#2305 Improve migration

TheHive-Project · Jan 13, 2022 · 80d49a9 · 80d49a9
1 parent 64cb9b0
commit 80d49a9
Show file tree

Hide file tree

Showing 9 changed files with 356 additions and 348 deletions.
diff --git a/conf/migration-logback.xml b/conf/migration-logback.xml
@@ -44,6 +44,7 @@
     <logger name="org.thp.thehive.services.LocalAuthSrv" level="TRACE" />
     <logger name="org.thp.scalligraph.graphql" level="TRACE" />
     -->
+    <logger name="org.thp.scalligraph.traversal" level="INFO"/>
     <logger name="org.janusgraph.graphdb.transaction.StandardJanusGraphTx" level="ERROR"/>
     <logger name="org.thp.thehive" level="INFO"/>
 

diff --git a/migration/src/main/resources/reference.conf b/migration/src/main/resources/reference.conf
@@ -17,9 +17,6 @@ input {
     randomFactor = 0.2
   }
   filter {
-    maxCaseAge: 0
-    maxAlertAge: 0
-    maxAuditAge: 0
     includeAlertTypes: []
     excludeAlertTypes: []
     includeAlertSources: []

diff --git a/migration/src/main/scala/org/thp/thehive/migration/Input.scala b/migration/src/main/scala/org/thp/thehive/migration/Input.scala
@@ -48,17 +48,15 @@ object Filter {
               new ParseException(s"Unparseable date: $s\nExpected format is ${dateFormats.map(_.toPattern).mkString("\"", "\" or \"", "\"")}", 0)
             )
         }
-    def readDate(dateConfigName: String, ageConfigName: String) =
+    def readDate(dateConfigName: String, ageConfigName: String): Option[Long] =
       Try(config.getString(dateConfigName))
         .flatMap(parseDate)
         .map(d => d.getTime)
-        .toOption
         .orElse {
-          Try {
-            val age = config.getDuration(ageConfigName)
-            if (age.isZero) None else Some(now - age.getSeconds * 1000)
-          }.toOption.flatten
+          Try(config.getDuration(ageConfigName))
+            .map(d => now - d.getSeconds * 1000)
         }
+        .toOption
     val caseFromDate    = readDate("caseFromDate", "maxCaseAge")
     val caseUntilDate   = readDate("caseUntilDate", "minCaseAge")
     val caseFromNumber  = Try(config.getInt("caseFromNumber")).toOption

diff --git a/migration/src/main/scala/org/thp/thehive/migration/Migrate.scala b/migration/src/main/scala/org/thp/thehive/migration/Migrate.scala
@@ -11,7 +11,7 @@ import java.io.File
 import java.nio.file.{Files, Paths}
 import scala.collection.JavaConverters._
 import scala.concurrent.duration.{Duration, DurationInt}
-import scala.concurrent.{Await, ExecutionContext}
+import scala.concurrent.{blocking, Await, ExecutionContext, Future}
 
 object Migrate extends App with MigrationOps {
   val defaultLoggerConfigFile = "/etc/thehive/logback-migration.xml"
@@ -205,11 +205,17 @@ object Migrate extends App with MigrationOps {
     implicit val mat: Materializer        = Materializer(actorSystem)
     transactionPageSize = config.getInt("transactionPageSize")
     threadCount = config.getInt("threadCount")
+    var stop = false
 
     try {
-      val timer = actorSystem.scheduler.scheduleAtFixedRate(10.seconds, 10.seconds) { () =>
-        logger.info(migrationStats.showStats())
-        migrationStats.flush()
+      Future {
+        blocking {
+          while (!stop) {
+            logger.info(migrationStats.showStats())
+            migrationStats.flush()
+            Thread.sleep(10000) // 10 seconds
+          }
+        }
       }
 
       val returnStatus =
@@ -219,16 +225,15 @@ object Migrate extends App with MigrationOps {
           val filter = Filter.fromConfig(config.getConfig("input.filter"))
 
           val process = migrate(input, output, filter)
-
-          Await.result(process, Duration.Inf)
+          blocking(Await.result(process, Duration.Inf))
           logger.info("Migration finished")
           0
         } catch {
           case e: Throwable =>
             logger.error(s"Migration failed", e)
             1
         } finally {
-          timer.cancel()
+          stop = true
           Await.ready(actorSystem.terminate(), 1.minute)
           ()
         }