From fb901447a40e44a3e962a5a654cb20e86c966055 Mon Sep 17 00:00:00 2001 From: Bob De Schutter Date: Thu, 18 Jun 2020 17:31:31 +0200 Subject: [PATCH] Added integration test and cleaned up code --- .../bigquery/dbt/adapters/bigquery/impl.py | 28 ++++++-- .../models-bigquery-nested/schema.yml | 19 ++++++ .../table_model_nested.sql | 8 +++ .../view_model_nested.sql | 8 +++ .../test_persist_docs.py | 67 +++++++++++++++++++ 5 files changed, 126 insertions(+), 4 deletions(-) create mode 100644 test/integration/060_persist_docs_tests/models-bigquery-nested/schema.yml create mode 100644 test/integration/060_persist_docs_tests/models-bigquery-nested/table_model_nested.sql create mode 100644 test/integration/060_persist_docs_tests/models-bigquery-nested/view_model_nested.sql diff --git a/plugins/bigquery/dbt/adapters/bigquery/impl.py b/plugins/bigquery/dbt/adapters/bigquery/impl.py index adb4b894957..225f30576da 100644 --- a/plugins/bigquery/dbt/adapters/bigquery/impl.py +++ b/plugins/bigquery/dbt/adapters/bigquery/impl.py @@ -596,16 +596,33 @@ def get_table_ref_from_relation(self, conn, relation): conn) def _update_column_dict(self, bq_column_dict, dbt_columns, parent=''): - # bq_colunmn_dict generated by SchemaField.to_api_repr() - dotted_column_name = '{}.{}'.format(parent, bq_column_dict['name']) if parent else bq_column_dict['name'] + """ + Helper function to recursively traverse the schema of a table in the + update_column_descriptions function below. + + bq_column_dict should be a dict as obtained by the to_api_repr() + function of a SchemaField object. + """ + if parent: + dotted_column_name = '{}.{}'.format(parent, bq_column_dict['name']) + else: + dotted_column_name = bq_column_dict['name'] + if dotted_column_name in dbt_columns: column_config = dbt_columns[dotted_column_name] bq_column_dict['description'] = column_config.get('description') + new_fields = [] for child_col_dict in bq_column_dict.get('fields', list()): - new_child_column_dict = self._update_column_dict(child_col_dict, dbt_columns, parent=dotted_column_name) + new_child_column_dict = self._update_column_dict( + child_col_dict, + dbt_columns, + parent=dotted_column_name + ) new_fields.append(new_child_column_dict) + bq_column_dict['fields'] = new_fields + return bq_column_dict @available.parse_none @@ -620,7 +637,10 @@ def update_column_descriptions(self, relation, columns): new_schema = [] for bq_column in table.schema: bq_column_dict = bq_column.to_api_repr() - new_bq_column_dict = self._update_column_dict(bq_column_dict, columns) + new_bq_column_dict = self._update_column_dict( + bq_column_dict, + columns + ) new_schema.append(SchemaField.from_api_repr(new_bq_column_dict)) new_table = google.cloud.bigquery.Table(table_ref, schema=new_schema) diff --git a/test/integration/060_persist_docs_tests/models-bigquery-nested/schema.yml b/test/integration/060_persist_docs_tests/models-bigquery-nested/schema.yml new file mode 100644 index 00000000000..0311dcb1449 --- /dev/null +++ b/test/integration/060_persist_docs_tests/models-bigquery-nested/schema.yml @@ -0,0 +1,19 @@ +version: 2 + +models: + - name: table_model_nested + columns: + - name: level_1 + description: level_1 column description + - name: level_1.level_2 + description: level_2 column description + - name: level_1.level_2.level_3_a + description: level_3 column description + - name: view_model_nested + columns: + - name: level_1 + description: level_1 column description + - name: level_1.level_2 + description: level_2 column description + - name: level_1.level_2.level_3_a + description: level_3 column description \ No newline at end of file diff --git a/test/integration/060_persist_docs_tests/models-bigquery-nested/table_model_nested.sql b/test/integration/060_persist_docs_tests/models-bigquery-nested/table_model_nested.sql new file mode 100644 index 00000000000..c2936d4f186 --- /dev/null +++ b/test/integration/060_persist_docs_tests/models-bigquery-nested/table_model_nested.sql @@ -0,0 +1,8 @@ +{{ config(materialized='table') }} +SELECT + STRUCT( + STRUCT( + 1 AS level_3_a, + 2 AS level_3_b + ) AS level_2 + ) AS level_1 \ No newline at end of file diff --git a/test/integration/060_persist_docs_tests/models-bigquery-nested/view_model_nested.sql b/test/integration/060_persist_docs_tests/models-bigquery-nested/view_model_nested.sql new file mode 100644 index 00000000000..e3323ddf4e6 --- /dev/null +++ b/test/integration/060_persist_docs_tests/models-bigquery-nested/view_model_nested.sql @@ -0,0 +1,8 @@ +{{ config(materialized='view') }} +SELECT + STRUCT( + STRUCT( + 1 AS level_3_a, + 2 AS level_3_b + ) AS level_2 + ) AS level_1 \ No newline at end of file diff --git a/test/integration/060_persist_docs_tests/test_persist_docs.py b/test/integration/060_persist_docs_tests/test_persist_docs.py index c5a6cc1689e..6b7dc783778 100644 --- a/test/integration/060_persist_docs_tests/test_persist_docs.py +++ b/test/integration/060_persist_docs_tests/test_persist_docs.py @@ -156,3 +156,70 @@ def test_snowflake_persist_docs(self): @use_profile('bigquery') def test_bigquery_persist_docs(self): self.run_dbt() + + +class TestPersistDocsNested(BasePersistDocsTest): + @property + def project_config(self): + return { + 'config-version': 2, + 'models': { + 'test': { + '+persist_docs': { + "relation": True, + "columns": True, + }, + } + } + } + + @property + def models(self): + return 'models-bigquery-nested' + + @use_profile('bigquery') + def test_bigquery_persist_docs(self): + """ + run dbt and use the bigquery client from the adapter to check if the + colunmn descriptions are persisted on the test model table and view. + """ + import sys + self.run_dbt() + + with self.adapter.connection_named('_test'): + client = self.adapter.connections.get_thread_connection().handle + + # test for table model + table_id = "{}.{}.{}".format( + self.default_database, + self.unique_schema(), + "table_model_nested" + ) + bq_table_schema = client.get_table(table_id).schema + + level_1_field = bq_table_schema[0] + assert level_1_field.description == "level_1 column description" + + level_2_field = level_1_field.fields[0] + assert level_2_field.description == "level_2 column description" + + level_3_field = level_2_field.fields[0] + assert level_3_field.description == "level_3 column description" + + # test for view model + view_id = "{}.{}.{}".format( + self.default_database, + self.unique_schema(), + "view_model_nested" + ) + bq_view_schema = client.get_table(view_id).schema + + level_1_field = bq_view_schema[0] + assert level_1_field.description == "level_1 column description" + + level_2_field = level_1_field.fields[0] + assert level_2_field.description == "level_2 column description" + + level_3_field = level_2_field.fields[0] + assert level_3_field.description == "level_3 column description" +