Skip to content

Commit

Permalink
Set query headers when manifest is passed in to dbtRunner (#9556)
Browse files Browse the repository at this point in the history
  • Loading branch information
gshank authored Feb 13, 2024
1 parent 03a4d11 commit 5841d52
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 0 deletions.
6 changes: 6 additions & 0 deletions .changes/unreleased/Fixes-20240212-144733.yaml
Original file line number Diff line number Diff line change
@@ -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"
5 changes: 5 additions & 0 deletions core/dbt/cli/requires.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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)
Expand Down
30 changes: 30 additions & 0 deletions tests/functional/dbt_runner/test_dbt_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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

0 comments on commit 5841d52

Please sign in to comment.