Update compatibility with google-cloud-kms>=2.0 (#13124)

apache · Dec 22, 2020 · b26b0df · b26b0df
1 parent 9a1d382
commit b26b0df
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 25 deletions.
diff --git a/airflow/providers/google/cloud/hooks/kms.py b/airflow/providers/google/cloud/hooks/kms.py
@@ -118,12 +118,14 @@ def encrypt(
  :rtype: str
  """
  response = self.get_conn().encrypt(
- name=key_name,
- plaintext=plaintext,
- additional_authenticated_data=authenticated_data,
+ request={
+ 'name': key_name,
+ 'plaintext': plaintext,
+ 'additional_authenticated_data': authenticated_data,
+ },
  retry=retry,
  timeout=timeout,
- metadata=metadata,
+ metadata=metadata or (),
  )
 
  ciphertext = _b64encode(response.ciphertext)
@@ -161,12 +163,14 @@ def decrypt(
  :rtype: bytes
  """
  response = self.get_conn().decrypt(
- name=key_name,
- ciphertext=_b64decode(ciphertext),
- additional_authenticated_data=authenticated_data,
+ request={
+ 'name': key_name,
+ 'ciphertext': _b64decode(ciphertext),
+ 'additional_authenticated_data': authenticated_data,
+ },
  retry=retry,
  timeout=timeout,
- metadata=metadata,
+ metadata=metadata or (),
  )
 
  return response.plaintext
diff --git a/setup.py b/setup.py
@@ -260,7 +260,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version
  'google-cloud-datacatalog>=1.0.0,<2.0.0',
  'google-cloud-dataproc>=1.0.1,<2.0.0',
  'google-cloud-dlp>=0.11.0,<2.0.0',
- 'google-cloud-kms>=1.2.1,<2.0.0',
+ 'google-cloud-kms>=2.0.0,<3.0.0',
  'google-cloud-language>=1.1.1,<2.0.0',
  'google-cloud-logging>=1.14.0,<2.0.0',
  'google-cloud-memcache>=0.2.0',

diff --git a/tests/providers/google/cloud/hooks/test_kms.py b/tests/providers/google/cloud/hooks/test_kms.py
@@ -82,12 +82,14 @@ def test_encrypt(self, mock_get_conn):
  result = self.kms_hook.encrypt(TEST_KEY_ID, PLAINTEXT)
  mock_get_conn.assert_called_once_with()
  mock_get_conn.return_value.encrypt.assert_called_once_with(
- name=TEST_KEY_ID,
- plaintext=PLAINTEXT,
- additional_authenticated_data=None,
+ request=dict(
+ name=TEST_KEY_ID,
+ plaintext=PLAINTEXT,
+ additional_authenticated_data=None,
+ ),
  retry=None,
  timeout=None,
- metadata=None,
+ metadata=(),
  )
  self.assertEqual(PLAINTEXT_b64, result)
 
@@ -97,12 +99,14 @@ def test_encrypt_with_auth_data(self, mock_get_conn):
  result = self.kms_hook.encrypt(TEST_KEY_ID, PLAINTEXT, AUTH_DATA)
  mock_get_conn.assert_called_once_with()
  mock_get_conn.return_value.encrypt.assert_called_once_with(
- name=TEST_KEY_ID,
- plaintext=PLAINTEXT,
- additional_authenticated_data=AUTH_DATA,
+ request=dict(
+ name=TEST_KEY_ID,
+ plaintext=PLAINTEXT,
+ additional_authenticated_data=AUTH_DATA,
+ ),
  retry=None,
  timeout=None,
- metadata=None,
+ metadata=(),
  )
  self.assertEqual(PLAINTEXT_b64, result)
 
@@ -112,12 +116,14 @@ def test_decrypt(self, mock_get_conn):
  result = self.kms_hook.decrypt(TEST_KEY_ID, CIPHERTEXT_b64)
  mock_get_conn.assert_called_once_with()
  mock_get_conn.return_value.decrypt.assert_called_once_with(
- name=TEST_KEY_ID,
- ciphertext=CIPHERTEXT,
- additional_authenticated_data=None,
+ request=dict(
+ name=TEST_KEY_ID,
+ ciphertext=CIPHERTEXT,
+ additional_authenticated_data=None,
+ ),
  retry=None,
  timeout=None,
- metadata=None,
+ metadata=(),
  )
  self.assertEqual(PLAINTEXT, result)
 
@@ -127,11 +133,13 @@ def test_decrypt_with_auth_data(self, mock_get_conn):
  result = self.kms_hook.decrypt(TEST_KEY_ID, CIPHERTEXT_b64, AUTH_DATA)
  mock_get_conn.assert_called_once_with()
  mock_get_conn.return_value.decrypt.assert_called_once_with(
- name=TEST_KEY_ID,
- ciphertext=CIPHERTEXT,
- additional_authenticated_data=AUTH_DATA,
+ request=dict(
+ name=TEST_KEY_ID,
+ ciphertext=CIPHERTEXT,
+ additional_authenticated_data=AUTH_DATA,
+ ),
  retry=None,
  timeout=None,
- metadata=None,
+ metadata=(),
  )
  self.assertEqual(PLAINTEXT, result)