diff --git a/cosmos/profiles/__init__.py b/cosmos/profiles/__init__.py index 1399c0a1d..e006f85ee 100644 --- a/cosmos/profiles/__init__.py +++ b/cosmos/profiles/__init__.py @@ -8,6 +8,7 @@ from .base import BaseProfileMapping from .bigquery.service_account_file import GoogleCloudServiceAccountFileProfileMapping +from .bigquery.service_account_keyfile_dict import GoogleCloudServiceAccountDictProfileMapping from .databricks.token import DatabricksTokenProfileMapping from .exasol.user_pass import ExasolUserPasswordProfileMapping from .postgres.user_pass import PostgresUserPasswordProfileMapping @@ -20,6 +21,7 @@ profile_mappings: list[Type[BaseProfileMapping]] = [ GoogleCloudServiceAccountFileProfileMapping, + GoogleCloudServiceAccountDictProfileMapping, DatabricksTokenProfileMapping, PostgresUserPasswordProfileMapping, RedshiftUserPasswordProfileMapping, diff --git a/cosmos/profiles/bigquery/__init__.py b/cosmos/profiles/bigquery/__init__.py index 3416eb300..e322c3af5 100644 --- a/cosmos/profiles/bigquery/__init__.py +++ b/cosmos/profiles/bigquery/__init__.py @@ -1,5 +1,9 @@ "BigQuery Airflow connection -> dbt profile mappings" from .service_account_file import GoogleCloudServiceAccountFileProfileMapping +from .service_account_keyfile_dict import GoogleCloudServiceAccountDictProfileMapping -__all__ = ["GoogleCloudServiceAccountFileProfileMapping"] +__all__ = [ + "GoogleCloudServiceAccountFileProfileMapping", + "GoogleCloudServiceAccountDictProfileMapping", +] diff --git a/cosmos/profiles/bigquery/service_account_keyfile_dict.py b/cosmos/profiles/bigquery/service_account_keyfile_dict.py new file mode 100644 index 000000000..e94c9dab6 --- /dev/null +++ b/cosmos/profiles/bigquery/service_account_keyfile_dict.py @@ -0,0 +1,43 @@ +"Maps Airflow GCP connections to dbt BigQuery profiles if they use a service account keyfile dict/json." +from __future__ import annotations + +from typing import Any + +from cosmos.profiles.base import BaseProfileMapping + + +class GoogleCloudServiceAccountDictProfileMapping(BaseProfileMapping): + """ + Maps Airflow GCP connections to dbt BigQuery profiles if they use a service account keyfile dict/json. + + https://docs.getdbt.com/reference/warehouse-setups/bigquery-setup#service-account-file + https://airflow.apache.org/docs/apache-airflow-providers-google/stable/connections/gcp.html + """ + + airflow_connection_type: str = "google_cloud_platform" + + required_fields = [ + "project", + "dataset", + "keyfile_dict", + ] + + airflow_param_mapping = { + "project": "extra.project", + "dataset": "dataset", + # multiple options for keyfile_dict param name because of older Airflow versions + "keyfile_dict": ["extra.keyfile_dict", "keyfile_dict", "extra__google_cloud_platform__keyfile_dict"], + } + + @property + def profile(self) -> dict[str, Any | None]: + """Generates profile. Defaults `threads` to 1.""" + return { + "type": "bigquery", + "method": "service-account-json", + "project": self.project, + "dataset": self.dataset, + "threads": self.profile_args.get("threads") or 1, + "keyfile_json": self.keyfile_dict, + **self.profile_args, + } diff --git a/docs/dbt/connections-profiles.rst b/docs/dbt/connections-profiles.rst index 59187b9d6..643625100 100644 --- a/docs/dbt/connections-profiles.rst +++ b/docs/dbt/connections-profiles.rst @@ -89,6 +89,14 @@ Service Account File :members: +Service Account Dict +~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: cosmos.profiles.bigquery.GoogleCloudServiceAccountDictProfileMapping + :undoc-members: + :members: + + Databricks ----------