Moved commit "Remove session from pool in case of BAD_SESSION and other errors which should trigger removing session from pool" from ydb repo

Author: dcherednik

include/ydb-cpp-sdk/client/query/client.h

@@ -130,6 +130,7 @@ class TTransaction;
 class TSession {
     friend class TQueryClient;
     friend class TTransaction;
+    friend class TExecuteQueryIterator;
 public:
     const std::string& GetId() const;
 

src/client/impl/ydb_internal/session_pool/session_pool.h

@@ -51,7 +51,9 @@ NThreading::TFuture<TResponse> InjectSessionStatusInterception(
         // Exclude CLIENT_RESOURCE_EXHAUSTED from transport errors which can cause to session disconnect
         // since we have guarantee this request wasn't been started to execute.
 
-        if (status.IsTransportError() && status.GetStatus() != EStatus::CLIENT_RESOURCE_EXHAUSTED) {
+        if (status.IsTransportError()
+            && status.GetStatus() != EStatus::CLIENT_RESOURCE_EXHAUSTED && status.GetStatus() != EStatus::CLIENT_OUT_OF_RANGE)
+        {
             impl->MarkBroken();
         } else if (status.GetStatus() == EStatus::SESSION_BUSY) {
             impl->MarkBroken();
@@ -70,6 +72,7 @@ NThreading::TFuture<TResponse> InjectSessionStatusInterception(
                 impl->ScheduleTimeToTouch(RandomizeThreshold(timeout), impl->GetState() == TKqpSessionCommon::EState::S_ACTIVE);
             }
         }
+
         if (cb) {
             cb(value, *impl);
         }

src/client/query/impl/exec_query.cpp

@@ -1,9 +1,11 @@
 #define INCLUDE_YDB_INTERNAL_H
 #include "exec_query.h"
+#include "client_session.h"
 
 #include <ydb-cpp-sdk/client/query/client.h>
 #include <src/client/impl/ydb_internal/make_request/make.h>
 #include <src/client/impl/ydb_internal/kqp_session_common/kqp_session_common.h>
+#include <src/client/impl/ydb_internal/session_pool/session_pool.h>
 #include <src/client/common_client/impl/client.h>
 #undef INCLUDE_YDB_INTERNAL_H
 
@@ -59,7 +61,7 @@ class TExecuteQueryIterator::TReaderImpl {
         return Finished_;
     }
 
-    TAsyncExecuteQueryPart ReadNext(std::shared_ptr<TSelf> self) {
+    TAsyncExecuteQueryPart DoReadNext(std::shared_ptr<TSelf> self) {
         auto promise = NThreading::NewPromise<TExecuteQueryPart>();
         // Capture self - guarantee no dtor call during the read
         auto readCb = [self, promise](TGRpcStatus&& grpcStatus) mutable {
@@ -100,6 +102,18 @@ class TExecuteQueryIterator::TReaderImpl {
         StreamProcessor_->Read(&Response_, readCb);
         return promise.GetFuture();
     }
+
+    TAsyncExecuteQueryPart ReadNext(std::shared_ptr<TSelf> self) {
+        if (!Session_)
+            return DoReadNext(std::move(self));
+
+        return NSessionPool::InjectSessionStatusInterception(
+            Session_->SessionImpl_,
+            DoReadNext(std::move(self)),
+            false, // no need to ping stream session
+            TDuration::Zero());
+    }
+
 private:
     TStreamProcessorPtr StreamProcessor_;
     TResponse Response_;