Add query_id of last executed statement to Adapter Response

[KeltonKarboviak]

resolves dbt-labs/dbt-adapters#630
docs dbt-labs/docs.getdbt.com/#

Problem

I want to be able to capture the query_id of the statement used to load a model.

Solution

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 realize it is different from Snowflake in that appears that the query ID is available directly on the cursor object, and with Redshift you need to use the pg_last_query_id() function in order to get the query ID.

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.

Please let me know what I'm missing here in the potential of this breaking something else due to injecting another another query into a model load flow.

Checklist

• I have read the contributing guide and understand what's expected of me
• I have run this code in development and it appears to resolve the stated issue
• This PR includes tests, or tests are not required/relevant for this PR
• This PR has no interface changes (e.g. macros, cli, logs, json artifacts, config files, adapter interface, etc) or this PR has already received feedback and approval from Product or DX

[KeltonKarboviak]

Well there is definitely a side effect to running the _get_last_query_id method. After adding that and running a single functional test, it blows up:

$ .tox/py38-redshift/bin/python -m pytest -vvv -s --dist=loadscope tests/functional/adapter/test_basic.py::TestSimpleMaterializationsRedshift::test_base
============================================================================= test session starts ==============================================================================
platform darwin -- Python 3.8.14, pytest-7.4.4, pluggy-1.5.0 -- /Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/bin/python
cachedir: .pytest_cache
rootdir: /Users/kelton.karboviak/Code/dbt-redshift
configfile: pytest.ini
plugins: ddtrace-2.3.0, csv-3.0.0, logbook-1.2.0, dotenv-0.5.2, xdist-3.6.1
collected 1 item

tests/functional/adapter/test_basic.py::TestSimpleMaterializationsRedshift::test_base
=== Test project_root: /private/var/folders/s8/7lghdz255_lb4khl4xpwvz3j50ydy1/T/pytest-of-kelton.karboviak/pytest-18/project0


Invoking dbt with ['seed']
19:53:50  Running with dbt=1.9.0-a1
19:53:50  Registered adapter: redshift=1.9.0-a1
19:53:50  Unable to do partial parsing because saved manifest not found. Starting full parse.
19:53:51  Found 3 models, 1 seed, 1 source, 484 macros
19:53:51
19:53:53  Concurrency: 1 threads (target='default')
19:53:53
19:53:53  1 of 1 START seed file test17224556293545641964_test_basic.base ................ [RUN]
19:53:54  1 of 1 OK loaded seed file test17224556293545641964_test_basic.base ............ [INSERT 10 in 1.02s]
19:53:55
19:53:55  Finished running 1 seed in 0 hours 0 minutes and 3.73 seconds (3.73s).
19:53:55
19:53:55  Completed successfully
19:53:55
19:53:55  Done. PASS=1 WARN=0 ERROR=0 SKIP=0 TOTAL=1


Invoking dbt with ['run']
19:53:55  Running with dbt=1.9.0-a1
19:53:55  Registered adapter: redshift=1.9.0-a1
19:53:55  Found 3 models, 1 seed, 1 source, 484 macros
19:53:55
19:53:57  Concurrency: 1 threads (target='default')
19:53:57
19:53:57  1 of 3 START sql table model test17224556293545641964_test_basic.swappable ..... [RUN]
19:53:59  1 of 3 OK created sql table model test17224556293545641964_test_basic.swappable  [SUCCESS in 1.36s]
19:53:59  2 of 3 START sql table model test17224556293545641964_test_basic.table_model ... [RUN]
19:54:00  2 of 3 OK created sql table model test17224556293545641964_test_basic.table_model  [SUCCESS in 1.38s]
19:54:00  3 of 3 START sql view model test17224556293545641964_test_basic.view_model ..... [RUN]
19:54:01  3 of 3 OK created sql view model test17224556293545641964_test_basic.view_model  [SUCCESS in 1.52s]
19:54:02
19:54:02  Finished running 2 table models, 1 view model in 0 hours 0 minutes and 6.98 seconds (6.98s).
19:54:02
19:54:02  Completed successfully
19:54:02
19:54:02  Done. PASS=3 WARN=0 ERROR=0 SKIP=0 TOTAL=3
BASIS "dbt_functional_tests"."test17224556293545641964_test_basic"."base"
COMPARES [<RedshiftRelation "dbt_functional_tests"."test17224556293545641964_test_basic"."view_model">, <RedshiftRelation "dbt_functional_tests"."test17224556293545641964_test_basic"."table_model">, <RedshiftRelation "dbt_functional_tests"."test17224556293545641964_test_basic"."swappable">]
FAILED

=================================================================================== FAILURES ===================================================================================
_________________________________________________________________ TestSimpleMaterializationsRedshift.test_base _________________________________________________________________

self = <redshift_connector.core.Connection object at 0x12497e640>, cursor = <redshift_connector.cursor.Cursor object at 0x12494d5b0>
operation = '/* {"app": "dbt", "dbt_version": "1.9.0a1", "profile_name": "test", "target_name": "default", "connection_name": "_te...s row_count_difference,\n    diff_count.num_missing as num_mismatched\nfrom row_count_diff\njoin diff_count using (id)'
vals = None

    def execute(self: "Connection", cursor: Cursor, operation: str, vals) -> None:
        """
        Executes a database operation. Parameters may be provided as a sequence, or as a mapping, depending upon the value of `redshift_connector.paramstyle`.

        Parameters
        ----------
        cursor : :class:`Cursor`
        operation : str The SQL statement to execute.
        vals : If `redshift_connector.paramstyle` is `qmark`, `numeric`, or `format` this argument should be an array of parameters to bind into the statement. If `redshift_connector.paramstyle` is `named` the argument should be a `dict` mapping of parameters. If `redshift_connector.paramstyle` is `pyformat`, the argument value may be either an array or mapping.

        Returns
        -------
        None:None
        """
        _logger.debug("Connection.execute()")

        # get the process ID of the calling process.
        pid: int = getpid()

        args: typing.Tuple[typing.Optional[typing.Tuple[str, typing.Any]], ...] = ()
        # transforms user provided bind parameters to server friendly bind parameters
        params: typing.Tuple[typing.Optional[typing.Tuple[int, int, typing.Callable]], ...] = ()
        has_bind_parameters: bool = False if vals is None else True
        # multi dimensional dictionary to store the data
        # cache = self._caches[cursor.paramstyle][pid]
        # cache = {'statement': {}, 'ps': {}}
        # statement stores the data of the statement, ps store the data of the prepared statement
        # statement = {operation(query): tuple from 'convert_paramstyle'(statement, make_args)}
        try:
            cache = self._caches[cursor.paramstyle][pid]
        except KeyError:
            try:
                param_cache = self._caches[cursor.paramstyle]
            except KeyError:
                param_cache = self._caches[cursor.paramstyle] = {}

            try:
                cache = param_cache[pid]
            except KeyError:
                cache = param_cache[pid] = {"statement": {}, "ps": {}}

        try:
            statement, make_args = cache["statement"][operation]
        except KeyError:
            if has_bind_parameters:
                statement, make_args = cache["statement"][operation] = convert_paramstyle(cursor.paramstyle, operation)
            else:
                # use a no-op make_args in lieu of parsing the sql statement
                statement, make_args = cache["statement"][operation] = operation, lambda p: ()
        if has_bind_parameters:
            args = make_args(vals)
            _logger.debug("User provided vals converted to %s args", len(args))
            # change the args to the format that the DB will identify
            # take reference from self.py_types
            params = self.make_params(args)
            _logger.debug("args converted to %s params", len(params))
        key = operation, params

        try:
>           ps = cache["ps"][key]
E           KeyError: ('/* {"app": "dbt", "dbt_version": "1.9.0a1", "profile_name": "test", "target_name": "default", "connection_name": "_test"} */\nwith diff_count as (\n    SELECT\n        1 as id,\n        COUNT(*) as num_missing FROM (\n            (SELECT  FROM "dbt_functional_tests"."test17224556293545641964_test_basic"."base" EXCEPT\n             SELECT  FROM "dbt_functional_tests"."test17224556293545641964_test_basic"."view_model")\n             UNION ALL\n            (SELECT  FROM "dbt_functional_tests"."test17224556293545641964_test_basic"."view_model" EXCEPT\n             SELECT  FROM "dbt_functional_tests"."test17224556293545641964_test_basic"."base")\n        ) as a\n), table_a as (\n    SELECT COUNT(*) as num_rows FROM "dbt_functional_tests"."test17224556293545641964_test_basic"."base"\n), table_b as (\n    SELECT COUNT(*) as num_rows FROM "dbt_functional_tests"."test17224556293545641964_test_basic"."view_model"\n), row_count_diff as (\n    select\n        1 as id,\n        table_a.num_rows - table_b.num_rows as difference\n    from table_a, table_b\n)\nselect\n    row_count_diff.difference as row_count_difference,\n    diff_count.num_missing as num_mismatched\nfrom row_count_diff\njoin diff_count using (id)', ())

/Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/lib/python3.8/site-packages/redshift_connector/core.py:1821: KeyError

During handling of the above exception, another exception occurred:

self = <dbt.adapters.redshift.connections.RedshiftConnectionManager object at 0x12448b700>
sql = '/* {"app": "dbt", "dbt_version": "1.9.0a1", "profile_name": "test", "target_name": "default", "connection_name": "_te...s row_count_difference,\n    diff_count.num_missing as num_mismatched\nfrom row_count_diff\njoin diff_count using (id)'

    @contextmanager
    def exception_handler(self, sql):
        try:
>           yield

/Users/kelton.karboviak/Code/dbt-redshift/dbt/adapters/redshift/connections.py:345:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/lib/python3.8/site-packages/dbt/adapters/sql/connections.py:93: in add_query
    cursor.execute(sql, bindings)
/Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/lib/python3.8/site-packages/redshift_connector/cursor.py:248: in execute
    raise e
/Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/lib/python3.8/site-packages/redshift_connector/cursor.py:241: in execute
    self._c.execute(self, operation, args)
/Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/lib/python3.8/site-packages/redshift_connector/core.py:1902: in execute
    self.handle_messages(cursor)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <redshift_connector.core.Connection object at 0x12497e640>, cursor = <redshift_connector.cursor.Cursor object at 0x12494d5b0>

    def handle_messages(self: "Connection", cursor: Cursor) -> None:
        """
        Reads messages formatted in ordinance with Amazon Redshift wire protocol, modifying the connection and cursor.

        Parameters
        ----------
        :param cursor: `Cursor`
            The `Cursor` object associated with the given connection object.

        Returns
        -------
        None:None
        """
        code = self.error = None

        while code != READY_FOR_QUERY:
            buffer = self._read(5)

            if len(buffer) == 0:
                if self._usock.timeout is not None:
                    raise InterfaceError(
                        "BrokenPipe: server socket closed. We noticed a timeout is set for this connection. Consider "
                        "raising the timeout or defaulting timeout to none."
                    )
                else:
                    raise InterfaceError(
                        "BrokenPipe: server socket closed. Please check that client side networking configurations such "
                        "as Proxies, firewalls, VPN, etc. are not affecting your network connection."
                    )

            code, data_len = ci_unpack(buffer)
            _logger.debug("Message received from BE with code %s length %s ", code, data_len)
            self.message_types[code](self._read(data_len - 4), cursor)

        if self.error is not None:
>           raise self.error
E           redshift_connector.error.ProgrammingError: {'S': 'ERROR', 'C': '42601', 'M': 'syntax error at or near "FROM"', 'P': '235', 'F': '/home/ec2-user/padb/src/pg/src/backend/parser/parser_scan.l', 'L': '828', 'R': 'yyerror'}

/Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/lib/python3.8/site-packages/redshift_connector/core.py:2194: ProgrammingError

The above exception was the direct cause of the following exception:

self = <test_basic.TestSimpleMaterializationsRedshift object at 0x1201258e0>, project = <dbt.tests.fixtures.project.TestProjInfo object at 0x120ecd940>

    @pytest.mark.flaky
    def test_base(self, project):
>       super().test_base(project)

/Users/kelton.karboviak/Code/dbt-redshift/tests/functional/adapter/test_basic.py:43:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/lib/python3.8/site-packages/dbt/tests/adapter/basic/test_base.py:75: in test_base
    check_relations_equal(project.adapter, ["base", "view_model", "table_model", "swappable"])
/Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/lib/python3.8/site-packages/dbt/tests/util.py:402: in check_relations_equal
    return check_relations_equal_with_relations(
/Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/lib/python3.8/site-packages/dbt/tests/util.py:442: in check_relations_equal_with_relations
    _, tbl = adapter.execute(sql, fetch=True)
/Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/lib/python3.8/site-packages/dbt/adapters/base/impl.py:356: in execute
    return self.connections.execute(sql=sql, auto_begin=auto_begin, fetch=fetch, limit=limit)
/Users/kelton.karboviak/Code/dbt-redshift/dbt/adapters/redshift/connections.py:423: in execute
    _, cursor = self.add_query(sql, auto_begin)
/Users/kelton.karboviak/Code/dbt-redshift/dbt/adapters/redshift/connections.py:451: in add_query
    connection, cursor = super().add_query(
/Users/kelton.karboviak/Code/dbt-redshift/.tox/py38-redshift/lib/python3.8/site-packages/dbt/adapters/sql/connections.py:103: in add_query
    return connection, cursor
/Users/kelton.karboviak/.asdf/installs/python/3.8.14/lib/python3.8/contextlib.py:131: in __exit__
    self.gen.throw(type, value, traceback)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <dbt.adapters.redshift.connections.RedshiftConnectionManager object at 0x12448b700>
sql = '/* {"app": "dbt", "dbt_version": "1.9.0a1", "profile_name": "test", "target_name": "default", "connection_name": "_te...s row_count_difference,\n    diff_count.num_missing as num_mismatched\nfrom row_count_diff\njoin diff_count using (id)'

    @contextmanager
    def exception_handler(self, sql):
        try:
            yield
        except redshift_connector.DatabaseError as e:
            try:
                err_msg = e.args[0]["M"]  # this is a type redshift sets, so we must use these keys
            except Exception:
                err_msg = str(e).strip()
            logger.debug(f"Redshift error: {err_msg}")
            self.rollback_if_open()
>           raise DbtDatabaseError(err_msg) from e
E           dbt_common.exceptions.base.DbtDatabaseError: Database Error
E             syntax error at or near "FROM"

/Users/kelton.karboviak/Code/dbt-redshift/dbt/adapters/redshift/connections.py:353: DbtDatabaseError
=========================================================================== short test summary info ============================================================================
FAILED tests/functional/adapter/test_basic.py::TestSimpleMaterializationsRedshift::test_base - dbt_common.exceptions.base.DbtDatabaseError: Database Error
============================================================================== 1 failed in 17.82s ==============================================================================

[github-actions]

This PR has been marked as Stale because it has been open with no activity as of late. If you would like the PR to remain open, please comment on the PR or else it will be closed in 7 days.

[github-actions]

Although we are closing this PR as stale, it can still be reopened to continue development. Just add a comment to notify the maintainers.