#1340 Add a start and end phase of migration

TheHive-Project · Jun 13, 2020 · b663339 · b663339
1 parent acd6c23
commit b663339
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/migration/src/main/scala/org/thp/thehive/migration/MigrationOps.scala b/migration/src/main/scala/org/thp/thehive/migration/MigrationOps.scala
@@ -342,6 +342,7 @@ trait MigrationOps {
         }
     }
 
+    output.startMigration()
     input.countOrganisations(filter).foreach(count => migrationStats.setTotal("Organisation", count))
     input.countCases(filter).foreach(count => migrationStats.setTotal("Case", count))
     input.countCaseObservables(filter).foreach(count => migrationStats.setTotal("Case/Observable", count))
@@ -372,6 +373,7 @@ trait MigrationOps {
       _ <- migrate("ObservableType", input.listObservableTypes(filter), output.createObservableTypes, output.observableTypeExists)
       _ <- migrateWholeCaseTemplates(input, output, filter)
       _ <- migrateCasesAndAlerts()
+      _ <- Future.fromTry(output.endMigration())
     } yield ()
   }
 }
diff --git a/migration/src/main/scala/org/thp/thehive/migration/Output.scala b/migration/src/main/scala/org/thp/thehive/migration/Output.scala
@@ -21,6 +21,8 @@ import org.thp.thehive.migration.dto.{
 }
 
 trait Output {
+  def startMigration(): Try[Unit]
+  def endMigration(): Try[Unit]
   def profileExists(inputProfile: InputProfile): Boolean
   def createProfile(inputProfile: InputProfile): Try[IdMapping]
   def organisationExists(inputOrganisation: InputOrganisation): Boolean