You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2021/03/03 09:32:42 UTC

[airflow] 18/41: Update compatibility with google-cloud-os-login>=2.0.0 (#13126)

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 <mi...@users.noreply.github.com>
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,