propagate redownload, use download_id in get download

Author: andrii-i

jupyter_scheduler/download_manager.py

@@ -17,9 +17,9 @@ def put(self, download: DescribeDownload):
             session.add(download)
             session.commit()
 
-    def get(self, job_id: str) -> Optional[DescribeDownload]:
+    def get(self, download_id: str) -> Optional[DescribeDownload]:
         with self.session() as session:
-            download = session.query(Download).filter(Download.job_id == job_id).first()
+            download = session.query(Download).filter(Download.download_id == download_id).first()
 
         if download:
             return DescribeDownload.from_orm(download)
@@ -46,13 +46,14 @@ def __init__(self, db_url: str):
         self.record_manager = DownloadRecordManager(db_url=db_url)
         self.queue = Queue()
 
-    def download_from_staging(self, job_id: str):
+    def download_from_staging(self, job_id: str, redownload: bool):
         download_initiated_time = get_utc_timestamp()
         download_id = generate_uuid()
         download = DescribeDownload(
             job_id=job_id,
             download_id=download_id,
             download_initiated_time=download_initiated_time,
+            redownload=redownload,
         )
         self.record_manager.put(download)
         self.queue.put(download)

jupyter_scheduler/download_runner.py

@@ -23,11 +23,11 @@ def __init__(
     async def process_download_queue(self):
         while not self.download_manager.queue.empty():
             download = self.download_manager.queue.get()
-            cache = self.download_manager.record_manager.get(download.job_id)
-            if not cache or not download:
+            download_record = self.download_manager.record_manager.get(download.download_id)
+            if not download_record:
                 continue
-            await self.job_files_manager.copy_from_staging(cache.job_id)
-            self.download_manager.record_manager.delete_download(cache.download_id)
+            await self.job_files_manager.copy_from_staging(download.job_id, download.redownload)
+            self.download_manager.delete_download(download.download_id)
 
     async def start(self):
         self.download_manager.populate_queue()

jupyter_scheduler/executors.py

@@ -163,6 +163,7 @@ def _download_from_staging(self, job_id: str):
             job_id=job_id,
             download_id=download_id,
             download_initiated_time=download_initiated_time,
+            redownload=True,
         )
         with self.db_session() as session:
             download_record = Download(**download.dict())

jupyter_scheduler/handlers.py

@@ -395,16 +395,8 @@ def get(self):
 
 
 class FilesDownloadHandler(ExtensionHandlerMixin, APIHandler):
-    # _job_files_manager = None
     _download_from_staging = None
 
-    # @property
-    # def job_files_manager(self):
-    #     if not self._job_files_manager:
-    #         self._job_files_manager = self.settings.get("job_files_manager", None)
-
-    #     return self._job_files_manager
-
     @property
     def download_from_staging(self):
         if not self._download_from_staging:
@@ -414,10 +406,9 @@ def download_from_staging(self):
 
     @authenticated
     async def get(self, job_id):
-        # redownload = self.get_query_argument("redownload", False)
+        redownload = self.get_query_argument("redownload", False)
         try:
-            # await self.job_files_manager.copy_from_staging(job_id=job_id, redownload=redownload)
-            self.download_from_staging(job_id)
+            self.download_from_staging(job_id, redownload)
         except Exception as e:
             self.log.exception(e)
             raise HTTPError(500, str(e)) from e

jupyter_scheduler/models.py

@@ -301,6 +301,7 @@ class DescribeDownload(BaseModel):
     job_id: str
     download_id: str
     download_initiated_time: int
+    redownload: bool
 
     class Config:
         orm_mode = True

jupyter_scheduler/orm.py

@@ -116,6 +116,7 @@ class Download(Base):
     job_id = Column(String(36), primary_key=True)
     download_id = Column(String(36), primary_key=True)
     download_initiated_time = Column(Integer)
+    redownload = Column(Boolean, default=False)
 
 
 def create_tables(db_url, drop_tables=False):