generic lookup retries: add tests

Author: yumkam

ydb/library/yql/providers/generic/actors/ut/yql_generic_lookup_actor_ut.cpp

@@ -226,4 +226,191 @@ 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::DEADLINE_EXCEEDED, grpc::StatusCode::OK }) {
+                for (grpc::StatusCode listStatus : { grpc::StatusCode::UNAVAILABLE, grpc::StatusCode::DEADLINE_EXCEEDED, 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)
+                    ;
+                    Cerr << "Simulate " << (int)listStatus << Endl;
+                    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()
+                ;
+                Cerr << "Simulate " << (int)readStatus << Endl;
+                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)

ydb/library/yql/providers/generic/actors/yql_generic_lookup_actor.cpp

@@ -397,8 +397,10 @@ namespace NYql::NDq {
         static void SendRetryOrError(NActors::TActorSystem* actorSystem, const NActors::TActorId& selfId, const NYdbGrpc::TGrpcStatus& status, const ui32 retriesRemaining) {
             if (NConnector::GrpcStatusNeedsRetry(status) || status.GRpcStatusCode == grpc::DEADLINE_EXCEEDED) {
                 if (retriesRemaining) {
+                    [[maybe_unused]]
                     const auto retry = RequestRetriesLimit - retriesRemaining;
                     // XXX FIXME tune/tweak
+                    [[maybe_unused]]
                     const auto delay = TDuration::MilliSeconds(1u << retry); // Exponential delay from 1ms to 1s
                     YQL_CLOG(WARN, ProviderGeneric) << "ActorId=" << selfId << " Got retrievable GRPC Error from Connector: " << status.ToDebugString() << ", retry " << (retry + 1) << " of " << RequestRetriesLimit << ", scheduled in " << delay;
                     actorSystem->Schedule(delay, new IEventHandle(selfId, selfId, new TEvRetry()));

ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.h

@@ -681,19 +681,25 @@ namespace NYql::NConnector::NTest {
                 return TListSplitsResultBuilder<TBuilder>(ResponseResult_, this);
             }
 
+            auto& Status(const NYdbGrpc::TGrpcStatus &status) {
+                ResponseStatus_ = status;
+                return *this;
+            }
+
             void FillWithDefaults() {
                 Result();
             }
 
         private:
             void SetExpectation() {
                 EXPECT_CALL(*Mock_, ListSplitsImpl(ProtobufRequestMatcher(*Result_)))
-                    .WillOnce(Return(TIteratorResult<IListSplitsStreamIterator>{NYdbGrpc::TGrpcStatus(), ResponseResult_}));
+                    .WillOnce(Return(TIteratorResult<IListSplitsStreamIterator>{ResponseStatus_, ResponseResult_}));
             }
 
         private:
             TConnectorClientMock* Mock_ = nullptr;
             TListSplitsStreamIteratorMock::TPtr ResponseResult_ = std::make_shared<TListSplitsStreamIteratorMock>();
+            NYdbGrpc::TGrpcStatus ResponseStatus_ {};
         };
 
         template <class TParent = void /* no parent by default */>
@@ -751,19 +757,25 @@ namespace NYql::NConnector::NTest {
                 return TReadSplitsResultBuilder<TBuilder>(ResponseResult_, this);
             }
 
+            auto& Status(const NYdbGrpc::TGrpcStatus &status) {
+                ResponseStatus_ = status;
+                return *this;
+            }
+
             void FillWithDefaults() {
                 Format(NApi::TReadSplitsRequest::ARROW_IPC_STREAMING);
             }
 
         private:
             void SetExpectation() {
                 EXPECT_CALL(*Mock_, ReadSplitsImpl(ProtobufRequestMatcher(*Result_)))
-                    .WillOnce(Return(TIteratorResult<IReadSplitsStreamIterator>{NYdbGrpc::TGrpcStatus(), ResponseResult_}));
+                    .WillOnce(Return(TIteratorResult<IReadSplitsStreamIterator>{ResponseStatus_, ResponseResult_}));
             }
 
         private:
             TConnectorClientMock* Mock_ = nullptr;
             TReadSplitsStreamIteratorMock::TPtr ResponseResult_ = std::make_shared<TReadSplitsStreamIteratorMock>();
+            NYdbGrpc::TGrpcStatus ResponseStatus_ {};
         };
 
         TDescribeTableExpectationBuilder ExpectDescribeTable() {