From 224682a8d55f401b6b16861d345384eb022ec31a Mon Sep 17 00:00:00 2001 From: robertlestak Date: Mon, 22 Jul 2024 22:18:23 -0700 Subject: [PATCH] add database interface to enable using external database as opposed to embedded tinydb --- README.md | 1 + backend/database.py | 116 ++++++++++++++++++++++++++++++++++++--- backend/requirements.txt | 2 + 3 files changed, 112 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 8665672..20490ef 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,7 @@ You will need to set the correct parameters for your setup: | `-e PAGE_SIZE=50` | (**optional**) To avoid plex timeouts, results are loaded in pages (or chunks). If you recieve Plex Timeout errors, try setting this parameter to a lower value. | | `-e DEBUG=0` | (**optional**) To enable debug logging set `DEBUG` to `1` | | `-e PLEX_TIMEOUT=7200` | (**optional**) modify the timeout for wrapper (Error : Failed to load content!) | +| `-e DATABASE_URL=postgresql://user:pass@database:5432/cleanarr` | (**optional**) To use a PostgreSQL database instead of the default TinyDB database. | #### Example running directly with docker (with make) diff --git a/backend/database.py b/backend/database.py index 1146356..de7fdfc 100644 --- a/backend/database.py +++ b/backend/database.py @@ -3,27 +3,55 @@ from tinydb import TinyDB, where from tinydb.table import Document +from sqlalchemy import create_engine, Column, String, Integer, BigInteger +from sqlalchemy.orm import sessionmaker, scoped_session +from sqlalchemy.ext.declarative import declarative_base from logger import get_logger DELETED_SIZE_DOC_ID = 1 IGNORED_ITEMS_TABLE = 'ignored' +logger = get_logger(__name__) +Base = declarative_base() -logger = get_logger(__name__) +class DeletedSize(Base): + __tablename__ = '_default' # named as such for backwards compatibility + id = Column(Integer, primary_key=True, default=DELETED_SIZE_DOC_ID) + library_name = Column(String, unique=True) + deleted_size = Column(BigInteger) + +class IgnoredItem(Base): + __tablename__ = IGNORED_ITEMS_TABLE + key = Column(String, primary_key=True) +class DatabaseInterface: + def set_deleted_size(self, library_name, deleted_size): + raise NotImplementedError + + def get_deleted_size(self, library_name): + raise NotImplementedError -class Database(object): + def get_ignored_item(self, content_key): + raise NotImplementedError + + def add_ignored_item(self, content_key): + raise NotImplementedError + + def remove_ignored_item(self, content_key): + raise NotImplementedError + +class TinyDBDatabase(DatabaseInterface): def __init__(self): - logger.debug("DB Init") - config_dir = os.environ.get("CONFIG_DIR", "") # Will be set by Dockerfile + logger.debug("TinyDB Init") + config_dir = os.environ.get("CONFIG_DIR", "") self.local = threading.local() - logger.debug("DB Init Success") + self.config_dir = config_dir + logger.debug("TinyDB Init Success") def get_db(self): if not hasattr(self.local, 'db'): - config_dir = os.environ.get("CONFIG_DIR", "") - self.local.db = TinyDB(os.path.join(config_dir, 'db.json')) + self.local.db = TinyDB(os.path.join(self.config_dir, 'db.json')) return self.local.db def set_deleted_size(self, library_name, deleted_size): @@ -57,3 +85,77 @@ def remove_ignored_item(self, content_key): logger.debug("content_key %s", content_key) table = self.get_db().table(IGNORED_ITEMS_TABLE) table.remove(where('key') == content_key) + +class SQLAlchemyDatabase(DatabaseInterface): + def __init__(self): + logger.debug("SQLAlchemy Init") + config_dir = os.environ.get("CONFIG_DIR", "") + db_url = os.environ.get("DATABASE_URL", f"sqlite:///{os.path.join(config_dir, 'db.sqlite')}") + self.engine = create_engine(db_url) + self.session_factory = sessionmaker(bind=self.engine) + self.Session = scoped_session(self.session_factory) + Base.metadata.create_all(self.engine) + logger.debug("SQLAlchemy Init Success") + + def set_deleted_size(self, library_name, deleted_size): + logger.debug("library_name %s, deleted_size %s", library_name, deleted_size) + session = self.Session() + obj = session.query(DeletedSize).filter_by(library_name=library_name).first() + if obj: + obj.deleted_size = deleted_size + else: + obj = DeletedSize(library_name=library_name, deleted_size=deleted_size) + session.add(obj) + session.commit() + session.close() + + def get_deleted_size(self, library_name): + logger.debug("library_name %s", library_name) + session = self.Session() + obj = session.query(DeletedSize).filter_by(library_name=library_name).first() + session.close() + return obj.deleted_size if obj else 0 + + def get_ignored_item(self, content_key): + logger.debug("content_key %s", content_key) + session = self.Session() + obj = session.query(IgnoredItem).get(content_key) + session.close() + return obj + + def add_ignored_item(self, content_key): + logger.debug("content_key %s", content_key) + session = self.Session() + obj = IgnoredItem(key=content_key) + session.add(obj) + session.commit() + session.close() + + def remove_ignored_item(self, content_key): + logger.debug("content_key %s", content_key) + session = self.Session() + session.query(IgnoredItem).filter_by(key=content_key).delete() + session.commit() + session.close() + +class Database(DatabaseInterface): + def __init__(self): + if os.environ.get("DATABASE_URL"): + self.db = SQLAlchemyDatabase() + else: + self.db = TinyDBDatabase() + + def set_deleted_size(self, library_name, deleted_size): + self.db.set_deleted_size(library_name, deleted_size) + + def get_deleted_size(self, library_name): + return self.db.get_deleted_size(library_name) + + def get_ignored_item(self, content_key): + return self.db.get_ignored_item(content_key) + + def add_ignored_item(self, content_key): + self.db.add_ignored_item(content_key) + + def remove_ignored_item(self, content_key): + self.db.remove_ignored_item(content_key) diff --git a/backend/requirements.txt b/backend/requirements.txt index d8033e6..cce5b3f 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -17,3 +17,5 @@ websocket-client==0.57.0 Werkzeug==3.0.1 python-dotenv==1.0.0 pytest-benchmark==4.0.0 +sqlalchemy==2.0.30 +psycopg2==2.9.9 \ No newline at end of file