From b26b0df5b03c4cd826fd7b2dff5771d64e18e6b7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= <mik-laj@users.noreply.github.com>
Date: Tue, 22 Dec 2020 12:59:27 +0100
Subject: [PATCH] Update compatibility with google-cloud-kms>=2.0 (#13124)

---
 airflow/providers/google/cloud/hooks/kms.py   | 20 ++++++----
 setup.py                                      |  2 +-
 .../providers/google/cloud/hooks/test_kms.py  | 40 +++++++++++--------
 3 files changed, 37 insertions(+), 25 deletions(-)

diff --git a/airflow/providers/google/cloud/hooks/kms.py b/airflow/providers/google/cloud/hooks/kms.py
index e63c2f1cb79f6..3fd14336d23d7 100644
--- 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
index bba626d648519..0475fd702d4b8 100644
--- 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
index 894532aaed41b..d2f4519df5534 100644
--- 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)