From 9a22ff937e58712949e93fc08aadbfbd7e821510 Mon Sep 17 00:00:00 2001
From: JoeSham <j.samanek@seznam.cz>
Date: Wed, 28 Jun 2023 14:14:01 +0200
Subject: [PATCH] Add support for BQ SA keyfile_dict

Related issue: https://github.com/astronomer/astronomer-cosmos/issues/350
---
 .../profiles/bigquery/service_account_file.py | 28 +++++++++++++++----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/cosmos/profiles/bigquery/service_account_file.py b/cosmos/profiles/bigquery/service_account_file.py
index 96fe56b00..2646aef36 100644
--- a/cosmos/profiles/bigquery/service_account_file.py
+++ b/cosmos/profiles/bigquery/service_account_file.py
@@ -19,24 +19,40 @@ class GoogleCloudServiceAccountFileProfileMapping(BaseProfileMapping):
     required_fields = [
         "project",
         "dataset",
-        "keyfile",
+        # plus one of: keyfile or keyfile_dict
     ]
 
     airflow_param_mapping = {
         "project": "extra.project",
         "dataset": "dataset",
-        "keyfile": "extra.key_path",
+        # multiple options for keyfile/keyfile_dict param name because of older Airflow versions
+        "keyfile": ["key_path", "extra__google_cloud_platform__key_path", "extra.key_path"],
+        "keyfile_dict": ["keyfile_dict", "extra__google_cloud_platform__keyfile_dict", "extra.keyfile_dict"],
     }
 
     @property
     def profile(self) -> dict[str, Any | None]:
-        "Generates profile. Defaults `threads` to 1."
-        return {
+        """
+            Generates profile. Defaults `threads` to 1.
+            Profile can either use keyfile as path to json file, or keyfile_dict as directly the json dict
+        """
+
+        profile_dict = {
             "type": "bigquery",
-            "method": "service-account",
             "project": self.project,
             "dataset": self.dataset,
             "threads": self.profile_args.get("threads") or 1,
-            "keyfile": self.keyfile,
             **self.profile_args,
         }
+
+        # use keyfile_dict if it exists, otherwise use keyfile
+        try:
+            profile_dict["keyfile_json"] = self.keyfile_dict
+            profile_dict["method"] = "service-account-json"
+        except AttributeError:
+            profile_dict["keyfile"] = self.keyfile
+            profile_dict["method"] = "service-account"
+
+        
+        return profile_dict
+