add query_id as optional to SQLQueryStatus

dbt-labs · Oct 30, 2024 · 26a01a8 · 26a01a8
1 parent b12c987
commit 26a01a8
Show file tree

Hide file tree

Showing 4 changed files with 110 additions and 104 deletions.
diff --git a/dbt/adapters/contracts/connection.py b/dbt/adapters/contracts/connection.py
@@ -41,6 +41,7 @@ class AdapterResponse(dbtClassMixin):
     _message: str
     code: Optional[str] = None
     rows_affected: Optional[int] = None
+    query_id: Optional[str] = None
 
     def __str__(self):
         return self._message

diff --git a/dbt/adapters/events/adapter_types.proto b/dbt/adapters/events/adapter_types.proto
@@ -266,6 +266,7 @@ message SQLQueryStatus {
     AdapterNodeInfo node_info = 1;
     string status = 2;
     float elapsed = 3;
+    optional string query_id = 4;
 }
 
 message SQLQueryStatusMsg {

diff --git a/dbt/adapters/events/adapter_types_pb2.py b/dbt/adapters/events/adapter_types_pb2.py
diff --git a/dbt/adapters/sql/connections.py b/dbt/adapters/sql/connections.py
@@ -92,11 +92,14 @@ def add_query(
             cursor = connection.handle.cursor()
             cursor.execute(sql, bindings)
 
+            result = self.get_response(cursor)
+
             fire_event(
                 SQLQueryStatus(
-                    status=str(self.get_response(cursor)),
+                    status=str(result),
                     elapsed=time.perf_counter() - pre,
                     node_info=get_node_info(),
+                    query_id=result.query_id,
                 )
             )