From b5edad51625418b537b552fc5121edbcdb68b5c7 Mon Sep 17 00:00:00 2001 From: Park Keunsoo Date: Mon, 27 Jan 2025 11:37:37 +0100 Subject: [PATCH 1/2] initial trial --- .../src/dbt/adapters/bigquery/impl.py | 5 +++++ .../bigquery/relation_configs/_options.py | 6 ++++++ .../adapter/simple_bigquery_view/fixtures.py | 11 ++++++++++ .../functional/adapter/test_simple_seed.py | 20 +++++++++++++++++++ 4 files changed, 42 insertions(+) diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/impl.py b/dbt-bigquery/src/dbt/adapters/bigquery/impl.py index 51c457129..a751dbd3c 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/impl.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/impl.py @@ -99,6 +99,7 @@ class BigqueryConfig(AdapterConfig): partition_by: Optional[Dict[str, Any]] = None kms_key_name: Optional[str] = None labels: Optional[Dict[str, str]] = None + tags: Optional[Dict[str, str]] = None partitions: Optional[List[str]] = None grant_access_to: Optional[List[Dict[str, str]]] = None hours_to_expiration: Optional[int] = None @@ -762,6 +763,10 @@ def get_common_options( labels = config.get("labels", {}) opts["labels"] = list(labels.items()) # type: ignore[assignment] + if config.get("tags"): + tags = config.get("tags", {}) + opts["tags"] = list(tags.items()) + return opts @available.parse(lambda *a, **k: {}) diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/relation_configs/_options.py b/dbt-bigquery/src/dbt/adapters/bigquery/relation_configs/_options.py index 7fd8797df..a1a050a71 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/relation_configs/_options.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/relation_configs/_options.py @@ -61,6 +61,7 @@ def array(x): "kms_key_name": string, "description": escaped_string, "labels": array, + "tags": array, } def formatted_option(name: str) -> Optional[Any]: @@ -88,6 +89,7 @@ def from_dict(cls, config_dict: Dict[str, Any]) -> Self: "kms_key_name": None, "description": None, "labels": None, + "tags": None, } def formatted_setting(name: str) -> Any: @@ -118,6 +120,7 @@ def parse_relation_config(cls, relation_config: RelationConfig) -> Dict[str, Any "kms_key_name", "description", "labels", + "tags", ] } @@ -145,6 +148,9 @@ def parse_bq_table(cls, table: BigQueryTable) -> Dict[str, Any]: if labels := table.labels: config_dict.update({"labels": labels}) + if tags := table.tags: + config_dict.update({"tags": tags}) + if encryption_configuration := table.encryption_configuration: config_dict.update({"kms_key_name": encryption_configuration.kms_key_name}) return config_dict diff --git a/dbt-bigquery/tests/functional/adapter/simple_bigquery_view/fixtures.py b/dbt-bigquery/tests/functional/adapter/simple_bigquery_view/fixtures.py index 9c13750f4..db7431cea 100644 --- a/dbt-bigquery/tests/functional/adapter/simple_bigquery_view/fixtures.py +++ b/dbt-bigquery/tests/functional/adapter/simple_bigquery_view/fixtures.py @@ -29,6 +29,17 @@ select * from {{ ref('view_model') }} """.lstrip() +tagged_model_sql = """ +{{ + config( + materialized = "table", + tags = {'environment': 'prod', 'component': 'batch', 'project': 'atlas'} + ) +}} + +select * from {{ ref('view_model') }} +""".lstrip() + multi_clustered_model_sql = """ {{ config( diff --git a/dbt-bigquery/tests/functional/adapter/test_simple_seed.py b/dbt-bigquery/tests/functional/adapter/test_simple_seed.py index 5ec19d420..bd5262345 100644 --- a/dbt-bigquery/tests/functional/adapter/test_simple_seed.py +++ b/dbt-bigquery/tests/functional/adapter/test_simple_seed.py @@ -66,6 +66,10 @@ labels: contains_pii: 'yes' contains_pie: 'no' + tags: + environment: 'prod' + component: 'batch' + project: 'atlas' """.lstrip() @@ -152,6 +156,22 @@ def test__bigquery_seed_table_with_labels_config_bigquery(self, project): assert bq_table.labels == self.table_labels() assert bq_table.expires + @staticmethod + def table_tags(): + return {"environment": "prod", "component": "batch", "project": "atlas"} + + def test__bigquery_seed_table_with_tags_config_bigquery(self, project): + seed_results = run_dbt(["seed"]) + assert len(seed_results) == 3 + with project.adapter.connection_named("_test"): + client = project.adapter.connections.get_thread_connection().handle + table_id = "{}.{}.{}".format(project.database, project.test_schema, "seed_configs") + bq_table = client.get_table(table_id) + + assert bq_table.tags + assert bq_table.tags == self.table_tags() + assert bq_table.expires + class TestBigQueryEmptySeed(BaseTestEmptySeed): pass From 46ca10a65e1838dbd89e4641e3d8c9d621d60f17 Mon Sep 17 00:00:00 2001 From: Park Keunsoo Date: Mon, 27 Jan 2025 13:53:30 +0100 Subject: [PATCH 2/2] fix: add tags on BigQueryOptionsConfig --- .../src/dbt/adapters/bigquery/relation_configs/_options.py | 1 + .../adapter/simple_bigquery_view/test_simple_bigquery_view.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/relation_configs/_options.py b/dbt-bigquery/src/dbt/adapters/bigquery/relation_configs/_options.py index a1a050a71..70d924609 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/relation_configs/_options.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/relation_configs/_options.py @@ -25,6 +25,7 @@ class BigQueryOptionsConfig(BigQueryBaseRelationConfig): kms_key_name: Optional[str] = None description: Optional[str] = None labels: Optional[Dict[str, str]] = None + tags: Optional[Dict[str, str]] = None def as_ddl_dict(self) -> Dict[str, Any]: """ diff --git a/dbt-bigquery/tests/functional/adapter/simple_bigquery_view/test_simple_bigquery_view.py b/dbt-bigquery/tests/functional/adapter/simple_bigquery_view/test_simple_bigquery_view.py index 6c7bc03e0..6376f7ce5 100644 --- a/dbt-bigquery/tests/functional/adapter/simple_bigquery_view/test_simple_bigquery_view.py +++ b/dbt-bigquery/tests/functional/adapter/simple_bigquery_view/test_simple_bigquery_view.py @@ -15,6 +15,7 @@ clustered_model_sql, funky_case_sql, labeled_model_sql, + tagged_model_sql, multi_clustered_model_sql, partitioned_model_sql, sql_header_model_sql, @@ -51,6 +52,7 @@ def models(self): "clustered_model.sql": clustered_model_sql, "fUnKyCaSe.sql": funky_case_sql, "labeled_model.sql": labeled_model_sql, + "tagged_model.sql": tagged_model_sql, "multi_clustered_model.sql": multi_clustered_model_sql, "partitioned_model.sql": partitioned_model_sql, "sql_header_model.sql": sql_header_model_sql,