Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 33 additions & 8 deletions src/datajoint/adapters/postgres.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"
)
9 changes: 7 additions & 2 deletions src/datajoint/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
121 changes: 121 additions & 0 deletions tests/integration/test_drop_shared_enum.py
Original file line number Diff line number Diff line change
@@ -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"]
12 changes: 12 additions & 0 deletions tests/unit/test_adapters.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down