[KM-533] add table level schema, differentiate with chunk level. expand retrieval result with FK exploration

src/pipeline/db_pipeline/db_pipeline_service.py

@@ -21,7 +21,13 @@ from src.db.postgres.connection import _pgvector_engine
 from src.db.postgres.vector_store import get_vector_store
 from src.middlewares.logging import get_logger
 from src.models.credentials import DbType
-from src.pipeline.db_pipeline.extractor import get_schema, profile_table
+from src.pipeline.db_pipeline.extractor import (
+    build_table_chunk,
+    fetch_sample_row,
+    get_row_count,
+    get_schema,
+    profile_table,
+)
 
 logger = get_logger("db_pipeline")
 
@@ -156,6 +162,7 @@ class DbPipelineService:
             metadata={
                 "user_id": user_id,
                 "source_type": "database",
+                "chunk_level": "column",
                 "database_client_id": client_id,
                 "updated_at": updated_at,
                 "data": {
@@ -168,6 +175,46 @@ class DbPipelineService:
             },
         )
 
+    def _to_table_document(
+        self,
+        user_id: str,
+        client_id: str,
+        table_name: str,
+        columns: list[dict],
+        row_count: int,
+        text: str,
+        updated_at: str,
+    ) -> LangChainDocument:
+        foreign_keys = []
+        for c in columns:
+            fk = c.get("foreign_key")
+            if not fk:
+                continue
+            target_table, _, target_column = fk.partition(".")
+            foreign_keys.append({
+                "column": c["name"],
+                "target_table": target_table,
+                "target_column": target_column,
+            })
+
+        return LangChainDocument(
+            page_content=text,
+            metadata={
+                "user_id": user_id,
+                "source_type": "database",
+                "chunk_level": "table",
+                "database_client_id": client_id,
+                "updated_at": updated_at,
+                "data": {
+                    "table_name": table_name,
+                    "row_count": row_count,
+                    "primary_key": [c["name"] for c in columns if c.get("is_primary_key")],
+                    "foreign_keys": foreign_keys,
+                    "column_names": [c["name"] for c in columns],
+                },
+            },
+        )
+
     async def run(
         self,
         user_id: str,
@@ -193,6 +240,29 @@ class DbPipelineService:
             entries = await asyncio.to_thread(profile_table, engine, table_name, columns)
             docs = [self._to_document(user_id, client_id, table_name, e, updated_at) for e in entries]
             all_docs.extend(docs)
+
+            # Table-level chunk. Failures here are logged and skipped — column
+            # chunks above are already in all_docs and will still be written.
+            try:
+                row_count = await asyncio.to_thread(get_row_count, engine, table_name)
+                sample_row = (
+                    await asyncio.to_thread(fetch_sample_row, engine, table_name)
+                    if row_count > 0
+                    else None
+                )
+                table_text = build_table_chunk(
+                    table_name, row_count, columns, entries, sample_row
+                )
+                all_docs.append(
+                    self._to_table_document(
+                        user_id, client_id, table_name, columns, row_count, table_text, updated_at
+                    )
+                )
+            except Exception as e:
+                logger.error(
+                    "table chunk generation failed", table=table_name, error=str(e)
+                )
+
             logger.info("profiled table", table=table_name, count=len(docs))
 
         # Insert new chunks first; only delete stale chunks after the insert succeeds.

src/pipeline/db_pipeline/extractor.py

@@ -85,7 +85,9 @@ def get_schema(
 
 
 def get_row_count(engine: Engine, table_name: str) -> int:
-    return pd.read_sql(f"SELECT COUNT(*) FROM {_qi(engine, table_name)}", engine).iloc[0, 0]
+    # Cast to plain int — pandas returns numpy.int64 which fails JSONB serialization
+    # when the value lands in PGVector cmetadata via the table-level chunk.
+    return int(pd.read_sql(f"SELECT COUNT(*) FROM {_qi(engine, table_name)}", engine).iloc[0, 0])
 
 
 def profile_column(
@@ -187,6 +189,74 @@ def profile_table(engine: Engine, table_name: str, columns: list[dict]) -> list[
     return results
 
 
+def fetch_sample_row(engine: Engine, table_name: str) -> Optional[dict]:
+    """First row of the table as a dict, or None if the table is empty.
+
+    Reuses _qi for dialect-correct quoting and _head_query for TOP/LIMIT.
+    """
+    qt = _qi(engine, table_name)
+    sql = _head_query(engine, "*", qt, 1)
+    df = pd.read_sql(sql, engine)
+    if df.empty:
+        return None
+    return df.iloc[0].to_dict()
+
+
+def build_table_chunk(
+    table_name: str,
+    row_count: int,
+    columns: list[dict],
+    column_profiles: list[dict],
+    sample_row: Optional[dict],
+) -> str:
+    """Build the table-level chunk text.
+
+    Format (lines omitted when not applicable):
+        Table: {name} ({row_count} rows)
+        Primary key: {pk_cols}
+        Foreign keys: {col} -> {target_table}.{target_col}, ...
+        Columns ({n}): {col1}, {col2}, ...
+        Numeric ranges: {col} [{min}-{max}], ...
+        Sample row: {dict}
+
+    Pure formatter — no DB I/O. column_profiles is the output of profile_table
+    and is reused so we don't re-introspect.
+    """
+    lines = [f"Table: {table_name} ({row_count} rows)"]
+
+    pk_cols = [c["name"] for c in columns if c.get("is_primary_key")]
+    if pk_cols:
+        lines.append(f"Primary key: {', '.join(pk_cols)}")
+
+    fk_parts = [
+        f"{c['name']} -> {c['foreign_key']}" for c in columns if c.get("foreign_key")
+    ]
+    if fk_parts:
+        lines.append(f"Foreign keys: {', '.join(fk_parts)}")
+
+    col_names = [c["name"] for c in columns]
+    lines.append(f"Columns ({len(col_names)}): {', '.join(col_names)}")
+
+    range_parts = []
+    for entry in column_profiles:
+        col = entry["col"]
+        profile = entry["profile"]
+        if not col.get("is_numeric"):
+            continue
+        mn = profile.get("min")
+        mx = profile.get("max")
+        if mn is None or mx is None:
+            continue
+        range_parts.append(f"{col['name']} [{mn}-{mx}]")
+    if range_parts:
+        lines.append(f"Numeric ranges: {', '.join(range_parts)}")
+
+    if sample_row is not None:
+        lines.append(f"Sample row: {sample_row}")
+
+    return "\n".join(lines)
+
+
 def build_text(table_name: str, row_count: int, col: dict, profile: dict) -> str:
     col_name = col["name"]
     col_type = col["type"]

src/query/executors/db_executor.py

@@ -41,6 +41,7 @@ _enc = tiktoken.get_encoding("cl100k_base")
 _SUPPORTED_DB_TYPES = {"postgres", "supabase", "mysql"}
 _MAX_RETRIES = 3
 _MAX_LIMIT = 500
+_FK_EXPANSION_MAX_TABLES = 5
 
 _SQL_SYSTEM_PROMPT = """\
 You are a SQL data analyst working with a user's database.
@@ -137,20 +138,31 @@ class DbExecutor(BaseExecutor):
             logger.warning("unsupported db_type for query execution", db_type=client.db_type)
             return None
 
-        # Distinct table names from retrieval results, expanded via FK relationships
-        table_names = list({
+        # Hit tables = tables retrieval pointed at directly. Get full per-column
+        # schema for these. Related tables (one FK hop away, both directions) are
+        # fetched separately in abbreviated form to give the LLM enough context
+        # to JOIN without paying the per-column profile token cost.
+        hit_tables = list({
             r.metadata.get("data", {}).get("table_name")
             for r in results
             if r.metadata.get("data", {}).get("table_name")
         })
-        table_names = await self._expand_with_fk_tables(client_id, user_id, table_names)
+        if not hit_tables:
+            logger.warning("no table_name on any retrieval result", client_id=client_id)
+            return None
 
-        full_schema = await self._fetch_full_schema(client_id, table_names, user_id)
+        full_schema = await self._fetch_full_schema(client_id, hit_tables, user_id)
         if not full_schema:
-            logger.warning("no schema found in vector store", client_id=client_id, tables=table_names)
+            logger.warning("no schema found in vector store", client_id=client_id, tables=hit_tables)
             return None
 
-        schema_ctx = self._build_schema_context(full_schema)
+        related_tables = await self._find_related_tables(client_id, user_id, hit_tables)
+        related_schema = (
+            await self._fetch_abbreviated_schema(client_id, user_id, related_tables)
+            if related_tables else {}
+        )
+
+        schema_ctx = self._build_schema_context(full_schema, related_schema)
         capped_limit = min(limit, _MAX_LIMIT)
         dialect = client.db_type
 
@@ -180,7 +192,8 @@ class DbExecutor(BaseExecutor):
                     "question": question,
                 })
                 sql = result.sql.strip()
-                validation_error = self._validate(sql, full_schema, capped_limit)
+                allowed_tables = set(full_schema) | set(related_schema)
+                validation_error = self._validate(sql, allowed_tables, capped_limit)
                 if validation_error:
                     prev_error = validation_error
                     prev_reasoning = result.reasoning
@@ -220,7 +233,7 @@ class DbExecutor(BaseExecutor):
         return QueryResult(
             source_type="database",
             source_id=client_id,
-            table_or_file=", ".join(table_names),
+            table_or_file=", ".join(hit_tables),
             columns=columns,
             rows=rows,
             row_count=len(rows),
@@ -236,57 +249,211 @@ class DbExecutor(BaseExecutor):
     # Schema helpers
     # ------------------------------------------------------------------
 
-    async def _expand_with_fk_tables(
+    async def _find_related_tables(
         self,
         client_id: str,
         user_id: str,
-        table_names: list[str],
+        hit_tables: list[str],
     ) -> list[str]:
-        """Expand table_names with any tables FK-referenced by the retrieved tables.
+        """One-hop FK neighbours of `hit_tables`, both directions, excluding hits.
 
-        Prevents SQL generation failures when a required table (e.g. orders) wasn't
-        returned by retrieval but is referenced via FK from a table that was
-        (e.g. order_items.order_id -> orders.id).
+        Prefers chunk_level='table' rows; if none exist for the client (legacy
+        ingest predating Phase 1), falls back to aggregating from column-chunk
+        metadata. Returns [] when no FK metadata is available.
+
+        Capped at _FK_EXPANSION_MAX_TABLES, ranked by edge count desc then
+        table name asc. A warning is logged when the cap kicks in.
         """
-        if not table_names:
-            return table_names
+        if not hit_tables:
+            return []
 
-        placeholders = ", ".join(f":t{i}" for i in range(len(table_names)))
-        sql = text(f"""
-            SELECT DISTINCT lpe.cmetadata->'data'->>'foreign_key' AS fk
+        hit_set = set(hit_tables)
+        # edge_counts[related_table] = number of FK edges connecting it to the hit set
+        edge_counts: dict[str, int] = defaultdict(int)
+
+        # ---- Primary path: table-level chunks ----
+        sql = text("""
+            SELECT lpe.cmetadata
             FROM langchain_pg_embedding lpe
             JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid
             WHERE lpc.name = 'document_embeddings'
               AND lpe.cmetadata->>'user_id' = :user_id
               AND lpe.cmetadata->>'source_type' = 'database'
               AND lpe.cmetadata->>'database_client_id' = :client_id
-              AND lpe.cmetadata->'data'->>'table_name' IN ({placeholders})
-              AND lpe.cmetadata->'data'->>'foreign_key' IS NOT NULL
+              AND lpe.cmetadata->>'chunk_level' = 'table'
         """)
+        async with _pgvector_engine.connect() as conn:
+            result = await conn.execute(sql, {"user_id": user_id, "client_id": client_id})
+            table_rows = result.fetchall()
+
+        if table_rows:
+            for row in table_rows:
+                data = row.cmetadata.get("data", {})
+                table = data.get("table_name")
+                fks = data.get("foreign_keys") or []
+                if not table:
+                    continue
+                if table in hit_set:
+                    # Outgoing: this hit's FKs point at related tables
+                    for fk in fks:
+                        target = fk.get("target_table")
+                        if target and target not in hit_set:
+                            edge_counts[target] += 1
+                else:
+                    # Incoming: this non-hit table's FKs point into the hit set
+                    for fk in fks:
+                        target = fk.get("target_table")
+                        if target in hit_set:
+                            edge_counts[table] += 1
+        else:
+            # ---- Fallback: aggregate from column chunks ----
+            sql = text("""
+                SELECT lpe.cmetadata->'data'->>'table_name' AS src_table,
+                       lpe.cmetadata->'data'->>'foreign_key' AS fk
+                FROM langchain_pg_embedding lpe
+                JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid
+                WHERE lpc.name = 'document_embeddings'
+                  AND lpe.cmetadata->>'user_id' = :user_id
+                  AND lpe.cmetadata->>'source_type' = 'database'
+                  AND lpe.cmetadata->>'database_client_id' = :client_id
+                  AND lpe.cmetadata->>'chunk_level' = 'column'
+                  AND lpe.cmetadata->'data'->>'foreign_key' IS NOT NULL
+            """)
+            async with _pgvector_engine.connect() as conn:
+                result = await conn.execute(sql, {"user_id": user_id, "client_id": client_id})
+                col_rows = result.fetchall()
+
+            for row in col_rows:
+                src = row.src_table
+                fk = row.fk
+                if not src or not fk:
+                    continue
+                target = fk.split(".", 1)[0]
+                if src in hit_set and target and target not in hit_set:
+                    edge_counts[target] += 1
+                elif src not in hit_set and target in hit_set:
+                    edge_counts[src] += 1
+
+        if not edge_counts:
+            return []
+
+        ranked = sorted(edge_counts.items(), key=lambda kv: (-kv[1], kv[0]))
+        if len(ranked) > _FK_EXPANSION_MAX_TABLES:
+            logger.warning(
+                "fk expansion cap hit",
+                client_id=client_id,
+                total=len(ranked),
+                cap=_FK_EXPANSION_MAX_TABLES,
+                dropped=[t for t, _ in ranked[_FK_EXPANSION_MAX_TABLES:]],
+            )
+            ranked = ranked[:_FK_EXPANSION_MAX_TABLES]
+
+        related = [t for t, _ in ranked]
+        logger.info("fk-related tables", hit=sorted(hit_set), related=related)
+        return related
+
+    async def _fetch_abbreviated_schema(
+        self,
+        client_id: str,
+        user_id: str,
+        table_names: list[str],
+    ) -> dict[str, dict[str, Any]]:
+        """Abbreviated schema: name, row_count, PK, FKs, column names — no profiles.
+
+        Prefers chunk_level='table' rows. Falls back to aggregating column-chunk
+        metadata when table chunks are missing for a given table_name.
 
+        Returns {table_name: {"row_count": int|None, "primary_key": [str],
+        "foreign_keys": [{column, target_table, target_column}],
+        "column_names": [str]}}.
+        """
+        if not table_names:
+            return {}
+
+        placeholders = ", ".join(f":t{i}" for i in range(len(table_names)))
         params: dict[str, Any] = {"user_id": user_id, "client_id": client_id}
         for i, name in enumerate(table_names):
             params[f"t{i}"] = name
 
+        # Primary path: one row per table from chunk_level='table'
+        sql_table = text(f"""
+            SELECT lpe.cmetadata
+            FROM langchain_pg_embedding lpe
+            JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid
+            WHERE lpc.name = 'document_embeddings'
+              AND lpe.cmetadata->>'user_id' = :user_id
+              AND lpe.cmetadata->>'source_type' = 'database'
+              AND lpe.cmetadata->>'database_client_id' = :client_id
+              AND lpe.cmetadata->>'chunk_level' = 'table'
+              AND lpe.cmetadata->'data'->>'table_name' IN ({placeholders})
+        """)
         async with _pgvector_engine.connect() as conn:
-            result = await conn.execute(sql, params)
-            rows = result.fetchall()
-
-        expanded = set(table_names)
-        for row in rows:
-            fk = row.fk  # format: "referred_table.referred_column"
-            if fk:
-                referred_table = fk.split(".")[0]
-                expanded.add(referred_table)
-
-        if expanded != set(table_names):
-            logger.info(
-                "expanded tables via FK",
-                original=sorted(table_names),
-                expanded=sorted(expanded),
-            )
+            result = await conn.execute(sql_table, params)
+            t_rows = result.fetchall()
 
-        return list(expanded)
+        out: dict[str, dict[str, Any]] = {}
+        for row in t_rows:
+            data = row.cmetadata.get("data", {})
+            tname = data.get("table_name")
+            if not tname:
+                continue
+            out[tname] = {
+                "row_count": data.get("row_count"),
+                "primary_key": list(data.get("primary_key") or []),
+                "foreign_keys": list(data.get("foreign_keys") or []),
+                "column_names": list(data.get("column_names") or []),
+            }
+
+        # Fallback for tables with no table-chunk: aggregate column chunks
+        missing = [t for t in table_names if t not in out]
+        if missing:
+            placeholders_m = ", ".join(f":m{i}" for i in range(len(missing)))
+            params_m: dict[str, Any] = {"user_id": user_id, "client_id": client_id}
+            for i, name in enumerate(missing):
+                params_m[f"m{i}"] = name
+            sql_col = text(f"""
+                SELECT lpe.cmetadata
+                FROM langchain_pg_embedding lpe
+                JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid
+                WHERE lpc.name = 'document_embeddings'
+                  AND lpe.cmetadata->>'user_id' = :user_id
+                  AND lpe.cmetadata->>'source_type' = 'database'
+                  AND lpe.cmetadata->>'database_client_id' = :client_id
+                  AND lpe.cmetadata->>'chunk_level' = 'column'
+                  AND lpe.cmetadata->'data'->>'table_name' IN ({placeholders_m})
+                ORDER BY lpe.cmetadata->'data'->>'table_name', lpe.cmetadata->'data'->>'column_name'
+            """)
+            async with _pgvector_engine.connect() as conn:
+                result = await conn.execute(sql_col, params_m)
+                c_rows = result.fetchall()
+
+            agg: dict[str, dict[str, Any]] = {
+                t: {"row_count": None, "primary_key": [], "foreign_keys": [], "column_names": []}
+                for t in missing
+            }
+            for row in c_rows:
+                data = row.cmetadata.get("data", {})
+                tname = data.get("table_name")
+                cname = data.get("column_name")
+                if not tname or tname not in agg or not cname:
+                    continue
+                bucket = agg[tname]
+                bucket["column_names"].append(cname)
+                if data.get("is_primary_key"):
+                    bucket["primary_key"].append(cname)
+                fk = data.get("foreign_key")
+                if fk:
+                    target_table, _, target_col = fk.partition(".")
+                    bucket["foreign_keys"].append({
+                        "column": cname,
+                        "target_table": target_table,
+                        "target_column": target_col,
+                    })
+            for t, v in agg.items():
+                if v["column_names"]:
+                    out[t] = v
+
+        return out
 
     async def _fetch_full_schema(
         self,
@@ -307,6 +474,7 @@ class DbExecutor(BaseExecutor):
             WHERE lpc.name = 'document_embeddings'
               AND lpe.cmetadata->>'user_id' = :user_id
               AND lpe.cmetadata->>'source_type' = 'database'
+              AND lpe.cmetadata->>'chunk_level' = 'column'
               AND lpe.cmetadata->>'database_client_id' = :client_id
               AND lpe.cmetadata->'data'->>'table_name' IN ({placeholders})
             ORDER BY lpe.cmetadata->'data'->>'table_name', lpe.cmetadata->'data'->>'column_name'
@@ -334,7 +502,11 @@ class DbExecutor(BaseExecutor):
                 })
         return dict(schema)
 
-    def _build_schema_context(self, schema: dict[str, list[dict[str, Any]]]) -> str:
+    def _build_schema_context(
+        self,
+        schema: dict[str, list[dict[str, Any]]],
+        related_schema: dict[str, dict[str, Any]] | None = None,
+    ) -> str:
         lines: list[str] = []
         for table, columns in schema.items():
             lines.append(f"Table: {table}")
@@ -352,14 +524,47 @@ class DbExecutor(BaseExecutor):
                         lines.append(f"    {line}")
                         break
             lines.append("")
+
+        related_block = self._build_related_schema_block(related_schema or {})
+        if related_block:
+            lines.append(related_block)
+
         return "\n".join(lines).strip()
 
+    def _build_related_schema_block(self, related_schema: dict[str, dict[str, Any]]) -> str:
+        """Format the abbreviated FK-related-tables section. Empty string when no related."""
+        if not related_schema:
+            return ""
+        lines: list[str] = ["Related tables (one hop via FK, abbreviated — use for JOINs only):"]
+        for table, info in related_schema.items():
+            row_count = info.get("row_count")
+            header = f"- {table} ({row_count} rows)" if row_count is not None else f"- {table}"
+            lines.append(header)
+            pk = info.get("primary_key") or []
+            lines.append(f"    Primary key: {', '.join(pk) if pk else '(none)'}")
+            fks = info.get("foreign_keys") or []
+            if fks:
+                fk_strs = [
+                    f"{fk.get('column')} -> {fk.get('target_table')}.{fk.get('target_column')}"
+                    for fk in fks
+                ]
+                lines.append(f"    Foreign keys: {', '.join(fk_strs)}")
+            else:
+                lines.append("    Foreign keys: (none)")
+            cols = info.get("column_names") or []
+            lines.append(f"    Columns: {', '.join(cols)}")
+        return "\n".join(lines)
+
     # ------------------------------------------------------------------
     # Guardrails
     # ------------------------------------------------------------------
 
-    def _validate(self, sql: str, schema: dict[str, list[dict]], limit: int) -> str:
-        """Return an error string if validation fails, empty string if OK."""
+    def _validate(self, sql: str, allowed_tables: set[str], limit: int) -> str:
+        """Return an error string if validation fails, empty string if OK.
+
+        `allowed_tables` is the union of hit-table names and FK-related table
+        names — both are legal targets for SELECT/JOIN.
+        """
         # Layer 1: sqlglot parse + SELECT-only check
         try:
             parsed = sqlglot.parse_one(sql)
@@ -374,7 +579,7 @@ class DbExecutor(BaseExecutor):
             return f"DML ({type(node).__name__}) is not allowed."
 
         # Layer 2: schema grounding — table names
-        known_tables = {t.lower() for t in schema}
+        known_tables = {t.lower() for t in allowed_tables}
         for tbl in parsed.find_all(exp.Table):
             name = tbl.name.lower()
             if name and name not in known_tables:

src/rag/retrievers/schema.py

@@ -47,6 +47,7 @@ class SchemaRetriever(BaseRetriever):
             WHERE lpc.name = 'document_embeddings'
               AND lpe.cmetadata->>'user_id' = :user_id
               AND lpe.cmetadata->>'source_type' = 'database'
+              AND lpe.cmetadata->>'chunk_level' = 'column'
             ORDER BY lpe.embedding <=> '{emb_str}'::vector ASC
             LIMIT :k
         """)
@@ -110,6 +111,7 @@ class SchemaRetriever(BaseRetriever):
             WHERE lpc.name = 'document_embeddings'
               AND lpe.cmetadata->>'user_id' = :user_id
               AND lpe.cmetadata->>'source_type' = 'database'
+              AND lpe.cmetadata->>'chunk_level' = 'column'
               AND to_tsvector('english', lpe.document) @@ plainto_tsquery('english', :query)
             ORDER BY rank DESC
             LIMIT :k