4 · Schema & Connectors

Two problems live here. Connectors are how you talk to six different databases through one interface. Schema handling is how you fit a 400-table warehouse into a context window — which is the retrieval problem, and the more interesting half.

The connector interface

# src/connectors/base.py
class DataConnector(ABC):
    @abstractmethod
    async def connect(self) -> None: ...
    @abstractmethod
    async def disconnect(self) -> None: ...
    @abstractmethod
    async def test_connection(self) -> ConnectionTestResult: ...
    @abstractmethod
    async def discover_schema(self) -> SchemaSnapshot: ...
    @abstractmethod
    async def execute_query(self, sql, params=None,
                            max_rows=1000, timeout=30) -> QueryResult: ...
    @abstractmethod
    async def execute_query_stream(self, sql, ...) -> AsyncIterator[list[Any]]: ...
 
    @property
    @abstractmethod
    def dialect(self) -> str: ...
 
    async def __aenter__(self):  await self.connect();  return self
    async def __aexit__(self, *a): await self.disconnect()

PostgreSQL · MySQL · SQLite · BigQuery · Snowflake · CSV/Excel (DuckDB + Polars).

Two details worth pointing at:

dialect is quietly load-bearing. It tells stage 4 which SQL flavour to ask for and stage 5 which grammar to validate against. Get it wrong and validation rejects valid SQL.

Async context manager. async with connector: guarantees the pool closes even if a stage raises. The pipeline opens the connector once and holds it across generation, validation and execution.

Normalised metadata

Every connector returns the same shapes regardless of vendor:

@dataclass
class ColumnInfo:
    name: str
    data_type: str
    nullable: bool
    is_primary_key: bool = False
    is_foreign_key: bool = False
    foreign_table: str | None = None
    foreign_column: str | None = None
    sample_values: list[Any] = field(default_factory=list)
 
@dataclass
class TableInfo:
    name: str
    schema: str
    columns: list[ColumnInfo]
    row_count_estimate: int = 0

foreign_table powers FK traversal in schema selection. sample_values help the model understand what a column actually contains — status is far more useful when the model can see 'shipped' | 'cancelled'. row_count_estimate feeds relevance scoring, since big tables are usually fact tables.

Credentials, encrypted

# src/connectors/encryption.py
def encrypt_credentials(credentials: dict) -> str:
    return _get_fernet().encrypt(json.dumps(credentials).encode()).decode()

Fernet (AES-128-CBC + HMAC) with a key from ENCRYPTION_KEY. Credentials are encrypted before they reach the database column, so a database dump alone does not yield connection strings.

Percent-encoding, a real bug class

user = quote_plus(str(c.get("username", "")))
password = quote_plus(str(c.get("password", "")))

Generated cloud passwords routinely contain @, /, : or #. Unencoded, they corrupt the DSN and produce a baffling “could not translate host name” — because everything after the stray @ is parsed as the host.

SSL, and a mismatch worth knowing

# asyncpg takes an `ssl` connect arg, NOT `sslmode`. Passing the libpq spelling
# raises: connect() got an unexpected keyword argument 'sslmode'.
# This bit every managed Postgres (Render, Neon, Supabase…).

libpq speaks sslmode=require; asyncpg speaks ssl=. Every managed provider hands you a libpq URL. The same translation exists in src/config.py for Retail Mind’s own database URL — see Deployment.

Schema discovery

discover_schema() walks the database. For SQLite:

result = await conn.execute(text("SELECT name FROM sqlite_master WHERE type='table'"))
for (table_name,) in result.fetchall():
    col_result = await conn.execute(text(f"PRAGMA table_info('{table_name}')"))
    fk_result  = await conn.execute(text(f"PRAGMA foreign_key_list('{table_name}')"))
    count      = await conn.execute(text(f"SELECT COUNT(*) FROM '{table_name}'"))

Postgres uses information_schema, BigQuery its INFORMATION_SCHEMA views, and so on — all returning the same SchemaSnapshot.

This runs as a background job because it takes minutes on a large warehouse. Snapshots are versioned rather than overwritten:

snap = SchemaSnapshot(
    source_id=source_id,
    version=(latest or 0) + 1,
    schema_data=schema_dict,
    compressed_ddl=compressed,
    table_count=len(snapshot.tables),
)

So a schema change is auditable and a bad sync can be traced. Before this, every row was written as version 1 and “latest” was ambiguous.

Compression

src/schema/compression.py turns verbose metadata into minimal DDL — names, types, keys, a few sample values — rather than dumping information_schema output. Fewer tokens, and less noise for the model to wade through.

Selection — the retrieval problem

def _score_table(table: dict, question_tokens: set[str]) -> float:
    score = 0.0
    score += len(_tokenize(table["name"]) & question_tokens) * 3.0    # table name
    for col in table.get("columns", []):
        score += len(_tokenize(col["name"]) & question_tokens) * 1.0  # column name
        for val in col.get("sample_values", []):
            if str(val).lower() in " ".join(question_tokens):
                score += 0.5                                          # sample value
    row_count = table.get("row_count_estimate", 0)
    if row_count > 1_000_000:  score += 1.0                           # fact-table bias
    elif row_count > 100_000:  score += 0.5
    return score

Then foreign-key traversal:

# A question about "revenue by product" scores orders and products, but needs
# order_items to join them — and nothing in the question mentions it.
selected = list(top_names)
seen = set(selected)
fk_additions: list[str] = []
for table in tables:
    if table["name"] not in seen:
        continue
    for col in table.get("columns", []):
        fk_target = col.get("foreign_table")
        if fk_target and fk_target not in seen and len(seen) + len(fk_additions) < MAX_TABLES_IN_PROMPT:
            fk_additions.append(fk_target)
            seen.add(fk_target)
return selected + sorted(fk_additions)

Determinism is a requirement, not a nicety

This used to end return list(top_set). Python randomises string hashing per interpreter run, so table order — and therefore the prompt — varied between processes. Two consequences:

• The same question produced different prompts across restarts.
• Eval runs were not reproducible, so you could not tell a prompt improvement from noise.

Now: score order, ties broken by name, FK additions appended sorted. tests/test_schema_search.py asserts identical output when the input table order is reversed, and it was verified across four PYTHONHASHSEED values.

Where this is weak — say it before you are asked

Keyword scoring fails when the question and the schema do not share vocabulary. Ask “how many buyers last month?” against a table called dim_cust_x and nothing scores.

The fix: embed each table’s name, columns and sample values; embed the question; retrieve by cosine similarity. pgvector is already installed and AIProvider.embed() already exists — the plumbing is there, the wiring is not. Best of all is hybrid: keyword for exact matches like a column literally named revenue, embeddings for paraphrase, then FK traversal over the union.

Reading the tests

tests/test_connectors_sqlite.py runs against the real fixture database and pins the two behaviours that are load-bearing for safety rather than correctness:

• execute_query enforces the row cap even when the SQL asks for more
• discover_schema finds foreign keys — miss one and schema selection cannot follow the join, silently