Spaces:
Running on Zero
Running on Zero
File size: 6,465 Bytes
8031356 1fcc069 8031356 1fcc069 8031356 1fcc069 8031356 1fcc069 8031356 1fcc069 8031356 1fcc069 8031356 1fcc069 8031356 1fcc069 8031356 1fcc069 | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | """
database/migrations.py
Versioned database migrations.
Migration files live in database/migrations/ and must be named NNN_description.sql
(0001_initial.sql, 0002_add_water_tracking.sql, ...). They are applied in version
order exactly once, and every applied version is recorded in the
schema_migrations table so the process is safe to run on every boot.
Unlike the previous "run schema.sql with CREATE TABLE IF NOT EXISTS" approach,
this supports real schema evolution: ALTER TABLE, new constraints, backfills,
data fixes, etc.
"""
import hashlib
import logging
import re
from pathlib import Path
from database.connection import get_db
logger = logging.getLogger(__name__)
MIGRATIONS_DIR = Path(__file__).parent / "migrations"
MIGRATION_TIMEOUT_SECONDS = 300
def split_sql_statements(sql: str) -> list[str]:
"""Split a migration file into individual statements.
Understands single-quoted string literals (with '' escapes),
double-quoted identifiers, line comments (--) and dollar-quoted bodies
($tag$ ... $tag$) so future migrations can contain PL/pgSQL functions.
"""
statements: list[str] = []
buffer: list[str] = []
index = 0
length = len(sql)
dollar_tag: str | None = None
while index < length:
char = sql[index]
# Inside a dollar-quoted body, copy verbatim until the closing tag.
if dollar_tag is not None:
closing = sql.find(dollar_tag, index)
if closing == -1:
buffer.append(sql[index:])
index = length
continue
buffer.append(sql[index:closing + len(dollar_tag)])
index = closing + len(dollar_tag)
dollar_tag = None
continue
# Comments: skip to end of line.
if char == "-" and index + 1 < length and sql[index + 1] == "-":
newline = sql.find("\n", index)
index = length if newline == -1 else newline
continue
# Single-quoted literal.
if char == "'":
buffer.append(char)
index += 1
while index < length:
buffer.append(sql[index])
if sql[index] == "'":
if index + 1 < length and sql[index + 1] == "'":
buffer.append(sql[index + 1])
index += 2
continue
index += 1
break
index += 1
continue
# Double-quoted identifier.
if char == '"':
buffer.append(char)
index += 1
while index < length:
buffer.append(sql[index])
if sql[index] == '"':
if index + 1 < length and sql[index + 1] == '"':
buffer.append(sql[index + 1])
index += 2
continue
index += 1
break
index += 1
continue
# Start of a dollar-quoted body.
if char == "$":
match = re.match(r"\$[A-Za-z_0-9]*\$", sql[index:])
if match:
dollar_tag = match.group(0)
buffer.append(dollar_tag)
index += len(dollar_tag)
continue
# Statement terminator.
if char == ";":
statement = "".join(buffer).strip()
if statement:
statements.append(statement)
buffer = []
index += 1
continue
buffer.append(char)
index += 1
tail = "".join(buffer).strip()
if tail:
statements.append(tail)
return statements
def _available_migrations() -> list[tuple[int, str, str, Path]]:
"""Return (version, name, checksum, path) for every migration on disk."""
migrations = []
for path in sorted(MIGRATIONS_DIR.glob("*.sql")):
match = re.match(r"^(\d+)_(.+)\.sql$", path.name)
if not match:
logger.warning("Skipping migration with invalid filename: %s", path.name)
continue
version = int(match.group(1))
statement = match.group(2)
checksum = hashlib.sha256(path.read_bytes()).hexdigest()
migrations.append((version, statement, checksum, path))
return migrations
async def run_migrations() -> list[str]:
"""Apply any pending migrations. Returns the names applied this run."""
db = await get_db()
await db.execute(
"""
CREATE TABLE IF NOT EXISTS schema_migrations (
version INTEGER PRIMARY KEY,
name TEXT NOT NULL,
checksum TEXT NOT NULL,
applied_at TIMESTAMPTZ NOT NULL DEFAULT now()
)
"""
)
applied = {
row["version"]: row["checksum"]
for row in await db.fetchall("SELECT version, checksum FROM schema_migrations")
}
applied_now: list[str] = []
for version, name, checksum, path in _available_migrations():
previous = applied.get(version)
if previous is not None:
if previous != checksum:
# Refuse to silently skip a changed migration — that is how you
# end up with schemas that differ per environment.
raise RuntimeError(
f"Migration {path.name} was modified after it was applied "
f"(expected checksum {previous[:12]}, found {checksum[:12]}). "
"Add a new migration instead of editing an applied one."
)
continue
logger.info("Applying migration %s ...", path.name)
statements = split_sql_statements(path.read_text(encoding="utf-8"))
async with db.transaction() as tx:
for statement in statements:
await tx.execute(statement)
await tx.execute(
"""
INSERT INTO schema_migrations (version, name, checksum)
VALUES ($1, $2, $3)
""",
(version, name, checksum),
)
applied_now.append(path.name)
logger.info("Applied migration %s (%d statements).", path.name, len(statements))
if not applied_now:
logger.info("Database schema is up to date.")
else:
logger.info("Applied %d migration(s): %s", len(applied_now), ", ".join(applied_now))
return applied_now
|