[
https://issues.apache.org/jira/browse/AIRFLOW-2062?page=com.atlassian.jira.plugin.system.issuetabpanels:comment-tabpanel=16723730#comment-16723730
]
ASF GitHub Bot commented on AIRFLOW-2062:
-
stale[bot] closed pull request #3805: [AIRFLOW-2062] Add per-connection KMS
encryption.
URL: https://github.com/apache/incubator-airflow/pull/3805
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/airflow/bin/cli.py b/airflow/bin/cli.py
index 1c5494ead1..15b061c94c 100644
--- a/airflow/bin/cli.py
+++ b/airflow/bin/cli.py
@@ -1133,7 +1133,8 @@ def version(args): # noqa
alternative_conn_specs = ['conn_type', 'conn_host',
- 'conn_login', 'conn_password', 'conn_schema',
'conn_port']
+ 'conn_login', 'conn_password', 'conn_schema',
'conn_port',
+ 'kms_conn_id', 'kms_extra']
@cli_utils.action_logging
@@ -1235,7 +1236,10 @@ def connections(args):
return
if args.conn_uri:
-new_conn = Connection(conn_id=args.conn_id, uri=args.conn_uri)
+new_conn = Connection(conn_id=args.conn_id,
+ uri=args.conn_uri,
+ kms_conn_id=args.kms_conn_id,
+ kms_extra=args.kms_extra)
else:
new_conn = Connection(conn_id=args.conn_id,
conn_type=args.conn_type,
@@ -1243,7 +1247,10 @@ def connections(args):
login=args.conn_login,
password=args.conn_password,
schema=args.conn_schema,
- port=args.conn_port)
+ port=args.conn_port,
+ kms_conn_id=args.kms_conn_id,
+ kms_extra=args.kms_extra
+ )
if args.conn_extra is not None:
new_conn.set_extra(args.conn_extra)
@@ -1883,6 +1890,15 @@ class CLIFactory(object):
('--conn_extra',),
help='Connection `Extra` field, optional when adding a connection',
type=str),
+'kms_conn_id': Arg(
+('--kms_conn_id',),
+help='An existing connection to use when encrypting this
connection with a '
+ 'KMS, optional when adding a connection',
+type=str),
+'kms_extra': Arg(
+('--kms_extra',),
+help='Connection `KMS Extra` field, optional when adding a
connection',
+type=str),
# users
'username': Arg(
('--username',),
diff --git a/airflow/contrib/hooks/gcp_kms_hook.py
b/airflow/contrib/hooks/gcp_kms_hook.py
index 6f2b3aedff..63e35fbe89 100644
--- a/airflow/contrib/hooks/gcp_kms_hook.py
+++ b/airflow/contrib/hooks/gcp_kms_hook.py
@@ -20,6 +20,7 @@
import base64
+from airflow.hooks.kmsapi_hook import KmsApiHook
from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook
from apiclient.discovery import build
@@ -35,7 +36,7 @@ def _b64decode(s):
return base64.b64decode(s.encode('utf-8'))
-class GoogleCloudKMSHook(GoogleCloudBaseHook):
+class GoogleCloudKMSHook(GoogleCloudBaseHook, KmsApiHook):
"""
Interact with Google Cloud KMS. This hook uses the Google Cloud Platform
connection.
@@ -106,3 +107,17 @@ def decrypt(self, key_name, ciphertext,
authenticated_data=None):
plaintext = _b64decode(response['plaintext'])
return plaintext
+
+def encrypt_conn_key(self, connection):
+kms_extras = connection.kms_extra_dejson
+key_name = kms_extras['kms_extra__google_cloud_platform__key_name']
+conn_key = connection._plain_conn_key
+
+connection.conn_key = self.encrypt(key_name, conn_key)
+
+def decrypt_conn_key(self, connection):
+kms_extras = connection.kms_extra_dejson
+key_name = kms_extras['kms_extra__google_cloud_platform__key_name']
+conn_key = connection.conn_key
+
+connection._plain_conn_key = self.decrypt(key_name, conn_key)
diff --git a/airflow/hooks/base_hook.py b/airflow/hooks/base_hook.py
index 103fa6260b..fe663f61c1 100644
--- a/airflow/hooks/base_hook.py
+++ b/airflow/hooks/base_hook.py
@@ -22,16 +22,9 @@
from __future__ import print_function
from __future__ import unicode_literals
-import os
-import random
-
from airflow.models import Connection
-from airflow.exceptions import AirflowException
-from airflow.utils.db import provide_session
from airflow.utils.log.logging_mixin import LoggingMixin
-CONN_ENV_PREFIX = 'AIRFLOW_CONN_'
-
class