🐛 aiida_computer: handle IntegrityError races (#7349)

Under `pytest-xdist`, the session-scoped `aiida_profile` is shared
across workers. When one worker calls `aiida_profile_clean` it wipes the
database mid-session, so subsequent calls to `aiida_computer` in other
workers hit a TOCTOU race: `get()` finds no row, but `store()` raises
`IntegrityError` because another worker already re-created the computer
in the meantime.

Band-aid fix for the UNIQUE-constraint failures reported in #7347:

- Catch `IntegrityError` from `store()` and fall back to `get()` on the
  same four-field match (label, hostname, scheduler_type,
  transport_type), consistent with the original lookup.
- If the fallback `get()` also raises `NotExistent` (rare: the racing
  worker's row already vanished), rebuild and store fresh.
- Apply caller-provided `minimum_job_poll_interval` and
  `default_mpiprocs_per_machine` on every recovery path so the caller's
  values are never silently dropped.

Add a regression test for each recovery branch.

Co-authored-by: Julian Geiger <julian.geiger@psi.ch>

Author: agoscinski

src/aiida/tools/pytest_fixtures/orm.py

@@ -91,24 +91,41 @@ def factory(
     ) -> 'Computer':
         import uuid
 
-        from aiida.common.exceptions import NotExistent
+        from aiida.common.exceptions import IntegrityError, NotExistent
         from aiida.orm import Computer
 
         label = label or f'test-computer-{uuid.uuid4().hex}'
 
-        try:
-            computer = Computer.collection.get(
-                label=label, hostname=hostname, scheduler_type=scheduler_type, transport_type=transport_type
-            )
-        except NotExistent:
-            computer = Computer(
+        def _build() -> 'Computer':
+            return Computer(
                 label=label,
                 hostname=hostname,
                 workdir=str(tmp_path),
                 transport_type=transport_type,
                 scheduler_type=scheduler_type,
             )
-            computer.store()
+
+        try:
+            computer = Computer.collection.get(
+                label=label, hostname=hostname, scheduler_type=scheduler_type, transport_type=transport_type
+            )
+        except NotExistent:
+            computer = _build()
+            try:
+                computer.store()
+            except IntegrityError:
+                # Another xdist worker concurrently re-created this computer after
+                # ``aiida_profile_clean`` reset the shared database. Re-querying is
+                # safe because the ``psql_dos`` storage backend rolls back our failed
+                # transaction before raising (see ``psql_dos/orm/utils.py``).
+                try:
+                    computer = Computer.collection.get(
+                        label=label, hostname=hostname, scheduler_type=scheduler_type, transport_type=transport_type
+                    )
+                except NotExistent:
+                    # Row no longer present (rare): rebuild and store fresh.
+                    computer = _build()
+                    computer.store()
             computer.set_minimum_job_poll_interval(minimum_job_poll_interval)
             computer.set_default_mpiprocs_per_machine(default_mpiprocs_per_machine)
 

tests/tools/pytest_fixtures/test_orm.py

@@ -1,9 +1,11 @@
 """Tests for the :mod:`aiida.tools.pytest_fixtures.orm` module."""
 
 import uuid
+from unittest.mock import patch
 
 import pytest
 
+from aiida.common.exceptions import IntegrityError, NotExistent
 from aiida.orm import Computer
 from aiida.transports import AsyncTransport, BlockingTransport
 from aiida.transports.plugins.async_backend import _AsyncSSH, _OpenSSH
@@ -50,6 +52,72 @@ def test_aiida_computer_fixtures(fixture_name, transport_cls, transport_type, re
     assert not computer_unconfigured.is_configured
 
 
+@pytest.mark.usefixtures('aiida_profile_clean')
+def test_aiida_computer_integrity_error_fallback(aiida_computer):
+    """The fallback ``get()`` returns the racing worker's row.
+
+    Simulates the xdist race where a row exists in the shared DB but our session
+    can't see it on the first ``get()`` (as if ``aiida_profile_clean`` just wiped
+    things from this worker's POV); ``store()`` then fails with ``IntegrityError``
+    on the UNIQUE label constraint, and the fallback ``get()`` finds the row.
+    """
+    label = f'test-fallback-{uuid.uuid4().hex}'
+    existing = aiida_computer(label=label)
+
+    with (
+        patch.object(
+            type(Computer.collection),
+            'get',
+            side_effect=[NotExistent('First get: DB was just cleaned'), existing],
+        ),
+        patch.object(Computer, 'store', side_effect=IntegrityError('UNIQUE constraint failed')),
+    ):
+        recovered = aiida_computer(label=label)
+
+    assert recovered.pk == existing.pk
+    assert recovered.uuid == existing.uuid
+
+
+@pytest.mark.usefixtures('aiida_profile_clean')
+def test_aiida_computer_integrity_error_rebuild(aiida_computer):
+    """The fallback ``get()`` also misses, so the fixture rebuilds and stores fresh.
+
+    Rare follow-up to the fallback case: ``store()`` raised ``IntegrityError`` (so
+    the racing worker's row should be there), but by the time we re-query it has
+    already vanished. The fixture must then build a new ``Computer`` and store it.
+    """
+    label = f'test-rebuild-{uuid.uuid4().hex}'
+
+    # store() fails the first time (simulating the race) and persists for real on
+    # the second call (the rebuild). ``autospec=True`` is required so the mock
+    # passes ``self`` to ``store_side_effect``, which we need to call through.
+    original_store = Computer.store
+    store_calls = {'n': 0}
+
+    def store_side_effect(self):
+        store_calls['n'] += 1
+        if store_calls['n'] == 1:
+            raise IntegrityError('UNIQUE constraint failed')
+        return original_store(self)
+
+    with (
+        patch.object(
+            type(Computer.collection),
+            'get',
+            side_effect=[
+                NotExistent('First get: DB was just cleaned'),
+                NotExistent('Second get: racing worker row vanished'),
+            ],
+        ),
+        patch.object(Computer, 'store', autospec=True, side_effect=store_side_effect),
+    ):
+        recovered = aiida_computer(label=label)
+
+    assert recovered.is_stored
+    assert recovered.label == label
+    assert store_calls['n'] == 2
+
+
 @pytest.mark.parametrize(
     'backend, backend_class',
     [