Merge pull request #32 from scalableminds/fix-columnFamily-options

Export To New DB

Author: fm3

src/main/protobuf/fossildbapi.proto

@@ -137,6 +137,16 @@ message CompactAllDataReply {
     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
+    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

@@ -3,7 +3,8 @@
  */
 package com.scalableminds.fossildb.db
 
-import java.nio.file.{Files, Path}
+import java.io.File
+import java.nio.file.{Files, Path, Paths}
 import java.util
 
 import com.typesafe.scalalogging.LazyLogging
@@ -21,6 +22,10 @@ class RocksDBManager(dataDir: Path, columnFamilies: List[String], optionsFilePat
 
   val (db: RocksDB, columnFamilyHandles) = {
     RocksDB.loadLibrary()
+    val columnOptions = new ColumnFamilyOptions()
+      .setArenaBlockSize(4 * 1024 * 1024) // 4MB
+      .setTargetFileSizeBase(1024 * 1024 * 1024) // 1GB
+      .setMaxBytesForLevelBase(10 * 1024 * 1024 * 1024) // 10GB
     val options = new DBOptions()
     val cfListRef: mutable.Buffer[ColumnFamilyDescriptor] = mutable.Buffer()
     optionsFilePathOpt.foreach { optionsFilePath =>
@@ -34,7 +39,7 @@ class RocksDBManager(dataDir: Path, columnFamilies: List[String], optionsFilePat
       }
     }
     options.setCreateIfMissing(true).setCreateMissingColumnFamilies(true)
-    val defaultColumnFamilyOptions = cfListRef.find(_.getName sameElements RocksDB.DEFAULT_COLUMN_FAMILY).map(_.getOptions).getOrElse(new ColumnFamilyOptions())
+    val defaultColumnFamilyOptions = cfListRef.find(_.getName sameElements RocksDB.DEFAULT_COLUMN_FAMILY).map(_.getOptions).getOrElse(columnOptions)
     val newColumnFamilyDescriptors = (columnFamilies.map(_.getBytes) :+ RocksDB.DEFAULT_COLUMN_FAMILY).diff(cfListRef.toList.map(_.getName)).map(new ColumnFamilyDescriptor(_, defaultColumnFamilyOptions))
     val columnFamilyDescriptors = cfListRef.toList ::: newColumnFamilyDescriptors
     logger.info("Opening RocksDB at " + dataDir.toAbsolutePath)
@@ -78,6 +83,19 @@ class RocksDBManager(dataDir: Path, columnFamilies: List[String], optionsFilePat
     logger.info("All data has been compacted to last level containing data")
   }
 
+  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))
+    }
+    logger.info("Writing data completed. Start compaction")
+    newManager.db.compactRange()
+    logger.info("Compaction finished")
+  }
+
   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,14 +66,21 @@ class StoreManager(dataDir: Path, backupDir: Path, columnFamilies: List[String],
     }
   }
 
-  def compactAllData = {
+  def compactAllData() = {
     failDuringBackup
     failDuringRestore
     try {
       rocksDBManager.get.compactAllData()
     }
   }
 
+  def exportDB(newDataDir: String, newOptionsFilePathOpt: Option[String]) = {
+    failDuringRestore
+    try {
+      rocksDBManager.get.exportToNewDB(Paths.get(newDataDir), newOptionsFilePathOpt)
+    }
+  }
+
   def close = {
     rocksDBManager.map(_.close)
   }