feat(subprocess_server): add fallback to Google Maven mirror

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: