Add simple schema migration system

cheeseshop/dbapi.py

@@ -1,6 +1,10 @@
 from enum import Enum
 
-# import asyncpg
+import asyncpg
+
+
+class SchemaError(Exception):
+    pass
 
 
 class NotFoundError(Exception):
@@ -433,6 +437,41 @@ def __init__(self, event_id, map_id):
         self.map_id = map_id
 
 
+class SystemConfig(object):
+    @staticmethod
+    async def create_schema(conn):
+        await conn.execute('''
+            CREATE TABLE system_config(
+                key text UNIQUE NOT NULL,
+                value json
+            )
+        ''')
+
+    @staticmethod
+    async def set(conn, key, value):
+        await conn.execute('''
+            INSERT INTO system_config (key, value)
+            VALUES ($1, $2)
+        ''', key, value)
+        return SystemConfig(key, value)
+
+    async def get(conn, key):
+        try:
+            row = await conn.fetchrow('''
+                SELECT * FROM system_config
+                WHERE key = $1
+            ''', key)
+        except asyncpg.exceptions.UndefinedTableError:
+            raise SchemaError()
+        if row is None:
+            raise NotFoundError
+        return SystemConfig(row['key'], row['value'])
+
+    def __init__(self, key, value):
+        self.key = key
+        self.value = value
+
+
 async def create_schema(conn):
     await Game.create_schema(conn)
     await Replay.create_schema(conn)
@@ -444,6 +483,7 @@ async def create_schema(conn):
     await CsGoGsiEvent.create_schema(conn)
     await CsGoMap.create_schema(conn)
     await CsGoEventMapRelation.create_schema(conn)
+    await SystemConfig.create_schema(conn)
 
 
 async def create_initial_records(conn):

cheeseshop/dbmigrations/__init__.py

@@ -0,0 +1,38 @@
+from cheeseshop import systemconfig
+from cheeseshop import dbapi
+from cheeseshop.dbmigrations import add_systemconfig
+
+# ONLY APPEND TO THIS LIST
+# Migrations are run in order and we start off based on the last-migration
+# property stored in SysemConfig.
+migrations = [
+    add_systemconfig
+]
+
+async def run_migrations(conn):
+    sc = systemconfig.SystemConfig(conn)
+    create_initial_schema = False
+    last_migration = None
+
+    try:
+        last_migration = await sc.last_migration()
+    except (dbapi.SchemaError, dbapi.NotFoundError):
+        create_initial_schema = True
+
+    if create_initial_schema:
+        await dbapi.create_schema(conn)
+        await dbapi.create_initial_records(conn)
+
+    start_migrating = False
+    if last_migration is None:
+        start_migrating = True
+    for migration in migrations:
+        cur_name = migration.__name__.rsplit('.', 1)[1]
+        if start_migrating is False:
+            # Check if were up to last_migration and if so start migrating
+            if last_migration == cur_name:
+                start_migrating = True
+                continue
+        else:
+            await migration.run(conn)
+            await sc.set_last_migration(cur_name)

cheeseshop/dbmigrations/add_systemconfig.py

@@ -0,0 +1,4 @@
+async def run(conn):
+    # This is a noop migration which is needed to let the migration system
+    # know we have created initial schemas (by setting the last-migration key)
+    pass

cheeseshop/main.py

@@ -11,6 +11,7 @@
 from cheeseshop import config as cs_config
 from cheeseshop import db
 from cheeseshop import dbapi
+from cheeseshop import dbmigrations
 from cheeseshop.games import csgo
 from cheeseshop import objectstoreapi
 from cheeseshop import swift
@@ -21,7 +22,7 @@ def parse_args(args):
     parser = argparse.ArgumentParser(description='cheeseshop webapp.')
     parser.add_argument('config_file', type=str,
                         help='Path to config file')
-    parser.add_argument('--create-schema', action='store_true')
+    parser.add_argument('--update-schema', action='store_true')
     return parser.parse_args(args)
 
 
@@ -165,11 +166,11 @@ def main():
     loop = asyncio.get_event_loop()
     pool = loop.run_until_complete(db.create_pool(config.sql))
 
-    if args.create_schema:
+    if args.update_schema:
+        print('Updating DB schema')
         conn = loop.run_until_complete(pool.acquire())
-        loop.run_until_complete(dbapi.create_schema(conn))
-        loop.run_until_complete(dbapi.create_initial_records(conn))
-        print('DB Schema created')
+        loop.run_until_complete(dbmigrations.run_migrations(conn))
+        print('DB Schema updated')
     else:
         app = App(config, pool)
         app.run()

cheeseshop/systemconfig.py

@@ -0,0 +1,15 @@
+import json
+
+from cheeseshop import dbapi
+
+class SystemConfig(object):
+    def __init__(self, conn):
+        self._conn = conn
+
+    async def last_migration(self):
+        record = await dbapi.SystemConfig.get(self._conn, 'last-migration')
+        return json.loads(record.value)
+
+    async def set_last_migration(self, value):
+        return await dbapi.SystemConfig.set(self._conn, 'last-migration',
+                                            json.dumps(value))