feat: 🎸 add concept of Resource

.github/workflows/_unit-tests-python.yml

@@ -49,6 +49,7 @@ jobs:
         env:
           CACHE_MONGO_URL: mongodb://localhost:${{ env.mongo-port }}
           QUEUE_MONGO_URL: mongodb://localhost:${{ env.mongo-port }}
+          MONGODB_MIGRATION_MONGO_URL: mongodb://localhost:${{ env.mongo-port }}
         run: poetry run python -m pytest -s --cov=./ --cov-report=xml:./coverage.xml --cov-report=term tests
       - name: Prepare codecov flag (slash "/" is not allowed)
         id: remove-slash

chart/templates/worker/config-names/_container.tpl

@@ -21,6 +21,7 @@
     # overridden
     value: {{ .Values.configNames.queue.maxJobsPerNamespace | quote }}
   volumeMounts:
+  {{ include "volumeMountAssetsRW" . | nindent 2 }}
   {{ include "volumeMountCache" . | nindent 2 }}
   securityContext:
     allowPrivilegeEscalation: false

chart/templates/worker/config-names/deployment.yaml

@@ -21,6 +21,7 @@ spec:
     spec:
       {{- include "image.imagePullSecrets" . | nindent 6 }}
       initContainers:
+        {{ include "initContainerAssets" . | nindent 8 }}
         {{ include "initContainerCache" . | nindent 8 }}
       containers: {{ include "containerWorkerConfigNames" . | nindent 8 }}
       nodeSelector: {{ toYaml .Values.configNames.nodeSelector | nindent 8 }}

chart/templates/worker/dataset-info/_container.tpl

@@ -17,6 +17,8 @@
     # value: {{ .Values.queue.maxJobsPerNamespace | quote }}
     # overridden
     value: {{ .Values.datasetInfo.queue.maxJobsPerNamespace | quote }}
+  volumeMounts:
+  {{ include "volumeMountAssetsRW" . | nindent 2 }}
   securityContext:
     allowPrivilegeEscalation: false
   resources: {{ toYaml .Values.datasetInfo.resources | nindent 4 }}

chart/templates/worker/dataset-info/deployment.yaml

@@ -20,6 +20,8 @@ spec:
       labels: {{ include "labels.datasetInfo" . | nindent 8 }}
     spec:
       {{- include "image.imagePullSecrets" . | nindent 6 }}
+      initContainers:
+        {{ include "initContainerAssets" . | nindent 8 }}
       containers: {{ include "containerWorkerDatasetInfo" . | nindent 8 }}
       nodeSelector: {{ toYaml .Values.datasetInfo.nodeSelector | nindent 8 }}
       tolerations: {{ toYaml .Values.datasetInfo.tolerations | nindent 8 }}

chart/templates/worker/parquet-and-dataset-info/_container.tpl

@@ -45,6 +45,7 @@
   - name: PARQUET_AND_DATASET_INFO_URL_TEMPLATE
     value: {{ .Values.parquetAndDatasetInfo.urlTemplate | quote }}
   volumeMounts:
+  {{ include "volumeMountAssetsRW" . | nindent 2 }}
   {{ include "volumeMountCache" . | nindent 2 }}
   securityContext:
     allowPrivilegeEscalation: false

chart/templates/worker/parquet-and-dataset-info/deployment.yaml

@@ -20,7 +20,9 @@ spec:
       labels: {{ include "labels.parquetAndDatasetInfo" . | nindent 8 }}
     spec:
       {{- include "image.imagePullSecrets" . | nindent 6 }}
-      initContainers: {{ include "initContainerCache" . | nindent 8 }}
+      initContainers:
+        {{ include "initContainerAssets" . | nindent 8 }}
+        {{ include "initContainerCache" . | nindent 8 }}
       containers: {{ include "containerWorkerParquetAndDatasetInfo" . | nindent 8 }}
       nodeSelector: {{ toYaml .Values.parquetAndDatasetInfo.nodeSelector | nindent 8 }}
       tolerations: {{ toYaml .Values.parquetAndDatasetInfo.tolerations | nindent 8 }}

chart/templates/worker/parquet/_container.tpl

@@ -17,6 +17,8 @@
     # value: {{ .Values.queue.maxJobsPerNamespace | quote }}
     # overridden
     value: {{ .Values.parquet.queue.maxJobsPerNamespace | quote }}
+  volumeMounts:
+  {{ include "volumeMountAssetsRW" . | nindent 2 }}
   securityContext:
     allowPrivilegeEscalation: false
   resources: {{ toYaml .Values.parquet.resources | nindent 4 }}

chart/templates/worker/parquet/deployment.yaml

@@ -20,6 +20,8 @@ spec:
       labels: {{ include "labels.parquet" . | nindent 8 }}
     spec:
       {{- include "image.imagePullSecrets" . | nindent 6 }}
+      initContainers:
+        {{ include "initContainerAssets" . | nindent 8 }}
       containers: {{ include "containerWorkerParquet" . | nindent 8 }}
       nodeSelector: {{ toYaml .Values.parquet.nodeSelector | nindent 8 }}
       tolerations: {{ toYaml .Values.parquet.tolerations | nindent 8 }}

chart/templates/worker/sizes/_container.tpl

@@ -17,6 +17,8 @@
     # value: {{ .Values.queue.maxJobsPerNamespace | quote }}
     # overridden
     value: {{ .Values.sizes.queue.maxJobsPerNamespace | quote }}
+  volumeMounts:
+  {{ include "volumeMountAssetsRW" . | nindent 2 }}
   securityContext:
     allowPrivilegeEscalation: false
   resources: {{ toYaml .Values.sizes.resources | nindent 4 }}

chart/templates/worker/sizes/deployment.yaml

@@ -20,6 +20,8 @@ spec:
       labels: {{ include "labels.sizes" . | nindent 8 }}
     spec:
       {{- include "image.imagePullSecrets" . | nindent 6 }}
+      initContainers:
+        {{ include "initContainerAssets" . | nindent 8 }}
       containers: {{ include "containerWorkerSizes" . | nindent 8 }}
       nodeSelector: {{ toYaml .Values.sizes.nodeSelector | nindent 8 }}
       tolerations: {{ toYaml .Values.sizes.tolerations | nindent 8 }}

chart/templates/worker/split-names/_container.tpl

@@ -21,6 +21,7 @@
     # overridden
     value: {{ .Values.splitNames.queue.maxJobsPerNamespace | quote }}
   volumeMounts:
+  {{ include "volumeMountAssetsRW" . | nindent 2 }}
   {{ include "volumeMountCache" . | nindent 2 }}
   securityContext:
     allowPrivilegeEscalation: false

chart/templates/worker/split-names/deployment.yaml

@@ -21,6 +21,7 @@ spec:
     spec:
       {{- include "image.imagePullSecrets" . | nindent 6 }}
       initContainers:
+        {{ include "initContainerAssets" . | nindent 8 }}
         {{ include "initContainerCache" . | nindent 8 }}
       containers: {{ include "containerWorkerSplitNames" . | nindent 8 }}
       nodeSelector: {{ toYaml .Values.splitNames.nodeSelector | nindent 8 }}

chart/templates/worker/splits/_container.tpl

@@ -21,6 +21,7 @@
     # overridden
     value: {{ .Values.splits.queue.maxJobsPerNamespace | quote }}
   volumeMounts:
+  {{ include "volumeMountAssetsRW" . | nindent 2 }}
   {{ include "volumeMountCache" . | nindent 2 }}
   securityContext:
     allowPrivilegeEscalation: false

chart/templates/worker/splits/deployment.yaml

@@ -20,7 +20,9 @@ spec:
       labels: {{ include "labels.splits" . | nindent 8 }}
     spec:
       {{- include "image.imagePullSecrets" . | nindent 6 }}
-      initContainers: {{ include "initContainerCache" . | nindent 8 }}
+      initContainers:
+        {{ include "initContainerAssets" . | nindent 8 }}
+        {{ include "initContainerCache" . | nindent 8 }}
       containers: {{ include "containerWorkerSplits" . | nindent 8 }}
       nodeSelector: {{ toYaml .Values.splits.nodeSelector | nindent 8 }}
       tolerations: {{ toYaml .Values.splits.tolerations | nindent 8 }}

jobs/mongodb_migration/src/mongodb_migration/config.py

@@ -6,42 +6,37 @@
 from environs import Env
 from libcommon.config import CacheConfig, CommonConfig, QueueConfig
 
-from mongodb_migration.database_migrations import connect_to_database
-
 MONGODB_MIGRATION_MONGO_DATABASE = "datasets_server_maintenance"
-MONGO_DATABASE_MONGO_URL = "mongodb://localhost:27017"
+MONGODB_MIGRATION_MONGO_URL = "mongodb://localhost:27017"
 
 
 @dataclass
-class MongodbMigrationConfig:
+class DatabaseMigrationsConfig:
     mongo_database: str = MONGODB_MIGRATION_MONGO_DATABASE
-    mongo_url: str = MONGO_DATABASE_MONGO_URL
-
-    def __post_init__(self):
-        connect_to_database(database=self.mongo_database, host=self.mongo_url)
+    mongo_url: str = MONGODB_MIGRATION_MONGO_URL
 
     @staticmethod
-    def from_env() -> "MongodbMigrationConfig":
+    def from_env() -> "DatabaseMigrationsConfig":
         env = Env(expand_vars=True)
         with env.prefixed("MONGODB_MIGRATION_"):
-            return MongodbMigrationConfig(
+            return DatabaseMigrationsConfig(
                 mongo_database=env.str(name="MONGO_DATABASE", default=MONGODB_MIGRATION_MONGO_DATABASE),
-                mongo_url=env.str(name="MONGO_URL", default=MONGO_DATABASE_MONGO_URL),
+                mongo_url=env.str(name="MONGO_URL", default=MONGODB_MIGRATION_MONGO_URL),
             )
 
 
 @dataclass
 class JobConfig:
     cache: CacheConfig = field(default_factory=CacheConfig)
     common: CommonConfig = field(default_factory=CommonConfig)
-    mongodb_migration: MongodbMigrationConfig = field(default_factory=MongodbMigrationConfig)
+    database_migrations: DatabaseMigrationsConfig = field(default_factory=DatabaseMigrationsConfig)
     queue: QueueConfig = field(default_factory=QueueConfig)
 
     @staticmethod
     def from_env() -> "JobConfig":
         return JobConfig(
             common=CommonConfig.from_env(),
             cache=CacheConfig.from_env(),
-            mongodb_migration=MongodbMigrationConfig.from_env(),
+            database_migrations=DatabaseMigrationsConfig.from_env(),
             queue=QueueConfig.from_env(),
         )

jobs/mongodb_migration/src/mongodb_migration/constants.py

@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2022 The HuggingFace Authors.
+
+DATABASE_MIGRATIONS_MONGOENGINE_ALIAS = "maintenance"

jobs/mongodb_migration/src/mongodb_migration/database_migrations.py

@@ -4,10 +4,12 @@
 import types
 from typing import Generic, Type, TypeVar
 
-from mongoengine import Document, DoesNotExist, connect
+from mongoengine import Document, DoesNotExist
 from mongoengine.fields import StringField
 from mongoengine.queryset.queryset import QuerySet
 
+from mongodb_migration.constants import DATABASE_MIGRATIONS_MONGOENGINE_ALIAS
+
 # START monkey patching ### hack ###
 # see https://github.com/sbdchd/mongo-types#install
 U = TypeVar("U", bound=Document)
@@ -27,12 +29,6 @@ def __get__(self, instance: object, cls: Type[U]) -> QuerySet[U]:
 
 # END monkey patching ### hack ###
 
-DATABASE_ALIAS = "maintenance"
-
-
-def connect_to_database(database: str, host: str) -> None:
-    connect(db=database, alias=DATABASE_ALIAS, host=host)
-
 
 class DatabaseMigration(Document):
     """A database migration that has already been executed.
@@ -44,7 +40,7 @@ class DatabaseMigration(Document):
 
     meta = {
         "collection": "databaseMigrations",
-        "db_alias": DATABASE_ALIAS,
+        "db_alias": DATABASE_MIGRATIONS_MONGOENGINE_ALIAS,
     }
     version = StringField(required=True)
     description = StringField(required=True)

jobs/mongodb_migration/src/mongodb_migration/main.py

@@ -3,18 +3,34 @@
 
 import sys
 
+from libcommon.resources import (
+    CacheDatabaseResource,
+    LogResource,
+    QueueDatabaseResource,
+)
+
 from mongodb_migration.collector import MigrationsCollector
 from mongodb_migration.config import JobConfig
 from mongodb_migration.plan import Plan
+from mongodb_migration.resources import MigrationsDatabaseResource
 
 if __name__ == "__main__":
     job_config = JobConfig.from_env()
-    collected_migrations = MigrationsCollector().get_migrations()
-    try:
-        Plan(collected_migrations=collected_migrations).execute()
-        sys.exit(0)
-    except Exception:
-        sys.exit(1)
+    with (
+        LogResource(log_level=job_config.common.log_level),
+        # ^ first resource to be acquired, in order to have logs as soon as possible
+        CacheDatabaseResource(database=job_config.cache.mongo_database, host=job_config.cache.mongo_url),
+        QueueDatabaseResource(database=job_config.queue.mongo_database, host=job_config.queue.mongo_url),
+        MigrationsDatabaseResource(
+            database=job_config.database_migrations.mongo_database, host=job_config.database_migrations.mongo_url
+        ),
+    ):
+        collected_migrations = MigrationsCollector().get_migrations()
+        try:
+            Plan(collected_migrations=collected_migrations).execute()
+            sys.exit(0)
+        except Exception:
+            sys.exit(1)
 
 # See:
 #  https://blog.appsignal.com/2020/04/14/dissecting-rails-migrationsl.html

jobs/mongodb_migration/src/mongodb_migration/resources.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2022 The HuggingFace Authors.
+
+from dataclasses import dataclass, field
+
+from libcommon.resources import DatabaseResource
+
+from mongodb_migration.constants import DATABASE_MIGRATIONS_MONGOENGINE_ALIAS
+
+
+class MigrationsDatabaseConnectionFailure(Exception):
+    pass
+
+
+@dataclass
+class MigrationsDatabaseResource(DatabaseResource):
+    """
+    A resource that represents a connection to the migrations database.
+
+    Args:
+        database (:obj:`str`): The name of the database.
+        host (:obj:`str`): The host of the database. It must start with ``mongodb://`` or ``mongodb+srv://``.
+    """
+
+    mongoengine_alias: str = field(default=DATABASE_MIGRATIONS_MONGOENGINE_ALIAS, init=False)

jobs/mongodb_migration/tests/conftest.py

@@ -1,29 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 The HuggingFace Authors.
 
-from pytest import MonkeyPatch, fixture
+from environs import Env
+from pytest import fixture
 
-from mongodb_migration.config import JobConfig
 
-
-# see https://github.com/pytest-dev/pytest/issues/363#issuecomment-406536200
 @fixture(scope="session")
-def monkeypatch_session():
-    monkeypatch_session = MonkeyPatch()
-    monkeypatch_session.setenv("CACHE_MONGO_DATABASE", "datasets_server_cache_test")
-    monkeypatch_session.setenv("QUEUE_MONGO_DATABASE", "datasets_server_queue_test")
-    monkeypatch_session.setenv("MONGODB_MIGRATION_MONGO_DATABASE", "datasets_server_maintenance_test")
-    yield monkeypatch_session
-    monkeypatch_session.undo()
+def env() -> Env:
+    return Env(expand_vars=True)
 
 
-@fixture(scope="session", autouse=True)
-def app_config(monkeypatch_session: MonkeyPatch) -> JobConfig:
-    job_config = JobConfig.from_env()
-    if (
-        "test" not in job_config.cache.mongo_database
-        or "test" not in job_config.queue.mongo_database
-        or "test" not in job_config.mongodb_migration.mongo_database
-    ):
-        raise ValueError("Test must be launched on a test mongo database")
-    return job_config
+@fixture(scope="session")
+def mongo_host(env: Env) -> str:
+    try:
+        return env.str(name="MONGODB_MIGRATION_MONGO_URL")
+    except Exception as e:
+        raise ValueError("MONGODB_MIGRATION_MONGO_URL is not set") from e

jobs/mongodb_migration/tests/test_plan.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 The HuggingFace Authors.
 
-from typing import List, Optional, Type
+from typing import Iterator, List, Optional, Type
 
 import pytest
 
@@ -11,10 +11,20 @@
 )
 from mongodb_migration.migration import IrreversibleMigration, Migration
 from mongodb_migration.plan import Plan, SavedMigrationsError
+from mongodb_migration.resources import MigrationsDatabaseResource
+
+
+@pytest.fixture(scope="module")
+def migrations_database_resource(mongo_host: str) -> Iterator[MigrationsDatabaseResource]:
+    database = "datasets_server_migrations_test"
+    if "test" not in database:
+        raise ValueError("Test must be launched on a test mongo database")
+    with MigrationsDatabaseResource(database=database, host=mongo_host) as resource:
+        yield resource
 
 
 @pytest.fixture(autouse=True)
-def clean_mongo_database() -> None:
+def clean_mongo_database(migrations_database_resource: MigrationsDatabaseResource) -> None:
     _clean_maintenance_database()
 
 

jobs/mongodb_migration/tests/test_resources.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2022 The HuggingFace Authors.
+
+from mongoengine import Document
+from mongoengine.fields import StringField
+
+from mongodb_migration.resources import MigrationsDatabaseResource
+
+
+def test_cache_database(mongo_host: str) -> None:
+    resource = MigrationsDatabaseResource(database="test_migrations_database", host=mongo_host)
+
+    class User(Document):
+        name = StringField()
+        meta = {"db_alias": resource.mongo_connection.mongoengine_alias}
+
+    assert len(User.objects()) == 0  # type: ignore
+    # clean
+    User.drop_collection()  # type: ignore
+    resource.release()