Update compatibility with google-cloud-os-login>=2.0.0 (#13126)

apache · Dec 17, 2020 · 1259c71 · 1259c71
1 parent 8529cb1
commit 1259c71
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 18 deletions.
diff --git 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 @@ def get_conn(self) -> OsLoginServiceClient:
  @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 @@ def import_ssh_public_key(
  :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
@@ -265,7 +265,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version
  '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
@@ -37,7 +37,7 @@
 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"
 
 
@@ -66,9 +66,11 @@ def test_import_ssh_public_key(self, mock_get_conn, mock_get_creds_and_project_i
  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,
@@ -100,9 +102,11 @@ def test_import_ssh_public_key(self, mock_get_conn, mock_get_creds_and_project_i
  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,
@@ -134,9 +138,7 @@ def test_import_ssh_public_key(self, mock_get_conn, mock_get_creds_and_project_i
  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,