Remove userInfo from Conda environment lists (apache#714)

* drop user info when ingesting conda list explicit uris * add test * special taste of scala * more test * fix test * more docs * 1 more docs
Alexis-D · Oct 16, 2020 · 386e92f · 386e92f
1 parent 73ce892
commit 386e92f
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 1 deletion.
diff --git a/core/src/main/scala/org/apache/spark/api/conda/CondaEnvironment.scala b/core/src/main/scala/org/apache/spark/api/conda/CondaEnvironment.scala
@@ -121,6 +121,8 @@ final class CondaEnvironment(
    *   <li>In {@code Solve} mode, list resolved packages into a specfile
    *    and use that on executors.</li>
    * </ul>
+   * Always using {@code File} mode for executors reduces conda init time by avoiding
+   * re-solving conda deps.
    */
   def buildSetupInstructions: CondaSetupInstructions = {
     bootstrapMode match {

diff --git a/core/src/main/scala/org/apache/spark/api/conda/CondaEnvironmentManager.scala b/core/src/main/scala/org/apache/spark/api/conda/CondaEnvironmentManager.scala
@@ -75,12 +75,24 @@ final class CondaEnvironmentManager(condaBinaryPath: String,
     0.until(verbosity).map(_ => "-v").toList
   }
 
+  /**
+   * List of exact uris of the packages in the solved environment, dropping any credential
+   * information (user info).
+   *
+   * This method is used by executors to obtain specfiles for repro-ing conda envs. File mode
+   * creation expects pkg uris without user info, so we need to drop them before returning.
+   * @param envDir
+   * @return List of uris
+   */
   def listPackagesExplicit(envDir: String): List[String] = {
     logInfo("Retrieving a conda environment's list of installed packages")
     val command = Process(List(condaBinaryPath, "list", "-p", envDir, "--explicit"), None)
 
     val out = runOrFail(command, "retrieving the conda installation's list of installed packages")
-    out.split("\n").filterNot(line => line.startsWith("#") || line.startsWith("@")).toList
+    out.split("\n")
+      .filterNot(line => line.startsWith("#") || line.startsWith("@"))
+      .map(CondaEnvironmentManager.dropUserInfo)
+      .toList
   }
 
   def createWithMode(
@@ -327,6 +339,14 @@ object CondaEnvironmentManager extends Logging {
     httpUrlToken.matcher(line).replaceAll("$1<password>")
   }
 
+  /**
+   * Safely dropping the <code>userInfo</code> component in URI via <code>UriBuilder</code>.
+   * <code>UriBuilder</code> can safely alter URI components without throwing exceptions.
+   */
+  private[conda] def dropUserInfo(uri: String): String = {
+    UriBuilder.fromUri(uri).userInfo(null).build().toString
+  }
+
   def fromConf(sparkConf: SparkConf): CondaEnvironmentManager = {
     val condaBinaryPath = sparkConf.get(CONDA_BINARY_PATH).getOrElse(
       sys.error(s"Expected config ${CONDA_BINARY_PATH.key} to be set"))

diff --git a/core/src/test/scala/org/apache/spark/api/conda/CondaEnvironmentManagerTest.scala b/core/src/test/scala/org/apache/spark/api/conda/CondaEnvironmentManagerTest.scala
@@ -71,4 +71,17 @@ class CondaEnvironmentManagerTest extends org.apache.spark.SparkFunSuite with Te
       "via spark.conda.bootstrapPackageUrlsUserInfo.")
         .equals(thrown.getMessage))
   }
+
+  test("CondaEnvironmentManager.dropUserInfo") {
+    val packageUrl = "https://x-5.bar/linux-64/package-0.0.1-py_0.tar.bz2"
+    assert(CondaEnvironmentManager.dropUserInfo(packageUrl) == packageUrl)
+    assert(CondaEnvironmentManager.dropUserInfo(
+      "https://a:b@x-5.bar/linux-64/package-0.0.1-py_0.tar.bz2") == packageUrl)
+    assert(CondaEnvironmentManager.dropUserInfo(
+      "https://:b@x-5.bar/linux-64/package-0.0.1-py_0.tar.bz2") == packageUrl)
+    assert(CondaEnvironmentManager.dropUserInfo(
+      "https://a:<password>@x-5.bar/linux-64/package-0.0.1-py_0.tar.bz2") == packageUrl)
+    assert(CondaEnvironmentManager.dropUserInfo(
+      "https://:Bearer bf.ghi@x-5.bar/linux-64/package-0.0.1-py_0.tar.bz2") == packageUrl)
+  }
 }