YQ-4867 Add streaming query name to checkpoint metrics (#28519)

Author: kardymonds

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> {

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;
 

ydb/core/kqp/executer_actor/kqp_data_executer.cpp

@@ -2887,7 +2887,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());

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);

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);

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;

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);

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

ydb/tests/fq/streaming/test_streaming.py

@@ -31,13 +31,13 @@ def get_sensors(self, kikimr, node_id, counters):
         url = self.monitoring_endpoint(kikimr, node_id) + "/counters/counters={}/json".format(counters)
         return load_metrics(url)
 
-    def get_checkpoint_coordinator_metric(self, kikimr, query_id, metric_name, expect_counters_exist=False):
+    def get_checkpoint_coordinator_metric(self, kikimr, path, metric_name, expect_counters_exist=False):
         sum = 0
         found = False
         for node_id in kikimr.Cluster.nodes:
             sensor = self.get_sensors(kikimr, node_id, "kqp").find_sensor(
                 {
-                    # "query_id": query_id,  # TODO
+                    "path": path,
                     "subsystem": "checkpoint_coordinator",
                     "sensor": metric_name
                 }
@@ -48,16 +48,16 @@ def get_checkpoint_coordinator_metric(self, kikimr, query_id, metric_name, expec
         assert found or not expect_counters_exist
         return sum
 
-    def get_completed_checkpoints(self, kikimr, query_id):
-        return self.get_checkpoint_coordinator_metric(kikimr, query_id, "CompletedCheckpoints")
+    def get_completed_checkpoints(self, kikimr, path):
+        return self.get_checkpoint_coordinator_metric(kikimr, path, "CompletedCheckpoints")
 
-    def wait_completed_checkpoints(self, kikimr, query_id,
+    def wait_completed_checkpoints(self, kikimr, path,
                                    timeout=plain_or_under_sanitizer_wrapper(120, 150)):
-        current = self.get_checkpoint_coordinator_metric(kikimr, query_id, "CompletedCheckpoints")
+        current = self.get_checkpoint_coordinator_metric(kikimr, path, "CompletedCheckpoints")
         checkpoints_count = current + 2
         deadline = time.time() + timeout
         while True:
-            completed = self.get_completed_checkpoints(kikimr, query_id)
+            completed = self.get_completed_checkpoints(kikimr, path)
             if completed >= checkpoints_count:
                 break
             assert time.time() < deadline, "Wait checkpoint failed, actual completed: " + str(completed)
@@ -125,16 +125,16 @@ def test_restart_query(self, kikimr):
                 INSERT INTO {source_name}.`{output_topic}` SELECT time FROM $in;
             END DO;'''
 
-        query_id = "query_id"  # TODO
+        path = f"/Root/{name}"
         kikimr.YdbClient.query(sql.format(query_name=name, source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
-        self.wait_completed_checkpoints(kikimr, query_id)
+        self.wait_completed_checkpoints(kikimr, path)
 
         data = ['{"time": "lunch time"}']
         expected_data = ['lunch time']
         self.write_stream(data)
 
         assert self.read_stream(len(expected_data), topic_path=self.output_topic) == expected_data
-        self.wait_completed_checkpoints(kikimr, query_id)
+        self.wait_completed_checkpoints(kikimr, path)
 
         kikimr.YdbClient.query(f"ALTER STREAMING QUERY `{name}` SET (RUN = FALSE);")
         time.sleep(0.5)
@@ -164,21 +164,22 @@ def test_read_topic_shared_reading_insert_to_topic(self, kikimr):
                 INSERT INTO {source_name}.`{output_topic}` SELECT time FROM $in;
             END DO;'''
 
-        kikimr.YdbClient.query(sql.format(query_name="query1", source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
-        kikimr.YdbClient.query(sql.format(query_name="query2", source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
-
-        query_id = "query_id"  # TODO
-        self.wait_completed_checkpoints(kikimr, query_id)
+        query_name1 = "test_read_topic_shared_reading_insert_to_topic1"
+        query_name2 = "test_read_topic_shared_reading_insert_to_topic2"
+        kikimr.YdbClient.query(sql.format(query_name=query_name1, source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
+        kikimr.YdbClient.query(sql.format(query_name=query_name2, source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
+        path1 = f"/Root/{query_name1}"
+        self.wait_completed_checkpoints(kikimr, path1)
 
         data = ['{"time": "lunch time"}']
         expected_data = ['lunch time', 'lunch time']
         self.write_stream(data)
         assert self.read_stream(len(expected_data), topic_path=self.output_topic) == expected_data
-        self.wait_completed_checkpoints(kikimr, query_id)
+        self.wait_completed_checkpoints(kikimr, path1)
 
         sql = R'''ALTER STREAMING QUERY `{query_name}` SET (RUN = FALSE);'''
-        kikimr.YdbClient.query(sql.format(query_name="query1"))
-        kikimr.YdbClient.query(sql.format(query_name="query2"))
+        kikimr.YdbClient.query(sql.format(query_name=query_name1))
+        kikimr.YdbClient.query(sql.format(query_name=query_name2))
 
         time.sleep(1)
 
@@ -187,13 +188,13 @@ def test_read_topic_shared_reading_insert_to_topic(self, kikimr):
         self.write_stream(data)
 
         sql = R'''ALTER STREAMING QUERY `{query_name}` SET (RUN = TRUE);'''
-        kikimr.YdbClient.query(sql.format(query_name="query1"))
-        kikimr.YdbClient.query(sql.format(query_name="query2"))
+        kikimr.YdbClient.query(sql.format(query_name=query_name1))
+        kikimr.YdbClient.query(sql.format(query_name=query_name2))
         assert self.read_stream(len(expected_data), topic_path=self.output_topic) == expected_data
 
         sql = R'''DROP STREAMING QUERY `{query_name}`;'''
-        kikimr.YdbClient.query(sql.format(query_name="query1"))
-        kikimr.YdbClient.query(sql.format(query_name="query2"))
+        kikimr.YdbClient.query(sql.format(query_name=query_name1))
+        kikimr.YdbClient.query(sql.format(query_name=query_name2))
 
     def test_read_topic_shared_reading_restart_nodes(self, kikimr):
         sourceName = "source_" + ''.join(random.choices(string.ascii_letters + string.digits, k=8))
@@ -211,14 +212,15 @@ def test_read_topic_shared_reading_restart_nodes(self, kikimr):
                 INSERT INTO {source_name}.`{output_topic}` SELECT value FROM $in;
             END DO;'''
 
-        kikimr.YdbClient.query(sql.format(query_name="query1", source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
-        query_id = "query_id"  # TODO
-        self.wait_completed_checkpoints(kikimr, query_id)
+        query_name = "test_read_topic_shared_reading_restart_nodes"
+        kikimr.YdbClient.query(sql.format(query_name=query_name, source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
+        path = f"/Root/{query_name}"
+        self.wait_completed_checkpoints(kikimr, path)
 
         self.write_stream(['{"value": "value1"}'])
         expected_data = ['value1']
         assert self.read_stream(len(expected_data), topic_path=self.output_topic) == expected_data
-        self.wait_completed_checkpoints(kikimr, query_id)
+        self.wait_completed_checkpoints(kikimr, path)
 
         restart_node_id = None
         for node_id in kikimr.Cluster.nodes:
@@ -234,7 +236,7 @@ def test_read_topic_shared_reading_restart_nodes(self, kikimr):
         self.write_stream(['{"value": "value2"}'])
         expected_data = ['value2']
         assert self.read_stream(len(expected_data), topic_path=self.output_topic) == expected_data
-        self.wait_completed_checkpoints(kikimr, query_id)
+        self.wait_completed_checkpoints(kikimr, path)
 
     def test_read_topic_restore_state(self, kikimr):
         sourceName = "source4_" + ''.join(random.choices(string.ascii_letters + string.digits, k=8))
@@ -267,9 +269,10 @@ def test_read_topic_restore_state(self, kikimr):
                     SELECT ToBytes(Unwrap(Json::SerializeJson(Yson::From(TableRow())))) FROM $mr;
             END DO;'''
 
-        kikimr.YdbClient.query(sql.format(query_name="query1", source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
-        query_id = "query_id"  # TODO
-        self.wait_completed_checkpoints(kikimr, query_id)
+        query_name = "test_read_topic_restore_state"
+        kikimr.YdbClient.query(sql.format(query_name=query_name, source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
+        path = f"/Root/{query_name}"
+        self.wait_completed_checkpoints(kikimr, path)
 
         data = [
             '{"dt": 1696849942000001, "str": "A" }',
@@ -278,7 +281,7 @@ def test_read_topic_restore_state(self, kikimr):
         self.write_stream(data)
         expected_data = ['{"a_time":1696849942000001,"b_time":1696849942500001,"c_time":null}']
         assert self.read_stream(len(expected_data), topic_path=self.output_topic) == expected_data
-        self.wait_completed_checkpoints(kikimr, query_id)
+        self.wait_completed_checkpoints(kikimr, path)
 
         restart_node_id = None
         for node_id in kikimr.Cluster.nodes:
@@ -313,9 +316,9 @@ def test_json_errors(self, kikimr):
                 INSERT INTO {source_name}.`{output_topic}` SELECT data FROM $in;
             END DO;'''
 
-        query_id = "query_id"  # TODO
+        path = f"/Root/{name}"
         kikimr.YdbClient.query(sql.format(query_name=name, source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
-        self.wait_completed_checkpoints(kikimr, query_id)
+        self.wait_completed_checkpoints(kikimr, path)
 
         data = [
             '{"time": 101, "data": "hello1"}',
@@ -328,7 +331,7 @@ def test_json_errors(self, kikimr):
         assert self.read_stream(len(expected), topic_path=self.output_topic) == expected
 
     def test_restart_query_by_rescaling(self, kikimr):
-        sourceName = ''.join(random.choices(string.ascii_letters + string.digits, k=8))
+        sourceName = 'source' + ''.join(random.choices(string.ascii_letters + string.digits, k=8))
         self.init_topics(sourceName, partitions_count=10)
         self.create_source(kikimr, sourceName, True)
 
@@ -347,15 +350,15 @@ def test_restart_query_by_rescaling(self, kikimr):
                 INSERT INTO `{source_name}`.`{output_topic}` SELECT time FROM $in;
             END DO;'''
 
-        query_id = "query_id"  # TODO
+        path = f"/Root/{name}"
         kikimr.YdbClient.query(sql.format(query_name=name, source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
-        self.wait_completed_checkpoints(kikimr, query_id)
+        self.wait_completed_checkpoints(kikimr, path)
 
         message_count = 20
         for i in range(message_count):
             self.write_stream(['{"time": "time to do it"}'], topic_path=None, partition_key=(''.join(random.choices(string.digits, k=8))))
         assert self.read_stream(message_count, topic_path=self.output_topic) == ["time to do it" for i in range(message_count)]
-        self.wait_completed_checkpoints(kikimr, query_id)
+        self.wait_completed_checkpoints(kikimr, path)
 
         logging.debug(f"stopping query {name}")
         kikimr.YdbClient.query(f"ALTER STREAMING QUERY `{name}` SET (RUN = FALSE);")