From d5b4114ca1ef28d3c2bd3c61812e3a2c901a2366 Mon Sep 17 00:00:00 2001
From: kodaho <23268819+kodaho@users.noreply.github.com>
Date: Fri, 6 Oct 2023 20:26:53 +0200
Subject: [PATCH 1/2] [fix] Use rendered query comment for job labels (#955)

* [fix] Use rendered query comment for job labels

* Add test from dbt-labs/dbt-query#872

* Valid JSON in test

* Add changie entry
---
 .../unreleased/Fixes-20231005-235950.yaml     |  6 +++
 dbt/adapters/bigquery/connections.py          |  5 +-
 .../query_comment_test/test_job_label.py      | 52 +++++++++++++++++++
 3 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 .changes/unreleased/Fixes-20231005-235950.yaml
 create mode 100644 tests/functional/adapter/query_comment_test/test_job_label.py

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 a5c7b9355..1e96ed5ef 100644
--- a/dbt/adapters/bigquery/connections.py
+++ b/dbt/adapters/bigquery/connections.py
@@ -451,9 +451,10 @@ def raw_execute(
             hasattr(self.profile, "query_comment")
             and self.profile.query_comment
             and self.profile.query_comment.job_label
+            and self.query_header
+            and (query_comment := self.query_header.comment.query_comment)
         ):
-            query_comment = self.profile.query_comment
-            labels = self._labels_from_query_comment(query_comment.comment)
+            labels = self._labels_from_query_comment(query_comment)
         else:
             labels = {}
 
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"

From 3e8d389e65a67bad0ae467761c4a4519507e367d Mon Sep 17 00:00:00 2001
From: Matthew McKnight <91097623+McKnight-42@users.noreply.github.com>
Date: Fri, 6 Oct 2023 16:08:24 -0500
Subject: [PATCH 2/2] a slight rework of 955 to keep a seperation of concerns
 (#957)

* a slight rework of 955 to keep a seperation of concerns

* move query_header back into inital if, save assignement for truthy, else return empty dict
---
 dbt/adapters/bigquery/connections.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/dbt/adapters/bigquery/connections.py b/dbt/adapters/bigquery/connections.py
index 1e96ed5ef..ff544f0d0 100644
--- a/dbt/adapters/bigquery/connections.py
+++ b/dbt/adapters/bigquery/connections.py
@@ -435,6 +435,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,
@@ -447,16 +459,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
-            and self.query_header
-            and (query_comment := self.query_header.comment.query_comment)
-        ):
-            labels = self._labels_from_query_comment(query_comment)
-        else:
-            labels = {}
+        labels = self.get_labels_from_query_comment()
 
         if active_user:
             labels["dbt_invocation_id"] = active_user.invocation_id