Supported join with external ydb

Author: GrigoriyPA

ydb/core/kqp/gateway/behaviour/streaming_query/optimization.cpp

@@ -39,7 +39,7 @@ bool ExploreStreamingQueryNode(TExprNode::TPtr node, TStreamingExploreCtx& res)
     const auto providerArg = node->ChildPtr(1);
     if (const auto maybeDataSource = TMaybeNode<TCoDataSource>(providerArg)) {
         const auto dataSourceCategory = maybeDataSource.Cast().Category().Value();
-        if (IsIn({NYql::PqProviderName, NYql::S3ProviderName}, dataSourceCategory)) {
+        if (IsIn({NYql::PqProviderName, NYql::S3ProviderName, NYql::GenericProviderName}, dataSourceCategory)) {
             res.StreamingReads += dataSourceCategory == NYql::PqProviderName;
             return true;
         }

ydb/core/kqp/ut/federated_query/datastreams/datastreams_ut.cpp

@@ -6,12 +6,12 @@
 #include <ydb/core/kqp/proxy_service/kqp_script_executions.h>
 #include <ydb/core/kqp/ut/common/kqp_ut_common.h>
 #include <ydb/core/kqp/ut/federated_query/common/common.h>
-
-#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/draft/ydb_scripting.h>
-
 #include <ydb/library/testlib/pq_helpers/mock_pq_gateway.h>
 #include <ydb/library/testlib/s3_recipe_helper/s3_recipe_helper.h>
+#include <ydb/library/yql/providers/generic/connector/libcpp/error.h>
+#include <ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.h>
 #include <ydb/library/yql/providers/s3/actors/yql_s3_actors_factory_impl.h>
+#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/draft/ydb_scripting.h>
 
 #include <fmt/format.h>
 
@@ -22,6 +22,8 @@ using namespace NYdb::NQuery;
 using namespace NKikimr::NKqp::NFederatedQueryTest;
 using namespace fmt::literals;
 using namespace NTestUtils;
+using namespace NYql::NConnector::NTest;
+using namespace NYql::NConnector::NApi;
 
 namespace {
 
@@ -65,6 +67,16 @@ class TStreamingTestFixture : public NUnitTest::TBaseFixture {
         return mockPqGateway;
     }
 
+    std::shared_ptr<TConnectorClientMock> SetupMockConnectorClient() {
+        UNIT_ASSERT_C(!ConnectorClient, "ConnectorClient is already initialized");
+        EnsureNotInitialized("ConnectorClient");
+
+        auto mockConnectorClient = std::make_shared<TConnectorClientMock>();
+        ConnectorClient = mockConnectorClient;
+
+        return mockConnectorClient;
+    }
+
     // Local kikimr test cluster
 
     std::shared_ptr<TKikimrRunner> GetKikimrRunner() {
@@ -79,7 +91,7 @@ class TStreamingTestFixture : public NUnitTest::TBaseFixture {
             queryServiceConfig.SetEnableMatchRecognize(true);
             queryServiceConfig.SetProgressStatsPeriodMs(1000);
 
-            Kikimr = MakeKikimrRunner(true, nullptr, nullptr, AppConfig, NYql::NDq::CreateS3ActorsFactory(), {
+            Kikimr = MakeKikimrRunner(true, ConnectorClient, nullptr, AppConfig, NYql::NDq::CreateS3ActorsFactory(), {
                 .PqGateway = PqGateway,
                 .CheckpointPeriod = CheckpointPeriod,
             });
@@ -608,6 +620,7 @@ class TStreamingTestFixture : public NUnitTest::TBaseFixture {
 private:
     std::optional<NKikimrConfig::TAppConfig> AppConfig;
     TIntrusivePtr<NYql::IPqGateway> PqGateway;
+    NYql::NConnector::IClient::TPtr ConnectorClient;
     std::shared_ptr<TKikimrRunner> Kikimr;
 
     std::shared_ptr<TDriver> InternalDriver;
@@ -1565,6 +1578,172 @@ Y_UNIT_TEST_SUITE(KqpStreamingQueriesDdl) {
         WaitMockPqReadSession(pqGateway, inputTopicName)->AddDataReceivedEvent(sampleMessages);
         ReadMockPqMessages(WaitMockPqWriteSession(pqGateway, outputTopicName), sampleResult);
     }
+
+    Y_UNIT_TEST_F(StreamingQueryWithYdbJoin, TStreamingTestFixture) {
+        // Test that defaults are overridden for streaming queries
+        auto& setting = *SetupAppConfig().MutableKQPConfig()->AddSettings();
+        setting.SetName("HashJoinMode");
+        setting.SetValue("grace");
+
+        const auto connectorClient = SetupMockConnectorClient();
+        const auto pqGateway = SetupMockPqGateway();
+
+        constexpr char inputTopicName[] = "inputTopicName";
+        constexpr char outputTopicName[] = "outputTopicName";
+        CreateTopic(inputTopicName);
+        CreateTopic(outputTopicName);
+
+        constexpr char pqSourceName[] = "pqSourceName";
+        CreatePqSource(pqSourceName);
+
+        constexpr char ydbSourceName[] = "ydbSourceName";
+        ExecQuery(fmt::format(R"(
+            CREATE OBJECT secret_name (TYPE SECRET) WITH (value = "{token}");
+            CREATE EXTERNAL DATA SOURCE `{ydb_source}` WITH (
+                SOURCE_TYPE = "Ydb",
+                LOCATION = "{ydb_location}",
+                DATABASE_NAME = "{ydb_database_name}",
+                AUTH_METHOD = "TOKEN",
+                TOKEN_SECRET_NAME = "secret_name",
+                USE_TLS = "FALSE"
+            );)",
+            "ydb_source"_a = ydbSourceName,
+            "ydb_location"_a = YDB_ENDPOINT,
+            "ydb_database_name"_a = YDB_DATABASE,
+            "token"_a = BUILTIN_ACL_ROOT
+        ));
+
+        constexpr char ydbTable[] = "lookup";
+        ExecExternalQuery(fmt::format(R"(
+            CREATE TABLE `{table}` (
+                fqdn String,
+                payload String,
+                PRIMARY KEY (fqdn)
+            ))",
+            "table"_a = ydbTable
+        ));
+
+        {   // Prepare connector mock
+            NYql::TGenericDataSourceInstance dataSourceInstance;
+            dataSourceInstance.set_kind(NYql::YDB);
+            dataSourceInstance.set_database(YDB_DATABASE);
+            dataSourceInstance.set_use_tls(false);
+            dataSourceInstance.set_protocol(NYql::NATIVE);
+
+            auto& endpoint = *dataSourceInstance.mutable_endpoint();
+            TIpPort port;
+            NHttp::CrackAddress(YDB_ENDPOINT, *endpoint.mutable_host(), port);
+            endpoint.set_port(port);
+
+            auto& iamToken = *dataSourceInstance.mutable_credentials()->mutable_token();
+            iamToken.set_type("IAM");
+            iamToken.set_value(BUILTIN_ACL_ROOT);
+
+            TTypeMappingSettings typeMappingSettings;
+            typeMappingSettings.set_date_time_format(STRING_FORMAT);
+
+            auto describeTableBuilder = connectorClient->ExpectDescribeTable();
+            describeTableBuilder
+                .Table(ydbTable)
+                .DataSourceInstance(dataSourceInstance)
+                .TypeMappingSettings(typeMappingSettings);
+
+            auto listSplitsBuilder = connectorClient->ExpectListSplits();
+            listSplitsBuilder.Select()
+                .DataSourceInstance(dataSourceInstance)
+                .Table(ydbTable);
+
+            const std::vector<std::string> fqdnColumn = {"host1.example.com", "host2.example.com", "host3.example.com"};
+            const std::vector<std::string> payloadColumn = {"P1", "P2", "P3"};
+            auto readSplitsBuilder = connectorClient->ExpectReadSplits();
+            readSplitsBuilder
+                .Filtering(TReadSplitsRequest::FILTERING_OPTIONAL)
+                .Split()
+                    .Description("some binary description")
+                    .Select()
+                        .Table(ydbTable)
+                        .DataSourceInstance(dataSourceInstance)
+                        .What()
+                            .Column("fqdn", Ydb::Type::STRING)
+                            .Column("payload", Ydb::Type::STRING);
+
+            const auto builtResults = [&]() {
+                describeTableBuilder.Response()
+                    .Column("fqdn", Ydb::Type::STRING)
+                    .Column("payload", Ydb::Type::STRING);
+
+                listSplitsBuilder.Result()
+                    .AddResponse(NYql::NConnector::NewSuccess())
+                        .Description("some binary description")
+                        .Select()
+                            .DataSourceInstance(dataSourceInstance)
+                            .What()
+                                .Column("fqdn", Ydb::Type::STRING)
+                                .Column("payload", Ydb::Type::STRING);
+
+                readSplitsBuilder.Result()
+                    .AddResponse(
+                        MakeRecordBatch(
+                            MakeArray<arrow::BinaryBuilder>("fqdn", fqdnColumn, arrow::binary()),
+                            MakeArray<arrow::BinaryBuilder>("payload", payloadColumn, arrow::binary())
+                        ),
+                        NYql::NConnector::NewSuccess()
+                    );
+            };
+
+            builtResults();
+            builtResults(); // Streaming queries compiled twice, also in test results requested twice due to retry
+        }
+
+        constexpr char queryName[] = "streamingQuery";
+        ExecQuery(fmt::format(R"(
+            CREATE STREAMING QUERY `{query_name}` AS
+            DO BEGIN
+                $ydb_lookup = SELECT * FROM `{ydb_source}`.`{ydb_table}`;
+
+                $pq_source = SELECT * FROM `{pq_source}`.`{input_topic}` WITH (
+                    FORMAT = "json_each_row",
+                    SCHEMA (
+                        time Int32 NOT NULL,
+                        event String,
+                        host String
+                    )
+                );
+
+                $joined = SELECT l.payload AS payload, p.* FROM $pq_source AS p
+                LEFT JOIN $ydb_lookup AS l
+                ON (l.fqdn = p.host);
+
+                INSERT INTO `{pq_source}`.`{output_topic}`
+                SELECT Unwrap(event || "-" || payload) FROM $joined
+            END DO;)",
+            "query_name"_a = queryName,
+            "pq_source"_a = pqSourceName,
+            "ydb_source"_a = ydbSourceName,
+            "ydb_table"_a = ydbTable,
+            "input_topic"_a = inputTopicName,
+            "output_topic"_a = outputTopicName
+        ));
+
+        CheckScriptExecutionsCount(1, 1);
+
+        auto readSession = WaitMockPqReadSession(pqGateway, inputTopicName);
+        const std::vector<IMockPqReadSession::TMessage> sampleMessages = {
+            {0, R"({"time": 0, "event": "A", "host": "host1.example.com"})"},
+            {1, R"({"time": 1, "event": "B", "host": "host3.example.com"})"},
+            {2, R"({"time": 2, "event": "A", "host": "host1.example.com"})"},
+        };
+        readSession->AddDataReceivedEvent(sampleMessages);
+
+        auto writeSession = WaitMockPqWriteSession(pqGateway, outputTopicName);
+        const std::vector<TString> sampleResult = {"A-P1", "B-P3", "A-P1"};
+        ReadMockPqMessages(writeSession, sampleResult);
+
+        readSession->AddCloseSessionEvent(EStatus::UNAVAILABLE, {NIssue::TIssue("Test pq session failure")});
+
+        WaitMockPqReadSession(pqGateway, inputTopicName)->AddDataReceivedEvent(sampleMessages);
+        ReadMockPqMessages(WaitMockPqWriteSession(pqGateway, outputTopicName), sampleResult);
+    }
 }
 
 } // namespace NKikimr::NKqp

ydb/core/kqp/ut/federated_query/datastreams/ya.make

@@ -16,6 +16,8 @@ PEERDIR(
     ydb/core/kqp/ut/federated_query/common
     ydb/library/testlib/pq_helpers
     ydb/library/testlib/s3_recipe_helper
+    ydb/library/yql/providers/generic/connector/libcpp
+    ydb/library/yql/providers/generic/connector/libcpp/ut_helpers
     yql/essentials/sql/pg
     yql/essentials/parser/pg_wrapper
 )

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

@@ -373,12 +373,11 @@ namespace NYql::NConnector::NTest {
             DATA_SOURCE_INSTANCE_SUBBUILDER();
 
             TDescribeTableResultBuilder<TBuilder> Response() {
-                return TDescribeTableResultBuilder<TBuilder>(ResponseResult_, this);
+                return TDescribeTableResultBuilder<TBuilder>(ResponseResults_.emplace_back(std::make_shared<NApi::TDescribeTableResponse>()), this);
             }
 
             void FillWithDefaults() {
                 Table(DEFAULT_TABLE);
-                Response();
             }
 
             TBuilder& TypeMappingSettings(const NApi::TTypeMappingSettings& proto) {
@@ -388,16 +387,19 @@ namespace NYql::NConnector::NTest {
 
         private:
             void SetExpectation() {
-                EXPECT_CALL(*Mock_, DescribeTableImpl(ProtobufRequestMatcher(*Result_)))
-                    .WillOnce(Return(
-                        TResult<NApi::TDescribeTableResponse>(
-                            {NYdbGrpc::TGrpcStatus(),
-                             *ResponseResult_})));
+                if (ResponseResults_.empty()) {
+                    Response();
+                }
+
+                auto& expectBuilder = EXPECT_CALL(*Mock_, DescribeTableImpl(ProtobufRequestMatcher(*Result_)));
+                for (auto result : ResponseResults_) {
+                    expectBuilder.WillOnce(Return(TResult<NApi::TDescribeTableResponse>({NYdbGrpc::TGrpcStatus(), *result})));
+                }
             }
 
         private:
             TConnectorClientMock* Mock_ = nullptr;
-            std::shared_ptr<NApi::TDescribeTableResponse> ResponseResult_ = std::make_shared<NApi::TDescribeTableResponse>();
+            std::vector<std::shared_ptr<NApi::TDescribeTableResponse>> ResponseResults_;
         };
 
         template <class TParent = void /* no parent by default */>
@@ -663,9 +665,7 @@ namespace NYql::NConnector::NTest {
             explicit TListSplitsExpectationBuilder(NApi::TListSplitsRequest* result = nullptr, TConnectorClientMock* mock = nullptr)
                 : TProtoBuilder<int, NApi::TListSplitsRequest>(result)
                 , Mock_(mock)
-
             {
-                FillWithDefaults();
             }
 
             explicit TListSplitsExpectationBuilder(TConnectorClientMock* mock)
@@ -681,27 +681,29 @@ namespace NYql::NConnector::NTest {
             SETTER(MaxSplitCount, max_split_count);
 
             TListSplitsResultBuilder<TBuilder> Result() {
-                return TListSplitsResultBuilder<TBuilder>(ResponseResult_, this);
+                return TListSplitsResultBuilder<TBuilder>(ResponseResults_.emplace_back(std::make_shared<TListSplitsStreamIteratorMock>()), 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>{ResponseStatus_, ResponseResult_}));
+                if (ResponseResults_.empty()) {
+                    Result();
+                }
+
+                auto& expectBuilder = EXPECT_CALL(*Mock_, ListSplitsImpl(ProtobufRequestMatcher(*Result_)));
+                for (auto response : ResponseResults_) {
+                    expectBuilder.WillOnce(Return(TIteratorResult<IListSplitsStreamIterator>{ResponseStatus_, response}));
+                }
             }
 
         private:
             TConnectorClientMock* Mock_ = nullptr;
-            TListSplitsStreamIteratorMock::TPtr ResponseResult_ = std::make_shared<TListSplitsStreamIteratorMock>();
+            std::vector<TListSplitsStreamIteratorMock::TPtr> ResponseResults_;
             NYdbGrpc::TGrpcStatus ResponseStatus_ {};
         };
 
@@ -757,7 +759,7 @@ namespace NYql::NConnector::NTest {
             SETTER(Filtering, filtering);
 
             TReadSplitsResultBuilder<TBuilder> Result() {
-                return TReadSplitsResultBuilder<TBuilder>(ResponseResult_, this);
+                return TReadSplitsResultBuilder<TBuilder>(ResponseResults_.emplace_back(std::make_shared<TReadSplitsStreamIteratorMock>()), this);
             }
 
             auto& Status(const NYdbGrpc::TGrpcStatus& status) {
@@ -771,13 +773,19 @@ namespace NYql::NConnector::NTest {
 
         private:
             void SetExpectation() {
-                EXPECT_CALL(*Mock_, ReadSplitsImpl(ProtobufRequestMatcher(*Result_)))
-                    .WillOnce(Return(TIteratorResult<IReadSplitsStreamIterator>{ResponseStatus_, ResponseResult_}));
+                if (ResponseResults_.empty()) {
+                    Result();
+                }
+
+                auto& expectBuilder = EXPECT_CALL(*Mock_, ReadSplitsImpl(ProtobufRequestMatcher(*Result_)));
+                for (auto response : ResponseResults_) {
+                    expectBuilder.WillOnce(Return(TIteratorResult<IReadSplitsStreamIterator>{ResponseStatus_, response}));
+                }
             }
 
         private:
             TConnectorClientMock* Mock_ = nullptr;
-            TReadSplitsStreamIteratorMock::TPtr ResponseResult_ = std::make_shared<TReadSplitsStreamIteratorMock>();
+            std::vector<TReadSplitsStreamIteratorMock::TPtr> ResponseResults_;
             NYdbGrpc::TGrpcStatus ResponseStatus_ {};
         };