Fun With Postgres: Custom Constraints

P
ostgres comes with a rich set of tool to help developers maintain data integrity. Many of them come in the form of Constraints. Foreign Key, Not Null, Exclusion, and Check constraint allow developers to off load what would usually be a lot of application code to the database server with very litte effort.

Between the Foreign key, which verifies values in another table, and the Check constraint, which verifies values in a specific column, you can typically accomplish just about everything you need to do rather easily. The main problem the you'll run into is that these kinds of constraints are restricted to a single column on a table. Additionally they only apply the the current value in that column. This restriction doesn't apply to custom constraint, which is a special kind of trigger. This effectively gives you the full power of SQL to implement much more complex constraint conditions.

CREATE TYPE states AS ENUM ('one', 'two', 'three');
create table example (
  id INT GENERATED ALWAYS AS IDENTITY
, state states NOT NULL  
);

A simple example to illustrate the utility of a constraint trigger is a simple state machine. Here, our table has a state field that has a custom ENUM column. This ensures that the column can only contain one, two, or three. However, I want to make sure that the initial value is always one. Additionally, I want to make sure that, during update operations, that the state value can only be set to two if the previous value was one and can only be set to three if the previous value was two. If any of those conditions are not met, that operation should fail and abort the transaction.

+ -- ONE --> -- TWO --> -- THREE -- +
Restricted Data Transitions

This kind of logic is generally pretty difficult to capture with traditional check constraints. At this point this is usually where most application developers will turn to doing this work at the application layer by pulling the current record, comparing the input value to the existing value and making a decision. Careful, of course, to account for the situation where there is no record and taking the appropriate action. The problem here is that doing it this way creates a data race between the running instances of the application in the time it take to read / check the record and updating or inserting value. It is an unnecessary bit of complexity when postgres can do all of that in a single step with a constraint

Trigger [tri'ger] -n., --noun

  1. special stored procedure that is run when specific actions occur within a database.
  2. cause an event or situation to happen or exist.

To do this we create a function like any other function that RETURNS TRIGGER. The function must do one of three things:

  1. Raise an exception to abort the transaction
  2. Return a record/row value having exactly the structure of the table the trigger was fired for
  3. Return NULL
CREATE OR REPLACE FUNCTION example_fsm_check()
RETURNS TRIGGER AS 
$$
  BEGIN
    RETURN NEW;
  END;
$$ LANGUAGE PLPGSQL VOLATILE;

Right now, this is a no-op. What we need to do is determine if the transaction is an INSERT and check for the value one or check for the transition values if it is an UPDATE. Postgres injects some additional information into trigger functions that you can inspect to make these kinds of determinations. In this case we want TG_OP.

CREATE OR REPLACE FUNCTION example_fsm_check()
RETURNS TRIGGER AS 
$$
  BEGIN
    IF TG_OP = 'INSERT' AND NEW.state != 'one' THEN
      RAISE EXCEPTION 'Invalid FSM State %s', NEW.state
    END IF;
    RETURN NEW;
  END;
$$ LANGUAGE PLPGSQL VOLATILE;

Now things are getting interesting. In the above snippet, if the operation is an insert, and the state column is NOT one, raise an exception. We can make use of  TG_OP to tell us that the operation is an insert, and the special NEW object to inspect the data being inserted. This prevents the initial state from being anything other than the value we want and raises an exception if it is.

INSERT INTO example (state)
VALUES ('two')
ERROR:  Invalid FSM State twos
CONTEXT:  PL/pgSQL function example_fsm() line 5 at RAISE
SQL state: P0001

The next thing to do is handle the allowable state transitions during updates. Like most other programming languages, plpgsql has some facilities for logic branching. Additionally, during an UPDATE operation, postgres gives trigger functions access to the OLD record as well as the NEW record

-- on update
IF TG_OP = 'UPDATE' THEN
  -- if state changes
  IF OLD.status IS DISTINCT FROM NEW.status THEN
    -- only from one -> two
    IF OLD.status = 'one' AND NEW.status = 'two' THEN
      RAISE EXCEPTION 'Invalid state transition';
    END IF;
		
    -- only from two -> three
    IF OLD.status = 'two' AND NEW.status != 'three' THEN
      RAISE EXCEPTION 'Invalid state transition';
    END IF;
  END IF;
END IF;
-- end update

Here, we add a couple of if checks. The main if condition is checking if status column has changed at all - if not, there is nothing to do. However, if it has changed, there are two additional conditions that make up the final logic of our little state machine. If the old status value was ONE, and the new value is not TWO, then raise an exception. Similarly if the value was TWO and the new value is not THREE, also raise an exception.

Lastly, we can attach the function to a table as a CONSTRAINT trigger like normal.

CREATE CONSTRAINT TRIGGER check_sample_fsm
AFTER INSERT OR UPDATE
ON example
FOR EACH ROW
EXECUTE procedure example_fsm_constraint();

That's it! we've defined our own constraint. Now unlike a regular trigger procedure, a constraint trigger can define when it is triggered as well as conditions that define when the trigger is executed. For example, if we want the constraint to fire at the end of a transaction and only when the state is set to one, we can easily add these instructions to the trigger using the DEFERRABLE option and the WHEN condition clause.

CREATE CONSTRAINT TRIGGER check_sample_fsm
AFTER INSERT OR UPDATE
ON example
DEFERRABLE INITIALLY DEFERRED 
FOR EACH ROW
WHEN NEW.state = 'one'
EXECUTE procedure example_fsm_constraint();

Trigger constraints give us the ability to define very complex triggers that extend well beyond what one could do with CHECK constraints. They could even be used, for example, to build out FOREIGN KEY like constraints that check multiple tables or other points of data that can't be accomplished otherwise.

Trigger Constraints - They are pretty cool.

postgres database trigger fun-with-postgres