diff --git a/core/src/main/scala/org/apache/spark/sql/rapids/tool/util/WebCrawlerUtil.scala b/core/src/main/scala/org/apache/spark/sql/rapids/tool/util/WebCrawlerUtil.scala index 2c31f3a2c..71ebab8c2 100644 --- a/core/src/main/scala/org/apache/spark/sql/rapids/tool/util/WebCrawlerUtil.scala +++ b/core/src/main/scala/org/apache/spark/sql/rapids/tool/util/WebCrawlerUtil.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2025, NVIDIA CORPORATION. + * Copyright (c) 2023-2026, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -16,8 +16,10 @@ package org.apache.spark.sql.rapids.tool.util -import java.io.IOException +import java.io.{InputStream, IOException} import java.net.URL +import java.nio.charset.StandardCharsets +import java.util.Base64 import scala.collection.mutable import scala.jdk.CollectionConverters._ @@ -34,7 +36,11 @@ import org.apache.spark.internal.Logging */ object WebCrawlerUtil extends Logging { private val MAX_CRAWLER_DEPTH = 1 - private val NV_MVN_BASE_URL = "https://repo1.maven.org/maven2/com/nvidia" + private val MAVEN_CENTRAL_BASE_URL = "https://repo1.maven.org/maven2" + private val RAPIDS_GROUP_PATH = "com/nvidia" + private val MAVEN_BASE_URL_ENV = "RAPIDS_TOOLS_MAVEN_BASE_URL" + private val MAVEN_USERNAME_ENV = "RAPIDS_TOOLS_MAVEN_USERNAME" + private val MAVEN_PASSWORD_ENV = "RAPIDS_TOOLS_MAVEN_PASSWORD" // defines the artifacts of the RAPIDS libraries private val NV_ARTIFACTS_LOOKUP = Map( "rapids.plugin" -> "rapids-4-spark_2.12", @@ -44,15 +50,57 @@ object WebCrawlerUtil extends Logging { // regular expression used to extract the version number from // the mvn repository url private val ARTIFACT_VERSION_REGEX = "\\d{2}\\.\\d{2}\\.\\d+/" + + def normalizeMavenBaseUrl(baseUrl: String): String = { + baseUrl.stripSuffix("/") + } + + def getMavenBaseUrl(env: Map[String, String] = sys.env): String = { + env.get(MAVEN_BASE_URL_ENV) + .map(_.trim) + .filter(_.nonEmpty) + .map(normalizeMavenBaseUrl) + .getOrElse(MAVEN_CENTRAL_BASE_URL) + } + + private def getMavenArtifactRootUrl(env: Map[String, String]): String = { + s"${getMavenBaseUrl(env)}/$RAPIDS_GROUP_PATH" + } + + private def getMavenBasicAuthHeader(env: Map[String, String]): Option[String] = { + for { + user <- env.get(MAVEN_USERNAME_ENV).map(_.trim).filter(_.nonEmpty) + password <- env.get(MAVEN_PASSWORD_ENV).filter(_.nonEmpty) + } yield { + val token = Base64.getEncoder.encodeToString( + s"$user:$password".getBytes(StandardCharsets.UTF_8)) + s"Basic $token" + } + } + + private def openMavenUrlStream( + mavenURL: String, + env: Map[String, String]): InputStream = { + val connection = new URL(mavenURL).openConnection() + getMavenBasicAuthHeader(env).foreach { authHeader => + connection.setRequestProperty("Authorization", authHeader) + } + connection.getInputStream + } + // given an artifactID returns the full mvn url that lists all the // releases - def getMVNArtifactURL(artifactID: String) : String = { + def getMVNArtifactURL( + artifactID: String, + env: Map[String, String] = sys.env) : String = { val artifactUrlPart = NV_ARTIFACTS_LOOKUP.getOrElse(artifactID, artifactID) - s"$NV_MVN_BASE_URL/$artifactUrlPart" + s"${getMavenArtifactRootUrl(env)}/$artifactUrlPart" } - def getMVNMetaURL(artifactID: String) : String = { - val artifactUrlPart = getMVNArtifactURL(artifactID) + def getMVNMetaURL( + artifactID: String, + env: Map[String, String] = sys.env) : String = { + val artifactUrlPart = getMVNArtifactURL(artifactID, env) s"$artifactUrlPart/$MAVEN_META_FILE" } @@ -67,7 +115,8 @@ object WebCrawlerUtil extends Logging { def getPageLinks( webURL: String, regEx: Option[String], - maxDepth: Int = MAX_CRAWLER_DEPTH): mutable.Set[String] = { + maxDepth: Int = MAX_CRAWLER_DEPTH, + env: Map[String, String] = sys.env): mutable.Set[String] = { def removeDefaultPorts(rawLink: String): String = { val jURL: URL = new URL(rawLink) jURL.getProtocol match { @@ -84,7 +133,11 @@ object WebCrawlerUtil extends Logging { allLinks: mutable.Set[String]): Unit = { if (currDepth < maxDepth && !allLinks.contains(currURL)) { try { - val doc = Jsoup.connect(currURL).get + val connection = Jsoup.connect(currURL) + getMavenBasicAuthHeader(env).foreach { authHeader => + connection.header("Authorization", authHeader) + } + val doc = connection.get val pageURLs = doc.select(cssQuery).asScala.toList val newDepth = currDepth + 1 for (page <- pageURLs) { @@ -111,19 +164,28 @@ object WebCrawlerUtil extends Logging { // given an artifactID, returns a list of strings containing all // available releases. - def getMvnReleasesForNVPackage(artifactID: String): Seq[String] = { - val mvnURL = getMVNArtifactURL(artifactID) - val definedLinks = getPageLinks(mvnURL, Some(ARTIFACT_VERSION_REGEX)).toSeq.sorted + def getMvnReleasesForNVPackage( + artifactID: String, + env: Map[String, String] = sys.env): Seq[String] = { + val mvnURL = getMVNArtifactURL(artifactID, env) + val definedLinks = getPageLinks(mvnURL, Some(ARTIFACT_VERSION_REGEX), env = env).toSeq.sorted definedLinks.map(_.split("/").last) } // given an artifactID, will return the latest version if any - def getLatestMvnReleaseForNVPackage(artifactID: String): Option[String] = { + def getLatestMvnReleaseForNVPackage( + artifactID: String, + env: Map[String, String] = sys.env): Option[String] = { // Reads maven-metadata.xml file to extract the latest version - val mvnMetaFile = getMVNMetaURL(artifactID) + val mvnMetaFile = getMVNMetaURL(artifactID, env) try { - val xml = XML.load(mvnMetaFile) - Some((xml \\ "metadata" \ "versioning" \ "latest").text) + val inputStream = openMavenUrlStream(mvnMetaFile, env) + try { + val xml = XML.load(inputStream) + Some((xml \\ "metadata" \ "versioning" \ "latest").text) + } finally { + inputStream.close() + } } catch { case NonFatal(e) => logWarning(s"Exception loading maven-metadata.xml: ${mvnMetaFile}", e) @@ -132,12 +194,17 @@ object WebCrawlerUtil extends Logging { } // given artifactID and release, returns the full mvn url to download the jar - def getMvnDownloadLink(artifactID: String, release: String): String = { - s"${getMVNArtifactURL(artifactID)}/$release/$artifactID-$release.jar" + def getMvnDownloadLink( + artifactID: String, + release: String, + env: Map[String, String] = sys.env): String = { + s"${getMVNArtifactURL(artifactID, env)}/$release/$artifactID-$release.jar" } - def getPluginMvnDownloadLink(release: String): String = { - getMvnDownloadLink(NV_ARTIFACTS_LOOKUP("rapids.plugin"), release) + def getPluginMvnDownloadLink( + release: String, + env: Map[String, String] = sys.env): String = { + getMvnDownloadLink(NV_ARTIFACTS_LOOKUP("rapids.plugin"), release, env) } // get the latest version available for rapids plugin diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuite.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuite.scala index 809a52b72..a467c3b5a 100644 --- a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuite.scala +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuite.scala @@ -121,7 +121,7 @@ abstract class ProfilingAutoTunerSuiteBase extends BaseAutoTunerSuite { case Some(v) => v case None => fail("Could not find pull the latest release successfully") } - ToolTestUtils.pluginMvnPrefix(latestRelease) + ".jar" + WebCrawlerUtil.getPluginMvnDownloadLink(latestRelease) } } diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/util/ToolUtilsSuite.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/util/ToolUtilsSuite.scala index 5558f3be6..d4fe75b6c 100644 --- a/core/src/test/scala/com/nvidia/spark/rapids/tool/util/ToolUtilsSuite.scala +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/util/ToolUtilsSuite.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2025, NVIDIA CORPORATION. + * Copyright (c) 2023-2026, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -17,7 +17,10 @@ package com.nvidia.spark.rapids.tool.util import java.io.File +import java.net.InetSocketAddress +import java.nio.charset.StandardCharsets import java.text.SimpleDateFormat +import java.util.Base64 import java.util.Calendar import scala.concurrent.duration._ @@ -25,6 +28,7 @@ import scala.xml.XML import com.nvidia.spark.rapids.tool.ToolTestUtils import com.nvidia.spark.rapids.tool.profiling.{ProfileOutputWriter, ProfileResult} +import com.sun.net.httpserver.HttpServer import org.scalatest.AppendedClues.convertToClueful import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers.{contain, convertToAnyShouldWrapper, equal, not} @@ -104,10 +108,10 @@ class ToolUtilsSuite extends AnyFunSuite with Logging { // test("list available mvn releases") { - // use mvn repo url got testing + // Use Maven Central explicitly because the expected values below read from that URL. val artifactID = "rapids-4-spark_2.12" val baseURL = ToolTestUtils.RAPIDS_MVN_BASE_URL - val nvReleases = WebCrawlerUtil.getMvnReleasesForNVPackage(artifactID) + val nvReleases = WebCrawlerUtil.getMvnReleasesForNVPackage(artifactID, Map.empty) // get all the links on the page val allLinks = WebCrawlerUtil.getPageLinks(baseURL, None).mkString("\n") val versionPattern = "(\\d{2}\\.\\d{2}\\.\\d+)/".r @@ -116,10 +120,11 @@ class ToolUtilsSuite extends AnyFunSuite with Logging { } test("get latest release") { - // use mvn repo url got testing + // Use Maven Central explicitly because the expected value below reads from that URL. val artifactID = "rapids-4-spark_2.12" val baseURL = ToolTestUtils.RAPIDS_MVN_BASE_URL - val latestRelease = WebCrawlerUtil.getLatestMvnReleaseForNVPackage(artifactID) match { + val latestRelease = WebCrawlerUtil.getLatestMvnReleaseForNVPackage( + artifactID, Map.empty) match { case Some(v) => v case None => fail("Could not find pull the latest release successfully") } @@ -132,6 +137,58 @@ class ToolUtilsSuite extends AnyFunSuite with Logging { latestRelease shouldBe actualRelease } + test("Maven base URL defaults to Maven Central") { + WebCrawlerUtil.getMavenBaseUrl(Map.empty) shouldBe "https://repo1.maven.org/maven2" + } + + test("Maven base URL can be overridden by environment") { + WebCrawlerUtil.getMavenBaseUrl( + Map("RAPIDS_TOOLS_MAVEN_BASE_URL" -> "https://mirror.example/maven/")) shouldBe + "https://mirror.example/maven" + } + + test("Maven artifact URL uses default Maven Central when environment is unset") { + assume(sys.env.get("RAPIDS_TOOLS_MAVEN_BASE_URL").forall(_.trim.isEmpty)) + WebCrawlerUtil.getMVNArtifactURL("rapids.plugin") shouldBe + "https://repo1.maven.org/maven2/com/nvidia/rapids-4-spark_2.12" + } + + test("Maven metadata request uses configured basic auth") { + val user = "maven-user" + val password = "maven-password" + val token = Base64.getEncoder.encodeToString( + s"$user:$password".getBytes(StandardCharsets.UTF_8)) + val expectedAuth = s"Basic $token" + val server = HttpServer.create(new InetSocketAddress("localhost", 0), 0) + + try { + server.createContext("/com/nvidia/test-artifact/maven-metadata.xml", exchange => { + val response = + if (exchange.getRequestHeaders.getFirst("Authorization") == expectedAuth) { + exchange.sendResponseHeaders(200, 0) + "26.04.0" + } else { + exchange.sendResponseHeaders(401, 0) + "unauthorized" + } + val bytes = response.getBytes(StandardCharsets.UTF_8) + exchange.getResponseBody.write(bytes) + exchange.close() + }) + server.start() + + val env = Map( + "RAPIDS_TOOLS_MAVEN_BASE_URL" -> s"http://localhost:${server.getAddress.getPort}", + "RAPIDS_TOOLS_MAVEN_USERNAME" -> user, + "RAPIDS_TOOLS_MAVEN_PASSWORD" -> password) + + WebCrawlerUtil.getLatestMvnReleaseForNVPackage("test-artifact", env) shouldBe + Some("26.04.0") + } finally { + server.stop(0) + } + } + test("Hadoop Configuration should load system properties") { // Tests that Hadoop configurations can load the system property passed to the // command line. i.e., "-Drapids.tools.hadoop.property.key=value" diff --git a/user_tools/src/spark_rapids_pytools/rapids/rapids_tool.py b/user_tools/src/spark_rapids_pytools/rapids/rapids_tool.py index d1c4555e1..9cbd77e2e 100644 --- a/user_tools/src/spark_rapids_pytools/rapids/rapids_tool.py +++ b/user_tools/src/spark_rapids_pytools/rapids/rapids_tool.py @@ -713,7 +713,8 @@ def cache_single_dependency(dep: RuntimeDependency) -> str: download_configs['timeOut'] = default_download_timeout if dep.verification is not None: verify_opts = dict(dep.verification) - download_task = DownloadTask(src_url=dep.uri, # pylint: disable=no-value-for-parameter) + src_url = Utilities.resolve_maven_url(dep.uri) + download_task = DownloadTask(src_url=src_url, # pylint: disable=no-value-for-parameter) dest_folder=dest_folder, verification=verify_opts, configs=download_configs) diff --git a/user_tools/src/spark_rapids_pytools/rapids/tool_ctxt.py b/user_tools/src/spark_rapids_pytools/rapids/tool_ctxt.py index 8d2a9a6ad..20f17205b 100644 --- a/user_tools/src/spark_rapids_pytools/rapids/tool_ctxt.py +++ b/user_tools/src/spark_rapids_pytools/rapids/tool_ctxt.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023-2025, NVIDIA CORPORATION. +# Copyright (c) 2023-2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -267,7 +267,7 @@ def get_rapids_jar_url(self) -> str: if self.use_local_tools_jar(): return self._get_tools_jar_from_local() self.logger.info('Tools JAR not found in local context. Downloading from Maven.') - mvn_base_url = self.get_value('sparkRapids', 'mvnUrl') + mvn_base_url = Utilities.resolve_maven_url(self.get_value('sparkRapids', 'mvnUrl')) jar_version = Utilities.get_latest_mvn_jar_from_metadata(mvn_base_url) rapids_url = self.get_value('sparkRapids', 'repoUrl').format(mvn_base_url, jar_version, jar_version) return rapids_url diff --git a/user_tools/src/spark_rapids_pytools/resources/dev/prepackage_mgr.py b/user_tools/src/spark_rapids_pytools/resources/dev/prepackage_mgr.py index 8bd902ae1..c0cae722f 100644 --- a/user_tools/src/spark_rapids_pytools/resources/dev/prepackage_mgr.py +++ b/user_tools/src/spark_rapids_pytools/resources/dev/prepackage_mgr.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023-2025, NVIDIA CORPORATION. +# Copyright (c) 2023-2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -91,8 +91,9 @@ def __init__(self, self.dest_dir = FSUtil.get_abs_path(self.dest_dir) def _get_spark_rapids_jar_url(self) -> str: - jar_version = Utilities.get_latest_mvn_jar_from_metadata(self._mvn_base_url) # pylint: disable=no-member - return (f'{self._mvn_base_url}/' # pylint: disable=no-member + mvn_base_url = Utilities.resolve_maven_url(self._mvn_base_url) # pylint: disable=no-member + jar_version = Utilities.get_latest_mvn_jar_from_metadata(mvn_base_url) + return (f'{mvn_base_url}/' f'{jar_version}/rapids-4-spark-tools_2.12-{jar_version}.jar') def _fetch_resources(self) -> dict: @@ -147,7 +148,7 @@ def _download_resources(self, resource_uris: dict): dest_folder = self.tools_resources_dir if is_tools_resource else self.dest_dir print(f'Creating download task: {resource_name}') # All the downloadTasks enforces download - download_tasks.append(DownloadTask(src_url=res_uri, # pylint: disable=no-value-for-parameter) + download_tasks.append(DownloadTask(src_url=Utilities.resolve_maven_url(res_uri), # pylint: disable=no-value-for-parameter dest_folder=dest_folder, configs={'forceDownload': True})) # Begin downloading the resources diff --git a/user_tools/src/spark_rapids_tools/utils/net_utils.py b/user_tools/src/spark_rapids_tools/utils/net_utils.py index 2a2b4ab8d..5ce514355 100644 --- a/user_tools/src/spark_rapids_tools/utils/net_utils.py +++ b/user_tools/src/spark_rapids_tools/utils/net_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -41,6 +41,7 @@ from spark_rapids_tools import CspPath from spark_rapids_tools.storagelib import CspFs from spark_rapids_tools.storagelib.tools.fs_utils import FileVerificationResult +from spark_rapids_tools.utils.util import Utilities def download_url_request(url: str, fpath: str, timeout: float = None, @@ -65,7 +66,8 @@ def download_url_request(url: str, fpath: str, timeout: float = None, # 3. This is why we need our own timeout check for the total download time # 4. We check timeout after writing each chunk, but allow the last chunk to complete start_time = time.time() - with requests.get(url, stream=True, timeout=timeout) as r: + with requests.get(url, stream=True, timeout=timeout, + headers=Utilities.get_maven_http_headers(url)) as r: r.raise_for_status() with open(fpath, 'wb') as f: # Set chunk size to 16 MB to lower the count of iterations. @@ -92,7 +94,10 @@ def download_url_urllib(url: str, fpath: str) -> str: logging.getLogger('urllib3').setLevel(logging.WARNING) # We create a context here to fix and issue with urlib requests issue. context = ssl.create_default_context(cafile=certifi.where()) - with urllib.request.urlopen(url, context=context) as resp: + request = urllib.request.Request(url) + for key, value in Utilities.get_maven_http_headers(url).items(): + request.add_header(key, value) + with urllib.request.urlopen(request, context=context) as resp: with open(fpath, 'wb') as f: shutil.copyfileobj(resp, f, 64 * 1024 * 1024) return fpath diff --git a/user_tools/src/spark_rapids_tools/utils/util.py b/user_tools/src/spark_rapids_tools/utils/util.py index cccdb2776..1a28c5017 100644 --- a/user_tools/src/spark_rapids_tools/utils/util.py +++ b/user_tools/src/spark_rapids_tools/utils/util.py @@ -14,6 +14,7 @@ """Utility and helper methods""" +import base64 import os from contextlib import contextmanager from pathlib import Path @@ -289,6 +290,61 @@ class Utilities: # Flag used to disable running tools in parallel. This is a temporary hack to reduce possibility # of OOME. Later we can re-enable it. conc_mode_enabled: ClassVar[bool] = False + maven_central_base_url: ClassVar[str] = 'https://repo1.maven.org/maven2' + maven_base_url_env: ClassVar[str] = 'RAPIDS_TOOLS_MAVEN_BASE_URL' + maven_username_env: ClassVar[str] = 'RAPIDS_TOOLS_MAVEN_USERNAME' + maven_password_env: ClassVar[str] = 'RAPIDS_TOOLS_MAVEN_PASSWORD' + + @classmethod + def get_maven_base_url(cls) -> str: + env_value = os.environ.get(cls.maven_base_url_env) + if env_value is None or env_value.strip() == '': + return cls.maven_central_base_url + return env_value.strip().rstrip('/') + + @classmethod + def resolve_maven_url(cls, url: str) -> str: + if not isinstance(url, str): + return url + default_base = cls.maven_central_base_url + configured_base = cls.get_maven_base_url() + if configured_base == default_base: + return url + if url == default_base: + return configured_base + if url.startswith(f'{default_base}/'): + return f'{configured_base}/{url[len(default_base) + 1:]}' + return url + + @classmethod + def is_maven_url(cls, url: str) -> bool: + if not isinstance(url, str): + return False + maven_base_url = cls.get_maven_base_url() + return url == maven_base_url or url.startswith(f'{maven_base_url}/') + + @classmethod + def get_maven_basic_auth_header(cls) -> Optional[str]: + user = os.environ.get(cls.maven_username_env) + password = os.environ.get(cls.maven_password_env) + if user is None or user.strip() == '' or password is None or password == '': + return None + token = base64.b64encode(f'{user.strip()}:{password}'.encode()).decode() + return f'Basic {token}' + + @classmethod + def get_maven_http_headers(cls, url: str) -> dict: + auth_header = cls.get_maven_basic_auth_header() + if auth_header is None or not cls.is_maven_url(url): + return {} + return {'Authorization': auth_header} + + @classmethod + def build_maven_url_request(cls, url: str): + request = urllib.request.Request(url) + for header, value in cls.get_maven_http_headers(url).items(): + request.add_header(header, value) + return request @classmethod def get_latest_mvn_jar_from_metadata(cls, url_base: str, @@ -314,7 +370,7 @@ def get_latest_mvn_jar_from_metadata(cls, url_base: str, defined_version = Version(loaded_version) jar_version = Version(loaded_version) xml_path = f'{url_base}/maven-metadata.xml' - with urllib.request.urlopen(xml_path, context=context) as resp: + with urllib.request.urlopen(cls.build_maven_url_request(xml_path), context=context) as resp: xml_content = resp.read() xml_root = elem_tree.fromstring(xml_content) for version_elem in xml_root.iter('version'): diff --git a/user_tools/tests/spark_rapids_tools_ut/test_maven_url.py b/user_tools/tests/spark_rapids_tools_ut/test_maven_url.py new file mode 100644 index 000000000..341cf9ef2 --- /dev/null +++ b/user_tools/tests/spark_rapids_tools_ut/test_maven_url.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for Maven URL overrides and Maven metadata authentication.""" + +import base64 +from unittest.mock import patch + +from spark_rapids_tools.utils import Utilities +from spark_rapids_pytools.resources.dev.prepackage_mgr import PrepackageMgr + + +def test_resolve_maven_url_keeps_default_when_env_missing(monkeypatch): + monkeypatch.delenv('RAPIDS_TOOLS_MAVEN_BASE_URL', raising=False) + assert Utilities.resolve_maven_url( + 'https://repo1.maven.org/maven2/com/nvidia/artifact') == ( + 'https://repo1.maven.org/maven2/com/nvidia/artifact') + + +def test_resolve_maven_url_uses_env_for_maven_central(monkeypatch): + monkeypatch.setenv('RAPIDS_TOOLS_MAVEN_BASE_URL', 'https://mirror.example/maven/') + assert Utilities.resolve_maven_url( + 'https://repo1.maven.org/maven2/com/nvidia/artifact') == ( + 'https://mirror.example/maven/com/nvidia/artifact') + + +def test_resolve_maven_url_ignores_non_maven_urls(monkeypatch): + monkeypatch.setenv('RAPIDS_TOOLS_MAVEN_BASE_URL', 'https://mirror.example/maven') + assert Utilities.resolve_maven_url( + 'https://archive.apache.org/dist/spark/spark.tgz') == ( + 'https://archive.apache.org/dist/spark/spark.tgz') + + +def test_prepackage_tools_jar_url_uses_maven_base_env(monkeypatch, tmp_path): + monkeypatch.setenv('RAPIDS_TOOLS_MAVEN_BASE_URL', 'https://mirror.example/maven/') + monkeypatch.setattr(Utilities, 'get_latest_mvn_jar_from_metadata', lambda url_base: '26.04.0') + mgr = PrepackageMgr(resource_dir=str(tmp_path), archive_enabled=False) + + assert mgr._get_spark_rapids_jar_url() == ( # pylint: disable=protected-access + 'https://mirror.example/maven/com/nvidia/rapids-4-spark-tools_2.12/' + '26.04.0/rapids-4-spark-tools_2.12-26.04.0.jar') + + +class _FakeMetadataResponse: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + del exc_type, exc_value, traceback + return False + + def read(self): + return b''' + + + + 26.02.0 + 26.04.0 + + + + ''' + + +def test_latest_mvn_jar_metadata_uses_configured_basic_auth(monkeypatch): + user = 'maven-user' + password = 'maven-password' + maven_base_url = 'https://mirror.example/maven' + monkeypatch.setenv('RAPIDS_TOOLS_MAVEN_BASE_URL', maven_base_url) + monkeypatch.setenv('RAPIDS_TOOLS_MAVEN_USERNAME', user) + monkeypatch.setenv('RAPIDS_TOOLS_MAVEN_PASSWORD', password) + expected_auth = 'Basic ' + base64.b64encode(f'{user}:{password}'.encode()).decode() + + def fake_urlopen(request, context): # pylint: disable=unused-argument + assert request.full_url == f'{maven_base_url}/com/nvidia/rapids-4-spark-tools_2.12/maven-metadata.xml' + assert request.get_header('Authorization') == expected_auth + return _FakeMetadataResponse() + + with patch('urllib.request.urlopen', side_effect=fake_urlopen): + assert Utilities.get_latest_mvn_jar_from_metadata( + f'{maven_base_url}/com/nvidia/rapids-4-spark-tools_2.12', + loaded_version='26.04.3') == '26.04.0' diff --git a/user_tools/tests/utils/test_net_utils.py b/user_tools/tests/utils/test_net_utils.py index c27006d13..8ae4f7b42 100644 --- a/user_tools/tests/utils/test_net_utils.py +++ b/user_tools/tests/utils/test_net_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,6 +14,7 @@ """This file include unit-test for utilities related to network/downloads""" +import base64 import time import pytest from unittest.mock import patch @@ -76,6 +77,27 @@ def test_normal_download(temp_file, mock_response): assert content == b'chunk1chunk2chunk3chunk4' +def test_maven_download_uses_configured_basic_auth(monkeypatch, temp_file, mock_response): + """Test Maven downloads include credentials when Maven auth env vars are configured.""" + user = 'maven-user' + password = 'maven-password' + monkeypatch.setenv('RAPIDS_TOOLS_MAVEN_BASE_URL', 'https://mirror.example/maven') + monkeypatch.setenv('RAPIDS_TOOLS_MAVEN_USERNAME', user) + monkeypatch.setenv('RAPIDS_TOOLS_MAVEN_PASSWORD', password) + expected_auth = 'Basic ' + base64.b64encode(f'{user}:{password}'.encode()).decode() + + def fake_get(url, stream, timeout, headers): # pylint: disable=unused-argument + assert headers == {'Authorization': expected_auth} + return mock_response + + with patch('requests.get', side_effect=fake_get): + result = download_url_request( + 'https://mirror.example/maven/com/nvidia/artifact.jar', + temp_file, + timeout=10) + assert result == temp_file + + def test_timeout_before_last_chunk(temp_file, slow_mock_response): """Test timeout before the last chunk.""" with patch('requests.get', return_value=slow_mock_response):