Move from sqlalchemy-migrate to alembic

This affects how columns are added and dropped from
tables. MetaData and SQLATable objects have to be
recreated to make sync changed columns.

dataset/persistence/database.py

@@ -10,12 +10,14 @@ except ImportError:
     from urlparse import parse_qs
 
 from sqlalchemy import create_engine
-from migrate.versioning.util import construct_engine
 from sqlalchemy.pool import NullPool
 from sqlalchemy.schema import MetaData, Column, Index
 from sqlalchemy.schema import Table as SQLATable
 from sqlalchemy import Integer, Text, String
 
+from alembic.migration import MigrationContext
+from alembic.operations import Operations
+
 from dataset.persistence.table import Table
 from dataset.persistence.util import ResultIter
 from dataset.util import DatasetException
@@ -42,9 +44,8 @@ class Database(object):
             if len(query):
                 url = url + '?' + urlencode(query, doseq=True)
         self.schema = schema
-        engine = create_engine(url, **kw)
+        self.engine = create_engine(url, **kw)
         self.url = url
-        self.engine = construct_engine(engine)
         self.metadata = MetaData(schema=schema)
         self.metadata.bind = self.engine
         if reflectMetadata:
@@ -59,6 +60,14 @@ class Database(object):
             return self.local.connection
         return self.engine
 
+    @property
+    def op(self):
+        if hasattr(self.local, 'connection'):
+            ctx = MigrationContext.configure(self.local.connection)
+        else:
+            ctx = MigrationContext.configure(self.engine)
+        return Operations(ctx)
+
     def _acquire(self):
         self.lock.acquire()
 
@@ -109,15 +118,15 @@ class Database(object):
 
     def create_table(self, table_name, primary_id='id', primary_type='Integer'):
         """
-        Creates a new table. The new table will automatically have an `id` column 
-        unless specified via optional parameter primary_id, which will be used 
-        as the primary key of the table. Automatic id is set to be an 
-        auto-incrementing integer, while the type of custom primary_id can be a 
-        String or an Integer as specified with primary_type flag. The default 
-        length of String is 256. The caller can specify the length. 
+        Creates a new table. The new table will automatically have an `id` column
+        unless specified via optional parameter primary_id, which will be used
+        as the primary key of the table. Automatic id is set to be an
+        auto-incrementing integer, while the type of custom primary_id can be a
+        String or an Integer as specified with primary_type flag. The default
+        length of String is 256. The caller can specify the length.
         The caller will be responsible for the uniqueness of manual primary_id.
 
-        This custom id feature is only available via direct create_table call. 
+        This custom id feature is only available via direct create_table call.
 
         Returns a :py:class:`Table <dataset.Table>` instance.
         ::
@@ -179,6 +188,12 @@ class Database(object):
         finally:
             self._release()
 
+    def update_table(self, table_name):
+        self.metadata = MetaData(schema=self.schema)
+        self.metadata.bind = self.engine
+        self.metadata.reflect(self.engine)
+        return SQLATable(table_name, self.metadata)
+
     def get_table(self, table_name, primary_id='id', primary_type='Integer'):
         """
         Smart wrapper around *load_table* and *create_table*. Either loads a table
@@ -222,7 +237,7 @@ class Database(object):
         iterator will yield each result sequentially.
 
         Any keyword arguments will be passed into the query to perform
-        parameter binding. 
+        parameter binding.
         ::
 
             res = db.query('SELECT user, COUNT(*) c FROM photos GROUP BY user')

dataset/persistence/table.py

@@ -218,9 +218,11 @@ class Table(object):
         self.database._acquire()
         try:
             if name not in self.table.columns.keys():
-                col = Column(name, type)
-                col.create(self.table,
-                           connection=self.database.executable)
+                self.database.op.add_column(
+                    self.table.name,
+                    Column(name, type)
+                )
+                self.table = self.database.update_table(self.table.name)
         finally:
             self.database._release()
 
@@ -235,8 +237,11 @@ class Table(object):
         self.database._acquire()
         try:
             if name in self.table.columns.keys():
-                col = self.table.columns[name]
-                col.drop()
+                self.database.op.drop_column(
+                    self.table.name,
+                    name
+                )
+                self.table = self.database.update_table(self.table.name)
         finally:
             self.database._release()
 

setup.py

@@ -23,7 +23,7 @@ setup(
     zip_safe=False,
     install_requires=[
         'sqlalchemy >= 0.8.1',
-        'sqlalchemy-migrate >= 0.7',
+        'alembic >= 0.6.1',
         "argparse >= 1.2.1",
         'python-slugify >= 0.0.6',
         "PyYAML >= 3.10"