Skip to content

Commit

Permalink
add query_id to SQLQueryStatus (demonstration only)
Browse files Browse the repository at this point in the history
  • Loading branch information
cmcarthur committed Oct 17, 2024
1 parent 7dea145 commit 6b55e7c
Showing 1 changed file with 16 additions and 2 deletions.
18 changes: 16 additions & 2 deletions dbt/adapters/snowflake/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
from contextlib import contextmanager
from dataclasses import dataclass
from io import StringIO
from time import sleep
from time import sleep, perf_counter

from typing import Optional, Tuple, Union, Any, List, Iterable, TYPE_CHECKING

Expand Down Expand Up @@ -43,8 +43,11 @@
DbtRuntimeError,
DbtConfigError,
)
from dbt_common.events.contextvars import get_node_info
from dbt_common.exceptions import DbtDatabaseError
from dbt_common.events.functions import fire_event
from dbt_common.record import get_record_mode_from_env, RecorderMode
from dbt.adapters.events.types import SQLQueryStatus
from dbt.adapters.exceptions.connection import FailedToConnectError
from dbt.adapters.contracts.connection import AdapterResponse, Connection, Credentials
from dbt.adapters.sql import SQLConnectionManager
Expand Down Expand Up @@ -86,7 +89,7 @@ def snowflake_private_key(private_key: RSAPrivateKey) -> bytes:

@dataclass
class SnowflakeAdapterResponse(AdapterResponse):
query_id: str = ""
pass


@dataclass
Expand Down Expand Up @@ -536,6 +539,8 @@ def add_query(
bindings: Optional[Any] = None,
abridge_sql_log: bool = False,
) -> Tuple[Connection, Any]:
pre = perf_counter()

if bindings:
# The snowflake connector is stricter than, e.g., psycopg2 -
# which allows any iterable thing to be passed as a binding.
Expand All @@ -561,6 +566,15 @@ def add_query(
if cursor is None:
self._raise_cursor_not_found_error(sql)

fire_event(
SQLQueryStatus(
status=str(self.get_response(cursor)),
elapsed=perf_counter() - pre,
node_info=get_node_info(),
query_id=cursor.sfqid,
)
)

return connection, cursor

def _stripped_queries(self, sql: str) -> List[str]:
Expand Down

0 comments on commit 6b55e7c

Please sign in to comment.