Foreign Key to partitioned table

One of the long standing limitations of partitions is that you can't have foreign keys pointing to them.

Let's see if I can make it possible to have some kind of constraint that would do the same thing as fkey.

First, obviously, we need partitioned table with some data:

=$ CREATE TABLE USERS (
    id             serial NOT NULL,
    username       text NOT NULL,
    UNIQUE (username)
)
PARTITION BY RANGE ( username );
 
=$ CREATE TABLE users_a
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM (minvalue) TO ('b');
=$ CREATE TABLE users_b
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('b') TO ('c');
=$ CREATE TABLE users_c
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('c') TO ('d');
=$ CREATE TABLE users_d
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('d') TO ('e');
=$ CREATE TABLE users_e
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('e') TO ('f');
=$ CREATE TABLE users_f
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('f') TO ('g');
=$ CREATE TABLE users_g
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('g') TO ('h');
=$ CREATE TABLE users_h
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('h') TO ('i');
=$ CREATE TABLE users_i
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('i') TO ('j');
=$ CREATE TABLE users_j
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('j') TO ('k');
=$ CREATE TABLE users_k
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('k') TO ('l');
=$ CREATE TABLE users_l
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('l') TO ('m');
=$ CREATE TABLE users_m
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('m') TO ('n');
=$ CREATE TABLE users_n
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('n') TO ('o');
=$ CREATE TABLE users_o
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('o') TO ('p');
=$ CREATE TABLE users_p
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('p') TO ('q');
=$ CREATE TABLE users_q
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('q') TO ('r');
=$ CREATE TABLE users_r
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('r') TO ('s');
=$ CREATE TABLE users_s
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('s') TO ('t');
=$ CREATE TABLE users_t
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('t') TO ('u');
=$ CREATE TABLE users_u
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('u') TO ('v');
=$ CREATE TABLE users_v
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('v') TO ('w');
=$ CREATE TABLE users_w
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('w') TO ('x');
=$ CREATE TABLE users_x
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('x') TO ('y');
=$ CREATE TABLE users_y
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('y') TO ('z');
=$ CREATE TABLE users_z
    partition OF users (PRIMARY KEY (id))
    FOR VALUES FROM ('z') TO (maxvalue);
 
=$ INSERT INTO users (username) VALUES
    ('atom'), ('aught'), ('bides'), ('blob'), ('count'),
    ('flags'), ('fleck'), ('leggy'), ('otter'), ('plan'),
    ('rice'), ('riots'), ('semis'), ('serf'), ('sped'),
    ('stab'), ('tics'), ('vocal'), ('ward'), ('wrist');

I purposely picked partition key that is not based on primary.

Now – let's make a table that should have foreign key to users:

=$ CREATE TABLE accounts (
    id serial PRIMARY KEY,
    user_id int4 NOT NULL REFERENCES users (id)
);

When I tried to run it, I got:

ERROR:  cannot reference partitioned TABLE "users"

OK. So let's make it without fkey, and add our own constraint:

=$ CREATE TABLE accounts (
    id serial PRIMARY KEY,
    user_id int4 NOT NULL
);

In sources we can see that fkeys, use select like:

=$ SELECT 1 FROM ONLY <pktable> x WHERE pkatt1 = $1 [AND ...] FOR KEY SHARE OF x

The “FOR KEY SHARE" part is important, so we need it there too.

Let's consider how many cases we have for a trigger:

  1. on insert to accounts we need to check if inserted user exists
  2. on update to accounts, if user_id has changed, we have to check new user_id if it exists
  3. on delete from users, we have to check if there are no rows in accounts with given user_id
  4. on update in users, if id column has changed, we have to check if there are no rows in accounts with old value of id in user_id

None of these looks overly complicated, so let's write them.

To test the code, we will need to know what are the working, and not working, user_id values, so:

=$ SELECT MIN(id), MAX(id), COUNT(*) FROM users;
 MIN | MAX | COUNT 
-----+-----+-------
   1 |  20 |    20
(1 ROW)

Great. So every user_id from range 1 to 20 (including) will be ok. Anything outside should be rejected.

First trigger:

=$ CREATE OR REPLACE FUNCTION fkey_fn_1() RETURNS TRIGGER AS $$
DECLARE
BEGIN
    perform 1 FROM users x WHERE id = NEW.user_id FOR KEY SHARE OF x;
    IF NOT FOUND THEN
        raise exception 'Constraint violation.'
            USING
                ERRCODE = 'foreign_key_violation',
                HINT = 'Value ' || NEW.user_id || ' does not exist in column id in table users.';
    END IF;
    RETURN NULL;
END;
$$ LANGUAGE plpgsql;
 
=$ CREATE CONSTRAINT TRIGGER fkey_fn_1
    AFTER INSERT ON accounts FOR EACH ROW EXECUTE PROCEDURE fkey_fn_1();

With this in place, I can run a simple test to make sure that it does work:

=$ INSERT INTO accounts (user_id) VALUES (20);
INSERT 0 1
 
=$ INSERT INTO accounts (user_id) VALUES (21);
ERROR:  CONSTRAINT violation.
HINT:  VALUE 21 does NOT exist IN COLUMN id IN TABLE users.
CONTEXT:  PL/pgSQL FUNCTION fkey_fn_1() line 6 at RAISE

Sweet. Correct value was approved, and incorrect one was rejected.

Second trigger will use identical function, but let's create separate one, for now:

=$ CREATE OR REPLACE FUNCTION fkey_fn_2() RETURNS TRIGGER AS $$
DECLARE
BEGIN
    perform 1 FROM users x WHERE id = NEW.user_id FOR KEY SHARE OF x;
    IF NOT FOUND THEN
        raise exception 'Constraint violation.'
            USING
                ERRCODE = 'foreign_key_violation',
                HINT = 'Value ' || NEW.user_id || ' does not exist in column id in table users.';
    END IF;
    RETURN NULL;
END;
$$ LANGUAGE plpgsql;
 
=$ CREATE CONSTRAINT TRIGGER fkey_fn_2
     AFTER UPDATE ON accounts
     FOR EACH ROW
     WHEN (OLD.user_id <> NEW.user_id)
     EXECUTE PROCEDURE fkey_fn_2();

Because of my prior calls, my data in accounts looks like:

=$ SELECT * FROM accounts;
┌────┬─────────┐
│ id │ user_id │
╞════╪═════════╡
│ 1320 │
└────┴─────────┘

So, let's test the updates:

=$ UPDATE accounts SET user_id = 19 WHERE user_id = 20;
UPDATE 1
 
=$ UPDATE accounts SET user_id = 21 WHERE user_id = 19;
ERROR:  CONSTRAINT violation.
HINT:  VALUE 21 does NOT exist IN COLUMN id IN TABLE users.
CONTEXT:  PL/pgSQL FUNCTION fkey_fn_2() line 6 at RAISE

Sweet. Now, let's move on checking users:

=$ CREATE OR REPLACE FUNCTION fkey_fn_3() RETURNS TRIGGER AS $$
DECLARE
BEGIN
    perform 1 FROM accounts x WHERE x.user_id = OLD.id FOR KEY SHARE OF x LIMIT 1;
    IF FOUND THEN
        raise exception 'Constraint violation.'
            USING
                ERRCODE = 'foreign_key_violation',
                HINT = 'Value ' || OLD.id || ' still exists in column user_id in table accounts.';
    END IF;
    RETURN NULL;
END;
$$ LANGUAGE plpgsql;
CREATE FUNCTION
 
=$ CREATE CONSTRAINT TRIGGER fkey_fn_3 AFTER DELETE ON users FOR EACH ROW EXECUTE PROCEDURE fkey_fn_3();
CREATE TRIGGER

OK. First question is: did it get propagated to partitions?

=$ \d users
                             TABLE "public.users"
  COLUMN  |  TYPE   | Collation | NULLABLE |              DEFAULT              
----------+---------+-----------+----------+-----------------------------------
 id       | INTEGER |           | NOT NULL | NEXTVAL('users_id_seq'::regclass)
 username | text    |           | NOT NULL | 
Partition KEY: RANGE (username)
Indexes:
    "users_username_key" UNIQUE CONSTRAINT, btree (username)
Triggers:
    fkey_fn_3 AFTER DELETE ON users NOT DEFERRABLE INITIALLY IMMEDIATE FOR EACH ROW EXECUTE PROCEDURE fkey_fn_3()
NUMBER OF partitions: 26 (USE \d+ TO list them.)
 
=$ \d users_a
                            TABLE "public.users_a"
  COLUMN  |  TYPE   | Collation | NULLABLE |              DEFAULT              
----------+---------+-----------+----------+-----------------------------------
 id       | INTEGER |           | NOT NULL | NEXTVAL('users_id_seq'::regclass)
 username | text    |           | NOT NULL | 
Partition OF: users FOR VALUES FROM (MINVALUE) TO ('b')
Indexes:
    "users_a_pkey" PRIMARY KEY, btree (id)
    "users_a_username_key" UNIQUE CONSTRAINT, btree (username)
Triggers:
    fkey_fn_3 AFTER DELETE ON users_a NOT DEFERRABLE INITIALLY IMMEDIATE FOR EACH ROW EXECUTE PROCEDURE fkey_fn_3()
 
=$ \d users_z
                            TABLE "public.users_z"
  COLUMN  |  TYPE   | Collation | NULLABLE |              DEFAULT              
----------+---------+-----------+----------+-----------------------------------
 id       | INTEGER |           | NOT NULL | NEXTVAL('users_id_seq'::regclass)
 username | text    |           | NOT NULL | 
Partition OF: users FOR VALUES FROM ('z') TO (MAXVALUE)
Indexes:
    "users_z_pkey" PRIMARY KEY, btree (id)
    "users_z_username_key" UNIQUE CONSTRAINT, btree (username)
Triggers:
    fkey_fn_3 AFTER DELETE ON users_z NOT DEFERRABLE INITIALLY IMMEDIATE FOR EACH ROW EXECUTE PROCEDURE fkey_fn_3()

Looks like it did, so let quickly recheck what account we have:

=$ SELECT * FROM accounts;
 id | user_id 
----+---------
 13 |      19
(1 ROW)

And now test the deletes:

=$ DELETE FROM users WHERE id = 1;
DELETE 1
 
=$ DELETE FROM users WHERE id = 19;
ERROR:  CONSTRAINT violation.
HINT:  VALUE 19 still EXISTS IN COLUMN user_id IN TABLE accounts.
CONTEXT:  PL/pgSQL FUNCTION fkey_fn_3() line 6 at RAISE

Sweet, worked and failed as expected, and as needed.

So, the last trigger is the one on users update …

=$ CREATE OR REPLACE FUNCTION fkey_fn_4() RETURNS TRIGGER AS $$
DECLARE
BEGIN
    perform 1 FROM accounts x WHERE x.user_id = OLD.id FOR KEY SHARE OF x LIMIT 1;
    IF FOUND THEN
        raise exception 'Constraint violation.'
            USING
                ERRCODE = 'foreign_key_violation',
                HINT = 'Value ' || OLD.id || ' still exists in column user_id in table accounts.';
    END IF;
    RETURN NULL;
END;
$$ LANGUAGE plpgsql;
CREATE FUNCTION
 
=$ CREATE CONSTRAINT TRIGGER fkey_fn_4 AFTER UPDATE ON users FOR EACH ROW WHEN (OLD.id <> NEW.id) EXECUTE PROCEDURE fkey_fn_4();
CREATE TRIGGER

We know that we have account for user 19, so first, let's try update that should work:

=$ UPDATE users SET id = 1 WHERE id = 2;
UPDATE 1

No surprises here. And how about one that shouldn't work?

=$ UPDATE users SET id = 2 WHERE id = 19;
ERROR:  CONSTRAINT violation.
HINT:  VALUE 19 still EXISTS IN COLUMN user_id IN TABLE accounts.
CONTEXT:  PL/pgSQL FUNCTION fkey_fn_4() line 6 at RAISE

Great. Looks like I got it to work. But it's not done yet – I should test performance, and perhaps make it in such a way that one could use the same functions, with no modifications, regardless of source and destination table and column names.

This will have to wait a bit, though …

Leave a Reply

Your email address will not be published.

This site uses Akismet to reduce spam. Learn how your comment data is processed.