From be0cf0caa8f34644d3191dcbba6eb0800b3a46f9 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Tue, 5 Jan 2021 20:07:49 +0200
Subject: [PATCH] Unify v0.17 migrations

---
 apiserver/mongo/migrations/0.17.0.py | 48 ++++++++++++++++++++++++++++
 apiserver/mongo/migrations/0.17.2.py | 45 --------------------------
 apiserver/mongo/migrations/0.17.3.py | 12 -------
 3 files changed, 48 insertions(+), 57 deletions(-)
 delete mode 100644 apiserver/mongo/migrations/0.17.2.py
 delete mode 100644 apiserver/mongo/migrations/0.17.3.py

diff --git a/apiserver/mongo/migrations/0.17.0.py b/apiserver/mongo/migrations/0.17.0.py
index 28d977d4..8f51feaa 100644
--- a/apiserver/mongo/migrations/0.17.0.py
+++ b/apiserver/mongo/migrations/0.17.0.py
@@ -1,9 +1,46 @@
+from datetime import datetime
+from typing import Optional
+
 from pymongo.database import Database, Collection
 
 from apiserver.bll.task.artifacts import get_artifact_id
 from apiserver.utilities.dicts import nested_get
 
 
+def _add_active_duration(db: Database):
+    active_duration_key = "active_duration"
+    query = {"$or": [{active_duration_key: {"$eq": None}}, {active_duration_key: {"$eq": 0}}]}
+    collection = db["task"]
+    for doc in collection.find(
+        filter=query, projection=[active_duration_key, "status", "started", "completed"]
+    ):
+        started = doc.get("started")
+        completed = doc.get("completed")
+        running = doc.get("status") == "running"
+        active_duration_value = doc.get(active_duration_key)
+        if active_duration_value == 0:
+            collection.update_one(
+                {"_id": doc["_id"]},
+                {"$set": {active_duration_key: None}},
+            )
+        elif started and active_duration_value is None:
+            collection.update_one(
+                {"_id": doc["_id"]},
+                {"$set": {active_duration_key: _get_active_duration(completed, running, started)}},
+            )
+
+
+def _get_active_duration(
+    completed: datetime, running: bool, started: datetime
+) -> Optional[float]:
+    if running:
+        return (datetime.utcnow() - started).total_seconds()
+    elif completed:
+        return (completed - started).total_seconds()
+    else:
+        return None
+
+
 def migrate_backend(db: Database):
     collection: Collection = db["task"]
     artifacts_field = "execution.artifacts"
@@ -17,3 +54,14 @@ def migrate_backend(db: Database):
         collection.update_one(
             {"_id": doc["_id"]}, {"$set": {artifacts_field: new_artifacts}}
         )
+
+    _add_active_duration(db)
+
+
+def migrate_auth(db: Database):
+    """
+    Remove the old indices from the user collections
+    to enable building of the updated user email index
+    """
+    collection: Collection = db["user"]
+    collection.drop_indexes()
diff --git a/apiserver/mongo/migrations/0.17.2.py b/apiserver/mongo/migrations/0.17.2.py
deleted file mode 100644
index 12011ee7..00000000
--- a/apiserver/mongo/migrations/0.17.2.py
+++ /dev/null
@@ -1,45 +0,0 @@
-from datetime import datetime
-from typing import Optional
-
-from pymongo.database import Database
-
-
-def _add_active_duration(db: Database):
-    active_duration_key = "active_duration"
-    query = {"$or": [{active_duration_key: {"$eq": None}}, {active_duration_key: {"$eq": 0}}]}
-    collection = db["task"]
-    for doc in collection.find(
-        filter=query, projection=[active_duration_key, "status", "started", "completed"]
-    ):
-        started = doc.get("started")
-        completed = doc.get("completed")
-        running = doc.get("status") == "running"
-        active_duration_value = doc.get(active_duration_key)
-        if active_duration_value == 0:
-            collection.update_one(
-                {"_id": doc["_id"]},
-                {"$set": {active_duration_key: None}},
-            )
-        elif started and active_duration_value is None:
-            collection.update_one(
-                {"_id": doc["_id"]},
-                {"$set": {active_duration_key: _get_active_duration(completed, running, started)}},
-            )
-
-
-def _get_active_duration(
-    completed: datetime, running: bool, started: datetime
-) -> Optional[float]:
-    if running:
-        return (datetime.utcnow() - started).total_seconds()
-    elif completed:
-        return (completed - started).total_seconds()
-    else:
-        return None
-
-
-def migrate_backend(db: Database):
-    """
-    Add active_duration field to tasks
-    """
-    _add_active_duration(db)
diff --git a/apiserver/mongo/migrations/0.17.3.py b/apiserver/mongo/migrations/0.17.3.py
deleted file mode 100644
index 2f4df4d0..00000000
--- a/apiserver/mongo/migrations/0.17.3.py
+++ /dev/null
@@ -1,12 +0,0 @@
-from collections import Collection
-
-from pymongo.database import Database, Collection
-
-
-def migrate_auth(db: Database):
-    """
-    Remove the old indices from the user collections
-    to enable building of the updated user email index
-    """
-    collection: Collection = db["user"]
-    collection.drop_indexes()