From: Mike Bayer Date: Thu, 21 Apr 2011 23:33:20 +0000 (-0400) Subject: - get more constraint construct to work, etc. X-Git-Tag: rel_0_1_0~75 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=2d7cdad5fb11a8c484331a6c13ba58025690134b;p=thirdparty%2Fsqlalchemy%2Falembic.git - get more constraint construct to work, etc. - move all the tests to test_op --- diff --git a/alembic/context.py b/alembic/context.py index e11f4d61..f4bb9cb6 100644 --- a/alembic/context.py +++ b/alembic/context.py @@ -122,7 +122,7 @@ class DefaultContext(object): nullable=util.NO_VALUE, server_default=util.NO_VALUE, name=util.NO_VALUE, - type=util.NO_VALUE, + type_=util.NO_VALUE, schema=None, ): @@ -133,7 +133,10 @@ class DefaultContext(object): table_name, column_name, server_default, schema=schema )) - + if type_ is not util.NO_VALUE: + self._exec(base.ColumnType( + table_name, column_name, type_, schema=schema + )) # ... etc def add_column(self, table_name, column): diff --git a/alembic/ddl/base.py b/alembic/ddl/base.py index 5f7bff44..8f1ec1c5 100644 --- a/alembic/ddl/base.py +++ b/alembic/ddl/base.py @@ -67,7 +67,7 @@ def visit_column_nullable(element, compiler, **kw): return "%s %s %s" % ( alter_table(compiler, element.table_name, element.schema), alter_column(compiler, element.column_name), - "NULL" if element.nullable else "NOT NULL" + "NULL" if element.nullable else "SET NOT NULL" ) def quote_dotted(name, quote): diff --git a/alembic/op.py b/alembic/op.py index 6d9f364f..a12aa644 100644 --- a/alembic/op.py +++ b/alembic/op.py @@ -6,6 +6,8 @@ from sqlalchemy import schema __all__ = [ 'alter_column', 'add_column', + 'drop_column', + 'add_constraint', 'create_foreign_key', 'create_table', 'drop_table', @@ -30,6 +32,17 @@ def _foreign_key_constraint(name, source, referent, local_cols, remote_cols): return f +def _ensure_table_for_constraint(name, constraint): + if getattr(constraint, 'parent', None) is not None: + return + if isinstance(constraint, schema.UniqueConstraint): + # TODO: what if constraint has Column objects already + columns = [schema.Column(n, NULLTYPE) for n in + constraint._pending_colargs] + else: + columns = [] + return schema.Table(name, schema.MetaData(), *(columns + [constraint]) ) + def _unique_constraint(name, source, local_cols): t = schema.Table(source, schema.MetaData(), *[schema.Column(n, NULLTYPE) for n in local_cols]) @@ -82,14 +95,14 @@ def alter_column(table_name, column_name, type_=type_ ) -def add_column(table_name, column_name, - type_, **kw): - c = _column(column_name, type_, **kw) - t = _table(table_name, c) +def add_column(table_name, column): + t = _table(table_name, column) get_context().add_column( table_name, - c + column ) + for constraint in [f.constraint for f in t.foreign_keys]: + get_context().add_constraint(constraint) def drop_column(table_name, column_name): get_context().drop_column( @@ -97,10 +110,15 @@ def drop_column(table_name, column_name): _column(column_name, NULLTYPE) ) +def add_constraint(table_name, constraint): + _ensure_table_for_constraint(table_name, constraint) + get_context().add_constraint( + constraint + ) def create_foreign_key(name, source, referent, local_cols, remote_cols): get_context().add_constraint( - _foreign_key_constraint(source, referent, + _foreign_key_constraint(name, source, referent, local_cols, remote_cols) ) diff --git a/tests/__init__.py b/tests/__init__.py index d88f45e4..3cd79ddf 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -5,6 +5,7 @@ import os import itertools from sqlalchemy import create_engine from alembic import context +import re staging_directory = os.path.join(os.path.dirname(__file__), 'scratch') @@ -45,14 +46,16 @@ class _op_fixture(context.DefaultContext): self.assertion = [] def _exec(self, construct): + sql = unicode(construct.compile()) + sql = re.sub(r'[\n\t]', '', sql) self.assertion.append( - unicode(construct.compile()) + sql ) - def assert_(self, sql): + def assert_(self, *sql): # TODO: make this more flexible about # whitespace and such - eq_("\n".join(self.assertion), sql) + eq_(self.assertion, list(sql)) def _sqlite_testing_config(): cfg = _testing_config() diff --git a/tests/test_ddl.py b/tests/test_ddl.py deleted file mode 100644 index 35967310..00000000 --- a/tests/test_ddl.py +++ /dev/null @@ -1,29 +0,0 @@ -from tests import assert_compiled -from sqlalchemy.schema import Column -from sqlalchemy.types import String, Integer, DateTime -from alembic.ddl.base import AddColumn, ColumnNullable, ColumnType, ColumnName - -# TODO: should these all just go to test_op ? - -def test_add_column(): - assert_compiled( - AddColumn("footable", Column("foocol", String(50), nullable=False)), - "ALTER TABLE footable ADD COLUMN foocol VARCHAR(50) NOT NULL" - ) - assert_compiled( - AddColumn("footable", Column("foocol", String(50), - server_default="12")), - "ALTER TABLE footable ADD COLUMN foocol VARCHAR(50) DEFAULT '12'" - ) - - -def test_column_nullable(): - assert_compiled( - ColumnNullable("footable", "foocol", True), - "ALTER TABLE footable ALTER COLUMN foocol NULL" - ) - - assert_compiled( - ColumnNullable("footable", "foocol", False), - "ALTER TABLE footable ALTER COLUMN foocol NOT NULL" - ) diff --git a/tests/test_op.py b/tests/test_op.py index 018ddbad..2ffd8ee6 100644 --- a/tests/test_op.py +++ b/tests/test_op.py @@ -2,15 +2,117 @@ from tests import _op_fixture from alembic import op -from sqlalchemy import Integer +from sqlalchemy import Integer, Column, ForeignKey, \ + UniqueConstraint, Table, MetaData def test_add_column(): context = _op_fixture() - op.add_column('t1', 'c1', Integer, nullable=False) + op.add_column('t1', Column('c1', Integer, nullable=False)) context.assert_("ALTER TABLE 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_fk(): + context = _op_fixture() + op.add_column('t1', Column('c1', Integer, ForeignKey('c2.id'), nullable=False)) + context.assert_( + "ALTER TABLE t1 ADD COLUMN c1 INTEGER NOT NULL", + "ALTER TABLE t1 ADD FOREIGN KEY(c1) REFERENCES c2 (id)" + ) def test_drop_column(): context = _op_fixture() op.drop_column('t1', 'c1') context.assert_("ALTER TABLE t1 DROP COLUMN c1") + +def test_alter_column_nullable(): + context = _op_fixture() + op.alter_column("t", "c", nullable=True) + context.assert_( + # TODO: not sure if this is supposed to be SET NULL + "ALTER TABLE t ALTER COLUMN c NULL" + ) + +def test_alter_column_not_nullable(): + context = _op_fixture() + op.alter_column("t", "c", nullable=False) + context.assert_( + # TODO: not sure if this is PG only or standard + # SQL + "ALTER TABLE t ALTER COLUMN c SET NOT NULL" + ) + +def test_add_foreign_key(): + context = _op_fixture() + op.create_foreign_key('fk_test', 't1', 't2', + ['foo', 'bar'], ['bat', 'hoho']) + context.assert_( + "ALTER TABLE t1 ADD CONSTRAINT fk_test FOREIGN KEY(foo, bar) " + "REFERENCES t2 (bat, hoho)" + ) + +def test_add_unique_constraint(): + context = _op_fixture() + op.create_unique_constraint('uk_test', 't1', ['foo', 'bar']) + context.assert_( + "ALTER TABLE t1 ADD CONSTRAINT uk_test UNIQUE (foo, bar)" + ) + +def test_add_unique_constraint_table_detached(): + context = _op_fixture() + op.add_constraint('t1', UniqueConstraint('foo', 'bar', name="uk_test")) + context.assert_( + "ALTER TABLE t1 ADD CONSTRAINT uk_test UNIQUE (foo, bar)" + ) + +def test_add_unique_constraint_table_attached(): + context = _op_fixture() + uq = UniqueConstraint('foo', 'bar', name="uk_test") + t1 = Table('t1', MetaData(), + Column('foo', Integer), + Column('bar', Integer), + uq + ) + op.add_constraint('t1', uq) + context.assert_( + "ALTER TABLE t1 ADD CONSTRAINT uk_test UNIQUE (foo, bar)" + ) + + +def test_create_table_fk_and_schema(): + context = _op_fixture() + op.create_table( + "some_table", + Column('id', Integer, primary_key=True), + Column('foo_id', Integer, ForeignKey('foo.id')), + schema='schema' + ) + context.assert_( + "CREATE TABLE schema.some_table (" + "id INTEGER NOT NULL, " + "foo_id INTEGER, " + "PRIMARY KEY (id), " + "FOREIGN KEY(foo_id) REFERENCES foo (id))" + ) + +def test_create_table_two_fk(): + context = _op_fixture() + op.create_table( + "some_table", + Column('id', Integer, primary_key=True), + Column('foo_id', Integer, ForeignKey('foo.id')), + Column('foo_bar', Integer, ForeignKey('foo.bar')), + ) + context.assert_( + "CREATE TABLE some_table (" + "id INTEGER NOT NULL, " + "foo_id INTEGER, " + "foo_bar INTEGER, " + "PRIMARY KEY (id), " + "FOREIGN KEY(foo_id) REFERENCES foo (id), " + "FOREIGN KEY(foo_bar) REFERENCES foo (bar))" + ) + diff --git a/tests/test_schema.py b/tests/test_schema.py deleted file mode 100644 index cd8e765f..00000000 --- a/tests/test_schema.py +++ /dev/null @@ -1,58 +0,0 @@ -from tests import assert_compiled -from alembic import op -from sqlalchemy.schema import AddConstraint, ForeignKeyConstraint, \ - CreateTable, Column, ForeignKey,\ - MetaData, Table -from sqlalchemy import Integer - -# TODO: should these all just go to test_op ? - -def test_foreign_key(): - fk = op._foreign_key_constraint('fk_test', 't1', 't2', - ['foo', 'bar'], ['bat', 'hoho']) - assert_compiled( - AddConstraint(fk), - "ALTER TABLE t1 ADD CONSTRAINT fk_test FOREIGN KEY(foo, bar) " - "REFERENCES t2 (bat, hoho)" - ) - -def test_unique_constraint(): - uc = op._unique_constraint('uk_test', 't1', ['foo', 'bar']) - assert_compiled( - AddConstraint(uc), - "ALTER TABLE t1 ADD CONSTRAINT uk_test UNIQUE (foo, bar)" - ) - - -def test_table_schema_fk(): - tb = op._table("some_table", - Column('id', Integer, primary_key=True), - Column('foo_id', Integer, ForeignKey('foo.id')), - schema='schema' - ) - assert_compiled( - CreateTable(tb), - "CREATE TABLE schema.some_table (" - "id INTEGER NOT NULL, " - "foo_id INTEGER, " - "PRIMARY KEY (id), " - "FOREIGN KEY(foo_id) REFERENCES foo (id))" - ) - -def test_table_two_fk(): - tb = op._table("some_table", - Column('id', Integer, primary_key=True), - Column('foo_id', Integer, ForeignKey('foo.id')), - Column('foo_bar', Integer, ForeignKey('foo.bar')), - ) - assert_compiled( - CreateTable(tb), - "CREATE TABLE some_table (" - "id INTEGER NOT NULL, " - "foo_id INTEGER, " - "foo_bar INTEGER, " - "PRIMARY KEY (id), " - "FOREIGN KEY(foo_id) REFERENCES foo (id), " - "FOREIGN KEY(foo_bar) REFERENCES foo (bar))" - ) -