Line data    Source code

       1             : /*
       2             :  * contrib/spi/autoinc.c
       3             :  */
       4             : #include "postgres.h"
       5             : 
       6             : #include "access/htup_details.h"
       7             : #include "catalog/pg_type.h"
       8             : #include "commands/sequence.h"
       9             : #include "commands/trigger.h"
      10             : #include "executor/spi.h"
      11             : #include "utils/builtins.h"
      12             : #include "utils/rel.h"
      13             : 
      14           2 : PG_MODULE_MAGIC;
      15             : 
      16           2 : PG_FUNCTION_INFO_V1(autoinc);
      17             : 
      18             : Datum
      19          12 : autoinc(PG_FUNCTION_ARGS)
      20             : {
      21          12 :     TriggerData *trigdata = (TriggerData *) fcinfo->context;
      22             :     Trigger    *trigger;        /* to get trigger name */
      23             :     int         nargs;          /* # of arguments */
      24             :     int        *chattrs;        /* attnums of attributes to change */
      25          12 :     int         chnattrs = 0;   /* # of above */
      26             :     Datum      *newvals;        /* vals of above */
      27             :     bool       *newnulls;       /* null flags for above */
      28             :     char      **args;           /* arguments */
      29             :     char       *relname;        /* triggered relation name */
      30             :     Relation    rel;            /* triggered relation */
      31          12 :     HeapTuple   rettuple = NULL;
      32             :     TupleDesc   tupdesc;        /* tuple description */
      33             :     bool        isnull;
      34             :     int         i;
      35             : 
      36          12 :     if (!CALLED_AS_TRIGGER(fcinfo))
      37             :         /* internal error */
      38           0 :         elog(ERROR, "not fired by trigger manager");
      39          12 :     if (!TRIGGER_FIRED_FOR_ROW(trigdata->tg_event))
      40             :         /* internal error */
      41           0 :         elog(ERROR, "must be fired for row");
      42          12 :     if (!TRIGGER_FIRED_BEFORE(trigdata->tg_event))
      43             :         /* internal error */
      44           0 :         elog(ERROR, "must be fired before event");
      45             : 
      46          12 :     if (TRIGGER_FIRED_BY_INSERT(trigdata->tg_event))
      47           6 :         rettuple = trigdata->tg_trigtuple;
      48           6 :     else if (TRIGGER_FIRED_BY_UPDATE(trigdata->tg_event))
      49           6 :         rettuple = trigdata->tg_newtuple;
      50             :     else
      51             :         /* internal error */
      52           0 :         elog(ERROR, "cannot process DELETE events");
      53             : 
      54          12 :     rel = trigdata->tg_relation;
      55          12 :     relname = SPI_getrelname(rel);
      56             : 
      57          12 :     trigger = trigdata->tg_trigger;
      58             : 
      59          12 :     nargs = trigger->tgnargs;
      60          12 :     if (nargs <= 0 || nargs % 2 != 0)
      61             :         /* internal error */
      62           0 :         elog(ERROR, "autoinc (%s): even number gt 0 of arguments was expected", relname);
      63             : 
      64          12 :     args = trigger->tgargs;
      65          12 :     tupdesc = rel->rd_att;
      66             : 
      67          12 :     chattrs = (int *) palloc(nargs / 2 * sizeof(int));
      68          12 :     newvals = (Datum *) palloc(nargs / 2 * sizeof(Datum));
      69          12 :     newnulls = (bool *) palloc(nargs / 2 * sizeof(bool));
      70             : 
      71          36 :     for (i = 0; i < nargs;)
      72             :     {
      73          12 :         int         attnum = SPI_fnumber(tupdesc, args[i]);
      74             :         int32       val;
      75             :         Datum       seqname;
      76             : 
      77          12 :         if (attnum <= 0)
      78           0 :             ereport(ERROR,
      79             :                     (errcode(ERRCODE_TRIGGERED_ACTION_EXCEPTION),
      80             :                      errmsg("\"%s\" has no attribute \"%s\"",
      81             :                             relname, args[i])));
      82             : 
      83          12 :         if (SPI_gettypeid(tupdesc, attnum) != INT4OID)
      84           0 :             ereport(ERROR,
      85             :                     (errcode(ERRCODE_TRIGGERED_ACTION_EXCEPTION),
      86             :                      errmsg("attribute \"%s\" of \"%s\" must be type INT4",
      87             :                             args[i], relname)));
      88             : 
      89          12 :         val = DatumGetInt32(SPI_getbinval(rettuple, tupdesc, attnum, &isnull));
      90             : 
      91          12 :         if (!isnull && val != 0)
      92             :         {
      93           9 :             i += 2;
      94           9 :             continue;
      95             :         }
      96             : 
      97           3 :         i++;
      98           3 :         chattrs[chnattrs] = attnum;
      99           3 :         seqname = CStringGetTextDatum(args[i]);
     100           3 :         newvals[chnattrs] = DirectFunctionCall1(nextval, seqname);
     101             :         /* nextval now returns int64; coerce down to int32 */
     102           3 :         newvals[chnattrs] = Int32GetDatum((int32) DatumGetInt64(newvals[chnattrs]));
     103           3 :         if (DatumGetInt32(newvals[chnattrs]) == 0)
     104             :         {
     105           1 :             newvals[chnattrs] = DirectFunctionCall1(nextval, seqname);
     106           1 :             newvals[chnattrs] = Int32GetDatum((int32) DatumGetInt64(newvals[chnattrs]));
     107             :         }
     108           3 :         newnulls[chnattrs] = false;
     109           3 :         pfree(DatumGetTextPP(seqname));
     110           3 :         chnattrs++;
     111           3 :         i++;
     112             :     }
     113             : 
     114          12 :     if (chnattrs > 0)
     115             :     {
     116           3 :         rettuple = heap_modify_tuple_by_cols(rettuple, tupdesc,
     117             :                                              chnattrs, chattrs,
     118             :                                              newvals, newnulls);
     119             :     }
     120             : 
     121          12 :     pfree(relname);
     122          12 :     pfree(chattrs);
     123          12 :     pfree(newvals);
     124          12 :     pfree(newnulls);
     125             : 
     126          12 :     return PointerGetDatum(rettuple);
     127             : }