テナント毎のシーケンスを生成する方法

Azure Cosmos DB for PostgreSQL (Citus)でCREATE SEQUENCEすると、当たり前だがクラスタ全体のシーケンスとなってしまう。また、bigserialを含むテーブルを作成し分散テーブルにしても、同じ結果になる。

一方で、テナント毎にシーケンスが必要になる場合ももちろんあるので、どうやるか、という話。

CREATE TABLE tenants_seq (tenant_id bigserial UNIQUE, seq bigint);

SELECT create_distributed_table('tenants_seq', 'tenant_id');

CREATE OR REPLACE FUNCTION get_sequence(id bigint)
RETURNS bigint AS $$
DECLARE
    res bigint;
BEGIN
    res := 1;
    INSERT INTO tenants_seq (tenant_id, seq) VALUES (id, 1)
    ON CONFLICT (tenant_id)
        DO UPDATE SET seq = tenants_seq.seq + 1
        WHERE tenants_seq.tenant_id = id
        RETURNING seq INTO res;
    RETURN res;
END;
$$ LANGUAGE plpgsql;

get_sequence() の引数はテナントIDのみで、呼び出すたびにテナント毎のシーケンスを返してくる。