diff --git a/airflow/providers/google/ads/hooks/ads.py b/airflow/providers/google/ads/hooks/ads.py index dcda4f77482b9..aa01df9bdebf1 100644 --- a/airflow/providers/google/ads/hooks/ads.py +++ b/airflow/providers/google/ads/hooks/ads.py @@ -32,6 +32,7 @@ from airflow import AirflowException from airflow.compat.functools import cached_property from airflow.hooks.base import BaseHook +from airflow.providers.google.common.hooks.base_google import get_field class GoogleAdsHook(BaseHook): @@ -200,8 +201,10 @@ def _update_config_with_secret(self, secrets_temp: IO[str]) -> None: Updates google ads config with file path of the temp file containing the secret Note, the secret must be passed as a file path for Google Ads API """ - secret_conn = self.get_connection(self.gcp_conn_id) - secret = secret_conn.extra_dejson["extra__google_cloud_platform__keyfile_dict"] + extras = self.get_connection(self.gcp_conn_id).extra_dejson + secret = get_field(extras, 'keyfile_dict') + if not secret: + raise KeyError("secret_conn.extra_dejson does not contain keyfile_dict") secrets_temp.write(secret) secrets_temp.flush() diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py index 19260c6ac6a10..5295aab078bf3 100644 --- a/airflow/providers/google/cloud/hooks/bigquery.py +++ b/airflow/providers/google/cloud/hooks/bigquery.py @@ -58,7 +58,7 @@ from airflow.exceptions import AirflowException from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.providers.google.common.consts import CLIENT_INFO -from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook +from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook, get_field from airflow.utils.helpers import convert_camel_to_snake from airflow.utils.log.logging_mixin import LoggingMixin @@ -156,15 +156,14 @@ def get_sqlalchemy_engine(self, engine_kwargs=None): """ if engine_kwargs is None: engine_kwargs = {} - connection = self.get_connection(self.gcp_conn_id) - if connection.extra_dejson.get("extra__google_cloud_platform__key_path"): - credentials_path = connection.extra_dejson['extra__google_cloud_platform__key_path'] + extras = self.get_connection(self.gcp_conn_id).extra_dejson + credentials_path = get_field(extras, 'key_path') + if credentials_path: return create_engine(self.get_uri(), credentials_path=credentials_path, **engine_kwargs) - elif connection.extra_dejson.get("extra__google_cloud_platform__keyfile_dict"): - credential_file_content = json.loads( - connection.extra_dejson["extra__google_cloud_platform__keyfile_dict"] - ) - return create_engine(self.get_uri(), credentials_info=credential_file_content, **engine_kwargs) + keyfile_dict = get_field(extras, 'keyfile_dict') + if keyfile_dict: + keyfile_content = keyfile_dict if isinstance(keyfile_dict, dict) else json.loads(keyfile_dict) + return create_engine(self.get_uri(), credentials_info=keyfile_content, **engine_kwargs) try: # 1. If the environment variable GOOGLE_APPLICATION_CREDENTIALS is set # ADC uses the service account key or configuration file that the variable points to. @@ -175,9 +174,7 @@ def get_sqlalchemy_engine(self, engine_kwargs=None): self.log.error(e) raise AirflowException( "For now, we only support instantiating SQLAlchemy engine by" - " using ADC" - ", extra__google_cloud_platform__key_path" - "and extra__google_cloud_platform__keyfile_dict" + " using ADC or extra fields `key_path` and `keyfile_dict`." ) def get_records(self, sql, parameters=None): diff --git a/airflow/providers/google/cloud/hooks/cloud_sql.py b/airflow/providers/google/cloud/hooks/cloud_sql.py index 4ea8cc8d11a3d..ba85baae853cf 100644 --- a/airflow/providers/google/cloud/hooks/cloud_sql.py +++ b/airflow/providers/google/cloud/hooks/cloud_sql.py @@ -48,7 +48,7 @@ # For requests that are "retriable" from airflow.hooks.base import BaseHook from airflow.models import Connection -from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +from airflow.providers.google.common.hooks.base_google import GoogleBaseHook, get_field from airflow.providers.mysql.hooks.mysql import MySqlHook from airflow.providers.postgres.hooks.postgres import PostgresHook from airflow.utils.log.logging_mixin import LoggingMixin @@ -375,9 +375,6 @@ def _wait_for_operation_to_complete(self, project_id: str, operation_name: str) "https://storage.googleapis.com/cloudsql-proxy/{}/cloud_sql_proxy.{}.{}" ) -GCP_CREDENTIALS_KEY_PATH = "extra__google_cloud_platform__key_path" -GCP_CREDENTIALS_KEYFILE_DICT = "extra__google_cloud_platform__keyfile_dict" - class CloudSqlProxyRunner(LoggingMixin): """ @@ -484,15 +481,16 @@ def _download_sql_proxy_if_needed(self) -> None: self.sql_proxy_was_downloaded = True def _get_credential_parameters(self) -> list[str]: - connection = GoogleBaseHook.get_connection(conn_id=self.gcp_conn_id) - - if connection.extra_dejson.get(GCP_CREDENTIALS_KEY_PATH): - credential_params = ['-credential_file', connection.extra_dejson[GCP_CREDENTIALS_KEY_PATH]] - elif connection.extra_dejson.get(GCP_CREDENTIALS_KEYFILE_DICT): - credential_file_content = json.loads(connection.extra_dejson[GCP_CREDENTIALS_KEYFILE_DICT]) + extras = GoogleBaseHook.get_connection(conn_id=self.gcp_conn_id).extra_dejson + key_path = get_field(extras, 'key_path') + keyfile_dict = get_field(extras, 'keyfile_dict') + if key_path: + credential_params = ['-credential_file', key_path] + elif keyfile_dict: + keyfile_content = keyfile_dict if isinstance(keyfile_dict, dict) else json.loads(keyfile_dict) self.log.info("Saving credentials to %s", self.credentials_path) with open(self.credentials_path, "w") as file: - json.dump(credential_file_content, file) + json.dump(keyfile_content, file) credential_params = ['-credential_file', self.credentials_path] else: self.log.info( @@ -504,7 +502,7 @@ def _get_credential_parameters(self) -> list[str]: credential_params = [] if not self.instance_specification: - project_id = connection.extra_dejson.get('extra__google_cloud_platform__project') + project_id = get_field(extras, 'project') if self.project_id: project_id = self.project_id if not project_id: diff --git a/airflow/providers/google/cloud/operators/cloud_sql.py b/airflow/providers/google/cloud/operators/cloud_sql.py index a291018f629fa..0193b849707ff 100644 --- a/airflow/providers/google/cloud/operators/cloud_sql.py +++ b/airflow/providers/google/cloud/operators/cloud_sql.py @@ -28,6 +28,7 @@ from airflow.providers.google.cloud.hooks.cloud_sql import CloudSQLDatabaseHook, CloudSQLHook from airflow.providers.google.cloud.links.cloud_sql import CloudSQLInstanceDatabaseLink, CloudSQLInstanceLink from airflow.providers.google.cloud.utils.field_validator import GcpBodyFieldValidator +from airflow.providers.google.common.hooks.base_google import get_field from airflow.providers.google.common.links.storage import FileDetailsLink from airflow.providers.mysql.hooks.mysql import MySqlHook from airflow.providers.postgres.hooks.postgres import PostgresHook @@ -1092,9 +1093,7 @@ def execute(self, context: Context): hook = CloudSQLDatabaseHook( gcp_cloudsql_conn_id=self.gcp_cloudsql_conn_id, gcp_conn_id=self.gcp_conn_id, - default_gcp_project_id=self.gcp_connection.extra_dejson.get( - 'extra__google_cloud_platform__project' - ), + default_gcp_project_id=get_field(self.gcp_connection.extra_dejson, 'project'), ) hook.validate_ssl_certs() connection = hook.create_connection() diff --git a/airflow/providers/google/cloud/utils/credentials_provider.py b/airflow/providers/google/cloud/utils/credentials_provider.py index 0041449b2517d..3eeea0f66fe9f 100644 --- a/airflow/providers/google/cloud/utils/credentials_provider.py +++ b/airflow/providers/google/cloud/utils/credentials_provider.py @@ -60,16 +60,14 @@ def build_gcp_conn( :return: String representing Airflow connection. """ conn = "google-cloud-platform://?{}" - extras = "extra__google_cloud_platform" - query_params = {} if key_file_path: - query_params[f"{extras}__key_path"] = key_file_path + query_params["key_path"] = key_file_path if scopes: scopes_string = ",".join(scopes) - query_params[f"{extras}__scope"] = scopes_string + query_params["scope"] = scopes_string if project_id: - query_params[f"{extras}__projects"] = project_id + query_params["projects"] = project_id query = urlencode(query_params) return conn.format(query) diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index de79f7887ce9e..98712ffe64505 100644 --- a/airflow/providers/google/common/hooks/base_google.py +++ b/airflow/providers/google/common/hooks/base_google.py @@ -127,6 +127,19 @@ def __init__(self): RT = TypeVar('RT') +def get_field(extras: dict, field_name: str): + """Get field from extra, first checking short name, then for backcompat we check for prefixed name.""" + if field_name.startswith('extra__'): + raise ValueError( + f"Got prefixed name {field_name}; please remove the 'extra__google_cloud_platform__' prefix " + "when using this method." + ) + if field_name in extras: + return extras[field_name] or None + prefixed_name = f"extra__google_cloud_platform__{field_name}" + return extras.get(prefixed_name) or None + + class GoogleBaseHook(BaseHook): """ A base hook for Google cloud-related hooks. Google cloud has a shared REST @@ -179,25 +192,17 @@ def get_connection_form_widgets() -> dict[str, Any]: from wtforms.validators import NumberRange return { - "extra__google_cloud_platform__project": StringField( - lazy_gettext('Project Id'), widget=BS3TextFieldWidget() - ), - "extra__google_cloud_platform__key_path": StringField( - lazy_gettext('Keyfile Path'), widget=BS3TextFieldWidget() - ), - "extra__google_cloud_platform__keyfile_dict": PasswordField( - lazy_gettext('Keyfile JSON'), widget=BS3PasswordFieldWidget() - ), - "extra__google_cloud_platform__scope": StringField( - lazy_gettext('Scopes (comma separated)'), widget=BS3TextFieldWidget() - ), - "extra__google_cloud_platform__key_secret_name": StringField( + "project": StringField(lazy_gettext('Project Id'), widget=BS3TextFieldWidget()), + "key_path": StringField(lazy_gettext('Keyfile Path'), widget=BS3TextFieldWidget()), + "keyfile_dict": PasswordField(lazy_gettext('Keyfile JSON'), widget=BS3PasswordFieldWidget()), + "scope": StringField(lazy_gettext('Scopes (comma separated)'), widget=BS3TextFieldWidget()), + "key_secret_name": StringField( lazy_gettext('Keyfile Secret Name (in GCP Secret Manager)'), widget=BS3TextFieldWidget() ), - "extra__google_cloud_platform__key_secret_project_id": StringField( + "key_secret_project_id": StringField( lazy_gettext('Keyfile Secret Project Id (in GCP Secret Manager)'), widget=BS3TextFieldWidget() ), - "extra__google_cloud_platform__num_retries": IntegerField( + "num_retries": IntegerField( lazy_gettext('Number of Retries'), validators=[NumberRange(min=0)], widget=BS3TextFieldWidget(), @@ -325,11 +330,7 @@ def _get_field(self, f: str, default: Any = None) -> Any: to the hook page, which allow admins to specify service_account, key_path, etc. They get formatted as shown below. """ - long_f = f'extra__google_cloud_platform__{f}' - if hasattr(self, 'extras') and long_f in self.extras: - return self.extras[long_f] - else: - return default + return hasattr(self, 'extras') and get_field(self.extras, f) or default @property def project_id(self) -> str | None: diff --git a/docs/apache-airflow-providers-amazon/secrets-backends/aws-secrets-manager.rst b/docs/apache-airflow-providers-amazon/secrets-backends/aws-secrets-manager.rst index b550e60e66be2..68127befa007a 100644 --- a/docs/apache-airflow-providers-amazon/secrets-backends/aws-secrets-manager.rst +++ b/docs/apache-airflow-providers-amazon/secrets-backends/aws-secrets-manager.rst @@ -147,5 +147,5 @@ For connecting to a google cloud conn, all the fields must be in the extra field .. code-block:: ini - {'extra__google_cloud_platform__key_path': '/opt/airflow/service_account.json', - 'extra__google_cloud_platform__scope': 'https://www.googleapis.com/auth/devstorage.read_only'} + {'key_path': '/opt/airflow/service_account.json', + 'scope': 'https://www.googleapis.com/auth/devstorage.read_only'} diff --git a/docs/apache-airflow-providers-google/connections/gcp.rst b/docs/apache-airflow-providers-google/connections/gcp.rst index ba87995e9eee8..92f1f37e8ac23 100644 --- a/docs/apache-airflow-providers-google/connections/gcp.rst +++ b/docs/apache-airflow-providers-google/connections/gcp.rst @@ -124,21 +124,27 @@ Number of Retries * query parameters contains information specific to this type of connection. The following keys are accepted: - * ``extra__google_cloud_platform__project`` - Project Id - * ``extra__google_cloud_platform__key_path`` - Keyfile Path - * ``extra__google_cloud_platform__keyfile_dict`` - Keyfile JSON - * ``extra__google_cloud_platform__key_secret_name`` - Secret name which holds Keyfile JSON - * ``extra__google_cloud_platform__key_secret_project_id`` - Project Id which holds Keyfile JSON - * ``extra__google_cloud_platform__scope`` - Scopes - * ``extra__google_cloud_platform__num_retries`` - Number of Retries + * ``project`` - Project Id + * ``key_path`` - Keyfile Path + * ``keyfile_dict`` - Keyfile JSON + * ``key_secret_name`` - Secret name which holds Keyfile JSON + * ``key_secret_project_id`` - Project Id which holds Keyfile JSON + * ``scope`` - Scopes + * ``num_retries`` - Number of Retries Note that all components of the URI should be URL-encoded. - For example: + For example, with URI format: .. code-block:: bash - export AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT='google-cloud-platform://?extra__google_cloud_platform__key_path=%2Fkeys%2Fkey.json&extra__google_cloud_platform__scope=https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fcloud-platform&extra__google_cloud_platform__project=airflow&extra__google_cloud_platform__num_retries=5' + export AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT='google-cloud-platform://?key_path=%2Fkeys%2Fkey.json&scope=https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fcloud-platform&project=airflow&num_retries=5' + + And using JSON format: + + .. code-block:: bash + + export AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT='{"conn_type": "google-cloud-platform", "key_path": "/keys/key.json", "scope": "https://www.googleapis.com/auth/cloud-platform", "project": "airflow", "num_retries": 5}' .. _howto/connection:gcp:impersonation: diff --git a/docs/apache-airflow-providers-google/connections/gcp_ssh.rst b/docs/apache-airflow-providers-google/connections/gcp_ssh.rst index 9391f2152b1a0..c23358d5b790b 100644 --- a/docs/apache-airflow-providers-google/connections/gcp_ssh.rst +++ b/docs/apache-airflow-providers-google/connections/gcp_ssh.rst @@ -46,12 +46,12 @@ Extra (optional) connection. The following parameters are supported in addition to those describing the Google Cloud connection. - * ``extra__google_cloud_platform__instance_name`` - The name of the Compute Engine instance. - * ``extra__google_cloud_platform__zone`` - The zone of the Compute Engine instance. - * ``extra__google_cloud_platform__use_internal_ip`` - Whether to connect using internal IP. - * ``extra__google_cloud_platform__use_iap_tunnel`` - Whether to connect through IAP tunnel. - * ``extra__google_cloud_platform__use_oslogin`` - Whether to manage keys using OsLogin API. If false, keys are managed using instance metadata. - * ``extra__google_cloud_platform__expire_time`` - The maximum amount of time in seconds before the private key expires. + * ``instance_name`` - The name of the Compute Engine instance. + * ``zone`` - The zone of the Compute Engine instance. + * ``use_internal_ip`` - Whether to connect using internal IP. + * ``use_iap_tunnel`` - Whether to connect through IAP tunnel. + * ``use_oslogin`` - Whether to manage keys using OsLogin API. If false, keys are managed using instance metadata. + * ``expire_time`` - The maximum amount of time in seconds before the private key expires. Environment variable @@ -64,9 +64,9 @@ For example: .. code-block:: bash export AIRFLOW_CONN_GOOGLE_CLOUD_SQL_DEFAULT="gcpssh://conn-user@conn-host?\ - extra__google_cloud_platform__instance_name=conn-instance-name&\ - extra__google_cloud_platform__zone=zone&\ - extra__google_cloud_platform__use_internal_ip=True&\ - extra__google_cloud_platform__use_iap_tunnel=True&\ - extra__google_cloud_platform__use_oslogin=False&\ - extra__google_cloud_platform__expire_time=4242" + instance_name=conn-instance-name&\ + zone=zone&\ + use_internal_ip=True&\ + use_iap_tunnel=True&\ + use_oslogin=False&\ + expire_time=4242" diff --git a/docs/apache-airflow/security/secrets/secrets-backend/local-filesystem-secrets-backend.rst b/docs/apache-airflow/security/secrets/secrets-backend/local-filesystem-secrets-backend.rst index b9b99250f1320..aaa71daf20af6 100644 --- a/docs/apache-airflow/security/secrets/secrets-backend/local-filesystem-secrets-backend.rst +++ b/docs/apache-airflow/security/secrets/secrets-backend/local-filesystem-secrets-backend.rst @@ -108,7 +108,7 @@ raise an exception. The following is a sample file. .. code-block:: text mysql_conn_id=mysql://log:password@13.1.21.1:3306/mysqldbrd - google_custom_key=google-cloud-platform://?extra__google_cloud_platform__key_path=%2Fkeys%2Fkey.json + google_custom_key=google-cloud-platform://?key_path=%2Fkeys%2Fkey.json Storing and Retrieving Variables """""""""""""""""""""""""""""""" diff --git a/tests/always/test_connection.py b/tests/always/test_connection.py index b55452731a620..95aa1e24be185 100644 --- a/tests/always/test_connection.py +++ b/tests/always/test_connection.py @@ -275,10 +275,8 @@ def test_connection_extra_with_encryption_rotate_fernet_key(self): description='no schema', ), UriTestCaseConfig( - test_conn_uri='google-cloud-platform://?extra__google_cloud_platform__key_' - 'path=%2Fkeys%2Fkey.json&extra__google_cloud_platform__scope=' - 'https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fcloud-platform&extra' - '__google_cloud_platform__project=airflow', + test_conn_uri='google-cloud-platform://?key_path=%2Fkeys%2Fkey.json&scope=' + 'https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fcloud-platform&project=airflow', test_conn_attributes=dict( conn_type='google_cloud_platform', host='', @@ -287,9 +285,9 @@ def test_connection_extra_with_encryption_rotate_fernet_key(self): password=None, port=None, extra_dejson=dict( - extra__google_cloud_platform__key_path='/keys/key.json', - extra__google_cloud_platform__scope='https://www.googleapis.com/auth/cloud-platform', - extra__google_cloud_platform__project='airflow', + key_path='/keys/key.json', + scope='https://www.googleapis.com/auth/cloud-platform', + project='airflow', ), ), description='with underscore', diff --git a/tests/always/test_secrets_local_filesystem.py b/tests/always/test_secrets_local_filesystem.py index bb2c40abf708a..2ed981413e22a 100644 --- a/tests/always/test_secrets_local_filesystem.py +++ b/tests/always/test_secrets_local_filesystem.py @@ -257,8 +257,8 @@ def test_missing_file(self, mock_exists): extra_dejson: arbitrary_dict: a: b - extra__google_cloud_platform__keyfile_dict: '{"a": "b"}' - extra__google_cloud_platform__keyfile_path: asaa""", + keyfile_dict: '{"a": "b"}' + keyfile_path: asaa""", { "conn_a": {'conn_type': 'mysql', 'host': 'hosta'}, "conn_b": { @@ -270,8 +270,8 @@ def test_missing_file(self, mock_exists): 'port': 1234, 'extra_dejson': { 'arbitrary_dict': {"a": "b"}, - 'extra__google_cloud_platform__keyfile_dict': '{"a": "b"}', - 'extra__google_cloud_platform__keyfile_path': 'asaa', + 'keyfile_dict': '{"a": "b"}', + 'keyfile_path': 'asaa', }, }, }, @@ -314,14 +314,14 @@ def test_yaml_file_should_load_connection(self, file_content, expected_attrs_dic password: None port: 1234 extra_dejson: - extra__google_cloud_platform__keyfile_dict: + keyfile_dict: a: b - extra__google_cloud_platform__key_path: xxx + key_path: xxx """, { "conn_d": { - "extra__google_cloud_platform__keyfile_dict": {"a": "b"}, - "extra__google_cloud_platform__key_path": "xxx", + "keyfile_dict": {"a": "b"}, + "key_path": "xxx", } }, ), @@ -334,9 +334,9 @@ def test_yaml_file_should_load_connection(self, file_content, expected_attrs_dic login: Login password: None port: 1234 - extra: '{\"extra__google_cloud_platform__keyfile_dict\": {\"a\": \"b\"}}' + extra: '{\"keyfile_dict\": {\"a\": \"b\"}}' """, - {"conn_d": {"extra__google_cloud_platform__keyfile_dict": {"a": "b"}}}, + {"conn_d": {"keyfile_dict": {"a": "b"}}}, ), ], ) diff --git a/tests/providers/google/ads/hooks/test_ads.py b/tests/providers/google/ads/hooks/test_ads.py index 6bea65a285b07..b06c2a63b6fa0 100644 --- a/tests/providers/google/ads/hooks/test_ads.py +++ b/tests/providers/google/ads/hooks/test_ads.py @@ -27,7 +27,7 @@ ADS_CLIENT = {"key": "value"} SECRET = "secret" EXTRAS = { - "extra__google_cloud_platform__keyfile_dict": SECRET, + "keyfile_dict": SECRET, "google_ads_client": ADS_CLIENT, } diff --git a/tests/providers/google/cloud/hooks/test_cloud_sql.py b/tests/providers/google/cloud/hooks/test_cloud_sql.py index 9e7bdaa5c79c8..3e349c5efc775 100644 --- a/tests/providers/google/cloud/hooks/test_cloud_sql.py +++ b/tests/providers/google/cloud/hooks/test_cloud_sql.py @@ -1015,9 +1015,9 @@ def setUp(self, m): "https://www.googleapis.com/auth/cloud-platform", ] conn_extra = { - "extra__google_cloud_platform__scope": ",".join(scopes), - "extra__google_cloud_platform__project": "your-gcp-project", - "extra__google_cloud_platform__key_path": '/var/local/google_cloud_default.json', + "scope": ",".join(scopes), + "project": "your-gcp-project", + "key_path": '/var/local/google_cloud_default.json', } conn_extra_json = json.dumps(conn_extra) self.connection.set_extra(conn_extra_json) diff --git a/tests/providers/google/cloud/hooks/test_compute_ssh.py b/tests/providers/google/cloud/hooks/test_compute_ssh.py index b9920a4aafaf0..a1d02828d28e8 100644 --- a/tests/providers/google/cloud/hooks/test_compute_ssh.py +++ b/tests/providers/google/cloud/hooks/test_compute_ssh.py @@ -297,12 +297,12 @@ def test_read_configuration_from_connection(self): host="conn-host", extra=json.dumps( { - "extra__google_cloud_platform__instance_name": "conn-instance-name", - "extra__google_cloud_platform__zone": "zone", - "extra__google_cloud_platform__use_internal_ip": True, - "extra__google_cloud_platform__use_iap_tunnel": True, - "extra__google_cloud_platform__use_oslogin": False, - "extra__google_cloud_platform__expire_time": 4242, + "instance_name": "conn-instance-name", + "zone": "zone", + "use_internal_ip": True, + "use_iap_tunnel": True, + "use_oslogin": False, + "expire_time": 4242, } ), ) diff --git a/tests/providers/google/cloud/operators/test_kubernetes_engine.py b/tests/providers/google/cloud/operators/test_kubernetes_engine.py index 3df68a210a289..5edacf8ab5554 100644 --- a/tests/providers/google/cloud/operators/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/operators/test_kubernetes_engine.py @@ -204,13 +204,7 @@ def test_template_fields(self): @mock.patch.dict(os.environ, {}) @mock.patch( "airflow.hooks.base.BaseHook.get_connections", - return_value=[ - Connection( - extra=json.dumps( - {"extra__google_cloud_platform__keyfile_dict": '{"private_key": "r4nd0m_k3y"}'} - ) - ) - ], + return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], ) @mock.patch('airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute') @mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook') @@ -244,13 +238,7 @@ def test_execute(self, file_mock, mock_execute_in_subprocess, mock_gcp_hook, exe @mock.patch.dict(os.environ, {}) @mock.patch( "airflow.hooks.base.BaseHook.get_connections", - return_value=[ - Connection( - extra=json.dumps( - {"extra__google_cloud_platform__keyfile_dict": '{"private_key": "r4nd0m_k3y"}'} - ) - ) - ], + return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], ) @mock.patch('airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute') @mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook') @@ -300,13 +288,7 @@ def test_config_file_throws_error(self): @mock.patch.dict(os.environ, {}) @mock.patch( "airflow.hooks.base.BaseHook.get_connections", - return_value=[ - Connection( - extra=json.dumps( - {"extra__google_cloud_platform__keyfile_dict": '{"private_key": "r4nd0m_k3y"}'} - ) - ) - ], + return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], ) @mock.patch('airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute') @mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook') @@ -344,13 +326,7 @@ def test_execute_with_internal_ip( @mock.patch.dict(os.environ, {}) @mock.patch( "airflow.hooks.base.BaseHook.get_connections", - return_value=[ - Connection( - extra=json.dumps( - {"extra__google_cloud_platform__keyfile_dict": '{"private_key": "r4nd0m_k3y"}'} - ) - ) - ], + return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], ) @mock.patch('airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute') @mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook') @@ -388,13 +364,7 @@ def test_execute_with_impersonation_service_account( @mock.patch.dict(os.environ, {}) @mock.patch( "airflow.hooks.base.BaseHook.get_connections", - return_value=[ - Connection( - extra=json.dumps( - {"extra__google_cloud_platform__keyfile_dict": '{"private_key": "r4nd0m_k3y"}'} - ) - ) - ], + return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], ) @mock.patch('airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute') @mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook') @@ -432,13 +402,7 @@ def test_execute_with_impersonation_service_chain_one_element( @mock.patch.dict(os.environ, {}) @mock.patch( "airflow.hooks.base.BaseHook.get_connections", - return_value=[ - Connection( - extra=json.dumps( - {"extra__google_cloud_platform__keyfile_dict": '{"private_key": "r4nd0m_k3y"}'} - ) - ) - ], + return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], ) @mock.patch('airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute') @mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook') diff --git a/tests/providers/google/cloud/utils/base_gcp_mock.py b/tests/providers/google/cloud/utils/base_gcp_mock.py index 6d9563a947197..9c8b9b60e7f78 100644 --- a/tests/providers/google/cloud/utils/base_gcp_mock.py +++ b/tests/providers/google/cloud/utils/base_gcp_mock.py @@ -31,7 +31,7 @@ def mock_base_gcp_hook_default_project_id( delegate_to=None, impersonation_chain=None, ): - self.extras_list = {'extra__google_cloud_platform__project': GCP_PROJECT_ID_HOOK_UNIT_TEST} + self.extras_list = {'project': GCP_PROJECT_ID_HOOK_UNIT_TEST} self._conn = gcp_conn_id self.delegate_to = delegate_to self.impersonation_chain = impersonation_chain @@ -57,9 +57,7 @@ def mock_base_gcp_hook_no_default_project_id( self._cached_project_id = None -GCP_CONNECTION_WITH_PROJECT_ID = Connection( - extra=json.dumps({'extra__google_cloud_platform__project': GCP_PROJECT_ID_HOOK_UNIT_TEST}) -) +GCP_CONNECTION_WITH_PROJECT_ID = Connection(extra=json.dumps({'project': GCP_PROJECT_ID_HOOK_UNIT_TEST})) GCP_CONNECTION_WITHOUT_PROJECT_ID = Connection(extra=json.dumps({})) diff --git a/tests/providers/google/cloud/utils/test_credentials_provider.py b/tests/providers/google/cloud/utils/test_credentials_provider.py index 0cc7b6bfa6de9..43e29789ef976 100644 --- a/tests/providers/google/cloud/utils/test_credentials_provider.py +++ b/tests/providers/google/cloud/utils/test_credentials_provider.py @@ -59,17 +59,17 @@ class TestHelper(unittest.TestCase): def test_build_gcp_conn_path(self): value = "test" conn = build_gcp_conn(key_file_path=value) - assert "google-cloud-platform://?extra__google_cloud_platform__key_path=test" == conn + assert "google-cloud-platform://?key_path=test" == conn def test_build_gcp_conn_scopes(self): value = ["test", "test2"] conn = build_gcp_conn(scopes=value) - assert "google-cloud-platform://?extra__google_cloud_platform__scope=test%2Ctest2" == conn + assert "google-cloud-platform://?scope=test%2Ctest2" == conn def test_build_gcp_conn_project(self): value = "test" conn = build_gcp_conn(project_id=value) - assert "google-cloud-platform://?extra__google_cloud_platform__projects=test" == conn + assert "google-cloud-platform://?projects=test" == conn class TestProvideGcpCredentials(unittest.TestCase): diff --git a/tests/providers/google/common/hooks/test_base_google.py b/tests/providers/google/common/hooks/test_base_google.py index 2c7a40a00e573..9d2866824d151 100644 --- a/tests/providers/google/common/hooks/test_base_google.py +++ b/tests/providers/google/common/hooks/test_base_google.py @@ -155,8 +155,8 @@ def setUp(self): def test_provide_gcp_credential_file_decorator_key_path_and_keyfile_dict(self): key_path = '/test/key-path' self.instance.extras = { - 'extra__google_cloud_platform__key_path': key_path, - 'extra__google_cloud_platform__keyfile_dict': '{"foo": "bar"}', + 'key_path': key_path, + 'keyfile_dict': '{"foo": "bar"}', } @hook.GoogleBaseHook.provide_gcp_credential_file @@ -172,7 +172,7 @@ def assert_gcp_credential_file_in_env(_): def test_provide_gcp_credential_file_decorator_key_path(self): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} @hook.GoogleBaseHook.provide_gcp_credential_file def assert_gcp_credential_file_in_env(_): @@ -185,7 +185,7 @@ def test_provide_gcp_credential_file_decorator_key_content(self, mock_file): string_file = StringIO() file_content = '{"foo": "bar"}' file_name = '/test/mock-file' - self.instance.extras = {'extra__google_cloud_platform__keyfile_dict': file_content} + self.instance.extras = {'keyfile_dict': file_content} mock_file_handler = mock_file.return_value.__enter__.return_value mock_file_handler.name = file_name mock_file_handler.write = string_file.write @@ -200,7 +200,7 @@ def assert_gcp_credential_file_in_env(_): @mock.patch.dict(os.environ, {CREDENTIALS: ENV_VALUE}) def test_provide_gcp_credential_keep_environment(self): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} @hook.GoogleBaseHook.provide_gcp_credential_file def assert_gcp_credential_file_in_env(_): @@ -212,7 +212,7 @@ def assert_gcp_credential_file_in_env(_): @mock.patch.dict(os.environ, {CREDENTIALS: ENV_VALUE}) def test_provide_gcp_credential_keep_environment_when_exception(self): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} @hook.GoogleBaseHook.provide_gcp_credential_file def assert_gcp_credential_file_in_env(_): @@ -226,7 +226,7 @@ def assert_gcp_credential_file_in_env(_): @mock.patch.dict(os.environ, clear=True) def test_provide_gcp_credential_keep_clear_environment(self): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} @hook.GoogleBaseHook.provide_gcp_credential_file def assert_gcp_credential_file_in_env(_): @@ -238,7 +238,7 @@ def assert_gcp_credential_file_in_env(_): @mock.patch.dict(os.environ, clear=True) def test_provide_gcp_credential_keep_clear_environment_when_exception(self): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} @hook.GoogleBaseHook.provide_gcp_credential_file def assert_gcp_credential_file_in_env(_): @@ -260,7 +260,7 @@ def setUp(self): def test_provide_gcp_credential_file_decorator_key_path(self): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} with self.instance.provide_gcp_credential_file_as_context(): assert os.environ[CREDENTIALS] == key_path @@ -270,7 +270,7 @@ def test_provide_gcp_credential_file_decorator_key_content(self, mock_file): string_file = StringIO() file_content = '{"foo": "bar"}' file_name = '/test/mock-file' - self.instance.extras = {'extra__google_cloud_platform__keyfile_dict': file_content} + self.instance.extras = {'keyfile_dict': file_content} mock_file_handler = mock_file.return_value.__enter__.return_value mock_file_handler.name = file_name mock_file_handler.write = string_file.write @@ -282,7 +282,7 @@ def test_provide_gcp_credential_file_decorator_key_content(self, mock_file): @mock.patch.dict(os.environ, {CREDENTIALS: ENV_VALUE}) def test_provide_gcp_credential_keep_environment(self): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} with self.instance.provide_gcp_credential_file_as_context(): assert os.environ[CREDENTIALS] == key_path @@ -292,7 +292,7 @@ def test_provide_gcp_credential_keep_environment(self): @mock.patch.dict(os.environ, {CREDENTIALS: ENV_VALUE}) def test_provide_gcp_credential_keep_environment_when_exception(self): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} with pytest.raises(Exception): with self.instance.provide_gcp_credential_file_as_context(): @@ -303,7 +303,7 @@ def test_provide_gcp_credential_keep_environment_when_exception(self): @mock.patch.dict(os.environ, clear=True) def test_provide_gcp_credential_keep_clear_environment(self): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} with self.instance.provide_gcp_credential_file_as_context(): assert os.environ[CREDENTIALS] == key_path @@ -313,7 +313,7 @@ def test_provide_gcp_credential_keep_clear_environment(self): @mock.patch.dict(os.environ, clear=True) def test_provide_gcp_credential_keep_clear_environment_when_exception(self): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} with pytest.raises(Exception): with self.instance.provide_gcp_credential_file_as_context(): @@ -364,7 +364,7 @@ def test_connection_failure(self, mock_get_creds_and_proj_id): def test_get_credentials_and_project_id_with_service_account_file(self, mock_get_creds_and_proj_id): mock_credentials = mock.MagicMock() mock_get_creds_and_proj_id.return_value = (mock_credentials, "PROJECT_ID") - self.instance.extras = {'extra__google_cloud_platform__key_path': "KEY_PATH.json"} + self.instance.extras = {'key_path': "KEY_PATH.json"} result = self.instance.get_credentials_and_project_id() mock_get_creds_and_proj_id.assert_called_once_with( key_path='KEY_PATH.json', @@ -379,12 +379,12 @@ def test_get_credentials_and_project_id_with_service_account_file(self, mock_get assert (mock_credentials, 'PROJECT_ID') == result def test_get_credentials_and_project_id_with_service_account_file_and_p12_key(self): - self.instance.extras = {'extra__google_cloud_platform__key_path': "KEY_PATH.p12"} + self.instance.extras = {'key_path': "KEY_PATH.p12"} with pytest.raises(AirflowException): self.instance.get_credentials_and_project_id() def test_get_credentials_and_project_id_with_service_account_file_and_unknown_key(self): - self.instance.extras = {'extra__google_cloud_platform__key_path': "KEY_PATH.unknown"} + self.instance.extras = {'key_path': "KEY_PATH.unknown"} with pytest.raises(AirflowException): self.instance.get_credentials_and_project_id() @@ -393,7 +393,7 @@ def test_get_credentials_and_project_id_with_service_account_info(self, mock_get mock_credentials = mock.MagicMock() mock_get_creds_and_proj_id.return_value = (mock_credentials, "PROJECT_ID") service_account = {'private_key': "PRIVATE_KEY"} - self.instance.extras = {'extra__google_cloud_platform__keyfile_dict': json.dumps(service_account)} + self.instance.extras = {'keyfile_dict': json.dumps(service_account)} result = self.instance.get_credentials_and_project_id() mock_get_creds_and_proj_id.assert_called_once_with( key_path=None, @@ -447,7 +447,7 @@ def test_get_credentials_and_project_id_with_default_auth_and_unsupported_delega def test_get_credentials_and_project_id_with_default_auth_and_overridden_project_id( self, mock_get_creds_and_proj_id ): - self.instance.extras = {'extra__google_cloud_platform__project': "SECOND_PROJECT_ID"} + self.instance.extras = {'project': "SECOND_PROJECT_ID"} result = self.instance.get_credentials_and_project_id() mock_get_creds_and_proj_id.assert_called_once_with( key_path=None, @@ -465,9 +465,9 @@ def test_get_credentials_and_project_id_with_mutually_exclusive_configuration( self, ): self.instance.extras = { - 'extra__google_cloud_platform__project': "PROJECT_ID", - 'extra__google_cloud_platform__key_path': "KEY_PATH", - 'extra__google_cloud_platform__keyfile_dict': '{"KEY": "VALUE"}', + 'project': "PROJECT_ID", + 'key_path': "KEY_PATH", + 'keyfile_dict': '{"KEY": "VALUE"}', } with pytest.raises( AirflowException, @@ -481,7 +481,7 @@ def test_get_credentials_and_project_id_with_invalid_keyfile_dict( self, ): self.instance.extras = { - 'extra__google_cloud_platform__keyfile_dict': 'INVALID_DICT', + 'keyfile_dict': 'INVALID_DICT', } with pytest.raises(AirflowException, match=re.escape('Invalid key JSON.')): self.instance.get_credentials_and_project_id() @@ -491,8 +491,8 @@ def test_get_credentials_and_project_id_with_invalid_keyfile_dict( ) def test_default_creds_with_scopes(self): self.instance.extras = { - 'extra__google_cloud_platform__project': default_project, - 'extra__google_cloud_platform__scope': ( + 'project': default_project, + 'scope': ( ','.join( ( 'https://www.googleapis.com/auth/bigquery', @@ -517,7 +517,7 @@ def test_default_creds_with_scopes(self): not default_creds_available, 'Default Google Cloud credentials not available to run tests' ) def test_default_creds_no_scopes(self): - self.instance.extras = {'extra__google_cloud_platform__project': default_project} + self.instance.extras = {'project': default_project} credentials = self.instance.get_credentials() @@ -531,7 +531,7 @@ def test_default_creds_no_scopes(self): def test_provide_gcp_credential_file_decorator_key_path(self): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} @hook.GoogleBaseHook.provide_gcp_credential_file def assert_gcp_credential_file_in_env(hook_instance): @@ -544,7 +544,7 @@ def test_provide_gcp_credential_file_decorator_key_content(self, mock_file): string_file = StringIO() file_content = '{"foo": "bar"}' file_name = '/test/mock-file' - self.instance.extras = {'extra__google_cloud_platform__keyfile_dict': file_content} + self.instance.extras = {'keyfile_dict': file_content} mock_file_handler = mock_file.return_value.__enter__.return_value mock_file_handler.name = file_name mock_file_handler.write = string_file.write @@ -558,8 +558,8 @@ def assert_gcp_credential_file_in_env(hook_instance): def test_provided_scopes(self): self.instance.extras = { - 'extra__google_cloud_platform__project': default_project, - 'extra__google_cloud_platform__scope': ( + 'project': default_project, + 'scope': ( ','.join( ( 'https://www.googleapis.com/auth/bigquery', @@ -575,7 +575,7 @@ def test_provided_scopes(self): ] def test_default_scopes(self): - self.instance.extras = {'extra__google_cloud_platform__project': default_project} + self.instance.extras = {'project': default_project} assert self.instance.scopes == ('https://www.googleapis.com/auth/cloud-platform',) @@ -585,7 +585,7 @@ def test_num_retries_is_not_none_by_default(self, get_con_mock): Verify that if 'num_retries' in extras is not set, the default value should not be None """ - get_con_mock.return_value.extra_dejson = {"extra__google_cloud_platform__num_retries": None} + get_con_mock.return_value.extra_dejson = {"num_retries": None} assert self.instance.num_retries == 5 @mock.patch("airflow.providers.google.common.hooks.base_google.build_http") @@ -684,8 +684,8 @@ def setUp(self): def test_provide_authorized_gcloud_key_path_and_keyfile_dict(self, mock_check_output, mock_default): key_path = '/test/key-path' self.instance.extras = { - 'extra__google_cloud_platform__key_path': key_path, - 'extra__google_cloud_platform__keyfile_dict': '{"foo": "bar"}', + 'key_path': key_path, + 'keyfile_dict': '{"foo": "bar"}', } with pytest.raises( @@ -704,7 +704,7 @@ def test_provide_authorized_gcloud_key_path_and_keyfile_dict(self, mock_check_ou @mock.patch(MODULE_NAME + '.check_output') def test_provide_authorized_gcloud_key_path(self, mock_check_output, mock_project_id): key_path = '/test/key-path' - self.instance.extras = {'extra__google_cloud_platform__key_path': key_path} + self.instance.extras = {'key_path': key_path} with self.instance.provide_authorized_gcloud(): assert os.environ[CREDENTIALS] == key_path @@ -726,7 +726,7 @@ def test_provide_authorized_gcloud_keyfile_dict(self, mock_file, mock_check_outp string_file = StringIO() file_content = '{"foo": "bar"}' file_name = '/test/mock-file' - self.instance.extras = {'extra__google_cloud_platform__keyfile_dict': file_content} + self.instance.extras = {'keyfile_dict': file_content} mock_file_handler = mock_file.return_value.__enter__.return_value mock_file_handler.name = file_name mock_file_handler.write = string_file.write @@ -782,7 +782,7 @@ class TestNumRetry(unittest.TestCase): def test_should_return_int_when_set_int_via_connection(self): instance = hook.GoogleBaseHook(gcp_conn_id="google_cloud_default") instance.extras = { - 'extra__google_cloud_platform__num_retries': 10, + 'num_retries': 10, } assert isinstance(instance.num_retries, int) @@ -790,9 +790,7 @@ def test_should_return_int_when_set_int_via_connection(self): @mock.patch.dict( 'os.environ', - AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT=( - 'google-cloud-platform://?extra__google_cloud_platform__num_retries=5' - ), + AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT=('google-cloud-platform://?num_retries=5'), ) def test_should_return_int_when_set_via_env_var(self): instance = hook.GoogleBaseHook(gcp_conn_id="google_cloud_default") @@ -800,9 +798,7 @@ def test_should_return_int_when_set_via_env_var(self): @mock.patch.dict( 'os.environ', - AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT=( - 'google-cloud-platform://?extra__google_cloud_platform__num_retries=cat' - ), + AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT=('google-cloud-platform://?num_retries=cat'), ) def test_should_raise_when_invalid_value_via_env_var(self): instance = hook.GoogleBaseHook(gcp_conn_id="google_cloud_default") @@ -811,9 +807,7 @@ def test_should_raise_when_invalid_value_via_env_var(self): @mock.patch.dict( 'os.environ', - AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT=( - 'google-cloud-platform://?extra__google_cloud_platform__num_retries=' - ), + AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT=('google-cloud-platform://?num_retries='), ) def test_should_fallback_when_empty_string_in_env_var(self): instance = hook.GoogleBaseHook(gcp_conn_id="google_cloud_default") diff --git a/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py b/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py index 0279badaf6f1e..59f4cb8ef1f27 100644 --- a/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py +++ b/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py @@ -83,9 +83,9 @@ def create_connection_gcp(conn_id_name: str, secret_arn: str): ) scopes = 'https://www.googleapis.com/auth/youtube.readonly' conn_extra = { - "extra__google_cloud_platform__scope": scopes, - "extra__google_cloud_platform__project": "aws-oss-airflow", - 'extra__google_cloud_platform__keyfile_dict': json_data, + "scope": scopes, + "project": "aws-oss-airflow", + 'keyfile_dict': json_data, } conn_extra_json = json.dumps(conn_extra) conn.set_extra(conn_extra_json)