enable extensions to use AutogenContext.run_name_filters Fixes: sqlal…

…chemy#813
olirice · Mar 5, 2021 · 31dccd2 · 31dccd2
1 parent c3fc47d
commit 31dccd2
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 9 deletions.
diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
@@ -330,15 +330,16 @@ def run_name_filters(self, name, type_, parent_names):
             if type_ == "table":
                 table_name = name
             else:
-                table_name = parent_names["table_name"]
-            schema_name = parent_names["schema_name"]
-            if schema_name:
-                parent_names["schema_qualified_table_name"] = "%s.%s" % (
-                    schema_name,
-                    table_name,
-                )
-            else:
-                parent_names["schema_qualified_table_name"] = table_name
+                table_name = parent_names.get("table_name", None)
+            if table_name:
+                schema_name = parent_names["schema_name"]
+                if schema_name:
+                    parent_names["schema_qualified_table_name"] = "%s.%s" % (
+                        schema_name,
+                        table_name,
+                    )
+                else:
+                    parent_names["schema_qualified_table_name"] = table_name
 
         for fn in self._name_filters:
 

diff --git a/tests/test_autogen_diffs.py b/tests/test_autogen_diffs.py
@@ -198,6 +198,36 @@ def include_name(name, type_, parent_names):
             },
         )
 
+    def test_run_name_filters_supports_extension_types(self):
+        def include_name(name, type_, parent_names):
+            # extension type successfully passed through
+            # `AutogenContext.run_name_filters`
+            assert True
+
+        self._update_context(name_filters=include_name, include_schemas=True)
+
+        self.autogen_context.run_name_filters(
+            name="some_function",
+            type_="function",
+            parent_names={"schema_name": "public"},
+        )
+
+    def test_run_object_filters_supports_extension_types(self):
+        def include_object(obj, name, type_, reflected, compare_to):
+            # extension type successfully passed through
+            # `AutogenContext.run_object_filters`
+            assert True
+
+        self._update_context(
+            object_filters=include_object, include_schemas=True
+        )
+
+        self.autogen_context.run_name_filters(
+            name="some_function",
+            type_="function",
+            parent_names={"schema_name": "public"},
+        )
+
     def test_default_schema_omitted_downgrade(self):
         def include_object(obj, name, type_, reflected, compare_to):
             if type_ == "table":