fix(job-server): JobSQLDAO: Prevent saving the same binary multiple (spark-jobserver#960)

times

Save hash of the binary in extra column in BINARIES and use it as a
primary key in BINARIES_CONTENTS table instead of binary id.

Author: ecandreev

job-server/src/main/resources/db/h2/migration/V0_7_6/V0_7_6__add_bin_hash_column.sql

@@ -0,0 +1,5 @@
+ALTER TABLE "BINARIES" ADD COLUMN "BIN_HASH" BINARY(32);
+UPDATE "BINARIES" SET "BIN_HASH" = (SELECT (CAST("BIN_ID" AS BINARY(32))));
+ALTER TABLE "BINARIES" ALTER COLUMN "BIN_HASH" SET NOT NULL;
+ALTER TABLE "BINARIES_CONTENTS" ALTER COLUMN "BIN_ID" RENAME TO "BIN_HASH";
+ALTER TABLE "BINARIES_CONTENTS" ALTER COLUMN "BIN_HASH" BINARY(32) NOT NULL;

job-server/src/main/resources/db/mysql/migration/V0_7_6/V0_7_6__add_bin_hash_column.sql

@@ -0,0 +1,3 @@
+ALTER TABLE `BINARIES` ADD COLUMN `BIN_HASH` VARBINARY(32) NOT NULL;
+UPDATE `BINARIES` SET `BIN_HASH` = (SELECT CAST(`BIN_ID` AS BINARY(32)));
+ALTER TABLE `BINARIES_CONTENTS` CHANGE COLUMN `BIN_ID` `BIN_HASH` VARBINARY(32) NOT NULL;

job-server/src/main/resources/db/postgresql/migration/V0_7_6/V0_7_6__add_bin_hash_column.sql

@@ -0,0 +1,8 @@
+ALTER TABLE "BINARIES" ADD COLUMN "BIN_HASH" BYTEA;
+UPDATE "BINARIES" SET "BIN_HASH" = (SELECT "BIN_ID"::VARCHAR(32)::BYTEA);
+ALTER TABLE "BINARIES" ALTER COLUMN "BIN_HASH" SET NOT NULL;
+ALTER TABLE "BINARIES_CONTENTS" ADD COLUMN "BIN_HASH" BYTEA;
+UPDATE "BINARIES_CONTENTS" SET "BIN_HASH" = (SELECT "BIN_ID"::VARCHAR(32)::BYTEA);
+ALTER TABLE "BINARIES_CONTENTS" ALTER COLUMN "BIN_HASH" SET NOT NULL;
+ALTER TABLE "BINARIES_CONTENTS" DROP COLUMN "BIN_ID";
+ALTER TABLE "BINARIES_CONTENTS" ADD PRIMARY KEY ("BIN_HASH");

job-server/src/main/scala/spark/jobserver/io/JobSqlDAO.scala

@@ -41,18 +41,19 @@ class JobSqlDAO(config: Config) extends JobDAO with FileCacher {
 
   // Definition of the tables
   //scalastyle:off
-  class Binaries(tag: Tag) extends Table[(Int, String, String, Timestamp)](tag, "BINARIES") {
+  class Binaries(tag: Tag) extends Table[(Int, String, String, Timestamp, Array[Byte])](tag, "BINARIES") {
     def binId = column[Int]("BIN_ID", O.PrimaryKey, O.AutoInc)
     def appName = column[String]("APP_NAME")
     def binaryType = column[String]("BINARY_TYPE")
     def uploadTime = column[Timestamp]("UPLOAD_TIME")
-    def * = (binId, appName, binaryType, uploadTime)
+    def binHash = column[Array[Byte]]("BIN_HASH")
+    def * = (binId, appName, binaryType, uploadTime, binHash)
   }
 
-  class BinariesContents(tag: Tag) extends Table[(Int, Blob)](tag, "BINARIES_CONTENTS") {
-    def binId = column[Int]("BIN_ID", O.PrimaryKey)
+  class BinariesContents(tag: Tag) extends Table[(Array[Byte], Blob)](tag, "BINARIES_CONTENTS") {
+    def binHash = column[Array[Byte]]("BIN_HASH", O.PrimaryKey)
     def binary = column[Blob]("BINARY")
-    def * = (binId, binary)
+    def * = (binHash, binary)
   }
 
   val binaries = TableQuery[Binaries]
@@ -183,12 +184,24 @@ class JobSqlDAO(config: Config) extends JobDAO with FileCacher {
     Await.result(db.run(query), 60 seconds)
   }
 
+  private def calculateBinaryHash(binBytes: Array[Byte]): Array[Byte] = {
+    import java.security.MessageDigest
+    val md = MessageDigest.getInstance("SHA-256");
+    md.digest(binBytes)
+  }
+
   // Insert JarInfo and its jar into db and return the primary key associated with that row
   private def insertBinaryInfo(binInfo: BinaryInfo, binBytes: Array[Byte]): Future[Int] = {
+    val hash = calculateBinaryHash(binBytes);
     val dbAction = (for {
       binId <- binaries.returning(binaries.map(_.binId)) +=
-        (-1, binInfo.appName, binInfo.binaryType.name, convertDateJodaToSql(binInfo.uploadTime))
-      _ <- binariesContents.map(bc => bc.*) += (binId, new SerialBlob(binBytes))
+        (-1, binInfo.appName, binInfo.binaryType.name, convertDateJodaToSql(binInfo.uploadTime), hash)
+      _ <- binariesContents.filter(_.binHash === hash).result.headOption.flatMap{
+            case None =>
+              binariesContents.map(bc => bc.*) += (hash, new SerialBlob(binBytes))
+            case Some(bc) =>
+              DBIO.successful(None) // no-op
+          }
     } yield binId).transactionally
     db.run(dbAction)
   }
@@ -200,9 +213,15 @@ class JobSqlDAO(config: Config) extends JobDAO with FileCacher {
 
   private def deleteBinaryInfo(appName: String): Future[Int] = {
     val deleteBinary = binaries.filter(_.appName === appName)
-    val deleteBinariesContents = binariesContents.filter(_.binId in deleteBinary.map(_.binId))
+    val hashUsed = binaries.filter(_.binHash in deleteBinary.map(_.binHash)).filter(_.appName =!= appName)
+    val deleteBinariesContents = binariesContents.filter(_.binHash in deleteBinary.map(_.binHash))
     val dbAction = (for {
-      _ <- deleteBinariesContents.delete
+      _ <- hashUsed.result.headOption.flatMap{
+        case None =>
+          deleteBinariesContents.delete
+        case Some(bc) =>
+          DBIO.successful(None) // no-op
+      }
       b <- deleteBinary.delete
     } yield b).transactionally
     db.run(dbAction).recover(logDeleteErrors)
@@ -231,7 +250,7 @@ class JobSqlDAO(config: Config) extends JobDAO with FileCacher {
       b <- binaries.filter { bin =>
         bin.appName === appName && bin.uploadTime === dateTime && bin.binaryType === binaryType.name
       }
-      bc <- binariesContents if b.binId === bc.binId
+      bc <- binariesContents if b.binHash === bc.binHash
     } yield bc.binary
     val dbAction = query.result
     db.run(dbAction.head.map { b => b.getBytes(1, b.length.toInt) }.transactionally)

job-server/src/test/scala/spark/jobserver/io/FlywayMigrationSpec.scala

@@ -80,13 +80,13 @@ class FlywayMigrationSpec extends FunSpec with Matchers {
       val descBinariesIt = ResultSetIterator(descBinaries){ r: ResultSet =>
         r.getString("FIELD")
       }
-      descBinariesIt.toList should be (List("BIN_ID", "APP_NAME", "UPLOAD_TIME", "BINARY_TYPE"))
+      descBinariesIt.toList should be (List("BIN_ID", "APP_NAME", "UPLOAD_TIME", "BINARY_TYPE", "BIN_HASH"))
 
       val descBinariesContents = sqlConn.createStatement().executeQuery("SHOW COLUMNS FROM BINARIES_CONTENTS")
       val descBinariesContentsIt = ResultSetIterator(descBinariesContents){ r: ResultSet =>
         r.getString("FIELD")
       }
-      descBinariesContentsIt.toList should be (List("BIN_ID", "BINARY"))
+      descBinariesContentsIt.toList should be (List("BIN_HASH", "BINARY"))
 
       val descJobs = sqlConn.createStatement().executeQuery("SHOW COLUMNS FROM JOBS")
       val descJobsIt = ResultSetIterator(descJobs){ r: ResultSet =>
@@ -115,14 +115,14 @@ class FlywayMigrationSpec extends FunSpec with Matchers {
       ))
 
       val migratedBinariesContents =
-        sqlConn.createStatement().executeQuery("SELECT BIN_ID, BINARY FROM BINARIES_CONTENTS")
+        sqlConn.createStatement().executeQuery("SELECT BIN_HASH, BINARY FROM BINARIES_CONTENTS")
       val migratedBinariesContentsIt = ResultSetIterator(migratedBinariesContents){ r: ResultSet =>
-        (r.getInt("BIN_ID"),
+        (r.getString("BIN_HASH"),
           r.getString("BINARY"))
       }
       migratedBinariesContentsIt.toList should be (List(
-        (1, "deadbeef"),
-        (2, "beadface")
+        ("0000000000000001", "deadbeef"),
+        ("0000000000000002", "beadface")
       ))
     }
   }