Add ability to propagate RetrySettings

Author: vgvoleg

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:

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