Merge pull request #2589 from azhard/dbt-2586-bigquery-policy-tags

Added support for setting policy tags for columns in BigQuery

CHANGELOG.md

@@ -8,6 +8,7 @@
 - Added new node selector methods (`config`, `test_type`, `test_name`, `package`) ([#2425](https://github.com/fishtown-analytics/dbt/issues/2425), [#2629](https://github.com/fishtown-analytics/dbt/pull/2629))
 - Added option to specify profile when connecting to Redshift via IAM ([#2437](https://github.com/fishtown-analytics/dbt/issues/2437), [#2581](https://github.com/fishtown-analytics/dbt/pull/2581))
 - Add more helpful error message for misconfiguration in profiles.yml ([#2569](https://github.com/fishtown-analytics/dbt/issues/2569), [#2627](https://github.com/fishtown-analytics/dbt/pull/2627))
+- Added support for setting policy tags for BigQuery columns ([#2586](https://github.com/fishtown-analytics/dbt/issues/2586), [#2589](https://github.com/fishtown-analytics/dbt/pull/2589))
 ### Fixes
 - Adapter plugins can once again override plugins defined in core ([#2548](https://github.com/fishtown-analytics/dbt/issues/2548), [#2590](https://github.com/fishtown-analytics/dbt/pull/2590))
 - Added `--selector` argument and support for `selectors.yml` file to define selection mechanisms. ([#2172](https://github.com/fishtown-analytics/dbt/issues/2172), [#2640](https://github.com/fishtown-analytics/dbt/pull/2640))
@@ -17,6 +18,7 @@ Contributors:
 - [@brunomurino](https://github.com/brunomurino) ([#2437](https://github.com/fishtown-analytics/dbt/pull/2581))
 - [@DrMcTaco](https://github.com/DrMcTaco) ([#1030](https://github.com/fishtown-analytics/dbt/issues/1030)),[#2555](https://github.com/fishtown-analytics/dbt/pull/2555/))
 - [@kning](https://github.com/kning) ([#2627](https://github.com/fishtown-analytics/dbt/pull/2627))
+- [@azhard](https://github.com/azhard) ([#2588](https://github.com/fishtown-analytics/dbt/pull/2588))
 
 
 ## dbt 0.18.0b1 (June 08, 2020)

core/dbt/contracts/graph/parsed.py

@@ -13,6 +13,7 @@
 )
 
 from hologram import JsonSchemaMixin
+from hologram.helpers import ExtensibleJsonSchemaMixin
 
 from dbt.clients.system import write_file
 from dbt.contracts.graph.unparsed import (
@@ -21,7 +22,7 @@
     HasYamlMetadata, MacroArgument, UnparsedSourceDefinition,
     UnparsedSourceTableDefinition, UnparsedColumn, TestDef
 )
-from dbt.contracts.util import Replaceable
+from dbt.contracts.util import Replaceable, AdditionalPropertiesMixin
 from dbt.logger import GLOBAL_LOGGER as logger  # noqa
 from dbt import flags
 from dbt.node_types import NodeType
@@ -44,12 +45,14 @@
 
 
 @dataclass
-class ColumnInfo(JsonSchemaMixin, Replaceable):
+class ColumnInfo(AdditionalPropertiesMixin, ExtensibleJsonSchemaMixin,
+                 Replaceable):
     name: str
     description: str = ''
     meta: Dict[str, Any] = field(default_factory=dict)
     data_type: Optional[str] = None
     tags: List[str] = field(default_factory=list)
+    _extra: Dict[str, Any] = field(default_factory=dict)
 
 
 @dataclass

core/dbt/contracts/graph/unparsed.py

@@ -71,12 +71,14 @@ class Docs(JsonSchemaMixin, Replaceable):
 
 
 @dataclass
-class HasDocs(JsonSchemaMixin, Replaceable):
+class HasDocs(AdditionalPropertiesMixin, ExtensibleJsonSchemaMixin,
+              Replaceable):
     name: str
     description: str = ''
     meta: Dict[str, Any] = field(default_factory=dict)
     data_type: Optional[str] = None
     docs: Docs = field(default_factory=Docs)
+    _extra: Dict[str, Any] = field(default_factory=dict)
 
 
 TestDef = Union[Dict[str, Any], str]

core/dbt/parser/schemas.py

@@ -100,6 +100,7 @@ def add(
             data_type=data_type,
             meta=meta,
             tags=tags,
+            _extra=column.extra
         )
 
     @classmethod

docker/requirements/requirements.0.18.0b1.txt

@@ -17,7 +17,7 @@ decorator==4.4.2
 docutils==0.15.2
 google-api-core==1.16.0
 google-auth==1.16.1
-google-cloud-bigquery==1.24.0
+google-cloud-bigquery==1.25.0
 google-cloud-core==1.3.0
 google-resumable-media==0.5.1
 googleapis-common-protos==1.6.0

plugins/bigquery/dbt/adapters/bigquery/impl.py

@@ -612,6 +612,10 @@ def _update_column_dict(self, bq_column_dict, dbt_columns, parent=''):
         if dotted_column_name in dbt_columns:
             column_config = dbt_columns[dotted_column_name]
             bq_column_dict['description'] = column_config.get('description')
+            if column_config.get('policy_tags'):
+                bq_column_dict['policyTags'] = {
+                    'names': column_config.get('policy_tags')
+                }
 
         new_fields = []
         for child_col_dict in bq_column_dict.get('fields', list()):
@@ -627,7 +631,7 @@ def _update_column_dict(self, bq_column_dict, dbt_columns, parent=''):
         return bq_column_dict
 
     @available.parse_none
-    def update_column_descriptions(self, relation, columns):
+    def update_columns(self, relation, columns):
         if len(columns) == 0:
             return
 

plugins/bigquery/dbt/include/bigquery/macros/adapters.sql

@@ -121,7 +121,7 @@
 {% endmacro %}
 
 {% macro bigquery__alter_column_comment(relation, column_dict) -%}
-  {% do adapter.update_column_descriptions(relation, column_dict) %}
+  {% do adapter.update_columns(relation, column_dict) %}
 {% endmacro %}
 
 {% macro bigquery__rename_relation(from_relation, to_relation) -%}

plugins/bigquery/setup.py

@@ -42,7 +42,7 @@
         'dbt-core=={}'.format(package_version),
         'protobuf>=3.6.0,<3.12',
         'google-cloud-core>=1.3.0,<1.4',
-        'google-cloud-bigquery>=1.24.0,<1.25.0',
+        'google-cloud-bigquery>=1.25.0,<1.26.0',
         'google-api-core>=1.16.0,<1.17.0',
         'googleapis-common-protos>=1.6.0,<1.7.0',
         'six>=1.14.0',

test.env.sample

@@ -14,6 +14,7 @@ SNOWFLAKE_TEST_OAUTH_CLIENT_SECRET=
 
 BIGQUERY_SERVICE_ACCOUNT_JSON=
 BIGQUERY_TEST_ALT_DATABASE=
+BIGQUERY_POLICY_TAG=
 
 REDSHIFT_TEST_HOST=
 REDSHIFT_TEST_USER=

test/integration/022_bigquery_test/test_bigquery_update_columns.py

@@ -0,0 +1,85 @@
+import os
+
+from test.integration.base import DBTIntegrationTest, use_profile
+
+
+class TestBigqueryUpdateColumnPolicyTag(DBTIntegrationTest):
+
+    @property
+    def schema(self):
+        return "bigquery_test_022"
+
+    @property
+    def models(self):
+        return "update-column-policy-tag"
+
+    @property
+    def project_config(self):
+        return {
+            'config-version': 2,
+            'vars': {
+                'policy_tag': self.policy_tag
+            }
+        }
+
+    @property
+    def policy_tag(self):
+        return os.environ.get('BIGQUERY_POLICY_TAG')
+
+    @use_profile('bigquery')
+    def test__bigquery_update_column_policy_tag(self):
+        if self.policy_tag:
+            results = self.run_dbt(['run', '--models', 'policy_tag_table'])
+            self.assertEqual(len(results), 1)
+
+            with self.get_connection() as conn:
+                client = conn.handle
+
+                table = client.get_table(
+                    self.adapter.connections.get_bq_table(
+                        self.default_database, self.unique_schema(), 'policy_tag_table')
+                )
+
+                for schema_field in table.schema:
+                    self.assertEquals(schema_field.policy_tags.names,
+                                      (self.policy_tag,))
+
+
+class TestBigqueryUpdateColumnDescription(DBTIntegrationTest):
+
+    @property
+    def schema(self):
+        return "bigquery_test_022"
+
+    @property
+    def models(self):
+        return "update-column-description"
+
+    @property
+    def project_config(self):
+        return {
+            'config-version': 2,
+            'vars': {
+                'field_description': self.field_description
+            }
+        }
+
+    @property
+    def field_description(self):
+        return 'this is a field'
+
+    @use_profile('bigquery')
+    def test__bigquery_update_column_description(self):
+        results = self.run_dbt(['run', '--models', 'description_table'])
+        self.assertEqual(len(results), 1)
+
+        with self.get_connection() as conn:
+            client = conn.handle
+
+            table = client.get_table(
+                self.adapter.connections.get_bq_table(
+                    self.default_database, self.unique_schema(), 'description_table')
+            )
+
+            for schema_field in table.schema:
+                self.assertEquals(schema_field.description, self.field_description)

test/integration/022_bigquery_test/update-column-description/description_table.sql

@@ -0,0 +1,9 @@
+{{
+  config(
+    materialized='table',
+    persist_docs={ 'columns': true }
+  )
+}}
+
+select
+  1 field

test/integration/022_bigquery_test/update-column-description/description_table.yml

@@ -0,0 +1,7 @@
+version: 2
+
+models:
+- name: description_table
+  columns:
+  - name: field
+    description: '{{ var("field_description") }}'

test/integration/022_bigquery_test/update-column-policy-tag/policy_tag_table.sql

@@ -0,0 +1,9 @@
+{{
+  config(
+    materialized='table',
+    persist_docs={ 'columns': true }
+  )
+}}
+
+select
+  1 field

test/integration/022_bigquery_test/update-column-policy-tag/policy_tag_table.yml

@@ -0,0 +1,8 @@
+version: 2
+
+models:
+- name: policy_tag_table
+  columns:
+  - name: field
+    policy_tags:
+      - '{{ var("policy_tag") }}'