switch to db migrator
This commit is contained in:
parent
8b17f133d6
commit
0b09a44cfb
SSH key fingerprint:
SHA256:oAMyvotlNFraMmZmr+p6AxnNfW/GioTs1pOn3V4tQ7A
3 changed files with 330 additions and 29 deletions
72
nomen/db.py
72
nomen/db.py
|
|
@ -1,43 +1,77 @@
|
|||
import logging
|
||||
import sqlite3
|
||||
|
||||
import aiosqlite
|
||||
|
||||
from .migrator import dumb_migrate_db
|
||||
from .utils import contains
|
||||
|
||||
log = logging.getLogger("nomen.db")
|
||||
log.setLevel(logging.INFO)
|
||||
|
||||
async def setup_db(db_file):
|
||||
db = await aiosqlite.connect(db_file)
|
||||
schema = """
|
||||
PRAGMA user_version = 1;
|
||||
|
||||
await db.executescript("""
|
||||
BEGIN;
|
||||
PRAGMA main.synchronous = NORMAL;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS keywords (
|
||||
CREATE TABLE keywords (
|
||||
guild_id INTEGER NOT NULL,
|
||||
keyword TEXT NOT NULL,
|
||||
user_id INTEGER NOT NULL,
|
||||
regex INTEGER NOT NULL DEFAULT 0 CHECK(regex IN (0, 1)),
|
||||
count INTEGER NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (guild_id, keyword, user_id)
|
||||
count INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
|
||||
CREATE TABLE guilds (
|
||||
guild_id INTEGER NOT NULL PRIMARY KEY,
|
||||
prefix TEXT NOT NULL DEFAULT ">"
|
||||
)
|
||||
WITHOUT ROWID;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS guilds (
|
||||
guild_id INTEGER NOT NULL PRIMARY KEY,
|
||||
prefix TEXT NOT NULL DEFAULT ">"
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
CREATE TABLE users (
|
||||
user_id INTEGER NOT NULL PRIMARY KEY,
|
||||
disabled INTEGER NOT NULL DEFAULT 0 CHECK(disabled IN (0, 1))
|
||||
disabled INTEGER NOT NULL DEFAULT 0 CHECK(disabled IN (0, 1)),
|
||||
use_embeds INTEGER NOT NULL DEFAULT 1 CHECK(use_embeds IN (0, 1)),
|
||||
notify_self INTEGER NOT NULL DEFAULT 0 CHECK(notify_self IN (0, 1)),
|
||||
bots_notify INTEGER NOT NULL DEFAULT 0 CHECK(bots_notify IN (0, 1))
|
||||
)
|
||||
WITHOUT ROWID;
|
||||
|
||||
CREATE TABLE user_ignores (
|
||||
user_id INTEGER NOT NULL,
|
||||
guild_id INTEGER NOT NULL,
|
||||
target INTEGER NOT NULL,
|
||||
PRIMARY KEY (user_id, guild_id, target)
|
||||
);
|
||||
|
||||
COMMIT;
|
||||
CREATE TABLE user_blocks (
|
||||
user_id INTEGER NOT NULL,
|
||||
target INTEGER NOT NULL,
|
||||
PRIMARY KEY (user_id, target)
|
||||
);
|
||||
"""
|
||||
|
||||
|
||||
def run_db_migrations(db_file):
|
||||
log.debug(f"Running automatic migration on {db_file}")
|
||||
with sqlite3.connect(db_file) as db:
|
||||
if dumb_migrate_db(db, schema):
|
||||
log.info(f"Migrated {db_file}")
|
||||
log.debug("Finished running automatic migration")
|
||||
|
||||
|
||||
async def setup_db(db_file):
|
||||
log.debug(f"Connecting to {db_file}")
|
||||
db = await aiosqlite.connect(db_file)
|
||||
|
||||
log.debug("Running start script")
|
||||
await db.executescript("""
|
||||
PRAGMA optimize(0x10002);
|
||||
PRAGMA main.synchronous = NORMAL;
|
||||
""")
|
||||
|
||||
log.debug("Adding contains function")
|
||||
await db.create_function("contains", 3, contains, deterministic=True)
|
||||
|
||||
log.debug("Done setting up DB")
|
||||
|
||||
return db
|
||||
|
||||
|
||||
# TODO: Database versioning and migrations
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from disnake.ext import commands
|
|||
from disnake.ext.commands import Bot
|
||||
from dotenv import find_dotenv, load_dotenv
|
||||
|
||||
from .db import setup_db
|
||||
from .db import run_db_migrations, setup_db
|
||||
from .notifications import Notifications
|
||||
from .settings import Settings
|
||||
|
||||
|
|
@ -153,4 +153,6 @@ async def prefix(ctx, prefix=None):
|
|||
|
||||
|
||||
def run():
|
||||
if run_db_migrations(DB_FILE):
|
||||
log.info(f"Migrated DB {DB_FILE}")
|
||||
bot.run(TOKEN)
|
||||
|
|
|
|||
265
nomen/migrator.py
Normal file
265
nomen/migrator.py
Normal file
|
|
@ -0,0 +1,265 @@
|
|||
# coding: utf-8
|
||||
|
||||
"""Simple declarative schema migration for SQLite.
|
||||
|
||||
See <https://david.rothlis.net/declarative-schema-migration-for-sqlite>.
|
||||
|
||||
Author: William Manley <will@stb-tester.com>.
|
||||
Copyright © 2019-2022 Stb-tester.com Ltd.
|
||||
License: MIT.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
import sqlite3
|
||||
from textwrap import dedent
|
||||
|
||||
|
||||
def dumb_migrate_db(db, schema, allow_deletions=False):
|
||||
"""
|
||||
Migrates a database to the new schema given by the SQL text `schema`
|
||||
preserving the data. We create any table that exists in schema, delete any
|
||||
old table that is no longer used and add/remove columns and indices as
|
||||
necessary.
|
||||
|
||||
Under this scheme there are a set of changes that we can make to the schema
|
||||
and this script will handle it fine:
|
||||
|
||||
1. Adding a new table
|
||||
2. Adding, deleting or modifying an index
|
||||
3. Adding a column to an existing table as long as the new column can be
|
||||
NULL or has a DEFAULT value specified.
|
||||
4. Changing a column to remove NULL or DEFAULT as long as all values in the
|
||||
database are not NULL
|
||||
5. Changing the type of a column
|
||||
6. Changing the user_version
|
||||
|
||||
In addition this function is capable of:
|
||||
|
||||
1. Deleting tables
|
||||
2. Deleting columns from tables
|
||||
|
||||
But only if allow_deletions=True. If the new schema requires a column/table
|
||||
to be deleted and allow_deletions=False this function will raise
|
||||
`RuntimeError`.
|
||||
|
||||
Note: When this function is called a transaction must not be held open on
|
||||
db. A transaction will be used internally. If you wish to perform
|
||||
additional migration steps as part of a migration use DBMigrator directly.
|
||||
|
||||
Any internally generated rowid columns by SQLite may change values by this
|
||||
migration.
|
||||
"""
|
||||
with DBMigrator(db, schema, allow_deletions) as migrator:
|
||||
migrator.migrate()
|
||||
return bool(migrator.n_changes)
|
||||
|
||||
|
||||
class DBMigrator:
|
||||
def __init__(self, db, schema, allow_deletions=False):
|
||||
self.db = db
|
||||
self.schema = schema
|
||||
self.allow_deletions = allow_deletions
|
||||
|
||||
self.pristine = sqlite3.connect(":memory:")
|
||||
self.pristine.executescript(schema)
|
||||
self.n_changes = 0
|
||||
|
||||
self.orig_foreign_keys = None
|
||||
|
||||
def log_execute(self, msg, sql, args=None):
|
||||
# It's important to log any changes we're making to the database for
|
||||
# forensics later
|
||||
msg_tmpl = "Database migration: %s with SQL:\n%s"
|
||||
msg_argv = (msg, _left_pad(dedent(sql)))
|
||||
if args:
|
||||
msg_tmpl += " args = %r"
|
||||
msg_argv += (args,)
|
||||
else:
|
||||
args = []
|
||||
logging.info(msg_tmpl, *msg_argv)
|
||||
self.db.execute(sql, args)
|
||||
self.n_changes += 1
|
||||
|
||||
def __enter__(self):
|
||||
self.orig_foreign_keys = (
|
||||
self.db.execute("PRAGMA foreign_keys").fetchone()[0])
|
||||
if self.orig_foreign_keys:
|
||||
self.log_execute("Disable foreign keys temporarily for migration",
|
||||
"PRAGMA foreign_keys = OFF")
|
||||
# This doesn't count as a change because we'll undo it at the end
|
||||
self.n_changes = 0
|
||||
|
||||
self.db.__enter__()
|
||||
self.db.execute('BEGIN')
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, exc_tb):
|
||||
self.db.__exit__(exc_type, exc_value, exc_tb)
|
||||
if exc_value is None:
|
||||
# The SQLite docs say:
|
||||
#
|
||||
# > This pragma is a no-op within a transaction; foreign key
|
||||
# > constraint enforcement may only be enabled or disabled when
|
||||
# > there is no pending BEGIN or SAVEPOINT.
|
||||
old_changes = self.n_changes
|
||||
new_val = self._migrate_pragma('foreign_keys')
|
||||
if new_val == self.orig_foreign_keys:
|
||||
self.n_changes = old_changes
|
||||
|
||||
# SQLite docs say:
|
||||
#
|
||||
# > A VACUUM will fail if there is an open transaction on the database
|
||||
# > connection that is attempting to run the VACUUM.
|
||||
if self.n_changes:
|
||||
self.db.execute("VACUUM")
|
||||
else:
|
||||
if self.orig_foreign_keys:
|
||||
self.log_execute(
|
||||
"Re-enable foreign keys after migration",
|
||||
"PRAGMA foreign_keys = ON")
|
||||
|
||||
def migrate(self):
|
||||
# In CI the database schema may be changing all the time. This checks
|
||||
# the current db and if it doesn't match database.sql we will
|
||||
# modify it so it does match where possible.
|
||||
pristine_tables = dict(self.pristine.execute("""\
|
||||
SELECT name, sql FROM sqlite_master
|
||||
WHERE type = \"table\" AND name != \"sqlite_sequence\"""").fetchall())
|
||||
pristine_indices = dict(self.pristine.execute("""\
|
||||
SELECT name, sql FROM sqlite_master
|
||||
WHERE type = \"index\"""").fetchall())
|
||||
|
||||
tables = dict(self.db.execute("""\
|
||||
SELECT name, sql FROM sqlite_master
|
||||
WHERE type = \"table\" AND name != \"sqlite_sequence\"""").fetchall())
|
||||
|
||||
new_tables = set(pristine_tables.keys()) - set(tables.keys())
|
||||
removed_tables = set(tables.keys()) - set(pristine_tables.keys())
|
||||
if removed_tables and not self.allow_deletions:
|
||||
raise RuntimeError(
|
||||
"Database migration: Refusing to delete tables %r" %
|
||||
removed_tables)
|
||||
|
||||
modified_tables = set(
|
||||
name for name, sql in pristine_tables.items()
|
||||
if normalise_sql(tables.get(name, "")) != normalise_sql(sql))
|
||||
|
||||
# This PRAGMA is automatically disabled when the db is committed
|
||||
self.db.execute("PRAGMA defer_foreign_keys = TRUE")
|
||||
|
||||
# New and removed tables are easy:
|
||||
for tbl_name in new_tables:
|
||||
self.log_execute("Create table %s" % tbl_name,
|
||||
pristine_tables[tbl_name])
|
||||
for tbl_name in removed_tables:
|
||||
self.log_execute("Drop table %s" % tbl_name,
|
||||
"DROP TABLE %s" % tbl_name)
|
||||
|
||||
for tbl_name in modified_tables:
|
||||
# The SQLite documentation insists that we create the new table and
|
||||
# rename it over the old rather than moving the old out of the way
|
||||
# and then creating the new
|
||||
create_table_sql = pristine_tables[tbl_name]
|
||||
create_table_sql = re.sub(r"\b%s\b" % re.escape(tbl_name),
|
||||
tbl_name + "_migration_new",
|
||||
create_table_sql)
|
||||
self.log_execute(
|
||||
"Columns change: Create table %s with updated schema" %
|
||||
tbl_name, create_table_sql)
|
||||
|
||||
cols = set([
|
||||
x[1] for x in self.db.execute(
|
||||
"PRAGMA table_info(%s)" % tbl_name)])
|
||||
pristine_cols = set([
|
||||
x[1] for x in
|
||||
self.pristine.execute("PRAGMA table_info(%s)" % tbl_name)])
|
||||
|
||||
removed_columns = cols - pristine_cols
|
||||
if not self.allow_deletions and removed_columns:
|
||||
logging.warning(
|
||||
"Database migration: Refusing to remove columns %r from "
|
||||
"table %s. Current cols are %r attempting migration to %r",
|
||||
removed_columns, tbl_name, cols, pristine_cols)
|
||||
raise RuntimeError(
|
||||
"Database migration: Refusing to remove columns %r from "
|
||||
"table %s" % (removed_columns, tbl_name))
|
||||
|
||||
logging.info("cols: %s, pristine_cols: %s", cols, pristine_cols)
|
||||
self.log_execute(
|
||||
"Migrate data for table %s" % tbl_name, """\
|
||||
INSERT INTO {tbl_name}_migration_new ({common})
|
||||
SELECT {common} FROM {tbl_name}""".format(
|
||||
tbl_name=tbl_name,
|
||||
common=", ".join(cols.intersection(pristine_cols))))
|
||||
|
||||
# Don't need the old table any more
|
||||
self.log_execute(
|
||||
"Drop old table %s now data has been migrated" % tbl_name,
|
||||
"DROP TABLE %s" % tbl_name)
|
||||
|
||||
self.log_execute(
|
||||
"Columns change: Move new table %s over old" % tbl_name,
|
||||
"ALTER TABLE %s_migration_new RENAME TO %s" % (
|
||||
tbl_name, tbl_name))
|
||||
|
||||
# Migrate the indices
|
||||
indices = dict(self.db.execute("""\
|
||||
SELECT name, sql FROM sqlite_master
|
||||
WHERE type = \"index\"""").fetchall())
|
||||
for name in set(indices.keys()) - set(pristine_indices.keys()):
|
||||
self.log_execute("Dropping obsolete index %s" % name,
|
||||
"DROP INDEX %s" % name)
|
||||
for name, sql in pristine_indices.items():
|
||||
if name not in indices:
|
||||
self.log_execute("Creating new index %s" % name, sql)
|
||||
elif sql != indices[name]:
|
||||
self.log_execute(
|
||||
"Index %s changed: Dropping old version" % name,
|
||||
"DROP INDEX %s" % name)
|
||||
self.log_execute(
|
||||
"Index %s changed: Creating updated version in its place" %
|
||||
name, sql)
|
||||
|
||||
self._migrate_pragma('user_version')
|
||||
|
||||
if self.pristine.execute("PRAGMA foreign_keys").fetchone()[0]:
|
||||
if self.db.execute("PRAGMA foreign_key_check").fetchall():
|
||||
raise RuntimeError(
|
||||
"Database migration: Would fail foreign_key_check")
|
||||
|
||||
def _migrate_pragma(self, pragma):
|
||||
pristine_val = self.pristine.execute(
|
||||
"PRAGMA %s" % pragma).fetchone()[0]
|
||||
val = self.db.execute("PRAGMA %s" % pragma).fetchone()[0]
|
||||
|
||||
if val != pristine_val:
|
||||
self.log_execute(
|
||||
"Set %s to %i from %i" % (pragma, pristine_val, val),
|
||||
"PRAGMA %s = %i" % (pragma, pristine_val))
|
||||
|
||||
return pristine_val
|
||||
|
||||
|
||||
def _left_pad(text, indent=" "):
|
||||
"""Maybe I can find a package in pypi for this?"""
|
||||
return "\n".join(indent + line for line in text.split('\n'))
|
||||
|
||||
|
||||
def normalise_sql(sql):
|
||||
# Remove comments:
|
||||
sql = re.sub(r'--[^\n]*\n', "", sql)
|
||||
# Normalise whitespace:
|
||||
sql = re.sub(r'\s+', " ", sql)
|
||||
sql = re.sub(r" *([(),]) *", r"\1", sql)
|
||||
# Remove unnecessary quotes
|
||||
sql = re.sub(r'"(\w+)"', r"\1", sql)
|
||||
return sql.strip()
|
||||
|
||||
|
||||
def test_normalise_sql():
|
||||
assert normalise_sql("""\
|
||||
CREATE TABLE "Node"( -- This is my table
|
||||
-- There are many like it but this one is mine
|
||||
A b, C D, "E F G", h)""") == \
|
||||
'CREATE TABLE Node(A b,C D,"E F G",h)'
|
||||
Loading…
Reference in a new issue