Migrations by yoyo

Change-Id: I6d06353462badac114c685a2f655bc41c6820aff

2 files changed, 30 insertions, 9 deletions

diff --git a/lorrycontroller/migrations/0001-add-gitlab_token-column.py b/lorrycontroller/migrations/0001-add-gitlab_token-column.py
new file mode 100644
index 0000000..e441809
--- /dev/null
+++ b/lorrycontroller/migrations/0001-add-gitlab_token-column.py
@@ -0,0 +1,19 @@
+# Copyright (C) 2016  Codethink Limited
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+
+import yoyo
+
+yoyo.step('ALTER TABLE troves ADD COLUMN gitlab_token')
diff --git a/lorrycontroller/statedb.py b/lorrycontroller/statedb.py
index 27e6fae..4b00db0 100644
--- a/lorrycontroller/statedb.py
+++ b/lorrycontroller/statedb.py
@@ -19,6 +19,8 @@ import os
 import sqlite3
 import time
 
+import yoyo
+
 import lorrycontroller
 
 
@@ -80,13 +82,14 @@ class StateDB(object):
             logging.debug('New connection is %r', self._conn)
             if not existed:
                 self._initialise_tables()
-            else:
-                c = self._conn.cursor()
-                c.execute('PRAGMA table_info(troves)')
-                columns = (info[1] for info in c.fetchall())
-                if 'gitlab_token' not in columns:
-                    c.execute('ALTER TABLE troves ADD COLUMN gitlab_token')
-                    self._conn.commit()
+
+            self.perform_any_migrations()
+
+    def perform_any_migrations(self):
+        backend = yoyo.get_backend('sqlite:///' + self._filename)
+        migrations_dir = os.path.join(os.path.dirname(__file__), 'migrations')
+        migrations = yoyo.read_migrations(migrations_dir)
+        backend.apply_migrations(backend.to_apply(migrations))
 
     def _initialise_tables(self):
         logging.debug('Initialising tables in database')
@@ -113,8 +116,7 @@ class StateDB(object):
             'ls_interval INT, '
             'ls_last_run INT, '
             'prefixmap TEXT, '
-            'ignore TEXT, '
-            'gitlab_token TEXT '
+            'ignore TEXT '
             ')')
 
         # Table for all the known lorries (the "run queue").