diff --git a/sdks/python/apache_beam/utils/subprocess_server.py b/sdks/python/apache_beam/utils/subprocess_server.py index 5637a1da575a..162f0f479754 100644 --- a/sdks/python/apache_beam/utils/subprocess_server.py +++ b/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 diff --git a/sdks/python/apache_beam/utils/subprocess_server_test.py b/sdks/python/apache_beam/utils/subprocess_server_test.py index b639e0e7cd6e..c848595db355 100644 --- a/sdks/python/apache_beam/utils/subprocess_server_test.py +++ b/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: