feat(subprocess_server): add fallback to Google Maven mirror (#36365)

* feat(subprocess_server): add fallback to Google Maven mirror when Maven Central fails

Implement fallback mechanism to use Google Maven mirror when downloads from Maven Central fail (e.g. 403 Forbidden). This improves reliability of jar downloads by providing an alternative source.

* lints

* refactor(subprocess_server): extract jar download logic into separate method

Move the jar download and caching logic from local_jar method to a new _download_jar_to_cache method to improve code reuse and maintainability

Author: liferoad

sdks/python/apache_beam/utils/subprocess_server.py

@@ -281,6 +281,7 @@ class JavaJarServer(SubprocessServer):
 
   MAVEN_CENTRAL_REPOSITORY = 'https://repo.maven.apache.org/maven2'
   MAVEN_STAGING_REPOSITORY = 'https://repository.apache.org/content/groups/staging'  # pylint: disable=line-too-long
+  GOOGLE_MAVEN_MIRROR = 'https://maven-central.storage-download.googleapis.com/maven2'  # pylint: disable=line-too-long
   BEAM_GROUP_ID = 'org.apache.beam'
   JAR_CACHE = os.path.expanduser("~/.apache_beam/cache/jars")
 
@@ -419,6 +420,32 @@ def path_to_beam_jar(
     return cls.path_to_maven_jar(
         artifact_id, cls.BEAM_GROUP_ID, version, maven_repo, appendix=appendix)
 
+  @classmethod
+  def _download_jar_to_cache(
+      cls, download_url, cached_jar_path, user_agent=None):
+    """Downloads a jar from the given URL to the specified cache path.
+    
+    Args:
+      download_url (str): The URL to download from.
+      cached_jar_path (str): The local path where the jar should be cached.
+      user_agent (str): The user agent to use when downloading.
+    """
+    try:
+      url_read = FileSystems.open(download_url)
+    except ValueError:
+      if user_agent is None:
+        user_agent = cls._DEFAULT_USER_AGENT
+      url_request = Request(download_url, headers={'User-Agent': user_agent})
+      url_read = urlopen(url_request)
+    with open(cached_jar_path + '.tmp', 'wb') as jar_write:
+      shutil.copyfileobj(url_read, jar_write, length=1 << 20)
+    try:
+      os.rename(cached_jar_path + '.tmp', cached_jar_path)
+    except FileNotFoundError:
+      # A race when multiple programs run in parallel and the cached_jar
+      # is already moved. Safe to ignore.
+      pass
+
   @classmethod
   def local_jar(cls, url, cache_dir=None, user_agent=None):
     """Returns a local path to the given jar, downloading it if necessary.
@@ -449,25 +476,31 @@ def local_jar(cls, url, cache_dir=None, user_agent=None):
           os.makedirs(cache_dir)
           # TODO: Clean up this cache according to some policy.
         try:
-          try:
-            url_read = FileSystems.open(url)
-          except ValueError:
-            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:
-            os.rename(cached_jar + '.tmp', cached_jar)
-          except FileNotFoundError:
-            # A race when multiple programs run in parallel and the cached_jar
-            # is already moved. Safe to ignore.
-            pass
+          cls._download_jar_to_cache(url, cached_jar, user_agent)
         except URLError as e:
-          raise RuntimeError(
-              f'Unable to fetch remote job server jar at {url}: {e}. If no '
-              f'Internet access at runtime, stage the jar at {cached_jar}')
+          # Try Google Maven mirror as fallback if the original URL is from
+          # Maven Central
+          if url.startswith(cls.MAVEN_CENTRAL_REPOSITORY):
+            fallback_url = url.replace(
+                cls.MAVEN_CENTRAL_REPOSITORY, cls.GOOGLE_MAVEN_MIRROR)
+            _LOGGER.info(
+                'Trying Google Maven mirror fallback: %s' % fallback_url)
+            try:
+              cls._download_jar_to_cache(fallback_url, cached_jar, user_agent)
+              _LOGGER.info(
+                  'Successfully downloaded from Google Maven mirror: %s' %
+                  fallback_url)
+            except URLError as fallback_e:
+              raise RuntimeError(
+                  f'Unable to fetch remote job server jar at {url}: {e}. '
+                  f'Also failed to fetch from Google Maven mirror at '
+                  f'{fallback_url}: {fallback_e}. '
+                  f'If no Internet access at runtime, stage the jar at '
+                  f'{cached_jar}')
+          else:
+            raise RuntimeError(
+                f'Unable to fetch remote job server jar at {url}: {e}. If no '
+                f'Internet access at runtime, stage the jar at {cached_jar}')
       return cached_jar
 
   @classmethod

sdks/python/apache_beam/utils/subprocess_server_test.py

@@ -126,6 +126,56 @@ def handle(self):
       with open(os.path.join(temp_dir, 'file.jar')) as fin:
         self.assertEqual(fin.read(), 'data')
 
+  def test_local_jar_fallback_to_google_maven_mirror(self):
+    """Test that Google Maven mirror is used as fallback 
+    when Maven Central fails."""
+    class MavenCentralHandler(socketserver.BaseRequestHandler):
+      timeout = 1
+
+      def handle(self):
+        # Simulate Maven Central returning 403 Forbidden
+        self.request.sendall(b'HTTP/1.1 403 Forbidden\n\n')
+
+    # Set up Maven Central server (will return 403)
+    maven_port, = subprocess_server.pick_port(None)
+    maven_server = socketserver.TCPServer(('localhost', maven_port),
+                                          MavenCentralHandler)
+    maven_thread = threading.Thread(target=maven_server.handle_request)
+    maven_thread.daemon = True
+    maven_thread.start()
+
+    # Temporarily replace the Maven Central constant to use our test server
+    original_maven_central = (
+        subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY)
+
+    try:
+      subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY = (
+          f'http://localhost:{maven_port}/maven2')
+
+      with tempfile.TemporaryDirectory() as temp_dir:
+        # Use a Maven Central URL that will trigger the fallback to real
+        # Google mirror
+        maven_url = (
+            f'http://localhost:{maven_port}/maven2/org/apache/beam/'
+            f'beam-sdks-java-extensions-schemaio-expansion-service/2.63.0/'
+            f'beam-sdks-java-extensions-schemaio-expansion-service-2.63.0.jar')
+
+        # This should fail on our mock Maven Central and fallback to the
+        # real Google mirror
+        jar_path = subprocess_server.JavaJarServer.local_jar(
+            maven_url, temp_dir)
+
+        # Verify the file was downloaded successfully (from the real Google
+        # mirror)
+        self.assertTrue(os.path.exists(jar_path))
+        jar_size = os.path.getsize(jar_path)
+        self.assertTrue(jar_size > 0)  # Should have actual content
+
+    finally:
+      # Restore original constants
+      subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY = (
+          original_maven_central)
+
   @unittest.skipUnless(shutil.which('javac'), 'missing java jdk')
   def test_classpath_jar(self):
     with tempfile.TemporaryDirectory() as temp_dir: