This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 29b1abe40a4 Adding CMEK support for temp_dataset for Python Bigquery
(#36118)
29b1abe40a4 is described below
commit 29b1abe40a4878803ca938b5723192b2e96bac2a
Author: Tanu Sharma <[email protected]>
AuthorDate: Mon Sep 15 19:04:36 2025 +0530
Adding CMEK support for temp_dataset for Python Bigquery (#36118)
* Adding CMEK support for temp_dataset
* Corrected formatting
* Resolved conflict
* formatting
* Formatting
* Fixing tests
---
sdks/python/apache_beam/io/gcp/bigquery.py | 8 +++++--
.../apache_beam/io/gcp/bigquery_read_internal.py | 3 ++-
.../io/gcp/bigquery_read_internal_test.py | 4 ++--
sdks/python/apache_beam/io/gcp/bigquery_tools.py | 18 +++++++++++---
.../apache_beam/io/gcp/bigquery_tools_test.py | 28 ++++++++++++++++++++++
5 files changed, 53 insertions(+), 8 deletions(-)
diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py
b/sdks/python/apache_beam/io/gcp/bigquery.py
index 4780f948be2..aa0ebc12ef1 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery.py
@@ -850,7 +850,8 @@ class _CustomBigQuerySource(BoundedSource):
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 @@ class _CustomBigQueryStorageSource(BoundedSource):
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):
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
index 8b8eb6eeb5c..6432f3b4eea 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
@@ -319,7 +319,8 @@ class _BigQueryReadSplit(beam.transforms.DoFn):
# 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,
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_internal_test.py
b/sdks/python/apache_beam/io/gcp/bigquery_read_internal_test.py
index 9d162457df5..46673b4ec2d 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_read_internal_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_read_internal_test.py
@@ -99,7 +99,7 @@ class BigQueryReadSplitTest(unittest.TestCase):
# 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 @@ class BigQueryReadSplitTest(unittest.TestCase):
# 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):
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
index 738d6e9c70f..889d3f1e96e 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
+++ b/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 @@ class BigQueryWrapper(object):
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 @@ class BigQueryWrapper(object):
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 @@ class BigQueryWrapper(object):
@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):
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
index 1320ced1dee..1101317439a 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
@@ -301,6 +301,34 @@ class TestBigQueryWrapper(unittest.TestCase):
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(
