Repository: astronomer/astronomer-cosmos (Python)

Issue:
GCP BigQuery service account json not supported in v0.7
We ran into this problem upgrading from v0.6.X -> v0.7.X using the default `google_cloud_platform` [airflow connection ](https://airflow.apache.org/docs/apache-airflow-providers-google/stable/connections/gcp.html#configuring-the-connection) that uses a keyfile json and not the keyfile path that is currently [supported](https://github.com/astronomer/astronomer-cosmos/blob/1ebee498b8ccc3a295b6c737a81d336ddcca6626/cosmos/profiles/bigquery/service_account_file.py#L9) in astronomer cosmos.

All of the profile var mapping of the json service account can be found in an older tag before the refactor of profiles in 0.7 here: https://github.com/astronomer/astronomer-cosmos/blob/astronomer-cosmos-v0.6.8/cosmos/providers/dbt/core/profiles/bigquery.py as a place to start from.

When I have time I'd be happy to take a stab at this, but opening up this issue in case others also run into this during an upgrade. 


Discussion: Hi, I was just going to open a similar issue :) We were not using the older version, but I guess we have the same problem: we store our Service Account JSON file as a secret in Google Secret Manager, so when we retrieve it, it is not a path to a file, but the actual file/JSON. I started testing with Cosmos 0.7.3 and so far did not manage getting it to somehow digest the SA keyfile as a json...
@tatiana Hi, I see you did the refactor probably, could you plz check this? Is there really no way to use SA keyfile_dict currently, or is there some other way? Thank you
@JoeSham @jbandoro I believe this may have been an issue introduced by the PR where we refactored profiles: https://github.com/astronomer/astronomer-cosmos/pull/271

To confirm if I understood the issue, Cosmos works as expected if you create a [Google Cloud Platform connection](https://airflow.apache.org/docs/apache-airflow-providers-google/stable/connections/gcp.html) specifying the `key_path` property but not if you specify `keyfile_dict` - is that the case?
@tatiana Hi, yes, correct. The [GoogleCloudServiceAccountFileProfileMapping](https://github.com/astronomer/astronomer-cosmos/blob/main/cosmos/profiles/bigquery/service_account_file.py) seems to only accept key_path as a path to SA keyfile, but there is no keyfile_dict or similar option. I also tested passing the SA keyfile_dict into the key_path param, in which case it passed in the ProfileMapping but then failed later on the fact that the keyfile (as given) does not exist
@JoeSham, the fix should be straightforward. The issue is that this:
https://github.com/astronomer/astronomer-cosmos/blob/astronomer-cosmos-v0.6.8/cosmos/providers/dbt/core/profiles/bigquery.py 

Was oversimplified as part of https://github.com/astronomer/astronomer-cosmos/pull/271:
https://github.com/astronomer/astronomer-cosmos/blob/main/cosmos/profiles/bigquery/service_account_file.py

We can add test cases to make sure we don't miss this feature in the future: https://github.com/astronomer/astronomer-cosmos/blob/main/tests/profiles/bigquery/test_bq_service_account_file.py

Would either of you have the capacity to work on the fix? Otherwise, we can take it
@JoeSham looked at the commits that are linked here, I think we should create a new profile mapping for the keyfile_dict/json. What do you think about that?
@jlaneve yea you are right, it was just a quick first draft. I now made another [quick 2nd draft](https://github.com/astronomer/astronomer-cosmos/compare/main...JoeSham:joe/bq_keyfile_dict?diff=unified), the way you recommended. Without tests so far. Not sure if I will have more time for it in the upcoming days though...
> @jlaneve yea you are right, it was just a quick first draft. I now made another [quick 2nd draft](https://github.com/astronomer/astronomer-cosmos/compare/main...JoeSham:joe/bq_keyfile_dict?diff=unified), the way you recommended. Without tests so far. Not sure if I will have more time for it in the upcoming days though...

That's fine! If you want to open a draft PR with what you have, I'm more than happy to jump in and add tests 🙂 
Ok great @jlaneve :) [Draft PR](https://github.com/astronomer/astronomer-cosmos/pull/352) opened



Fix (patch):
diff --git a/cosmos/operators/local.py b/cosmos/operators/local.py
index 8461bd07b..605fd4ce3 100644
--- a/cosmos/operators/local.py
+++ b/cosmos/operators/local.py
@@ -231,7 +231,6 @@ def run_command(
                 output_encoding=self.output_encoding,
                 cwd=tmp_project_dir,
             )
-
             self.exception_handling(result)
             self.store_compiled_sql(tmp_project_dir, context)
             if self.callback:
diff --git a/cosmos/profiles/__init__.py b/cosmos/profiles/__init__.py
index 48abdf482..2ebdd47cc 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
@@ -21,6 +22,7 @@
 
 profile_mappings: list[Type[BaseProfileMapping]] = [
     GoogleCloudServiceAccountFileProfileMapping,
+    GoogleCloudServiceAccountDictProfileMapping,
     DatabricksTokenProfileMapping,
     PostgresUserPasswordProfileMapping,
     RedshiftUserPasswordProfileMapping,
diff --git a/cosmos/profiles/base.py b/cosmos/profiles/base.py
index 081b9fb00..30ca8a8f1 100644
--- a/cosmos/profiles/base.py
+++ b/cosmos/profiles/base.py
@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-
 from logging import getLogger
 from typing import Any
 
@@ -42,18 +41,21 @@ def can_claim_connection(self) -> bool:
         if self.conn.conn_type != self.airflow_connection_type:
             return False
 
+        logger.info(dir(self.conn))
+        logger.info(self.conn.__dict__)
+
         for field in self.required_fields:
             try:
                 if not getattr(self, field):
                     logger.info(
-                        "Not using mapping %s because %s is not set",
+                        "1 Not using mapping %s because %s is not set",
                         self.__class__.__name__,
                         field,
                     )
                     return False
             except AttributeError:
                 logger.info(
-                    "Not using mapping %s because %s is not set",
+                    "2 Not using mapping %s because %s is not set",
                     self.__class__.__name__,
                     field,
                 )
@@ -97,7 +99,6 @@ def get_profile_file_contents(self, profile_name: str, target_name: str = "cosmo
                 "outputs": {target_name: profile_vars},
             }
         }
-
         return str(yaml.dump(profile_contents, indent=4))
 
     def get_dbt_value(self, name: str) -> Any:
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..f5bd87e20
--- /dev/null
+++ b/cosmos/profiles/bigquery/service_account_keyfile_dict.py
@@ -0,0 +1,48 @@
+"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",
+        # multiple options for dataset because of older Airflow versions
+        "dataset": ["extra.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 a GCP profile.
+        Even though the Airflow connection contains hard-coded Service account credentials,
+        we generate a temporary file and the DBT profile uses it.
+        """
+        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 7b62e9bdc..95f8d47b7 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
 ----------