fix: prevent silent page loss in store_entity parallel upload

Three linked defects caused store_entity() to drop a page during batch
uploads while reporting success.

Fix A (wtsite.py, WtPage.edit): the retry loop returned None after
exhausting all attempts, discarding the last exception. It now raises a
RuntimeError chained from the underlying error so callers see the failure.

Fix B1 (wtsite.py): parallel uploads share one mwclient session. Guard the
destructive session mutations (_clear_cookies, _relogin, and edit()'s
token-refresh recovery) with a per-WtSite RLock, and iterate a snapshot of
the cookie jar in _clear_cookies. A lazy _get_session_lock() accessor keeps
WtSite instances built via __new__ (tests) working.

Fix C (util.py, core.py): add return_exceptions to parallelize() (captures
per-task exceptions instead of aborting the batch on the first). store_entity
now collects per-entity outcomes for both serial and parallel paths, adds
StoreEntityResult.failed, and raises OSW.StoreEntityPartialError carrying the
partial result (stored + failed) on any failure.

Adds offline unit tests for all three fixes.

(cherry picked from commit 280838c)

Author: LukasGold

src/osw/core.py

@@ -1593,11 +1593,34 @@ class StoreEntityResult(OswBaseModel):
         change_id: str
         """The ID of the change"""
         pages: Dict[str, WtPage]
-        """The pages that have been stored"""
+        """The pages that have been successfully stored, keyed by full page title.
+        On partial failure this contains only the successfully-stored pages."""
+        failed: Dict[str, Exception] = {}
+        """Entities that could not be stored, keyed by full page title and mapped to
+        the exception that caused the failure. Empty on full success."""
 
         class Config:
             arbitrary_types_allowed = True
 
+    class StoreEntityPartialError(Exception):
+        """Raised by store_entity() when one or more entities could not be stored.
+
+        Carries the partial ``StoreEntityResult`` so callers can learn exactly which
+        entities were written (``stored`` / ``result.pages``) and which failed
+        (``failed`` / ``result.failed``) without a separate existence query.
+        """
+
+        def __init__(self, result: "OSW.StoreEntityResult"):
+            self.result = result
+            self.stored = list(result.pages.keys())
+            self.failed = result.failed
+            total = len(result.pages) + len(result.failed)
+            failed_titles = ", ".join(result.failed.keys())
+            super().__init__(
+                f"store_entity failed for {len(result.failed)} of {total} "
+                f"entities: {failed_titles}"
+            )
+
     def store_entity(
         self, param: Union[StoreEntityParam, OswBaseModel, List[OswBaseModel]]
     ) -> StoreEntityResult:
@@ -1790,27 +1813,63 @@ class UploadObject(BaseModel):
                 upload_index += 1
 
         def handle_upload_object_(upload_object: UploadObject) -> None:
+            # Let exceptions propagate: the caller collects them per entity below,
+            # so a single failure neither aborts the batch nor is silently
+            # swallowed (store_entity_ only records created_pages on success).
+            store_entity_(
+                upload_object.entity,
+                upload_object.namespace,
+                upload_object.index,
+                upload_object.overwrite_class_param,
+            )
+
+        def failure_title_(upload_object: UploadObject) -> str:
+            """Best-effort full page title of a failed entity, for error reporting."""
             try:
-                store_entity_(
-                    upload_object.entity,
-                    upload_object.namespace,
-                    upload_object.index,
-                    upload_object.overwrite_class_param,
+                namespace = upload_object.namespace or get_namespace(
+                    upload_object.entity
+                )
+                return f"{namespace}:{get_title(upload_object.entity)}"
+            except Exception:
+                return (
+                    getattr(upload_object.entity, "name", None)
+                    or getattr(upload_object.entity, "uuid", None)
+                    or "unknown"
                 )
-            except Exception as e:
-                entity_name = getattr(upload_object.entity, "name", None) or "unknown"
-                _logger.error(f"Error storing entity '{entity_name}': {e}")
 
         if param.parallel:
-            _ = parallelize(
-                handle_upload_object_, upload_object_list, flush_at_end=param.debug
+            # return_exceptions=True keeps results aligned with upload_object_list
+            # and lets every entity be attempted even if some fail.
+            results = parallelize(
+                handle_upload_object_,
+                upload_object_list,
+                flush_at_end=param.debug,
+                return_exceptions=True,
             )
         else:
-            _ = [
-                handle_upload_object_(upload_object)
-                for upload_object in upload_object_list
-            ]
-        return OSW.StoreEntityResult(change_id=param.change_id, pages=created_pages)
+            results = []
+            for upload_object in upload_object_list:
+                try:
+                    handle_upload_object_(upload_object)
+                    results.append(None)
+                except Exception as e:  # noqa: BLE001 - collected below
+                    results.append(e)
+
+        failed: Dict[str, Exception] = {}
+        for upload_object, result in zip(upload_object_list, results):
+            if isinstance(result, Exception):
+                title = failure_title_(upload_object)
+                _logger.error(f"Error storing entity '{title}': {result}")
+                failed[title] = result
+
+        store_result = OSW.StoreEntityResult(
+            change_id=param.change_id, pages=created_pages, failed=failed
+        )
+        if failed:
+            # Surface partial/total failure so callers cannot mistake a dropped
+            # page for a success. The result (successes + failures) rides along.
+            raise OSW.StoreEntityPartialError(store_result)
+        return store_result
 
     class DeleteEntityParam(OswBaseModel):
         entities: Union[OswBaseModel, List[OswBaseModel]]

src/osw/utils/util.py

@@ -274,6 +274,7 @@ def parallelize(
     iterable: Iterable,
     flush_at_end: bool = False,
     progress_bar: bool = True,
+    return_exceptions: bool = False,
     **kwargs,
 ):
     """A function to parallelize tasks with a progress bar and a message buffer.
@@ -289,6 +290,11 @@ def parallelize(
         execution.
     progress_bar:
         If True, a progress bar will be displayed.
+    return_exceptions:
+        If True, exceptions raised by ``func`` are returned in the result list at the
+        position of the corresponding item (instead of aborting the whole batch on the
+        first failure). Mirrors ``asyncio.gather(return_exceptions=...)``. Results stay
+        aligned with the input order, so callers can zip them back to ``iterable``.
     kwargs:
         Keyword arguments to be passed to the function.z
     """
@@ -324,6 +330,18 @@ async def _run_tasks():
             print(
                 f"Performing parallel execution of {func.__name__} ({len(tasks)} tasks)."
             )
+            if return_exceptions:
+                # tqdm.gather() re-raises the first exception (it has no
+                # return_exceptions kwarg), which would abandon the remaining
+                # results. Capture each task's exception instead so the batch
+                # completes and results stay aligned with the input order.
+                async def _capture(task):
+                    try:
+                        return await task
+                    except Exception as exc:  # noqa: BLE001 - returned to caller
+                        return exc
+
+                tasks = [_capture(task) for task in tasks]
             res = await tqdm.gather(
                 *tasks
             )  # like asyncio.gather, but with a progress bar

src/osw/wtsite.py

@@ -5,6 +5,7 @@
 import json
 import os
 import shutil
+import threading
 import urllib
 import warnings
 import xml.etree.ElementTree as et
@@ -88,6 +89,13 @@ def __init__(self, config: Union[WtSiteConfig, WtSiteLegacyConfig]):
 
         """
 
+        # Serializes destructive mutations of the shared mwclient session
+        # (cookie jar / tokens). Parallel uploads share one requests.Session, so
+        # concurrent clears/re-logins would otherwise corrupt each other's state.
+        # RLock (not Lock) because the edit() retry path may call _relogin() while
+        # already holding this lock.
+        self._session_lock = threading.RLock()
+
         scheme = "https"
 
         # Store credentials for potential re-login on session timeout
@@ -153,6 +161,20 @@ def __init__(self, config: Union[WtSiteConfig, WtSiteLegacyConfig]):
         self._page_cache = {}
         self._cache_enabled = False
 
+    def _get_session_lock(self) -> threading.RLock:
+        """Return the session lock, lazily creating it if absent.
+
+        In normal use the lock is created in ``__init__`` before any worker
+        threads exist. This fallback keeps WtSite instances that bypassed
+        ``__init__`` (e.g. tests using ``WtSite.__new__``) working; those are
+        single-threaded at construction, so the lazy creation is race-free.
+        """
+        lock = getattr(self, "_session_lock", None)
+        if lock is None:
+            lock = threading.RLock()
+            self._session_lock = lock
+        return lock
+
     def _relogin(self):
         """Re-login to the wiki site using stored credentials.
 
@@ -174,9 +196,12 @@ def _relogin(self):
             # Stale session cookies cause MediaWiki to abort the login flow with
             # "Unable to continue login. Your session most likely timed out."
             # Clear client-side session state so login starts from a clean slate.
-            self._site.connection.cookies.clear()
-            self._site.tokens.clear()
-            self._site.login(username=cred.username, password=cred.password)
+            # Hold the session lock so in-flight parallel uploads cannot read the
+            # cookie jar / tokens while they are being wiped and rebuilt.
+            with self._get_session_lock():
+                self._site.connection.cookies.clear()
+                self._site.tokens.clear()
+                self._site.login(username=cred.username, password=cred.password)
         else:
             raise RuntimeError(
                 "Re-login is only supported for username/password credentials."
@@ -457,11 +482,14 @@ def clear_cache(self):
 
     def _clear_cookies(self):
         # see https://github.com/mwclient/mwclient/issues/221
-        for cookie in self._site.connection.cookies:
-            if "PostEditRevision" in cookie.name:
-                self._site.connection.cookies.clear(
-                    cookie.domain, cookie.path, cookie.name
-                )
+        # Iterate a snapshot (list(...)) so mutating the jar mid-loop is safe, and
+        # hold the session lock so a peer thread cannot clear/re-login concurrently.
+        with self._get_session_lock():
+            for cookie in list(self._site.connection.cookies):
+                if "PostEditRevision" in cookie.name:
+                    self._site.connection.cookies.clear(
+                        cookie.domain, cookie.path, cookie.name
+                    )
 
     class SearchParam(wt.SearchParam):
         pass
@@ -1718,26 +1746,42 @@ def edit(self, comment: str = None, mode="action-multislot", bot_edit: bool = Tr
         mode:
             (optional) single API call ('action-multislot') or multiple (
             'action-singleslot'), by default 'action-multislot' (faster)
+
+        Raises
+        ------
+        RuntimeError
+            if the edit still fails after ``max_retry`` attempts. The last
+            underlying exception is chained via ``__cause__``. Previously this
+            method returned None on exhaustion, silently discarding the failure.
         """
-        retry = 0
         max_retry = 5
-        while retry < max_retry:
+        last_exc: Optional[Exception] = None
+        for attempt in range(max_retry):
             try:
                 return self._edit(comment, mode, bot_edit)
-            except Exception as e:
-                if retry < max_retry:
-                    retry += 1
-                    print(f"Page edit failed: {e}. Retry ({retry}/{max_retry})")
+            except Exception as e:  # noqa: BLE001 - re-raised below after retries
+                last_exc = e
+                print(f"Page edit failed: {e}. Retry ({attempt + 1}/{max_retry})")
+                if attempt + 1 < max_retry:
+                    # Attempt to recover the shared session before retrying.
+                    # Guard the whole block: a recovery failure must never mask
+                    # the original edit error we intend to re-raise below.
                     try:
-                        # refresh token for longer running processes
-                        self.wtSite._site.get_token("csrf", force=True)
+                        # Serialize session mutations so a peer thread's request
+                        # cannot read a half-cleared cookie jar / token set.
+                        with self.wtSite._get_session_lock():
+                            try:
+                                # refresh token for longer running processes
+                                self.wtSite._site.get_token("csrf", force=True)
+                            except Exception:
+                                # token refresh failed, attempt full re-login
+                                self.wtSite._relogin()
                     except Exception:
-                        # token refresh failed, attempt full re-login
-                        try:
-                            self.wtSite._relogin()
-                        except Exception:
-                            pass  # re-login failed, will retry anyway
+                        pass  # recovery failed, will retry / re-raise anyway
                     sleep(5)
+        raise RuntimeError(
+            f"Page edit for '{self.title}' failed after {max_retry} attempts"
+        ) from last_exc
 
     def _edit(
         self, comment: str = None, mode="action-multislot", bot_edit: bool = True

tests/test_store_entity_failure.py

@@ -0,0 +1,99 @@
+"""Unit tests for store_entity() per-entity failure reporting (Fix C).
+
+These run fully offline: WtPage.init and the overwrite policy are stubbed so no
+network is required, and WtPage.edit is replaced with a controllable stub that
+fails for selected page titles.
+"""
+
+import pytest
+
+import osw.model.entity as model
+from osw.core import OSW
+from osw.utils.wiki import get_namespace, get_title
+from osw.wtsite import WtPage
+
+
+def _title(entity):
+    return f"{get_namespace(entity)}:{get_title(entity)}"
+
+
+@pytest.fixture
+def offline_osw(monkeypatch):
+    # no network when a WtPage is constructed with do_init=True; mimic the
+    # do_init=False branch of WtPage.__init__ which sets .exists
+    monkeypatch.setattr(WtPage, "init", lambda self: setattr(self, "exists", False))
+    # bypass the overwrite policy: return the page that was built for the entity
+    monkeypatch.setattr(
+        OSW, "_apply_overwrite_policy", staticmethod(lambda param: param.page)
+    )
+    return OSW.construct(site=object())
+
+
+def _install_edit(monkeypatch, failing_titles):
+    def fake_edit(self, *args, **kwargs):
+        if self.title in failing_titles:
+            raise RuntimeError(f"edit failed for {self.title}")
+        return None
+
+    monkeypatch.setattr(WtPage, "edit", fake_edit)
+
+
+def test_store_entity_serial_single_failure_raises(offline_osw, monkeypatch):
+    item = model.Item(label=[model.Label(text="Solo")])
+    title = _title(item)
+    _install_edit(monkeypatch, {title})
+
+    with pytest.raises(OSW.StoreEntityPartialError) as exc_info:
+        offline_osw.store_entity(OSW.StoreEntityParam(entities=[item], parallel=False))
+
+    err = exc_info.value
+    assert title in err.failed
+    assert isinstance(err.failed[title], RuntimeError)
+    # the dropped page must NOT be reported as stored
+    assert title not in err.result.pages
+    assert err.stored == []
+
+
+def test_store_entity_parallel_middle_failure_reports_all(offline_osw, monkeypatch):
+    items = [model.Item(label=[model.Label(text=f"Item{i}")]) for i in range(3)]
+    titles = [_title(it) for it in items]
+    failing = titles[1]
+    _install_edit(monkeypatch, {failing})
+
+    with pytest.raises(OSW.StoreEntityPartialError) as exc_info:
+        offline_osw.store_entity(OSW.StoreEntityParam(entities=items, parallel=True))
+
+    err = exc_info.value
+    # the two good entities are recorded as stored (the first exception did not
+    # abandon the remaining tasks) ...
+    assert set(err.result.pages.keys()) == {titles[0], titles[2]}
+    assert set(err.stored) == {titles[0], titles[2]}
+    # ... and the failing one is reported instead of being silently dropped
+    assert set(err.failed.keys()) == {failing}
+    assert isinstance(err.failed[failing], RuntimeError)
+
+
+def test_store_entity_all_success_returns_result(offline_osw, monkeypatch):
+    items = [model.Item(label=[model.Label(text=f"Ok{i}")]) for i in range(3)]
+    titles = [_title(it) for it in items]
+    _install_edit(monkeypatch, set())  # nothing fails
+
+    result = offline_osw.store_entity(
+        OSW.StoreEntityParam(entities=items, parallel=True)
+    )
+
+    assert set(result.pages.keys()) == set(titles)
+    assert result.failed == {}
+
+
+def test_store_entity_serial_all_success_returns_result(offline_osw, monkeypatch):
+    items = [model.Item(label=[model.Label(text=f"Ser{i}")]) for i in range(2)]
+    titles = [_title(it) for it in items]
+    _install_edit(monkeypatch, set())
+
+    result = offline_osw.store_entity(
+        OSW.StoreEntityParam(entities=items, parallel=False)
+    )
+
+    assert set(result.pages.keys()) == set(titles)
+    assert result.failed == {}