Commits

bbinet committed f730920

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

Comments (0)

Files changed (3)

                                 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 \

alembic/operations.py

                         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.
         :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(
 
         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):
             server_default=server_default,
             name=name,
             type_=type_,
+            schema=schema,
             autoincrement=autoincrement,
             existing_type=existing_type,
             existing_server_default=existing_server_default,
         )
 
         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.
 
         :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):
             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.
 
 
             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):
     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))
         "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()
         '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."""
         "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))
     '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)
         "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)
         "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")
         "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))
         '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")
         "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)
         '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')
         '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"))
         '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"))
         '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())
         '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())
         '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',
         "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(
         "for the Alembic 'Operations' class.  "
         "Try placing this code inside a callable.",
         op.inline_literal, "asdf"
-    )
+    )