Merge pull request #854 from AbsaOSS/pr/759

Pr/759

Author: wajda

README.md

@@ -698,10 +698,19 @@ When one of these commands occurs spline will let you know by logging a warning.
 ### Plugin API
 
 Using a plugin API you can capture lineage from a 3rd party data source provider.
-Spline discover plugins automatically by scanning a classpath, so no special steps required to register and configure a plugin.
+By default, Spline discover plugins automatically by scanning a classpath, so no special steps required to register and configure a plugin.
 All you need is to create a class extending the `za.co.absa.spline.harvester.plugin.Plugin` marker trait
 mixed with one or more `*Processing` traits, depending on your intention.
 
+To disable automatic plugin discovery and speed up initialization, set `spline.scanClasspath` to `false` in your configuration file.
+Then, you will need to register all necessary plugins one by one, using `spline.plugins.{className}.enabled` property, e.g.:
+```properties
+# Disable automatic plugin discovery to save on bootstrap time
+spline.scanClasspath=false
+# This explicitly registers and enables a plugin with the class name 'com.example.MyPlugin'
+spline.plugins.com.example.MyPlugin.enabled=true
+```
+
 There are three general processing traits:
 
 - `DataSourceFormatNameResolving` - returns a name of a data provider/format in use.

core/src/main/resources/spline.default.yaml

@@ -242,7 +242,17 @@ spline:
       # can be a json string or url to json file
       rules: { }
 
+  # ===========================================
   # Plugins configuration.
+  # ===========================================
+
+  # Classpath scanning is required for auto-discovery of Spline agent plugins.
+  # It might however take up some of the job's startup time depending on the size
+  # of your Spark application classpath.
+  # If you are experiencing long job startup times, try to disable classpath
+  # scanning, and instead register your custom plugins explicitly in the configuration.
+  scanClasspath: true # true | false
+
   # The `key` is the fully specified plugin class name.
   # For example:
   # ---
@@ -251,6 +261,58 @@ spline:
   #      prop1: foo
   #      prop2: bar
   plugins:
+
+    za.co.absa.spline.harvester.plugin.composite.LogicalRelationPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.composite.SaveIntoDataSourceCommandPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.CobrixPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.ExcelPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.BigQueryPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.DataSourceV2Plugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.RDDPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.MongoPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.AvroPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.CassandraPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.DatabricksPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.SQLPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.JDBCPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.KafkaPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.XMLPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.DeltaPlugin:
+      enabled: true
+
+    za.co.absa.spline.harvester.plugin.embedded.ElasticSearchPlugin:
+      enabled: true
+
     za.co.absa.spline.harvester.plugin.embedded.NonPersistentActionsCapturePlugin:
       enabled: false
       funcNames:

core/src/main/scala/za/co/absa/spline/agent/AgentBOM.scala

@@ -24,6 +24,7 @@ import za.co.absa.spline.harvester.IdGenerator.UUIDVersion
 import za.co.absa.spline.harvester.conf.{SQLFailureCaptureMode, SplineMode}
 import za.co.absa.spline.harvester.dispatcher.{CompositeLineageDispatcher, LineageDispatcher}
 import za.co.absa.spline.harvester.iwd.IgnoredWriteDetectionStrategy
+import za.co.absa.spline.harvester.plugin.PluginsConfiguration
 import za.co.absa.spline.harvester.postprocessing.{CompositePostProcessingFilter, PostProcessingFilter}
 
 import scala.collection.JavaConverters._
@@ -37,7 +38,7 @@ private[spline] trait AgentBOM {
   def lineageDispatcher: LineageDispatcher
   def iwdStrategy: IgnoredWriteDetectionStrategy
   def execPlanUUIDVersion: UUIDVersion
-  def pluginsConfig: Configuration
+  def pluginsConfig: PluginsConfiguration
 }
 
 object AgentBOM {
@@ -62,8 +63,11 @@ object AgentBOM {
       mergedConfig.getRequiredInt(ConfProperty.ExecPlanUUIDVersion)
     }
 
-    override def pluginsConfig: Configuration = {
-      mergedConfig.subset(ConfProperty.PluginsConfigNamespace)
+    override def pluginsConfig: PluginsConfiguration = {
+      PluginsConfiguration(
+        mergedConfig.getRequiredBoolean(ConfProperty.ScanClasspath),
+        mergedConfig.subset(ConfProperty.PluginsConfigNamespace)
+      )
     }
 
     override lazy val postProcessingFilter: Option[PostProcessingFilter] = {

core/src/main/scala/za/co/absa/spline/agent/AgentConfig.scala

@@ -81,6 +81,16 @@ object AgentConfig {
       this
     }
 
+    def scanClasspath(enabled: Boolean): this.type = synchronized {
+      options += ConfProperty.ScanClasspath -> enabled
+      this
+    }
+
+    def enablePlugin(name: String): this.type = synchronized {
+      options += s"${ConfProperty.PluginsConfigNamespace}.$name.enabled" -> true
+      this
+    }
+
     def build(): AgentConfig = new AgentConfig {
       options.foreach(tupled(addProperty))
     }
@@ -123,6 +133,11 @@ object AgentConfig {
      */
     val IgnoreWriteDetectionStrategy = "spline.IWDStrategy"
 
+    /**
+     * Should the classpath be scanned at startup
+     */
+    val ScanClasspath = "spline.scanClasspath"
+
     val PluginsConfigNamespace = "spline.plugins"
 
     def dispatcherClassName(logicalName: String): String = s"$RootLineageDispatcher.$logicalName.${HierarchicalObjectFactory.ClassName}"

core/src/main/scala/za/co/absa/spline/agent/SplineAgent.scala

@@ -16,7 +16,6 @@
 
 package za.co.absa.spline.agent
 
-import org.apache.commons.configuration.Configuration
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.execution.QueryExecution
@@ -30,6 +29,7 @@ import za.co.absa.spline.harvester.builder.write.PluggableWriteCommandExtractor
 import za.co.absa.spline.harvester.converter.{DataConverter, DataTypeConverter}
 import za.co.absa.spline.harvester.dispatcher.LineageDispatcher
 import za.co.absa.spline.harvester.iwd.IgnoredWriteDetectionStrategy
+import za.co.absa.spline.harvester.plugin.PluginsConfiguration
 import za.co.absa.spline.harvester.plugin.registry.AutoDiscoveryPluginRegistry
 import za.co.absa.spline.harvester.postprocessing._
 import za.co.absa.spline.harvester.qualifier.HDFSPathQualifier
@@ -54,7 +54,7 @@ object SplineAgent extends Logging {
   )
 
   def create(
-    pluginsConfig: Configuration,
+    pluginsConfig: PluginsConfiguration,
     session: SparkSession,
     lineageDispatcher: LineageDispatcher,
     userPostProcessingFilter: Option[PostProcessingFilter],

core/src/main/scala/za/co/absa/spline/harvester/builder/dsformat/PluggableDataSourceFormatResolver.scala

@@ -25,7 +25,8 @@ class PluggableDataSourceFormatResolver(pluginRegistry: PluginRegistry) extends
   private val processFn =
     pluginRegistry.plugins[DataSourceFormatNameResolving]
       .map(_.formatNameResolver)
-      .reduce(_ orElse _)
+      .reduceOption(_ orElse _)
+      .getOrElse(PartialFunction.empty)
       .orElse[AnyRef, String] {
         case dsr: DataSourceRegister => dsr.shortName
         case o => o.toString

core/src/main/scala/za/co/absa/spline/harvester/builder/read/PluggableReadCommandExtractor.scala

@@ -34,12 +34,14 @@ class PluggableReadCommandExtractor(
   private val planProcessFn =
     pluginRegistry.plugins[ReadNodeProcessing]
       .map(_.readNodeProcessor)
-      .reduce(_ orElse _)
+      .reduceOption(_ orElse _)
+      .getOrElse(PartialFunction.empty)
 
   private val rddProcessFn =
     pluginRegistry.plugins[RddReadNodeProcessing]
       .map(_.rddReadNodeProcessor)
-      .reduce(_ orElse _)
+      .reduceOption(_ orElse _)
+      .getOrElse(PartialFunction.empty)
 
   override def asReadCommand(planOrRdd: PlanOrRdd): Option[ReadCommand] = {
     val res = planOrRdd match {

core/src/main/scala/za/co/absa/spline/harvester/builder/write/PluggableWriteCommandExtractor.scala

@@ -38,7 +38,8 @@ class PluggableWriteCommandExtractor(
   private val processFn: ((FuncName, LogicalPlan)) => Option[WriteNodeInfo] =
     pluginRegistry.plugins[WriteNodeProcessing]
       .map(_.writeNodeProcessor)
-      .reduce(_ orElse _)
+      .reduceOption(_ orElse _)
+      .getOrElse(PartialFunction.empty)
       .lift
 
   def asWriteCommand(funcName: FuncName, logicalPlan: LogicalPlan): Option[WriteCommand] = {

core/src/main/scala/za/co/absa/spline/harvester/plugin/PluginsConfiguration.scala

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2023 ABSA Group Limited
+ *
+ * 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.
+ */
+
+package za.co.absa.spline.harvester.plugin
+
+import org.apache.commons.configuration.Configuration
+
+case class PluginsConfiguration(
+  classpathScanEnabled: Boolean,
+  config: Configuration
+)

core/src/main/scala/za/co/absa/spline/harvester/plugin/composite/LogicalRelationPlugin.scala

@@ -30,7 +30,8 @@ class LogicalRelationPlugin(pluginRegistry: PluginRegistry) extends Plugin with
   private lazy val baseRelProcessor =
     pluginRegistry.plugins[BaseRelationProcessing]
       .map(_.baseRelationProcessor)
-      .reduce(_ orElse _)
+      .reduceOption(_ orElse _)
+      .getOrElse(PartialFunction.empty)
 
   override val readNodeProcessor: PartialFunction[LogicalPlan, ReadNodeInfo] = {
     case lr: LogicalRelation

core/src/main/scala/za/co/absa/spline/harvester/plugin/composite/SaveIntoDataSourceCommandPlugin.scala

@@ -39,8 +39,8 @@ class SaveIntoDataSourceCommandPlugin(
   private lazy val rpProcessor =
     pluginRegistry.plugins[RelationProviderProcessing]
       .map(_.relationProviderProcessor)
-      .reduce(_ orElse _)
-
+      .reduceOption(_ orElse _)
+      .getOrElse(PartialFunction.empty)
 
   override def writeNodeProcessor: PartialFunction[(FuncName, LogicalPlan), WriteNodeInfo] = {
     case (_, cmd: SaveIntoDataSourceCommand) => cmd match {

core/src/main/scala/za/co/absa/spline/harvester/plugin/registry/AutoDiscoveryPluginRegistry.scala

@@ -21,9 +21,8 @@ import org.apache.commons.configuration.Configuration
 import org.apache.commons.lang.ClassUtils.{getAllInterfaces, getAllSuperclasses}
 import org.apache.spark.internal.Logging
 import za.co.absa.spline.commons.lang.ARM
-import za.co.absa.spline.harvester.plugin.Plugin
 import za.co.absa.spline.harvester.plugin.Plugin.Precedence
-import za.co.absa.spline.harvester.plugin.registry.AutoDiscoveryPluginRegistry.{EnabledByDefault, EnabledConfProperty, PluginClasses, getOnlyOrThrow}
+import za.co.absa.spline.harvester.plugin.{Plugin, PluginsConfiguration}
 
 import javax.annotation.Priority
 import scala.collection.JavaConverters._
@@ -32,11 +31,13 @@ import scala.util.Try
 import scala.util.control.NonFatal
 
 class AutoDiscoveryPluginRegistry(
-  conf: Configuration,
+  pluginsConf: PluginsConfiguration,
   injectables: AnyRef*
 ) extends PluginRegistry
   with Logging {
 
+  import za.co.absa.spline.harvester.plugin.registry.AutoDiscoveryPluginRegistry._
+
   private val injectablesByType: Map[Class[_], Seq[_ <: AnyRef]] = {
     val typedInjectables =
       for {
@@ -47,13 +48,30 @@ class AutoDiscoveryPluginRegistry(
     typedInjectables.groupBy(_._1).mapValues(_.map(_._2))
   }
 
-  private val allPlugins: Seq[Plugin] =
-    for (pc <- PluginClasses if isPluginEnabled(pc)) yield {
-      logInfo(s"Loading plugin: $pc")
-      instantiatePlugin(pc)
-        .recover({ case NonFatal(e) => throw new RuntimeException(s"Plugin instantiation failure: $pc", e) })
+  private val allPlugins: Seq[Plugin] = {
+    val discoveredClasses: Seq[Class[Plugin]] =
+      if (pluginsConf.classpathScanEnabled) scanForPluginClasses()
+      else {
+        logInfo(s"Classpath scanning is DISABLED. Only explicitly configured plugins will be loaded.")
+        Seq.empty
+      }
+
+    val configuredClasses: Seq[Class[Plugin]] = getRegisteredPluginClasses(pluginsConf.config)
+
+    val allFoundPluginClasses: Seq[Class[Plugin]] = (discoveredClasses ++ configuredClasses).distinct
+
+    val allSortedPluginClasses = allFoundPluginClasses
+      .map(c => c -> priorityOf(c))
+      .sortBy({ case (_, p) => p })
+      .map({ case (c, _) => c })
+
+    for (cls <- allSortedPluginClasses if isPluginEnabled(cls)) yield {
+      logInfo(s"Loading plugin: $cls")
+      instantiatePlugin(cls)
+        .recover({ case NonFatal(e) => throw new RuntimeException(s"Plugin instantiation failure: $cls", e) })
         .get
     }
+  }
 
   override def plugins[A: ClassTag]: Seq[Plugin with A] = {
     val ct = implicitly[ClassTag[A]]
@@ -65,7 +83,7 @@ class AutoDiscoveryPluginRegistry(
     val constr = getOnlyOrThrow(constrs, s"Plugin class must have a single public constructor: ${constrs.mkString(", ")}")
     val args = constr.getParameterTypes.map {
       case ct if classOf[Configuration].isAssignableFrom(ct) =>
-        conf.subset(pluginClass.getName)
+        pluginsConf.config.subset(pluginClass.getName)
       case pt =>
         val candidates = injectablesByType.getOrElse(pt, sys.error(s"Cannot bind $pt. No value found"))
         getOnlyOrThrow(candidates, s"Ambiguous constructor parameter binding. Multiple values found for $pt: ${candidates.length}")
@@ -74,7 +92,7 @@ class AutoDiscoveryPluginRegistry(
   }
 
   private def isPluginEnabled(pc: Class[Plugin]): Boolean = {
-    val pluginConf = conf.subset(pc.getName)
+    val pluginConf = pluginsConf.config.subset(pc.getName)
     val isEnabled = pluginConf.getBoolean(EnabledConfProperty, EnabledByDefault)
     if (!isEnabled) {
       logWarning(s"Plugin ${pc.getName} is disabled in the configuration.")
@@ -89,22 +107,33 @@ object AutoDiscoveryPluginRegistry extends Logging {
   private val EnabledConfProperty = "enabled"
   private val EnabledByDefault = true
 
-  private val PluginClasses: Seq[Class[Plugin]] = {
+  private def scanForPluginClasses(): Seq[Class[Plugin]] = {
     logDebug("Scanning for plugins")
     val classGraph = new ClassGraph().enableClassInfo
     for {
       scanResult <- ARM.managed(classGraph.scan)
-      (cls, prt) <- scanResult
+      cls <- scanResult
         .getClassesImplementing(classOf[Plugin].getName)
         .loadClasses.asScala.asInstanceOf[Seq[Class[Plugin]]]
-        .map(c => c -> priorityOf(c))
-        .sortBy(_._2)
     } yield {
-      logDebug(s"Found plugin [priority=$prt]\t: $cls")
+      logDebug(s"Discovered plugin: $cls")
       cls
     }
   }
 
+  private def getRegisteredPluginClasses(conf: Configuration): Seq[Class[Plugin]] = {
+    for {
+      key <- conf.getKeys.asScala.toSeq
+      if key.endsWith(s".$EnabledConfProperty") // Looking for keys ending with ".enabled", since plugins must be explicitly enabled
+      className = key.dropRight(EnabledConfProperty.length + 1) // Dropping ".enabled" to get plugin class name
+      cls = Class.forName(className)
+      if classOf[Plugin].isAssignableFrom(cls)
+    } yield {
+      logDebug(s"Found registered plugin: $cls")
+      cls.asInstanceOf[Class[Plugin]]
+    }
+  }
+
   private def priorityOf(c: Class[Plugin]): Int =
     Option(c.getAnnotation(classOf[Priority]))
       .map(_.value)