casualjim
2/9/2012 - 11:27 PM

assembly.scala

import collection.mutable.ListBuffer
import sbt._
import Keys._
import Project.Initialize
import collection.JavaConversions._
import com.amazonaws.auth.{BasicAWSCredentials, AWSCredentials}
import com.amazonaws.services.s3.AmazonS3Client
import com.amazonaws.services.s3.transfer.TransferManager

object S3 {

  val S3 = config("s3") extend(Runtime) hide

  val amazonAccessKeyId = SettingKey[String]("amazon-access-key-id")
  val amazonSecretKeyId = SettingKey[String]("amazon-secret-key-id")
  val s3Bucket          = SettingKey[String]("amazon-s3-bucket")
  val s3Region          = SettingKey[String]("amazon-region")
  val s3Settings        = SettingKey[S3Settings]("amazon-s3-settings")
  val s3Credentials     = SettingKey[AWSCredentials]("amazon-s3-credentials")
  val filesToUploadToS3 = SettingKey[Seq[File]]("amazon-files-to-upload")

  val listBucket        = TaskKey[Seq[String]]("amazon-list-bucket")
  val uploadFilesToS3   = TaskKey[Seq[File]]("amazon-upload-files")

  case class S3Settings(credentials: AWSCredentials, bucket: String, region: String)

  private def uploadToS3: Initialize[Task[Seq[File]]] = {
    (s3Settings, filesToUploadToS3, streams) map uploadFiles
  }

  def uploadFiles(set: S3Settings, files: Traversable[File], s: TaskStreams) = {
    val tm = new TransferManager(set.credentials)
    val uploadedFiles = ListBuffer[File]()
    files foreach { f =>
      try {
        val curr = tm.upload(set.bucket, f.getName, f)
        s.log.info("Uploading: " + f.getAbsolutePath)
        curr.waitForCompletion()
        uploadedFiles += f
      } catch {
        case e => {
          s.log error (e.getMessage + "\n" + e.getStackTraceString)
        }
      }
    }
    uploadedFiles.toSeq
  }

  private def listBucketTask: Initialize[Task[Seq[String]]] = {
    (s3Settings, streams) map { (s, st) =>
      try {
        val client = new AmazonS3Client(s.credentials)
        client setEndpoint (s.region + ".amazonaws.com")
        val listing = client listObjects s.bucket
        listing.getObjectSummaries map (_.getKey) toSeq
      } catch {
        case e => {
          st.log error (e.getMessage + "\n" + e.getStackTraceString)
          throw e
        }
      }
    }
  }
  
  val awsAccessKeyId = {
    val aa = System.getenv("AMAZON_ACCESS_KEY_ID")
    if(aa == null) "notset" else aa
  }
  
  val awsSecretKeyId = {
    val aa = System.getenv("AMAZON_SECRET_ACCESS_KEY")
    if(aa == null) "notset" else aa
  }
  
  val defaultSettings = inConfig(S3)(Seq(
    amazonAccessKeyId := awsAccessKeyId,
    amazonSecretKeyId := awsSecretKeyId,
    s3Bucket          := "backchat-deploys",
    s3Region          := "s3-eu-west-1",
    s3Credentials     <<= (amazonAccessKeyId, amazonSecretKeyId) apply { (ak, sk) => new BasicAWSCredentials(ak, sk) },
    s3Settings        <<= (s3Credentials, s3Bucket, s3Region) apply S3Settings.apply,
    filesToUploadToS3 := Seq.empty,
    listBucket        <<= listBucketTask,
    uploadFilesToS3   <<= uploadToS3
  ))


}
import java.util.{TimeZone, Calendar}
import sbt._
import Keys._
import Assembly._
import S3._
import sbt.Project.Initialize
import util.matching.Regex
import io.backchat.sbt._

object Dist extends sbt.Plugin {


  private def interpolate(text: String, vars: Map[String, String]) =
    """\#\{([^}]+)\}""".r.replaceAllIn(text, (_: Regex.Match) match {
      case Regex.Groups(v) => vars.getOrElse(v, "")
    })

  private implicit def stringWithInterPolate(s: String) = new {
    def fill(values: (String, String)*) = interpolate(s, Map(values:_*))
  }

  val Dist                  = config("dist") extend(Runtime)

  val dist                  = TaskKey[Seq[File]]("dist", "Builds an archive with a reference config and launcher script, and uploads it to S3")
  val generateDistConfigs   = TaskKey[File]("generate-dist-configs", "Generates the configuration files for this application.")
  val createDistLauncher    = TaskKey[File]("generate-dist-launcher", "Generates a launcher script for this application.")
  val createBZip2Archive    = TaskKey[Option[File]]("create-bzip2-archive", "Creates a BZip2 archive from the generated dist")
  val createDeb             = TaskKey[(File, Option[File])]("create-deb", "Creates a debian package for the project")
  val createDist            = TaskKey[File]("create-dist", "Creates the layout to be archived")
  val createTimestamp       = TaskKey[File]("create-dist-timestamp", "Creates a file with the build date of this release")
  val createVersionFile     = TaskKey[File]("create-version-file", "Creates a file with the version number of this release")
  val copyPublicAssets      = TaskKey[File]("copy-public-assets", "Copies the public assets of the website to the dist folder")
  val createLatestFile      = TaskKey[File]("create-latest-file", "Creates a file with the name of the most recent backchat version")

  val distPath              = SettingKey[File]("dist-path")
  val distConfigPath        = SettingKey[File]("dist-config-path")
  val distLibPath           = SettingKey[File]("dist-lib-path")
  val distBinPath           = SettingKey[File]("dist-bin-path")
  val distPublicPath        = SettingKey[File]("dist-public-path")
  val distName              = SettingKey[String]("dist-name")
  val distTimestampPath     = SettingKey[File]("dist-timestamp")
  val distLatestPath        = SettingKey[File]("dist-latest-path")
  val versionFilePath       = SettingKey[File]("version-file-path")


  val archiveName           = SettingKey[String]("archive-name")
  val configSourcePath      = SettingKey[File]("config-source-path")
  val configsToInclude      = SettingKey[Map[String, String]]("configs-to-include")
  val webAppSourcePath      = SettingKey[File]("web-app-source-path")
  val templateContext       = SettingKey[Map[String, String]]("template-context")
  val elasticMappingsPath   = SettingKey[File]("elasticsearch-mappings-path")

  val defaultSettings       = S3.defaultSettings ++ Assembly.defaultSettings ++ inConfig(Dist)(Seq(
    dist <<= distTask,
    generateDistConfigs <<= generateConfigsTask,
//    createDistLauncher <<= generateLauncherTask,
    createBZip2Archive <<= createBZip2ArchiveTask,
    createDeb <<= generateDebTask,
    createTimestamp <<= createTimestampTask,
    createVersionFile <<= createVersionFileTask,
    createDist <<= createDistTask,
    copyPublicAssets <<= copyPublicAssetsTask,
    aggregate in dist := false,
    distName := "backchat",
    distPath <<= (target, distName) apply { _ / _ },
    distConfigPath <<= (distPath) apply { _ / "etc" },
    distBinPath <<=  (distPath) apply { _ / "bin" },
    distLibPath <<= (distPath) apply { _ / "lib" },
    distPublicPath <<= (distPath) apply { _ / "webapp" },
    distTimestampPath <<= (distPath) apply { _ / "build_date" },
    versionFilePath <<= (distPath) apply { _ / "VERSION" },
    distLatestPath <<= (target) apply { _ / "LATEST" },
    s3Settings <<= (s3Settings in S3.S3),
    timestamp <<= (timestamp in Assembly.Assembly),
    archiveName <<= (distName, timestamp) apply { (n, v) => "%s-%s.tar.bz2" format (n, v) },
    (jarPath in Assembly.Assembly) <<= (distLibPath, jarName) apply { _ / _ },
    jarName <<= (jarName in Assembly.Assembly),
    configSourcePath <<= (baseDirectory) apply { _ / ".." / "configs" },
    elasticMappingsPath <<= (configSourcePath) apply { _ / "mappings" },
    configsToInclude := defaultConfigsToInclude,
    webAppSourcePath <<= (sourceDirectory in Compile) apply { _ / "webapp" },
    templateContext := Map.empty
  )) ++ Seq(
    dist <<= (dist in Dist),
    aggregate in dist := false)

  private def distTask: Initialize[Task[Seq[File]]] = (createDeb, target, distLatestPath, s3Settings, streams) map { (debZip, dip, latest, set, s) =>
    val (_: File, fi: Option[File]) = debZip
    if (fi.isEmpty) sys.error("Couldn't create the bzip archive") // bail this is wrong, plain wrong
    val ar = fi.get
    val md5 = dip / (ar.getName + ".md5")
    IO.write(md5, "%s %s".format(md5sum, ar.getAbsolutePath) lines_! Git.devnull mkString)
    IO.write(latest, ar.getName)
    val files = Seq(ar, latest, md5)
    s.log.info("Files to upload: " + files.map(_.getName).mkString("[", ", ", "]"))
    S3.uploadFiles(set, files, s)
    files
  }
  
  def generateDebTask = (createBZip2Archive, target, timestamp, distPath, streams) map { (fi, tgt, vers, dip, s) =>
    val cmd = "bash scripts/gen-deb.sh %s %s %s".format(dip.getAbsolutePath, vers, tgt.getAbsolutePath)
    s.log.info("Generating deb with: %s" format cmd)
    val fp = (cmd lines_! s.log).mkString
    val re = "^Created.(.*)".r
    val deb = re findFirstMatchIn fp map (_ group 1) getOrElse "-"
    if (! re.findFirstIn(fp).isDefined) fail("Couldn't create debian archive")
    s.log.info("Generated deb with: %s" format deb)
    (file(deb), fi)
  }

  private def md5sum = System.getProperty("os.name", "generic").toLowerCase match {
    case s if (s.startsWith("mac") || s.contains("darwin")) => "md5"
    case _ => "md5sum"
  }

  private def tar = System.getProperty("os.name", "generic").toLowerCase match {
    case s if (s.startsWith("mac") || s.contains("darwin")) => "gtar"
    case _ => "tar"
  }

  private def copyPublicAssetsTask: Initialize[Task[File]] = {
    (webAppSourcePath, distPublicPath, streams) map { (pa, pp, s) =>
      s.log info "Copying public assets from %s to dist [%s] folder".format(pa.getAbsolutePath, pp.getAbsolutePath)
      IO.copyDirectory(pa, pp, overwrite = true, preserveLastModified = true)
      pp
    }
  }

  private def createTimestampTask: Initialize[Task[File]] = {
    (distTimestampPath, streams) map { (timest, s) =>
      s.log info ("Generating timestamp")
      val cal = Calendar.getInstance(TimeZone.getTimeZone("GMT"))
      IO.write(timest, cal.getTime.toString, append = false)
      timest
    }
  }

  private def createVersionFileTask: Initialize[Task[File]] = {
    (versionFilePath, timestamp, streams) map { (pth, vers, s) =>
      s.log info "Generating version file"
      IO.write(pth, "%s" format (vers), append = false)
      pth
    }
  }

  private def generateConfigsTask: Initialize[Task[File]] = {
    (configSourcePath, elasticMappingsPath, configsToInclude, distConfigPath, templateContext, streams) map {
      (csp, emp, tti, dcp, tec, s) =>
        tti map {
          case (src, tgt) => {
            s.log.info("Copying %s to dist".format(tgt))
            val f = dcp / tgt
            IO.writeLines(f, IO.readLines(csp / src) map (_.fill(tec.toSeq:_*)))
            f
          }
        }
        s.log.info("Copying elastic search mappings to dist")
        IO.copyDirectory(csp / "mappings", dcp / "mappings", true, true)
        IO.delete((dcp / "mappings" * "*.disabled").get)
        dcp
    }
  }

  private def createDistTask: Initialize[Task[File]] = {
    (target, assembly, generateDistConfigs, createTimestamp, createVersionFile, copyPublicAssets, streams) map {
      (pth, _, _, _, _, _, s) =>
        s.log.info("Built dist directory")
        pth
    }
  }

  def createBZip2ArchiveTask: Initialize[Task[Option[File]]] = {
    (createDist, distName, archiveName, streams) map { (pth, dn, ar, s) =>
      val ap = pth / ar
      IO.delete(ap)
      s.log.info("Creating archive: %s" format ap.getAbsolutePath)
      s.log.info("%s --exclude '.DS_STORE' --use-compress-prog=pbzip2 -cf %s -C %s %s".format(tar, ap, pth.getAbsolutePath, dn))
      val exitCode = "%s --exclude '.DS_STORE' --use-compress-prog=pbzip2 -cf %s -C %s %s".format(tar, ap, pth.getAbsolutePath, dn) ! s.log
      if(exitCode > 0) None else Some(ap)
    }
  }

  private val defaultConfigsToInclude = Map(
    "reference-config.conf" -> "application-reference.conf",
    "lb-connfu.xml" -> "logback.reference.xml",
    "es-empty.yml" -> "elasticsearch.reference.yml"
  )

}

import sbt._
import Keys._
import java.io.PrintWriter
import collection.mutable
import scala.io.Source
import Project.Initialize
import sbt.classpath.ClasspathUtilities
import io.backchat.sbt._

object Assembly extends sbt.Plugin {
  val Assembly = config("assembly") extend(Runtime)
  val assembly = TaskKey[File]("assembly", "Builds a single-file deployable jar.")

  val jarName           = SettingKey[String]("jar-name")
  val jarPath           = SettingKey[File]("jar-path")
//  val gitDeployPath     = SettingKey[File]("git-deploy-path")
  val excludedFiles     = SettingKey[Seq[File] => Seq[File]]("excluded-files")
  val conflictingFiles  = SettingKey[Seq[File] => Seq[File]]("conflicting-files")
  val context           = SettingKey[AssemblyContext]("assembly-context")
  val gitShortRev       = SettingKey[String]("git-short-rev")
  val timestamp         = SettingKey[String]("assembly-timestamp")

  case class AssemblyContext(
                excludedFiles: Seq[File] => Seq[File],
                conflictingFiles: Seq[File] => Seq[File],
                includeAsJar: Seq[String])

  def createJar(srcs: scala.Seq[(File, String)], jarFile: Types.Id[sbt.File], jars: Seq[File], options: Types.Id[scala.Seq[PackageOption]], cacheDir: Types.Id[File], s: Types.Id[Keys.TaskStreams]): Types.Id[sbt.File] = {
    val libp = jarFile.getParentFile
    val extra = libp.getParentFile / "libs"
    IO.delete(Seq(libp, extra))
    IO.createDirectories(Seq(libp, extra))
    IO.copy(jars map { j => (j, extra / j.getName) }, true, false)
    val config = new Package.Configuration(srcs, jarFile, options)
    Package(config, cacheDir, s.log)
    jarFile
  }

  def copyToGitDeploy(srcs: Seq[(File, String)], jars: Seq[File], gd: File, s: Types.Id[Keys.TaskStreams]) {
    val libDir = gd / "lib"
    s.log info "Copying assembly files to %s".format(libDir.getAbsolutePath)
    val extraLibs = gd / "libs"
    IO.delete(Seq(libDir, extraLibs))
    IO.createDirectories(Seq(libDir, extraLibs))
    IO.copy(jars map { j => (j, extraLibs / j.getName) }, true, false)
    IO.copy(srcs map { case (src, tgt) => (src, libDir / tgt) }, true, false)
  }

  private def assemblyTask: Initialize[Task[File]] =
    (packageOptions, baseDirectory, cacheDirectory, jarPath, dependencyClasspath in Runtime,
        fullClasspath, context, streams) map {
      (options, baseDir, cacheDir, jarFile, dcp, cp, ctx, s) =>

        val (jars, directories) = cp.map(_.data).partition(ClasspathUtilities.isArchive)
        val libDir = jarFile.getParentFile //gd / "lib"
        val bdfiles = "*.conf" | "*.yml" | "logback.xml" | "lb-staging.xml"
        println(libDir.getAbsolutePath)
        IO.delete(Seq(libDir))
        IO.createDirectories(Seq(libDir))
        IO.copy(jars map { j => (j, libDir / j.getName) }, true, false)
        val onlyFiles = (directories ** (-DirectoryFilter) --- (directories ** bdfiles)).get
        val fs = (onlyFiles x relativeTo(directories))
        IO.copy(fs map { case (f, p) => (f, libDir / p) }, true, false)
        jarFile
//        if (Git.currentIsDirty) error("The current branch %s is dirty. Commit your changes before releasing" format Git.currentBranch)
//        IO.withTemporaryDirectory { tempDir =>
//          val (srcs, jars) = assemblyPaths(tempDir, cp, dcp, ctx.excludedFiles, ctx.conflictingFiles, ctx.includeAsJar, s.log)
//          s.log.info("Preparing to copy to git repository.")
//
//          copyToGitDeploy(srcs, jars, gd, s)
//          s.log info "Creating %s".format(jarFile.getName)
//          
//          createJar(srcs, jarFile, jars, options, cacheDir, s)
//        }
    }



  private def assemblyPackageOptionsTask: Initialize[Task[Seq[PackageOption]]] =
    (packageOptions in Compile, mainClass in Assembly) map { (os, mainClass) =>
      mainClass map { s =>
        os find { o => o.isInstanceOf[Package.MainClass] } map { _ => os
        } getOrElse { Package.MainClass(s) +: os }
      } getOrElse {os}
    }

  private def assemblyExcludedFiles(base: Seq[File]): Seq[File] =
    ((base / "com" / "sun" / "syndication" / "io" / "impl" * "Atom10Parser.class" ) +++
      (base / "META-INF" * "*")).get collect {
        case f if f.getName.toLowerCase == "license" => f
        case f if f.getName.toLowerCase == "manifest.mf" => f
        case f if f.getName.toLowerCase == "atom10parser.class" => f
      }

  private def assemblyPaths(tempDir: File, classpath: Classpath, dependencies: Classpath,
      exclude: Seq[File] => Seq[File], conflicting: Seq[File] => Seq[File], toIncludeAsJar: Seq[String], log: Logger) = {

    val (libs, directories) = classpath.map(_.data).partition(ClasspathUtilities.isArchive)
//    val (depLibs, depDirs) = dependencies.map(_.data).partition(ClasspathUtilities.isArchive)
    val services = mutable.Map[String, mutable.ListBuffer[String]]()
    val includedAsJar = new mutable.ListBuffer[File]
    for(jar <- libs) {
      val jarName = jar.asFile.getName
//      if (toIncludeAsJar contains jarName) {
//        log.info("Including as jar: %s" format jarName)
//        includedAsJar += jar
//      } else {
        log.info("Including %s".format(jarName))
        val licenseFilter: NameFilter = "META-INF/LICENSE" | "licence" | "LICENSE" | "LICENSE.txt" | "mime.cache"
        IO.unzip(jar, tempDir, -licenseFilter)
  //      IO.unzip(jar, tempDir / jar.getName, licenseFilter)
        IO.delete(conflicting(Seq(tempDir)))
        val servicesDir = tempDir / "META-INF" / "services"
        if (servicesDir.asFile.exists) {
          for (service <- (servicesDir ** "*").get) {
            val serviceFile = service.asFile
            if (serviceFile.exists && serviceFile.isFile) {
              val entries = services.getOrElseUpdate(serviceFile.getName, new mutable.ListBuffer[String]())
              for (provider <- Source.fromFile(serviceFile).getLines) {
                if (!entries.contains(provider)) {
                  entries += provider
                }
              }
            }
          }
        }


        for ((service, providers) <- services) {
          log.info("Merging providers for %s".format(service))
          val serviceFile = (tempDir / "META-INF" / "services" / service).asFile
          val writer = new PrintWriter(serviceFile)
          for (provider <- providers.map { _.trim }.filter { !_.isEmpty }) {
            log.info(" - %s".format(provider))
            writer.println(provider)
          }
          writer.close()
        }
//      }
    }

    log info "preparing the final directory structure for the assembly"
    val base = tempDir +: directories
    val bdfiles = "*.conf" | "*.yml" | "logback.xml" | "lb-staging.xml"
    val descendants = ((base ** (-DirectoryFilter)) --- exclude(base) --- (base ** bdfiles)).get //++ includedAsJar.toSeq
    (descendants x relativeTo(base), includedAsJar.toSeq)
  }

  val defaultSettings = inConfig(Assembly)(Seq(
    assembly <<= packageBin,
    mainClass <<= (mainClass in Runtime),
    fullClasspath <<= (fullClasspath in Runtime),
    (aggregate in assembly) := false,
    packageBin <<= assemblyTask,
    gitShortRev := Git.shortRev,
    //gitDeployPath := new File(System.getenv("BACKCHAT_HOME")) / ".." / "deploys",
    timestamp := new java.text.SimpleDateFormat("yyyyMMddHHmmss").format(new java.util.Date),
    jarName <<= (timestamp) apply { ts => "backchat-%s.jar" format ts },
    jarPath <<= (target, jarName) apply { (t, s) => t / s },
    packageOptions <<= assemblyPackageOptionsTask,
//    dependencyClasspath in assembly <<= dependencyClasspath or (dependencyClasspath in Runtime),
    context := AssemblyContext(
      assemblyExcludedFiles _,
      assemblyExcludedFiles _,
      Seq("elasticsearch-cloud-aws-%s.jar".format(Dependencies.elasticSearchVersion)))
  )) ++
  Seq(
    assembly <<= (assembly in Assembly)
  )
}