# This file is part of AOMO, "Archive Of My Own", a collection of Python and PHP # scripts designed act as an improved local backup system # for works published on https://archiveofourown.org # # Copyright (c) 2025 Cyberpro123, except where otherwise noted. # # This project is available under the GNU General Public License (GPL) v2.0, # available at https://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html # or in the 'LICENSE.md' file that should be distributed alongside this one. # # Report issues to https://codeberg.org/Cyberpro123/AOMO # Contact author at cyberpro123@posteo.com import datetime import logging import common class UniqueConstraintFailed(Exception): def __init__(self, message): super().__init__(message) self.message = message def execute( commands: list, config: dict, logger: logging.Logger, ) -> list: config["sqlType"] config["sqlLocation"] output = [] match config["sqlType"]: case "sqlite": import contextlib import os.path import random import sqlite3 import time dbPath = os.path.join(config["dirWebUi"], config["sqlLocation"]) with contextlib.closing(sqlite3.connect(dbPath, timeout=10)) as con: cur = con.cursor() for command in commands: logger.debug(command[0]) logger.debug(str(command[1])[:100]) loopNo = 1 while loopNo <= common.loopRetries: try: cur.execute(command[0], command[1]) except sqlite3.OperationalError as ex: if "database is locked" in str(ex): random.seed() pauseMult = 0.5 + random.random() pauseTime = loopNo * 5 * pauseMult logger.log( (10 + 20 * int(2 > loopNo)), common.loopErrorTemplate.format( "Doing database operation", pauseTime, type(ex).__name__, ex.args, ), ) loopNo += 1 time.sleep(pauseTime) else: raise ex except sqlite3.IntegrityError as ex: if "UNIQUE constraint failed" in str(ex): raise UniqueConstraintFailed(str(ex)) else: raise ex else: loopNo += common.loopRetries output.append(cur.fetchall()) con.commit() case "postgres": import psycopg with psycopg.connect( f"host={config['sqlLocation']} user={config['sqlUsername']} password={config['sqlPassword']} dbname='AOMO'" ) as con: with con.cursor() as cur: for command in commands: logger.debug(command[0]) logger.debug(str(command[1])[:100]) cur.execute(str(command[0]).replace("?", "%s"), command[1]) try: output.append(cur.fetchall()) except psycopg.ProgrammingError: output.append([]) # psycopg is a picky little piece of shit that will throw an error # if it doesn't have a result for you (when it should just return an empty fucking list) con.commit() case "mysql": import MySQLdb if "localhost" in config["sqlLocation"]: config["sqlLocation"] = config["sqlLocation"].replace( "localhost", "127.0.0.1" ) con = MySQLdb.connect( host=config["sqlLocation"], user=config["sqlUsername"], password=config["sqlPassword"], database="AOMO", connect_timeout=15, ) with con.cursor() as cur: for command in commands: logger.debug(command[0]) logger.debug(str(command[1])[:100]) cur.execute( command[0] .replace("?", "%s") .replace( "ON CONFLICT DO NOTHING", "ON DUPLICATE KEY UPDATE id = id" ), command[1], ) logger.debug("Executed") output.append(cur.fetchall()) con.commit() case _: exceptionStr = ( f"Invalid sqlType given to db.execute() [{config['sqlType']}]" ) logger.critical(exceptionStr) raise common.InvalidConfigurationError(exceptionStr) return output def init( config: dict, logger: logging.Logger, ) -> None: commandsRaw = [] commandsRaw.append( """ CREATE TABLE IF NOT EXISTS works( id INTEGER PRIMARY KEY, chapters_count INTEGER, chapters_expected INTEGER, word_count_ao3 INTEGER, date_last_downloaded INTEGER, date_last_edited INTEGER, date_last_updated INTEGER, date_published INTEGER, date_processed INTEGER);""" ) commandsRaw.append( """ CREATE TABLE IF NOT EXISTS tags( id INTEGER, tag VARCHAR(150), category INTEGER, sequence INTEGER)""" ) # values for tags table category column: # 1: Fandom # 2: Ships # 3: Characters # 4: Freeform # 5: Languages # 6: Ratings # 7: Warnings # 8: Categories # 9: Authors commandsRaw.append( """ CREATE TABLE IF NOT EXISTS series_meta( id INTEGER PRIMARY KEY, title VARCHAR(255), date_begun INTEGER, date_updated INTEGER, date_checked INTEGER, description VARCHAR(1250), notes VARCHAR(5000), complete INTEGER )""" ) commandsRaw.append( """ CREATE TABLE IF NOT EXISTS series_works( series_id INTEGER, id INTEGER, pos INTEGER )""" ) commandsRaw.append( """ CREATE TABLE IF NOT EXISTS text( id INTEGER, chapter INTEGER, title VARCHAR(255), summary VARCHAR(1250), notes_start VARCHAR(5000), notes_end VARCHAR(5000), body TEXT, PRIMARY KEY (id, chapter) ); """ ) commandsRaw.append( """ CREATE TABLE IF NOT EXISTS users( id INTEGER PRIMARY KEY, username VARCHAR(100) UNIQUE, password_hash VARCHAR(255), site_skin INTEGER, can_see_explicit_works INTEGER )""" ) commandsRaw.append( """ CREATE TABLE IF NOT EXISTS login_tokens( token VARCHAR(1024) PRIMARY KEY, id INTEGER, expire_timestamp INTEGER )""" ) commandsRaw.append( """ CREATE TABLE IF NOT EXISTS progress_tracker( id INTEGER, user_id INTEGER, chapter INTEGER, timestamp INTEGER, PRIMARY KEY (id, user_id) )""" ) commandsRaw.append( """ CREATE TABLE IF NOT EXISTS site_skins( id INTEGER PRIMARY KEY, body TEXT, name VARCHAR(100), preview_img_link VARCHAR(300) )""" ) commands = [] for command in commandsRaw: commands.append((command, ())) commands.append( ( "INSERT INTO site_skins VALUES(?, ?, ?, ?) ON CONFLICT DO NOTHING", (0, "", "None", ""), ) ) execute( commands=commands, config=config, logger=logger, ) logger.info("Initialized Database") def addWork( id: int, info: dict, config: dict, logger: logging.Logger, ) -> None: try: for i in ( ("nchapters", int), ("chaptersExpected", int), ("dateDl", int), ("dateEd", int), ("dateUp", int), ("datePb", int), ("rating", str), ("warnings", list), ("categories", list), ("authors", list), ("tagsFandom", list), ("tagsShips", list), ("tagsChara", list), ("tagsOther", list), ("serieses", list), ("text", dict), ): if not isinstance(info[i[0]], i[1]): raise KeyError except KeyError: exceptionStr = "Invalid info dict given to addWork() in db.py" logger.critical(exceptionStr) raise KeyError(exceptionStr) commands = [] commands.append(("INSERT INTO works(id) VALUES(?) ON CONFLICT DO NOTHING", (id,))) commands.append( ( "UPDATE works SET chapters_count = ?, chapters_expected = ?, word_count_ao3 = ?, date_last_downloaded = ?," + " date_last_edited = ?, date_last_updated = ?, date_published = ?, date_processed = ? WHERE id = ?", ( info["nchapters"], info["chaptersExpected"], int(str(info["wordCount"]).replace(",", "")), info["dateDl"], info["dateEd"], info["dateUp"], info["datePb"], int( datetime.datetime.now() .astimezone() .replace(microsecond=0) .timestamp() ), id, ), ) ) def listToDesiredTagSets( enumerable: list, catNo: int, ) -> set: outList = set(()) for num, i in enumerate(enumerable): outList.add((i, catNo, (num + 1))) return outList desiredTags = set(()) desiredTags.add((info["rating"], 6, 1)) desiredTags.add((info["language"], 5, 1)) desiredTags.update(listToDesiredTagSets(info["warnings"], 7)) desiredTags.update(listToDesiredTagSets(info["categories"], 8)) desiredTags.update(listToDesiredTagSets(info["tagsFandom"], 1)) desiredTags.update(listToDesiredTagSets(info["tagsShips"], 2)) desiredTags.update(listToDesiredTagSets(info["tagsChara"], 3)) desiredTags.update(listToDesiredTagSets(info["tagsOther"], 4)) desiredTags.update(listToDesiredTagSets(info["authors"], 9)) existingTags = getTags(id, config, logger) tagsToInsert = desiredTags.difference(existingTags) tagsToDelete = existingTags.difference(desiredTags) for i in tagsToInsert: commands.append( ("INSERT INTO tags VALUES (?, ?, ?, ?)", (id, i[0], i[1], i[2])) ) for i in tagsToDelete: commands.append( ( "DELETE FROM tags WHERE id = ? AND tag = ? AND category = ? AND sequence = ?", (id, i[0], i[1], i[2]), ) ) # series[0] = series ID # series[1] = work's position in that series # series[2] = series title for series in info["serieses"]: commands.extend( ( ( "INSERT INTO series_meta(id) VALUES (?) ON CONFLICT DO NOTHING", (series[0],), ), ( "UPDATE series_meta SET title = ? WHERE id = ?", (series[2], series[0]), ), ) ) existingPos = execute( commands=[ ( "SELECT pos FROM series_works WHERE series_id = ? AND id = ?", (series[0], id), ) ], config=config, logger=logger, )[0] if existingPos: if existingPos[0][0] != series[1]: commands.append( ( "UPDATE series_works SET pos = ? WHERE series_id = ? AND id = ?", (series[1], series[0], id), ) ) else: commands.append( ( "INSERT INTO series_works VALUES (?, ?, ?)", (series[0], id, series[1]), ) ) for chaptNum in info["text"]: sqlStrTextInit = ( "INSERT INTO text(id, chapter) VALUES (?, ?) ON CONFLICT DO NOTHING" ) sqlStrTextUpdt = "UPDATE text SET title = ?, summary = ?, notes_start = ?, notes_end = ?, body = ? WHERE id = ? AND chapter = ?" for i in info["text"][chaptNum]: info["text"][chaptNum][i].replace( common.imagePrefix, common.imagePrefixSubstitute ) sqlTplTextUpdt = ( info["text"][chaptNum]["title"], info["text"][chaptNum]["summary"], info["text"][chaptNum]["notesStart"], info["text"][chaptNum]["notesEnd"], info["text"][chaptNum]["body"], id, chaptNum, ) commands.append((sqlStrTextInit, (id, chaptNum))) commands.append((sqlStrTextUpdt, sqlTplTextUpdt)) execute( commands=commands, config=config, logger=logger, ) def getWorkIDs( config: dict, logger: common.logging.Logger, ) -> set: output = set(()) for i in execute( commands=[ ("SELECT id FROM works", ()), ], config=config, logger=logger, )[0]: output.add(i[0]) return output def getCompleteWorkIDs( config: dict, logger: common.logging.Logger, ) -> set: output = set(()) for i in execute( commands=[ ("SELECT id FROM works WHERE chapters_count = chapters_expected", ()), ], config=config, logger=logger, )[0]: output.add(i[0]) return output def getSeriesIDs( config: dict, logger: common.logging.Logger, ) -> dict: output = {} for i in execute( commands=[("SELECT id, title FROM series_meta", ())], config=config, logger=logger, )[0]: output[i[0]] = bool(i[1]) return output def getWorksInSeriesWithPos( seriesID: int, config: dict, logger: common.logging.Logger, ) -> list: output = [] sqlStr = "SELECT id, pos FROM series_works WHERE series_id = ? ORDER BY pos" for i in execute( commands=[(sqlStr, (seriesID,))], config=config, logger=logger, )[0]: output.append(i) return output def getWorksInSeriesWithoutPos( seriesID: int, config: dict, logger: common.logging.Logger, ) -> list: output = [] sqlStr = "SELECT id FROM series_works WHERE series_id = ? ORDER BY pos" for i in execute( commands=[(sqlStr, (seriesID,))], config=config, logger=logger, )[0]: output.append(i[0]) return output def getAllEditTimePerWorkID( config: dict, logger: common.logging.Logger, ) -> dict: output = {} for i in execute( commands=[("SELECT id, date_last_edited FROM works", ())], config=config, logger=logger, )[0]: output[i[0]] = i[1] return output def getEditTimeFromWorkID( id: int, config: dict, logger: common.logging.Logger, ) -> int: output = execute( commands=[("SELECT date_last_edited FROM works WHERE id = ?;", (id,))], config=config, logger=logger, ) return int(output[0][0][0]) def getTags( id: int, config: dict, logger: common.logging.Logger, ) -> set: return set( execute( commands=[ ("SELECT tag, category, sequence FROM tags WHERE id = ?;", (id,)) ], config=config, logger=logger, )[0] )