AO3-DL/python/sql.py

# 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)
)"""
    )
    commandsRaw.append(
        """
    CREATE TABLE IF NOT EXISTS inspirations(
    id INTEGER,
    parent_id INTEGER,
    translation INTEGER,
    PRIMARY KEY (id, parent_id)
    )"""
    )
    commandsRaw.append(
        """
    CREATE TABLE IF NOT EXISTS blacklist(
    id INTEGER PRIMARY KEY,
    timestamp_added INTEGER,
    reason VARCHAR(1000)
    )"""
    )
    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))

    existingParents = set(
        execute(
            commands=[
                ("SELECT * FROM inspirations WHERE id = ? OR parent_id = ?", (id, id))
            ],
            config=config,
            logger=logger,
        )[0]
    )
    desiredParents = set(info["parents"])
    parentsToAdd = list(desiredParents.difference(existingParents))
    parentsToDelete = list(existingParents.difference(desiredParents))
    cancelChange = []
    for i in parentsToAdd:
        if i[0] == 0:
            for j in parentsToDelete:
                if j[1] == i[1]:
                    cancelChange.append((i, j))
        if i[1] == 0:
            for j in parentsToDelete:
                if j[0] == i[0]:
                    cancelChange.append((i, j))
    for i in cancelChange:
        parentsToDelete.remove(i[1])
        parentsToAdd.remove(i[0])

    for i in parentsToAdd:
        commands.append(("INSERT INTO inspirations VALUES (?, ?, ?)", tuple(i)))
    for i in parentsToDelete:
        commands.append(
            (
                "DELETE FROM inspirations WHERE id = ? AND parent_id = ? AND translation = ?",
                tuple(i),
            )
        )

    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 getExcluded() -> set:
    output = set(())
    for i in execute(
        commands=[
            ("SELECT id FROM blacklist", ()),
        ],
        config=common.config,
        logger=common.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]
    )