Skip to content

Commit cf6324e

Browse files
thejensJens Larsson
thejens
and
Jens Larsson
authored
Implement BigQuery Table Schema Update Operator (#15367)
Co-authored-by: Jens Larsson <jens.larsson@c02cv73mml85.lan>
1 parent 59be278 commit cf6324e

6 files changed

Lines changed: 441 additions & 0 deletions

File tree

‎airflow/providers/google/cloud/example_dags/example_bigquery_operations.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
BigQueryPatchDatasetOperator,
3737
BigQueryUpdateDatasetOperator,
3838
BigQueryUpdateTableOperator,
39+
BigQueryUpdateTableSchemaOperator,
3940
BigQueryUpsertTableOperator,
4041
)
4142
from airflow.utils.dates import days_ago
@@ -73,6 +74,18 @@
7374
)
7475
# [END howto_operator_bigquery_create_table]
7576

77+
# [START howto_operator_bigquery_update_table_schema]
78+
update_table_schema = BigQueryUpdateTableSchemaOperator(
79+
task_id="update_table_schema",
80+
dataset_id=DATASET_NAME,
81+
table_id="test_table",
82+
schema_fields_updates=[
83+
{"name": "emp_name", "description": "Name of employee"},
84+
{"name": "salary", "description": "Monthly salary in USD"},
85+
],
86+
)
87+
# [END howto_operator_bigquery_update_table_schema]
88+
7689
# [START howto_operator_bigquery_delete_table]
7790
delete_table = BigQueryDeleteTableOperator(
7891
task_id="delete_table",
@@ -216,6 +229,7 @@
216229
delete_view,
217230
]
218231
>> upsert_table
232+
>> update_table_schema
219233
>> delete_materialized_view
220234
>> delete_table
221235
>> delete_dataset

‎airflow/providers/google/cloud/hooks/bigquery.py‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1373,6 +1373,101 @@ def get_schema(self, dataset_id: str, table_id: str, project_id: Optional[str] =
13731373
table = self.get_client(project_id=project_id).get_table(table_ref)
13741374
return {"fields": [s.to_api_repr() for s in table.schema]}
13751375

1376+
@GoogleBaseHook.fallback_to_default_project_id
1377+
def update_table_schema(
1378+
self,
1379+
schema_fields_updates: List[Dict[str, Any]],
1380+
include_policy_tags: bool,
1381+
dataset_id: str,
1382+
table_id: str,
1383+
project_id: Optional[str] = None,
1384+
) -> None:
1385+
"""
1386+
Update fields within a schema for a given dataset and table. Note that
1387+
some fields in schemas are immutable and trying to change them will cause
1388+
an exception.
1389+
If a new field is included it will be inserted which requires all required fields to be set.
1390+
See https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#TableSchema
1391+
1392+
:param include_policy_tags: If set to True policy tags will be included in
1393+
the update request which requires special permissions even if unchanged
1394+
see https://cloud.google.com/bigquery/docs/column-level-security#roles
1395+
:type include_policy_tags: bool
1396+
:param dataset_id: the dataset ID of the requested table to be updated
1397+
:type dataset_id: str
1398+
:param table_id: the table ID of the table to be updated
1399+
:type table_id: str
1400+
:param schema_fields_updates: a partial schema resource. see
1401+
https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#TableSchema
1402+
1403+
**Example**: ::
1404+
1405+
schema_fields_updates=[
1406+
{"name": "emp_name", "description": "Some New Description"},
1407+
{"name": "salary", "description": "Some New Description"},
1408+
{"name": "departments", "fields": [
1409+
{"name": "name", "description": "Some New Description"},
1410+
{"name": "type", "description": "Some New Description"}
1411+
]},
1412+
]
1413+
1414+
:type schema_fields_updates: List[dict]
1415+
:param project_id: The name of the project where we want to update the table.
1416+
:type project_id: str
1417+
"""
1418+
1419+
def _build_new_schema(
1420+
current_schema: List[Dict[str, Any]], schema_fields_updates: List[Dict[str, Any]]
1421+
) -> List[Dict[str, Any]]:
1422+
1423+
# Turn schema_field_updates into a dict keyed on field names
1424+
schema_fields_updates = {field["name"]: field for field in deepcopy(schema_fields_updates)}
1425+
1426+
# Create a new dict for storing the new schema, initated based on the current_schema
1427+
# as of Python 3.6, dicts retain order.
1428+
new_schema = {field["name"]: field for field in deepcopy(current_schema)}
1429+
1430+
# Each item in schema_fields_updates contains a potential patch
1431+
# to a schema field, iterate over them
1432+
for field_name, patched_value in schema_fields_updates.items():
1433+
# If this field already exists, update it
1434+
if field_name in new_schema:
1435+
# If this field is of type RECORD and has a fields key we need to patch it recursively
1436+
if "fields" in patched_value:
1437+
patched_value["fields"] = _build_new_schema(
1438+
new_schema[field_name]["fields"], patched_value["fields"]
1439+
)
1440+
# Update the new_schema with the patched value
1441+
new_schema[field_name].update(patched_value)
1442+
# This is a new field, just include the whole configuration for it
1443+
else:
1444+
new_schema[field_name] = patched_value
1445+
1446+
return list(new_schema.values())
1447+
1448+
def _remove_policy_tags(schema: List[Dict[str, Any]]):
1449+
for field in schema:
1450+
if "policyTags" in field:
1451+
del field["policyTags"]
1452+
if "fields" in field:
1453+
_remove_policy_tags(field["fields"])
1454+
1455+
current_table_schema = self.get_schema(
1456+
dataset_id=dataset_id, table_id=table_id, project_id=project_id
1457+
)["fields"]
1458+
new_schema = _build_new_schema(current_table_schema, schema_fields_updates)
1459+
1460+
if not include_policy_tags:
1461+
_remove_policy_tags(new_schema)
1462+
1463+
self.update_table(
1464+
table_resource={"schema": {"fields": new_schema}},
1465+
fields=["schema"],
1466+
project_id=project_id,
1467+
dataset_id=dataset_id,
1468+
table_id=table_id,
1469+
)
1470+
13761471
@GoogleBaseHook.fallback_to_default_project_id
13771472
def poll_job_complete(
13781473
self,

‎airflow/providers/google/cloud/operators/bigquery.py‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2039,6 +2039,118 @@ def execute(self, context) -> None:
20392039
)
20402040

20412041

2042+
class BigQueryUpdateTableSchemaOperator(BaseOperator):
2043+
"""
2044+
Update BigQuery Table Schema
2045+
Updates fields on a table schema based on contents of the supplied schema_fields_updates
2046+
parameter. The supplied schema does not need to be complete, if the field
2047+
already exists in the schema you only need to supply keys & values for the
2048+
items you want to patch, just ensure the "name" key is set.
2049+
2050+
.. seealso::
2051+
For more information on how to use this operator, take a look at the guide:
2052+
:ref:`howto/operator:BigQueryUpdateTableSchemaOperator`
2053+
2054+
:param schema_fields_updates: a partial schema resource. see
2055+
https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#TableSchema
2056+
2057+
**Example**: ::
2058+
2059+
schema_fields_updates=[
2060+
{"name": "emp_name", "description": "Some New Description"},
2061+
{"name": "salary", "policyTags": {'names': ['some_new_policy_tag']},},
2062+
{"name": "departments", "fields": [
2063+
{"name": "name", "description": "Some New Description"},
2064+
{"name": "type", "description": "Some New Description"}
2065+
]},
2066+
]
2067+
2068+
:type schema_fields_updates: List[dict]
2069+
:param include_policy_tags: (Optional) If set to True policy tags will be included in
2070+
the update request which requires special permissions even if unchanged (default False)
2071+
see https://cloud.google.com/bigquery/docs/column-level-security#roles
2072+
:type include_policy_tags: bool
2073+
:param dataset_id: A dotted
2074+
``(<project>.|<project>:)<dataset>`` that indicates which dataset
2075+
will be updated. (templated)
2076+
:type dataset_id: str
2077+
:param table_id: The table ID of the requested table. (templated)
2078+
:type table_id: str
2079+
:param project_id: The name of the project where we want to update the dataset.
2080+
Don't need to provide, if projectId in dataset_reference.
2081+
:type project_id: str
2082+
:param gcp_conn_id: (Optional) The connection ID used to connect to Google Cloud.
2083+
:type gcp_conn_id: str
2084+
:param bigquery_conn_id: (Deprecated) The connection ID used to connect to Google Cloud.
2085+
This parameter has been deprecated. You should pass the gcp_conn_id parameter instead.
2086+
:type bigquery_conn_id: str
2087+
:param delegate_to: The account to impersonate, if any.
2088+
For this to work, the service account making the request must have domain-wide
2089+
delegation enabled.
2090+
:type delegate_to: str
2091+
:param location: The location used for the operation.
2092+
:type location: str
2093+
:param impersonation_chain: Optional service account to impersonate using short-term
2094+
credentials, or chained list of accounts required to get the access_token
2095+
of the last account in the list, which will be impersonated in the request.
2096+
If set as a string, the account must grant the originating account
2097+
the Service Account Token Creator IAM role.
2098+
If set as a sequence, the identities from the list must grant
2099+
Service Account Token Creator IAM role to the directly preceding identity, with first
2100+
account from the list granting this role to the originating account (templated).
2101+
:type impersonation_chain: Union[str, Sequence[str]]
2102+
"""
2103+
2104+
template_fields = (
2105+
'schema_fields_updates',
2106+
'dataset_id',
2107+
'table_id',
2108+
'project_id',
2109+
'impersonation_chain',
2110+
)
2111+
template_fields_renderers = {"schema_fields_updates": "json"}
2112+
ui_color = BigQueryUIColors.TABLE.value
2113+
2114+
@apply_defaults
2115+
def __init__(
2116+
self,
2117+
*,
2118+
schema_fields_updates: List[Dict[str, Any]],
2119+
include_policy_tags: Optional[bool] = False,
2120+
dataset_id: Optional[str] = None,
2121+
table_id: Optional[str] = None,
2122+
project_id: Optional[str] = None,
2123+
gcp_conn_id: str = 'google_cloud_default',
2124+
delegate_to: Optional[str] = None,
2125+
impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
2126+
**kwargs,
2127+
) -> None:
2128+
self.schema_fields_updates = schema_fields_updates
2129+
self.include_policy_tags = include_policy_tags
2130+
self.table_id = table_id
2131+
self.dataset_id = dataset_id
2132+
self.project_id = project_id
2133+
self.gcp_conn_id = gcp_conn_id
2134+
self.delegate_to = delegate_to
2135+
self.impersonation_chain = impersonation_chain
2136+
super().__init__(**kwargs)
2137+
2138+
def execute(self, context):
2139+
bq_hook = BigQueryHook(
2140+
gcp_conn_id=self.gcp_conn_id,
2141+
delegate_to=self.delegate_to,
2142+
impersonation_chain=self.impersonation_chain,
2143+
)
2144+
2145+
return bq_hook.update_table_schema(
2146+
schema_fields_updates=self.schema_fields_updates,
2147+
include_policy_tags=self.include_policy_tags,
2148+
dataset_id=self.dataset_id,
2149+
table_id=self.table_id,
2150+
project_id=self.project_id,
2151+
)
2152+
2153+
20422154
# pylint: disable=too-many-arguments
20432155
class BigQueryInsertJobOperator(BaseOperator):
20442156
"""

‎docs/apache-airflow-providers-google/operators/cloud/bigquery.rst‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,23 @@ in the given dataset.
245245
:start-after: [START howto_operator_bigquery_upsert_table]
246246
:end-before: [END howto_operator_bigquery_upsert_table]
247247

248+
.. _howto/operator:BigQueryUpdateTableSchemaOperator:
249+
250+
Update table schema
251+
"""""""""""""""""""
252+
253+
To update the schema of a table you can use
254+
:class:`~airflow.providers.google.cloud.operators.bigquery.BigQueryUpdateTableSchemaOperator`.
255+
256+
This operator updates the schema field values supplied, while leaving the rest unchanged. This is useful
257+
for instance to set new field descriptions on an existing table schema.
258+
259+
.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_bigquery_operations.py
260+
:language: python
261+
:dedent: 4
262+
:start-after: [START howto_operator_bigquery_update_table_schema]
263+
:end-before: [END howto_operator_bigquery_update_table_schema]
264+
248265
.. _howto/operator:BigQueryDeleteTableOperator:
249266

250267
Delete table

0 commit comments

Comments
 (0)