Add ability to propagate RetrySettings

Author: vgvoleg

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -99,6 +99,21 @@ def get_characteristic(
         return dialect.get_ydb_request_settings(dbapi_connection)
 
 
+class YdbRetrySettingsCharacteristic(characteristics.ConnectionCharacteristic):
+    def reset_characteristic(self, dialect: "YqlDialect", dbapi_connection: ydb_dbapi.Connection) -> None:
+        dialect.reset_ydb_retry_settings(dbapi_connection)
+
+    def set_characteristic(
+        self, dialect: "YqlDialect", dbapi_connection: ydb_dbapi.Connection, value: ydb.RetrySettings
+    ) -> None:
+        dialect.set_ydb_retry_settings(dbapi_connection, value)
+
+    def get_characteristic(
+        self, dialect: "YqlDialect", dbapi_connection: ydb_dbapi.Connection
+    ) -> ydb.RetrySettings:
+        return dialect.get_ydb_retry_settings(dbapi_connection)
+
+
 class YqlDialect(StrCompileDialect):
     name = "yql"
     driver = "ydb"
@@ -150,6 +165,7 @@ class YqlDialect(StrCompileDialect):
         {
             "isolation_level": characteristics.IsolationLevelCharacteristic(),
             "ydb_request_settings": YdbRequestSettingsCharacteristic(),
+            "ydb_retry_settings": YdbRetrySettingsCharacteristic(),
         }
     )
 
@@ -312,6 +328,19 @@ def reset_ydb_request_settings(self, dbapi_connection: ydb_dbapi.Connection):
     def get_ydb_request_settings(self, dbapi_connection: ydb_dbapi.Connection) -> ydb.BaseRequestSettings:
         return dbapi_connection.get_ydb_request_settings()
 
+    def set_ydb_retry_settings(
+        self,
+        dbapi_connection: ydb_dbapi.Connection,
+        value: ydb.RetrySettings,
+    ) -> None:
+        dbapi_connection.set_ydb_retry_settings(value)
+
+    def reset_ydb_retry_settings(self, dbapi_connection: ydb_dbapi.Connection):
+        self.set_ydb_retry_settings(dbapi_connection, ydb.RetrySettings())
+
+    def get_ydb_retry_settings(self, dbapi_connection: ydb_dbapi.Connection) -> ydb.RetrySettings:
+        return dbapi_connection.get_ydb_retry_settings()
+
     def create_connect_args(self, url):
         args, kwargs = super().create_connect_args(url)
         # YDB database name should start with '/'

ydb_sqlalchemy/sqlalchemy/dbapi_adapter.py

@@ -56,6 +56,12 @@ def set_ydb_request_settings(self, value: ydb.BaseRequestSettings) -> None:
     def get_ydb_request_settings(self) -> ydb.BaseRequestSettings:
         return self._connection.get_ydb_request_settings()
 
+    def set_ydb_retry_settings(self, value: ydb.RetrySettings) -> None:
+        self._connection.set_ydb_retry_settings(value)
+
+    def get_ydb_retry_settings(self) -> ydb.RetrySettings:
+        return self._connection.get_ydb_retry_settings()
+
     def describe(self, table_path: str):
         return await_only(self._connection.describe(table_path))