This is an automated email from the ASF dual-hosted git repository.
turbaszek pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/master by this push:
new 3994030 Refactor BigQuery operators (#8858)
3994030 is described below
commit 3994030ea678727daaf9c2bfed0ca94a096f8d2a
Author: Tomek Urbaszek <[email protected]>
AuthorDate: Tue May 26 10:03:23 2020 +0200
Refactor BigQuery operators (#8858)
* Refactor BigQueryCreateEmptyTableOperator
* Refactor BigQueryCreateExternalTableOperator
* Refactor BigQueryDeleteDatasetOperator
* Refactor BigQueryCreateEmptyDatasetOperator
* Refactor BigQueryGetDataOperator
* BigQueryGetDatasetTablesOperator
* Refactor BigQueryPatchDatasetOperator
* Refactor BigQueryUpdateDatasetOperator
* Refactor BigQueryDeleteTableOperator
* Refactor BigQueryUpsertTableOperator
* Apply cr suggestions
* fixup! Apply cr suggestions
---
airflow/providers/google/cloud/hooks/bigquery.py | 44 +-
.../providers/google/cloud/operators/bigquery.py | 483 ++++++++++++---------
.../providers/google/cloud/hooks/test_bigquery.py | 2 +-
.../google/cloud/operators/test_bigquery.py | 100 ++---
4 files changed, 363 insertions(+), 266 deletions(-)
diff --git a/airflow/providers/google/cloud/hooks/bigquery.py
b/airflow/providers/google/cloud/hooks/bigquery.py
index 73cba36..6155bc4 100644
--- a/airflow/providers/google/cloud/hooks/bigquery.py
+++ b/airflow/providers/google/cloud/hooks/bigquery.py
@@ -233,6 +233,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
retry: Optional[Retry] = DEFAULT_RETRY,
num_retries: Optional[int] = None,
location: Optional[str] = None,
+ exists_ok: bool = True
) -> Table:
"""
Creates a new, empty table in the dataset.
@@ -293,6 +294,8 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
:type encryption_configuration: dict
:param num_retries: Maximum number of retries in case of connection
problems.
:type num_retries: int
+ :param exists_ok: If ``True``, ignore "already exists" errors when
creating the table.
+ :type exists_ok: bool
:return: Created table
"""
if num_retries:
@@ -333,16 +336,19 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
table = Table.from_api_repr(table_resource)
return self.get_client(project_id=project_id,
location=location).create_table(
table=table,
- exists_ok=True,
+ exists_ok=exists_ok,
retry=retry
)
@GoogleBaseHook.fallback_to_default_project_id
- def create_empty_dataset(self,
- dataset_id: Optional[str] = None,
- project_id: Optional[str] = None,
- location: Optional[str] = None,
- dataset_reference: Optional[Dict[str, Any]] =
None) -> None:
+ def create_empty_dataset(
+ self,
+ dataset_id: Optional[str] = None,
+ project_id: Optional[str] = None,
+ location: Optional[str] = None,
+ dataset_reference: Optional[Dict[str, Any]] = None,
+ exists_ok: bool = True,
+ ) -> None:
"""
Create a new empty dataset:
https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/insert
@@ -358,6 +364,8 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
:param dataset_reference: Dataset reference that could be provided
with request body. More info:
https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#resource
:type dataset_reference: dict
+ :param exists_ok: If ``True``, ignore "already exists" errors when
creating the DATASET.
+ :type exists_ok: bool
"""
dataset_reference = dataset_reference or {"datasetReference": {}}
@@ -388,8 +396,10 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
if location:
dataset_reference["location"] = dataset_reference.get("location",
location)
- dataset = Dataset.from_api_repr(dataset_reference)
- self.get_client(location=location).create_dataset(dataset=dataset,
exists_ok=True)
+ dataset: Dataset = Dataset.from_api_repr(dataset_reference)
+ self.log.info('Creating dataset: %s in project: %s ',
dataset.dataset_id, dataset.project)
+ self.get_client(location=location).create_dataset(dataset=dataset,
exists_ok=exists_ok)
+ self.log.info('Dataset created successfully.')
@GoogleBaseHook.fallback_to_default_project_id
def get_dataset_tables(
@@ -397,7 +407,6 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
dataset_id: str,
project_id: Optional[str] = None,
max_results: Optional[int] = None,
- page_token: Optional[str] = None,
retry: Retry = DEFAULT_RETRY,
) -> List[Dict[str, Any]]:
"""
@@ -413,19 +422,17 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
:type project_id: str
:param max_results: (Optional) the maximum number of tables to return.
:type max_results: int
- :param page_token: (Optional) page token, returned from a previous
call,
- identifying the result set.
- :type page_token: str
:param retry: How to retry the RPC.
:type retry: google.api_core.retry.Retry
:return: List of tables associated with the dataset.
"""
+ self.log.info('Start getting tables list from dataset: %s.%s',
project_id, dataset_id)
tables = self.get_client().list_tables(
dataset=DatasetReference(project=project_id,
dataset_id=dataset_id),
max_results=max_results,
- page_token=page_token,
retry=retry,
)
+ # Convert to a list (consumes all values)
return [t.reference.to_api_repr() for t in tables]
@GoogleBaseHook.fallback_to_default_project_id
@@ -480,7 +487,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
project_id: Optional[str] = None,
) -> None:
"""
- Creates a new external table in the dataset with the data in Google
+ Creates a new external table in the dataset with the data from Google
Cloud Storage. See here:
https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#resource
@@ -619,7 +626,12 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
table.encryption_configuration =
EncryptionConfiguration.from_api_repr(encryption_configuration)
self.log.info('Creating external table: %s',
external_project_dataset_table)
- self.create_empty_table(table_resource=table.to_api_repr(),
project_id=project_id, location=location)
+ self.create_empty_table(
+ table_resource=table.to_api_repr(),
+ project_id=project_id,
+ location=location,
+ exists_ok=True
+ )
self.log.info('External table created successfully: %s',
external_project_dataset_table)
@GoogleBaseHook.fallback_to_default_project_id
@@ -894,6 +906,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
if value and not spec_value:
dataset_resource["datasetReference"][key] = value
+ self.log.info('Start updating dataset')
dataset = self.get_client(project_id=project_id).update_dataset(
dataset=Dataset.from_api_repr(dataset_resource),
fields=fields,
@@ -937,6 +950,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
service = self.get_service()
dataset_project_id = project_id or self.project_id
+ self.log.info('Start patching dataset: %s:%s', dataset_project_id,
dataset_id)
dataset = (
service.datasets() # pylint: disable=no-member
.patch(
diff --git a/airflow/providers/google/cloud/operators/bigquery.py
b/airflow/providers/google/cloud/operators/bigquery.py
index b93b757..8d36d25 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -26,7 +26,8 @@ import warnings
from typing import Any, Dict, Iterable, List, Optional, SupportsAbs, Union
import attr
-from googleapiclient.errors import HttpError
+from google.api_core.exceptions import Conflict
+from google.cloud.bigquery import TableReference
from airflow.exceptions import AirflowException
from airflow.models import BaseOperator, BaseOperatorLink
@@ -36,7 +37,7 @@ from airflow.providers.google.cloud.hooks.bigquery import
BigQueryHook
from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
from airflow.utils.decorators import apply_defaults
-BIGQUERY_JOB_DETAILS_LINK_FMT =
'https://console.cloud.google.com/bigquery?j={job_id}'
+BIGQUERY_JOB_DETAILS_LINK_FMT =
"https://console.cloud.google.com/bigquery?j={job_id}";
_DEPRECATION_MSG = "The bigquery_conn_id parameter has been deprecated. " \
"You should pass the gcp_conn_id parameter."
@@ -72,7 +73,7 @@ class BigQueryConsoleIndexableLink(BaseOperatorLink):
@property
def name(self) -> str:
- return 'BigQuery Console #{index}'.format(index=self.index + 1)
+ return f'BigQuery Console #{self.index + 1}'
def get_link(self, operator, dttm):
ti = TaskInstance(task=operator, execution_date=dttm)
@@ -343,21 +344,23 @@ class BigQueryGetDataOperator(BaseOperator):
:param location: The location used for the operation.
:type location: str
"""
- template_fields = ('dataset_id', 'table_id', 'max_results')
+ template_fields = ('dataset_id', 'table_id', 'max_results',
'selected_fields')
ui_color = BigQueryUIColors.QUERY.value
@apply_defaults
- def __init__(self,
- dataset_id: str,
- table_id: str,
- max_results: int = 100,
- selected_fields: Optional[str] = None,
- gcp_conn_id: str = 'google_cloud_default',
- bigquery_conn_id: Optional[str] = None,
- delegate_to: Optional[str] = None,
- location: Optional[str] = None,
- *args,
- **kwargs) -> None:
+ def __init__(
+ self,
+ dataset_id: str,
+ table_id: str,
+ max_results: int = 100,
+ selected_fields: Optional[str] = None,
+ gcp_conn_id: str = 'google_cloud_default',
+ bigquery_conn_id: Optional[str] = None,
+ delegate_to: Optional[str] = None,
+ location: Optional[str] = None,
+ *args,
+ **kwargs
+ ) -> None:
super().__init__(*args, **kwargs)
if bigquery_conn_id:
@@ -375,8 +378,7 @@ class BigQueryGetDataOperator(BaseOperator):
self.location = location
def execute(self, context):
- self.log.info('Fetching Data from:')
- self.log.info('Dataset: %s ; Table: %s ; Max Results: %s',
+ self.log.info('Fetching Data from %s.%s max results: %s',
self.dataset_id, self.table_id, self.max_results)
hook = BigQueryHook(
@@ -392,7 +394,7 @@ class BigQueryGetDataOperator(BaseOperator):
location=self.location
)
- self.log.info('Total Extracted rows: %s', len(rows))
+ self.log.info('Total extracted rows: %s', len(rows))
table_data = [row.values() for row in rows]
return table_data
@@ -647,6 +649,10 @@ class BigQueryCreateEmptyTableOperator(BaseOperator):
:type dataset_id: str
:param table_id: The Name of the table to be created. (templated)
:type table_id: str
+ :param table_resource: Table resource as described in documentation:
+ https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#Table
+ If provided all other parameters are ignored.
+ :type table_resource: Dict[str, Any]
:param schema_fields: If set, the schema field list as defined here:
https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.load.schema
@@ -741,28 +747,38 @@ class BigQueryCreateEmptyTableOperator(BaseOperator):
https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#clustering.fields
:type cluster_fields: list
"""
- template_fields = ('dataset_id', 'table_id', 'project_id',
- 'gcs_schema_object', 'labels', 'view')
+ template_fields = (
+ 'dataset_id',
+ 'table_id',
+ 'project_id',
+ 'gcs_schema_object',
+ 'labels',
+ 'view'
+ )
+
ui_color = BigQueryUIColors.TABLE.value
# pylint: disable=too-many-arguments
@apply_defaults
- def __init__(self,
- dataset_id: str,
- table_id: str,
- project_id: Optional[str] = None,
- schema_fields: Optional[List] = None,
- gcs_schema_object: Optional[str] = None,
- time_partitioning: Optional[Dict] = None,
- bigquery_conn_id: str = 'google_cloud_default',
- google_cloud_storage_conn_id: str = 'google_cloud_default',
- delegate_to: Optional[str] = None,
- labels: Optional[Dict] = None,
- view: Optional[Dict] = None,
- encryption_configuration: Optional[Dict] = None,
- location: Optional[str] = None,
- cluster_fields: Optional[List[str]] = None,
- *args, **kwargs) -> None:
+ def __init__(
+ self,
+ dataset_id: str,
+ table_id: str,
+ table_resource: Optional[Dict[str, Any]] = None,
+ project_id: Optional[str] = None,
+ schema_fields: Optional[List] = None,
+ gcs_schema_object: Optional[str] = None,
+ time_partitioning: Optional[Dict] = None,
+ bigquery_conn_id: str = 'google_cloud_default',
+ google_cloud_storage_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None,
+ labels: Optional[Dict] = None,
+ view: Optional[Dict] = None,
+ encryption_configuration: Optional[Dict] = None,
+ location: Optional[str] = None,
+ cluster_fields: Optional[List[str]] = None,
+ *args, **kwargs
+ ) -> None:
super().__init__(*args, **kwargs)
self.project_id = project_id
@@ -779,16 +795,17 @@ class BigQueryCreateEmptyTableOperator(BaseOperator):
self.encryption_configuration = encryption_configuration
self.location = location
self.cluster_fields = cluster_fields
+ self.table_resource = table_resource
def execute(self, context):
- bq_hook = BigQueryHook(gcp_conn_id=self.bigquery_conn_id,
- delegate_to=self.delegate_to,
- location=self.location)
+ bq_hook = BigQueryHook(
+ gcp_conn_id=self.bigquery_conn_id,
+ delegate_to=self.delegate_to,
+ location=self.location
+ )
if not self.schema_fields and self.gcs_schema_object:
-
gcs_bucket, gcs_object = _parse_gcs_url(self.gcs_schema_object)
-
gcs_hook = GCSHook(
google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
delegate_to=self.delegate_to)
@@ -799,9 +816,8 @@ class BigQueryCreateEmptyTableOperator(BaseOperator):
schema_fields = self.schema_fields
try:
- self.log.info('Creating Table %s:%s.%s',
- self.project_id, self.dataset_id, self.table_id)
- bq_hook.create_empty_table(
+ self.log.info('Creating table')
+ table = bq_hook.create_empty_table(
project_id=self.project_id,
dataset_id=self.dataset_id,
table_id=self.table_id,
@@ -810,22 +826,20 @@ class BigQueryCreateEmptyTableOperator(BaseOperator):
cluster_fields=self.cluster_fields,
labels=self.labels,
view=self.view,
- encryption_configuration=self.encryption_configuration
+ encryption_configuration=self.encryption_configuration,
+ table_resource=self.table_resource,
+ exists_ok=False,
)
- self.log.info('Table created successfully: %s:%s.%s',
- self.project_id, self.dataset_id, self.table_id)
- except HttpError as err:
- if err.resp.status != 409:
- raise
- else:
- self.log.info('Table %s:%s.%s already exists.',
self.project_id,
- self.dataset_id, self.table_id)
+ self.log.info('Table %s.%s.%s created successfully',
+ table.project, table.dataset_id, table.table_id)
+ except Conflict:
+ self.log.info('Table %s.%s already exists.', self.dataset_id,
self.table_id)
# pylint: disable=too-many-instance-attributes
class BigQueryCreateExternalTableOperator(BaseOperator):
"""
- Creates a new external table in the dataset with the data in Google Cloud
+ Creates a new external table in the dataset with the data from Google Cloud
Storage.
The schema to be used for the BigQuery table may be specified in one of
@@ -836,8 +850,7 @@ class BigQueryCreateExternalTableOperator(BaseOperator):
:param bucket: The bucket to point the external table to. (templated)
:type bucket: str
:param source_objects: List of Google Cloud Storage URIs to point
- table to. (templated)
- If source_format is 'DATASTORE_BACKUP', the list must only contain a
single URI.
+ table to. If source_format is 'DATASTORE_BACKUP', the list must only
contain a single URI.
:type source_objects: list
:param destination_project_dataset_table: The dotted
``(<project>.)<dataset>.<table>``
BigQuery table to load data into (templated). If ``<project>`` is not
included,
@@ -852,6 +865,10 @@ class BigQueryCreateExternalTableOperator(BaseOperator):
{"name": "salary", "type": "INTEGER", "mode":
"NULLABLE"}]
Should not be set when source_format is 'DATASTORE_BACKUP'.
+ :param table_resource: Table resource as described in documentation:
+ https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#Table
+ If provided all other parameters are ignored. External schema from
object will be resolved.
+ :type table_resource: Dict[str, Any]
:type schema_fields: list
:param schema_object: If set, a GCS object path pointing to a .json file
that
contains the schema for the table. (templated)
@@ -906,34 +923,43 @@ class BigQueryCreateExternalTableOperator(BaseOperator):
:param location: The location used for the operation.
:type location: str
"""
- template_fields = ('bucket', 'source_objects',
- 'schema_object', 'destination_project_dataset_table',
'labels')
+ template_fields = (
+ 'bucket',
+ 'source_objects',
+ 'schema_object',
+ 'destination_project_dataset_table',
+ 'labels',
+ 'table_resource',
+ )
ui_color = BigQueryUIColors.TABLE.value
- # pylint: disable=too-many-arguments
+ # pylint: disable=too-many-arguments,too-many-locals
@apply_defaults
- def __init__(self,
- bucket: str,
- source_objects: List,
- destination_project_dataset_table: str,
- schema_fields: Optional[List] = None,
- schema_object: Optional[str] = None,
- source_format: str = 'CSV',
- compression: str = 'NONE',
- skip_leading_rows: int = 0,
- field_delimiter: str = ',',
- max_bad_records: int = 0,
- quote_character: Optional[str] = None,
- allow_quoted_newlines: bool = False,
- allow_jagged_rows: bool = False,
- bigquery_conn_id: str = 'google_cloud_default',
- google_cloud_storage_conn_id: str = 'google_cloud_default',
- delegate_to: Optional[str] = None,
- src_fmt_configs: Optional[dict] = None,
- labels: Optional[Dict] = None,
- encryption_configuration: Optional[Dict] = None,
- location: Optional[str] = None,
- *args, **kwargs) -> None:
+ def __init__(
+ self,
+ bucket: str,
+ source_objects: List,
+ destination_project_dataset_table: str,
+ table_resource: Optional[Dict[str, Any]] = None,
+ schema_fields: Optional[List] = None,
+ schema_object: Optional[str] = None,
+ source_format: str = 'CSV',
+ compression: str = 'NONE',
+ skip_leading_rows: int = 0,
+ field_delimiter: str = ',',
+ max_bad_records: int = 0,
+ quote_character: Optional[str] = None,
+ allow_quoted_newlines: bool = False,
+ allow_jagged_rows: bool = False,
+ bigquery_conn_id: str = 'google_cloud_default',
+ google_cloud_storage_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None,
+ src_fmt_configs: Optional[dict] = None,
+ labels: Optional[Dict] = None,
+ encryption_configuration: Optional[Dict] = None,
+ location: Optional[str] = None,
+ *args, **kwargs
+ ) -> None:
super().__init__(*args, **kwargs)
# GCS config
@@ -942,6 +968,34 @@ class BigQueryCreateExternalTableOperator(BaseOperator):
self.schema_object = schema_object
# BQ config
+ kwargs_passed = any([
+ destination_project_dataset_table,
+ schema_fields,
+ source_format,
+ compression,
+ skip_leading_rows,
+ field_delimiter,
+ max_bad_records,
+ quote_character,
+ allow_quoted_newlines,
+ allow_jagged_rows,
+ src_fmt_configs,
+ labels,
+ encryption_configuration,
+ ])
+
+ if not table_resource:
+ warnings.warn(
+ "Passing table parameters via keywords arguments will be
deprecated. "
+ "Please use provide table definition using `table_resource`
parameter."
+ "You can still use external `schema_object`. ",
+ DeprecationWarning, stacklevel=2
+ )
+
+ if table_resource and kwargs_passed:
+ raise ValueError("You provided both `table_resource` and exclusive
keywords arguments.")
+
+ self.table_resource = table_resource
self.destination_project_dataset_table =
destination_project_dataset_table
self.schema_fields = schema_fields
self.source_format = source_format
@@ -957,30 +1011,41 @@ class BigQueryCreateExternalTableOperator(BaseOperator):
self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
self.delegate_to = delegate_to
- self.src_fmt_configs = src_fmt_configs if src_fmt_configs is not None
else dict()
+ self.src_fmt_configs = src_fmt_configs or dict()
self.labels = labels
self.encryption_configuration = encryption_configuration
self.location = location
def execute(self, context):
- bq_hook = BigQueryHook(gcp_conn_id=self.bigquery_conn_id,
- delegate_to=self.delegate_to,
- location=self.location)
+ bq_hook = BigQueryHook(
+ gcp_conn_id=self.bigquery_conn_id,
+ delegate_to=self.delegate_to,
+ location=self.location
+ )
if not self.schema_fields and self.schema_object and
self.source_format != 'DATASTORE_BACKUP':
gcs_hook = GCSHook(
google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
delegate_to=self.delegate_to)
- schema_fields = json.loads(gcs_hook.download(
- self.bucket,
- self.schema_object).decode("utf-8"))
+ schema_object = gcs_hook.download(self.bucket, self.schema_object)
+ schema_fields = json.loads(schema_object.decode("utf-8"))
else:
schema_fields = self.schema_fields
- source_uris = ['gs://{}/{}'.format(self.bucket, source_object)
- for source_object in self.source_objects]
+ if schema_fields and self.table_resource:
+ self.table_resource["externalDataConfiguration"]["schema"] =
schema_fields
+
+ if self.table_resource:
+ tab_ref =
TableReference.from_string(self.destination_project_dataset_table)
+ bq_hook.create_empty_table(
+ table_resource=self.table_resource,
+ project_id=tab_ref.project,
+ table_id=tab_ref.table_id,
+ dataset_id=tab_ref.dataset_id,
+ )
+ else:
+ source_uris = [f"gs://{self.bucket}/{source_object}" for
source_object in self.source_objects]
- try:
bq_hook.create_external_table(
external_project_dataset_table=self.destination_project_dataset_table,
schema_fields=schema_fields,
@@ -997,9 +1062,6 @@ class BigQueryCreateExternalTableOperator(BaseOperator):
labels=self.labels,
encryption_configuration=self.encryption_configuration
)
- except HttpError as err:
- if err.resp.status != 409:
- raise
class BigQueryDeleteDatasetOperator(BaseOperator):
@@ -1037,14 +1099,16 @@ class BigQueryDeleteDatasetOperator(BaseOperator):
ui_color = BigQueryUIColors.DATASET.value
@apply_defaults
- def __init__(self,
- dataset_id: str,
- project_id: Optional[str] = None,
- delete_contents: bool = False,
- gcp_conn_id: str = 'google_cloud_default',
- bigquery_conn_id: Optional[str] = None,
- delegate_to: Optional[str] = None,
- *args, **kwargs) -> None:
+ def __init__(
+ self,
+ dataset_id: str,
+ project_id: Optional[str] = None,
+ delete_contents: bool = False,
+ gcp_conn_id: str = 'google_cloud_default',
+ bigquery_conn_id: Optional[str] = None,
+ delegate_to: Optional[str] = None,
+ *args, **kwargs
+ ) -> None:
if bigquery_conn_id:
warnings.warn(
"The bigquery_conn_id parameter has been deprecated. You
should pass "
@@ -1062,8 +1126,10 @@ class BigQueryDeleteDatasetOperator(BaseOperator):
def execute(self, context):
self.log.info('Dataset id: %s Project id: %s', self.dataset_id,
self.project_id)
- bq_hook = BigQueryHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
+ bq_hook = BigQueryHook(
+ gcp_conn_id=self.gcp_conn_id,
+ delegate_to=self.delegate_to
+ )
bq_hook.delete_dataset(
project_id=self.project_id,
@@ -1074,17 +1140,14 @@ class BigQueryDeleteDatasetOperator(BaseOperator):
class BigQueryCreateEmptyDatasetOperator(BaseOperator):
"""
- This operator is used to create new dataset for your Project in Big query.
+ This operator is used to create new dataset for your Project in BigQuery.
https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#resource
:param project_id: The name of the project where we want to create the
dataset.
- Don't need to provide, if projectId in dataset_reference.
:type project_id: str
- :param dataset_id: The id of dataset. Don't need to provide,
- if datasetId in dataset_reference.
+ :param dataset_id: The id of dataset. Don't need to provide, if datasetId
in dataset_reference.
:type dataset_id: str
- :param location: (Optional) The geographic location where the dataset
should reside.
- There is no default value but the dataset will be created in US if
nothing is provided.
+ :param location: The geographic location where the dataset should reside.
:type location: str
:param dataset_reference: Dataset reference that could be provided with
request body.
More info:
@@ -1108,12 +1171,12 @@ class BigQueryCreateEmptyDatasetOperator(BaseOperator):
:type location: str
"""
- template_fields = ('dataset_id', 'project_id')
+ template_fields = ('dataset_id', 'project_id', 'dataset_reference')
ui_color = BigQueryUIColors.DATASET.value
@apply_defaults
def __init__(self,
- dataset_id: str,
+ dataset_id: Optional[str] = None,
project_id: Optional[str] = None,
dataset_reference: Optional[Dict] = None,
location: Optional[str] = None,
@@ -1138,24 +1201,23 @@ class BigQueryCreateEmptyDatasetOperator(BaseOperator):
super().__init__(*args, **kwargs)
def execute(self, context):
- self.log.info('Dataset id: %s Project id: %s', self.dataset_id,
self.project_id)
-
- bq_hook = BigQueryHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to,
- location=self.location)
+ bq_hook = BigQueryHook(
+ gcp_conn_id=self.gcp_conn_id,
+ delegate_to=self.delegate_to,
+ location=self.location
+ )
try:
- self.log.info('Creating Dataset: %s in project: %s ',
self.dataset_id, self.project_id)
bq_hook.create_empty_dataset(
project_id=self.project_id,
dataset_id=self.dataset_id,
dataset_reference=self.dataset_reference,
- location=self.location)
- self.log.info('Dataset created successfully.')
- except HttpError as err:
- if err.resp.status != 409:
- raise
- self.log.info('Dataset %s already exists.', self.dataset_id)
+ location=self.location,
+ exists_ok=False,
+ )
+ except Conflict:
+ dataset_id = self.dataset_reference.get("datasetReference",
{}).get("datasetId", self.dataset_id)
+ self.log.info('Dataset %s already exists.', dataset_id)
class BigQueryGetDatasetOperator(BaseOperator):
@@ -1213,50 +1275,44 @@ class BigQueryGetDatasetTablesOperator(BaseOperator):
:type project_id: str
:param max_results: (Optional) the maximum number of tables to return.
:type max_results: int
- :param page_token: (Optional) page token, returned from a previous call,
- identifying the result set.
- :type page_token: str
:param gcp_conn_id: (Optional) The connection ID used to connect to Google
Cloud Platform.
:type gcp_conn_id: str
:param delegate_to: (Optional) The account to impersonate, if any.
For this to work, the service account making the request must have
domain-wide
delegation enabled.
:type delegate_to: str
-
- :rtype: dict
- .. seealso::
https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/list#response-body
"""
template_fields = ('dataset_id', 'project_id')
ui_color = BigQueryUIColors.DATASET.value
@apply_defaults
- def __init__(self,
- dataset_id: str,
- project_id: Optional[str] = None,
- max_results: Optional[int] = None,
- page_token: Optional[str] = None,
- gcp_conn_id: Optional[str] = 'google_cloud_default',
- delegate_to: Optional[str] = None,
- *args, **kwargs) -> None:
+ def __init__(
+ self,
+ dataset_id: str,
+ project_id: Optional[str] = None,
+ max_results: Optional[int] = None,
+ gcp_conn_id: Optional[str] = 'google_cloud_default',
+ delegate_to: Optional[str] = None,
+ *args, **kwargs
+ ) -> None:
self.dataset_id = dataset_id
self.project_id = project_id
self.max_results = max_results
- self.page_token = page_token
self.gcp_conn_id = gcp_conn_id
self.delegate_to = delegate_to
super().__init__(*args, **kwargs)
def execute(self, context):
- bq_hook = BigQueryHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
-
- self.log.info('Start getting tables list from dataset: %s:%s',
self.project_id, self.dataset_id)
+ bq_hook = BigQueryHook(
+ gcp_conn_id=self.gcp_conn_id,
+ delegate_to=self.delegate_to,
+ )
return bq_hook.get_dataset_tables(
dataset_id=self.dataset_id,
project_id=self.project_id,
max_results=self.max_results,
- page_token=self.page_token)
+ )
class BigQueryPatchDatasetOperator(BaseOperator):
@@ -1283,13 +1339,20 @@ class BigQueryPatchDatasetOperator(BaseOperator):
ui_color = BigQueryUIColors.DATASET.value
@apply_defaults
- def __init__(self,
- dataset_id: str,
- dataset_resource: dict,
- project_id: Optional[str] = None,
- gcp_conn_id: str = 'google_cloud_default',
- delegate_to: Optional[str] = None,
- *args, **kwargs) -> None:
+ def __init__(
+ self,
+ dataset_id: str,
+ dataset_resource: dict,
+ project_id: Optional[str] = None,
+ gcp_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None,
+ *args, **kwargs,
+ ) -> None:
+
+ warnings.warn(
+ "This operator is deprecated. Please use
BigQueryUpdateDatasetOperator.",
+ DeprecationWarning, stacklevel=3
+ )
self.dataset_id = dataset_id
self.project_id = project_id
self.gcp_conn_id = gcp_conn_id
@@ -1298,22 +1361,25 @@ class BigQueryPatchDatasetOperator(BaseOperator):
super().__init__(*args, **kwargs)
def execute(self, context):
- bq_hook = BigQueryHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
-
- self.log.info('Start patching dataset: %s:%s', self.project_id,
self.dataset_id)
+ bq_hook = BigQueryHook(
+ gcp_conn_id=self.gcp_conn_id,
+ delegate_to=self.delegate_to,
+ )
return bq_hook.patch_dataset(
dataset_id=self.dataset_id,
dataset_resource=self.dataset_resource,
- project_id=self.project_id)
+ project_id=self.project_id,
+ )
class BigQueryUpdateDatasetOperator(BaseOperator):
"""
This operator is used to update dataset for your Project in BigQuery.
- The update method replaces the entire dataset resource, whereas the patch
- method only replaces fields that are provided in the submitted dataset
resource.
+ Use ``fields`` to specify which fields of dataset to update. If a field
+ is listed in ``fields`` and is ``None`` in dataset, it will be deleted.
+ If no ``fields`` are provided then all fields of provided
``dataset_reources``
+ will be used.
:param dataset_id: The id of dataset. Don't need to provide,
if datasetId in dataset_reference.
@@ -1321,6 +1387,8 @@ class BigQueryUpdateDatasetOperator(BaseOperator):
:param dataset_resource: Dataset resource that will be provided with
request body.
https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#resource
:type dataset_resource: dict
+ :param fields: The properties of dataset to change (e.g. "friendly_name").
+ :type fields: Sequence[str]
:param project_id: The name of the project where we want to create the
dataset.
Don't need to provide, if projectId in dataset_reference.
:type project_id: str
@@ -1334,27 +1402,30 @@ class BigQueryUpdateDatasetOperator(BaseOperator):
ui_color = BigQueryUIColors.DATASET.value
@apply_defaults
- def __init__(self,
- dataset_resource: dict,
- dataset_id: Optional[str] = None,
- project_id: Optional[str] = None,
- gcp_conn_id: str = 'google_cloud_default',
- delegate_to: Optional[str] = None,
- *args, **kwargs) -> None:
+ def __init__(
+ self,
+ dataset_resource: dict,
+ fields: Optional[List[str]] = None,
+ dataset_id: Optional[str] = None,
+ project_id: Optional[str] = None,
+ gcp_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None,
+ *args, **kwargs
+ ) -> None:
self.dataset_id = dataset_id
self.project_id = project_id
+ self.fields = fields
self.gcp_conn_id = gcp_conn_id
self.dataset_resource = dataset_resource
self.delegate_to = delegate_to
super().__init__(*args, **kwargs)
def execute(self, context):
- bq_hook = BigQueryHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to)
-
- self.log.info('Start updating dataset: %s:%s', self.project_id,
self.dataset_id)
-
- fields = list(self.dataset_resource.keys())
+ bq_hook = BigQueryHook(
+ gcp_conn_id=self.gcp_conn_id,
+ delegate_to=self.delegate_to,
+ )
+ fields = self.fields or list(self.dataset_resource.keys())
dataset = bq_hook.update_dataset(
dataset_resource=self.dataset_resource,
@@ -1392,15 +1463,17 @@ class BigQueryDeleteTableOperator(BaseOperator):
ui_color = BigQueryUIColors.TABLE.value
@apply_defaults
- def __init__(self,
- deletion_dataset_table: str,
- gcp_conn_id: str = 'google_cloud_default',
- bigquery_conn_id: Optional[str] = None,
- delegate_to: Optional[str] = None,
- ignore_if_missing: bool = False,
- location: Optional[str] = None,
- *args,
- **kwargs) -> None:
+ def __init__(
+ self,
+ deletion_dataset_table: str,
+ gcp_conn_id: str = 'google_cloud_default',
+ bigquery_conn_id: Optional[str] = None,
+ delegate_to: Optional[str] = None,
+ ignore_if_missing: bool = False,
+ location: Optional[str] = None,
+ *args,
+ **kwargs,
+ ) -> None:
super().__init__(*args, **kwargs)
if bigquery_conn_id:
@@ -1417,12 +1490,15 @@ class BigQueryDeleteTableOperator(BaseOperator):
def execute(self, context):
self.log.info('Deleting: %s', self.deletion_dataset_table)
- hook = BigQueryHook(gcp_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to,
- location=self.location)
- hook.run_table_delete(
- deletion_dataset_table=self.deletion_dataset_table,
- ignore_if_missing=self.ignore_if_missing)
+ hook = BigQueryHook(
+ gcp_conn_id=self.gcp_conn_id,
+ delegate_to=self.delegate_to,
+ location=self.location
+ )
+ hook.delete_table(
+ table_id=self.deletion_dataset_table,
+ not_found_ok=self.ignore_if_missing
+ )
class BigQueryUpsertTableOperator(BaseOperator):
@@ -1451,20 +1527,22 @@ class BigQueryUpsertTableOperator(BaseOperator):
:param location: The location used for the operation.
:type location: str
"""
- template_fields = ('dataset_id', 'table_resource',)
+ template_fields = ('dataset_id', 'table_resource')
ui_color = BigQueryUIColors.TABLE.value
@apply_defaults
- def __init__(self,
- dataset_id: str,
- table_resource: dict,
- project_id: Optional[str] = None,
- gcp_conn_id: str = 'google_cloud_default',
- bigquery_conn_id: Optional[str] = None,
- delegate_to: Optional[str] = None,
- location: Optional[str] = None,
- *args,
- **kwargs) -> None:
+ def __init__(
+ self,
+ dataset_id: str,
+ table_resource: dict,
+ project_id: Optional[str] = None,
+ gcp_conn_id: str = 'google_cloud_default',
+ bigquery_conn_id: Optional[str] = None,
+ delegate_to: Optional[str] = None,
+ location: Optional[str] = None,
+ *args,
+ **kwargs,
+ ) -> None:
super().__init__(*args, **kwargs)
if bigquery_conn_id:
@@ -1482,10 +1560,13 @@ class BigQueryUpsertTableOperator(BaseOperator):
def execute(self, context):
self.log.info('Upserting Dataset: %s with table_resource: %s',
self.dataset_id, self.table_resource)
- hook = BigQueryHook(bigquery_conn_id=self.gcp_conn_id,
- delegate_to=self.delegate_to,
- location=self.location)
+ hook = BigQueryHook(
+ bigquery_conn_id=self.gcp_conn_id,
+ delegate_to=self.delegate_to,
+ location=self.location,
+ )
hook.run_table_upsert(
dataset_id=self.dataset_id,
table_resource=self.table_resource,
- project_id=self.project_id)
+ project_id=self.project_id,
+ )
diff --git a/tests/providers/google/cloud/hooks/test_bigquery.py
b/tests/providers/google/cloud/hooks/test_bigquery.py
index bdb9c15..8c34690 100644
--- a/tests/providers/google/cloud/hooks/test_bigquery.py
+++ b/tests/providers/google/cloud/hooks/test_bigquery.py
@@ -891,7 +891,6 @@ class TestTableOperations(_BigQueryBaseTestClass):
mock_client.return_value.list_tables.assert_called_once_with(
dataset=dataset_reference,
max_results=None,
- page_token=None,
retry=DEFAULT_RETRY,
)
for res, exp in zip(result, table_list):
@@ -1598,6 +1597,7 @@ class TestBigQueryWithKMS(_BigQueryBaseTestClass):
table_resource=body,
project_id=PROJECT_ID,
location=None,
+ exists_ok=True,
)
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.insert_job")
diff --git a/tests/providers/google/cloud/operators/test_bigquery.py
b/tests/providers/google/cloud/operators/test_bigquery.py
index 8a019bb..6c032eb 100644
--- a/tests/providers/google/cloud/operators/test_bigquery.py
+++ b/tests/providers/google/cloud/operators/test_bigquery.py
@@ -65,14 +65,15 @@ class
TestBigQueryCreateEmptyTableOperator(unittest.TestCase):
@mock.patch('airflow.providers.google.cloud.operators.bigquery.BigQueryHook')
def test_execute(self, mock_hook):
- operator = BigQueryCreateEmptyTableOperator(task_id=TASK_ID,
- dataset_id=TEST_DATASET,
-
project_id=TEST_GCP_PROJECT_ID,
- table_id=TEST_TABLE_ID)
+ operator = BigQueryCreateEmptyTableOperator(
+ task_id=TASK_ID,
+ dataset_id=TEST_DATASET,
+ project_id=TEST_GCP_PROJECT_ID,
+ table_id=TEST_TABLE_ID
+ )
operator.execute(None)
- mock_hook.return_value \
- .create_empty_table \
+ mock_hook.return_value.create_empty_table \
.assert_called_once_with(
dataset_id=TEST_DATASET,
project_id=TEST_GCP_PROJECT_ID,
@@ -82,20 +83,23 @@ class
TestBigQueryCreateEmptyTableOperator(unittest.TestCase):
cluster_fields=None,
labels=None,
view=None,
- encryption_configuration=None
+ encryption_configuration=None,
+ table_resource=None,
+ exists_ok=False,
)
@mock.patch('airflow.providers.google.cloud.operators.bigquery.BigQueryHook')
def test_create_view(self, mock_hook):
- operator = BigQueryCreateEmptyTableOperator(task_id=TASK_ID,
- dataset_id=TEST_DATASET,
-
project_id=TEST_GCP_PROJECT_ID,
- table_id=TEST_TABLE_ID,
- view=VIEW_DEFINITION)
+ operator = BigQueryCreateEmptyTableOperator(
+ task_id=TASK_ID,
+ dataset_id=TEST_DATASET,
+ project_id=TEST_GCP_PROJECT_ID,
+ table_id=TEST_TABLE_ID,
+ view=VIEW_DEFINITION
+ )
operator.execute(None)
- mock_hook.return_value \
- .create_empty_table \
+ mock_hook.return_value.create_empty_table \
.assert_called_once_with(
dataset_id=TEST_DATASET,
project_id=TEST_GCP_PROJECT_ID,
@@ -105,7 +109,9 @@ class
TestBigQueryCreateEmptyTableOperator(unittest.TestCase):
cluster_fields=None,
labels=None,
view=VIEW_DEFINITION,
- encryption_configuration=None
+ encryption_configuration=None,
+ table_resource=None,
+ exists_ok=False,
)
@mock.patch('airflow.providers.google.cloud.operators.bigquery.BigQueryHook')
@@ -133,18 +139,18 @@ class
TestBigQueryCreateEmptyTableOperator(unittest.TestCase):
"field": "date_hired"
}
cluster_fields = ["date_birth"]
- operator = BigQueryCreateEmptyTableOperator(task_id=TASK_ID,
- dataset_id=TEST_DATASET,
-
project_id=TEST_GCP_PROJECT_ID,
- table_id=TEST_TABLE_ID,
-
schema_fields=schema_fields,
-
time_partitioning=time_partitioning,
-
cluster_fields=cluster_fields
- )
+ operator = BigQueryCreateEmptyTableOperator(
+ task_id=TASK_ID,
+ dataset_id=TEST_DATASET,
+ project_id=TEST_GCP_PROJECT_ID,
+ table_id=TEST_TABLE_ID,
+ schema_fields=schema_fields,
+ time_partitioning=time_partitioning,
+ cluster_fields=cluster_fields
+ )
operator.execute(None)
- mock_hook.return_value \
- .create_empty_table \
+ mock_hook.return_value.create_empty_table \
.assert_called_once_with(
dataset_id=TEST_DATASET,
project_id=TEST_GCP_PROJECT_ID,
@@ -154,7 +160,9 @@ class
TestBigQueryCreateEmptyTableOperator(unittest.TestCase):
cluster_fields=cluster_fields,
labels=None,
view=None,
- encryption_configuration=None
+ encryption_configuration=None,
+ table_resource=None,
+ exists_ok=False,
)
@@ -228,14 +236,13 @@ class
TestBigQueryCreateEmptyDatasetOperator(unittest.TestCase):
)
operator.execute(None)
- mock_hook.return_value \
- .create_empty_dataset \
- .assert_called_once_with(
- dataset_id=TEST_DATASET,
- project_id=TEST_GCP_PROJECT_ID,
- location=TEST_DATASET_LOCATION,
- dataset_reference={}
- )
+ mock_hook.return_value.create_empty_dataset.assert_called_once_with(
+ dataset_id=TEST_DATASET,
+ project_id=TEST_GCP_PROJECT_ID,
+ location=TEST_DATASET_LOCATION,
+ dataset_reference={},
+ exists_ok=False,
+ )
class TestBigQueryGetDatasetOperator(unittest.TestCase):
@@ -705,7 +712,7 @@ class TestBigQueryTableDeleteOperator(unittest.TestCase):
@mock.patch('airflow.providers.google.cloud.operators.bigquery.BigQueryHook')
def test_execute(self, mock_hook):
ignore_if_missing = True
- deletion_dataset_table = '{}.{}'.format(TEST_DATASET, TEST_TABLE_ID)
+ deletion_dataset_table = f'{TEST_DATASET}.{TEST_TABLE_ID}'
operator = BigQueryDeleteTableOperator(
task_id=TASK_ID,
@@ -714,12 +721,10 @@ class TestBigQueryTableDeleteOperator(unittest.TestCase):
)
operator.execute(None)
- mock_hook.return_value \
- .run_table_delete \
- .assert_called_once_with(
- deletion_dataset_table=deletion_dataset_table,
- ignore_if_missing=ignore_if_missing
- )
+ mock_hook.return_value.delete_table.assert_called_once_with(
+ table_id=deletion_dataset_table,
+ not_found_ok=ignore_if_missing
+ )
class TestBigQueryGetDatasetTablesOperator(unittest.TestCase):
@@ -733,14 +738,11 @@ class
TestBigQueryGetDatasetTablesOperator(unittest.TestCase):
)
operator.execute(None)
- mock_hook.return_value \
- .get_dataset_tables \
- .assert_called_once_with(
- dataset_id=TEST_DATASET,
- project_id=TEST_GCP_PROJECT_ID,
- max_results=2,
- page_token=None
- )
+ mock_hook.return_value.get_dataset_tables.assert_called_once_with(
+ dataset_id=TEST_DATASET,
+ project_id=TEST_GCP_PROJECT_ID,
+ max_results=2,
+ )
class TestBigQueryConnIdDeprecationWarning(unittest.TestCase):
