diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index 6595d683911b..c6435c8a6f4b 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/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 = [] diff --git a/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py b/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py index 3b302e334a5f..45ca6f92bd5d 100644 --- a/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py +++ b/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( diff --git a/sdks/python/apache_beam/runners/portability/job_server.py b/sdks/python/apache_beam/runners/portability/job_server.py index 0d98de6bdf3d..9fdaabd1a177 100644 --- a/sdks/python/apache_beam/runners/portability/job_server.py +++ b/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) diff --git a/sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server.py b/sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server.py index f754b4c330ad..e38da2d46f5e 100644 --- a/sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server.py +++ b/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( diff --git a/sdks/python/apache_beam/transforms/external.py b/sdks/python/apache_beam/transforms/external.py index 3f9f56a54139..c90291192411 100644 --- a/sdks/python/apache_beam/transforms/external.py +++ b/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,9 +1030,15 @@ 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 @@ -1040,12 +1046,13 @@ def __init__( 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 ' diff --git a/sdks/python/apache_beam/transforms/external_test.py b/sdks/python/apache_beam/transforms/external_test.py index 2ed7d622ecd6..c59058a6e62b 100644 --- a/sdks/python/apache_beam/transforms/external_test.py +++ b/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 diff --git a/sdks/python/apache_beam/utils/subprocess_server.py b/sdks/python/apache_beam/utils/subprocess_server.py index 84848479430b..5637a1da575a 100644 --- a/sdks/python/apache_beam/utils/subprocess_server.py +++ b/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: diff --git a/sdks/python/apache_beam/utils/subprocess_server_test.py b/sdks/python/apache_beam/utils/subprocess_server_test.py index ddf8b3498001..b639e0e7cd6e 100644 --- a/sdks/python/apache_beam/utils/subprocess_server_test.py +++ b/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) diff --git a/sdks/python/apache_beam/utils/transform_service_launcher.py b/sdks/python/apache_beam/utils/transform_service_launcher.py index 8742efd1573a..7b2ad53e8e22 100644 --- a/sdks/python/apache_beam/utils/transform_service_launcher.py +++ b/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)