[AIRFLOW-4931] Add KMS Encryption Configuration to BigQuery Hook and Operators (apache#5567)

Add KMS to BigQuery

(cherry picked from commit 682aea2)

Author: ryanyuan

airflow/contrib/hooks/bigquery_hook.py

@@ -226,6 +226,7 @@ def __init__(self,
         self.location = location
         self.num_retries = num_retries
 
+    # pylint: disable=too-many-arguments
     def create_empty_table(self,
                            project_id,
                            dataset_id,
@@ -235,6 +236,7 @@ def create_empty_table(self,
                            cluster_fields=None,
                            labels=None,
                            view=None,
+                           encryption_configuration=None,
                            num_retries=None):
         """
         Creates a new, empty table in the dataset.
@@ -280,6 +282,13 @@ def create_empty_table(self,
                 "useLegacySql": False
             }
 
+        :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys).
+            **Example**: ::
+
+                encryption_configuration = {
+                    "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key"
+                }
+        :type encryption_configuration: dict
         :return: None
         """
 
@@ -308,6 +317,9 @@ def create_empty_table(self,
         if view:
             table_resource['view'] = view
 
+        if encryption_configuration:
+            table_resource["encryptionConfiguration"] = encryption_configuration
+
         num_retries = num_retries if num_retries else self.num_retries
 
         self.log.info('Creating Table %s:%s.%s',
@@ -342,7 +354,8 @@ def create_external_table(self,
                               allow_quoted_newlines=False,
                               allow_jagged_rows=False,
                               src_fmt_configs=None,
-                              labels=None
+                              labels=None,
+                              encryption_configuration=None
                               ):
         """
         Creates a new external table in the dataset with the data in Google
@@ -405,6 +418,13 @@ def create_external_table(self,
         :type src_fmt_configs: dict
         :param labels: a dictionary containing labels for the table, passed to BigQuery
         :type labels: dict
+        :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys).
+            **Example**: ::
+
+                encryption_configuration = {
+                    "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key"
+                }
+        :type encryption_configuration: dict
         """
 
         if src_fmt_configs is None:
@@ -508,6 +528,9 @@ def create_external_table(self,
         if labels:
             table_resource['labels'] = labels
 
+        if encryption_configuration:
+            table_resource["encryptionConfiguration"] = encryption_configuration
+
         try:
             self.service.tables().insert(
                 projectId=project_id,
@@ -535,7 +558,8 @@ def patch_table(self,
                     schema=None,
                     time_partitioning=None,
                     view=None,
-                    require_partition_filter=None):
+                    require_partition_filter=None,
+                    encryption_configuration=None):
         """
         Patch information in an existing table.
         It only updates fileds that are provided in the request object.
@@ -587,6 +611,13 @@ def patch_table(self,
         :param require_partition_filter: [Optional] If true, queries over the this table require a
             partition filter. If false, queries over the table
         :type require_partition_filter: bool
+        :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys).
+            **Example**: ::
+
+                encryption_configuration = {
+                    "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key"
+                }
+        :type encryption_configuration: dict
 
         """
 
@@ -612,6 +643,8 @@ def patch_table(self,
             table_resource['view'] = view
         if require_partition_filter is not None:
             table_resource['requirePartitionFilter'] = require_partition_filter
+        if encryption_configuration:
+            table_resource["encryptionConfiguration"] = encryption_configuration
 
         self.log.info('Patching Table %s:%s.%s',
                       project_id, dataset_id, table_id)
@@ -631,7 +664,7 @@ def patch_table(self,
                 'BigQuery job failed. Error was: {}'.format(err.content)
             )
 
-    def run_query(self,
+    def run_query(self,  # pylint: disable=too-many-locals,too-many-arguments
                   bql=None,
                   sql=None,
                   destination_dataset_table=None,
@@ -650,7 +683,8 @@ def run_query(self,
                   time_partitioning=None,
                   api_resource_configs=None,
                   cluster_fields=None,
-                  location=None):
+                  location=None,
+                  encryption_configuration=None):
         """
         Executes a BigQuery SQL query. Optionally persists results in a BigQuery
         table. See here:
@@ -725,6 +759,13 @@ def run_query(self,
             US and EU. See details at
             https://cloud.google.com/bigquery/docs/locations#specifying_your_location
         :type location: str
+        :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys).
+            **Example**: ::
+
+                encryption_configuration = {
+                    "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key"
+                }
+        :type encryption_configuration: dict
         """
 
         if time_partitioning is None:
@@ -867,6 +908,11 @@ def run_query(self,
                 'labels', labels, configuration)
             configuration['labels'] = labels
 
+        if encryption_configuration:
+            configuration["query"][
+                "destinationEncryptionConfiguration"
+            ] = encryption_configuration
+
         return self.run_with_configuration(configuration)
 
     def run_extract(  # noqa
@@ -942,7 +988,8 @@ def run_copy(self,
                  destination_project_dataset_table,
                  write_disposition='WRITE_EMPTY',
                  create_disposition='CREATE_IF_NEEDED',
-                 labels=None):
+                 labels=None,
+                 encryption_configuration=None):
         """
         Executes a BigQuery copy command to copy data from one BigQuery table
         to another. See here:
@@ -968,6 +1015,13 @@ def run_copy(self,
         :param labels: a dictionary containing labels for the job/query,
             passed to BigQuery
         :type labels: dict
+        :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys).
+            **Example**: ::
+
+                encryption_configuration = {
+                    "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key"
+                }
+        :type encryption_configuration: dict
         """
         source_project_dataset_tables = ([
             source_project_dataset_tables
@@ -1008,6 +1062,11 @@ def run_copy(self,
         if labels:
             configuration['labels'] = labels
 
+        if encryption_configuration:
+            configuration["copy"][
+                "destinationEncryptionConfiguration"
+            ] = encryption_configuration
+
         return self.run_with_configuration(configuration)
 
     def run_load(self,
@@ -1028,7 +1087,8 @@ def run_load(self,
                  src_fmt_configs=None,
                  time_partitioning=None,
                  cluster_fields=None,
-                 autodetect=False):
+                 autodetect=False,
+                 encryption_configuration=None):
         """
         Executes a BigQuery load command to load data from Google Cloud Storage
         to BigQuery. See here:
@@ -1098,6 +1158,13 @@ def run_load(self,
             by one or more columns. This is only available in combination with
             time_partitioning. The order of columns given determines the sort order.
         :type cluster_fields: list[str]
+        :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys).
+            **Example**: ::
+
+                encryption_configuration = {
+                    "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key"
+                }
+        :type encryption_configuration: dict
         """
 
         # bigquery only allows certain source formats
@@ -1189,6 +1256,11 @@ def run_load(self,
         if max_bad_records:
             configuration['load']['maxBadRecords'] = max_bad_records
 
+        if encryption_configuration:
+            configuration["load"][
+                "destinationEncryptionConfiguration"
+            ] = encryption_configuration
+
         # if following fields are not specified in src_fmt_configs,
         # honor the top-level params for backward-compatibility
         if 'skipLeadingRows' not in src_fmt_configs:

airflow/contrib/operators/bigquery_operator.py

@@ -130,6 +130,13 @@ class BigQueryOperator(BaseOperator):
         US and EU. See details at
         https://cloud.google.com/bigquery/docs/locations#specifying_your_location
     :type location: str
+    :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys).
+        **Example**: ::
+
+            encryption_configuration = {
+                "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key"
+            }
+    :type encryption_configuration: dict
     """
 
     template_fields = ('bql', 'sql', 'destination_dataset_table', 'labels')
@@ -140,7 +147,7 @@ class BigQueryOperator(BaseOperator):
         BigQueryConsoleLink(),
     )
 
-    # pylint: disable=too-many-arguments
+    # pylint: disable=too-many-arguments, too-many-locals
     @apply_defaults
     def __init__(self,
                  bql=None,
@@ -164,6 +171,7 @@ def __init__(self,
                  api_resource_configs=None,
                  cluster_fields=None,
                  location=None,
+                 encryption_configuration=None,
                  *args,
                  **kwargs):
         super(BigQueryOperator, self).__init__(*args, **kwargs)
@@ -189,6 +197,7 @@ def __init__(self,
         self.api_resource_configs = api_resource_configs
         self.cluster_fields = cluster_fields
         self.location = location
+        self.encryption_configuration = encryption_configuration
 
         # TODO remove `bql` in Airflow 2.0
         if self.bql:
@@ -233,6 +242,7 @@ def execute(self, context):
                 time_partitioning=self.time_partitioning,
                 api_resource_configs=self.api_resource_configs,
                 cluster_fields=self.cluster_fields,
+                encryption_configuration=self.encryption_configuration
             )
         elif isinstance(self.sql, Iterable):
             job_id = [
@@ -253,6 +263,7 @@ def execute(self, context):
                     time_partitioning=self.time_partitioning,
                     api_resource_configs=self.api_resource_configs,
                     cluster_fields=self.cluster_fields,
+                    encryption_configuration=self.encryption_configuration
                 )
                 for s in self.sql]
         else:
@@ -354,7 +365,13 @@ class BigQueryCreateEmptyTableOperator(BaseOperator):
                 google_cloud_storage_conn_id='airflow-service-account'
             )
     :type labels: dict
+    :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys).
+        **Example**: ::
 
+            encryption_configuration = {
+                "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key"
+            }
+    :type encryption_configuration: dict
     """
     template_fields = ('dataset_id', 'table_id', 'project_id',
                        'gcs_schema_object', 'labels')
@@ -373,6 +390,7 @@ def __init__(self,
                  google_cloud_storage_conn_id='google_cloud_default',
                  delegate_to=None,
                  labels=None,
+                 encryption_configuration=None,
                  *args, **kwargs):
 
         super(BigQueryCreateEmptyTableOperator, self).__init__(*args, **kwargs)
@@ -387,6 +405,7 @@ def __init__(self,
         self.delegate_to = delegate_to
         self.time_partitioning = {} if time_partitioning is None else time_partitioning
         self.labels = labels
+        self.encryption_configuration = encryption_configuration
 
     def execute(self, context):
         bq_hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id,
@@ -414,7 +433,8 @@ def execute(self, context):
             table_id=self.table_id,
             schema_fields=schema_fields,
             time_partitioning=self.time_partitioning,
-            labels=self.labels
+            labels=self.labels,
+            encryption_configuration=self.encryption_configuration
         )
 
 
@@ -490,6 +510,13 @@ class BigQueryCreateExternalTableOperator(BaseOperator):
     :type src_fmt_configs: dict
     :param labels: a dictionary containing labels for the table, passed to BigQuery
     :type labels: dict
+    :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys).
+        **Example**: ::
+
+            encryption_configuration = {
+                "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key"
+            }
+    :type encryption_configuration: dict
     """
     template_fields = ('bucket', 'source_objects',
                        'schema_object', 'destination_project_dataset_table', 'labels')
@@ -516,6 +543,7 @@ def __init__(self,
                  delegate_to=None,
                  src_fmt_configs=None,
                  labels=None,
+                 encryption_configuration=None,
                  *args, **kwargs):
 
         super(BigQueryCreateExternalTableOperator, self).__init__(*args, **kwargs)
@@ -543,6 +571,7 @@ def __init__(self,
 
         self.src_fmt_configs = src_fmt_configs if src_fmt_configs is not None else dict()
         self.labels = labels
+        self.encryption_configuration = encryption_configuration
 
     def execute(self, context):
         bq_hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id,
@@ -577,7 +606,8 @@ def execute(self, context):
             allow_quoted_newlines=self.allow_quoted_newlines,
             allow_jagged_rows=self.allow_jagged_rows,
             src_fmt_configs=self.src_fmt_configs,
-            labels=self.labels
+            labels=self.labels,
+            encryption_configuration=self.encryption_configuration
         )
 
 

airflow/contrib/operators/bigquery_to_bigquery.py

@@ -52,6 +52,13 @@ class BigQueryToBigQueryOperator(BaseOperator):
     :param labels: a dictionary containing labels for the job/query,
         passed to BigQuery
     :type labels: dict
+    :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys).
+        **Example**: ::
+
+            encryption_configuration = {
+                "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key"
+            }
+    :type encryption_configuration: dict
     """
     template_fields = ('source_project_dataset_tables',
                        'destination_project_dataset_table', 'labels')
@@ -67,6 +74,7 @@ def __init__(self,
                  bigquery_conn_id='bigquery_default',
                  delegate_to=None,
                  labels=None,
+                 encryption_configuration=None,
                  *args,
                  **kwargs):
         super(BigQueryToBigQueryOperator, self).__init__(*args, **kwargs)
@@ -77,6 +85,7 @@ def __init__(self,
         self.bigquery_conn_id = bigquery_conn_id
         self.delegate_to = delegate_to
         self.labels = labels
+        self.encryption_configuration = encryption_configuration
 
     def execute(self, context):
         self.log.info(
@@ -92,4 +101,5 @@ def execute(self, context):
             destination_project_dataset_table=self.destination_project_dataset_table,
             write_disposition=self.write_disposition,
             create_disposition=self.create_disposition,
-            labels=self.labels)
+            labels=self.labels,
+            encryption_configuration=self.encryption_configuration)