From e45994c31a8d89687070d0341b7c513c8781720a Mon Sep 17 00:00:00 2001 From: Kelton Karboviak Date: Wed, 31 Jul 2024 12:40:11 -0500 Subject: [PATCH 1/2] Add `query_id` of last executed statement to Adapter Response I want to be able to capture the `query_id` of the statement used to load a model. I ripped this off of the dbt-snowflake adapter and how they are setting `query_id` on the adapter response: https://github.com/dbt-labs/dbt-snowflake/blob/f95b9192f6eec9af4e30eaab87f9e3412febf7d1/dbt/adapters/snowflake/connections.py#L456-L461 I'm sure this doesn't handle 100% of accurately capturing the correct query ID for a particular load (I'm sure there are cases where multiple statements are run in order to perform the model or snapshot load), but I think this is a good start. --- dbt/adapters/redshift/connections.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/dbt/adapters/redshift/connections.py b/dbt/adapters/redshift/connections.py index d3fbcafea..e96d9896b 100644 --- a/dbt/adapters/redshift/connections.py +++ b/dbt/adapters/redshift/connections.py @@ -72,6 +72,11 @@ class RedshiftSSLMode(StrEnum): } +@dataclass +class RedshiftAdapterResponse(AdapterResponse): + query_id: int + + @dataclass class RedshiftSSLConfig(dbtClassMixin, Replaceable): # type: ignore ssl: bool = True @@ -314,13 +319,24 @@ def _get_backend_pid(cls, connection): res = c.execute(sql).fetchone() return res[0] + @classmethod + def _get_last_query_id(cls, cursor): + sql = "select pg_last_query_id();" + res = cursor.execute(sql).fetchone() + return res[0] + @classmethod def get_response(cls, cursor: redshift_connector.Cursor) -> AdapterResponse: # redshift_connector.Cursor doesn't have a status message attribute but # this function is only used for successful run, so we can just return a dummy rows = cursor.rowcount message = "SUCCESS" - return AdapterResponse(_message=message, rows_affected=rows) + query_id = cls._get_last_query_id(cursor) + return RedshiftAdapterResponse( + _message=message, + rows_affected=rows, + query_id=query_id, + ) @contextmanager def exception_handler(self, sql): From 2c52d79a77554def52342b3d350e12e0addb9f93 Mon Sep 17 00:00:00 2001 From: Kelton Karboviak Date: Wed, 31 Jul 2024 13:27:54 -0500 Subject: [PATCH 2/2] Add default of `-1` for `query_id` --- dbt/adapters/redshift/connections.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dbt/adapters/redshift/connections.py b/dbt/adapters/redshift/connections.py index e96d9896b..a2ead4022 100644 --- a/dbt/adapters/redshift/connections.py +++ b/dbt/adapters/redshift/connections.py @@ -74,7 +74,7 @@ class RedshiftSSLMode(StrEnum): @dataclass class RedshiftAdapterResponse(AdapterResponse): - query_id: int + query_id: int = -1 @dataclass @@ -333,8 +333,8 @@ def get_response(cls, cursor: redshift_connector.Cursor) -> AdapterResponse: message = "SUCCESS" query_id = cls._get_last_query_id(cursor) return RedshiftAdapterResponse( - _message=message, - rows_affected=rows, + _message=message, + rows_affected=rows, query_id=query_id, )