Moved commit "Added new versions of RetryQuery" from ydb repo

ydb-platform · Aug 26, 2024 · ce4222c · ce4222c
1 parent f48930a
commit ce4222c
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 7 deletions.
diff --git a/include/ydb-cpp-sdk/client/query/client.h b/include/ydb-cpp-sdk/client/query/client.h
@@ -14,7 +14,11 @@ namespace NYdb {
     namespace NRetry::Async {
         template <typename TClient, typename TAsyncStatusType>
         class TRetryContext;
-    }
+    } // namespace NRetry::Async
+    namespace NRetry::Sync {
+        template <typename TClient, typename TStatusType>
+        class TRetryContext;
+    } // namespace NRetry::Sync
 }
 
 namespace NYdb::NQuery {
@@ -55,10 +59,15 @@ class TSession;
 class TQueryClient {
     friend class TSession;
     friend class NRetry::Async::TRetryContext<TQueryClient, TAsyncExecuteQueryResult>;
+    friend class NRetry::Async::TRetryContext<TQueryClient, TAsyncStatus>;
+    friend class NRetry::Sync::TRetryContext<TQueryClient, TStatus>;
 
 public:
-    using TQueryFunc = std::function<TAsyncExecuteQueryResult(TSession session)>;
-    using TQueryWithoutSessionFunc = std::function<TAsyncExecuteQueryResult(TQueryClient& client)>;
+    using TQueryResultFunc = std::function<TAsyncExecuteQueryResult(TSession session)>;
+    using TQueryFunc = std::function<TAsyncStatus(TSession session)>;
+    using TQuerySyncFunc = std::function<TStatus(TSession session)>;
+    using TQueryWithoutSessionFunc = std::function<TAsyncStatus(TQueryClient& client)>;
+    using TQueryWithoutSessionSyncFunc = std::function<TStatus(TQueryClient& client)>;
     using TSettings = TClientSettings;
     using TSession = TSession;
     using TCreateSessionSettings = TCreateSessionSettings;
@@ -79,7 +88,15 @@ class TQueryClient {
     TAsyncExecuteQueryIterator StreamExecuteQuery(const std::string& query, const TTxControl& txControl,
         const TParams& params, const TExecuteQuerySettings& settings = TExecuteQuerySettings());
 
-    TAsyncExecuteQueryResult RetryQuery(TQueryFunc&& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings());
+    TAsyncExecuteQueryResult RetryQuery(TQueryResultFunc&& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings());
+
+    TAsyncStatus RetryQuery(TQueryFunc&& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings());
+
+    TAsyncStatus RetryQuery(TQueryWithoutSessionFunc&& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings());
+
+    TStatus RetryQuery(const TQuerySyncFunc& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings());
+
+    TStatus RetryQuery(const TQueryWithoutSessionSyncFunc& queryFunc, TRetryOperationSettings settings = TRetryOperationSettings());
 
     TAsyncExecuteQueryResult RetryQuery(const std::string& query, const TTxControl& txControl,
         TDuration timeout, bool isIndempotent);

diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp
@@ -21,7 +21,8 @@
 
 namespace NYdb::NQuery {
 
-using TRetryContextAsync = NRetry::Async::TRetryContext<TQueryClient, TAsyncExecuteQueryResult>;
+using TRetryContextResultAsync = NRetry::Async::TRetryContext<TQueryClient, TAsyncExecuteQueryResult>;
+using TRetryContextAsync = NRetry::Async::TRetryContext<TQueryClient, TAsyncStatus>;
 
 NYdb::NRetry::TRetryOperationSettings GetRetrySettings(TDuration timeout, bool isIndempotent) {
     return NYdb::NRetry::TRetryOperationSettings()
@@ -577,20 +578,40 @@ int64_t TQueryClient::GetCurrentPoolSize() const {
     return Impl_->GetCurrentPoolSize();
 }
 
-TAsyncExecuteQueryResult TQueryClient::RetryQuery(TQueryFunc&& queryFunc, TRetryOperationSettings settings)
+TAsyncExecuteQueryResult TQueryClient::RetryQuery(TQueryResultFunc&& queryFunc, TRetryOperationSettings settings)
 {
+    TRetryContextResultAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings));
+    return ctx->Execute();
+}
+
+TAsyncStatus TQueryClient::RetryQuery(TQueryFunc&& queryFunc, TRetryOperationSettings settings) {
     TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings));
     return ctx->Execute();
 }
 
+TAsyncStatus TQueryClient::RetryQuery(TQueryWithoutSessionFunc&& queryFunc, TRetryOperationSettings settings) {
+    TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithoutSession(*this, std::move(queryFunc), settings));
+    return ctx->Execute();
+}
+
+TStatus TQueryClient::RetryQuery(const TQuerySyncFunc& queryFunc, TRetryOperationSettings settings) {
+    NRetry::Sync::TRetryWithSession ctx(*this, queryFunc, settings);
+    return ctx.Execute();
+}
+
+TStatus TQueryClient::RetryQuery(const TQueryWithoutSessionSyncFunc& queryFunc, TRetryOperationSettings settings) {
+    NRetry::Sync::TRetryWithoutSession ctx(*this, queryFunc, settings);
+    return ctx.Execute();
+}
+
 TAsyncExecuteQueryResult TQueryClient::RetryQuery(const std::string& query, const TTxControl& txControl,
     TDuration timeout, bool isIndempotent)
 {
     auto settings = GetRetrySettings(timeout, isIndempotent);
     auto queryFunc = [&query, &txControl](TSession session, TDuration duration) -> TAsyncExecuteQueryResult {
         return session.ExecuteQuery(query, txControl, TExecuteQuerySettings().ClientTimeout(duration));
     };
-    TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings));
+    TRetryContextResultAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings));
     return ctx->Execute();
 }