Skip to Content
Internal docs are powered by Nextra Docs Theme.
ProjectsMove query_data to API2Plan

query_data to api2: Phased Migration

Date: 2026-05-15 Status: All three phases shipped. Scope: Move the agent-facing query_data tool out of Rails and onto api2 (Node) in three incremental phases, ending with a database-enforced read-only role for agent queries.


Why

Before Phase 1, the api2 query_data tool was a thin shell that did one extra HTTP round-trip into Rails for every agent SQL query:

agent → api2 (queryData tool) → POST /internal/tools/call (Rails) → ToolCalls::ComplianceAgent::QueryData → Postgres

All validation (PgQuery), session config (SET LOCAL statement_timeout + read-only + app.current_organization_id), execution, and a 5000-row persistence step into stored_query_results happened in Ruby. The HTTP hop added latency and made tool spans harder to reason about in Braintrust traces.

Phase 1 collapses that to a single in-process call against the shared Postgres. Phases 2 and 3 then move the q_ views and the access-control story onto api2 so the eventual end state is “agent runs as a least-privileged Postgres role against Drizzle-defined views — Ruby is no longer involved in agent reads.”


Phase 1 — Port validator + executor to api2 (shipped 2026-05-15)

What changed

Two new modules in api2:

  • rulebase-web/rulebase-api2/src/lib/query-data-validator.tslibpg-query (WASM) parses the agent’s SQL. Rejects empty input, multi-statement queries, non-SelectStmt, and any reference whose RangeVar.relname does not start with q_. Recursively walks the parse tree so subqueries, CTE bodies, lateral joins, set ops, and schema-qualified refs are all covered. CommonTableExpr.ctename is excluded so a CTE alias like WITH recent AS (...) doesn’t trip the prefix check. Returns { ok: true } | { ok: false, reason } — no throws.
  • rulebase-web/rulebase-api2/src/lib/query-data-runner.ts — opens a Drizzle transaction against the existing api2 pool and, inside it:
    SET LOCAL statement_timeout = 20000; SET LOCAL transaction_read_only = true; SELECT set_config('app.current_organization_id', $orgId, true); -- then the user's query (raw)
    Caps at STORAGE_ROW_LIMIT = 5000, PREVIEW_ROW_LIMIT = 200. Returns { columns, row_count, truncated, preview_rows } | { error }. One small upgrade over the Ruby version: set_config(..., $1, true) uses a bind parameter instead of the Ruby tool’s "SET LOCAL ... = '#{org_id}'" string interpolation, removing an unnecessary interpolation surface.

The runner accepts an injectable database parameter so tests pass a pglite instance instead of the prod pool.

Wire-up + cleanup

  • rulebase-web/rulebase-api2/src/lib/agent-query-data-tools.ts — swapped the queryData() HTTP call for runQueryData(). The data: z.unknown() output schema is structurally compatible; the agent only reads columns / row_count / truncated / preview_rows, all of which are preserved.
  • rulebase-web/rulebase-api2/src/lib/rails-client.ts — deleted the queryData() export. No other callers existed in the repo.
  • rulebase-web/rulebase-api2/package.json — added libpg-query@^17.7.3.

Response shape change

The Rails tool returned { query_result_id, columns, row_count, truncated, preview_rows, message }. query_result_id referenced a row in Rails’s stored_query_results table, used only for Rails-side retrospective inspection. Nothing in api2 or rulebase-ui consumed it (verified via repo-wide grep). The Node executor returns { columns, row_count, truncated, preview_rows } and does not persist anything — Braintrust already logs the tool input and output, and if a “view full result” UI is needed later we’ll add an api2.stored_query_results Drizzle table at that point.

Risk that disappeared

The plan originally flagged libpg-query’s native-C-addon build as a deploy compatibility risk, with pgsql-parser as a fallback. libpg-query@17.x is now WASM-only — single prebuilt artifact, no node-gyp, no per-platform builds. No fallback needed.

Tests

22 new tests, all green; full api2 suite (168/168) and tsc --noEmit clean.

  • query-data-validator.test.ts (16 tests) — empty/whitespace, multi-statement, UPDATE/INSERT/DELETE, single q_ select, multi-q_ join, non-q_ reject, mixed q_ + non-q_ reject, schema-qualified public.users reject, CTE-with-q_-source accept (and CTE name not checked against prefix), CTE-with-non-q_-source reject, subquery referencing non-q_ reject, malformed SQL (parse error) reject, SELECT 1 (no FROM) reject.
  • query-data-runner.db.test.ts (6 pglite tests) — tenant isolation via current_setting('app.current_organization_id'), SET LOCAL scope across consecutive calls (org 1 / org 2 don’t leak), validator-rejection short-circuit (no transaction opened), underlying SQL errors, row cap + preview cap with generate_series(1, 5250), and a direct assertion of the three GUC values from inside the transaction (statement_timeout = 20s, transaction_read_only = on, app.current_organization_id = '42').

Important pglite caveat for future tests: pglite is single-process WASM, so statement_timeout semantics are weaker than prod. We assert the GUC is set; we don’t try to fire a long-running query and assert it gets cancelled.

What stayed on Rails (intentional)

  • The Ruby ToolCalls::ComplianceAgent::QueryData tool and the Internal::ToolsController#call query_data case both remain. They’re still used by the workflow-node and sentinel-conversation-message paths (workflow_node/run/agent/tools/query_data.rb, sentinel_conversation_message/tools/query_data.rb), neither of which goes through api2.
  • q_* Scenic views remain in Rails-managed public.*.
  • The stored_query_results table keeps getting written by Ruby callers. The api2 path stops writing it.

If anything Rails-side (audit, eval, retrospective tooling) reads stored_query_results rows from agent runs that came through api2, those rows will disappear after this PR. Repo-wide grep found no such consumer — flagging here in case an out-of-tree dashboard depends on it.


Phase 2 — Drizzle-owned views in agent_api schema (shipped 2026-05-15)

The drift problem behind this phase: until Phase 2, q_tables_schema.json (the static doc served by get_schema) was hand-edited every time a Rails Scenic view changed. Two recent incidents came from that drift — actor_type was missing from q_conversation_parts for weeks, and ticket_assignee_id / name / actor_type were never documented at all until 2026-05-14. The Ownership Watcher agent was over-counting “agents on a ticket” partly because of it.

What changed

  • agent_api Postgres schema, owned by Drizzle. rulebase-web/rulebase-api2/src/db/agent-api-schema.ts declares all 26 views via pgSchema("agent_api").view(...).as(...) with the body in a Drizzle sql template literal. Each body is copy-pasted verbatim from the matching rulebase-api/db/views/q_<name>_v<latest>.sql, with two cosmetic adjustments: q_ is dropped from the view name (the schema namespaces) and base-table references are prefixed with public. so the resolved DDL is search-path-independent. Column declarations use real Drizzle types (bigint, integer, text, text[], timestamp with withTimezone: true, boolean, numeric, jsonb, interval) — column types are not part of the CREATE VIEW DDL, but they’re the source from which the agent’s get_schema derives type strings, so accuracy here ≡ accuracy in the agent prompt.
  • agent-api-descriptions.ts — pure derivation layer. Owns only the view-level descriptions and per-column human descriptions; everything else (column names, column types) is read off the Drizzle view at module load via view[Symbol.for("drizzle:ViewBaseConfig")].selectedFields[col].getSQLType(). Has a drift guard that throws at module load if any description references a column the matching view doesn’t have. One-time migration from the now-deleted q_tables_schema.json.
  • Drizzle-generated migration (drizzle/0023_shallow_the_watchers.sql) — creates the schema unconditionally and wraps all 26 CREATE VIEW statements in a DO block guarded on to_regclass('public.conversations') IS NOT NULL. Real Postgres dbs (Rails-shared) create all 26; pglite — which lacks the Rails base tables — silently skips them so the test harness still boots.
  • Validator flip. validator.ts replaced the relname.startsWith("q_") check with a schema check: schemaname must be null or "agent_api". CTE-name exclusion logic unchanged. Hard cutover — public.q_* references from agent SQL are now rejected before the runner ever opens a transaction.
  • Runner pins search_path. runner.ts adds SET LOCAL search_path = agent_api to the transaction prologue, so agents can write SELECT * FROM conversations without typing agent_api. every time. Anything that resolves outside agent_api fails with relation does not exist; the validator already rejects explicit public.* / api2.* refs upstream.
  • Prompt + snapshots. managed-agent.ts buildToolDocumentation and managed_agent_system.md lost all q_* references; the 12 managed-agent snapshots were regenerated.

Key design decisions

  • Duplicate view bodies, not pass-through. Each agent_api.<name> body is a copy of the Rails Scenic q_<name> body, not a SELECT * FROM public.q_<name>. The price: a Scenic-side column change must be replicated into agent-api-schema.ts or the two surfaces diverge. Accepted trade-off so api2 owns its agent-facing surface end-to-end and Phase 3 can grant a least-privileged role on just agent_api.*.
  • Hard cutover. Validator stops accepting q_* (in public) the day this lands. The Ruby query_data tool (workflow-node + sentinel paths) is unaffected and keeps reading public.q_*.
  • Drop q_ prefix in agent_api. agent_api.conversations, not agent_api.q_conversations. The schema does the namespacing.

Drift surface (the new “we accepted this” cost)

The JSON drift problem is gone. agent-api-descriptions.ts is now a derivation — it can’t drift on column names or types because it reads both straight off the Drizzle view; the only fields it owns are the human-readable descriptions. If a description references a column the view doesn’t have, the module throws at load.

The remaining drift surface is between agent-api-schema.ts and the Rails Scenic source files in rulebase-api/db/views/q_*.sql. There is no automated guard. When a Scenic view changes:

  1. Update the matching db/views/q_<name>_v<N+1>.sql.
  2. Re-paste the body into the matching agent_api.<name> view in agent-api-schema.ts (drop the q_ aliases or rename to taste — the body is yours now). Add/remove column declarations to match.
  3. Add/update column descriptions in agent-api-descriptions.ts. The drift guard catches stale entries (description for a removed column → throw at boot); columns added without a description are allowed (agent sees the type but an empty description).

Phase 3 (below) hardens this by enforcing read-only at the DB role level so a misaligned view body can’t accidentally surface base tables — but a wrong column in an agent_api view body is still a logic bug the role can’t catch.

Tests

  • Validator (19 tests, rewritten): unqualified accept, agent_api.x accept, public.x reject, public.q_* legacy reject, api2.x reject, mixed-schema reject, CTE-with-agent_api accept, CTE-with-forbidden-schema reject, subquery-with-forbidden-schema reject, plus all the existing empty/DML/multi-statement/parse-error cases.
  • Runner pglite (7 tests): agent_api.test_rows view set up in the harness; org isolation, search_path resolution, row/preview cap, GUC assertion (now includes search_path = agent_api).
  • Managed-agent snapshots: 12 regenerated. All 24 tests green.

Phase 3 — DB role + SET LOCAL ROLE enforcement (shipped 2026-05-15)

End state: agent queries run as a least-privileged Postgres role (agent_query_reader) that has SELECT-only on the agent_api schema and nothing else, even if the validator is ever bypassed.

What changed

  • Custom Drizzle migration (drizzle/0024_agent-query-reader-role.sql, scaffolded via drizzle-kit generate --custom). One DO $do$ ... $do$; block that:
    1. Creates agent_query_reader with NOLOGIN if it doesn’t already exist (idempotent across cluster-scoped re-runs).
    2. GRANT USAGE ON SCHEMA agent_api and GRANT SELECT ON ALL TABLES IN SCHEMA agent_api to the new role.
    3. ALTER DEFAULT PRIVILEGES IN SCHEMA agent_api GRANT SELECT ON TABLES TO agent_query_reader so views added in future Drizzle migrations are auto-readable.
    4. EXECUTE format('GRANT agent_query_reader TO %I', current_user) so the api2 runtime (which connects with the same DATABASE_URL credentials as the migrator — verified via flightcontrol.cue) can SET LOCAL ROLE into it.
  • Runner (runner.ts) now ends the transaction prologue with SET LOCAL ROLE agent_query_reader, after all GUC sets (the GUCs would be unsettable once dropped to the restricted role). The user query then executes as agent_query_reader — direct SELECT against public.* or api2.* dies with permission denied for table ... even if the validator is bypassed.
  • No GRANT on public.* needed. Postgres views run with the owner’s permissions, not the caller’s. The view owner is the master user (same role that runs Rails migrations + api2 migrations + api2 runtime in prod, per flightcontrol.cue), and it already has SELECT on the Rails-managed base tables. agent_query_reader reads agent_api.conversations; Postgres checks the role’s SELECT on the view (granted), then the view’s body executes under the owner’s identity to read public.conversations (granted to the owner). No cross-schema grant juggling.

Why pglite-friendliness fell out for free

The earlier “Phase 3 will need real Postgres in CI” assumption was wrong. Pglite’s electric-sql/pglite build implements CREATE ROLE, SET ROLE, GRANT, ALTER DEFAULT PRIVILEGES, and view-owner ACL semantics — verified by probe against @electric-sql/pglite@latest. The Phase 3 migration applies cleanly under pglite, the test-time agent_api.test_rows view inherits SELECT via the default-privilege grant, and runQueryData actually runs as agent_query_reader in the pglite suite the same way it does in prod.

Tests

  • runner.db.test.ts (8 tests, 1 added, 1 expanded).
    • The GUC-assertion test now also asserts current_user = 'agent_query_reader' inside the runner’s transaction.
    • New “denies direct reads of public.* base tables even via the validator-bypassing role” test: creates agent_api.public_passthrough as a deliberate escape hatch, confirms the agent can read through it (views run as owner) but a manually-issued SELECT FROM public.test_rows_base under SET LOCAL ROLE agent_query_reader rejects with permission denied. Drizzle wraps the pg error, so the assertion reads .cause.message.

Operational notes

  • The role is cluster-scoped, not database-scoped. pnpm db:reset:local drops the agent_api schema (which revokes grants) but leaves the role itself in place. The migration’s IF to_regrole(...) IS NULL guard makes that fine across resets. To fully recycle locally: DROP OWNED BY agent_query_reader CASCADE; DROP ROLE agent_query_reader; before db:reset:local.
  • Future agent-readable views must live in agent_api. Views added there post-migration auto-inherit SELECT for agent_query_reader via the ALTER DEFAULT PRIVILEGES step; explicit grants are unnecessary.

Files (Phase 1)

New:

  • rulebase-web/rulebase-api2/src/lib/query-data/validator.ts (initially src/lib/query-data-validator.ts; relocated to the query-data/ subdirectory mid-PR)
  • rulebase-web/rulebase-api2/src/lib/query-data/validator.test.ts
  • rulebase-web/rulebase-api2/src/lib/query-data/runner.ts
  • rulebase-web/rulebase-api2/src/lib/query-data/runner.db.test.ts

Edited:

  • rulebase-web/rulebase-api2/src/lib/query-data/tools.tsqueryData()runQueryData()
  • rulebase-web/rulebase-api2/src/lib/rails-client.ts — deleted queryData() export
  • rulebase-web/rulebase-api2/package.json — added libpg-query@^17.7.3

Files (Phase 2)

New:

  • rulebase-web/rulebase-api2/src/db/agent-api-schema.tspgSchema("agent_api") + 26 view declarations
  • rulebase-web/rulebase-api2/src/db/agent-api-descriptions.ts — TS column-descriptions registry served by get_schema
  • rulebase-web/rulebase-api2/drizzle/0023_shallow_the_watchers.sqlCREATE SCHEMA agent_api; + 26 views wrapped in a DO block guarded on to_regclass('public.conversations')

Edited:

  • rulebase-web/rulebase-api2/drizzle.config.tsschema: ["./src/db/schema.ts", "./src/db/agent-api-schema.ts"], schemaFilter: ["api2", "agent_api"]
  • rulebase-web/rulebase-api2/src/lib/query-data/validator.ts — schema check (agent_api or null) instead of prefix check
  • rulebase-web/rulebase-api2/src/lib/query-data/runner.ts — added SET LOCAL search_path = agent_api
  • rulebase-web/rulebase-api2/src/lib/query-data/tools.ts — serve TS descriptions registry, drop JSON file read
  • rulebase-web/rulebase-api2/src/lib/managed-agent.ts + src/prompts/managed_agent_system.mdq_* references → unqualified or agent_api.*; snapshots regenerated
  • rulebase-web/rulebase-api2/src/lib/query-data/validator.test.ts + runner.db.test.ts — rewritten for the schema contract

Deleted:

  • rulebase-web/rulebase-api2/src/lib/query-data/q_tables_schema.json

Files (Phase 3)

New:

  • rulebase-web/rulebase-api2/drizzle/0024_agent-query-reader-role.sql — custom Drizzle migration: creates agent_query_reader (NOLOGIN), grants USAGE + SELECT on agent_api.*, ALTER DEFAULT PRIVILEGES, grants the role to the connecting user

Edited:

  • rulebase-web/rulebase-api2/src/lib/query-data/runner.ts — appended SET LOCAL ROLE agent_query_reader to the transaction prologue
  • rulebase-web/rulebase-api2/src/lib/query-data/runner.db.test.ts — added current_user assertion + a defense-in-depth test that direct public.* reads are refused under the role
Last updated on