From 9ebda052a3737c77c718354c113b1a402eaf7a5d Mon Sep 17 00:00:00 2001
From: Tatiana Al-Chueyr <tatiana.alchueyr@gmail.com>
Date: Tue, 18 Jul 2023 13:49:41 +0100
Subject: [PATCH] Store BQ profile keyfile_dict as an envvar

---
 cosmos/profiles/base.py                        |  4 +++-
 .../bigquery/service_account_keyfile_dict.py   |  6 +++++-
 .../test_bq_service_account_keyfile_dict.py    | 18 +++++++++++++++++-
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/cosmos/profiles/base.py b/cosmos/profiles/base.py
index 081b9fb00..bb5ef28d4 100644
--- a/cosmos/profiles/base.py
+++ b/cosmos/profiles/base.py
@@ -4,8 +4,8 @@
 """
 from __future__ import annotations
 
+import json
 from abc import ABC, abstractmethod
-
 from logging import getLogger
 from typing import Any
 
@@ -77,6 +77,8 @@ def env_vars(self) -> dict[str, str]:
         for field in self.secret_fields:
             env_var_name = self.get_env_var_name(field)
             value = self.get_dbt_value(field)
+            if isinstance(value, dict):
+                env_vars[env_var_name] = json.dumps(value)
             if value is not None:
                 env_vars[env_var_name] = str(value)
 
diff --git a/cosmos/profiles/bigquery/service_account_keyfile_dict.py b/cosmos/profiles/bigquery/service_account_keyfile_dict.py
index 597c92231..aba1e7f48 100644
--- a/cosmos/profiles/bigquery/service_account_keyfile_dict.py
+++ b/cosmos/profiles/bigquery/service_account_keyfile_dict.py
@@ -22,6 +22,10 @@ class GoogleCloudServiceAccountDictProfileMapping(BaseProfileMapping):
         "keyfile_dict",
     ]
 
+    secret_fields = [
+        "keyfile_dict",
+    ]
+
     airflow_param_mapping = {
         "project": "extra.project",
         # multiple options for dataset because of older Airflow versions
@@ -39,6 +43,6 @@ def profile(self) -> dict[str, Any | None]:
             "project": self.project,
             "dataset": self.dataset,
             "threads": self.profile_args.get("threads") or 1,
-            "keyfile_json": self.keyfile_dict,
+            "keyfile_json": self.get_env_var_format("keyfile_dict"),
             **self.profile_args,
         }
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 a763a91a2..091da3099 100644
--- a/tests/profiles/bigquery/test_bq_service_account_keyfile_dict.py
+++ b/tests/profiles/bigquery/test_bq_service_account_keyfile_dict.py
@@ -42,7 +42,23 @@ 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
+    # Remove the `dataset` key, which is mandatory
     mock_bigquery_conn_with_dict.extra = json.dumps({"project": "my_project", "keyfile_dict": {"key": "value"}})
     profile_mapping = GoogleCloudServiceAccountDictProfileMapping(mock_bigquery_conn_with_dict, {})
     assert not profile_mapping.can_claim_connection()
+
+
+def test_profile_env_vars(
+    mock_bigquery_conn_with_dict: Connection,
+) -> None:
+    """
+    Tests that the environment variables get set correctly.
+    """
+    profile_mapping = get_profile_mapping(
+        mock_bigquery_conn_with_dict.conn_id,
+    )
+    assert profile_mapping.env_vars == {
+        "COSMOS_CONN_GOOGLE_CLOUD_PLATFORM_KEYFILE_DICT": str(
+            mock_bigquery_conn_with_dict.extra_dejson["keyfile_dict"]
+        ),
+    }