From: Mike Bayer Date: Thu, 21 Apr 2011 21:15:48 +0000 (-0400) Subject: - implement add_column, drop_column, start thinking about scaling up how we do the... X-Git-Tag: rel_0_1_0~76 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=177b7de4ee0952dde61636c4a8c00b6568a576b4;p=thirdparty%2Fsqlalchemy%2Falembic.git - implement add_column, drop_column, start thinking about scaling up how we do the tests --- diff --git a/alembic/context.py b/alembic/context.py index 6ea301b2..e11f4d61 100644 --- a/alembic/context.py +++ b/alembic/context.py @@ -2,8 +2,10 @@ from alembic import util from sqlalchemy import MetaData, Table, Column, String, literal_column, \ text from sqlalchemy import schema, create_engine -import logging +from sqlalchemy.util import importlater +import logging +base = importlater("alembic.ddl", "base") log = logging.getLogger(__name__) class ContextMeta(type): @@ -120,18 +122,26 @@ 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, ): if nullable is not util.NO_VALUE: - self._exec(base.ColumnNullable(table_name, column_name, nullable)) + self._exec(base.ColumnNullable(table_name, column_name, nullable, schema=schema)) if server_default is not util.NO_VALUE: self._exec(base.ColumnDefault( - table_name, column_name, server_default + table_name, column_name, server_default, + schema=schema )) # ... etc + def add_column(self, table_name, column): + self._exec(base.AddColumn(table_name, column)) + + def drop_column(self, table_name, column): + self._exec(base.DropColumn(table_name, column)) + def add_constraint(self, const): self._exec(schema.AddConstraint(const)) diff --git a/alembic/ddl/base.py b/alembic/ddl/base.py index 3c50bf12..5f7bff44 100644 --- a/alembic/ddl/base.py +++ b/alembic/ddl/base.py @@ -43,6 +43,11 @@ class AddColumn(AlterTable): super(AddColumn, self).__init__(name, schema=schema) self.column = column +class DropColumn(AlterTable): + def __init__(self, name, column, schema=None): + super(DropColumn, self).__init__(name, schema=schema) + self.column = column + @compiles(AddColumn) def visit_add_column(element, compiler, **kw): return "%s %s" % ( @@ -50,6 +55,13 @@ def visit_add_column(element, compiler, **kw): add_column(compiler, element.column, **kw) ) +@compiles(DropColumn) +def visit_drop_column(element, compiler, **kw): + return "%s %s" % ( + alter_table(compiler, element.table_name, element.schema), + drop_column(compiler, element.column.name, **kw) + ) + @compiles(ColumnNullable) def visit_column_nullable(element, compiler, **kw): return "%s %s %s" % ( diff --git a/alembic/op.py b/alembic/op.py index c402d7cb..6d9f364f 100644 --- a/alembic/op.py +++ b/alembic/op.py @@ -5,6 +5,7 @@ from sqlalchemy import schema __all__ = [ 'alter_column', + 'add_column', 'create_foreign_key', 'create_table', 'drop_table', @@ -13,22 +14,6 @@ __all__ = [ 'get_bind', 'execute'] -def alter_column(table_name, column_name, - nullable=util.NO_VALUE, - server_default=util.NO_VALUE, - name=util.NO_VALUE, - type_=util.NO_VALUE -): - """Issue ALTER COLUMN using the current change context.""" - - context.alter_column(table_name, column_name, - nullable=nullable, - server_default=server_default, - name=name, - type_=type_ - ) - - def _foreign_key_constraint(name, source, referent, local_cols, remote_cols): m = schema.MetaData() t1 = schema.Table(source, m, @@ -57,6 +42,9 @@ def _table(name, *columns, **kw): _ensure_table_for_fk(m, f) return t +def _column(name, type_, **kw): + return schema.Column(name, type_, **kw) + def _ensure_table_for_fk(metadata, fk): """create a placeholder Table object for the referent of a ForeignKey. @@ -78,6 +66,38 @@ def _ensure_table_for_fk(metadata, fk): if not rel_t.c.contains_column(cname): rel_t.append_column(schema.Column(cname, NULLTYPE)) + +def alter_column(table_name, column_name, + nullable=util.NO_VALUE, + server_default=util.NO_VALUE, + name=util.NO_VALUE, + type_=util.NO_VALUE +): + """Issue ALTER COLUMN using the current change context.""" + + get_context().alter_column(table_name, column_name, + nullable=nullable, + server_default=server_default, + name=name, + type_=type_ + ) + +def add_column(table_name, column_name, + type_, **kw): + c = _column(column_name, type_, **kw) + t = _table(table_name, c) + get_context().add_column( + table_name, + c + ) + +def drop_column(table_name, column_name): + get_context().drop_column( + table_name, + _column(column_name, NULLTYPE) + ) + + def create_foreign_key(name, source, referent, local_cols, remote_cols): get_context().add_constraint( _foreign_key_constraint(source, referent, diff --git a/tests/__init__.py b/tests/__init__.py index fea29fd0..d88f45e4 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -4,6 +4,7 @@ import shutil import os import itertools from sqlalchemy import create_engine +from alembic import context staging_directory = os.path.join(os.path.dirname(__file__), 'scratch') @@ -37,6 +38,22 @@ def _testing_config(): os.mkdir(staging_directory) return Config(os.path.join(staging_directory, 'test_alembic.ini')) +class _op_fixture(context.DefaultContext): + def __init__(self): + # TODO: accept dialect here. + context._context = self + self.assertion = [] + + def _exec(self, construct): + self.assertion.append( + unicode(construct.compile()) + ) + + def assert_(self, sql): + # TODO: make this more flexible about + # whitespace and such + eq_("\n".join(self.assertion), sql) + def _sqlite_testing_config(): cfg = _testing_config() dir_ = os.path.join(staging_directory, 'scripts') diff --git a/tests/test_ddl.py b/tests/test_ddl.py index 4e0d19fa..35967310 100644 --- a/tests/test_ddl.py +++ b/tests/test_ddl.py @@ -3,6 +3,7 @@ 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( diff --git a/tests/test_op.py b/tests/test_op.py new file mode 100644 index 00000000..018ddbad --- /dev/null +++ b/tests/test_op.py @@ -0,0 +1,16 @@ +"""Test against the builders in the op.* module.""" + +from tests import _op_fixture +from alembic import op +from sqlalchemy import Integer + +def test_add_column(): + context = _op_fixture() + op.add_column('t1', 'c1', Integer, nullable=False) + context.assert_("ALTER TABLE t1 ADD COLUMN c1 INTEGER NOT NULL") + + +def test_drop_column(): + context = _op_fixture() + op.drop_column('t1', 'c1') + context.assert_("ALTER TABLE t1 DROP COLUMN c1") diff --git a/tests/test_schema.py b/tests/test_schema.py index bc68de26..cd8e765f 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -5,6 +5,8 @@ from sqlalchemy.schema import AddConstraint, ForeignKeyConstraint, \ 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']) @@ -22,7 +24,7 @@ def test_unique_constraint(): ) -def test_table(): +def test_table_schema_fk(): tb = op._table("some_table", Column('id', Integer, primary_key=True), Column('foo_id', Integer, ForeignKey('foo.id')), @@ -37,6 +39,7 @@ def test_table(): "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')), @@ -53,17 +56,3 @@ def test_table(): "FOREIGN KEY(foo_bar) REFERENCES foo (bar))" ) - m = MetaData() - foo = Table('foo', m, Column('id', Integer, primary_key=True)) - tb = op._table("some_table", - Column('id', Integer, primary_key=True), - Column('foo_id', Integer, ForeignKey('foo.id')), - ) - assert_compiled( - CreateTable(tb), - "CREATE TABLE some_table (" - "id INTEGER NOT NULL, " - "foo_id INTEGER, " - "PRIMARY KEY (id), " - "FOREIGN KEY(foo_id) REFERENCES foo (id))" - ) \ No newline at end of file