Keyfile dict can be dict not str (#29135)

dstandish · jedcunningham · web-flow · commit cf90a1a5673d · 2023-01-25T10:39:08.000-08:00
* Keyfile dict can be dict not str

This makes it much easier to define a connection with json.  We can do this:

```json
{"extra": {"keyfile_dict": {"foo": "bar", "private_key": "hi"}}}
```

Instead of this:

```json
{"extra": {"keyfile_dict": "{\"foo\": \"bar\", \"private_key\": \"hi\"}"}}
```

I.e. instead of json-dumping the contents of the key first, we can just paste the key as is into the "keyfile_dict" field in extra.

Co-authored-by: Jed Cunningham &lt;66968678+jedcunningham@users.noreply.github.com&gt;
diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py
@@ -242,10 +242,13 @@ def get_credentials_and_project_id(self) -> tuple[google.auth.credentials.Creden
 
         key_path: str | None = self._get_field("key_path", None)
         try:
-            keyfile_dict: str | None = self._get_field("keyfile_dict", None)
+            keyfile_dict: str | dict[str, str] | None = self._get_field("keyfile_dict", None)
             keyfile_dict_json: dict[str, str] | None = None
             if keyfile_dict:
-                keyfile_dict_json = json.loads(keyfile_dict)
+                if isinstance(keyfile_dict, dict):
+                    keyfile_dict_json = keyfile_dict
+                else:
+                    keyfile_dict_json = json.loads(keyfile_dict)
         except json.decoder.JSONDecodeError:
             raise AirflowException("Invalid key JSON.")
         key_secret_name: str | None = self._get_field("key_secret_name", None)
@@ -494,7 +497,7 @@ def provide_gcp_credential_file_as_context(self) -> Generator[str | None, None,
         file in ``GOOGLE_APPLICATION_CREDENTIALS`` environment variable.
         """
         key_path: str | None = self._get_field("key_path", None)
-        keyfile_dict: str | None = self._get_field("keyfile_dict", None)
+        keyfile_dict: str | dict[str, str] | None = self._get_field("keyfile_dict", None)
         if key_path and keyfile_dict:
             raise AirflowException(
                 "The `keyfile_dict` and `key_path` fields are mutually exclusive. "
@@ -507,6 +510,8 @@ def provide_gcp_credential_file_as_context(self) -> Generator[str | None, None,
                 yield key_path
         elif keyfile_dict:
             with tempfile.NamedTemporaryFile(mode="w+t") as conf_file:
+                if isinstance(keyfile_dict, dict):
+                    keyfile_dict = json.dumps(keyfile_dict)
                 conf_file.write(keyfile_dict)
                 conf_file.flush()
                 with patch_environ({CREDENTIALS: conf_file.name}):
diff --git a/tests/providers/google/common/hooks/test_base_google.py b/tests/providers/google/common/hooks/test_base_google.py
@@ -21,7 +21,9 @@
 import os
 import re
 from io import StringIO
+from pathlib import Path
 from unittest import mock
+from unittest.mock import patch
 
 import google.auth
 import pytest
@@ -34,6 +36,7 @@
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.utils.credentials_provider import _DEFAULT_SCOPES
 from airflow.providers.google.common.hooks import base_google as hook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from tests.providers.google.cloud.utils.base_gcp_mock import mock_base_gcp_hook_default_project_id
 
 default_creds_available = True
@@ -166,6 +169,37 @@ def assert_gcp_credential_file_in_env(_):
         ):
             assert_gcp_credential_file_in_env(self.instance)
 
+    def test_provide_gcp_credential_keyfile_dict_json(self):
+        """
+        Historically, keyfile_dict had to be str in the conn extra.  Now it
+        can be dict and this is verified here.
+        """
+        conn_dict = {
+            "extra": {
+                "keyfile_dict": {"foo": "bar", "private_key": "hi"},  # notice keyfile_dict is dict not str
+            }
+        }
+
+        @GoogleBaseHook.provide_gcp_credential_file
+        def assert_gcp_credential_file_in_env(instance):
+            assert Path(os.environ[CREDENTIALS]).read_text() == json.dumps(conn_dict["extra"]["keyfile_dict"])
+
+        with patch.dict("os.environ", AIRFLOW_CONN_MY_GOOGLE=json.dumps(conn_dict)):
+            # keyfile dict is handled in two different areas
+
+            hook = GoogleBaseHook("my_google")
+
+            # the first is in provide_gcp_credential_file
+            assert_gcp_credential_file_in_env(hook)
+
+            with patch("google.oauth2.service_account.Credentials.from_service_account_info") as m:
+                # the second is in get_credentials_and_project_id
+                hook.get_credentials_and_project_id()
+                m.assert_called_once_with(
+                    conn_dict["extra"]["keyfile_dict"],
+                    scopes=("https://www.googleapis.com/auth/cloud-platform",),
+                )
+
     def test_provide_gcp_credential_file_decorator_key_path(self):
         key_path = "/test/key-path"
         self.instance.extras = {"key_path": key_path}
