Implement column comments for batch

Added missing "create comment" feature for columns that are altered in
batch migrations.

Change-Id: I7468937702cefc8f581e27b1c4ba670e4f8b7d17
Fixes: #761

diff --git a/alembic/operations/batch.py b/alembic/operations/batch.py
index b0b04157c22b32c5453b544a72e0d9af10767f6c..a5b31c9a090742a485e13202a6f57e9a97bd0e9a 100644 (file)
--- a/alembic/operations/batch.py
+++ b/alembic/operations/batch.py
@@ -152,6 +152,9 @@ class BatchOperationsImpl(object):
     def drop_table(self, table):
         raise NotImplementedError("Can't drop table in batch mode")
 
+    def create_column_comment(self, column):
+        self.batch.append(("create_column_comment", (column,), {}))
+
 
 class ApplyBatchImpl(object):
     def __init__(
@@ -398,6 +401,7 @@ class ApplyBatchImpl(object):
         name=None,
         type_=None,
         autoincrement=None,
+        comment=False,
         **kw
     ):
         existing = self.columns[column_name]
@@ -441,6 +445,9 @@ class ApplyBatchImpl(object):
         if autoincrement is not None:
             existing.autoincrement = bool(autoincrement)
 
+        if comment is not False:
+            existing.comment = comment
+
     def _setup_dependencies_for_add_column(
         self, colname, insert_before, insert_after
     ):
@@ -506,6 +513,16 @@ class ApplyBatchImpl(object):
         del self.column_transfers[column.name]
         self.existing_ordering.remove(column.name)
 
+    def create_column_comment(self, column):
+        """the batch table creation function will issue create_column_comment
+        on the real "impl" as part of the create table process.
+
+        That is, the Column object will have the comment on it already,
+        so when it is received by add_column() it will be a normal part of
+        the CREATE TABLE and doesn't need an extra step here.
+
+        """
+
     def add_constraint(self, const):
         if not const.name:
             raise ValueError("Constraint must have a name")

diff --git a/docs/build/unreleased/761.rst b/docs/build/unreleased/761.rst
new file mode 100644 (file)
index 0000000..b4cc018
--- /dev/null
+++ b/docs/build/unreleased/761.rst
@@ -0,0 +1,7 @@
+.. change::
+    :tags: bug, batch
+    :tickets: 761
+
+    Added missing "create comment" feature for columns that are altered in
+    batch migrations.
+

diff --git a/tests/test_batch.py b/tests/test_batch.py
index 9501e3234ac0702653c8395b79a7387839c88562..d10955258da0fc7b8e77363950deafc30f239dbd 100644 (file)
--- a/tests/test_batch.py
+++ b/tests/test_batch.py
@@ -348,6 +348,18 @@ class BatchApplyTest(TestBase):
         new_table = self._assert_impl(impl)
         eq_(new_table.c.x.name, "q")
 
+    def test_alter_column_comment(self):
+        impl = self._simple_fixture()
+        impl.alter_column("tname", "x", comment="some comment")
+        new_table = self._assert_impl(impl)
+        eq_(new_table.c.x.comment, "some comment")
+
+    def test_add_column_comment(self):
+        impl = self._simple_fixture()
+        impl.add_column("tname", Column("q", Integer, comment="some comment"))
+        new_table = self._assert_impl(impl, colnames=["id", "x", "y", "q"])
+        eq_(new_table.c.q.comment, "some comment")
+
     def test_rename_col_boolean(self):
         impl = self._boolean_fixture()
         impl.alter_column("tname", "flag", name="bflag")
@@ -1663,6 +1675,82 @@ class BatchRoundTripTest(TestBase):
             ]
         )
 
+    def _assert_column_comment(self, tname, cname, comment):
+        insp = inspect(config.db)
+
+        cols = {col["name"]: col for col in insp.get_columns(tname)}
+        eq_(cols[cname]["comment"], comment)
+
+    @config.requirements.comments
+    def test_add_column_comment(self):
+        with self.op.batch_alter_table("foo") as batch_op:
+            batch_op.add_column(Column("y", Integer, comment="some comment"))
+
+        self._assert_column_comment("foo", "y", "some comment")
+
+        self._assert_data(
+            [
+                {"id": 1, "data": "d1", "x": 5, "y": None},
+                {"id": 2, "data": "22", "x": 6, "y": None},
+                {"id": 3, "data": "8.5", "x": 7, "y": None},
+                {"id": 4, "data": "9.46", "x": 8, "y": None},
+                {"id": 5, "data": "d5", "x": 9, "y": None},
+            ]
+        )
+
+    @config.requirements.comments
+    def test_add_column_comment_recreate(self):
+        with self.op.batch_alter_table("foo", recreate="always") as batch_op:
+            batch_op.add_column(Column("y", Integer, comment="some comment"))
+
+        self._assert_column_comment("foo", "y", "some comment")
+
+        self._assert_data(
+            [
+                {"id": 1, "data": "d1", "x": 5, "y": None},
+                {"id": 2, "data": "22", "x": 6, "y": None},
+                {"id": 3, "data": "8.5", "x": 7, "y": None},
+                {"id": 4, "data": "9.46", "x": 8, "y": None},
+                {"id": 5, "data": "d5", "x": 9, "y": None},
+            ]
+        )
+
+    @config.requirements.comments
+    def test_alter_column_comment(self):
+        with self.op.batch_alter_table("foo") as batch_op:
+            batch_op.alter_column(
+                "x", existing_type=Integer(), comment="some comment"
+            )
+
+        self._assert_column_comment("foo", "x", "some comment")
+
+        self._assert_data(
+            [
+                {"id": 1, "data": "d1", "x": 5},
+                {"id": 2, "data": "22", "x": 6},
+                {"id": 3, "data": "8.5", "x": 7},
+                {"id": 4, "data": "9.46", "x": 8},
+                {"id": 5, "data": "d5", "x": 9},
+            ]
+        )
+
+    @config.requirements.comments
+    def test_alter_column_comment_recreate(self):
+        with self.op.batch_alter_table("foo", recreate="always") as batch_op:
+            batch_op.alter_column("x", comment="some comment")
+
+        self._assert_column_comment("foo", "x", "some comment")
+
+        self._assert_data(
+            [
+                {"id": 1, "data": "d1", "x": 5},
+                {"id": 2, "data": "22", "x": 6},
+                {"id": 3, "data": "8.5", "x": 7},
+                {"id": 4, "data": "9.46", "x": 8},
+                {"id": 5, "data": "d5", "x": 9},
+            ]
+        )
+
     def test_rename_column(self):
         with self.op.batch_alter_table("foo") as batch_op:
             batch_op.alter_column("x", new_column_name="y")