fix(state-pg): reclaim expired set_if_not_exists rows, support migration-owned schemas (#240)

[patrick-chinchill]

Summary

Ports two upstream Postgres state changes:

• Expired claims are reclaimed. PostgresStateAdapter.set_if_not_exists used ON CONFLICT DO NOTHING, so an expired row in chat_state_cache blocked every new claim until a get() of that exact key deleted it. Dedupe keys and leases built on set_if_not_exists (for example Telegram update_id claims from [4.41/TG0] Telegram: require webhook verification by default, dedupe repeated updates #224) refused work they should accept. The query is now upstream's conditional upsert:

  ON CONFLICT (key_prefix, cache_key) DO UPDATE SET value = EXCLUDED.value, expires_at = EXCLUDED.expires_at, updated_at = now()
    WHERE chat_state_cache.expires_at IS NOT NULL AND chat_state_cache.expires_at <= now()
  RETURNING cache_key

  The result is read with fetchval(...) is not None. Live and permanent (no-TTL) rows are never overwritten. ttl_ms falsy still means permanent, as upstream.

• Migration-owned schemas. New keyword-only auto_create_schema: bool | None = True (None means True, as upstream's ?? true) on PostgresStateAdapter and create_postgres_state. With False, connect() runs SELECT 1 and then one read-only probe, and never runs DDL. The probe checks per-table DELETE, per-column SELECT/INSERT/UPDATE, and the list/queue seq sequence (either an identity column or USAGE, UPDATE). If anything is missing, connect() raises StateSchemaError("PostgreSQL state schema is not ready: …") naming the problem (the first relation PostgreSQL cannot resolve, or every object whose grants are missing), followed by the hint. A probe failure is chained as __cause__.

• New public API. POSTGRES_SCHEMA_STATEMENTS (tuple[str, ...], the nine statements in execution order), exported from chat_sdk.state.postgres and chat_sdk.state. Also chat_sdk.StateSchemaError(ChatError). The README has a new "PostgreSQL state" section with the migration SQL and grants, and a unit test keeps that SQL identical to POSTGRES_SCHEMA_STATEMENTS.

Verify first (from the issue)

I ran the new query through asyncpg against a real PostgreSQL 16.0, an embedded pgserver instance. Docker was not running.

Existing row	execute tag	fetchval result	Row after
absent	INSERT 0 1	the key	new value
live (expires in 1h)	INSERT 0 0	None	unchanged
expired (1h ago)	INSERT 0 1	the key	new value
permanent (NULL)	INSERT 0 0	None	unchanged

The tags agree with the results, but RETURNING is the less ambiguous signal, so I implemented the recommended fetchval(... RETURNING cache_key) is not None.

Upstream commits mapped

Upstream	Python
d88789c9 fix(state-pg): setIfNotExists TTL expiry (vercel/chat#636, chat@4.35.0)	PostgresStateAdapter.set_if_not_exists
ea025af7 feat(postgres): allow migration-managed schemas (vercel/chat#913, chat@4.41.0)	auto_create_schema, POSTGRES_SCHEMA_STATEMENTS, _TABLE_PRIVILEGES / _SCHEMA_PROBE / _verify_schema, StateSchemaError, README section

I re-derived _TABLE_PRIVILEGES from Python's SQL in postgres.py, covering conflict targets, WHERE/RETURNING columns, and the UPDATE sets for set, set_if_not_exists, acquire_lock, extend_lock and the list TTL refresh. It matches upstream's map exactly. The live column-grant test below confirms that every adapter statement works with exactly those column grants.

Tests ported (tests/test_state_postgres.py)

• index.test.ts, "should allow setIfNotExists to replace expired keys": test_should_allow_set_if_not_exists_to_replace_expired_keys. It replaces the phantom test_succeeds_after_expired_key, which passed only because the mock's DO NOTHING branch reclaimed expired rows. Real Postgres does not.
• describe("schema initialization"), all 8 ported into TestPostgresStateSchemaInitialization:
  • creates every table and index when autoCreateSchema is %s
  • keeps the published migration in sync …
  • probes instead of creating the schema for an external pool via %s
  • probes only the privileges each table needs
  • rejects connect() when a migration-owned table is missing
  • rejects connect() naming every object …
  • forwards opt-out for a URL from %s and closes the owned pool
  • retries a failed connectivity check without DDL (adapted; see Divergences)
• postgres.integration.test.ts: all 9 tests ported into TestPostgresMigrationOwnedSchemaIntegration. This is a superset of the 5 the issue named. It adds the column-grant, the naming, the expired-renew (TTL / no TTL) and the 8-way concurrent one-winner tests. The suite is skipped unless POSTGRES_TEST_URL is set, and never reads POSTGRES_URL. It needs asyncpg, which the dev group does not include. Locally, against PG 16: POSTGRES_TEST_URL=… uv run --with asyncpg pytest tests/test_state_postgres.py → 89 passed (at 5f2c6eb).
• Python-specific tests: test_permanent_key_is_never_reclaimed, test_zero_ttl_claim_is_permanent, test_concurrent_connect_retries_after_a_failed_connectivity_check. test_fails_when_key_exists now uses a live TTL row and asserts that its expiry is unchanged.
• Mock fixes: the DO NOTHING branch never overwrites. A conditional-upsert branch is dispatched before the generic DO UPDATE branch, in execute, fetchrow and fetchval. The clock is injectable (MockAsyncpgPool.advance), so every asyncio.sleep in the file is gone, including a 150 ms one. The weak test_connect_creates_tables (>= 5 CREATEs) was replaced by the exact-sequence upstream port.
• Mutation-checked, each caught by at least one test, most by unit and live tests:
  • reverting to DO NOTHING
  • dropping the IS NOT NULL guard
  • always running DDL
  • treating a NULL probe value as missing
  • dropping updated_at from the cache UPDATE map
  • losing the __cause__ chain
  • the factory dropping the flag

Fidelity

TS_ROOT=<chat@4.41.1> uv run python scripts/verify_test_fidelity.py --report-target:

Delta vs committed report (HEAD): missing 282 -> 282 (+0)

packages/state-pg is not in MAPPING/TARGET_MAPPING, so scripts/fidelity_target.json is unchanged and not part of this diff. Strict at the pin is still 733/733.

Divergences (3, within budget; all recorded in docs/UPSTREAM_SYNC.md under Known Non-Parity, plus a new "Postgres state: expired claims and migration-owned schemas" section)

1. Error type. Python raises StateSchemaError(ChatError); upstream throws a plain Error. The message prefix and structure are identical. The hint names auto_create_schema=True.
2. Concurrent connect(). This behavior predates this PR and is kept. Python serializes connect() on an asyncio.Lock, so a caller queued behind a failed attempt retries it. Upstream shares one in-flight promise, so all concurrent callers reject together. The upstream test retries a failed connectivity check without DDL is ported sequentially: it asserts that no DDL or probe runs after the failed SELECT 1, get() raises not-connected, and a later connect() succeeds. test_concurrent_connect_retries_after_a_failed_connectivity_check pins the Python behavior, and a breadcrumb sits at the lock.
3. Owned pool closed on a failed connect(). asyncpg opens min_size (10) connections eagerly, and disconnect() is a no-op until connect() succeeds, so a startup retry loop leaked a full pool per failed attempt (live on PG 16: 0 → 10 → 20 → 30 client backends). Upstream's lazy pg.Pool keeps at most one idle client. Python now closes and resets a pool that the failed (or cancelled) attempt created; the next attempt builds a fresh one (same live script: 0 → 0 → 0 → 0). An injected pool is never closed. Pinned by test_closes_the_owned_pool_when_connect_fails.

Consumer impact

• Postgres state users only. Slack, Teams and the other adapters are unaffected unless they use Postgres state.
• Strict fix: dedupe and leases on Postgres now recover from expired rows without cleanup. No schema change, no migration.
• Opt-in: auto_create_schema=False. The default behavior is unchanged.
• New exports: POSTGRES_SCHEMA_STATEMENTS (in chat_sdk.state) and StateSchemaError (in chat_sdk).

Validation

• ruff check, ruff format --check, audit_test_quality.py (0 hard failures), verify_test_fidelity.py --check-docs, and --strict at chat@4.31.0 (all TS tests have Python equivalents): all pass.
• pytest tests/: 5580 passed, 24 skipped (main baseline: 5561 passed).
• pyrefly check: 0 errors.

Merge gate

Final HEAD: 5f2c6eb (branch already contains origin/main 4817dc0; no merge needed).

Independent review findings (4), all fixed in 5f2c6eb:

1. Owned asyncpg pool leaked on a failed connect() (minor, real). Fixed: a pool created by the failed or cancelled attempt is closed and reset; recorded as divergence 3 above. test_closes_the_owned_pool_when_connect_fails[schema|connectivity|cancelled] fails on the old code; the injected-pool test now asserts close_calls == 0. Verified live with the reviewer's script (0/0/0 backends instead of 10/20/30).
2. README/CHANGELOG said "naming every missing table or grant" (minor, real). Reworded: the first unresolvable relation is named for missing tables; all objects are named for missing grants. docs/UPSTREAM_SYNC.md updated the same way.
3. Most per-column _TABLE_PRIVILEGES entries unpinned (minor, real). test_probes_only_the_privileges_each_table_needs now asserts the exact set of has_column_privilege terms against a literal copy of upstream's tablePrivileges map; the reviewer's drop_token_select mutation now fails it. The live revoke loop also revokes SELECT (token) ON chat_state_locks (live suite 89 passed).
4. auto_create_schema=None switched to verify-only (minor, real). Now None means True, matching upstream's ?? true; pinned by the new none case of test_creates_every_table_and_index_when_auto_create_schema_is (fails on the old code).

Declined: none.

gpt-6-astra: 1 round on 5f2c6eb, verdict: no actionable findings.

Bots: CodeRabbit skipped the draft and was rate-limited after ready-for-review; no inline comments or reviews. No gemini comments.

CI on 5f2c6eb: test (3.12), test (3.13), CodeQL, Analyze (actions/python) and Lint & Type Check all pass.

Local full validation on 5f2c6eb: ruff check + format, audit (0 hard failures), --check-docs, --strict at chat@4.31.0 (all TS tests have Python equivalents), pytest 5580 passed / 24 skipped, pyrefly 0 errors. Fidelity report-target: Delta vs committed report (HEAD): missing 282 -> 282 (+0).

Closes #240
Part of #184

[coderabbitai]

Warning

Review limit reached

Next included review available in 50 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used the included review currently available.

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

Learn how review limits work.

Review configuration:

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Advanced

Run ID: 8f9752f1-72a6-4164-a87e-3b653dd23e61

📥 Commits

Reviewing files that changed from the base of the PR and between 1d0eb49 and 1c32ba9.

📒 Files selected for processing (9)

• CHANGELOG.md
• README.md
• docs/TESTING.md
• docs/UPSTREAM_SYNC.md
• src/chat_sdk/__init__.py
• src/chat_sdk/errors.py
• src/chat_sdk/state/__init__.py
• src/chat_sdk/state/postgres.py
• tests/test_state_postgres.py

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[patrick-chinchill]

Merge gate: CI green (Lint & Type Check, test (3.12), test (3.13), Analyze (python), Analyze (actions), CodeQL) on 1c32ba9; local Codex review (gpt-6-astra, xhigh, --base origin/main) on 5f2c6eb: "No actionable regressions found relative to the specified merge base. PostgreSQL unit tests passed (78 passed); lint and formatting checks also passed." (11 live PostgreSQL tests skipped without POSTGRES_TEST_URL); converged pre-merge astra rounds, no new round needed after merging origin/main (#202): only CHANGELOG.md / docs/UPSTREAM_SYNC.md conflicted (additive, both sides kept), and the PR's src/tests/scripts diff vs main is line-for-line identical to the reviewed diff (1220/1220 changed lines; #202 touches cards/modals, this PR touches state/postgres + errors); fidelity_target.json regenerated, delta +0; local full validation 5623 passed, strict 733/733, pyrefly 0 errors; CodeRabbit rate-limited (no findings). Merging with --admin (Protect Main requires a code-owner approval).