Fix Migrator behavior... again!

I got pretty close last time, but the one table that I rebuilt manually
inside the with block was, of course, still hanging on to the table_old
when all the others got renamed. Grrr. This new format breaks the whole
thing into separate steps for rename, transfer, drop, all tables
in lockstep.

utilities/database_upgrader.py

@@ -6,7 +6,7 @@ from voussoirkit import sqlhelpers
 
 import etiquette
 
-class Regenerator:
+class Migrator:
     '''
     Many of the upgraders involve adding columns. ALTER TABLE ADD COLUMN only
     allows adding at the end, which I usually don't prefer. In order to add a
@@ -14,48 +14,59 @@ class Regenerator:
     the data, and drop the old one. But, foreign keys and indices will still
     point to the old table, which causes broken foreign keys and dropped
     indices. So, the only way to prevent all that is to regenerate all affected
-    tables and indices. Rathern than parsing relationships to determine the
+    tables and indices. Rather than parsing relationships to determine the
     affected tables, this implementation just regenerates everything.
 
     It's kind of horrible but it allows me to have the columns in the order I
     want instead of just always appending. Besides, modifying collations cannot
     be done in-place either.
-    '''
-    def __init__(self, photodb, except_tables=[]):
-        self.photodb = photodb
-        if isinstance(except_tables, str):
-            except_tables = [except_tables]
-        self.except_tables = except_tables
 
-    def __enter__(self):
+    If you want to truly remove a table or index and not have it get
+    regenerated, just do that before instantiating the Migrator.
+    '''
+    def __init__(self, photodb):
+        self.photodb = photodb
+
         query = 'SELECT name, sql FROM sqlite_master WHERE type == "table"'
-        if self.except_tables:
+        self.tables = {
-            query += ' AND name NOT IN ' + sqlhelpers.listify(self.except_tables)
+            name: {'create': sql, 'transfer': f'INSERT INTO {name} SELECT * FROM {name}_old'}
-        self.tables = list(self.photodb.sql_select(query))
+            for (name, sql) in self.photodb.sql_select(query)
+        }
+
+        # The user may be adding entirely new tables derived from the data of
+        # old ones. We'll need to skip new tables for the rename and drop_old
+        # steps. So we track which tables already existed at the beginning.
+        self.existing_tables = set(self.tables)
 
         query = 'SELECT name, sql FROM sqlite_master WHERE type == "index" AND name NOT LIKE "sqlite_%"'
         self.indices = list(self.photodb.sql_select(query))
 
-    def __exit__(self, exc_type, exc, exc_traceback):
+    def go(self):
-        if exc:
+        # This loop is split in many parts, because otherwise if table A
-            raise exc
-
-        # This loop is split in two parts, because otherwise if table A
         # references table B and table A is completely reconstructed, it will
         # be pointing to the version of B which has not been reconstructed yet,
         # which is about to get renamed to B_old and then A's reference will be
         # broken.
-        for (name, query) in self.tables:
+        self.photodb.sql_execute('PRAGMA foreign_keys = OFF')
+        self.photodb.sql_execute('BEGIN')
+        for (name, table) in self.tables.items():
+            if name not in self.existing_tables:
+                continue
             self.photodb.sql_execute(f'ALTER TABLE {name} RENAME TO {name}_old')
 
-        for (name, query) in self.tables:
+        for (name, table) in self.tables.items():
-            self.photodb.sql_execute(query)
+            self.photodb.sql_execute(table['create'])
-            self.photodb.sql_execute(f'INSERT INTO {name} SELECT * FROM {name}_old')
+
+        for (name, table) in self.tables.items():
+            self.photodb.sql_execute(table['transfer'])
+
+        for (name, query) in self.tables.items():
+            if name not in self.existing_tables:
+                continue
             self.photodb.sql_execute(f'DROP TABLE {name}_old')
 
         for (name, query) in self.indices:
             self.photodb.sql_execute(query)
-        # self.photodb.sql_execute('REINDEX')
 
 def upgrade_1_to_2(photodb):
     '''
@@ -184,35 +195,39 @@ def upgrade_6_to_7(photodb):
     for index in indices:
         photodb.sql_execute(f'DROP INDEX {index}')
 
-    with Regenerator(photodb, except_tables='albums'):
+    m = Migrator(photodb)
-        photodb.sql_executescript('''
+    m.tables['album_associated_directories']['create'] = '''
     CREATE TABLE album_associated_directories(
         albumid TEXT,
         directory TEXT COLLATE NOCASE
     );
+    '''
+    m.tables['album_associated_directories']['transfer'] = '''
+    INSERT INTO album_associated_directories SELECT
+        id,
+        associated_directory
+    FROM albums_old
+    WHERE associated_directory IS NOT NULL;
+    '''
 
-        ALTER TABLE albums RENAME TO deleting_albums;
+    m.tables['albums']['create'] = '''
-
     CREATE TABLE albums(
         id TEXT,
         title TEXT,
         description TEXT
     );
-
+    '''
+    m.tables['albums']['transfer'] = '''
     INSERT INTO albums SELECT
         id,
         title,
         description
-        FROM deleting_albums;
+    FROM albums_old;
+    '''
 
-        INSERT INTO album_associated_directories SELECT
+    m.go()
-            id,
-            associated_directory
-        FROM deleting_albums
-        WHERE associated_directory IS NOT NULL;
-
-        DROP TABLE deleting_albums;
 
+    photodb.sql_executescript('''
     CREATE INDEX IF NOT EXISTS index_album_associated_directories_albumid on
         album_associated_directories(albumid);
     CREATE INDEX IF NOT EXISTS index_album_associated_directories_directory on
@@ -292,13 +307,9 @@ def upgrade_10_to_11(photodb):
     Added Primary keys, Foreign keys, and NOT NULL constraints.
     Added author_id column to Album and Tag tables.
     '''
-    with Regenerator(photodb, except_tables=['albums', 'tags']):
+    m = Migrator(photodb)
-        photodb.sql_executescript('''
-        PRAGMA foreign_keys = OFF;
-        BEGIN;
-
-        ALTER TABLE albums RENAME TO albums_old;
 
+    m.tables['albums']['create'] = '''
     CREATE TABLE albums(
         id TEXT PRIMARY KEY NOT NULL,
         title TEXT,
@@ -306,18 +317,17 @@ def upgrade_10_to_11(photodb):
         author_id TEXT,
         FOREIGN KEY(author_id) REFERENCES users(id)
     );
-
+    '''
+    m.tables['albums']['transfer'] = '''
     INSERT INTO albums SELECT
         id,
         title,
         description,
         NULL
     FROM albums_old;
+    '''
 
-        DROP TABLE albums_old;
+    m.tables['tags']['create'] = '''
-
-        ALTER_TABLE tags RENAME TO tags_old;
-
     CREATE TABLE tags(
         id TEXT PRIMARY KEY NOT NULL,
         name TEXT NOT NULL,
@@ -325,16 +335,17 @@ def upgrade_10_to_11(photodb):
         author_id TEXT,
         FOREIGN KEY(author_id) REFERENCES users(id)
     );
-
+    '''
+    m.tables['tags']['transfer'] = '''
     INSERT INTO tags SELECT
         id,
         name,
         description,
         NULL
     FROM tags_old;
+    '''
 
-        DROP TABLE tags_old;
+    m.go()
-        ''')
 
 def upgrade_11_to_12(photodb):
     '''
@@ -352,14 +363,8 @@ def upgrade_12_to_13(photodb):
     '''
     Added display_name column to the User table.
     '''
-    with Regenerator(photodb, except_tables='users'):
+    m = Migrator(photodb)
-        photodb.sql_executescript('''
+    m.tables['users']['create'] = '''
-        PRAGMA foreign_keys = OFF;
-
-        BEGIN;
-
-        ALTER TABLE users RENAME TO users_old;
-
     CREATE TABLE users(
         id TEXT PRIMARY KEY NOT NULL,
         username TEXT NOT NULL COLLATE NOCASE,
@@ -367,7 +372,8 @@ def upgrade_12_to_13(photodb):
         display_name TEXT,
         created INT
     );
-
+    '''
+    m.tables['users']['transfer'] = '''
     INSERT INTO users SELECT
         id,
         username,
@@ -375,9 +381,7 @@ def upgrade_12_to_13(photodb):
         NULL,
         created
     FROM users_old;
-
+    '''
-        DROP TABLE users_old;
-        ''')
 
 def upgrade_13_to_14(photodb):
     '''
@@ -391,14 +395,8 @@ def upgrade_14_to_15(photodb):
     '''
     Added the dev_ino column to photos.
     '''
-    with Regenerator(photodb, except_tables='photos'):
+    m = Migrator(photodb)
-        photodb.sql_executescript('''
+    m.tables['photos']['create'] = '''
-        PRAGMA foreign_keys = OFF;
-
-        BEGIN;
-
-        ALTER TABLE photos RENAME TO photos_old;
-
     CREATE TABLE photos(
         id TEXT PRIMARY KEY NOT NULL,
         filepath TEXT COLLATE NOCASE,
@@ -418,7 +416,8 @@ def upgrade_14_to_15(photodb):
         searchhidden INT,
         FOREIGN KEY(author_id) REFERENCES users(id)
     );
-
+    '''
+    m.tables['photos']['transfer'] = '''
     INSERT INTO photos SELECT
         id,
         filepath,
@@ -437,11 +436,11 @@ def upgrade_14_to_15(photodb):
         author_id,
         searchhidden
     FROM photos_old;
+    '''
 
-        DROP TABLE photos_old;
+    m.go()
 
-        CREATE INDEX index_photos_dev_ino ON photos(dev_ino);
+    photodb.sql_execute('CREATE INDEX index_photos_dev_ino ON photos(dev_ino);')
-        ''')
 
     for photo in photodb.get_photos_by_recent():
         if not photo.real_path.is_file:
@@ -457,14 +456,8 @@ def upgrade_15_to_16(photodb):
     '''
     Added the basename column to photos. Added collate nocase to extension.
     '''
-    with Regenerator(photodb, except_tables='photos'):
+    m = Migrator(photodb)
-        photodb.sql_executescript('''
+    m.tables['photos']['create'] = '''
-        PRAGMA foreign_keys = OFF;
-
-        BEGIN;
-
-        ALTER TABLE photos RENAME TO photos_old;
-
     CREATE TABLE photos(
         id TEXT PRIMARY KEY NOT NULL,
         filepath TEXT COLLATE NOCASE,
@@ -485,7 +478,8 @@ def upgrade_15_to_16(photodb):
         searchhidden INT,
         FOREIGN KEY(author_id) REFERENCES users(id)
     );
-
+    '''
+    m.tables['photos']['transfer'] = '''
     INSERT INTO photos SELECT
         id,
         filepath,
@@ -505,9 +499,9 @@ def upgrade_15_to_16(photodb):
         author_id,
         searchhidden
     FROM photos_old;
+    '''
 
-        DROP TABLE photos_old;
+    m.go()
-        ''')
 
     for (id, filepath) in photodb.sql_select('SELECT id, filepath FROM photos'):
         basename = os.path.basename(filepath)