[theluckystrike]

schema.rebuild_lineage() failed on PostgreSQL because rebuild_schema_lineage read keys from information_schema with MySQL only names (REFERENCED_TABLE_SCHEMA and friends, and CONSTRAINT_NAME = 'PRIMARY'). It had already run DELETE FROM ~lineage by then, so the failed call left the schema with an empty ~lineage table.

The function now reads primary and foreign keys through adapter.load_primary_keys_sql and adapter.load_foreign_keys_sql, the same queries Dependencies.load uses, and splits the returned names with adapter.split_full_table_name. Those foreign key queries also return keys that other schemas hold on this one, so rows whose referencing table is in another schema are ignored. The DELETE moved to just before the insert, after every lineage has resolved, so a rebuild that cannot resolve a key leaves ~lineage as it was. That also covers the existing MySQL failure where an upstream schema has no lineage.

Two tests were added to tests/integration/test_multi_backend.py, and both run on MySQL and PostgreSQL. test_rebuild_lineage empties ~lineage and checks the rebuilt entries for a parent and child table. test_failed_rebuild_lineage_keeps_entries removes the upstream lineage of a cross schema foreign key and checks that the downstream entries survive the error.

I ran the tests against postgres:15 and datajoint/mysql:8.0 containers with DJ_USE_EXTERNAL_CONTAINERS=1. On master with the new tests, pytest tests/integration/test_multi_backend.py -k rebuild gave 3 failed and 1 passed (the MySQL happy path passes on master). With this change it gives 4 passed. pytest tests/integration/test_multi_backend.py tests/integration/test_semantic_matching.py tests/integration/test_dependencies.py tests/unit gave 512 passed. ruff check and ruff format --check pass on the two changed files. I did not run the full pixi run -e test test suite or the S3 backed tests.

Fixes #1641

[dimitri-yatsenko]

Thanks for this. The full suite passes locally on MySQL and PostgreSQL, and the two new tests fail on master as you describe (3 failed, 1 passed).

One gap in keeping ~lineage on failure, noted inline.

[dimitri-yatsenko]

Moving the DELETE here covers a key that cannot be resolved. But the DELETE and the insert below still run as separate statements outside a transaction, so if the insert fails (a dropped connection, for example), ~lineage is left empty, which is the state #1641 is about. Could the two run in one transaction?

One thing to watch: insert_lineages calls ensure_lineage_table, which runs CREATE TABLE IF NOT EXISTS. MySQL commits any open transaction before a CREATE TABLE, even when the table already exists, so calling insert_lineages inside with connection.transaction: would still commit the DELETE on its own. The table must already exist here, since the DELETE succeeded, so one way is to split the upsert out of insert_lineages into a helper that does not create the table, and run the DELETE and that upsert together:

with connection.transaction:
    connection.query(f"DELETE FROM {lineage_table}")
    if entries:
        _upsert_lineages(connection, lineage_table, entries)

A test could make the insert fail (for example with monkeypatch on the upsert) and check that the entries survive.

[theluckystrike]

Done in e9d2686, the way you sketched it. _upsert_lineages does the upsert without ensure_lineage_table, insert_lineages calls it after creating the table, and rebuild_schema_lineage runs the DELETE and the upsert inside with connection.transaction.

One addition. start_transaction raises on a nested transaction, so calling schema.rebuild_lineage() inside a transaction the caller already opened would start failing. In that case the rebuild reuses the open transaction instead of starting one.

Two new tests, on MySQL and PostgreSQL.

• test_failed_lineage_insert_keeps_entries patches _upsert_lineages to raise LostConnectionError and checks that ~lineage still holds the old entries and no transaction is left open. With the transaction removed it fails on both backends.
• test_rebuild_lineage_in_open_transaction runs a rebuild inside with connection.transaction. Without the reuse it fails on both backends.

I ran the full suite locally against MySQL 8.4 and PostgreSQL 15 (1229 passed). The 12 failures there are 11 diagram tests (no Graphviz on this machine) and test_activate_unauthorized, and master fails the same 12 in the same setup.

[dimitri-yatsenko]

Thanks @theluckystrike. _upsert_lineages and the transaction around the DELETE and the upsert fix the case where the rebuild opens its own transaction: nothing inside it can commit early.

The new path that reuses a transaction the caller already opened still has the problem from my earlier review, on a different line. rebuild_schema_lineage calls ensure_lineage_table (CREATE TABLE IF NOT EXISTS) before it checks connection.in_transaction. Inside the caller's transaction, MySQL commits at that statement, even when the table exists. DataJoint still thinks it is in a transaction, so the DELETE and the upsert each commit on their own. For example:

with conn.transaction:
    Animal.insert1(row)
    schema.rebuild_lineage()   # the upsert fails here

On MySQL, the insert is already committed by the CREATE, ~lineage is left empty (the state #1641 describes), and the caller's rollback has nothing left to undo. test_rebuild_lineage_in_open_transaction checks only that no error is raised and the entries are unchanged, so it passes on MySQL anyway. PostgreSQL is fine, because its DDL runs inside the transaction.

The simplest fix is to refuse, as declare and alter already do inside a transaction (table.py, "Cannot declare new tables inside a transaction"). At the top of rebuild_schema_lineage, before ensure_lineage_table:

if connection.in_transaction:
    raise DataJointError("Cannot rebuild lineage inside a transaction.")

Then use a plain with connection.transaction: around the DELETE and the upsert, drop the contextlib import, and change test_rebuild_lineage_in_open_transaction to expect DataJointError inside with connection.transaction: and check that schema.lineage is unchanged afterward. On master, a rebuild inside a caller's transaction also committed that transaction on MySQL, so raising is no loss, only explicit.

If you'd rather keep rebuild working inside a caller's transaction, run the CREATE only when the table is missing, and refuse when it is missing and a transaction is open. That version needs a test proving the caller's rollback still works on MySQL: insert a row, rebuild, raise inside with conn.transaction:, then check that the row is gone and ~lineage is intact.

[theluckystrike]

Thanks @dimitri-yatsenko, you're right. The CREATE TABLE IF NOT EXISTS ran before the in_transaction check, so on MySQL it committed the caller's transaction and my new test couldn't see it.

I took the simpler option in d59d4ec. rebuild_schema_lineage now raises DataJointError("Cannot rebuild lineage inside a transaction.") before ensure_lineage_table, the DELETE and the upsert run in a plain with connection.transaction:, and the contextlib import is gone.

test_rebuild_lineage_in_open_transaction now inserts a row inside with connection.transaction:, calls schema.rebuild_lineage(), and expects DataJointError. Afterwards it checks that no transaction is open, that the row was rolled back (len(Animal()) == 0) and that schema.lineage is unchanged. Against the previous commit it fails on both MySQL and PostgreSQL with "DID NOT RAISE DataJointError". With the guard disabled it fails too, because the plain transaction then raises "Nested connections are not supported." The lineage, transaction and migrate tests pass on both backends (41 passed), and ruff check and format are clean.