From ea44e7f3398aa5001e86e6a7e30f61f6a86d9674 Mon Sep 17 00:00:00 2001
From: Adrien Berchet <adrien.berchet@gmail.com>
Date: Mon, 16 May 2022 10:04:35 +0200
Subject: [PATCH] Trigger events before and after drop table statements

---
 alembic/ddl/impl.py |  6 ++++++
 tests/test_op.py    | 49 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py
index 8c9e0b91..070c124b 100644
--- a/alembic/ddl/impl.py
+++ b/alembic/ddl/impl.py
@@ -371,7 +371,13 @@ def create_table(self, table: "Table") -> None:
                 self.create_column_comment(column)
 
     def drop_table(self, table: "Table") -> None:
+        table.dispatch.before_drop(
+            table, self.connection, checkfirst=False, _ddl_runner=self
+        )
         self._exec(schema.DropTable(table))
+        table.dispatch.after_drop(
+            table, self.connection, checkfirst=False, _ddl_runner=self
+        )
 
     def create_index(self, index: "Index") -> None:
         self._exec(schema.CreateIndex(index))
diff --git a/tests/test_op.py b/tests/test_op.py
index 677b7bc0..7ed7e45a 100644
--- a/tests/test_op.py
+++ b/tests/test_op.py
@@ -3,6 +3,7 @@
 from sqlalchemy import Boolean
 from sqlalchemy import CheckConstraint
 from sqlalchemy import Column
+from sqlalchemy import event
 from sqlalchemy import exc
 from sqlalchemy import ForeignKey
 from sqlalchemy import Index
@@ -1087,6 +1088,54 @@ def test_drop_table_comment_op(self):
 
         context.assert_("COMMENT ON TABLE some_table IS NULL")
 
+    def test_create_table_event(self):
+        context = op_fixture()
+
+        events_triggered = []
+
+        TestTable = Table("tb_test", MetaData(), Column("c1", Integer, nullable=False))
+
+        @event.listens_for(Table, "before_create")
+        def record_before_event(table, conn, **kwargs):
+            events_triggered.append(("before_create", table.name))
+
+        @event.listens_for(Table, "after_create")
+        def record_after_event(table, conn, **kwargs):
+            events_triggered.append(("after_create", table.name))
+
+        op.create_table(TestTable)
+        op.drop_table(TestTable)
+        context.assert_("CREATE TABLE tb_test ()", "DROP TABLE tb_test")
+
+        assert events_triggered == [
+            ("before_create", "tb_test"),
+            ("after_create", "tb_test"),
+        ]
+
+    def test_drop_table_event(self):
+        context = op_fixture()
+
+        events_triggered = []
+
+        TestTable = Table("tb_test", MetaData(), Column("c1", Integer, nullable=False))
+
+        @event.listens_for(Table, "before_drop")
+        def record_before_event(table, conn, **kwargs):
+            events_triggered.append(("before_drop", table.name))
+
+        @event.listens_for(Table, "after_drop")
+        def record_after_event(table, conn, **kwargs):
+            events_triggered.append(("after_drop", table.name))
+
+        op.create_table(TestTable)
+        op.drop_table(TestTable)
+        context.assert_("CREATE TABLE tb_test ()", "DROP TABLE tb_test")
+
+        assert events_triggered == [
+            ("before_drop", "tb_test"),
+            ("after_drop", "tb_test"),
+        ]
+
 
 class SQLModeOpTest(TestBase):
     def test_auto_literals(self):