batch_claim_otk_keys.sql

\echo nuke table\\
DROP TABLE IF EXISTS e2e_one_time_keys_json;

\echo Make table\\
CREATE TABLE e2e_one_time_keys_json (
    user_id     text   NOT NULL,
    device_id   text   NOT NULL,
    algorithm   text   NOT NULL,
    key_id      text   NOT NULL,
    ts_added_ms bigint NOT NULL,
    key_json    text   NOT NULL,
    UNIQUE (user_id, device_id, algorithm, key_id)
);

\echo Dummy data. Multiple users, devices, and algorithms.\\
INSERT INTO e2e_one_time_keys_json (user_id, device_id, algorithm, key_id, ts_added_ms, key_json)
SELECT '@alice:test', 'alice_dev_1', 'algA', concat('key', id.id), 0, CONCAT('json', id.id)
FROM generate_series(1, 100) as id(id)
UNION
SELECT '@alice:test', 'alice_dev_2', 'algA', concat('key', id.id), 0, CONCAT('json', id.id)
FROM generate_series(101, 150) as id(id)
UNION
SELECT '@alice:test', 'alice_dev_2', 'algB', concat('key', id.id), 0, CONCAT('json', id.id)
FROM generate_series(151, 200) as id(id)
UNION
SELECT '@bob:test', 'bob_dev', 'algA', concat('key', id.id), 0, CONCAT('json', id.id)
FROM generate_series(201, 250) as id(id)
UNION
SELECT '@bob:test', 'bob_dev', 'algB', concat('key', id.id), 0, CONCAT('json', id.id)
FROM generate_series(251, 300) as id(id)
;

\echo How many keys did we insert? \\

SELECT user_id, device_id, algorithm, count(*)
FROM e2e_one_time_keys_json
GROUP BY (user_id, device_id, algorithm)
ORDER BY user_id, device_id, algorithm;


\echo The query of doom. \\

WITH
    claims(user_id, device_id, algorithm, claim_count) AS (
        VALUES ('@alice:test', 'alice_dev_1', 'algA', 4),
               ('@alice:test', 'alice_dev_2', 'algA', 2),
               ('@bob:test', 'bob_dev', 'algB', 1)
    ),
    ranked_keys AS (
        SELECT user_id,
               device_id,
               algorithm,
               key_id,
               claims.claim_count,
               ROW_NUMBER() OVER (PARTITION BY (user_id, device_id, algorithm)) AS r
        FROM e2e_one_time_keys_json
            JOIN claims USING (user_id, device_id, algorithm)
    )
DELETE
FROM e2e_one_time_keys_json k
WHERE (user_id, device_id, algorithm, key_id) IN (
    SELECT user_id, device_id, algorithm, key_id
    FROM ranked_keys
    WHERE r <= claim_count
)
RETURNING user_id, device_id, algorithm, key_id, key_json;

\echo How many keys remain. Did we claim the right amount? \\

SELECT user_id, device_id, algorithm, count(*)
FROM e2e_one_time_keys_json
GROUP BY (user_id, device_id, algorithm)
ORDER BY user_id, device_id, algorithm;

psql output

nuke table
DROP TABLE
Time: 17.736 ms
Make table
CREATE TABLE
Time: 22.177 ms
Dummy data. Multiple users, devices, and algorithms.
INSERT 0 300
Time: 14.532 ms
How many keys did we insert?
   user_id   │  device_id  │ algorithm │ count 
═════════════╪═════════════╪═══════════╪═══════
 @alice:test │ alice_dev_1 │ algA      │   100
 @alice:test │ alice_dev_2 │ algA      │    50
 @alice:test │ alice_dev_2 │ algB      │    50
 @bob:test   │ bob_dev     │ algA      │    50
 @bob:test   │ bob_dev     │ algB      │    50
(5 rows)

Time: 0.227 ms
The query of doom.
   user_id   │  device_id  │ algorithm │ key_id │ key_json 
═════════════╪═════════════╪═══════════╪════════╪══════════
 @alice:test │ alice_dev_2 │ algA      │ key105 │ json105
 @alice:test │ alice_dev_1 │ algA      │ key61  │ json61
 @alice:test │ alice_dev_1 │ algA      │ key94  │ json94
 @bob:test   │ bob_dev     │ algB      │ key279 │ json279
 @alice:test │ alice_dev_2 │ algA      │ key125 │ json125
 @alice:test │ alice_dev_1 │ algA      │ key6   │ json6
 @alice:test │ alice_dev_1 │ algA      │ key82  │ json82
(7 rows)

DELETE 7
Time: 7.486 ms
How many keys remain. Did we claim the right amount?
   user_id   │  device_id  │ algorithm │ count 
═════════════╪═════════════╪═══════════╪═══════
 @alice:test │ alice_dev_1 │ algA      │    96
 @alice:test │ alice_dev_2 │ algA      │    48
 @alice:test │ alice_dev_2 │ algB      │    50
 @bob:test   │ bob_dev     │ algA      │    50
 @bob:test   │ bob_dev     │ algB      │    49
(5 rows)

Time: 0.279 ms