contrib/spi/autoinc.c

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