assign "datasetReference" attribute to dataset_reference dict. by def…

…ault if not already set in create_empty_dataset method of bigquery hook (#28782)
apache · Jan 9, 2023 · c0b2fcf · c0b2fcf
1 parent 71306b3
commit c0b2fcf
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py
@@ -425,7 +425,10 @@ def create_empty_dataset(
             https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#resource
         :param exists_ok: If ``True``, ignore "already exists" errors when creating the dataset.
         """
-        dataset_reference = dataset_reference or {"datasetReference": {}}
+        dataset_reference = dataset_reference or {}
+
+        if "datasetReference" not in dataset_reference:
+            dataset_reference["datasetReference"] = {}
 
         for param, value in zip(["datasetId", "projectId"], [dataset_id, project_id]):
             specified_param = dataset_reference["datasetReference"].get(param)

diff --git a/tests/providers/google/cloud/hooks/test_bigquery.py b/tests/providers/google/cloud/hooks/test_bigquery.py
@@ -1487,6 +1487,22 @@ def test_create_empty_dataset_use_values_from_object(self, mock_client, mock_dat
             dataset=api_repr.return_value, exists_ok=True
         )
 
+    @mock.patch("airflow.providers.google.cloud.hooks.bigquery.Dataset")
+    @mock.patch("airflow.providers.google.cloud.hooks.bigquery.Client")
+    def test_create_empty_dataset_without_datasetreference_key(self, mock_client, mock_dataset):
+        dataset = {"defaultTableExpirationMs": str(1000 * 60 * 60 * 24 * 30)}
+        dataset_copy = dataset.copy()
+        self.hook.create_empty_dataset(
+            dataset_reference=dataset, dataset_id="DATASET_ID", project_id="PROJECT_ID"
+        )
+        assert dataset["defaultTableExpirationMs"] == dataset_copy["defaultTableExpirationMs"]
+        assert dataset["datasetReference"] == {"datasetId": "DATASET_ID", "projectId": "PROJECT_ID"}
+        api_repr = mock_dataset.from_api_repr
+        api_repr.assert_called_once_with(dataset)
+        mock_client.return_value.create_dataset.assert_called_once_with(
+            dataset=api_repr.return_value, exists_ok=True
+        )
+
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.Client")
     def test_get_dataset(self, mock_client):
         _expected_result = {