Add exists_ok flag to BigQueryCreateEmptyTable(Dataset)Operator (#14026)

ysktir · ephraimbuddy · web-flow · commit 59c94c679e99 · 2021-02-13T14:31:27.000+01:00
Co-authored-by: uma6 &lt;&gt;
Co-authored-by: Ephraim Anierobi &lt;splendidzigy24@gmail.com&gt;
diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py
@@ -411,7 +411,7 @@ def create_empty_dataset(
         :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.
+        :param exists_ok: If ``True``, ignore "already exists" errors when creating the dataset.
         :type exists_ok: bool
         """
         dataset_reference = dataset_reference or {"datasetReference": {}}
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
@@ -868,6 +868,8 @@ class BigQueryCreateEmptyTableOperator(BaseOperator):
         Service Account Token Creator IAM role to the directly preceding identity, with first
         account from the list granting this role to the originating account (templated).
     :type impersonation_chain: Union[str, Sequence[str]]
+    :param exists_ok: If ``True``, ignore "already exists" errors when creating the table.
+    :type exists_ok: bool
     """
 
     template_fields = (
@@ -905,6 +907,7 @@ def __init__(
         location: Optional[str] = None,
         cluster_fields: Optional[List[str]] = None,
         impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
+        exists_ok: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -926,6 +929,7 @@ def __init__(
         self.cluster_fields = cluster_fields
         self.table_resource = table_resource
         self.impersonation_chain = impersonation_chain
+        self.exists_ok = exists_ok
 
     def execute(self, context) -> None:
         bq_hook = BigQueryHook(
@@ -960,7 +964,7 @@ def execute(self, context) -> None:
                 materialized_view=self.materialized_view,
                 encryption_configuration=self.encryption_configuration,
                 table_resource=self.table_resource,
-                exists_ok=False,
+                exists_ok=self.exists_ok,
             )
             self.log.info(
                 'Table %s.%s.%s created successfully', table.project, table.dataset_id, table.table_id
@@ -1357,6 +1361,8 @@ class BigQueryCreateEmptyDatasetOperator(BaseOperator):
         Service Account Token Creator IAM role to the directly preceding identity, with first
         account from the list granting this role to the originating account (templated).
     :type impersonation_chain: Union[str, Sequence[str]]
+    :param exists_ok: If ``True``, ignore "already exists" errors when creating the dataset.
+    :type exists_ok: bool
         **Example**: ::
 
             create_new_dataset = BigQueryCreateEmptyDatasetOperator(
@@ -1389,6 +1395,7 @@ def __init__(
         bigquery_conn_id: Optional[str] = None,
         delegate_to: Optional[str] = None,
         impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
+        exists_ok: bool = False,
         **kwargs,
     ) -> None:
 
@@ -1408,6 +1415,7 @@ def __init__(
         self.dataset_reference = dataset_reference if dataset_reference else {}
         self.delegate_to = delegate_to
         self.impersonation_chain = impersonation_chain
+        self.exists_ok = exists_ok
 
         super().__init__(**kwargs)
 
@@ -1425,7 +1433,7 @@ def execute(self, context) -> None:
                 dataset_id=self.dataset_id,
                 dataset_reference=self.dataset_reference,
                 location=self.location,
-                exists_ok=False,
+                exists_ok=self.exists_ok,
             )
         except Conflict:
             dataset_id = self.dataset_reference.get("datasetReference", {}).get("datasetId", self.dataset_id)
