fix: rebuild lineage on PostgreSQL and keep entries on failure

src/datajoint/lineage.py

@@ -207,12 +207,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,
@@ -265,7 +273,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
     ----------
@@ -277,17 +285,19 @@ 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)
 
     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 +308,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,7 +373,10 @@ def resolve_lineage(table, col):
             if not col.startswith("_"):
                 resolve_lineage(table, col)
 
-    # 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()]
+    with connection.transaction:
+        connection.query(f"DELETE FROM {lineage_table}")
+        if entries:
+            _upsert_lineages(connection, lineage_table, entries)

tests/integration/test_multi_backend.py

@@ -86,6 +86,155 @@ 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_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 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,
+    )
+
+    @schema
+    class Animal(dj.Manual):
+        definition = """
+        animal_id : int
+        """
+
+    before = schema.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
+    schema.drop()
+
+
 @pytest.mark.backend_agnostic
 def test_data_types(connection_by_backend, backend, prefix):
     """Test that core data types work on both backends."""