From 849ea5d9467d1f0476a385c6e529b749168125df Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Thu, 11 Dec 2025 21:50:13 +0300 Subject: [PATCH] Add ability to propagate RetrySettings --- requirements.txt | 2 +- ydb_sqlalchemy/sqlalchemy/__init__.py | 27 ++++++++++++++++++++++ ydb_sqlalchemy/sqlalchemy/dbapi_adapter.py | 6 +++++ 3 files changed, 34 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index bef9443..46c275a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,3 @@ sqlalchemy >= 1.4.0, < 3.0.0 ydb >= 3.21.6 -ydb-dbapi >= 0.1.10 +ydb-dbapi >= 0.1.16 diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index 70738d8..62875cb 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -99,6 +99,19 @@ 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 +163,7 @@ class YqlDialect(StrCompileDialect): { "isolation_level": characteristics.IsolationLevelCharacteristic(), "ydb_request_settings": YdbRequestSettingsCharacteristic(), + "ydb_retry_settings": YdbRetrySettingsCharacteristic(), } ) @@ -312,6 +326,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 '/' diff --git a/ydb_sqlalchemy/sqlalchemy/dbapi_adapter.py b/ydb_sqlalchemy/sqlalchemy/dbapi_adapter.py index e7a3f47..068fd43 100644 --- a/ydb_sqlalchemy/sqlalchemy/dbapi_adapter.py +++ b/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))