Remove DownloadTask data class, use DescribeDownload for both queue and db records

Author: andrii-i

jupyter_scheduler/download_manager.py

@@ -1,68 +1,43 @@
-from dataclasses import dataclass
-from datetime import datetime
 from multiprocessing import Queue
 from typing import List, Optional
 
-from jupyter_scheduler.orm import Downloads, create_session, generate_uuid
-from jupyter_scheduler.pydantic_v1 import BaseModel
+from jupyter_scheduler.models import DescribeDownload
+from jupyter_scheduler.orm import Download, create_session, generate_uuid
 from jupyter_scheduler.utils import get_utc_timestamp
 from jupyter_scheduler.pydantic_v1 import BaseModel
 
 
-class DescribeDownload(BaseModel):
-    job_id: str
-    download_id: str
-    download_initiated_time: int
-
-    class Config:
-        orm_mode = True
-
-
-@dataclass
-class DownloadTask:
-    job_id: str
-    download_id: str
-    download_initiated_time: int
-
-    def __lt__(self, other):
-        return self.download_initiated_time < other.download_initiated_time
-
-    def __str__(self):
-        download_initiated_time = datetime.fromtimestamp(self.download_initiated_time / 1e3)
-        return f"Id: {self.job_id}, Download initiated: {download_initiated_time}"
-
-
 class DownloadRecordManager:
     def __init__(self, db_url):
         self.session = create_session(db_url)
 
     def put(self, download: DescribeDownload):
         with self.session() as session:
-            new_download = Downloads(**download.dict())
-            session.add(new_download)
+            download = Download(**download.dict())
+            session.add(download)
             session.commit()
 
     def get(self, job_id: str) -> Optional[DescribeDownload]:
         with self.session() as session:
-            download = session.query(Downloads).filter(Downloads.job_id == job_id).first()
+            download = session.query(Download).filter(Download.job_id == job_id).first()
 
         if download:
             return DescribeDownload.from_orm(download)
         else:
             return None
 
-    def get_tasks(self) -> List[DescribeDownload]:
+    def get_downloads(self) -> List[DescribeDownload]:
         with self.session() as session:
-            return session.query(Downloads).order_by(Downloads.download_initiated_time).all()
+            return session.query(Download).order_by(Download.download_initiated_time).all()
 
     def delete_download(self, download_id: str):
         with self.session() as session:
-            session.query(Downloads).filter(Downloads.download_id == download_id).delete()
+            session.query(Download).filter(Download.download_id == download_id).delete()
             session.commit()
 
     def delete_job_downloads(self, job_id: str):
         with self.session() as session:
-            session.query(Downloads).filter(Downloads.job_id == job_id).delete()
+            session.query(Download).filter(Download.job_id == job_id).delete()
             session.commit()
 
 
@@ -74,18 +49,13 @@ def __init__(self, db_url: str):
     def download_from_staging(self, job_id: str):
         download_initiated_time = get_utc_timestamp()
         download_id = generate_uuid()
-        download_cache = DescribeDownload(
-            job_id=job_id,
-            download_id=download_id,
-            download_initiated_time=download_initiated_time,
-        )
-        self.record_manager.put(download_cache)
-        download_task = DownloadTask(
+        download = DescribeDownload(
             job_id=job_id,
             download_id=download_id,
             download_initiated_time=download_initiated_time,
         )
-        self.queue.put(download_task)
+        self.record_manager.put(download)
+        self.queue.put(download)
 
     def delete_download(self, download_id: str):
         self.record_manager.delete_download(download_id)
@@ -94,11 +64,6 @@ def delete_job_downloads(self, job_id: str):
         self.record_manager.delete_job_downloads(job_id)
 
     def populate_queue(self):
-        tasks = self.record_manager.get_tasks()
-        for task in tasks:
-            download_task = DownloadTask(
-                job_id=task.job_id,
-                download_id=task.download_id,
-                download_initiated_time=task.download_initiated_time,
-            )
-            self.queue.put(download_task)
+        downloads = self.record_manager.get_downloads()
+        for download in downloads:
+            self.queue.put(download)

jupyter_scheduler/executors.py

@@ -1,5 +1,4 @@
 import io
-import multiprocessing
 import os
 import shutil
 import tarfile
@@ -14,11 +13,7 @@
 
 from jupyter_scheduler.models import DescribeJob, JobFeature, Status
 from jupyter_scheduler.orm import Job, create_session
-from jupyter_scheduler.download_manager import (
-    DescribeDownload,
-    Downloads,
-    DownloadTask,
-)
+from jupyter_scheduler.download_manager import DescribeDownload, Download
 from jupyter_scheduler.models import DescribeJob, JobFeature, Status
 from jupyter_scheduler.orm import Job, create_session, generate_uuid
 from jupyter_scheduler.parameterize import add_parameters
@@ -159,25 +154,21 @@ def execute(self):
         finally:
             self.add_side_effects_files(staging_dir)
             self.create_output_files(job, nb)
+            self._download_from_staging(job.job_id)
 
-    def download_from_staging(self, job_id: str):
+    def _download_from_staging(self, job_id: str):
         download_initiated_time = get_utc_timestamp()
         download_id = generate_uuid()
-        download_cache = DescribeDownload(
+        download = DescribeDownload(
             job_id=job_id,
             download_id=download_id,
             download_initiated_time=download_initiated_time,
         )
         with self.db_session() as session:
-            new_download = Downloads(**download_cache.dict())
-            session.add(new_download)
+            download_record = Download(**download.dict())
+            session.add(download_record)
             session.commit()
-        download_task = DownloadTask(
-            job_id=job_id,
-            download_id=download_id,
-            download_initiated_time=download_initiated_time,
-        )
-        self.download_queue.put(download_task)
+        self.download_queue.put(download)
 
     def add_side_effects_files(self, staging_dir: str):
         """Scan for side effect files potentially created after input file execution and update the job's packaged_files with these files"""

jupyter_scheduler/models.py

@@ -295,3 +295,15 @@ class JobFeature(str, Enum):
     output_filename_template = "output_filename_template"
     stop_job = "stop_job"
     delete_job = "delete_job"
+
+
+class DescribeDownload(BaseModel):
+    job_id: str
+    download_id: str
+    download_initiated_time: int
+
+    class Config:
+        orm_mode = True
+
+    def __str__(self) -> str:
+        return self.json()

jupyter_scheduler/orm.py

@@ -111,7 +111,7 @@ class JobDefinition(CommonColumns, Base):
     active = Column(Boolean, default=True)
 
 
-class Downloads(Base):
+class Download(Base):
     __tablename__ = "downloads"
     job_id = Column(String(36), primary_key=True)
     download_id = Column(String(36), primary_key=True)