From 5841d52792c2acfa27b299c330dd4e5c573d2a3c Mon Sep 17 00:00:00 2001
From: Gerda Shank <gerda@dbtlabs.com>
Date: Tue, 13 Feb 2024 11:58:42 -0500
Subject: [PATCH] Set query headers when manifest is passed in to dbtRunner
 (#9556)

---
 .../unreleased/Fixes-20240212-144733.yaml     |  6 ++++
 core/dbt/cli/requires.py                      |  5 ++++
 .../functional/dbt_runner/test_dbt_runner.py  | 30 +++++++++++++++++++
 3 files changed, 41 insertions(+)
 create mode 100644 .changes/unreleased/Fixes-20240212-144733.yaml

diff --git a/.changes/unreleased/Fixes-20240212-144733.yaml b/.changes/unreleased/Fixes-20240212-144733.yaml
new file mode 100644
index 00000000000..cdf9df56b78
--- /dev/null
+++ b/.changes/unreleased/Fixes-20240212-144733.yaml
@@ -0,0 +1,6 @@
+kind: Fixes
+body: Set query headers when manifest is passed in to dbtRunner
+time: 2024-02-12T14:47:33.092877-05:00
+custom:
+  Author: gshank
+  Issue: "9546"
diff --git a/core/dbt/cli/requires.py b/core/dbt/cli/requires.py
index 101a63c137d..ccd5ffc7150 100644
--- a/core/dbt/cli/requires.py
+++ b/core/dbt/cli/requires.py
@@ -15,6 +15,7 @@
 from dbt.cli.flags import Flags
 from dbt.config import RuntimeConfig
 from dbt.config.runtime import load_project, load_profile, UnsetProfile
+from dbt.context.manifest import generate_query_header_context
 
 from dbt_common.events.base_types import EventLevel
 from dbt_common.events.functions import (
@@ -290,6 +291,10 @@ def wrapper(*args, **kwargs):
                 adapter = get_adapter(runtime_config)
                 adapter.set_macro_context_generator(generate_runtime_macro_context)
                 adapter.set_macro_resolver(ctx.obj["manifest"])
+                query_header_context = generate_query_header_context(
+                    adapter.config, ctx.obj["manifest"]
+                )
+                adapter.connections.set_query_header(query_header_context)
             return func(*args, **kwargs)
 
         return update_wrapper(wrapper, func)
diff --git a/tests/functional/dbt_runner/test_dbt_runner.py b/tests/functional/dbt_runner/test_dbt_runner.py
index 85e3f281751..b951820d4e8 100644
--- a/tests/functional/dbt_runner/test_dbt_runner.py
+++ b/tests/functional/dbt_runner/test_dbt_runner.py
@@ -6,6 +6,8 @@
 from dbt.cli.main import dbtRunner
 from dbt.exceptions import DbtProjectError
 from dbt.adapters.factory import reset_adapters, FACTORY
+from dbt.tests.util import read_file, write_file
+from dbt.version import __version__ as dbt_version
 
 
 class TestDbtRunner:
@@ -90,3 +92,31 @@ def test_pass_in_manifest(self, project, dbt):
         # Check that the adapters are registered again.
         assert result.success
         assert len(FACTORY.adapters) == 1
+
+
+class TestDbtRunnerQueryComments:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "models.sql": "select 1 as id",
+        }
+
+    @pytest.fixture(scope="class")
+    def project_config_update(self):
+        return {
+            "query-comment": {
+                "comment": f"comment: {dbt_version}",
+                "append": True,
+            }
+        }
+
+    def test_query_comment_saved_manifest(self, project, logs_dir):
+        dbt = dbtRunner()
+        dbt.invoke(["build", "--select", "models"])
+        result = dbt.invoke(["parse"])
+        write_file("", logs_dir, "dbt.log")
+        # pass in manifest from parse command
+        dbt = dbtRunner(result.result)
+        dbt.invoke(["build", "--select", "models"])
+        log_file = read_file(logs_dir, "dbt.log")
+        assert f"comment: {dbt_version}" in log_file