diff options
| -rw-r--r-- | alembic/autogenerate/render.py | 67 | ||||
| -rw-r--r-- | alembic/batch.py | 22 | ||||
| -rw-r--r-- | alembic/operations.py | 30 | ||||
| -rw-r--r-- | tests/test_autogenerate.py | 4 | ||||
| -rw-r--r-- | tests/test_batch.py | 55 |
5 files changed, 129 insertions, 49 deletions
diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py index d14842e..3bfa9c0 100644 --- a/alembic/autogenerate/render.py +++ b/alembic/autogenerate/render.py @@ -137,24 +137,32 @@ def _add_index(index, autogen_context): :class:`~sqlalchemy.schema.Index` instance. """ - text = "%(prefix)screate_index(%(name)r, %(table)r, [%(columns)s], "\ - "unique=%(unique)r%(schema)s%(kwargs)s)" % { - 'prefix': _alembic_autogenerate_prefix(autogen_context), - 'name': _render_gen_name(autogen_context, index.name), - 'table': _ident(index.table.name), - 'columns': ", ".join( - _get_index_rendered_expressions(index, autogen_context)), - 'unique': index.unique or False, - 'schema': (", schema=%r" % _ident(index.table.schema)) - if index.table.schema else '', - 'kwargs': ( - ', ' + - ', '.join( - ["%s=%s" % - (key, _render_potential_expr(val, autogen_context)) - for key, val in index.kwargs.items()])) - if len(index.kwargs) else '' - } + has_batch = 'batch_prefix' in autogen_context + + if has_batch: + tmpl = "%(prefix)screate_index(%(name)r, [%(columns)s], "\ + "unique=%(unique)r%(kwargs)s)" + else: + tmpl = "%(prefix)screate_index(%(name)r, %(table)r, [%(columns)s], "\ + "unique=%(unique)r%(schema)s%(kwargs)s)" + + text = tmpl % { + 'prefix': _alembic_autogenerate_prefix(autogen_context), + 'name': _render_gen_name(autogen_context, index.name), + 'table': _ident(index.table.name), + 'columns': ", ".join( + _get_index_rendered_expressions(index, autogen_context)), + 'unique': index.unique or False, + 'schema': (", schema=%r" % _ident(index.table.schema)) + if index.table.schema else '', + 'kwargs': ( + ', ' + + ', '.join( + ["%s=%s" % + (key, _render_potential_expr(val, autogen_context)) + for key, val in index.kwargs.items()])) + if len(index.kwargs) else '' + } return text @@ -163,14 +171,21 @@ def _drop_index(index, autogen_context): Generate Alembic operations for the DROP INDEX of an :class:`~sqlalchemy.schema.Index` instance. """ - text = "%(prefix)sdrop_index(%(name)r, "\ - "table_name=%(table_name)r%(schema)s)" % { - 'prefix': _alembic_autogenerate_prefix(autogen_context), - 'name': _render_gen_name(autogen_context, index.name), - 'table_name': _ident(index.table.name), - 'schema': ((", schema=%r" % _ident(index.table.schema)) - if index.table.schema else '') - } + has_batch = 'batch_prefix' in autogen_context + + if has_batch: + tmpl = "%(prefix)sdrop_index(%(name)r)" + else: + tmpl = "%(prefix)sdrop_index(%(name)r, "\ + "table_name=%(table_name)r%(schema)s)" + + text = tmpl % { + 'prefix': _alembic_autogenerate_prefix(autogen_context), + 'name': _render_gen_name(autogen_context, index.name), + 'table_name': _ident(index.table.name), + 'schema': ((", schema=%r" % _ident(index.table.schema)) + if index.table.schema else '') + } return text diff --git a/alembic/batch.py b/alembic/batch.py index 7bfcf22..f5957ca 100644 --- a/alembic/batch.py +++ b/alembic/batch.py @@ -78,18 +78,18 @@ class BatchOperationsImpl(object): def rename_table(self, *arg, **kw): self.batch.append(("rename_table", arg, kw)) + def create_index(self, idx): + self.batch.append(("create_index", (idx,), {})) + + def drop_index(self, idx): + self.batch.append(("drop_index", (idx,), {})) + def create_table(self, table): raise NotImplementedError("Can't create table in batch mode") def drop_table(self, table): raise NotImplementedError("Can't drop table in batch mode") - def create_index(self, index): - raise NotImplementedError("Can't create index in batch mode") - - def drop_index(self, index): - raise NotImplementedError("Can't drop index in batch mode") - class ApplyBatchImpl(object): def __init__(self, table, table_args, table_kwargs): @@ -117,6 +117,7 @@ class ApplyBatchImpl(object): self.named_constraints[const.name] = const else: self.unnamed_constraints.append(const) + for idx in self.table.indexes: self.indexes[idx.name] = idx @@ -251,5 +252,14 @@ class ApplyBatchImpl(object): except KeyError: raise ValueError("No such constraint: '%s'" % const.name) + def add_index(self, idx): + self.indexes[idx.name] = idx + + def drop_index(self, idx): + try: + del self.indexes[idx.name] + except KeyError: + raise ValueError("No such index: '%s'" % idx.name) + def rename_table(self, *arg, **kw): raise NotImplementedError("TODO") diff --git a/alembic/operations.py b/alembic/operations.py index 4665505..fc52e03 100644 --- a/alembic/operations.py +++ b/alembic/operations.py @@ -1309,6 +1309,7 @@ class BatchOperations(Operations): :meth:`.Operations.create_primary_key` """ + raise NotImplementedError("not yet implemented") def create_foreign_key( self, name, referent, local_cols, remote_cols, **kw): @@ -1331,7 +1332,8 @@ class BatchOperations(Operations): """ return super(BatchOperations, self).create_foreign_key( - name, self.impl.table_name, referent, local_cols, remote_cols,) + name, self.impl.table_name, referent, local_cols, remote_cols, + schema=self.impl.schema) def create_unique_constraint(self, name, local_cols, **kw): """Issue a "create unique constraint" instruction using the @@ -1345,6 +1347,7 @@ class BatchOperations(Operations): :meth:`.Operations.create_unique_constraint` """ + kw['schema'] = self.impl.schema return super(BatchOperations, self).create_unique_constraint( name, self.impl.table_name, local_cols, **kw) @@ -1360,6 +1363,7 @@ class BatchOperations(Operations): :meth:`.Operations.create_check_constraint` """ + raise NotImplementedError("not yet implemented") def drop_constraint(self, name, type_=None): """Issue a "drop constraint" instruction using the @@ -1374,13 +1378,23 @@ class BatchOperations(Operations): """ return super(BatchOperations, self).drop_constraint( - name, self.impl.table_name, type_=type_) + name, self.impl.table_name, type_=type_, + schema=self.impl.schema) + + def create_index(self, name, columns, **kw): + """Issue a "create index" instruction using the + current batch migration context.""" + + kw['schema'] = self.impl.schema + + return super(BatchOperations, self).create_index( + name, self.impl.table_name, columns, **kw) - def create_index(self, *arg, **kw): - """Not implemented for batch table operations.""" - self._noop("create_index") + def drop_index(self, name, **kw): + """Issue a "drop index" instruction using the + current batch migration context.""" - def drop_index(self, name): - """Not implemented for batch table operations.""" - self._noop("drop_index") + kw['schema'] = self.impl.schema + return super(BatchOperations, self).drop_index( + name, self.impl.table_name, **kw) diff --git a/tests/test_autogenerate.py b/tests/test_autogenerate.py index 60b3a3e..51cae27 100644 --- a/tests/test_autogenerate.py +++ b/tests/test_autogenerate.py @@ -594,7 +594,7 @@ nullable=True)) batch_op.alter_column('name', existing_type=sa.VARCHAR(length=50), nullable=False) - batch_op.drop_index('pw_idx', table_name='user') + batch_op.drop_index('pw_idx') batch_op.drop_column('pw') ### end Alembic commands ###""") @@ -603,7 +603,7 @@ nullable=True)) """### commands auto generated by Alembic - please adjust! ### with op.batch_alter_table('user', schema=None) as batch_op: batch_op.add_column(sa.Column('pw', sa.VARCHAR(length=50), nullable=True)) - batch_op.create_index('pw_idx', 'user', ['pw'], unique=False) + batch_op.create_index('pw_idx', ['pw'], unique=False) batch_op.alter_column('name', existing_type=sa.VARCHAR(length=50), nullable=True) diff --git a/tests/test_batch.py b/tests/test_batch.py index a387792..3892951 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -10,9 +10,9 @@ from alembic.batch import ApplyBatchImpl from alembic.migration import MigrationContext from sqlalchemy import Integer, Table, Column, String, MetaData, ForeignKey, \ - UniqueConstraint, ForeignKeyConstraint + UniqueConstraint, ForeignKeyConstraint, Index from sqlalchemy.sql import column -from sqlalchemy.schema import CreateTable +from sqlalchemy.schema import CreateTable, CreateIndex class BatchApplyTest(TestBase): @@ -42,6 +42,17 @@ class BatchApplyTest(TestBase): ) return ApplyBatchImpl(t, table_args, table_kwargs) + def _ix_fixture(self, table_args=(), table_kwargs={}): + m = MetaData() + t = Table( + 'tname', m, + Column('id', Integer, primary_key=True), + Column('x', String()), + Column('y', Integer), + Index('ix1', 'y') + ) + return ApplyBatchImpl(t, table_args, table_kwargs) + def _fk_fixture(self, table_args=(), table_kwargs={}): m = MetaData() t = Table( @@ -90,13 +101,23 @@ class BatchApplyTest(TestBase): CreateTable(impl.new_table).compile(dialect=context.dialect)) create_stmt = re.sub(r'[\n\t]', '', create_stmt) + idx_stmt = "" + for idx in impl.new_table.indexes: + idx_stmt += str(CreateIndex(idx).compile(dialect=context.dialect)) + idx_stmt = re.sub(r'[\n\t]', '', idx_stmt) + if ddl_contains: - assert ddl_contains in create_stmt + assert ddl_contains in create_stmt + idx_stmt if ddl_not_contains: - assert ddl_not_contains not in create_stmt + assert ddl_not_contains not in create_stmt + idx_stmt - context.assert_( + expected = [ create_stmt, + ] + if impl.new_table.indexes: + expected.append(idx_stmt) + + expected.extend([ 'INSERT INTO _alembic_batch_temp (%(colnames)s) ' 'SELECT %(tname_colnames)s FROM tname' % { "colnames": ", ".join([ @@ -116,7 +137,8 @@ class BatchApplyTest(TestBase): }, 'DROP TABLE tname', 'ALTER TABLE _alembic_batch_temp RENAME TO tname' - ) + ]) + context.assert_(*expected) return impl.new_table def test_change_type(self): @@ -250,6 +272,24 @@ class BatchApplyTest(TestBase): impl, colnames=['id', 'x', 'y'], ddl_not_contains="CONSTRAINT uq1 UNIQUE") + def test_add_index(self): + impl = self._simple_fixture() + ix = self.op._index('ix1', 'tname', ['y']) + + impl.add_index(ix) + self._assert_impl( + impl, colnames=['id', 'x', 'y'], + ddl_contains="CREATE INDEX ix1") + + def test_drop_index(self): + impl = self._ix_fixture() + + ix = self.op._index('ix1', 'tname', ['y']) + impl.drop_index(ix) + self._assert_impl( + impl, colnames=['id', 'x', 'y'], + ddl_not_contains="CONSTRAINT uq1 UNIQUE") + def test_add_table_opts(self): impl = self._simple_fixture(table_kwargs={'mysql_engine': 'InnoDB'}) self._assert_impl( @@ -304,7 +344,8 @@ class BatchAPITest(TestBase): mock.call( ['x'], ['user.y'], onupdate=None, ondelete=None, name='myfk', - initially=None, deferrable=None, match=None) + initially=None, deferrable=None, match=None, + schema=None) ] ) eq_( |