Set user agent when downloading jars with local_jar (#36342)

* Set user agent when downloading jars with local_jar

* Fix mock

* Fix typo

Author: Abacn

sdks/python/apache_beam/options/pipeline_options.py

@@ -1716,6 +1716,12 @@ def _add_argparse_args(cls, parser):
         help=(
             'Docker registry url to use for tagging and pushing the prebuilt '
             'sdk worker container image.'))
+    parser.add_argument(
+        '--user_agent',
+        default=None,
+        help=(
+            'A user agent string describing the pipeline to external services. '
+            'The format should follow RFC2616.'))
 
   def validate(self, validator):
     errors = []

sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py

@@ -48,6 +48,7 @@ def __init__(self, master_url, options):
     self._executable_jar = (
         options.view_as(
             pipeline_options.FlinkRunnerOptions).flink_job_server_jar)
+    self._user_agent = options.view_as(pipeline_options.SetupOptions).user_agent
     self._artifact_port = (
         options.view_as(pipeline_options.JobServerOptions).artifact_port)
     self._temp_dir = tempfile.mkdtemp(prefix='apache-beam-flink')
@@ -77,7 +78,8 @@ def executable_jar(self):
     else:
       url = job_server.JavaJarJobServer.path_to_beam_jar(
           ':runners:flink:%s:job-server:shadowJar' % self.flink_version())
-    return job_server.JavaJarJobServer.local_jar(url)
+    return job_server.JavaJarJobServer.local_jar(
+        url, user_agent=self._user_agent)
 
   def flink_version(self):
     full_version = requests.get(

sdks/python/apache_beam/runners/portability/job_server.py

@@ -155,8 +155,9 @@ def path_to_beam_jar(gradle_target, artifact_id=None):
         gradle_target, artifact_id=artifact_id)
 
   @staticmethod
-  def local_jar(url, jar_cache_dir=None):
-    return subprocess_server.JavaJarServer.local_jar(url, jar_cache_dir)
+  def local_jar(url, jar_cache_dir=None, user_agent=None):
+    return subprocess_server.JavaJarServer.local_jar(
+        url, jar_cache_dir, user_agent)
 
   def subprocess_cmd_and_endpoint(self):
     jar_path = self.local_jar(self.path_to_jar(), self._jar_cache_dir)

sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server.py

@@ -53,6 +53,7 @@ def __init__(self, rest_url, options):
     spark_options = options.view_as(pipeline_options.SparkRunnerOptions)
     self._executable_jar = spark_options.spark_job_server_jar
     self._spark_version = spark_options.spark_version
+    self._user_agent = options.view_as(pipeline_options.SetupOptions).user_agent
 
   def start(self):
     return self
@@ -78,7 +79,8 @@ def executable_jar(self):
       else:
         url = job_server.JavaJarJobServer.path_to_beam_jar(
             ':runners:spark:3:job-server:shadowJar')
-    return job_server.JavaJarJobServer.local_jar(url)
+    return job_server.JavaJarJobServer.local_jar(
+        url, user_agent=self._user_agent)
 
   def create_beam_job(self, job_id, job_name, pipeline, options):
     return SparkBeamJob(

sdks/python/apache_beam/transforms/external.py

@@ -35,7 +35,7 @@
 
 from apache_beam import pvalue
 from apache_beam.coders import RowCoder
-from apache_beam.options.pipeline_options import CrossLanguageOptions
+from apache_beam.options import pipeline_options
 from apache_beam.portability import common_urns
 from apache_beam.portability.api import beam_artifact_api_pb2_grpc
 from apache_beam.portability.api import beam_expansion_api_pb2
@@ -1030,22 +1030,29 @@ class JavaJarExpansionService(object):
     append_args: arguments to be provided when starting up the
       expansion service using the jar file. These arguments will be appended to
       the default arguments.
+    user_agent: the user agent to use when downloading the jar.
   """
   def __init__(
-      self, path_to_jar, extra_args=None, classpath=None, append_args=None):
+      self,
+      path_to_jar,
+      extra_args=None,
+      classpath=None,
+      append_args=None,
+      user_agent=None):
     if extra_args and append_args:
       raise ValueError('Only one of extra_args or append_args may be provided')
     self.path_to_jar = path_to_jar
     self._extra_args = extra_args
     self._classpath = classpath or []
     self._service_count = 0
     self._append_args = append_args or []
+    self._user_agent = user_agent
 
   def is_existing_service(self):
     return subprocess_server.is_service_endpoint(self.path_to_jar)
 
   @staticmethod
-  def _expand_jars(jar):
+  def _expand_jars(jar, user_agent=None):
     if glob.glob(jar):
       return glob.glob(jar)
     elif isinstance(jar, str) and (jar.startswith('http://') or
@@ -1064,14 +1071,15 @@ def _expand_jars(jar):
         return [jar]
       path = subprocess_server.JavaJarServer.local_jar(
           subprocess_server.JavaJarServer.path_to_maven_jar(
-              artifact_id, group_id, version))
+              artifact_id, group_id, version),
+          user_agent=user_agent)
       return [path]
 
   def _default_args(self):
     """Default arguments to be used by `JavaJarExpansionService`."""
 
     to_stage = ','.join([self.path_to_jar] + sum((
-        JavaJarExpansionService._expand_jars(jar)
+        JavaJarExpansionService._expand_jars(jar, self._user_agent)
         for jar in self._classpath or []), []))
     args = ['{{PORT}}', f'--filesToStage={to_stage}']
     # TODO(robertwb): See if it's possible to scope this per pipeline.
@@ -1080,10 +1088,14 @@ def _default_args(self):
       args.append('--alsoStartLoopbackWorker')
     return args
 
+  def with_user_agent(self, user_agent: str):
+    self._user_agent = user_agent
+    return self
+
   def __enter__(self):
     if self._service_count == 0:
       self.path_to_jar = subprocess_server.JavaJarServer.local_jar(
-          self.path_to_jar)
+          self.path_to_jar, user_agent=self._user_agent)
       if self._extra_args is None:
         self._extra_args = self._default_args() + self._append_args
       # Consider memoizing these servers (with some timeout).
@@ -1095,7 +1107,8 @@ def __enter__(self):
       classpath_urls = [
           subprocess_server.JavaJarServer.local_jar(path)
           for jar in self._classpath
-          for path in JavaJarExpansionService._expand_jars(jar)
+          for path in JavaJarExpansionService._expand_jars(
+              jar, user_agent=self._user_agent)
       ]
       self._service_provider = subprocess_server.JavaJarServer(
           ExpansionAndArtifactRetrievalStub,
@@ -1138,12 +1151,17 @@ def __init__(
       extra_args=None,
       gradle_appendix=None,
       classpath=None,
-      append_args=None):
+      append_args=None,
+      user_agent=None):
     path_to_jar = subprocess_server.JavaJarServer.path_to_beam_jar(
         gradle_target, gradle_appendix)
     self.gradle_target = gradle_target
     super().__init__(
-        path_to_jar, extra_args, classpath=classpath, append_args=append_args)
+        path_to_jar,
+        extra_args,
+        classpath=classpath,
+        append_args=append_args,
+        user_agent=user_agent)
 
 
 def _maybe_use_transform_service(provided_service=None, options=None):
@@ -1185,10 +1203,11 @@ def is_docker_available():
   docker_available = is_docker_available()
 
   use_transform_service = options.view_as(
-      CrossLanguageOptions).use_transform_service
+      pipeline_options.CrossLanguageOptions).use_transform_service
+  user_agent = options.view_as(pipeline_options.SetupOptions).user_agent
 
   if (java_available and provided_service and not use_transform_service):
-    return provided_service
+    return provided_service.with_user_agent(user_agent)
   elif docker_available:
     if use_transform_service:
       error_append = 'it was explicitly requested'
@@ -1210,7 +1229,7 @@ def is_docker_available():
     beam_version = beam_version.__version__
 
     return transform_service_launcher.TransformServiceLauncher(
-        project_name, port, beam_version)
+        project_name, port, beam_version, user_agent)
   else:
     raise ValueError(
         'Cannot start an expansion service since neither Java nor '

sdks/python/apache_beam/transforms/external_test.py

@@ -829,7 +829,7 @@ def _side_effect_fn(path):
 
   @mock.patch.object(JavaJarServer, 'local_jar')
   def test_classpath_with_gradle_artifact(self, local_jar):
-    def _side_effect_fn(path):
+    def _side_effect_fn(path, user_agent=None):
       return path[path.rindex('/') + 1:]
 
     local_jar.side_effect = _side_effect_fn

sdks/python/apache_beam/utils/subprocess_server.py

@@ -34,11 +34,13 @@
 from typing import Any
 from typing import Set
 from urllib.error import URLError
+from urllib.request import Request
 from urllib.request import urlopen
 
 import grpc
 
 from apache_beam.io.filesystems import FileSystems
+from apache_beam.runners.internal.names import BEAM_SDK_NAME
 from apache_beam.version import __version__ as beam_version
 
 _LOGGER = logging.getLogger(__name__)
@@ -286,6 +288,8 @@ class JavaJarServer(SubprocessServer):
       'local', (threading.local, ),
       dict(__init__=lambda self: setattr(self, 'replacements', {})))()
 
+  _DEFAULT_USER_AGENT = f'{BEAM_SDK_NAME}/{beam_version}'
+
   def __init__(
       self,
       stub_class,
@@ -416,7 +420,18 @@ def path_to_beam_jar(
         artifact_id, cls.BEAM_GROUP_ID, version, maven_repo, appendix=appendix)
 
   @classmethod
-  def local_jar(cls, url, cache_dir=None):
+  def local_jar(cls, url, cache_dir=None, user_agent=None):
+    """Returns a local path to the given jar, downloading it if necessary.
+
+    Args:
+      url (str): A URL or local path to a jar file.
+      cache_dir (str): The directory to use for caching downloaded jars. If not
+        specified, a default temporary directory will be used.
+      user_agent (str): The user agent to use when downloading the jar.
+
+    Returns:
+      str: The local path to the jar file.
+    """
     if cache_dir is None:
       cache_dir = cls.JAR_CACHE
     # TODO: Verify checksum?
@@ -437,7 +452,10 @@ def local_jar(cls, url, cache_dir=None):
           try:
             url_read = FileSystems.open(url)
           except ValueError:
-            url_read = urlopen(url)
+            if user_agent is None:
+              user_agent = cls._DEFAULT_USER_AGENT
+            url_request = Request(url, headers={'User-Agent': user_agent})
+            url_read = urlopen(url_request)
           with open(cached_jar + '.tmp', 'wb') as jar_write:
             shutil.copyfileobj(url_read, jar_write, length=1 << 20)
           try:

sdks/python/apache_beam/utils/subprocess_server_test.py

@@ -108,8 +108,11 @@ class Handler(socketserver.BaseRequestHandler):
       timeout = 1
 
       def handle(self):
-        self.request.recv(1024)
-        self.request.sendall(b'HTTP/1.1 200 OK\n\ndata')
+        data = self.request.recv(1024)
+        if 'User-Agent: Apache Beam SDK for Python' in str(data):
+          self.request.sendall(b'HTTP/1.1 200 OK\n\ndata')
+        else:
+          self.request.sendall(b'HTTP/1.1 400 BAD REQUEST\n\n')
 
     port, = subprocess_server.pick_port(None)
     server = socketserver.TCPServer(('localhost', port), Handler)

sdks/python/apache_beam/utils/transform_service_launcher.py

@@ -57,13 +57,13 @@ class TransformServiceLauncher(object):
 
   # Maintaining a static list of launchers to prevent temporary resources
   # from being created unnecessarily.
-  def __new__(cls, project_name, port, beam_version=None):
+  def __new__(cls, project_name, port, beam_version=None, user_agent=None):
     if project_name not in TransformServiceLauncher._launchers:
       TransformServiceLauncher._launchers[project_name] = super(
           TransformServiceLauncher, cls).__new__(cls)
     return TransformServiceLauncher._launchers[project_name]
 
-  def __init__(self, project_name, port, beam_version=None):
+  def __init__(self, project_name, port, beam_version=None, user_agent=None):
     logging.info('Initializing the Beam Transform Service %s.' % project_name)
 
     self._project_name = project_name
@@ -85,7 +85,8 @@ def __init__(self, project_name, port, beam_version=None):
     # Get the jar with configs
     path_to_local_jar = subprocess_server.JavaJarServer.local_jar(
         subprocess_server.JavaJarServer.path_to_beam_jar(
-            _EXPANSION_SERVICE_LAUNCHER_JAR))
+            _EXPANSION_SERVICE_LAUNCHER_JAR),
+        user_agent=user_agent)
 
     with zipfile.ZipFile(path_to_local_jar) as launcher_jar:
       launcher_jar.extract('docker-compose.yml', path=temp_dir)