diff --git a/cosmos/profiles/bigquery/service_account_keyfile_dict.py b/cosmos/profiles/bigquery/service_account_keyfile_dict.py index 038b34153..89b92e812 100644 --- 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,11 +44,15 @@ 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 = { **self.mapped_params, "threads": 1, **self.profile_args, } + if "dataset" in self.profile_args: + profile["dataset"] = self.profile_args["dataset"] + + return profile @property def mock_profile(self) -> dict[str, Any | None]: diff --git a/tests/profiles/bigquery/test_bq_service_account_keyfile_dict.py b/tests/profiles/bigquery/test_bq_service_account_keyfile_dict.py index 4e56f5ba1..992fe9eb0 100755 --- 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({"keyfile_dict": sample_keyfile_dict}) profile_mapping = GoogleCloudServiceAccountDictProfileMapping(mock_bigquery_conn_with_dict, {}) assert not profile_mapping.can_claim_connection() @@ -87,6 +87,33 @@ def test_profile(mock_bigquery_conn_with_dict: Connection): assert profile_mapping.profile == expected +def test_profile_without_dataset(mock_bigquery_conn_with_dict: Connection): + mock_bigquery_conn_with_dict.extra = json.dumps( + { + "project": "my_project", + "keyfile_dict": json.dumps(sample_keyfile_dict), + } + ) + + profile_mapping = GoogleCloudServiceAccountDictProfileMapping(mock_bigquery_conn_with_dict, {}) + + # Expected profile does not include 'dataset' + expected = { + "type": "bigquery", + "method": "service-account-json", + "project": "my_project", + "dataset": None, + "threads": 1, + "keyfile_json": { + "type": "service_account", + "private_key_id": "{{ env_var('COSMOS_CONN_GOOGLE_CLOUD_PLATFORM_PRIVATE_KEY_ID') }}", + "private_key": "{{ env_var('COSMOS_CONN_GOOGLE_CLOUD_PLATFORM_PRIVATE_KEY') }}", + }, + } + + assert profile_mapping.profile == expected + + def test_mock_profile(mock_bigquery_conn_with_dict: Connection): """ Tests profile mock for keyfile_json is None. Giving keyfile_json a value will crash dbt ls. @@ -96,7 +123,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, }