fix(job-server): JobSQLDAO: Move blobs to a separate table (spark-jobserver#876) (spark-jobserver#932)

msohn · noorul · commit c9269b94b0a0 · 2017-10-20T07:13:38.000+05:30
* Fix postgres database migration for schema version v0_7_4 Add missing quotes for table and column identifiers. Without these quotes the migration v0_7_4 fails on postgres. Also quote table and column identifiers for h2 and mysql. * fix(job-server): JobSQLDAO: Move blobs to a separate table (spark-jobserver#876) * Separate the BINARIES table into BINARIES and BINARIES_CONTENTS, so that only the BINARIES_CONTENTS table contains the "BINARY" (Blob) column. Both tables can be joined on the same BIN_ID. * Optimize the JobDAO.getLastUploadTimeAndType(...) method - make it abstract and implement an optimized version in the JobFileDAO and JobSqlDAO. This change is based on previous work done by Michał Januszewski in pull request spark-jobserver#886.
diff --git a/job-server/src/main/resources/db/h2/migration/V0_7_4/V0_7_4__create_error_class_and_stacktrace.sql b/job-server/src/main/resources/db/h2/migration/V0_7_4/V0_7_4__create_error_class_and_stacktrace.sql
@@ -1,2 +1,2 @@
-ALTER TABLE JOBS ADD COLUMN ERROR_CLASS VARCHAR(255);
-ALTER TABLE JOBS ADD COLUMN ERROR_STACK_TRACE CLOB;
+ALTER TABLE "JOBS" ADD COLUMN "ERROR_CLASS" VARCHAR(255);
+ALTER TABLE "JOBS" ADD COLUMN "ERROR_STACK_TRACE" CLOB;
diff --git a/job-server/src/main/resources/db/mysql/migration/V0_7_4/V0_7_4__create_error_class_and_stacktrace.sql b/job-server/src/main/resources/db/mysql/migration/V0_7_4/V0_7_4__create_error_class_and_stacktrace.sql
@@ -1,2 +1,2 @@
-ALTER TABLE JOBS ADD COLUMN ERROR_CLASS VARCHAR(255);
-ALTER TABLE JOBS ADD COLUMN ERROR_STACK_TRACE TEXT;
+ALTER TABLE `JOBS` ADD COLUMN `ERROR_CLASS` VARCHAR(255);
+ALTER TABLE `JOBS` ADD COLUMN `ERROR_STACK_TRACE` TEXT;
diff --git a/job-server/src/main/resources/db/postgresql/migration/V0_7_4/V0_7_4__create_error_class_and_stacktrace.sql b/job-server/src/main/resources/db/postgresql/migration/V0_7_4/V0_7_4__create_error_class_and_stacktrace.sql
@@ -1,2 +1,2 @@
-ALTER TABLE JOBS ADD COLUMN ERROR_CLASS VARCHAR(255);
-ALTER TABLE JOBS ADD COLUMN ERROR_STACK_TRACE TEXT;
+ALTER TABLE "JOBS" ADD COLUMN "ERROR_CLASS" VARCHAR(255);
+ALTER TABLE "JOBS" ADD COLUMN "ERROR_STACK_TRACE" TEXT;
diff --git a/job-server/src/main/scala/db/h2/migration/V0_7_5/V0_7_5__Migrate_Blobs.scala b/job-server/src/main/scala/db/h2/migration/V0_7_5/V0_7_5__Migrate_Blobs.scala
@@ -0,0 +1,38 @@
+package db.h2.migration.V0_7_5
+
+import java.sql.Blob
+import java.sql.Connection
+import javax.sql.rowset.serial.SerialBlob
+
+import scala.concurrent.Await
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.duration.DurationInt
+import scala.util.control.NonFatal
+
+import db.migration.V0_7_5.Migration
+import org.flywaydb.core.api.migration.jdbc.JdbcMigration
+import org.slf4j.LoggerFactory
+import slick.dbio.DBIO
+import slick.dbio.Effect
+import slick.dbio.NoStream
+import slick.driver.H2Driver.api.actionBasedSQLInterpolation
+import slick.jdbc.GetResult
+import slick.jdbc.PositionedParameters
+import slick.jdbc.SetParameter
+import slick.profile.SqlAction
+import spark.jobserver.slick.unmanaged.UnmanagedDatabase
+
+class V0_7_5__Migrate_Blobs extends Migration {
+  val logger = LoggerFactory.getLogger(getClass)
+
+  protected def insertBlob(id: Int, blob: SerialBlob): SqlAction[Int, NoStream, Effect] = {
+    sqlu"""INSERT INTO "BINARIES_CONTENTS" ("BIN_ID", "BINARY") VALUES (${id}, ${blob})"""
+  }
+  val createContentsTable = sqlu"""CREATE TABLE "BINARIES_CONTENTS" (
+    "BIN_ID"  BIGINT  NOT NULL PRIMARY KEY,
+    "BINARY"  BLOB
+  );"""
+  val getBinaryContents = sql"""SELECT "BIN_ID", "BINARY" FROM "BINARIES"""".as[BinaryContent]
+  val dropColumn = sqlu"""ALTER TABLE "BINARIES" DROP COLUMN "BINARY""""
+}
+
diff --git a/job-server/src/main/scala/db/migration/V0_7_5/Migration.scala b/job-server/src/main/scala/db/migration/V0_7_5/Migration.scala
@@ -0,0 +1,72 @@
+package db.migration.V0_7_5
+
+import java.sql.Blob
+import java.sql.Connection
+import javax.sql.rowset.serial.SerialBlob
+
+import scala.concurrent.Await
+import scala.concurrent.duration.DurationInt
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.util.control.NonFatal
+
+import org.flywaydb.core.api.migration.jdbc.JdbcMigration
+import org.slf4j.Logger
+import slick.dbio.DBIO
+import slick.dbio.Effect
+import slick.dbio.NoStream
+import slick.jdbc.GetResult
+import slick.jdbc.PositionedParameters
+import slick.jdbc.SetParameter
+import slick.profile.SqlAction
+import spark.jobserver.slick.unmanaged.UnmanagedDatabase
+import slick.dbio.DBIOAction
+import slick.dbio.Streaming
+
+trait Migration extends JdbcMigration {
+  protected val Timeout = 10 minutes
+  protected val logger: Logger
+
+  protected case class BinaryContent(id: Int, binary: Blob)
+
+  protected implicit object SetSerialBlob extends SetParameter[SerialBlob] {
+    def apply(v: SerialBlob, pp: PositionedParameters) {
+      pp.setBlob(v)
+    }
+  }
+  protected def insertBlob(id: Int, blob: SerialBlob): SqlAction[Int, NoStream, Effect]
+
+  protected val createContentsTable: SqlAction[Int, NoStream, Effect]
+
+  protected implicit val getBinaryResult = GetResult[BinaryContent](
+      r => BinaryContent(r.nextInt(), r.nextBlob()))
+  protected val getBinaryContents: DBIOAction[Seq[(BinaryContent)], Streaming[BinaryContent], Effect]
+
+  protected val dropColumn: SqlAction[Int, NoStream, Effect]
+
+  protected def logErrors = PartialFunction[Throwable, Unit] {
+    case e: Throwable => logger.error(e.getMessage, e)
+  }
+
+  protected def insertBlob(db: UnmanagedDatabase, b: BinaryContent): Unit = {
+    val blob = new SerialBlob(b.binary.getBytes(1, b.binary.length().toInt))
+    Await.ready(db.run(insertBlob(b.id, blob)).recover{logErrors}, Timeout)
+  }
+
+  def migrate(c: Connection): Unit = {
+    val db = new UnmanagedDatabase(c)
+    c.setAutoCommit(false);
+    try {
+      Await.ready(
+          for {
+            _ <- db.run(createContentsTable)
+            _ <- db.stream(getBinaryContents).foreach(b => insertBlob(db, b))
+            _ <- db.run(dropColumn)
+          } yield Unit, Timeout
+      ).recover{logErrors}
+      c.commit()
+    } catch {
+      case NonFatal(e) => { c.rollback() }
+    }
+  }
+}
+
diff --git a/job-server/src/main/scala/db/mysql/migration/V0_7_5/V0_7_5__Migrate_Blobs.scala b/job-server/src/main/scala/db/mysql/migration/V0_7_5/V0_7_5__Migrate_Blobs.scala
@@ -0,0 +1,38 @@
+package db.mysql.migration.V0_7_5
+
+import java.sql.Blob
+import java.sql.Connection
+import javax.sql.rowset.serial.SerialBlob
+
+import scala.concurrent.Await
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.duration.DurationInt
+import scala.util.control.NonFatal
+
+import db.migration.V0_7_5.Migration
+import org.flywaydb.core.api.migration.jdbc.JdbcMigration
+import org.slf4j.LoggerFactory
+import slick.dbio.DBIO
+import slick.dbio.Effect
+import slick.dbio.NoStream
+import slick.driver.MySQLDriver.api.actionBasedSQLInterpolation
+import slick.jdbc.GetResult
+import slick.jdbc.PositionedParameters
+import slick.jdbc.SetParameter
+import slick.profile.SqlAction
+import spark.jobserver.slick.unmanaged.UnmanagedDatabase
+
+class V0_7_5__Migrate_Blobs extends Migration {
+  val logger = LoggerFactory.getLogger(getClass)
+
+  protected def insertBlob(id: Int, blob: SerialBlob): SqlAction[Int, NoStream, Effect] = {
+    sqlu"""INSERT INTO `BINARIES_CONTENTS` (`BIN_ID`, `BINARY`) VALUES (${id}, ${blob})"""
+  }
+  val createContentsTable = sqlu"""CREATE TABLE `BINARIES_CONTENTS` (
+    `BIN_ID`  SERIAL  NOT NULL PRIMARY KEY,
+    `BINARY`  LONGBLOB
+  );"""
+  val getBinaryContents = sql"""SELECT `BIN_ID`, `BINARY` FROM `BINARIES`""".as[BinaryContent]
+  val dropColumn = sqlu"""ALTER TABLE `BINARIES` DROP COLUMN `BINARY`"""
+}
+
diff --git a/job-server/src/main/scala/db/postgresql/migration/V0_7_5/V0_7_5__Migrate_Blobs.scala b/job-server/src/main/scala/db/postgresql/migration/V0_7_5/V0_7_5__Migrate_Blobs.scala
@@ -0,0 +1,61 @@
+package db.postgresql.migration.V0_7_5
+
+import java.sql.Blob
+import java.sql.Connection
+import javax.sql.rowset.serial.SerialBlob
+
+import scala.concurrent.Await
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.duration.DurationInt
+import scala.util.control.NonFatal
+
+import db.migration.V0_7_5.Migration
+import org.flywaydb.core.api.migration.jdbc.JdbcMigration
+import org.slf4j.LoggerFactory
+import slick.dbio.DBIO
+import slick.dbio.Effect
+import slick.dbio.NoStream
+import slick.driver.PostgresDriver.api.actionBasedSQLInterpolation
+import slick.jdbc.GetResult
+import slick.jdbc.PositionedParameters
+import slick.jdbc.SetParameter
+import slick.profile.SqlAction
+import spark.jobserver.slick.unmanaged.UnmanagedDatabase
+
+class V0_7_5__Migrate_Blobs extends Migration {
+  val logger = LoggerFactory.getLogger(getClass)
+
+  protected def insertBlob(id: Int, blob: SerialBlob): SqlAction[Int, NoStream, Effect] = {
+    sqlu"""INSERT INTO "BINARIES_CONTENTS" ("BIN_ID", "BINARY") VALUES (${id}, ${blob})"""
+  }
+  val createContentsTable = sqlu"""CREATE TABLE "BINARIES_CONTENTS" (
+    "BIN_ID"  SERIAL  NOT NULL PRIMARY KEY,
+    "BINARY"  OID
+  );"""
+  val getBinaryContents = sql"""SELECT "BIN_ID", "BINARY" FROM "BINARIES"""".as[BinaryContent]
+  val dropColumn = sqlu"""ALTER TABLE "BINARIES" DROP COLUMN "BINARY""""
+
+  override def migrate(c: Connection): Unit = {
+    val createTriggerBinariesContents = sqlu"""CREATE TRIGGER t_binary
+            BEFORE UPDATE OR DELETE ON "BINARIES_CONTENTS"
+            FOR EACH ROW EXECUTE PROCEDURE lo_manage("BINARY")"""
+    val dropTriggerBinaries = sqlu"""DROP TRIGGER t_binary ON "BINARIES""""
+    val db = new UnmanagedDatabase(c)
+    c.setAutoCommit(false);
+    try {
+      Await.ready(
+          for {
+            _ <- db.run(createContentsTable)
+            _ <- db.run(createTriggerBinariesContents)
+            _ <- db.stream(getBinaryContents).foreach(b => insertBlob(db, b))
+            _ <- db.run(dropColumn)
+            _ <- db.run(dropTriggerBinaries)
+          } yield Unit, Timeout
+      ).recover{logErrors}
+      c.commit()
+    } catch {
+      case NonFatal(e) => { c.rollback() }
+    }
+  }
+}
+
diff --git a/job-server/src/main/scala/spark/jobserver/io/JobCassandraDAO.scala b/job-server/src/main/scala/spark/jobserver/io/JobCassandraDAO.scala
@@ -6,23 +6,23 @@ import java.nio.ByteBuffer
 import java.nio.file.{Files, Paths}
 import java.util.UUID
 
-import com.datastax.driver.core.querybuilder.{Insert, QueryBuilder => QB}
-import com.datastax.driver.core.querybuilder.QueryBuilder._
+import scala.collection.convert.WrapAsJava
+import scala.collection.convert.Wrappers.JListWrapper
+import scala.concurrent.{Await, Future}
+import scala.concurrent.duration.DurationInt
+import scala.util.Try
+
 import com.datastax.driver.core._
-import com.datastax.driver.core.schemabuilder.SchemaBuilder.Direction
+import com.datastax.driver.core.querybuilder.{Insert, QueryBuilder => QB }
+import com.datastax.driver.core.querybuilder.QueryBuilder._
 import com.datastax.driver.core.schemabuilder.{Create, SchemaBuilder}
+import com.datastax.driver.core.schemabuilder.SchemaBuilder.Direction
 import com.typesafe.config.{Config, ConfigFactory, ConfigRenderOptions}
 import org.joda.time.DateTime
 import org.slf4j.LoggerFactory
 
-import scala.concurrent.duration._
-import scala.collection.convert.WrapAsJava
-import scala.collection.convert.Wrappers.JListWrapper
-import scala.concurrent.{Await, Future}
 import spark.jobserver.cassandra.Cassandra.Resultset.toFuture
 
-import scala.util.Try
-
 object Metadata {
   val BinariesTable = "binaries"
   val BinariesChronologicalTable = "binaries_chronological"
@@ -160,7 +160,6 @@ class JobCassandraDAO(config: Config) extends JobDAO with FileCacher {
   }
 
   override def getJobInfos(limit: Int, status: Option[String] = None): Future[Seq[JobInfo]] = {
-    import Metadata._
     val query = QB.select(
       JobId, ContextName, AppName, BType, UploadTime, Classpath, StartTime, EndTime,
       Error, ErrorClass, ErrorStackTrace
@@ -181,7 +180,6 @@ class JobCassandraDAO(config: Config) extends JobDAO with FileCacher {
   }
 
   override def getRunningJobInfosForContextName(contextName: String): Future[Seq[JobInfo]] = {
-    import Metadata._
     val query = QB.select(
       JobId, ContextName, AppName, BType, UploadTime, Classpath, StartTime, EndTime,
       Error, ErrorClass, ErrorStackTrace
@@ -307,6 +305,11 @@ class JobCassandraDAO(config: Config) extends JobDAO with FileCacher {
     }
   }
 
+  override def getLastUploadTimeAndType(appName: String): Option[(DateTime, BinaryType)] = {
+    // Copied from the base JobDAO, feel free to optimize this (having in mind this specific storage type)
+    Await.result(getApps, 60 seconds).get(appName).map(t => (t._2, t._1))
+  }
+
   private def setup(config: Config): Session = {
     val cassandraConfig = config.getConfig("spark.jobserver.cassandra")
     val hosts = JListWrapper(cassandraConfig.getStringList("hosts"))
@@ -337,7 +340,6 @@ class JobCassandraDAO(config: Config) extends JobDAO with FileCacher {
   }
 
   private def setupSchema() = {
-    import Metadata._
 
     val binariesTable: Create = SchemaBuilder.createTable(BinariesTable).ifNotExists.
       addPartitionKey(AppName, DataType.text).
diff --git a/job-server/src/main/scala/spark/jobserver/io/JobDAO.scala b/job-server/src/main/scala/spark/jobserver/io/JobDAO.scala
@@ -104,6 +104,7 @@ trait JobDAO {
     * @param appName
     */
   def deleteBinary(appName: String)
+
   /**
    * Return all applications name and their last upload times.
    *
@@ -182,8 +183,7 @@ trait JobDAO {
    * Returns the last upload time for a given app name.
    * @return Some(lastUploadedTime) if the app exists and the list of times is nonempty, None otherwise
    */
-  def getLastUploadTimeAndType(appName: String): Option[(DateTime, BinaryType)] =
-    Await.result(getApps, 60 seconds).get(appName).map(t => (t._2, t._1))
+  def getLastUploadTimeAndType(appName: String): Option[(DateTime, BinaryType)]
 
   /**
     * Fetch submited jar or egg content for remote driver and JobManagerActor to cache in local
diff --git a/job-server/src/main/scala/spark/jobserver/io/JobFileDAO.scala b/job-server/src/main/scala/spark/jobserver/io/JobFileDAO.scala
@@ -1,15 +1,15 @@
 package spark.jobserver.io
 
-import com.typesafe.config._
 import java.io._
 import java.nio.file.{Files, Paths}
 
+import com.typesafe.config._
 import org.joda.time.DateTime
 import org.slf4j.LoggerFactory
 
 import scala.collection.mutable
-import scala.concurrent.Future
 import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.Future
 
 /**
   * NB This class does NOT support persisting binary types
@@ -224,6 +224,10 @@ class JobFileDAO(config: Config) extends JobDAO {
     configs.get(jobId)
   }
 
+  override def getLastUploadTimeAndType(appName: String): Option[(DateTime, BinaryType)] = {
+    apps(appName).headOption.map(uploadTime => (uploadTime, BinaryType.Jar))
+  }
+
   private def writeJobConfig(out: DataOutputStream, jobId: String, jobConfig: Config) {
     out.writeUTF(jobId)
     out.writeUTF(jobConfig.root().render(ConfigRenderOptions.concise()))
diff --git a/job-server/src/main/scala/spark/jobserver/io/JobSqlDAO.scala b/job-server/src/main/scala/spark/jobserver/io/JobSqlDAO.scala
diff --git a/job-server/src/test/scala/spark/jobserver/InMemoryDAO.scala b/job-server/src/test/scala/spark/jobserver/InMemoryDAO.scala
diff --git a/job-server/src/test/scala/spark/jobserver/io/FlywayMigrationSpec.scala b/job-server/src/test/scala/spark/jobserver/io/FlywayMigrationSpec.scala
diff --git a/job-server/src/test/scala/spark/jobserver/io/JobDAOActorSpec.scala b/job-server/src/test/scala/spark/jobserver/io/JobDAOActorSpec.scala
