- implement add_column, drop_column, start thinking about scaling up how we do the...

diff --git a/alembic/context.py b/alembic/context.py
index 6ea301b2f7427d0abb385386df8eff4acd58f0ce..e11f4d61cc30798b601805e6c800aa1441ffb7d2 100644 (file)
--- 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 3c50bf12f0c6b3cc488bb183fa29a0e432ee5545..5f7bff44ddbacbab23ad9e0f7892e1ff0197782a 100644 (file)
--- 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 c402d7cb31f06e46ad6a0da44c0f87aee5072c4d..6d9f364fb610732a5e727118a1cf5ed13f477c1a 100644 (file)
--- 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 fea29fd08375c4aa85701853d2a8a87a9bdc1eef..d88f45e4faab2986db71cab16f248328a61f329c 100644 (file)
--- 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 4e0d19fac0e3c70c38ef26e52e54efb836b9a3a9..359673103e886f4cab1fac80be28f0e4019ca805 100644 (file)
--- 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 (file)
index 0000000..018ddba
--- /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 bc68de26b24107ff58fb452cc6be80c00464be59..cd8e765fa08b82acbff42f04027983d43c263f6a 100644 (file)
--- 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