summaryrefslogtreecommitdiff
path: root/personal_agent/database/database.py
blob: 7119c7ea8df946c8834e32bf292f93aef6d69c12 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import sqlite3
from datetime import UTC, datetime
from pathlib import Path


class Database:
    def __init__(self, path: Path) -> None:
        self.path = path
        self.migrations_dir = Path(__file__).parent / "migrations"

    def connect(self) -> sqlite3.Connection:
        return sqlite3.connect(self.path)

    def migrate(self) -> None:
        with self.connect() as connection:
            connection.execute(
                """
                CREATE TABLE IF NOT EXISTS schema_migrations (
                    version INTEGER PRIMARY KEY,
                    applied_at TEXT NOT NULL
                )
                """
            )

            applied_versions = {
                row[0]
                for row in connection.execute("SELECT version FROM schema_migrations")
            }

            for migration_version, path in self._get_migrations():
                if migration_version in applied_versions:
                    continue

                sql = path.read_text(encoding="utf-8")

                connection.execute("BEGIN")

                try:
                    connection.execute(sql)

                    applied_at = datetime.now(UTC).isoformat().replace("+00:00", "Z")

                    connection.execute(
                        """
                        INSERT INTO schema_migrations (version, applied_at)
                        VALUES (?, ?)
                        """,
                        (migration_version, applied_at),
                    )

                    connection.commit()
                except Exception:
                    connection.rollback()
                    raise

    def _get_migrations(self) -> list[tuple[int, Path]]:
        migrations = []

        for path in self.migrations_dir.glob("*.sql"):
            version_text, _, _ = path.stem.partition("_")

            if not version_text.isdigit():
                continue

            migrations.append((int(version_text), path))

        migrations.sort()

        versions = [version for version, _ in migrations]
        if len(versions) != len(set(versions)):
            raise ValueError("Duplicate migration version")

        return migrations