Make GoogleCloudServiceAccountDictProfileMapping dataset profile argu… (

astronomer#1017) This PR follows the methodology in astronomer#683 by modifying the GCP `GoogleCloudServiceAccountDictProfileMapping()` `profile_args` not require a dataset as a required argument. DAG RUN <img width="1473" alt="Screenshot 2024-06-05 at 6 50 04 PM" src="https://github.com/astronomer/astronomer-cosmos/assets/98807258/81b127b9-c5e5-4983-9efe-bbf00d81914f"> Co-authored-by: "Ollie Ma" <[email protected]> Original PR by @oliverrmaa: astronomer#839 Closes astronomer#837
arojasb3 · Jul 14, 2024 · 40de74a · 40de74a
1 parent 07ebfe3
commit 40de74a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/cosmos/profiles/bigquery/service_account_keyfile_dict.py b/cosmos/profiles/bigquery/service_account_keyfile_dict.py
@@ -22,7 +22,6 @@ class GoogleCloudServiceAccountDictProfileMapping(BaseProfileMapping):
 
     required_fields = [
         "project",
-        "dataset",
         "keyfile_json",
     ]
 
@@ -45,12 +44,14 @@ def profile(self) -> dict[str, Any | None]:
         Even though the Airflow connection contains hard-coded Service account credentials,
         we generate a temporary file and the DBT profile uses it.
         """
-        return {
+        profile_dict = {
             **self.mapped_params,
             "threads": 1,
             **self.profile_args,
         }
 
+        return self.filter_null(profile_dict)
+
     @property
     def mock_profile(self) -> dict[str, Any | None]:
         "Generates mock profile. Defaults `threads` to 1."

diff --git a/tests/profiles/bigquery/test_bq_service_account_keyfile_dict.py b/tests/profiles/bigquery/test_bq_service_account_keyfile_dict.py
@@ -64,8 +64,8 @@ def test_connection_claiming_succeeds(mock_bigquery_conn_with_dict: Connection):
 
 
 def test_connection_claiming_fails(mock_bigquery_conn_with_dict: Connection):
-    # Remove the `dataset` key, which is mandatory
-    mock_bigquery_conn_with_dict.extra = json.dumps({"project": "my_project", "keyfile_dict": sample_keyfile_dict})
+    # Remove the `project` key, which is mandatory
+    mock_bigquery_conn_with_dict.extra = json.dumps({"dataset": "my_dataset", "keyfile_dict": sample_keyfile_dict})
     profile_mapping = GoogleCloudServiceAccountDictProfileMapping(mock_bigquery_conn_with_dict, {})
     assert not profile_mapping.can_claim_connection()
 
@@ -96,7 +96,6 @@ def test_mock_profile(mock_bigquery_conn_with_dict: Connection):
         "type": "bigquery",
         "method": "service-account-json",
         "project": "mock_value",
-        "dataset": "mock_value",
         "threads": 1,
         "keyfile_json": None,
     }