From 919ff6c60befdb3011841e19b92dbb591eebe907 Mon Sep 17 00:00:00 2001
From: Oleg Ovcharuk <vgvoleg@gmail.com>
Date: Thu, 11 Dec 2025 20:57:55 +0300
Subject: [PATCH] Add ability to propagate RetrySettings

---
 tests/test_cursors.py    |  2 ++
 ydb_dbapi/connections.py |  9 +++++++++
 ydb_dbapi/cursors.py     | 26 ++++++++++++++++++++------
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/tests/test_cursors.py b/tests/test_cursors.py
index f951121..177b394 100644
--- a/tests/test_cursors.py
+++ b/tests/test_cursors.py
@@ -166,6 +166,7 @@ def sync_cursor(
             session_pool_sync,
             ydb.QuerySerializableReadWrite(),
             request_settings=ydb.BaseRequestSettings(),
+            retry_settings=ydb.RetrySettings(),
         )
         yield cursor
         cursor.close()
@@ -208,6 +209,7 @@ async def async_cursor(
             session_pool,
             ydb.QuerySerializableReadWrite(),
             request_settings=ydb.BaseRequestSettings(),
+            retry_settings=ydb.RetrySettings(),
         )
         yield cursor
         cursor.close()
diff --git a/ydb_dbapi/connections.py b/ydb_dbapi/connections.py
index 61c692f..6bd9a25 100644
--- a/ydb_dbapi/connections.py
+++ b/ydb_dbapi/connections.py
@@ -117,6 +117,7 @@ def __init__(
         self.request_settings: ydb.BaseRequestSettings = (
             ydb.BaseRequestSettings()
         )
+        self.retry_settings: ydb.RetrySettings = ydb.RetrySettings()
 
     def set_isolation_level(self, isolation_level: IsolationLevel) -> None:
         if self._tx_context and self._tx_context.tx_id:
@@ -151,6 +152,12 @@ def set_ydb_request_settings(self, value: ydb.BaseRequestSettings) -> None:
     def get_ydb_request_settings(self) -> ydb.BaseRequestSettings:
         return self.request_settings
 
+    def set_ydb_retry_settings(self, value: ydb.RetrySettings) -> None:
+        self.retry_settings = value
+
+    def get_ydb_retry_settings(self) -> ydb.RetrySettings:
+        return self.retry_settings
+
     def _get_request_settings(self) -> ydb.BaseRequestSettings:
         settings = self.request_settings.make_copy()
 
@@ -210,6 +217,7 @@ def cursor(self) -> Cursor:
             tx_context=self._tx_context,
             table_path_prefix=self.table_path_prefix,
             request_settings=self.request_settings,
+            retry_settings=self.retry_settings,
         )
 
     def wait_ready(self, timeout: int = 10) -> None:
@@ -402,6 +410,7 @@ def cursor(self) -> AsyncCursor:
             tx_context=self._tx_context,
             table_path_prefix=self.table_path_prefix,
             request_settings=self.request_settings,
+            retry_settings=self.retry_settings,
         )
 
     async def wait_ready(self, timeout: int = 10) -> None:
diff --git a/ydb_dbapi/cursors.py b/ydb_dbapi/cursors.py
index 41706f9..7060401 100644
--- a/ydb_dbapi/cursors.py
+++ b/ydb_dbapi/cursors.py
@@ -199,6 +199,7 @@ def __init__(
         session_pool: ydb.QuerySessionPool,
         tx_mode: ydb.BaseQueryTxMode,
         request_settings: ydb.BaseRequestSettings,
+        retry_settings: ydb.RetrySettings,
         tx_context: ydb.QueryTxContext | None = None,
         table_path_prefix: str = "",
     ) -> None:
@@ -207,9 +208,9 @@ def __init__(
         self._session_pool = session_pool
         self._tx_mode = tx_mode
         self._request_settings = request_settings
+        self._retry_settings = retry_settings
         self._tx_context = tx_context
         self._table_path_prefix = table_path_prefix
-
         self._stream: Iterator | None = None
 
     def fetchone(self) -> tuple | None:
@@ -253,7 +254,10 @@ def callee(
                 )
             )
 
-        return self._session_pool.retry_operation_sync(callee)
+        return self._session_pool.retry_operation_sync(
+            callee,
+            retry_settings=self._retry_settings,
+        )
 
     @handle_ydb_errors
     @invalidate_cursor_on_ydb_error
@@ -276,7 +280,10 @@ def callee(
                 )
             )
 
-        return self._session_pool.retry_operation_sync(callee)
+        return self._session_pool.retry_operation_sync(
+            callee,
+            retry_settings=self._retry_settings,
+        )
 
     @handle_ydb_errors
     @invalidate_cursor_on_ydb_error
@@ -369,6 +376,7 @@ def __init__(
         session_pool: ydb.aio.QuerySessionPool,
         tx_mode: ydb.BaseQueryTxMode,
         request_settings: ydb.BaseRequestSettings,
+        retry_settings: ydb.RetrySettings,
         tx_context: ydb.aio.QueryTxContext | None = None,
         table_path_prefix: str = "",
     ) -> None:
@@ -377,9 +385,9 @@ def __init__(
         self._session_pool = session_pool
         self._tx_mode = tx_mode
         self._request_settings = request_settings
+        self._retry_settings = retry_settings
         self._tx_context = tx_context
         self._table_path_prefix = table_path_prefix
-
         self._stream: AsyncIterator | None = None
 
     def fetchone(self) -> tuple | None:
@@ -423,7 +431,10 @@ async def callee(
                 )
             )
 
-        return await self._session_pool.retry_operation_async(callee)
+        return await self._session_pool.retry_operation_async(
+            callee,
+            retry_settings=self._retry_settings,
+        )
 
     @handle_ydb_errors
     @invalidate_cursor_on_ydb_error
@@ -446,7 +457,10 @@ async def callee(
                 )
             )
 
-        return await self._session_pool.retry_operation_async(callee)
+        return await self._session_pool.retry_operation_async(
+            callee,
+            retry_settings=self._retry_settings,
+        )
 
     @handle_ydb_errors
     @invalidate_cursor_on_ydb_error