| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -1591,10 +1591,21 @@ def drop_enum_types_for_table(self, schema_name: str, table_name: str) -> list[s | |||
| 1591 | 1591 | # to get actual enum types, then call this with results | |
| 1592 | 1592 | return [] # Placeholder - actual implementation requires query execution | |
| 1593 | 1593 | ||
| 1594 | + def _quote_enum_type_name(self, enum_type_name: str) -> str: | ||
| 1595 | + """Quote a ``schema.typename`` enum type name, each part separately.""" | ||
| 1596 | + parts = enum_type_name.split(".") | ||
| 1597 | + if len(parts) == 2: | ||
| 1598 | + return f"{self.quote_identifier(parts[0])}.{self.quote_identifier(parts[1])}" | ||
| 1599 | + return self.quote_identifier(enum_type_name) | ||
| 1600 | + | ||
| 1594 | 1601 | def drop_enum_type_ddl(self, enum_type_name: str) -> str: | |
| 1595 | 1602 | """ | |
| 1596 | 1603 | Generate DROP TYPE IF EXISTS statement for a PostgreSQL enum. | |
| 1597 | 1604 | ||
| 1605 | + The statement has no CASCADE: enum types are named by their values and | ||
| 1606 | + shared by every column in the schema with the same values, and CASCADE | ||
| 1607 | + would drop those columns too. A type still in use is not dropped. | ||
| 1608 | + | ||
| 1598 | 1609 | Parameters | |
| 1599 | 1610 | ---------- | |
| 1600 | 1611 | enum_type_name : str | |
@@ -1603,12 +1614,26 @@ def drop_enum_type_ddl(self, enum_type_name: str) -> str: | |||
| 1603 | 1614 | Returns | |
| 1604 | 1615 | ------- | |
| 1605 | 1616 | str | |
| 1606 | - DROP TYPE IF EXISTS statement with CASCADE. | ||
| 1617 | + DROP TYPE IF EXISTS statement. | ||
| 1607 | 1618 | """ | |
| 1608 | - # Split schema.typename and quote each part | ||
| 1609 | - parts = enum_type_name.split(".") | ||
| 1610 | - if len(parts) == 2: | ||
| 1611 | - qualified_name = f"{self.quote_identifier(parts[0])}.{self.quote_identifier(parts[1])}" | ||
| 1612 | - else: | ||
| 1613 | - qualified_name = self.quote_identifier(enum_type_name) | ||
| 1614 | - return f"DROP TYPE IF EXISTS {qualified_name} CASCADE" | ||
| 1619 | + return f"DROP TYPE IF EXISTS {self._quote_enum_type_name(enum_type_name)}" | ||
| 1620 | + | ||
| 1621 | + def enum_type_in_use_sql(self, enum_type_name: str) -> str: | ||
| 1622 | + """ | ||
| 1623 | + Query that returns a row if any table column uses a PostgreSQL enum type. | ||
| 1624 | + | ||
| 1625 | + Parameters | ||
| 1626 | + ---------- | ||
| 1627 | + enum_type_name : str | ||
| 1628 | + Fully qualified enum type name (schema.typename). | ||
| 1629 | + | ||
| 1630 | + Returns | ||
| 1631 | + ------- | ||
| 1632 | + str | ||
| 1633 | + SQL query returning at most one row. | ||
| 1634 | + """ | ||
| 1635 | + return ( | ||
| 1636 | + "SELECT 1 FROM pg_catalog.pg_attribute " | ||
| 1637 | + f"WHERE atttypid = {self.quote_string(self._quote_enum_type_name(enum_type_name))}::regtype " | ||
| 1638 | + "AND attnum > 0 AND NOT attisdropped LIMIT 1" | ||
| 1639 | + ) | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -1267,15 +1267,20 @@ def drop_quick(self): | |||
| 1267 | 1267 | self.connection.query(query) | |
| 1268 | 1268 | logger.info("Dropped table %s" % self.full_table_name) | |
| 1269 | 1269 | ||
| 1270 | - # For PostgreSQL, clean up enum types after dropping the table | ||
| 1270 | + # For PostgreSQL, clean up enum types after dropping the table. An enum | ||
| 1271 | + # type is named by its values and shared by every column in the schema | ||
| 1272 | + # with the same values (every jobs table's status, for one), so it is | ||
| 1273 | + # dropped only when no column uses it any more. | ||
| 1271 | 1274 | if enum_types_to_drop and hasattr(adapter, "drop_enum_type_ddl"): | |
| 1272 | 1275 | for enum_type in enum_types_to_drop: | |
| 1273 | 1276 | try: | |
| 1277 | + if self.connection.query(adapter.enum_type_in_use_sql(enum_type)).fetchone(): | ||
| 1278 | + continue | ||
| 1274 | 1279 | drop_ddl = adapter.drop_enum_type_ddl(enum_type) | |
| 1275 | 1280 | self.connection.query(drop_ddl) | |
| 1276 | 1281 | logger.debug("Dropped enum type %s" % enum_type) | |
| 1277 | 1282 | except Exception: | |
| 1278 | - pass # Ignore errors - type may be used by other tables | ||
| 1283 | + pass # Ignore errors - enum cleanup is best-effort | ||
| 1279 | 1284 | else: | |
| 1280 | 1285 | logger.info("Nothing to drop: table %s is not declared" % self.full_table_name) | |
| 1281 | 1286 | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -0,0 +1,121 @@ | |||
| 1 | + """ | ||
| 2 | + Dropping a table must not drop columns of other tables. | ||
| 3 | + | ||
| 4 | + On PostgreSQL, an enum attribute is declared with an enum type named by its | ||
| 5 | + values, so every column in the schema with the same values shares one type. | ||
| 6 | + Dropping a table drops its enum types only when no other column uses them. | ||
| 7 | + """ | ||
| 8 | + | ||
| 9 | + import time | ||
| 10 | + | ||
| 11 | + import pytest | ||
| 12 | + | ||
| 13 | + import datajoint as dj | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + @pytest.fixture(scope="function") | ||
| 17 | + def schema_by_backend(connection_by_backend, db_creds_by_backend): | ||
| 18 | + backend = db_creds_by_backend["backend"] | ||
| 19 | + schema_name = f"djtest_shared_enum_{backend}_{str(int(time.time() * 1000))[-8:]}" | ||
| 20 | + schema = dj.Schema(schema_name, connection=connection_by_backend) | ||
| 21 | + yield schema | ||
| 22 | + if connection_by_backend.is_connected: | ||
| 23 | + connection_by_backend.query( | ||
| 24 | + f"DROP DATABASE IF EXISTS {connection_by_backend.adapter.quote_identifier(schema_name)}" | ||
| 25 | + if backend == "mysql" | ||
| 26 | + else f"DROP SCHEMA IF EXISTS {connection_by_backend.adapter.quote_identifier(schema_name)} CASCADE" | ||
| 27 | + ) | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + def _columns(schema, table_name): | ||
| 31 | + rows = schema.connection.query( | ||
| 32 | + "SELECT column_name FROM information_schema.columns " | ||
| 33 | + "WHERE table_schema = %s AND table_name = %s ORDER BY ordinal_position", | ||
| 34 | + args=(schema.database, table_name), | ||
| 35 | + ).fetchall() | ||
| 36 | + return [row[0] for row in rows] | ||
| 37 | + | ||
| 38 | + | ||
| 39 | + def _enum_types(schema): | ||
| 40 | + if schema.connection.adapter.backend != "postgresql": | ||
| 41 | + return [] | ||
| 42 | + rows = schema.connection.query( | ||
| 43 | + "SELECT t.typname FROM pg_catalog.pg_type t " | ||
| 44 | + "JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace " | ||
| 45 | + "WHERE t.typtype = 'e' AND n.nspname = %s", | ||
| 46 | + args=(schema.database,), | ||
| 47 | + ).fetchall() | ||
| 48 | + return [row[0] for row in rows] | ||
| 49 | + | ||
| 50 | + | ||
| 51 | + def test_drop_keeps_a_shared_enum_column(schema_by_backend): | ||
| 52 | + """Two tables declare the same enum. Dropping one keeps the other's column | ||
| 53 | + and its values; dropping the second removes the type.""" | ||
| 54 | + | ||
| 55 | + @schema_by_backend | ||
| 56 | + class Animal(dj.Manual): | ||
| 57 | + definition = """ | ||
| 58 | + animal_id : int32 | ||
| 59 | + --- | ||
| 60 | + species : enum('mouse', 'rat') | ||
| 61 | + """ | ||
| 62 | + | ||
| 63 | + @schema_by_backend | ||
| 64 | + class Cage(dj.Manual): | ||
| 65 | + definition = """ | ||
| 66 | + cage_id : int32 | ||
| 67 | + --- | ||
| 68 | + species : enum('mouse', 'rat') | ||
| 69 | + """ | ||
| 70 | + | ||
| 71 | + Animal.insert([(1, "mouse"), (2, "rat")]) | ||
| 72 | + Cage.insert([(1, "rat")]) | ||
| 73 | + types_before = _enum_types(schema_by_backend) | ||
| 74 | + | ||
| 75 | + Cage().drop_quick() | ||
| 76 | + assert _columns(schema_by_backend, "animal") == ["animal_id", "species"] | ||
| 77 | + assert Animal.to_dicts(order_by="animal_id") == [ | ||
| 78 | + {"animal_id": 1, "species": "mouse"}, | ||
| 79 | + {"animal_id": 2, "species": "rat"}, | ||
| 80 | + ] | ||
| 81 | + assert _enum_types(schema_by_backend) == types_before | ||
| 82 | + | ||
| 83 | + Animal().drop_quick() | ||
| 84 | + assert _enum_types(schema_by_backend) == [] | ||
| 85 | + | ||
| 86 | + | ||
| 87 | + def test_dropping_a_jobs_table_keeps_other_jobs_tables_whole(schema_by_backend): | ||
| 88 | + """Every jobs table's status column has the same enum. Dropping one jobs | ||
| 89 | + table leaves the others' status columns, and their jobs, in place.""" | ||
| 90 | + | ||
| 91 | + @schema_by_backend | ||
| 92 | + class Item(dj.Manual): | ||
| 93 | + definition = """ | ||
| 94 | + item : int32 | ||
| 95 | + """ | ||
| 96 | + | ||
| 97 | + @schema_by_backend | ||
| 98 | + class First(dj.Computed): | ||
| 99 | + definition = """ | ||
| 100 | + -> Item | ||
| 101 | + """ | ||
| 102 | + | ||
| 103 | + def make(self, key): | ||
| 104 | + self.insert1(key) | ||
| 105 | + | ||
| 106 | + @schema_by_backend | ||
| 107 | + class Second(dj.Computed): | ||
| 108 | + definition = """ | ||
| 109 | + -> Item | ||
| 110 | + """ | ||
| 111 | + | ||
| 112 | + def make(self, key): | ||
| 113 | + raise RuntimeError("not ready") | ||
| 114 | + | ||
| 115 | + Item.insert([(1,), (2,)]) | ||
| 116 | + First.populate(reserve_jobs=True) | ||
| 117 | + Second.populate(reserve_jobs=True, suppress_errors=True) | ||
| 118 | + | ||
| 119 | + First.jobs.drop_quick() | ||
| 120 | + assert "status" in _columns(schema_by_backend, "~~second") | ||
| 121 | + assert sorted(Second.jobs.to_arrays("status").tolist()) == ["error", "error"] | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -352,6 +352,18 @@ def test_enum_type_sql(self, adapter): | |||
| 352 | 352 | assert "'pending'" in sql | |
| 353 | 353 | assert "'complete'" in sql | |
| 354 | 354 | ||
| 355 | + def test_drop_enum_type_ddl_does_not_cascade(self, adapter): | ||
| 356 | + """An enum type shared by other columns must not take them with it.""" | ||
| 357 | + sql = adapter.drop_enum_type_ddl("myschema.enum_abc") | ||
| 358 | + assert sql == 'DROP TYPE IF EXISTS "myschema"."enum_abc"' | ||
| 359 | + | ||
| 360 | + def test_enum_type_in_use_sql(self, adapter): | ||
| 361 | + """The usage check looks the type up by its quoted, qualified name.""" | ||
| 362 | + sql = adapter.enum_type_in_use_sql("myschema.enum_abc") | ||
| 363 | + assert "pg_attribute" in sql | ||
| 364 | + assert """'"myschema"."enum_abc"'::regtype""" in sql | ||
| 365 | + assert "NOT attisdropped" in sql | ||
| 366 | + | ||
| 355 | 367 | def test_drop_enum_type_sql(self, adapter): | |
| 356 | 368 | """Test PostgreSQL enum type dropping.""" | |
| 357 | 369 | sql = adapter.drop_enum_type_sql("myschema", "mytable", "status") | |
| Back | FazBrowse Home | New Git URL |
0 commit comments