add schema support for add_column, drop_column, alter_column, drop_table operations

diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py
index 220fbef03c5309fd681241c54244f86990b45ab5..60cf4d54f0dba1b95c1f88b60b66e01ffd52d582 100644 (file)
--- a/alembic/ddl/impl.py
+++ b/alembic/ddl/impl.py
@@ -122,11 +122,11 @@ class DefaultImpl(object):
                                 existing_nullable=existing_nullable,
                             ))
 
-    def add_column(self, table_name, column):
-        self._exec(base.AddColumn(table_name, column))
+    def add_column(self, table_name, column, schema=None):
+        self._exec(base.AddColumn(table_name, column, schema=schema))
 
-    def drop_column(self, table_name, column, **kw):
-        self._exec(base.DropColumn(table_name, column))
+    def drop_column(self, table_name, column, schema=None, **kw):
+        self._exec(base.DropColumn(table_name, column, schema=schema))
 
     def add_constraint(self, const):
         if const._create_rule is None or \

diff --git a/alembic/operations.py b/alembic/operations.py
index 791a4c2465299cb6df1ee7a2adfbb6ebc0dc046d..5c1d27b9b2adb1870d9e85c5ecccfb3057902391 100644 (file)
--- a/alembic/operations.py
+++ b/alembic/operations.py
@@ -160,7 +160,8 @@ class Operations(object):
                         existing_type=None,
                         existing_server_default=False,
                         existing_nullable=None,
-                        existing_autoincrement=None
+                        existing_autoincrement=None,
+                        schema=None
     ):
         """Issue an "alter column" instruction using the
         current migration context.
@@ -228,6 +229,7 @@ class Operations(object):
         :param existing_autoincrement: Optional; the existing autoincrement
          of the column.  Used for MySQL's system of altering a column
          that specifies ``AUTO_INCREMENT``.
+        :param schema: Optional, name of schema to operate within.
         """
 
         compiler = self.impl.dialect.statement_compiler(
@@ -241,7 +243,8 @@ class Operations(object):
 
         if existing_type and type_:
             t = self._table(table_name,
-                        sa_schema.Column(column_name, existing_type)
+                        sa_schema.Column(column_name, existing_type),
+                        schema=schema
                     )
             for constraint in t.constraints:
                 if _count_constraint(constraint):
@@ -252,6 +255,7 @@ class Operations(object):
             server_default=server_default,
             name=name,
             type_=type_,
+            schema=schema,
             autoincrement=autoincrement,
             existing_type=existing_type,
             existing_server_default=existing_server_default,
@@ -260,12 +264,15 @@ class Operations(object):
         )
 
         if type_:
-            t = self._table(table_name, sa_schema.Column(column_name, type_))
+            t = self._table(table_name,
+                        sa_schema.Column(column_name, type_),
+                        schema=schema
+                    )
             for constraint in t.constraints:
                 if _count_constraint(constraint):
                     self.impl.add_constraint(constraint)
 
-    def add_column(self, table_name, column):
+    def add_column(self, table_name, column, schema=None):
         """Issue an "add column" instruction using the current
         migration context.
 
@@ -308,13 +315,15 @@ class Operations(object):
         :param table_name: String name of the parent table.
         :param column: a :class:`sqlalchemy.schema.Column` object
          representing the new column.
+        :param schema: Optional, name of schema to operate within.
 
         """
 
-        t = self._table(table_name, column)
+        t = self._table(table_name, column, schema=schema)
         self.impl.add_column(
             table_name,
-            column
+            column,
+            schema=schema
         )
         for constraint in t.constraints:
             if not isinstance(constraint, sa_schema.PrimaryKeyConstraint):
@@ -534,7 +543,7 @@ class Operations(object):
             self._table(name, *columns, **kw)
         )
 
-    def drop_table(self, name):
+    def drop_table(self, name, **kw):
         """Issue a "drop table" instruction using the current
         migration context.
 
@@ -543,9 +552,13 @@ class Operations(object):
 
             drop_table("accounts")
 
+        :param name: Name of the table
+        :param \**kw: Other keyword arguments are passed to the underlying
+         :class:`.Table` object created for the command.
+
         """
         self.impl.drop_table(
-            self._table(name)
+            self._table(name, **kw)
         )
 
     def create_index(self, name, tablename, *columns, **kw):

diff --git a/tests/test_op.py b/tests/test_op.py
index b5dbc26dcf736c3d2c34709b758f5497c3eb9f62..7bf58a1ac3689b0f1e8fb1b64fea239d3959e9c6 100644 (file)
--- a/tests/test_op.py
+++ b/tests/test_op.py
@@ -23,11 +23,23 @@ def test_add_column():
     op.add_column('t1', Column('c1', Integer, nullable=False))
     context.assert_("ALTER TABLE t1 ADD COLUMN c1 INTEGER NOT NULL")
 
+def test_add_column_schema():
+    context = op_fixture()
+    op.add_column('t1', Column('c1', Integer, nullable=False), schema="foo")
+    context.assert_("ALTER TABLE foo.t1 ADD COLUMN c1 INTEGER NOT NULL")
+
 def test_add_column_with_default():
     context = op_fixture()
     op.add_column('t1', Column('c1', Integer, nullable=False, server_default="12"))
     context.assert_("ALTER TABLE t1 ADD COLUMN c1 INTEGER DEFAULT '12' NOT NULL")
 
+def test_add_column_schema_with_default():
+    context = op_fixture()
+    op.add_column('t1',
+            Column('c1', Integer, nullable=False, server_default="12"),
+            schema='foo')
+    context.assert_("ALTER TABLE foo.t1 ADD COLUMN c1 INTEGER DEFAULT '12' NOT NULL")
+
 def test_add_column_fk():
     context = op_fixture()
     op.add_column('t1', Column('c1', Integer, ForeignKey('c2.id'), nullable=False))
@@ -36,6 +48,16 @@ def test_add_column_fk():
         "ALTER TABLE t1 ADD FOREIGN KEY(c1) REFERENCES c2 (id)"
     )
 
+def test_add_column_schema_fk():
+    context = op_fixture()
+    op.add_column('t1',
+            Column('c1', Integer, ForeignKey('c2.id'), nullable=False),
+            schema='foo')
+    context.assert_(
+        "ALTER TABLE foo.t1 ADD COLUMN c1 INTEGER NOT NULL",
+        "ALTER TABLE foo.t1 ADD FOREIGN KEY(c1) REFERENCES c2 (id)"
+    )
+
 def test_add_column_schema_type():
     """Test that a schema type generates its constraints...."""
     context = op_fixture()
@@ -45,6 +67,15 @@ def test_add_column_schema_type():
         'ALTER TABLE t1 ADD CHECK (c1 IN (0, 1))'
     )
 
+def test_add_column_schema_schema_type():
+    """Test that a schema type generates its constraints...."""
+    context = op_fixture()
+    op.add_column('t1', Column('c1', Boolean, nullable=False), schema='foo')
+    context.assert_(
+        'ALTER TABLE foo.t1 ADD COLUMN c1 BOOLEAN NOT NULL',
+        'ALTER TABLE foo.t1 ADD CHECK (c1 IN (0, 1))'
+    )
+
 def test_add_column_schema_type_checks_rule():
     """Test that a schema type doesn't generate a
     constraint based on check rule."""
@@ -62,6 +93,16 @@ def test_add_column_fk_self_referential():
         "ALTER TABLE t1 ADD FOREIGN KEY(c1) REFERENCES t1 (c2)"
     )
 
+def test_add_column_schema_fk_self_referential():
+    context = op_fixture()
+    op.add_column('t1',
+            Column('c1', Integer, ForeignKey('foo.t1.c2'), nullable=False),
+            schema='foo')
+    context.assert_(
+        "ALTER TABLE foo.t1 ADD COLUMN c1 INTEGER NOT NULL",
+        "ALTER TABLE foo.t1 ADD FOREIGN KEY(c1) REFERENCES foo.t1 (c2)"
+    )
+
 def test_add_column_fk_schema():
     context = op_fixture()
     op.add_column('t1', Column('c1', Integer, ForeignKey('remote.t2.c2'), nullable=False))
@@ -70,11 +111,26 @@ def test_add_column_fk_schema():
     'ALTER TABLE t1 ADD FOREIGN KEY(c1) REFERENCES remote.t2 (c2)'
     )
 
+def test_add_column_schema_fk_schema():
+    context = op_fixture()
+    op.add_column('t1',
+            Column('c1', Integer, ForeignKey('remote.t2.c2'), nullable=False),
+            schema='foo')
+    context.assert_(
+    'ALTER TABLE foo.t1 ADD COLUMN c1 INTEGER NOT NULL',
+    'ALTER TABLE foo.t1 ADD FOREIGN KEY(c1) REFERENCES remote.t2 (c2)'
+    )
+
 def test_drop_column():
     context = op_fixture()
     op.drop_column('t1', 'c1')
     context.assert_("ALTER TABLE t1 DROP COLUMN c1")
 
+def test_drop_column_schema():
+    context = op_fixture()
+    op.drop_column('t1', 'c1', schema='foo')
+    context.assert_("ALTER TABLE foo.t1 DROP COLUMN c1")
+
 def test_alter_column_nullable():
     context = op_fixture()
     op.alter_column("t", "c", nullable=True)
@@ -84,6 +140,15 @@ def test_alter_column_nullable():
         "ALTER TABLE t ALTER COLUMN c DROP NOT NULL"
     )
 
+def test_alter_column_schema_nullable():
+    context = op_fixture()
+    op.alter_column("t", "c", nullable=True, schema='foo')
+    context.assert_(
+        # TODO: not sure if this is PG only or standard
+        # SQL
+        "ALTER TABLE foo.t ALTER COLUMN c DROP NOT NULL"
+    )
+
 def test_alter_column_not_nullable():
     context = op_fixture()
     op.alter_column("t", "c", nullable=False)
@@ -93,6 +158,15 @@ def test_alter_column_not_nullable():
         "ALTER TABLE t ALTER COLUMN c SET NOT NULL"
     )
 
+def test_alter_column_schema_not_nullable():
+    context = op_fixture()
+    op.alter_column("t", "c", nullable=False, schema='foo')
+    context.assert_(
+        # TODO: not sure if this is PG only or standard
+        # SQL
+        "ALTER TABLE foo.t ALTER COLUMN c SET NOT NULL"
+    )
+
 def test_alter_column_rename():
     context = op_fixture()
     op.alter_column("t", "c", name="x")
@@ -100,6 +174,13 @@ def test_alter_column_rename():
         "ALTER TABLE t RENAME c TO x"
     )
 
+def test_alter_column_schema_rename():
+    context = op_fixture()
+    op.alter_column("t", "c", name="x", schema='foo')
+    context.assert_(
+        "ALTER TABLE foo.t RENAME c TO x"
+    )
+
 def test_alter_column_type():
     context = op_fixture()
     op.alter_column("t", "c", type_=String(50))
@@ -107,6 +188,13 @@ def test_alter_column_type():
         'ALTER TABLE t ALTER COLUMN c TYPE VARCHAR(50)'
     )
 
+def test_alter_column_schema_type():
+    context = op_fixture()
+    op.alter_column("t", "c", type_=String(50), schema='foo')
+    context.assert_(
+        'ALTER TABLE foo.t ALTER COLUMN c TYPE VARCHAR(50)'
+    )
+
 def test_alter_column_set_default():
     context = op_fixture()
     op.alter_column("t", "c", server_default="q")
@@ -114,13 +202,30 @@ def test_alter_column_set_default():
         "ALTER TABLE t ALTER COLUMN c SET DEFAULT 'q'"
     )
 
+def test_alter_column_schema_set_default():
+    context = op_fixture()
+    op.alter_column("t", "c", server_default="q", schema='foo')
+    context.assert_(
+        "ALTER TABLE foo.t ALTER COLUMN c SET DEFAULT 'q'"
+    )
+
 def test_alter_column_set_compiled_default():
     context = op_fixture()
-    op.alter_column("t", "c", server_default=func.utc_thing(func.current_timestamp()))
+    op.alter_column("t", "c",
+            server_default=func.utc_thing(func.current_timestamp()))
     context.assert_(
         "ALTER TABLE t ALTER COLUMN c SET DEFAULT utc_thing(CURRENT_TIMESTAMP)"
     )
 
+def test_alter_column_schema_set_compiled_default():
+    context = op_fixture()
+    op.alter_column("t", "c",
+            server_default=func.utc_thing(func.current_timestamp()),
+            schema='foo')
+    context.assert_(
+        "ALTER TABLE foo.t ALTER COLUMN c SET DEFAULT utc_thing(CURRENT_TIMESTAMP)"
+    )
+
 def test_alter_column_drop_default():
     context = op_fixture()
     op.alter_column("t", "c", server_default=None)
@@ -128,6 +233,13 @@ def test_alter_column_drop_default():
         'ALTER TABLE t ALTER COLUMN c DROP DEFAULT'
     )
 
+def test_alter_column_schema_drop_default():
+    context = op_fixture()
+    op.alter_column("t", "c", server_default=None, schema='foo')
+    context.assert_(
+        'ALTER TABLE foo.t ALTER COLUMN c DROP DEFAULT'
+    )
+
 
 def test_alter_column_schema_type_unnamed():
     context = op_fixture('mssql')
@@ -137,6 +249,14 @@ def test_alter_column_schema_type_unnamed():
         'ALTER TABLE t ADD CHECK (c IN (0, 1))'
     )
 
+def test_alter_column_schema_schema_type_unnamed():
+    context = op_fixture('mssql')
+    op.alter_column("t", "c", type_=Boolean(), schema='foo')
+    context.assert_(
+        'ALTER TABLE foo.t ALTER COLUMN c BIT',
+        'ALTER TABLE foo.t ADD CHECK (c IN (0, 1))'
+    )
+
 def test_alter_column_schema_type_named():
     context = op_fixture('mssql')
     op.alter_column("t", "c", type_=Boolean(name="xyz"))
@@ -145,6 +265,14 @@ def test_alter_column_schema_type_named():
         'ALTER TABLE t ADD CONSTRAINT xyz CHECK (c IN (0, 1))'
     )
 
+def test_alter_column_schema_schema_type_named():
+    context = op_fixture('mssql')
+    op.alter_column("t", "c", type_=Boolean(name="xyz"), schema='foo')
+    context.assert_(
+        'ALTER TABLE foo.t ALTER COLUMN c BIT',
+        'ALTER TABLE foo.t ADD CONSTRAINT xyz CHECK (c IN (0, 1))'
+    )
+
 def test_alter_column_schema_type_existing_type():
     context = op_fixture('mssql')
     op.alter_column("t", "c", type_=String(10), existing_type=Boolean(name="xyz"))
@@ -153,6 +281,14 @@ def test_alter_column_schema_type_existing_type():
         'ALTER TABLE t ALTER COLUMN c VARCHAR(10)'
     )
 
+def test_alter_column_schema_schema_type_existing_type():
+    context = op_fixture('mssql')
+    op.alter_column("t", "c", type_=String(10), existing_type=Boolean(name="xyz"), schema='foo')
+    context.assert_(
+        'ALTER TABLE foo.t DROP CONSTRAINT xyz',
+        'ALTER TABLE foo.t ALTER COLUMN c VARCHAR(10)'
+    )
+
 def test_alter_column_schema_type_existing_type_no_const():
     context = op_fixture('postgresql')
     op.alter_column("t", "c", type_=String(10), existing_type=Boolean())
@@ -160,6 +296,13 @@ def test_alter_column_schema_type_existing_type_no_const():
         'ALTER TABLE t ALTER COLUMN c TYPE VARCHAR(10)'
     )
 
+def test_alter_column_schema_schema_type_existing_type_no_const():
+    context = op_fixture('postgresql')
+    op.alter_column("t", "c", type_=String(10), existing_type=Boolean(), schema='foo')
+    context.assert_(
+        'ALTER TABLE foo.t ALTER COLUMN c TYPE VARCHAR(10)'
+    )
+
 def test_alter_column_schema_type_existing_type_no_new_type():
     context = op_fixture('postgresql')
     op.alter_column("t", "c", nullable=False, existing_type=Boolean())
@@ -167,6 +310,13 @@ def test_alter_column_schema_type_existing_type_no_new_type():
         'ALTER TABLE t ALTER COLUMN c SET NOT NULL'
     )
 
+def test_alter_column_schema_schema_type_existing_type_no_new_type():
+    context = op_fixture('postgresql')
+    op.alter_column("t", "c", nullable=False, existing_type=Boolean(), schema='foo')
+    context.assert_(
+        'ALTER TABLE foo.t ALTER COLUMN c SET NOT NULL'
+    )
+
 def test_add_foreign_key():
     context = op_fixture()
     op.create_foreign_key('fk_test', 't1', 't2',
@@ -268,6 +418,13 @@ def test_drop_table():
         "DROP TABLE tb_test"
     )
 
+def test_drop_table_schema():
+    context = op_fixture()
+    op.drop_table('tb_test', schema='foo')
+    context.assert_(
+        "DROP TABLE foo.tb_test"
+    )
+
 def test_create_table_selfref():
     context = op_fixture()
     op.create_table(
@@ -351,4 +508,4 @@ def test_cant_op():
         "for the Alembic 'Operations' class.  "
         "Try placing this code inside a callable.",
         op.inline_literal, "asdf"
-    )
\ No newline at end of file
+    )