Update compatibility with google-cloud-kms>=2.0

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,13 +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 (),
         )
 
-        plaintext = response.plaintext
-        return plaintext
+        return response.plaintext

setup.py

@@ -260,7 +260,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version
     'google-cloud-datacatalog>=0.5.0, <0.8',  # TODO: we should migrate to 1.0 likely and add <2.0.0 then
     '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',

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)

tests/providers/google/cloud/hooks/test_kms_system.py

@@ -0,0 +1,104 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import base64
+import os
+from tempfile import TemporaryDirectory
+
+import pytest
+
+from airflow.providers.google.cloud.hooks.kms import CloudKMSHook
+from tests.providers.google.cloud.utils.gcp_authenticator import GCP_KMS_KEY
+from tests.test_utils.gcp_system_helpers import GoogleSystemTest, provide_gcp_context
+
+# To prevent resource name collisions, key ring and key resources CANNOT be deleted, so
+# to avoid cluttering the project, we only create the key once during project initialization.
+# See: https://cloud.google.com/kms/docs/faq#cannot_delete
+GCP_KMS_KEYRING_NAME = os.environ.get('GCP_KMS_KEYRING_NAME', 'test-airflow-system-tests-keyring')
+GCP_KMS_KEY_NAME = os.environ.get('GCP_KMS_KEY_NAME', 'test-airflow-system-tests-key')
+
+
+@pytest.mark.credential_file(GCP_KMS_KEY)
+class TestKmsHook(GoogleSystemTest):
+    @provide_gcp_context(GCP_KMS_KEY)
+    def test_encrypt(self):
+        with TemporaryDirectory() as tmp_dir:
+            kms_hook = CloudKMSHook()
+            content = kms_hook.encrypt(
+                key_name=(
+                    f"projects/{kms_hook.project_id}/locations/global/keyRings/"
+                    f"{GCP_KMS_KEYRING_NAME}/cryptoKeys/{GCP_KMS_KEY_NAME}"
+                ),
+                plaintext=b"TEST-SECRET",
+            )
+            with open(f"{tmp_dir}/mysecret.txt.encrypted", "wb") as encrypted_file:
+                encrypted_file.write(base64.b64decode(content))
+            self.execute_cmd(
+                [
+                    "gcloud",
+                    "kms",
+                    "decrypt",
+                    "--location",
+                    "global",
+                    "--keyring",
+                    GCP_KMS_KEYRING_NAME,
+                    "--key",
+                    GCP_KMS_KEY_NAME,
+                    "--ciphertext-file",
+                    f"{tmp_dir}/mysecret.txt.encrypted",
+                    "--plaintext-file",
+                    f"{tmp_dir}/mysecret.txt",
+                ]
+            )
+            with open(f"{tmp_dir}/mysecret.txt", "rb") as secret_file:
+                secret = secret_file.read()
+            self.assertEqual(secret, b"TEST-SECRET")
+
+    @provide_gcp_context(GCP_KMS_KEY)
+    def test_decrypt(self):
+        with TemporaryDirectory() as tmp_dir:
+            with open(f"{tmp_dir}/mysecret.txt", "w") as secret_file:
+                secret_file.write("TEST-SECRET")
+            self.execute_cmd(
+                [
+                    "gcloud",
+                    "kms",
+                    "encrypt",
+                    "--location",
+                    "global",
+                    "--keyring",
+                    GCP_KMS_KEYRING_NAME,
+                    "--key",
+                    GCP_KMS_KEY_NAME,
+                    "--plaintext-file",
+                    f"{tmp_dir}/mysecret.txt",
+                    "--ciphertext-file",
+                    f"{tmp_dir}/mysecret.txt.encrypted",
+                ]
+            )
+            with open(f"{tmp_dir}/mysecret.txt.encrypted", "rb") as encrypted_file:
+                encrypted_secret = base64.b64encode(encrypted_file.read()).decode()
+
+            kms_hook = CloudKMSHook()
+            content = kms_hook.decrypt(
+                key_name=(
+                    f"projects/{kms_hook.project_id}/locations/global/keyRings/"
+                    f"{GCP_KMS_KEYRING_NAME}/cryptoKeys/{GCP_KMS_KEY_NAME}"
+                ),
+                ciphertext=encrypted_secret,
+            )
+            self.assertEqual(content, b"TEST-SECRET")

tests/providers/google/cloud/utils/gcp_authenticator.py

@@ -46,6 +46,7 @@
 GCP_GCS_KEY = 'gcp_gcs.json'
 GCP_GCS_TRANSFER_KEY = 'gcp_gcs_transfer.json'
 GCP_GKE_KEY = "gcp_gke.json"
+GCP_KMS_KEY = "gcp_kms.json"
 GCP_LIFE_SCIENCES_KEY = 'gcp_life_sciences.json'
 GCP_MEMORYSTORE = 'gcp_memorystore.json'
 GCP_PUBSUB_KEY = "gcp_pubsub.json"

tests/test_utils/gcp_system_helpers.py

@@ -83,9 +83,12 @@ def provide_gcp_context(
     :param scopes: OAuth scopes for the connection
     :type scopes: Sequence
     :param project_id: The id of GCP project for the connection.
+        Default: ``os.environ["GCP_PROJECT_ID"]`` or None
     :type project_id: str
     """
     key_file_path = resolve_full_gcp_key_path(key_file_path)  # type: ignore
+    if project_id is None:
+        project_id = os.environ.get("GCP_PROJECT_ID")
     with provide_gcp_conn_and_credentials(
         key_file_path, scopes, project_id
     ), tempfile.TemporaryDirectory() as gcloud_config_tmp, mock.patch.dict(