diff --git a/.changes/unreleased/Fixes-20231005-235950.yaml b/.changes/unreleased/Fixes-20231005-235950.yaml new file mode 100644 index 000000000..bf0bf6fa6 --- /dev/null +++ b/.changes/unreleased/Fixes-20231005-235950.yaml @@ -0,0 +1,6 @@ +kind: Fixes +body: Fix issue where job labels are not rendered when using macro for query comment +time: 2023-10-05T23:59:50.077842+02:00 +custom: + Author: kodaho mikealfare + Issue: "863" diff --git a/dbt/adapters/bigquery/connections.py b/dbt/adapters/bigquery/connections.py index 39deedbba..6ebd0477d 100644 --- a/dbt/adapters/bigquery/connections.py +++ b/dbt/adapters/bigquery/connections.py @@ -428,6 +428,18 @@ def get_table_from_response(cls, resp): column_names = [field.name for field in resp.schema] return agate_helper.table_from_data_flat(resp, column_names) + def get_labels_from_query_comment(cls): + if ( + hasattr(cls.profile, "query_comment") + and cls.profile.query_comment + and cls.profile.query_comment.job_label + and cls.query_header + ): + query_comment = cls.query_header.comment.query_comment + return cls._labels_from_query_comment(query_comment) + + return {} + def raw_execute( self, sql, @@ -440,15 +452,7 @@ def raw_execute( fire_event(SQLQuery(conn_name=conn.name, sql=sql, node_info=get_node_info())) - if ( - hasattr(self.profile, "query_comment") - and self.profile.query_comment - and self.profile.query_comment.job_label - ): - query_comment = self.profile.query_comment - labels = self._labels_from_query_comment(query_comment.comment) - else: - labels = {} + labels = self.get_labels_from_query_comment() if active_user: labels["dbt_invocation_id"] = active_user.invocation_id diff --git a/tests/functional/adapter/query_comment_test/test_job_label.py b/tests/functional/adapter/query_comment_test/test_job_label.py new file mode 100644 index 000000000..af984a8c4 --- /dev/null +++ b/tests/functional/adapter/query_comment_test/test_job_label.py @@ -0,0 +1,52 @@ +import pytest + +from google.cloud.bigquery.client import Client + +from dbt.tests.util import run_dbt + + +_MACRO__BQ_LABELS = """ +{% macro bq_labels() %}{ + "system": "{{ env_var('LABEL_SYSTEM', 'my_system') }}", + "env_type": "{{ env_var('LABEL_ENV', 'dev') }}" +}{% endmacro %} +""" +_MODEL__MY_TABLE = """ +{{ config(materialized= "table") }} +select 1 as id +""" + + +class TestQueryCommentJobLabel: + @pytest.fixture(scope="class") + def models(self): + return {"my_table.sql": _MODEL__MY_TABLE} + + @pytest.fixture(scope="class") + def macros(self): + return {"bq_labels.sql": _MACRO__BQ_LABELS} + + @pytest.fixture(scope="class") + def project_config_update(self): + return { + "query-comment": { + "comment": "{{ bq_labels() }}", + "job-label": True, + "append": True, + } + } + + def test_query_comments_displays_as_job_labels(self, project): + """ + Addresses this regression in dbt-bigquery 1.6: + https://github.com/dbt-labs/dbt-bigquery/issues/863 + """ + results = run_dbt(["run"]) + job_id = results.results[0].adapter_response.get("job_id") + with project.adapter.connection_named("_test"): + client: Client = project.adapter.connections.get_thread_connection().handle + job = client.get_job(job_id=job_id) + + # this is what should happen + assert job.labels.get("system") == "my_system" + assert job.labels.get("env_type") == "dev"