Upgrade database migrations (#180)

.gitignore

@@ -178,4 +178,7 @@ inv.yml
 test-scripts/
 
 # Ruff cache
-.ruff_cache/
+.ruff_cache/
+
+# Old stuff
+old/

bot/migrations.py

@@ -0,0 +1,261 @@
+import asyncio
+import re
+import traceback
+from functools import wraps
+from pathlib import Path
+from types import TracebackType
+from typing import Optional, TypeVar
+
+import asyncpg
+import click
+from discord.utils import utcnow
+from libs.utils.config import CatherineConfig
+from typing_extensions import Self
+
+path = Path(__file__).parent / "config.yml"
+config = CatherineConfig(path)
+
+BE = TypeVar("BE", bound=BaseException)
+
+REVISION_FILE = re.compile(r"(?P<kind>V)(?P<version>\d+)__(?P<description>.+).sql")
+POSTGRES_URI = config["postgres"]["uri"]
+
+CREATE_MIGRATIONS_TABLE = """
+CREATE TABLE IF NOT EXISTS migrations (
+    id SERIAL PRIMARY KEY,
+    description TEXT,
+    created_at TIMESTAMP DEFAULT (NOW() AT TIME ZONE 'utc')
+);
+"""
+
+GET_LATEST_VERSION = """
+SELECT id FROM migrations
+ORDER BY id DESC
+LIMIT 1;
+"""
+
+INSERT_VERSION = """
+INSERT INTO migrations (id, description)
+VALUES ($1, $2);
+"""
+
+
+def coro(f):
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        return asyncio.run(f(*args, **kwargs))
+
+    return wrapper
+
+
+class Revision:
+    __slots__ = ("kind", "version", "description", "file")
+
+    def __init__(
+        self, *, kind: str, version: int, description: str, file: Path
+    ) -> None:
+        self.kind: str = kind
+        self.version: int = version
+        self.description: str = description
+        self.file: Path = file
+
+    @classmethod
+    def from_match(cls, match: re.Match[str], file: Path):
+        return cls(
+            kind=match.group("kind"),
+            version=int(match.group("version")),
+            description=match.group("description"),
+            file=file,
+        )
+
+
+class Migrations:
+    def __init__(self, *, no_conn: bool = False, migrations_path: str = "migrations"):
+        self.no_conn = no_conn
+        self.migrations_path = migrations_path
+        self.root: Path = Path(__file__).parent
+        self.revisions: dict[int, Revision] = self.get_revisions()
+        self.ensure_path()
+
+    async def __aenter__(self) -> Self:
+        if self.no_conn is False:
+            self.conn = await asyncpg.connect(POSTGRES_URI)
+            self.version = await self.get_latest_version()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BE]],
+        exc: Optional[BE],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        await self.conn.close()
+
+    async def get_latest_version(self):
+        record_version = await self.conn.fetchval(GET_LATEST_VERSION)
+        if record_version is None:
+            return 0
+        return record_version
+
+    def ensure_path(self) -> None:
+        migrations_path = self.root / self.migrations_path
+        migrations_path.mkdir(exist_ok=True)
+
+    def get_revisions(self) -> dict[int, Revision]:
+        result: dict[int, Revision] = {}
+        for file in self.root.glob("migrations/*.sql"):
+            match = REVISION_FILE.match(file.name)
+            if match is not None:
+                rev = Revision.from_match(match, file)
+                result[rev.version] = rev
+
+        return result
+
+    def is_next_revision_taken(self) -> bool:
+        return self.version + 1 in self.revisions
+
+    @property
+    def ordered_revisions(self) -> list[Revision]:
+        return sorted(self.revisions.values(), key=lambda r: r.version)
+
+    def create_revision(self, reason: str, *, kind: str = "V") -> Revision:
+        cleaned = re.sub(r"\s", "_", reason)
+        filename = f"{kind}{self.version + 1}__{cleaned}.sql"
+        path = self.root / self.migrations_path / filename
+
+        stub = (
+            f"-- Revision Version: V{self.version + 1}\n"
+            f"-- Revises: V{self.version}\n"
+            f"-- Creation Date: {utcnow()} UTC\n"
+            f"-- Reason: {reason}\n\n"
+        )
+
+        with open(path, "w", encoding="utf-8", newline="\n") as fp:
+            fp.write(stub)
+
+        return Revision(
+            kind=kind, description=reason, version=self.version + 1, file=path
+        )
+
+    async def upgrade(self) -> int:
+        ordered = self.ordered_revisions
+        successes = 0
+        async with self.conn.transaction():
+            for revision in ordered:
+                if revision.version > self.version:
+                    sql = revision.file.read_text("utf-8")
+                    await self.conn.execute(sql)
+                    await self.conn.execute(
+                        INSERT_VERSION, revision.version, revision.description
+                    )
+                    successes += 1
+
+        self.version += successes
+        return successes
+
+    def display(self) -> None:
+        ordered = self.ordered_revisions
+        for revision in ordered:
+            if revision.version > self.version:
+                sql = revision.file.read_text("utf-8")
+                click.echo(sql)
+
+
+async def create_migrations_table() -> None:
+    conn = await asyncpg.connect(POSTGRES_URI)
+    await conn.execute(CREATE_MIGRATIONS_TABLE)
+    await conn.close()
+
+
+@click.group(short_help="database migrations util", options_metavar="[options]")
+def main():
+    # grouped database commands
+    pass
+
+
+@main.command()
+@coro
+async def init():
+    """Initializes the database and runs all the current migrations"""
+    await create_migrations_table()
+    async with Migrations() as mg:
+        try:
+            applied = await mg.upgrade()
+            click.secho(
+                f"Successfully initialized and applied {applied} revisions(s)",
+                fg="green",
+            )
+        except Exception:
+            traceback.print_exc()
+            click.secho(
+                "failed to initialize and apply migrations due to error", fg="red"
+            )
+
+
+@main.command()
+@click.option("--reason", "-r", help="The reason for this revision.", required=True)
+@coro
+async def migrate(reason: str):
+    """Creates a new revision for you to edit"""
+    async with Migrations() as mg:
+        if mg.is_next_revision_taken():
+            click.echo(
+                "an unapplied migration already exists for the next version, exiting"
+            )
+            click.secho(
+                "hint: apply pending migrations with the `upgrade` command", bold=True
+            )
+            return
+        revision = mg.create_revision(reason)
+        click.echo(f"Created revision V{revision.version!r}")
+
+
+@main.command()
+@coro
+async def current():
+    """Shows the current version"""
+    async with Migrations() as mg:
+        click.echo(f"Version {mg.version}")
+
+
+@main.command()
+@click.option("--sql", help="Print the SQL instead of executing it", is_flag=True)
+@coro
+async def upgrade(sql):
+    """Upgrade to the latest version"""
+    async with Migrations() as mg:
+        if sql:
+            mg.display()
+            return
+
+        try:
+            applied = await mg.upgrade()
+            click.secho(
+                f"Applied {applied} revisions(s) (Current: V{mg.version})", fg="green"
+            )
+        except Exception:
+            traceback.print_exc()
+            click.secho("failed to apply migrations due to error", fg="red")
+
+
+@main.command()
+@click.option("--reverse", help="Print in reverse order (oldest first).", is_flag=True)
+@coro
+async def log(reverse):
+    """Displays the revision history"""
+    # We don't need to make an connection in ths case
+    migrations = Migrations(no_conn=True)
+
+    # Revisions is oldest first already
+    revs = (
+        reversed(migrations.ordered_revisions)
+        if not reverse
+        else migrations.ordered_revisions
+    )
+    for rev in revs:
+        as_yellow = click.style(f"V{rev.version:>03}", fg="yellow")
+        click.echo(f'{as_yellow} {rev.description.replace("_", " ")}')
+
+
+if __name__ == "__main__":
+    main()

bot/migrations/V1__initial_migrations.sql

@@ -0,0 +1,63 @@
+-- Revision Version: V1
+-- Revises: V0
+-- Creation Date: 2024-05-20 19:06:40.005347 UTC
+-- Reason: initial_migrations
+
+-- These migrations are intended to continue from the previous migrations
+CREATE TABLE IF NOT EXISTS catherine_users (
+    id BIGINT PRIMARY KEY,
+    created_at timestamp WITHOUT TIME ZONE DEFAULT (NOW() AT TIME ZONE 'utc')
+);
+
+CREATE TABLE IF NOT EXISTS profiles (
+    id SERIAL PRIMARY KEY,
+    views INT DEFAULT 0,
+    name VARCHAR(50),
+    pronouns VARCHAR(50),
+    gender_identity VARCHAR(50),
+    sexual_orientation VARCHAR(50),
+    romantic_orientation VARCHAR(50),
+    user_id BIGINT REFERENCES catherine_users (id) ON DELETE CASCADE ON UPDATE NO ACTION
+);
+
+CREATE INDEX IF NOT EXISTS profiles_name_idx ON profiles (name);
+CREATE INDEX IF NOT EXISTS profiles_name_trgm_idx ON profiles USING GIN (name gin_trgm_ops);
+CREATE INDEX IF NOT EXISTS profiles_name_lower_idx ON profiles (LOWER(name));
+CREATE UNIQUE INDEX IF NOT EXISTS profiles_user_idx ON profiles (user_id);
+
+
+CREATE TABLE IF NOT EXISTS pronouns_examples (
+    id SERIAL PRIMARY KEY,
+    sentence TEXT,
+    created_at timestamp WITHOUT TIME ZONE DEFAULT (NOW() AT TIME ZONE 'utc'),
+    approved BOOLEAN DEFAULT FALSE,
+    user_id BIGINT REFERENCES catherine_users (id) ON DELETE NO ACTION ON UPDATE NO ACTION
+);
+
+CREATE INDEX IF NOT EXISTS pronouns_examples_user_idx ON pronouns_examples (user_id);
+
+-- Keeping the tonetags for consistency with previous migrations
+CREATE TABLE IF NOT EXISTS tonetags (
+    id SERIAL PRIMARY KEY,
+    indicator VARCHAR(255),
+    definition TEXT,
+    created_at TIMESTAMP WITHOUT TIME ZONE DEFAULT (NOW() AT TIME ZONE 'utc'),
+    author_id BIGINT REFERENCES catherine_users (id) ON DELETE CASCADE ON UPDATE NO ACTION
+);
+
+CREATE TABLE IF NOT EXISTS tonetags_lookup (
+    id SERIAL PRIMARY KEY,
+    indicator VARCHAR (255),
+    author_id BIGINT,
+    tonetags_id INTEGER REFERENCES tonetags (id) ON DELETE CASCADE ON UPDATE NO ACTION
+);
+
+CREATE INDEX IF NOT EXISTS tonetags_indicator_idx ON tonetags (indicator);
+CREATE INDEX IF NOT EXISTS tonetags_indicator_trgm_idx ON tonetags USING GIN (indicator gin_trgm_ops);
+CREATE INDEX IF NOT EXISTS tonetags_indicator_lower_idx ON tonetags (LOWER(indicator));
+CREATE UNIQUE INDEX IF NOT EXISTS tonetags_uniq_idx ON tonetags (LOWER(indicator), author_id);
+
+CREATE INDEX IF NOT EXISTS tonetags_lookup_indicator_idx ON tonetags_lookup (indicator);
+CREATE INDEX IF NOT EXISTS tonetags_lookup_indicator_trgm_idx ON tonetags_lookup USING GIN (indicator gin_trgm_ops);
+CREATE INDEX IF NOT EXISTS tonetags_lookup_indicator_lower_idx ON tonetags_lookup (LOWER(indicator));
+CREATE UNIQUE INDEX IF NOT EXISTS tonetags_lookup_uniq_idx ON tonetags_lookup (LOWER(indicator), author_id);

bot/migrations/V2__pride_profiles.sql

@@ -0,0 +1,26 @@
+-- Revision Version: V2
+-- Revises: V1
+-- Creation Date: 2024-06-16 17:49:08.154360+00:00 UTC
+-- Reason: pride_profiles
+
+-- This is an replacement table for the old profiles table
+-- Also we need to reduce the dependency on the catherine_users table
+CREATE TABLE IF NOT EXISTS pride_profiles (
+    id SERIAL PRIMARY KEY,
+    user_id BIGINT,
+    name TEXT,
+    pronouns TEXT,
+    gender_identity TEXT,
+    sexual_orientation TEXT,
+    romantic_orientation TEXT,
+    views INT DEFAULT 0,
+    UNIQUE (user_id, name)
+);
+
+CREATE INDEX IF NOT EXISTS pride_profiles_name_idx ON pride_profiles (name);
+CREATE INDEX IF NOT EXISTS pride_profiles_name_trgm_idx ON pride_profiles USING GIN (name gin_trgm_ops);
+CREATE INDEX IF NOT EXISTS pride_profiles_name_lower_idx ON pride_profiles (LOWER(name));
+CREATE INDEX IF NOT EXISTS pride_profiles_user_idx ON pride_profiles (user_id);
+
+-- There isn't that much data in production so this is justify (and has been tested to work)
+INSERT INTO pride_profiles (SELECT id, user_id, name, pronouns, gender_identity, sexual_orientation, romantic_orientation, views FROM profiles);

config-example.yml

@@ -10,5 +10,4 @@ bot:
 
 # Postgres Configuration
 postgres:
-  revision: "rev5"
   uri: "postgres://username:[email protected]:5432/postgres"

docker/Dockerfile

@@ -15,8 +15,6 @@ RUN apt-get update \
 
 WORKDIR /catherine
 COPY /bot/ /catherine/bot
-COPY /migrations /catherine/migrations
-COPY /migrations-runner.py /catherine/migrations-runner.py
 COPY /requirements.txt /catherine/requirements.txt
 COPY /docker/start.sh /catherine/start.sh
 COPY /docker/wait-for /catherine/wait-for

docker/start.sh

@@ -5,7 +5,7 @@ CATHERINE_FIRST_START_CHECK="CATHERINE_FIRST_START"
 if [ ! -f $CATHERINE_FIRST_START_CHECK ]; then
     touch $CATHERINE_FIRST_START_CHECK
     echo 'DO NOT EDIT THIS FILE! THIS IS USED WHEN YOU FIRST RUN CATHERINE USING DOCKER!' >> $CATHERINE_FIRST_START_CHECK
-    python3 /catherine/migrations-runner.py
+    python3 /catherine/bot/migrations.py init
 fi
 
 exec python3 /catherine/bot/catherinebot.py

docs/source/guides/dev/setup.rst

@@ -57,7 +57,7 @@ Local Setup
 
     .. code-block:: bash
 
-        python3 migrations-runner.py
+        python3 bot/migrations.py
 
 Special Configuration Variables
 ---------------------