implement new rpc route for db export and remove unnecessary code

Author: Youri K

src/main/protobuf/fossildbapi.proto

@@ -130,13 +130,23 @@ message RestoreFromBackupReply {
     optional string errorMessage = 2;
 }
 
-message CompactAllDataRequest {optional uint32 type = 1;}
+message CompactAllDataRequest {}
 
 message CompactAllDataReply {
     required bool success = 1;
     optional string errorMessage = 2;
 }
 
+message ExportDBRequest {
+    required string newDataDir = 1;
+    optional string optionsFile = 2;
+}
+
+message ExportDBReply {
+    required bool success = 1;
+    optional string errorMessage = 2;
+}
+
 
 service FossilDB {
     rpc Health (HealthRequest) returns (HealthReply) {}
@@ -151,4 +161,5 @@ service FossilDB {
     rpc Backup (BackupRequest) returns (BackupReply) {}
     rpc RestoreFromBackup (RestoreFromBackupRequest) returns (RestoreFromBackupReply) {}
     rpc CompactAllData (CompactAllDataRequest) returns (CompactAllDataReply) {}
+    rpc ExportDB (ExportDBRequest) returns (ExportDBReply) {}
 }

src/main/scala/com/scalableminds/fossildb/FossilDBGrpcImpl.scala

@@ -15,11 +15,11 @@ import scala.concurrent.Future
 
 class FossilDBGrpcImpl(storeManager: StoreManager)
   extends FossilDBGrpc.FossilDB
-  with LazyLogging {
+    with LazyLogging {
 
   override def health(req: HealthRequest) = withExceptionHandler(req) {
     HealthReply(true)
-  } {errorMsg => HealthReply(false, errorMsg)}
+  } { errorMsg => HealthReply(false, errorMsg) }
 
   override def get(req: GetRequest) = withExceptionHandler(req) {
     val store = storeManager.getStore(req.collection)
@@ -31,71 +31,76 @@ class FossilDBGrpcImpl(storeManager: StoreManager)
         GetReply(false, Some("No such element"), ByteString.EMPTY, 0)
       }
     }
-  } {errorMsg => GetReply(false, errorMsg, ByteString.EMPTY, 0)}
+  } { errorMsg => GetReply(false, errorMsg, ByteString.EMPTY, 0) }
 
   override def put(req: PutRequest) = withExceptionHandler(req) {
     val store = storeManager.getStore(req.collection)
     val version = req.version.getOrElse(store.get(req.key, None).map(_.version + 1).getOrElse(0L))
     require(version >= 0, "Version numbers must be non-negative")
     store.put(req.key, version, req.value.toByteArray)
     PutReply(true)
-  } {errorMsg => PutReply(false, errorMsg)}
+  } { errorMsg => PutReply(false, errorMsg) }
 
   override def delete(req: DeleteRequest) = withExceptionHandler(req) {
     val store = storeManager.getStore(req.collection)
     store.delete(req.key, req.version)
     DeleteReply(true)
-  } {errorMsg => DeleteReply(false, errorMsg)}
+  } { errorMsg => DeleteReply(false, errorMsg) }
 
   override def getMultipleVersions(req: GetMultipleVersionsRequest) = withExceptionHandler(req) {
     val store = storeManager.getStore(req.collection)
     val (values, versions) = store.getMultipleVersions(req.key, req.oldestVersion, req.newestVersion)
     GetMultipleVersionsReply(true, None, values.map(ByteString.copyFrom(_)), versions)
-  } {errorMsg => GetMultipleVersionsReply(false, errorMsg)}
+  } { errorMsg => GetMultipleVersionsReply(false, errorMsg) }
 
   override def getMultipleKeys(req: GetMultipleKeysRequest) = withExceptionHandler(req) {
     val store = storeManager.getStore(req.collection)
     val (keys, values, versions) = store.getMultipleKeys(req.key, req.prefix, req.version, req.limit)
     GetMultipleKeysReply(true, None, keys, values.map(ByteString.copyFrom(_)), versions)
-  } {errorMsg => GetMultipleKeysReply(false, errorMsg)}
+  } { errorMsg => GetMultipleKeysReply(false, errorMsg) }
 
   override def deleteMultipleVersions(req: DeleteMultipleVersionsRequest) = withExceptionHandler(req) {
     val store = storeManager.getStore(req.collection)
     store.deleteMultipleVersions(req.key, req.oldestVersion, req.newestVersion)
     DeleteMultipleVersionsReply(true)
-  } {errorMsg => DeleteMultipleVersionsReply(false, errorMsg)}
+  } { errorMsg => DeleteMultipleVersionsReply(false, errorMsg) }
 
   override def listKeys(req: ListKeysRequest) = withExceptionHandler(req) {
     val store = storeManager.getStore(req.collection)
     val keys = store.listKeys(req.limit, req.startAfterKey)
     ListKeysReply(true, None, keys)
-  } {errorMsg => ListKeysReply(false, errorMsg)}
+  } { errorMsg => ListKeysReply(false, errorMsg) }
 
   override def listVersions(req: ListVersionsRequest) = withExceptionHandler(req) {
     val store = storeManager.getStore(req.collection)
     val versions = store.listVersions(req.key, req.limit, req.offset)
     ListVersionsReply(true, None, versions)
-  } {errorMsg => ListVersionsReply(false, errorMsg)}
+  } { errorMsg => ListVersionsReply(false, errorMsg) }
 
   override def backup(req: BackupRequest) = withExceptionHandler(req) {
     val backupInfoOpt = storeManager.backup
     backupInfoOpt match {
       case Some(backupInfo) => BackupReply(true, None, backupInfo.id, backupInfo.timestamp, backupInfo.size)
       case _ => throw new Exception("Backup did not return valid BackupInfo")
     }
-  } {errorMsg => BackupReply(false, errorMsg, 0, 0, 0)}
+  } { errorMsg => BackupReply(false, errorMsg, 0, 0, 0) }
 
   override def restoreFromBackup(req: RestoreFromBackupRequest) = withExceptionHandler(req) {
     storeManager.restoreFromBackup
     RestoreFromBackupReply(true)
-  } {errorMsg => RestoreFromBackupReply(false, errorMsg)}
+  } { errorMsg => RestoreFromBackupReply(false, errorMsg) }
 
   override def compactAllData(req: CompactAllDataRequest) = withExceptionHandler(req) {
-    storeManager.compactAllData(req.`type`)
+    storeManager.compactAllData()
     CompactAllDataReply(true)
-  } {errorMsg => CompactAllDataReply(false, errorMsg)}
+  } { errorMsg => CompactAllDataReply(false, errorMsg) }
 
-  private def withExceptionHandler [T, R <: GeneratedMessage](request: R)(tryBlock: => T)(onErrorBlock: Option[String] => T): Future[T] = {
+  override def exportDB(req: ExportDBRequest) = withExceptionHandler(req) {
+    storeManager.exportDB(req.newDataDir, req.optionsFile)
+    ExportDBReply(true)
+  } { errorMsg => ExportDBReply(false, errorMsg) }
+
+  private def withExceptionHandler[T, R <: GeneratedMessage](request: R)(tryBlock: => T)(onErrorBlock: Option[String] => T): Future[T] = {
     try {
       logger.debug("received " + requestToString(request))
       Future.successful(tryBlock)
@@ -112,7 +117,7 @@ class FossilDBGrpcImpl(storeManager: StoreManager)
   }
 
   private def requestToString[R <: GeneratedMessage](request: R) =
-    request.getClass.getSimpleName + "(" + request.toString.replaceAll("\n"," ") + ")"
+    request.getClass.getSimpleName + "(" + request.toString.replaceAll("\n", " ") + ")"
 
   private def getStackTraceAsString(t: Throwable) = {
     val sw = new StringWriter

src/main/scala/com/scalableminds/fossildb/db/RocksDBStore.scala

@@ -76,77 +76,29 @@ class RocksDBManager(dataDir: Path, columnFamilies: List[String], optionsFilePat
     logger.info("Restoring from backup complete. Reopening RocksDB")
   }
 
-  def compactAllData(idx: Option[Int]) = {
+  def compactAllData() = {
     logger.info("Compacting all data")
     RocksDB.loadLibrary()
-    idx.getOrElse(0) match {
-      case 0 => db.compactRange()
-      case 1 => writeAllSSts()
-      case 2 => ingestFiles()
-      case 3 => writeToNewDB()
-    }
+    db.compactRange()
     logger.info("All data has been compacted to last level containing data")
   }
 
-  def close(): Future[Unit] = {
-    logger.info("Closing RocksDB handle")
-    Future.successful(db.close())
-  }
-
-  def ingestFiles() = {
-    val ifo = new IngestExternalFileOptions()
-    ifo.setMoveFiles(true)
-    val fileNames = new File("toIngest").listFiles.filter(_.isFile).filter(_.getName.endsWith("sst")).map(_.getPath)
-    val asd: mutable.Buffer[String] = fileNames.toBuffer
-    val handle = columnFamilyHandles("skeletons")
-    db.ingestExternalFile(handle, asd.asJava, ifo)
-  }
-
-  def writeAllSSts() = {
-    val (dbOptions, columnFamilyDescriptors) = loadOptions("config/options.ini")
-    val descriptor = columnFamilyDescriptors.find(_.getName sameElements "skeletons".getBytes)
-    val options = new Options(dbOptions, descriptor.get.getOptions)
-    val writer = new SstFileWriter(new EnvOptions(), options)
-    val store = getStoreForColumnFamily("skeletons")
-    val it = store.get.scan("", None)
-    var idx = 0
-    writer.open(s"data/export/test${idx}.sst")
-    it.foreach { el =>
-      if (new File(s"data/export/test${idx}.sst").length() + el.key.getBytes.length + el.value.length > options.targetFileSizeBase()) {
-        writer.finish()
-        idx += 1
-        writer.open(s"data/export/test${idx}.sst")
-      }
-      writer.put(el.key.getBytes, el.value)
+  def exportToNewDB(newDataDir: Path, newOptionsFilePathOpt: Option[String]) = {
+    RocksDB.loadLibrary()
+    logger.info(s"Exporting to new DB at ${newDataDir.toString} with options file ${newOptionsFilePathOpt}")
+    val newManager = new RocksDBManager(newDataDir, columnFamilies, newOptionsFilePathOpt)
+    newManager.columnFamilyHandles.foreach { case (name, handle) =>
+      val dataIterator = getStoreForColumnFamily(name).get.scan("", None)
+      dataIterator.foreach(el => newManager.db.put(handle, el.key.getBytes, el.value))
     }
-    writer.finish()
-  }
-
-  def writeToNewDB() = {
-    val manager = new RocksDBManager(Paths.get("data_new"), columnFamilies, Some("config/options.ini"))
-    val skeletonHandle = manager.columnFamilyHandles("skeletons")
-    val it = getStoreForColumnFamily("skeletons").get.scan("", None)
-    it.foreach { el => manager.db.put(skeletonHandle, el.key.getBytes, el.value) }
+    logger.info("Writing data completed. Start compaction")
+    newManager.db.compactRange()
+    logger.info("Compaction finished")
   }
 
-  def loadOptions(optionFilepath: String) = {
-    val options = new DBOptions()
-    val cfListRef: mutable.Buffer[ColumnFamilyDescriptor] = mutable.Buffer()
-    optionsFilePathOpt.foreach { optionsFilePath =>
-      try {
-        org.rocksdb.OptionsUtil.loadOptionsFromFile(optionsFilePath, Env.getDefault, options, cfListRef.asJava)
-        logger.info("successfully loaded rocksdb options from " + optionsFilePath)
-      } catch {
-        case e: Exception => {
-          throw new Exception("Failed to load rocksdb options from file " + optionsFilePath, e)
-        }
-      }
-    }
-    options.setCreateIfMissing(true).setCreateMissingColumnFamilies(true)
-    val defaultColumnFamilyOptions = cfListRef.find(_.getName sameElements RocksDB.DEFAULT_COLUMN_FAMILY).map(_.getOptions).getOrElse(new ColumnFamilyOptions())
-    val newColumnFamilyDescriptors = (columnFamilies.map(_.getBytes) :+ RocksDB.DEFAULT_COLUMN_FAMILY).diff(cfListRef.toList.map(_.getName)).map(new ColumnFamilyDescriptor(_, defaultColumnFamilyOptions))
-    val columnFamilyDescriptors = cfListRef.toList ::: newColumnFamilyDescriptors
-    (options, columnFamilyDescriptors)
+  def close(): Future[Unit] = {
+    logger.info("Closing RocksDB handle")
+    Future.successful(db.close())
   }
 }
 

src/main/scala/com/scalableminds/fossildb/db/StoreManager.scala

@@ -3,7 +3,7 @@
  */
 package com.scalableminds.fossildb.db
 
-import java.nio.file.Path
+import java.nio.file.{Path, Paths}
 import java.util.concurrent.atomic.AtomicBoolean
 
 class StoreManager(dataDir: Path, backupDir: Path, columnFamilies: List[String], rocksdbOptions: Option[String]) {
@@ -66,11 +66,18 @@ class StoreManager(dataDir: Path, backupDir: Path, columnFamilies: List[String],
     }
   }
 
-  def compactAllData(idx: Option[Int]) = {
+  def compactAllData() = {
     failDuringBackup
     failDuringRestore
     try {
-      rocksDBManager.get.compactAllData(idx)
+      rocksDBManager.get.compactAllData()
+    }
+  }
+
+  def exportDB(newDataDir: String, newOptionsFilePathOpt: Option[String]) = {
+    failDuringRestore
+    try {
+      rocksDBManager.get.exportToNewDB(Paths.get(newDataDir), newOptionsFilePathOpt)
     }
   }