Skip to content

Commit

Permalink
Make GoogleCloudServiceAccountDictProfileMapping dataset profile argu… (
Browse files Browse the repository at this point in the history
#1017)

This PR follows the methodology in
#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" <oliver.zheyi.ma@gmail.com> 
Original PR by @oliverrmaa: #839

Closes #837
  • Loading branch information
pankajastro authored and tatiana committed Jun 6, 2024
1 parent df09ea9 commit 09b6cbb
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 5 deletions.
5 changes: 3 additions & 2 deletions cosmos/profiles/bigquery/service_account_keyfile_dict.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ class GoogleCloudServiceAccountDictProfileMapping(BaseProfileMapping):

required_fields = [
"project",
"dataset",
"keyfile_json",
]

Expand All @@ -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."
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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,
}
Expand Down

0 comments on commit 09b6cbb

Please sign in to comment.