How to limit rows to at most N per category

The question was asked relatively recently on irc. And it proved to be non-trivial.

Surely, if you want to have one row per category (one address per user), it's trivial – add user_id column to addresses, make it unique, and we're done. But what if we want to allow five addresses? Or five thousands?

Let's see.

Obviously we have to start with some table. Let's stick with addresses. Since I don't really care about other fields, it will be simple:

CREATE TABLE users (
    id       text NOT NULL,
    PRIMARY KEY (id)
);
 
INSERT INTO users (id) VALUES ('depesz');
 
CREATE TABLE addresses (
    id      int8 generated BY DEFAULT AS IDENTITY,
    user_id text NOT NULL REFERENCES users (id),
    PRIMARY KEY (id)
);
 
CREATE INDEX addresses_user_id ON addresses (user_id);

All in place.

Now, let's say we want to have at most 3 addresses per user.

To do it, we can simply use a trigger:

CREATE FUNCTION trg_check_addresses_per_user() RETURNS TRIGGER AS $$
DECLARE
    v_count int4;
BEGIN
    SELECT COUNT(*) INTO v_count FROM addresses WHERE user_id = NEW.user_id;
    IF v_count >= 3 THEN
        raise exception
            'User % would have % addresses, but only 3 allowed.', NEW.user_id, v_count+1
            USING ERRCODE = 'check_violation';
    END IF;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;
 
CREATE TRIGGER trg_check_addresses_per_user
    BEFORE INSERT OR UPDATE ON addresses
    FOR EACH ROW EXECUTE FUNCTION trg_check_addresses_per_user();

And now we can test it:

=$ INSERT INTO addresses (user_id) VALUES ('depesz');
INSERT 0 1
 
=$ INSERT INTO addresses (user_id) VALUES ('depesz');
INSERT 0 1
 
=$ INSERT INTO addresses (user_id) VALUES ('depesz');
INSERT 0 1
 
=$ INSERT INTO addresses (user_id) VALUES ('depesz');
ERROR:  USER depesz would have 4 addresses, but ONLY 3 allowed.
CONTEXT:  PL/pgSQL FUNCTION trg_check_addresses_per_user() line 7 at RAISE

Great. Works. Thanks to index on user_id, and low limit, it will be fast enough for most cases.

But what about case where we want to limit to 5000 rows? Getting count() of them will be slow. Is there a better way?

Yes – precached counts.

To do it, let's drop existing trigger, as it will not server any purpose:

DROP TRIGGER trg_check_addresses_per_user ON addresses;
DROP FUNCTION trg_check_addresses_per_user();

Now – we need a place to store existing address count. Since it's per user, the logical place for this is to store it in users table, so:

ALTER TABLE users ADD COLUMN address_count int4;

With this in place, make sure current counts are correct:

UPDATE users AS u
    SET address_count = x.count
    FROM ( SELECT user_id, COUNT(*) FROM addresses GROUP BY user_id ) AS x
    WHERE u.id = x.user_id;

And now we will need three triggers – one for insert, one for delete, and one for update, in case someone would change user_id in address. Let's start with the INSERT one, as it's simpler:

CREATE FUNCTION trg_update_address_count_i() RETURNS TRIGGER AS $$
DECLARE
BEGIN
    UPDATE users SET address_count = COALESCE(address_count, 0) + 1 WHERE id = NEW.user_id;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;
 
CREATE TRIGGER trg_update_address_count_i
    AFTER INSERT ON ADDRESSES
    FOR EACH ROW EXECUTE FUNCTION trg_update_address_count_i();

Please note that since this should happen after all potential data changes in this row has happened, this trigger changed from BEFORE to AFTER

Delete trigger is very similar:

CREATE FUNCTION trg_update_address_count_d() RETURNS TRIGGER AS $$
DECLARE
BEGIN
    UPDATE users SET address_count = COALESCE(address_count, 0) - 1 WHERE id = OLD.user_id;
    RETURN OLD;
END;
$$ LANGUAGE plpgsql;
 
CREATE TRIGGER trg_update_address_count_d
    AFTER DELETE ON ADDRESSES
    FOR EACH ROW EXECUTE FUNCTION trg_update_address_count_d();

Now, the other trigger. Theoretically changing user_id in address data can be forbidden, but let's assume it's not.

CREATE FUNCTION trg_update_address_count_u() RETURNS TRIGGER AS $$
DECLARE
BEGIN
    IF NEW.user_id < OLD.user_id THEN
        UPDATE users SET address_count = COALESCE(address_count, 0) + 1 WHERE id = NEW.user_id;
        UPDATE users SET address_count = COALESCE(address_count, 0) - 1 WHERE id = OLD.user_id;
    ELSE
        UPDATE users SET address_count = COALESCE(address_count, 0) - 1 WHERE id = OLD.user_id;
        UPDATE users SET address_count = COALESCE(address_count, 0) + 1 WHERE id = NEW.user_id;
    END IF;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;
 
CREATE TRIGGER trg_update_address_count_u
    AFTER UPDATE OF user_id ON ADDRESSES
    FOR EACH ROW EXECUTE FUNCTION trg_update_address_count_u();

You might ask – why the if, and reordering of updates. This is so updates will always happen in alphabetical order, so we will not end up with deadlock. I wrote more about it earlier (search for deadlock).

With all this in place, let's see if it will really work. I will need to add some more users to allow updates:

INSERT INTO users (id) VALUES ('test');

and change the data a bit:

DELETE FROM addresses WHERE id = ( SELECT MIN(id) FROM addresses );
UPDATE addresses SET user_id = 'test' WHERE id = ( SELECT MIN(id) FROM addresses );
INSERT INTO addresses (user_id) VALUES ('depesz'), ('test');

Afterwards I got these addresses:

SELECT * FROM addresses;
 id | user_id
----+---------
  3 | depesz
  2 | test
  7 | depesz
  8 | test
(4 ROWS)

and users address_count is correct:

SELECT * FROM users;
   id   | address_count
--------+---------------
 depesz |             2
 test   |             2
(2 ROWS)

With this all in place, the last thing to do, is to add actual check:

ALTER TABLE users ADD CHECK ( address_count <= 5000 );

Given that we have 2 addresses per user, let's insert 5000 for one of users:

=$ \timing
Timing IS ON.
 
=$ INSERT INTO addresses (user_id)
    SELECT 'depesz' FROM generate_series(1,5000);
ERROR:  NEW ROW FOR relation "users" violates CHECK CONSTRAINT "users_address_count_check"
DETAIL:  Failing ROW contains (depesz, 5001).
CONTEXT:  SQL statement "UPDATE users SET address_count = coalesce(address_count, 0) + 1 WHERE id = NEW.user_id"
PL/pgSQL FUNCTION trg_update_address_count_i() line 4 at SQL statement
TIME: 690.912 ms

Failed as expected. And what is the state of our tables?

SELECT * FROM addresses;
 id | user_id
----+---------
  3 | depesz
  2 | test
  7 | depesz
  8 | test
(4 ROWS)
 
SELECT * FROM users;
   id   | address_count
--------+---------------
 depesz |             2
 test   |             2
(2 ROWS)

Nice. Of course updating so many times rows in users, in case of such massive update might not be the best idea – and that's why it took 690ms. Can it be done better?

Well, yeah. We'll need to use statement triggers for this, though. But let's see what we can do.

The three triggers that we have will stay there, but their body will be changed:

CREATE OR REPLACE FUNCTION trg_update_address_count_i() RETURNS TRIGGER AS $$
DECLARE
    v_counts   jsonb := current_setting('trg.update_address_count');
    v_previous INT4  := COALESCE( (v_counts ->> NEW.user_id )::int4, 0 );
BEGIN
    v_counts := jsonb_set(v_counts, ARRAY[ NEW.user_id ], to_jsonb( v_previous + 1 ) );
    perform set_config( 'trg.update_address_count', v_counts::text, TRUE );
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;
 
CREATE OR REPLACE FUNCTION trg_update_address_count_d() RETURNS TRIGGER AS $$
DECLARE
    v_counts   jsonb := current_setting('trg.update_address_count');
    v_previous INT4  := COALESCE( (v_counts ->> OLD.user_id )::int4, 0 );
BEGIN
    v_counts := jsonb_set(v_counts, ARRAY[ OLD.user_id ], to_jsonb( v_previous - 1 ) );
    perform set_config( 'trg.update_address_count', v_counts::text, TRUE );
    RETURN OLD;
END;
$$ LANGUAGE plpgsql;
 
CREATE OR REPLACE FUNCTION trg_update_address_count_u() RETURNS TRIGGER AS $$
DECLARE
    v_counts   jsonb := current_setting('trg.update_address_count');
    v_prev_new INT4  := COALESCE( (v_counts ->> NEW.user_id )::int4, 0 );
    v_prev_old INT4  := COALESCE( (v_counts ->> OLD.user_id )::int4, 0 );
BEGIN
    IF NEW.user_id <> OLD.user_id THEN
        v_counts := jsonb_set(v_counts, ARRAY[ NEW.user_id ], to_jsonb( v_prev_new + 1 ) );
        v_counts := jsonb_set(v_counts, ARRAY[ OLD.user_id ], to_jsonb( v_prev_old - 1 ) );
        perform set_config( 'trg.update_address_count', v_counts::text, TRUE );
     END IF;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

Now, a bit of explanation what these do.

Instead of updating users on each row, triggers will keep track of how many addresses were changed for each user in variable (trg.update_address_count). This variable will use JSONB to keep many counters.

Name of the variable might need explanation, too. Trg is a variable class. I can't set random variables without class, because these are normal PostgreSQL config variables (a.k.a. GUCs). Class can be basically any string. I picked trg as class to somehow make it obvious that this variable is used in triggers, and make sure it will not clash with some other variable named update_address_count in class used by something else.

INSERT INTO addresses (user_id) VALUES
    ( 'depesz' ),
    ( 'depesz' ),
    ( 'depesz' ),
    ( 'test' ),
    ( 'test' ),
    ( 'depesz' );

I'd expect the trg.update_address_count to be:

{"test": 2, "depesz": 4}

But, we need two more things – setting the variable to sensible state at start, and then finally doing the updates.

First trigger is simple:

CREATE FUNCTION trg_check_addresses_per_user_startup() RETURNS TRIGGER AS $$
DECLARE
BEGIN
    PERFORM set_config('trg.update_address_count', '{}', TRUE );
    IF TG_OP = 'DELETE' THEN
        RETURN OLD;
    END IF;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;
 
CREATE TRIGGER trg_check_addresses_per_user_startup
    BEFORE INSERT OR UPDATE OR DELETE ON addresses
    FOR EACH STATEMENT EXECUTE FUNCTION trg_check_addresses_per_user_startup();

It just sets the variable to empty jsonb object, and returns whatever is needed.

Trigger that will do the actual finalizes is more complicated:

CREATE FUNCTION trg_check_addresses_per_user_finalize() RETURNS TRIGGER AS $$
DECLARE
    v_counts   jsonb := current_setting('trg.update_address_count');
    v_row      record;
BEGIN
    FOR v_row IN SELECT KEY, VALUE FROM jsonb_each(v_counts) ORDER BY KEY LOOP
        UPDATE users SET address_count = COALESCE(address_count, 0) + v_row.value::int4 WHERE id = v_row.key;
    END loop;
    RETURN NULL;
END;
$$ LANGUAGE plpgsql;
 
CREATE TRIGGER trg_check_addresses_per_user_finalize
    after INSERT OR UPDATE OR DELETE ON addresses
    FOR EACH STATEMENT EXECUTE FUNCTION trg_check_addresses_per_user_finalize();

In the for loop it extracts information from the variable, and sorts it by username. And then updates all of them, in order to avoid deadlocks.

Previously – with update on each row, failing 5000 row insert failed after almost 700 ms.

Now?

=$ \timing ON
Timing IS ON.
 
=$ INSERT INTO addresses (user_id)
    SELECT 'depesz' FROM generate_series(1,5000);
ERROR:  NEW ROW FOR relation "users" violates CHECK CONSTRAINT "users_address_count_check"
DETAIL:  Failing ROW contains (depesz, 5002).
CONTEXT:  SQL statement "UPDATE users SET address_count = coalesce(address_count, 0) + v_row.value::int4 WHERE id = v_row.key"
PL/pgSQL FUNCTION trg_check_addresses_per_user_finalize() line 7 at SQL statement
TIME: 103.904 ms

Better. Of course without triggers it would be even faster:

=$ BEGIN;
BEGIN
 
=$ ALTER TABLE addresses disable TRIGGER USER;
ALTER TABLE
 
=$ INSERT INTO addresses (user_id)
    SELECT 'depesz' FROM generate_series(1,5000);
INSERT 0 5000
TIME: 52.992 ms
 
=$ ROLLBACK;
ROLLBACK

I think this slowdown is acceptable, especially since the table in question is very narrow, and it's changes are not really IO intensive.

To finalize, let's make some random things, and see if all triggers work OK.

Starting with sanity check:

WITH REAL AS (
    SELECT user_id, COUNT(*)
    FROM addresses
    GROUP BY user_id
)
SELECT
    u.id AS USER,
    u.address_count AS cached_count,
    r.count AS real_count,
    CASE
        WHEN u.address_count IS DISTINCT FROM r.count THEN '✗'
        ELSE '✓'
    END AS is_ok
FROM
    users u
    LEFT OUTER JOIN REAL r ON u.id = r.user_id
ORDER BY
    u.id;
  USER  | cached_count | real_count | is_ok
--------+--------------+------------+-------
 depesz |            2 |          2 | ✓
 test   |            2 |          2 |(2 ROWS)

To make inserting/updating easier, let's make a macro that will return random username:

\SET random_user '(ARRAY[''depesz'', ''test''])[floor(1 + random() * 2)]'

And test it:

SELECT :random_user, COUNT(*)  FROM generate_series(1,100) i GROUP BY 1 ORDER BY 1;
 array  | COUNT
--------+-------
 depesz |    52
 test   |    48
(2 ROWS)

Looks sane. Now, insert 2000 random addresses:

INSERT INTO addresses (user_id)
    SELECT :random_user FROM generate_series(1,2000) i;

, update ~ 50% of them:

UPDATE addresses
     SET user_id =  :random_user
     WHERE random() < 0.5;

, and finally, delete ~ 50% of them:

DELETE FROM  addresses WHERE random() < 0.5;

How does the count look like now?

WITH REAL AS (
    SELECT user_id, COUNT(*)
    FROM addresses
    GROUP BY user_id
)
SELECT
    u.id AS USER,
    u.address_count AS cached_count,
    r.count AS real_count,
    CASE
        WHEN u.address_count IS DISTINCT FROM r.count THEN '✗'
        ELSE '✓'
    END AS is_ok
FROM
    users u
    LEFT OUTER JOIN REAL r ON u.id = r.user_id
ORDER BY
    u.id;
  USER  | cached_count | real_count | is_ok
--------+--------------+------------+-------
 depesz |          494 |        494 | ✓
 test   |          517 |        517 |(2 ROWS)

All looks good. Hope you'll find it useful.

6 thoughts on “How to limit rows to at most N per category”

  1. Hi,

    > But what if we want to allow five addresses?

    I checked only the first approach. It is designed for a single-user system.
    If there are multiple sessions inserting stuff there, it will not work.

    For instance, I deleted the third row to have only two rows in addresses:
    postgres=# delete from addresses where id=3;
    DELETE 1

    Then, I started a transaction in session 1:
    postgres=# begin;
    BEGIN
    postgres=# INSERT INTO addresses (user_id) VALUES (‘depesz’);
    INSERT 0 1

    The same in session 2.
    Committed in both sessions and got 4 rows in addresses.

    postgres=# commit;
    COMMIT
    postgres=# select * from addresses;
    id | user_id
    —-+———
    1 | depesz
    2 | depesz
    5 | depesz
    6 | depesz
    (4 rows)

    Even without explicit transactions, we can inject a sleep call after the SELECT COUNT in your trigger. Then run multiple INSERT’s simultaneously. It should leave us with more than 3 rows in addresses.

    If such an issue happens in Oracle, which I am more familiar with, I will serialize the modification of a user so as to guarantee that two sessions do not try to change addresses of one user at the same time.
    Basically, it would be:
    begin;
    select id from users where id=’depesz’ for update;
    INSERT INTO addresses (user_id) VALUES (‘depesz’);
    commit;

    The same works fine in PostgreSQL too.

    Regards,
    Mikhail.

  2. This is cool, but just curious: What is the real-world use case?

    Would it be so bad if I simply deleted old rows where the count > 5000 periodically?

  3. @John:

    No idea. This was a question on irc, and interesting thing to write. Usecase is a problem for someone that will use it 🙂

  4. Interesting, the “Counter cache” approach is very popular in the Rails framework and easy to enabled, but I’d always avoided it due to performance concerns and race conditions. If the `addresses` table already has many inserts/deletes, if you also need to update users with the count of addresses, you’ve essentially doubled the IO needed to make that change. Also if the current count is `5`, and you run two simultaneous transactions, one with an insert tries to increment the count to `6`, and one with a delete decrements the count to `4`. So the count ends up wrong no matter which one “wins”, or you have to lock the user row, which brings back the potential for performance problems.

  5. This is an interesting problem. Without assertions it isn’t easy to express rules that span multiple tables. I’ve used a similar approach before. But I prefer to encode the invalid state detection into a view. This can simplify the logic inside of a trigger as it boils down to something like:

    SELECT * FROM users_with_too_many_addresses WHERE id = NEW.id;

    If any row is found the constraint has been violated. It also allows more flexibility as to when the check actually runs. For example, in some cases an hourly or nightly check for an error condition may be sufficient. I have more details and some tooling for this technique at https://github.com/jackc/doublecheck.

    How does using a GUC to optimize the counter cache compare to using the transition relation for the statement trigger? I would expect the transition relation to perform pretty well.

Comments are closed.