Adding CMEK support for temp_dataset for Python Bigquery (#36118)

* Adding CMEK support for temp_dataset

* Corrected formatting

* Resolved conflict

* formatting

* Formatting

* Fixing tests

Author: TanuSharma2511

sdks/python/apache_beam/io/gcp/bigquery.py

@@ -850,7 +850,8 @@ def _setup_temporary_dataset(self, bq):
       return
     location = bq.get_query_location(
         self._get_project(), self.query.get(), self.use_legacy_sql)
-    bq.create_temporary_dataset(self._get_project(), location)
+    bq.create_temporary_dataset(
+        self._get_project(), location, kms_key=self.kms_key)
 
   @check_accessible(['query'])
   def _execute_query(self, bq):
@@ -1062,7 +1063,10 @@ def _setup_temporary_dataset(self, bq):
         self._get_parent_project(), self.query.get(), self.use_legacy_sql)
     _LOGGER.warning("### Labels: %s", str(self.bigquery_dataset_labels))
     bq.create_temporary_dataset(
-        self._get_parent_project(), location, self.bigquery_dataset_labels)
+        self._get_parent_project(),
+        location,
+        self.bigquery_dataset_labels,
+        kms_key=self.kms_key)
 
   @check_accessible(['query'])
   def _execute_query(self, bq):

sdks/python/apache_beam/io/gcp/bigquery_read_internal.py

@@ -319,7 +319,8 @@ def _setup_temporary_dataset(
     # Use the project from temp_dataset if it's a DatasetReference,
     # otherwise use the pipeline project
     temp_dataset_project = self._get_temp_dataset_project()
-    bq.create_temporary_dataset(temp_dataset_project, location)
+    bq.create_temporary_dataset(
+        temp_dataset_project, location, kms_key=self.kms_key)
 
   def _execute_query(
       self,

sdks/python/apache_beam/io/gcp/bigquery_read_internal_test.py

@@ -99,7 +99,7 @@ def test_setup_temporary_dataset_uses_correct_project(self, mock_bq_wrapper):
 
     # Verify that create_temporary_dataset was called with the custom project
     mock_bq.create_temporary_dataset.assert_called_once_with(
-        'custom-project', 'US')
+        'custom-project', 'US', kms_key=None)
     # Verify that get_query_location was called with the pipeline project
     mock_bq.get_query_location.assert_called_once_with(
         'test-project', 'SELECT * FROM table', False)
@@ -145,7 +145,7 @@ def test_setup_temporary_dataset_with_string_temp_dataset(
 
     # Verify that create_temporary_dataset was called with the pipeline project
     mock_bq.create_temporary_dataset.assert_called_once_with(
-        'test-project', 'US')
+        'test-project', 'US', kms_key=None)
 
   @mock.patch('apache_beam.io.gcp.bigquery_tools.BigQueryWrapper')
   def test_finish_bundle_with_string_temp_dataset(self, mock_bq_wrapper):

sdks/python/apache_beam/io/gcp/bigquery_tools.py

@@ -333,6 +333,10 @@ def _build_filter_from_labels(labels):
   return filter_str
 
 
+def _build_dataset_encryption_config(kms_key):
+  return bigquery.EncryptionConfiguration(kmsKeyName=kms_key)
+
+
 class BigQueryWrapper(object):
   """BigQuery client wrapper with utilities for querying.
 
@@ -835,7 +839,7 @@ def _create_table(
       num_retries=MAX_RETRIES,
       retry_filter=retry.retry_on_server_errors_and_timeout_filter)
   def get_or_create_dataset(
-      self, project_id, dataset_id, location=None, labels=None):
+      self, project_id, dataset_id, location=None, labels=None, kms_key=None):
     # Check if dataset already exists otherwise create it
     try:
       dataset = self.client.datasets.Get(
@@ -858,6 +862,9 @@ def get_or_create_dataset(
           dataset.location = location
         if labels is not None:
           dataset.labels = _build_dataset_labels(labels)
+        if kms_key is not None:
+          dataset.defaultEncryptionConfiguration = (
+              _build_dataset_encryption_config(kms_key))
         request = bigquery.BigqueryDatasetsInsertRequest(
             projectId=project_id, dataset=dataset)
         response = self.client.datasets.Insert(request)
@@ -929,9 +936,14 @@ def is_user_configured_dataset(self):
   @retry.with_exponential_backoff(
       num_retries=MAX_RETRIES,
       retry_filter=retry.retry_on_server_errors_and_timeout_filter)
-  def create_temporary_dataset(self, project_id, location, labels=None):
+  def create_temporary_dataset(
+      self, project_id, location, labels=None, kms_key=None):
     self.get_or_create_dataset(
-        project_id, self.temp_dataset_id, location=location, labels=labels)
+        project_id,
+        self.temp_dataset_id,
+        location=location,
+        labels=labels,
+        kms_key=kms_key)
 
     if (project_id is not None and not self.is_user_configured_dataset() and
         not self.created_temp_dataset):

sdks/python/apache_beam/io/gcp/bigquery_tools_test.py

@@ -301,6 +301,34 @@ def test_get_or_create_dataset_created(self):
     new_dataset = wrapper.get_or_create_dataset('project-id', 'dataset_id')
     self.assertEqual(new_dataset.datasetReference.datasetId, 'dataset_id')
 
+  def test_create_temporary_dataset_with_kms_key(self):
+    kms_key = (
+        'projects/my-project/locations/global/keyRings/my-kr/'
+        'cryptoKeys/my-key')
+    client = mock.Mock()
+    client.datasets.Get.side_effect = HttpError(
+        response={'status': '404'}, url='', content='')
+
+    client.datasets.Insert.return_value = bigquery.Dataset(
+        datasetReference=bigquery.DatasetReference(
+            projectId='project-id', datasetId='temp_dataset'))
+    wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client)
+
+    try:
+      wrapper.create_temporary_dataset(
+          'project-id', 'location', kms_key=kms_key)
+    except Exception:
+      pass
+
+    args, _ = client.datasets.Insert.call_args
+    insert_request = args[0]  # BigqueryDatasetsInsertRequest
+    inserted_dataset = insert_request.dataset  # Actual Dataset object
+
+    # Assertions
+    self.assertIsNotNone(inserted_dataset.defaultEncryptionConfiguration)
+    self.assertEqual(
+        inserted_dataset.defaultEncryptionConfiguration.kmsKeyName, kms_key)
+
   def test_get_or_create_dataset_fetched(self):
     client = mock.Mock()
     client.datasets.Get.return_value = bigquery.Dataset(