[SPARK-59553][CONNECT] Make SparkSession.addArtifact honor spark.jars.ivySettings

common/utils/src/main/scala/org/apache/spark/util/MavenUtils.scala

@@ -22,17 +22,19 @@ import java.net.URI
 import java.text.ParseException
 import java.util.UUID
 
+import scala.jdk.CollectionConverters._
+
 import org.apache.ivy.Ivy
 import org.apache.ivy.core.LogOptions
 import org.apache.ivy.core.module.descriptor.{Artifact, DefaultDependencyDescriptor, DefaultExcludeRule, DefaultModuleDescriptor, ExcludeRule}
 import org.apache.ivy.core.module.id.{ArtifactId, ModuleId, ModuleRevisionId}
 import org.apache.ivy.core.report.{DownloadStatus, ResolveReport}
 import org.apache.ivy.core.resolve.ResolveOptions
 import org.apache.ivy.core.retrieve.RetrieveOptions
-import org.apache.ivy.core.settings.IvySettings
+import org.apache.ivy.core.settings.{IvySettings, NamedTimeoutConstraint}
 import org.apache.ivy.plugins.matcher.GlobPatternMatcher
 import org.apache.ivy.plugins.repository.file.FileRepository
-import org.apache.ivy.plugins.resolver.{ChainResolver, FileSystemResolver, IBiblioResolver}
+import org.apache.ivy.plugins.resolver.{AbstractResolver, ChainResolver, FileSystemResolver, IBiblioResolver}
 
 import org.apache.spark.SparkException
 import org.apache.spark.internal.{Logging, LogKeys}
@@ -340,6 +342,28 @@ private[spark] object MavenUtils extends Logging {
     ivySettings
   }
 
+  /** Apply bounded network timeouts to every resolver in an Ivy settings graph. */
+  private[spark] def setResolverTimeouts(
+      ivySettings: IvySettings,
+      connectTimeoutMs: Int,
+      readTimeoutMs: Int): Unit = {
+    require(connectTimeoutMs > 0, "The Ivy connection timeout must be positive")
+    require(readTimeoutMs > 0, "The Ivy read timeout must be positive")
+
+    val name = s"spark-runtime-${UUID.randomUUID()}"
+    val timeout = new NamedTimeoutConstraint(name)
+    timeout.setConnectionTimeout(connectTimeoutMs)
+    timeout.setReadTimeout(readTimeoutMs)
+    ivySettings.addConfigured(timeout)
+
+    ivySettings.getResolvers.asScala.foreach {
+      case resolver: AbstractResolver =>
+        resolver.setTimeoutConstraint(name)
+        resolver.validate()
+      case _ =>
+    }
+  }
+
   /* Set ivy settings for location of cache, if option is supplied */
   private[util] def processIvyPathArg(ivySettings: IvySettings, ivyPath: Option[String]): Unit = {
     val alternateIvyDir = ivyPath.filterNot(_.trim.isEmpty).getOrElse {

common/utils/src/main/scala/org/apache/spark/util/RuntimeDependencyResolver.scala

@@ -0,0 +1,138 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.
+ */
+
+package org.apache.spark.util
+
+import java.io.PrintStream
+import java.net.URI
+import java.nio.file.{Path, Paths}
+import java.util.concurrent.CancellationException
+import java.util.concurrent.locks.ReentrantLock
+
+import org.apache.spark.util.ArrayImplicits._
+
+/** Resolves runtime Ivy dependencies using one immutable Spark configuration snapshot. */
+private[spark] final class RuntimeDependencyResolver(
+    ivySettingsPath: Option[String],
+    configuredRepositories: Seq[String],
+    ivyPath: Option[String]) {
+  import RuntimeDependencyResolver._
+
+  /** Resolve an ivy URI. Calls are serialized because Ivy mutates process-wide state. */
+  def resolve(
+      uri: URI,
+      repositoryPolicy: RepositoryPolicy = AllowRequestedRepositories,
+      connectTimeoutMs: Int = DefaultConnectTimeoutMs,
+      readTimeoutMs: Int = DefaultReadTimeoutMs,
+      isCancelled: () => Boolean = () => false): Seq[Path] = {
+    checkCancelled(isCancelled)
+    try {
+      ivyLock.lockInterruptibly()
+    } catch {
+      case _: InterruptedException =>
+        Thread.currentThread().interrupt()
+        throw new CancellationException("Runtime Maven dependency resolution was cancelled")
+    }
+
+    try {
+      require(uri.getScheme == "ivy", s"Expected an ivy URI, found: $uri")
+      val authority = Option(uri.getAuthority).getOrElse {
+        throw new IllegalArgumentException(
+          s"Invalid Ivy URI authority in uri $uri: Expected 'org:module:version', found null.")
+      }
+      if (authority.split(":").length != 3) {
+        throw new IllegalArgumentException(
+          s"Invalid Ivy URI authority in uri $uri: " +
+            s"Expected 'org:module:version', found $authority.")
+      }
+
+      checkCancelled(isCancelled)
+      val (transitive, exclusions, requestedRepositories) = MavenUtils.parseQueryParams(uri)
+      val requested = requestedRepositories
+        .split(",")
+        .iterator
+        .map(_.trim)
+        .filter(_.nonEmpty)
+        .toImmutableArraySeq
+      val repositories = (configuredRepositories ++ repositoryPolicy.validate(requested))
+        .iterator
+        .map(_.trim)
+        .filter(_.nonEmpty)
+        .toSeq
+        .distinct
+
+      implicit val printStream: PrintStream = System.err
+      val ivySettings = ivySettingsPath.filter(_.trim.nonEmpty) match {
+        case Some(path) =>
+          MavenUtils.loadIvySettings(path, repositoriesOption(repositories), ivyPath)
+        case None => MavenUtils.buildIvySettings(repositoriesOption(repositories), ivyPath)
+      }
+      MavenUtils.setResolverTimeouts(ivySettings, connectTimeoutMs, readTimeoutMs)
+
+      checkCancelled(isCancelled)
+      val exclusionsList = exclusions
+        .split(",")
+        .iterator
+        .map(_.trim)
+        .filter(_.nonEmpty)
+        .toImmutableArraySeq
+      val result = MavenUtils.resolveMavenCoordinates(
+        authority,
+        ivySettings,
+        transitive = transitive,
+        exclusions = exclusionsList)
+      checkCancelled(isCancelled)
+      result.map(Paths.get(_))
+    } finally {
+      ivyLock.unlock()
+    }
+  }
+}
+
+private[spark] object RuntimeDependencyResolver {
+  val DefaultConnectTimeoutMs: Int = 30 * 1000
+  val DefaultReadTimeoutMs: Int = 5 * 60 * 1000
+
+  private val ivyLock = new ReentrantLock()
+
+  trait RepositoryPolicy {
+    def validate(requestedRepositories: Seq[String]): Seq[String]
+  }
+
+  object AllowRequestedRepositories extends RepositoryPolicy {
+    override def validate(requestedRepositories: Seq[String]): Seq[String] = requestedRepositories
+  }
+
+  object RejectRequestedRepositories extends RepositoryPolicy {
+    override def validate(requestedRepositories: Seq[String]): Seq[String] = {
+      if (requestedRepositories.nonEmpty) {
+        throw new IllegalArgumentException(
+          "Server-side Maven dependencies do not allow repositories from the ivy URI")
+      }
+      Nil
+    }
+  }
+
+  private def repositoriesOption(repositories: Seq[String]): Option[String] =
+    Option(repositories.mkString(",")).filter(_.nonEmpty)
+
+  private def checkCancelled(isCancelled: () => Boolean): Unit = {
+    if (isCancelled() || Thread.currentThread().isInterrupted) {
+      throw new CancellationException("Runtime Maven dependency resolution was cancelled")
+    }
+  }
+}

common/utils/src/test/scala/org/apache/spark/util/MavenUtilsSuite.scala

@@ -21,6 +21,7 @@ import java.io.{File, OutputStream, PrintStream}
 import java.net.URI
 import java.nio.charset.StandardCharsets
 import java.nio.file.{Files, Paths}
+import java.util.concurrent.CancellationException
 
 import scala.collection.mutable.ArrayBuffer
 import scala.jdk.CollectionConverters._
@@ -100,6 +101,21 @@ class MavenUtilsSuite
     }
   }
 
+  test("runtime resolver configures Ivy network timeouts") {
+    val settings = MavenUtils.buildIvySettings(None, Some(tempIvyPath))
+
+    MavenUtils.setResolverTimeouts(settings, connectTimeoutMs = 1234, readTimeoutMs = 5678)
+
+    val resolvers = settings.getDefaultResolver
+      .asInstanceOf[ChainResolver]
+      .getResolvers
+      .asScala
+      .collect { case resolver: AbstractResolver => resolver }
+    assert(resolvers.nonEmpty)
+    assert(resolvers.forall(_.getTimeoutConstraint.getConnectionTimeout == 1234))
+    assert(resolvers.forall(_.getTimeoutConstraint.getReadTimeout == 5678))
+  }
+
   test("add dependencies works correctly") {
     val md = MavenUtils.getModuleDescriptor
     val artifacts = MavenUtils.extractMavenCoordinates("com.databricks:spark-csv_2.12:0.1," +
@@ -145,6 +161,51 @@ class MavenUtilsSuite
     }
   }
 
+  test("runtime dependency resolver loads custom Ivy settings") {
+    val main = MavenCoordinate("my.runtime.lib", "mylib", "0.1")
+    IvyTestUtils.withRepository(main, None, None) { repo =>
+      val settings = Paths.get(tempIvyPath, "ivysettings.xml")
+      Files.writeString(
+        settings,
+        s"""<ivysettings>
+           |  <settings defaultResolver="runtime"/>
+           |  <resolvers>
+           |    <ibiblio name="runtime" m2compatible="true" root="$repo"/>
+           |  </resolvers>
+           |</ivysettings>""".stripMargin)
+      val resolver = new RuntimeDependencyResolver(
+        ivySettingsPath = Some(settings.toString),
+        configuredRepositories = Nil,
+        ivyPath = Some(tempIvyPath))
+
+      val resolved = resolver.resolve(URI.create(s"ivy://${main.toString}"))
+
+      assert(resolved.exists(_.getFileName.toString.contains("my.runtime.lib_mylib-0.1")))
+      assert(resolved.forall(_.startsWith(Paths.get(tempIvyPath))))
+    }
+  }
+
+  test("runtime dependency resolver applies the repository policy") {
+    val resolver = new RuntimeDependencyResolver(None, Nil, Some(tempIvyPath))
+    val error = intercept[IllegalArgumentException] {
+      resolver.resolve(
+        URI.create("ivy://my.runtime.lib:mylib:0.1?repos=https://example.com/repository"),
+        RuntimeDependencyResolver.RejectRequestedRepositories)
+    }
+
+    assert(error.getMessage.contains("do not allow repositories"))
+  }
+
+  test("runtime dependency resolver honors cancellation before resolution") {
+    val resolver = new RuntimeDependencyResolver(None, Nil, Some(tempIvyPath))
+
+    intercept[CancellationException] {
+      resolver.resolve(
+        URI.create("ivy://my.runtime.lib:mylib:0.1"),
+        isCancelled = () => true)
+    }
+  }
+
   test("search for artifact at local repositories") {
     val main = new MavenCoordinate("my.great.lib", "mylib", "0.1")
     val dep = "my.great.dep:mydep:0.5"

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -3042,7 +3042,8 @@ package object config {
   private[spark] val JAR_IVY_SETTING_PATH =
     ConfigBuilder(MavenUtils.JAR_IVY_SETTING_PATH_KEY)
       .doc("Path to an Ivy settings file to customize resolution of jars specified " +
-        "using spark.jars.packages instead of the built-in defaults, such as maven central. " +
+        "using spark.jars.packages or ivy:// URIs passed to SparkSession.addArtifact instead " +
+        "of the built-in defaults, such as maven central. " +
         "Additional repositories given by the command-line option --repositories " +
         "or spark.jars.repositories will also be included. " +
         "Useful for allowing Spark to resolve artifacts from behind a firewall " +

python/pyspark/sql/connect/proto/base_pb2.pyi

@@ -896,14 +896,26 @@ class AnalyzePlanResponse(google.protobuf.message.Message):
         DESCRIPTOR: google.protobuf.descriptor.Descriptor
 
         VERSION_FIELD_NUMBER: builtins.int
+        CAPABILITIES_FIELD_NUMBER: builtins.int
         version: builtins.str
+        @property
+        def capabilities(
+            self,
+        ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]:
+            """Capabilities supported by the server. Capability names are versioned so clients can
+            negotiate behavior without inferring it from the Spark version or a configuration value.
+            """
         def __init__(
             self,
             *,
             version: builtins.str = ...,
+            capabilities: collections.abc.Iterable[builtins.str] | None = ...,
         ) -> None: ...
         def ClearField(
-            self, field_name: typing_extensions.Literal["version", b"version"]
+            self,
+            field_name: typing_extensions.Literal[
+                "capabilities", b"capabilities", "version", b"version"
+            ],
         ) -> None: ...
 
     class DDLParse(google.protobuf.message.Message):
@@ -2488,26 +2500,87 @@ class AddArtifactsRequest(google.protobuf.message.Message):
             self, field_name: typing_extensions.Literal["data", b"data", "name", b"name"]
         ) -> None: ...
 
+    class MavenDependency(google.protobuf.message.Message):
+        """A Maven dependency that must be resolved by the server. The URI uses Spark's existing
+        ivy://group:module:version syntax.
+        """
+
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        URI_FIELD_NUMBER: builtins.int
+        uri: builtins.str
+        def __init__(
+            self,
+            *,
+            uri: builtins.str = ...,
+        ) -> None: ...
+        def ClearField(self, field_name: typing_extensions.Literal["uri", b"uri"]) -> None: ...
+
+    class ArtifactEntry(google.protobuf.message.Message):
+        """An ordered artifact batch entry used by clients that support server-side Maven resolution."""
+
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        ARTIFACT_FIELD_NUMBER: builtins.int
+        MAVEN_DEPENDENCY_FIELD_NUMBER: builtins.int
+        @property
+        def artifact(self) -> global___AddArtifactsRequest.SingleChunkArtifact: ...
+        @property
+        def maven_dependency(self) -> global___AddArtifactsRequest.MavenDependency: ...
+        def __init__(
+            self,
+            *,
+            artifact: global___AddArtifactsRequest.SingleChunkArtifact | None = ...,
+            maven_dependency: global___AddArtifactsRequest.MavenDependency | None = ...,
+        ) -> None: ...
+        def HasField(
+            self,
+            field_name: typing_extensions.Literal[
+                "artifact", b"artifact", "maven_dependency", b"maven_dependency", "value", b"value"
+            ],
+        ) -> builtins.bool: ...
+        def ClearField(
+            self,
+            field_name: typing_extensions.Literal[
+                "artifact", b"artifact", "maven_dependency", b"maven_dependency", "value", b"value"
+            ],
+        ) -> None: ...
+        def WhichOneof(
+            self, oneof_group: typing_extensions.Literal["value", b"value"]
+        ) -> typing_extensions.Literal["artifact", "maven_dependency"] | None: ...
+
     class Batch(google.protobuf.message.Message):
         """A number of `SingleChunkArtifact` batched into a single RPC."""
 
         DESCRIPTOR: google.protobuf.descriptor.Descriptor
 
         ARTIFACTS_FIELD_NUMBER: builtins.int
+        ENTRIES_FIELD_NUMBER: builtins.int
         @property
         def artifacts(
             self,
         ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
             global___AddArtifactsRequest.SingleChunkArtifact
-        ]: ...
+        ]:
+            """Legacy artifact list used by clients that do not negotiate server-side Maven resolution."""
+        @property
+        def entries(
+            self,
+        ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
+            global___AddArtifactsRequest.ArtifactEntry
+        ]:
+            """Ordered mixture of uploaded artifacts and server-resolved Maven dependencies."""
         def __init__(
             self,
             *,
             artifacts: collections.abc.Iterable[global___AddArtifactsRequest.SingleChunkArtifact]
             | None = ...,
+            entries: collections.abc.Iterable[global___AddArtifactsRequest.ArtifactEntry]
+            | None = ...,
         ) -> None: ...
         def ClearField(
-            self, field_name: typing_extensions.Literal["artifacts", b"artifacts"]
+            self,
+            field_name: typing_extensions.Literal["artifacts", b"artifacts", "entries", b"entries"],
         ) -> None: ...
 
     class BeginChunkedArtifact(google.protobuf.message.Message):