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            1 : PG_MODULE_MAGIC_EXT(
      15              :                     .name = "autoinc",
      16              :                     .version = PG_VERSION
      17              : );
      18              : 
      19            2 : PG_FUNCTION_INFO_V1(autoinc);
      20              : 
      21              : Datum
      22           12 : autoinc(PG_FUNCTION_ARGS)
      23              : {
      24           12 :     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           12 :     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           12 :     HeapTuple   rettuple = NULL;
      35              :     TupleDesc   tupdesc;        /* tuple description */
      36              :     bool        isnull;
      37              :     int         i;
      38              : 
      39           12 :     if (!CALLED_AS_TRIGGER(fcinfo))
      40              :         /* internal error */
      41            0 :         elog(ERROR, "not fired by trigger manager");
      42           12 :     if (!TRIGGER_FIRED_FOR_ROW(trigdata->tg_event))
      43              :         /* internal error */
      44            0 :         elog(ERROR, "must be fired for row");
      45           12 :     if (!TRIGGER_FIRED_BEFORE(trigdata->tg_event))
      46              :         /* internal error */
      47            0 :         elog(ERROR, "must be fired before event");
      48              : 
      49           12 :     if (TRIGGER_FIRED_BY_INSERT(trigdata->tg_event))
      50            3 :         rettuple = trigdata->tg_trigtuple;
      51            9 :     else if (TRIGGER_FIRED_BY_UPDATE(trigdata->tg_event))
      52            9 :         rettuple = trigdata->tg_newtuple;
      53              :     else
      54              :         /* internal error */
      55            0 :         elog(ERROR, "cannot process DELETE events");
      56              : 
      57           12 :     rel = trigdata->tg_relation;
      58           12 :     relname = SPI_getrelname(rel);
      59              : 
      60           12 :     trigger = trigdata->tg_trigger;
      61              : 
      62           12 :     nargs = trigger->tgnargs;
      63           12 :     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           12 :     args = trigger->tgargs;
      68           12 :     tupdesc = rel->rd_att;
      69              : 
      70           12 :     chattrs = (int *) palloc(nargs / 2 * sizeof(int));
      71           12 :     newvals = (Datum *) palloc(nargs / 2 * sizeof(Datum));
      72           12 :     newnulls = (bool *) palloc(nargs / 2 * sizeof(bool));
      73              : 
      74           24 :     for (i = 0; i < nargs;)
      75              :     {
      76           12 :         int         attnum = SPI_fnumber(tupdesc, args[i]);
      77              :         int32       val;
      78              :         Datum       seqname;
      79              : 
      80           12 :         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           12 :         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           12 :         val = DatumGetInt32(SPI_getbinval(rettuple, tupdesc, attnum, &isnull));
      93              : 
      94           12 :         if (!isnull && val != 0)
      95              :         {
      96            4 :             i += 2;
      97            4 :             continue;
      98              :         }
      99              : 
     100            8 :         i++;
     101            8 :         chattrs[chnattrs] = attnum;
     102            8 :         seqname = CStringGetTextDatum(args[i]);
     103            8 :         newvals[chnattrs] = DirectFunctionCall1(nextval, seqname);
     104              :         /* nextval now returns int64; coerce down to int32 */
     105            8 :         newvals[chnattrs] = Int32GetDatum((int32) DatumGetInt64(newvals[chnattrs]));
     106            8 :         if (DatumGetInt32(newvals[chnattrs]) == 0)
     107              :         {
     108            1 :             newvals[chnattrs] = DirectFunctionCall1(nextval, seqname);
     109            1 :             newvals[chnattrs] = Int32GetDatum((int32) DatumGetInt64(newvals[chnattrs]));
     110              :         }
     111            8 :         newnulls[chnattrs] = false;
     112            8 :         pfree(DatumGetTextPP(seqname));
     113            8 :         chnattrs++;
     114            8 :         i++;
     115              :     }
     116              : 
     117           12 :     if (chnattrs > 0)
     118              :     {
     119            8 :         rettuple = heap_modify_tuple_by_cols(rettuple, tupdesc,
     120              :                                              chnattrs, chattrs,
     121              :                                              newvals, newnulls);
     122              :     }
     123              : 
     124           12 :     pfree(relname);
     125           12 :     pfree(chattrs);
     126           12 :     pfree(newvals);
     127           12 :     pfree(newnulls);
     128              : 
     129           12 :     return PointerGetDatum(rettuple);
     130              : }