Skip to content

Commit

Permalink
Add functionality to persist descriptions to nested columns in bigquery
Browse files Browse the repository at this point in the history
  • Loading branch information
bodschut committed Jun 16, 2020
1 parent 74c78ef commit 33b4e78
Showing 1 changed file with 17 additions and 7 deletions.
24 changes: 17 additions & 7 deletions plugins/bigquery/dbt/adapters/bigquery/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -595,6 +595,19 @@ def get_table_ref_from_relation(self, conn, relation):
relation.identifier,
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']
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_fields.append(new_child_column_dict)
bq_column_dict['fields'] = new_fields
return bq_column_dict

@available.parse_none
def update_column_descriptions(self, relation, columns):
if len(columns) == 0:
Expand All @@ -605,13 +618,10 @@ def update_column_descriptions(self, relation, columns):
table = conn.handle.get_table(table_ref)

new_schema = []
for column in table.schema:
if column.name in columns:
column_config = columns[column.name]
column_dict = column.to_api_repr()
column_dict['description'] = column_config.get('description')
column = SchemaField.from_api_repr(column_dict)
new_schema.append(column)
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_schema.append(SchemaField.from_api_repr(new_bq_column_dict))

new_table = google.cloud.bigquery.Table(table_ref, schema=new_schema)
conn.handle.update_table(new_table, ['schema'])
Expand Down

0 comments on commit 33b4e78

Please sign in to comment.