From 0b2f7dbcf57623e2a66c908f11704e04dab9ac68 Mon Sep 17 00:00:00 2001 From: theluckystrike <51033404+theluckystrike@users.noreply.github.com> Date: Mon, 9 Nov 2026 01:58:42 +0100 Subject: [PATCH 1/3] fix: rebuild lineage on PostgreSQL and keep entries on failure --- src/datajoint/lineage.py | 43 ++++++------- tests/integration/test_multi_backend.py | 82 +++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 24 deletions(-) diff --git a/src/datajoint/lineage.py b/src/datajoint/lineage.py index a7d8a272c..697298503 100644 --- a/src/datajoint/lineage.py +++ b/src/datajoint/lineage.py @@ -265,7 +265,7 @@ def rebuild_schema_lineage(connection, database): This function assumes that any upstream schemas (referenced via cross-schema foreign keys) have already had their lineage rebuilt. If a referenced attribute in another schema has no lineage entry, - a DataJointError is raised. + a DataJointError is raised and the existing ~lineage entries are kept. Parameters ---------- @@ -285,9 +285,6 @@ def rebuild_schema_lineage(connection, database): adapter = connection.adapter lineage_table = f"{adapter.quote_identifier(database)}.{adapter.quote_identifier('~lineage')}" - # Clear all existing lineage entries for this schema - connection.query(f"DELETE FROM {lineage_table}") - # Get all tables in the schema (excluding hidden tables) tables_result = connection.query( """ @@ -298,35 +295,30 @@ def rebuild_schema_lineage(connection, database): ).fetchall() all_tables = {row[0] for row in tables_result} - if not all_tables: - return + # Read keys through the adapter, as Dependencies.load does: the catalog differs by backend + # Note: Both PyMySQL and psycopg use %s placeholders, so escape % as %% + schemas_list = adapter.quote_string(database) + like_pattern = "'~%%'" # Get all primary key columns for all tables - pk_result = connection.query( - """ - SELECT TABLE_NAME, COLUMN_NAME FROM information_schema.KEY_COLUMN_USAGE - WHERE TABLE_SCHEMA = %s AND CONSTRAINT_NAME = 'PRIMARY' - """, - args=(database,), - ).fetchall() # table -> set of PK columns pk_columns = {} - for table, col in pk_result: + for full_name, col in connection.query(adapter.load_primary_keys_sql(schemas_list, like_pattern)): + _, table = adapter.split_full_table_name(full_name) pk_columns.setdefault(table, set()).add(col) # Get all FK relationships within and across schemas - fk_result = connection.query( - """ - SELECT TABLE_NAME, COLUMN_NAME, - REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME - FROM information_schema.KEY_COLUMN_USAGE - WHERE TABLE_SCHEMA = %s AND REFERENCED_TABLE_NAME IS NOT NULL - """, - args=(database,), - ).fetchall() + fk_result = connection.query(adapter.load_foreign_keys_sql(schemas_list, like_pattern), as_dict=True) # Build FK map: (table, column) -> (parent_schema, parent_table, parent_column) - fk_map = {(table, col): (ref_schema, ref_table, ref_col) for table, col, ref_schema, ref_table, ref_col in fk_result} + fk_map = {} + for row in fk_result: + key = {k.lower(): v for k, v in row.items()} + child_schema, child_table = adapter.split_full_table_name(key["referencing_table"]) + # The adapter query also returns foreign keys that other schemas hold on this one + if child_schema == database: + ref_schema, ref_table = adapter.split_full_table_name(key["referenced_table"]) + fk_map[(child_table, key["column_name"])] = (ref_schema, ref_table, key["referenced_column_name"]) # Lineage cache: (table, column) -> lineage string (for this schema) lineage_cache = {} @@ -368,6 +360,9 @@ def resolve_lineage(table, col): if not col.startswith("_"): resolve_lineage(table, col) + # Replace the existing entries only now that every lineage has resolved + connection.query(f"DELETE FROM {lineage_table}") + # Insert all lineages in one batch if lineage_cache: entries = [(table, col, lineage) for (table, col), lineage in lineage_cache.items()] diff --git a/tests/integration/test_multi_backend.py b/tests/integration/test_multi_backend.py index 707b05d20..cf6ee5df9 100644 --- a/tests/integration/test_multi_backend.py +++ b/tests/integration/test_multi_backend.py @@ -86,6 +86,88 @@ class Observation(dj.Manual): schema.drop() +def _lineage_table(schema): + adapter = schema.connection.adapter + return f"{adapter.quote_identifier(schema.database)}.{adapter.quote_identifier('~lineage')}" + + +@pytest.mark.backend_agnostic +def test_rebuild_lineage(connection_by_backend, backend, prefix): + """Test that rebuild_lineage restores lineage from the keys on both backends.""" + schema = dj.Schema( + f"{prefix}_multi_backend_{backend}_rebuild", + connection=connection_by_backend, + ) + + @schema + class Animal(dj.Manual): + definition = """ + animal_id : int + """ + + @schema + class Observation(dj.Manual): + definition = """ + -> Animal + obs_id : int + """ + + connection_by_backend.query(f"DELETE FROM {_lineage_table(schema)}") + assert schema.lineage == {} + + schema.rebuild_lineage() + + db = schema.database + assert schema.lineage == { + f"{db}.animal.animal_id": f"{db}.animal.animal_id", + f"{db}.observation.animal_id": f"{db}.animal.animal_id", + f"{db}.observation.obs_id": f"{db}.observation.obs_id", + } + + # Cleanup + schema.drop() + + +@pytest.mark.backend_agnostic +def test_failed_rebuild_lineage_keeps_entries(connection_by_backend, backend, prefix): + """Test that a rebuild that cannot resolve a key leaves ~lineage unchanged.""" + upstream = dj.Schema( + f"{prefix}_multi_backend_{backend}_rebuild_up", + connection=connection_by_backend, + ) + downstream = dj.Schema( + f"{prefix}_multi_backend_{backend}_rebuild_down", + connection=connection_by_backend, + ) + + @upstream + class Animal(dj.Manual): + definition = """ + animal_id : int + """ + + @downstream + class Observation(dj.Manual): + definition = """ + -> Animal + obs_id : int + """ + + before = downstream.lineage + assert before + + connection_by_backend.query(f"DELETE FROM {_lineage_table(upstream)}") + + with pytest.raises(dj.errors.DataJointError, match="has no lineage"): + downstream.rebuild_lineage() + + assert downstream.lineage == before + + # Cleanup + downstream.drop() + upstream.drop() + + @pytest.mark.backend_agnostic def test_data_types(connection_by_backend, backend, prefix): """Test that core data types work on both backends.""" From e9d2686799b77220dfd59a49b91ea2d9cf60b4d3 Mon Sep 17 00:00:00 2001 From: theluckystrike <51033404+theluckystrike@users.noreply.github.com> Date: Tue, 10 Nov 2026 17:32:41 +0100 Subject: [PATCH 2/3] fix: replace ~lineage entries in one transaction rebuild_schema_lineage now runs the DELETE and the upsert together, so an insert that fails rolls the DELETE back. The upsert is split out of insert_lineages because ensure_lineage_table's CREATE TABLE would commit the transaction on MySQL. A rebuild inside a caller's open transaction reuses it. --- src/datajoint/lineage.py | 28 ++++++++---- tests/integration/test_multi_backend.py | 58 +++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 9 deletions(-) diff --git a/src/datajoint/lineage.py b/src/datajoint/lineage.py index 697298503..8f935de0f 100644 --- a/src/datajoint/lineage.py +++ b/src/datajoint/lineage.py @@ -20,6 +20,7 @@ declaration time. Use schema.rebuild_lineage() to restore lineage for legacy schemas. """ +import contextlib import logging from .errors import DataJointError @@ -207,12 +208,20 @@ def insert_lineages(connection, database, entries): adapter = connection.adapter lineage_table = f"{adapter.quote_identifier(database)}.{adapter.quote_identifier('~lineage')}" + _upsert_lineages(connection, lineage_table, entries) - # Build backend-agnostic upsert statement + +def _upsert_lineages(connection, lineage_table, entries): + """ + Upsert lineage entries into an existing ~lineage table. + + Unlike insert_lineages, this does not create the table, so it can run inside + a transaction (MySQL commits an open transaction before any CREATE TABLE). + """ columns = ["table_name", "attribute_name", "lineage"] primary_key = ["table_name", "attribute_name"] - sql = adapter.upsert_on_duplicate_sql( + sql = connection.adapter.upsert_on_duplicate_sql( lineage_table, columns, primary_key, @@ -360,10 +369,11 @@ def resolve_lineage(table, col): if not col.startswith("_"): resolve_lineage(table, col) - # Replace the existing entries only now that every lineage has resolved - connection.query(f"DELETE FROM {lineage_table}") - - # Insert all lineages in one batch - if lineage_cache: - entries = [(table, col, lineage) for (table, col), lineage in lineage_cache.items()] - insert_lineages(connection, database, entries) + # Replace the existing entries only now that every lineage has resolved, in one + # transaction so a failed insert rolls the DELETE back and keeps the old entries + entries = [(table, col, lineage) for (table, col), lineage in lineage_cache.items()] + # An open transaction already covers both statements; nesting one would raise + with contextlib.nullcontext() if connection.in_transaction else connection.transaction: + connection.query(f"DELETE FROM {lineage_table}") + if entries: + _upsert_lineages(connection, lineage_table, entries) diff --git a/tests/integration/test_multi_backend.py b/tests/integration/test_multi_backend.py index cf6ee5df9..2be8432dd 100644 --- a/tests/integration/test_multi_backend.py +++ b/tests/integration/test_multi_backend.py @@ -168,6 +168,64 @@ class Observation(dj.Manual): upstream.drop() +@pytest.mark.backend_agnostic +def test_failed_lineage_insert_keeps_entries(connection_by_backend, backend, prefix, monkeypatch): + """Test that a rebuild whose insert fails rolls back the DELETE and keeps ~lineage.""" + import datajoint.lineage + + schema = dj.Schema( + f"{prefix}_multi_backend_{backend}_rebuild_insert", + connection=connection_by_backend, + ) + + @schema + class Animal(dj.Manual): + definition = """ + animal_id : int + """ + + before = schema.lineage + assert before + + def failing_upsert(connection, lineage_table, entries): + raise dj.errors.LostConnectionError("connection dropped during the insert") + + monkeypatch.setattr(datajoint.lineage, "_upsert_lineages", failing_upsert) + + with pytest.raises(dj.errors.LostConnectionError): + schema.rebuild_lineage() + + monkeypatch.undo() + assert not connection_by_backend.in_transaction + assert schema.lineage == before + + # Cleanup + schema.drop() + + +@pytest.mark.backend_agnostic +def test_rebuild_lineage_in_open_transaction(connection_by_backend, backend, prefix): + """Test that rebuild_lineage works inside a transaction the caller already opened.""" + schema = dj.Schema( + f"{prefix}_multi_backend_{backend}_rebuild_txn", + connection=connection_by_backend, + ) + + @schema + class Animal(dj.Manual): + definition = """ + animal_id : int + """ + + before = schema.lineage + with connection_by_backend.transaction: + schema.rebuild_lineage() + assert schema.lineage == before + + # Cleanup + schema.drop() + + @pytest.mark.backend_agnostic def test_data_types(connection_by_backend, backend, prefix): """Test that core data types work on both backends.""" From d59d4ecce987f4e52cfefade0a6f49ac5b2d93c9 Mon Sep 17 00:00:00 2001 From: theluckystrike <51033404+theluckystrike@users.noreply.github.com> Date: Tue, 10 Nov 2026 22:17:18 +0100 Subject: [PATCH 3/3] fix: refuse to rebuild lineage inside a transaction CREATE TABLE IF NOT EXISTS for ~lineage commits an open transaction on MySQL, so a rebuild inside the caller's transaction committed the caller's earlier writes and left the DELETE and upsert unprotected. Raise before any DDL, as declare and alter do. Signed-off-by: theluckystrike <51033404+theluckystrike@users.noreply.github.com> --- src/datajoint/lineage.py | 11 +++++++---- tests/integration/test_multi_backend.py | 15 ++++++++++++--- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/src/datajoint/lineage.py b/src/datajoint/lineage.py index 8f935de0f..e2d95bf37 100644 --- a/src/datajoint/lineage.py +++ b/src/datajoint/lineage.py @@ -20,7 +20,6 @@ declaration time. Use schema.rebuild_lineage() to restore lineage for legacy schemas. """ -import contextlib import logging from .errors import DataJointError @@ -286,8 +285,13 @@ def rebuild_schema_lineage(connection, database): Raises ------ DataJointError - If a referenced attribute in another schema has no lineage entry. + If a referenced attribute in another schema has no lineage entry, + or if called inside a transaction. """ + # CREATE TABLE IF NOT EXISTS below commits an open transaction on MySQL, as declare does + if connection.in_transaction: + raise DataJointError("Cannot rebuild lineage inside a transaction.") + # Ensure the lineage table exists ensure_lineage_table(connection, database) @@ -372,8 +376,7 @@ def resolve_lineage(table, col): # Replace the existing entries only now that every lineage has resolved, in one # transaction so a failed insert rolls the DELETE back and keeps the old entries entries = [(table, col, lineage) for (table, col), lineage in lineage_cache.items()] - # An open transaction already covers both statements; nesting one would raise - with contextlib.nullcontext() if connection.in_transaction else connection.transaction: + with connection.transaction: connection.query(f"DELETE FROM {lineage_table}") if entries: _upsert_lineages(connection, lineage_table, entries) diff --git a/tests/integration/test_multi_backend.py b/tests/integration/test_multi_backend.py index 2be8432dd..dcab28b1d 100644 --- a/tests/integration/test_multi_backend.py +++ b/tests/integration/test_multi_backend.py @@ -205,7 +205,11 @@ def failing_upsert(connection, lineage_table, entries): @pytest.mark.backend_agnostic def test_rebuild_lineage_in_open_transaction(connection_by_backend, backend, prefix): - """Test that rebuild_lineage works inside a transaction the caller already opened.""" + """Test that rebuild_lineage refuses to run inside a transaction and leaves it intact. + + On MySQL the CREATE TABLE IF NOT EXISTS for ~lineage would commit the caller's + transaction, so the rebuild raises before it, as declare does. + """ schema = dj.Schema( f"{prefix}_multi_backend_{backend}_rebuild_txn", connection=connection_by_backend, @@ -218,8 +222,13 @@ class Animal(dj.Manual): """ before = schema.lineage - with connection_by_backend.transaction: - schema.rebuild_lineage() + with pytest.raises(dj.DataJointError, match="inside a transaction"): + with connection_by_backend.transaction: + Animal.insert1(dict(animal_id=1)) + schema.rebuild_lineage() + # The caller's rollback still undid the insert, and ~lineage is unchanged + assert not connection_by_backend.in_transaction + assert len(Animal()) == 0 assert schema.lineage == before # Cleanup