[KM 556] add columns field to QueryResult + update DbExecutor

src/query/executor/base.py

@@ -11,6 +11,7 @@ from ..ir.models import QueryIR
 class QueryResult:
     source_id: str
     backend: str  # "sql" | "tabular"
+    columns: list[str] = field(default_factory=list)
     rows: list[dict[str, Any]] = field(default_factory=list)
     row_count: int = 0
     truncated: bool = False

src/query/executor/db.py

@@ -79,7 +79,7 @@ class DbExecutor(BaseExecutor):
                 )
             creds = decrypt_credentials_dict(client.credentials)
 
-            rows = await asyncio.wait_for(
+            columns, rows = await asyncio.wait_for(
                 asyncio.to_thread(self._run_sync, client.db_type, creds, compiled),
                 timeout=_QUERY_TIMEOUT_SECONDS,
             )
@@ -97,6 +97,7 @@ class DbExecutor(BaseExecutor):
             return QueryResult(
                 source_id=ir.source_id,
                 backend="sql",
+                columns=columns,
                 rows=capped,
                 row_count=len(capped),
                 truncated=truncated,
@@ -175,7 +176,7 @@ class DbExecutor(BaseExecutor):
             )
 
     @staticmethod
-    def _run_sync(db_type: str, creds: dict, compiled: CompiledSql) -> list[dict]:
+    def _run_sync(db_type: str, creds: dict, compiled: CompiledSql) -> tuple[list[str], list[dict]]:
         with db_pipeline_service.engine_scope(db_type, creds) as engine:
             with engine.connect() as conn:
                 if db_type in _POSTGRES_LIKE:
@@ -185,4 +186,6 @@ class DbExecutor(BaseExecutor):
                         text(f"SET statement_timeout = {_QUERY_TIMEOUT_SECONDS * 1000}")
                     )
                 result = conn.execute(text(compiled.sql), compiled.params)
-                return [dict(row) for row in result.mappings()]
+                columns = list(result.keys())
+                rows = [dict(row) for row in result.mappings()]
+                return columns, rows