YQ-4699 Check shared_ptr in metadata_loader (#25580)

5 files changed, 43 insertions, 39 deletions

diff --git a/ydb/core/fq/libs/row_dispatcher/coordinator.cpp b/ydb/core/fq/libs/row_dispatcher/coordinator.cpp
index 102154b596f..382a5296462 100644
--- a/ydb/core/fq/libs/row_dispatcher/coordinator.cpp
+++ b/ydb/core/fq/libs/row_dispatcher/coordinator.cpp
@@ -515,7 +515,7 @@ void TActorCoordinator::Handle(TEvPrivate::TEvPrintState::TPtr&) {
 }
 
 void TActorCoordinator::Handle(NKikimr::TEvTenantNodeEnumerator::TEvLookupResult::TPtr& ev) {
-     if (!ev->Get()->Success) {
+    if (!ev->Get()->Success) {
         LOG_ROW_DISPATCHER_ERROR("Failed to get TEvLookupResult, try later...");
         ScheduleNodeInfoRequest();
         return;
diff --git a/ydb/core/kqp/gateway/kqp_metadata_loader.cpp b/ydb/core/kqp/gateway/kqp_metadata_loader.cpp
index da4511e1777..c48bdef6136 100644
--- a/ydb/core/kqp/gateway/kqp_metadata_loader.cpp
+++ b/ydb/core/kqp/gateway/kqp_metadata_loader.cpp
@@ -920,11 +920,12 @@ NThreading::TFuture<TTableMetadataResult> TKqpTableMetadataLoader::LoadTableMeta
 
     const auto schemeCacheId = MakeSchemeCacheID();
 
+    auto ptr = weak_from_base();
     auto future = SendActorRequest<TRequest, TResponse, TResult>(
         ActorSystem,
         schemeCacheId,
         ev.Release(),
-        [userToken, database, cluster, mainCluster = Cluster, table, settings, expectedSchemaVersion, this, queryName, externalPath]
+        [userToken, database, cluster, mainCluster = Cluster, table, settings, expectedSchemaVersion, ptr, queryName, externalPath]
             (TPromise<TResult> promise, TResponse&& response) mutable
         {
             try {
@@ -955,7 +956,12 @@ NThreading::TFuture<TTableMetadataResult> TKqpTableMetadataLoader::LoadTableMeta
                     }
                 }
 
-                const bool resolveEntityInsideDataSource = (cluster != Cluster);
+                auto locked = ptr.lock();
+                if (!locked) {
+                    promise.SetValue(ResultFromError<TResult>(YqlIssue({}, TIssuesIds::KIKIMR_INDEX_METADATA_LOAD_FAILED, "lock failed")));
+                    return;
+                }
+                const bool resolveEntityInsideDataSource = (cluster != locked->Cluster);
                 // resolveEntityInsideDataSource => entry.Kind == EKind::KindExternalDataSource
                 if (resolveEntityInsideDataSource && entry.Kind != EKind::KindExternalDataSource) {
                     const auto message = TStringBuilder()
@@ -976,8 +982,8 @@ NThreading::TFuture<TTableMetadataResult> TKqpTableMetadataLoader::LoadTableMeta
                         if (externalPath) {
                             externalDataSourceMetadata.Metadata->ExternalSource.TableLocation = *externalPath;
                         }
-                        LoadExternalDataSourceSecretValues(entry, userToken, ActorSystem)
-                            .Subscribe([promise, externalDataSourceMetadata, settings, table, database, externalPath, this](const TFuture<TEvDescribeSecretsResponse::TDescription>& result) mutable
+                        LoadExternalDataSourceSecretValues(entry, userToken, locked->ActorSystem)
+                            .Subscribe([promise, externalDataSourceMetadata, settings, table, database, externalPath, locked](const TFuture<TEvDescribeSecretsResponse::TDescription>& result) mutable
                         {
                             UpdateExternalDataSourceSecretsValue(externalDataSourceMetadata, result.GetValue());
                             if (!externalDataSourceMetadata.Success()) {
@@ -1036,7 +1042,7 @@ NThreading::TFuture<TTableMetadataResult> TKqpTableMetadataLoader::LoadTableMeta
                                 auto path = databaseName + "/" + *externalPath;
 
                                 GetSchemeEntryType(
-                                    FederatedQuerySetup,
+                                    locked->FederatedQuerySetup,
                                     source.DataSourceLocation,
                                     databaseName,
                                     useTls,
@@ -1067,7 +1073,7 @@ NThreading::TFuture<TTableMetadataResult> TKqpTableMetadataLoader::LoadTableMeta
                             return;
                         }
                         settings.WithExternalDatasources_ = true;
-                        LoadTableMetadataCache(cluster, dataSourcePath, settings, database, userToken)
+                        locked->LoadTableMetadataCache(cluster, dataSourcePath, settings, database, userToken)
                             .Apply([promise, externalTableMetadata](const TFuture<TTableMetadataResult>& result) mutable
                         {
                             auto externalDataSourceMetadata = result.GetValue();
@@ -1084,7 +1090,7 @@ NThreading::TFuture<TTableMetadataResult> TKqpTableMetadataLoader::LoadTableMeta
                             }
                             TIndexId pathId = TIndexId(child.PathId, child.SchemaVersion);
 
-                            LoadTableMetadataCache(cluster, std::make_pair(pathId, table), settings, database, userToken)
+                            locked->LoadTableMetadataCache(cluster, std::make_pair(pathId, table), settings, database, userToken)
                                 .Apply([promise](const TFuture<TTableMetadataResult>& result) mutable
                             {
                                 promise.SetValue(result.GetValue());
diff --git a/ydb/core/kqp/provider/yql_kikimr_gateway_ut.cpp b/ydb/core/kqp/provider/yql_kikimr_gateway_ut.cpp
index eac754a219b..316ba5a6cfd 100644
--- a/ydb/core/kqp/provider/yql_kikimr_gateway_ut.cpp
+++ b/ydb/core/kqp/provider/yql_kikimr_gateway_ut.cpp
@@ -404,7 +404,8 @@ Y_UNIT_TEST_SUITE(KikimrIcGateway) {
             );)";
         auto result = session.ExecuteSchemeQuery(query).GetValueSync();
         UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
-        auto responseFuture = GetIcGateway(kikimr.GetTestServer())->LoadTableMetadata(TestCluster, externalTableName, IKikimrGateway::TLoadTableMetadataSettings());
+        auto gateway = GetIcGateway(kikimr.GetTestServer());
+        auto responseFuture = gateway->LoadTableMetadata(TestCluster, externalTableName, IKikimrGateway::TLoadTableMetadataSettings());
         responseFuture.Wait();
         auto response = responseFuture.GetValue();
         response.Issues().PrintTo(Cerr);
@@ -513,7 +514,8 @@ Y_UNIT_TEST_SUITE(KikimrIcGateway) {
         auto result = session.ExecuteSchemeQuery(query).GetValueSync();
         UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
 
-        auto responseFuture = GetIcGateway(kikimr.GetTestServer())->LoadTableMetadata(TestCluster, externalDataSourceName, IKikimrGateway::TLoadTableMetadataSettings());
+        auto gateway = GetIcGateway(kikimr.GetTestServer());
+        auto responseFuture = gateway->LoadTableMetadata(TestCluster, externalDataSourceName, IKikimrGateway::TLoadTableMetadataSettings());
         responseFuture.Wait();
 
         auto response = responseFuture.GetValue();
@@ -554,8 +556,8 @@ Y_UNIT_TEST_SUITE(KikimrIcGateway) {
             );)";
         auto result = session.ExecuteSchemeQuery(query).GetValueSync();
         UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
-
-        auto responseFuture = GetIcGateway(kikimr.GetTestServer())->LoadTableMetadata(TestCluster, externalDataSourceName, IKikimrGateway::TLoadTableMetadataSettings());
+        auto gateway = GetIcGateway(kikimr.GetTestServer());
+        auto responseFuture = gateway->LoadTableMetadata(TestCluster, externalDataSourceName, IKikimrGateway::TLoadTableMetadataSettings());
         responseFuture.Wait();
 
         auto response = responseFuture.GetValue();
@@ -596,7 +598,8 @@ Y_UNIT_TEST_SUITE(KikimrIcGateway) {
         auto result = session.ExecuteSchemeQuery(query).GetValueSync();
         UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
 
-        auto responseFuture = GetIcGateway(kikimr.GetTestServer())->LoadTableMetadata(TestCluster, externalDataSourceName, IKikimrGateway::TLoadTableMetadataSettings());
+        auto gateway = GetIcGateway(kikimr.GetTestServer());
+        auto responseFuture = gateway->LoadTableMetadata(TestCluster, externalDataSourceName, IKikimrGateway::TLoadTableMetadataSettings());
         responseFuture.Wait();
 
         auto response = responseFuture.GetValue();
@@ -644,7 +647,8 @@ Y_UNIT_TEST_SUITE(KikimrIcGateway) {
         auto result = session.ExecuteSchemeQuery(query).GetValueSync();
         UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
 
-        auto responseFuture = GetIcGateway(kikimr.GetTestServer())->LoadTableMetadata(TestCluster, externalDataSourceName, IKikimrGateway::TLoadTableMetadataSettings());
+        auto gateway = GetIcGateway(kikimr.GetTestServer());
+        auto responseFuture = gateway->LoadTableMetadata(TestCluster, externalDataSourceName, IKikimrGateway::TLoadTableMetadataSettings());
         responseFuture.Wait();
 
         auto response = responseFuture.GetValue();
@@ -685,7 +689,8 @@ Y_UNIT_TEST_SUITE(KikimrIcGateway) {
         auto result = session.ExecuteSchemeQuery(query).GetValueSync();
         UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
 
-        auto responseFuture = GetIcGateway(kikimr.GetTestServer())->LoadTableMetadata(TestCluster, externalDataSourceName, IKikimrGateway::TLoadTableMetadataSettings());
+        auto gateway = GetIcGateway(kikimr.GetTestServer());
+        auto responseFuture = gateway->LoadTableMetadata(TestCluster, externalDataSourceName, IKikimrGateway::TLoadTableMetadataSettings());
         responseFuture.Wait();
 
         auto response = responseFuture.GetValue();
diff --git a/ydb/tests/fq/streaming/base.py b/ydb/tests/fq/streaming/base.py
index 567cfa164f0..3b95bfa7068 100644
--- a/ydb/tests/fq/streaming/base.py
+++ b/ydb/tests/fq/streaming/base.py
@@ -12,10 +12,8 @@ logger = logging.getLogger(__name__)
 
 class YdbClient:
     def __init__(self, endpoint: str, database: str):
-        self.endpoint = endpoint
-        self.database = database
-
-        self.driver = ydb.Driver(endpoint=endpoint, database=database, oauth=None)
+        driver_config = ydb.DriverConfig(endpoint, database, auth_token='root@builtin')
+        self.driver = ydb.Driver(driver_config)
         self.session_pool = ydb.QuerySessionPool(self.driver)
 
     def stop(self):
@@ -47,11 +45,12 @@ class StreamingImportTestBase(object):
         config = KikimrConfigGenerator(
             extra_feature_flags={
                 "enable_external_data_sources": True,
-                "enable_move_column_table": True
+                "enable_streaming_queries": True
             },
             query_service_config={"available_external_data_sources": ["Ydb"]},
             table_service_config={},
-            nodes=2
+            nodes=2,
+            default_clusteradmin="root@builtin"
         )
         config.yaml_config["query_service_config"] = {}
         config.yaml_config["query_service_config"]["available_external_data_sources"] = ["ObjectStorage", "Ydb", "YdbTopics"]
@@ -78,6 +77,7 @@ class StreamingImportTestBase(object):
         config.yaml_config["query_service_config"]["checkpoints_config"]["checkpoint_garbage_config"] = {}
         config.yaml_config["query_service_config"]["checkpoints_config"]["checkpoint_garbage_config"]["enabled"] = True
         config.yaml_config["query_service_config"]["checkpoints_config"]["max_inflight"] = 1
+        config.yaml_config["query_service_config"]["checkpoints_config"]["checkpointing_period_millis"] = 500
         return config
 
     @classmethod
diff --git a/ydb/tests/fq/streaming/test_streaming.py b/ydb/tests/fq/streaming/test_streaming.py
index da32939e6d7..958aca16409 100644
--- a/ydb/tests/fq/streaming/test_streaming.py
+++ b/ydb/tests/fq/streaming/test_streaming.py
@@ -2,8 +2,6 @@ import logging
 import os
 import time
 
-import ydb
-
 from ydb.tests.fq.streaming.base import StreamingImportTestBase
 from ydb.tests.tools.datastreams_helpers.test_yds_base import TestYdsBase
 
@@ -75,30 +73,25 @@ class TestStreamingInYdb(StreamingImportTestBase, TestYdsBase):
                 SHARED_READING="TRUE",
                 AUTH_METHOD="NONE");""")
 
-        sql = f"""
-            $in = SELECT time FROM {sourceName}.`{self.input_topic}`
+        sql = R'''
+            CREATE STREAMING QUERY `{query_name}` AS
+            DO BEGIN
+                $in = SELECT time FROM {source_name}.`{input_topic}`
                 WITH (
                     FORMAT="json_each_row",
                     SCHEMA=(time String NOT NULL))
                 WHERE time like "%lunch%";
+                INSERT INTO {source_name}.`{output_topic}` SELECT time FROM $in;
+            END DO;'''
 
-            INSERT INTO {sourceName}.`{self.output_topic}` SELECT time FROM $in;"""
-
-        session1 = ydb.QuerySession(self.ydb_client.driver)
-        session1.create()
-        it1 = session1.execute(sql)
-
-        session2 = ydb.QuerySession(self.ydb_client.driver)
-        session2.create()
-        it2 = session2.execute(sql)
-
+        self.ydb_client.query(sql.format(query_name="query1", source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
+        self.ydb_client.query(sql.format(query_name="query2", source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
         time.sleep(3)
         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
 
-        it1.cancel()
-        session1.delete()
-        it2.cancel()
-        session2.delete()
+        sql = R'''DROP STREAMING QUERY `{query_name}`;'''
+        self.ydb_client.query(sql.format(query_name="query1"))
+        self.ydb_client.query(sql.format(query_name="query2"))