YQ-4084 Add AllExternalDataSourcesAreAvailable to QueryServiceConfig (#18042)

kardymonds · kardymonds · commit 7343abbe9068 · 2025-05-19T12:06:20.000Z
diff --git a/ydb/core/external_sources/external_source_factory.cpp b/ydb/core/external_sources/external_source_factory.cpp
@@ -16,8 +16,10 @@ namespace {
 struct TExternalSourceFactory : public IExternalSourceFactory {
     TExternalSourceFactory(
         const TMap<TString, IExternalSource::TPtr>& sources,
+        bool allExternalDataSourcesAreAvailable,
         const std::set<TString>& availableExternalDataSources)
         : Sources(sources)
+        , AllExternalDataSourcesAreAvailable(allExternalDataSourcesAreAvailable)
         , AvailableExternalDataSources(availableExternalDataSources)
     {}
 
@@ -26,14 +28,15 @@ struct TExternalSourceFactory : public IExternalSourceFactory {
         if (it == Sources.end()) {
             throw TExternalSourceException() << "External source with type " << type << " was not found";
         }
-        if (!AvailableExternalDataSources.contains(type)) {
+        if (!AllExternalDataSourcesAreAvailable && !AvailableExternalDataSources.contains(type)) {
             throw TExternalSourceException() << "External source with type " << type << " is disabled. Please contact your system administrator to enable it";
         }
         return it->second;
     }
 
 private:
     const TMap<TString, IExternalSource::TPtr> Sources;
+    bool AllExternalDataSourcesAreAvailable;
     const std::set<TString> AvailableExternalDataSources;
 };
 
@@ -45,6 +48,7 @@ IExternalSourceFactory::TPtr CreateExternalSourceFactory(const std::vector<TStri
                                                          std::shared_ptr<NYql::ISecuredServiceAccountCredentialsFactory> credentialsFactory,
                                                          bool enableInfer,
                                                          bool allowLocalFiles,
+                                                         bool allExternalDataSourcesAreAvailable,
                                                          const std::set<TString>& availableExternalDataSources) {
     std::vector<TRegExMatch> hostnamePatternsRegEx(hostnamePatterns.begin(), hostnamePatterns.end());
     return MakeIntrusive<TExternalSourceFactory>(TMap<TString, IExternalSource::TPtr>{
@@ -93,6 +97,7 @@ IExternalSourceFactory::TPtr CreateExternalSourceFactory(const std::vector<TStri
             CreateExternalDataSource(TString{NYql::SolomonProviderName}, {"NONE", "TOKEN"}, {"use_ssl", "grpc_port"}, hostnamePatternsRegEx)
         }
     },
+    allExternalDataSourcesAreAvailable,
     availableExternalDataSources); 
 }
 
diff --git a/ydb/core/external_sources/external_source_factory.h b/ydb/core/external_sources/external_source_factory.h
@@ -18,6 +18,7 @@ IExternalSourceFactory::TPtr CreateExternalSourceFactory(const std::vector<TStri
                                                          std::shared_ptr<NYql::ISecuredServiceAccountCredentialsFactory> credentialsFactory = nullptr,
                                                          bool enableInfer = false,
                                                          bool allowLocalFiles = false,
+                                                         bool allExternalDataSourcesAreAvailable = true,
                                                          const std::set<TString>& availableExternalDataSources = {});
 
 }
diff --git a/ydb/core/grpc_services/rpc_describe_external_table.cpp b/ydb/core/grpc_services/rpc_describe_external_table.cpp
@@ -36,7 +36,7 @@ bool ConvertContent(
     google::protobuf::Map<TProtoStringType, TProtoStringType>& out,
     TIssues& issues
 ) {
-    const auto externalSourceFactory = NExternalSource::CreateExternalSourceFactory({}, nullptr, 50000, nullptr, false, false, NYql::GetAllExternalDataSourceTypes());
+    const auto externalSourceFactory = NExternalSource::CreateExternalSourceFactory({}, nullptr, 50000, nullptr, false, false, true, NYql::GetAllExternalDataSourceTypes());
     try {
         const auto source = externalSourceFactory->GetOrCreate(sourceType);
         for (const auto& [key, items] : source->GetParameters(in)) {
diff --git a/ydb/core/kqp/host/kqp_host.cpp b/ydb/core/kqp/host/kqp_host.cpp
@@ -1172,6 +1172,7 @@ class TKqpHost : public IKqpHost {
                                                                                  FederatedQuerySetup ? FederatedQuerySetup->CredentialsFactory : nullptr,
                                                                                  Config->FeatureFlags.GetEnableExternalSourceSchemaInference(),
                                                                                  FederatedQuerySetup->S3GatewayConfig.GetAllowLocalFiles(),
+                                                                                 QueryServiceConfig.GetAllExternalDataSourcesAreAvailable(),
                                                                                  std::set<TString>(availableExternalDataSources.cbegin(), availableExternalDataSources.cend()));
         }
     }
diff --git a/ydb/core/kqp/ut/federated_query/common/common.cpp b/ydb/core/kqp/ut/federated_query/common/common.cpp
@@ -41,11 +41,7 @@ namespace NKikimr::NKqp::NFederatedQueryTest {
         if (!appConfig) {
             appConfig.emplace();
         }
-        appConfig->MutableQueryServiceConfig()->AddAvailableExternalDataSources("ObjectStorage");
-        appConfig->MutableQueryServiceConfig()->AddAvailableExternalDataSources("ClickHouse");
-        appConfig->MutableQueryServiceConfig()->AddAvailableExternalDataSources("PostgreSQL");
-        appConfig->MutableQueryServiceConfig()->AddAvailableExternalDataSources("MySQL");
-        appConfig->MutableQueryServiceConfig()->AddAvailableExternalDataSources("Ydb");
+        appConfig->MutableQueryServiceConfig()->SetAllExternalDataSourcesAreAvailable(true);
 
         auto settings = TKikimrSettings();
 
diff --git a/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp b/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp
@@ -6613,7 +6613,6 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
     Y_UNIT_TEST(CreateExternalDataSource) {
         NKikimrConfig::TAppConfig appCfg;
         appCfg.MutableQueryServiceConfig()->AddHostnamePatterns("my-bucket");
-        appCfg.MutableQueryServiceConfig()->AddAvailableExternalDataSources("ObjectStorage");
 
         TKikimrRunner kikimr(appCfg);
         kikimr.GetTestServer().GetRuntime()->GetAppData(0).FeatureFlags.SetEnableExternalDataSources(true);
@@ -6696,6 +6695,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
 
     Y_UNIT_TEST(DisableS3ExternalDataSource) {
         NKikimrConfig::TAppConfig appCfg;
+        appCfg.MutableQueryServiceConfig()->SetAllExternalDataSourcesAreAvailable(false);
         appCfg.MutableQueryServiceConfig()->AddAvailableExternalDataSources("PostgreSQL");
         TKikimrRunner kikimr(appCfg);
         kikimr.GetTestServer().GetRuntime()->GetAppData(0).FeatureFlags.SetEnableExternalDataSources(true);
diff --git a/ydb/core/protos/config.proto b/ydb/core/protos/config.proto
@@ -1077,7 +1077,8 @@ message TQueryServiceConfig {
     optional uint64 ProgressStatsPeriodMs = 14 [default = 0]; // 0 = disabled
     optional uint32 QueryTimeoutDefaultSeconds = 19 [default = 1800];
     optional bool EnableMatchRecognize = 20 [default = false];
-    repeated string AvailableExternalDataSources = 22;
+    repeated string AvailableExternalDataSources = 22;                  // Ignored if AllExternalDataSourcesAreAvailable is true
+    optional bool AllExternalDataSourcesAreAvailable = 23 [default = true];   
 }
 
 // Config describes immediate controls and allows
diff --git a/ydb/core/tx/schemeshard/schemeshard_impl.cpp b/ydb/core/tx/schemeshard/schemeshard_impl.cpp
@@ -7287,6 +7287,7 @@ void TSchemeShard::ApplyConsoleConfigs(const NKikimrConfig::TAppConfig& appConfi
             nullptr,
             appConfig.GetFeatureFlags().GetEnableExternalSourceSchemaInference(),
             appConfig.GetQueryServiceConfig().GetS3().GetAllowLocalFiles(),
+            appConfig.GetQueryServiceConfig().GetAllExternalDataSourcesAreAvailable(),
             std::set<TString>(availableExternalDataSources.cbegin(), availableExternalDataSources.cend())
         );
     }
diff --git a/ydb/core/tx/schemeshard/ut_helpers/test_env.cpp b/ydb/core/tx/schemeshard/ut_helpers/test_env.cpp
@@ -1026,17 +1026,7 @@ void NSchemeShardUT_Private::TTestEnv::BootTxAllocator(NActors::TTestActorRuntim
 NKikimrConfig::TAppConfig NSchemeShardUT_Private::TTestEnv::GetAppConfig() const {
     NKikimrConfig::TAppConfig appConfig;
     auto* queryServiceConfig = appConfig.MutableQueryServiceConfig();
-    queryServiceConfig->AddAvailableExternalDataSources("ObjectStorage");
-    queryServiceConfig->AddAvailableExternalDataSources("ClickHouse");
-    queryServiceConfig->AddAvailableExternalDataSources("PostgreSQL");
-    queryServiceConfig->AddAvailableExternalDataSources("MySQL");
-    queryServiceConfig->AddAvailableExternalDataSources("Ydb");
-    queryServiceConfig->AddAvailableExternalDataSources("YT");
-    queryServiceConfig->AddAvailableExternalDataSources("Greenplum");
-    queryServiceConfig->AddAvailableExternalDataSources("MsSQLServer");
-    queryServiceConfig->AddAvailableExternalDataSources("Oracle");
-    queryServiceConfig->AddAvailableExternalDataSources("Logging");
-    queryServiceConfig->AddAvailableExternalDataSources("Solomon");
+    queryServiceConfig->SetAllExternalDataSourcesAreAvailable(true);
     return appConfig;
 }
 
diff --git a/ydb/core/viewer/viewer_describe.h b/ydb/core/viewer/viewer_describe.h
@@ -314,7 +314,7 @@ class TJsonDescribe : public TViewerPipeClient {
             return;
         }
 
-        NExternalSource::IExternalSourceFactory::TPtr externalSourceFactory{NExternalSource::CreateExternalSourceFactory({}, nullptr, 50000, nullptr, false, false, NYql::GetAllExternalDataSourceTypes())};
+        NExternalSource::IExternalSourceFactory::TPtr externalSourceFactory{NExternalSource::CreateExternalSourceFactory({}, nullptr, 50000, nullptr, false, false, true, NYql::GetAllExternalDataSourceTypes())};
         NJson::TJsonValue root;
         const auto& sourceType = DescribeResult->GetPathDescription().GetExternalTableDescription().GetSourceType();
         try {
diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/run/kqprun.py b/ydb/library/yql/providers/generic/connector/tests/utils/run/kqprun.py
@@ -186,6 +186,8 @@ class AppConfigRenderer:
 }
 
 QueryServiceConfig {
+  AllExternalDataSourcesAreAvailable: true
+
   Generic {
     Connector {
         Endpoint {
diff --git a/ydb/tests/tools/fq_runner/kikimr_utils.py b/ydb/tests/tools/fq_runner/kikimr_utils.py
@@ -75,18 +75,7 @@ def apply_to_kikimr(self, request, kikimr):
 
 
 def enable_external_data_sources(qs_config):
-    qs_config['available_external_data_sources'] = []
-    qs_config['available_external_data_sources'].append("ObjectStorage")
-    qs_config['available_external_data_sources'].append("ClickHouse")
-    qs_config['available_external_data_sources'].append("PostgreSQL")
-    qs_config['available_external_data_sources'].append("MySQL")
-    qs_config['available_external_data_sources'].append("Ydb")
-    qs_config['available_external_data_sources'].append("YT")
-    qs_config['available_external_data_sources'].append("Greenplum")
-    qs_config['available_external_data_sources'].append("MsSQLServer")
-    qs_config['available_external_data_sources'].append("Oracle")
-    qs_config['available_external_data_sources'].append("Logging")
-    qs_config['available_external_data_sources'].append("Solomon")
+    qs_config['all_external_data_sources_are_available'] = True
 
 
 class AddFormatSizeLimitExtension(ExtensionPoint):
diff --git a/ydb/tests/tools/kqprun/configuration/app_config.conf b/ydb/tests/tools/kqprun/configuration/app_config.conf
@@ -83,17 +83,7 @@ QueryServiceConfig {
   QueryArtifactsCompressionMethod: "zstd_6"
   ScriptResultRowsLimit: 0
   ScriptResultSizeLimit: 10485760
-  AvailableExternalDataSources: "ObjectStorage"
-  AvailableExternalDataSources: "ClickHouse"
-  AvailableExternalDataSources: "PostgreSQL"
-  AvailableExternalDataSources: "MySQL"
-  AvailableExternalDataSources: "Ydb"
-  AvailableExternalDataSources: "YT"
-  AvailableExternalDataSources: "Greenplum"
-  AvailableExternalDataSources: "MsSQLServer"
-  AvailableExternalDataSources: "Oracle"
-  AvailableExternalDataSources: "Logging"
-  AvailableExternalDataSources: "Solomon"
+  AllExternalDataSourcesAreAvailable: true
 
   FileStorage {
     MaxFiles: 1000

Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,7 @@ IExternalSourceFactory::TPtr CreateExternalSourceFactory(const std::vector<TStri`
`18`	`18`	`std::shared_ptr<NYql::ISecuredServiceAccountCredentialsFactory> credentialsFactory = nullptr,`
`19`	`19`	`bool enableInfer = false,`
`20`	`20`	`bool allowLocalFiles = false,`
	`21`	`+ bool allExternalDataSourcesAreAvailable = true,`
`21`	`22`	`const std::set<TString>& availableExternalDataSources = {});`
`22`	`23`
`23`	`24`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1172,6 +1172,7 @@ class TKqpHost : public IKqpHost {`
`1172`	`1172`	`FederatedQuerySetup ? FederatedQuerySetup->CredentialsFactory : nullptr,`
`1173`	`1173`	`Config->FeatureFlags.GetEnableExternalSourceSchemaInference(),`
`1174`	`1174`	`FederatedQuerySetup->S3GatewayConfig.GetAllowLocalFiles(),`
	`1175`	`+ QueryServiceConfig.GetAllExternalDataSourcesAreAvailable(),`
`1175`	`1176`	`std::set<TString>(availableExternalDataSources.cbegin(), availableExternalDataSources.cend()));`
`1176`	`1177`	`}`
`1177`	`1178`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1077,7 +1077,8 @@ message TQueryServiceConfig {`
`1077`	`1077`	`optional uint64 ProgressStatsPeriodMs = 14 [default = 0]; // 0 = disabled`
`1078`	`1078`	`optional uint32 QueryTimeoutDefaultSeconds = 19 [default = 1800];`
`1079`	`1079`	`optional bool EnableMatchRecognize = 20 [default = false];`
`1080`		`- repeated string AvailableExternalDataSources = 22;`
	`1080`	`+ repeated string AvailableExternalDataSources = 22; // Ignored if AllExternalDataSourcesAreAvailable is true`
	`1081`	`+ optional bool AllExternalDataSourcesAreAvailable = 23 [default = true];`
`1081`	`1082`	`}`
`1082`	`1083`
`1083`	`1084`	`// Config describes immediate controls and allows`
Original file line number	Diff line number	Diff line change
`@@ -7287,6 +7287,7 @@ void TSchemeShard::ApplyConsoleConfigs(const NKikimrConfig::TAppConfig& appConfi`
`7287`	`7287`	`nullptr,`
`7288`	`7288`	`appConfig.GetFeatureFlags().GetEnableExternalSourceSchemaInference(),`
`7289`	`7289`	`appConfig.GetQueryServiceConfig().GetS3().GetAllowLocalFiles(),`
	`7290`	`+ appConfig.GetQueryServiceConfig().GetAllExternalDataSourcesAreAvailable(),`
`7290`	`7291`	`std::set<TString>(availableExternalDataSources.cbegin(), availableExternalDataSources.cend())`
`7291`	`7292`	`);`
`7292`	`7293`	`}`
Original file line number	Diff line number	Diff line change
`@@ -314,7 +314,7 @@ class TJsonDescribe : public TViewerPipeClient {`
`314`	`314`	`return;`
`315`	`315`	`}`
`316`	`316`
`317`		`- NExternalSource::IExternalSourceFactory::TPtr externalSourceFactory{NExternalSource::CreateExternalSourceFactory({}, nullptr, 50000, nullptr, false, false, NYql::GetAllExternalDataSourceTypes())};`
	`317`	`+ NExternalSource::IExternalSourceFactory::TPtr externalSourceFactory{NExternalSource::CreateExternalSourceFactory({}, nullptr, 50000, nullptr, false, false, true, NYql::GetAllExternalDataSourceTypes())};`
`318`	`318`	`NJson::TJsonValue root;`
`319`	`319`	`const auto& sourceType = DescribeResult->GetPathDescription().GetExternalTableDescription().GetSourceType();`
`320`	`320`	`try {`