Migrations by yoyo

Change-Id: I6d06353462badac114c685a2f655bc41c6820aff

3 files changed, 32 insertions, 11 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").
diff --git a/setup.py b/setup.py
index 5fc49d2..f5b7838 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,6 @@
 #!/usr/bin/python
 #
-# Copyright (C) 2012-2014  Codethink Limited
+# Copyright (C) 2012-2016  Codethink Limited
 
 
 from distutils.core import setup
@@ -14,7 +14,7 @@ setup(name='lorry-controller',
       url='http://www.baserock.com/',
       scripts=['lorry-controller-webapp', 'lorry-controller-minion',
                'lorry-controller-remove-old-jobs'],
-      packages=['lorrycontroller'],
+      packages=['lorrycontroller', 'lorrycontroller.migrations'],
       data_files=[
         ('share/lorry-controller/templates', glob.glob('templates/*')),
         ('share/lorry-controller/static', glob.glob('static/*')),