add query_id as optional to SQLQueryStatus (#342)

dbt-labs · Nov 5, 2024 · da2a380 · da2a380
1 parent b12c987
commit da2a380
Show file tree

Hide file tree

Showing 5 changed files with 115 additions and 103 deletions.
diff --git a/.changes/unreleased/Under the Hood-20241105-121440.yaml b/.changes/unreleased/Under the Hood-20241105-121440.yaml
@@ -0,0 +1,6 @@
+kind: Under the Hood
+body: Add `query_id` to SQLQueryStatus
+time: 2024-11-05T12:14:40.181931-05:00
+custom:
+    Author: cmcarthur
+    Issue: "342"
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;
+    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,
                 )
             )