Keyfile dict can be dict not str (#29135)

* 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 <[email protected]>
apache · Jan 25, 2023 · cf90a1a · cf90a1a
1 parent bfbe2cb
commit cf90a1a
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 3 deletions.
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}