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