diff --git a/dataset/persistence/database.py b/dataset/persistence/database.py index 5a99470..b1c07bf 100644 --- a/dataset/persistence/database.py +++ b/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 ` 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') diff --git a/dataset/persistence/table.py b/dataset/persistence/table.py index 7edb475..9241e78 100644 --- a/dataset/persistence/table.py +++ b/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() diff --git a/setup.py b/setup.py index 490a922..bc6a221 100644 --- a/setup.py +++ b/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"