Commits

bbinet committed 8d5ef85

add schema support for add/drop constraints and indexes

Comments (0)

Files changed (2)

alembic/operations.py

             t1_cols = local_cols + remote_cols
         else:
             t1_cols = local_cols
-            sa_schema.Table(referent, m,
-                *[sa_schema.Column(n, NULLTYPE) for n in remote_cols])
+            sname, tname = self._parse_table_key(referent)
+            sa_schema.Table(tname, m,
+                    *[sa_schema.Column(n, NULLTYPE) for n in remote_cols],
+                    schema=sname)
 
-        t1 = sa_schema.Table(source, m,
-                *[sa_schema.Column(n, NULLTYPE) for n in t1_cols])
+        sname, tname = self._parse_table_key(source)
+        t1 = sa_schema.Table(tname, m,
+                *[sa_schema.Column(n, NULLTYPE) for n in t1_cols],
+                schema=sname)
 
         f = sa_schema.ForeignKeyConstraint(local_cols,
                                             ["%s.%s" % (referent, n)
         return f
 
     def _unique_constraint(self, name, source, local_cols, **kw):
-        t = sa_schema.Table(source, sa_schema.MetaData(),
-                    *[sa_schema.Column(n, NULLTYPE) for n in local_cols])
+        sname, tname = self._parse_table_key(source)
+        t = sa_schema.Table(tname, sa_schema.MetaData(),
+                    *[sa_schema.Column(n, NULLTYPE) for n in local_cols],
+                    schema=sname)
         kw['name'] = name
         uq = sa_schema.UniqueConstraint(*[t.c[n] for n in local_cols], **kw)
         # TODO: need event tests to ensure the event
         return uq
 
     def _check_constraint(self, name, source, condition, **kw):
-        t = sa_schema.Table(source, sa_schema.MetaData(),
-                    sa_schema.Column('x', Integer))
+        sname, tname = self._parse_table_key(source)
+        t = sa_schema.Table(tname, sa_schema.MetaData(),
+                    sa_schema.Column('x', Integer), schema=sname)
         ck = sa_schema.CheckConstraint(condition, name=name, **kw)
         t.append_constraint(ck)
         return ck
     def _column(self, name, type_, **kw):
         return sa_schema.Column(name, type_, **kw)
 
-    def _index(self, name, tablename, columns, **kw):
+    def _index(self, name, tablename, columns, schema=None, **kw):
         t = sa_schema.Table(tablename or 'no_table', sa_schema.MetaData(),
-            *[sa_schema.Column(n, NULLTYPE) for n in columns]
+            *[sa_schema.Column(n, NULLTYPE) for n in columns],
+            schema=schema
         )
         return sa_schema.Index(name, *list(t.c), **kw)
 
+    def _parse_table_key(self, table_key):
+        if '.' in table_key:
+            tokens = table_key.split('.')
+            sname = ".".join(tokens[0:-1])
+            tname = tokens[-1]
+        else:
+            tname = table_key
+            sname = None
+        return (sname, tname)
+
     def _ensure_table_for_fk(self, metadata, fk):
         """create a placeholder Table object for the referent of a
         ForeignKey.
         """
         if isinstance(fk._colspec, basestring):
             table_key, cname = fk._colspec.rsplit('.', 1)
-            if '.' in table_key:
-                tokens = table_key.split('.')
-                sname = ".".join(tokens[0:-1])
-                tname = tokens[-1]
-            else:
-                tname = table_key
-                sname = None
+            sname, tname = self._parse_table_key(table_key)
             if table_key not in metadata.tables:
                 rel_t = sa_schema.Table(tname, metadata, schema=sname)
             else:
          ``name`` here can be ``None``, as the event listener will
          apply the name to the constraint object when it is associated
          with the table.
-        :param source: String name of the source table.  Currently
-         there is no support for dotted schema names.
-        :param referent: String name of the destination table. Currently
-         there is no support for dotted schema names.
+        :param source: String name of the source table. Dotted schema names
+         are supported.
+        :param referent: String name of the destination table. Dotted schema
+         names are supported.
         :param local_cols: a list of string column names in the
          source table.
         :param remote_cols: a list of string column names in the
          ``name`` here can be ``None``, as the event listener will
          apply the name to the constraint object when it is associated
          with the table.
-        :param source: String name of the source table.  Currently
-         there is no support for dotted schema names.
+        :param source: String name of the source table. Dotted schema names are
+         supported.
         :param local_cols: a list of string column names in the
          source table.
         :param deferrable: optional bool. If set, emit DEFERRABLE or NOT DEFERRABLE when
          ``name`` here can be ``None``, as the event listener will
          apply the name to the constraint object when it is associated
          with the table.
-        :param source: String name of the source table.  Currently
-         there is no support for dotted schema names.
+        :param source: String name of the source table. Dotted schema names
+         are supported.
         :param condition: SQL expression that's the condition of the constraint.
          Can be a string or SQLAlchemy expression language structure.
         :param deferrable: optional bool. If set, emit DEFERRABLE or NOT DEFERRABLE when
             self._table(name, **kw)
         )
 
-    def create_index(self, name, tablename, *columns, **kw):
+    def create_index(self, name, tablename, columns, schema=None, **kw):
         """Issue a "create index" instruction using the current
         migration context.
 
             from alembic import op
             op.create_index('ik_test', 't1', ['foo', 'bar'])
 
+        :param name: name of the index.
+        :param tablename: name of the owning table.
+        :param columns: a list of string column names in the
+         table.
+        :param schema: Optional, name of schema to operate within.
+
         """
 
         self.impl.create_index(
-            self._index(name, tablename, *columns, **kw)
+            self._index(name, tablename, columns, schema=schema, **kw)
         )
 
-    def drop_index(self, name, tablename=None):
+    def drop_index(self, name, tablename=None, schema=None):
         """Issue a "drop index" instruction using the current
         migration context.
 
 
             drop_index("accounts")
 
+        :param name: name of the index.
         :param tablename: name of the owning table.  Some
          backends such as Microsoft SQL Server require this.
+        :param schema: Optional, name of schema to operate within.
 
         """
         # need a dummy column name here since SQLAlchemy
         # 0.7.6 and further raises on Index with no columns
-        self.impl.drop_index(self._index(name, tablename, ['x']))
+        self.impl.drop_index(
+            self._index(name, tablename, ['x'], schema=schema)
+        )
 
-    def drop_constraint(self, name, tablename, type=None):
+    def drop_constraint(self, name, tablename, type=None, schema=None):
         """Drop a constraint of the given name, typically via DROP CONSTRAINT.
 
         :param name: name of the constraint.
         .. versionadded:: 0.3.6 'primary' qualfier to enable
            dropping of MySQL primary key constraints.
 
+        :param schema: Optional, name of schema to operate within.
+
         """
-        t = self._table(tablename)
+        t = self._table(tablename, schema=schema)
         types = {
             'foreignkey':lambda name:sa_schema.ForeignKeyConstraint(
                                 [], [], name=name),
 
 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')
+    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_schema_type_existing_type_no_const():
     context = op_fixture('postgresql')
-    op.alter_column("t", "c", type_=String(10), existing_type=Boolean(), schema='foo')
+    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_schema_type_existing_type_no_new_type():
     context = op_fixture('postgresql')
-    op.alter_column("t", "c", nullable=False, existing_type=Boolean(), schema='foo')
+    op.alter_column("t", "c", nullable=False, existing_type=Boolean(),
+            schema='foo')
     context.assert_(
         'ALTER TABLE foo.t ALTER COLUMN c SET NOT NULL'
     )
             "REFERENCES t2 (bat, hoho)"
     )
 
+def test_add_foreign_key_schema():
+    context = op_fixture()
+    op.create_foreign_key('fk_test', 'foo.t1', 'bar.t2',
+                    ['foo', 'bar'], ['bat', 'hoho'])
+    context.assert_(
+        "ALTER TABLE foo.t1 ADD CONSTRAINT fk_test FOREIGN KEY(foo, bar) "
+            "REFERENCES bar.t2 (bat, hoho)"
+    )
+
 def test_add_foreign_key_onupdate():
     context = op_fixture()
     op.create_foreign_key('fk_test', 't1', 't2',
         "CHECK (len(name) > 5)"
     )
 
+def test_add_check_constraint_schema():
+    context = op_fixture()
+    op.create_check_constraint(
+        "ck_user_name_len",
+        "foo.user_table",
+        func.len(column('name')) > 5
+    )
+    context.assert_(
+        "ALTER TABLE foo.user_table ADD CONSTRAINT ck_user_name_len "
+        "CHECK (len(name) > 5)"
+    )
+
 def test_add_unique_constraint():
     context = op_fixture()
     op.create_unique_constraint('uk_test', 't1', ['foo', 'bar'])
         "ALTER TABLE t1 ADD CONSTRAINT uk_test UNIQUE (foo, bar)"
     )
 
+def test_add_unique_constraint_schema():
+    context = op_fixture()
+    op.create_unique_constraint('uk_test', 'foo.t1', ['foo', 'bar'])
+    context.assert_(
+        "ALTER TABLE foo.t1 ADD CONSTRAINT uk_test UNIQUE (foo, bar)"
+    )
+
 def test_add_unique_constraint_auto_cols():
     context = op_fixture()
     from sqlalchemy import event, DateTime
         "ALTER TABLE t1 DROP CONSTRAINT foo_bar_bat"
     )
 
+def test_drop_constraint_schema():
+    context = op_fixture()
+    op.drop_constraint('foo_bar_bat', 't1', schema='foo')
+    context.assert_(
+        "ALTER TABLE foo.t1 DROP CONSTRAINT foo_bar_bat"
+    )
+
 def test_create_index():
     context = op_fixture()
     op.create_index('ik_test', 't1', ['foo', 'bar'])
         "CREATE INDEX ik_test ON t1 (foo, bar)"
     )
 
+def test_create_index_schema():
+    context = op_fixture()
+    op.create_index('ik_test', 't1', ['foo', 'bar'], schema='foo')
+    context.assert_(
+        "CREATE INDEX ik_test ON foo.t1 (foo, bar)"
+    )
 
 def test_drop_index():
     context = op_fixture()
         "DROP INDEX ik_test"
     )
 
+def test_drop_index_schema():
+    context = op_fixture()
+    op.drop_index('ik_test', schema='foo')
+    context.assert_(
+        "DROP INDEX foo.ik_test"
+    )
+
 def test_drop_table():
     context = op_fixture()
     op.drop_table('tb_test')