Add in-job Step DSL (Source/Transform/Sink) for Spark jobs

pipeline/src/main/scala/com/kakao/actionbase/pipeline/dsl/Job.scala

@@ -0,0 +1,71 @@
+package com.kakao.actionbase.pipeline.dsl
+
+import com.fasterxml.jackson.databind.{DeserializationFeature, ObjectMapper}
+import com.fasterxml.jackson.module.scala.{ClassTagExtensions, DefaultScalaModule}
+import org.apache.spark.sql.SparkSession
+
+import scala.reflect.ClassTag
+
+// spark-submit entry point. Subclass and implement `plan(cfg)`. `main` binds `--key=value` argv into Cfg;
+// `planFromMap` binds a YAML-shaped Map (for in-process runners that need nested fields argv can't express).
+abstract class Job[C <: Product: ClassTag] {
+
+  def plan(cfg: C): Plan.Closed
+
+  protected def configure(builder: SparkSession.Builder): SparkSession.Builder = builder
+
+  def planFromMap(args: Map[String, Any]): Plan.Closed = {
+    val cls = implicitly[ClassTag[C]].runtimeClass.asInstanceOf[Class[C]]
+    val cfg = Job.mapper.convertValue(args, cls)
+    plan(cfg)
+  }
+
+  def main(argv: Array[String]): Unit = {
+    val args = Job.parseArgv(argv)
+    println(s"Running ${getClass.getSimpleName}: $args")
+
+    val spark = configure(
+      SparkSession.builder().appName(getClass.getCanonicalName.stripSuffix("$"))
+    ).getOrCreate()
+
+    try planFromMap(args).run()(spark)
+    finally {
+      println("Stopping Spark session...")
+      spark.stop()
+    }
+  }
+}
+
+object Job {
+
+  // Lax — Spark adds its own `--spark.*` flags we don't model.
+  private[pipeline] lazy val mapper: ObjectMapper with ClassTagExtensions =
+    buildMapper(failOnUnknownProperties = false)
+
+  // Strict — YAML-authored step args; typos must surface.
+  private[pipeline] lazy val stepMapper: ObjectMapper with ClassTagExtensions =
+    buildMapper(failOnUnknownProperties = true)
+
+  private def buildMapper(failOnUnknownProperties: Boolean): ObjectMapper with ClassTagExtensions = {
+    val m = new ObjectMapper() with ClassTagExtensions
+    m.registerModule(DefaultScalaModule)
+    m.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, failOnUnknownProperties)
+    m.configure(DeserializationFeature.FAIL_ON_NULL_FOR_PRIMITIVES, true)
+    m
+  }
+
+  private[pipeline] def parseArgv(argv: Array[String]): Map[String, Any] =
+    argv.iterator.flatMap { s =>
+      if (!s.startsWith("--") || !s.contains('=')) {
+        System.err.println(s"[Job] ignoring unrecognized arg: '$s' (expected --key=value)")
+        None
+      } else {
+        s.drop(2).split("=", 2) match {
+          case Array(k, v) => Some(k -> v)
+          case _ =>
+            System.err.println(s"[Job] ignoring malformed arg: '$s' (expected --key=value)")
+            None
+        }
+      }
+    }.toMap
+}

pipeline/src/main/scala/com/kakao/actionbase/pipeline/dsl/Plan.scala

@@ -0,0 +1,135 @@
+package com.kakao.actionbase.pipeline.dsl
+
+import org.apache.spark.sql.{DataFrame, SparkSession}
+
+import java.util.IdentityHashMap
+
+// Type-state: Open = at a single output. Forked = at a Split's M outputs. Closed = ends in Sink/Fork/Group (runnable).
+sealed trait Plan {
+  private[dsl] def ast: Ast
+}
+
+object Plan {
+
+  final class Open private[dsl] (
+      private[dsl] val ast: Ast,
+      private[dsl] val edgeLabel: String = DefaultLabel
+  ) extends Plan {
+
+    // Label this output for downstream `as:` references (e.g., a SqlTransform temp view). Defaults to "_0".
+    def as(label: String): Open = new Open(ast, label)
+
+    def ~>(f: Flow): Open     = new Open(Ast.Fl(ast, f))
+    def ~>(m: Merge): Open    = new Open(Ast.Mg(Seq(edgeLabel -> ast), m))
+    def ~>(sp: Split): Forked = new Forked(Ast.Sp(ast, sp))
+    def ~>(s: Sink): Closed   = new Closed(Ast.Snk(ast, s))
+
+    def +(other: Open): MultiOpen =
+      new MultiOpen(Seq(edgeLabel -> ast, other.edgeLabel -> other.ast))
+
+    def fanOut(branches: (Open => Closed)*): Closed = {
+      require(branches.nonEmpty, "fanOut requires at least one branch")
+      val branchAsts = branches.map { b =>
+        val closed = b(new Open(Ast.Ref))
+        PlanValidation.validateBranch(closed.ast)
+        closed.ast
+      }
+      new Closed(Ast.Fork(ast, branchAsts))
+    }
+  }
+
+  final class MultiOpen private[dsl] (private[dsl] val parts: Seq[(String, Ast)]) {
+    def +(other: Open): MultiOpen = new MultiOpen(parts :+ (other.edgeLabel -> other.ast))
+    def ~>(m: Merge): Open        = new Open(Ast.Mg(parts, m))
+  }
+
+  final class Forked private[dsl] (private[dsl] val splitAst: Ast.Sp) {
+    def apply(port: String): Open = {
+      require(splitAst.sp.ports.contains(port), s"unknown port '$port'; known: ${splitAst.sp.ports.mkString(", ")}")
+      new Open(Ast.Port(splitAst, port))
+    }
+    def ports: Seq[String] = splitAst.sp.ports
+  }
+
+  final class Closed private[dsl] (private[dsl] val ast: Ast) extends Plan {
+    def run()(implicit spark: SparkSession): Unit = {
+      PlanValidation.validate(ast)
+      Executor.run(ast)
+    }
+  }
+
+  // Multiple Closeds share an execution memo, so a common upstream materializes once.
+  def bundle(closeds: Closed*): Closed = {
+    require(closeds.nonEmpty, "bundle requires at least one Closed")
+    if (closeds.size == 1) closeds.head
+    else new Closed(Ast.Group(closeds.map(_.ast)))
+  }
+
+  // Escape hatches for runners assembling the AST directly (e.g., StepsBuilder).
+  private[pipeline] def closed(ast: Ast): Closed = new Closed(ast)
+  private[pipeline] def open(ast: Ast): Open     = new Open(ast)
+}
+
+private[pipeline] sealed trait Ast
+private[pipeline] object Ast {
+  case class Src(s: Source)                              extends Ast
+  case class Fl(upstream: Ast, f: Flow)                  extends Ast
+  case class Mg(upstreams: Seq[(String, Ast)], m: Merge) extends Ast
+  case class Sp(upstream: Ast, sp: Split)                extends Ast
+  case class Port(splitNode: Sp, port: String)           extends Ast
+  case class Snk(upstream: Ast, s: Sink)                 extends Ast
+  case class Fork(upstream: Ast, branches: Seq[Ast])     extends Ast
+  case class Group(roots: Seq[Ast])                      extends Ast
+  // Placeholder for a fanOut branch's shared upstream; valid only inside a branch, where the Executor seeds the
+  // memo with Ref → the cached root. PlanValidation.validateBranch keeps it from escaping that context.
+  case object Ref extends Ast
+}
+
+private[dsl] object Executor {
+
+  // Identity-keyed: a shared upstream node materializes once; two separately-built Sources stay independent.
+  private type Memo      = IdentityHashMap[Ast, DataFrame]
+  private type SplitMemo = IdentityHashMap[Ast.Sp, Map[String, DataFrame]]
+
+  def run(ast: Ast)(implicit spark: SparkSession): Unit = runRoot(ast, new Memo(), new SplitMemo())
+
+  private def runRoot(ast: Ast, memo: Memo, splitMemo: SplitMemo)(implicit spark: SparkSession): Unit = ast match {
+    case Ast.Snk(up, sink) => sink.write(materialize(up, memo, splitMemo))
+    case Ast.Group(roots)  => roots.foreach(r => runRoot(r, memo, splitMemo))
+    case Ast.Fork(up, branches) =>
+      val df = materialize(up, memo, splitMemo).cache()
+      // Non-blocking: blocking would serialize against executor RPCs for no real benefit at job-termination.
+      try branches.foreach(runBranch(_, df))
+      finally df.unpersist(blocking = false)
+    case other => throw new IllegalStateException(s"Top-level Plan must end in a Sink, Fork, or Group, got: $other")
+  }
+
+  private def materialize(ast: Ast, memo: Memo, splitMemo: SplitMemo)(implicit spark: SparkSession): DataFrame =
+    memo.computeIfAbsent(
+      ast,
+      _ =>
+        ast match {
+          case Ast.Src(s)    => s.read()
+          case Ast.Fl(up, f) => f.apply(materialize(up, memo, splitMemo))
+          case Ast.Mg(ups, m) =>
+            m.apply(ups.map { case (label, up) => label -> materialize(up, memo, splitMemo) })
+          case Ast.Port(sp, port) =>
+            val splitMap = splitMemo.computeIfAbsent(sp, _ => sp.sp.split(materialize(sp.upstream, memo, splitMemo)))
+            splitMap.getOrElse(
+              port,
+              throw new IllegalStateException(s"unknown port '$port'; got: ${splitMap.keys.mkString(", ")}")
+            )
+          case Ast.Ref => throw new IllegalStateException("Ast.Ref outside fanOut branch — invalid Plan")
+          case other   => throw new IllegalStateException(s"materialize cannot handle: $other")
+        }
+    )
+
+  // Branches share the fork's cached root via a pre-seeded memo (Ast.Ref → root).
+  private def runBranch(ast: Ast, root: DataFrame)(implicit spark: SparkSession): Unit = ast match {
+    case Ast.Snk(up, sink) =>
+      val memo = new Memo()
+      memo.put(Ast.Ref, root)
+      sink.write(materialize(up, memo, new SplitMemo()))
+    case other => throw new IllegalStateException(s"fanOut branch must end in Sink, got: $other")
+  }
+}

pipeline/src/main/scala/com/kakao/actionbase/pipeline/dsl/PlanValidation.scala

@@ -0,0 +1,41 @@
+package com.kakao.actionbase.pipeline.dsl
+
+// Structural plan-time checks shared by both front-ends: the DSL runs `validate` from `Closed.run` and reuses
+// `validateBranch` for eager `fanOut` errors; the YAML builder produces a `Closed` that hits the same `validate`.
+// A single definition here keeps the two paths from drifting (their other checks are type- vs. string-specific).
+private[pipeline] object PlanValidation {
+
+  // Walk the finished AST and apply every whole-plan invariant before execution.
+  def validate(ast: Ast): Unit = ast match {
+    case Ast.Src(_)       =>
+    case Ast.Ref          =>
+    case Ast.Fl(up, _)    => validate(up)
+    case Ast.Mg(ups, _)   => ups.foreach { case (_, up) => validate(up) }
+    case Ast.Sp(up, _)    => validate(up)
+    case Ast.Port(sp, _)  => validate(sp)
+    case Ast.Snk(up, _)   => validate(up)
+    case Ast.Group(roots) => roots.foreach(validate)
+    case Ast.Fork(up, branches) =>
+      validate(up)
+      branches.foreach(validateBranch)
+  }
+
+  // A fanOut branch consumes the shared upstream and terminates in a Sink; it cannot start a new Source.
+  def validateBranch(ast: Ast): Unit = ast match {
+    case Ast.Snk(up, _) => validateBranchBody(up)
+    case other =>
+      throw new IllegalArgumentException(s"fanOut branch must end in a Sink, got ${other.getClass.getSimpleName}")
+  }
+
+  private def validateBranchBody(ast: Ast): Unit = ast match {
+    case Ast.Ref        =>
+    case Ast.Fl(up, _)  => validateBranchBody(up)
+    case Ast.Mg(ups, _) => ups.foreach { case (_, up) => validateBranchBody(up) }
+    case Ast.Src(_) =>
+      throw new IllegalArgumentException(
+        "fanOut branch cannot introduce a new Source; consume the shared upstream only"
+      )
+    case other =>
+      throw new IllegalArgumentException(s"fanOut branch cannot contain ${other.getClass.getSimpleName}")
+  }
+}

pipeline/src/main/scala/com/kakao/actionbase/pipeline/dsl/Step.scala

@@ -0,0 +1,33 @@
+package com.kakao.actionbase.pipeline.dsl
+
+import org.apache.spark.sql.{DataFrame, SparkSession}
+
+// Shapes: Source 0→1, Flow 1→1, Merge N→1, Split 1→M, Sink 1→0.
+sealed trait Step
+
+trait Source extends Step {
+  def read()(implicit spark: SparkSession): DataFrame
+}
+
+// DataFrame-to-DataFrame middle of the pipeline, varying only by arity (Flow/Merge/Split).
+sealed trait Transform extends Step
+
+trait Flow extends Transform {
+  def apply(in: DataFrame)(implicit spark: SparkSession): DataFrame
+}
+
+// Inputs arrive as (label, df) so consumers (e.g., SqlTransform) can address each by its producer's `as:` label.
+trait Merge extends Transform {
+  def apply(inputs: Seq[(String, DataFrame)])(implicit spark: SparkSession): DataFrame
+  def apply(in: DataFrame)(implicit spark: SparkSession): DataFrame = apply(Seq(DefaultLabel -> in))
+}
+
+// Ports declared up-front so the DSL validates references at plan time. Body runs once; ports share the result.
+trait Split extends Transform {
+  def ports: Seq[String]
+  def split(in: DataFrame)(implicit spark: SparkSession): Map[String, DataFrame]
+}
+
+trait Sink extends Step {
+  def write(in: DataFrame)(implicit spark: SparkSession): Unit
+}

pipeline/src/main/scala/com/kakao/actionbase/pipeline/dsl/package.scala

@@ -0,0 +1,13 @@
+package com.kakao.actionbase.pipeline
+
+import scala.language.implicitConversions
+
+package object dsl {
+
+  // The implicit edge label for an unlabeled output: the linear-chain default and a Merge's single-input view name.
+  private[pipeline] val DefaultLabel: String = "_0"
+
+  // Lets a Source start a `~>` chain (the operator is defined on Plan.Open).
+  implicit def sourceToOpen(s: Source): Plan.Open =
+    new Plan.Open(Ast.Src(s))
+}

pipeline/src/main/scala/com/kakao/actionbase/pipeline/jobs/StepsRunnerJob.scala

@@ -0,0 +1,12 @@
+package com.kakao.actionbase.pipeline.jobs
+
+import com.kakao.actionbase.pipeline.dsl.{Job, Plan}
+import com.kakao.actionbase.pipeline.runner.StepsBuilder
+import com.kakao.actionbase.pipeline.workflow.StepSpec
+
+// Job whose Plan is assembled from a YAML `steps:` list — express a Source~>...~>Sink chain without writing a Job class.
+case class StepsRunnerCfg(steps: Seq[StepSpec])
+
+object StepsRunnerJob extends Job[StepsRunnerCfg] {
+  override def plan(cfg: StepsRunnerCfg): Plan.Closed = StepsBuilder.build(cfg.steps)
+}

pipeline/src/main/scala/com/kakao/actionbase/pipeline/runner/ClassResolver.scala

@@ -0,0 +1,83 @@
+package com.kakao.actionbase.pipeline.runner
+
+import java.io.File
+import java.net.{JarURLConnection, URL}
+import java.util.concurrent.ConcurrentHashMap
+import scala.collection.JavaConverters._
+
+// Resolves a YAML class name to a Class, loaded without running `<clinit>` so callers can reject non-Job/non-Step
+// candidates (e.g., `java.lang.Runtime`) before side-effecting init. A name resolves as (1) a fully-qualified name,
+// (2) `<root>.<name>` per root, or (3) failing those, by recursively scanning the classpath beneath each root for a
+// unique simple name. The scan is a convenience for the built-ins shipped under `…pipeline.steps.*`; third-party
+// steps live outside that package and so are reachable by FQN only. Results are memoized per (roots, name).
+object ClassResolver {
+
+  val JobRoots: Seq[String] = Seq("com.kakao.actionbase.pipeline.jobs")
+
+  // Single root; built-ins live in subpackages (source/transform/sink) found by the recursive scan.
+  val StepRoots: Seq[String] = Seq("com.kakao.actionbase.pipeline.steps")
+
+  private val cache = new ConcurrentHashMap[String, Class[_]]()
+
+  def resolve(name: String, roots: Seq[String]): Class[_] =
+    cache.computeIfAbsent(s"${roots.mkString(",")}::$name", _ => doResolve(name, roots))
+
+  private def doResolve(name: String, roots: Seq[String]): Class[_] = {
+    val direct = (name +: roots.map(r => s"$r.$name")).view.flatMap(tryLoad).headOption
+    direct.getOrElse(scanForSimpleName(name, roots))
+  }
+
+  // Distinct FQNs sharing the requested simple name. Zero → not found; one → load it; many → ambiguous, so refuse
+  // rather than pick by scan order (which varies across classpath layouts and would be silently nondeterministic).
+  private def scanForSimpleName(name: String, roots: Seq[String]): Class[_] = {
+    val matches = roots.flatMap(classNamesUnder).filter(fqn => fqn.substring(fqn.lastIndexOf('.') + 1) == name).distinct
+    matches match {
+      case Seq(only) =>
+        tryLoad(only).getOrElse(throw new ClassNotFoundException(only))
+      case Seq() =>
+        throw new ClassNotFoundException(
+          s"Cannot resolve '$name' as a full FQN or under any of: ${roots.mkString(", ")}"
+        )
+      case many =>
+        throw new IllegalArgumentException(
+          s"ambiguous step name '$name' matches ${many.mkString(", ")}; reference it by fully-qualified name"
+        )
+    }
+  }
+
+  private def loader: ClassLoader = Thread.currentThread().getContextClassLoader
+
+  private def tryLoad(fqn: String): Option[Class[_]] =
+    try Some(Class.forName(fqn, false, loader))
+    catch { case _: ClassNotFoundException => None }
+
+  // Top-level class FQNs (companion/anonymous `$` entries skipped) anywhere beneath `pkg` on the classpath.
+  private def classNamesUnder(pkg: String): Seq[String] = {
+    val path = pkg.replace('.', '/')
+    loader.getResources(path).asScala.toSeq.flatMap { url =>
+      url.getProtocol match {
+        case "file" => scanDir(new File(url.toURI), pkg)
+        case "jar"  => scanJar(url, path)
+        case _      => Nil
+      }
+    }
+  }
+
+  private def scanDir(dir: File, pkg: String): Seq[String] =
+    Option(dir.listFiles()).getOrElse(Array.empty[File]).toSeq.flatMap { f =>
+      if (f.isDirectory) scanDir(f, s"$pkg.${f.getName}")
+      else if (f.getName.endsWith(".class") && !f.getName.contains('$')) Seq(s"$pkg.${f.getName.stripSuffix(".class")}")
+      else Nil
+    }
+
+  private def scanJar(url: URL, pathPrefix: String): Seq[String] = {
+    val jar = url.openConnection().asInstanceOf[JarURLConnection].getJarFile
+    jar
+      .entries()
+      .asScala
+      .map(_.getName)
+      .filter(n => n.startsWith(pathPrefix) && n.endsWith(".class") && !n.contains('$'))
+      .map(_.stripSuffix(".class").replace('/', '.'))
+      .toSeq
+  }
+}