Merge 81552aa into 27eb560

yumkam · web-flow · commit ad4950ad3f22 · 2025-01-16T09:36:51.000Z
diff --git a/ydb/library/yql/providers/generic/actors/ut/yql_generic_lookup_actor_ut.cpp b/ydb/library/yql/providers/generic/actors/ut/yql_generic_lookup_actor_ut.cpp
@@ -226,4 +226,189 @@ Y_UNIT_TEST_SUITE(GenericProviderLookupActor) {
         }
     }
 
+    Y_UNIT_TEST(LookupWithErrors) {
+        auto alloc = std::make_shared<NKikimr::NMiniKQL::TScopedAlloc>(__LOCATION__, NKikimr::TAlignedPagePoolCounters(), true, false);
+        NKikimr::NMiniKQL::TMemoryUsageInfo memUsage("TestMemUsage");
+        NKikimr::NMiniKQL::THolderFactory holderFactory(alloc->Ref(), memUsage);
+        NKikimr::NMiniKQL::TTypeEnvironment typeEnv(*alloc);
+        NKikimr::NMiniKQL::TTypeBuilder typeBuilder(typeEnv);
+
+        auto loggerConfig = NYql::NProto::TLoggingConfig();
+        loggerConfig.set_allcomponentslevel(::NYql::NProto::TLoggingConfig_ELevel::TLoggingConfig_ELevel_TRACE);
+        NYql::NLog::InitLogger(loggerConfig, false);
+
+        TTestActorRuntimeBase runtime(1, 1, true);
+        runtime.Initialize();
+        auto edge = runtime.AllocateEdgeActor();
+
+        NYql::TGenericDataSourceInstance dsi;
+        dsi.Setkind(NYql::EGenericDataSourceKind::YDB);
+        dsi.mutable_endpoint()->Sethost("some_host");
+        dsi.mutable_endpoint()->Setport(2135);
+        dsi.Setdatabase("some_db");
+        dsi.Setuse_tls(true);
+        dsi.set_protocol(::NYql::EGenericProtocol::NATIVE);
+        auto token = dsi.mutable_credentials()->mutable_token();
+        token->Settype("IAM");
+        token->Setvalue("TEST_TOKEN");
+
+        auto connectorMock = std::make_shared<NYql::NConnector::NTest::TConnectorClientMock>();
+
+        // clang-format off
+        // step 1: ListSplits
+        {
+            ::testing::InSequence seq;
+            for (grpc::StatusCode readStatus : { grpc::StatusCode::UNAVAILABLE, grpc::StatusCode::OK }) {
+                for (grpc::StatusCode listStatus : { grpc::StatusCode::UNAVAILABLE, readStatus == grpc::StatusCode::OK ? grpc::StatusCode::DEADLINE_EXCEEDED : grpc::StatusCode::UNAVAILABLE, grpc::StatusCode::OK }) {
+                    auto listBuilder = connectorMock->ExpectListSplits();
+                    listBuilder
+                        .Select()
+                            .DataSourceInstance(dsi)
+                            .What()
+                                .Column("id", Ydb::Type::UINT64)
+                                .NullableColumn("optional_id", Ydb::Type::UINT64)
+                                .NullableColumn("string_value", Ydb::Type::STRING)
+                                .Done()
+                            .Table("lookup_test")
+                            .Where()
+                                .Filter()
+                                    .Disjunction()
+                                        .Operand()
+                                            .Conjunction()
+                                                .Operand().Equal().Column("id").Value<ui64>(2).Done().Done()
+                                                .Operand().Equal().Column("optional_id").OptionalValue<ui64>(102).Done().Done()
+                                                .Done()
+                                            .Done()
+                                        .Operand()
+                                            .Conjunction()
+                                                .Operand().Equal().Column("id").Value<ui64>(1).Done().Done()
+                                                .Operand().Equal().Column("optional_id").OptionalValue<ui64>(101).Done().Done()
+                                                .Done()
+                                            .Done()
+                                        .Operand()
+                                            .Conjunction()
+                                                .Operand().Equal().Column("id").Value<ui64>(0).Done().Done()
+                                                .Operand().Equal().Column("optional_id").OptionalValue<ui64>(100).Done().Done()
+                                                .Done()
+                                            .Done()
+                                        .Done()
+                                    .Done()
+                                .Done()
+                            .Done()
+                        .MaxSplitCount(1)
+                    ;
+                    if (listStatus != grpc::StatusCode::OK) {
+                        listBuilder
+                            .Status(NYdbGrpc::TGrpcStatus(listStatus, "Mocked Error"))
+                        ;
+                        continue;
+                    }
+                    listBuilder
+                        .Result()
+                            .AddResponse(NewSuccess())
+                                .Description("Actual split info is not important")
+                    ;
+                }
+
+                auto readBuilder = connectorMock->ExpectReadSplits();
+                readBuilder
+                    .DataSourceInstance(dsi)
+                    .Filtering(NYql::NConnector::NApi::TReadSplitsRequest::FILTERING_MANDATORY)
+                    .Split()
+                        .Description("Actual split info is not important")
+                        .Done()
+                ;
+                if (readStatus != grpc::StatusCode::OK) {
+                    readBuilder
+                        .Status(NYdbGrpc::TGrpcStatus(readStatus, "Mocked Error"))
+                    ;
+                    continue;
+                }
+                readBuilder
+                    .Result()
+                        .AddResponse(
+                            MakeRecordBatch(
+                                MakeArray<arrow::UInt64Builder, uint64_t>("id", {0, 1, 2}, arrow::uint64()),
+                                MakeArray<arrow::UInt64Builder, uint64_t>("optional_id", {100, 101, 103}, arrow::uint64()), // the last value is intentially wrong
+                                MakeArray<arrow::StringBuilder, std::string>("string_value", {"a", "b", "c"}, arrow::utf8())
+                            ),
+                            NewSuccess()
+                        )
+                ;
+            }
+        }
+        // clang-format on
+
+        NYql::Generic::TLookupSource lookupSourceSettings;
+        *lookupSourceSettings.mutable_data_source_instance() = dsi;
+        lookupSourceSettings.Settable("lookup_test");
+        lookupSourceSettings.SetServiceAccountId("testsaid");
+        lookupSourceSettings.SetServiceAccountIdSignature("fake_signature");
+
+        google::protobuf::Any packedLookupSource;
+        Y_ABORT_UNLESS(packedLookupSource.PackFrom(lookupSourceSettings));
+
+        NKikimr::NMiniKQL::TStructTypeBuilder keyTypeBuilder{typeEnv};
+        keyTypeBuilder.Add("id", typeBuilder.NewDataType(NYql::NUdf::EDataSlot::Uint64, false));
+        keyTypeBuilder.Add("optional_id", typeBuilder.NewDataType(NYql::NUdf::EDataSlot::Uint64, true));
+        NKikimr::NMiniKQL::TStructTypeBuilder outputypeBuilder{typeEnv};
+        outputypeBuilder.Add("string_value", typeBuilder.NewDataType(NYql::NUdf::EDataSlot::String, true));
+
+        auto guard = Guard(*alloc.get());
+        auto keyTypeHelper = std::make_shared<NYql::NDq::IDqAsyncLookupSource::TKeyTypeHelper>(keyTypeBuilder.Build());
+
+        auto [lookupSource, actor] = NYql::NDq::CreateGenericLookupActor(
+            connectorMock,
+            std::make_shared<NYql::NTestCreds::TSecuredServiceAccountCredentialsFactory>(),
+            edge,
+            nullptr,
+            alloc,
+            keyTypeHelper,
+            std::move(lookupSourceSettings),
+            keyTypeBuilder.Build(),
+            outputypeBuilder.Build(),
+            typeEnv,
+            holderFactory,
+            1'000'000);
+        auto lookupActor = runtime.Register(actor);
+
+        auto request = std::make_shared<NYql::NDq::IDqAsyncLookupSource::TUnboxedValueMap>(3, keyTypeHelper->GetValueHash(), keyTypeHelper->GetValueEqual());
+        for (size_t i = 0; i != 3; ++i) {
+            NYql::NUdf::TUnboxedValue* keyItems;
+            auto key = holderFactory.CreateDirectArrayHolder(2, keyItems);
+            keyItems[0] = NYql::NUdf::TUnboxedValuePod(ui64(i));
+            keyItems[1] = NYql::NUdf::TUnboxedValuePod(ui64(100 + i));
+            request->emplace(std::move(key), NYql::NUdf::TUnboxedValue{});
+        }
+
+        guard.Release(); // let actors use alloc
+
+        auto callLookupActor = new TCallLookupActor(alloc, lookupActor, request);
+        runtime.Register(callLookupActor);
+
+        auto ev = runtime.GrabEdgeEventRethrow<NYql::NDq::IDqAsyncLookupSource::TEvLookupResult>(edge);
+        auto guard2 = Guard(*alloc.get());
+        auto lookupResult = ev->Get()->Result.lock();
+        UNIT_ASSERT(lookupResult);
+
+        UNIT_ASSERT_EQUAL(3, lookupResult->size());
+        {
+            const auto* v = lookupResult->FindPtr(CreateStructValue(holderFactory, {0, 100}));
+            UNIT_ASSERT(v);
+            NYql::NUdf::TUnboxedValue val = v->GetElement(0);
+            UNIT_ASSERT(val.AsStringRef() == TStringBuf("a"));
+        }
+        {
+            const auto* v = lookupResult->FindPtr(CreateStructValue(holderFactory, {1, 101}));
+            UNIT_ASSERT(v);
+            NYql::NUdf::TUnboxedValue val = v->GetElement(0);
+            UNIT_ASSERT(val.AsStringRef() == TStringBuf("b"));
+        }
+        {
+            const auto* v = lookupResult->FindPtr(CreateStructValue(holderFactory, {2, 102}));
+            UNIT_ASSERT(v);
+            UNIT_ASSERT(!*v);
+        }
+    }
+
 } // Y_UNIT_TEST_SUITE(GenericProviderLookupActor)
diff --git a/ydb/library/yql/providers/generic/actors/yql_generic_lookup_actor.cpp b/ydb/library/yql/providers/generic/actors/yql_generic_lookup_actor.cpp
@@ -29,6 +29,7 @@ namespace NYql::NDq {
 
     namespace {
         constexpr ui32 RequestRetriesLimit = 10; // TODO lookup parameters or PRAGMA?
+        constexpr TDuration RequestTimeout = TDuration::Minutes(3); // TODO lookup parameters or PRAGMA?
 
         const NKikimr::NMiniKQL::TStructType* MergeStructTypes(const NKikimr::NMiniKQL::TTypeEnvironment& env, const NKikimr::NMiniKQL::TStructType* t1, const NKikimr::NMiniKQL::TStructType* t2) {
             Y_ABORT_UNLESS(t1);
@@ -195,7 +196,7 @@ namespace NYql::NDq {
             *readRequest.add_splits() = split;
             readRequest.Setformat(NConnector::NApi::TReadSplitsRequest_EFormat::TReadSplitsRequest_EFormat_ARROW_IPC_STREAMING);
             readRequest.set_filtering(NConnector::NApi::TReadSplitsRequest::FILTERING_MANDATORY);
-            Connector->ReadSplits(readRequest).Subscribe([
+            Connector->ReadSplits(readRequest, RequestTimeout).Subscribe([
                     actorSystem = TActivationContext::ActorSystem(),
                     selfId = SelfId(),
                     retriesRemaining = RetriesRemaining
@@ -284,7 +285,7 @@ namespace NYql::NDq {
             };
 
             splitRequest.Setmax_split_count(1);
-            Connector->ListSplits(splitRequest).Subscribe([
+            Connector->ListSplits(splitRequest, RequestTimeout).Subscribe([
                     actorSystem = TActivationContext::ActorSystem(),
                     selfId = SelfId(),
                     retriesRemaining = RetriesRemaining
@@ -395,13 +396,17 @@ namespace NYql::NDq {
         }
 
         static void SendRetryOrError(NActors::TActorSystem* actorSystem, const NActors::TActorId& selfId, const NYdbGrpc::TGrpcStatus& status, ui32 retriesRemaining) {
-            if (NConnector::GrpcStatusNeedsRetry(status)) {
+            if (NConnector::GrpcStatusNeedsRetry(status) || status.GRpcStatusCode == grpc::DEADLINE_EXCEEDED) {
                 if (retriesRemaining) {
                     const auto retry = RequestRetriesLimit - retriesRemaining;
                     const auto delay = TDuration::MilliSeconds(1u << retry); // Exponential delay from 1ms to ~0.5s
                     // <<< TODO tune/tweak
                     YQL_CLOG(WARN, ProviderGeneric) << "ActorId=" << selfId << " Got retrievable GRPC Error from Connector: " << status.ToDebugString() << ", retry " << (retry + 1) << " of " << RequestRetriesLimit << ", scheduled in " << delay;
                     --retriesRemaining;
+                    if (status.GRpcStatusCode == grpc::DEADLINE_EXCEEDED) {
+                        // if error was deadline, retry only once
+                        retriesRemaining = 0; // TODO tune/tweak
+                    }
                     actorSystem->Schedule(delay, new IEventHandle(selfId, selfId, new TEvRetry(retriesRemaining)));
                     return;
                 }
diff --git a/ydb/library/yql/providers/generic/connector/libcpp/client.cpp b/ydb/library/yql/providers/generic/connector/libcpp/client.cpp
@@ -47,16 +47,16 @@ namespace NYql::NConnector {
             GrpcConnection_ = GrpcClient_->CreateGRpcServiceConnection<NApi::Connector>(GrpcConfig_);
         }
 
-        virtual TDescribeTableAsyncResult DescribeTable(const NApi::TDescribeTableRequest& request) override {
-            return UnaryCall<NApi::TDescribeTableRequest, NApi::TDescribeTableResponse>(request, &NApi::Connector::Stub::AsyncDescribeTable);
+        virtual TDescribeTableAsyncResult DescribeTable(const NApi::TDescribeTableRequest& request, TDuration timeout = {}) override {
+            return UnaryCall<NApi::TDescribeTableRequest, NApi::TDescribeTableResponse>(request, &NApi::Connector::Stub::AsyncDescribeTable, timeout);
         }
 
-        virtual TListSplitsStreamIteratorAsyncResult ListSplits(const NApi::TListSplitsRequest& request) override {
-            return ServerSideStreamingCall<NApi::TListSplitsRequest, NApi::TListSplitsResponse>(request, &NApi::Connector::Stub::AsyncListSplits);
+        virtual TListSplitsStreamIteratorAsyncResult ListSplits(const NApi::TListSplitsRequest& request, TDuration timeout = {}) override {
+            return ServerSideStreamingCall<NApi::TListSplitsRequest, NApi::TListSplitsResponse>(request, &NApi::Connector::Stub::AsyncListSplits, timeout);
         }
 
-        virtual TReadSplitsStreamIteratorAsyncResult ReadSplits(const NApi::TReadSplitsRequest& request) override {
-            return ServerSideStreamingCall<NApi::TReadSplitsRequest, NApi::TReadSplitsResponse>(request, &NApi::Connector::Stub::AsyncReadSplits);
+        virtual TReadSplitsStreamIteratorAsyncResult ReadSplits(const NApi::TReadSplitsRequest& request, TDuration timeout = {}) override {
+            return ServerSideStreamingCall<NApi::TReadSplitsRequest, NApi::TReadSplitsResponse>(request, &NApi::Connector::Stub::AsyncReadSplits, timeout);
         }
 
         ~TClientGRPC() {
@@ -74,7 +74,7 @@ namespace NYql::NConnector {
         template <class TRequest, class TResponse>
         TAsyncResult<TResponse> UnaryCall(
             const TRequest& request,
-            typename NYdbGrpc::TSimpleRequestProcessor<NApi::Connector::Stub, TRequest, TResponse>::TAsyncRequest rpc) {
+            typename NYdbGrpc::TSimpleRequestProcessor<NApi::Connector::Stub, TRequest, TResponse>::TAsyncRequest rpc, TDuration timeout = {}) {
             auto context = GrpcClient_->CreateContext();
             if (!context) {
                 throw yexception() << "Client is being shutted down";
@@ -89,7 +89,7 @@ namespace NYql::NConnector {
                 std::move(request),
                 std::move(callback),
                 rpc,
-                {},
+                { .Timeout = timeout },
                 context.get());
 
             return promise.GetFuture();
@@ -98,7 +98,8 @@ namespace NYql::NConnector {
         template <class TRequest, class TResponse>
         TIteratorAsyncResult<IStreamIterator<TResponse>> ServerSideStreamingCall(
             const TRequest& request,
-            TStreamRpc<NApi::Connector::Stub, TRequest, TResponse, NYdbGrpc::TStreamRequestReadProcessor> rpc) {
+            TStreamRpc<NApi::Connector::Stub, TRequest, TResponse, NYdbGrpc::TStreamRequestReadProcessor> rpc,
+            TDuration timeout = {}) {
             using TStreamProcessorPtr = typename NYdbGrpc::IStreamRequestReadProcessor<TResponse>::TPtr;
             using TStreamInitResult = std::pair<NYdbGrpc::TGrpcStatus, TStreamProcessorPtr>;
 
@@ -115,7 +116,7 @@ namespace NYql::NConnector {
                     promise.SetValue({std::move(status), streamProcessor});
                 },
                 rpc,
-                {},
+                { .Timeout = timeout },
                 context.get());
 
             // TODO: async handling YQ-2513
diff --git a/ydb/library/yql/providers/generic/connector/libcpp/client.h b/ydb/library/yql/providers/generic/connector/libcpp/client.h
@@ -89,9 +89,9 @@ namespace NYql::NConnector {
     public:
         using TPtr = std::shared_ptr<IClient>;
 
-        virtual TDescribeTableAsyncResult DescribeTable(const NApi::TDescribeTableRequest& request) = 0;
-        virtual TListSplitsStreamIteratorAsyncResult ListSplits(const NApi::TListSplitsRequest& request) = 0;
-        virtual TReadSplitsStreamIteratorAsyncResult ReadSplits(const NApi::TReadSplitsRequest& request) = 0;
+        virtual TDescribeTableAsyncResult DescribeTable(const NApi::TDescribeTableRequest& request, TDuration timeout = {}) = 0;
+        virtual TListSplitsStreamIteratorAsyncResult ListSplits(const NApi::TListSplitsRequest& request, TDuration timeout = {}) = 0;
+        virtual TReadSplitsStreamIteratorAsyncResult ReadSplits(const NApi::TReadSplitsRequest& request, TDuration timeout = {}) = 0;
         virtual ~IClient() = default;
     };
 
diff --git a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.h b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.h
diff --git a/ydb/library/yql/providers/generic/provider/ut/pushdown/pushdown_ut.cpp b/ydb/library/yql/providers/generic/provider/ut/pushdown/pushdown_ut.cpp
