Add user-agent string for Beam in the GCS request header. (apache#29765)

iht · Dec 14, 2023 · 1ce707c · 1ce707c
1 parent 00462b1
commit 1ce707c
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 5 deletions.
diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py b/sdks/python/apache_beam/io/gcp/gcsio.py
@@ -40,6 +40,7 @@
 from google.cloud.storage.fileio import BlobReader
 from google.cloud.storage.fileio import BlobWriter
 
+from apache_beam import version as beam_version
 from apache_beam.internal.gcp import auth
 from apache_beam.options.pipeline_options import GoogleCloudOptions
 from apache_beam.options.pipeline_options import PipelineOptions
@@ -111,7 +112,11 @@ def __init__(self, storage_client=None, pipeline_options=None):
       if credentials:
         storage_client = storage.Client(
             credentials=credentials.get_google_auth_credentials(),
-            project=pipeline_options.view_as(GoogleCloudOptions).project)
+            project=pipeline_options.view_as(GoogleCloudOptions).project,
+            extra_headers={
+                "User-Agent": "apache-beam/%s (GPN:Beam)" %
+                beam_version.__version__
+            })
       else:
         storage_client = storage.Client.create_anonymous_client()
     self.client = storage_client

diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
@@ -514,7 +514,11 @@ def __init__(self, options, root_staging_location=None):
       # client in terms of which GCP project to use.
       self._storage_client = storage.Client(
           credentials=storage_credentials,
-          project=self.google_cloud_options.project)
+          project=self.google_cloud_options.project,
+          extra_headers={
+              "User-Agent": "apache-beam/%s (GPN:Beam)" %
+              beam_version.__version__
+          })
     else:
       self._storage_client = storage.Client.create_anonymous_client()
     self._sdk_image_overrides = self._get_sdk_image_overrides(options)

diff --git a/sdks/python/apache_beam/runners/interactive/utils.py b/sdks/python/apache_beam/runners/interactive/utils.py
@@ -30,6 +30,7 @@
 import pandas as pd
 
 import apache_beam as beam
+from apache_beam import version as beam_version
 from apache_beam.dataframe.convert import to_pcollection
 from apache_beam.dataframe.frame_base import DeferredBase
 from apache_beam.internal.gcp import auth
@@ -457,7 +458,12 @@ def assert_bucket_exists(bucket_name):
       # We set project to None, so it will not try to use project id from
       # the environment (ADC).
       storage_client = storage.Client(
-          credentials=credentials.get_google_auth_credentials(), project=None)
+          credentials=credentials.get_google_auth_credentials(),
+          project=None,
+          extra_headers={
+              "User-Agent": "apache-beam/%s (GPN:Beam)" %
+              beam_version.__version__
+          })
     else:
       storage_client = storage.Client.create_anonymous_client()
     storage_client.get_bucket(bucket_name)

diff --git a/sdks/python/apache_beam/runners/portability/sdk_container_builder.py b/sdks/python/apache_beam/runners/portability/sdk_container_builder.py
@@ -213,7 +213,11 @@ def __init__(self, options):
     if credentials:
       self._storage_client = storage.Client(
           credentials=credentials.get_google_auth_credentials(),
-          project=self._google_cloud_options.project)
+          project=self._google_cloud_options.project,
+          extra_headers={
+              "User-Agent": "apache-beam/%s (GPN:Beam)" %
+              beam_version.__version__
+          })
     else:
       self._storage_client = storage.Client.create_anonymous_client()
     self._cloudbuild_client = cloudbuild.CloudbuildV1(

diff --git a/sdks/python/setup.py b/sdks/python/setup.py
@@ -373,7 +373,7 @@ def get_portability_package_data():
               'google-cloud-datastore>=2.0.0,<3',
               'google-cloud-pubsub>=2.1.0,<3',
               'google-cloud-pubsublite>=1.2.0,<2',
-              'google-cloud-storage>=2.10.0,<3',
+              'google-cloud-storage>=2.14.0,<3',
               # GCP packages required by tests
               'google-cloud-bigquery>=2.0.0,<4',
               'google-cloud-bigquery-storage>=2.6.3,<3',