Repository: dbt-labs/dbt-core

Issue:
Support BigQuery Job Tags and Labels
### Describe the feature

I would like to be able to control [tagging and labeling of BigQuery Jobs](https://cloud.google.com/bigquery/docs/adding-labels#job-label) as I run dbt on BigQuery.

A similar (but not the same) issue is #1947, for labeling BigQuery Tables and Datasets. This issue focuses on BigQuery Jobs (such as Insert Jobs or Query Jobs).

### Describe alternatives you've considered

It's not possible to label or tag jobs after they have started. From [the docs](https://cloud.google.com/bigquery/docs/adding-labels#job-label)

> You cannot add labels to or update labels on pending, running, or completed jobs.

### Additional context

The main reason why one would tag and label their BigQuery Job is to analyze BigQuery spend. For example, if one were able to link a BigQuery Job to a certain Airflow operator run (or similar -- in my case a python script run by a cron! :-D) then a real dollar value can be put on running that operator over time.

I think it's important to give the developer control on what tags and labels can be added, so it supports their data processing setup. And so I think tags and labels should be able to be set at launch-time. (In my case, I run a python script that calls `dbt run` -- I would want my python script to be able to set the BigQuery Job tags and labels, while the Jobs are ultimately launched by `dbt run`.)

### Who will this benefit?

Folks who are responsible for their BigQuery spend should benefit by using relevant Job tags and labels.

Thanks!


Discussion:
This is interesting! Thanks for the detailed writeup @boxysean.

More so than labels on tables and columns, this idea feels like a close relative of [query comments](https://docs.getdbt.com/reference/project-configs/query-comment/), which we introduced in v0.15. Parsing out and aggregating query comments is the best way to calculate usage on Postgres/Redshift.

### Clarifying questions

A single `dbt run` may execute dozens or hundreds query jobs. Is your thought that all query jobs executed by the same dbt command should get the same label(s)? Should each query job label include information about the specific node (model/seed/snapshot/test/etc) it's running?

You mention that you'd want to have control of the label when invoking dbt from, say, Airflow or a cron script. Is this something you expect as a CLI arg, e.g.
```
dbt run --job-label '{"name": "nightly_run"}'
```
Or could it be version-controlled code that, depending on situational variables (`target`, `var`, `env_var`) outputs a different label?

### My proposal

Here's what makes sense to me:
* We could add a `query-comment` attribute, `job-label: true`. On BigQuery, in addition to prepending the query comment as a SQL comment (`/* */`), dbt would add the comment's key-value components as labels to each query jobs it runs.
* Using the existing `query-comment` functionality, you could write a macro that returns a different comment (or comment dict attributes) depending on the values of `target`/`var`/`env_var`
The proposal is reasonable and I experimented with it. However, I found there is a maximum length of a label value in GCP of 128 bytes, and the default comment exceeded it. I'm not sure it makes sense to use the same query-comment functionality.

Alternative:
- On every DBT BigQuery job it passes in the label dbt_invocation_id matching with the invocation_id in DBT.
- For cases where more detailed profiling is needed, a DBT step can be added to insert into a tracking table an association between the invocation_id and whatever information is needed. This can include whatever macros you need/want (eg coming from Airflow perhaps, or Argo).
- The INFORMATION_SCHEMA job information can be used to extract performance stats and left join with the tracking table for more detailed information/analysis.

The code changes in DBT are tiny -- I'm putting together a pullrequest for it.

@mescanne Thanks for checking out this issue, and for opening up the PR. I'll take a look there in a moment.

I had envisioned splitting out the arguments currently passed in the query comment to each be their own label, so one for each of `app`, `dbt_version`, `profile_name`, `target_name`, and `node_id`. That would enable us to avoid the 128 byte limit. A few of those values would require more plumbing than what's currently available to `raw_execute`, however, so the work required is a bit more involved.

I quite like your suggestion to query `INFORMATION_SCHEMA.JOBS_BY_*` and filter / join on the `invocation_id` label as a way to get metadata about a given run, e.g. total spend.
Hi @jtcohen6 coming back to this. I think it's helpful to add as much information as possible to the query jobs, so your `query-comment` suggestion makes sense.

I would still prefer to specify the BigQuery Job label at dbt run-time, such as you suggested with `dbt run --job-label '{"name": "nightly_run"}'`. That way it would be controlled by the executor (e.g., Airflow, cron) and I could customize how I slice & dice the executions, which indeed `INFORMATION_SCHEMA.JOBS_BY_* ` could be used for!
I would also like to upvote this feature request.

In our use case, we would like to pass other environment variables to the bq-job label for ease of monitoring and benchmarking. For example, we would like to pass the environment (prod, stage), pipeline release version, and our orchestration workflow run id (prefect run id) to BQ job label. Currently we set these info in environment variable, and it's quite easy to add them into the query-comment dict.

I think @jtcohen6 's proposal below would make it work for us if implemented. 
> I had envisioned splitting out the arguments currently passed in the query comment to each be their own label, so one for each of app, dbt_version, profile_name, target_name, and node_id. That would enable us to avoid the 128 byte limit. A few of those values would require more plumbing than what's currently available to raw_execute, however, so the work required is a bit more involved.


We are in the exact same situation than those described by @boxysean and @hui-zheng: we'd like to stamp our dbt's big query jobs with labels/tags so that we can get better visibility on how we spend our BQ budget/quota. 

Although per `dbt run` custom labels/tags would be enough in our case (eg. customer, environment, etc), having additional (generic?) context per query jobs would be great too (eg. dbt model/node the query job is running for as suggested by @jtcohen6, or model/node's own [dbt tags](https://docs.getdbt.com/reference/resource-configs/tags/)).

This looks to be supported and achievable through BigQuery's [adding job label](https://cloud.google.com/bigquery/docs/adding-labels#job-label) feature. It is also documented in the [JobConfiguration](https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#jobconfiguration) API doc.

> Currently we set these info in environment variable, and it's quite easy to add them into the query-comment dict.

@hui-zheng, could you describe how you do it?
My team was looking for exactly this feature. Are you planning to work on it, or would you review a PR that adds a `--job-label` flag to `dbt run`?
This isn't something we're prioritizing now. FYI #2809 did add `invocation_id` as a label to all dbt-bigquery jobs, starting in v0.19.0. That `invocation_id` can be used to query `INFORMATION_SCHEMA.JOBS_BY_*` and calculate total time/spend per invocation; it can also be used to associate BigQuery query history with dbt run artifacts ([docs](https://docs.getdbt.com/reference/artifacts/dbt-artifacts#common-metadata)), namely `run_results.json`. Those run artifacts will contain lots of useful metadata, including (e.g.) any environment variables prefixed with `DBT_ENV_CUSTOM_ENV_`.

I agree that dbt should be able to pass more information than just the `invocation_id`, though I think that's a strong start. A `--job-label` flag that allows the user / orchestration tool to set one value for all nodes in an invocation should be straightforward to implement. The more I think about it, though, I find it functionally limiting but also one-off as an implementation, not well integrated with existing dbt constructs that seek to accomplish the same goal.

I do think the best version of this would make available the full query comment context as per-node job labels. That context, available to the `query_comment` macro, is defined in [`query_headers.py`](https://github.com/fishtown-analytics/dbt/blob/344a14416d22f0cfbeb56b9904092c8a4f38b1fc/core/dbt/adapters/base/query_headers.py). I agree with @hui-zheng, it's quite easy to pass environment variables or `--vars` into the `query-comment` config or `query_comment` macro, so this approach would solve for both use cases we've been discussing.

The string version of this comment—the default value, the string passed to the config, or the value returned by the custom macro—is available to the connection manager, via [`set_query_header`](https://github.com/fishtown-analytics/dbt/blob/344a14416d22f0cfbeb56b9904092c8a4f38b1fc/core/dbt/adapters/base/connections.py#L46-L47) and [`_add_query_comment`](https://github.com/fishtown-analytics/dbt/blob/344a14416d22f0cfbeb56b9904092c8a4f38b1fc/core/dbt/adapters/base/connections.py#L286-L289). The `execute` method already calls `_add_query_comment` to prepend the comment to SQL before execution:

https://github.com/fishtown-analytics/dbt/blob/344a14416d22f0cfbeb56b9904092c8a4f38b1fc/plugins/bigquery/dbt/adapters/bigquery/connections.py#L333-L336

So here's what I'm thinking about:
- How would we enable dbt-bigquery users to turn this on? I'm thinking an additional option, nested under `query_comment` in `dbt_project.yml`, called `job_label: true | false`.
- If `query_comment.job_label` is turned on, and the query comment config/macro returns a dict / JSON string (such as the [advanced usage example in the docs](https://docs.getdbt.com/reference/project-configs/query-comment#advanced-use-a-macro-to-generate-a-comment)), should dbt try to parse the returned value into a python dict, and pass each key-value pair as a separate label? I think yes; this should even work for the default query comment value.
- If `query_comment.job_label` is turned on, and the query-comment returns an unstructured string, should dbt still try to pass the first 128 bytes (truncated if needed) as the value to a single label called `query_comment`? I still think yes, but I'm open to your thoughts on this point (and every point above).

Having written all that out, acknowledging that there are a few tricky pieces, I do think the requisite changes would be relatively self-contained in the codebase. Would anyone be interested in giving it a go?
I'm not very familiar with the dbt internals, so it would probably take me some time to figure out, but I'd be happy to give this a try if nobody picks it up first.
Happy to help along the way @jmcarp! I think the methods I linked above would be the right places to start.

In particular, `raw_execute` already passes job labels here:

https://github.com/fishtown-analytics/dbt/blob/344a14416d22f0cfbeb56b9904092c8a4f38b1fc/plugins/bigquery/dbt/adapters/bigquery/connections.py#L302-L313

And [`_add_query_comment`](https://github.com/fishtown-analytics/dbt/blob/344a14416d22f0cfbeb56b9904092c8a4f38b1fc/core/dbt/adapters/base/connections.py#L286-L289) is already a method available to the Connection Manager. So this may be as simple as adding within `raw_execute`:
```python
set query_comment = self._add_query_comment('')
```
Then, try to parse `query_comment` is a dict. If it's a dict, pass each key/value pair as a separate item to `job_params['labels']`. If it's a non-empty string, pass the entire `query_comment` string (truncated to 128 characters) as one item. If it's empty, do nothing.

There would be some additional work to make this an optional config, and to turn it on/off based on `query_comment.job_params`, but that feels eminently doable. Probably it looks like creating a special dbt-bigquery subclass of `MacroQueryStringSetter` and having the `BigQueryConnectionManager` override [`set_query_header`](https://github.com/fishtown-analytics/dbt/blob/344a14416d22f0cfbeb56b9904092c8a4f38b1fc/core/dbt/adapters/base/connections.py#L46-L47), unless we can think of comparable functionality on other adapters. (Snowflake has [query tags](https://docs.getdbt.com/reference/resource-configs/snowflake-configs#query-tags), but those are accessible in SQL, so there's no need to wedge them into the connection.)
Thanks for the detailed explanations @jtcohen6, I can give it a try too 🖐️  

Resolving patch:
diff --git a/core/dbt/contracts/connection.py b/core/dbt/contracts/connection.py
--- a/core/dbt/contracts/connection.py
+++ b/core/dbt/contracts/connection.py
@@ -9,7 +9,7 @@
 from dbt.logger import GLOBAL_LOGGER as logger
 from typing_extensions import Protocol
 from dbt.dataclass_schema import (
-    dbtClassMixin, StrEnum, ExtensibleDbtClassMixin,
+    dbtClassMixin, StrEnum, ExtensibleDbtClassMixin, HyphenatedDbtClassMixin,
     ValidatedStringMixin, register_pattern
 )
 from dbt.contracts.util import Replaceable
@@ -212,9 +212,10 @@ def to_target_dict(self):
 
 
 @dataclass
-class QueryComment(dbtClassMixin):
+class QueryComment(HyphenatedDbtClassMixin):
     comment: str = DEFAULT_QUERY_COMMENT
     append: bool = False
+    job_label: bool = False
 
 
 class AdapterRequiredConfig(HasCredentials, Protocol):
diff --git a/core/dbt/dataclass_schema.py b/core/dbt/dataclass_schema.py
--- a/core/dbt/dataclass_schema.py
+++ b/core/dbt/dataclass_schema.py
@@ -73,7 +73,9 @@ def __post_serialize__(self, dct):
     # performing the conversion to a dict
     @classmethod
     def __pre_deserialize__(cls, data):
-        if cls._hyphenated:
+        # `data` might not be a dict, e.g. for `query_comment`, which accepts
+        # a dict or a string; only snake-case for dict values.
+        if cls._hyphenated and isinstance(data, dict):
             new_dict = {}
             for key in data:
                 if '-' in key:
diff --git a/plugins/bigquery/dbt/adapters/bigquery/connections.py b/plugins/bigquery/dbt/adapters/bigquery/connections.py
--- a/plugins/bigquery/dbt/adapters/bigquery/connections.py
+++ b/plugins/bigquery/dbt/adapters/bigquery/connections.py
@@ -1,3 +1,5 @@
+import json
+import re
 from contextlib import contextmanager
 from dataclasses import dataclass
 from functools import lru_cache
@@ -305,12 +307,16 @@ def raw_execute(self, sql, fetch=False, *, use_legacy_sql=False):
 
         logger.debug('On {}: {}', conn.name, sql)
 
-        job_params = {'use_legacy_sql': use_legacy_sql}
+        if self.profile.query_comment.job_label:
+            query_comment = self.query_header.comment.query_comment
+            labels = self._labels_from_query_comment(query_comment)
+        else:
+            labels = {}
 
         if active_user:
-            job_params['labels'] = {
-                'dbt_invocation_id': active_user.invocation_id
-            }
+            labels['dbt_invocation_id'] = active_user.invocation_id
+
+        job_params = {'use_legacy_sql': use_legacy_sql, 'labels': labels}
 
         priority = conn.credentials.priority
         if priority == Priority.Batch:
@@ -544,6 +550,16 @@ def _retry_generator(self):
             initial=self.DEFAULT_INITIAL_DELAY,
             maximum=self.DEFAULT_MAXIMUM_DELAY)
 
+    def _labels_from_query_comment(self, comment: str) -> Dict:
+        try:
+            comment_labels = json.loads(comment)
+        except (TypeError, ValueError):
+            return {'query_comment': _sanitize_label(comment)}
+        return {
+            _sanitize_label(key): _sanitize_label(str(value))
+            for key, value in comment_labels.items()
+        }
+
 
 class _ErrorCounter(object):
     """Counts errors seen up to a threshold then raises the next error."""
@@ -573,3 +589,13 @@ def _is_retryable(error):
             e['reason'] == 'rateLimitExceeded' for e in error.errors):
         return True
     return False
+
+
+_SANITIZE_LABEL_PATTERN = re.compile(r"[^a-z0-9_-]")
+
+
+def _sanitize_label(value: str) -> str:
+    """Return a legal value for a BigQuery label."""
+    value = value.strip().lower()
+    value = _SANITIZE_LABEL_PATTERN.sub("_", value)
+    return value


Test patch:
diff --git a/test/unit/test_bigquery_adapter.py b/test/unit/test_bigquery_adapter.py
--- a/test/unit/test_bigquery_adapter.py
+++ b/test/unit/test_bigquery_adapter.py
@@ -1,6 +1,8 @@
 import agate
 import decimal
+import json
 import re
+import pytest
 import unittest
 from contextlib import contextmanager
 from requests.exceptions import ConnectionError
@@ -15,6 +17,7 @@
 from dbt.adapters.bigquery import BigQueryRelation
 from dbt.adapters.bigquery import Plugin as BigQueryPlugin
 from dbt.adapters.bigquery.connections import BigQueryConnectionManager
+from dbt.adapters.bigquery.connections import _sanitize_label
 from dbt.adapters.base.query_headers import MacroQueryStringSetter
 from dbt.clients import agate_helper
 import dbt.exceptions
@@ -588,7 +591,6 @@ def test_query_and_results(self, mock_bq):
         self.mock_client.query.assert_called_once_with(
           'sql', job_config=mock_bq.QueryJobConfig())
 
-
     def test_copy_bq_table_appends(self):
         self._copy_table(
             write_disposition=dbt.adapters.bigquery.impl.WRITE_APPEND)
@@ -615,12 +617,20 @@ def test_copy_bq_table_truncates(self):
             kwargs['job_config'].write_disposition,
             dbt.adapters.bigquery.impl.WRITE_TRUNCATE)
 
+    def test_job_labels_valid_json(self):
+        expected = {"key": "value"}
+        labels = self.connections._labels_from_query_comment(json.dumps(expected))
+        self.assertEqual(labels, expected)
+
+    def test_job_labels_invalid_json(self):
+        labels = self.connections._labels_from_query_comment("not json")
+        self.assertEqual(labels, {"query_comment": "not_json"})
+
     def _table_ref(self, proj, ds, table, conn):
         return google.cloud.bigquery.table.TableReference.from_string(
             '{}.{}.{}'.format(proj, ds, table))
 
     def _copy_table(self, write_disposition):
-
         self.connections.table_ref = self._table_ref
         source = BigQueryRelation.create(
             database='project', schema='dataset', identifier='table1')
@@ -931,3 +941,15 @@ def test_convert_time_type(self):
         expected = ['time', 'time', 'time']
         for col_idx, expect in enumerate(expected):
             assert BigQueryAdapter.convert_time_type(agate_table, col_idx) == expect
+
+
+@pytest.mark.parametrize(
+    ["input", "output"],
+    [
+        ("ABC", "abc"),
+        ("a c", "a_c"),
+        ("a ", "a"),
+    ],
+)
+def test_sanitize_label(input, output):
+    assert _sanitize_label(input) == output