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")