feat(datastore): add namespace_encryption_keys migration and NamespaceKeyStore

Suleimi Ahmedrequested to merge
sahmed/s04a-column-level-encryption-step-3 into main

Summary

  • Add the namespace_encryption_keys database migration: PARTITION BY HASH (namespace_id) into 64 partitions, composite (id, namespace_id) primary key (the partition key must appear in every unique constraint), FK to namespaces, a partial unique index enforcing single-active-key-per-namespace, a partial root_key_uri index (WHERE shredded_at IS NULL), and the shredded_at tombstone column for crypto-shredding.
  • Implement NamespaceEncryptionKeyStore — the PostgreSQL-backed crypto.NamespaceKeyStore with the S04-local FOR UPDATE barrier protocol, deactivate-before-insert, fresh version computation, a shred-guarded conditional insert (INSERT ... SELECT ... WHERE NOT EXISTS (tombstone) RETURNING), SQLSTATE 23505→ErrKeyConflict / 23503→ErrNamespaceNotFound mapping, set-based crypto-shredding, a 2x cache-TTL cooldown for re-enable, and WithTx for caller-controlled transaction composition.
  • UpdateWrappedKey takes namespaceID so the rotation re-wrap prunes to the single partition holding the row instead of probing all 64.
  • Crypto-shred fix: on an empty (never-seen) namespace the set-based tombstone sweep always runs after the version-1 tombstone upsert (no early return), so a concurrent first-write that commits between the barrier and the upsert cannot leave a live key.
  • Add integration tests covering all store methods, error mapping (including 23503→ErrNamespaceNotFound), mixed-state reads, lifecycle transitions, concurrent first-write, and the empty-namespace shred/insert race.

S04-A Step 3 — depends on Step 1 (crypto foundation).

Test plan

  • Migration applies and rolls back cleanly across all supported PostgreSQL versions (CI db:migrate matrix). The migration now creates a partitioned parent plus 64 child partitions, so timings differ from the pre-partition run and must be refreshed by CI.
  • Integration tests pass: go test -tags integration ./internal/datastore/...
  • Jet types match schema (CI jet:generate-check) — the composite PK adds sql:"primary_key" to namespace_id in the model and drops it from mutableColumns.
  • structure.sql regenerated (mise run db:dump-structure) — parent is PARTITION BY HASH, 64 child partitions in the partitions schema.
  • go vet and go build pass; golangci-lint reports 0 issues on ./internal/datastore/ ./internal/crypto/.
  • Squawk lint passes.

Related to https://gitlab.com/gitlab-org/ops/artifact-registry/-/work_items/384

Database Review Evidence

Migrations

Note

The migration was rewritten to PARTITION BY HASH (namespace_id) with 64 child partitions, so the apply/rollback cost changed materially from the pre-partition run and the CI db:migrate matrix numbers below must be re-captured. Interim local data point (ephemeral postgres:17-alpine, goose verbose): apply ≈ 0.14s, rollback ≈ 0.11s; up/down/up verified reversible. Each partition CREATE (Up) and DROP (Down) runs in its own implicit transaction under -- +goose NO TRANSACTION to stay within max_locks_per_transaction.

Migration PG 16 PG 17 PG 18
20260724120000_create_namespace_encryption_keys.sql refresh in CI refresh in CI (local apply/rollback ≈ 0.14s / 0.11s) refresh in CI

Queries

Note

Plans are from EXPLAIN (ANALYZE, BUFFERS) against an ephemeral PostgreSQL 17 container on the partitioned schema, seeded with namespaces=200, namespace_encryption_keys=5000 (25 versions per namespace, highest version active, all live, a single root_key_uri). Mutating statements ran inside a transaction rolled back per query. With genuinely multi-namespace data the planner prunes every namespace_id-scoped statement to exactly one of the 64 partitions — the target namespace maps to namespace_encryption_keys_p47, the empty namespace used for the first-write/tombstone probes maps to p38. Numbers reflect moderate cardinality and do not capture production-scale effects.

Method Plan node Index / partition Rows (plan / actual) Cost Time Buffers (hit / read) Partitions scanned
datastore.GetActiveKey.shredCheck Result + InitPlan Seq Scan p47 1 / 0 3.88 0.043ms 2 / 0 1 of 64 (pruned)
datastore.GetActiveKey.activeLookup Limit → Sort → Seq Scan p47 1 / 1 3.89 0.112ms 5 / 0 1 of 64 (pruned)
datastore.GetKeyByVersion Limit → Seq Scan p47 1 / 1 4.25 0.058ms 2 / 0 1 of 64 (pruned)
datastore.ListKeysForRootKeyURI Limit → Append (64 scans) index_ns_enc_keys_on_root_key_uri 1000 / 1000 34.28 0.422ms 19 / 0 64 of 64 (no prune)
datastore.UpdateWrappedKey Update → Seq Scan p47 1 / 1 4.25 0.257ms 23 / 2 1 of 64 (pruned)
datastore.insertKeyTx.deactivate Update → Seq Scan p47 1 / 1 3.88 0.191ms 25 / 0 1 of 64 (pruned)
datastore.insertKeyTx.insert Insert → Result (One-Time Filter on InitPlan) p38 1 / 1 1.96 0.698ms 7 / 0 1 of 64 (pruned)
datastore.lockAndCheckShredded LockRows → Seq Scan p47 25 / 25 4.12 0.062ms 27 / 0 1 of 64 (pruned)
datastore.nextVersion Result + InitPlan Index Only Scan Backward p47_namespace_id_version_key 1 / 1 0.65 0.042ms 2 / 0 1 of 64 (pruned)
datastore.shredNamespaceTx.lock LockRows → Seq Scan p47 25 / 25 4.12 0.062ms 27 / 0 1 of 64 (pruned)
datastore.shredNamespaceTx.tombstoneInsert Insert (ON CONFLICT) unique_ns_enc_keys_ns_version 1 / 1 0.02 1.460ms 157 / 0 1 of 64 (pruned)
datastore.shredNamespaceTx.tombstoneUpdate Update → Seq Scan p47 25 / 25 3.94 0.383ms 210 / 0 1 of 64 (pruned)
datastore.deleteNamespaceKeysTx Delete → Seq Scan p47 25 / 25 5.81 0.137ms 28 / 0 1 of 64 (pruned)

Query notes:

  • Partition pruning is the headline change. Every statement that carries a namespace_id equality prunes to exactly one of the 64 partitions (the single namespace_encryption_keys_pNN child in each plan). This replaces the previous evidence's "Seq Scan is a seed artifact; the index would be picked with multi-namespace data" caveats: with 200 namespaces the planner prunes to one small partition, and the within-partition Seq Scan over ~25 rows is optimal — a secondary index buys nothing on a 25-row relation.
  • Removed indexes. index_ns_enc_keys_on_namespace_id and the non-partial index_ns_enc_keys_root_key_uri no longer exist. namespace_id-equality lookups are served by partition pruning plus the per-partition UNIQUE(namespace_id, version) (which leads with namespace_id); root_key_uri lookups are served by the partial index_ns_enc_keys_on_root_key_uri (root_key_uri) WHERE shredded_at IS NULL.
  • ListKeysForRootKeyURI is the one statement without a namespace_id predicate, so it cannot prune: it fans out across all 64 partitions via an Append (each child a scan under the partial root_key_uri index; here a Seq Scan per partition because the single-root_key_uri seed makes every partition match). LIMIT 1000 stops the Append early — partitions past the 1000th row show never executed. This fan-out is inherent to a rotation sweep keyed only on root_key_uri and is bounded by the batch contract (repeat-until-empty; the predicate shrinks as rows are re-wrapped).
  • UpdateWrappedKey now takes namespace_id and filters on it, so the rotation re-wrap prunes to the single partition holding the row instead of scanning all 64.
  • insertKeyTx.insert is the shred-guarded conditional insert (INSERT ... SELECT ... WHERE NOT EXISTS (tombstone) RETURNING), planned as an Insert over a Result gated by a One-Time Filter: (NOT (InitPlan 1).col1); the InitPlan is the tombstone EXISTS probe, pruned to the namespace's partition (p38). The FK to namespaces fires as an insert-time trigger. On a committed tombstone the NOT EXISTS yields no row, RETURNING is empty (qrm.ErrNoRows), and the store maps it to ErrNamespaceShredded — the fail-closed first-write path.
  • shredNamespaceTx empty-namespace path now always sweeps. The version-1 tombstone Insert ... ON CONFLICT (arbiter unique_ns_enc_keys_ns_version) DO UPDATE is followed unconditionally by the set-based tombstone UPDATE ... WHERE namespace_id = N (the early return was removed). A concurrent first-write that commits between the barrier and the upsert is caught by the always-run UPDATE, which takes a fresh snapshot and tombstones every committed row (COALESCE preserves the first stamp). Covered by TestNamespaceEncryptionKeyStore_EmptyNamespaceShredSweep.
  • GetActiveKey.activeLookup now carries ORDER BY version DESC (the Sort node above the scan). The partial unique active index already guarantees a single active row; the ORDER BY ... LIMIT 1 makes the guard real — if that invariant ever regressed, the query returns the highest active version rather than an arbitrary row.
datastore.GetActiveKey.shredCheck

Summary: SELECT EXISTS(...) tombstone probe. Pruned to p47; the InitPlan short-circuits and execution stays well under 1ms. No anomalies.

Seed shape: namespaces=200, namespace_encryption_keys=5000

Rendered SQL:

SELECT (EXISTS (
          SELECT 1
          FROM public.namespace_encryption_keys
          WHERE (namespace_id = $1::uuid) AND (shredded_at IS NOT NULL)
     )) AS "is_shredded"

Plan:

Result  (cost=3.88..3.88 rows=1 width=1) (actual time=0.018..0.018 rows=1 loops=1)
  Buffers: shared hit=2
  InitPlan 1
    ->  Seq Scan on namespace_encryption_keys_p47 namespace_encryption_keys  (cost=0.00..3.88 rows=1 width=0) (actual time=0.017..0.017 rows=0 loops=1)
          Filter: ((shredded_at IS NOT NULL) AND (namespace_id = '...042'::uuid))
          Rows Removed by Filter: 150
          Buffers: shared hit=2
Planning Time: 1.014 ms
Execution Time: 0.043 ms
datastore.GetActiveKey.activeLookup

Summary: Pruned to p47, then Sort on version DESC + LIMIT 1. The ORDER BY version DESC is the deterministic-winner guard over the partial unique active index. 1 / 1 rows. No anomalies.

Rendered SQL:

SELECT id, namespace_id, version, wrapped_key, root_key_uri, active, created_at
FROM public.namespace_encryption_keys
WHERE (namespace_id = $1::uuid) AND (active = $2::boolean)
ORDER BY version DESC
LIMIT $3

Plan:

Limit  (cost=3.88..3.89 rows=1 width=61) (actual time=0.053..0.053 rows=1 loops=1)
  Buffers: shared hit=5
  ->  Sort  (cost=3.88..3.89 rows=1 width=61) (actual time=0.053..0.053 rows=1 loops=1)
        Sort Key: version DESC
        Sort Method: quicksort  Memory: 25kB
        ->  Seq Scan on namespace_encryption_keys_p47 namespace_encryption_keys  (cost=0.00..3.88 rows=1 width=61) (actual time=0.013..0.014 rows=1 loops=1)
              Filter: (active AND (namespace_id = '...042'::uuid))
              Rows Removed by Filter: 149
Execution Time: 0.112 ms
datastore.GetKeyByVersion

Summary: Pruned to p47; both namespace_id and version are equality filters. 1 / 1 rows. No anomalies.

Rendered SQL:

SELECT id, namespace_id, version, wrapped_key, root_key_uri, active, created_at
FROM public.namespace_encryption_keys
WHERE (namespace_id = $1::uuid) AND (version = $2::integer)
LIMIT $3

Plan:

Limit  (cost=0.00..4.25 rows=1 width=61) (actual time=0.011..0.012 rows=1 loops=1)
  ->  Seq Scan on namespace_encryption_keys_p47 namespace_encryption_keys  (cost=0.00..4.25 rows=1 width=61) (actual time=0.011..0.011 rows=1 loops=1)
        Filter: ((namespace_id = '...042'::uuid) AND (version = 13))
        Rows Removed by Filter: 75
        Buffers: shared hit=2
Execution Time: 0.058 ms
datastore.ListKeysForRootKeyURI

Summary: No namespace_id predicate, so no pruning: Limit over an Append across all 64 partitions (each child a scan under the partial root_key_uri index; Seq Scan here because the single-root_key_uri seed matches every partition). LIMIT 1000 stops the Append early — later partitions are never executed. Inherent to a root-key-scoped rotation sweep; bounded by the batch contract.

Rendered SQL:

SELECT id, namespace_id, version, wrapped_key, root_key_uri, active, created_at
FROM public.namespace_encryption_keys
WHERE (root_key_uri = $1::text) AND (shredded_at IS NULL)
LIMIT $2

Plan (head; 64-way Append elided):

Limit  (cost=0.00..34.28 rows=1000 width=61) (actual time=0.010..0.228 rows=1000 loops=1)
  Buffers: shared hit=19
  ->  Append  (cost=0.00..171.52 rows=5004 width=61) (actual time=0.009..0.179 rows=1000 loops=1)
        ->  Seq Scan on namespace_encryption_keys_p00 ...  (actual rows=50 loops=1)
              Filter: ((shredded_at IS NULL) AND (root_key_uri = 'local:root-v1'::text))
        ->  Seq Scan on namespace_encryption_keys_p01 ...  (actual rows=75 loops=1)
        ... (partitions p02..p14 executed; p15..p63 "never executed" once LIMIT 1000 is met) ...
Execution Time: 0.422 ms
datastore.UpdateWrappedKey

Summary: Now scoped by namespace_id, so the UPDATE prunes to p47 (one partition) rather than probing all 64. The shredded_at IS NULL guard makes a re-wrap racing a shred a no-op (0 rows → RotateAll skips the row). 1 / 1 rows.

Rendered SQL:

UPDATE public.namespace_encryption_keys
SET wrapped_key = $1, root_key_uri = $2
WHERE (namespace_id = $3::uuid) AND (id = $4::uuid) AND (shredded_at IS NULL)

Plan:

Update on namespace_encryption_keys  (actual time=0.174..0.174 rows=0 loops=1)
  Update on namespace_encryption_keys_p47 namespace_encryption_keys_1
  Buffers: shared hit=23 read=2
  ->  Seq Scan on namespace_encryption_keys_p47 namespace_encryption_keys_1  (cost=0.00..4.25 rows=1 width=74) (actual time=0.016..0.016 rows=1 loops=1)
        Filter: ((shredded_at IS NULL) AND (namespace_id = '...042'::uuid) AND (id = '...'::uuid))
        Rows Removed by Filter: 149
Execution Time: 0.257 ms
datastore.insertKeyTx.deactivate

Summary: Clears active on the namespace's current active row before the new active row is inserted (the partial unique active index is non-deferrable — insert-first self-violates). Pruned to p47. 1 / 1 rows.

Rendered SQL:

UPDATE public.namespace_encryption_keys
SET active = $1
WHERE (namespace_id = $2::uuid) AND (active = $3::boolean)

Plan:

Update on namespace_encryption_keys  (actual time=0.109..0.109 rows=0 loops=1)
  Update on namespace_encryption_keys_p47 namespace_encryption_keys_1
  ->  Seq Scan on namespace_encryption_keys_p47 namespace_encryption_keys_1  (cost=0.00..3.88 rows=1 width=11) (actual time=0.018..0.018 rows=1 loops=1)
        Filter: (active AND (namespace_id = '...042'::uuid))
        Rows Removed by Filter: 149
Execution Time: 0.191 ms
datastore.insertKeyTx.insert

Summary: The shred-guarded conditional insert. Insert over a Result gated by One-Time Filter: (NOT (InitPlan 1).col1); the InitPlan is the tombstone EXISTS probe, pruned to the namespace's partition (p38 for the empty namespace). The FK to namespaces fires as an insert-time trigger. A committed tombstone makes the filter drop the row → empty RETURNING → ErrNamespaceShredded.

Rendered SQL:

INSERT INTO public.namespace_encryption_keys (id, namespace_id, version, wrapped_key, root_key_uri, active)
SELECT $1::uuid, $2::uuid, $3, $4, $5, $6
WHERE NOT EXISTS (
    SELECT 1 FROM public.namespace_encryption_keys
    WHERE (namespace_id = $2::uuid) AND (shredded_at IS NOT NULL)
)
RETURNING id, namespace_id, version, wrapped_key, root_key_uri, active, created_at

Plan:

Insert on namespace_encryption_keys  (cost=1.94..1.96 rows=1 width=117) (actual time=0.142..0.143 rows=1 loops=1)
  ->  Subquery Scan on "*SELECT*"  (actual time=0.025..0.026 rows=1 loops=1)
        ->  Result  (cost=1.94..1.95 rows=1 width=101) (actual time=0.024..0.024 rows=1 loops=1)
              One-Time Filter: (NOT (InitPlan 1).col1)
              InitPlan 1
                ->  Seq Scan on namespace_encryption_keys_p38 namespace_encryption_keys_1  (cost=0.00..1.94 rows=1 width=0) (actual time=0.017..0.018 rows=0 loops=1)
                      Filter: ((shredded_at IS NOT NULL) AND (namespace_id = '...9999'::uuid))
                      Rows Removed by Filter: 75
Trigger for constraint fk_ns_enc_keys_namespace_id_namespaces on namespace_encryption_keys_p38: time=0.490 calls=1
Execution Time: 0.698 ms
datastore.lockAndCheckShredded / shredNamespaceTx.lock

Summary: The S04-local FOR UPDATE barrier — LockRows over a scan pruned to p47. Projects the primary key alongside shredded_at so every locked row maps to a result element (qrm iterates the full set, acquiring the row locks). 25 / 25 rows.

Rendered SQL:

SELECT id, shredded_at
FROM public.namespace_encryption_keys
WHERE (namespace_id = $1::uuid)
FOR UPDATE

Plan:

LockRows  (cost=0.00..4.12 rows=25 width=34) (actual time=0.018..0.038 rows=25 loops=1)
  ->  Seq Scan on namespace_encryption_keys_p47 namespace_encryption_keys  (cost=0.00..3.88 rows=25 width=34) (actual time=0.008..0.023 rows=25 loops=1)
        Filter: (namespace_id = '...042'::uuid)
        Rows Removed by Filter: 125
Execution Time: 0.062 ms
datastore.nextVersion

Summary: MAX(version) computed after the barrier from a fresh snapshot. Pruned to p47; served by an Index Only Scan Backward on the partition's (namespace_id, version) unique index. 1 / 1 rows.

Rendered SQL:

SELECT MAX(version) AS "max_version"
FROM public.namespace_encryption_keys
WHERE (namespace_id = $1::uuid)

Plan:

Result  (cost=0.64..0.65 rows=1 width=4) (actual time=0.022..0.022 rows=1 loops=1)
  InitPlan 1
    ->  Limit  (actual time=0.020..0.020 rows=1 loops=1)
          ->  Index Only Scan Backward using namespace_encryption_keys_p47_namespace_id_version_key on namespace_encryption_keys_p47 namespace_encryption_keys  (cost=0.14..12.49 rows=25 width=4) (actual time=0.019..0.019 rows=1 loops=1)
                Index Cond: (namespace_id = '...042'::uuid)
                Heap Fetches: 1
Execution Time: 0.042 ms
datastore.shredNamespaceTx.tombstoneInsert (empty-namespace path)

Summary: Version-1 tombstone Insert ... ON CONFLICT DO UPDATE; the conflict arbiter is unique_ns_enc_keys_ns_version, which serializes the empty-namespace shred against a concurrent first-write auto-create. Followed unconditionally by the set-based sweep (see notes). FK trigger on p38.

Rendered SQL:

INSERT INTO public.namespace_encryption_keys (id, namespace_id, version, wrapped_key, root_key_uri, active, shredded_at)
VALUES ($1::uuid, $2::uuid, 1, ''::bytea, '', false, clock_timestamp())
ON CONFLICT (namespace_id, version)
DO UPDATE SET wrapped_key = ''::bytea, active = false,
    shredded_at = COALESCE(namespace_encryption_keys.shredded_at, clock_timestamp())

Plan:

Insert on namespace_encryption_keys  (cost=0.00..0.02 rows=0 width=0) (actual time=0.460..0.460 rows=0 loops=1)
  Conflict Resolution: UPDATE
  Conflict Arbiter Indexes: unique_ns_enc_keys_ns_version
  Tuples Inserted: 1
  Conflicting Tuples: 0
  ->  Result  (actual time=0.006..0.006 rows=1 loops=1)
Trigger for constraint fk_ns_enc_keys_namespace_id_namespaces on namespace_encryption_keys_p38: time=0.621 calls=1
Execution Time: 1.460 ms
datastore.shredNamespaceTx.tombstoneUpdate

Summary: Set-based tombstone sweep — zero wrapped_key, clear active, COALESCE(shredded_at, clock_timestamp()). Pruned to p47. Now always runs on the empty-namespace path too, so a row a concurrent insert committed between the barrier and the upsert is swept. 25 / 25 rows.

Rendered SQL:

UPDATE public.namespace_encryption_keys
SET wrapped_key = $1, active = $2, shredded_at = COALESCE(shredded_at, clock_timestamp())
WHERE (namespace_id = $3::uuid)

Plan:

Update on namespace_encryption_keys  (actual time=0.287..0.287 rows=0 loops=1)
  Update on namespace_encryption_keys_p47 namespace_encryption_keys_1
  Buffers: shared hit=210 dirtied=1 written=1
  ->  Seq Scan on namespace_encryption_keys_p47 namespace_encryption_keys_1  (cost=0.00..3.94 rows=25 width=51) (actual time=0.008..0.021 rows=25 loops=1)
        Filter: (namespace_id = '...042'::uuid)
        Rows Removed by Filter: 125
Execution Time: 0.383 ms
datastore.deleteNamespaceKeysTx

Summary: Re-enable hard-delete of the namespace's key rows (after the shred + 2x cache-TTL cooldown). Pruned to p47. 25 / 25 rows.

Rendered SQL:

DELETE FROM public.namespace_encryption_keys
WHERE (namespace_id = $1::uuid)

Plan:

Delete on namespace_encryption_keys  (actual time=0.062..0.062 rows=0 loops=1)
  Delete on namespace_encryption_keys_p47 namespace_encryption_keys_1
  ->  Seq Scan on namespace_encryption_keys_p47 namespace_encryption_keys_1  (cost=0.00..5.81 rows=38 width=10) (actual time=0.016..0.044 rows=25 loops=1)
        Filter: (namespace_id = '...042'::uuid)
        Rows Removed by Filter: 125
Execution Time: 0.137 ms

e2e scenarios (Guardrail 11)

No docs/testing/ change — the store has no user-facing surface. It exposes no protocol handler, management API, or CLI; it is exercised only through internal/crypto and (later) the S13 credential tables, so no e2e scenario catalog is added or affected.

Edited by Suleimi Ahmed

Merge request reports