testing

requirements.txt

@@ -1 +1,2 @@
 python-qbittorrent==0.4.3
+sqlalchemy==2.0.27

tarc/main.py

@@ -13,145 +13,66 @@ import sys
 import re
 import uuid
 import argparse
-import sqlite3
-
 from datetime import datetime, timezone
 
 import qbittorrent
+from sqlalchemy import create_engine, event
+from sqlalchemy.orm import sessionmaker
+
+from .models import Base, Client
 
 # SCHEMA format is YYYYMMDDX
 SCHEMA = 202410060
 
-
+def init_db(engine):
-def init_db(conn):
     """
     Initialize database
     """
-
+    # Create all tables first
-    c = conn.cursor()
+    Base.metadata.create_all(engine)
-    c.executescript(
+    
-        f"""
+    # Set the schema version using SQLAlchemy primitives
-        PRAGMA user_version = {SCHEMA};
+    @event.listens_for(engine, 'connect')
-
+    def set_sqlite_pragma(dbapi_connection, connection_record):
-        CREATE TABLE IF NOT EXISTS clients (
+        cursor = dbapi_connection.cursor()
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
+        cursor.execute(f"PRAGMA user_version = {SCHEMA}")
-            name TEXT NOT NULL UNIQUE,
+        cursor.close()
-            uuid TEXT NOT NULL UNIQUE,
-            endpoint TEXT NOT NULL,
-            last_seen DATETIME NOT NULL
-        );
-
-        CREATE TABLE IF NOT EXISTS torrents (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            info_hash_v1 TEXT NOT NULL UNIQUE,
-            info_hash_v2 TEXT UNIQUE,
-            file_count INTEGER NOT NULL,
-            completed_on DATETIME NOT NULL
-        );
-
-        CREATE TABLE IF NOT EXISTS torrent_clients (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            torrent_id INTEGER NOT NULL,
-            client_id INTEGER NOT NULL,
-            name TEXT NOT NULL,
-            content_path TEXT NOT NULL,
-            last_seen DATETIME NOT NULL,
-            FOREIGN KEY (torrent_id) REFERENCES torrents(id),
-            FOREIGN KEY (client_id) REFERENCES clients(id),
-            UNIQUE (torrent_id, client_id)
-        );
-
-        CREATE TABLE IF NOT EXISTS trackers (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            url TEXT NOT NULL UNIQUE,
-            last_seen DATETIME NOT NULL
-        );
-
-        CREATE TABLE IF NOT EXISTS torrent_trackers (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            client_id INTEGER NOT NULL,
-            torrent_id INTEGER NOT NULL,
-            tracker_id INTEGER NOT NULL,
-            last_seen DATETIME NOT NULL,
-            FOREIGN KEY (client_id) REFERENCES clients(id),
-            FOREIGN KEY (torrent_id) REFERENCES torrents(id),
-            FOREIGN KEY (tracker_id) REFERENCES trackers(id),
-            UNIQUE (client_id, torrent_id, tracker_id)
-        );
-
-        CREATE TABLE IF NOT EXISTS files (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            size INTEGER NOT NULL,
-            oshash TEXT NOT NULL UNIQUE,
-            hash TEXT UNIQUE
-        );
-
-        CREATE TABLE IF NOT EXISTS torrent_files (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            file_id INTEGER NOT NULL,
-            torrent_id INTEGER NOT NULL,
-            client_id INTEGER NOT NULL,
-            file_index INTEGER NOT NULL,
-            file_path TEXT NOT NULL,
-            is_downloaded BOOLEAN NOT NULL,
-            last_checked DATETIME NOT NULL,
-            FOREIGN KEY (file_id) REFERENCES files(id),
-            FOREIGN KEY (torrent_id) REFERENCES torrents(id),
-            FOREIGN KEY (client_id) REFERENCES clients(id),
-            UNIQUE (file_id, torrent_id, client_id, file_index)
-        );
-        """
-    )
-    conn.commit()
-    c.close()
 
 
-def list_tables(conn):
+def list_tables(engine):
     """
     List all tables in database
     """
-    c = conn.cursor()
+    return Base.metadata.tables.keys()
-    c.execute("SELECT name FROM sqlite_master WHERE type='table';")
-    table_list = c.fetchall()
-    c.close()
-    return [table[0] for table in table_list]
 
 
-def add_client(conn, name, endpoint, last_seen):
+def add_client(session, name, endpoint, last_seen):
     """
     Add a new client endpoint to database
     """
-    c = conn.cursor()
+    new_client = Client(
-    c.execute(
+        uuid=str(uuid.uuid4()),
-        f"""
+        name=name,
-        INSERT INTO clients (uuid, name, endpoint, last_seen)
+        endpoint=endpoint,
-        VALUES ("{uuid.uuid4()}", "{name}", "{endpoint}", "{last_seen}");
+        last_seen=last_seen
-        """
     )
-    conn.commit()
+    session.add(new_client)
-    c.close()
+    session.commit()
 
 
-def find_client(conn, endpoint):
+def find_client(session, endpoint):
     """
     Find existing client
     """
-    c = conn.cursor()
+    clients = session.query(Client.id, Client.name, Client.uuid).filter(Client.endpoint == endpoint).all()
-    c.execute(f'SELECT id, name, uuid FROM clients WHERE endpoint="{endpoint}";')
+    return clients
-    response = c.fetchall()
-    c.close()
-    return response
 
 
-def list_clients(conn):
+def list_clients(session):
     """
     List all stored clients
     """
-    c = conn.cursor()
+    return session.query(Client).all()
-    c.execute("SELECT * FROM clients;")
-    rows = c.fetchall()
-    c.close()
-    return rows
 
 
 def main():
@@ -185,26 +106,42 @@ def main():
         else:
             STORAGE = args.storage
         try:
-            sqlitedb = sqlite3.connect(STORAGE)
+            engine = create_engine(f"sqlite:///{STORAGE}")
-            tables = list_tables(sqlitedb)
+            tables = list_tables(engine)
-        except sqlite3.DatabaseError as e:
+            Session = sessionmaker(bind=engine)
+            session = Session()
+        except Exception as e:
             print(f'[ERROR]: Database Error "{STORAGE}" ({str(e)})')
             sys.exit(1)
         if len(tables) == 0:
             print(f"[INFO]: Initializing database at {STORAGE}")
-            init_db(sqlitedb)
+            init_db(engine)
-        cursor = sqlitedb.cursor()
+
-        cursor.execute("PRAGMA user_version;")
+        # Check schema version using SQLAlchemy primitives
-        SCHEMA_FOUND = cursor.fetchone()[0]
+        schema_found = None
-        cursor.close()
+        
-        if not SCHEMA == SCHEMA_FOUND:
+        @event.listens_for(engine, 'connect')
-            print(f"[ERROR]: SCHEMA {SCHEMA_FOUND}, expected {SCHEMA}")
+        def get_sqlite_pragma(dbapi_connection, connection_record):
+            nonlocal schema_found
+            cursor = dbapi_connection.cursor()
+            cursor.execute("PRAGMA user_version")
+            schema_found = cursor.fetchone()[0]
+            cursor.close()
+        
+        # Force a connection to trigger the event
+        with engine.connect() as conn:
+            pass
+        
+        if not SCHEMA == schema_found:
+            print(f"[ERROR]: SCHEMA {schema_found}, expected {SCHEMA}")
             sys.exit(1)
         if not args.directory is None:
             print("[INFO]: --directory is not implemented")
             sys.exit(0)
         elif not args.endpoint is None:
             qb = qbittorrent.Client(args.endpoint)
+            if args.username and args.password:
+                qb.login(args.username, args.password)
             if qb.qbittorrent_version is None:
                 print(f'[ERROR]: Couldn\'t find client version at "{args.endpoint}"')
                 sys.exit(1)
@@ -217,14 +154,12 @@ def main():
                 print(
                     f'[INFO]: Found qbittorrent {qb.qbittorrent_version} at "{args.endpoint}"'
                 )
-            clients = find_client(sqlitedb, args.endpoint)
+            clients = find_client(session, args.endpoint)
             if args.confirm_add:
                 if len(clients) == 0:
                     if not args.name is None:
-                        now = datetime.now(timezone.utc).isoformat(
+                        now = datetime.now(timezone.utc)
-                            sep=" ", timespec="seconds"
+                        add_client(session, args.name, args.endpoint, now)
-                        )
-                        add_client(sqlitedb, args.name, args.endpoint, now)
                         print(f"[INFO]: Added client {args.name} ({args.endpoint})")
                     else:
                         print("[ERROR]: Must specify --name for a new client")
@@ -244,6 +179,7 @@ def main():
             elif len(clients) == 1:
                 torrents = qb.torrents()
                 print(f"[INFO]: There are {len(torrents)} torrents\n")
+                
                 for torrent in torrents[:2]:
                     files = qb.get_torrent_files(torrent["hash"])
                     trackers = qb.get_torrent_trackers(torrent["hash"])

tarc/models.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+"""
+Database models for Torrent Archiver
+"""
+
+from sqlalchemy import Column, Integer, String, Boolean, DateTime, ForeignKey, UniqueConstraint
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import relationship
+
+Base = declarative_base()
+
+class Client(Base):
+    __tablename__ = 'clients'
+    
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    name = Column(String, nullable=False, unique=True)
+    uuid = Column(String, nullable=False, unique=True)
+    endpoint = Column(String, nullable=False)
+    last_seen = Column(DateTime, nullable=False)
+    
+    torrent_clients = relationship("TorrentClient", back_populates="client")
+    torrent_trackers = relationship("TorrentTracker", back_populates="client")
+    torrent_files = relationship("TorrentFile", back_populates="client")
+
+
+class Torrent(Base):
+    __tablename__ = 'torrents'
+    
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    info_hash_v1 = Column(String, nullable=False, unique=True)
+    info_hash_v2 = Column(String, unique=True)
+    file_count = Column(Integer, nullable=False)
+    completed_on = Column(DateTime, nullable=False)
+    
+    torrent_clients = relationship("TorrentClient", back_populates="torrent")
+    torrent_trackers = relationship("TorrentTracker", back_populates="torrent")
+    torrent_files = relationship("TorrentFile", back_populates="torrent")
+
+
+class TorrentClient(Base):
+    __tablename__ = 'torrent_clients'
+    
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    torrent_id = Column(Integer, ForeignKey('torrents.id'), nullable=False)
+    client_id = Column(Integer, ForeignKey('clients.id'), nullable=False)
+    name = Column(String, nullable=False)
+    content_path = Column(String, nullable=False)
+    last_seen = Column(DateTime, nullable=False)
+    
+    torrent = relationship("Torrent", back_populates="torrent_clients")
+    client = relationship("Client", back_populates="torrent_clients")
+    
+    __table_args__ = (UniqueConstraint('torrent_id', 'client_id'),)
+
+
+class Tracker(Base):
+    __tablename__ = 'trackers'
+    
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    url = Column(String, nullable=False, unique=True)
+    last_seen = Column(DateTime, nullable=False)
+    
+    torrent_trackers = relationship("TorrentTracker", back_populates="tracker")
+
+
+class TorrentTracker(Base):
+    __tablename__ = 'torrent_trackers'
+    
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    client_id = Column(Integer, ForeignKey('clients.id'), nullable=False)
+    torrent_id = Column(Integer, ForeignKey('torrents.id'), nullable=False)
+    tracker_id = Column(Integer, ForeignKey('trackers.id'), nullable=False)
+    last_seen = Column(DateTime, nullable=False)
+    
+    client = relationship("Client", back_populates="torrent_trackers")
+    torrent = relationship("Torrent", back_populates="torrent_trackers")
+    tracker = relationship("Tracker", back_populates="torrent_trackers")
+    
+    __table_args__ = (UniqueConstraint('client_id', 'torrent_id', 'tracker_id'),)
+
+
+class File(Base):
+    __tablename__ = 'files'
+    
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    size = Column(Integer, nullable=False)
+    oshash = Column(String, nullable=False, unique=True)
+    hash = Column(String, unique=True)
+    
+    torrent_files = relationship("TorrentFile", back_populates="file")
+
+
+class TorrentFile(Base):
+    __tablename__ = 'torrent_files'
+    
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    file_id = Column(Integer, ForeignKey('files.id'), nullable=False)
+    torrent_id = Column(Integer, ForeignKey('torrents.id'), nullable=False)
+    client_id = Column(Integer, ForeignKey('clients.id'), nullable=False)
+    file_index = Column(Integer, nullable=False)
+    file_path = Column(String, nullable=False)
+    is_downloaded = Column(Boolean, nullable=False)
+    last_checked = Column(DateTime, nullable=False)
+    
+    file = relationship("File", back_populates="torrent_files")
+    torrent = relationship("Torrent", back_populates="torrent_files")
+    client = relationship("Client", back_populates="torrent_files")
+    
+    __table_args__ = (UniqueConstraint('file_id', 'torrent_id', 'client_id', 'file_index'),)