From e35803ce1a019282c5540204024df4b6e0f79ec9 Mon Sep 17 00:00:00 2001
From: Yannik Schmidt <yannik.schmidt@potaris.de>
Date: Sat, 15 Feb 2025 19:43:20 +0100
Subject: [PATCH] feat: basic download tracking in db

---
 db.py          | 34 ++++++++++++++++++++++++++++++++++
 software.py    |  4 ++++
 statekeeper.py | 34 +++++++++++++++++++++++++++++++++-
 3 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 db.py

diff --git a/db.py b/db.py
new file mode 100644
index 0000000..1c3fff9
--- /dev/null
+++ b/db.py
@@ -0,0 +1,34 @@
+from sqlalchemy import Column, String, Integer, Boolean, create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker, scoped_session
+
+Base = declarative_base()
+
+class Download(Base):
+
+    __tablename__ = 'files'
+
+    path = Column(String, primary_key=True)
+    size = Column(Integer)
+    type = Column(String)
+    finished = Column(Boolean)
+
+class Database:
+    def __init__(self, db_url="sqlite:///database.db"):
+        self.engine = create_engine(db_url, echo=True)
+        self.session_factory = sessionmaker(bind=self.engine)
+        self.Session = scoped_session(self.session_factory)  # Thread-safe sessions
+
+        # Automatically create tables
+        Base.metadata.create_all(self.engine)
+
+    def session(self):
+        """Returns a new session (or an existing one if in the same thread)."""
+        return self.Session()
+
+    def close_session(self):
+        """Closes the current session."""
+        self.Session.remove()
+
+# Singleton instance of Database
+db = Database()
\ No newline at end of file
diff --git a/software.py b/software.py
index f3a9288..b50acc4 100644
--- a/software.py
+++ b/software.py
@@ -11,6 +11,7 @@ import jinja_helper
 import threading
 import sys
 import tkinter
+import statekeeper
 
 class Software:
 
@@ -131,7 +132,10 @@ class Software:
         except IndexError:
             print("No main_dir:", path)
             raise AssertionError("No main_dir for this software")
+        
+        statekeeper.log_begin_download(remote_file)
         local_file = self.backend.get(remote_file, self.cache_dir, wait=True)
+        statekeeper.log_end_download(remote_file)
 
         # execute or unpack #
         if local_file.endswith(".exe"):
diff --git a/statekeeper.py b/statekeeper.py
index 04fb35d..21587bf 100644
--- a/statekeeper.py
+++ b/statekeeper.py
@@ -1,6 +1,9 @@
 import requests
 import os
+import sqlalchemy
 import threading
+from db import db, Download
+from sqlalchemy import or_, and_
 
 def add_to_download_queue(url, path):
     '''The download is added to the global queue and downloaded eventually'''
@@ -33,4 +36,33 @@ def _download(url, path):
     
     else:
 
-        raise AssertionError("Non-200 Response for:", url, path, response.status_code, response.text)
\ No newline at end of file
+        raise AssertionError("Non-200 Response for:", url, path, response.status_code, response.text)
+
+def log_begin_download(path):
+
+    session = db.session()
+    print("Current path", path)
+    path_exists = session.query(Download).filter(and_(Download.path==path, Download.finished==False)).first()
+    if path_exists:
+        print("DAFUG", path_exists)
+        print("WTF", path_exists.path)
+        raise AssertionError("ERROR: {} is already downloading.".format(path))
+    else:
+        print("Adding to download log:", path)
+        session.merge(Download(path=path, size=0, type="download", finished=False))
+        session.commit()
+
+    db.close_session()
+
+def log_end_download(path):
+
+    session = db.session()
+    path_exists = session.query(Download).filter(Download.path==path).first()
+    if not path_exists:
+        raise AssertionError("ERROR: {} is not downloading/cannot remove.".format(path))
+    else:
+        print("Removing from download log:", path)
+        session.merge(Download(path=path, size=0, type="download", finished=True))
+        session.commit()
+
+    db.close_session()
\ No newline at end of file