This is an automated email from the ASF dual-hosted git repository.
Abacn 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 63bcd3dd2b5 Validate Project of staging bucket (#39008)
63bcd3dd2b5 is described below
commit 63bcd3dd2b5d54261de386ad388e2d576913e525
Author: Tarun Annapareddy <[email protected]>
AuthorDate: Thu Jun 18 11:32:28 2026 -0700
Validate Project of staging bucket (#39008)
---
sdks/python/apache_beam/io/gcp/gcsio.py | 50 ++++++++++++++-
sdks/python/apache_beam/io/gcp/gcsio_test.py | 92 ++++++++++++++++++++++++++++
sdks/python/setup.py | 1 +
3 files changed, 140 insertions(+), 3 deletions(-)
diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py
b/sdks/python/apache_beam/io/gcp/gcsio.py
index da5c20aa0e7..7cc68fd1e4a 100644
--- a/sdks/python/apache_beam/io/gcp/gcsio.py
+++ b/sdks/python/apache_beam/io/gcp/gcsio.py
@@ -77,6 +77,48 @@ def default_gcs_bucket_name(project, region):
region, md5(project.encode('utf8')).hexdigest())
+def _get_project_number(project_id, credentials=None):
+ """Resolves a project ID to its project number using Cloud Resource Manager
API."""
+ from google.cloud import resourcemanager_v3
+ client = resourcemanager_v3.ProjectsClient(credentials=credentials)
+ project_info = client.get_project(name=f"projects/{project_id}")
+ # project_info.name is of the form "projects/PROJECT_NUMBER"
+ return int(project_info.name.split('/')[-1])
+
+
+def _validate_bucket_project(bucket, project_id, credentials=None):
+ """Verifies that the GCS bucket is owned by the executing project."""
+ bucket_project_number = bucket.project_number
+
+ # Skip validation if the bucket project number is a mock object
+ if (type(bucket_project_number).__name__.endswith('Mock') or
+ hasattr(bucket_project_number, '_mock_self')):
+ _LOGGER.warning(
+ 'Bucket project number is a mock object. Skipping ownership
validation.'
+ )
+ return
+
+ if bucket_project_number is None:
+ _LOGGER.warning(
+ 'Bucket gs://%s does not have a project number. Skipping ownership
validation.',
+ bucket.name)
+ return
+
+ try:
+ project_number = _get_project_number(project_id, credentials=credentials)
+ except Exception as e:
+ _LOGGER.warning(
+ 'Failed to resolve project number for project %s: %s. '
+ 'Skipping bucket ownership validation.',
+ project_id,
+ e)
+ return
+
+ if bucket_project_number != project_number:
+ raise ValueError(
+ f'Bucket gs://{bucket.name} is not owned by project {project_id}.')
+
+
def get_or_create_default_gcs_bucket(options):
"""Create a default GCS bucket for this project."""
if getattr(options, 'dataflow_kms_key', None):
@@ -90,16 +132,18 @@ def get_or_create_default_gcs_bucket(options):
return None
bucket_name = default_gcs_bucket_name(project, region)
- bucket = GcsIO(pipeline_options=options).get_bucket(bucket_name)
+ gcs = GcsIO(pipeline_options=options)
+ bucket = gcs.get_bucket(bucket_name)
if bucket:
+ _validate_bucket_project(
+ bucket, project, credentials=getattr(gcs.client, '_credentials', None))
return bucket
else:
_LOGGER.warning(
'Creating default GCS bucket for project %s: gs://%s',
project,
bucket_name)
- return GcsIO(pipeline_options=options).create_bucket(
- bucket_name, project, location=region)
+ return gcs.create_bucket(bucket_name, project, location=region)
def create_storage_client(pipeline_options, use_credentials=True):
diff --git a/sdks/python/apache_beam/io/gcp/gcsio_test.py
b/sdks/python/apache_beam/io/gcp/gcsio_test.py
index ec4ccbf1cf5..9c4414175e4 100644
--- a/sdks/python/apache_beam/io/gcp/gcsio_test.py
+++ b/sdks/python/apache_beam/io/gcp/gcsio_test.py
@@ -395,6 +395,98 @@ class TestGCSIO(unittest.TestCase):
DEFAULT_GCP_PROJECT, "us-central1", kms_key="kmskey!")),
None)
+ @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient')
+ @mock.patch('apache_beam.io.gcp.gcsio.GcsIO')
+ def test_get_or_create_default_gcs_bucket_ownership_match(
+ self, mock_gcsio_class, mock_crm_class):
+ mock_gcsio = mock_gcsio_class.return_value
+ mock_bucket = mock.Mock()
+ mock_bucket.project_number = 123456789
+ mock_gcsio.get_bucket.return_value = mock_bucket
+
+ mock_crm_client = mock_crm_class.return_value
+ mock_project_info = mock.Mock()
+ mock_project_info.name = 'projects/123456789'
+ mock_crm_client.get_project.return_value = mock_project_info
+
+ options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1')
+ bucket = gcsio.get_or_create_default_gcs_bucket(options)
+
+ self.assertEqual(bucket, mock_bucket)
+ mock_gcsio.get_bucket.assert_called_once_with(
+ 'dataflow-staging-us-central1-77b801c0838aee13391c0d1885860494')
+
+ @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient')
+ @mock.patch('apache_beam.io.gcp.gcsio.GcsIO')
+ def test_get_or_create_default_gcs_bucket_ownership_mismatch(
+ self, mock_gcsio_class, mock_crm_class):
+ mock_gcsio = mock_gcsio_class.return_value
+ mock_bucket = mock.Mock()
+ mock_bucket.project_number = 999999999
+ mock_gcsio.get_bucket.return_value = mock_bucket
+
+ mock_crm_client = mock_crm_class.return_value
+ mock_project_info = mock.Mock()
+ mock_project_info.name = 'projects/123456789'
+ mock_crm_client.get_project.return_value = mock_project_info
+
+ options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1')
+ with self.assertRaises(ValueError) as ctx:
+ gcsio.get_or_create_default_gcs_bucket(options)
+
+ self.assertIn(
+ f'is not owned by project {DEFAULT_GCP_PROJECT}', str(ctx.exception))
+
+ @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient')
+ @mock.patch('apache_beam.io.gcp.gcsio.GcsIO')
+ def test_get_or_create_default_gcs_bucket_ownership_no_bucket_number(
+ self, mock_gcsio_class, mock_crm_class):
+ mock_gcsio = mock_gcsio_class.return_value
+ mock_bucket = mock.Mock()
+ mock_bucket.project_number = None
+ mock_bucket.name =
'dataflow-staging-us-central1-77b801c0838aee13391c0d1885860494'
+ mock_gcsio.get_bucket.return_value = mock_bucket
+
+ options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1')
+ bucket = gcsio.get_or_create_default_gcs_bucket(options)
+
+ self.assertEqual(bucket, mock_bucket)
+ mock_crm_class.assert_not_called()
+
+ @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient')
+ @mock.patch('apache_beam.io.gcp.gcsio.GcsIO')
+ def test_get_or_create_default_gcs_bucket_ownership_crm_error(
+ self, mock_gcsio_class, mock_crm_class):
+ mock_gcsio = mock_gcsio_class.return_value
+ mock_bucket = mock.Mock()
+ mock_bucket.project_number = 123456789
+ mock_gcsio.get_bucket.return_value = mock_bucket
+
+ mock_crm_client = mock_crm_class.return_value
+ mock_crm_client.get_project.side_effect = Exception("API Disabled")
+
+ options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1')
+ bucket = gcsio.get_or_create_default_gcs_bucket(options)
+
+ # Should fall back to success (warn only)
+ self.assertEqual(bucket, mock_bucket)
+
+ @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient')
+ @mock.patch('apache_beam.io.gcp.gcsio.GcsIO')
+ def test_get_or_create_default_gcs_bucket_ownership_mock_project_number(
+ self, mock_gcsio_class, mock_crm_class):
+ mock_gcsio = mock_gcsio_class.return_value
+ # Creating a mock bucket without setting project_number (it returns
another mock object)
+ mock_bucket = mock.Mock()
+ mock_gcsio.get_bucket.return_value = mock_bucket
+
+ options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1')
+ bucket = gcsio.get_or_create_default_gcs_bucket(options)
+
+ # Verification should be skipped, returning the mock bucket and never
calling CRM
+ self.assertEqual(bucket, mock_bucket)
+ mock_crm_class.assert_not_called()
+
def test_exists(self):
file_name = 'gs://gcsio-test/dummy_file'
file_size = 1234
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index 027b3f65039..c3543995a85 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -531,6 +531,7 @@ if __name__ == '__main__':
'google-cloud-datastore>=2.0.0,<3',
'google-cloud-pubsub>=2.1.0,<3',
'google-cloud-storage>=2.18.2,<4',
+ 'google-cloud-resource-manager>=1.12.0,<2',
'google-cloud-dataflow-client>=0.13.0,<0.14.0',
# GCP packages required by tests
'google-cloud-bigquery>=2.0.0,<4',
