Merge pull request #488 from iiasa/enh/avoid-lock

Avoid locking on failed operations

RELEASE_NOTES.rst

@@ -1,9 +1,11 @@
-.. Next release
-.. ============
+Next release
+============
 
 .. All changes
 .. -----------
 
+- New :func:`.utils.discard_on_error` and matching argument to :meth:`.TimeSeries.transact` to avoid locking :class:`.TimeSeries` / :class:`.Scenario` on failed operations with :class:`.JDBCBackend` (:pull:`488`).
+
 .. _v3.7.0:
 
 v3.7.0 (2023-05-17)

doc/api.rst

@@ -205,6 +205,7 @@ Utilities
 
    .. autosummary::
       diff
+      discard_on_error
       format_scenario_list
       maybe_check_out
       maybe_commit

ixmp/core/timeseries.py

@@ -1,5 +1,5 @@
 import logging
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from os import PathLike
 from pathlib import Path
 from typing import Any, Dict, Optional, Sequence, Tuple, Union
@@ -203,13 +203,27 @@ def discard_changes(self) -> None:
         self._backend("discard_changes")
 
     @contextmanager
-    def transact(self, message: str = "", condition: bool = True):
+    def transact(
+        self, message: str = "", condition: bool = True, discard_on_error: bool = False
+    ):
         """Context manager to wrap code in a 'transaction'.
 
-        If `condition` is :obj:`True`, the TimeSeries (or :class:`.Scenario`) is
-        checked out *before* the block begins. When the block ends, the object is
-        committed with `message`. If `condition` is :obj:`False`, nothing occurs before
-        or after the block.
+        Parameters
+        ----------
+        message : str
+            Commit message to use, if any commit is performed.
+        condition : bool
+            If :obj:`True` (the default):
+
+            - Before entering the code block, the TimeSeries (or :class:`.Scenario`) is
+              checked out.
+            - On exiting the code block normally (without an exception), changes are
+              committed with `message`.
+
+            If :obj:`False`, nothing occurs on entry or exit.
+        discard_on_error : bool
+            If :obj:`True` (default :obj:`False`), then the anti-locking behaviour of
+            :func:`.discard_on_error` also applies to any exception raised in the block.
 
         Example
         -------
@@ -221,10 +235,15 @@ def transact(self, message: str = "", condition: bool = True):
         >>> # Changes to `ts` have been committed
         """
         # TODO implement __enter__ and __exit__ to allow simpler "with ts: …"
+        from ixmp.utils import discard_on_error as discard_on_error_cm
+
         if condition:
             maybe_check_out(self)
         try:
-            yield
+            # Use the discard_on_error context manager (cm) if the parameter of the same
+            # name is True
+            with discard_on_error_cm(self) if discard_on_error else nullcontext():
+                yield
         finally:
             maybe_commit(self, condition, message)
 

ixmp/testing/__init__.py

@@ -86,7 +86,12 @@
 def pytest_addoption(parser):
     """Add the ``--user-config`` command-line option to pytest."""
     parser.addoption(
-        "--user-config",
+        "--ixmp-jvm-mem",
+        action="store",
+        help="Heap space to allocated for the JDBCBackend/JVM.",
+    )
+    parser.addoption(
+        "--ixmp-user-config",
         action="store_true",
         help="Use the user's existing config/'local' platform.",
     )
@@ -96,7 +101,7 @@ def pytest_sessionstart(session):
     """Unset any configuration read from the user's directory."""
     from ixmp.backend import jdbc
 
-    if not session.config.option.user_config:
+    if not session.config.option.ixmp_user_config:
         ixmp_config.clear()
         # Further clear an automatic reference to the user's home directory. See fixture
         # tmp_env below.
@@ -164,7 +169,7 @@ def tmp_env(pytestconfig, tmp_path_factory):
     base_temp = tmp_path_factory.getbasetemp()
     os.environ["IXMP_DATA"] = str(base_temp)
 
-    if not pytestconfig.option.user_config:
+    if not pytestconfig.option.ixmp_user_config:
         # Set the path for the default/local platform in the test directory
         localdb = base_temp.joinpath("localdb", "default")
         ixmp_config.values["platform"]["local"]["path"] = localdb

ixmp/tests/core/test_timeseries.py

@@ -1,3 +1,5 @@
+import logging
+import re
 from datetime import datetime, timedelta
 
 import pandas as pd
@@ -313,6 +315,34 @@ def test_remove(self, mp, ts, commit):
         # Result is empty
         assert ts.timeseries().empty
 
+    def test_transact_discard(self, caplog, mp, ts):
+        caplog.set_level(logging.INFO, "ixmp.utils")
+
+        df = expected(DATA[2050], ts)
+
+        ts.add_timeseries(DATA[2050])
+        ts.commit("")
+
+        # Catch the deliberately-raised exception so that the test passes
+        with pytest.raises(AttributeError):
+            with ts.transact(discard_on_error=True):
+                # Remove a single data point; this operation will not be committed
+                ts.remove_timeseries(df[df.year == 2010])
+
+                # Trigger AttributeError
+                ts.foo_bar()
+
+        # Reopen the database connection
+        mp.open_db()
+
+        # Exception was caught and logged
+        assert caplog.messages[-4].startswith("Avoid locking ")
+        assert re.match("Discard (timeseries|scenario) changes", caplog.messages[-3])
+        assert "Close database connection" == caplog.messages[-2]
+
+        # Data are unchanged
+        assert_frame_equal(expected(DATA[2050], ts), ts.timeseries())
+
     # Geodata
 
     def test_add_geodata(self, ts):

ixmp/tests/test_utils.py

@@ -134,6 +134,47 @@ def test_diff_items(test_mp):
         pass  # No check of the contents
 
 
+def test_discard_on_error(caplog, test_mp):
+    caplog.set_level(logging.INFO, "ixmp.utils")
+
+    # Create a test scenario, checked-in state
+    s = make_dantzig(test_mp)
+    url = s.url
+
+    # Some actions that don't trigger exceptions
+    assert dict(value=90, unit="USD/km") == s.scalar("f")
+    s.check_out()
+    s.change_scalar("f", 100, "USD/km")
+
+    # Catch the deliberately-raised exception so that the test passes
+    with pytest.raises(KeyError):
+        # Trigger KeyError and the discard_on_error() behaviour
+        with utils.discard_on_error(s):
+            s.add_par("d", pd.DataFrame([["foo", "bar", 1.0, "kg"]]))
+
+    # Exception was caught and logged
+    assert caplog.messages[-3].startswith("Avoid locking ")
+    assert [
+        "Discard scenario changes",
+        "Close database connection",
+    ] == caplog.messages[-2:]
+
+    # Re-load the mp and the scenario
+    with pytest.raises(RuntimeError):
+        # Fails because the connection to test_mp was closed by discard_on_error()
+        s2 = Scenario(test_mp, **utils.parse_url(url)[1])
+
+    # Reopen the connection
+    test_mp.open_db()
+
+    # Now the scenario can be reloaded
+    s2 = Scenario(test_mp, **utils.parse_url(url)[1])
+    assert s2 is not s  # Different object instance than above
+
+    # Data modification above was discarded by discard_on_error()
+    assert dict(value=90, unit="USD/km") == s.scalar("f")
+
+
 def test_filtered():
     df = pd.DataFrame()
     assert df is utils.filtered(df, filters=None)

ixmp/utils/__init__.py

@@ -1,8 +1,9 @@
 import logging
 import re
 import sys
+from contextlib import contextmanager
 from pathlib import Path
-from typing import Dict, Iterable, Iterator, List, Tuple
+from typing import TYPE_CHECKING, Dict, Iterable, Iterator, List, Tuple
 from urllib.parse import urlparse
 from warnings import warn
 
@@ -11,6 +12,9 @@
 
 from ixmp.backend import ItemType
 
+if TYPE_CHECKING:  # pragma: no cover
+    from ixmp import TimeSeries
+
 log = logging.getLogger(__name__)
 
 # Globally accessible logger.
@@ -159,6 +163,51 @@ def diff(a, b, filters=None) -> Iterator[Tuple[str, pd.DataFrame]]:
             break
 
 
+@contextmanager
+def discard_on_error(ts: "TimeSeries"):
+    """Context manager to discard changes to `ts` and close the DB on any exception.
+
+    For :mod:`JDBCBackend`, this can avoid leaving `ts` in a "locked" state in the
+    database.
+
+    Examples
+    --------
+    >>> mp = ixmp.Platform()
+    >>> s = ixmp.Scenario(mp, ...)
+    >>> with discard_on_error(s):
+    ...     s.add_par(...)  # Any code
+    ...     s.not_a_method()  # Code that raises some exception
+
+    Before the the exception in the final line is raised (and possibly handled by
+    surrounding code):
+
+    - Any changes—for example, here changes due to the call to :meth:`.add_par`—are
+      discarded/not committed;
+    - ``s`` is guaranteed to be in a non-locked state; and
+    - :meth:`.close_db` is called on ``mp``.
+    """
+    mp = ts.platform
+    try:
+        yield
+    except Exception as e:
+        log.info(
+            f"Avoid locking {ts!r} before raising {e.__class__.__name__}: "
+            + str(e).splitlines()[0].strip('"')
+        )
+
+        try:
+            ts.discard_changes()
+        except Exception:  # pragma: no cover
+            pass  # Some exception trying to discard changes()
+        else:
+            log.info(f"Discard {ts.__class__.__name__.lower()} changes")
+
+        mp.close_db()
+        log.info("Close database connection")
+
+        raise
+
+
 def maybe_check_out(timeseries, state=None):
     """Check out `timeseries` depending on `state`.
 

ixmp/utils/sphinx_linkcode_github.py

@@ -10,7 +10,7 @@
 
 from sphinx.util import logging
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     import sphinx.application
 
 log = logging.getLogger(__name__)