Unify v0.17 migrations

This commit is contained in:
allegroai 2021-01-05 20:07:49 +02:00
parent a8d90887e2
commit be0cf0caa8
3 changed files with 48 additions and 57 deletions

View File

@ -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()

View File

@ -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)

View File

@ -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()