Repository: huggingface/dataset-viewer

Issue:
Migration job fails on a new database
When run on a new database, where the migrations have never been run, the migration script that adds a "force" field fails with:

```
ERROR: 2023-01-31 14:36:28,825 - root - Migration failed: The fields "{'priority'}" do not exist on the document "JobSnapshot"
```

Reported by @rtrompier 


Discussion:
hot fix: https://github.com/huggingface/datasets-server/pull/743 to disable it for now on prod.


Resolving patch:
diff --git a/jobs/mongodb_migration/src/mongodb_migration/migrations/_20221116133500_queue_job_add_force.py b/jobs/mongodb_migration/src/mongodb_migration/migrations/_20221116133500_queue_job_add_force.py
--- a/jobs/mongodb_migration/src/mongodb_migration/migrations/_20221116133500_queue_job_add_force.py
+++ b/jobs/mongodb_migration/src/mongodb_migration/migrations/_20221116133500_queue_job_add_force.py
@@ -1,15 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 The HuggingFace Authors.
 
-import enum
 import logging
-import types
-from typing import Generic, Type, TypeVar
 
-from mongoengine import Document
+from libcommon.queue import Job
 from mongoengine.connection import get_db
-from mongoengine.fields import BooleanField, DateTimeField, EnumField, StringField
-from mongoengine.queryset.queryset import QuerySet
 
 from mongodb_migration.check import check_documents
 from mongodb_migration.migration import Migration
@@ -19,9 +14,9 @@
 class MigrationAddForceToJob(Migration):
     def up(self) -> None:
         # See https://docs.mongoengine.org/guide/migration.html#example-1-addition-of-a-field
-        logging.info("Add the force field, with the default value (False), to all the jobs")
+        logging.info("If missing, add the force field with the default value (False) to  the jobs")
         db = get_db("queue")
-        db["jobsBlue"].update_many({}, {"$set": {"force": False}})
+        db["jobsBlue"].update_many({"force": {"$exists": False}}, {"$set": {"force": False}})
 
     def down(self) -> None:
         logging.info("Remove the force field from all the jobs")
@@ -29,88 +24,6 @@ def down(self) -> None:
         db["jobsBlue"].update_many({}, {"$unset": {"force": ""}})
 
     def validate(self) -> None:
-        logging.info("Ensure that a random selection of jobs have the 'force' field set to False")
+        logging.info("Ensure that a random selection of jobs have the 'force' field")
 
-        def custom_validation(doc: JobSnapshot) -> None:
-            if doc.force:
-                raise ValueError("force should be False")
-
-        check_documents(DocCls=JobSnapshot, sample_size=10, custom_validation=custom_validation)
-        if JobSnapshot.objects(force=False).count() != JobSnapshot.objects.count():
-            raise ValueError('All the objects should have the "force" field, set to False')
-
-
-# --- JobSnapshot ---
-# copied from libcommon.queue.Job, as a snapshot of when the migration was created
-class Status(enum.Enum):
-    WAITING = "waiting"
-    STARTED = "started"
-    SUCCESS = "success"
-    ERROR = "error"
-    CANCELLED = "cancelled"
-    SKIPPED = "skipped"
-
-
-# START monkey patching ### hack ###
-# see https://github.com/sbdchd/mongo-types#install
-U = TypeVar("U", bound=Document)
-
-
-def no_op(self, x):  # type: ignore
-    return self
-
-
-QuerySet.__class_getitem__ = types.MethodType(no_op, QuerySet)
-
-
-class QuerySetManager(Generic[U]):
-    def __get__(self, instance: object, cls: Type[U]) -> QuerySet[U]:
-        return QuerySet(cls, cls._get_collection())
-
-
-# END monkey patching ### hack ###
-
-
-class JobSnapshot(Document):
-    """A job in the mongoDB database
-
-    Args:
-        type (`str`): The type of the job, identifies the queue
-        dataset (`str`): The dataset on which to apply the job.
-        config (`str`, optional): The config on which to apply the job.
-        split (`str`, optional): The config on which to apply the job.
-        unicity_id (`str`): A string that identifies the job uniquely. Only one job with the same unicity_id can be in
-          the started state.
-        namespace (`str`): The dataset namespace (user or organization) if any, else the dataset name (canonical name).
-        force (`bool`, optional): If True, the job SHOULD not be skipped. Defaults to False.
-        status (`Status`, optional): The status of the job. Defaults to Status.WAITING.
-        created_at (`datetime`): The creation date of the job.
-        started_at (`datetime`, optional): When the job has started.
-        finished_at (`datetime`, optional): When the job has finished.
-    """
-
-    meta = {
-        "collection": "jobsBlue",
-        "db_alias": "queue",
-        "indexes": [
-            "status",
-            ("type", "status"),
-            ("type", "dataset", "status"),
-            ("type", "dataset", "config", "split", "status"),
-            ("status", "type", "created_at", "namespace"),
-            "-created_at",
-        ],
-    }
-    type = StringField(required=True)
-    dataset = StringField(required=True)
-    config = StringField()
-    split = StringField()
-    unicity_id = StringField(required=True)
-    namespace = StringField(required=True)
-    force = BooleanField(default=False)
-    status = EnumField(Status, default=Status.WAITING)
-    created_at = DateTimeField(required=True)
-    started_at = DateTimeField()
-    finished_at = DateTimeField()
-
-    objects = QuerySetManager["JobSnapshot"]()
+        check_documents(DocCls=Job, sample_size=10)
diff --git a/jobs/mongodb_migration/src/mongodb_migration/migrations/_20221117223000_cache_generic_response.py b/jobs/mongodb_migration/src/mongodb_migration/migrations/_20221117223000_cache_generic_response.py
--- a/jobs/mongodb_migration/src/mongodb_migration/migrations/_20221117223000_cache_generic_response.py
+++ b/jobs/mongodb_migration/src/mongodb_migration/migrations/_20221117223000_cache_generic_response.py
@@ -1,38 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 The HuggingFace Authors.
 
+import contextlib
 import logging
-import types
-from datetime import datetime, timezone
-from enum import Enum
-from http import HTTPStatus
-from typing import Generic, Type, TypeVar
 
-from bson import ObjectId
-from mongoengine import Document
+from libcommon.simple_cache import CachedResponse
 from mongoengine.connection import get_db
-from mongoengine.fields import (
-    DateTimeField,
-    DictField,
-    EnumField,
-    ObjectIdField,
-    StringField,
-)
-from mongoengine.queryset.queryset import QuerySet
+from pymongo.errors import InvalidName
 
 from mongodb_migration.check import check_documents
 from mongodb_migration.migration import Migration
 
-
-class CacheKind(Enum):
-    SPLITS = "/splits"
-    FIRST_ROWS = "/first-rows"
-
-
 db_name = "cache"
 splitsResponseCollection = "splitsResponse"
 firstRowsResponseCollection = "firstRowsResponse"
 cachedResponseCollection = "cachedResponsesBlue"
+SPLITS_KIND = "/splits"
+FIRST_ROWS_KIND = "/first-rows"
 
 
 # connection already occurred in the main.py (caveat: we use globals)
@@ -45,54 +29,56 @@ def up(self) -> None:
         db = get_db(db_name)
         # Copy the data from the previous collections (splitsResponse, firstRowsResponse) to
         # the new generic collection (cachedResponse)
-        for splits_response in db[splitsResponseCollection].find():
-            if not isinstance(splits_response, dict):
-                # for mypy
-                raise ValueError("splits_response should be a dict")
-            db[cachedResponseCollection].insert_one(
-                {
-                    "_id": splits_response.get("_id"),
-                    "kind": CacheKind.SPLITS.value,
-                    # ^ "kind" is a new field
-                    "dataset": splits_response.get("dataset_name"),
-                    "config": None,
-                    "split": None,
-                    # ^ "config" and "split" are None for kind=/splits
-                    "http_status": splits_response.get("http_status"),
-                    "error_code": splits_response.get("error_code"),
-                    "content": splits_response.get("response"),
-                    # ^ "response" field has been renamed to "content"
-                    "worker_version": splits_response.get("worker_version"),
-                    "dataset_git_revision": splits_response.get("dataset_git_revision"),
-                    "details": splits_response.get("details"),
-                    "updated_at": splits_response.get("updated_at"),
-                    # "stale" field is not used anymore
-                }
-            )
-        for first_rows_response in db[firstRowsResponseCollection].find():
-            if not isinstance(first_rows_response, dict):
-                # for mypy
-                raise ValueError("first_rows_response should be a dict")
-            db[cachedResponseCollection].insert_one(
-                {
-                    "_id": first_rows_response.get("_id"),
-                    "kind": CacheKind.FIRST_ROWS.value,
-                    # ^ "kind" is a new field
-                    "dataset": first_rows_response.get("dataset_name"),
-                    "config": first_rows_response.get("config_name"),
-                    "split": first_rows_response.get("split_name"),
-                    # ^ "config" and "split" are None for kind=/splits
-                    "http_status": first_rows_response.get("http_status"),
-                    "error_code": first_rows_response.get("error_code"),
-                    "content": first_rows_response.get("response"),
-                    # ^ "response" field has been renamed to "content"
-                    "worker_version": first_rows_response.get("worker_version"),
-                    "dataset_git_revision": first_rows_response.get("dataset_git_revision"),
-                    "details": first_rows_response.get("details"),
-                    "updated_at": first_rows_response.get("updated_at"),
-                    # "stale" field is not used anymore
-                }
-            )
+        with contextlib.suppress(InvalidName):
+            for splits_response in db[splitsResponseCollection].find():
+                if not isinstance(splits_response, dict):
+                    # for mypy
+                    raise ValueError("splits_response should be a dict")
+                db[cachedResponseCollection].insert_one(
+                    {
+                        "_id": splits_response.get("_id"),
+                        "kind": SPLITS_KIND,
+                        # ^ "kind" is a new field
+                        "dataset": splits_response.get("dataset_name"),
+                        "config": None,
+                        "split": None,
+                        # ^ "config" and "split" are None for kind=/splits
+                        "http_status": splits_response.get("http_status"),
+                        "error_code": splits_response.get("error_code"),
+                        "content": splits_response.get("response"),
+                        # ^ "response" field has been renamed to "content"
+                        "worker_version": splits_response.get("worker_version"),
+                        "dataset_git_revision": splits_response.get("dataset_git_revision"),
+                        "details": splits_response.get("details"),
+                        "updated_at": splits_response.get("updated_at"),
+                        # "stale" field is not used anymore
+                    }
+                )
+        with contextlib.suppress(InvalidName):
+            for first_rows_response in db[firstRowsResponseCollection].find():
+                if not isinstance(first_rows_response, dict):
+                    # for mypy
+                    raise ValueError("first_rows_response should be a dict")
+                db[cachedResponseCollection].insert_one(
+                    {
+                        "_id": first_rows_response.get("_id"),
+                        "kind": FIRST_ROWS_KIND,
+                        # ^ "kind" is a new field
+                        "dataset": first_rows_response.get("dataset_name"),
+                        "config": first_rows_response.get("config_name"),
+                        "split": first_rows_response.get("split_name"),
+                        # ^ "config" and "split" are None for kind=/splits
+                        "http_status": first_rows_response.get("http_status"),
+                        "error_code": first_rows_response.get("error_code"),
+                        "content": first_rows_response.get("response"),
+                        # ^ "response" field has been renamed to "content"
+                        "worker_version": first_rows_response.get("worker_version"),
+                        "dataset_git_revision": first_rows_response.get("dataset_git_revision"),
+                        "details": first_rows_response.get("details"),
+                        "updated_at": first_rows_response.get("updated_at"),
+                        # "stale" field is not used anymore
+                    }
+                )
         # We will not delete the old collections for now. It will be made in a later migration.
         # Also: no need to create indexes on the new collection, mongoengine will do it automatically on the next
         # request.
@@ -100,105 +86,26 @@ def up(self) -> None:
     def down(self) -> None:
         logging.info(f"Delete the {cachedResponseCollection} collection")
         db = get_db(db_name)
-        db[cachedResponseCollection].drop()
+        with contextlib.suppress(InvalidName):
+            db[cachedResponseCollection].drop()
 
     def validate(self) -> None:
         logging.info("Validate the migrated documents")
 
-        def custom_validation(doc: CachedResponseSnapshot) -> None:
-            if doc.kind not in (CacheKind.SPLITS.value, CacheKind.FIRST_ROWS.value):
-                raise ValueError("kind should be /splits or /first-rows")
-
-        check_documents(DocCls=CachedResponseSnapshot, sample_size=10, custom_validation=custom_validation)
+        check_documents(DocCls=CachedResponse, sample_size=10)
 
         db = get_db(db_name)
-        splits_responses_count = db[splitsResponseCollection].count_documents({})
-        first_rows_responses_count = db[firstRowsResponseCollection].count_documents({})
-        cached_responses_count = CachedResponseSnapshot.objects.count()
+        try:
+            splits_responses_count = db[splitsResponseCollection].count_documents({})
+        except InvalidName:
+            splits_responses_count = 0
+        try:
+            first_rows_responses_count = db[firstRowsResponseCollection].count_documents({})
+        except InvalidName:
+            first_rows_responses_count = 0
+        cached_responses_count = CachedResponse.objects.count()
         if splits_responses_count + first_rows_responses_count > cached_responses_count:
             raise ValueError(
                 f"Some documents are missing in the new collection: splitsResponse ({splits_responses_count}),"
                 f" firstRowsResponse ({first_rows_responses_count}), cachedResponseBlue ({cached_responses_count})"
             )
-
-
-# --- CachedResponseSnapshot ---
-
-# START monkey patching ### hack ###
-# see https://github.com/sbdchd/mongo-types#install
-U = TypeVar("U", bound=Document)
-
-
-def no_op(self, x):  # type: ignore
-    return self
-
-
-QuerySet.__class_getitem__ = types.MethodType(no_op, QuerySet)
-
-
-class QuerySetManager(Generic[U]):
-    def __get__(self, instance: object, cls: Type[U]) -> QuerySet[U]:
-        return QuerySet(cls, cls._get_collection())
-
-
-# END monkey patching ### hack ###
-
-
-def get_datetime() -> datetime:
-    return datetime.now(timezone.utc)
-
-
-# cache of any endpoint
-class CachedResponseSnapshot(Document):
-    """A response to an endpoint request, cached in the mongoDB database
-
-    Args:
-        kind (`str`): The kind of the cached response, identifies the endpoint
-        dataset (`str`): The requested dataset.
-        config (`str`, optional): The requested config, if any.
-        split (`str`, optional): The requested split, if any.
-        http_status (`HTTPStatus`): The HTTP status code.
-        error_code (`str`, optional): The error code, if any.
-        content (`dict`): The content of the cached response. Can be an error or a valid content.
-        details (`dict`, optional): Additional details, eg. a detailed error that we don't want to send as a response.
-        updated_at (`datetime`): When the cache entry has been last updated.
-        worker_version (`str`): The semver version of the worker that cached the response.
-        dataset_git_revision (`str`): The commit (of the git dataset repo) used to generate the response.
-    """
-
-    id = ObjectIdField(db_field="_id", primary_key=True, default=ObjectId)
-
-    kind = StringField(required=True, unique_with=["dataset", "config", "split"])
-    dataset = StringField(required=True)
-    config = StringField()
-    split = StringField()
-
-    http_status = EnumField(HTTPStatus, required=True)
-    error_code = StringField()
-    content = DictField(required=True)
-    worker_version = StringField()
-    dataset_git_revision = StringField()
-
-    details = DictField()
-    updated_at = DateTimeField(default=get_datetime)
-
-    meta = {
-        "collection": cachedResponseCollection,
-        "db_alias": db_name,
-        "indexes": [
-            ("kind", "dataset", "config", "split"),
-            ("dataset", "kind", "http_status"),
-            ("kind", "http_status", "dataset"),
-            ("kind", "http_status", "error_code"),
-            ("kind", "id"),
-        ],
-    }
-    objects = QuerySetManager["CachedResponseSnapshot"]()
-
-
-# Fix issue with mongoengine: https://github.com/MongoEngine/mongoengine/issues/1242#issuecomment-810501601
-# mongoengine automatically sets "config" and "splits" as required fields, because they are listed in the unique_with
-# field of the "kind" field. But it's an error, since unique indexes (which are used to enforce unique_with) accept
-# null values, see https://www.mongodb.com/docs/v5.0/core/index-unique/#unique-index-and-missing-field.
-CachedResponseSnapshot.config.required = False  # type: ignore
-CachedResponseSnapshot.split.required = False  # type: ignore
diff --git a/jobs/mongodb_migration/src/mongodb_migration/migrations/_20230126164900_queue_job_add_priority.py b/jobs/mongodb_migration/src/mongodb_migration/migrations/_20230126164900_queue_job_add_priority.py
--- a/jobs/mongodb_migration/src/mongodb_migration/migrations/_20230126164900_queue_job_add_priority.py
+++ b/jobs/mongodb_migration/src/mongodb_migration/migrations/_20230126164900_queue_job_add_priority.py
@@ -1,15 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 The HuggingFace Authors.
 
-import enum
 import logging
-import types
-from typing import Generic, Type, TypeVar
 
-from mongoengine import Document
+from libcommon.queue import Job
 from mongoengine.connection import get_db
-from mongoengine.fields import BooleanField, DateTimeField, EnumField, StringField
-from mongoengine.queryset.queryset import QuerySet
 
 from mongodb_migration.check import check_documents
 from mongodb_migration.migration import Migration
@@ -19,9 +14,9 @@
 class MigrationAddPriorityToJob(Migration):
     def up(self) -> None:
         # See https://docs.mongoengine.org/guide/migration.html#example-1-addition-of-a-field
-        logging.info("Add the priority field, with the default value ('normal'), to all the jobs")
+        logging.info("If missing, add the priority field with the default value ('normal') to the jobs")
         db = get_db("queue")
-        db["jobsBlue"].update_many({}, {"$set": {"priority": "normal"}})
+        db["jobsBlue"].update_many({"priority": {"$exists": False}}, {"$set": {"priority": "normal"}})
 
     def down(self) -> None:
         logging.info("Remove the priority field from all the jobs")
@@ -29,95 +24,6 @@ def down(self) -> None:
         db["jobsBlue"].update_many({}, {"$unset": {"priority": ""}})
 
     def validate(self) -> None:
-        logging.info("Ensure that a random selection of jobs have the 'priority' field set to 'normal'")
+        logging.info("Ensure that a random selection of jobs have the 'priority' field set")
 
-        def custom_validation(doc: JobSnapshot) -> None:
-            if doc.priority != Priority.NORMAL:
-                raise ValueError("priority should be 'normal'")
-
-        check_documents(DocCls=JobSnapshot, sample_size=10, custom_validation=custom_validation)
-        if JobSnapshot.objects(priority=Priority.NORMAL).count() != JobSnapshot.objects.count():
-            raise ValueError('All the objects should have the "priority" field, set to "normal"')
-
-
-# --- JobSnapshot ---
-# copied from libcommon.queue.Job, as a snapshot of when the migration was created
-class Status(enum.Enum):
-    WAITING = "waiting"
-    STARTED = "started"
-    SUCCESS = "success"
-    ERROR = "error"
-    CANCELLED = "cancelled"
-    SKIPPED = "skipped"
-
-
-class Priority(enum.Enum):
-    NORMAL = "normal"
-    LOW = "low"
-
-
-# START monkey patching ### hack ###
-# see https://github.com/sbdchd/mongo-types#install
-U = TypeVar("U", bound=Document)
-
-
-def no_op(self, x):  # type: ignore
-    return self
-
-
-QuerySet.__class_getitem__ = types.MethodType(no_op, QuerySet)
-
-
-class QuerySetManager(Generic[U]):
-    def __get__(self, instance: object, cls: Type[U]) -> QuerySet[U]:
-        return QuerySet(cls, cls._get_collection())
-
-
-# END monkey patching ### hack ###
-
-
-class JobSnapshot(Document):
-    """A job in the mongoDB database
-
-    Args:
-        type (`str`): The type of the job, identifies the queue
-        dataset (`str`): The dataset on which to apply the job.
-        config (`str`, optional): The config on which to apply the job.
-        split (`str`, optional): The config on which to apply the job.
-        unicity_id (`str`): A string that identifies the job uniquely. Only one job with the same unicity_id can be in
-          the started state.
-        namespace (`str`): The dataset namespace (user or organization) if any, else the dataset name (canonical name).
-        force (`bool`, optional): If True, the job SHOULD not be skipped. Defaults to False.
-        priority (`Priority`, optional): The priority of the job. Defaults to Priority.NORMAL.
-        status (`Status`, optional): The status of the job. Defaults to Status.WAITING.
-        created_at (`datetime`): The creation date of the job.
-        started_at (`datetime`, optional): When the job has started.
-        finished_at (`datetime`, optional): When the job has finished.
-    """
-
-    meta = {
-        "collection": "jobsBlue",
-        "db_alias": "queue",
-        "indexes": [
-            "status",
-            ("type", "status"),
-            ("type", "dataset", "status"),
-            ("type", "dataset", "config", "split", "status", "force", "priority"),
-            ("status", "type", "created_at", "namespace", "unicity_id", "priority"),
-            "-created_at",
-        ],
-    }
-    type = StringField(required=True)
-    dataset = StringField(required=True)
-    config = StringField()
-    split = StringField()
-    unicity_id = StringField(required=True)
-    namespace = StringField(required=True)
-    force = BooleanField(default=False)
-    priority = EnumField(Priority, default=Priority.NORMAL)
-    status = EnumField(Status, default=Status.WAITING)
-    created_at = DateTimeField(required=True)
-    started_at = DateTimeField()
-    finished_at = DateTimeField()
-
-    objects = QuerySetManager["JobSnapshot"]()
+        check_documents(DocCls=Job, sample_size=10)


Test patch: