Merge pull request #1632 from datajoint/fix/postgres-shared-enum-drop

fix(drop): drop a PostgreSQL enum type only when no column uses it

Author: dimitri-yatsenko

src/datajoint/adapters/postgres.py

@@ -1591,10 +1591,21 @@ def drop_enum_types_for_table(self, schema_name: str, table_name: str) -> list[s
         # to get actual enum types, then call this with results
         return []  # Placeholder - actual implementation requires query execution
 
+    def _quote_enum_type_name(self, enum_type_name: str) -> str:
+        """Quote a ``schema.typename`` enum type name, each part separately."""
+        parts = enum_type_name.split(".")
+        if len(parts) == 2:
+            return f"{self.quote_identifier(parts[0])}.{self.quote_identifier(parts[1])}"
+        return self.quote_identifier(enum_type_name)
+
     def drop_enum_type_ddl(self, enum_type_name: str) -> str:
         """
         Generate DROP TYPE IF EXISTS statement for a PostgreSQL enum.
 
+        The statement has no CASCADE: enum types are named by their values and
+        shared by every column in the schema with the same values, and CASCADE
+        would drop those columns too. A type still in use is not dropped.
+
         Parameters
         ----------
         enum_type_name : str
@@ -1603,12 +1614,26 @@ def drop_enum_type_ddl(self, enum_type_name: str) -> str:
         Returns
         -------
         str
-            DROP TYPE IF EXISTS statement with CASCADE.
+            DROP TYPE IF EXISTS statement.
         """
-        # Split schema.typename and quote each part
-        parts = enum_type_name.split(".")
-        if len(parts) == 2:
-            qualified_name = f"{self.quote_identifier(parts[0])}.{self.quote_identifier(parts[1])}"
-        else:
-            qualified_name = self.quote_identifier(enum_type_name)
-        return f"DROP TYPE IF EXISTS {qualified_name} CASCADE"
+        return f"DROP TYPE IF EXISTS {self._quote_enum_type_name(enum_type_name)}"
+
+    def enum_type_in_use_sql(self, enum_type_name: str) -> str:
+        """
+        Query that returns a row if any table column uses a PostgreSQL enum type.
+
+        Parameters
+        ----------
+        enum_type_name : str
+            Fully qualified enum type name (schema.typename).
+
+        Returns
+        -------
+        str
+            SQL query returning at most one row.
+        """
+        return (
+            "SELECT 1 FROM pg_catalog.pg_attribute "
+            f"WHERE atttypid = {self.quote_string(self._quote_enum_type_name(enum_type_name))}::regtype "
+            "AND attnum > 0 AND NOT attisdropped LIMIT 1"
+        )

src/datajoint/table.py

@@ -1267,15 +1267,20 @@ def drop_quick(self):
             self.connection.query(query)
             logger.info("Dropped table %s" % self.full_table_name)
 
-            # For PostgreSQL, clean up enum types after dropping the table
+            # For PostgreSQL, clean up enum types after dropping the table. An enum
+            # type is named by its values and shared by every column in the schema
+            # with the same values (every jobs table's status, for one), so it is
+            # dropped only when no column uses it any more.
             if enum_types_to_drop and hasattr(adapter, "drop_enum_type_ddl"):
                 for enum_type in enum_types_to_drop:
                     try:
+                        if self.connection.query(adapter.enum_type_in_use_sql(enum_type)).fetchone():
+                            continue
                         drop_ddl = adapter.drop_enum_type_ddl(enum_type)
                         self.connection.query(drop_ddl)
                         logger.debug("Dropped enum type %s" % enum_type)
                     except Exception:
-                        pass  # Ignore errors - type may be used by other tables
+                        pass  # Ignore errors - enum cleanup is best-effort
         else:
             logger.info("Nothing to drop: table %s is not declared" % self.full_table_name)
 

tests/integration/test_drop_shared_enum.py

@@ -0,0 +1,121 @@
+"""
+Dropping a table must not drop columns of other tables.
+
+On PostgreSQL, an enum attribute is declared with an enum type named by its
+values, so every column in the schema with the same values shares one type.
+Dropping a table drops its enum types only when no other column uses them.
+"""
+
+import time
+
+import pytest
+
+import datajoint as dj
+
+
+@pytest.fixture(scope="function")
+def schema_by_backend(connection_by_backend, db_creds_by_backend):
+    backend = db_creds_by_backend["backend"]
+    schema_name = f"djtest_shared_enum_{backend}_{str(int(time.time() * 1000))[-8:]}"
+    schema = dj.Schema(schema_name, connection=connection_by_backend)
+    yield schema
+    if connection_by_backend.is_connected:
+        connection_by_backend.query(
+            f"DROP DATABASE IF EXISTS {connection_by_backend.adapter.quote_identifier(schema_name)}"
+            if backend == "mysql"
+            else f"DROP SCHEMA IF EXISTS {connection_by_backend.adapter.quote_identifier(schema_name)} CASCADE"
+        )
+
+
+def _columns(schema, table_name):
+    rows = schema.connection.query(
+        "SELECT column_name FROM information_schema.columns "
+        "WHERE table_schema = %s AND table_name = %s ORDER BY ordinal_position",
+        args=(schema.database, table_name),
+    ).fetchall()
+    return [row[0] for row in rows]
+
+
+def _enum_types(schema):
+    if schema.connection.adapter.backend != "postgresql":
+        return []
+    rows = schema.connection.query(
+        "SELECT t.typname FROM pg_catalog.pg_type t "
+        "JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace "
+        "WHERE t.typtype = 'e' AND n.nspname = %s",
+        args=(schema.database,),
+    ).fetchall()
+    return [row[0] for row in rows]
+
+
+def test_drop_keeps_a_shared_enum_column(schema_by_backend):
+    """Two tables declare the same enum. Dropping one keeps the other's column
+    and its values; dropping the second removes the type."""
+
+    @schema_by_backend
+    class Animal(dj.Manual):
+        definition = """
+        animal_id : int32
+        ---
+        species : enum('mouse', 'rat')
+        """
+
+    @schema_by_backend
+    class Cage(dj.Manual):
+        definition = """
+        cage_id : int32
+        ---
+        species : enum('mouse', 'rat')
+        """
+
+    Animal.insert([(1, "mouse"), (2, "rat")])
+    Cage.insert([(1, "rat")])
+    types_before = _enum_types(schema_by_backend)
+
+    Cage().drop_quick()
+    assert _columns(schema_by_backend, "animal") == ["animal_id", "species"]
+    assert Animal.to_dicts(order_by="animal_id") == [
+        {"animal_id": 1, "species": "mouse"},
+        {"animal_id": 2, "species": "rat"},
+    ]
+    assert _enum_types(schema_by_backend) == types_before
+
+    Animal().drop_quick()
+    assert _enum_types(schema_by_backend) == []
+
+
+def test_dropping_a_jobs_table_keeps_other_jobs_tables_whole(schema_by_backend):
+    """Every jobs table's status column has the same enum. Dropping one jobs
+    table leaves the others' status columns, and their jobs, in place."""
+
+    @schema_by_backend
+    class Item(dj.Manual):
+        definition = """
+        item : int32
+        """
+
+    @schema_by_backend
+    class First(dj.Computed):
+        definition = """
+        -> Item
+        """
+
+        def make(self, key):
+            self.insert1(key)
+
+    @schema_by_backend
+    class Second(dj.Computed):
+        definition = """
+        -> Item
+        """
+
+        def make(self, key):
+            raise RuntimeError("not ready")
+
+    Item.insert([(1,), (2,)])
+    First.populate(reserve_jobs=True)
+    Second.populate(reserve_jobs=True, suppress_errors=True)
+
+    First.jobs.drop_quick()
+    assert "status" in _columns(schema_by_backend, "~~second")
+    assert sorted(Second.jobs.to_arrays("status").tolist()) == ["error", "error"]

tests/unit/test_adapters.py

@@ -352,6 +352,18 @@ def test_enum_type_sql(self, adapter):
         assert "'pending'" in sql
         assert "'complete'" in sql
 
+    def test_drop_enum_type_ddl_does_not_cascade(self, adapter):
+        """An enum type shared by other columns must not take them with it."""
+        sql = adapter.drop_enum_type_ddl("myschema.enum_abc")
+        assert sql == 'DROP TYPE IF EXISTS "myschema"."enum_abc"'
+
+    def test_enum_type_in_use_sql(self, adapter):
+        """The usage check looks the type up by its quoted, qualified name."""
+        sql = adapter.enum_type_in_use_sql("myschema.enum_abc")
+        assert "pg_attribute" in sql
+        assert """'"myschema"."enum_abc"'::regtype""" in sql
+        assert "NOT attisdropped" in sql
+
     def test_drop_enum_type_sql(self, adapter):
         """Test PostgreSQL enum type dropping."""
         sql = adapter.drop_enum_type_sql("myschema", "mytable", "status")