feat(bb-distributed-data): atomic Counter primitive for DSQL

.changeset/atomic-counter-primitive.md

@@ -0,0 +1,5 @@
+---
+"@aws-blocks/bb-distributed-data": minor
+---
+
+Add an atomic `Counter` primitive to `DistributedDatabase`. `db.counter(name).next()` returns a race-free monotonic sequence value (with `.current()` and `.reset()`), backed by a single `INSERT ... ON CONFLICT DO UPDATE ... RETURNING` upsert against a framework-managed `_blocks_counters` table (created automatically on deploy). This replaces the racy `SELECT MAX(seq) + 1` read-modify-write pattern that DSQL otherwise forces, since it has no sequences (`SERIAL` / `BIGSERIAL`).

packages/bb-distributed-data/API.md

@@ -6,18 +6,29 @@
 
 import type { ChildLogger } from '@aws-blocks/bb-logger';
 import { createKyselyAdapter } from '@aws-blocks/data-common';
+import type { DatabaseBase } from '@aws-blocks/data-common';
 import { DatabaseEngine } from '@aws-blocks/data-common';
 import { Scope } from '@aws-blocks/core';
 import type { ScopeParent } from '@aws-blocks/core';
 import { sql } from '@aws-blocks/data-common';
 import { SqlQuery } from '@aws-blocks/data-common';
 import { Transaction } from '@aws-blocks/data-common';
 
+// @public
+export class Counter {
+    // @internal
+    constructor(resolveBase: () => Promise<DatabaseBase>, name: string);
+    current(): Promise<number>;
+    next(delta?: number): Promise<number>;
+    reset(value?: number): Promise<void>;
+}
+
 export { createKyselyAdapter }
 
 // @public (undocumented)
 export class DistributedDatabase extends Scope {
     constructor(scope: ScopeParent, id: string, _options?: DistributedDatabaseOptions);
+    counter(name: string): Counter;
     // (undocumented)
     execute(query: SqlQuery): Promise<{
         rowCount: number;

packages/bb-distributed-data/DESIGN.md

@@ -120,6 +120,28 @@ async transaction<T>(fn, options?) {
 
 Default: no retry (honest, predictable). `retryOnConflict: true` is explicit opt-in with JSDoc warning about side effects.
 
+## Atomic Counters
+
+DSQL has no sequences (`SERIAL` / `BIGSERIAL` are rejected by the validator), so code that needs a monotonic sequence number is otherwise pushed into a racy `SELECT MAX(seq) + 1` read-modify-write: two concurrent callers read the same maximum and both write the same next value, producing duplicate or skipped numbers.
+
+`db.counter(name)` replaces that pattern with a single atomic upsert against a framework-managed table:
+
+```sql
+CREATE TABLE IF NOT EXISTS _blocks_counters (name TEXT PRIMARY KEY, value BIGINT NOT NULL DEFAULT 0);
+
+-- Counter.next(delta) — atomic read-and-increment in one statement:
+INSERT INTO _blocks_counters (name, value) VALUES ($name, $delta)
+  ON CONFLICT (name) DO UPDATE SET value = _blocks_counters.value + $delta
+  RETURNING value;
+```
+
+- **Atomicity** — the increment is one `INSERT ... ON CONFLICT DO UPDATE ... RETURNING` statement, so there is no read-then-write window for a competing writer. It runs inside `transaction({ retryOnConflict: true })` so a DSQL OCC conflict on the row retries transparently.
+- **Managed table** — `_blocks_counters` is created by the migration Lambda (admin) on every deploy, like `_migrations`. The app runtime is DML-only and never issues the DDL. Because the table is admin-owned and not named `_migrations`, `provisionAppRole` automatically grants the app role `SELECT/INSERT/UPDATE/DELETE` on it.
+- **Mock parity** — `DsqlMockEngine` creates the same table inside its `withDdl` init window, so `db.counter()` behaves identically against PGlite locally.
+- **Range** — values are stored as `BIGINT` and returned as `number` (exact up to 2^53 − 1).
+
+API: `Counter.next(delta = 1)`, `Counter.current()`, `Counter.reset(value = 0)`.
+
 ## Error Translation
 
 Happens in the engine layer (same pattern as bb-data engines):
@@ -155,6 +177,7 @@ The `DistributedDatabase` class does not wrap errors — engines handle translat
 - Retries with exponential backoff on transient connection errors (cluster may take a moment after creation)
 - `.sql` files bundled into the Lambda package via CDK `commandHooks`
 - `migrationsHash` property triggers re-invocation when files change
+- **Ensures the `_blocks_counters` table** exists (see [Atomic Counters](#atomic-counters)) — created here as admin so the DML grants below cover it
 - **Provisions custom DB role** on every deploy (idempotent):
   1. `CREATE ROLE "app_role" WITH LOGIN` (if not exists)
   2. `AWS IAM GRANT "app_role" TO 'arn:aws:iam::...:role/...'` (maps IAM to DB role)

packages/bb-distributed-data/README.md

@@ -61,6 +61,25 @@ interface TransactionOptions {
 }
 ```
 
+## Atomic Counters
+
+DSQL has no sequences (`SERIAL` / `BIGSERIAL`), so a monotonic sequence number would otherwise force a racy `SELECT MAX(seq) + 1` read-modify-write — two concurrent callers read the same maximum and both write the same next value, producing duplicate or skipped numbers. `db.counter(name)` gives you a race-free named counter instead: each operation is a single atomic upsert wrapped in OCC retry.
+
+```typescript
+// Per-user monotonic sequence, safe under concurrency:
+const seq = await db.counter(`notes:${userId}`).next();
+await db.execute(
+  sql`INSERT INTO notes (id, user_id, seq, body) VALUES (${id}, ${userId}, ${seq}, ${body})`
+);
+
+await db.counter('signups').next();      // increment by 1, returns the new value
+await db.counter('score').next(10);      // increment by a custom (integer) delta
+await db.counter('signups').current();   // read without incrementing (0 if never seen)
+await db.counter('signups').reset();     // set back to 0
+```
+
+Counters are stored in a framework-managed `_blocks_counters` table that is created automatically on deploy — no migration required. Values are stored as `BIGINT` and returned as `number` (exact up to `Number.MAX_SAFE_INTEGER`).
+
 ## Migrations
 
 One DDL statement per file. DML in separate files. This matches DSQL's transaction constraints.
@@ -102,7 +121,7 @@ DSQL is a subset of PostgreSQL. The local mock enforces these restrictions so co
 | Triggers | Event-driven logic (EventBridge, Lambda) |
 | Views | CTEs or application-layer query composition |
 | PL/pgSQL functions | `LANGUAGE SQL` functions or app logic |
-| SERIAL / BIGSERIAL | UUIDs (`gen_random_uuid()`) |
+| SERIAL / BIGSERIAL | UUIDs (`gen_random_uuid()`) for IDs; `db.counter()` for monotonic sequences |
 | TRUNCATE | `DELETE FROM` |
 | Temporary tables | CTEs or subqueries |
 | LISTEN / NOTIFY | AppSync Events, EventBridge, or polling |

packages/bb-distributed-data/src/counter.test.ts

@@ -0,0 +1,158 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Tests for the atomic Counter primitive (DSQL upsert semantics).
+ *
+ * Runs against the mock engine (PGlite). The atomic upsert replaces the racy
+ * `SELECT MAX(seq) + 1` read-modify-write that DSQL otherwise forces (no
+ * sequences), so the headline test asserts no duplicate values under
+ * concurrent increments.
+ */
+
+import { describe, it, before, after, beforeEach } from 'node:test';
+import assert from 'node:assert/strict';
+import { rmSync } from 'node:fs';
+import { DatabaseBase, sql } from '@aws-blocks/data-common';
+import { DsqlMockEngine } from './engines/dsql-mock-engine.js';
+import { Counter, ensureCounterTable, COUNTER_TABLE } from './counter.js';
+import { DistributedDatabase } from './index.mock.js';
+
+const DIR = '.bb-data/__test_counter__';
+
+describe('Counter (mock)', () => {
+  let engine: DsqlMockEngine;
+  let base: DatabaseBase;
+  const counter = (name: string) => new Counter(async () => base, name);
+
+  before(async () => {
+    rmSync(DIR, { recursive: true, force: true });
+    engine = new DsqlMockEngine(DIR);
+    base = new DatabaseBase(engine);
+    // The counter table is DDL — created by admin/migrations in prod. Locally
+    // it goes through the same withDdl escape hatch the migration runner uses.
+    await engine.withDdl(() => ensureCounterTable(engine));
+  });
+
+  after(async () => {
+    await engine.destroy();
+    rmSync(DIR, { recursive: true, force: true });
+  });
+
+  beforeEach(async () => {
+    await base.execute(sql`DELETE FROM _blocks_counters`);
+  });
+
+  it('next() returns a monotonic sequence starting at 1', async () => {
+    const c = counter('seq');
+    assert.equal(await c.next(), 1);
+    assert.equal(await c.next(), 2);
+    assert.equal(await c.next(), 3);
+  });
+
+  it('next(delta) adds an arbitrary integer, including negatives', async () => {
+    const c = counter('by-tens');
+    assert.equal(await c.next(10), 10);
+    assert.equal(await c.next(5), 15);
+    assert.equal(await c.next(-3), 12);
+  });
+
+  it('current() reads without incrementing; unknown counter is 0', async () => {
+    const c = counter('reads');
+    assert.equal(await c.current(), 0);
+    await c.next();
+    await c.next();
+    assert.equal(await c.current(), 2);
+    assert.equal(await c.current(), 2);
+  });
+
+  it('reset() sets an explicit value', async () => {
+    const c = counter('resettable');
+    await c.next();
+    await c.next();
+    await c.reset(100);
+    assert.equal(await c.current(), 100);
+    assert.equal(await c.next(), 101);
+    await c.reset();
+    assert.equal(await c.current(), 0);
+  });
+
+  it('named counters are independent', async () => {
+    const a = counter('a');
+    const b = counter('b');
+    await a.next();
+    await a.next();
+    await a.next();
+    await b.next();
+    assert.equal(await a.current(), 3);
+    assert.equal(await b.current(), 1);
+  });
+
+  it('recovers from an OCC conflict via retry without double counting', async () => {
+    const c = counter('occ');
+    assert.equal(await c.next(), 1);
+    engine.simulateConflict();
+    // First commit throws 40001 and rolls back; retry succeeds. The rolled-back
+    // attempt must NOT leak its increment.
+    assert.equal(await c.next(), 2);
+    assert.equal(await c.current(), 2);
+  });
+
+  it('concurrent next() calls never produce a duplicate value', async () => {
+    const c = counter('race');
+    const N = 25;
+    const results = await Promise.all(Array.from({ length: N }, () => c.next()));
+    const unique = new Set(results);
+    assert.equal(
+      unique.size,
+      N,
+      `expected ${N} unique values, got ${unique.size}: ${[...results].sort((x, y) => x - y).join(',')}`,
+    );
+    assert.equal(Math.max(...results), N);
+    assert.equal(Math.min(...results), 1);
+    assert.equal(await c.current(), N);
+  });
+
+  it('rejects an empty counter name', () => {
+    assert.throws(() => counter(''), /non-empty string/);
+  });
+
+  it('rejects a non-integer delta or reset value', async () => {
+    const c = counter('validate');
+    await assert.rejects(() => c.next(1.5), /must be an integer/);
+    await assert.rejects(() => c.reset(0.5), /must be an integer/);
+  });
+});
+
+describe('DistributedDatabase.counter (mock integration)', () => {
+  const scope = { id: 'test' };
+  let db: InstanceType<typeof DistributedDatabase>;
+
+  before(async () => {
+    db = new DistributedDatabase(scope as any, 'counterint');
+    // Clear any state persisted from a previous run of this data dir.
+    await db.execute(sql`DELETE FROM _blocks_counters`);
+  });
+
+  after(async () => {
+    await (db as any).mockEngine.destroy();
+  });
+
+  it('auto-creates the counter table with no migrations configured', async () => {
+    // No migrationsPath was provided, yet counter() works — proving the table
+    // is created during init (mirrors the migration Lambda creating it in prod).
+    const seq = db.counter('notes:user-1');
+    assert.equal(await seq.next(), 1);
+    assert.equal(await seq.next(), 2);
+    assert.equal(await db.counter('notes:user-2').next(), 1);
+  });
+
+  it('persists to the framework-managed _blocks_counters table', async () => {
+    assert.equal(COUNTER_TABLE, '_blocks_counters');
+    await db.counter('direct-check').next(7);
+    const row = await db.queryOne<{ value: string | number }>(
+      sql`SELECT value FROM _blocks_counters WHERE name = ${'direct-check'}`,
+    );
+    assert.equal(Number(row?.value), 7);
+  });
+});