feat(pipeline_options): add support for custom maven repository url (#36390)

* feat(pipeline_options): add support for custom maven repository url

Add --maven_repository_url flag to SetupOptions to allow specifying custom Maven repository
Modify JavaJarServer, JavaJarExpansionService and BeamJarExpansionService to use custom repository
Update _resolve_expansion_service to pass maven_repository_url from pipeline options
Add unit tests for maven_repository_url functionality

* comments

* retvert sdks/python/apache_beam/utils/subprocess_server.py

* need to support maven_repository_url for path_to_beam_jar

* use MAVEN_STAGING_REPOSITORY for rc

Author: liferoad

sdks/python/apache_beam/options/pipeline_options.py

@@ -1741,6 +1741,13 @@ def _add_argparse_args(cls, parser):
         help=(
             'A user agent string describing the pipeline to external services. '
             'The format should follow RFC2616.'))
+    parser.add_argument(
+        '--maven_repository_url',
+        default=None,
+        help=(
+            'Custom Maven repository URL to use for downloading JAR files. '
+            'If not specified, the default Maven Central repository will be '
+            'used.'))
 
   def validate(self, validator):
     errors = []

sdks/python/apache_beam/transforms/external.py

@@ -1032,15 +1032,21 @@ 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.
+    user_agent: HTTP user agent string used when downloading jars via
+      `JavaJarServer.local_jar`, including the main jar and any classpath
+      dependencies.
+    maven_repository_url: Maven repository base URL to resolve artifacts when
+      classpath entries or jars are specified as Maven coordinates
+      (`group:artifact:version`). Defaults to Maven Central if not provided.
   """
   def __init__(
       self,
       path_to_jar,
       extra_args=None,
       classpath=None,
       append_args=None,
-      user_agent=None):
+      user_agent=None,
+      maven_repository_url=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
@@ -1049,12 +1055,13 @@ def __init__(
     self._service_count = 0
     self._append_args = append_args or []
     self._user_agent = user_agent
+    self._maven_repository_url = maven_repository_url
 
   def is_existing_service(self):
     return subprocess_server.is_service_endpoint(self.path_to_jar)
 
   @staticmethod
-  def _expand_jars(jar, user_agent=None):
+  def _expand_jars(jar, user_agent=None, maven_repository_url=None):
     if glob.glob(jar):
       return glob.glob(jar)
     elif isinstance(jar, str) and (jar.startswith('http://') or
@@ -1073,15 +1080,21 @@ def _expand_jars(jar, user_agent=None):
         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,
+              repository=(
+                  maven_repository_url or
+                  subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY)),
           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, self._user_agent)
+        JavaJarExpansionService._expand_jars(
+            jar, self._user_agent, self._maven_repository_url)
         for jar in self._classpath or []), []))
     args = ['{{PORT}}', f'--filesToStage={to_stage}']
     # TODO(robertwb): See if it's possible to scope this per pipeline.
@@ -1110,7 +1123,8 @@ def __enter__(self):
           subprocess_server.JavaJarServer.local_jar(path)
           for jar in self._classpath
           for path in JavaJarExpansionService._expand_jars(
-              jar, user_agent=self._user_agent)
+              jar, user_agent=self._user_agent,
+              maven_repository_url=self._maven_repository_url)
       ]
       self._service_provider = subprocess_server.JavaJarServer(
           ExpansionAndArtifactRetrievalStub,
@@ -1146,6 +1160,11 @@ class BeamJarExpansionService(JavaJarExpansionService):
     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: HTTP user agent string used when downloading the Beam jar and
+      any classpath dependencies.
+    maven_repository_url: Maven repository base URL to resolve the Beam jar
+      for the provided Gradle target. Defaults to Maven Central if not
+      provided.
   """
   def __init__(
       self,
@@ -1154,16 +1173,20 @@ def __init__(
       gradle_appendix=None,
       classpath=None,
       append_args=None,
-      user_agent=None):
+      user_agent=None,
+      maven_repository_url=None):
     path_to_jar = subprocess_server.JavaJarServer.path_to_beam_jar(
-        gradle_target, gradle_appendix)
+        gradle_target,
+        gradle_appendix,
+        maven_repository_url=maven_repository_url)
     self.gradle_target = gradle_target
     super().__init__(
         path_to_jar,
         extra_args,
         classpath=classpath,
         append_args=append_args,
-        user_agent=user_agent)
+        user_agent=user_agent,
+        maven_repository_url=maven_repository_url)
 
 
 def _maybe_use_transform_service(provided_service=None, options=None):

sdks/python/apache_beam/transforms/managed.py

@@ -118,16 +118,24 @@ def __init__(
           f"An unsupported source was specified: '{source}'. Please specify "
           f"one of the following sources: {list(self._READ_TRANSFORMS.keys())}")
 
-    self._expansion_service = _resolve_expansion_service(
-        source, identifier, expansion_service)
+    # Store parameters for deferred expansion service creation
+    self._identifier = identifier
+    self._provided_expansion_service = expansion_service
     self._underlying_identifier = identifier
     self._yaml_config = yaml.dump(config)
     self._config_url = config_url
 
   def expand(self, input):
+    # Create expansion service with access to pipeline options
+    expansion_service = _resolve_expansion_service(
+        self._source,
+        self._identifier,
+        self._provided_expansion_service,
+        pipeline_options=input.pipeline._options)
+
     return input | SchemaAwareExternalTransform(
         identifier=MANAGED_SCHEMA_TRANSFORM_IDENTIFIER,
-        expansion_service=self._expansion_service,
+        expansion_service=expansion_service,
         rearrange_based_on_discovery=True,
         transform_identifier=self._underlying_identifier,
         config=self._yaml_config,
@@ -162,16 +170,24 @@ def __init__(
           f"An unsupported sink was specified: '{sink}'. Please specify "
           f"one of the following sinks: {list(self._WRITE_TRANSFORMS.keys())}")
 
-    self._expansion_service = _resolve_expansion_service(
-        sink, identifier, expansion_service)
+    # Store parameters for deferred expansion service creation
+    self._identifier = identifier
+    self._provided_expansion_service = expansion_service
     self._underlying_identifier = identifier
     self._yaml_config = yaml.dump(config)
     self._config_url = config_url
 
   def expand(self, input):
+    # Create expansion service with access to pipeline options
+    expansion_service = _resolve_expansion_service(
+        self._sink,
+        self._identifier,
+        self._provided_expansion_service,
+        pipeline_options=input.pipeline._options)
+
     return input | SchemaAwareExternalTransform(
         identifier=MANAGED_SCHEMA_TRANSFORM_IDENTIFIER,
-        expansion_service=self._expansion_service,
+        expansion_service=expansion_service,
         rearrange_based_on_discovery=True,
         transform_identifier=self._underlying_identifier,
         config=self._yaml_config,
@@ -182,7 +198,10 @@ def default_label(self) -> str:
 
 
 def _resolve_expansion_service(
-    transform_name: str, identifier: str, expansion_service):
+    transform_name: str,
+    identifier: str,
+    expansion_service,
+    pipeline_options=None):
   if expansion_service:
     return expansion_service
 
@@ -193,4 +212,18 @@ def _resolve_expansion_service(
     raise ValueError(
         "No expansion service was specified and could not find a "
         f"default expansion service for {transform_name}: '{identifier}'.")
-  return BeamJarExpansionService(gradle_target)
+
+  # Extract maven_repository_url and user_agent from pipeline options if
+  # available
+  maven_repository_url = None
+  user_agent = None
+  if pipeline_options:
+    from apache_beam.options import pipeline_options as po
+    setup_options = pipeline_options.view_as(po.SetupOptions)
+    maven_repository_url = setup_options.maven_repository_url
+    user_agent = setup_options.user_agent
+
+  return BeamJarExpansionService(
+      gradle_target,
+      maven_repository_url=maven_repository_url,
+      user_agent=user_agent)