[KM-516][KM-517] add new feature; ai can now see table & column names that have fk relationship with retrieved result

src/query/executors/db_executor.py

@@ -137,12 +137,13 @@ class DbExecutor(BaseExecutor):
             logger.warning("unsupported db_type for query execution", db_type=client.db_type)
             return None
 
-        # Distinct table names from retrieval results
+        # Distinct table names from retrieval results, expanded via FK relationships
         table_names = 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)
 
         full_schema = await self._fetch_full_schema(client_id, table_names, user_id)
         if not full_schema:
@@ -235,6 +236,58 @@ class DbExecutor(BaseExecutor):
     # Schema helpers
     # ------------------------------------------------------------------
 
+    async def _expand_with_fk_tables(
+        self,
+        client_id: str,
+        user_id: str,
+        table_names: list[str],
+    ) -> list[str]:
+        """Expand table_names with any tables FK-referenced by the retrieved tables.
+
+        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).
+        """
+        if not table_names:
+            return table_names
+
+        placeholders = ", ".join(f":t{i}" for i in range(len(table_names)))
+        sql = text(f"""
+            SELECT DISTINCT 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->'data'->>'table_name' IN ({placeholders})
+              AND lpe.cmetadata->'data'->>'foreign_key' IS NOT NULL
+        """)
+
+        params: dict[str, Any] = {"user_id": user_id, "client_id": client_id}
+        for i, name in enumerate(table_names):
+            params[f"t{i}"] = name
+
+        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),
+            )
+
+        return list(expanded)
+
     async def _fetch_full_schema(
         self,
         client_id: str,