DataMate — Database Schema Selection Feature · Implementation Plan

Add a sidebar dropdown to pick a source-DB schema; DataMate then answers questions grounded in that schema's structure (writes correct SQL, explains tables/columns/relationships). Schema-only reasoning — no live row execution.

Date: 2026-06-10 · Scope: asktge (UI + backend) + datamate/agent + new S3 store

---

1. Locked Decisions

The 6 files in asktge/DB_Schema_Files/ are MySQL/MariaDB DDL dumps (CREATE TABLE only, no rows): Goldhopper_DEV2, MYTEAMGE, Rockhopper_NSW, TAPS, TED-EDS, TED-RDS. They are source-system schemas, not the Redshift the agent's MCP tools already reach.

---

2. What Already Exists (reuse, don't rebuild)

Investigation found partial plumbing already in place:

• asktge/routes/chat_routes.py already carries a selected_schema through the session and into the invoke flow. The SQLMatic branch already does session["selected_schema"] = session.get("selectedSchema", "Goldhopper") — i.e. a schema-selection slot exists and already defaults to Goldhopper.
• asktge/utils/session_utils.py::get_agent_id_for_schema() and asktge/config/knowledge_bases.py::get_kb_id_for_schema() map a schema name → agent/KB.
• asktge/invoke_client.py::invoke_datamate_agent() posts to the DataMate AgentCore runtime (bedrock-agentcore…/runtimes/…datamate_agent-VHS3PoFgPB/invocations) with a Bearer (Azure AD JWT) and SSE handling.
• The datamate agent already uses S3 (tools/kb_search.py::_generate_presigned_url, boto3 S3) — so S3 read IAM + client are present. No new infra category.
• datamate/agent/api/models.py::InvokeRequest is the payload contract; ai/graphs.py builds the LangGraph and binds tools; ai/prompts.py::seed_system_prompt_withTools is DataMate's system prompt.

Implication: we extend an existing selected_schema channel and add one payload field (databaseId) + one loader + a prompt block. We are not inventing the pathway.

---

3. Chosen Delivery Approach & Why

S3 canonical DDL + build-time condensation + agent-side cached loader + system-prompt injection, with an on-demand drill-down tool.

Why not the plain "S3 + read-tool" as first described: feeding a raw 9,000-line MySQL dump (MYTEAMGE) into context every turn is heavy and noisy (SET statements, charset clauses, ENGINE/AUTO_INCREMENT boilerplate). Two refinements fix it:

1. Condense at upload time. A preprocessing step strips dump noise and emits a compact schema view per DB: tables → columns (name, type, nullable, PK/FK), plus foreign-key edges. ~9k DDL lines → roughly 1.5–3k tokens. Store both condensed and raw on S3.
2. Inject condensed schema into the system prompt for the selected DB (deterministic — the agent always has structure, so "ask anything" reliably works), and expose a get_table_ddl(databaseId, table) tool for the exact raw DDL of a specific table when precise constraints/types are needed. Use Bedrock prompt caching on the schema block to amortize tokens across turns.

This keeps requests tiny (only databaseId travels), lets schemas update without an agent redeploy (re-upload to S3), and bounds token cost.

S3:  s3://<bucket>/datamate/schemas/<databaseId>/
        raw.sql              (original dump, for drill-down)
        condensed.json       (tables/columns/keys — injected)
        manifest.json        (display name, table count, updated_at)

---

4. Data Flow (target)

User opens DataMate ──► sidebar dropdown lists schemas (from /api/datamate/schemas)
        │ picks "Goldhopper"
        ▼
chat.js stores selectedSchema ──► sends selectedSchema in /send_message body
        ▼
chat_routes.py: session["selected_schema"]=...; resolve databaseId
        ▼
invoke_client.invoke_datamate_agent(..., databaseId="goldhopper")
        ▼  payload += {"databaseId": "..."}
datamate InvokeRequest.databaseId ──► graph state
        ▼
schema_loader.load_condensed(databaseId)  [S3 GET + TTL cache]
        ▼
system prompt = base + "<DB SCHEMA: Goldhopper>…condensed…</>"  (prompt-cached)
        ▼
LLM answers / writes SQL.  Needs exact DDL? → tool get_table_ddl(databaseId, table)

---

5. Implementation Phases

Phase 0 — Schema prep & S3 (foundation)

• Create S3 prefix datamate/schemas/ in the existing datamate bucket (same bucket as KB docs; reuse IAM).
• Write datamate/agent/scripts/build_schema_store.py: parse each .sql dump → condensed.json (tables, columns w/ type+nullable, PK, FK edges) + manifest.json; upload raw + condensed + manifest. Parser can use sqlparse or regex over CREATE TABLE blocks (dumps are uniform mysqldump output).
• Define a schema registry datamate/agent/config/schemas.py: databaseId → {display_name, s3_prefix, enabled}. Seed the 6 DBs. This is the single source of truth the UI list is derived from.
• One-time run: upload the 6 current files. (Move ownership of DB_Schema_Files/ here, or keep asktge copy as the editing source and add a CI step that re-runs the builder.)

Phase 1 — asktge UI (sidebar dropdown)

• asktge/templates/chat.html: add a <select id="datamate-schema-select"> block inside #chat-sidebar, rendered only when currentAgent === "DataMate".
• asktge/static/js/chat.js: on DataMate load, GET /api/datamate/schemas → populate dropdown; persist choice to ChatState.selectedSchema + localStorage; include selectedSchema in the /send_message request body (chat.js already tracks currentAgent).
• asktge/static/css/agent_select.css (or chat.css): style the selector to match the sidebar.
• UX: default to last-used (localStorage) else first enabled schema; show the active schema as a chip above the input.

Phase 2 — asktge backend (list endpoint + payload)

• New route (e.g. utility_routes.py or a small datamate_routes.py): GET /api/datamate/schemas → returns [{databaseId, display_name}] from a config that mirrors datamate/agent/config/schemas.py (keep one canonical list; asktge can fetch it from the agent at startup or hold a static copy).
• chat_routes.py: in the DataMate branch, read selectedSchema from request → set session["selected_schema"] → resolve to databaseId (slugify / map). Validate against the allowed list; reject unknown.
• invoke_client.py::invoke_datamate_agent(): add database_id param; payload["databaseId"] = database_id.
• Pass databaseId through chat_service / wherever invoke_datamate_agent is called.

Phase 3 — datamate agent (ingest + load)

• api/models.py::InvokeRequest: add databaseId: Optional[str] = None (+ optional validator against config/schemas.py).
• app.py: read request.databaseId, put it in the LangGraph initial state (e.g. state["database_id"]).
• New datamate/agent/tools/schema_loader.py:
  • load_condensed(database_id) -> dict — S3 GET condensed.json, cachetools.TTLCache (e.g. 30 min) keyed by databaseId; reuse the existing boto3 S3 client from kb_search.
  • load_table_ddl(database_id, table) -> str — slice the table's block from raw.sql (cache raw per DB).
  • list_schemas() -> list — from config/schemas.py (used by the asktge endpoint too, or exposed via a tiny agent route).

Phase 4 — prompt + tool wiring

• ai/prompts.py: when state.get("database_id"), append a ## Selected Database Schema block built from load_condensed(...), plus rules: "Answer using ONLY this schema; if a question needs another DB, say it isn't selected; for exact column constraints call get_table_ddl." Mark the block for prompt caching.
• ai/graphs.py: add get_table_ddl(databaseId, table) to the bound tool list (alongside the existing redshift_* / semantic_schema_* tools). Keep the live redshift_* tools available but instruct the prompt that for the selected source-schema mode they are for reference only (schema-only reasoning — no execution expected).
• Ensure database_id survives across turns: persist it in session metadata (storage/dynamodb.py::_store_session_metadata) so restored conversations keep their schema.

---

6. Config / Registry (single source of truth)

datamate/agent/config/schemas.py:

SCHEMAS = {
    "goldhopper":    {"display_name": "Goldhopper (DEV2)", "s3_prefix": "datamate/schemas/goldhopper/",    "enabled": True},
    "myteamge":      {"display_name": "MyTeamGE",          "s3_prefix": "datamate/schemas/myteamge/",      "enabled": True},
    "rockhopper_nsw":{"display_name": "Rockhopper (NSW)",  "s3_prefix": "datamate/schemas/rockhopper_nsw/","enabled": True},
    "taps":          {"display_name": "TAPS",              "s3_prefix": "datamate/schemas/taps/",          "enabled": True},
    "ted_eds":       {"display_name": "TED-EDS",           "s3_prefix": "datamate/schemas/ted_eds/",       "enabled": True},
    "ted_rds":       {"display_name": "TED-RDS",           "s3_prefix": "datamate/schemas/ted_rds/",       "enabled": True},
}

UI dropdown, asktge validation, and agent loader all derive from this. Adding a DB = add a row + run the builder upload.

---

7. Token / Performance

• Condensed schema ~1.5–3k tokens; inject once per turn, prompt-cache the block (Sonnet 4 on Bedrock supports caching) → near-zero marginal cost after first turn.
• schema_loader TTLCache avoids repeat S3 GETs within a session.
• Drill-down get_table_ddl keeps the big raw DDL out of context unless a specific table is needed.

8. Security / Compliance

• Schemas may expose table/column names that are themselves sensitive (PII-revealing). Treat the S3 prefix as internal; agent role gets read-only s3:GetObject on datamate/schemas/* (extend the existing policy).
• No credentials to the source DBs are involved (schema-only). If this later becomes live-query, that's a new auth+secrets workstream (the earlier "new DB connections" option).
• Keep schema selection server-validated against the registry — never let an arbitrary databaseId reach S3 path construction (prevent path traversal: map via registry, don't interpolate raw input).

9. Testing

• Unit: build_schema_store parser on each of the 6 dumps (table/column/FK counts vs. eyeball).
• Unit: schema_loader cache hit/miss, unknown databaseId → safe error.
• Integration: invoke datamate with databaseId="goldhopper" + "list the tables and the key columns of X" → answer grounded in condensed schema; get_table_ddl returns exact block.
• E2E: UI dropdown → ask a structural question → correct, schema-scoped answer; switch schema mid-session → context updates.
• Negative: ask about a table from a non-selected DB → agent declines / notes it isn't selected.

10. Rollout

1. Phase 0 offline (S3 + builder) — no user impact.
2. Ship agent changes behind tolerance for missing databaseId (absent → current generic behavior).
3. Ship asktge backend (sends databaseId).
4. Ship UI dropdown last. Optionally feature-flag via the agent registry enabled flags + a DATAMATE_SCHEMA_SELECTION env toggle.

11. Effort (rough)

12. Open Items to Confirm

• Bucket: reuse the existing datamate KB bucket, or a dedicated one? (Plan assumes reuse.)
• Editing source of truth: keep asktge/DB_Schema_Files/ as the human-edited copy with a CI step that condenses+uploads, or move ownership into the datamate repo? Recommend: keep asktge copy as source, add a GitLab CI job running build_schema_store.py on change.
• Redshift relationship: confirm these 6 source schemas are intentionally separate from the agent's live Redshift, so "schema-only" is the right long-term model and we won't surprise users who expect real numbers back.
• Naming: surface as "Database" or "Schema" in the UI? (Affects label only.)