From 7e9cae956fd7538a54d990841e515fd782293648 Mon Sep 17 00:00:00 2001 From: Dimitri Yatsenko Date: Sat, 17 Oct 2026 19:22:43 -0500 Subject: [PATCH] fix(drop): drop a PostgreSQL enum type only when no column uses it On PostgreSQL, dropping a table could drop columns of other tables, with their data. 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. `drop_quick()` dropped the table's enum types afterwards with `DROP TYPE ... CASCADE`, and CASCADE also dropped every other column of that type. Two tables that both declare `species : enum('mouse', 'rat')` share a type, and dropping one took the other's `species` column. Every jobs table has the same `status` enum, so dropping one jobs table removed `status` from all the others in the schema. MySQL has no enum types and was not affected. An enum type is now dropped only when no column uses it any more, and the statement no longer cascades, so a type still in use is never dropped with its columns. The new tests, on both backends, drop one of two tables that share an enum and check that the other keeps its column and values, that the type is removed once neither table remains, and that dropping one jobs table leaves another's status column and jobs in place. Both fail on PostgreSQL before this change. --- src/datajoint/adapters/postgres.py | 41 +++++-- src/datajoint/table.py | 9 +- tests/integration/test_drop_shared_enum.py | 121 +++++++++++++++++++++ tests/unit/test_adapters.py | 12 ++ 4 files changed, 173 insertions(+), 10 deletions(-) create mode 100644 tests/integration/test_drop_shared_enum.py diff --git a/src/datajoint/adapters/postgres.py b/src/datajoint/adapters/postgres.py index 3a55d8eb8..70d37a2da 100644 --- a/src/datajoint/adapters/postgres.py +++ b/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" + ) diff --git a/src/datajoint/table.py b/src/datajoint/table.py index 9f6784fb3..b69b6678e 100644 --- a/src/datajoint/table.py +++ b/src/datajoint/table.py @@ -1247,15 +1247,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) diff --git a/tests/integration/test_drop_shared_enum.py b/tests/integration/test_drop_shared_enum.py new file mode 100644 index 000000000..979f7f0b9 --- /dev/null +++ b/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"] diff --git a/tests/unit/test_adapters.py b/tests/unit/test_adapters.py index bd0afa9e0..b5db84318 100644 --- a/tests/unit/test_adapters.py +++ b/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")