Merge pull request #304 from alexarchambault/add-keep-jar-method

Add keepJar method, don't send source JARs by default

build.sc

@@ -19,10 +19,12 @@ implicit def millModuleBasePath: define.BasePath =
 trait AmmSparkPublishModule extends PublishModule {
   import mill.scalalib.publish._
   def publishVersion = T {
-    val v = VcsVersion.vcsState().format()
+    val v        = VcsVersion.vcsState().format()
     val dirtyIdx = v.indexOf("-DIRTY")
+    def endsWithCommitHash =
+      v.length > 6 && v.substring(v.length - 6).forall(c => c.isDigit || (c >= 'a' && c <= 'f'))
     if (dirtyIdx >= 0) v.take(dirtyIdx) + "-SNAPSHOT"
-    else if (v.length > 6 && v.substring(v.length - 6).forall(c => c.isDigit || (c >= 'a' && c <= 'f'))) v + "-SNAPSHOT"
+    else if (endsWithCommitHash) v + "-SNAPSHOT"
     else v
   }
   def pomSettings = PomSettings(

modules/core/src/main/scala/org/apache/spark/sql/ammonitesparkinternals/AmmoniteSparkSessionBuilder.scala

@@ -203,6 +203,8 @@ class AmmoniteSparkSessionBuilder(implicit
   private var sendSparkYarnJars0 = true
   private var sendSparkJars0     = true
   private var ignoreJars0        = Set.empty[URI]
+  private var sendSourceJars0    = false
+  private var keepJars0          = Seq.empty[URI => Boolean]
 
   def sendSparkYarnJars(force: Boolean = true): this.type = {
     sendSparkYarnJars0 = force
@@ -219,6 +221,16 @@ class AmmoniteSparkSessionBuilder(implicit
     this
   }
 
+  def sendSourceJars(send: Boolean): this.type = {
+    sendSourceJars0 = send
+    this
+  }
+
+  def keepJars(keep: URI => Boolean): this.type = {
+    keepJars0 = keepJars0 :+ keep
+    this
+  }
+
   var classServerOpt = Option.empty[AmmoniteClassServer]
 
   private def isYarn(): Boolean =
@@ -349,7 +361,11 @@ class AmmoniteSparkSessionBuilder(implicit
           .map(normalize)
           .toSet
       val nonSparkJars = jars.filter(uri => !sparkJarFileSet.contains(normalize(uri)))
-      config("spark.jars", nonSparkJars.map(_.toASCIIString).mkString(","))
+      val nonSparkJars0 =
+        if (sendSourceJars0) nonSparkJars
+        else nonSparkJars.filter(uri => !uri.toASCIIString.endsWith("-sources.jar"))
+      val finalNonSparkJars = keepJars0.foldLeft(nonSparkJars0)(_.filter(_))
+      config("spark.jars", finalNonSparkJars.map(_.toASCIIString).mkString(","))
     }
 
     if (interpApi != null)