contrib/spi/autoinc.c

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