Support Maven mirror configuration for tools downloads

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

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)
   }
 }
 

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,14 +17,18 @@
 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._
 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)
+            "<metadata><versioning><latest>26.04.0</latest></versioning></metadata>"
+          } 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"

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)

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

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

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