From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Tue, 29 Nov 2011 21:09:13 +0000 (-0500)
Subject: - implement "start migrations" event for impls
X-Git-Tag: rel_0_1_0~10
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=2370cdf0f8decab38419aa28175e6be5cf35ac25;p=thirdparty%2Fsqlalchemy%2Falembic.git

- implement "start migrations" event for impls
- implement counter logic for SQL server constraint/default
drop so that variables are declared uniquely within
a full migration run, #12
---

diff --git a/alembic/context.py b/alembic/context.py
index e737ed96..89bf606b 100644
--- a/alembic/context.py
+++ b/alembic/context.py
@@ -88,6 +88,7 @@ class Context(object):
     def run_migrations(self, **kw):
 
         current_rev = rev = False
+        self.impl.start_migrations()
         for change, prev_rev, rev in self._migrations_fn(
                                         self._current_rev()):
             if current_rev is False:
diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py
index 8306fcb6..a109ed4f 100644
--- a/alembic/ddl/impl.py
+++ b/alembic/ddl/impl.py
@@ -187,6 +187,14 @@ class DefaultImpl(object):
         conn_col_default = inspector_column['default']
         return conn_col_default != rendered_metadata_default
 
+    def start_migrations(self):
+        """A hook called when :meth:`.Context.run_migrations`
+        is called.
+        
+        Implementations can set up per-migration-run state here.
+        
+        """
+
     def emit_begin(self):
         """Emit the string ``BEGIN``, or the backend-specific
         equivalent, on the current connection context.
diff --git a/alembic/ddl/mssql.py b/alembic/ddl/mssql.py
index 400f6147..44d72f1e 100644
--- a/alembic/ddl/mssql.py
+++ b/alembic/ddl/mssql.py
@@ -1,12 +1,20 @@
 from alembic.ddl.impl import DefaultImpl
 from alembic.ddl.base import alter_table, AddColumn, ColumnName, \
     format_table_name, format_column_name, ColumnNullable, alter_column
+from alembic import util
 from sqlalchemy.ext.compiler import compiles
 
 class MSSQLImpl(DefaultImpl):
     __dialect__ = 'mssql'
     transactional_ddl = True
 
+    def start_migrations(self):
+        self.__dict__.pop('const_sym_counter', None)
+
+    @util.memoized_property
+    def const_sym_counter(self):
+        return 1
+
     def emit_begin(self):
         self._exec("BEGIN TRANSACTION")
 
@@ -29,30 +37,36 @@ class MSSQLImpl(DefaultImpl):
         drop_default = kw.pop('mssql_drop_default', False)
         if drop_default:
             self._exec(
-                _exec_drop_col_constraint(table_name, column, 'sys.default_constraints')
+                _exec_drop_col_constraint(self, 
+                        table_name, column, 
+                        'sys.default_constraints')
             )
         drop_check = kw.pop('mssql_drop_check', False)
         if drop_check:
             self._exec(
-                _exec_drop_col_constraint(table_name, column, 'sys.check_constraints')
+                _exec_drop_col_constraint(self, 
+                        table_name, column, 
+                        'sys.check_constraints')
             )
         super(MSSQLImpl, self).drop_column(table_name, column)
 
-
-def _exec_drop_col_constraint(tname, colname, type_):
+def _exec_drop_col_constraint(impl, tname, colname, type_):
     # from http://www.mssqltips.com/sqlservertip/1425/working-with-default-constraints-in-sql-server/
     # TODO: needs table formatting, etc.
-    return """declare @const_name varchar(256)
-select @const_name = [name] from %(type)s
+    counter = impl.const_sym_counter
+    impl.const_sym_counter += 1
+
+    return """declare @const_name_%(sym)s varchar(256)
+select @const_name_%(sym)s = [name] from %(type)s
 where parent_object_id = object_id('%(tname)s')
 and col_name(parent_object_id, parent_column_id) = '%(colname)s'
-exec('alter table %(tname)s drop constraint ' + @const_name)""" % {
+exec('alter table %(tname)s drop constraint ' + @const_name_%(sym)s)""" % {
         'type':type_,
         'tname':tname,
-        'colname':colname
+        'colname':colname,
+        'sym':counter
     }
 
-
 @compiles(AddColumn, 'mssql')
 def visit_add_column(element, compiler, **kw):
     return "%s %s" % (
diff --git a/tests/test_mssql.py b/tests/test_mssql.py
index 9392c6c9..5a653188 100644
--- a/tests/test_mssql.py
+++ b/tests/test_mssql.py
@@ -49,16 +49,25 @@ class OpTest(TestCase):
     def test_drop_column_w_default(self):
         context = op_fixture('mssql')
         op.drop_column('t1', 'c1', mssql_drop_default=True)
-        context.assert_contains("exec('alter table t1 drop constraint ' + @const_name)")
+        op.drop_column('t1', 'c2', mssql_drop_default=True)
+        context.assert_contains("exec('alter table t1 drop constraint ' + @const_name_1)")
         context.assert_contains("ALTER TABLE t1 DROP COLUMN c1")
 
+        # counter increments
+        context.assert_contains("exec('alter table t1 drop constraint ' + @const_name_2)")
+        context.assert_contains("ALTER TABLE t1 DROP COLUMN c2")
 
     def test_drop_column_w_check(self):
         context = op_fixture('mssql')
         op.drop_column('t1', 'c1', mssql_drop_check=True)
-        context.assert_contains("exec('alter table t1 drop constraint ' + @const_name)")
+        op.drop_column('t1', 'c2', mssql_drop_check=True)
+        context.assert_contains("exec('alter table t1 drop constraint ' + @const_name_1)")
         context.assert_contains("ALTER TABLE t1 DROP COLUMN c1")
 
+        # counter increments
+        context.assert_contains("exec('alter table t1 drop constraint ' + @const_name_2)")
+        context.assert_contains("ALTER TABLE t1 DROP COLUMN c2")
+
     def test_alter_column_nullable(self):
         context = op_fixture('mssql')
         op.alter_column("t", "c", nullable=True)