Use APScheduler

Use APScheduler

Use APScheduler

Use APScheduler

Use APScheduler

Author: BenjaminCharmes

pydatalab/pyproject.toml

@@ -53,6 +53,7 @@ server = [
     "Flask-PyMongo ~= 2.3",
     "Flask-Mail ~= 0.10",
     "Flask-Compress ~= 1.15",
+    "APScheduler ~= 3.10",
     "Werkzeug ~= 3.0",
     "python-dotenv ~= 1.0",
     "pillow ~= 11.0",

pydatalab/src/pydatalab/export_utils.py

@@ -146,36 +146,21 @@ def create_eln_file(collection_id: str, output_path: str) -> None:
 
             item_metadata = {k: v for k, v in item.items() if k not in ["_id", "file_ObjectIds"]}
 
-            for key, value in item_metadata.items():
-                if isinstance(value, ObjectId):
-                    item_metadata[key] = str(value)
-                elif isinstance(value, datetime):
-                    item_metadata[key] = value.isoformat()
-                elif isinstance(value, dict):
-                    item_metadata[key] = _convert_objectids_in_dict(value)
-                elif isinstance(value, list):
-                    item_metadata[key] = [
-                        str(v)
-                        if isinstance(v, ObjectId)
-                        else v.isoformat()
-                        if isinstance(v, datetime)
-                        else _convert_objectids_in_dict(v)
-                        if isinstance(v, dict)
-                        else v
-                        for v in value
-                    ]
+            item_metadata = _convert_objectids_in_dict(item_metadata)
 
             with open(item_folder / "metadata.json", "w", encoding="utf-8") as f:
                 json.dump(item_metadata, f, indent=2, ensure_ascii=False)
 
             if item.get("file_ObjectIds"):
-                for file_id in item["file_ObjectIds"]:
+                for file_id in item.get("file_ObjectIds", []):
                     file_data = flask_mongo.db.files.find_one({"_id": ObjectId(file_id)})
-                    if file_data and file_data.get("location"):
+                    if file_data:
                         source_path = Path(file_data["location"])
                         if source_path.exists():
-                            dest_path = item_folder / file_data["name"]
-                            shutil.copy2(source_path, dest_path)
+                            dest_file = item_folder / file_data["name"]
+                            shutil.copy2(source_path, dest_file)
+                        else:
+                            print(f"Warning: File not found on disk: {file_data['location']}")
 
         with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zipf:
             for file_path in root_folder.rglob("*"):

pydatalab/src/pydatalab/models/export_task.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timezone
 from enum import Enum
 
 from pydantic import BaseModel, Field
@@ -23,7 +23,8 @@ class ExportTask(BaseModel):
     )
     creator_id: str = Field(..., description="ID of the user who created the export")
     created_at: datetime = Field(
-        default_factory=datetime.now, description="When the task was created"
+        default_factory=lambda: datetime.now(tz=timezone.utc),
+        description="When the task was created",
     )
     completed_at: datetime | None = Field(None, description="When the task was completed")
     file_path: str | None = Field(None, description="Path to the generated .eln file")

pydatalab/src/pydatalab/mongo.py

@@ -217,5 +217,6 @@ def create_user_fts():
     ret += db.export_tasks.create_index(
         "created_at", name="export task created at", background=background
     )
+    ret += db.export_tasks.create_index("status", name="export task status", background=background)
 
     return ret

pydatalab/src/pydatalab/routes/v0_1/export.py

@@ -1,17 +1,17 @@
 import os
-import threading
 import uuid
 from datetime import datetime, timezone
 from pathlib import Path
 
-from flask import Blueprint, jsonify, send_file
+from flask import Blueprint, current_app, jsonify, send_file
 from flask_login import current_user
 
 from pydatalab.config import CONFIG
 from pydatalab.export_utils import create_eln_file
 from pydatalab.models.export_task import ExportStatus, ExportTask
 from pydatalab.mongo import flask_mongo
 from pydatalab.permissions import active_users_or_get_only
+from pydatalab.scheduler import export_scheduler
 
 EXPORT = Blueprint("export", __name__)
 
@@ -21,62 +21,63 @@
 def _(): ...
 
 
-def _generate_export_in_background(task_id: str, collection_id: str):
+def _generate_export_in_background(task_id: str, collection_id: str, app):
     """Background function to generate the .eln file.
 
     Parameters:
         task_id: ID of the export task
         collection_id: ID of the collection to export
+        app: Flask application instance
     """
-    try:
-        flask_mongo.db.export_tasks.update_one(
-            {"task_id": task_id}, {"$set": {"status": ExportStatus.PROCESSING}}
-        )
-
-        export_dir = Path(CONFIG.FILE_DIRECTORY) / "exports"
-        export_dir.mkdir(exist_ok=True)
-
-        output_path = export_dir / f"{task_id}.eln"
-        create_eln_file(collection_id, str(output_path))
-
-        flask_mongo.db.export_tasks.update_one(
-            {"task_id": task_id},
-            {
-                "$set": {
-                    "status": ExportStatus.READY,
-                    "file_path": str(output_path),
-                    "completed_at": datetime.now(tz=timezone.utc),
-                }
-            },
-        )
-
-    except Exception as e:
-        flask_mongo.db.export_tasks.update_one(
-            {"task_id": task_id},
-            {
-                "$set": {
-                    "status": ExportStatus.ERROR,
-                    "error_message": str(e),
-                    "completed_at": datetime.now(tz=timezone.utc),
-                }
-            },
-        )
+    with app.app_context():
+        try:
+            flask_mongo.db.export_tasks.update_one(
+                {"task_id": task_id}, {"$set": {"status": ExportStatus.PROCESSING}}
+            )
+
+            export_dir = Path(CONFIG.FILE_DIRECTORY) / "exports"
+            export_dir.mkdir(exist_ok=True)
+
+            output_path = export_dir / f"{task_id}.eln"
+            create_eln_file(collection_id, str(output_path))
+
+            flask_mongo.db.export_tasks.update_one(
+                {"task_id": task_id},
+                {
+                    "$set": {
+                        "status": ExportStatus.READY,
+                        "file_path": str(output_path),
+                        "completed_at": datetime.now(tz=timezone.utc),
+                    }
+                },
+            )
+
+        except Exception as e:
+            flask_mongo.db.export_tasks.update_one(
+                {"task_id": task_id},
+                {
+                    "$set": {
+                        "status": ExportStatus.ERROR,
+                        "error_message": str(e),
+                        "completed_at": datetime.now(tz=timezone.utc),
+                    }
+                },
+            )
 
 
 @EXPORT.route("/collections/<string:collection_id>/export", methods=["POST"])
 def start_collection_export(collection_id: str):
-    """Start exporting a collection to .eln format.
-
-    Parameters:
-        collection_id: The collection ID to export
+    from pydatalab.permissions import get_default_permissions
 
-    Returns:
-        JSON response with task_id and status_url
-    """
+    collection_exists = flask_mongo.db.collections.find_one({"collection_id": collection_id})
+    if not collection_exists:
+        return jsonify({"status": "error", "message": "Collection not found"}), 404
 
-    collection = flask_mongo.db.collections.find_one({"collection_id": collection_id})
-    if not collection:
-        return jsonify({"status": "error", "message": f"Collection {collection_id} not found"}), 404
+    collection_with_perms = flask_mongo.db.collections.find_one(
+        {"collection_id": collection_id, **get_default_permissions(user_only=True)}
+    )
+    if not collection_with_perms:
+        return jsonify({"status": "error", "message": "Access denied"}), 403
 
     task_id = str(uuid.uuid4())
 
@@ -94,9 +95,13 @@ def start_collection_export(collection_id: str):
 
     flask_mongo.db.export_tasks.insert_one(export_task.dict())
 
-    thread = threading.Thread(target=_generate_export_in_background, args=(task_id, collection_id))
-    thread.daemon = True
-    thread.start()
+    app = current_app._get_current_object()
+
+    export_scheduler.add_job(
+        func=_generate_export_in_background,
+        args=[task_id, collection_id, app],
+        job_id=f"export_{task_id}",
+    )
 
     return jsonify(
         {"status": "success", "task_id": task_id, "status_url": f"/exports/{task_id}/status"}

pydatalab/src/pydatalab/scheduler.py

@@ -0,0 +1,49 @@
+from apscheduler.executors.pool import ThreadPoolExecutor
+from apscheduler.jobstores.memory import MemoryJobStore
+from apscheduler.schedulers.background import BackgroundScheduler
+
+
+class ExportScheduler:
+    _instance = None
+    _scheduler = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def init_scheduler(self):
+        if self._scheduler is None:
+            jobstores = {"default": MemoryJobStore()}
+            executors = {"default": ThreadPoolExecutor(10)}
+            job_defaults = {"coalesce": False, "max_instances": 3}
+
+            self._scheduler = BackgroundScheduler(
+                jobstores=jobstores, executors=executors, job_defaults=job_defaults, timezone="UTC"
+            )
+            self._scheduler.start()
+        return self._scheduler
+
+    def get_scheduler(self):
+        if self._scheduler is None:
+            self.init_scheduler()
+        return self._scheduler
+
+    def add_job(self, func, args, job_id=None):
+        """Add a job to the scheduler or run it in a thread if APScheduler not available."""
+        if self._scheduler:
+            return self._scheduler.add_job(func=func, args=args, id=job_id, replace_existing=True)
+        else:
+            import threading
+
+            thread = threading.Thread(target=func, args=args)
+            thread.daemon = True
+            thread.start()
+            return None
+
+    def shutdown(self):
+        if self._scheduler and self._scheduler.running:
+            self._scheduler.shutdown()
+
+
+export_scheduler = ExportScheduler()