merge to stable-25-3 YQ-4312 fixes for streaming queries (#28776)

GrigoriyPA · web-flow · commit 7d1a81bd900c · 2025-11-17T16:45:33.000+05:00
diff --git a/ydb/core/kqp/common/events/events.h b/ydb/core/kqp/common/events/events.h
@@ -140,6 +140,7 @@ struct TEvKqp {
         bool DisableDefaultTimeout = false;
         i64 Generation = 1;
         TString CheckpointId;
+        TString StreamingQueryPath;
     };
 
     struct TEvScriptResponse : public TEventLocal<TEvScriptResponse, TKqpEvents::EvScriptResponse> {
diff --git a/ydb/core/kqp/common/kqp_user_request_context.h b/ydb/core/kqp/common/kqp_user_request_context.h
@@ -21,6 +21,7 @@ namespace NKikimr::NKqp {
         std::optional<NResourcePool::TPoolSettings> PoolConfig;
         bool IsStreamingQuery = false;
         TString CheckpointId;
+        TString StreamingQueryPath;
 
         TUserRequestContext() = default;
 
diff --git a/ydb/core/kqp/executer_actor/kqp_data_executer.cpp b/ydb/core/kqp/executer_actor/kqp_data_executer.cpp
@@ -2822,7 +2822,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
             NYql::NDq::MakeCheckpointStorageID(),
             SelfId(),
             {},
-            Counters->Counters->GetKqpCounters(),
+            Counters->Counters->GetKqpCounters()->GetSubgroup("path", context->StreamingQueryPath),
             graphParams,
             stateLoadMode,
             streamingDisposition).Release());
diff --git a/ydb/core/kqp/executer_actor/kqp_tasks_graph.cpp b/ydb/core/kqp/executer_actor/kqp_tasks_graph.cpp
@@ -1559,7 +1559,6 @@ void TKqpTasksGraph::PersistTasksGraphInfo(NKikimrKqp::TQueryPhysicalGraph& resu
 
 void TKqpTasksGraph::RestoreTasksGraphInfo(const TVector<IKqpGateway::TPhysicalTxData>& transactions, const TVector<NKikimrKqp::TKqpNodeResources>& resourcesSnapshot, const NKikimrKqp::TQueryPhysicalGraph& graphInfo) {
     GetMeta().IsRestored = true;
-    GetMeta().AllowWithSpilling = false;
 
     const auto restoreDqTransform = [](const auto& protoInfo) -> TMaybe<TTransform> {
         if (!protoInfo.HasTransform()) {
@@ -1755,6 +1754,8 @@ void TKqpTasksGraph::RestoreTasksGraphInfo(const TVector<IKqpGateway::TPhysicalT
                 const auto it = scheduledTaskCount.find(stageIdx);
                 BuildReadTasksFromSource(stageInfo, resourcesSnapshot, it != scheduledTaskCount.end() ? it->second.TaskCount : 0);
             }
+
+            GetMeta().AllowWithSpilling |= stage.GetAllowWithSpilling();
         }
     }
 }
diff --git a/ydb/core/kqp/gateway/behaviour/streaming_query/queries.cpp b/ydb/core/kqp/gateway/behaviour/streaming_query/queries.cpp
@@ -1746,6 +1746,7 @@ class TStartStreamingQueryTableActor final : public TActionActorBase<TStartStrea
         ev->ForgetAfter = TDuration::Max();
         ev->Generation = PreviousGeneration + 1;
         ev->CheckpointId = State.GetCheckpointId();
+        ev->StreamingQueryPath = QueryPath;
 
         if (const auto statsPeriod = AppData()->QueryServiceConfig.GetProgressStatsPeriodMs()) {
             ev->ProgressStatsPeriod = TDuration::MilliSeconds(statsPeriod);
diff --git a/ydb/core/kqp/provider/yql_kikimr_datasource.cpp b/ydb/core/kqp/provider/yql_kikimr_datasource.cpp
@@ -116,9 +116,9 @@ class TKiSourceIntentDeterminationTransformer: public TKiSourceVisitorTransforme
         TExprBase currentNode(node);
         if (auto maybeReadTable = currentNode.Maybe<TKiReadTable>()) {
             auto readTable = maybeReadTable.Cast();
-            for (auto setting : readTable.Settings()) {
-                auto name = setting.Name().Value();
-                if (name == "sysViewRewritten") {
+            for (auto setting : readTable.Settings().Ref().ChildrenList()) {
+                auto maybeTuple = TMaybeNode<TCoNameValueTuple>(setting);
+                if (maybeTuple && maybeTuple.Cast().Name().Value() == "sysViewRewritten"sv) {
                     sysViewRewritten = true;
                 }
             }
diff --git a/ydb/core/kqp/proxy_service/kqp_script_executions.cpp b/ydb/core/kqp/proxy_service/kqp_script_executions.cpp
@@ -562,6 +562,7 @@ class TCreateScriptExecutionActor : public TActorBootstrapped<TCreateScriptExecu
             .PhysicalGraph = ev.QueryPhysicalGraph,
             .DisableDefaultTimeout = ev.DisableDefaultTimeout,
             .CheckpointId = ev.CheckpointId,
+            .StreamingQueryPath = ev.StreamingQueryPath
         }, QueryServiceConfig));
 
         const auto& creatorId = Register(new TCreateScriptOperationQuery(ExecutionId, RunScriptActorId, ev.Record, meta, MaxRunTime, GetRetryState(), ev.QueryPhysicalGraph, QueryServiceConfig, ev.Generation));
@@ -615,6 +616,7 @@ class TCreateScriptExecutionActor : public TActorBootstrapped<TCreateScriptExecu
         meta.SetTraceId(eventProto.GetTraceId());
         meta.SetResourcePoolId(request.GetPoolId());
         meta.SetCheckpointId(ev.CheckpointId);
+        meta.SetStreamingQueryPath(ev.StreamingQueryPath);
         meta.SetClientAddress(request.GetClientAddress());
         meta.SetCollectStats(request.GetCollectStats());
         meta.SetSaveQueryPhysicalGraph(ev.SaveQueryPhysicalGraph);
@@ -1092,6 +1094,7 @@ class TRestartScriptOperationQuery : public TQueryBase {
             .PhysicalGraph = std::move(physicalGraph),
             .DisableDefaultTimeout = meta.GetDisableDefaultTimeout(),
             .CheckpointId = meta.GetCheckpointId(),
+            .StreamingQueryPath = meta.GetStreamingQueryPath()
         }, QueryServiceConfig));
 
         KQP_PROXY_LOG_D("Restart with RunScriptActorId: " << RunScriptActorId << ", has PhysicalGraph: " << hasPhysicalGraph);
diff --git a/ydb/core/kqp/run_script_actor/kqp_run_script_actor.cpp b/ydb/core/kqp/run_script_actor/kqp_run_script_actor.cpp
@@ -127,6 +127,7 @@ class TRunScriptActor : public NActors::TActorBootstrapped<TRunScriptActor> {
         , DisableDefaultTimeout(settings.DisableDefaultTimeout)
         , CheckpointId(settings.CheckpointId)
         , PhysicalGraph(std::move(settings.PhysicalGraph))
+        , StreamingQueryPath(settings.StreamingQueryPath)
         , Counters(settings.Counters)
     {}
 
@@ -144,6 +145,7 @@ class TRunScriptActor : public NActors::TActorBootstrapped<TRunScriptActor> {
         );
         UserRequestContext->IsStreamingQuery = SaveQueryPhysicalGraph;
         UserRequestContext->CheckpointId = CheckpointId;
+        UserRequestContext->StreamingQueryPath = StreamingQueryPath;
 
         LOG_I("Bootstrap");
 
@@ -969,6 +971,7 @@ class TRunScriptActor : public NActors::TActorBootstrapped<TRunScriptActor> {
     const bool DisableDefaultTimeout = false;
     const TString CheckpointId;
     std::optional<NKikimrKqp::TQueryPhysicalGraph> PhysicalGraph;
+    const TString StreamingQueryPath;
     std::optional<TActorId> PhysicalGraphSender;
     TIntrusivePtr<TKqpCounters> Counters;
     TString SessionId;
diff --git a/ydb/core/kqp/run_script_actor/kqp_run_script_actor.h b/ydb/core/kqp/run_script_actor/kqp_run_script_actor.h
@@ -25,6 +25,7 @@ struct TKqpRunScriptActorSettings {
     std::optional<NKikimrKqp::TQueryPhysicalGraph> PhysicalGraph;
     bool DisableDefaultTimeout = false;
     TString CheckpointId;
+    TString StreamingQueryPath;
 };
 
 NActors::IActor* CreateRunScriptActor(const NKikimrKqp::TEvQueryRequest& request, TKqpRunScriptActorSettings&& settings, NKikimrConfig::TQueryServiceConfig queryServiceConfig);
diff --git a/ydb/core/kqp/ut/federated_query/datastreams/datastreams_ut.cpp b/ydb/core/kqp/ut/federated_query/datastreams/datastreams_ut.cpp
@@ -660,35 +660,39 @@ class TStreamingTestFixture : public NUnitTest::TBaseFixture {
         TTypeMappingSettings typeMappingSettings;
         typeMappingSettings.set_date_time_format(STRING_FORMAT);
 
-        auto describeTableBuilder = mockClient->ExpectDescribeTable();
-        describeTableBuilder
-            .Table(settings.TableName)
-            .DataSourceInstance(GetMockConnectorSourceInstance())
-            .TypeMappingSettings(typeMappingSettings);
-
-        auto listSplitsBuilder = mockClient->ExpectListSplits();
-        auto fillListSplitExpectation = listSplitsBuilder
-            .ValidateArgs(settings.ValidateListSplitsArgs ? TConnectorClientMock::EArgsValidation::Strict : TConnectorClientMock::EArgsValidation::DataSourceInstance)
-            .Select()
-                .DataSourceInstance(GetMockConnectorSourceInstance())
+        if (settings.DescribeCount) {
+            auto describeTableBuilder = mockClient->ExpectDescribeTable();
+            describeTableBuilder
                 .Table(settings.TableName)
-                .What();
-
-        FillMockConnectorRequestColumns(fillListSplitExpectation, settings.Columns);
+                .DataSourceInstance(GetMockConnectorSourceInstance())
+                .TypeMappingSettings(typeMappingSettings);
 
-        for (ui64 i = 0; i < settings.DescribeCount; ++i) {
-            auto responseBuilder = describeTableBuilder.Response();
-            FillMockConnectorRequestColumns(responseBuilder, settings.Columns);
+            for (ui64 i = 0; i < settings.DescribeCount; ++i) {
+                auto responseBuilder = describeTableBuilder.Response();
+                FillMockConnectorRequestColumns(responseBuilder, settings.Columns);
+            }
         }
 
-        for (ui64 i = 0; i < settings.ListSplitsCount; ++i) {
-            auto responseBuilder = listSplitsBuilder.Result()
-                .AddResponse(NYql::NConnector::NewSuccess())
-                    .Description("some binary description")
-                    .Select()
-                        .DataSourceInstance(GetMockConnectorSourceInstance())
-                        .What();
-            FillMockConnectorRequestColumns(responseBuilder, settings.Columns);
+        if (settings.ListSplitsCount) {
+            auto listSplitsBuilder = mockClient->ExpectListSplits();
+            auto fillListSplitExpectation = listSplitsBuilder
+                .ValidateArgs(settings.ValidateListSplitsArgs ? TConnectorClientMock::EArgsValidation::Strict : TConnectorClientMock::EArgsValidation::DataSourceInstance)
+                .Select()
+                    .DataSourceInstance(GetMockConnectorSourceInstance())
+                    .Table(settings.TableName)
+                    .What();
+
+            FillMockConnectorRequestColumns(fillListSplitExpectation, settings.Columns);
+
+            for (ui64 i = 0; i < settings.ListSplitsCount; ++i) {
+                auto responseBuilder = listSplitsBuilder.Result()
+                    .AddResponse(NYql::NConnector::NewSuccess())
+                        .Description("some binary description")
+                        .Select()
+                            .DataSourceInstance(GetMockConnectorSourceInstance())
+                            .What();
+                FillMockConnectorRequestColumns(responseBuilder, settings.Columns);
+            }
         }
     }
 
@@ -1561,6 +1565,51 @@ Y_UNIT_TEST_SUITE(KqpFederatedQueryDatastreams) {
         UNIT_ASSERT_STRING_CONTAINS(readyOp.Status().GetIssues().ToString(), "Runtime listing is not supported for streaming queries, pragma value was ignored");
         UNIT_ASSERT_VALUES_EQUAL(GetAllObjects(sourceBucket), "{\"data\":\"x\"}\n{\"data\": \"x\"}");
     }
+
+    Y_UNIT_TEST_F(CrossJoinWithNotExistingDataSource, TStreamingTestFixture) {
+        const auto connectorClient = SetupMockConnectorClient();
+
+        constexpr char ydbSourceName[] = "ydbSourceName";
+        CreateYdbSource(ydbSourceName);
+
+        constexpr char ydbTable[] = "unknownSourceLookup";
+        ExecExternalQuery(fmt::format(R"(
+            CREATE TABLE `{table}` (
+                fqdn String,
+                payload String,
+                PRIMARY KEY (fqdn)
+            ))",
+            "table"_a = ydbTable
+        ));
+
+        {   // Prepare connector mock
+            const std::vector<TColumn> columns = {
+                {"fqdn", Ydb::Type::STRING},
+                {"payload", Ydb::Type::STRING}
+            };
+            SetupMockConnectorTableDescription(connectorClient, {
+                .TableName = ydbTable,
+                .Columns = columns,
+                .DescribeCount = 1,
+                .ListSplitsCount = 0
+            });
+        }
+
+        ExecQuery(fmt::format(R"(
+                SELECT
+                    *
+                FROM `unknown-datasource`.`unknown-topic` WITH (
+                    FORMAT = raw,
+                    SCHEMA (Data String NOT NULL)
+                ) AS p
+                CROSS JOIN (
+                    SELECT * FROM `{ydb_source}`.`{table}`
+                ) AS l
+            )",
+            "ydb_source"_a = ydbSourceName,
+            "table"_a = ydbTable
+        ), EStatus::SCHEME_ERROR, "Cannot find table '/Root/unknown-datasource.[unknown-topic]' because it does not exist or you do not have access permissions");
+    }
 }
 
 Y_UNIT_TEST_SUITE(KqpStreamingQueriesDdl) {
@@ -2298,6 +2347,91 @@ Y_UNIT_TEST_SUITE(KqpStreamingQueriesDdl) {
         pqGateway->WaitWriteSession(outputTopicName)->ExpectMessages(sampleResult);
     }
 
+    Y_UNIT_TEST_F(StreamingQueryWithDoubleYdbJoin, TStreamingTestFixture) {
+        const auto connectorClient = SetupMockConnectorClient();
+        const auto pqGateway = SetupMockPqGateway();
+
+        constexpr char inputTopicName[] = "doubleYdbJoinInputTopicName";
+        constexpr char outputTopicName[] = "doubleYdbJoinOutputTopicName";
+        CreateTopic(inputTopicName);
+        CreateTopic(outputTopicName);
+
+        constexpr char pqSourceName[] = "pqSourceName";
+        constexpr char ydbSourceName[] = "ydbSourceName";
+        CreatePqSource(pqSourceName);
+        CreateYdbSource(ydbSourceName);
+
+        constexpr char ydbTable[] = "doubleYdbJoinLookup";
+        ExecExternalQuery(fmt::format(R"(
+            CREATE TABLE `{table}` (
+                fqdn String,
+                PRIMARY KEY (fqdn)
+            ))",
+            "table"_a = ydbTable
+        ));
+
+        {   // Prepare connector mock
+            const std::vector<TColumn> columns = {{"fqdn", Ydb::Type::STRING}};
+            SetupMockConnectorTableDescription(connectorClient, {
+                .TableName = ydbTable,
+                .Columns = columns,
+                .DescribeCount = 2,
+                .ListSplitsCount = 1
+            });
+
+            const std::vector<std::string> fqdnColumn = {"host1", "host2"};
+            SetupMockConnectorTableData(connectorClient, {
+                .TableName = ydbTable,
+                .Columns = columns,
+                .NumberReadSplits = 2,
+                .ResultFactory = [&]() {
+                    return MakeRecordBatch(MakeArray<arrow::BinaryBuilder>("fqdn", fqdnColumn, arrow::binary()));
+                }
+            });
+        }
+
+        constexpr char queryName[] = "streamingQuery";
+        ExecQuery(fmt::format(R"(
+            CREATE STREAMING QUERY `{query_name}` AS
+            DO BEGIN
+                INSERT INTO `{pq_source}`.`{output_topic}`
+                SELECT
+                    p.Data || "-" || la.fqdn || "-" || lb.fqdn
+                FROM `{pq_source}`.`{input_topic}` AS p
+                CROSS JOIN `{ydb_source}`.`{ydb_table}` AS la
+                CROSS JOIN `{ydb_source}`.`{ydb_table}` AS lb
+            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 = pqGateway->WaitReadSession(inputTopicName);
+        readSession->AddDataReceivedEvent(0, "data1");
+
+        pqGateway->WaitWriteSession(outputTopicName)->ExpectMessages({
+            "data1-host1-host2",
+            "data1-host2-host1",
+            "data1-host1-host1",
+            "data1-host2-host2"
+        }, /* sort  */ true);
+
+        readSession->AddCloseSessionEvent(EStatus::UNAVAILABLE, {NIssue::TIssue("Test pq session failure")});
+
+        pqGateway->WaitReadSession(inputTopicName)->AddDataReceivedEvent(1, "data2");
+        pqGateway->WaitWriteSession(outputTopicName)->ExpectMessages({
+            "data2-host1-host2",
+            "data2-host2-host1",
+            "data2-host1-host1",
+            "data2-host2-host2"
+        }, /* sort  */ true);
+    }
+
     Y_UNIT_TEST_F(StreamingQueryWithStreamLookupJoin, TStreamingTestFixture) {
         SetupAppConfig().MutableQueryServiceConfig()->SetProgressStatsPeriodMs(0);
 
diff --git a/ydb/core/protos/kqp.proto b/ydb/core/protos/kqp.proto
@@ -925,6 +925,7 @@ message TScriptExecutionOperationMeta {
     optional bool SaveQueryPhysicalGraph = 12;
     optional bool DisableDefaultTimeout = 13;
     optional string CheckpointId = 14;
+    optional string StreamingQueryPath = 15;
 }
 
 // stored in column "retry_state" of .metadata/script_executions table
diff --git a/ydb/tests/fq/streaming/conftest.py b/ydb/tests/fq/streaming/conftest.py
@@ -32,7 +32,7 @@ def query_async(self, statement):
         return self.session_pool.execute_with_retries_async(statement)
 
 
-@pytest.fixture(scope="function")
+@pytest.fixture(scope="module")
 def kikimr(request):
 
     class Kikimr:
diff --git a/ydb/tests/fq/streaming/test_streaming.py b/ydb/tests/fq/streaming/test_streaming.py
diff --git a/ydb/tests/fq/yds/test_row_dispatcher.py b/ydb/tests/fq/yds/test_row_dispatcher.py

Original file line number	Diff line number	Diff line change
`@@ -1559,7 +1559,6 @@ void TKqpTasksGraph::PersistTasksGraphInfo(NKikimrKqp::TQueryPhysicalGraph& resu`
`1559`	`1559`
`1560`	`1560`	`void TKqpTasksGraph::RestoreTasksGraphInfo(const TVector<IKqpGateway::TPhysicalTxData>& transactions, const TVector<NKikimrKqp::TKqpNodeResources>& resourcesSnapshot, const NKikimrKqp::TQueryPhysicalGraph& graphInfo) {`
`1561`	`1561`	`GetMeta().IsRestored = true;`
`1562`		`- GetMeta().AllowWithSpilling = false;`
`1563`	`1562`
`1564`	`1563`	`const auto restoreDqTransform = [](const auto& protoInfo) -> TMaybe<TTransform> {`
`1565`	`1564`	`if (!protoInfo.HasTransform()) {`
`@@ -1755,6 +1754,8 @@ void TKqpTasksGraph::RestoreTasksGraphInfo(const TVector<IKqpGateway::TPhysicalT`
`1755`	`1754`	`const auto it = scheduledTaskCount.find(stageIdx);`
`1756`	`1755`	`BuildReadTasksFromSource(stageInfo, resourcesSnapshot, it != scheduledTaskCount.end() ? it->second.TaskCount : 0);`
`1757`	`1756`	`}`
	`1757`	`+`
	`1758`	`+ GetMeta().AllowWithSpilling \|= stage.GetAllowWithSpilling();`
`1758`	`1759`	`}`
`1759`	`1760`	`}`
`1760`	`1761`	`}`
Original file line number	Diff line number	Diff line change
`@@ -116,9 +116,9 @@ class TKiSourceIntentDeterminationTransformer: public TKiSourceVisitorTransforme`
`116`	`116`	`TExprBase currentNode(node);`
`117`	`117`	`if (auto maybeReadTable = currentNode.Maybe<TKiReadTable>()) {`
`118`	`118`	`auto readTable = maybeReadTable.Cast();`
`119`		`- for (auto setting : readTable.Settings()) {`
`120`		`- auto name = setting.Name().Value();`
`121`		`- if (name == "sysViewRewritten") {`
	`119`	`+ for (auto setting : readTable.Settings().Ref().ChildrenList()) {`
	`120`	`+ auto maybeTuple = TMaybeNode<TCoNameValueTuple>(setting);`
	`121`	`+ if (maybeTuple && maybeTuple.Cast().Name().Value() == "sysViewRewritten"sv) {`
`122`	`122`	`sysViewRewritten = true;`
`123`	`123`	`}`
`124`	`124`	`}`