Merge pull request #11789 from freqtrade/fix/customdata_sessions

fix: improved session handling for custom_data sessions

freqtrade/persistence/models.py

@@ -2,6 +2,7 @@
 This module contains the class to persist trades into SQLite
 """
 
+import functools
 import logging
 import threading
 from contextvars import ContextVar
@@ -94,3 +95,22 @@ def init_db(db_url: str) -> None:
     previous_tables = inspect(engine).get_table_names()
     ModelBase.metadata.create_all(engine)
     check_migrate(engine, decl_base=ModelBase, previous_tables=previous_tables)
+
+
+def custom_data_rpc_wrapper(func):
+    """
+    Wrapper for RPC methods when using custom_data
+    Similar behavior to deps.get_rpc() - but limited to custom_data.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            _CustomData.session.rollback()
+            return func(*args, **kwargs)
+        finally:
+            _CustomData.session.rollback()
+            # Ensure the session is removed after use
+            _CustomData.session.remove()
+
+    return wrapper

freqtrade/rpc/rpc.py

@@ -35,7 +35,7 @@ from freqtrade.exchange.exchange_utils import price_to_precision
 from freqtrade.ft_types import AnnotationType
 from freqtrade.loggers import bufferHandler
 from freqtrade.persistence import CustomDataWrapper, KeyValueStore, PairLocks, Trade
-from freqtrade.persistence.models import PairLock
+from freqtrade.persistence.models import PairLock, custom_data_rpc_wrapper
 from freqtrade.plugins.pairlist.pairlist_helpers import expand_pairlist
 from freqtrade.rpc.fiat_convert import CryptoToFiatConverter
 from freqtrade.rpc.rpc_types import RPCSendMsg
@@ -1125,6 +1125,7 @@ class RPC:
                 "cancel_order_count": c_count,
             }
 
+    @custom_data_rpc_wrapper
     def _rpc_list_custom_data(
         self, trade_id: int | None = None, key: str | None = None, limit: int = 100, offset: int = 0
     ) -> list[dict[str, Any]]:
@@ -1137,6 +1138,7 @@ class RPC:
         - "custom_data": a list of custom data dicts, each with the fields:
                 "id", "key", "type", "value", "created_at", "updated_at"
         """
+
         trades: Sequence[Trade]
         if trade_id is None:
             # Get all open trades