overflowdb.formats.ExporterMain.scala Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of overflowdb-formats_3 Show documentation
overflowdb-formats
The newest version!
package overflowdb.formats

import org.slf4j.LoggerFactory
import overflowdb.formats.dot.DotExporter
import overflowdb.formats.graphml.GraphMLExporter
import overflowdb.formats.graphson.GraphSONExporter
import overflowdb.formats.neo4jcsv.Neo4jCsvExporter
import overflowdb.{EdgeFactory, Graph, NodeFactory}
import scopt.OParser

import java.io.File
import java.nio.file.{Files, Path, Paths}
import scala.jdk.CollectionConverters.SeqHasAsJava
import scala.util.Using

/** Base functionality to export a given OverflowDB graph to various export formats. Because ODB relies on domain
  * specific implementations, specifically the NodeFactories and EdgeFactories from the domain-specific generated
  * classes (typically generated by by https://github.com/ShiftLeftSecurity/overflowdb-codegen) need to be passed in.
  */
object ExporterMain {
  lazy val logger = LoggerFactory.getLogger(getClass)

  def apply(nodeFactories: Seq[NodeFactory[_]], edgeFactories: Seq[EdgeFactory[_]]): Array[String] => Unit = { args =>
    OParser
      .parse(parser, args, Config(Paths.get("/dev/null"), null, Paths.get("/dev/null")))
      .map { case Config(inputFile, format, outputFile) =>
        if (Files.notExists(inputFile))
          throw new AssertionError(s"given input file $inputFile does not exist")
        if (Files.exists(outputFile)) {
          if (Files.isRegularFile(outputFile))
            throw new AssertionError(s"output file $outputFile already exists and is not a directory")
        } else {
          Files.createDirectories(outputFile)
        }

        val exporter: Exporter = format match {
          case Format.Neo4jCsv => Neo4jCsvExporter
          case Format.GraphML  => GraphMLExporter
          case Format.GraphSON => GraphSONExporter
          case Format.Dot      => DotExporter
        }
        val odbConfig = overflowdb.Config.withoutOverflow.withStorageLocation(inputFile)
        logger.info(s"starting export of graph in $inputFile to storagePath=$outputFile in format=$format")
        val ExportResult(nodeCount, edgeCount, files, additionalInfo) =
          Using.resource(Graph.open(odbConfig, nodeFactories.asJava, edgeFactories.asJava)) { graph =>
            exporter.runExport(graph, outputFile)
          }
        logger.info(s"export completed successfully: $nodeCount nodes, $edgeCount edges in ${files.size} files")
        additionalInfo.foreach(logger.info)
      }
  }

  private lazy val builder = OParser.builder[Config]
  private lazy val parser = {
    import builder._
    OParser.sequence(
      programName("odb-export"),
      help("help").text("prints this usage text"),
      opt[String]('f', "format")
        .required()
        .action((x, c) => c.copy(format = Format.byNameLowercase(x)))
        .text(s"export format, one of [${Format.valuesAsStringLowercase.mkString("|")}]"),
      opt[File]('o', "out") // will be able to read a `Path` with scopt 4.0.2+ (once released)
        .required()
        .action((x, c) => c.copy(outputFile = x.toPath))
        .text("output file or directory - must exist and be writable"),
      arg[File]("odbBinaryFile")
        .required()
        .action((x, c) => c.copy(inputFile = x.toPath))
        .text("input overflowdb graph file - must exist and be readable")
    )
  }

  private case class Config(inputFile: Path, format: Format.Value, outputFile: Path)
}