-
Notifications
You must be signed in to change notification settings - Fork 159
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support OAuth authentication for Big Query (#431)
## Description Support OAuth authentication for Big Query ## Related Issue(s) closes #420 ## Breaking Change? None ## Checklist - [ ] I have made corresponding changes to the documentation (if required) - [X] I have added tests that prove my fix is effective or that my feature works --------- Co-authored-by: Monideep De <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
- Loading branch information
Showing
5 changed files
with
107 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,39 @@ | ||
"Maps Airflow GCP connections to dbt BigQuery profiles that uses oauth via gcloud, if they don't use key file or JSON." | ||
from __future__ import annotations | ||
|
||
from typing import Any | ||
|
||
from cosmos.profiles.base import BaseProfileMapping | ||
|
||
|
||
class GoogleCloudOauthProfileMapping(BaseProfileMapping): | ||
""" | ||
Maps Airflow GCP connections to dbt BigQuery profiles that uses oauth via gcloud, | ||
if they don't use key file or JSON. | ||
https://docs.getdbt.com/docs/core/connect-data-platform/bigquery-setup#oauth-via-gcloud | ||
https://airflow.apache.org/docs/apache-airflow-providers-google/stable/connections/gcp.html | ||
""" | ||
|
||
airflow_connection_type: str = "google_cloud_platform" | ||
|
||
required_fields = [ | ||
"project", | ||
"dataset", | ||
] | ||
|
||
airflow_param_mapping = { | ||
"project": "extra.project", | ||
"dataset": "extra.dataset", | ||
} | ||
|
||
@property | ||
def profile(self) -> dict[str, Any | None]: | ||
"Generates profile. Defaults `threads` to 1." | ||
return { | ||
**self.mapped_params, | ||
"type": "bigquery", | ||
"method": "oauth", | ||
"threads": 1, | ||
**self.profile_args, | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,61 @@ | ||
"Tests for the BigQuery profile." | ||
|
||
import json | ||
from unittest.mock import patch | ||
|
||
import pytest | ||
from airflow.models.connection import Connection | ||
|
||
from cosmos.profiles import get_automatic_profile_mapping | ||
from cosmos.profiles.bigquery.oauth import ( | ||
GoogleCloudOauthProfileMapping, | ||
) | ||
|
||
|
||
@pytest.fixture() | ||
def mock_bigquery_conn(request): | ||
""" | ||
Mocks and returns an Airflow BigQuery connection. | ||
""" | ||
extra = {"project": "my_project", "dataset": "my_dataset"} if not hasattr(request, "param") else request.param | ||
conn = Connection( | ||
conn_id="my_bigquery_connection", | ||
conn_type="google_cloud_platform", | ||
extra=json.dumps(extra), | ||
) | ||
|
||
with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn): | ||
yield conn | ||
|
||
|
||
def test_bigquery_mapping_selected(mock_bigquery_conn: Connection): | ||
profile_mapping = get_automatic_profile_mapping(mock_bigquery_conn.conn_id, {}) | ||
assert isinstance(profile_mapping, GoogleCloudOauthProfileMapping) | ||
|
||
|
||
@pytest.mark.parametrize( | ||
"mock_bigquery_conn", [{"project": "my_project"}, {"dataset": "my_dataset"}, {}], indirect=True | ||
) | ||
def test_connection_claiming_fails(mock_bigquery_conn: Connection) -> None: | ||
""" | ||
Tests that the BigQuery profile mapping claims the correct connection type. | ||
""" | ||
profile_mapping = GoogleCloudOauthProfileMapping(mock_bigquery_conn) | ||
assert not profile_mapping.can_claim_connection() | ||
|
||
|
||
def test_connection_claiming_succeeds(mock_bigquery_conn: Connection): | ||
profile_mapping = GoogleCloudOauthProfileMapping(mock_bigquery_conn, {}) | ||
assert profile_mapping.can_claim_connection() | ||
|
||
|
||
def test_profile(mock_bigquery_conn: Connection): | ||
profile_mapping = GoogleCloudOauthProfileMapping(mock_bigquery_conn, {}) | ||
expected = { | ||
"type": "bigquery", | ||
"method": "oauth", | ||
"project": "my_project", | ||
"dataset": "my_dataset", | ||
"threads": 1, | ||
} | ||
assert profile_mapping.profile == expected |