From b4108e4b36333d7a45289b73781a48a45b8b876b Mon Sep 17 00:00:00 2001 From: WorldTeacher Date: Sat, 1 Nov 2025 21:32:28 +0100 Subject: [PATCH] add __all__ --- src/komcache/__init__.py | 1 + src/komcache/cache.py | 202 +++++++++++++++++++++----------- src/komcache/schemas/mariadb.py | 28 +++-- 3 files changed, 152 insertions(+), 79 deletions(-) diff --git a/src/komcache/__init__.py b/src/komcache/__init__.py index 917efa6..e355816 100644 --- a/src/komcache/__init__.py +++ b/src/komcache/__init__.py @@ -1 +1,2 @@ +__all__ = ["KomCache"] from .cache import KomCache diff --git a/src/komcache/cache.py b/src/komcache/cache.py index 51d9a07..0cf14ec 100644 --- a/src/komcache/cache.py +++ b/src/komcache/cache.py @@ -1,11 +1,14 @@ -from typing import Any -from sqlalchemy import create_engine, Column, String, Integer, Date, text -from sqlalchemy.orm import sessionmaker, declarative_base -from sqlalchemy.exc import SQLAlchemyError -from komcache.schemas.sqlite import CREATE_SQLITE_TABLES -from komcache.schemas.mariadb import CREATE_MARIADB_TABLES -from komconfig import KomConfig +import time +from typing import Any, Tuple, Union + import loguru +from komconfig import KomConfig +from sqlalchemy import Column, Date, Integer, String, create_engine, text +from sqlalchemy.exc import OperationalError, SQLAlchemyError +from sqlalchemy.orm import declarative_base, sessionmaker + +from komcache.schemas.mariadb import CREATE_MARIADB_TABLES +from komcache.schemas.sqlite import CREATE_SQLITE_TABLES log = loguru.logger log.remove() @@ -28,6 +31,10 @@ class KomGrabber(Base): complete = Column(Integer, nullable=False) +def protect(url: str) -> str: + return "mysql+pymysql://user:pass@host:3306/dbname" + + class KomCache: def __init__(self, db_path: str = ""): # Default to empty string if not provided self.db_path = db_path or config.cache.path @@ -35,15 +42,24 @@ class KomCache: if config.cache.mode == "local": self.db_path = db_path or config.cache.path log.debug(f"Cache path: {self.db_path}") - self.engine = create_engine(f"sqlite:///{self.db_path}") + self.engine = create_engine(f"sqlite:///{self.db_path}", pool_pre_ping=True) elif config.cache.mode == "remote": db_url = ( config.cache.url ) # e.g., "mysql+pymysql://user:pass@host:3306/dbname" - log.debug(f"Using remote DB URL: {db_url}") - self.engine = create_engine(db_url) - self.Session = sessionmaker(bind=self.engine) + log.debug(f"Using remote DB URL: {protect(db_url)}") + self.engine = create_engine( + db_url, + pool_pre_ping=True, + pool_recycle=1800, + pool_size=5, + max_overflow=10, + pool_timeout=30, + connect_args={"connect_timeout": 10}, + ) + + self.Session = sessionmaker(bind=self.engine, expire_on_commit=False) # if tables do not exist, create them if config.cache.mode == "local": if not self.query( @@ -54,6 +70,29 @@ class KomCache: if not self.query("SHOW TABLES LIKE 'komgrabber'"): self.create_table() + def _run(self, fn_desc: str, callable_, retries: int = 2, *args, **kwargs): + attempt = 0 + while True: + try: + return callable_(*args, **kwargs) + except OperationalError as e: + # MySQL server has gone away (2006) / Lost connection (2013) + if attempt < retries and any( + code in str(e.orig) for code in ("2006", "2013") + ): + attempt += 1 + wait = 1 * attempt + log.warning( + f"{fn_desc} failed due to connection loss (attempt {attempt}/{retries}). Retrying in {wait}s." + ) + time.sleep(wait) + continue + log.error(f"{fn_desc} failed with OperationalError: {e}") + return None + except SQLAlchemyError as e: + log.error(f"{fn_desc} failed with SQLAlchemyError: {e}") + return None + def create_table(self): """Ensure all tables are created in the database.""" if config.cache.mode == "local": @@ -90,87 +129,110 @@ class KomCache: return False def query(self, query: str, args: dict[str, Any] = None): + """Run an arbitrary SQL statement. + For SELECT (or other row‑returning) statements: returns list of rows. + For non‑SELECT: executes and commits, returns []. + """ if args is None: args = {} - try: + + def _do(): session = self.Session() - result = session.execute(text(query), args).fetchall() - session.close() - return result - except SQLAlchemyError as e: - log.error(f"Error executing query: {e}") - return [] + try: + result = session.execute(text(query), args) + # SQLAlchemy 1.4/2.0: result.returns_rows tells us if rows are present + if getattr(result, "returns_rows", False): + rows = result.fetchall() + return rows + # Non-row statements: commit if they mutate + first = query.lstrip().split(None, 1)[0].upper() + if first in { + "INSERT", + "UPDATE", + "DELETE", + "REPLACE", + "ALTER", + "CREATE", + "DROP", + "TRUNCATE", + }: + session.commit() + return [] + finally: + session.close() + + result = self._run("query", _do) + return result or [] def insert(self, query: str, args: dict[str, Any]) -> bool: - try: + # (Optionally you can now just call self.query(query, args)) + def _do(): session = self.Session() - session.execute(text(query), args) - session.commit() - session.close() - return True - except SQLAlchemyError as e: - log.error(f"Error inserting data: {e}") - return False + try: + session.execute(text(query), args) + session.commit() + return True + finally: + session.close() + + return bool(self._run("insert", _do)) def update(self, query: str, args: dict[str, Any]) -> bool: - try: + def _do(): session = self.Session() - session.execute(text(query), args) - session.commit() - session.close() - return True - except SQLAlchemyError as e: - log.error(f"Error updating data: {e}") - return False + try: + session.execute(text(query), args) + session.commit() + return True + finally: + session.close() + + return bool(self._run("update", _do)) def delete(self, query: str, args: dict[str, Any]) -> bool: - try: + def _do(): session = self.Session() - session.execute(text(query), args) - session.commit() - session.close() - return True - except SQLAlchemyError as e: - log.error(f"Error deleting data: {e}") - return False + try: + session.execute(text(query), args) + session.commit() + return True + finally: + session.close() - def get_last_update_date(self, series_name: str) -> str: - try: - session = self.Session() - result = ( - session.query(KomGrabber.last_update_date) - .filter_by(series_name=series_name) - .first() - ) - session.close() - return result[0] if result else "" - except SQLAlchemyError as e: - log.error(f"Error fetching last update date: {e}") - return "" + return bool(self._run("delete", _do)) - def fetch_one(self, query: str, args: dict[str, Any] = None): + def fetch_one(self, query: str, args: dict[str, Any] = None) -> Union[Tuple, None]: if args is None: args = {} - try: + + def _do(): session = self.Session() - result = session.execute(text(query), args).fetchone() - session.close() - return result - except SQLAlchemyError as e: - log.error(f"Error executing query: {e}") - return None + try: + result = session.execute(text(query), args) + if getattr(result, "returns_rows", False): + return result.fetchone() + return None + finally: + session.close() + + return self._run("fetch_one", _do) def fetch_all(self, query: str, args: dict[str, Any] | None = None): if args is None: args = {} - try: + + def _do(): session = self.Session() - result = session.execute(text(query), args).fetchall() - session.close() - return result - except SQLAlchemyError as e: - log.error(f"Error executing query: {e}") - return [] + try: + result = session.execute(text(query), args) + if getattr(result, "returns_rows", False): + return result.fetchall() + return [] + finally: + session.close() + + result = self._run("fetch_all", _do) + return result or [] if __name__ == "__main__": diff --git a/src/komcache/schemas/mariadb.py b/src/komcache/schemas/mariadb.py index 6008f9e..d1d2330 100644 --- a/src/komcache/schemas/mariadb.py +++ b/src/komcache/schemas/mariadb.py @@ -1,9 +1,13 @@ -CREATE_MARIADB_TABLES = [""" +CREATE_MARIADB_TABLES = [ + """ CREATE TABLE IF NOT EXISTS manga_requests ( id INT AUTO_INCREMENT PRIMARY KEY, manga_id INT, - grabbed TINYINT(1) DEFAULT 0 -);""","""CREATE TABLE IF NOT EXISTS komgrabber ( + grabbed TINYINT(1) DEFAULT 0, + image TEXT NOT NULL, + title TEXT NOT NULL +);""", + """CREATE TABLE IF NOT EXISTS komgrabber ( id INT AUTO_INCREMENT PRIMARY KEY, name TEXT NOT NULL, series_id TEXT NOT NULL, @@ -12,12 +16,18 @@ CREATE TABLE IF NOT EXISTS manga_requests ( updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, last_checked TIMESTAMP DEFAULT NULL, completed TINYINT(1) DEFAULT 0 -);""","""CREATE TABLE IF NOT EXISTS komtagger ( +);""", + """CREATE TABLE IF NOT EXISTS komtagger ( id INT AUTO_INCREMENT PRIMARY KEY, series_id TEXT NOT NULL, title TEXT NOT NULL, - updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - last_checked TIMESTAMP DEFAULT NULL, - status TEXT NOT NULL -); -"""] \ No newline at end of file + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + last_checked DATETIME NOT NULL DEFAULT '1970-01-01 00:00:00', + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + UNIQUE KEY unique_series_id (series_id(255)), + status TEXT NOT NULL, + tag_status TEXT NOT NULL DEFAULT 'untagged', + anilist_id INT DEFAULT NULL + ); +""", +]