implement table comments for batch

Added missing ``batch_op.create_table_comment()``,
``batch_op.drop_table_comment()`` directives to batch ops.

Change-Id: Ia7779619d150a2fe26abb8a8cc89d147a8432f8c
Fixes: #799

diff --git a/alembic/operations/batch.py b/alembic/operations/batch.py
index 990b3a887e1b6d67a83681f274a53fcb4889c251..b9268603ffa0837ea3c181648eb3438d2407a823 100644 (file)
--- a/alembic/operations/batch.py
+++ b/alembic/operations/batch.py
@@ -149,6 +149,12 @@ class BatchOperationsImpl(object):
     def drop_index(self, idx):
         self.batch.append(("drop_index", (idx,), {}))
 
+    def create_table_comment(self, table):
+        self.batch.append(("create_table_comment", (table,), {}))
+
+    def drop_table_comment(self, table):
+        self.batch.append(("drop_table_comment", (table,), {}))
+
     def create_table(self, table):
         raise NotImplementedError("Can't create table in batch mode")
 
@@ -528,6 +534,18 @@ class ApplyBatchImpl(object):
 
         """
 
+    def create_table_comment(self, table):
+        """the batch table creation function will issue create_table_comment
+        on the real "impl" as part of the create table process.
+
+        """
+
+    def drop_table_comment(self, table):
+        """the batch table creation function will issue drop_table_comment
+        on the real "impl" as part of the create table process.
+
+        """
+
     def add_constraint(self, const):
         if not const.name:
             raise ValueError("Constraint must have a name")

diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py
index 7ef219020e362a3650bf38ff04a06eb21cb258cb..9283a2d651d4cc94f0128773adf3c5d12c5173b6 100644 (file)
--- a/alembic/operations/ops.py
+++ b/alembic/operations/ops.py
@@ -1172,6 +1172,9 @@ class RenameTableOp(AlterTableOp):
 
 
 @Operations.register_operation("create_table_comment")
+@BatchOperations.register_operation(
+    "create_table_comment", "batch_create_table_comment"
+)
 class CreateTableCommentOp(AlterTableOp):
     """Represent a COMMENT ON `table` operation."""
 
@@ -1220,6 +1223,35 @@ class CreateTableCommentOp(AlterTableOp):
         )
         return operations.invoke(op)
 
+    @classmethod
+    def batch_create_table_comment(
+        cls,
+        operations,
+        comment,
+        existing_comment=None,
+    ):
+        """Emit a COMMENT ON operation to set the comment for a table
+        using the current batch migration context.
+
+        .. versionadded:: 1.6.0
+
+        :param comment: string value of the comment being registered against
+         the specified table.
+        :param existing_comment: String value of a comment
+         already registered on the specified table, used within autogenerate
+         so that the operation is reversible, but not required for direct
+         use.
+
+        """
+
+        op = cls(
+            operations.impl.table_name,
+            comment,
+            existing_comment=existing_comment,
+            schema=operations.impl.schema,
+        )
+        return operations.invoke(op)
+
     def reverse(self):
         """Reverses the COMMENT ON operation against a table."""
         if self.existing_comment is None:
@@ -1248,6 +1280,9 @@ class CreateTableCommentOp(AlterTableOp):
 
 
 @Operations.register_operation("drop_table_comment")
+@BatchOperations.register_operation(
+    "drop_table_comment", "batch_drop_table_comment"
+)
 class DropTableCommentOp(AlterTableOp):
     """Represent an operation to remove the comment from a table."""
 
@@ -1280,6 +1315,26 @@ class DropTableCommentOp(AlterTableOp):
         op = cls(table_name, existing_comment=existing_comment, schema=schema)
         return operations.invoke(op)
 
+    @classmethod
+    def batch_drop_table_comment(cls, operations, existing_comment=None):
+        """Issue a "drop table comment" operation to
+        remove an existing comment set on a table using the current
+        batch operations context.
+
+        .. versionadded:: 1.6.0
+
+        :param existing_comment: An optional string value of a comment already
+         registered on the specified table.
+
+        """
+
+        op = cls(
+            operations.impl.table_name,
+            existing_comment=existing_comment,
+            schema=operations.impl.schema,
+        )
+        return operations.invoke(op)
+
     def reverse(self):
         """Reverses the COMMENT ON operation against a table."""
         return CreateTableCommentOp(

diff --git a/docs/build/unreleased/799.rst b/docs/build/unreleased/799.rst
new file mode 100644 (file)
index 0000000..8f296ee
--- /dev/null
+++ b/docs/build/unreleased/799.rst
@@ -0,0 +1,6 @@
+.. change::
+    :tags: bug, batch
+    :tickets: 799
+
+    Added missing ``batch_op.create_table_comment()``,
+    ``batch_op.drop_table_comment()`` directives to batch ops.
\ No newline at end of file

diff --git a/tests/test_batch.py b/tests/test_batch.py
index 23ab364c1bde3b905938a319f05b7d2394e8a8ea..9e0491d7aaaf4f9679487d4c7c7c5f9a88c2f066 100644 (file)
--- a/tests/test_batch.py
+++ b/tests/test_batch.py
@@ -1703,6 +1703,36 @@ class BatchRoundTripTest(TestBase):
             ]
         )
 
+    def _assert_table_comment(self, tname, comment):
+        insp = inspect(config.db)
+
+        tcomment = insp.get_table_comment(tname)
+        eq_(tcomment, {"text": comment})
+
+    @config.requirements.comments
+    def test_add_table_comment(self):
+        with self.op.batch_alter_table("foo") as batch_op:
+            batch_op.create_table_comment("some comment")
+
+        self._assert_table_comment("foo", "some comment")
+
+        with self.op.batch_alter_table("foo") as batch_op:
+            batch_op.create_table_comment(
+                "some new comment", existing_comment="some comment"
+            )
+
+        self._assert_table_comment("foo", "some new comment")
+
+    @config.requirements.comments
+    def test_drop_table_comment(self):
+        with self.op.batch_alter_table("foo") as batch_op:
+            batch_op.create_table_comment("some comment")
+
+        with self.op.batch_alter_table("foo") as batch_op:
+            batch_op.drop_table_comment(existing_comment="some comment")
+
+        self._assert_table_comment("foo", None)
+
     def _assert_column_comment(self, tname, cname, comment):
         insp = inspect(config.db)