This is an automated email from the ASF dual-hosted git repository. potiuk pushed a commit to branch v2-0-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit e2eceb5bc8ea8c6094b3fdc5057564d284661811 Author: Kamil BreguĊa <[email protected]> AuthorDate: Thu Dec 17 11:00:59 2020 +0100 Update compatibility with google-cloud-os-login>=2.0.0 (#13126) (cherry picked from commit 1259c712a42d69135dc389de88f79942c70079a3) --- airflow/providers/google/cloud/hooks/os_login.py | 16 +++++++++------- setup.py | 2 +- .../providers/google/cloud/hooks/test_os_login.py | 22 ++++++++++++---------- 3 files changed, 22 insertions(+), 18 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/os_login.py b/airflow/providers/google/cloud/hooks/os_login.py index c7a4234..361ea60 100644 --- a/airflow/providers/google/cloud/hooks/os_login.py +++ b/airflow/providers/google/cloud/hooks/os_login.py @@ -17,7 +17,7 @@ from typing import Dict, Optional, Sequence, Union -from google.cloud.oslogin_v1 import OsLoginServiceClient +from google.cloud.oslogin_v1 import ImportSshPublicKeyResponse, OsLoginServiceClient from airflow.providers.google.common.hooks.base_google import GoogleBaseHook @@ -54,7 +54,7 @@ class OSLoginHook(GoogleBaseHook): @GoogleBaseHook.fallback_to_default_project_id def import_ssh_public_key( self, user: str, ssh_public_key: Dict, project_id: str, retry=None, timeout=None, metadata=None - ): + ) -> ImportSshPublicKeyResponse: """ Adds an SSH public key and returns the profile information. Default POSIX account information is set when no username and UID exist as part of the @@ -74,14 +74,16 @@ class OSLoginHook(GoogleBaseHook): :type timeout: Optional[float] :param metadata: Additional metadata that is provided to the method. :type metadata: Optional[Sequence[Tuple[str, str]]] - :return: A :class:`~google.cloud.oslogin_v1.types.ImportSshPublicKeyResponse` instance. + :return: A :class:`~google.cloud.oslogin_v1.ImportSshPublicKeyResponse` instance. """ conn = self.get_conn() return conn.import_ssh_public_key( - parent=OsLoginServiceClient.user_path(user=user), - ssh_public_key=ssh_public_key, - project_id=project_id, + request=dict( + parent=f"users/{user}", + ssh_public_key=ssh_public_key, + project_id=project_id, + ), retry=retry, timeout=timeout, - metadata=metadata, + metadata=metadata or (), ) diff --git a/setup.py b/setup.py index 7071795..0586bf3 100644 --- a/setup.py +++ b/setup.py @@ -295,7 +295,7 @@ google = [ 'google-cloud-logging>=1.14.0,<2.0.0', 'google-cloud-memcache>=0.2.0', 'google-cloud-monitoring>=0.34.0,<2.0.0', - 'google-cloud-os-login>=1.0.0,<2.0.0', + 'google-cloud-os-login>=2.0.0,<3.0.0', 'google-cloud-pubsub>=1.0.0,<2.0.0', 'google-cloud-redis>=0.3.0,<2.0.0', 'google-cloud-secret-manager>=0.2.0,<2.0.0', diff --git a/tests/providers/google/cloud/hooks/test_os_login.py b/tests/providers/google/cloud/hooks/test_os_login.py index 303f1ea..d2b88e4 100644 --- a/tests/providers/google/cloud/hooks/test_os_login.py +++ b/tests/providers/google/cloud/hooks/test_os_login.py @@ -38,7 +38,7 @@ TEST_CREDENTIALS = mock.MagicMock() TEST_BODY: Dict = mock.MagicMock() TEST_RETRY: Retry = mock.MagicMock() TEST_TIMEOUT: float = 4 -TEST_METADATA: Sequence[Tuple[str, str]] = [] +TEST_METADATA: Sequence[Tuple[str, str]] = () TEST_PARENT: str = "users/test-user" @@ -67,9 +67,11 @@ class TestOSLoginHook(TestCase): metadata=TEST_METADATA, ) mock_get_conn.return_value.import_ssh_public_key.assert_called_once_with( - parent=TEST_PARENT, - ssh_public_key=TEST_BODY, - project_id=TEST_PROJECT_ID, + request=dict( + parent=TEST_PARENT, + ssh_public_key=TEST_BODY, + project_id=TEST_PROJECT_ID, + ), retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, @@ -101,9 +103,11 @@ class TestOSLoginHookWithDefaultProjectIdHook(TestCase): metadata=TEST_METADATA, ) mock_get_conn.return_value.import_ssh_public_key.assert_called_once_with( - parent=TEST_PARENT, - ssh_public_key=TEST_BODY, - project_id=TEST_PROJECT_ID_2, + request=dict( + parent=TEST_PARENT, + ssh_public_key=TEST_BODY, + project_id=TEST_PROJECT_ID_2, + ), retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, @@ -135,9 +139,7 @@ class TestOSLoginHookWithoutDefaultProjectIdHook(TestCase): metadata=TEST_METADATA, ) mock_get_conn.return_value.import_ssh_public_key.assert_called_once_with( - parent=TEST_PARENT, - ssh_public_key=TEST_BODY, - project_id=TEST_PROJECT_ID, + request=dict(parent=TEST_PARENT, ssh_public_key=TEST_BODY, project_id=TEST_PROJECT_ID), retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA,
