diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 84c2890e575c..3bf486866579 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1,54 +1 @@ -/*.md @ydb-platform/docs - -/build @ydb-platform/ci -/certs @ydb-platform/ci -/contrib @ydb-platform/ci -/devtools @ydb-platform/ci -/library @ydb-platform/ci -/scripts @ydb-platform/ci -/tools @ydb-platform/ci -/util @ydb-platform/ci -/vendor @ydb-platform/ci -/yql @ydb-platform/ci -/yt @ydb-platform/ci - -/ydb/core/fq/ @ydb-platform/fq -/ydb/core/public_http/ @ydb-platform/fq - -/ydb/docs/ @ydb-platform/docs - -/ydb/library/yql/ @ydb-platform/yql -/ydb/library/yql/dq @ydb-platform/yql @ydb-platform/qp -/ydb/library/yql/dq/actors/common @ydb-platform/fq -/ydb/library/yql/providers/common/http_gateway @ydb-platform/fq -/ydb/library/yql/providers/common/db_id_async_resolver @ydb-platform/fq -/ydb/library/yql/providers/common/pushdown @ydb-platform/fq -/ydb/library/yql/providers/generic @ydb-platform/fq -/ydb/library/yql/providers/pq @ydb-platform/fq -/ydb/library/yql/providers/s3 @ydb-platform/fq -/ydb/library/yql/providers/solomon @ydb-platform/fq -/ydb/library/yql/tests/sql/solomon @ydb-platform/fq -/ydb/library/yql/tests/sql/suites/solomon @ydb-platform/fq -/ydb/library/yql/udfs/common/clickhouse/client @ydb-platform/fq - -/ydb/library/yql/yt @Krock21 @Krisha11 @zlobober @gritukan - -/ydb/services/fq/ @ydb-platform/fq - -/ydb/core/kafka_proxy @ydb-platform/Topics -/ydb/core/persqueue @ydb-platform/Topics -/ydb/services/datastreams @ydb-platform/Topics -/ydb/services/deprecated/persqueue_v0 @ydb-platform/Topics -/ydb/services/persqueue_v1 @ydb-platform/Topics - -/ydb/core/config/ut @ydb-platform/core - -/ydb/core/viewer @ydb-platform/ui-backend -/ydb/core/protos/node_whiteboard.proto @ydb-platform/ui-backend - -/ydb/core/formats/arrow @ydb-platform/cs -/ydb/core/tx/columnshard @ydb-platform/cs - -/ydb/apps/ydb @ydb-platform/cli -/ydb/public/lib/ydb_cli @ydb-platform/cli - +* @ydb-platform/fq diff --git a/ydb/core/driver_lib/run/config_helpers.cpp b/ydb/core/driver_lib/run/config_helpers.cpp index 2d0b834eb713..66a37eee15f1 100644 --- a/ydb/core/driver_lib/run/config_helpers.cpp +++ b/ydb/core/driver_lib/run/config_helpers.cpp @@ -1,5 +1,9 @@ #include "config_helpers.h" +#include +#include +#include + #include @@ -114,4 +118,31 @@ NActors::TSchedulerConfig CreateSchedulerConfig(const NKikimrConfig::TActorSyste } // namespace NActorSystemConfigHelpers +namespace NKikimrConfigHelpers { + +NMemory::TResourceBrokerConfig CreateMemoryControllerResourceBrokerConfig(const NKikimrConfig::TAppConfig& config) { + NMemory::TResourceBrokerConfig resourceBrokerSelfConfig; // for backward compatibility + auto mergeResourceBrokerConfigs = [&](const NKikimrResourceBroker::TResourceBrokerConfig& resourceBrokerConfig) { + if (resourceBrokerConfig.HasResourceLimit() && resourceBrokerConfig.GetResourceLimit().HasMemory()) { + resourceBrokerSelfConfig.LimitBytes = resourceBrokerConfig.GetResourceLimit().GetMemory(); + } + for (const auto& queue : resourceBrokerConfig.GetQueues()) { + if (queue.GetName() == NLocalDb::KqpResourceManagerQueue) { + if (queue.HasLimit() && queue.GetLimit().HasMemory()) { + resourceBrokerSelfConfig.QueryExecutionLimitBytes = queue.GetLimit().GetMemory(); + } + } + } + }; + if (config.HasBootstrapConfig() && config.GetBootstrapConfig().HasResourceBroker()) { + mergeResourceBrokerConfigs(config.GetBootstrapConfig().GetResourceBroker()); + } + if (config.HasResourceBrokerConfig()) { + mergeResourceBrokerConfigs(config.GetResourceBrokerConfig()); + } + return resourceBrokerSelfConfig; +} + +} // namespace NKikimrConfigHelpers + } // namespace NKikimr diff --git a/ydb/core/driver_lib/run/config_helpers.h b/ydb/core/driver_lib/run/config_helpers.h index d38e416b6f5b..39a0c6cdb053 100644 --- a/ydb/core/driver_lib/run/config_helpers.h +++ b/ydb/core/driver_lib/run/config_helpers.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -15,4 +16,10 @@ NActors::TSchedulerConfig CreateSchedulerConfig(const NKikimrConfig::TActorSyste } // namespace NActorSystemConfigHelpers +namespace NKikimrConfigHelpers { + +NMemory::TResourceBrokerConfig CreateMemoryControllerResourceBrokerConfig(const NKikimrConfig::TAppConfig& config); + +} // namespace NKikimrConfigHelpers + } // namespace NKikimr diff --git a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp index 3f073ab52be3..e3d4ef977249 100644 --- a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp +++ b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp @@ -2049,28 +2049,8 @@ void TMemoryControllerInitializer::InitializeServices( NActors::TActorSystemSetup* setup, const NKikimr::TAppData* appData) { - NMemory::TResourceBrokerConfig resourceBrokerSelfConfig; // for backward compatibility - auto mergeResourceBrokerConfigs = [&](const NKikimrResourceBroker::TResourceBrokerConfig& resourceBrokerConfig) { - if (resourceBrokerConfig.HasResourceLimit() && resourceBrokerConfig.GetResourceLimit().HasMemory()) { - resourceBrokerSelfConfig.LimitBytes = resourceBrokerConfig.GetResourceLimit().GetMemory(); - } - for (const auto& queue : resourceBrokerConfig.GetQueues()) { - if (queue.GetName() == NLocalDb::KqpResourceManagerQueue) { - if (queue.HasLimit() && queue.GetLimit().HasMemory()) { - resourceBrokerSelfConfig.QueryExecutionLimitBytes = queue.GetLimit().GetMemory(); - } - } - } - }; - if (Config.HasBootstrapConfig() && Config.GetBootstrapConfig().HasResourceBroker()) { - mergeResourceBrokerConfigs(Config.GetBootstrapConfig().GetResourceBroker()); - } - if (Config.HasResourceBrokerConfig()) { - mergeResourceBrokerConfigs(Config.GetResourceBrokerConfig()); - } - auto* actor = NMemory::CreateMemoryController(TDuration::Seconds(1), ProcessMemoryInfoProvider, - Config.GetMemoryControllerConfig(), resourceBrokerSelfConfig, + Config.GetMemoryControllerConfig(), NKikimrConfigHelpers::CreateMemoryControllerResourceBrokerConfig(Config), appData->Counters); setup->LocalServices.emplace_back( NMemory::MakeMemoryControllerId(0), diff --git a/ydb/core/external_sources/s3/ut/s3_aws_credentials_ut.cpp b/ydb/core/external_sources/s3/ut/s3_aws_credentials_ut.cpp index bf2b1d0f2b2b..1d892504d002 100644 --- a/ydb/core/external_sources/s3/ut/s3_aws_credentials_ut.cpp +++ b/ydb/core/external_sources/s3/ut/s3_aws_credentials_ut.cpp @@ -251,6 +251,88 @@ Y_UNIT_TEST_SUITE(S3AwsCredentials) { } } + + Y_UNIT_TEST(TestInsertEscaping) { + const TString externalDataSourceName = "/Root/external_data_source"; + auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory(); + auto kikimr = MakeKikimrRunner(true, nullptr, nullptr, std::nullopt, s3ActorsFactory); + + auto tc = kikimr->GetTableClient(); + auto session = tc.CreateSession().GetValueSync().GetSession(); + const TString query = fmt::format(R"( + CREATE OBJECT id (TYPE SECRET) WITH (value=`minio`); + CREATE OBJECT key (TYPE SECRET) WITH (value=`minio123`); + CREATE EXTERNAL DATA SOURCE `{external_source}` WITH ( + SOURCE_TYPE="ObjectStorage", + LOCATION="{location}", + AUTH_METHOD="AWS", + AWS_ACCESS_KEY_ID_SECRET_NAME="id", + AWS_SECRET_ACCESS_KEY_SECRET_NAME="key", + AWS_REGION="ru-central-1" + ); + )", + "external_source"_a = externalDataSourceName, + "location"_a = "localhost:" + GetExternalPort("minio", "9000") + "/datalake/" + ); + + auto result = session.ExecuteSchemeQuery(query).GetValueSync(); + UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString()); + + WaitBucket(kikimr, externalDataSourceName); + + auto db = kikimr->GetQueryClient(); + + TString path = TStringBuilder() << "exp_folder/some_" << EscapeC(GetSymbolsString(' ', '~', "*?{}`")) << "\\`"; + + { + // NB: AtomicUploadCommit = "false" because in minio ListMultipartUploads by prefix is not supported + auto scriptExecutionOperation = db.ExecuteScript(fmt::format(R"( + PRAGMA s3.AtomicUploadCommit = "false"; + INSERT INTO `{external_source}`.`{path}/` WITH (FORMAT = "csv_with_names") + SELECT * FROM `{external_source}`.`/a/` WITH ( + format="json_each_row", + schema( + key Utf8 NOT NULL, + value Utf8 NOT NULL + ) + ) + )", "external_source"_a = externalDataSourceName, "path"_a = path)).ExtractValueSync(); + UNIT_ASSERT_VALUES_EQUAL_C(scriptExecutionOperation.Status().GetStatus(), EStatus::SUCCESS, scriptExecutionOperation.Status().GetIssues().ToString()); + UNIT_ASSERT(!scriptExecutionOperation.Metadata().ExecutionId.empty()); + + NYdb::NQuery::TScriptExecutionOperation readyOp = WaitScriptExecutionOperation(scriptExecutionOperation.Id(), kikimr->GetDriver()); + UNIT_ASSERT_EQUAL_C(readyOp.Metadata().ExecStatus, EExecStatus::Completed, readyOp.Status().GetIssues().ToString()); + } + + { + auto scriptExecutionOperation = db.ExecuteScript(fmt::format(R"( + SELECT * FROM `{external_source}`.`{path}/` WITH ( + format="csv_with_names", + schema( + key Utf8 NOT NULL, + value Utf8 NOT NULL + ) + ) + )", "external_source"_a = externalDataSourceName, "path"_a = path)).ExtractValueSync(); + UNIT_ASSERT_VALUES_EQUAL_C(scriptExecutionOperation.Status().GetStatus(), EStatus::SUCCESS, scriptExecutionOperation.Status().GetIssues().ToString()); + UNIT_ASSERT(!scriptExecutionOperation.Metadata().ExecutionId.empty()); + + NYdb::NQuery::TScriptExecutionOperation readyOp = WaitScriptExecutionOperation(scriptExecutionOperation.Id(), kikimr->GetDriver()); + UNIT_ASSERT_EQUAL_C(readyOp.Metadata().ExecStatus, EExecStatus::Completed, readyOp.Status().GetIssues().ToString()); + TFetchScriptResultsResult results = db.FetchScriptResults(scriptExecutionOperation.Id(), 0).ExtractValueSync(); + UNIT_ASSERT_C(results.IsSuccess(), results.GetIssues().ToString()); + + TResultSetParser resultSet(results.ExtractResultSet()); + UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnsCount(), 2); + UNIT_ASSERT_VALUES_EQUAL(resultSet.RowsCount(), 2); + UNIT_ASSERT(resultSet.TryNextRow()); + UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(0).GetUtf8(), "1"); + UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(1).GetUtf8(), "trololo"); + UNIT_ASSERT(resultSet.TryNextRow()); + UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(0).GetUtf8(), "2"); + UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(1).GetUtf8(), "hello world"); + } + } } } // namespace NKikimr::NKqp diff --git a/ydb/core/fq/libs/actors/clusters_from_connections.cpp b/ydb/core/fq/libs/actors/clusters_from_connections.cpp index 90434706f219..3438e85fd5ca 100644 --- a/ydb/core/fq/libs/actors/clusters_from_connections.cpp +++ b/ydb/core/fq/libs/actors/clusters_from_connections.cpp @@ -1,7 +1,7 @@ #include "clusters_from_connections.h" #include -#include +#include #include #include #include @@ -88,6 +88,19 @@ std::pair ParseHttpEndpoint(const TString& endpoint) { return std::make_pair(ToString(host), scheme != "http"); } +std::pair ParseGrpcEndpoint(const TString& endpoint) { + TStringBuf scheme; + TStringBuf address; + TStringBuf uri; + NHttp::CrackURL(endpoint, scheme, address, uri); + + TString hostname; + TIpPort port; + NHttp::CrackAddress(TString(address), hostname, port); + + return {hostname, port}; +} + void FillSolomonClusterConfig(NYql::TSolomonClusterConfig& clusterConfig, const TString& name, const TString& authToken, @@ -112,7 +125,7 @@ void FillGenericClusterConfigBase( TGenericClusterConfig& clusterCfg, const TConnection& connection, const TString& connectionName, - NConnector::NApi::EDataSourceKind dataSourceKind, + NYql::EGenericDataSourceKind dataSourceKind, const TString& authToken, const THashMap& accountIdSignatures ) { @@ -128,21 +141,21 @@ void FillGenericClusterConfigBase( // In YQv1 we just hardcode the appropriate protocols here. // In YQv2 protocol can be configured via `CREATE EXTERNAL DATA SOURCE` params. switch (dataSourceKind) { - case NYql::NConnector::NApi::CLICKHOUSE: - clusterCfg.SetProtocol(common.GetUseNativeProtocolForClickHouse() ? NYql::NConnector::NApi::EProtocol::NATIVE : NYql::NConnector::NApi::EProtocol::HTTP); + case NYql::EGenericDataSourceKind::CLICKHOUSE: + clusterCfg.SetProtocol(common.GetUseNativeProtocolForClickHouse() ? NYql::EGenericProtocol::NATIVE : NYql::EGenericProtocol::HTTP); break; - case NYql::NConnector::NApi::GREENPLUM: - clusterCfg.SetProtocol(NYql::NConnector::NApi::EProtocol::NATIVE); + case NYql::EGenericDataSourceKind::GREENPLUM: + clusterCfg.SetProtocol(NYql::EGenericProtocol::NATIVE); break; - case NYql::NConnector::NApi::MYSQL: - clusterCfg.SetProtocol(NYql::NConnector::NApi::EProtocol::NATIVE); + case NYql::EGenericDataSourceKind::MYSQL: + clusterCfg.SetProtocol(NYql::EGenericProtocol::NATIVE); break; - case NYql::NConnector::NApi::POSTGRESQL: - clusterCfg.SetProtocol(NYql::NConnector::NApi::EProtocol::NATIVE); + case NYql::EGenericDataSourceKind::POSTGRESQL: + clusterCfg.SetProtocol(NYql::EGenericProtocol::NATIVE); break; default: ythrow yexception() << "Unexpected data source kind: '" - << NYql::NConnector::NApi::EDataSourceKind_Name(dataSourceKind) << "'"; + << NYql::EGenericDataSourceKind_Name(dataSourceKind) << "'"; } ValidateGenericClusterConfig(clusterCfg, "NFq::FillGenericClusterFromConfig"); @@ -154,7 +167,7 @@ void FillGenericClusterConfig( TGenericClusterConfig& clusterCfg, const TConnection& connection, const TString& connectionName, - NConnector::NApi::EDataSourceKind dataSourceKind, + NYql::EGenericDataSourceKind dataSourceKind, const TString& authToken, const THashMap& accountIdSignatures ) { @@ -167,7 +180,7 @@ void FillGenericClusterConfig( TGenericClusterConfig& clusterCfg, const FederatedQuery::PostgreSQLCluster& connection, const TString& connectionName, - NConnector::NApi::EDataSourceKind dataSourceKind, + NYql::EGenericDataSourceKind dataSourceKind, const TString& authToken, const THashMap& accountIdSignatures ){ @@ -227,11 +240,21 @@ void AddClustersFromConnections( case FederatedQuery::ConnectionSetting::kYdbDatabase: { const auto& db = conn.content().setting().ydb_database(); auto* clusterCfg = gatewaysConfig.MutableGeneric()->AddClusterMapping(); - clusterCfg->SetKind(NYql::NConnector::NApi::EDataSourceKind::YDB); - clusterCfg->SetProtocol(NYql::NConnector::NApi::EProtocol::NATIVE); + clusterCfg->SetKind(NYql::EGenericDataSourceKind::YDB); + clusterCfg->SetProtocol(NYql::EGenericProtocol::NATIVE); clusterCfg->SetName(connectionName); - clusterCfg->SetDatabaseId(db.database_id()); - clusterCfg->SetUseSsl(!common.GetDisableSslForGenericDataSources()); + if (const auto& databaseId = db.database_id()) { + clusterCfg->SetDatabaseId(databaseId); + clusterCfg->SetUseSsl(!common.GetDisableSslForGenericDataSources()); + } else { + const auto& [host, port] = ParseGrpcEndpoint(db.endpoint()); + + auto& endpoint = *clusterCfg->MutableEndpoint(); + endpoint.set_host(host); + endpoint.set_port(port); + clusterCfg->SetUseSsl(db.secure()); + clusterCfg->SetDatabaseName(db.database()); + } FillClusterAuth(*clusterCfg, db.auth(), authToken, accountIdSignatures); clusters.emplace(connectionName, GenericProviderName); break; @@ -242,7 +265,7 @@ void AddClustersFromConnections( *gatewaysConfig.MutableGeneric()->AddClusterMapping(), conn.content().setting().clickhouse_cluster(), connectionName, - NYql::NConnector::NApi::EDataSourceKind::CLICKHOUSE, + NYql::EGenericDataSourceKind::CLICKHOUSE, authToken, accountIdSignatures); clusters.emplace(connectionName, GenericProviderName); @@ -275,7 +298,7 @@ void AddClustersFromConnections( *gatewaysConfig.MutableGeneric()->AddClusterMapping(), conn.content().setting().postgresql_cluster(), connectionName, - NYql::NConnector::NApi::EDataSourceKind::POSTGRESQL, + NYql::EGenericDataSourceKind::POSTGRESQL, authToken, accountIdSignatures); clusters.emplace(connectionName, GenericProviderName); @@ -287,7 +310,7 @@ void AddClustersFromConnections( *gatewaysConfig.MutableGeneric()->AddClusterMapping(), conn.content().setting().greenplum_cluster(), connectionName, - NYql::NConnector::NApi::EDataSourceKind::GREENPLUM, + NYql::EGenericDataSourceKind::GREENPLUM, authToken, accountIdSignatures); clusters.emplace(connectionName, GenericProviderName); @@ -299,7 +322,7 @@ void AddClustersFromConnections( *gatewaysConfig.MutableGeneric()->AddClusterMapping(), conn.content().setting().mysql_cluster(), connectionName, - NYql::NConnector::NApi::EDataSourceKind::MYSQL, + NYql::EGenericDataSourceKind::MYSQL, authToken, accountIdSignatures); clusters.emplace(connectionName, GenericProviderName); @@ -308,7 +331,7 @@ void AddClustersFromConnections( case FederatedQuery::ConnectionSetting::kLogging: { const auto& connection = conn.content().setting().logging(); auto* clusterCfg = gatewaysConfig.MutableGeneric()->AddClusterMapping(); - clusterCfg->SetKind(NYql::NConnector::NApi::EDataSourceKind::LOGGING); + clusterCfg->SetKind(NYql::EGenericDataSourceKind::LOGGING); clusterCfg->SetName(connectionName); clusterCfg->mutable_datasourceoptions()->insert({"folder_id", connection.folder_id()}); FillClusterAuth(*clusterCfg, connection.auth(), authToken, accountIdSignatures); diff --git a/ydb/core/fq/libs/actors/database_resolver.cpp b/ydb/core/fq/libs/actors/database_resolver.cpp index 98be09d7335a..ef033e7047e4 100644 --- a/ydb/core/fq/libs/actors/database_resolver.cpp +++ b/ydb/core/fq/libs/actors/database_resolver.cpp @@ -29,7 +29,7 @@ using TParser = std::function; using TParsers = THashMap; @@ -292,7 +292,7 @@ class TDatabaseResolver: public TActor .SetErrorTtl(TDuration::Minutes(1)) .SetMaxSize(1000000)) { - auto ydbParser = [](NJson::TJsonValue& databaseInfo, const NYql::IMdbEndpointGenerator::TPtr&, bool, NConnector::NApi::EProtocol) { + auto ydbParser = [](NJson::TJsonValue& databaseInfo, const NYql::IMdbEndpointGenerator::TPtr&, bool, NYql::EGenericProtocol) { bool secure = false; TString endpoint = databaseInfo.GetMap().at("endpoint").GetStringRobust(); TString prefix("/?database="); @@ -333,7 +333,7 @@ class TDatabaseResolver: public TActor NJson::TJsonValue& databaseInfo, const NYql::IMdbEndpointGenerator::TPtr& mdbEndpointGenerator, bool useTls, - NConnector::NApi::EProtocol protocol) + NYql::EGenericProtocol protocol) { auto ret = ydbParser(databaseInfo, mdbEndpointGenerator, useTls, protocol); // TODO: Take explicit field from MVP @@ -349,7 +349,7 @@ class TDatabaseResolver: public TActor NJson::TJsonValue& databaseInfo, const NYql::IMdbEndpointGenerator::TPtr& mdbEndpointGenerator, bool useTls, - NConnector::NApi::EProtocol protocol + NYql::EGenericProtocol protocol ) { NYql::IMdbEndpointGenerator::TEndpoint endpoint; TVector aliveHosts; @@ -380,7 +380,7 @@ class TDatabaseResolver: public TActor NJson::TJsonValue& databaseInfo, const NYql::IMdbEndpointGenerator::TPtr& mdbEndpointGenerator, bool useTls, - NConnector::NApi::EProtocol protocol + NYql::EGenericProtocol protocol ) { NYql::IMdbEndpointGenerator::TEndpoint endpoint; TVector aliveHosts; @@ -427,7 +427,7 @@ class TDatabaseResolver: public TActor NJson::TJsonValue& databaseInfo, const NYql::IMdbEndpointGenerator::TPtr& mdbEndpointGenerator, bool useTls, - NConnector::NApi::EProtocol protocol + NYql::EGenericProtocol protocol ) { NYql::IMdbEndpointGenerator::TEndpoint endpoint; TString aliveHost; @@ -465,7 +465,7 @@ class TDatabaseResolver: public TActor NJson::TJsonValue& databaseInfo, const NYql::IMdbEndpointGenerator::TPtr& mdbEndpointGenerator, bool useTls, - NConnector::NApi::EProtocol protocol + NYql::EGenericProtocol protocol ) { NYql::IMdbEndpointGenerator::TEndpoint endpoint; TVector aliveHosts; diff --git a/ydb/core/fq/libs/actors/pending_fetcher.cpp b/ydb/core/fq/libs/actors/pending_fetcher.cpp index a91a01498ebd..99cef5053e0c 100644 --- a/ydb/core/fq/libs/actors/pending_fetcher.cpp +++ b/ydb/core/fq/libs/actors/pending_fetcher.cpp @@ -53,6 +53,7 @@ #include #include #include +#include #include #include @@ -135,7 +136,7 @@ class TPendingFetcher : public NActors::TActorBootstrapped { private: static ::NMonitoring::IHistogramCollectorPtr GetLatencyHistogramBuckets() { - return ::NMonitoring::ExplicitHistogram({0, 1, 2, 5, 10, 20, 50, 100, 500, 1000, 2000, 5000, 10000, 30000, 50000, 500000}); + return ::NMonitoring::ExplicitHistogram({0, 10, 100, 1000, 10000}); } }; @@ -156,7 +157,8 @@ class TPendingFetcher : public NActors::TActorBootstrapped { const ::NMonitoring::TDynamicCounterPtr& clientCounters, const TString& tenantName, NActors::TMon* monitoring, - std::shared_ptr s3ActorsFactory + std::shared_ptr s3ActorsFactory, + NYql::IPqGatewayFactory::TPtr pqGatewayFactory ) : YqSharedResources(yqSharedResources) , CredentialsProviderFactory(credentialsProviderFactory) @@ -179,6 +181,7 @@ class TPendingFetcher : public NActors::TActorBootstrapped { , Monitoring(monitoring) , ComputeConfig(config.GetCompute()) , S3ActorsFactory(std::move(s3ActorsFactory)) + , PqGatewayFactory(std::move(pqGatewayFactory)) { Y_ENSURE(GetYqlDefaultModuleResolverWithContext(ModuleResolver)); } @@ -358,6 +361,16 @@ class TPendingFetcher : public NActors::TActorBootstrapped { const TString folderId = NYdb::NFq::TScope(task.scope()).ParseFolder(); const TString cloudId = task.sensor_labels().at("cloud_id"); const TString queryId = task.query_id().value(); + const bool isStreaming = task.query_type() == FederatedQuery::QueryContent::STREAMING; + TString queryIdLabel; + TString queryNameLabel = SanitizeLabel(task.query_name()); + if (task.automatic()) { + queryIdLabel = isStreaming ? "streaming" : "analytics"; + } else if (isStreaming) { + queryIdLabel = queryId; + } else { + queryIdLabel = "manual"; + } ::NYql::NCommon::TServiceCounters queryCounters(ServiceCounters); auto publicCountersParent = ServiceCounters.PublicCounters; @@ -365,8 +378,15 @@ class TPendingFetcher : public NActors::TActorBootstrapped { if (cloudId && folderId) { publicCountersParent = publicCountersParent->GetSubgroup("cloud_id", cloudId)->GetSubgroup("folder_id", folderId); } - queryCounters.PublicCounters = publicCountersParent->GetSubgroup("query_id", - task.automatic() ? (task.query_name() ? task.query_name() : "automatic") : queryId); + + ::NMonitoring::TDynamicCounterPtr queryPublicCounters = publicCountersParent; + // use original query id here + queryPublicCounters = queryPublicCounters->GetSubgroup("query_id", queryId); + + if (queryNameLabel) { + queryPublicCounters = queryPublicCounters->GetSubgroup("query_name", queryNameLabel); + } + queryCounters.PublicCounters = queryPublicCounters; auto rootCountersParent = ServiceCounters.RootCounters; std::set> sensorLabels(task.sensor_labels().begin(), task.sensor_labels().end()); @@ -374,9 +394,17 @@ class TPendingFetcher : public NActors::TActorBootstrapped { rootCountersParent = rootCountersParent->GetSubgroup(label, item); } - queryCounters.RootCounters = rootCountersParent->GetSubgroup("query_id", - task.automatic() ? (folderId ? "automatic_" + folderId : "automatic") : queryId); - queryCounters.Counters = queryCounters.RootCounters; + ::NMonitoring::TDynamicCounterPtr queryRootCounters = rootCountersParent; + if (queryIdLabel) { + queryRootCounters = queryRootCounters->GetSubgroup("query_id", queryIdLabel); + } + + if (!task.automatic() && isStreaming && queryNameLabel) { + queryRootCounters = queryRootCounters->GetSubgroup("query_name", queryNameLabel); + } + + queryCounters.RootCounters = queryRootCounters; + queryCounters.Counters = queryRootCounters; queryCounters.InitUptimeCounter(); const auto createdAt = TInstant::Now(); @@ -446,7 +474,8 @@ class TPendingFetcher : public NActors::TActorBootstrapped { NProtoInterop::CastFromProto(task.result_ttl()), std::map(task.parameters().begin(), task.parameters().end()), S3ActorsFactory, - ComputeConfig.GetWorkloadManagerConfig(task.scope()) + ComputeConfig.GetWorkloadManagerConfig(task.scope()), + PqGatewayFactory ); auto runActorId = @@ -455,9 +484,7 @@ class TPendingFetcher : public NActors::TActorBootstrapped { : Register(CreateRunActor(SelfId(), queryCounters, std::move(params))); RunActorMap[runActorId] = TRunActorInfo { .QueryId = queryId, .QueryName = task.query_name() }; - if (!task.automatic()) { - CountersMap[queryId] = { rootCountersParent, publicCountersParent, runActorId }; - } + CountersMap[queryId] = { rootCountersParent, publicCountersParent, runActorId }; } NActors::IActor* CreateYdbRunActor(TRunActorParams&& params, const ::NYql::NCommon::TServiceCounters& queryCounters) const { @@ -524,6 +551,7 @@ class TPendingFetcher : public NActors::TActorBootstrapped { NActors::TMon* Monitoring; TComputeConfig ComputeConfig; std::shared_ptr S3ActorsFactory; + NYql::IPqGatewayFactory::TPtr PqGatewayFactory; }; @@ -543,7 +571,8 @@ NActors::IActor* CreatePendingFetcher( const ::NMonitoring::TDynamicCounterPtr& clientCounters, const TString& tenantName, NActors::TMon* monitoring, - std::shared_ptr s3ActorsFactory) + std::shared_ptr s3ActorsFactory, + NYql::IPqGatewayFactory::TPtr pqGatewayFactory) { return new TPendingFetcher( yqSharedResources, @@ -561,7 +590,8 @@ NActors::IActor* CreatePendingFetcher( clientCounters, tenantName, monitoring, - std::move(s3ActorsFactory)); + std::move(s3ActorsFactory), + std::move(pqGatewayFactory)); } TActorId MakePendingFetcherId(ui32 nodeId) { diff --git a/ydb/core/fq/libs/actors/proxy.h b/ydb/core/fq/libs/actors/proxy.h index 92cc9dd13faa..7ebe1f378bf6 100644 --- a/ydb/core/fq/libs/actors/proxy.h +++ b/ydb/core/fq/libs/actors/proxy.h @@ -16,6 +16,7 @@ #include #include #include +#include #include #include @@ -53,7 +54,8 @@ NActors::IActor* CreatePendingFetcher( const ::NMonitoring::TDynamicCounterPtr& clientCounters, const TString& tenantName, NActors::TMon* monitoring, - std::shared_ptr s3ActorsFactory + std::shared_ptr s3ActorsFactory, + NYql::IPqGatewayFactory::TPtr pqGatewayFactory ); NActors::IActor* CreateRunActor( diff --git a/ydb/core/fq/libs/actors/run_actor.cpp b/ydb/core/fq/libs/actors/run_actor.cpp index 950f935c03a1..c1bba77dee44 100644 --- a/ydb/core/fq/libs/actors/run_actor.cpp +++ b/ydb/core/fq/libs/actors/run_actor.cpp @@ -571,6 +571,7 @@ class TRunActor : public NActors::TActorBootstrapped { SelfId(), Params.QueryId, Params.YqSharedResources->UserSpaceYdbDriver, + Params.PqGatewayFactory->CreatePqGateway(), Params.Resources.topic_consumers(), PrepareReadRuleCredentials() ) @@ -1435,6 +1436,7 @@ class TRunActor : public NActors::TActorBootstrapped { SelfId(), Params.QueryId, Params.YqSharedResources->UserSpaceYdbDriver, + Params.PqGatewayFactory->CreatePqGateway(), Params.Resources.topic_consumers(), PrepareReadRuleCredentials() ) @@ -1970,14 +1972,8 @@ class TRunActor : public NActors::TActorBootstrapped { } { - NYql::TPqGatewayServices pqServices( - Params.YqSharedResources->UserSpaceYdbDriver, - Params.PqCmConnections, - Params.CredentialsFactory, - std::make_shared(gatewaysConfig.GetPq()), - Params.FunctionRegistry - ); - const auto pqGateway = NYql::CreatePqNativeGateway(pqServices); + auto pqGateway = Params.PqGatewayFactory->CreatePqGateway(); + pqGateway->UpdateClusterConfigs(std::make_shared(gatewaysConfig.GetPq())); dataProvidersInit.push_back(GetPqDataProviderInitializer(pqGateway, false, dbResolver)); } diff --git a/ydb/core/fq/libs/actors/ut/database_resolver_ut.cpp b/ydb/core/fq/libs/actors/ut/database_resolver_ut.cpp index 7295d7beb804..59a4cb47ae07 100644 --- a/ydb/core/fq/libs/actors/ut/database_resolver_ut.cpp +++ b/ydb/core/fq/libs/actors/ut/database_resolver_ut.cpp @@ -124,7 +124,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { void Test( NYql::EDatabaseType databaseType, - NYql::NConnector::NApi::EProtocol protocol, + NYql::EGenericProtocol protocol, const TString& getUrl, const TString& status, const TString& responseBody, @@ -186,7 +186,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Y_UNIT_TEST(Ydb_Serverless) { Test( NYql::EDatabaseType::Ydb, - NYql::NConnector::NApi::EProtocol::PROTOCOL_UNSPECIFIED, + NYql::EGenericProtocol::PROTOCOL_UNSPECIFIED, "https://ydbc.ydb.cloud.yandex.net:8789/ydbc/cloud-prod/database?databaseId=etn021us5r9rhld1vgbh", "200", R"( @@ -218,7 +218,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Test( NYql::EDatabaseType::Ydb, - NYql::NConnector::NApi::EProtocol::PROTOCOL_UNSPECIFIED, + NYql::EGenericProtocol::PROTOCOL_UNSPECIFIED, "https://ydbc.ydb.cloud.yandex.net:8789/ydbc/cloud-prod/database?databaseId=etn021us5r9rhld1vgbh", "", "", @@ -237,7 +237,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Y_UNIT_TEST(Ydb_Dedicated) { Test( NYql::EDatabaseType::Ydb, - NYql::NConnector::NApi::EProtocol::PROTOCOL_UNSPECIFIED, + NYql::EGenericProtocol::PROTOCOL_UNSPECIFIED, "https://ydbc.ydb.cloud.yandex.net:8789/ydbc/cloud-prod/database?databaseId=etn021us5r9rhld1vgbh", "200", R"( @@ -259,7 +259,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Y_UNIT_TEST(DataStreams_Serverless) { Test( NYql::EDatabaseType::DataStreams, - NYql::NConnector::NApi::EProtocol::PROTOCOL_UNSPECIFIED, + NYql::EGenericProtocol::PROTOCOL_UNSPECIFIED, "https://ydbc.ydb.cloud.yandex.net:8789/ydbc/cloud-prod/database?databaseId=etn021us5r9rhld1vgbh", "200", R"( @@ -280,7 +280,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Y_UNIT_TEST(DataStreams_Dedicated) { Test( NYql::EDatabaseType::DataStreams, - NYql::NConnector::NApi::EProtocol::PROTOCOL_UNSPECIFIED, + NYql::EGenericProtocol::PROTOCOL_UNSPECIFIED, "https://ydbc.ydb.cloud.yandex.net:8789/ydbc/cloud-prod/database?databaseId=etn021us5r9rhld1vgbh", "200", R"( @@ -302,7 +302,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Y_UNIT_TEST(ClickHouseNative) { Test( NYql::EDatabaseType::ClickHouse, - NYql::NConnector::NApi::EProtocol::NATIVE, + NYql::EGenericProtocol::NATIVE, "https://mdb.api.cloud.yandex.net:443/managed-clickhouse/v1/clusters/etn021us5r9rhld1vgbh/hosts", "200", R"({ @@ -336,7 +336,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Y_UNIT_TEST(ClickHouseHttp) { Test( NYql::EDatabaseType::ClickHouse, - NYql::NConnector::NApi::EProtocol::HTTP, + NYql::EGenericProtocol::HTTP, "https://mdb.api.cloud.yandex.net:443/managed-clickhouse/v1/clusters/etn021us5r9rhld1vgbh/hosts", "200", R"({ @@ -381,7 +381,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Test( NYql::EDatabaseType::ClickHouse, - NYql::NConnector::NApi::EProtocol::HTTP, + NYql::EGenericProtocol::HTTP, "https://mdb.api.cloud.yandex.net:443/managed-clickhouse/v1/clusters/etn021us5r9rhld1vgbh/hosts", "403", R"( @@ -405,7 +405,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Y_UNIT_TEST(PostgreSQL) { Test( NYql::EDatabaseType::PostgreSQL, - NYql::NConnector::NApi::EProtocol::NATIVE, + NYql::EGenericProtocol::NATIVE, "https://mdb.api.cloud.yandex.net:443/managed-postgresql/v1/clusters/etn021us5r9rhld1vgbh/hosts", "200", R"({ @@ -454,7 +454,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Test( NYql::EDatabaseType::PostgreSQL, - NYql::NConnector::NApi::EProtocol::NATIVE, + NYql::EGenericProtocol::NATIVE, "https://mdb.api.cloud.yandex.net:443/managed-postgresql/v1/clusters/etn021us5r9rhld1vgbh/hosts", "403", R"( @@ -478,7 +478,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Y_UNIT_TEST(Greenplum_MasterNode) { Test( NYql::EDatabaseType::Greenplum, - NYql::NConnector::NApi::EProtocol::NATIVE, + NYql::EGenericProtocol::NATIVE, "https://mdb.api.cloud.yandex.net:443/managed-greenplum/v1/clusters/etn021us5r9rhld1vgbh/master-hosts", "200", R"({ @@ -520,7 +520,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Test( NYql::EDatabaseType::Greenplum, - NYql::NConnector::NApi::EProtocol::NATIVE, + NYql::EGenericProtocol::NATIVE, "https://mdb.api.cloud.yandex.net:443/managed-greenplum/v1/clusters/etn021us5r9rhld1vgbh/master-hosts", "403", R"( @@ -542,7 +542,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Y_UNIT_TEST(MySQL) { Test( NYql::EDatabaseType::MySQL, - NYql::NConnector::NApi::EProtocol::NATIVE, + NYql::EGenericProtocol::NATIVE, "https://mdb.api.cloud.yandex.net:443/managed-mysql/v1/clusters/etn021us5r9rhld1vgbh/hosts", "200", R"({ @@ -590,7 +590,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { Test( NYql::EDatabaseType::MySQL, - NYql::NConnector::NApi::EProtocol::NATIVE, + NYql::EGenericProtocol::NATIVE, "https://mdb.api.cloud.yandex.net:443/managed-mysql/v1/clusters/etn021us5r9rhld1vgbh/hosts", "403", R"( @@ -624,7 +624,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { }; Test( NYql::EDatabaseType::DataStreams, - NYql::NConnector::NApi::EProtocol::PROTOCOL_UNSPECIFIED, + NYql::EGenericProtocol::PROTOCOL_UNSPECIFIED, "https://ydbc.ydb.cloud.yandex.net:8789/ydbc/cloud-prod/database?databaseId=etn021us5r9rhld1vgbh", "403", R"( @@ -642,7 +642,7 @@ Y_UNIT_TEST_SUITE(TDatabaseResolverTests) { NYql::TDatabaseAuth databaseAuth; databaseAuth.UseTls = true; - databaseAuth.Protocol = NYql::NConnector::NApi::EProtocol::PROTOCOL_UNSPECIFIED; + databaseAuth.Protocol = NYql::EGenericProtocol::PROTOCOL_UNSPECIFIED; TString databaseId1{"etn021us5r9rhld1vgb1"}; TString databaseId2{"etn021us5r9rhld1vgb2"}; diff --git a/ydb/core/fq/libs/actors/ya.make b/ydb/core/fq/libs/actors/ya.make index 1b238dbfd9b1..6aff3c94bc9e 100644 --- a/ydb/core/fq/libs/actors/ya.make +++ b/ydb/core/fq/libs/actors/ya.make @@ -44,6 +44,7 @@ PEERDIR( ydb/core/fq/libs/db_schema ydb/core/fq/libs/events ydb/core/fq/libs/grpc + ydb/core/fq/libs/metrics ydb/core/fq/libs/private_client ydb/core/fq/libs/rate_limiter/utils ydb/core/fq/libs/result_formatter @@ -72,7 +73,6 @@ PEERDIR( ydb/library/yql/providers/dq/provider ydb/library/yql/providers/dq/provider/exec ydb/library/yql/providers/dq/worker_manager/interface - ydb/library/yql/providers/generic/connector/api/common ydb/library/yql/providers/generic/connector/libcpp ydb/library/yql/providers/generic/provider ydb/library/yql/providers/pq/cm_client diff --git a/ydb/core/fq/libs/checkpoint_storage/storage_proxy.cpp b/ydb/core/fq/libs/checkpoint_storage/storage_proxy.cpp index 7037b2734ea0..52d9ab6edfb5 100644 --- a/ydb/core/fq/libs/checkpoint_storage/storage_proxy.cpp +++ b/ydb/core/fq/libs/checkpoint_storage/storage_proxy.cpp @@ -168,40 +168,18 @@ void TStorageProxy::Handle(TEvCheckpointStorage::TEvCreateCheckpointRequest::TPt cookie = ev->Cookie, sender = ev->Sender, totalGraphCheckpointsSizeLimit = Config.GetStateStorageLimits().GetMaxGraphCheckpointsSizeBytes(), - actorSystem = TActivationContext::ActorSystem()] + graphDesc = std::move(event->GraphDescription), + storage = CheckpointStorage] (const NThreading::TFuture& resultFuture) { - auto result = resultFuture.GetValue(); - auto issues = result.second; - - if (issues) { - LOG_STREAMS_STORAGE_SERVICE_AS_WARN(*actorSystem, "[" << coordinatorId << "] [" << checkpointId << "] Failed to fetch total graph checkpoints size: " << issues.ToString()); - actorSystem->Send(sender, new TEvCheckpointStorage::TEvCreateCheckpointResponse(checkpointId, std::move(issues), TString()), 0, cookie); - return false; - } + auto [totalGraphCheckpointsSize, issues] = resultFuture.GetValue(); - auto totalGraphCheckpointsSize = result.first; - - if (totalGraphCheckpointsSize > totalGraphCheckpointsSizeLimit) { + if (!issues && totalGraphCheckpointsSize > totalGraphCheckpointsSizeLimit) { TStringStream ss; - ss << "[" << coordinatorId << "] [" << checkpointId << "] Graph checkpoints size limit exceeded: limit " << totalGraphCheckpointsSizeLimit << ", current checkpoints size: " << totalGraphCheckpointsSize; - auto message = ss.Str(); - LOG_STREAMS_STORAGE_SERVICE_AS_WARN(*actorSystem, message) - issues.AddIssue(message); - LOG_STREAMS_STORAGE_SERVICE_AS_DEBUG(*actorSystem, "[" << coordinatorId << "] [" << checkpointId << "] Send TEvCreateCheckpointResponse"); - actorSystem->Send(sender, new TEvCheckpointStorage::TEvCreateCheckpointResponse(checkpointId, std::move(issues), TString()), 0, cookie); - return false; + ss << "Graph checkpoints size limit exceeded: limit " << totalGraphCheckpointsSizeLimit << ", current checkpoints size: " << totalGraphCheckpointsSize; + issues.AddIssue(std::move(ss.Str())); } - return true; - }) - .Apply([checkpointId = event->CheckpointId, - coordinatorId = event->CoordinatorId, - cookie = ev->Cookie, - sender = ev->Sender, - graphDesc = event->GraphDescription, - storage = CheckpointStorage] - (const NThreading::TFuture& passedSizeLimitCheckFuture) { - if (!passedSizeLimitCheckFuture.GetValue()) { - return NThreading::TFuture(); + if (issues) { + return NThreading::MakeFuture(ICheckpointStorage::TCreateCheckpointResult {TString(), std::move(issues) } ); } if (std::holds_alternative(graphDesc)) { return storage->CreateCheckpoint(coordinatorId, checkpointId, std::get(graphDesc), ECheckpointStatus::Pending); @@ -215,12 +193,8 @@ void TStorageProxy::Handle(TEvCheckpointStorage::TEvCreateCheckpointRequest::TPt sender = ev->Sender, actorSystem = TActivationContext::ActorSystem()] (const NThreading::TFuture& resultFuture) { - if (!resultFuture.Initialized()) { // didn't pass the size limit check - return; - } - auto result = resultFuture.GetValue(); - auto issues = result.second; - auto response = std::make_unique(checkpointId, std::move(issues), result.first); + auto [graphDescId, issues] = resultFuture.GetValue(); + auto response = std::make_unique(checkpointId, std::move(issues), std::move(graphDescId)); if (response->Issues) { LOG_STREAMS_STORAGE_SERVICE_AS_WARN(*actorSystem, "[" << coordinatorId << "] [" << checkpointId << "] Failed to create checkpoint: " << response->Issues.ToString()); } else { diff --git a/ydb/core/fq/libs/checkpoint_storage/ydb_checkpoint_storage.cpp b/ydb/core/fq/libs/checkpoint_storage/ydb_checkpoint_storage.cpp index c270f8e7201f..661db83eee28 100644 --- a/ydb/core/fq/libs/checkpoint_storage/ydb_checkpoint_storage.cpp +++ b/ydb/core/fq/libs/checkpoint_storage/ydb_checkpoint_storage.cpp @@ -57,15 +57,18 @@ struct TCheckpointContext : public TThrRefBase { TGenerationContextPtr GenerationContext; TCheckpointGraphDescriptionContextPtr CheckpointGraphDescriptionContext; IEntityIdGenerator::TPtr EntityIdGenerator; + TExecDataQuerySettings Settings; TCheckpointContext(const TCheckpointId& id, ECheckpointStatus status, ECheckpointStatus expected, - ui64 stateSizeBytes) + ui64 stateSizeBytes, + TExecDataQuerySettings settings) : CheckpointId(id) , Status(status) , ExpectedStatus(expected) , StateSizeBytes(stateSizeBytes) + , Settings(settings) { } }; @@ -218,7 +221,7 @@ TFuture CreateCheckpoint(const TCheckpointContextPtr& context) { } auto ttxControl = TTxControl::Tx(*generationContext->Transaction).CommitTx(); - return generationContext->Session.ExecuteDataQuery(query, ttxControl, params.Build()).Apply( + return generationContext->Session.ExecuteDataQuery(query, ttxControl, params.Build(), context->Settings).Apply( [] (const TFuture& future) { TStatus status = future.GetValue(); return status; @@ -234,21 +237,41 @@ TFuture UpdateCheckpoint(const TCheckpointContextPtr& context) { auto query = Sprintf(R"( --!syntax_v1 PRAGMA TablePathPrefix("%s"); - $ts = cast(%lu as Timestamp); + DECLARE $graph_id AS String; + DECLARE $coordinator_generation AS Uint64; + DECLARE $seq_no AS Uint64; + DECLARE $status AS Uint8; + DECLARE $state_size AS Uint64; + DECLARE $ts AS Timestamp; UPSERT INTO %s (graph_id, coordinator_generation, seq_no, status, state_size, modified_by) VALUES - ("%s", %lu, %lu, %u, %lu, $ts); + ($graph_id, $coordinator_generation, $seq_no, $status, $state_size, $ts); )", generationContext->TablePathPrefix.c_str(), - TInstant::Now().MicroSeconds(), - CheckpointsMetadataTable, - generationContext->PrimaryKey.c_str(), - context->CheckpointId.CoordinatorGeneration, - context->CheckpointId.SeqNo, - (ui32)context->Status, - context->StateSizeBytes); + CheckpointsMetadataTable); + + NYdb::TParamsBuilder params; + params + .AddParam("$graph_id") + .String(generationContext->PrimaryKey) + .Build() + .AddParam("$coordinator_generation") + .Uint64(context->CheckpointId.CoordinatorGeneration) + .Build() + .AddParam("$seq_no") + .Uint64(context->CheckpointId.SeqNo) + .Build() + .AddParam("$status") + .Uint8((ui8)context->Status) + .Build() + .AddParam("$state_size") + .Uint64(context->StateSizeBytes) + .Build() + .AddParam("$ts") + .Timestamp(TInstant::Now()) + .Build(); auto ttxControl = TTxControl::Tx(*generationContext->Transaction).CommitTx(); - return generationContext->Session.ExecuteDataQuery(query, ttxControl).Apply( + return generationContext->Session.ExecuteDataQuery(query, ttxControl, params.Build(), context->Settings).Apply( [] (const TFuture& future) { TStatus status = future.GetValue(); return status; @@ -262,15 +285,20 @@ TFuture SelectGraphDescId(const TCheckpointContextPtr& context auto query = Sprintf(R"( --!syntax_v1 PRAGMA TablePathPrefix("%s"); + DECLARE $graph_desc_id AS String; SELECT ref_count FROM %s - WHERE id = "%s"; + WHERE id = $graph_desc_id; )", generationContext->TablePathPrefix.c_str(), - CheckpointsGraphsDescriptionTable, - graphDescContext->GraphDescId.c_str()); + CheckpointsGraphsDescriptionTable); + NYdb::TParamsBuilder params; + params + .AddParam("$graph_desc_id") + .String(graphDescContext->GraphDescId) + .Build(); - return generationContext->Session.ExecuteDataQuery(query, TTxControl::Tx(*generationContext->Transaction)); + return generationContext->Session.ExecuteDataQuery(query, TTxControl::Tx(*generationContext->Transaction), params.Build(), context->Settings); } bool GraphDescIdExists(const TFuture& result) { @@ -290,6 +318,7 @@ TFuture GenerateGraphDescId(const TCheckpointContextPtr& context) { if (!result.GetValue().IsSuccess()) { return MakeFuture(result.GetValue()); } + // TODO racing! if (!GraphDescIdExists(result)) { return MakeFuture(TStatus(EStatus::SUCCESS, NYql::TIssues())); } else { @@ -441,19 +470,33 @@ TFuture SelectCheckpoint(const TCheckpointContextPtr& context) auto query = Sprintf(R"( --!syntax_v1 PRAGMA TablePathPrefix("%s"); + DECLARE $graph_id AS String; + DECLARE $coordinator_generation AS Uint64; + DECLARE $seq_no AS Uint64; SELECT status FROM %s - WHERE graph_id = "%s" AND coordinator_generation = %lu AND seq_no = %lu; + WHERE graph_id = $graph_id AND coordinator_generation = $coordinator_generation AND seq_no = $seq_no; )", generationContext->TablePathPrefix.c_str(), - CheckpointsMetadataTable, - generationContext->PrimaryKey.c_str(), - context->CheckpointId.CoordinatorGeneration, - context->CheckpointId.SeqNo); + CheckpointsMetadataTable); + + NYdb::TParamsBuilder params; + params + .AddParam("$graph_id") + .String(generationContext->PrimaryKey) + .Build() + .AddParam("$coordinator_generation") + .Uint64(context->CheckpointId.CoordinatorGeneration) + .Build() + .AddParam("$seq_no") + .Uint64(context->CheckpointId.SeqNo) + .Build(); return generationContext->Session.ExecuteDataQuery( query, - TTxControl::Tx(*generationContext->Transaction)); + TTxControl::Tx(*generationContext->Transaction), + params.Build(), + context->Settings); } TFuture CheckCheckpoint( @@ -766,7 +809,7 @@ TFuture TCheckpointStorage::CreateC ECheckpointStatus status) { Y_ABORT_UNLESS(graphDescId); - auto checkpointContext = MakeIntrusive(checkpointId, status, ECheckpointStatus::Pending, 0ul); + auto checkpointContext = MakeIntrusive(checkpointId, status, ECheckpointStatus::Pending, 0ul, DefaultExecDataQuerySettings()); checkpointContext->CheckpointGraphDescriptionContext = MakeIntrusive(graphDescId); return CreateCheckpointImpl(coordinator, checkpointContext); } @@ -777,7 +820,7 @@ TFuture TCheckpointStorage::CreateC const NProto::TCheckpointGraphDescription& graphDesc, ECheckpointStatus status) { - auto checkpointContext = MakeIntrusive(checkpointId, status, ECheckpointStatus::Pending, 0ul); + auto checkpointContext = MakeIntrusive(checkpointId, status, ECheckpointStatus::Pending, 0ul, DefaultExecDataQuerySettings()); checkpointContext->CheckpointGraphDescriptionContext = MakeIntrusive(graphDesc); checkpointContext->EntityIdGenerator = EntityIdGenerator; return CreateCheckpointImpl(coordinator, checkpointContext); @@ -818,7 +861,7 @@ TFuture TCheckpointStorage::UpdateCheckpointStatus( ECheckpointStatus prevStatus, ui64 stateSizeBytes) { - auto checkpointContext = MakeIntrusive(checkpointId, newStatus, prevStatus, stateSizeBytes); + auto checkpointContext = MakeIntrusive(checkpointId, newStatus, prevStatus, stateSizeBytes, DefaultExecDataQuerySettings()); auto future = YdbConnection->TableClient.RetryOperation( [prefix = YdbConnection->TablePathPrefix, coordinator, checkpointContext] (TSession session) { auto generationContext = MakeIntrusive( @@ -844,7 +887,7 @@ TFuture TCheckpointStorage::AbortCheckpoint( const TCoordinatorId& coordinator, const TCheckpointId& checkpointId) { - auto checkpointContext = MakeIntrusive(checkpointId, ECheckpointStatus::Aborted, ECheckpointStatus::Pending, 0ul); + auto checkpointContext = MakeIntrusive(checkpointId, ECheckpointStatus::Aborted, ECheckpointStatus::Pending, 0ul, DefaultExecDataQuerySettings()); auto future = YdbConnection->TableClient.RetryOperation( [prefix = YdbConnection->TablePathPrefix, coordinator, checkpointContext] (TSession session) { auto generationContext = MakeIntrusive( @@ -903,28 +946,35 @@ TFuture TCheckpointStorage::GetCheckp TFuture TCheckpointStorage::DeleteGraph(const TString& graphId) { auto future = YdbConnection->TableClient.RetryOperation( - [prefix = YdbConnection->TablePathPrefix, graphId] (TSession session) { + [prefix = YdbConnection->TablePathPrefix, graphId, settings = DefaultExecDataQuerySettings()] (TSession session) { // TODO: use prepared queries auto query = Sprintf(R"( --!syntax_v1 PRAGMA TablePathPrefix("%s"); + DECLARE $graph_id AS String; DELETE FROM %s - WHERE graph_id = "%s"; + WHERE graph_id = $graph_id; DELETE FROM %s - WHERE graph_id = "%s"; + WHERE graph_id = $graph_id; )", prefix.c_str(), CoordinatorsSyncTable, - graphId.c_str(), - CheckpointsMetadataTable, - graphId.c_str()); + CheckpointsMetadataTable); + + NYdb::TParamsBuilder params; + params + .AddParam("$graph_id") + .String(graphId) + .Build(); auto future = session.ExecuteDataQuery( query, - TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx()); + TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), + params.Build(), + settings); return future.Apply( [] (const TFuture& future) { @@ -941,30 +991,48 @@ TFuture TCheckpointStorage::MarkCheckpointsGC( const TCheckpointId& checkpointUpperBound) { auto future = YdbConnection->TableClient.RetryOperation( - [prefix = YdbConnection->TablePathPrefix, graphId, checkpointUpperBound] (TSession session) { + [prefix = YdbConnection->TablePathPrefix, graphId, checkpointUpperBound, thisPtr = TIntrusivePtr(this)] (TSession session) { // TODO: use prepared queries auto query = Sprintf(R"( --!syntax_v1 PRAGMA TablePathPrefix("%s"); - $ts = cast(%lu as Timestamp); + DECLARE $ts AS Timestamp; + DECLARE $status AS Uint8; + DECLARE $graph_id AS String; + DECLARE $coordinator_generation AS Uint64; + DECLARE $seq_no AS Uint64; UPDATE %s - SET status = %u, modified_by = $ts - WHERE graph_id = "%s" AND - (coordinator_generation < %lu OR - (coordinator_generation = %lu AND seq_no < %lu)); + SET status = $status, modified_by = $ts + WHERE graph_id = $graph_id AND + (coordinator_generation < $coordinator_generation OR + (coordinator_generation = $coordinator_generation AND seq_no < $seq_no)); )", prefix.c_str(), - TInstant::Now().MicroSeconds(), - CheckpointsMetadataTable, - (ui32)ECheckpointStatus::GC, - graphId.c_str(), - checkpointUpperBound.CoordinatorGeneration, - checkpointUpperBound.CoordinatorGeneration, - checkpointUpperBound.SeqNo); + CheckpointsMetadataTable); + + NYdb::TParamsBuilder params; + params + .AddParam("$graph_id") + .String(graphId) + .Build() + .AddParam("$coordinator_generation") + .Uint64(checkpointUpperBound.CoordinatorGeneration) + .Build() + .AddParam("$seq_no") + .Uint64(checkpointUpperBound.SeqNo) + .Build() + .AddParam("$status") + .Uint8((ui8)ECheckpointStatus::GC) + .Build() + .AddParam("$ts") + .Timestamp(TInstant::Now()) + .Build(); auto future = session.ExecuteDataQuery( query, - TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx()); + TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), + params.Build(), + thisPtr->DefaultExecDataQuerySettings()); return future.Apply( [] (const TFuture& future) { @@ -981,7 +1049,7 @@ TFuture TCheckpointStorage::DeleteMarkedCheckpoints( const TCheckpointId& checkpointUpperBound) { auto future = YdbConnection->TableClient.RetryOperation( - [prefix = YdbConnection->TablePathPrefix, graphId, checkpointUpperBound] (TSession session) { + [prefix = YdbConnection->TablePathPrefix, graphId, checkpointUpperBound, settings = DefaultExecDataQuerySettings()] (TSession session) { // TODO: use prepared queries using namespace fmt::literals; const TString query = fmt::format(R"sql( @@ -1040,7 +1108,7 @@ TFuture TCheckpointStorage::DeleteMarkedCheckpoints( auto future = session.ExecuteDataQuery( query, - TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), params.Build()); + TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), params.Build(), settings); return future.Apply( [] (const TFuture& future) { diff --git a/ydb/core/fq/libs/checkpoint_storage/ydb_state_storage.cpp b/ydb/core/fq/libs/checkpoint_storage/ydb_state_storage.cpp index c18e0127ce57..0d9a104fdd92 100644 --- a/ydb/core/fq/libs/checkpoint_storage/ydb_state_storage.cpp +++ b/ydb/core/fq/libs/checkpoint_storage/ydb_state_storage.cpp @@ -725,8 +725,8 @@ TFuture TStateStorage::DeleteGraph(const TString& graphId) { DELETE FROM %s - WHERE graph_id = "%s"; - )", prefix.c_str(), StatesTable, graphId.c_str()); + WHERE graph_id = $graph_id; + )", prefix.c_str(), StatesTable); auto future = session.ExecuteDataQuery( query, diff --git a/ydb/core/fq/libs/checkpointing/checkpoint_coordinator.cpp b/ydb/core/fq/libs/checkpointing/checkpoint_coordinator.cpp index 6c2b5d9afd31..c19abbc3f464 100644 --- a/ydb/core/fq/libs/checkpointing/checkpoint_coordinator.cpp +++ b/ydb/core/fq/libs/checkpointing/checkpoint_coordinator.cpp @@ -372,11 +372,12 @@ void TCheckpointCoordinator::Handle(const TEvCheckpointCoordinator::TEvScheduleC CC_LOG_D("Got TEvScheduleCheckpointing"); ScheduleNextCheckpoint(); const auto checkpointsInFly = PendingCheckpoints.size() + PendingCommitCheckpoints.size(); - if (checkpointsInFly >= Settings.GetMaxInflight() || InitingZeroCheckpoint) { + if (checkpointsInFly >= Settings.GetMaxInflight() || (InitingZeroCheckpoint && !FailedZeroCheckpoint)) { CC_LOG_W("Skip schedule checkpoint event since inflight checkpoint limit exceeded: current: " << checkpointsInFly << ", limit: " << Settings.GetMaxInflight()); Metrics.SkippedDueToInFlightLimit->Inc(); return; } + FailedZeroCheckpoint = false; Metrics.SkippedDueToInFlightLimit->Set(0); InitCheckpoint(); } @@ -386,13 +387,18 @@ void TCheckpointCoordinator::Handle(const TEvCheckpointStorage::TEvCreateCheckpo const auto& issues = ev->Get()->Issues; CC_LOG_D("[" << checkpointId << "] Got TEvCreateCheckpointResponse"); - if (issues) { - CC_LOG_E("[" << checkpointId << "] StorageError: can't create checkpoint: " << issues.ToOneLineString()); + auto cancelCheckpoint = [&](const TString& str) { + CC_LOG_E("[" << checkpointId << "] " << str); PendingCheckpoints.erase(checkpointId); + FailedZeroCheckpoint = InitingZeroCheckpoint; UpdateInProgressMetric(); ++*Metrics.FailedToCreate; ++*Metrics.StorageError; CheckpointingSnapshotRotationIndex = CheckpointingSnapshotRotationPeriod; // Next checkpoint is snapshot. + }; + + if (issues) { + cancelCheckpoint("StorageError: can't create checkpoint: " + issues.ToOneLineString()); return; } @@ -400,7 +406,10 @@ void TCheckpointCoordinator::Handle(const TEvCheckpointStorage::TEvCreateCheckpo Y_ABORT_UNLESS(GraphDescId == ev->Get()->GraphDescId); } else { GraphDescId = ev->Get()->GraphDescId; - Y_ABORT_UNLESS(GraphDescId); + if (!GraphDescId) { + cancelCheckpoint("StorageError (internal error), empty GraphDescId"); + return; + } } if (PendingInit) { @@ -470,6 +479,7 @@ void TCheckpointCoordinator::Handle(const NYql::NDq::TEvDqCompute::TEvSaveTaskSt CC_LOG_E("[" << checkpointId << "] Got all acks for aborted checkpoint, aborting in storage"); CheckpointingSnapshotRotationIndex = CheckpointingSnapshotRotationPeriod; // Next checkpoint is snapshot. Send(StorageProxy, new TEvCheckpointStorage::TEvAbortCheckpointRequest(CoordinatorId, checkpointId, "Can't save node state"), IEventHandle::FlagTrackDelivery); + FailedZeroCheckpoint = InitingZeroCheckpoint; } else { CC_LOG_I("[" << checkpointId << "] Got all acks, changing checkpoint status to 'PendingCommit'"); Send(StorageProxy, new TEvCheckpointStorage::TEvSetCheckpointPendingCommitStatusRequest(CoordinatorId, checkpointId, checkpoint.GetStats().StateSize), IEventHandle::FlagTrackDelivery); @@ -494,6 +504,7 @@ void TCheckpointCoordinator::Handle(const TEvCheckpointStorage::TEvSetCheckpoint CC_LOG_E("[" << checkpointId << "] StorageError: can't change checkpoint status to 'PendingCommit': " << issues.ToString()); ++*Metrics.StorageError; PendingCheckpoints.erase(it); + FailedZeroCheckpoint = InitingZeroCheckpoint; return; } @@ -571,6 +582,7 @@ void TCheckpointCoordinator::Handle(const TEvCheckpointStorage::TEvAbortCheckpoi ++*Metrics.Aborted; } PendingCheckpoints.erase(checkpointId); + FailedZeroCheckpoint = InitingZeroCheckpoint; PendingCommitCheckpoints.erase(checkpointId); UpdateInProgressMetric(); } @@ -616,6 +628,8 @@ void TCheckpointCoordinator::Handle(NActors::TEvents::TEvPoison::TPtr& ev) { } void TCheckpointCoordinator::Handle(const TEvCheckpointCoordinator::TEvRunGraph::TPtr&) { + Y_DEBUG_ABORT_UNLESS(InitingZeroCheckpoint); + Y_DEBUG_ABORT_UNLESS(!FailedZeroCheckpoint); InitingZeroCheckpoint = false; // TODO: run graph only now, not before zero checkpoint inited } diff --git a/ydb/core/fq/libs/checkpointing/checkpoint_coordinator.h b/ydb/core/fq/libs/checkpointing/checkpoint_coordinator.h index 280130f38163..b5740178dadd 100644 --- a/ydb/core/fq/libs/checkpointing/checkpoint_coordinator.h +++ b/ydb/core/fq/libs/checkpointing/checkpoint_coordinator.h @@ -193,6 +193,7 @@ class TCheckpointCoordinator : public NYql::TTaskControllerImpl PendingInit; bool GraphIsRunning = false; bool InitingZeroCheckpoint = false; + bool FailedZeroCheckpoint = false; bool RestoringFromForeignCheckpoint = false; TCheckpointCoordinatorMetrics Metrics; diff --git a/ydb/core/fq/libs/checkpointing/ut/checkpoint_coordinator_ut.cpp b/ydb/core/fq/libs/checkpointing/ut/checkpoint_coordinator_ut.cpp index bab3c2aec960..77cc91fbff37 100644 --- a/ydb/core/fq/libs/checkpointing/ut/checkpoint_coordinator_ut.cpp +++ b/ydb/core/fq/libs/checkpointing/ut/checkpoint_coordinator_ut.cpp @@ -411,7 +411,6 @@ Y_UNIT_TEST_SUITE(TCheckpointCoordinatorTests) { TEvCheckpointStorage::TEvCompleteCheckpointRequest(CoordinatorId, checkpointId, 300, type)); MockCompleteCheckpointResponse(checkpointId); - MockRunGraph(); } void SaveFailed(TCheckpointId checkpointId) { @@ -423,7 +422,6 @@ Y_UNIT_TEST_SUITE(TCheckpointCoordinatorTests) { ExpectEvent(StorageProxy, TEvCheckpointStorage::TEvAbortCheckpointRequest( CoordinatorId, checkpointId, "Can't save node state")); MockAbortCheckpointResponse(checkpointId); - MockRunGraph(); } void ScheduleCheckpointing() { @@ -436,6 +434,7 @@ Y_UNIT_TEST_SUITE(TCheckpointCoordinatorTests) { test.RegisterCoordinator(); test.InjectCheckpoint(test.CheckpointId1); test.AllSavedAndCommited(test.CheckpointId1); + test.MockRunGraph(); } Y_UNIT_TEST(ShouldTriggerCheckpointWithSourcesAndWithChannel) { @@ -443,6 +442,7 @@ Y_UNIT_TEST_SUITE(TCheckpointCoordinatorTests) { test.RegisterCoordinator(); test.InjectCheckpoint(test.CheckpointId1); test.AllSavedAndCommited(test.CheckpointId1); + test.MockRunGraph(); } Y_UNIT_TEST(ShouldAllSnapshots) { @@ -450,6 +450,7 @@ Y_UNIT_TEST_SUITE(TCheckpointCoordinatorTests) { test.RegisterCoordinator(); test.InjectCheckpoint(test.CheckpointId1); test.AllSavedAndCommited(test.CheckpointId1); + test.MockRunGraph(); test.ScheduleCheckpointing(); test.InjectCheckpoint(test.CheckpointId2, test.GraphDescId, NYql::NDqProto::CHECKPOINT_TYPE_SNAPSHOT); @@ -461,6 +462,7 @@ Y_UNIT_TEST_SUITE(TCheckpointCoordinatorTests) { test.RegisterCoordinator(); test.InjectCheckpoint(test.CheckpointId1); test.AllSavedAndCommited(test.CheckpointId1); + test.MockRunGraph(); test.ScheduleCheckpointing(); test.InjectCheckpoint(test.CheckpointId2, test.GraphDescId, NYql::NDqProto::CHECKPOINT_TYPE_INCREMENT_OR_SNAPSHOT); diff --git a/ydb/core/fq/libs/compute/common/config.h b/ydb/core/fq/libs/compute/common/config.h index 8598dc9dc4de..579869553e12 100644 --- a/ydb/core/fq/libs/compute/common/config.h +++ b/ydb/core/fq/libs/compute/common/config.h @@ -68,6 +68,25 @@ class TComputeConfig { } } + TVector GetExternalSourcesAccessSIDs(const TString& scope) const { + const auto& controlPlane = ComputeConfig.GetYdb().GetControlPlane(); + switch (controlPlane.type_case()) { + case NConfig::TYdbComputeControlPlane::TYPE_NOT_SET: + return {}; + case NConfig::TYdbComputeControlPlane::kSingle: + return {}; + case NConfig::TYdbComputeControlPlane::kCms: + return GetExternalSourcesAccessSIDs(scope, controlPlane.GetCms().GetDatabaseMapping()); + case NConfig::TYdbComputeControlPlane::kYdbcp: + return GetExternalSourcesAccessSIDs(scope, controlPlane.GetYdbcp().GetDatabaseMapping()); + } + } + + TVector GetExternalSourcesAccessSIDs(const TString& scope, const ::NFq::NConfig::TDatabaseMapping& databaseMapping) const { + const auto protoExternalSourcesAccessSIDs = GetComputeDatabaseConfig(scope, databaseMapping).GetAccessConfig().GetExternalSourcesAccessSID(); + return TVector{protoExternalSourcesAccessSIDs.begin(), protoExternalSourcesAccessSIDs.end()}; + } + NFq::NConfig::TYdbStorageConfig GetControlPlaneConnection(const TString& scope) const { const auto& controlPlane = ComputeConfig.GetYdb().GetControlPlane(); switch (controlPlane.type_case()) { @@ -83,15 +102,8 @@ class TComputeConfig { } NFq::NConfig::TYdbStorageConfig GetControlPlaneConnection(const TString& scope, const ::NFq::NConfig::TDatabaseMapping& databaseMapping) const { - auto it = databaseMapping.GetScopeToComputeDatabase().find(scope); - if (it != databaseMapping.GetScopeToComputeDatabase().end()) { - return it->second.GetControlPlaneConnection(); - } - return databaseMapping.GetCommon().empty() - ? NFq::NConfig::TYdbStorageConfig{} - : databaseMapping - .GetCommon(MultiHash(scope) % databaseMapping.GetCommon().size()) - .GetControlPlaneConnection(); + auto computeDatabaseConfig = GetComputeDatabaseConfig(scope, databaseMapping); + return computeDatabaseConfig.GetControlPlaneConnection(); } NFq::NConfig::TYdbStorageConfig GetExecutionConnection(const TString& scope) const { @@ -137,18 +149,11 @@ class TComputeConfig { } NFq::NConfig::TYdbStorageConfig GetExecutionConnection(const TString& scope, const ::NFq::NConfig::TDatabaseMapping& databaseMapping) const { - auto it = databaseMapping.GetScopeToComputeDatabase().find(scope); - if (it != databaseMapping.GetScopeToComputeDatabase().end()) { - return it->second.GetExecutionConnection(); - } - return databaseMapping.GetCommon().empty() - ? NFq::NConfig::TYdbStorageConfig{} - : databaseMapping - .GetCommon(MultiHash(scope) % databaseMapping.GetCommon().size()) - .GetExecutionConnection(); + auto computeDatabaseConfig = GetComputeDatabaseConfig(scope, databaseMapping); + return computeDatabaseConfig.GetExecutionConnection(); } - NFq::NConfig::TWorkloadManagerConfig GetWorkloadManagerConfig(const TString& scope) { + NFq::NConfig::TWorkloadManagerConfig GetWorkloadManagerConfig(const TString& scope) const { const auto& controlPlane = ComputeConfig.GetYdb().GetControlPlane(); switch (controlPlane.type_case()) { case NConfig::TYdbComputeControlPlane::TYPE_NOT_SET: @@ -163,24 +168,33 @@ class TComputeConfig { } NFq::NConfig::TWorkloadManagerConfig GetWorkloadManagerConfig(const TString& scope, const ::NFq::NConfig::TDatabaseMapping& databaseMapping) const { + auto computeDatabaseConfig = GetComputeDatabaseConfig(scope, databaseMapping); + return computeDatabaseConfig.GetWorkloadManagerConfig(); + } + + NFq::NConfig::TComputeDatabaseConfig GetComputeDatabaseConfig(const TString& scope, const ::NFq::NConfig::TDatabaseMapping& databaseMapping) const { auto it = databaseMapping.GetScopeToComputeDatabase().find(scope); if (it != databaseMapping.GetScopeToComputeDatabase().end()) { - return it->second.GetWorkloadManagerConfig(); + return it->second; + } + if (databaseMapping.GetCommon().empty()) { + return NFq::NConfig::TComputeDatabaseConfig{}; } - return databaseMapping.GetCommon().empty() - ? NFq::NConfig::TWorkloadManagerConfig{} - : databaseMapping - .GetCommon(MultiHash(scope) % databaseMapping.GetCommon().size()) - .GetWorkloadManagerConfig(); + return databaseMapping.GetCommon(MultiHash(scope) % databaseMapping.GetCommon().size()); } bool YdbComputeControlPlaneEnabled(const TString& scope) const { + return YdbComputeControlPlaneEnabled(scope, FederatedQuery::QueryContent::ANALYTICS) || + YdbComputeControlPlaneEnabled(scope, FederatedQuery::QueryContent::STREAMING); + } + + bool YdbComputeControlPlaneEnabled(const TString& scope, FederatedQuery::QueryContent::QueryType queryType) const { + if (queryType == FederatedQuery::QueryContent::QUERY_TYPE_UNSPECIFIED) { + return YdbComputeControlPlaneEnabled(scope); + } return ComputeConfig.GetYdb().GetEnable() && ComputeConfig.GetYdb().GetControlPlane().GetEnable() && - (GetComputeType(FederatedQuery::QueryContent::ANALYTICS, scope) == - NFq::NConfig::EComputeType::YDB || - GetComputeType(FederatedQuery::QueryContent::STREAMING, scope) == - NFq::NConfig::EComputeType::YDB); + GetComputeType(queryType, scope) == NFq::NConfig::EComputeType::YDB; } bool IsYDBSchemaOperationsEnabled( diff --git a/ydb/core/fq/libs/compute/common/run_actor_params.cpp b/ydb/core/fq/libs/compute/common/run_actor_params.cpp index a6a26c0f9cea..2af6a14c8293 100644 --- a/ydb/core/fq/libs/compute/common/run_actor_params.cpp +++ b/ydb/core/fq/libs/compute/common/run_actor_params.cpp @@ -60,7 +60,8 @@ TRunActorParams::TRunActorParams( TDuration resultTtl, std::map&& queryParameters, std::shared_ptr s3ActorsFactory, - const ::NFq::NConfig::TWorkloadManagerConfig& workloadManager + const ::NFq::NConfig::TWorkloadManagerConfig& workloadManager, + NYql::IPqGatewayFactory::TPtr pqGatewayFactory ) : YqSharedResources(yqSharedResources) , CredentialsProviderFactory(credentialsProviderFactory) @@ -117,6 +118,7 @@ TRunActorParams::TRunActorParams( , QueryParameters(std::move(queryParameters)) , S3ActorsFactory(std::move(s3ActorsFactory)) , WorkloadManager(workloadManager) + , PqGatewayFactory(std::move(pqGatewayFactory)) { } diff --git a/ydb/core/fq/libs/compute/common/run_actor_params.h b/ydb/core/fq/libs/compute/common/run_actor_params.h index 7ae4eda044ff..f49419a6334d 100644 --- a/ydb/core/fq/libs/compute/common/run_actor_params.h +++ b/ydb/core/fq/libs/compute/common/run_actor_params.h @@ -13,6 +13,7 @@ #include #include #include +#include #include #include @@ -79,7 +80,8 @@ struct TRunActorParams { // TODO2 : Change name TDuration resultTtl, std::map&& queryParameters, std::shared_ptr s3ActorsFactory, - const ::NFq::NConfig::TWorkloadManagerConfig& workloadManager + const ::NFq::NConfig::TWorkloadManagerConfig& workloadManager, + NYql::IPqGatewayFactory::TPtr pqGatewayFactory ); TRunActorParams(const TRunActorParams& params) = default; @@ -145,6 +147,7 @@ struct TRunActorParams { // TODO2 : Change name std::map QueryParameters; std::shared_ptr S3ActorsFactory; ::NFq::NConfig::TWorkloadManagerConfig WorkloadManager; + NYql::IPqGatewayFactory::TPtr PqGatewayFactory; }; } /* NFq */ diff --git a/ydb/core/fq/libs/compute/common/ut/utils_ut.cpp b/ydb/core/fq/libs/compute/common/ut/utils_ut.cpp index 0f1b9e58bd06..35644850ea66 100644 --- a/ydb/core/fq/libs/compute/common/ut/utils_ut.cpp +++ b/ydb/core/fq/libs/compute/common/ut/utils_ut.cpp @@ -63,7 +63,7 @@ Y_UNIT_TEST_SUITE(StatsFormat) { Y_UNIT_TEST(AggregateStat) { auto res = NFq::AggregateStats(NResource::Find("plan.json")); - UNIT_ASSERT_VALUES_EQUAL(res.size(), 14); + UNIT_ASSERT_VALUES_EQUAL(res.size(), 18); UNIT_ASSERT_VALUES_EQUAL(res["IngressBytes"], 6333256); UNIT_ASSERT_VALUES_EQUAL(res["EgressBytes"], 0); UNIT_ASSERT_VALUES_EQUAL(res["InputBytes"], 1044); diff --git a/ydb/core/fq/libs/compute/common/utils.cpp b/ydb/core/fq/libs/compute/common/utils.cpp index 7abd6da7492f..ba2dcde4cc51 100644 --- a/ydb/core/fq/libs/compute/common/utils.cpp +++ b/ydb/core/fq/libs/compute/common/utils.cpp @@ -80,6 +80,10 @@ struct TTotalStatistics { TAggregate EgressRows; TAggregate Tasks; TAggregates Aggregates; + TAggregate IngressFilteredBytes; + TAggregate IngressFilteredRows; + TAggregate IngressQueuedBytes; + TAggregate IngressQueuedRows; }; TString FormatDurationMs(ui64 durationMs) { @@ -309,6 +313,14 @@ void WriteNamedNode(NYson::TYsonWriter& writer, NJson::TJsonValue& node, const T totals.SourceCpuTimeUs.Add(*sum); } else if (name == "Tasks") { totals.Tasks.Add(*sum); + } else if (name == "IngressFilteredBytes") { + totals.IngressFilteredBytes.Add(*sum); + } else if (name == "IngressFilteredRows") { + totals.IngressFilteredRows.Add(*sum); + } else if (name == "IngressQueuedBytes") { + totals.IngressQueuedBytes.Add(*sum); + } else if (name == "IngressQueuedRows") { + totals.IngressQueuedRows.Add(*sum); } } } @@ -468,6 +480,10 @@ TString GetV1StatFromV2Plan(const TString& plan, double* cpuUsage, TString* time totals.IngressRows.Write(writer, "IngressRows"); totals.EgressBytes.Write(writer, "EgressBytes"); totals.EgressRows.Write(writer, "EgressRows"); + totals.IngressFilteredBytes.Write(writer, "IngressFilteredBytes"); + totals.IngressFilteredRows.Write(writer, "IngressFilteredRows"); + totals.IngressQueuedBytes.Write(writer, "IngressQueuedBytes"); + totals.IngressQueuedRows.Write(writer, "IngressQueuedRows"); totals.Tasks.Write(writer, "Tasks"); writer.OnEndMap(); } @@ -532,6 +548,26 @@ struct TStatsAggregator { Aggregates[source + ".Splits"] += ingress->GetIntegerSafe(); success = true; } + if (auto ingress = node.GetValueByPath("Ingress.FilteredBytes.Sum")) { + auto source = name.substr(prefix.size()); + Aggregates[source + ".FilteredBytes"] += ingress->GetIntegerSafe(); + success = true; + } + if (auto ingress = node.GetValueByPath("Ingress.FilteredRows.Sum")) { + auto source = name.substr(prefix.size()); + Aggregates[source + ".FilteredRows"] += ingress->GetIntegerSafe(); + success = true; + } + if (auto ingress = node.GetValueByPath("Ingress.QueuedBytes.Sum")) { + auto source = name.substr(prefix.size()); + Aggregates[source + ".QueuedBytes"] += ingress->GetIntegerSafe(); + success = true; + } + if (auto ingress = node.GetValueByPath("Ingress.QueuedRows.Sum")) { + auto source = name.substr(prefix.size()); + Aggregates[source + ".QueuedRows"] += ingress->GetIntegerSafe(); + success = true; + } return success; } @@ -543,7 +579,11 @@ struct TStatsAggregator { {"EgressRows", 0}, {"InputBytes", 0}, {"OutputBytes", 0}, - {"CpuTimeUs", 0} + {"CpuTimeUs", 0}, + {"IngressFilteredBytes", 0}, + {"IngressFilteredRows", 0}, + {"IngressQueuedBytes", 0}, + {"IngressQueuedRows", 0} }; }; @@ -989,6 +1029,10 @@ TString GetPrettyStatistics(const TString& statistics) { RemapNode(writer, p.second, "TaskRunner.Stage=Total.EgressBytes", "EgressBytes"); RemapNode(writer, p.second, "TaskRunner.Stage=Total.EgressRows", "EgressRows"); RemapNode(writer, p.second, "TaskRunner.Stage=Total.MkqlMaxMemoryUsage", "MaxMemoryUsage"); + RemapNode(writer, p.second, "TaskRunner.Stage=Total.IngressFilteredBytes", "IngressFilteredBytes"); + RemapNode(writer, p.second, "TaskRunner.Stage=Total.IngressFilteredRows", "IngressFilteredRows"); + RemapNode(writer, p.second, "TaskRunner.Stage=Total.IngressQueuedBytes", "IngressQueuedBytes"); + RemapNode(writer, p.second, "TaskRunner.Stage=Total.IngressQueuedRows", "IngressQueuedRows"); writer.OnEndMap(); } // YQv2 @@ -1010,6 +1054,10 @@ TString GetPrettyStatistics(const TString& statistics) { RemapNode(writer, p.second, "EgressBytes", "EgressBytes"); RemapNode(writer, p.second, "EgressRows", "EgressRows"); RemapNode(writer, p.second, "MaxMemoryUsage", "MaxMemoryUsage"); + RemapNode(writer, p.second, "IngressFilteredBytes", "IngressFilteredBytes"); + RemapNode(writer, p.second, "IngressFilteredRows", "IngressFilteredRows"); + RemapNode(writer, p.second, "IngressQueuedBytes", "IngressQueuedBytes"); + RemapNode(writer, p.second, "IngressQueuedRows", "IngressQueuedRows"); writer.OnEndMap(); } } diff --git a/ydb/core/fq/libs/compute/ydb/control_plane/cms_grpc_client_actor.cpp b/ydb/core/fq/libs/compute/ydb/control_plane/cms_grpc_client_actor.cpp index 0f19e61a467a..970f88fe281a 100644 --- a/ydb/core/fq/libs/compute/ydb/control_plane/cms_grpc_client_actor.cpp +++ b/ydb/core/fq/libs/compute/ydb/control_plane/cms_grpc_client_actor.cpp @@ -1,3 +1,5 @@ +#include "ydb_grpc_helpers.h" + #include #include @@ -78,7 +80,7 @@ class TCmsGrpcServiceActor : public NActors::TActor, NGrpc auto forwardRequest = std::make_unique(); forwardRequest->Request.mutable_serverless_resources()->set_shared_database_path(request.BasePath); forwardRequest->Request.set_path(request.Path); - forwardRequest->Token = CredentialsProvider->GetAuthInfo(); + SetYdbRequestToken(*forwardRequest, CredentialsProvider->GetAuthInfo()); TEvPrivate::TEvCreateDatabaseRequest::TPtr forwardEvent = (NActors::TEventHandle*)new IEventHandle(SelfId(), SelfId(), forwardRequest.release(), 0, Cookie); MakeCall(std::move(forwardEvent)); Requests[Cookie++] = ev; @@ -119,7 +121,7 @@ class TCmsGrpcServiceActor : public NActors::TActor, NGrpc void Handle(TEvYdbCompute::TEvListDatabasesRequest::TPtr& ev) { auto forwardRequest = std::make_unique(); - forwardRequest->Token = CredentialsProvider->GetAuthInfo(); + SetYdbRequestToken(*forwardRequest, CredentialsProvider->GetAuthInfo()); TEvPrivate::TEvListDatabasesRequest::TPtr forwardEvent = (NActors::TEventHandle*)new IEventHandle(SelfId(), SelfId(), forwardRequest.release(), 0, Cookie); MakeCall(std::move(forwardEvent)); Requests[Cookie++] = ev; diff --git a/ydb/core/fq/libs/compute/ydb/control_plane/compute_database_control_plane_service.cpp b/ydb/core/fq/libs/compute/ydb/control_plane/compute_database_control_plane_service.cpp index b1e77602de93..e6cb95298d41 100644 --- a/ydb/core/fq/libs/compute/ydb/control_plane/compute_database_control_plane_service.cpp +++ b/ydb/core/fq/libs/compute/ydb/control_plane/compute_database_control_plane_service.cpp @@ -11,7 +11,6 @@ #include #include -#include #include #include @@ -309,18 +308,6 @@ class TCreateDatabaseRequestActor : public NActors::TActorBootstrapped { @@ -355,7 +342,7 @@ class TComputeDatabaseControlPlaneServiceActor : public NActors::TActorBootstrap switch (controlPlane.type_case()) { case NConfig::TYdbComputeControlPlane::TYPE_NOT_SET: case NConfig::TYdbComputeControlPlane::kSingle: - CreateSingleClientActors(); + CreateSingleClientActors(controlPlane.GetSingle()); break; case NConfig::TYdbComputeControlPlane::kCms: CreateCmsClientActors(controlPlane.GetCms(), controlPlane.GetDatabasesCacheReloadPeriod()); @@ -374,13 +361,11 @@ class TComputeDatabaseControlPlaneServiceActor : public NActors::TActorBootstrap if (connection.GetCertificateFile()) { settings.CertificateRootCA = StripString(TFileInput(connection.GetCertificateFile()).ReadAll()); } - settings.Headers[NYdb::YDB_DATABASE_HEADER] = connection.GetDatabase(); return settings; } - static NGrpcActorClient::TGrpcClientSettings CreateGrpcClientSettings(const NConfig::TComputeDatabaseConfig& config) { + static NGrpcActorClient::TGrpcClientSettings CreateGrpcClientSettings(const auto& connection) { NGrpcActorClient::TGrpcClientSettings settings; - const auto& connection = config.GetExecutionConnection(); settings.Endpoint = connection.GetEndpoint(); settings.EnableSsl = connection.GetUseSsl(); if (connection.GetCertificateFile()) { @@ -390,21 +375,23 @@ class TComputeDatabaseControlPlaneServiceActor : public NActors::TActorBootstrap return settings; } - void CreateSingleClientActors() { + static NGrpcActorClient::TGrpcClientSettings CreateGrpcClientSettings(const NConfig::TComputeDatabaseConfig& config) { + return CreateGrpcClientSettings(config.GetControlPlaneConnection()); + } + + void CreateSingleClientActors(const NConfig::TYdbComputeControlPlane::TSingle& singleConfig) { auto globalLoadConfig = Config.GetYdb().GetLoadControlConfig(); - if (!globalLoadConfig.GetEnable() || !IsValidLoadControlConfig(globalLoadConfig)) { - return; - } - TActorId clientActor; - auto monitoringEndpoint = globalLoadConfig.GetMonitoringEndpoint(); - const auto& databaseConnection = globalLoadConfig.GetDatabaseConnection(); - auto credentialsProvider = CredentialsProviderFactory(GetYdbCredentialSettings(databaseConnection))->CreateProvider(); - if (monitoringEndpoint) { - clientActor = Register(CreateMonitoringRestClientActor(monitoringEndpoint, databaseConnection.GetDatabase(), credentialsProvider).release()); - } else { - clientActor = Register(CreateMonitoringGrpcClientActor(CreateGrpcClientSettings(databaseConnection), credentialsProvider).release()); + if (globalLoadConfig.GetEnable()) { + TActorId clientActor; + auto monitoringEndpoint = globalLoadConfig.GetMonitoringEndpoint(); + auto credentialsProvider = CredentialsProviderFactory(GetYdbCredentialSettings(singleConfig.GetConnection()))->CreateProvider(); + if (monitoringEndpoint) { + clientActor = Register(CreateMonitoringRestClientActor(monitoringEndpoint, singleConfig.GetConnection().GetDatabase(), credentialsProvider).release()); + } else { + clientActor = Register(CreateMonitoringGrpcClientActor(CreateGrpcClientSettings(singleConfig.GetConnection()), credentialsProvider).release()); + } + MonitoringActorId = Register(CreateDatabaseMonitoringActor(clientActor, globalLoadConfig, Counters).release()); } - MonitoringActorId = Register(CreateDatabaseMonitoringActor(clientActor, globalLoadConfig, Counters).release()); } void CreateCmsClientActors(const NConfig::TYdbComputeControlPlane::TCms& cmsConfig, const TString& databasesCacheReloadPeriod) { @@ -418,15 +405,14 @@ class TComputeDatabaseControlPlaneServiceActor : public NActors::TActorBootstrap const NConfig::TLoadControlConfig& loadConfig = config.GetLoadControlConfig().GetEnable() ? config.GetLoadControlConfig() : globalLoadConfig; - if (loadConfig.GetEnable() && IsValidLoadControlConfig(loadConfig)) { + if (loadConfig.GetEnable()) { TActorId clientActor; auto monitoringEndpoint = loadConfig.GetMonitoringEndpoint(); - const auto& databaseConnection = loadConfig.GetDatabaseConnection(); - auto credentialsProvider = CredentialsProviderFactory(GetYdbCredentialSettings(databaseConnection))->CreateProvider(); + auto credentialsProvider = CredentialsProviderFactory(GetYdbCredentialSettings(config.GetControlPlaneConnection()))->CreateProvider(); if (monitoringEndpoint) { - clientActor = Register(CreateMonitoringRestClientActor(monitoringEndpoint, databaseConnection.GetDatabase(), credentialsProvider).release()); + clientActor = Register(CreateMonitoringRestClientActor(monitoringEndpoint, config.GetControlPlaneConnection().GetDatabase(), credentialsProvider).release()); } else { - clientActor = Register(CreateMonitoringGrpcClientActor(CreateGrpcClientSettings(databaseConnection), credentialsProvider).release()); + clientActor = Register(CreateMonitoringGrpcClientActor(CreateGrpcClientSettings(config), credentialsProvider).release()); } databaseMonitoringActor = Register(CreateDatabaseMonitoringActor(clientActor, loadConfig, databaseCounters).release()); } @@ -443,15 +429,14 @@ class TComputeDatabaseControlPlaneServiceActor : public NActors::TActorBootstrap const NConfig::TLoadControlConfig& loadConfig = config.GetLoadControlConfig().GetEnable() ? config.GetLoadControlConfig() : globalLoadConfig; - if (loadConfig.GetEnable() && IsValidLoadControlConfig(loadConfig)) { + if (loadConfig.GetEnable()) { TActorId clientActor; auto monitoringEndpoint = loadConfig.GetMonitoringEndpoint(); - const auto& databaseConnection = loadConfig.GetDatabaseConnection(); - auto credentialsProvider = CredentialsProviderFactory(GetYdbCredentialSettings(databaseConnection))->CreateProvider(); + auto credentialsProvider = CredentialsProviderFactory(GetYdbCredentialSettings(config.GetControlPlaneConnection()))->CreateProvider(); if (monitoringEndpoint) { - clientActor = Register(CreateMonitoringRestClientActor(monitoringEndpoint, databaseConnection.GetDatabase(), credentialsProvider).release()); + clientActor = Register(CreateMonitoringRestClientActor(monitoringEndpoint, config.GetControlPlaneConnection().GetDatabase(), credentialsProvider).release()); } else { - clientActor = Register(CreateMonitoringGrpcClientActor(CreateGrpcClientSettings(databaseConnection), credentialsProvider).release()); + clientActor = Register(CreateMonitoringGrpcClientActor(CreateGrpcClientSettings(config), credentialsProvider).release()); } databaseMonitoringActor = Register(CreateDatabaseMonitoringActor(clientActor, loadConfig, databaseCounters).release()); } diff --git a/ydb/core/fq/libs/compute/ydb/control_plane/compute_databases_cache.cpp b/ydb/core/fq/libs/compute/ydb/control_plane/compute_databases_cache.cpp index 2b760d0d52cd..5ca8774c8c20 100644 --- a/ydb/core/fq/libs/compute/ydb/control_plane/compute_databases_cache.cpp +++ b/ydb/core/fq/libs/compute/ydb/control_plane/compute_databases_cache.cpp @@ -55,7 +55,7 @@ class TComputeDatabasesCacheActor : public NActors::TActorBootstrapped #include @@ -64,7 +66,7 @@ class TMonitoringGrpcServiceActor : public NActors::TActor(); forwardRequest->Request.set_return_verbose_status(true); - forwardRequest->Token = CredentialsProvider->GetAuthInfo(); + SetYdbRequestToken(*forwardRequest, CredentialsProvider->GetAuthInfo()); TEvPrivate::TEvSelfCheckRequest::TPtr forwardEvent = (NActors::TEventHandle*)new IEventHandle(SelfId(), SelfId(), forwardRequest.release(), 0, Cookie); MakeCall(std::move(forwardEvent)); Requests[Cookie++] = ev; diff --git a/ydb/core/fq/libs/compute/ydb/control_plane/ya.make b/ydb/core/fq/libs/compute/ydb/control_plane/ya.make index 26fef7d2d404..45a0ee2cca50 100644 --- a/ydb/core/fq/libs/compute/ydb/control_plane/ya.make +++ b/ydb/core/fq/libs/compute/ydb/control_plane/ya.make @@ -21,15 +21,13 @@ PEERDIR( ydb/library/actors/protos ydb/library/db_pool/protos ydb/library/grpc/actor_client - yql/essentials/public/issue - yql/essentials/utils ydb/library/yql/utils/actors ydb/public/api/grpc ydb/public/api/grpc/draft ydb/public/lib/operation_id/protos ydb/public/sdk/cpp/client/resources - ydb/public/sdk/cpp/client/ydb_operation - ydb/public/sdk/cpp/client/ydb_query + yql/essentials/public/issue + yql/essentials/utils ) YQL_LAST_ABI_VERSION() diff --git a/ydb/core/fq/libs/compute/ydb/control_plane/ydb_grpc_helpers.h b/ydb/core/fq/libs/compute/ydb/control_plane/ydb_grpc_helpers.h new file mode 100644 index 000000000000..8dba7e7c2185 --- /dev/null +++ b/ydb/core/fq/libs/compute/ydb/control_plane/ydb_grpc_helpers.h @@ -0,0 +1,16 @@ +#pragma once + +#include +#include + +namespace NFq { + +template +void SetYdbRequestToken(NCloud::TEvGrpcProtoRequest& event, const TString& token) { + if (token) { + event.Token = token; + event.Headers.emplace(NYdb::YDB_AUTH_TICKET_HEADER, token); + } +} + +} // namespace NFq diff --git a/ydb/core/fq/libs/compute/ydb/synchronization_service/synchronization_service.cpp b/ydb/core/fq/libs/compute/ydb/synchronization_service/synchronization_service.cpp index fe8df83ebd6c..fef5777b69fb 100644 --- a/ydb/core/fq/libs/compute/ydb/synchronization_service/synchronization_service.cpp +++ b/ydb/core/fq/libs/compute/ydb/synchronization_service/synchronization_service.cpp @@ -627,13 +627,15 @@ class TSynchronizeScopeActor : public NActors::TActorBootstrappedSend(self, new TEvYdbCompute::TEvCreateResourcePoolResponse(ExtractStatus(future)), 0, i); @@ -663,13 +665,16 @@ class TSynchronizeScopeActor : public NActors::TActorBootstrappedSend(self, new TEvYdbCompute::TEvCreateResourcePoolResponse(ExtractStatus(future)), 0, index); diff --git a/ydb/core/fq/libs/config/protos/common.proto b/ydb/core/fq/libs/config/protos/common.proto index 4d89d11a35aa..b4b7a974b7c6 100644 --- a/ydb/core/fq/libs/config/protos/common.proto +++ b/ydb/core/fq/libs/config/protos/common.proto @@ -32,4 +32,6 @@ message TCommonConfig { bool ShowQueryTimeline = 16; uint64 MaxQueryTimelineSize = 17; // default: 200KB string PqReconnectPeriod = 18; // default: disabled + uint32 TopicClientHandlersExecutorThreadsNum = 19; // default: 0 that means use default from TopicClientSettings (1) + uint32 TopicClientCompressionExecutorThreadsNum = 20; // default: 0 that means use default from TopicClientSettings (2) } diff --git a/ydb/core/fq/libs/config/protos/compute.proto b/ydb/core/fq/libs/config/protos/compute.proto index 38ca12b987dd..172a6092b1cc 100644 --- a/ydb/core/fq/libs/config/protos/compute.proto +++ b/ydb/core/fq/libs/config/protos/compute.proto @@ -27,7 +27,6 @@ message TLoadControlConfig { bool Strict = 7; // default false, whether to deny execution in load level unavailable uint32 CpuNumber = 8; string MonitoringEndpoint = 9; // if defined, will be used as REST API instead of default GRPC - TYdbStorageConfig DatabaseConnection = 10; } message TWorkloadManagerConfig { @@ -36,6 +35,7 @@ message TWorkloadManagerConfig { int32 ConcurrentQueryLimit = 2; int32 QueueSize = 3; double DatabaseLoadCpuThreshold = 4; + double TotalCpuLimitPercentPerNode = 5; } bool Enable = 1; @@ -43,6 +43,10 @@ message TWorkloadManagerConfig { repeated TResourcePool ResourcePool = 3; } +message TAccessConfig { + repeated string ExternalSourcesAccessSID = 1; +} + message TComputeDatabaseConfig { string Id = 7; TYdbStorageConfig ControlPlaneConnection = 1; @@ -51,6 +55,7 @@ message TComputeDatabaseConfig { string Tenant = 2; TLoadControlConfig LoadControlConfig = 4; TWorkloadManagerConfig WorkloadManagerConfig = 5; + TAccessConfig AccessConfig = 8; } message TDatabaseMapping { diff --git a/ydb/core/fq/libs/config/protos/row_dispatcher.proto b/ydb/core/fq/libs/config/protos/row_dispatcher.proto index 2ee16c93ddef..e05dd9b82edc 100644 --- a/ydb/core/fq/libs/config/protos/row_dispatcher.proto +++ b/ydb/core/fq/libs/config/protos/row_dispatcher.proto @@ -25,6 +25,10 @@ message TJsonParserConfig { uint64 BufferCellCount = 3; // (number rows) * (number columns) limit, default 10^6 } +message TCompileServiceConfig { + uint64 ParallelCompilationLimit = 1; // 1 by default +} + message TRowDispatcherConfig { bool Enabled = 1; uint64 TimeoutBeforeStartSessionSec = 2; @@ -32,5 +36,6 @@ message TRowDispatcherConfig { uint64 MaxSessionUsedMemory = 4; bool WithoutConsumer = 5; TJsonParserConfig JsonParser = 7; + TCompileServiceConfig CompileService = 8; TRowDispatcherCoordinatorConfig Coordinator = 6; } diff --git a/ydb/core/fq/libs/control_plane_proxy/actors/counters.h b/ydb/core/fq/libs/control_plane_proxy/actors/counters.h index 32dec96a725e..cb73258c1121 100644 --- a/ydb/core/fq/libs/control_plane_proxy/actors/counters.h +++ b/ydb/core/fq/libs/control_plane_proxy/actors/counters.h @@ -171,6 +171,7 @@ enum ERequestTypeCommon { RTC_DELETE_BINDING_IN_YDB, RTC_CREATE_COMPUTE_DATABASE, RTC_LIST_CPS_ENTITY, + RTC_RATE_LIMITER, RTC_MAX, }; @@ -228,6 +229,7 @@ class TCounters : public virtual TThrRefBase { {MakeIntrusive("DeleteBindingInYDB")}, {MakeIntrusive("CreateComputeDatabase")}, {MakeIntrusive("ListCPSEntities")}, + {MakeIntrusive("RateLimiter")}, }); TTtlCache ScopeCounters{TTtlCacheSettings{}.SetTtl(TDuration::Days(1))}; diff --git a/ydb/core/fq/libs/control_plane_proxy/actors/query_utils.cpp b/ydb/core/fq/libs/control_plane_proxy/actors/query_utils.cpp index f842ccf6a363..db162d81f5fa 100644 --- a/ydb/core/fq/libs/control_plane_proxy/actors/query_utils.cpp +++ b/ydb/core/fq/libs/control_plane_proxy/actors/query_utils.cpp @@ -17,6 +17,61 @@ TString MakeSecretKeyName(const TString& prefix, const TString& folderId, const return TStringBuilder{} << prefix << "_" << folderId << "_" << name; } +TString MakeCreateSecretObjectSql(const TString& name, const TString& value) { + using namespace fmt::literals; + return fmt::format( + R"( + UPSERT OBJECT {name} (TYPE SECRET) WITH value={value}; + )", + "name"_a = EncloseAndEscapeString(name, '`'), + "value"_a = EncloseSecret(EncloseAndEscapeString(value, '"'))); +} + +TString MakeCreateSecretAccessObjectSql(const TString& name, const TString& sid) { + using namespace fmt::literals; + TString accessName = TStringBuilder{} << name << ":" << sid; + return fmt::format( + R"( + UPSERT OBJECT {name} (TYPE SECRET_ACCESS); + )", + "name"_a = EncloseAndEscapeString(accessName, '`')); +} + +TString MakeCreateSecretAccessObjectsSql(const TString& name, const TVector& externalSourcesAccessSIDs) { + TStringBuilder result; + for (const auto& sid : externalSourcesAccessSIDs) { + result << MakeCreateSecretAccessObjectSql(name, sid); + } + return result; +} + +TString MakeDropSecretObjectSql(const TString& name) { + using namespace fmt::literals; + return fmt::format( + R"( + DROP OBJECT {name} (TYPE SECRET); + )", + "name"_a = EncloseAndEscapeString(name, '`')); +} + +TString MakeDropAccessSecretObjectSql(const TString& name, const TString& sid) { + using namespace fmt::literals; + TString accessName = TStringBuilder{} << name << ":" << sid; + return fmt::format( + R"( + DROP OBJECT {name} (TYPE SECRET_ACCESS); + )", + "name"_a = EncloseAndEscapeString(accessName, '`')); +} + +TString MakeDropSecretAccessObjectsSql(const TString& name, const TVector& externalSourcesAccessSIDs) { + TStringBuilder result; + for (const auto& sid : externalSourcesAccessSIDs) { + result << MakeDropAccessSecretObjectSql(name, sid); + } + return result; +} + } TString MakeCreateExternalDataTableQuery(const FederatedQuery::BindingContent& content, @@ -100,33 +155,30 @@ TString SignAccountId(const TString& id, const TSigner::TPtr& signer) { return signer ? signer->SignAccountId(id) : TString{}; } + + TMaybe CreateSecretObjectQuery(const FederatedQuery::ConnectionSetting& setting, - const TString& name, - const TSigner::TPtr& signer, - const TString& folderId) { - using namespace fmt::literals; - TString secretObjects; + const TString& name, + const TSigner::TPtr& signer, + const TString& folderId, + const TVector& externalSourcesAccessSIDs) { + TStringBuilder result; auto serviceAccountId = ExtractServiceAccountId(setting); - if (serviceAccountId) { - secretObjects = signer ? fmt::format( - R"( - UPSERT OBJECT {sa_secret_name} (TYPE SECRET) WITH value={signature}; - )", - "sa_secret_name"_a = EncloseAndEscapeString(MakeSecretKeyName("f1", folderId, name), '`'), - "signature"_a = EncloseSecret(EncloseAndEscapeString(SignAccountId(serviceAccountId, signer), '"'))) : std::string{}; + if (serviceAccountId && signer) { + const TString saSecretName = MakeSecretKeyName("f1", folderId, name); + const TString signature = SignAccountId(serviceAccountId, signer); + result << MakeCreateSecretObjectSql(saSecretName, signature); + result << MakeCreateSecretAccessObjectsSql(saSecretName, externalSourcesAccessSIDs); } auto password = GetPassword(setting); if (password) { - secretObjects += fmt::format( - R"( - UPSERT OBJECT {password_secret_name} (TYPE SECRET) WITH value={password}; - )", - "password_secret_name"_a = EncloseAndEscapeString(MakeSecretKeyName("f2", folderId, name), '`'), - "password"_a = EncloseSecret(EncloseAndEscapeString(*password, '"'))); + const TString passwordSecretName = MakeSecretKeyName("f2", folderId, name); + result << MakeCreateSecretObjectSql(passwordSecretName, *password); + result << MakeCreateSecretAccessObjectsSql(passwordSecretName, externalSourcesAccessSIDs); } - return secretObjects ? secretObjects : TMaybe{}; + return result ? result : TMaybe{}; } TString CreateAuthParamsQuery(const FederatedQuery::ConnectionSetting& setting, @@ -301,21 +353,14 @@ TString MakeCreateExternalDataSourceQuery( folderId)); } -TMaybe DropSecretObjectQuery(const TString& name, const TString& folderId) { - using namespace fmt::literals; - return fmt::format( - R"( - DROP OBJECT {secret_name1} (TYPE SECRET); - DROP OBJECT {secret_name2} (TYPE SECRET); - DROP OBJECT {secret_name3} (TYPE SECRET); -- for backward compatibility - DROP OBJECT {secret_name4} (TYPE SECRET); -- for backward compatibility - DROP OBJECT {secret_name5} (TYPE SECRET); -- for backward compatibility - )", - "secret_name1"_a = EncloseAndEscapeString(MakeSecretKeyName("f1", folderId, name), '`'), - "secret_name2"_a = EncloseAndEscapeString(MakeSecretKeyName("f2", folderId, name), '`'), - "secret_name3"_a = EncloseAndEscapeString(TStringBuilder{} << "k1" << name, '`'), - "secret_name4"_a = EncloseAndEscapeString(TStringBuilder{} << "k2" << name, '`'), - "secret_name5"_a = EncloseAndEscapeString(name, '`')); +TMaybe DropSecretObjectQuery(const TString& name, const TString& folderId, const TVector& externalSourcesAccessSIDs) { + const TString secretName1 = MakeSecretKeyName("f1", folderId, name); + const TString secretName2 = MakeSecretKeyName("f2", folderId, name); + return TStringBuilder{} + << MakeDropSecretAccessObjectsSql(secretName1, externalSourcesAccessSIDs) + << MakeDropSecretObjectSql(secretName1) + << MakeDropSecretAccessObjectsSql(secretName2, externalSourcesAccessSIDs) + << MakeDropSecretObjectSql(secretName2); } TString MakeDeleteExternalDataTableQuery(const TString& tableName) { diff --git a/ydb/core/fq/libs/control_plane_proxy/actors/query_utils.h b/ydb/core/fq/libs/control_plane_proxy/actors/query_utils.h index 92ed74341c43..5d8e86df1389 100644 --- a/ydb/core/fq/libs/control_plane_proxy/actors/query_utils.h +++ b/ydb/core/fq/libs/control_plane_proxy/actors/query_utils.h @@ -11,9 +11,10 @@ namespace NPrivate { TMaybe CreateSecretObjectQuery(const FederatedQuery::ConnectionSetting& setting, const TString& name, const TSigner::TPtr& signer, - const TString& folderId); + const TString& folderId, + const TVector& externalSourcesAccessSIDs); -TMaybe DropSecretObjectQuery(const TString& name, const TString& folderId); +TMaybe DropSecretObjectQuery(const TString& name, const TString& folderId, const TVector& externalSourcesAccessSIDs); TString MakeCreateExternalDataSourceQuery( const FederatedQuery::ConnectionContent& connectionContent, diff --git a/ydb/core/fq/libs/control_plane_proxy/actors/request_actor.h b/ydb/core/fq/libs/control_plane_proxy/actors/request_actor.h index 2ceab95ae510..6925f9e217ad 100644 --- a/ydb/core/fq/libs/control_plane_proxy/actors/request_actor.h +++ b/ydb/core/fq/libs/control_plane_proxy/actors/request_actor.h @@ -192,11 +192,23 @@ class TCreateQueryRequestActor : TEvControlPlaneStorage::TEvCreateQueryResponse, TEvControlPlaneProxy::TEvCreateQueryRequest, TEvControlPlaneProxy::TEvCreateQueryResponse>; - using TBaseRequestActor::TBaseRequestActor; + + TCreateQueryRequestActor(typename TEvControlPlaneProxy::TEvCreateQueryRequest::TPtr requestProxy, + const TControlPlaneProxyConfig& config, + const TActorId& serviceId, + const TRequestCounters& counters, + const TRequestCommonCountersPtr& rateLimiterCounters, + const std::function& probe, + const TPermissions& availablePermissions, + bool replyWithResponseOnSuccess = true) + : TBaseRequestActor(requestProxy, config, serviceId, counters, probe, availablePermissions, replyWithResponseOnSuccess) + , RateLimiterCounters(rateLimiterCounters) { + } STFUNC(StateFunc) { switch (ev->GetTypeRewrite()) { hFunc(TEvRateLimiter::TEvCreateResourceResponse, Handle); + cFunc(NActors::TEvents::TSystem::Wakeup, HandleTimeout); default: return TBaseRequestActor::StateFunc(ev); } @@ -208,6 +220,16 @@ class TCreateQueryRequestActor : || !Config.ComputeConfig.YdbComputeControlPlaneEnabled(RequestProxy->Get()->Scope)); } + void HandleTimeout() { + // Don't need to set the RateLimiterCreationInProgress = false + // because of the PassAway will be called in this callback + if (RateLimiterCreationInProgress) { + RateLimiterCounters->Timeout->Inc(); + RateLimiterCounters->InFly->Dec(); + } + TBaseRequestActor::HandleTimeout(); + } + void OnBootstrap() override { this->UnsafeBecome(&TCreateQueryRequestActor::StateFunc); if (ShouldCreateRateLimiter()) { @@ -223,9 +245,13 @@ class TCreateQueryRequestActor : 10); // percent -> milliseconds CPP_LOG_T("Create rate limiter resource for cloud with limit " << cloudLimit << "ms"); + RateLimiterCreationInProgress = true; + RateLimiterCounters->InFly->Inc(); + StartRateLimiterCreation = TInstant::Now(); Send(RateLimiterControlPlaneServiceId(), new TEvRateLimiter::TEvCreateResource(RequestProxy->Get()->CloudId, cloudLimit)); } else { + RateLimiterCounters->Error->Inc(); NYql::TIssues issues; NYql::TIssue issue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, @@ -238,12 +264,17 @@ class TCreateQueryRequestActor : } void Handle(TEvRateLimiter::TEvCreateResourceResponse::TPtr& ev) { + RateLimiterCreationInProgress = false; + RateLimiterCounters->InFly->Dec(); + RateLimiterCounters->LatencyMs->Collect((TInstant::Now() - StartRateLimiterCreation).MilliSeconds()); CPP_LOG_D( "Create response from rate limiter service. Success: " << ev->Get()->Success); if (ev->Get()->Success) { + RateLimiterCounters->Ok->Inc(); QuoterResourceCreated = true; SendRequestIfCan(); } else { + RateLimiterCounters->Error->Inc(); NYql::TIssue issue("Failed to create rate limiter resource"); for (const NYql::TIssue& i : ev->Get()->Issues) { issue.AddSubIssue(MakeIntrusive(i)); @@ -257,6 +288,11 @@ class TCreateQueryRequestActor : bool CanSendRequest() const override { return (QuoterResourceCreated || !ShouldCreateRateLimiter()) && TBaseRequestActor::CanSendRequest(); } + +private: + TInstant StartRateLimiterCreation; + bool RateLimiterCreationInProgress = false; + TRequestCommonCountersPtr RateLimiterCounters; }; } // namespace NFq::NPrivate diff --git a/ydb/core/fq/libs/control_plane_proxy/actors/ydb_schema_query_actor.cpp b/ydb/core/fq/libs/control_plane_proxy/actors/ydb_schema_query_actor.cpp index 26eae1ef7670..e62849f7daea 100644 --- a/ydb/core/fq/libs/control_plane_proxy/actors/ydb_schema_query_actor.cpp +++ b/ydb/core/fq/libs/control_plane_proxy/actors/ydb_schema_query_actor.cpp @@ -564,7 +564,8 @@ IActor* MakeCreateConnectionActor( auto createSecretStatement = CreateSecretObjectQuery(connectionContent.setting(), connectionContent.name(), signer, - folderId); + folderId, + computeConfig.GetExternalSourcesAccessSIDs(scope)); std::vector statements; if (createSecretStatement) { @@ -665,14 +666,16 @@ IActor* MakeModifyConnectionActor( auto& newConnectionContent = request->Get()->Request.content(); const auto& scope = request->Get()->Scope; const TString folderId = NYdb::NFq::TScope{scope}.ParseFolder(); + const auto externalSourcesAccessSIDs = computeConfig.GetExternalSourcesAccessSIDs(scope); auto dropOldSecret = - DropSecretObjectQuery(oldConnectionContent.name(), folderId); + DropSecretObjectQuery(oldConnectionContent.name(), folderId, externalSourcesAccessSIDs); auto createNewSecret = CreateSecretObjectQuery(newConnectionContent.setting(), newConnectionContent.name(), signer, - folderId); + folderId, + externalSourcesAccessSIDs); bool replaceSupported = computeConfig.IsReplaceIfExistsSyntaxSupported(); if (replaceSupported && @@ -680,7 +683,8 @@ IActor* MakeModifyConnectionActor( // CREATE OR REPLACE auto createSecretStatement = CreateSecretObjectQuery(newConnectionContent.setting(), - newConnectionContent.name(), signer, folderId); + newConnectionContent.name(), signer, folderId, + externalSourcesAccessSIDs); std::vector statements; if (createSecretStatement) { @@ -727,13 +731,14 @@ IActor* MakeModifyConnectionActor( .SQL = *dropOldSecret, .RollbackSQL = CreateSecretObjectQuery(oldConnectionContent.setting(), oldConnectionContent.name(), - signer, folderId), + signer, folderId, + externalSourcesAccessSIDs), .ShouldSkipStepOnError = IsPathDoesNotExistIssue}); } if (createNewSecret) { statements.push_back(TSchemaQueryTask{.SQL = *createNewSecret, .RollbackSQL = DropSecretObjectQuery( - newConnectionContent.name(), folderId)}); + newConnectionContent.name(), folderId, externalSourcesAccessSIDs)}); } statements.push_back( @@ -787,18 +792,20 @@ IActor* MakeDeleteConnectionActor( TDuration requestTimeout, TCounters& counters, const TCommonConfig& commonConfig, + const ::NFq::TComputeConfig& computeConfig, TSigner::TPtr signer) { auto queryFactoryMethod = [signer = std::move(signer), - commonConfig]( + commonConfig, computeConfig]( const TEvControlPlaneProxy::TEvDeleteConnectionRequest::TPtr& request) -> std::vector { auto& connectionContent = *request->Get()->ConnectionContent; const auto& scope = request->Get()->Scope; const TString folderId = NYdb::NFq::TScope{scope}.ParseFolder(); + const auto externalSourcesAccessSIDs = computeConfig.GetExternalSourcesAccessSIDs(scope); auto dropSecret = - DropSecretObjectQuery(connectionContent.name(), folderId); + DropSecretObjectQuery(connectionContent.name(), folderId, externalSourcesAccessSIDs); std::vector statements = { TSchemaQueryTask{.SQL = TString{MakeDeleteExternalDataSourceQuery( @@ -812,7 +819,8 @@ IActor* MakeDeleteConnectionActor( .RollbackSQL = CreateSecretObjectQuery(connectionContent.setting(), connectionContent.name(), - signer, folderId), + signer, folderId, + externalSourcesAccessSIDs), .ShouldSkipStepOnError = IsPathDoesNotExistIssue}); } return statements; diff --git a/ydb/core/fq/libs/control_plane_proxy/actors/ydb_schema_query_actor.h b/ydb/core/fq/libs/control_plane_proxy/actors/ydb_schema_query_actor.h index 2be77b844048..260c8ecaaf96 100644 --- a/ydb/core/fq/libs/control_plane_proxy/actors/ydb_schema_query_actor.h +++ b/ydb/core/fq/libs/control_plane_proxy/actors/ydb_schema_query_actor.h @@ -45,6 +45,7 @@ NActors::IActor* MakeDeleteConnectionActor( TDuration requestTimeout, TCounters& counters, const NConfig::TCommonConfig& commonConfig, + const TComputeConfig& computeConfig, TSigner::TPtr signer); /// Binding manipulation actors diff --git a/ydb/core/fq/libs/control_plane_proxy/control_plane_proxy.cpp b/ydb/core/fq/libs/control_plane_proxy/control_plane_proxy.cpp index 971e77048ed9..e79fee68adaa 100644 --- a/ydb/core/fq/libs/control_plane_proxy/control_plane_proxy.cpp +++ b/ydb/core/fq/libs/control_plane_proxy/control_plane_proxy.cpp @@ -402,6 +402,7 @@ class TCreateComputeDatabaseActor : public NActors::TActorBootstrapped Probe; TEventRequest Event; ui32 Cookie; + FederatedQuery::QueryContent::QueryType QueryType; TInstant StartTime; public: @@ -413,7 +414,8 @@ class TCreateComputeDatabaseActor : public NActors::TActorBootstrapped& probe, TEventRequest event, - ui32 cookie) + ui32 cookie, + FederatedQuery::QueryContent::QueryType queryType = FederatedQuery::QueryContent::QUERY_TYPE_UNSPECIFIED) : Config(config) , ComputeConfig(computeConfig) , Sender(sender) @@ -423,13 +425,14 @@ class TCreateComputeDatabaseActor : public NActors::TActorBootstrappedGet()->ComputeDatabase = FederatedQuery::Internal::ComputeDatabaseInternal{}; TActivationContext::Send(Event->Forward(ControlPlaneProxyActorId())); PassAway(); @@ -649,6 +652,7 @@ class TControlPlaneProxyActor : public NActors::TActorBootstrappedSender; ui64 cookie = ev->Cookie; + FederatedQuery::QueryContent::QueryType queryType = request.content().type(); auto probe = [=](const TDuration& delta, bool isSuccess, bool isTimeout) { LWPROBE(CreateQueryRequest, scope, user, delta, byteSize, isSuccess, isTimeout); @@ -689,7 +693,7 @@ class TControlPlaneProxyActor : public NActors::TActorBootstrapped (Counters.GetCommonCounters(RTC_CREATE_COMPUTE_DATABASE), sender, Config, Config.ComputeConfig, cloudId, - scope, probe, ev, cookie)); + scope, probe, ev, cookie, queryType)); return; } @@ -702,6 +706,7 @@ class TControlPlaneProxyActor : public NActors::TActorBootstrappedSender; ui64 cookie = ev->Cookie; + FederatedQuery::QueryContent::QueryType queryType = request.content().type(); auto probe = [=](const TDuration& delta, bool isSuccess, bool isTimeout) { LWPROBE(ModifyQueryRequest, scope, user, queryId, delta, byteSize, isSuccess, isTimeout); @@ -958,7 +964,7 @@ class TControlPlaneProxyActor : public NActors::TActorBootstrapped (Counters.GetCommonCounters(RTC_CREATE_COMPUTE_DATABASE), sender, Config, Config.ComputeConfig, cloudId, - scope, probe, ev, cookie)); + scope, probe, ev, cookie, queryType)); return; } @@ -1866,6 +1872,7 @@ class TControlPlaneProxyActor : public NActors::TActorBootstrapped { + using TProto = Fq::Private::WriteTaskResultRequest; TEvWriteResultDataRequest() = default; @@ -411,6 +412,8 @@ struct TEvControlPlaneStorage { struct TEvWriteResultDataResponse : NActors::TEventLocal { static constexpr bool Auditable = false; + using TProto = Fq::Private::WriteTaskResultResult; + explicit TEvWriteResultDataResponse( const Fq::Private::WriteTaskResultResult& record) : Record(record) @@ -434,6 +437,7 @@ struct TEvControlPlaneStorage { }; struct TEvGetTaskRequest : NActors::TEventLocal { + using TProto = Fq::Private::GetTaskRequest; TEvGetTaskRequest() = default; @@ -454,6 +458,8 @@ struct TEvControlPlaneStorage { struct TEvGetTaskResponse : NActors::TEventLocal { static constexpr bool Auditable = false; + using TProto = Fq::Private::GetTaskResult; + explicit TEvGetTaskResponse( const Fq::Private::GetTaskResult& record) : Record(record) @@ -499,6 +505,7 @@ struct TEvControlPlaneStorage { }; struct TEvPingTaskRequest : NActors::TEventLocal { + using TProto = Fq::Private::PingTaskRequest; TEvPingTaskRequest() = default; @@ -519,6 +526,8 @@ struct TEvControlPlaneStorage { struct TEvPingTaskResponse : NActors::TEventLocal { static constexpr bool Auditable = false; + using TProto = Fq::Private::PingTaskResult; + explicit TEvPingTaskResponse( const Fq::Private::PingTaskResult& record) : Record(record) @@ -542,6 +551,7 @@ struct TEvControlPlaneStorage { }; struct TEvNodesHealthCheckRequest : NActors::TEventLocal { + using TProto = Fq::Private::NodesHealthCheckRequest; TEvNodesHealthCheckRequest() = default; @@ -561,6 +571,8 @@ struct TEvControlPlaneStorage { struct TEvNodesHealthCheckResponse : NActors::TEventLocal { static constexpr bool Auditable = false; + using TProto = Fq::Private::NodesHealthCheckResult; + explicit TEvNodesHealthCheckResponse( const Fq::Private::NodesHealthCheckResult& record) : Record(record) diff --git a/ydb/core/fq/libs/control_plane_storage/in_memory_control_plane_storage.cpp b/ydb/core/fq/libs/control_plane_storage/in_memory_control_plane_storage.cpp index f9089b0cac4b..7f715d2ec9f0 100644 --- a/ydb/core/fq/libs/control_plane_storage/in_memory_control_plane_storage.cpp +++ b/ydb/core/fq/libs/control_plane_storage/in_memory_control_plane_storage.cpp @@ -1,59 +1,173 @@ #include "control_plane_storage.h" -#include "util.h" +#include "ydb_control_plane_storage_impl.h" -#include -#include +#include -#include - -#include -#include +#include namespace NFq { -class TInMemoryControlPlaneStorageActor : public NActors::TActor { - struct TKey { +class TInMemoryControlPlaneStorageActor : public NActors::TActor, + public TControlPlaneStorageBase { + struct TScopeKey { TString Scope; TString Id; - bool operator<(const TKey& other) const - { - return tie(Scope, Id) < tie(other.Scope, other.Id); - } + std::strong_ordering operator<=>(const TScopeKey& other) const = default; }; - struct TConfig { - NConfig::TControlPlaneStorageConfig Proto; - TDuration IdempotencyKeyTtl; - TDuration AutomaticQueriesTtl; - TDuration ResultSetsTtl; - TDuration AnalyticsRetryCounterUpdateTime; - TDuration StreamingRetryCounterUpdateTime; - TDuration TaskLeaseTtl; - - TConfig(const NConfig::TControlPlaneStorageConfig& config) - : Proto(FillDefaultParameters(config)) - , IdempotencyKeyTtl(GetDuration(Proto.GetIdempotencyKeysTtl(), TDuration::Minutes(10))) - , AutomaticQueriesTtl(GetDuration(Proto.GetAutomaticQueriesTtl(), TDuration::Days(1))) - , ResultSetsTtl(GetDuration(Proto.GetResultSetsTtl(), TDuration::Days(1))) - , TaskLeaseTtl(GetDuration(Proto.GetTaskLeaseTtl(), TDuration::Seconds(30))) - { - } + struct TQueries { + using TKey = TScopeKey; + + struct TValue { + FederatedQuery::Query Query; + FederatedQuery::Internal::QueryInternal QueryInternal; + TString LastJobId; + TString User; + TString ResultId; + TInstant ResultExpireAt; + ui64 Generation = 0; + TInstant ExpireAt = TInstant::Zero(); + }; + + TMap Values; }; - TConfig Config; - TMap Queries; - TMap Connections; - TMap IdempotencyKeys; // idempotency_key -> created_at + struct TPendingQueries { + struct TKey { + TString Tenant; + TString Scope; + TString QueryId; + + std::strong_ordering operator<=>(const TKey& other) const = default; + }; + + struct TValue { + TRetryLimiter RetryLimiter; + TString Owner; + TInstant AssignedUntil; + TInstant LastSeenAt; + }; + + TMap Values; + }; + + struct TJobs { + struct TKey { + TString Scope; + TString QueryId; + TString JobId; + + std::strong_ordering operator<=>(const TKey& other) const = default; + }; + + struct TValue { + FederatedQuery::Job Job; + TInstant ExpireAt = TInstant::Zero(); + }; + + TMap Values; + }; - static constexpr int64_t InitialRevision = 1; + struct TConnections { + using TKey = TScopeKey; + using TEntity = FederatedQuery::Connection; + + struct TValue { + TEntity Connection; + TString User; + + const TEntity& GetEntity() const { + return Connection; + } + }; + + TMap Values; + }; + + struct TBindings { + using TKey = TScopeKey; + using TEntity = FederatedQuery::Binding; + + struct TValue { + TEntity Binding; + TString User; + + const TEntity& GetEntity() const { + return Binding; + } + }; + + TMap Values; + }; + + struct TIdempotencyKeys { + using TKey = TScopeKey; + + struct TValue { + TString Response; + TInstant ExpireAt = TInstant::Zero(); + }; + + TMap Values; + }; + + struct TResultSets { + struct TKey { + TString ResultId; + i32 ResultSetId; + + std::strong_ordering operator<=>(const TKey& other) const = default; + }; + + struct TValue { + TVector Rows; + TInstant ExpireAt = TInstant::Zero(); + }; + + TMap Values; + }; + + struct TNodes { + static constexpr TDuration TTL = TDuration::Seconds(15); + + struct TKey { + TString Tenant; + ui32 NodeId; + + std::strong_ordering operator<=>(const TKey& other) const = default; + }; + + struct TValue { + Fq::Private::NodeInfo Node; + TInstant ExpireAt = TInstant::Zero(); + }; + + TMap Values; + }; + + using TBase = TControlPlaneStorageBase; + + TQueries Queries; + TPendingQueries PendingQueries; + TJobs Jobs; + TConnections Connections; + TBindings Bindings; + TResultSets ResultSets; + TIdempotencyKeys IdempotencyKeys; + TNodes Nodes; public: - TInMemoryControlPlaneStorageActor(const NConfig::TControlPlaneStorageConfig& config) + TInMemoryControlPlaneStorageActor( + const NConfig::TControlPlaneStorageConfig& config, + const NYql::TS3GatewayConfig& s3Config, + const NConfig::TCommonConfig& common, + const NConfig::TComputeConfig& computeConfig, + const ::NMonitoring::TDynamicCounterPtr& counters, + const TString& tenantName) : TActor(&TThis::StateFunc) - , Config(config) - { - } + , TBase(config, s3Config, common, computeConfig, counters, tenantName) + {} static constexpr char ActorName[] = "YQ_CONTROL_PLANE_STORAGE"; @@ -83,75 +197,258 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActorGet()->Request; - CPS_LOG_D("CreateQueryRequest: " << request.DebugString()); - CleanupIndempotencyKeys(); - auto now = TInstant::Now(); - const TString idempotencyKey = request.idempotency_key(); - if (idempotencyKey && IdempotencyKeys.contains(idempotencyKey)) { - CPS_LOG_D("CreateQueryRequest, idempotency key already exist: " << request.DebugString()); - NYql::TIssue issue = MakeErrorIssue(TIssuesIds::BAD_REQUEST, "idempotency key already exist"); - Send(ev->Sender, new TEvControlPlaneStorage::TEvCreateQueryResponse(NYql::TIssues{issue}), 0, ev->Cookie); - return; - } - - NYql::TIssues issues = ValidateCreateQueryRequest(ev); - if (issues) { - CPS_LOG_D("CreateQueryRequest, validation failed: " << request.DebugString() << " error: " << issues.ToString()); - Send(ev->Sender, new TEvControlPlaneStorage::TEvCreateQueryResponse(issues), 0, ev->Cookie); - return; - } - - const TString user = ev->Get()->User; - const TString scope = ev->Get()->Scope; - const TString queryId = CreateGuidAsString(); - FederatedQuery::Query query; - FederatedQuery::QueryContent& content = *query.mutable_content(); - content = request.content(); - FederatedQuery::QueryMeta& meta = *query.mutable_meta(); - FederatedQuery::CommonMeta& common = *meta.mutable_common(); - common.set_id(queryId); - common.set_created_by(user); - auto timestamp = NProtoInterop::CastToProto(now); - *common.mutable_created_at() = timestamp; - common.set_revision(InitialRevision); - - Queries[{scope, queryId}] = query; - - if (!idempotencyKey) { - IdempotencyKeys[idempotencyKey] = now; - } - - CPS_LOG_D("CreateQueryRequest, success: " << request.DebugString() << " query_id: " << queryId); - FederatedQuery::CreateQueryResult result; - result.set_query_id(queryId); - Send(ev->Sender, new TEvControlPlaneStorage::TEvCreateQueryResponse(result, TAuditDetails{}), 0, ev->Cookie); - } - - void Handle(TEvControlPlaneStorage::TEvListQueriesRequest::TPtr& ev) - { + template + class TCommonRequestContext { + public: + using TResultType = TPrepareResponseResultType; + using TResponse = TResultType::Type; + using TAuditDetails = TResultType::TResponseAuditDetails; + + TCommonRequestContext(TInMemoryControlPlaneStorageActor& self, TEvRequest::TPtr& ev, const TString& cloudId, + const TString& scope, const TString& logPrefix, const TString& requestStr, const TString& responseStr) + : StartTime(TInstant::Now()) + , Event(*ev->Get()) + , Request(Event.Request) + , LogPrefix(logPrefix) + , RequestCounters(self.Counters.GetCounters(cloudId, scope, TYPE_SCOPE, TYPE_COMMON)) + , Self(self) + , EventPtr(ev) + , RequestStr(requestStr) + , ResponseStr(responseStr) + { + Self.Cleanup(); + CPS_LOG_I(RequestStr); + CPS_LOG_T(RequestStr << ":" << LogPrefix); + + RequestCounters.IncInFly(); + RequestCounters.Common->RequestBytes->Add(Event.GetByteSize()); + } + + TCommonRequestContext(TInMemoryControlPlaneStorageActor& self, TEvRequest::TPtr& ev, const TString& requestStr, const TString& responseStr) + : TCommonRequestContext(self, ev, "", "", GetLogPrefix(ev), requestStr, responseStr) + {} + + virtual bool Validate() { + if (const auto& issues = Self.ValidateRequest(EventPtr)) { + Fail("query validation", issues); + return false; + } + return true; + } + + bool IsFailed() const { + return Failed; + } + + void Fail(const TString& logInfo, const NYql::TIssues& issues) const { + Y_ABORT_UNLESS(!Failed, "Can not fail twice"); + CPS_LOG_W(RequestStr << ":" << LogPrefix << logInfo << " FAILED: " << issues.ToOneLineString()); + Self.SendResponseIssues(EventPtr->Sender, issues, EventPtr->Cookie, TInstant::Now() - StartTime, RequestCounters); + } + + virtual ~TCommonRequestContext() { + if (Failed) { + return; + } + + Self.SendResponse( + TStringBuilder() << RequestStr << " - " << ResponseStr, + NActors::TActivationContext::ActorSystem(), + NThreading::MakeFuture(NYdb::TStatus(NYdb::EStatus::SUCCESS, {})), + Self.SelfId(), + EventPtr, + StartTime, + RequestCounters, + [response = Response] { return response; }, + Self.Config->Proto.GetEnableDebugMode() ? std::make_shared() : TDebugInfoPtr{}); + } + + static TString GetLogPrefix(TEvRequest::TPtr& ev) { + return TStringBuilder() << "{" << ev->Get()->Request.DebugString() << "} "; + } + + public: + const TInstant StartTime; + const TEvRequest& Event; + const TEvRequest::TProto Request; + const TString LogPrefix; + TRequestCounters RequestCounters; + TResponse Response; + + protected: + TInMemoryControlPlaneStorageActor& Self; + TEvRequest::TPtr& EventPtr; + const TString RequestStr; + const TString ResponseStr; + bool Failed = false; + }; + + template + class TRequestContext : public TCommonRequestContext { + using TBase = TCommonRequestContext; + + Y_HAS_MEMBER(idempotency_key); + static constexpr bool HasIdempotencyKey = THasidempotency_key::value; + + public: + TRequestContext(TInMemoryControlPlaneStorageActor& self, TEvRequest::TPtr& ev, const TString& requestStr, const TString& responseStr) + : TBase( + self, ev, ev->Get()->CloudId, ev->Get()->Scope, + TStringBuilder() << TBase::GetLogPrefix(ev) << MakeUserInfo(ev->Get()->User, ev->Get()->Token), + requestStr, responseStr + ) + , CloudId(TBase::Event.CloudId) + , Scope(TBase::Event.Scope) + , User(TBase::Event.User) + , Token(TBase::Event.Token) + , Permissions(TBase::Event.Permissions) + , Quotas(TBase::Event.Quotas) + , TenantInfo(TBase::Event.TenantInfo) + , ComputeDatabase(TBase::Event.ComputeDatabase) + { + if constexpr (!std::is_same_v) { + TBase::Response.second.CloudId = CloudId; + } + } + + bool Validate() override { + if (!TBase::Validate()) { + return false; + } + if constexpr (HasIdempotencyKey) { + if (const TString& idempotencyKey = TBase::Request.idempotency_key()) { + if (const auto& value = TBase::Self.GetEntity(TBase::Self.IdempotencyKeys, {Scope, idempotencyKey})) { + if (!TBase::Response.first.ParseFromString(value->Response)) { + TBase::RequestCounters.Common->ParseProtobufError->Inc(); + TBase::Fail("idempotency key parse", {NYql::TIssue("INTERNAL ERROR. Error parsing proto message for idempotency key request. Please contact internal support")}); + } else { + TBase::Response.second.IdempotencyResult = true; + } + return false; + } + } + } + return true; + } + + ~TRequestContext() override { + if (TBase::Failed) { + return; + } + + if constexpr (HasIdempotencyKey) { + if (const TString& idempotencyKey = TBase::Request.idempotency_key()) { + this->Self.AddEntity(this->Self.IdempotencyKeys, {this->Scope, idempotencyKey}, { + .Response = TBase::Response.first.SerializeAsString(), + .ExpireAt = TBase::StartTime + this->Self.Config->IdempotencyKeyTtl + }); + } + } + } + + public: + const TString CloudId; + const TString Scope; + const TString User; + const TString Token; + const TPermissions Permissions; + const TMaybe Quotas; + const TTenantInfo::TPtr TenantInfo; + const FederatedQuery::Internal::ComputeDatabaseInternal ComputeDatabase; + }; + +#define HANDLE_CPS_REQUEST_IMPL(TEvRequest, TEvResponse, TContext, RTS_COUNTERS_ENUM, RTC_COUNTERS_ENUM) \ + using TContext##TEvRequest = TContext< \ + TEvControlPlaneStorage::TEvRequest, TEvControlPlaneStorage::TEvResponse, \ + RTS_COUNTERS_ENUM, RTC_COUNTERS_ENUM>; \ + void Handle(TEvControlPlaneStorage::TEvRequest::TPtr& ev) { \ + TContext##TEvRequest ctx(*this, ev, #TEvRequest, #TEvResponse); \ + if (!ctx.Validate()) { \ + return; \ + } \ + try { \ + Process##TRequest(ctx); \ + } catch (...) { \ + const auto& backtrace = TBackTrace::FromCurrentException().PrintToString(); \ + const auto logError = TStringBuilder() << "pocess "#TEvRequest" call, back trace:\n" << backtrace; \ + ctx.Fail(logError, {NYql::TIssue(CurrentExceptionMessage())}); \ + } \ + } \ + void Process##TRequest(TContext##TEvRequest& ctx) + +#define HANDLE_CPS_REQUEST(TEvRequest, TEvResponse, COUNTERS_ENUM) HANDLE_CPS_REQUEST_IMPL(TEvRequest, TEvResponse, TRequestContext, RTS_##COUNTERS_ENUM, RTC_##COUNTERS_ENUM) + + HANDLE_CPS_REQUEST(TEvCreateQueryRequest, TEvCreateQueryResponse, CREATE_QUERY) { + const auto& [query, job] = GetCreateQueryProtos(ctx.Request, ctx.User, ctx.StartTime); + if (query.ByteSizeLong() > Config->Proto.GetMaxRequestSize()) { + return ctx.Fail("query size validation", {NYql::TIssue(TStringBuilder() << "incoming request exceeded the size limit: " << query.ByteSizeLong() << " of " << Config->Proto.GetMaxRequestSize() << ". Please shorten your request")}); + } + ctx.Response.second.After = query; + + const TString& queryId = query.meta().common().id(); + ctx.Response.first.set_query_id(queryId); + + auto queryInternal = GetQueryInternalProto(ctx.Request, ctx.CloudId, ctx.Token, ctx.Quotas); + const auto queryType = ctx.Request.content().type(); + if (ctx.Request.execute_mode() != FederatedQuery::SAVE) { + *queryInternal.mutable_compute_connection() = ctx.ComputeDatabase.connection(); + FillConnectionsAndBindings( + queryInternal, + queryType, + GetEntities(Connections, ctx.Scope, ctx.User), + GetEntitiesWithVisibilityPriority(Connections, ctx.Scope, ctx.User), + GetEntitiesWithVisibilityPriority(Bindings, ctx.Scope, ctx.User) + ); + } + if (queryInternal.ByteSizeLong() > Config->Proto.GetMaxRequestSize()) { + return ctx.Fail("query internal size validation", {NYql::TIssue(TStringBuilder() << "the size of all connections and bindings in the project exceeded the limit: " << queryInternal.ByteSizeLong() << " of " << Config->Proto.GetMaxRequestSize() << ". Please reduce the number of connections and bindings")}); + } + + const auto& jobId = job.meta().id(); + if (ctx.Request.execute_mode() != FederatedQuery::SAVE) { + AddEntity(Jobs, {ctx.Scope, queryId, jobId}, {job}); + + TRetryLimiter retryLimiter; + retryLimiter.Assign(0, ctx.StartTime, 0.0); + + AddEntity(PendingQueries, { + .Tenant = ctx.TenantInfo->Assign(ctx.CloudId, ctx.Scope, queryType, TenantName), + .Scope = ctx.Scope, + .QueryId = queryId + }, {.RetryLimiter = retryLimiter}); + } + + AddEntity(Queries, {ctx.Scope, queryId}, { + .Query = query, + .QueryInternal = queryInternal, + .LastJobId = jobId, + .User = ctx.User + }); + } + + void Handle(TEvControlPlaneStorage::TEvListQueriesRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvListQueriesRequest::TPtr, FederatedQuery::ListQueriesResult, TEvControlPlaneStorage::TEvListQueriesResponse>(ev, "ListQueriesRequest"); } - void Handle(TEvControlPlaneStorage::TEvDescribeQueryRequest::TPtr& ev) - { - SendEmptyResponse< - TEvControlPlaneStorage::TEvDescribeQueryRequest::TPtr, - FederatedQuery::DescribeQueryResult, - TEvControlPlaneStorage::TEvDescribeQueryResponse>(ev, "DescribeQueryRequest"); + HANDLE_CPS_REQUEST(TEvDescribeQueryRequest, TEvDescribeQueryResponse, DESCRIBE_QUERY) { + const auto& query = GetEntity(Queries, {ctx.Scope, ctx.Request.query_id()}); + if (!query) { + return ctx.Fail("find query", {NYql::TIssue("Query does not exist")}); + } + + *ctx.Response.mutable_query() = query->Query; + FillDescribeQueryResult(ctx.Response, query->QueryInternal, ctx.User, ctx.Permissions); } - void Handle(TEvControlPlaneStorage::TEvModifyQueryRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvModifyQueryRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvModifyQueryRequest::TPtr, FederatedQuery::ModifyQueryResult, @@ -159,8 +456,8 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "ModifyQueryRequest"); } - void Handle(TEvControlPlaneStorage::TEvDeleteQueryRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDeleteQueryRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvDeleteQueryRequest::TPtr, FederatedQuery::DeleteQueryResult, @@ -168,8 +465,8 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "DeleteQueryRequest"); } - void Handle(TEvControlPlaneStorage::TEvControlQueryRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvControlQueryRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvControlQueryRequest::TPtr, FederatedQuery::ControlQueryResult, @@ -177,49 +474,98 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "ControlQueryRequest"); } - void Handle(TEvControlPlaneStorage::TEvGetResultDataRequest::TPtr& ev) - { - SendEmptyResponse< - TEvControlPlaneStorage::TEvGetResultDataRequest::TPtr, - FederatedQuery::GetResultDataResult, - TEvControlPlaneStorage::TEvGetResultDataResponse>(ev, "GetResultDataRequest"); + HANDLE_CPS_REQUEST(TEvGetResultDataRequest, TEvGetResultDataResponse, GET_RESULT_DATA) { + const auto& query = GetEntity(Queries, {ctx.Scope, ctx.Request.query_id()}); + if (!query) { + return ctx.Fail("find query", {NYql::TIssue("Query does not exist")}); + } + + if (!HasViewAccess(GetResultDataReadPerimssions(ctx.Event), query->Query.content().acl().visibility(), query->User, ctx.User)) { + return ctx.Fail("check ACL", {NYql::TIssue("Permission denied")}); + } + + const auto resultSetIndex = ctx.Request.result_set_index(); + const auto& resultSetMeta = query->Query.result_set_meta(); + if (resultSetIndex >= resultSetMeta.size()) { + return ctx.Fail("check result set index", {NYql::TIssue(TStringBuilder() << "Result set index out of bound: " << resultSetIndex << " >= " << resultSetMeta.size())}); + } + + const auto expireAt = query->ResultExpireAt; + if (query->Query.meta().status() != FederatedQuery::QueryMeta::COMPLETED || !expireAt) { + return ctx.Fail("check status", {NYql::TIssue("Result doesn't exist")}); + } + + if (expireAt < TInstant::Now()) { + return ctx.Fail("check expiration", {NYql::TIssue("Result removed by TTL")}); + } + + const auto& result = GetEntity(ResultSets, {query->ResultId, ctx.Request.result_set_index()}); + if (!result) { + return ctx.Fail("get result", {NYql::TIssue("INTERNAL ERROR. Failed to find result set")}); + } + + auto& resultSet = *ctx.Response.mutable_result_set(); + *resultSet.mutable_columns() = resultSetMeta[resultSetIndex].column(); + + const i64 offset = ctx.Request.offset(); + const i64 numberRows = result->Rows.size(); + for (i64 rowId = offset; rowId < offset + ctx.Request.limit() && rowId < numberRows; ++rowId) { + *resultSet.add_rows() = result->Rows[rowId]; + } } - void Handle(TEvControlPlaneStorage::TEvListJobsRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvListJobsRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvListJobsRequest::TPtr, FederatedQuery::ListJobsResult, TEvControlPlaneStorage::TEvListJobsResponse>(ev, "ListJobsRequest"); } - void Handle(TEvControlPlaneStorage::TEvCreateConnectionRequest::TPtr& ev) - { - SendEmptyAuditResponse< - TEvControlPlaneStorage::TEvCreateConnectionRequest::TPtr, - FederatedQuery::CreateConnectionResult, - TEvControlPlaneStorage::TEvCreateConnectionResponse, - TAuditDetails>(ev, "CreateConnectionRequest"); + HANDLE_CPS_REQUEST(TEvCreateConnectionRequest, TEvCreateConnectionResponse, CREATE_CONNECTION) { + const auto& content = ctx.Request.content(); + const auto visibility = content.acl().visibility(); + const auto& name = content.name(); + if (!CheckConnectionOrBindingName(Connections, ctx.Scope, ctx.User, visibility, name)) { + return ctx.Fail("check name", {NYql::TIssue("Connection with the same name already exists. Please choose another name")}); + } + if (!CheckConnectionOrBindingName(Bindings, ctx.Scope, ctx.User, visibility, name)) { + return ctx.Fail("check name", {NYql::TIssue("Binding with the same name already exists. Please choose another name")}); + } + if (GetNumberEntitiesByScope(Connections, ctx.Scope) >= Config->Proto.GetMaxCountConnections()) { + return ctx.Fail("check number", {NYql::TIssue(TStringBuilder() << "Too many connections in folder: " << Config->Proto.GetMaxCountConnections() << ". Please remove unused connections")}); + } + + const auto& [connection, _] = GetCreateConnectionProtos(ctx.Request, ctx.CloudId, ctx.User, ctx.StartTime); + ctx.Response.second.After = connection; + + const TString& connectionId = connection.meta().id(); + ctx.Response.first.set_connection_id(connectionId); + + AddEntity(Connections, {ctx.Scope, connectionId}, { + .Connection = connection, + .User = ctx.User + }); } - void Handle(TEvControlPlaneStorage::TEvListConnectionsRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvListConnectionsRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvListConnectionsRequest::TPtr, FederatedQuery::ListConnectionsResult, TEvControlPlaneStorage::TEvListConnectionsResponse>(ev, "ListConnectionsRequest"); } - void Handle(TEvControlPlaneStorage::TEvDescribeConnectionRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDescribeConnectionRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvDescribeConnectionRequest::TPtr, FederatedQuery::DescribeConnectionResult, TEvControlPlaneStorage::TEvDescribeConnectionResponse>(ev, "DescribeConnectionRequest"); } - void Handle(TEvControlPlaneStorage::TEvModifyConnectionRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvModifyConnectionRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvModifyConnectionRequest::TPtr, FederatedQuery::ModifyConnectionResult, @@ -227,8 +573,8 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "ModifyConnectionRequest"); } - void Handle(TEvControlPlaneStorage::TEvDeleteConnectionRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDeleteConnectionRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvDeleteConnectionRequest::TPtr, FederatedQuery::DeleteConnectionResult, @@ -236,33 +582,64 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "DeleteConnectionRequest"); } - void Handle(TEvControlPlaneStorage::TEvCreateBindingRequest::TPtr& ev) - { - SendEmptyAuditResponse< - TEvControlPlaneStorage::TEvCreateBindingRequest::TPtr, - FederatedQuery::CreateBindingResult, - TEvControlPlaneStorage::TEvCreateBindingResponse, - TAuditDetails>(ev, "CreateBindingRequest"); + HANDLE_CPS_REQUEST(TEvCreateBindingRequest, TEvCreateBindingResponse, CREATE_BINDING) { + const auto& content = ctx.Request.content(); + const auto visibility = content.acl().visibility(); + const auto& name = content.name(); + if (!CheckConnectionOrBindingName(Connections, ctx.Scope, ctx.User, visibility, name)) { + return ctx.Fail("check name", {NYql::TIssue("Connection with the same name already exists. Please choose another name")}); + } + if (!CheckConnectionOrBindingName(Bindings, ctx.Scope, ctx.User, visibility, name)) { + return ctx.Fail("check name", {NYql::TIssue("Binding with the same name already exists. Please choose another name")}); + } + if (GetNumberEntitiesByScope(Bindings, ctx.Scope) >= Config->Proto.GetMaxCountBindings()) { + return ctx.Fail("check number", {NYql::TIssue(TStringBuilder() << "Too many bindings in folder: " << Config->Proto.GetMaxCountBindings() << ". Please remove unused bindings")}); + } + + const auto& connection = GetEntity(Connections, {ctx.Scope, content.connection_id()}); + if (!connection) { + return ctx.Fail("check connection", {NYql::TIssue("Connection for binding not found")}); + } + + const auto connectionVisibility = connection->Connection.content().acl().visibility(); + if (content.acl().visibility() == FederatedQuery::Acl::SCOPE && connectionVisibility == FederatedQuery::Acl::PRIVATE) { + return ctx.Fail("check connection ACL", {NYql::TIssue("Binding with SCOPE visibility cannot refer to connection with PRIVATE visibility")}); + } + + if (!HasManageAccess(GetCreateBindingPerimssions(ctx.Event), connectionVisibility, connection->User, ctx.User)) { + return ctx.Fail("check connection ACL", {NYql::TIssue("Permission denied for binding connection")}); + } + + const auto& [binding, _] = GetCreateBindingProtos(ctx.Request, ctx.CloudId, ctx.User, ctx.StartTime); + ctx.Response.second.After = binding; + + const TString& bindingId = binding.meta().id(); + ctx.Response.first.set_binding_id(bindingId); + + AddEntity(Bindings, {ctx.Scope, bindingId}, { + .Binding = binding, + .User = ctx.User + }); } - void Handle(TEvControlPlaneStorage::TEvListBindingsRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvListBindingsRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvListBindingsRequest::TPtr, FederatedQuery::ListBindingsResult, TEvControlPlaneStorage::TEvListBindingsResponse>(ev, "ListBindingsRequest"); } - void Handle(TEvControlPlaneStorage::TEvDescribeBindingRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDescribeBindingRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvDescribeBindingRequest::TPtr, FederatedQuery::DescribeBindingResult, TEvControlPlaneStorage::TEvDescribeBindingResponse>(ev, "DescribeBindingRequest"); } - void Handle(TEvControlPlaneStorage::TEvModifyBindingRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvModifyBindingRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvModifyBindingRequest::TPtr, FederatedQuery::ModifyBindingResult, @@ -270,8 +647,8 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "ModifyBindingRequest"); } - void Handle(TEvControlPlaneStorage::TEvDeleteBindingRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDeleteBindingRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvDeleteBindingRequest::TPtr, FederatedQuery::DeleteBindingResult, @@ -279,53 +656,109 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "DeleteBindingRequest"); } - void Handle(TEvControlPlaneStorage::TEvDescribeJobRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDescribeJobRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvDescribeJobRequest::TPtr, FederatedQuery::DescribeJobResult, TEvControlPlaneStorage::TEvDescribeJobResponse>(ev, "DescribeJobRequest"); } - void Handle(TEvControlPlaneStorage::TEvWriteResultDataRequest::TPtr& ev) - { - SendEmptyResponse< - TEvControlPlaneStorage::TEvWriteResultDataRequest::TPtr, - NYql::TIssues, - TEvControlPlaneStorage::TEvWriteResultDataResponse>(ev, "WriteResultDataRequest"); + HANDLE_CPS_REQUEST_IMPL(TEvWriteResultDataRequest, TEvWriteResultDataResponse, TCommonRequestContext, RTS_MAX, RTC_WRITE_RESULT_DATA) { + ctx.Response.set_request_id(ctx.Request.request_id()); + + const auto offset = ctx.Request.offset(); + const auto& newRows = ctx.Request.result_set().rows(); + + auto& [resultRows, expireAt] = ResultSets.Values[{ctx.Request.result_id().value(), static_cast(ctx.Request.result_set_id())}]; + expireAt = NProtoInterop::CastFromProto(ctx.Request.deadline()); + + resultRows.resize(std::max(offset + newRows.size(), resultRows.size())); + for (size_t i = offset; const auto& row : newRows) { + resultRows[i++] = row; + } } - void Handle(TEvControlPlaneStorage::TEvGetTaskRequest::TPtr& ev) - { - CPS_LOG_I("GetTaskRequest"); - Fq::Private::GetTaskResult result; - auto event = std::make_unique(result); - NActors::TActivationContext::ActorSystem()->Send(new IEventHandle(ev->Sender, SelfId(), event.release(), 0, ev->Cookie)); + HANDLE_CPS_REQUEST_IMPL(TEvGetTaskRequest, TEvGetTaskResponse, TCommonRequestContext, RTS_MAX, RTC_GET_TASK) { + const auto& tasksInternal = GetActiveTasks(ctx); + + TVector tasks; + tasks.reserve(tasksInternal.size()); + for (const auto& taskInternal : tasksInternal) { + if (const auto& task = AssignTask(ctx, taskInternal)) { + tasks.emplace_back(*task); + } + if (ctx.IsFailed()) { + return; + } + } + + FillGetTaskResult(ctx.Response, tasks); } - void Handle(TEvControlPlaneStorage::TEvPingTaskRequest::TPtr& ev) - { - SendEmptyResponse< - TEvControlPlaneStorage::TEvPingTaskRequest::TPtr, - Fq::Private::PingTaskResult, - TEvControlPlaneStorage::TEvPingTaskResponse>(ev, "PingTaskRequest"); + HANDLE_CPS_REQUEST_IMPL(TEvPingTaskRequest, TEvPingTaskResponse, TCommonRequestContext, RTS_MAX, RTC_PING_TASK) { + const auto& scope = ctx.Request.scope(); + const auto& queryId = ctx.Request.query_id().value(); + + auto query = GetEntity(Queries, {scope, queryId}); + if (!query) { + return ctx.Fail("get query", {NYql::TIssue("INTERNAL ERROR. Query for ping task not found")}); + } + + auto job = GetEntity(Jobs, {scope, queryId, query->LastJobId}); + if (!job) { + return ctx.Fail("get job", {NYql::TIssue("INTERNAL ERROR. Job for ping task not found")}); + } + + auto pendingQuery = GetEntity(PendingQueries, {ctx.Request.tenant(), scope, queryId}); + if (!pendingQuery) { + return ctx.Fail("get pending query", {NYql::TIssue("INTERNAL ERROR. Pending query for ping task not found")}); + } + + auto resuest = ctx.Request; + auto finalStatus = std::make_shared(); + + TDuration backoff = Config->TaskLeaseTtl; + TInstant expireAt = ctx.StartTime + Config->AutomaticQueriesTtl; + UpdateTaskInfo(TActivationContext::ActorSystem(), resuest, finalStatus, query->Query, query->QueryInternal, job->Job, pendingQuery->Owner, pendingQuery->RetryLimiter, backoff, expireAt); + PingTask(ctx, *query, *job, *pendingQuery, backoff, expireAt); + + if (IsTerminalStatus(ctx.Request.status())) { + FillQueryStatistics(finalStatus, query->Query, query->QueryInternal, pendingQuery->RetryLimiter); + } + Send(SelfId(), new TEvControlPlaneStorage::TEvFinalStatusReport( + queryId, finalStatus->JobId, finalStatus->CloudId, scope, std::move(finalStatus->FinalStatistics), + finalStatus->Status, finalStatus->StatusCode, finalStatus->QueryType, finalStatus->Issues, finalStatus->TransientIssues)); } - void Handle(TEvControlPlaneStorage::TEvNodesHealthCheckRequest::TPtr& ev) - { - SendEmptyResponse< - TEvControlPlaneStorage::TEvNodesHealthCheckRequest::TPtr, - Fq::Private::NodesHealthCheckResult, - TEvControlPlaneStorage::TEvNodesHealthCheckResponse>(ev, "NodesHealthCheckRequest"); + HANDLE_CPS_REQUEST_IMPL(TEvNodesHealthCheckRequest, TEvNodesHealthCheckResponse, TCommonRequestContext, RTS_MAX, RTC_NODES_HEALTH_CHECK) { + const auto& tenant = ctx.Request.tenant(); + const auto& node = ctx.Request.node(); + + AddEntity(Nodes, {tenant, node.node_id()}, { + .Node = node, + .ExpireAt = ctx.StartTime + TNodes::TTL + }); + + for (const auto& [key, value] : Nodes.Values) { + if (key.Tenant == tenant) { + *ctx.Response.add_nodes() = value.Node; + } + } } +#undef HANDLE_CPS_REQUEST +#undef HANDLE_CPS_REQUEST_IMPL + void Handle(NActors::NMon::TEvHttpInfo::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); TStringStream str; Send(ev->Sender, new NActors::NMon::TEvHttpInfoRes(str.Str())); } template void SendEmptyResponse(TRequest& ev, std::string logText) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("SendEmptyResponse"); CPS_LOG_I(logText); TResult result = {}; @@ -335,6 +768,7 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor void SendEmptyAuditResponse(TRequest& ev, std::string logText) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("SendEmptyAuditResponse"); CPS_LOG_I(logText); TResult result = {}; @@ -343,50 +777,250 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActorSend(new IEventHandle(ev->Sender, SelfId(), event.release(), 0, ev->Cookie)); } - NYql::TIssues ValidateCreateQueryRequest(TEvControlPlaneStorage::TEvCreateQueryRequest::TPtr& ev) - { - NYql::TIssues issues; - const FederatedQuery::CreateQueryRequest& request = ev->Get()->Request; - const TString user = ev->Get()->User; - const TString scope = ev->Get()->Scope; - const FederatedQuery::QueryContent& query = request.content(); +private: + template + static std::optional GetEntity(const TTable& table, const TTable::TKey& key) { + const auto it = table.Values.find(key); + if (it == table.Values.end()) { + return std::nullopt; + } + return it->second; + } + + template + static bool AddEntity(TTable& table, const TTable::TKey& key, const TTable::TValue& value) { + return table.Values.emplace(key, value).second; + } + + template + auto GetScopeRange(const TMap& table, const TString& scope) const { + const auto startIt = table.lower_bound({scope, ""}); + + std::ranges::subrange range(startIt, table.end()); + return range | std::views::take_while([scope](auto element) { + return element.first.Scope == scope; + }); + } + + template + auto GetVisibleRange(const TMap& table, const TString& scope, const TString& user, std::optional visibility = std::nullopt) const { + auto range = GetScopeRange(table, scope); + return range | std::views::filter([user, visibility, ignorePrivate = Config->Proto.GetIgnorePrivateSources()](const auto& element) { + const auto entityVisibility = element.second.GetEntity().content().acl().visibility(); + if (ignorePrivate && entityVisibility == FederatedQuery::Acl::PRIVATE) { + return false; + } + if (visibility && entityVisibility != *visibility) { + return false; + } + return entityVisibility == FederatedQuery::Acl::SCOPE || element.second.User == user; + }); + } + + template + TVector GetEntities(const TTable& table, const TString& scope, const TString& user) const { + auto range = GetVisibleRange(table.Values, scope, user) | std::views::transform([](const auto& element) { + return element.second.GetEntity(); + }) | std::views::common; + return {range.begin(), range.end()}; + } + + template + THashMap GetEntitiesWithVisibilityPriority(const TTable& table, const TString& scope, const TString& user) const { + THashMap entities; + for (const auto& [_, value] : GetVisibleRange(table.Values, scope, user)) { + const auto& entity = value.GetEntity(); + const auto visibility = entity.content().acl().visibility(); + const TString& name = entity.content().name(); + if (auto it = entities.find(name); it != entities.end()) { + if (visibility == FederatedQuery::Acl::PRIVATE) { + it->second = entity; + } + } else { + entities[name] = entity; + } + } + return entities; + } + + template + bool CheckConnectionOrBindingName(const TTable& table, const TString& scope, const TString& user, FederatedQuery::Acl::Visibility visibility, const TString& name) const { + auto range = GetVisibleRange(table.Values, scope, user, visibility) | std::views::transform([](const auto& element) { + return element.second.GetEntity().content().name(); + }); + return std::ranges::find(range, name) == range.end(); + } + + template + ui64 GetNumberEntitiesByScope(const TTable& table, const TString& scope) const { + auto range = GetScopeRange(table.Values, scope) | std::views::common; + return std::distance(range.begin(), range.end()); + } + +private: + void Cleanup() { + CleanupTable(Queries); + CleanupTable(Jobs); + CleanupTable(ResultSets); + CleanupTable(IdempotencyKeys); + CleanupTable(Nodes); + } + + template + static void CleanupTable(TTable& table) { + const auto now = TInstant::Now(); + std::erase_if(table.Values, [now](const auto& item) { + const auto expireAt = item.second.ExpireAt; + return expireAt && expireAt < now; + }); + } + +private: + // Get / Ping task utils + + struct TTaskInternal { + TTask Task; + TString Owner; + TRetryLimiter RetryLimiter; + TString TenantName; + bool ShouldAbortTask; + }; + + TVector GetActiveTasks(const TCommonRequestContextTEvGetTaskRequest& ctx) const { + const ui64 tasksBatchSize = Config->Proto.GetTasksBatchSize(); + const TString& tenantName = ctx.Request.tenant(); + + TVector tasks; + tasks.reserve(std::min(tasksBatchSize, PendingQueries.Values.size())); + for (const auto& [key, query] : PendingQueries.Values) { + if (key.Tenant != tenantName || query.AssignedUntil >= ctx.StartTime) { + continue; + } + + TTaskInternal& taskInternal = tasks.emplace_back(); + taskInternal.Owner = ctx.Request.owner_id(); + taskInternal.TenantName = tenantName; + taskInternal.RetryLimiter = query.RetryLimiter; + + auto& task = taskInternal.Task; + task.Scope = key.Scope; + task.QueryId = key.QueryId; + + if (query.Owner) { + CPS_LOG_T("Task (Query): " << task.QueryId << " Lease TIMEOUT, RetryCounterUpdatedAt " << taskInternal.RetryLimiter.RetryCounterUpdatedAt << " LastSeenAt: " << query.LastSeenAt); + taskInternal.ShouldAbortTask = !taskInternal.RetryLimiter.UpdateOnRetry(query.LastSeenAt, Config->TaskLeaseRetryPolicy, ctx.StartTime); + } + task.RetryCount = taskInternal.RetryLimiter.RetryCount; + + CPS_LOG_T("Task (Query): " << task.QueryId << " RetryRate: " << taskInternal.RetryLimiter.RetryRate << " RetryCounter: " << taskInternal.RetryLimiter.RetryCount << " At: " << taskInternal.RetryLimiter.RetryCounterUpdatedAt << (taskInternal.ShouldAbortTask ? " ABORTED" : "")); + + if (tasks.size() >= tasksBatchSize) { + break; + } + } + + std::shuffle(tasks.begin(), tasks.end(), std::default_random_engine(TInstant::Now().MicroSeconds())); + const ui64 numTasksProportion = Config->Proto.GetNumTasksProportion(); + tasks.resize((tasks.size() + numTasksProportion - 1) / numTasksProportion); + + return tasks; + } + + std::optional AssignTask(const TCommonRequestContextTEvGetTaskRequest& ctx, TTaskInternal taskInternal) { + auto& task = taskInternal.Task; - TString error; - if (!request.validate(error)) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, error)); + const auto& query = GetEntity(Queries, {task.Scope, task.QueryId}); + if (!query) { + ctx.Fail("task build", {NYql::TIssue(TStringBuilder() << "INTERNAL ERROR. Not found query for task in scope " << task.Scope << " with query id " << task.QueryId)}); + return std::nullopt; } - if (query.type() == FederatedQuery::QueryContent::QUERY_TYPE_UNSPECIFIED) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "type field is not specified")); + task.Generation = query->Generation + 1; + task.Query = query->Query; + task.Internal = query->QueryInternal; + task.Deadline = TInstant::Now() + (task.Query.content().automatic() ? std::min(Config->AutomaticQueriesTtl, Config->ResultSetsTtl) : Config->ResultSetsTtl); + *task.Internal.mutable_result_ttl() = NProtoInterop::CastToProto(Config->ResultSetsTtl); + + if (Config->Proto.GetDisableCurrentIam()) { + task.Internal.clear_token(); } - if (query.acl().visibility() == FederatedQuery::Acl::VISIBILITY_UNSPECIFIED) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "acl.visibility field is not specified")); + if (taskInternal.ShouldAbortTask) { + AddTransientIssues(task.Query.mutable_transient_issue(), {NYql::TIssue("Query was aborted by system due to high failure rate")}); + task.Query.mutable_meta()->set_status(FederatedQuery::QueryMeta::ABORTING_BY_SYSTEM); } - if (request.ByteSize() > static_cast(Config.Proto.GetMaxRequestSize())) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "Request size exceeded " + ToString(request.ByteSize()) + " out of " + ToString(Config.Proto.GetMaxRequestSize()) + " bytes")); + if (const auto tenantInfo = ctx.Event.TenantInfo) { + const TString& newTenant = tenantInfo->Assign(task.Internal.cloud_id(), task.Scope, task.Query.content().type(), taskInternal.TenantName); + if (newTenant != taskInternal.TenantName) { + UpdateTaskState(ctx, taskInternal, newTenant); + return std::nullopt; + } + if (tenantInfo->TenantState.Value(taskInternal.TenantName, TenantState::Active) != TenantState::Active) { + return std::nullopt; + } } - const uint64_t countQueries = count_if(Queries.begin(), Queries.end(), [scope](const auto& item) { - const auto& [key, value] = item; - return key.Scope == scope; - }); + UpdateTaskState(ctx, taskInternal); + + return task; + } - if (countQueries > Config.Proto.GetMaxCountQueries()) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "The count of the queries exceeds the limit of " + ToString(countQueries) + " out of " + ToString(Config.Proto.GetMaxCountQueries()))); + void UpdateTaskState(const TCommonRequestContextTEvGetTaskRequest& ctx, const TTaskInternal& taskInternal, const TString& newTenant = "") { + const auto& task = taskInternal.Task; + + const auto queryIt = Queries.Values.find({task.Scope, task.QueryId}); + if (queryIt != Queries.Values.end()) { + queryIt->second.Query = task.Query; + queryIt->second.QueryInternal = task.Internal; + queryIt->second.Generation = task.Generation; } - return issues; + const auto pendingIt = PendingQueries.Values.find({taskInternal.TenantName, task.Scope, task.QueryId}); + if (pendingIt != PendingQueries.Values.end()) { + pendingIt->second.Owner = taskInternal.Owner; + pendingIt->second.RetryLimiter = taskInternal.RetryLimiter; + pendingIt->second.AssignedUntil = ctx.StartTime + Config->TaskLeaseTtl; + pendingIt->second.LastSeenAt = ctx.StartTime; + if (newTenant) { + const auto value = pendingIt->second; + PendingQueries.Values.erase(pendingIt); + AddEntity(PendingQueries, {newTenant, task.Scope, task.QueryId}, value); + } + } } - void CleanupIndempotencyKeys() - { - auto now = TInstant::Now(); - erase_if(IdempotencyKeys, [this, now](const auto& item) { - auto const& [idempotencyKey, timestamp] = item; - return timestamp + Config.IdempotencyKeyTtl < now; - }); + void PingTask(const TCommonRequestContextTEvPingTaskRequest& ctx, const TQueries::TValue& query, const TJobs::TValue& job, const TPendingQueries::TValue& pendingQuery, TDuration backoff, TInstant expireAt) { + const auto& scope = ctx.Request.scope(); + const auto& queryId = ctx.Request.query_id().value(); + const auto status = query.Query.meta().status(); + + const auto pendingIt = PendingQueries.Values.find({ctx.Request.tenant(), scope, queryId}); + if (pendingIt != PendingQueries.Values.end()) { + if (IsTerminalStatus(status)) { + PendingQueries.Values.erase(pendingIt); + } else { + pendingIt->second = pendingQuery; + pendingIt->second.AssignedUntil = ctx.StartTime + backoff; + } + } + + const bool automaticFinished = IsTerminalStatus(status) && query.Query.content().automatic(); + const auto jobIt = Jobs.Values.find({scope, queryId, job.Job.meta().id()}); + if (jobIt != Jobs.Values.end()) { + jobIt->second = job; + jobIt->second.ExpireAt = automaticFinished ? expireAt : TInstant::Zero(); + } + + const auto queryIt = Queries.Values.find({scope, queryId}); + if (queryIt != Queries.Values.end()) { + queryIt->second = query; + if (ctx.Request.has_result_id()) { + queryIt->second.ResultId = ctx.Request.result_id().value(); + } + queryIt->second.ResultExpireAt = status == FederatedQuery::QueryMeta::COMPLETED ? NProtoInterop::CastFromProto(ctx.Request.deadline()) : TInstant::Zero(); + queryIt->second.ExpireAt = automaticFinished ? expireAt : TInstant::Zero(); + } } }; @@ -395,8 +1029,14 @@ NActors::TActorId ControlPlaneStorageServiceActorId(ui32 nodeId) { return NActors::TActorId(nodeId, name); } -NActors::IActor* CreateInMemoryControlPlaneStorageServiceActor(const NConfig::TControlPlaneStorageConfig& config) { - return new TInMemoryControlPlaneStorageActor(config); +NActors::IActor* CreateInMemoryControlPlaneStorageServiceActor( + const NConfig::TControlPlaneStorageConfig& config, + const NYql::TS3GatewayConfig& s3Config, + const NConfig::TCommonConfig& common, + const NConfig::TComputeConfig& computeConfig, + const ::NMonitoring::TDynamicCounterPtr& counters, + const TString& tenantName) { + return new TInMemoryControlPlaneStorageActor(config, s3Config, common, computeConfig, counters, tenantName); } } // NFq diff --git a/ydb/core/fq/libs/control_plane_storage/internal/nodes_health_check.cpp b/ydb/core/fq/libs/control_plane_storage/internal/nodes_health_check.cpp index 652098e00de3..ac815ff3b37d 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/nodes_health_check.cpp +++ b/ydb/core/fq/libs/control_plane_storage/internal/nodes_health_check.cpp @@ -1,9 +1,17 @@ #include "utils.h" +#include #include namespace NFq { +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvNodesHealthCheckRequest::TPtr& ev) const { + const auto& request = ev->Get()->Request; + const auto& node = request.node(); + + return ValidateNodesHealthCheck(request.tenant(), node.instance_id(), node.hostname()); +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvNodesHealthCheckRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -29,8 +37,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvNodesHealth CPS_LOG_T("NodesHealthCheckRequest: {" << request.DebugString() << "}"); - NYql::TIssues issues = ValidateNodesHealthCheck(tenant, instanceId, hostName); - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_W("NodesHealthCheckRequest: {" << request.DebugString() << "} validation FAILED: " << issues.ToOneLineString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); diff --git a/ydb/core/fq/libs/control_plane_storage/internal/task_get.cpp b/ydb/core/fq/libs/control_plane_storage/internal/task_get.cpp index 41bad0bb970d..3031cac8df01 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/task_get.cpp +++ b/ydb/core/fq/libs/control_plane_storage/internal/task_get.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include @@ -247,6 +248,85 @@ TDuration ExtractLimit(const TTask& task) { return executionLimit; } +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvGetTaskRequest::TPtr& ev) const +{ + const auto& request = ev->Get()->Request; + NYql::TIssues issues = ValidateGetTask(request.owner_id(), request.host()); + + if (!ev->Get()->TenantInfo) { + issues.AddIssue(MakeErrorIssue(TIssuesIds::NOT_READY, "Control Plane is not ready yet. Please retry later.")); + } + + return issues; +} + +void TControlPlaneStorageBase::FillGetTaskResult(Fq::Private::GetTaskResult& result, const TVector& tasks) const +{ + for (const auto& task : tasks) { + const auto& queryType = task.Query.content().type(); + if (queryType != FederatedQuery::QueryContent::ANALYTICS && queryType != FederatedQuery::QueryContent::STREAMING) { //TODO: fix + ythrow yexception() + << "query type " + << FederatedQuery::QueryContent::QueryType_Name(queryType) + << " unsupported"; + } + + auto* newTask = result.add_tasks(); + newTask->set_query_type(queryType); + newTask->set_query_syntax(task.Query.content().syntax()); + newTask->set_execute_mode(task.Query.meta().execute_mode()); + newTask->set_state_load_mode(task.Internal.state_load_mode()); + auto* queryId = newTask->mutable_query_id(); + queryId->set_value(task.Query.meta().common().id()); + newTask->set_streaming(queryType == FederatedQuery::QueryContent::STREAMING); + newTask->set_text(task.Query.content().text()); + *newTask->mutable_connection() = task.Internal.connection(); + *newTask->mutable_binding() = task.Internal.binding(); + newTask->set_user_token(task.Internal.token()); + newTask->set_user_id(task.Query.meta().common().created_by()); + newTask->set_generation(task.Generation); + newTask->set_status(task.Query.meta().status()); + *newTask->mutable_created_topic_consumers() = task.Internal.created_topic_consumers(); + newTask->mutable_sensor_labels()->insert({"cloud_id", task.Internal.cloud_id()}); + newTask->mutable_sensor_labels()->insert({"scope", task.Scope}); + newTask->set_automatic(task.Query.content().automatic()); + newTask->set_query_name(task.Query.content().name()); + *newTask->mutable_deadline() = NProtoInterop::CastToProto(task.Deadline); + newTask->mutable_disposition()->CopyFrom(task.Internal.disposition()); + newTask->set_result_limit(task.Internal.result_limit()); + *newTask->mutable_execution_limit() = NProtoInterop::CastToProto(ExtractLimit(task)); + *newTask->mutable_request_started_at() = task.Query.meta().started_at(); + *newTask->mutable_request_submitted_at() = task.Query.meta().submitted_at(); + + newTask->set_restart_count(task.RetryCount); + auto* jobId = newTask->mutable_job_id(); + jobId->set_value(task.Query.meta().last_job_id()); + + for (const auto& connection: task.Internal.connection()) { + const auto serviceAccountId = ExtractServiceAccountId(connection); + if (!serviceAccountId) { + continue; + } + auto* account = newTask->add_service_accounts(); + account->set_value(serviceAccountId); + } + + *newTask->mutable_dq_graph() = task.Internal.dq_graph(); + newTask->set_dq_graph_index(task.Internal.dq_graph_index()); + *newTask->mutable_dq_graph_compressed() = task.Internal.dq_graph_compressed(); + + *newTask->mutable_result_set_meta() = task.Query.result_set_meta(); + newTask->set_scope(task.Scope); + *newTask->mutable_resources() = task.Internal.resources(); + + newTask->set_execution_id(task.Internal.execution_id()); + newTask->set_operation_id(task.Internal.operation_id()); + *newTask->mutable_compute_connection() = task.Internal.compute_connection(); + *newTask->mutable_result_ttl() = task.Internal.result_ttl(); + *newTask->mutable_parameters() = task.Query.content().parameters(); + } +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetTaskRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -262,13 +342,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetTaskRequ CPS_LOG_T("GetTaskRequest: {" << request.DebugString() << "}"); - NYql::TIssues issues = ValidateGetTask(owner, hostName); - - if (!ev->Get()->TenantInfo) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::NOT_READY, "Control Plane is not ready yet. Please retry later.")); - } - - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_W("GetTaskRequest: {" << request.DebugString() << "} FAILED: " << issues.ToOneLineString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); @@ -364,15 +438,17 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetTaskRequ responseTasks=responseTasks] (const auto& readFuture) mutable { try { - if (!readFuture.GetValue().IsSuccess()) + if (!readFuture.GetValue().IsSuccess()) { return readFuture; + } } catch (...) { return readFuture; } auto pickTaskParams = prepareParams(*resultSets); - if (pickTaskParams.empty()) + if (pickTaskParams.empty()) { return readFuture; + } auto debugInfos = std::make_shared>(pickTaskParams.size()); if (Config->Proto.GetEnableDebugMode()) { @@ -410,74 +486,9 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetTaskRequ }); }); - auto prepare = [response] { + auto prepare = [this, response] { Fq::Private::GetTaskResult result; - const auto& tasks = std::get<0>(*response); - - for (const auto& task : tasks) { - const auto& queryType = task.Query.content().type(); - if (queryType != FederatedQuery::QueryContent::ANALYTICS && queryType != FederatedQuery::QueryContent::STREAMING) { //TODO: fix - ythrow yexception() - << "query type " - << FederatedQuery::QueryContent::QueryType_Name(queryType) - << " unsupported"; - } - - auto* newTask = result.add_tasks(); - newTask->set_query_type(queryType); - newTask->set_query_syntax(task.Query.content().syntax()); - newTask->set_execute_mode(task.Query.meta().execute_mode()); - newTask->set_state_load_mode(task.Internal.state_load_mode()); - auto* queryId = newTask->mutable_query_id(); - queryId->set_value(task.Query.meta().common().id()); - newTask->set_streaming(queryType == FederatedQuery::QueryContent::STREAMING); - newTask->set_text(task.Query.content().text()); - *newTask->mutable_connection() = task.Internal.connection(); - *newTask->mutable_binding() = task.Internal.binding(); - newTask->set_user_token(task.Internal.token()); - newTask->set_user_id(task.Query.meta().common().created_by()); - newTask->set_generation(task.Generation); - newTask->set_status(task.Query.meta().status()); - *newTask->mutable_created_topic_consumers() = task.Internal.created_topic_consumers(); - newTask->mutable_sensor_labels()->insert({"cloud_id", task.Internal.cloud_id()}); - newTask->mutable_sensor_labels()->insert({"scope", task.Scope}); - newTask->set_automatic(task.Query.content().automatic()); - newTask->set_query_name(task.Query.content().name()); - *newTask->mutable_deadline() = NProtoInterop::CastToProto(task.Deadline); - newTask->mutable_disposition()->CopyFrom(task.Internal.disposition()); - newTask->set_result_limit(task.Internal.result_limit()); - *newTask->mutable_execution_limit() = NProtoInterop::CastToProto(ExtractLimit(task)); - *newTask->mutable_request_started_at() = task.Query.meta().started_at(); - *newTask->mutable_request_submitted_at() = task.Query.meta().submitted_at(); - - newTask->set_restart_count(task.RetryCount); - auto* jobId = newTask->mutable_job_id(); - jobId->set_value(task.Query.meta().last_job_id()); - - for (const auto& connection: task.Internal.connection()) { - const auto serviceAccountId = ExtractServiceAccountId(connection); - if (!serviceAccountId) { - continue; - } - auto* account = newTask->add_service_accounts(); - account->set_value(serviceAccountId); - } - - *newTask->mutable_dq_graph() = task.Internal.dq_graph(); - newTask->set_dq_graph_index(task.Internal.dq_graph_index()); - *newTask->mutable_dq_graph_compressed() = task.Internal.dq_graph_compressed(); - - *newTask->mutable_result_set_meta() = task.Query.result_set_meta(); - newTask->set_scope(task.Scope); - *newTask->mutable_resources() = task.Internal.resources(); - - newTask->set_execution_id(task.Internal.execution_id()); - newTask->set_operation_id(task.Internal.operation_id()); - *newTask->mutable_compute_connection() = task.Internal.compute_connection(); - *newTask->mutable_result_ttl() = task.Internal.result_ttl(); - *newTask->mutable_parameters() = task.Query.content().parameters(); - } - + FillGetTaskResult(result, std::get<0>(*response)); return result; }; auto success = SendResponse diff --git a/ydb/core/fq/libs/control_plane_storage/internal/task_ping.cpp b/ydb/core/fq/libs/control_plane_storage/internal/task_ping.cpp index 976cceeca87b..b98a0e6b549d 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/task_ping.cpp +++ b/ydb/core/fq/libs/control_plane_storage/internal/task_ping.cpp @@ -4,6 +4,7 @@ #include #include +#include #include #include @@ -45,36 +46,15 @@ THashMap DeserializeFlatStats(const google::protobuf::RepeatedPtrF } -struct TPingTaskParams { - TString Query; - TParams Params; - const std::function(const TVector&)> Prepare; - std::shared_ptr> MeteringRecords; -}; - -struct TFinalStatus { - FederatedQuery::QueryMeta::ComputeStatus Status = FederatedQuery::QueryMeta::COMPUTE_STATUS_UNSPECIFIED; - NYql::NDqProto::StatusIds::StatusCode StatusCode = NYql::NDqProto::StatusIds::UNSPECIFIED; - FederatedQuery::QueryContent::QueryType QueryType = FederatedQuery::QueryContent::QUERY_TYPE_UNSPECIFIED; - NYql::TIssues Issues; - NYql::TIssues TransientIssues; - StatsValuesList FinalStatistics; - TString CloudId; - TString JobId; -}; - -TPingTaskParams ConstructHardPingTask( +TYdbControlPlaneStorageActor::TPingTaskParams TYdbControlPlaneStorageActor::ConstructHardPingTask( const Fq::Private::PingTaskRequest& request, std::shared_ptr response, - const TString& tablePathPrefix, const TDuration& automaticQueriesTtl, const TDuration& taskLeaseTtl, - const THashMap& retryPolicies, ::NMonitoring::TDynamicCounterPtr rootCounters, - uint64_t maxRequestSize, bool dumpRawStatistics, const std::shared_ptr& finalStatus, - const TRequestCommonCountersPtr& commonCounters) { + const std::shared_ptr& finalStatus, const TRequestCommonCountersPtr& commonCounters) const { auto scope = request.scope(); auto query_id = request.query_id().value(); - auto counters = rootCounters->GetSubgroup("scope", scope)->GetSubgroup("query_id", query_id); + auto counters = Counters.Counters->GetSubgroup("scope", scope)->GetSubgroup("query_id", query_id); - TSqlQueryBuilder readQueryBuilder(tablePathPrefix, "HardPingTask(read)"); + TSqlQueryBuilder readQueryBuilder(YdbConnection->TablePathPrefix, "HardPingTask(read)"); readQueryBuilder.AddString("tenant", request.tenant()); readQueryBuilder.AddString("scope", scope); readQueryBuilder.AddString("query_id", query_id); @@ -146,297 +126,12 @@ TPingTaskParams ConstructHardPingTask( ); } - TMaybe queryStatus; - if (request.status() != FederatedQuery::QueryMeta::COMPUTE_STATUS_UNSPECIFIED) { - queryStatus = request.status(); - } - TMaybe issues; - if (request.issues().size() > 0) { - NYql::TIssues requestIssues; - NYql::IssuesFromMessage(request.issues(), requestIssues); - issues = requestIssues; - } - TMaybe transientIssues; - if (request.transient_issues().size() > 0) { - NYql::TIssues requestTransientIssues; - NYql::IssuesFromMessage(request.transient_issues(), requestTransientIssues); - transientIssues = requestTransientIssues; - } // running query us locked for lease period - TDuration backoff = taskLeaseTtl; - - if (request.resign_query()) { - if (request.status_code() == NYql::NDqProto::StatusIds::UNSPECIFIED && internal.pending_status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { - request.set_status_code(internal.pending_status_code()); - internal.clear_pending_status_code(); - internal.clear_execution_id(); - internal.clear_operation_id(); - } - - TRetryPolicyItem policy(0, 0, TDuration::Seconds(1), TDuration::Zero()); - auto it = retryPolicies.find(request.status_code()); - auto policyFound = it != retryPolicies.end(); - if (policyFound) { - policy = it->second; - } - - auto now = TInstant::Now(); - auto executionDeadline = TInstant::Max(); - - auto submittedAt = NProtoInterop::CastFromProto(query.meta().submitted_at()); - auto executionTtl = NProtoInterop::CastFromProto(internal.execution_ttl()); - if (submittedAt && executionTtl) { - executionDeadline = submittedAt + executionTtl; - } - - if (retryLimiter.UpdateOnRetry(now, policy) && now < executionDeadline) { - queryStatus.Clear(); - // failing query is throttled for backoff period - backoff = policy.BackoffPeriod * (retryLimiter.RetryRate + 1); - owner = ""; - if (!transientIssues) { - transientIssues.ConstructInPlace(); - } - TStringBuilder builder; - builder << "Query failed with code " << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()) - << " and will be restarted (RetryCount: " << retryLimiter.RetryCount << ")" - << " at " << now; - transientIssues->AddIssue(NYql::TIssue(builder)); - } else { - // failure query should be processed instantly - queryStatus = FederatedQuery::QueryMeta::FAILING; - backoff = TDuration::Zero(); - TStringBuilder builder; - builder << "Query failed with code " << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()); - if (policy.RetryCount) { - builder << " (" << retryLimiter.LastError << ")"; - } - builder << " at " << now; - - // in case of problems with finalization, do not change the issues - if (query.meta().status() == FederatedQuery::QueryMeta::FAILING || query.meta().status() == FederatedQuery::QueryMeta::ABORTING_BY_SYSTEM || query.meta().status() == FederatedQuery::QueryMeta::ABORTING_BY_USER) { - if (issues) { - transientIssues->AddIssues(*issues); - } - transientIssues->AddIssue(NYql::TIssue(builder)); - } else { - if (!issues) { - issues.ConstructInPlace(); - } - auto issue = NYql::TIssue(builder); - if (query.issue().size() > 0 && request.issues().empty()) { - NYql::TIssues queryIssues; - NYql::IssuesFromMessage(query.issue(), queryIssues); - for (auto& subIssue : queryIssues) { - issue.AddSubIssue(MakeIntrusive(subIssue)); - } - } - if (transientIssues) { - for (auto& subIssue : *transientIssues) { - issue.AddSubIssue(MakeIntrusive(subIssue)); - } - transientIssues.Clear(); - } - issues->AddIssue(issue); - } - } - CPS_LOG_AS_D(*actorSystem, "PingTaskRequest (resign): " << (!policyFound ? " DEFAULT POLICY" : "") << (owner ? " FAILURE " : " ") << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()) << " " << retryLimiter.RetryCount << " " << retryLimiter.RetryCounterUpdatedAt << " " << backoff); - } - - if (queryStatus) { - query.mutable_meta()->set_status(*queryStatus); - job.mutable_query_meta()->set_status(*queryStatus); - } - - if (request.status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { - internal.set_status_code(request.status_code()); - } - - if (request.pending_status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { - internal.set_pending_status_code(request.pending_status_code()); - } - - if (issues) { - NYql::IssuesToMessage(*issues, query.mutable_issue()); - NYql::IssuesToMessage(*issues, job.mutable_issue()); - } - - if (transientIssues) { - AddTransientIssues(query.mutable_transient_issue(), std::move(*transientIssues)); - } - - if (request.internal_issues().size()) { - *internal.mutable_internal_issue() = request.internal_issues(); - } - - if (request.statistics()) { - TString statistics = request.statistics(); - if (request.flat_stats_size() == 0) { - internal.clear_statistics(); - // TODO: remove once V1 and V2 stats go the same way - PackStatisticsToProtobuf(*internal.mutable_statistics(), statistics, TInstant::Now() - NProtoInterop::CastFromProto(job.meta().created_at())); - } - - // global dumpRawStatistics will be removed with YQv1 - if (!dumpRawStatistics && !request.dump_raw_statistics()) { - try { - statistics = GetPrettyStatistics(statistics); - } catch (const std::exception&) { - // LOG_AS? - CPS_LOG_E("Error on statistics prettification: " << CurrentExceptionMessage()); - } - } - *query.mutable_statistics()->mutable_json() = statistics; - *job.mutable_statistics()->mutable_json() = statistics; - } - - if (request.current_load()) { - internal.set_current_load(request.current_load()); - } - - if (request.timeline()) { - internal.set_timeline(request.timeline()); - } - - if (request.flat_stats_size() != 0) { - internal.clear_statistics(); - auto stats = DeserializeFlatStats(request.flat_stats()); - PackStatisticsToProtobuf(*internal.mutable_statistics(), stats, TInstant::Now() - NProtoInterop::CastFromProto(job.meta().created_at())); - } - - if (!request.result_set_meta().empty()) { - // we will overwrite result_set_meta's COMPLETELY - *query.mutable_result_set_meta() = request.result_set_meta(); - *job.mutable_result_set_meta() = request.result_set_meta(); - } - - if (request.ast()) { - query.mutable_ast()->set_data(request.ast()); - job.mutable_ast()->set_data(request.ast()); - } - - if (request.plan()) { - query.mutable_plan()->set_json(request.plan()); - job.mutable_plan()->set_json(request.plan()); - } - - if (request.ast_compressed().data()) { - internal.mutable_ast_compressed()->set_method(request.ast_compressed().method()); - internal.mutable_ast_compressed()->set_data(request.ast_compressed().data()); - // todo: keep AST compressed in JobInternal - // job.mutable_ast()->set_data(request.ast()); - } - - if (request.plan_compressed().data()) { - internal.mutable_plan_compressed()->set_method(request.plan_compressed().method()); - internal.mutable_plan_compressed()->set_data(request.plan_compressed().data()); - // todo: keep plan compressed in JobInternal - // job.mutable_plan()->set_json(request.plan()); - } - - if (request.has_started_at()) { - *query.mutable_meta()->mutable_started_at() = request.started_at(); - *job.mutable_query_meta()->mutable_started_at() = request.started_at(); - } - - if (request.has_finished_at()) { - *query.mutable_meta()->mutable_finished_at() = request.finished_at(); - *job.mutable_query_meta()->mutable_finished_at() = request.finished_at(); - if (!query.meta().has_started_at()) { - *query.mutable_meta()->mutable_started_at() = request.finished_at(); - *job.mutable_query_meta()->mutable_started_at() = request.finished_at(); - } - } - - TInstant expireAt = TInstant::Now() + automaticQueriesTtl; - if (IsTerminalStatus(query.meta().status()) && query.content().automatic()) { - *query.mutable_meta()->mutable_expire_at() = NProtoInterop::CastToProto(expireAt); - *job.mutable_query_meta()->mutable_expire_at() = NProtoInterop::CastToProto(expireAt); - *job.mutable_expire_at() = NProtoInterop::CastToProto(expireAt); - } - - if (query.meta().status() == FederatedQuery::QueryMeta::COMPLETED) { - *query.mutable_meta()->mutable_result_expire_at() = request.deadline(); - } - - if (request.state_load_mode()) { - internal.set_state_load_mode(request.state_load_mode()); - if (request.state_load_mode() == FederatedQuery::FROM_LAST_CHECKPOINT) { // Saved checkpoint - query.mutable_meta()->set_has_saved_checkpoints(true); - } - } - - if (request.has_disposition()) { - *internal.mutable_disposition() = request.disposition(); - } - - if (request.status() && IsTerminalStatus(request.status())) { - internal.clear_created_topic_consumers(); - // internal.clear_dq_graph(); keep for debug - internal.clear_dq_graph_index(); - // internal.clear_execution_id(); keep for debug - // internal.clear_operation_id(); keep for debug - } - - if (!request.created_topic_consumers().empty()) { - std::set mergedConsumers; - for (auto&& c : *internal.mutable_created_topic_consumers()) { - mergedConsumers.emplace(std::move(c)); - } - for (const auto& c : request.created_topic_consumers()) { - mergedConsumers.emplace(c); - } - internal.clear_created_topic_consumers(); - for (auto&& c : mergedConsumers) { - *internal.add_created_topic_consumers() = std::move(c); - } - } - - if (!request.execution_id().empty()) { - internal.set_execution_id(request.execution_id()); - } - - if (!request.operation_id().empty()) { - internal.set_operation_id(request.operation_id()); - } + TDuration backoff = Config->TaskLeaseTtl; + TInstant expireAt = TInstant::Now() + Config->AutomaticQueriesTtl; + UpdateTaskInfo(actorSystem, request, finalStatus, query, internal, job, owner, retryLimiter, backoff, expireAt); - if (!request.dq_graph().empty()) { - *internal.mutable_dq_graph() = request.dq_graph(); - } - - if (!request.dq_graph_compressed().empty()) { - *internal.mutable_dq_graph_compressed() = request.dq_graph_compressed(); - } - - if (request.dq_graph_index()) { - internal.set_dq_graph_index(request.dq_graph_index()); - } - - if (request.has_resources()) { - *internal.mutable_resources() = request.resources(); - } - - if (job.ByteSizeLong() > maxRequestSize) { - ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Job proto exceeded the size limit: " << job.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(job).ToString(); - } - - if (query.ByteSizeLong() > maxRequestSize) { - ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Query proto exceeded the size limit: " << query.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(query).ToString(); - } - - if (internal.ByteSizeLong() > maxRequestSize) { - ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "QueryInternal proto exceeded the size limit: " << internal.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(internal).ToString(); - } - - finalStatus->Status = query.meta().status(); - finalStatus->QueryType = query.content().type(); - finalStatus->StatusCode = internal.status_code(); - finalStatus->CloudId = internal.cloud_id(); - finalStatus->JobId = jobId; - NYql::IssuesFromMessage(query.issue(), finalStatus->Issues); - NYql::IssuesFromMessage(query.transient_issue(), finalStatus->TransientIssues); - - TSqlQueryBuilder writeQueryBuilder(tablePathPrefix, "HardPingTask(write)"); + TSqlQueryBuilder writeQueryBuilder(YdbConnection->TablePathPrefix, "HardPingTask(write)"); writeQueryBuilder.AddString("tenant", request.tenant()); writeQueryBuilder.AddString("scope", request.scope()); writeQueryBuilder.AddString("job_id", jobId); @@ -531,18 +226,7 @@ TPingTaskParams ConstructHardPingTask( // YQv2 may not provide statistics with terminal status, use saved one statistics = query.statistics().json(); } - finalStatus->FinalStatistics = ExtractStatisticsFromProtobuf(internal.statistics()); - finalStatus->FinalStatistics.push_back(std::make_pair("IsAutomatic", query.content().automatic())); - if (query.content().name().Contains("DataLens YQ query")) { - finalStatus->FinalStatistics.push_back(std::make_pair("IsDataLens", 1)); - } else if (query.content().name().Contains("Audit-trails")) { - finalStatus->FinalStatistics.push_back(std::make_pair("IsAuditTrails", 1)); - } else if (query.content().name().Contains("Query from YDB SDK")) { - finalStatus->FinalStatistics.push_back(std::make_pair("IsSDK", 1)); - } - finalStatus->FinalStatistics.push_back(std::make_pair("RetryCount", retryLimiter.RetryCount)); - finalStatus->FinalStatistics.push_back(std::make_pair("RetryRate", retryLimiter.RetryRate * 100)); - finalStatus->FinalStatistics.push_back(std::make_pair("Load", internal.current_load())); + FillQueryStatistics(finalStatus, query, internal, retryLimiter); auto records = GetMeteringRecords(statistics, isBillable, jobId, request.scope(), HostName()); meteringRecords->swap(records); @@ -558,10 +242,10 @@ TPingTaskParams ConstructHardPingTask( return {readQuery.Sql, readQuery.Params, prepareParams, meteringRecords}; } -TPingTaskParams ConstructSoftPingTask( +TYdbControlPlaneStorageActor::TPingTaskParams TYdbControlPlaneStorageActor::ConstructSoftPingTask( const Fq::Private::PingTaskRequest& request, std::shared_ptr response, - const TString& tablePathPrefix, const TDuration& taskLeaseTtl, const TRequestCommonCountersPtr& commonCounters) { - TSqlQueryBuilder readQueryBuilder(tablePathPrefix, "SoftPingTask(read)"); + const TRequestCommonCountersPtr& commonCounters) const { + TSqlQueryBuilder readQueryBuilder(YdbConnection->TablePathPrefix, "SoftPingTask(read)"); readQueryBuilder.AddString("tenant", request.tenant()); readQueryBuilder.AddString("scope", request.scope()); readQueryBuilder.AddString("query_id", request.query_id().value()); @@ -603,11 +287,11 @@ TPingTaskParams ConstructSoftPingTask( } } - TInstant ttl = TInstant::Now() + taskLeaseTtl; + TInstant ttl = TInstant::Now() + Config->TaskLeaseTtl; response->set_action(internal.action()); *response->mutable_expired_at() = google::protobuf::util::TimeUtil::MillisecondsToTimestamp(ttl.MilliSeconds()); - TSqlQueryBuilder writeQueryBuilder(tablePathPrefix, "SoftPingTask(write)"); + TSqlQueryBuilder writeQueryBuilder(YdbConnection->TablePathPrefix, "SoftPingTask(write)"); writeQueryBuilder.AddTimestamp("now", TInstant::Now()); writeQueryBuilder.AddTimestamp("ttl", ttl); writeQueryBuilder.AddString("tenant", request.tenant()); @@ -626,6 +310,330 @@ TPingTaskParams ConstructSoftPingTask( return {readQuery.Sql, readQuery.Params, prepareParams, std::shared_ptr>{}}; } +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvPingTaskRequest::TPtr& ev) const { + const Fq::Private::PingTaskRequest& request = ev->Get()->Request; + + NYql::TIssues issues = ValidatePingTask(request.scope(), request.query_id().value(), request.owner_id(), NProtoInterop::CastFromProto(request.deadline()), Config->ResultSetsTtl); + + const auto tenantInfo = ev->Get()->TenantInfo; + if (tenantInfo && tenantInfo->TenantState.Value(request.tenant(), TenantState::Active) == TenantState::Idle) { + issues.AddIssue("Tenant is idle, no processing is allowed"); + } + + return issues; +} + +void TControlPlaneStorageBase::UpdateTaskInfo( + NActors::TActorSystem* actorSystem, Fq::Private::PingTaskRequest& request, const std::shared_ptr& finalStatus, FederatedQuery::Query& query, + FederatedQuery::Internal::QueryInternal& internal, FederatedQuery::Job& job, TString& owner, + TRetryLimiter& retryLimiter, TDuration& backoff, TInstant& expireAt) const +{ + TMaybe queryStatus; + if (request.status() != FederatedQuery::QueryMeta::COMPUTE_STATUS_UNSPECIFIED) { + queryStatus = request.status(); + } + TMaybe issues; + if (request.issues().size() > 0) { + NYql::TIssues requestIssues; + NYql::IssuesFromMessage(request.issues(), requestIssues); + issues = requestIssues; + } + TMaybe transientIssues; + if (request.transient_issues().size() > 0) { + NYql::TIssues requestTransientIssues; + NYql::IssuesFromMessage(request.transient_issues(), requestTransientIssues); + transientIssues = requestTransientIssues; + } + + if (request.resign_query()) { + if (request.status_code() == NYql::NDqProto::StatusIds::UNSPECIFIED && internal.pending_status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { + request.set_status_code(internal.pending_status_code()); + internal.clear_pending_status_code(); + internal.clear_execution_id(); + internal.clear_operation_id(); + } + + TRetryPolicyItem policy(0, 0, TDuration::Seconds(1), TDuration::Zero()); + auto it = Config->RetryPolicies.find(request.status_code()); + auto policyFound = it != Config->RetryPolicies.end(); + if (policyFound) { + policy = it->second; + } + + auto now = TInstant::Now(); + auto executionDeadline = TInstant::Max(); + + auto submittedAt = NProtoInterop::CastFromProto(query.meta().submitted_at()); + auto executionTtl = NProtoInterop::CastFromProto(internal.execution_ttl()); + if (submittedAt && executionTtl) { + executionDeadline = submittedAt + executionTtl; + } + + if (retryLimiter.UpdateOnRetry(now, policy) && now < executionDeadline) { + queryStatus.Clear(); + // failing query is throttled for backoff period + backoff = policy.BackoffPeriod * (retryLimiter.RetryRate + 1); + owner = ""; + if (!transientIssues) { + transientIssues.ConstructInPlace(); + } + TStringBuilder builder; + builder << "Query failed with code " << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()) + << " and will be restarted (RetryCount: " << retryLimiter.RetryCount << ")" + << " at " << now; + transientIssues->AddIssue(NYql::TIssue(builder)); + } else { + // failure query should be processed instantly + queryStatus = FederatedQuery::QueryMeta::FAILING; + backoff = TDuration::Zero(); + TStringBuilder builder; + builder << "Query failed with code " << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()); + if (policy.RetryCount) { + builder << " (" << retryLimiter.LastError << ")"; + } + builder << " at " << now; + + // in case of problems with finalization, do not change the issues + if (query.meta().status() == FederatedQuery::QueryMeta::FAILING || query.meta().status() == FederatedQuery::QueryMeta::ABORTING_BY_SYSTEM || query.meta().status() == FederatedQuery::QueryMeta::ABORTING_BY_USER) { + if (issues) { + transientIssues->AddIssues(*issues); + } + transientIssues->AddIssue(NYql::TIssue(builder)); + } else { + if (!issues) { + issues.ConstructInPlace(); + } + auto issue = NYql::TIssue(builder); + if (query.issue().size() > 0 && request.issues().empty()) { + NYql::TIssues queryIssues; + NYql::IssuesFromMessage(query.issue(), queryIssues); + for (auto& subIssue : queryIssues) { + issue.AddSubIssue(MakeIntrusive(subIssue)); + } + } + if (transientIssues) { + for (auto& subIssue : *transientIssues) { + issue.AddSubIssue(MakeIntrusive(subIssue)); + } + transientIssues.Clear(); + } + issues->AddIssue(issue); + } + } + CPS_LOG_AS_D(*actorSystem, "PingTaskRequest (resign): " << (!policyFound ? " DEFAULT POLICY" : "") << (owner ? " FAILURE " : " ") << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()) << " " << retryLimiter.RetryCount << " " << retryLimiter.RetryCounterUpdatedAt << " " << backoff); + } + + if (queryStatus) { + query.mutable_meta()->set_status(*queryStatus); + job.mutable_query_meta()->set_status(*queryStatus); + } + + if (request.status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { + internal.set_status_code(request.status_code()); + } + + if (request.pending_status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { + internal.set_pending_status_code(request.pending_status_code()); + } + + if (issues) { + NYql::IssuesToMessage(*issues, query.mutable_issue()); + NYql::IssuesToMessage(*issues, job.mutable_issue()); + } + + if (transientIssues) { + AddTransientIssues(query.mutable_transient_issue(), std::move(*transientIssues)); + } + + if (request.internal_issues().size()) { + *internal.mutable_internal_issue() = request.internal_issues(); + } + + if (request.statistics()) { + TString statistics = request.statistics(); + if (request.flat_stats_size() == 0) { + internal.clear_statistics(); + // TODO: remove once V1 and V2 stats go the same way + PackStatisticsToProtobuf(*internal.mutable_statistics(), statistics, TInstant::Now() - NProtoInterop::CastFromProto(job.meta().created_at())); + } + + // global dumpRawStatistics will be removed with YQv1 + if (!Config->Proto.GetDumpRawStatistics() && !request.dump_raw_statistics()) { + try { + statistics = GetPrettyStatistics(statistics); + } catch (const std::exception&) { + CPS_LOG_AS_E(*actorSystem, "Error on statistics prettification: " << CurrentExceptionMessage()); + } + } + *query.mutable_statistics()->mutable_json() = statistics; + *job.mutable_statistics()->mutable_json() = statistics; + } + + if (request.current_load()) { + internal.set_current_load(request.current_load()); + } + + if (request.timeline()) { + internal.set_timeline(request.timeline()); + } + + if (request.flat_stats_size() != 0) { + internal.clear_statistics(); + auto stats = DeserializeFlatStats(request.flat_stats()); + PackStatisticsToProtobuf(*internal.mutable_statistics(), stats, TInstant::Now() - NProtoInterop::CastFromProto(job.meta().created_at())); + } + + if (!request.result_set_meta().empty()) { + // we will overwrite result_set_meta's COMPLETELY + *query.mutable_result_set_meta() = request.result_set_meta(); + *job.mutable_result_set_meta() = request.result_set_meta(); + } + + if (request.ast()) { + query.mutable_ast()->set_data(request.ast()); + job.mutable_ast()->set_data(request.ast()); + } + + if (request.plan()) { + query.mutable_plan()->set_json(request.plan()); + job.mutable_plan()->set_json(request.plan()); + } + + if (request.ast_compressed().data()) { + internal.mutable_ast_compressed()->set_method(request.ast_compressed().method()); + internal.mutable_ast_compressed()->set_data(request.ast_compressed().data()); + // todo: keep AST compressed in JobInternal + // job.mutable_ast()->set_data(request.ast()); + } + + if (request.plan_compressed().data()) { + internal.mutable_plan_compressed()->set_method(request.plan_compressed().method()); + internal.mutable_plan_compressed()->set_data(request.plan_compressed().data()); + // todo: keep plan compressed in JobInternal + // job.mutable_plan()->set_json(request.plan()); + } + + if (request.has_started_at()) { + *query.mutable_meta()->mutable_started_at() = request.started_at(); + *job.mutable_query_meta()->mutable_started_at() = request.started_at(); + } + + if (request.has_finished_at()) { + *query.mutable_meta()->mutable_finished_at() = request.finished_at(); + *job.mutable_query_meta()->mutable_finished_at() = request.finished_at(); + if (!query.meta().has_started_at()) { + *query.mutable_meta()->mutable_started_at() = request.finished_at(); + *job.mutable_query_meta()->mutable_started_at() = request.finished_at(); + } + } + + if (IsTerminalStatus(query.meta().status()) && query.content().automatic()) { + *query.mutable_meta()->mutable_expire_at() = NProtoInterop::CastToProto(expireAt); + *job.mutable_query_meta()->mutable_expire_at() = NProtoInterop::CastToProto(expireAt); + *job.mutable_expire_at() = NProtoInterop::CastToProto(expireAt); + } + + if (query.meta().status() == FederatedQuery::QueryMeta::COMPLETED) { + *query.mutable_meta()->mutable_result_expire_at() = request.deadline(); + } + + if (request.state_load_mode()) { + internal.set_state_load_mode(request.state_load_mode()); + if (request.state_load_mode() == FederatedQuery::FROM_LAST_CHECKPOINT) { // Saved checkpoint + query.mutable_meta()->set_has_saved_checkpoints(true); + } + } + + if (request.has_disposition()) { + *internal.mutable_disposition() = request.disposition(); + } + + if (request.status() && IsTerminalStatus(request.status())) { + internal.clear_created_topic_consumers(); + // internal.clear_dq_graph(); keep for debug + internal.clear_dq_graph_index(); + // internal.clear_execution_id(); keep for debug + // internal.clear_operation_id(); keep for debug + } + + if (!request.created_topic_consumers().empty()) { + std::set mergedConsumers; + for (auto&& c : *internal.mutable_created_topic_consumers()) { + mergedConsumers.emplace(std::move(c)); + } + for (const auto& c : request.created_topic_consumers()) { + mergedConsumers.emplace(c); + } + internal.clear_created_topic_consumers(); + for (auto&& c : mergedConsumers) { + *internal.add_created_topic_consumers() = std::move(c); + } + } + + if (!request.execution_id().empty()) { + internal.set_execution_id(request.execution_id()); + } + + if (!request.operation_id().empty()) { + internal.set_operation_id(request.operation_id()); + } + + if (!request.dq_graph().empty()) { + *internal.mutable_dq_graph() = request.dq_graph(); + } + + if (!request.dq_graph_compressed().empty()) { + *internal.mutable_dq_graph_compressed() = request.dq_graph_compressed(); + } + + if (request.dq_graph_index()) { + internal.set_dq_graph_index(request.dq_graph_index()); + } + + if (request.has_resources()) { + *internal.mutable_resources() = request.resources(); + } + + const auto maxRequestSize = Config->Proto.GetMaxRequestSize(); + if (job.ByteSizeLong() > maxRequestSize) { + ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Job proto exceeded the size limit: " << job.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(job).ToString(); + } + + if (query.ByteSizeLong() > maxRequestSize) { + ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Query proto exceeded the size limit: " << query.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(query).ToString(); + } + + if (internal.ByteSizeLong() > maxRequestSize) { + ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "QueryInternal proto exceeded the size limit: " << internal.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(internal).ToString(); + } + + finalStatus->Status = query.meta().status(); + finalStatus->QueryType = query.content().type(); + finalStatus->StatusCode = internal.status_code(); + finalStatus->CloudId = internal.cloud_id(); + finalStatus->JobId = job.meta().id(); + NYql::IssuesFromMessage(query.issue(), finalStatus->Issues); + NYql::IssuesFromMessage(query.transient_issue(), finalStatus->TransientIssues); +} + +void TControlPlaneStorageBase::FillQueryStatistics( + const std::shared_ptr& finalStatus, const FederatedQuery::Query& query, + const FederatedQuery::Internal::QueryInternal& internal, const TRetryLimiter& retryLimiter) const +{ + finalStatus->FinalStatistics = ExtractStatisticsFromProtobuf(internal.statistics()); + finalStatus->FinalStatistics.push_back(std::make_pair("IsAutomatic", query.content().automatic())); + if (query.content().name().Contains("DataLens YQ query")) { + finalStatus->FinalStatistics.push_back(std::make_pair("IsDataLens", 1)); + } else if (query.content().name().Contains("Audit-trails")) { + finalStatus->FinalStatistics.push_back(std::make_pair("IsAuditTrails", 1)); + } else if (query.content().name().Contains("Query from YDB SDK")) { + finalStatus->FinalStatistics.push_back(std::make_pair("IsSDK", 1)); + } + finalStatus->FinalStatistics.push_back(std::make_pair("RetryCount", retryLimiter.RetryCount)); + finalStatus->FinalStatistics.push_back(std::make_pair("RetryRate", retryLimiter.RetryRate * 100)); + finalStatus->FinalStatistics.push_back(std::make_pair("Load", internal.current_load())); +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvPingTaskRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -635,20 +643,10 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvPingTaskReq requestCounters.IncInFly(); requestCounters.Common->RequestBytes->Add(ev->Get()->GetByteSize()); const TString queryId = request.query_id().value(); - const TString owner = request.owner_id(); - const TInstant deadline = NProtoInterop::CastFromProto(request.deadline()); - const TString tenant = request.tenant(); CPS_LOG_T("PingTaskRequest: {" << request.DebugString() << "}"); - NYql::TIssues issues = ValidatePingTask(scope, queryId, owner, deadline, Config->ResultSetsTtl); - - auto tenantInfo = ev->Get()->TenantInfo; - if (tenantInfo && tenantInfo->TenantState.Value(tenant, TenantState::Active) == TenantState::Idle) { - issues.AddIssue("Tenant is idle, no processing is allowed"); - } - - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_W("PingTaskRequest: {" << request.DebugString() << "} validation FAILED: " << issues.ToOneLineString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); @@ -660,10 +658,8 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvPingTaskReq std::shared_ptr finalStatus = std::make_shared(); auto pingTaskParams = DoesPingTaskUpdateQueriesTable(request) ? - ConstructHardPingTask(request, response, YdbConnection->TablePathPrefix, Config->AutomaticQueriesTtl, - Config->TaskLeaseTtl, Config->RetryPolicies, Counters.Counters, Config->Proto.GetMaxRequestSize(), - Config->Proto.GetDumpRawStatistics(), finalStatus, requestCounters.Common) : - ConstructSoftPingTask(request, response, YdbConnection->TablePathPrefix, Config->TaskLeaseTtl, requestCounters.Common); + ConstructHardPingTask(request, response, finalStatus, requestCounters.Common) : + ConstructSoftPingTask(request, response, requestCounters.Common); auto debugInfo = Config->Proto.GetEnableDebugMode() ? std::make_shared() : TDebugInfoPtr{}; auto result = ReadModifyWrite(pingTaskParams.Query, pingTaskParams.Params, pingTaskParams.Prepare, requestCounters, debugInfo); auto prepare = [response] { return *response; }; @@ -696,7 +692,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvPingTaskReq }); } -void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvFinalStatusReport::TPtr& ev) { +void TControlPlaneStorageBase::Handle(TEvControlPlaneStorage::TEvFinalStatusReport::TPtr& ev) { const auto& event = *ev->Get(); if (!IsTerminalStatus(event.Status)) { return; @@ -719,11 +715,10 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvFinalStatus Counters.GetFinalStatusCounters(event.CloudId, event.Scope)->IncByStatus(event.Status); - Statistics statistics{event.Statistics}; + TStatistics statistics{event.Statistics}; LOG_YQ_AUDIT_SERVICE_INFO("FinalStatus: cloud id: [" << event.CloudId << "], scope: [" << event.Scope << "], query id: [" << event.QueryId << "], job id: [" << event.JobId << "], query type: [" << FederatedQuery::QueryContent::QueryType_Name(event.QueryType) << "], " << statistics << ", " << "status: " << FederatedQuery::QueryMeta::ComputeStatus_Name(event.Status)); } - } // NFq diff --git a/ydb/core/fq/libs/control_plane_storage/internal/task_result_write.cpp b/ydb/core/fq/libs/control_plane_storage/internal/task_result_write.cpp index c2a39dbfe05b..469f8eab3230 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/task_result_write.cpp +++ b/ydb/core/fq/libs/control_plane_storage/internal/task_result_write.cpp @@ -1,7 +1,19 @@ #include "utils.h" +#include + namespace NFq { +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvWriteResultDataRequest::TPtr& ev) const { + const auto& request = ev->Get()->Request; + return ValidateWriteResultData( + request.result_id().value(), + request.result_set(), + NProtoInterop::CastFromProto(request.deadline()), + Config->ResultSetsTtl + ); +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvWriteResultDataRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -17,11 +29,9 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvWriteResult const Ydb::ResultSet& resultSet = request.result_set(); const int byteSize = resultSet.ByteSize(); - CPS_LOG_T("WriteResultDataRequest: " << resultId << " " << resultSetId << " " << startRowId << " " << resultSet.ByteSize() << " " << deadline); - NYql::TIssues issues = ValidateWriteResultData(resultId, resultSet, deadline, Config->ResultSetsTtl); - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_D("WriteResultDataRequest, validation failed: " << resultId << " " << resultSetId << " " << startRowId << " " << resultSet.DebugString() << " " << deadline << " error: " << issues.ToString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); diff --git a/ydb/core/fq/libs/control_plane_storage/internal/utils.cpp b/ydb/core/fq/libs/control_plane_storage/internal/utils.cpp index 3c9281e86eda..ab7c8a3617e0 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/utils.cpp +++ b/ydb/core/fq/libs/control_plane_storage/internal/utils.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -143,15 +144,15 @@ std::vector GetMeteringRecords(const TString& statistics, bool billable } auto now = Now(); - result.emplace_back(TBillRecord() + result.emplace_back(NKikimr::TBillRecord() .Id(jobId + "_i") .Schema("yq.ingress.v1") .FolderId(TScope(scope).ParseFolder()) .SourceWt(now) .SourceId(sourceId) - .Usage(TBillRecord::TUsage() - .Type(TBillRecord::TUsage::EType::Delta) - .Unit(TBillRecord::TUsage::EUnit::MByte) + .Usage(NKikimr::TBillRecord::TUsage() + .Type(NKikimr::TBillRecord::TUsage::EType::Delta) + .Unit(NKikimr::TBillRecord::TUsage::EUnit::MByte) .Quantity(ingressMBytes) .Start(now) .Finish(now) @@ -311,8 +312,8 @@ void PackStatisticsToProtobuf(google::protobuf::RepeatedPtrField& statsProto) { - StatsValuesList statPairs; +TStatsValuesList ExtractStatisticsFromProtobuf(const google::protobuf::RepeatedPtrField& statsProto) { + TStatsValuesList statPairs; statPairs.reserve(statsProto.size()); for (const auto& stat : statsProto) { statPairs.emplace_back(stat.name(), stat.value()); @@ -320,7 +321,7 @@ StatsValuesList ExtractStatisticsFromProtobuf(const google::protobuf::RepeatedPt return statPairs; } -TStringBuilder& operator<<(TStringBuilder& builder, const Statistics& statistics) { +TStringBuilder& operator<<(TStringBuilder& builder, const TStatistics& statistics) { bool first = true; builder << '{'; for (const auto& [field, value] : statistics.Stats) { diff --git a/ydb/core/fq/libs/control_plane_storage/internal/utils.h b/ydb/core/fq/libs/control_plane_storage/internal/utils.h index 8a35170c094c..d52cf84ec42f 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/utils.h +++ b/ydb/core/fq/libs/control_plane_storage/internal/utils.h @@ -3,11 +3,13 @@ #include #include +#include #include #include -#include +#include +#include #include namespace NFq { @@ -41,17 +43,17 @@ void PackStatisticsToProtobuf(google::protobuf::RepeatedPtrField>; +using TStatsValuesList = std::vector>; -StatsValuesList ExtractStatisticsFromProtobuf(const google::protobuf::RepeatedPtrField& statsProto); +TStatsValuesList ExtractStatisticsFromProtobuf(const google::protobuf::RepeatedPtrField& statsProto); -struct Statistics { +struct TStatistics { operator bool() const noexcept { return !Stats.empty(); } - const StatsValuesList& Stats; + const TStatsValuesList& Stats; }; -TStringBuilder& operator<<(TStringBuilder& builder, const Statistics& statistics); +TStringBuilder& operator<<(TStringBuilder& builder, const TStatistics& statistics); void AddTransientIssues(::google::protobuf::RepeatedPtrField< ::Ydb::Issue::IssueMessage>* protoIssues, NYql::TIssues&& issues); diff --git a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage.cpp b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage.cpp index 362cb5d7ec6d..a5057a16804c 100644 --- a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage.cpp +++ b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage.cpp @@ -32,6 +32,20 @@ TYdbSdkRetryPolicy::TPtr MakeCreateSchemaRetryPolicy() { } // namespace +TControlPlaneStorageBase::TControlPlaneStorageBase( + const NConfig::TControlPlaneStorageConfig& config, + const NYql::TS3GatewayConfig& s3Config, + const NConfig::TCommonConfig& common, + const NConfig::TComputeConfig& computeConfig, + const ::NMonitoring::TDynamicCounterPtr& counters, + const TString& tenantName) + : TBase(config, s3Config, common, computeConfig) + , Counters(counters, *Config) + , FailedStatusCodeCounters(MakeIntrusive("FinalFailedStatusCode", counters->GetSubgroup("component", "QueryDiagnostic"))) + , TenantName(tenantName) +{ +} + void TYdbControlPlaneStorageActor::Bootstrap() { CPS_LOG_I("Starting ydb control plane storage service. Actor id: " << SelfId()); NLwTraceMonPage::ProbeRegistry().AddProbesList(LWTRACE_GET_PROBES(YQ_CONTROL_PLANE_STORAGE_PROVIDER)); @@ -323,7 +337,7 @@ void TYdbControlPlaneStorageActor::AfterTablesCreated() { // Schedule(TDuration::Zero(), new NActors::TEvents::TEvWakeup()); } -bool TControlPlaneStorageUtils::IsSuperUser(const TString& user) +bool TControlPlaneStorageUtils::IsSuperUser(const TString& user) const { return AnyOf(Config->Proto.GetSuperUsers(), [&user](const auto& superUser) { return superUser == user; diff --git a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_bindings.cpp b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_bindings.cpp index 75f5557ea197..2bc787fc4522 100644 --- a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_bindings.cpp +++ b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_bindings.cpp @@ -8,6 +8,46 @@ namespace NFq { +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvCreateBindingRequest::TPtr& ev) const +{ + NYql::TIssues issues = ValidateBinding(ev); + + const auto& event = *ev->Get(); + const auto& permissions = GetCreateBindingPerimssions(event); + if (event.Request.content().acl().visibility() == FederatedQuery::Acl::SCOPE && !permissions.Check(TPermissions::MANAGE_PUBLIC)) { + issues.AddIssue(MakeErrorIssue(TIssuesIds::ACCESS_DENIED, "Permission denied to create a binding with these parameters. Please receive a permission yq.resources.managePublic")); + } + + return issues; +} + +TPermissions TControlPlaneStorageBase::GetCreateBindingPerimssions(const TEvControlPlaneStorage::TEvCreateBindingRequest& event) const +{ + TPermissions permissions = Config->Proto.GetEnablePermissions() + ? event.Permissions + : TPermissions{TPermissions::MANAGE_PUBLIC}; + if (IsSuperUser(event.User)) { + permissions.SetAll(); + } + return permissions; +} + +std::pair TControlPlaneStorageBase::GetCreateBindingProtos( + const FederatedQuery::CreateBindingRequest& request, const TString& cloudId, const TString& user, TInstant startTime) const +{ + const TString& bindingId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::BINDING); + + FederatedQuery::Binding binding; + FederatedQuery::BindingContent& content = *binding.mutable_content(); + content = request.content(); + *binding.mutable_meta() = CreateCommonMeta(bindingId, user, startTime, InitialRevision); + + FederatedQuery::Internal::BindingInternal bindingInternal; + bindingInternal.set_cloud_id(cloudId); + + return {binding, bindingInternal}; +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateBindingRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -19,28 +59,21 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateBindi requestCounters.Common->RequestBytes->Add(event.GetByteSize()); const TString user = event.User; const TString token = event.Token; - TPermissions permissions = Config->Proto.GetEnablePermissions() - ? event.Permissions - : TPermissions{TPermissions::MANAGE_PUBLIC}; - if (IsSuperUser(user)) { - permissions.SetAll(); - } const FederatedQuery::CreateBindingRequest& request = event.Request; - const TString bindingId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::BINDING); int byteSize = request.ByteSize(); const TString connectionId = request.content().connection_id(); const TString idempotencyKey = request.idempotency_key(); + const auto [binding, bindingInternal] = GetCreateBindingProtos(request, cloudId, user, startTime); + const auto& content = binding.content(); + const TString& bindingId = binding.meta().id(); + CPS_LOG_T(MakeLogPrefix(scope, user, bindingId) << "CreateBindingRequest: " << NKikimr::MaskTicket(token) << " " << request.DebugString()); - NYql::TIssues issues = ValidateBinding(ev); - if (request.content().acl().visibility() == FederatedQuery::Acl::SCOPE && !permissions.Check(TPermissions::MANAGE_PUBLIC)) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::ACCESS_DENIED, "Permission denied to create a binding with these parameters. Please receive a permission yq.resources.managePublic")); - } - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_D(MakeLogPrefix(scope, user, bindingId) << "CreateBindingRequest, validation failed: " << NKikimr::MaskTicket(token) << " " @@ -52,14 +85,6 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateBindi return; } - FederatedQuery::Binding binding; - FederatedQuery::BindingContent& content = *binding.mutable_content(); - content = request.content(); - *binding.mutable_meta() = CreateCommonMeta(bindingId, user, startTime, InitialRevision); - - FederatedQuery::Internal::BindingInternal bindingInternal; - bindingInternal.set_cloud_id(cloudId); - std::shared_ptr>> response = std::make_shared>>(); response->first.set_binding_id(bindingId); response->second.After.ConstructInPlace().CopyFrom(binding); @@ -111,7 +136,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateBindi scope, connectionId, "Connection " + connectionId + " does not exist or permission denied. Please check the id connection or your access rights", - permissions, + GetCreateBindingPerimssions(event), user, content.acl().visibility(), YdbConnection->TablePathPrefix); @@ -122,7 +147,6 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateBindi user, YdbConnection->TablePathPrefix); - TVector validators; if (idempotencyKey) { validators.push_back(CreateIdempotencyKeyValidator(scope, idempotencyKey, response, YdbConnection->TablePathPrefix, requestCounters.Common->ParseProtobufError)); diff --git a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_connections.cpp b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_connections.cpp index 426ce7594725..76ebf8ddd451 100644 --- a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_connections.cpp +++ b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_connections.cpp @@ -30,6 +30,46 @@ void PrepareSensitiveFields(::FederatedQuery::Connection& connection, bool extra } +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvCreateConnectionRequest::TPtr& ev) const +{ + NYql::TIssues issues = ValidateConnection(ev); + + const auto& event = *ev->Get(); + const auto& permissions = GetCreateConnectionPerimssions(event); + if (event.Request.content().acl().visibility() == FederatedQuery::Acl::SCOPE && !permissions.Check(TPermissions::MANAGE_PUBLIC)) { + issues.AddIssue(MakeErrorIssue(TIssuesIds::ACCESS_DENIED, "Permission denied to create a connection with these parameters. Please receive a permission yq.resources.managePublic")); + } + + return issues; +} + +TPermissions TControlPlaneStorageBase::GetCreateConnectionPerimssions(const TEvControlPlaneStorage::TEvCreateConnectionRequest& event) const +{ + TPermissions permissions = Config->Proto.GetEnablePermissions() + ? event.Permissions + : TPermissions{TPermissions::MANAGE_PUBLIC}; + if (IsSuperUser(event.User)) { + permissions.SetAll(); + } + return permissions; +} + +std::pair TControlPlaneStorageBase::GetCreateConnectionProtos( + const FederatedQuery::CreateConnectionRequest& request, const TString& cloudId, const TString& user, TInstant startTime) const +{ + const TString& connectionId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::CONNECTION); + + FederatedQuery::Connection connection; + FederatedQuery::ConnectionContent& content = *connection.mutable_content(); + content = request.content(); + *connection.mutable_meta() = CreateCommonMeta(connectionId, user, startTime, InitialRevision); + + FederatedQuery::Internal::ConnectionInternal connectionInternal; + connectionInternal.set_cloud_id(cloudId); + + return {connection, connectionInternal}; +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateConnectionRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -43,25 +83,18 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateConne const TString user = event.User; const TString token = event.Token; const int byteSize = request.ByteSize(); - TPermissions permissions = Config->Proto.GetEnablePermissions() - ? event.Permissions - : TPermissions{TPermissions::MANAGE_PUBLIC}; - if (IsSuperUser(user)) { - permissions.SetAll(); - } const TString idempotencyKey = request.idempotency_key(); - const TString connectionId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::CONNECTION); + + const auto [connection, connectionInternal] = GetCreateConnectionProtos(request, cloudId, user, startTime); + const auto& content = connection.content(); + const TString& connectionId = connection.meta().id(); CPS_LOG_T(MakeLogPrefix(scope, user, connectionId) << "CreateConnectionRequest: " << NKikimr::MaskTicket(token) << " " << request.DebugString()); - NYql::TIssues issues = ValidateConnection(ev); - if (request.content().acl().visibility() == FederatedQuery::Acl::SCOPE && !permissions.Check(TPermissions::MANAGE_PUBLIC)) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::ACCESS_DENIED, "Permission denied to create a connection with these parameters. Please receive a permission yq.resources.managePublic")); - } - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_D(MakeLogPrefix(scope, user, connectionId) << "CreateConnectionRequest, validation failed: " << NKikimr::MaskTicket(token) << " " @@ -73,14 +106,6 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateConne return; } - FederatedQuery::Connection connection; - FederatedQuery::ConnectionContent& content = *connection.mutable_content(); - content = request.content(); - *connection.mutable_meta() = CreateCommonMeta(connectionId, user, startTime, InitialRevision); - - FederatedQuery::Internal::ConnectionInternal connectionInternal; - connectionInternal.set_cloud_id(cloudId); - std::shared_ptr>> response = std::make_shared>>(); response->first.set_connection_id(connectionId); response->second.After.ConstructInPlace().CopyFrom(connection); @@ -141,7 +166,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateConne auto overridBindingValidator = CreateConnectionOverrideBindingValidator( scope, content.name(), - permissions, + GetCreateConnectionPerimssions(event), user, YdbConnection->TablePathPrefix); validators.push_back(overridBindingValidator); diff --git a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h index dc212439650d..434c9405446b 100644 --- a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h +++ b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h @@ -8,6 +8,7 @@ #include "request_validators.h" #include "util.h" #include "validators.h" +#include #include #include @@ -312,10 +313,10 @@ class TControlPlaneStorageUtils { /* * Utility */ - bool IsSuperUser(const TString& user); + bool IsSuperUser(const TString& user) const; template - NYql::TIssues ValidateConnection(T& ev, bool passwordRequired = true) + NYql::TIssues ValidateConnection(T& ev, bool passwordRequired = true) const { return ::NFq::ValidateConnection(ev, Config->Proto.GetMaxRequestSize(), Config->AvailableConnections, Config->Proto.GetDisableCurrentIam(), @@ -323,19 +324,19 @@ class TControlPlaneStorageUtils { } template - NYql::TIssues ValidateBinding(T& ev) + NYql::TIssues ValidateBinding(T& ev) const { return ::NFq::ValidateBinding(ev, Config->Proto.GetMaxRequestSize(), Config->AvailableBindings, Config->GeneratorPathsLimit); } template - NYql::TIssues ValidateQuery(const T& ev) + NYql::TIssues ValidateQuery(const T& ev) const { return ::NFq::ValidateQuery(ev, Config->Proto.GetMaxRequestSize()); } template - NYql::TIssues ValidateEvent(const P& ev) + NYql::TIssues ValidateEvent(const P& ev) const { return ::NFq::ValidateEvent

(ev, Config->Proto.GetMaxRequestSize()); } @@ -370,10 +371,8 @@ class TControlPlaneStorageUtils { static constexpr int64_t InitialRevision = 1; }; -class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped, - public TDbRequester, - public TControlPlaneStorageUtils -{ +class TControlPlaneStorageBase : public TControlPlaneStorageUtils { +protected: enum ERequestTypeScope { RTS_CREATE_QUERY, RTS_LIST_QUERIES, @@ -558,6 +557,10 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped cacheVal; ScopeCounters.Get(key, &cacheVal); @@ -580,13 +583,217 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped GetCreateQueryProtos( + const FederatedQuery::CreateQueryRequest& request, const TString& user, TInstant startTime) const; + + FederatedQuery::Internal::QueryInternal GetQueryInternalProto( + const FederatedQuery::CreateQueryRequest& request, const TString& cloudId, const TString& token, + const TMaybe& quotas) const; + + void FillConnectionsAndBindings( + FederatedQuery::Internal::QueryInternal& queryInternal, FederatedQuery::QueryContent::QueryType queryType, + const TVector& allConnections, + const THashMap& visibleConnections, + const THashMap& visibleBindings) const; + + // Describe query request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvDescribeQueryRequest::TPtr& ev) const; + + void FillDescribeQueryResult( + FederatedQuery::DescribeQueryResult& result, FederatedQuery::Internal::QueryInternal internal, + const TString& user, TPermissions permissions) const; + + // Get result data request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvGetResultDataRequest::TPtr& ev) const; + + TPermissions GetResultDataReadPerimssions(const TEvControlPlaneStorage::TEvGetResultDataRequest& event) const; + + // Create connection request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvCreateConnectionRequest::TPtr& ev) const; + + TPermissions GetCreateConnectionPerimssions(const TEvControlPlaneStorage::TEvCreateConnectionRequest& event) const; + + std::pair GetCreateConnectionProtos( + const FederatedQuery::CreateConnectionRequest& request, const TString& cloudId, const TString& user, TInstant startTime) const; + + // Create binding request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvCreateBindingRequest::TPtr& ev) const; + + TPermissions GetCreateBindingPerimssions(const TEvControlPlaneStorage::TEvCreateBindingRequest& event) const; + + std::pair GetCreateBindingProtos( + const FederatedQuery::CreateBindingRequest& request, const TString& cloudId, const TString& user, TInstant startTime) const; + + // Write result data request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvWriteResultDataRequest::TPtr& ev) const; + + // Get task request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvGetTaskRequest::TPtr& ev) const; + + void FillGetTaskResult(Fq::Private::GetTaskResult& result, const TVector& tasks) const; + + // Ping task request + + struct TFinalStatus { + FederatedQuery::QueryMeta::ComputeStatus Status = FederatedQuery::QueryMeta::COMPUTE_STATUS_UNSPECIFIED; + NYql::NDqProto::StatusIds::StatusCode StatusCode = NYql::NDqProto::StatusIds::UNSPECIFIED; + FederatedQuery::QueryContent::QueryType QueryType = FederatedQuery::QueryContent::QUERY_TYPE_UNSPECIFIED; + NYql::TIssues Issues; + NYql::TIssues TransientIssues; + TStatsValuesList FinalStatistics; + TString CloudId; + TString JobId; + }; + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvPingTaskRequest::TPtr& ev) const; + + void UpdateTaskInfo( + NActors::TActorSystem* actorSystem, Fq::Private::PingTaskRequest& request, const std::shared_ptr& finalStatus, FederatedQuery::Query& query, + FederatedQuery::Internal::QueryInternal& internal, FederatedQuery::Job& job, TString& owner, + TRetryLimiter& retryLimiter, TDuration& backoff, TInstant& expireAt) const; + + void FillQueryStatistics( + const std::shared_ptr& finalStatus, const FederatedQuery::Query& query, + const FederatedQuery::Internal::QueryInternal& internal, const TRetryLimiter& retryLimiter) const; + + void Handle(TEvControlPlaneStorage::TEvFinalStatusReport::TPtr& ev); + + // Node health check + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvNodesHealthCheckRequest::TPtr& ev) const; + +protected: + // Should not be used from callbacks + template + void SendResponseIssues(TActorId sender, const NYql::TIssues& issues, ui64 cookie, const TDuration& delta, TRequestCounters requestCounters) { + std::unique_ptr event(new T{issues}); + requestCounters.Common->ResponseBytes->Add(event->GetByteSize()); + TActivationContext::Send(sender, std::move(event), 0, cookie); + requestCounters.DecInFly(); + requestCounters.IncError(); + requestCounters.Common->LatencyMs->Collect(delta.MilliSeconds()); + } + + template + TFuture SendResponse(const TString& name, + NActors::TActorSystem* actorSystem, + const TAsyncStatus& status, + TActorId self, + const RequestEventPtr& ev, + const TInstant& startTime, + const TRequestCounters& requestCounters, + const std::function::Type()>& prepare, + TDebugInfoPtr debugInfo) + { + return status.Apply([=, requestCounters=requestCounters](const auto& future) mutable { + NYql::TIssues internalIssues; + NYql::TIssues issues; + Result result; + typename TPrepareResponseResultType::TResponseAuditDetails auditDetails; // void* for nonauditable events + + try { + TStatus status = future.GetValue(); + if (status.IsSuccess()) { + if constexpr (ResponseEvent::Auditable) { + auto p = prepare(); + result = std::move(p.first); + auditDetails = std::move(p.second); + } else { + result = prepare(); + } + } else { + issues.AddIssues(status.GetIssues()); + internalIssues.AddIssues(status.GetIssues()); + } + } catch (const NYql::TCodeLineException& exception) { + NYql::TIssue issue = MakeErrorIssue(exception.Code, exception.GetRawMessage()); + issues.AddIssue(issue); + NYql::TIssue internalIssue = MakeErrorIssue(exception.Code, CurrentExceptionMessage()); + internalIssues.AddIssue(internalIssue); + } catch (const std::exception& exception) { + NYql::TIssue issue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, exception.what()); + issues.AddIssue(issue); + NYql::TIssue internalIssue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); + internalIssues.AddIssue(internalIssue); + } catch (...) { + NYql::TIssue issue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); + issues.AddIssue(issue); + NYql::TIssue internalIssue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); + internalIssues.AddIssue(internalIssue); + } + + const auto& request = ev->Get()->Request; + size_t responseByteSize = 0; + if (issues) { + CPS_LOG_AS_W(*actorSystem, name << ": {" << TrimForLogs(request.DebugString()) << "} ERROR: " << internalIssues.ToOneLineString()); + auto event = std::make_unique(issues); + event->DebugInfo = debugInfo; + responseByteSize = event->GetByteSize(); + actorSystem->Send(new IEventHandle(ev->Sender, self, event.release(), 0, ev->Cookie)); + requestCounters.IncError(); + for (const auto& issue : issues) { + NYql::WalkThroughIssues(issue, true, [&requestCounters](const NYql::TIssue& err, ui16 level) { + Y_UNUSED(level); + requestCounters.Common->Issues->GetCounter(ToString(err.GetCode()), true)->Inc(); + }); + } + } else { + CPS_LOG_AS_T(*actorSystem, name << ": {" << TrimForLogs(result.DebugString()) << "} SUCCESS"); + std::unique_ptr event; + if constexpr (ResponseEvent::Auditable) { + event = std::make_unique(result, auditDetails); + } else { + event = std::make_unique(result); + } + event->DebugInfo = debugInfo; + responseByteSize = event->GetByteSize(); + actorSystem->Send(new IEventHandle(ev->Sender, self, event.release(), 0, ev->Cookie)); + requestCounters.IncOk(); + } + requestCounters.DecInFly(); + requestCounters.Common->ResponseBytes->Add(responseByteSize); + TDuration delta = TInstant::Now() - startTime; + requestCounters.Common->LatencyMs->Collect(delta.MilliSeconds()); + return MakeFuture(!issues); + }); + } + + static ui64 GetExecutionLimitMills(FederatedQuery::QueryContent::QueryType queryType, const TMaybe& quotas); +}; + +class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped, + public TDbRequester, + public TControlPlaneStorageBase +{ + using TBase = TControlPlaneStorageBase; ::NFq::TYqSharedResources::TPtr YqSharedResources; NKikimr::TYdbCredentialsProviderFactory CredProviderFactory; - TString TenantName; // Query Quota THashMap QueryQuotas; @@ -606,12 +813,9 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped("FinalFailedStatusCode", counters->GetSubgroup("component", "QueryDiagnostic"))) + : TBase(config, s3Config, common, computeConfig, counters, tenantName) , YqSharedResources(yqSharedResources) , CredProviderFactory(credProviderFactory) - , TenantName(tenantName) { } @@ -654,7 +858,7 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped void HandleRateLimiterImpl(TEventPtr& ev); @@ -736,90 +938,6 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped - TFuture SendResponse(const TString& name, - NActors::TActorSystem* actorSystem, - const TAsyncStatus& status, - TActorId self, - const RequestEventPtr& ev, - const TInstant& startTime, - const TRequestCounters& requestCounters, - const std::function::Type()>& prepare, - TDebugInfoPtr debugInfo) - { - return status.Apply([=, requestCounters=requestCounters](const auto& future) mutable { - NYql::TIssues internalIssues; - NYql::TIssues issues; - Result result; - typename TPrepareResponseResultType::TResponseAuditDetails auditDetails; // void* for nonauditable events - - try { - TStatus status = future.GetValue(); - if (status.IsSuccess()) { - if constexpr (ResponseEvent::Auditable) { - auto p = prepare(); - result = std::move(p.first); - auditDetails = std::move(p.second); - } else { - result = prepare(); - } - } else { - issues.AddIssues(status.GetIssues()); - internalIssues.AddIssues(status.GetIssues()); - } - } catch (const NYql::TCodeLineException& exception) { - NYql::TIssue issue = MakeErrorIssue(exception.Code, exception.GetRawMessage()); - issues.AddIssue(issue); - NYql::TIssue internalIssue = MakeErrorIssue(exception.Code, CurrentExceptionMessage()); - internalIssues.AddIssue(internalIssue); - } catch (const std::exception& exception) { - NYql::TIssue issue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, exception.what()); - issues.AddIssue(issue); - NYql::TIssue internalIssue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); - internalIssues.AddIssue(internalIssue); - } catch (...) { - NYql::TIssue issue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); - issues.AddIssue(issue); - NYql::TIssue internalIssue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); - internalIssues.AddIssue(internalIssue); - } - - const auto& request = ev->Get()->Request; - size_t responseByteSize = 0; - if (issues) { - CPS_LOG_AS_W(*actorSystem, name << ": {" << TrimForLogs(request.DebugString()) << "} ERROR: " << internalIssues.ToOneLineString()); - auto event = std::make_unique(issues); - event->DebugInfo = debugInfo; - responseByteSize = event->GetByteSize(); - actorSystem->Send(new IEventHandle(ev->Sender, self, event.release(), 0, ev->Cookie)); - requestCounters.IncError(); - for (const auto& issue : issues) { - NYql::WalkThroughIssues(issue, true, [&requestCounters](const NYql::TIssue& err, ui16 level) { - Y_UNUSED(level); - requestCounters.Common->Issues->GetCounter(ToString(err.GetCode()), true)->Inc(); - }); - } - } else { - CPS_LOG_AS_T(*actorSystem, name << ": {" << TrimForLogs(result.DebugString()) << "} SUCCESS"); - std::unique_ptr event; - if constexpr (ResponseEvent::Auditable) { - event = std::make_unique(result, auditDetails); - } else { - event = std::make_unique(result); - } - event->DebugInfo = debugInfo; - responseByteSize = event->GetByteSize(); - actorSystem->Send(new IEventHandle(ev->Sender, self, event.release(), 0, ev->Cookie)); - requestCounters.IncOk(); - } - requestCounters.DecInFly(); - requestCounters.Common->ResponseBytes->Add(responseByteSize); - TDuration delta = TInstant::Now() - startTime; - requestCounters.Common->LatencyMs->Collect(delta.MilliSeconds()); - return MakeFuture(!issues); - }); - } - template TFuture SendResponseTuple(const TString& name, NActors::TActorSystem* actorSystem, @@ -890,20 +1008,6 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped - void SendResponseIssues(const TActorId sender, - const NYql::TIssues& issues, - ui64 cookie, - const TDuration& delta, - TRequestCounters requestCounters) { - std::unique_ptr event(new T{issues}); - requestCounters.Common->ResponseBytes->Add(event->GetByteSize()); - Send(sender, event.release(), 0, cookie); - requestCounters.DecInFly(); - requestCounters.IncError(); - requestCounters.Common->LatencyMs->Collect(delta.MilliSeconds()); - } - struct TPickTaskParams { TString ReadQuery; TParams ReadParams; @@ -920,9 +1024,20 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped& validators = {}, TTxSettings transactionMode = TTxSettings::SerializableRW()); - ui64 GetExecutionLimitMills( - FederatedQuery::QueryContent_QueryType queryType, - const TMaybe& quotas); + struct TPingTaskParams { + TString Query; + TParams Params; + const std::function(const TVector&)> Prepare; + std::shared_ptr> MeteringRecords; + }; + + TPingTaskParams ConstructHardPingTask( + const Fq::Private::PingTaskRequest& request, std::shared_ptr response, + const std::shared_ptr& finalStatus, const TRequestCommonCountersPtr& commonCounters) const; + + TPingTaskParams ConstructSoftPingTask( + const Fq::Private::PingTaskRequest& request, std::shared_ptr response, + const TRequestCommonCountersPtr& commonCounters) const; }; } diff --git a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp index fa2e67f2e27c..94b2353ba082 100644 --- a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp +++ b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp @@ -54,36 +54,17 @@ FederatedQuery::IamAuth::IdentityCase GetIamAuth(const FederatedQuery::Connectio namespace NFq { -void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQueryRequest::TPtr& ev) +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvCreateQueryRequest::TPtr& ev) const { - TInstant startTime = TInstant::Now(); const TEvControlPlaneStorage::TEvCreateQueryRequest& event = *ev->Get(); - const TString cloudId = event.CloudId; const FederatedQuery::CreateQueryRequest& request = event.Request; - const FederatedQuery::Internal::ComputeDatabaseInternal& computeDatabase = event.ComputeDatabase; - ui64 resultLimit = 0; - if (event.Quotas) { - if (auto it = event.Quotas->find(QUOTA_QUERY_RESULT_LIMIT); it != event.Quotas->end()) { - resultLimit = it->second.Limit.Value; - } - } - auto queryType = request.content().type(); - ui64 executionLimitMills = GetExecutionLimitMills(queryType, event.Quotas); - const TString scope = event.Scope; - TRequestCounters requestCounters = Counters.GetCounters(cloudId, scope, RTS_CREATE_QUERY, RTC_CREATE_QUERY); - requestCounters.IncInFly(); - requestCounters.Common->RequestBytes->Add(event.GetByteSize()); - const TString user = event.User; - const TString token = event.Token; + TPermissions permissions = Config->Proto.GetEnablePermissions() ? event.Permissions : TPermissions{TPermissions::QUERY_INVOKE | TPermissions::MANAGE_PUBLIC}; - if (IsSuperUser(user)) { + if (IsSuperUser(event.User)) { permissions.SetAll(); } - const size_t byteSize = request.ByteSizeLong(); - const TString queryId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::QUERY); - CPS_LOG_T("CreateQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token)); NYql::TIssues issues = ValidateQuery(ev); if (request.execute_mode() != FederatedQuery::SAVE && !permissions.Check(TPermissions::QUERY_INVOKE)) { @@ -97,10 +78,9 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQuery issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "Streaming disposition \"from_last_checkpoint\" is not allowed in CreateQuery request")); } - auto tenant = ev->Get()->TenantInfo->Assign(cloudId, scope, request.content().type(), TenantName); - if (event.Quotas) { TQuotaMap::const_iterator it = event.Quotas->end(); + const auto queryType = request.content().type(); if (queryType == FederatedQuery::QueryContent::ANALYTICS) { it = event.Quotas->find(QUOTA_ANALYTICS_COUNT_LIMIT); } else if (queryType == FederatedQuery::QueryContent::STREAMING) { @@ -120,24 +100,22 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQuery issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "VCPU rate limit can't be less than zero")); } - if (issues) { - CPS_LOG_W("CreateQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token) << "validation FAILED: " << issues.ToOneLineString()); - const TDuration delta = TInstant::Now() - startTime; - SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); - LWPROBE(CreateQueryRequest, scope, user, delta, byteSize, false); - return; - } - - const TString idempotencyKey = request.idempotency_key(); - const TString jobId = request.execute_mode() == FederatedQuery::SAVE ? "" : GetEntityIdAsString(Config->IdsPrefix, EEntityType::JOB); + return issues; +} +std::pair TControlPlaneStorageBase::GetCreateQueryProtos( + const FederatedQuery::CreateQueryRequest& request, const TString& user, TInstant startTime) const +{ FederatedQuery::Query query; FederatedQuery::QueryContent& content = *query.mutable_content() = request.content(); FederatedQuery::QueryMeta& meta = *query.mutable_meta(); + + const TString queryId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::QUERY); FederatedQuery::CommonMeta& common = *meta.mutable_common() = CreateCommonMeta(queryId, user, startTime, InitialRevision); meta.set_execute_mode(request.execute_mode()); meta.set_status(request.execute_mode() == FederatedQuery::SAVE ? FederatedQuery::QueryMeta::COMPLETED : FederatedQuery::QueryMeta::STARTING); + const TString jobId = request.execute_mode() == FederatedQuery::SAVE ? "" : GetEntityIdAsString(Config->IdsPrefix, EEntityType::JOB); FederatedQuery::Job job; if (request.execute_mode() != FederatedQuery::SAVE) { meta.set_last_job_query_revision(InitialRevision); @@ -154,6 +132,112 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQuery *job.mutable_parameters() = content.parameters(); } + return {query, job}; +} + +FederatedQuery::Internal::QueryInternal TControlPlaneStorageBase::GetQueryInternalProto( + const FederatedQuery::CreateQueryRequest& request, const TString& cloudId, const TString& token, + const TMaybe& quotas) const +{ + FederatedQuery::Internal::QueryInternal queryInternal; + if (!Config->Proto.GetDisableCurrentIam()) { + queryInternal.set_token(token); + } + + queryInternal.set_cloud_id(cloudId); + queryInternal.set_state_load_mode(FederatedQuery::StateLoadMode::EMPTY); + queryInternal.mutable_disposition()->CopyFrom(request.disposition()); + ui64 resultLimit = 0; + if (quotas) { + if (const auto it = quotas->find(QUOTA_QUERY_RESULT_LIMIT); it != quotas->end()) { + resultLimit = it->second.Limit.Value; + } + } + queryInternal.set_result_limit(resultLimit); + *queryInternal.mutable_execution_ttl() = NProtoInterop::CastToProto(TDuration::MilliSeconds(GetExecutionLimitMills(request.content().type(), quotas))); + + return queryInternal; +} + +void TControlPlaneStorageBase::FillConnectionsAndBindings( + FederatedQuery::Internal::QueryInternal& queryInternal, FederatedQuery::QueryContent::QueryType queryType, + const TVector& allConnections, + const THashMap& visibleConnections, + const THashMap& visibleBindings) const +{ + TSet disabledConnections; + for (const auto& connection : allConnections) { + const auto connectionCase = connection.content().setting().connection_case(); + if (!Config->AvailableConnections.contains(connectionCase)) { + disabledConnections.insert(connection.meta().id()); + continue; + } + if ((queryType == FederatedQuery::QueryContent::STREAMING) && !Config->AvailableStreamingConnections.contains(connectionCase)) { + disabledConnections.insert(connection.meta().id()); + continue; + } + if (GetIamAuth(connection) == FederatedQuery::IamAuth::kCurrentIam && Config->Proto.GetDisableCurrentIam()) { + disabledConnections.insert(connection.meta().id()); + continue; + } + } + + TSet connectionIds; + for (const auto& [_, connection] : visibleConnections) { + if (disabledConnections.contains(connection.meta().id())) { + continue; + } + *queryInternal.add_connection() = connection; + connectionIds.insert(connection.meta().id()); + } + + for (const auto& [_, binding] : visibleBindings) { + if (!Config->AvailableBindings.contains(binding.content().setting().binding_case())) { + continue; + } + if (disabledConnections.contains(binding.content().connection_id())) { + continue; + } + + *queryInternal.add_binding() = binding; + if (!connectionIds.contains(binding.content().connection_id())) { + ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Unable to resolve connection for binding " << binding.meta().id() << ", name " << binding.content().name() << ", connection id " << binding.content().connection_id(); + } + } +} + +void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQueryRequest::TPtr& ev) +{ + TInstant startTime = TInstant::Now(); + const TEvControlPlaneStorage::TEvCreateQueryRequest& event = *ev->Get(); + const TString cloudId = event.CloudId; + const FederatedQuery::CreateQueryRequest& request = event.Request; + const FederatedQuery::Internal::ComputeDatabaseInternal& computeDatabase = event.ComputeDatabase; + const auto queryType = request.content().type(); + const TString scope = event.Scope; + TRequestCounters requestCounters = Counters.GetCounters(cloudId, scope, RTS_CREATE_QUERY, RTC_CREATE_QUERY); + requestCounters.IncInFly(); + requestCounters.Common->RequestBytes->Add(event.GetByteSize()); + const TString user = event.User; + const TString token = event.Token; + const size_t byteSize = request.ByteSizeLong(); + CPS_LOG_T("CreateQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token)); + + auto tenant = ev->Get()->TenantInfo->Assign(cloudId, scope, queryType, TenantName); + + if (const auto& issues = ValidateRequest(ev)) { + CPS_LOG_W("CreateQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token) << "validation FAILED: " << issues.ToOneLineString()); + const TDuration delta = TInstant::Now() - startTime; + SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); + LWPROBE(CreateQueryRequest, scope, user, delta, byteSize, false); + return; + } + + const TString idempotencyKey = request.idempotency_key(); + const auto [query, job] = GetCreateQueryProtos(request, user, startTime); + const TString queryId = query.meta().common().id(); + const TString jobId = job.meta().id(); + std::shared_ptr>> response = std::make_shared>>(); response->first.set_query_id(queryId); response->second.CloudId = cloudId; @@ -179,7 +263,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQuery ); } - auto prepareParams = [=, as=TActivationContext::ActorSystem(), commonCounters=requestCounters.Common](const TVector& resultSets) mutable { + auto prepareParams = [=, as=TActivationContext::ActorSystem(), commonCounters=requestCounters.Common, quotas=event.Quotas](const TVector& resultSets) mutable { const size_t countSets = (idempotencyKey ? 1 : 0) + (request.execute_mode() != FederatedQuery::SAVE ? 2 : 0); if (resultSets.size() != countSets) { ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Result set size is not equal to " << countSets << " but equal " << resultSets.size() << ". Please contact internal support"; @@ -197,63 +281,18 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQuery } } - FederatedQuery::Internal::QueryInternal queryInternal; - if (!Config->Proto.GetDisableCurrentIam()) { - queryInternal.set_token(token); - } - - queryInternal.set_cloud_id(cloudId); - queryInternal.set_state_load_mode(FederatedQuery::StateLoadMode::EMPTY); - queryInternal.mutable_disposition()->CopyFrom(request.disposition()); - queryInternal.set_result_limit(resultLimit); - *queryInternal.mutable_execution_ttl() = NProtoInterop::CastToProto(TDuration::MilliSeconds(executionLimitMills)); + auto queryInternal = GetQueryInternalProto(request, cloudId, token, quotas); if (request.execute_mode() != FederatedQuery::SAVE) { // TODO: move to run actor priority selection *queryInternal.mutable_compute_connection() = computeDatabase.connection(); - TSet disabledConnections; - for (const auto& connection: GetEntities(resultSets[resultSets.size() - 2], CONNECTION_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters)) { - auto connectionCase = connection.content().setting().connection_case(); - if (!Config->AvailableConnections.contains(connectionCase)) { - disabledConnections.insert(connection.meta().id()); - continue; - } - if ((queryType == FederatedQuery::QueryContent::STREAMING) && !Config->AvailableStreamingConnections.contains(connectionCase)) { - disabledConnections.insert(connection.meta().id()); - continue; - } - - if (GetIamAuth(connection) == FederatedQuery::IamAuth::kCurrentIam && Config->Proto.GetDisableCurrentIam()) { - disabledConnections.insert(connection.meta().id()); - continue; - } - } - - TSet connectionIds; - auto connections = GetEntitiesWithVisibilityPriority(resultSets[resultSets.size() - 2], CONNECTION_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters); - for (const auto& [_, connection]: connections) { - if (disabledConnections.contains(connection.meta().id())) { - continue; - } - *queryInternal.add_connection() = connection; - connectionIds.insert(connection.meta().id()); - } - - auto bindings = GetEntitiesWithVisibilityPriority(resultSets[resultSets.size() - 1], BINDING_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters); - for (const auto& [_, binding]: bindings) { - if (!Config->AvailableBindings.contains(binding.content().setting().binding_case())) { - continue; - } - - if (disabledConnections.contains(binding.content().connection_id())) { - continue; - } - - *queryInternal.add_binding() = binding; - if (!connectionIds.contains(binding.content().connection_id())) { - ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Unable to resolve connection for binding " << binding.meta().id() << ", name " << binding.content().name() << ", connection id " << binding.content().connection_id(); - } - } + FillConnectionsAndBindings( + queryInternal, + queryType, + GetEntities(resultSets[resultSets.size() - 2], CONNECTION_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters), + GetEntitiesWithVisibilityPriority(resultSets[resultSets.size() - 2], CONNECTION_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters), + GetEntitiesWithVisibilityPriority(resultSets[resultSets.size() - 1], BINDING_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters) + ); } if (query.ByteSizeLong() > Config->Proto.GetMaxRequestSize()) { @@ -510,6 +549,73 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvListQueries }); } +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvDescribeQueryRequest::TPtr& ev) const +{ + return ValidateEvent(ev); +} + +void TControlPlaneStorageBase::FillDescribeQueryResult( + FederatedQuery::DescribeQueryResult& result, FederatedQuery::Internal::QueryInternal internal, + const TString& user, TPermissions permissions) const +{ + const auto lastJobId = result.query().meta().last_job_id(); + result.mutable_query()->mutable_meta()->set_last_job_id(lastJobId + "-" + result.query().meta().common().id()); + + permissions = Config->Proto.GetEnablePermissions() + ? permissions + : TPermissions{TPermissions::VIEW_PUBLIC | TPermissions::VIEW_AST | TPermissions::VIEW_QUERY_TEXT}; + if (IsSuperUser(user)) { + permissions.SetAll(); + } + + const auto queryVisibility = result.query().content().acl().visibility(); + const auto queryUser = result.query().meta().common().created_by(); + const bool hasViewAccess = HasViewAccess(permissions, queryVisibility, queryUser, user); + if (!hasViewAccess) { + ythrow NYql::TCodeLineException(TIssuesIds::ACCESS_DENIED) << "Query does not exist or permission denied. Please check the id of the query or your access rights"; + } + + // decompress plan + if (internal.plan_compressed().data()) { // todo: remove this if after migration + TCompressor compressor(internal.plan_compressed().method()); + result.mutable_query()->mutable_plan()->set_json(compressor.Decompress(internal.plan_compressed().data())); + if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { + if (result.query().plan().json().size() > 1000) { + // modifing plan this way should definitely reduce query msg size + result.mutable_query()->mutable_plan()->set_json(TStringBuilder() << "Message is too big: " << result.query().ByteSizeLong() << " bytes, dropping plan of size " << result.query().plan().json().size() << " bytes"); + } + } + } + + auto timeline = internal.timeline(); + if (timeline) { + result.mutable_query()->mutable_timeline()->set_svg(timeline); + } + + if (!permissions.Check(TPermissions::VIEW_AST)) { + result.mutable_query()->clear_ast(); + } else { + // decompress AST + if (internal.ast_compressed().data()) { // todo: remove this if after migration + TCompressor compressor(internal.ast_compressed().method()); + result.mutable_query()->mutable_ast()->set_data(compressor.Decompress(internal.ast_compressed().data())); + } + if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { + if (result.query().ast().data().size() > 1000) { + // modifing AST this way should definitely reduce query msg size + result.mutable_query()->mutable_ast()->set_data(TStringBuilder() << "Message is too big: " << result.query().ByteSizeLong() << " bytes, dropping AST of size " << result.query().ast().data().size() << " bytes"); + } + } + } + if (!permissions.Check(TPermissions::VIEW_QUERY_TEXT)) { + result.mutable_query()->mutable_content()->clear_text(); + } + + if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { + ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Resulting query of size " << result.query().ByteSizeLong() << " bytes is too big"; + } +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvDescribeQueryRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -521,20 +627,13 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvDescribeQue requestCounters.Common->RequestBytes->Add(event.GetByteSize()); const TString user = event.User; const TString token = event.Token; - TPermissions permissions = Config->Proto.GetEnablePermissions() - ? event.Permissions - : TPermissions{TPermissions::VIEW_PUBLIC | TPermissions::VIEW_AST | TPermissions::VIEW_QUERY_TEXT}; - if (IsSuperUser(user)) { - permissions.SetAll(); - } const FederatedQuery::DescribeQueryRequest& request = event.Request; const TString queryId = request.query_id(); const int byteSize = request.ByteSize(); CPS_LOG_T("DescribeQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token)); - NYql::TIssues issues = ValidateEvent(ev); - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_W("DescribeQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token) << "validation FAILED: " << issues.ToOneLineString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); @@ -553,7 +652,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvDescribeQue const auto query = queryBuilder.Build(); auto debugInfo = Config->Proto.GetEnableDebugMode() ? std::make_shared() : TDebugInfoPtr{}; auto [result, resultSets] = Read(query.Sql, query.Params, requestCounters, debugInfo); - auto prepare = [resultSets=resultSets, user, permissions, commonCounters=requestCounters.Common] { + auto prepare = [this, resultSets=resultSets, user, permissions=event.Permissions, commonCounters=requestCounters.Common] { if (resultSets->size() != 1) { ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Result set size is not equal to 1 but equal " << resultSets->size() << ". Please contact internal support"; } @@ -569,61 +668,13 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvDescribeQue ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Error parsing proto message for query. Please contact internal support"; } - const auto lastJobId = result.query().meta().last_job_id(); - result.mutable_query()->mutable_meta()->set_last_job_id(lastJobId + "-" + result.query().meta().common().id()); - - const auto queryVisibility = result.query().content().acl().visibility(); - const auto queryUser = result.query().meta().common().created_by(); - const bool hasViewAccess = HasViewAccess(permissions, queryVisibility, queryUser, user); - if (!hasViewAccess) { - ythrow NYql::TCodeLineException(TIssuesIds::ACCESS_DENIED) << "Query does not exist or permission denied. Please check the id of the query or your access rights"; - } - FederatedQuery::Internal::QueryInternal internal; if (!internal.ParseFromString(*parser.ColumnParser(INTERNAL_COLUMN_NAME).GetOptionalString())) { commonCounters->ParseProtobufError->Inc(); ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Error parsing proto message for query internal. Please contact internal support"; } - // decompress plan - if (internal.plan_compressed().data()) { // todo: remove this if after migration - TCompressor compressor(internal.plan_compressed().method()); - result.mutable_query()->mutable_plan()->set_json(compressor.Decompress(internal.plan_compressed().data())); - if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { - if (result.query().plan().json().size() > 1000) { - // modifing plan this way should definitely reduce query msg size - result.mutable_query()->mutable_plan()->set_json(TStringBuilder() << "Message is too big: " << result.query().ByteSizeLong() << " bytes, dropping plan of size " << result.query().plan().json().size() << " bytes"); - } - } - } - - auto timeline = internal.timeline(); - if (timeline) { - result.mutable_query()->mutable_timeline()->set_svg(timeline); - } - - if (!permissions.Check(TPermissions::VIEW_AST)) { - result.mutable_query()->clear_ast(); - } else { - // decompress AST - if (internal.ast_compressed().data()) { // todo: remove this if after migration - TCompressor compressor(internal.ast_compressed().method()); - result.mutable_query()->mutable_ast()->set_data(compressor.Decompress(internal.ast_compressed().data())); - } - if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { - if (result.query().ast().data().size() > 1000) { - // modifing AST this way should definitely reduce query msg size - result.mutable_query()->mutable_ast()->set_data(TStringBuilder() << "Message is too big: " << result.query().ByteSizeLong() << " bytes, dropping AST of size " << result.query().ast().data().size() << " bytes"); - } - } - } - if (!permissions.Check(TPermissions::VIEW_QUERY_TEXT)) { - result.mutable_query()->mutable_content()->clear_text(); - } - - if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { - ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Resulting query of size " << result.query().ByteSizeLong() << " bytes is too big"; - } + FillDescribeQueryResult(result, internal, user, permissions); return result; }; @@ -1484,6 +1535,21 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvControlQuer }); } +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvGetResultDataRequest::TPtr& ev) const +{ + return ValidateEvent(ev); +} + +TPermissions TControlPlaneStorageBase::GetResultDataReadPerimssions(const TEvControlPlaneStorage::TEvGetResultDataRequest& event) const { + TPermissions permissions = Config->Proto.GetEnablePermissions() + ? event.Permissions + : TPermissions{TPermissions::VIEW_PUBLIC}; + if (IsSuperUser(event.User)) { + permissions.SetAll(); + } + return permissions; +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetResultDataRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -1501,17 +1567,11 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetResultDa const TString queryId = request.query_id(); const int byteSize = event.Request.ByteSize(); const TString token = event.Token; - TPermissions permissions = Config->Proto.GetEnablePermissions() - ? event.Permissions - : TPermissions{TPermissions::VIEW_PUBLIC}; - if (IsSuperUser(user)) { - permissions.SetAll(); - } + const TPermissions permissions = GetResultDataReadPerimssions(event); const int64_t limit = request.limit(); CPS_LOG_T("GetResultDataRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token)); - NYql::TIssues issues = ValidateEvent(ev); - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_W("GetResultDataRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token) << "validation FAILED: " << issues.ToOneLineString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); @@ -1851,10 +1911,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvDescribeJob }); } -ui64 TYdbControlPlaneStorageActor::GetExecutionLimitMills( - FederatedQuery::QueryContent_QueryType queryType, - const TMaybe& quotas) { - +ui64 TControlPlaneStorageBase::GetExecutionLimitMills(FederatedQuery::QueryContent::QueryType queryType, const TMaybe& quotas) { if (!quotas) { return 0; } diff --git a/ydb/core/fq/libs/db_id_async_resolver_impl/mdb_endpoint_generator.cpp b/ydb/core/fq/libs/db_id_async_resolver_impl/mdb_endpoint_generator.cpp index 8c0523189eb1..15d0c6522c3d 100644 --- a/ydb/core/fq/libs/db_id_async_resolver_impl/mdb_endpoint_generator.cpp +++ b/ydb/core/fq/libs/db_id_async_resolver_impl/mdb_endpoint_generator.cpp @@ -56,14 +56,14 @@ namespace NFq { ui32 port; switch (params.Protocol) { - case NYql::NConnector::NApi::EProtocol::NATIVE: + case NYql::EGenericProtocol::NATIVE: port = params.UseTls ? CLICKHOUSE_NATIVE_SECURE_PORT : CLICKHOUSE_NATIVE_INSECURE_PORT; break; - case NYql::NConnector::NApi::EProtocol::HTTP: + case NYql::EGenericProtocol::HTTP: port = params.UseTls ? CLICKHOUSE_HTTP_SECURE_PORT : CLICKHOUSE_HTTP_INSECURE_PORT; break; default: - ythrow yexception() << "Unexpected protocol for ClickHouse: " << NYql::NConnector::NApi::EProtocol_Name(params.Protocol); + ythrow yexception() << "Unexpected protocol for ClickHouse: " << NYql::EGenericProtocol_Name(params.Protocol); } return TEndpoint(fixedHost, port); @@ -71,26 +71,26 @@ namespace NFq { case NYql::EDatabaseType::PostgreSQL: // https://cloud.yandex.ru/docs/managed-postgresql/operations/connect switch (params.Protocol) { - case NYql::NConnector::NApi::EProtocol::NATIVE: + case NYql::EGenericProtocol::NATIVE: return TEndpoint(fixedHost, POSTGRESQL_PORT); default: - ythrow yexception() << "Unexpected protocol for PostgreSQL " << NYql::NConnector::NApi::EProtocol_Name(params.Protocol); + ythrow yexception() << "Unexpected protocol for PostgreSQL " << NYql::EGenericProtocol_Name(params.Protocol); } case NYql::EDatabaseType::Greenplum: // https://cloud.yandex.ru/docs/managed-greenplum/operations/connect switch (params.Protocol) { - case NYql::NConnector::NApi::EProtocol::NATIVE: + case NYql::EGenericProtocol::NATIVE: return TEndpoint(fixedHost, GREENPLUM_PORT); default: - ythrow yexception() << "Unexpected protocol for Greenplum: " << NYql::NConnector::NApi::EProtocol_Name(params.Protocol); + ythrow yexception() << "Unexpected protocol for Greenplum: " << NYql::EGenericProtocol_Name(params.Protocol); } case NYql::EDatabaseType::MySQL: // https://cloud.yandex.ru/docs/managed-mysql/operations/connect switch (params.Protocol) { - case NYql::NConnector::NApi::EProtocol::NATIVE: + case NYql::EGenericProtocol::NATIVE: return TEndpoint(fixedHost, MYSQL_PORT); default: - ythrow yexception() << "Unexpected protocol for MySQL: " << NYql::NConnector::NApi::EProtocol_Name(params.Protocol); + ythrow yexception() << "Unexpected protocol for MySQL: " << NYql::EGenericProtocol_Name(params.Protocol); } default: ythrow yexception() << "Unexpected database type: " << ToString(params.DatabaseType); diff --git a/ydb/core/fq/libs/db_id_async_resolver_impl/ut/mdb_endpoint_generator_ut.cpp b/ydb/core/fq/libs/db_id_async_resolver_impl/ut/mdb_endpoint_generator_ut.cpp index f02d3a4e08ac..f05a6e49aa7c 100644 --- a/ydb/core/fq/libs/db_id_async_resolver_impl/ut/mdb_endpoint_generator_ut.cpp +++ b/ydb/core/fq/libs/db_id_async_resolver_impl/ut/mdb_endpoint_generator_ut.cpp @@ -11,7 +11,7 @@ Y_UNIT_TEST_SUITE(MdbEndpoingGenerator) { .DatabaseType = NYql::EDatabaseType::ClickHouse, .MdbHost = "rc1c-p5waby2y5y1kb5ue.db.yandex.net", .UseTls = true, - .Protocol = NYql::NConnector::NApi::EProtocol::HTTP, + .Protocol = NYql::EGenericProtocol::HTTP, }; UNIT_ASSERT_VALUES_EQUAL( @@ -22,7 +22,7 @@ Y_UNIT_TEST_SUITE(MdbEndpoingGenerator) { .DatabaseType = NYql::EDatabaseType::ClickHouse, .MdbHost = "ya.ru", .UseTls = false, - .Protocol = NYql::NConnector::NApi::EProtocol::HTTP, + .Protocol = NYql::EGenericProtocol::HTTP, }; UNIT_ASSERT_VALUES_EQUAL( @@ -37,7 +37,7 @@ Y_UNIT_TEST_SUITE(MdbEndpoingGenerator) { .DatabaseType = NYql::EDatabaseType::ClickHouse, .MdbHost = "rc1a-d6dv17lv47v5mcop.mdb.yandexcloud.net", .UseTls = true, - .Protocol = NYql::NConnector::NApi::EProtocol::HTTP, + .Protocol = NYql::EGenericProtocol::HTTP, }; UNIT_ASSERT_VALUES_EQUAL( @@ -48,7 +48,7 @@ Y_UNIT_TEST_SUITE(MdbEndpoingGenerator) { .DatabaseType = NYql::EDatabaseType::PostgreSQL, .MdbHost = "rc1b-eyt6dtobu96rwydq.mdb.yandexcloud.net", .UseTls = false, - .Protocol = NYql::NConnector::NApi::EProtocol::NATIVE, + .Protocol = NYql::EGenericProtocol::NATIVE, }; UNIT_ASSERT_VALUES_EQUAL( @@ -65,7 +65,7 @@ Y_UNIT_TEST_SUITE(MdbEndpoingGenerator) { .DatabaseType = NYql::EDatabaseType::ClickHouse, .MdbHost = "rc1a-d6dv17lv47v5mcop.mdb.yandexcloud.net", .UseTls = false, - .Protocol = NYql::NConnector::NApi::EProtocol::HTTP, + .Protocol = NYql::EGenericProtocol::HTTP, }; UNIT_ASSERT_VALUES_EQUAL( @@ -76,7 +76,7 @@ Y_UNIT_TEST_SUITE(MdbEndpoingGenerator) { .DatabaseType = NYql::EDatabaseType::ClickHouse, .MdbHost = "rc1a-d6dv17lv47v5mcop.mdb.yandexcloud.net", .UseTls = false, - .Protocol = NYql::NConnector::NApi::EProtocol::NATIVE, + .Protocol = NYql::EGenericProtocol::NATIVE, }; UNIT_ASSERT_VALUES_EQUAL( @@ -87,7 +87,7 @@ Y_UNIT_TEST_SUITE(MdbEndpoingGenerator) { .DatabaseType = NYql::EDatabaseType::ClickHouse, .MdbHost = "rc1a-d6dv17lv47v5mcop.mdb.yandexcloud.net", .UseTls = true, - .Protocol = NYql::NConnector::NApi::EProtocol::HTTP, + .Protocol = NYql::EGenericProtocol::HTTP, }; UNIT_ASSERT_VALUES_EQUAL( @@ -98,7 +98,7 @@ Y_UNIT_TEST_SUITE(MdbEndpoingGenerator) { .DatabaseType = NYql::EDatabaseType::ClickHouse, .MdbHost = "rc1a-d6dv17lv47v5mcop.mdb.yandexcloud.net", .UseTls = true, - .Protocol = NYql::NConnector::NApi::EProtocol::NATIVE, + .Protocol = NYql::EGenericProtocol::NATIVE, }; UNIT_ASSERT_VALUES_EQUAL( @@ -111,7 +111,7 @@ Y_UNIT_TEST_SUITE(MdbEndpoingGenerator) { .DatabaseType = NYql::EDatabaseType::PostgreSQL, .MdbHost = "rc1b-eyt6dtobu96rwydq.mdb.yandexcloud.net", .UseTls = true, - .Protocol = NYql::NConnector::NApi::EProtocol::NATIVE, + .Protocol = NYql::EGenericProtocol::NATIVE, }; UNIT_ASSERT_VALUES_EQUAL( diff --git a/ydb/core/fq/libs/http_api_client/http_client.py b/ydb/core/fq/libs/http_api_client/http_client.py index 88b9e781c52e..e5770ce9adde 100644 --- a/ydb/core/fq/libs/http_api_client/http_client.py +++ b/ydb/core/fq/libs/http_api_client/http_client.py @@ -11,7 +11,7 @@ from .query_results import YQResults -MAX_RETRY_FOR_SESSION = 4 +MAX_RETRY_FOR_SESSION = 100 BACK_OFF_FACTOR = 0.3 TIME_BETWEEN_RETRIES = 1000 ERROR_CODES = (500, 502, 504) @@ -150,6 +150,22 @@ def create_query( self._validate_http_error(response, expected_code=expected_code) return response.json()["id"] + def start_query( + self, + query_id: str, + request_id=None, + idempotency_key: str | None = None, + expected_code: int = 204, + ): + response = self.session.post( + self._compose_api_url(f"/api/fq/v1/queries/{query_id}/start"), + headers=self._build_headers(idempotency_key=idempotency_key, request_id=request_id), + params=self._build_params(), + ) + + self._validate_http_error(response, expected_code) + return response + def get_query_status(self, query_id, request_id=None, expected_code=200) -> Any: response = self.session.get( self._compose_api_url(f"/api/fq/v1/queries/{query_id}/status"), @@ -272,9 +288,7 @@ def get_query_result_set(self, query_id: str, result_set_index: int, raw_format: return YQResults(result).results - def get_query_all_result_sets( - self, query_id: str, result_set_count: int, raw_format: bool = False - ) -> Any: + def get_query_all_result_sets(self, query_id: str, result_set_count: int, raw_format: bool = False) -> Any: result = [] for i in range(0, result_set_count): r = self.get_query_result_set(query_id, result_set_index=i, raw_format=raw_format) diff --git a/ydb/core/fq/libs/init/init.cpp b/ydb/core/fq/libs/init/init.cpp index e365ec04cf84..cd9afe4dcbc0 100644 --- a/ydb/core/fq/libs/init/init.cpp +++ b/ydb/core/fq/libs/init/init.cpp @@ -56,6 +56,17 @@ namespace NFq { using namespace NKikimr; +NYdb::NTopic::TTopicClientSettings GetCommonTopicClientSettings(const NFq::NConfig::TCommonConfig& config) { + NYdb::NTopic::TTopicClientSettings settings; + if (config.GetTopicClientHandlersExecutorThreadsNum()) { + settings.DefaultHandlersExecutor(NYdb::NTopic::CreateThreadPoolExecutor(config.GetTopicClientHandlersExecutorThreadsNum())); + } + if (config.GetTopicClientCompressionExecutorThreadsNum()) { + settings.DefaultCompressionExecutor(NYdb::NTopic::CreateThreadPoolExecutor(config.GetTopicClientCompressionExecutorThreadsNum())); + } + return settings; +} + void Init( const NFq::NConfig::TConfig& protoConfig, ui32 nodeId, @@ -66,7 +77,8 @@ void Init( const IYqSharedResources::TPtr& iyqSharedResources, const std::function& folderServiceFactory, ui32 icPort, - const std::vector& additionalCompNodeFactories + const std::vector& additionalCompNodeFactories, + NYql::IPqGatewayFactory::TPtr pqGatewayFactory ) { Y_ABORT_UNLESS(iyqSharedResources, "No YQ shared resources created"); @@ -76,14 +88,21 @@ void Init( const auto clientCounters = yqCounters->GetSubgroup("subsystem", "ClientMetrics"); if (protoConfig.GetControlPlaneStorage().GetEnabled()) { + const auto counters = yqCounters->GetSubgroup("subsystem", "ControlPlaneStorage"); auto controlPlaneStorage = protoConfig.GetControlPlaneStorage().GetUseInMemory() - ? NFq::CreateInMemoryControlPlaneStorageServiceActor(protoConfig.GetControlPlaneStorage()) + ? NFq::CreateInMemoryControlPlaneStorageServiceActor( + protoConfig.GetControlPlaneStorage(), + protoConfig.GetGateways().GetS3(), + protoConfig.GetCommon(), + protoConfig.GetCompute(), + counters, + tenant) : NFq::CreateYdbControlPlaneStorageServiceActor( protoConfig.GetControlPlaneStorage(), protoConfig.GetGateways().GetS3(), protoConfig.GetCommon(), protoConfig.GetCompute(), - yqCounters->GetSubgroup("subsystem", "ControlPlaneStorage"), + counters, yqSharedResources, NKikimr::CreateYdbCredentialsProviderFactory, tenant); @@ -189,14 +208,18 @@ void Init( credentialsFactory = NYql::CreateSecuredServiceAccountCredentialsOverTokenAccessorFactory(tokenAccessorConfig.GetEndpoint(), tokenAccessorConfig.GetUseSsl(), caContent, tokenAccessorConfig.GetConnectionPoolSize()); } + auto commonTopicClientSettings = GetCommonTopicClientSettings(protoConfig.GetCommon()); + if (protoConfig.GetRowDispatcher().GetEnabled()) { NYql::TPqGatewayServices pqServices( yqSharedResources->UserSpaceYdbDriver, nullptr, nullptr, std::make_shared(), - nullptr); - + nullptr, + nullptr, + commonTopicClientSettings + ); auto rowDispatcher = NFq::NewRowDispatcherService( protoConfig.GetRowDispatcher(), NKikimr::CreateYdbCredentialsProviderFactory, @@ -204,8 +227,9 @@ void Init( credentialsFactory, tenant, yqCounters->GetSubgroup("subsystem", "row_dispatcher"), - CreatePqNativeGateway(pqServices), - appData->Mon); + pqGatewayFactory ? pqGatewayFactory->CreatePqGateway() : CreatePqNativeGateway(pqServices), + appData->Mon, + appData->Counters); actorRegistrator(NFq::RowDispatcherServiceActorId(), rowDispatcher.release()); } @@ -222,9 +246,12 @@ void Init( pqCmConnections, credentialsFactory, std::make_shared(protoConfig.GetGateways().GetPq()), - appData->FunctionRegistry + appData->FunctionRegistry, + nullptr, + commonTopicClientSettings ); - RegisterDqPqReadActorFactory(*asyncIoFactory, yqSharedResources->UserSpaceYdbDriver, credentialsFactory, NYql::CreatePqNativeGateway(std::move(pqServices)), + auto pqGateway = pqGatewayFactory ? pqGatewayFactory->CreatePqGateway() : NYql::CreatePqNativeGateway(std::move(pqServices)); + RegisterDqPqReadActorFactory(*asyncIoFactory, yqSharedResources->UserSpaceYdbDriver, credentialsFactory, pqGateway, yqCounters->GetSubgroup("subsystem", "DqSourceTracker"), protoConfig.GetCommon().GetPqReconnectPeriod()); s3ActorsFactory->RegisterS3ReadActorFactory(*asyncIoFactory, credentialsFactory, httpGateway, s3HttpRetryPolicy, readActorFactoryCfg, @@ -233,7 +260,7 @@ void Init( httpGateway, s3HttpRetryPolicy); RegisterGenericProviderFactories(*asyncIoFactory, credentialsFactory, connectorClient); - RegisterDqPqWriteActorFactory(*asyncIoFactory, yqSharedResources->UserSpaceYdbDriver, credentialsFactory, yqCounters->GetSubgroup("subsystem", "DqSinkTracker")); + RegisterDqPqWriteActorFactory(*asyncIoFactory, yqSharedResources->UserSpaceYdbDriver, credentialsFactory, pqGateway, yqCounters->GetSubgroup("subsystem", "DqSinkTracker")); RegisterDQSolomonWriteActorFactory(*asyncIoFactory, credentialsFactory); } @@ -327,6 +354,15 @@ void Init( } if (protoConfig.GetPendingFetcher().GetEnabled()) { + NYql::TPqGatewayServices pqServices( + yqSharedResources->UserSpaceYdbDriver, + pqCmConnections, + credentialsFactory, + std::make_shared(protoConfig.GetGateways().GetPq()), + appData->FunctionRegistry, + nullptr, + commonTopicClientSettings + ); auto fetcher = CreatePendingFetcher( yqSharedResources, NKikimr::CreateYdbCredentialsProviderFactory, @@ -343,7 +379,8 @@ void Init( clientCounters, tenant, appData->Mon, - s3ActorsFactory + s3ActorsFactory, + pqGatewayFactory ? pqGatewayFactory : NYql::CreatePqNativeGatewayFactory(pqServices) ); actorRegistrator(MakePendingFetcherId(nodeId), fetcher); diff --git a/ydb/core/fq/libs/init/init.h b/ydb/core/fq/libs/init/init.h index 6dcf35afd6d1..ad1bca9a431f 100644 --- a/ydb/core/fq/libs/init/init.h +++ b/ydb/core/fq/libs/init/init.h @@ -12,6 +12,7 @@ #include #include +#include #include @@ -36,7 +37,8 @@ void Init( const IYqSharedResources::TPtr& yqSharedResources, const std::function& folderServiceFactory, ui32 icPort, - const std::vector& additionalCompNodeFactories + const std::vector& additionalCompNodeFactories, + NYql::IPqGatewayFactory::TPtr pqGatewayFactory = nullptr ); } // NFq diff --git a/ydb/core/fq/libs/row_dispatcher/common.cpp b/ydb/core/fq/libs/metrics/sanitize_label.cpp similarity index 50% rename from ydb/core/fq/libs/row_dispatcher/common.cpp rename to ydb/core/fq/libs/metrics/sanitize_label.cpp index 2f2696c23d33..53104099f735 100644 --- a/ydb/core/fq/libs/row_dispatcher/common.cpp +++ b/ydb/core/fq/libs/metrics/sanitize_label.cpp @@ -1,12 +1,13 @@ -#include "common.h" +#include "sanitize_label.h" #include namespace NFq { -TString CleanupCounterValueString(const TString& value) { - TString clean; - constexpr auto valueLenghtLimit = 200; +TString SanitizeLabel(const TString& value) { + TString result; + result.reserve(value.size()); + constexpr auto labelLengthLimit = 200; for (auto c : value) { switch (c) { @@ -19,13 +20,13 @@ TString CleanupCounterValueString(const TString& value) { case '\\': continue; default: - clean.push_back(c); - if (clean.size() == valueLenghtLimit) { - break; + result.push_back(c); + if (result.size() == labelLengthLimit) { + return result; } } } - return clean; + return result; } } // namespace NFq diff --git a/ydb/core/fq/libs/metrics/sanitize_label.h b/ydb/core/fq/libs/metrics/sanitize_label.h new file mode 100644 index 000000000000..99ade02ef0ba --- /dev/null +++ b/ydb/core/fq/libs/metrics/sanitize_label.h @@ -0,0 +1,9 @@ +#pragma once + +#include + +namespace NFq { + +TString SanitizeLabel(const TString& value); + +} // namespace NFq diff --git a/ydb/core/fq/libs/metrics/ut/sanitize_label_ut.cpp b/ydb/core/fq/libs/metrics/ut/sanitize_label_ut.cpp new file mode 100644 index 000000000000..5170d59ce416 --- /dev/null +++ b/ydb/core/fq/libs/metrics/ut/sanitize_label_ut.cpp @@ -0,0 +1,29 @@ +#include + +#include + +Y_UNIT_TEST_SUITE(SanitizeLable) { + Y_UNIT_TEST(Empty) { + UNIT_ASSERT_VALUES_EQUAL(NFq::SanitizeLabel(""), ""); + } + + Y_UNIT_TEST(SkipSingleBadSymbol) { + UNIT_ASSERT_VALUES_EQUAL(NFq::SanitizeLabel("|"), ""); + UNIT_ASSERT_VALUES_EQUAL(NFq::SanitizeLabel("*"), ""); + UNIT_ASSERT_VALUES_EQUAL(NFq::SanitizeLabel("?"), ""); + UNIT_ASSERT_VALUES_EQUAL(NFq::SanitizeLabel("\""), ""); + UNIT_ASSERT_VALUES_EQUAL(NFq::SanitizeLabel("'"), ""); + UNIT_ASSERT_VALUES_EQUAL(NFq::SanitizeLabel("`"), ""); + UNIT_ASSERT_VALUES_EQUAL(NFq::SanitizeLabel("\\"), ""); + } + + Y_UNIT_TEST(SkipBadSymbols) { + UNIT_ASSERT_VALUES_EQUAL(NFq::SanitizeLabel("a|b*c?d\"e'f`g\\h"), "abcdefgh"); + } + + Y_UNIT_TEST(Truncate200) { + TString s1(400, 'a'); + TString s2(200, 'a'); + UNIT_ASSERT_VALUES_EQUAL(NFq::SanitizeLabel(s1), s2); + } +} diff --git a/ydb/core/fq/libs/metrics/ut/ya.make b/ydb/core/fq/libs/metrics/ut/ya.make index cb0311da95e6..829ad1258ff1 100644 --- a/ydb/core/fq/libs/metrics/ut/ya.make +++ b/ydb/core/fq/libs/metrics/ut/ya.make @@ -1,13 +1,12 @@ UNITTEST_FOR(ydb/core/fq/libs/metrics) -FORK_SUBTESTS() - IF (SANITIZER_TYPE OR WITH_VALGRIND) SIZE(MEDIUM) ENDIF() SRCS( metrics_ut.cpp + sanitize_label_ut.cpp ) YQL_LAST_ABI_VERSION() diff --git a/ydb/core/fq/libs/metrics/ya.make b/ydb/core/fq/libs/metrics/ya.make index fa7443622a6d..e1bafbd6d295 100644 --- a/ydb/core/fq/libs/metrics/ya.make +++ b/ydb/core/fq/libs/metrics/ya.make @@ -1,13 +1,14 @@ LIBRARY() SRCS( + sanitize_label.cpp status_code_counters.cpp ) PEERDIR( library/cpp/monlib/dynamic_counters - ydb/library/yql/dq/actors/protos yql/essentials/public/issue + ydb/library/yql/dq/actors/protos ) YQL_LAST_ABI_VERSION() diff --git a/ydb/core/fq/libs/read_rule/read_rule_creator.cpp b/ydb/core/fq/libs/read_rule/read_rule_creator.cpp index e926018a2be5..fb1f111ad9e0 100644 --- a/ydb/core/fq/libs/read_rule/read_rule_creator.cpp +++ b/ydb/core/fq/libs/read_rule/read_rule_creator.cpp @@ -71,6 +71,7 @@ class TSingleReadRuleCreator : public TActorBootstrapped NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const Fq::Private::TopicConsumer& topicConsumer, std::shared_ptr credentialsProvider, ui64 index @@ -79,6 +80,7 @@ class TSingleReadRuleCreator : public TActorBootstrapped , QueryId(std::move(queryId)) , TopicConsumer(topicConsumer) , YdbDriver(std::move(ydbDriver)) + , PqGateway(pqGateway) , TopicClient(YdbDriver, GetTopicClientSettings(std::move(credentialsProvider))) , Index(index) { @@ -183,7 +185,7 @@ class TSingleReadRuleCreator : public TActorBootstrapped private: NYdb::NTopic::TTopicClientSettings GetTopicClientSettings(std::shared_ptr credentialsProvider) { - return NYdb::NTopic::TTopicClientSettings() + return PqGateway->GetTopicClientSettings() .Database(TopicConsumer.database()) .DiscoveryEndpoint(TopicConsumer.cluster_endpoint()) .CredentialsProviderFactory(std::move(credentialsProvider)) @@ -196,6 +198,7 @@ class TSingleReadRuleCreator : public TActorBootstrapped const TString QueryId; const Fq::Private::TopicConsumer TopicConsumer; NYdb::TDriver YdbDriver; + NYql::IPqGateway::TPtr PqGateway; NYdb::NTopic::TTopicClient TopicClient; ui64 Index = 0; NYdb::NTopic::IRetryPolicy::IRetryState::TPtr RetryState; @@ -210,12 +213,14 @@ class TReadRuleCreator : public TActorBootstrapped { NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials ) : Owner(owner) , QueryId(std::move(queryId)) , YdbDriver(std::move(ydbDriver)) + , PqGateway(pqGateway) , TopicConsumers(VectorFromProto(topicConsumers)) , Credentials(std::move(credentials)) { @@ -232,7 +237,7 @@ class TReadRuleCreator : public TActorBootstrapped { Results.reserve(TopicConsumers.size()); for (size_t i = 0; i < TopicConsumers.size(); ++i) { LOG_D("Create read rule creation actor for `" << TopicConsumers[i].topic_path() << "` [" << i << "]"); - Children.push_back(Register(new TSingleReadRuleCreator(SelfId(), QueryId, YdbDriver, TopicConsumers[i], Credentials[i], i))); + Children.push_back(Register(new TSingleReadRuleCreator(SelfId(), QueryId, YdbDriver, PqGateway, TopicConsumers[i], Credentials[i], i))); } } @@ -281,6 +286,7 @@ class TReadRuleCreator : public TActorBootstrapped { const NActors::TActorId Owner; const TString QueryId; NYdb::TDriver YdbDriver; + NYql::IPqGateway::TPtr PqGateway; const TVector TopicConsumers; const TVector> Credentials; size_t ResultsGot = 0; @@ -295,6 +301,7 @@ NActors::IActor* MakeReadRuleCreatorActor( NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials ) @@ -303,6 +310,7 @@ NActors::IActor* MakeReadRuleCreatorActor( owner, std::move(queryId), std::move(ydbDriver), + pqGateway, topicConsumers, std::move(credentials) ); diff --git a/ydb/core/fq/libs/read_rule/read_rule_creator.h b/ydb/core/fq/libs/read_rule/read_rule_creator.h index 51873fb98272..87b4b4ede5d2 100644 --- a/ydb/core/fq/libs/read_rule/read_rule_creator.h +++ b/ydb/core/fq/libs/read_rule/read_rule_creator.h @@ -4,6 +4,7 @@ #include #include +#include namespace NFq { @@ -11,6 +12,7 @@ NActors::IActor* MakeReadRuleCreatorActor( NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials // For each topic ); diff --git a/ydb/core/fq/libs/read_rule/read_rule_deleter.cpp b/ydb/core/fq/libs/read_rule/read_rule_deleter.cpp index 7a3837fbbe02..970d87f838a0 100644 --- a/ydb/core/fq/libs/read_rule/read_rule_deleter.cpp +++ b/ydb/core/fq/libs/read_rule/read_rule_deleter.cpp @@ -67,6 +67,7 @@ class TSingleReadRuleDeleter : public TActorBootstrapped NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, Fq::Private::TopicConsumer topic, std::shared_ptr credentialsProvider, ui64 index, @@ -76,6 +77,7 @@ class TSingleReadRuleDeleter : public TActorBootstrapped , QueryId(std::move(queryId)) , Topic(std::move(topic)) , YdbDriver(std::move(ydbDriver)) + , PqGateway(pqGateway) , TopicClient(YdbDriver, GetTopicClientSettings(std::move(credentialsProvider))) , Index(index) , MaxRetries(maxRetries) @@ -158,7 +160,7 @@ class TSingleReadRuleDeleter : public TActorBootstrapped private: NYdb::NTopic::TTopicClientSettings GetTopicClientSettings(std::shared_ptr credentialsProvider) { - return NYdb::NTopic::TTopicClientSettings() + return PqGateway->GetTopicClientSettings() .Database(Topic.database()) .DiscoveryEndpoint(Topic.cluster_endpoint()) .CredentialsProviderFactory(std::move(credentialsProvider)) @@ -171,6 +173,7 @@ class TSingleReadRuleDeleter : public TActorBootstrapped const TString QueryId; const Fq::Private::TopicConsumer Topic; NYdb::TDriver YdbDriver; + NYql::IPqGateway::TPtr PqGateway; NYdb::NTopic::TTopicClient TopicClient; ui64 Index = 0; const size_t MaxRetries; @@ -184,6 +187,7 @@ class TReadRuleDeleter : public TActorBootstrapped { NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials, size_t maxRetries @@ -191,6 +195,7 @@ class TReadRuleDeleter : public TActorBootstrapped { : Owner(owner) , QueryId(std::move(queryId)) , YdbDriver(std::move(ydbDriver)) + , PqGateway(pqGateway) , Topics(VectorFromProto(topicConsumers)) , Credentials(std::move(credentials)) , MaxRetries(maxRetries) @@ -206,7 +211,7 @@ class TReadRuleDeleter : public TActorBootstrapped { Results.reserve(Topics.size()); for (size_t i = 0; i < Topics.size(); ++i) { LOG_D("Create read rule deleter actor for `" << Topics[i].topic_path() << "` [" << i << "]"); - Children.push_back(Register(new TSingleReadRuleDeleter(SelfId(), QueryId, YdbDriver, Topics[i], Credentials[i], i, MaxRetries))); + Children.push_back(Register(new TSingleReadRuleDeleter(SelfId(), QueryId, YdbDriver, PqGateway, Topics[i], Credentials[i], i, MaxRetries))); } } @@ -257,6 +262,7 @@ class TReadRuleDeleter : public TActorBootstrapped { const NActors::TActorId Owner; const TString QueryId; NYdb::TDriver YdbDriver; + NYql::IPqGateway::TPtr PqGateway; const TVector Topics; const TVector> Credentials; const size_t MaxRetries; @@ -272,6 +278,7 @@ NActors::IActor* MakeReadRuleDeleterActor( NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials, // For each topic size_t maxRetries @@ -281,6 +288,7 @@ NActors::IActor* MakeReadRuleDeleterActor( owner, std::move(queryId), std::move(ydbDriver), + pqGateway, topicConsumers, std::move(credentials), maxRetries diff --git a/ydb/core/fq/libs/read_rule/read_rule_deleter.h b/ydb/core/fq/libs/read_rule/read_rule_deleter.h index bcf332b0f81e..a23e23260248 100644 --- a/ydb/core/fq/libs/read_rule/read_rule_deleter.h +++ b/ydb/core/fq/libs/read_rule/read_rule_deleter.h @@ -4,6 +4,7 @@ #include #include +#include namespace NFq { @@ -11,6 +12,7 @@ NActors::IActor* MakeReadRuleDeleterActor( NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials, // For each topic size_t maxRetries = 15 diff --git a/ydb/core/fq/libs/row_dispatcher/actors_factory.cpp b/ydb/core/fq/libs/row_dispatcher/actors_factory.cpp index 6e1669e92e8e..e77be57b9107 100644 --- a/ydb/core/fq/libs/row_dispatcher/actors_factory.cpp +++ b/ydb/core/fq/libs/row_dispatcher/actors_factory.cpp @@ -9,6 +9,7 @@ struct TActorFactory : public IActorFactory { TActorFactory() {} NActors::TActorId RegisterTopicSession( + const TString& readGroup, const TString& topicPath, const TString& endpoint, const TString& database, @@ -19,10 +20,12 @@ struct TActorFactory : public IActorFactory { NYdb::TDriver driver, std::shared_ptr credentialsProviderFactory, const ::NMonitoring::TDynamicCounterPtr& counters, + const ::NMonitoring::TDynamicCounterPtr& countersRoot, const NYql::IPqGateway::TPtr& pqGateway, ui64 maxBufferSize) const override { auto actorPtr = NFq::NewTopicSession( + readGroup, topicPath, endpoint, database, @@ -33,6 +36,7 @@ struct TActorFactory : public IActorFactory { std::move(driver), credentialsProviderFactory, counters, + countersRoot, pqGateway, maxBufferSize ); diff --git a/ydb/core/fq/libs/row_dispatcher/actors_factory.h b/ydb/core/fq/libs/row_dispatcher/actors_factory.h index 0c0836949f89..9c7dc0fac19f 100644 --- a/ydb/core/fq/libs/row_dispatcher/actors_factory.h +++ b/ydb/core/fq/libs/row_dispatcher/actors_factory.h @@ -12,6 +12,7 @@ struct IActorFactory : public TThrRefBase { using TPtr = TIntrusivePtr; virtual NActors::TActorId RegisterTopicSession( + const TString& readGroup, const TString& topicPath, const TString& endpoint, const TString& database, @@ -22,6 +23,7 @@ struct IActorFactory : public TThrRefBase { NYdb::TDriver driver, std::shared_ptr credentialsProviderFactory, const ::NMonitoring::TDynamicCounterPtr& counters, + const ::NMonitoring::TDynamicCounterPtr& countersRoot, const NYql::IPqGateway::TPtr& pqGateway, ui64 maxBufferSize) const = 0; }; diff --git a/ydb/core/fq/libs/row_dispatcher/common.h b/ydb/core/fq/libs/row_dispatcher/common.h deleted file mode 100644 index dd9b18c6c9a1..000000000000 --- a/ydb/core/fq/libs/row_dispatcher/common.h +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include - -namespace NFq { - -TString CleanupCounterValueString(const TString& value); - -} // namespace NFq diff --git a/ydb/core/fq/libs/row_dispatcher/coordinator.cpp b/ydb/core/fq/libs/row_dispatcher/coordinator.cpp index 385fad09b763..62e2d734945a 100644 --- a/ydb/core/fq/libs/row_dispatcher/coordinator.cpp +++ b/ydb/core/fq/libs/row_dispatcher/coordinator.cpp @@ -51,23 +51,42 @@ struct TEvPrivate { class TActorCoordinator : public TActorBootstrapped { const ui64 PrintStatePeriodSec = 300; + const ui64 PrintStateToLogSplitSize = 64000; - struct TPartitionKey { + struct TTopicKey { TString Endpoint; TString Database; TString TopicName; - ui64 PartitionId; size_t Hash() const noexcept { ui64 hash = std::hash()(Endpoint); hash = CombineHashes(hash, std::hash()(Database)); hash = CombineHashes(hash, std::hash()(TopicName)); + return hash; + } + bool operator==(const TTopicKey& other) const { + return Endpoint == other.Endpoint && Database == other.Database + && TopicName == other.TopicName; + } + }; + + struct TTopicKeyHash { + int operator()(const TTopicKey& k) const { + return k.Hash(); + } + }; + + struct TPartitionKey { + TTopicKey Topic; + ui64 PartitionId; + + size_t Hash() const noexcept { + ui64 hash = Topic.Hash(); hash = CombineHashes(hash, std::hash()(PartitionId)); return hash; } bool operator==(const TPartitionKey& other) const { - return Endpoint == other.Endpoint && Database == other.Database - && TopicName == other.TopicName && PartitionId == other.PartitionId; + return Topic == other.Topic && PartitionId == other.PartitionId; } }; @@ -156,7 +175,7 @@ class TActorCoordinator : public TActorBootstrapped { const TString Tenant; TMap RowDispatchers; THashMap PartitionLocations; - THashMap TopicsInfo; + THashMap TopicsInfo; std::unordered_map PendingReadActors; TCoordinatorMetrics Metrics; THashSet InterconnectSessions; @@ -196,11 +215,12 @@ class TActorCoordinator : public TActorBootstrapped { void AddRowDispatcher(NActors::TActorId actorId, bool isLocal); void PrintInternalState(); - TTopicInfo& GetOrCreateTopicInfo(const TString& topicName); + TTopicInfo& GetOrCreateTopicInfo(const TTopicKey& topic); std::optional GetAndUpdateLocation(const TPartitionKey& key); // std::nullopt if TopicPartitionsLimitPerNode reached bool ComputeCoordinatorRequest(TActorId readActorId, const TCoordinatorRequest& request); void UpdatePendingReadActors(); void UpdateInterconnectSessions(const NActors::TActorId& interconnectSession); + TString GetInternalState(); }; TActorCoordinator::TActorCoordinator( @@ -278,7 +298,7 @@ void TActorCoordinator::Handle(NActors::TEvents::TEvPing::TPtr& ev) { Send(ev->Sender, new NActors::TEvents::TEvPong(), IEventHandle::FlagTrackDelivery); } -void TActorCoordinator::PrintInternalState() { +TString TActorCoordinator::GetInternalState() { TStringStream str; str << "Known row dispatchers:\n"; @@ -288,15 +308,22 @@ void TActorCoordinator::PrintInternalState() { str << "\nLocations:\n"; for (auto& [key, actorId] : PartitionLocations) { - str << " " << key.Endpoint << " / " << key.Database << " / " << key.TopicName << ", partId " << key.PartitionId << ", row dispatcher actor id: " << actorId << "\n"; + str << " " << key.Topic.Endpoint << " / " << key.Topic.Database << " / " << key.Topic.TopicName << ", partId " << key.PartitionId << ", row dispatcher actor id: " << actorId << "\n"; } str << "\nPending partitions:\n"; - for (const auto& [topicName, topicInfo] : TopicsInfo) { - str << " " << topicName << ", pending partitions: " << topicInfo.PendingPartitions.size() << "\n"; + for (const auto& [topic, topicInfo] : TopicsInfo) { + str << " " << topic.TopicName << " (" << topic.Endpoint << "), pending partitions: " << topicInfo.PendingPartitions.size() << "\n"; } + return str.Str(); +} - LOG_ROW_DISPATCHER_DEBUG(str.Str()); +void TActorCoordinator::PrintInternalState() { + auto str = GetInternalState(); + auto buf = TStringBuf(str); + for (ui64 offset = 0; offset < buf.size(); offset += PrintStateToLogSplitSize) { + LOG_ROW_DISPATCHER_DEBUG(buf.SubString(offset, PrintStateToLogSplitSize)); + } } void TActorCoordinator::HandleConnected(TEvInterconnect::TEvNodeConnected::TPtr& ev) { @@ -337,18 +364,18 @@ void TActorCoordinator::Handle(NFq::TEvRowDispatcher::TEvCoordinatorChanged::TPt Metrics.IsActive->Set(isActive); } -TActorCoordinator::TTopicInfo& TActorCoordinator::GetOrCreateTopicInfo(const TString& topicName) { - const auto it = TopicsInfo.find(topicName); +TActorCoordinator::TTopicInfo& TActorCoordinator::GetOrCreateTopicInfo(const TTopicKey& topic) { + const auto it = TopicsInfo.find(topic); if (it != TopicsInfo.end()) { return it->second; } - return TopicsInfo.insert({topicName, TTopicInfo(Metrics, topicName)}).first->second; + return TopicsInfo.insert({topic, TTopicInfo(Metrics, topic.TopicName)}).first->second; } std::optional TActorCoordinator::GetAndUpdateLocation(const TPartitionKey& key) { Y_ENSURE(!PartitionLocations.contains(key)); - auto& topicInfo = GetOrCreateTopicInfo(key.TopicName); + auto& topicInfo = GetOrCreateTopicInfo(key.Topic); TActorId bestLocation; ui64 bestNumberPartitions = std::numeric_limits::max(); @@ -391,17 +418,14 @@ void TActorCoordinator::Handle(NFq::TEvRowDispatcher::TEvCoordinatorRequest::TPt UpdateInterconnectSessions(ev->InterconnectSession); TStringStream str; - str << "TEvCoordinatorRequest from " << ev->Sender.ToString() << ", " << source.GetTopicPath() << ", partIds: "; - for (auto& partitionId : ev->Get()->Record.GetPartitionId()) { - str << partitionId << ", "; - } - LOG_ROW_DISPATCHER_DEBUG(str.Str()); + LOG_ROW_DISPATCHER_INFO("TEvCoordinatorRequest from " << ev->Sender.ToString() << ", " << source.GetTopicPath() << ", partIds: " << JoinSeq(", ", ev->Get()->Record.GetPartitionIds())); Metrics.IncomingRequests->Inc(); TCoordinatorRequest request = {.Cookie = ev->Cookie, .Record = ev->Get()->Record}; if (ComputeCoordinatorRequest(ev->Sender, request)) { PendingReadActors.erase(ev->Sender); } else { + LOG_ROW_DISPATCHER_INFO("All nodes are overloaded, add request into pending queue"); // All nodes are overloaded, add request into pending queue // We save only last request from each read actor PendingReadActors[ev->Sender] = request; @@ -415,8 +439,9 @@ bool TActorCoordinator::ComputeCoordinatorRequest(TActorId readActorId, const TC bool hasPendingPartitions = false; TMap> tmpResult; - for (auto& partitionId : request.Record.GetPartitionId()) { - TPartitionKey key{source.GetEndpoint(), source.GetDatabase(), source.GetTopicPath(), partitionId}; + for (auto& partitionId : request.Record.GetPartitionIds()) { + TTopicKey topicKey{source.GetEndpoint(), source.GetDatabase(), source.GetTopicPath()}; + TPartitionKey key {topicKey, partitionId}; auto locationIt = PartitionLocations.find(key); NActors::TActorId rowDispatcherId; if (locationIt != PartitionLocations.end()) { @@ -441,7 +466,7 @@ bool TActorCoordinator::ComputeCoordinatorRequest(TActorId readActorId, const TC auto* partitionsProto = response->Record.AddPartitions(); ActorIdToProto(actorId, partitionsProto->MutableActorId()); for (auto partitionId : partitions) { - partitionsProto->AddPartitionId(partitionId); + partitionsProto->AddPartitionIds(partitionId); } } diff --git a/ydb/core/fq/libs/row_dispatcher/events/data_plane.h b/ydb/core/fq/libs/row_dispatcher/events/data_plane.h index 88953d424432..593a3ed1cc07 100644 --- a/ydb/core/fq/libs/row_dispatcher/events/data_plane.h +++ b/ydb/core/fq/libs/row_dispatcher/events/data_plane.h @@ -3,13 +3,16 @@ #include #include #include - #include #include #include +#include #include +#include +#include + namespace NFq { NActors::TActorId RowDispatcherServiceActorId(); @@ -46,10 +49,12 @@ struct TEvRowDispatcher { EvCoordinatorResult, EvSessionStatistic, EvHeartbeat, + EvNoSession, EvGetInternalStateRequest, EvGetInternalStateResponse, EvPurecalcCompileRequest, EvPurecalcCompileResponse, + EvPurecalcCompileAbort, EvEnd, }; @@ -72,7 +77,7 @@ struct TEvRowDispatcher { const std::vector& partitionIds) { *Record.MutableSource() = sourceParams; for (const auto& id : partitionIds) { - Record.AddPartitionId(id); + Record.AddPartitionIds(id); } } }; @@ -82,22 +87,28 @@ struct TEvRowDispatcher { TEvCoordinatorResult() = default; }; +// Session events (with seqNo checks) + struct TEvStartSession : public NActors::TEventPB { TEvStartSession() = default; TEvStartSession( const NYql::NPq::NProto::TDqPqTopicSource& sourceParams, - ui64 partitionId, + const std::set& partitionIds, const TString token, - TMaybe readOffset, + const std::map& readOffsets, ui64 startingMessageTimestampMs, const TString& queryId) { *Record.MutableSource() = sourceParams; - Record.SetPartitionId(partitionId); + for (auto partitionId : partitionIds) { + Record.AddPartitionIds(partitionId); + } Record.SetToken(token); - if (readOffset) { - Record.SetOffset(*readOffset); + for (const auto& [partitionId, offset] : readOffsets) { + auto* partitionOffset = Record.AddOffsets(); + partitionOffset->SetPartitionId(partitionId); + partitionOffset->SetOffset(offset); } Record.SetStartingMessageTimestampMs(startingMessageTimestampMs); Record.SetQueryId(queryId); @@ -138,7 +149,6 @@ struct TEvRowDispatcher { struct TEvStatistics : public NActors::TEventPB { TEvStatistics() = default; - NActors::TActorId ReadActorId; }; struct TEvSessionError : public NActors::TEventPB { - TEvSessionStatistic(const TopicSessionStatistic& stat) + TEvSessionStatistic(const TTopicSessionStatistic& stat) : Stat(stat) {} - TopicSessionStatistic Stat; + TTopicSessionStatistic Stat; }; + // two purposes: confirm seqNo and check the availability of the recipient actor (wait TEvUndelivered) struct TEvHeartbeat : public NActors::TEventPB { TEvHeartbeat() = default; - TEvHeartbeat(ui32 partitionId) { - Record.SetPartitionId(partitionId); - } + }; + +// Network events (without seqNo checks) + + struct TEvNoSession : public NActors::TEventPB { + TEvNoSession() = default; }; struct TEvGetInternalStateRequest : public NActors::TEventPB { - explicit TEvPurecalcCompileResponse(const TString& error) - : Error(error) + TEvPurecalcCompileResponse(NYql::NDqProto::StatusIds::StatusCode status, NYql::TIssues issues) + : Status(status) + , Issues(std::move(issues)) {} explicit TEvPurecalcCompileResponse(IProgramHolder::TPtr programHolder) : ProgramHolder(std::move(programHolder)) + , Status(NYql::NDqProto::StatusIds::SUCCESS) {} IProgramHolder::TPtr ProgramHolder; // Same holder that passed into TEvPurecalcCompileRequest - TString Error; + NYql::NDqProto::StatusIds::StatusCode Status; + NYql::TIssues Issues; }; + + struct TEvPurecalcCompileAbort : public NActors::TEventLocal {}; }; } // namespace NFq diff --git a/ydb/core/fq/libs/row_dispatcher/events/topic_session_stats.h b/ydb/core/fq/libs/row_dispatcher/events/topic_session_stats.h index 68075b4a92e0..5aa71084b316 100644 --- a/ydb/core/fq/libs/row_dispatcher/events/topic_session_stats.h +++ b/ydb/core/fq/libs/row_dispatcher/events/topic_session_stats.h @@ -5,62 +5,105 @@ namespace NFq { -struct TopicSessionClientStatistic { +struct TTopicSessionClientStatistic { NActors::TActorId ReadActorId; ui32 PartitionId = 0; - i64 UnreadRows = 0; // Current value - i64 UnreadBytes = 0; // Current value + i64 QueuedRows = 0; // Current value + i64 QueuedBytes = 0; // Current value ui64 Offset = 0; // Current value - ui64 ReadBytes = 0; // Increment / filtered + ui64 FilteredBytes = 0; // Increment / filtered + ui64 FilteredRows = 0; // Increment / filtered + ui64 ReadBytes = 0; // Increment bool IsWaiting = false; // Current value i64 ReadLagMessages = 0; // Current value ui64 InitialOffset = 0; - void Add(const TopicSessionClientStatistic& stat) { - UnreadRows = stat.UnreadRows; - UnreadBytes = stat.UnreadBytes; + void Add(const TTopicSessionClientStatistic& stat) { + QueuedRows = stat.QueuedRows; + QueuedBytes = stat.QueuedBytes; Offset = stat.Offset; + FilteredBytes += stat.FilteredBytes; + FilteredRows += stat.FilteredRows; ReadBytes += stat.ReadBytes; IsWaiting = stat.IsWaiting; ReadLagMessages = stat.ReadLagMessages; InitialOffset = stat.InitialOffset; } void Clear() { + FilteredBytes = 0; + FilteredRows = 0; ReadBytes = 0; } }; -struct TopicSessionCommonStatistic { - ui64 UnreadBytes = 0; // Current value +struct TParserStatistic { + TDuration ParserLatency; + + void Add(const TParserStatistic& stat) { + ParserLatency = stat.ParserLatency != TDuration::Zero() ? stat.ParserLatency : ParserLatency; + } +}; + +struct TFiltersStatistic { + TDuration FilterLatency; + + void Add(const TFiltersStatistic& stat) { + FilterLatency = stat.FilterLatency != TDuration::Zero() ? stat.FilterLatency : FilterLatency; + } +}; + +struct TFormatHandlerStatistic { + TDuration ParseAndFilterLatency; + + TParserStatistic ParserStats; + TFiltersStatistic FilterStats; + + void Add(const TFormatHandlerStatistic& stat) { + ParseAndFilterLatency = stat.ParseAndFilterLatency != TDuration::Zero() ? stat.ParseAndFilterLatency : ParseAndFilterLatency; + + ParserStats.Add(stat.ParserStats); + FilterStats.Add(stat.FilterStats); + } +}; + +struct TTopicSessionCommonStatistic { + ui64 QueuedBytes = 0; // Current value ui64 RestartSessionByOffsets = 0; ui64 ReadBytes = 0; // Increment ui64 ReadEvents = 0; // Increment ui64 LastReadedOffset = 0; - TDuration ParseAndFilterLatency; - void Add(const TopicSessionCommonStatistic& stat) { - UnreadBytes = stat.UnreadBytes; + + std::unordered_map FormatHandlers; + + void Add(const TTopicSessionCommonStatistic& stat) { + QueuedBytes = stat.QueuedBytes; RestartSessionByOffsets = stat.RestartSessionByOffsets; ReadBytes += stat.ReadBytes; ReadEvents += stat.ReadEvents; LastReadedOffset = stat.LastReadedOffset; - ParseAndFilterLatency = stat.ParseAndFilterLatency != TDuration::Zero() ? stat.ParseAndFilterLatency : ParseAndFilterLatency; + + for (const auto& [formatName, foramtStats] : stat.FormatHandlers) { + FormatHandlers[formatName].Add(foramtStats); + } } + void Clear() { ReadBytes = 0; ReadEvents = 0; } }; -struct TopicSessionParams { +struct TTopicSessionParams { + TString ReadGroup; TString Endpoint; TString Database; TString TopicPath; ui64 PartitionId = 0; }; -struct TopicSessionStatistic { - TopicSessionParams SessionKey; - TVector Clients; - TopicSessionCommonStatistic Common; +struct TTopicSessionStatistic { + TTopicSessionParams SessionKey; + std::vector Clients; + TTopicSessionCommonStatistic Common; }; } // namespace NFq diff --git a/ydb/core/fq/libs/row_dispatcher/events/ya.make b/ydb/core/fq/libs/row_dispatcher/events/ya.make index 60f0b00e7e90..2968bf590ebe 100644 --- a/ydb/core/fq/libs/row_dispatcher/events/ya.make +++ b/ydb/core/fq/libs/row_dispatcher/events/ya.make @@ -7,8 +7,11 @@ SRCS( PEERDIR( ydb/core/fq/libs/events ydb/core/fq/libs/row_dispatcher/protos + ydb/library/actors/core ydb/library/yql/providers/pq/provider + + yql/essentials/public/issue ) END() diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/common/common.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/common/common.cpp new file mode 100644 index 000000000000..cfb7526bdcad --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/common/common.cpp @@ -0,0 +1,21 @@ +#include "common.h" + +#include + +namespace NFq::NRowDispatcher { + +//// TSchemaColumn + +TString TSchemaColumn::ToString() const { + return TStringBuilder() << "'" << Name << "' : " << TypeYson; +} + +//// TCountersDesc + +TCountersDesc TCountersDesc::CopyWithNewMkqlCountersName(const TString& mkqlCountersName) const { + TCountersDesc result(*this); + result.MkqlCountersName = mkqlCountersName; + return result; +} + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/common/common.h b/ydb/core/fq/libs/row_dispatcher/format_handler/common/common.h new file mode 100644 index 000000000000..dbd5e10bc249 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/common/common.h @@ -0,0 +1,35 @@ +#pragma once + +#include + +#include +#include +#include + +namespace NFq::NRowDispatcher { + +using EStatusId = NYql::NDqProto::StatusIds; +using TStatusCode = EStatusId::StatusCode; +using TStatus = NKikimr::TYQLConclusionSpecialStatus; + +template +using TValueStatus = NKikimr::TConclusionImpl; + +struct TSchemaColumn { + TString Name; + TString TypeYson; + + bool operator==(const TSchemaColumn& other) const = default; + + TString ToString() const; +}; + +struct TCountersDesc { + NMonitoring::TDynamicCounterPtr CountersRoot = MakeIntrusive(); + NMonitoring::TDynamicCounterPtr CountersSubgroup = MakeIntrusive(); + TString MkqlCountersName; // Used for TAlignedPagePoolCounters created from CountersRoot + + [[nodiscard]] TCountersDesc CopyWithNewMkqlCountersName(const TString& mkqlCountersName) const; +}; + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/common/ya.make b/ydb/core/fq/libs/row_dispatcher/format_handler/common/ya.make new file mode 100644 index 000000000000..1ac30eda68a9 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/common/ya.make @@ -0,0 +1,18 @@ +LIBRARY() + +SRCS( + common.cpp +) + +PEERDIR( + library/cpp/monlib/dynamic_counters + + ydb/library/conclusion + ydb/library/yql/dq/actors/protos + + yql/essentials/public/issue +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/filters/filters_set.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/filters/filters_set.cpp new file mode 100644 index 000000000000..571820fd6a59 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/filters/filters_set.cpp @@ -0,0 +1,287 @@ +#include "filters_set.h" + +#include + +namespace NFq::NRowDispatcher { + +namespace { + +class TTopicFilters : public ITopicFilters { + struct TCounters { + const NMonitoring::TDynamicCounterPtr Counters; + + NMonitoring::TDynamicCounters::TCounterPtr ActiveFilters; + NMonitoring::TDynamicCounters::TCounterPtr InFlightCompileRequests; + NMonitoring::TDynamicCounters::TCounterPtr CompileErrors; + + explicit TCounters(NMonitoring::TDynamicCounterPtr counters) + : Counters(counters) + { + Register(); + } + + private: + void Register() { + ActiveFilters = Counters->GetCounter("ActiveFilters", false); + InFlightCompileRequests = Counters->GetCounter("InFlightCompileRequests", false); + CompileErrors = Counters->GetCounter("CompileErrors", true); + } + }; + + struct TStats { + void AddFilterLatency(TDuration filterLatency) { + FilterLatency = std::max(FilterLatency, filterLatency); + } + + void Clear() { + FilterLatency = TDuration::Zero(); + } + + TDuration FilterLatency; + }; + + class TFilterHandler { + public: + TFilterHandler(TTopicFilters& self, IFilteredDataConsumer::TPtr consumer, IPurecalcFilter::TPtr purecalcFilter) + : Self(self) + , Consumer(std::move(consumer)) + , PurecalcFilter(std::move(purecalcFilter)) + , LogPrefix(TStringBuilder() << Self.LogPrefix << "TFilterHandler " << Consumer->GetFilterId() << " : ") + {} + + ~TFilterHandler() { + if (InFlightCompilationId) { + Self.Counters.InFlightCompileRequests->Dec(); + } else if (FilterStarted) { + Self.Counters.ActiveFilters->Dec(); + } + } + + IFilteredDataConsumer::TPtr GetConsumer() const { + return Consumer; + } + + IPurecalcFilter::TPtr GetPurecalcFilter() const { + return PurecalcFilter; + } + + bool IsStarted() const { + return FilterStarted; + } + + void CompileFilter() { + if (!PurecalcFilter) { + StartFilter(); + return; + } + + InFlightCompilationId = Self.FreeCompileId++; + Self.Counters.InFlightCompileRequests->Inc(); + Y_ENSURE(Self.InFlightCompilations.emplace(InFlightCompilationId, Consumer->GetFilterId()).second, "Got duplicated compilation event id"); + + LOG_ROW_DISPATCHER_TRACE("Send compile request with id " << InFlightCompilationId); + NActors::TActivationContext::ActorSystem()->Send(new NActors::IEventHandle(Self.Config.CompileServiceId, Self.Owner, PurecalcFilter->GetCompileRequest().release(), 0, InFlightCompilationId)); + } + + void AbortCompilation() { + if (!InFlightCompilationId) { + return; + } + + LOG_ROW_DISPATCHER_TRACE("Send abort compile request with id " << InFlightCompilationId); + NActors::TActivationContext::ActorSystem()->Send(new NActors::IEventHandle(Self.Config.CompileServiceId, Self.Owner, new TEvRowDispatcher::TEvPurecalcCompileAbort(), 0, InFlightCompilationId)); + + InFlightCompilationId = 0; + Self.Counters.InFlightCompileRequests->Dec(); + } + + void OnCompileResponse(TEvRowDispatcher::TEvPurecalcCompileResponse::TPtr ev) { + if (ev->Cookie != InFlightCompilationId) { + LOG_ROW_DISPATCHER_DEBUG("Outdated compiler response ignored for id " << ev->Cookie << ", current compile id " << InFlightCompilationId); + return; + } + + Y_ENSURE(InFlightCompilationId, "Unexpected compilation response"); + InFlightCompilationId = 0; + Self.Counters.InFlightCompileRequests->Dec(); + + if (!ev->Get()->ProgramHolder) { + auto status = TStatus::Fail(ev->Get()->Status, std::move(ev->Get()->Issues)); + LOG_ROW_DISPATCHER_ERROR("Filter compilation error: " << status.GetErrorMessage()); + + Self.Counters.CompileErrors->Inc(); + Consumer->OnFilteringError(status.AddParentIssue("Failed to compile client filter")); + return; + } + + Y_ENSURE(PurecalcFilter, "Unexpected compilation response for client without filter"); + PurecalcFilter->OnCompileResponse(std::move(ev)); + + LOG_ROW_DISPATCHER_TRACE("Filter compilation finished"); + StartFilter(); + } + + private: + void StartFilter() { + FilterStarted = true; + Self.Counters.ActiveFilters->Inc(); + Consumer->OnFilterStarted(); + } + + private: + TTopicFilters& Self; + const IFilteredDataConsumer::TPtr Consumer; + const IPurecalcFilter::TPtr PurecalcFilter; + const TString LogPrefix; + + ui64 InFlightCompilationId = 0; + bool FilterStarted = false; + }; + +public: + explicit TTopicFilters(NActors::TActorId owner, const TTopicFiltersConfig& config, NMonitoring::TDynamicCounterPtr counters) + : Config(config) + , Owner(owner) + , LogPrefix("TTopicFilters: ") + , Counters(counters) + {} + + ~TTopicFilters() { + Filters.clear(); + } + +public: + void FilterData(const TVector& columnIndex, const TVector& offsets, const TVector*>& values, ui64 numberRows) override { + LOG_ROW_DISPATCHER_TRACE("FilterData for " << Filters.size() << " clients, number rows: " << numberRows); + + const TInstant startFilter = TInstant::Now(); + for (const auto& [_, filterHandler] : Filters) { + const auto consumer = filterHandler.GetConsumer(); + if (const auto nextOffset = consumer->GetNextMessageOffset(); !numberRows || (nextOffset && offsets[numberRows - 1] < *nextOffset)) { + LOG_ROW_DISPATCHER_TRACE("Ignore filtering for " << consumer->GetFilterId() << ", historical offset"); + continue; + } + if (!filterHandler.IsStarted()) { + LOG_ROW_DISPATCHER_TRACE("Ignore filtering for " << consumer->GetFilterId() << ", client filter is not compiled"); + continue; + } + + PushToFilter(filterHandler, offsets, columnIndex, values, numberRows); + } + Stats.AddFilterLatency(TInstant::Now() - startFilter); + } + + void OnCompileResponse(TEvRowDispatcher::TEvPurecalcCompileResponse::TPtr ev) override { + LOG_ROW_DISPATCHER_TRACE("Got compile response for request with id " << ev->Cookie); + + const auto requestIt = InFlightCompilations.find(ev->Cookie); + if (requestIt == InFlightCompilations.end()) { + LOG_ROW_DISPATCHER_DEBUG("Compile response ignored for id " << ev->Cookie); + return; + } + + const auto filterId = requestIt->second; + InFlightCompilations.erase(requestIt); + + const auto filterIt = Filters.find(filterId); + if (filterIt == Filters.end()) { + LOG_ROW_DISPATCHER_DEBUG("Compile response ignored for id " << ev->Cookie << ", filter with id " << filterId << " not found"); + return; + } + + filterIt->second.OnCompileResponse(std::move(ev)); + } + + TStatus AddFilter(IFilteredDataConsumer::TPtr filter) override { + LOG_ROW_DISPATCHER_TRACE("Create filter with id " << filter->GetFilterId()); + + IPurecalcFilter::TPtr purecalcFilter; + if (const auto& predicate = filter->GetWhereFilter()) { + LOG_ROW_DISPATCHER_TRACE("Create purecalc filter for predicate '" << predicate << "' (filter id: " << filter->GetFilterId() << ")"); + + auto filterStatus = CreatePurecalcFilter(filter); + if (filterStatus.IsFail()) { + return filterStatus; + } + purecalcFilter = filterStatus.DetachResult(); + } + + const auto [it, inserted] = Filters.insert({filter->GetFilterId(), TFilterHandler(*this, filter, std::move(purecalcFilter))}); + if (!inserted) { + return TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Failed to create new filter, filter with id " << filter->GetFilterId() << " already exists"); + } + + it->second.CompileFilter(); + return TStatus::Success(); + } + + void RemoveFilter(NActors::TActorId filterId) override { + LOG_ROW_DISPATCHER_TRACE("Remove filter with id " << filterId); + + const auto it = Filters.find(filterId); + if (it == Filters.end()) { + return; + } + + it->second.AbortCompilation(); + Filters.erase(it); + } + + TFiltersStatistic GetStatistics() override { + TFiltersStatistic statistics; + statistics.FilterLatency = Stats.FilterLatency; + Stats.Clear(); + + return statistics; + } + +private: + void PushToFilter(const TFilterHandler& filterHandler, const TVector& offsets, const TVector& columnIndex, const TVector*>& values, ui64 numberRows) { + const auto consumer = filterHandler.GetConsumer(); + const auto& columnIds = consumer->GetColumnIds(); + + TVector*> result; + result.reserve(columnIds.size()); + for (ui64 columnId : columnIds) { + Y_ENSURE(columnId < columnIndex.size(), "Unexpected column id " << columnId << ", it is larger than index array size " << columnIndex.size()); + const ui64 index = columnIndex[columnId]; + + Y_ENSURE(index < values.size(), "Unexpected column index " << index << ", it is larger than values array size " << values.size()); + if (const auto value = values[index]) { + result.emplace_back(value); + } else { + LOG_ROW_DISPATCHER_TRACE("Ignore filtering for " << consumer->GetFilterId() << ", client got parsing error for column " << columnId); + return; + } + } + + if (const auto filter = filterHandler.GetPurecalcFilter()) { + LOG_ROW_DISPATCHER_TRACE("Pass " << numberRows << " rows to purecalc filter (filter id: " << consumer->GetFilterId() << ")"); + filter->FilterData(result, numberRows); + } else if (numberRows) { + LOG_ROW_DISPATCHER_TRACE("Add " << numberRows << " rows to client " << consumer->GetFilterId() << " without filtering"); + consumer->OnFilteredBatch(0, numberRows - 1); + } + } + +private: + const TTopicFiltersConfig Config; + const NActors::TActorId Owner; + const TString LogPrefix; + + ui64 FreeCompileId = 1; // 0 <=> compilation is not started + std::unordered_map InFlightCompilations; + std::unordered_map Filters; + + // Metrics + const TCounters Counters; + TStats Stats; +}; + +} // anonymous namespace + +ITopicFilters::TPtr CreateTopicFilters(NActors::TActorId owner, const TTopicFiltersConfig& config, NMonitoring::TDynamicCounterPtr counters) { + return MakeIntrusive(owner, config, counters); +} + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/filters/filters_set.h b/ydb/core/fq/libs/row_dispatcher/format_handler/filters/filters_set.h new file mode 100644 index 000000000000..57da103b1a76 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/filters/filters_set.h @@ -0,0 +1,47 @@ +#pragma once + +#include "purecalc_filter.h" + +#include + +#include + +namespace NFq::NRowDispatcher { + +class IFilteredDataConsumer : public IPurecalcFilterConsumer { +public: + using TPtr = TIntrusivePtr; + +public: + virtual NActors::TActorId GetFilterId() const = 0; + virtual const TVector& GetColumnIds() const = 0; + virtual TMaybe GetNextMessageOffset() const = 0; + + virtual void OnFilteredBatch(ui64 firstRow, ui64 lastRow) = 0; // inclusive interval [firstRow, lastRow] + + virtual void OnFilterStarted() = 0; + virtual void OnFilteringError(TStatus status) = 0; +}; + +class ITopicFilters : public TThrRefBase, public TNonCopyable { +public: + using TPtr = TIntrusivePtr; + +public: + // columnIndex - mapping from stable column id to index in values array + virtual void FilterData(const TVector& columnIndex, const TVector& offsets, const TVector*>& values, ui64 numberRows) = 0; + virtual void OnCompileResponse(TEvRowDispatcher::TEvPurecalcCompileResponse::TPtr ev) = 0; + + virtual TStatus AddFilter(IFilteredDataConsumer::TPtr filter) = 0; + virtual void RemoveFilter(NActors::TActorId filterId) = 0; + + virtual TFiltersStatistic GetStatistics() = 0; +}; + +struct TTopicFiltersConfig { + NActors::TActorId CompileServiceId; +}; + +ITopicFilters::TPtr CreateTopicFilters(NActors::TActorId owner, const TTopicFiltersConfig& config, NMonitoring::TDynamicCounterPtr counters); + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/filters/purecalc_filter.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/filters/purecalc_filter.cpp new file mode 100644 index 000000000000..b9b5bfd7c04a --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/filters/purecalc_filter.cpp @@ -0,0 +1,346 @@ +#include "purecalc_filter.h" + +#include + +#include +#include +#include + +namespace NFq::NRowDispatcher { + +namespace { + +constexpr const char* OFFSET_FIELD_NAME = "_offset"; + +NYT::TNode CreateTypeNode(const TString& fieldType) { + return NYT::TNode::CreateList() + .Add("DataType") + .Add(fieldType); +} + +void AddField(NYT::TNode& node, const TString& fieldName, const TString& fieldType) { + node.Add( + NYT::TNode::CreateList() + .Add(fieldName) + .Add(CreateTypeNode(fieldType)) + ); +} + +void AddColumn(NYT::TNode& node, const TSchemaColumn& column) { + TString parseTypeError; + TStringOutput errorStream(parseTypeError); + NYT::TNode parsedType; + if (!NYql::NCommon::ParseYson(parsedType, column.TypeYson, errorStream)) { + throw yexception() << "Failed to parse column '" << column.Name << "' type yson " << column.TypeYson << ", error: " << parseTypeError; + } + + node.Add( + NYT::TNode::CreateList() + .Add(column.Name) + .Add(parsedType) + ); +} + +NYT::TNode MakeInputSchema(const TVector& columns) { + auto structMembers = NYT::TNode::CreateList(); + AddField(structMembers, OFFSET_FIELD_NAME, "Uint64"); + for (const auto& column : columns) { + AddColumn(structMembers, column); + } + return NYT::TNode::CreateList().Add("StructType").Add(std::move(structMembers)); +} + +NYT::TNode MakeOutputSchema() { + auto structMembers = NYT::TNode::CreateList(); + AddField(structMembers, OFFSET_FIELD_NAME, "Uint64"); + return NYT::TNode::CreateList().Add("StructType").Add(std::move(structMembers)); +} + +struct TInputType { + const TVector*>& Values; + const ui64 NumberRows; +}; + +class TFilterInputSpec : public NYql::NPureCalc::TInputSpecBase { +public: + explicit TFilterInputSpec(const NYT::TNode& schema) + : Schemas({schema}) + {} + +public: + const TVector& GetSchemas() const override { + return Schemas; + } + +private: + const TVector Schemas; +}; + +class TFilterInputConsumer : public NYql::NPureCalc::IConsumer { +public: + TFilterInputConsumer(const TFilterInputSpec& spec, NYql::NPureCalc::TWorkerHolder worker) + : Worker(std::move(worker)) + { + const NKikimr::NMiniKQL::TStructType* structType = Worker->GetInputType(); + const ui64 count = structType->GetMembersCount(); + + THashMap schemaPositions; + for (ui64 i = 0; i < count; ++i) { + const auto name = structType->GetMemberName(i); + if (name == OFFSET_FIELD_NAME) { + OffsetPosition = i; + } else { + schemaPositions[name] = i; + } + } + + const auto& fields = spec.GetSchemas()[0][1]; + Y_ENSURE(fields.IsList(), "Unexpected input spec type"); + Y_ENSURE(count == fields.Size(), "Unexpected purecalc schema size"); + + FieldsPositions.reserve(count); + for (const auto& field : fields.AsList()) { + const auto& name = field[0].AsString(); + if (name != OFFSET_FIELD_NAME) { + FieldsPositions.emplace_back(schemaPositions[name]); + } + } + } + + ~TFilterInputConsumer() override { + with_lock(Worker->GetScopedAlloc()) { + Cache.Clear(); + } + } + +public: + void OnObject(TInputType input) override { + Y_ENSURE(FieldsPositions.size() == input.Values.size(), "Unexpected input scheme size"); + + NKikimr::NMiniKQL::TThrowingBindTerminator bind; + with_lock (Worker->GetScopedAlloc()) { + Y_DEFER { + // Clear cache after each object because + // values allocated on another allocator and should be released + Cache.Clear(); + Worker->Invalidate(); + }; + + auto& holderFactory = Worker->GetGraph().GetHolderFactory(); + + for (ui64 rowId = 0; rowId < input.NumberRows; ++rowId) { + NYql::NUdf::TUnboxedValue* items = nullptr; + NYql::NUdf::TUnboxedValue result = Cache.NewArray(holderFactory, static_cast(input.Values.size() + 1), items); + + items[OffsetPosition] = NYql::NUdf::TUnboxedValuePod(rowId); + + for (ui64 fieldId = 0; const auto column : input.Values) { + items[FieldsPositions[fieldId++]] = column->at(rowId); + } + + Worker->Push(std::move(result)); + } + } + } + + void OnFinish() override { + NKikimr::NMiniKQL::TBindTerminator bind(Worker->GetGraph().GetTerminator()); + with_lock(Worker->GetScopedAlloc()) { + Worker->OnFinish(); + } + } + +private: + const NYql::NPureCalc::TWorkerHolder Worker; + NKikimr::NMiniKQL::TPlainContainerCache Cache; + + ui64 OffsetPosition = 0; + TVector FieldsPositions; +}; + +class TFilterOutputSpec : public NYql::NPureCalc::TOutputSpecBase { +public: + explicit TFilterOutputSpec(const NYT::TNode& schema) + : Schema(schema) + {} + +public: + const NYT::TNode& GetSchema() const override { + return Schema; + } + +private: + const NYT::TNode Schema; +}; + +class TFilterOutputConsumer : public NYql::NPureCalc::IConsumer { +public: + explicit TFilterOutputConsumer(IPurecalcFilterConsumer::TPtr consumer) + : Consumer(std::move(consumer)) + {} + +public: + void OnObject(ui64 value) override { + Consumer->OnFilteredData(value); + } + + void OnFinish() override { + } + +private: + const IPurecalcFilterConsumer::TPtr Consumer; +}; + +class TFilterPushRelayImpl : public NYql::NPureCalc::IConsumer { +public: + TFilterPushRelayImpl(const TFilterOutputSpec& outputSpec, NYql::NPureCalc::IPushStreamWorker* worker, THolder> underlying) + : Underlying(std::move(underlying)) + , Worker(worker) + { + Y_UNUSED(outputSpec); + } + +public: + void OnObject(const NYql::NUdf::TUnboxedValue* value) override { + Y_ENSURE(value->GetListLength() == 1, "Unexpected output schema size"); + + auto unguard = Unguard(Worker->GetScopedAlloc()); + Underlying->OnObject(value->GetElement(0).Get()); + } + + void OnFinish() override { + auto unguard = Unguard(Worker->GetScopedAlloc()); + Underlying->OnFinish(); + } + +private: + const THolder> Underlying; + NYql::NPureCalc::IWorker* Worker; +}; + +} // anonymous namespace + +} // namespace NFq::NRowDispatcher + +template <> +struct NYql::NPureCalc::TInputSpecTraits { + static constexpr bool IsPartial = false; + static constexpr bool SupportPushStreamMode = true; + + using TConsumerType = THolder>; + + static TConsumerType MakeConsumer(const NFq::NRowDispatcher::TFilterInputSpec& spec, NYql::NPureCalc::TWorkerHolder worker) { + return MakeHolder(spec, std::move(worker)); + } +}; + +template <> +struct NYql::NPureCalc::TOutputSpecTraits { + static const constexpr bool IsPartial = false; + static const constexpr bool SupportPushStreamMode = true; + + static void SetConsumerToWorker(const NFq::NRowDispatcher::TFilterOutputSpec& outputSpec, NYql::NPureCalc::IPushStreamWorker* worker, THolder> consumer) { + worker->SetConsumer(MakeHolder(outputSpec, worker, std::move(consumer))); + } +}; + +namespace NFq::NRowDispatcher { + +namespace { + +class TProgramHolder : public IProgramHolder { +public: + using TPtr = TIntrusivePtr; + +public: + TProgramHolder(IPurecalcFilterConsumer::TPtr consumer, const TString& sql) + : Consumer(consumer) + , Sql(sql) + {} + + NYql::NPureCalc::IConsumer& GetConsumer() { + Y_ENSURE(InputConsumer, "Program is not compiled"); + return *InputConsumer; + } + +public: + void CreateProgram(NYql::NPureCalc::IProgramFactoryPtr programFactory) override { + // Program should be stateless because input values + // allocated on another allocator and should be released + Program = programFactory->MakePushStreamProgram( + TFilterInputSpec(MakeInputSchema(Consumer->GetColumns())), + TFilterOutputSpec(MakeOutputSchema()), + Sql, + NYql::NPureCalc::ETranslationMode::SQL + ); + InputConsumer = Program->Apply(MakeHolder(Consumer)); + } + +private: + const IPurecalcFilterConsumer::TPtr Consumer; + const TString Sql; + + THolder> Program; + THolder> InputConsumer; +}; + +class TPurecalcFilter : public IPurecalcFilter { +public: + TPurecalcFilter(IPurecalcFilterConsumer::TPtr consumer) + : Consumer(consumer) + , PurecalcSettings(consumer->GetPurecalcSettings()) + , LogPrefix("TPurecalcFilter: ") + , ProgramHolder(MakeIntrusive(consumer, GenerateSql())) + {} + +public: + void FilterData(const TVector*>& values, ui64 numberRows) override { + LOG_ROW_DISPATCHER_TRACE("Do filtering for " << numberRows << " rows"); + ProgramHolder->GetConsumer().OnObject({.Values = values, .NumberRows = numberRows}); + } + + std::unique_ptr GetCompileRequest() override { + Y_ENSURE(ProgramHolder, "Can not create compile request twice"); + auto result = std::make_unique(std::move(ProgramHolder), PurecalcSettings); + ProgramHolder = nullptr; + return result; + } + + void OnCompileResponse(TEvRowDispatcher::TEvPurecalcCompileResponse::TPtr ev) override { + Y_ENSURE(!ProgramHolder, "Can not handle compile response twice"); + + auto result = static_cast(ev->Get()->ProgramHolder.Release()); + Y_ENSURE(result, "Unexpected compile response"); + + ProgramHolder = TIntrusivePtr(result); + } + +private: + TString GenerateSql() { + TStringStream str; + str << "PRAGMA config.flags(\"LLVM\", \"" << (PurecalcSettings.EnabledLLVM ? "ON" : "OFF") << "\");\n"; + str << "SELECT " << OFFSET_FIELD_NAME << " FROM Input " << Consumer->GetWhereFilter() << ";\n"; + + LOG_ROW_DISPATCHER_DEBUG("Generated sql:\n" << str.Str()); + return str.Str(); + } + +private: + const IPurecalcFilterConsumer::TPtr Consumer; + const TPurecalcCompileSettings PurecalcSettings; + const TString LogPrefix; + + TProgramHolder::TPtr ProgramHolder; +}; + +} // anonymous namespace + +TValueStatus CreatePurecalcFilter(IPurecalcFilterConsumer::TPtr consumer) { + try { + return IPurecalcFilter::TPtr(MakeIntrusive(consumer)); + } catch (...) { + return TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Failed to create purecalc filter with predicate '" << consumer->GetWhereFilter() << "', got unexpected exception: " << CurrentExceptionMessage()); + } +} + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/filters/purecalc_filter.h b/ydb/core/fq/libs/row_dispatcher/format_handler/filters/purecalc_filter.h new file mode 100644 index 000000000000..62ad0b7899e4 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/filters/purecalc_filter.h @@ -0,0 +1,35 @@ +#pragma once + +#include +#include + +#include + +namespace NFq::NRowDispatcher { + +class IPurecalcFilterConsumer : public TThrRefBase { +public: + using TPtr = TIntrusivePtr; + +public: + virtual const TVector& GetColumns() const = 0; + virtual const TString& GetWhereFilter() const = 0; + virtual TPurecalcCompileSettings GetPurecalcSettings() const = 0; + + virtual void OnFilteredData(ui64 rowId) = 0; +}; + +class IPurecalcFilter : public TThrRefBase, public TNonCopyable { +public: + using TPtr = TIntrusivePtr; + +public: + virtual void FilterData(const TVector*>& values, ui64 numberRows) = 0; + + virtual std::unique_ptr GetCompileRequest() = 0; // Should be called exactly once + virtual void OnCompileResponse(TEvRowDispatcher::TEvPurecalcCompileResponse::TPtr ev) = 0; +}; + +TValueStatus CreatePurecalcFilter(IPurecalcFilterConsumer::TPtr consumer); + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/filters/ya.make b/ydb/core/fq/libs/row_dispatcher/format_handler/filters/ya.make new file mode 100644 index 000000000000..629c468519f3 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/filters/ya.make @@ -0,0 +1,26 @@ +LIBRARY() + +SRCS( + purecalc_filter.cpp + filters_set.cpp +) + +PEERDIR( + ydb/core/fq/libs/actors/logging + ydb/core/fq/libs/row_dispatcher/events + ydb/core/fq/libs/row_dispatcher/format_handler/common + ydb/core/fq/libs/row_dispatcher/purecalc_no_pg_wrapper + + ydb/library/actors/core + + yql/essentials/minikql + yql/essentials/minikql/computation + yql/essentials/minikql/comp_nodes + yql/essentials/minikql/invoke_builtins + yql/essentials/providers/common/schema/parser + yql/essentials/public/udf +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/format_handler.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/format_handler.cpp new file mode 100644 index 000000000000..8994f61d8404 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/format_handler.cpp @@ -0,0 +1,628 @@ +#include "format_handler.h" + +#include + +#include +#include +#include + +#include + +#include + +namespace NFq::NRowDispatcher { + +namespace { + +class TTopicFormatHandler : public NActors::TActor, public ITopicFormatHandler, public TTypeParser { + using TBase = NActors::TActor; + +public: + static constexpr char ActorName[] = "FQ_ROW_DISPATCHER_FORMAT_HANDLER"; + +private: + struct TCounters { + TCountersDesc Desc; + + NMonitoring::TDynamicCounters::TCounterPtr ActiveFormatHandlers; + NMonitoring::TDynamicCounters::TCounterPtr ActiveClients; + + TCounters(const TCountersDesc& counters, const TSettings& settings) + : Desc(counters) + { + Desc.CountersSubgroup = Desc.CountersSubgroup->GetSubgroup("format", settings.ParsingFormat); + + Register(); + } + + private: + void Register() { + ActiveFormatHandlers = Desc.CountersRoot->GetCounter("ActiveFormatHandlers", false); + + ActiveClients = Desc.CountersSubgroup->GetCounter("ActiveClients", false); + } + }; + + struct TColumnDesc { + ui64 ColumnId; // Stable column id, used in TClientHandler + TString TypeYson; + std::unordered_set Clients; // Set of clients with this column + }; + + class TParserHandler : public IParsedDataConsumer { + public: + using TPtr = TIntrusivePtr; + + public: + TParserHandler(TTopicFormatHandler& self, TVector parerSchema) + : Self(self) + , ParerSchema(parerSchema) + , LogPrefix(TStringBuilder() << Self.LogPrefix << "TParserHandler: ") + {} + + public: + const TVector& GetColumns() const override { + return ParerSchema; + } + + void OnParsingError(TStatus status) override { + LOG_ROW_DISPATCHER_ERROR("Got parsing error: " << status.GetErrorMessage()); + Self.FatalError(status); + } + + void OnParsedData(ui64 numberRows) override { + LOG_ROW_DISPATCHER_TRACE("Got parsed data, number rows: " << numberRows); + + Self.ParsedData.assign(ParerSchema.size(), nullptr); + for (size_t i = 0; i < ParerSchema.size(); ++i) { + auto columnStatus = Self.Parser->GetParsedColumn(i); + if (Y_LIKELY(columnStatus.IsSuccess())) { + Self.ParsedData[i] = columnStatus.DetachResult(); + } else { + OnColumnError(i, columnStatus); + } + } + + Self.Offsets = &Self.Parser->GetOffsets(); + Self.FilterData(numberRows); + } + + private: + void OnColumnError(ui64 columnIndex, TStatus status) { + const auto& column = ParerSchema[columnIndex]; + LOG_ROW_DISPATCHER_WARN("Failed to parse column " << column.ToString() << ", " << status.GetErrorMessage()); + + const auto columnIt = Self.ColumnsDesc.find(column.Name); + if (columnIt == Self.ColumnsDesc.end()) { + return; + } + + for (const auto clientId : columnIt->second.Clients) { + const auto clientIt = Self.Clients.find(clientId); + if (clientIt != Self.Clients.end()) { + clientIt->second->OnClientError(status); + } + } + } + + private: + TTopicFormatHandler& Self; + const TVector ParerSchema; + const TString LogPrefix; + }; + + class TClientHandler : public IFilteredDataConsumer { + public: + using TPtr = TIntrusivePtr; + + public: + TClientHandler(TTopicFormatHandler& self, IClientDataConsumer::TPtr client) + : Self(self) + , Client(client) + , Columns(Client->GetColumns()) + , LogPrefix(TStringBuilder() << Self.LogPrefix << "TClientHandler " << Client->GetClientId() << ": ") + , FilteredRow(Columns.size()) + { + ColumnsIds.reserve(Columns.size()); + } + + IClientDataConsumer::TPtr GetClient() const { + return Client; + } + + bool IsClientStarted() const { + return ClientStarted; + } + + TStatus SetupColumns() { + if (Columns.empty()) { + return TStatus::Fail(EStatusId::INTERNAL_ERROR, "Client should have at least one column in schema"); + } + + for (const auto& column : Columns) { + const auto it = Self.ColumnsDesc.find(column.Name); + if (it != Self.ColumnsDesc.end()) { + if (it->second.TypeYson != column.TypeYson) { + return TStatus::Fail(EStatusId::SCHEME_ERROR, TStringBuilder() << "Use the same column type in all queries via RD, current type for column `" << column.Name << "` is " << it->second.TypeYson << " (requested type is " << column.TypeYson <<")"); + } + + it->second.Clients.emplace(Client->GetClientId()); + ColumnsIds.emplace_back(it->second.ColumnId); + } else { + ui64 columnId = Self.MaxColumnId; + if (Self.FreeColumnIds) { + columnId = Self.FreeColumnIds.back(); + Self.FreeColumnIds.pop_back(); + } + + Self.ColumnsDesc[column.Name] = TColumnDesc{.ColumnId = columnId, .TypeYson = column.TypeYson, .Clients = {Client->GetClientId()}}; + Self.MaxColumnId = std::max(Self.MaxColumnId, columnId + 1); + ColumnsIds.emplace_back(columnId); + } + } + + return SetupPacker(); + } + + TQueue>> ExtractClientData() { + FinishPacking(); + TQueue>> result; + result.swap(ClientData); + LOG_ROW_DISPATCHER_TRACE("ExtractClientData, number batches: " << result.size()); + return result; + } + + void OnClientError(TStatus status) { + LOG_ROW_DISPATCHER_WARN("OnClientError, " << status.GetErrorMessage()); + Client->OnClientError(std::move(status)); + } + + public: + NActors::TActorId GetFilterId() const override { + return Client->GetClientId(); + } + + const TVector& GetColumns() const override { + return Columns; + } + + const TVector& GetColumnIds() const override { + return ColumnsIds; + } + + TMaybe GetNextMessageOffset() const override { + return Client->GetNextMessageOffset(); + } + + const TString& GetWhereFilter() const override { + return Client->GetWhereFilter(); + } + + TPurecalcCompileSettings GetPurecalcSettings() const override { + return Client->GetPurecalcSettings(); + } + + void OnFilteringError(TStatus status) override { + Client->OnClientError(status); + } + + void OnFilterStarted() override { + ClientStarted = true; + Client->StartClientSession(); + } + + void OnFilteredBatch(ui64 firstRow, ui64 lastRow) override { + LOG_ROW_DISPATCHER_TRACE("OnFilteredBatch, rows [" << firstRow << ", " << lastRow << "]"); + for (ui64 rowId = firstRow; rowId <= lastRow; ++rowId) { + OnFilteredData(rowId); + } + } + + void OnFilteredData(ui64 rowId) override { + const ui64 offset = Self.Offsets->at(rowId); + if (const auto nextOffset = Client->GetNextMessageOffset(); nextOffset && offset < *nextOffset) { + LOG_ROW_DISPATCHER_TRACE("OnFilteredData, skip historical offset: " << offset << ", next message offset: " << *nextOffset); + return; + } + + Y_DEFER { + // Values allocated on parser allocator and should be released + FilteredRow.assign(Columns.size(), NYql::NUdf::TUnboxedValue()); + }; + + for (size_t i = 0; const ui64 columnId : ColumnsIds) { + // All data was locked in parser, so copy is safe + FilteredRow[i++] = Self.ParsedData[Self.ParserSchemaIndex[columnId]]->at(rowId); + } + DataPacker->AddWideItem(FilteredRow.data(), FilteredRow.size()); + FilteredOffsets.emplace_back(offset); + + const ui64 newPackerSize = DataPacker->PackedSizeEstimate(); + LOG_ROW_DISPATCHER_TRACE("OnFilteredData, row id: " << rowId << ", offset: " << offset << ", new packer size: " << newPackerSize); + Client->AddDataToClient(offset, newPackerSize - DataPackerSize); + + DataPackerSize = newPackerSize; + if (DataPackerSize > MAX_BATCH_SIZE) { + FinishPacking(); + } + } + + private: + TStatus SetupPacker() { + TVector columnTypes; + columnTypes.reserve(Columns.size()); + for (const auto& column : Columns) { + auto status = Self.ParseTypeYson(column.TypeYson); + if (status.IsFail()) { + return status; + } + columnTypes.emplace_back(status.DetachResult()); + } + + with_lock(Self.Alloc) { + const auto rowType = Self.ProgramBuilder->NewMultiType(columnTypes); + DataPacker = std::make_unique>(rowType); + } + return TStatus::Success(); + } + + void FinishPacking() { + if (!DataPacker->IsEmpty()) { + LOG_ROW_DISPATCHER_TRACE("FinishPacking, batch size: " << DataPackerSize << ", number rows: " << FilteredOffsets.size()); + ClientData.emplace(NYql::MakeReadOnlyRope(DataPacker->Finish()), FilteredOffsets); + DataPackerSize = 0; + FilteredOffsets.clear(); + } + } + + private: + TTopicFormatHandler& Self; + const IClientDataConsumer::TPtr Client; + const TVector Columns; + const TString LogPrefix; + + TVector ColumnsIds; + bool ClientStarted = false; + + // Filtered data + ui64 DataPackerSize = 0; + TVector FilteredRow; // Temporary value holder for DataPacket + std::unique_ptr> DataPacker; + TVector FilteredOffsets; // Offsets of current batch in DataPacker + TQueue>> ClientData; // vector of (messages batch, [offsets]) + }; + +public: + TTopicFormatHandler(const TFormatHandlerConfig& config, const TSettings& settings, const TCountersDesc& counters) + : TBase(&TTopicFormatHandler::StateFunc) + , TTypeParser(__LOCATION__, counters.CopyWithNewMkqlCountersName("row_dispatcher")) + , Config(config) + , Settings(settings) + , LogPrefix(TStringBuilder() << "TTopicFormatHandler [" << Settings.ParsingFormat << "]: ") + , Counters(counters, settings) + { + Counters.ActiveFormatHandlers->Inc(); + } + + ~TTopicFormatHandler() { + Counters.ActiveFormatHandlers->Dec(); + Counters.ActiveClients->Set(0); + + with_lock(Alloc) { + Clients.clear(); + } + } + + STRICT_STFUNC_EXC(StateFunc, + hFunc(TEvRowDispatcher::TEvPurecalcCompileResponse, Handle); + hFunc(NActors::TEvents::TEvWakeup, Handle); + hFunc(NActors::TEvents::TEvPoison, Handle);, + ExceptionFunc(std::exception, HandleException) + ) + + void Handle(TEvRowDispatcher::TEvPurecalcCompileResponse::TPtr& ev) { + if (Filters) { + Filters->OnCompileResponse(std::move(ev)); + } + } + + void Handle(NActors::TEvents::TEvWakeup::TPtr&) { + RefreshScheduled = false; + + if (Parser) { + LOG_ROW_DISPATCHER_TRACE("Refresh parser"); + Parser->Refresh(); + ScheduleRefresh(); + } + } + + void Handle(NActors::TEvents::TEvPoison::TPtr&) { + if (Filters) { + for (const auto& [clientId, _] : Clients) { + Filters->RemoveFilter(clientId); + } + Filters.Reset(); + } + PassAway(); + } + + void HandleException(const std::exception& error) { + LOG_ROW_DISPATCHER_ERROR("Got unexpected exception: " << error.what()); + FatalError(TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Format handler error, got unexpected exception: " << error.what())); + } + +public: + void ParseMessages(const TVector& messages) override { + LOG_ROW_DISPATCHER_TRACE("Send " << messages.size() << " messages to parser"); + + if (messages) { + CurrentOffset = messages.back().GetOffset(); + } + + if (Parser) { + Parser->ParseMessages(messages); + ScheduleRefresh(); + } else if (!Clients.empty()) { + FatalError(TStatus::Fail(EStatusId::INTERNAL_ERROR, "Failed to parse messages, expected empty clients set without parser")); + } + } + + TQueue>> ExtractClientData(NActors::TActorId clientId) override { + const auto it = Clients.find(clientId); + if (it == Clients.end()) { + return {}; + } + return it->second->ExtractClientData(); + } + + TStatus AddClient(IClientDataConsumer::TPtr client) override { + LOG_ROW_DISPATCHER_DEBUG("Add client with id " << client->GetClientId()); + + if (const auto clientOffset = client->GetNextMessageOffset()) { + if (Parser && CurrentOffset && *CurrentOffset > *clientOffset) { + LOG_ROW_DISPATCHER_DEBUG("Parser was flushed due to new historical offset " << *clientOffset << "(previous parser offset: " << *CurrentOffset << ")"); + Parser->Refresh(true); + } + } + + auto clientHandler = MakeIntrusive(*this, client); + if (!Clients.emplace(client->GetClientId(), clientHandler).second) { + return TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Failed to create new client, client with id " << client->GetClientId() << " already exists"); + } + Counters.ActiveClients->Inc(); + + if (auto status = clientHandler->SetupColumns(); status.IsFail()) { + RemoveClient(client->GetClientId()); + return status.AddParentIssue("Failed to modify common parsing schema"); + } + + if (auto status = UpdateParser(); status.IsFail()) { + RemoveClient(client->GetClientId()); + return status.AddParentIssue("Failed to update parser with new client"); + } + + CreateFilters(); + if (auto status = Filters->AddFilter(clientHandler); status.IsFail()) { + RemoveClient(client->GetClientId()); + return status.AddParentIssue("Failed to create filter for new client"); + } + + return TStatus::Success(); + } + + void RemoveClient(NActors::TActorId clientId) override { + LOG_ROW_DISPATCHER_DEBUG("Remove client with id " << clientId); + + if (Filters) { + Filters->RemoveFilter(clientId); + } + + const auto it = Clients.find(clientId); + if (it == Clients.end()) { + return; + } + + const auto client = it->second->GetClient(); + Counters.ActiveClients->Dec(); + Clients.erase(it); + + for (const auto& column : client->GetColumns()) { + const auto columnIt = ColumnsDesc.find(column.Name); + if (columnIt == ColumnsDesc.end()) { + continue; + } + + columnIt->second.Clients.erase(client->GetClientId()); + if (columnIt->second.Clients.empty()) { + FreeColumnIds.emplace_back(columnIt->second.ColumnId); + ColumnsDesc.erase(columnIt); + } + } + + if (auto status = UpdateParser(); status.IsFail()) { + FatalError(status.AddParentIssue("Failed to update parser after removing client")); + } + } + + bool HasClients() const override { + return !Clients.empty(); + } + + TFormatHandlerStatistic GetStatistics() override { + TFormatHandlerStatistic statistics; + if (Parser) { + Parser->FillStatistics(statistics); + } + if (Filters) { + statistics.FilterStats = Filters->GetStatistics(); + } + return statistics; + } + + void ForceRefresh() override { + if (Parser) { + Parser->Refresh(true); + } + } + +protected: + NActors::TActorId GetSelfId() const override { + return SelfId(); + } + +private: + void ScheduleRefresh() { + if (const auto refreshPeriod = Config.JsonParserConfig.LatencyLimit; !RefreshScheduled && refreshPeriod) { + RefreshScheduled = true; + Schedule(refreshPeriod, new NActors::TEvents::TEvWakeup()); + } + } + + TStatus UpdateParser() { + TVector parerSchema; + parerSchema.reserve(ColumnsDesc.size()); + for (const auto& [columnName, columnDesc] : ColumnsDesc) { + parerSchema.emplace_back(TSchemaColumn{.Name = columnName, .TypeYson = columnDesc.TypeYson}); + } + + if (ParserHandler && parerSchema == ParserHandler->GetColumns()) { + return TStatus::Success(); + } + + if (Parser) { + Parser->Refresh(true); + Parser.Reset(); + } + + LOG_ROW_DISPATCHER_DEBUG("UpdateParser to new schema with size " << parerSchema.size()); + ParserHandler = MakeIntrusive(*this, std::move(parerSchema)); + + if (const ui64 schemaSize = ParserHandler->GetColumns().size()) { + auto newParser = CreateParserForFormat(); + if (newParser.IsFail()) { + return newParser; + } + + LOG_ROW_DISPATCHER_DEBUG("Parser was updated on new schema with " << schemaSize << " columns"); + + Parser = newParser.DetachResult(); + ParserSchemaIndex.resize(MaxColumnId, std::numeric_limits::max()); + for (ui64 i = 0; const auto& [_, columnDesc] : ColumnsDesc) { + ParserSchemaIndex[columnDesc.ColumnId] = i++; + } + } else { + LOG_ROW_DISPATCHER_INFO("No columns to parse, reset parser"); + } + + return TStatus::Success(); + } + + TValueStatus CreateParserForFormat() const { + const auto& counters = Counters.Desc.CopyWithNewMkqlCountersName("row_dispatcher_parser"); + if (Settings.ParsingFormat == "raw") { + return CreateRawParser(ParserHandler, counters); + } + if (Settings.ParsingFormat == "json_each_row") { + return CreateJsonParser(ParserHandler, Config.JsonParserConfig, counters); + } + return TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Unsupported parsing format: " << Settings.ParsingFormat); + } + + void CreateFilters() { + if (!Filters) { + Filters = CreateTopicFilters(SelfId(), Config.FiltersConfig, Counters.Desc.CountersSubgroup); + } + } + + void FilterData(ui64 numberRows) { + if (!numberRows) { + return; + } + + const ui64 lastOffset = Offsets->at(numberRows - 1); + LOG_ROW_DISPATCHER_TRACE("Send " << numberRows << " messages to filters, first offset: " << Offsets->front() << ", last offset: " << lastOffset); + + if (Filters) { + Filters->FilterData(ParserSchemaIndex, *Offsets, ParsedData, numberRows); + } + + for (const auto& [_, client] : Clients) { + if (client->IsClientStarted()) { + LOG_ROW_DISPATCHER_TRACE("Commit client " << client->GetClient()->GetClientId() << " offset " << lastOffset); + client->GetClient()->UpdateClientOffset(lastOffset); + } + } + } + + void FatalError(TStatus status) const { + LOG_ROW_DISPATCHER_ERROR("Got fatal error: " << status.GetErrorMessage()); + for (const auto& [_, client] : Clients) { + client->OnClientError(status); + } + } + +private: + const TFormatHandlerConfig Config; + const TSettings Settings; + const TString LogPrefix; + + // Columns indexes + ui64 MaxColumnId = 0; + TVector FreeColumnIds; + TVector ParserSchemaIndex; // Column id to index in parser schema + std::map ColumnsDesc; + std::unordered_map Clients; + + // Perser and filters + ITopicParser::TPtr Parser; + TParserHandler::TPtr ParserHandler; + ITopicFilters::TPtr Filters; + std::optional CurrentOffset; + + // Parsed data + const TVector* Offsets; + TVector*> ParsedData; + bool RefreshScheduled = false; + + // Metrics + const TCounters Counters; +}; + +} // anonymous namespace + +//// ITopicFormatHandler::TDestroy + +void ITopicFormatHandler::TDestroy::Destroy(ITopicFormatHandler* handler) { + if (NActors::TlsActivationContext) { + NActors::TActivationContext::ActorSystem()->Send(handler->GetSelfId(), new NActors::TEvents::TEvPoison()); + } else { + // Destroy from not AS thread my be caused only in case AS destruction (so handler will be deleted) + } +} + +ITopicFormatHandler::TPtr CreateTopicFormatHandler(const NActors::TActorContext& owner, const TFormatHandlerConfig& config, const ITopicFormatHandler::TSettings& settings, const TCountersDesc& counters) { + const auto handler = new TTopicFormatHandler(config, settings, counters); + owner.RegisterWithSameMailbox(handler); + return ITopicFormatHandler::TPtr(handler); +} + +TFormatHandlerConfig CreateFormatHandlerConfig(const NConfig::TRowDispatcherConfig& rowDispatcherConfig, NActors::TActorId compileServiceId) { + return { + .JsonParserConfig = CreateJsonParserConfig(rowDispatcherConfig.GetJsonParser()), + .FiltersConfig = { + .CompileServiceId = compileServiceId + } + }; +} + +namespace NTests { + +ITopicFormatHandler::TPtr CreateTestFormatHandler(const TFormatHandlerConfig& config, const ITopicFormatHandler::TSettings& settings) { + const auto handler = new TTopicFormatHandler(config, settings, {}); + NActors::TActivationContext::ActorSystem()->Register(handler); + return ITopicFormatHandler::TPtr(handler); +} + +} // namespace NTests + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/format_handler.h b/ydb/core/fq/libs/row_dispatcher/format_handler/format_handler.h new file mode 100644 index 000000000000..38325515f6f4 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/format_handler.h @@ -0,0 +1,79 @@ +#pragma once + +#include +#include + +#include +#include + +namespace NFq::NRowDispatcher { + +static constexpr ui64 MAX_BATCH_SIZE = 10_MB; + +class IClientDataConsumer : public TThrRefBase { +public: + using TPtr = TIntrusivePtr; + +public: + virtual TVector GetColumns() const = 0; + virtual const TString& GetWhereFilter() const = 0; + virtual TPurecalcCompileSettings GetPurecalcSettings() const = 0; + virtual NActors::TActorId GetClientId() const = 0; + virtual TMaybe GetNextMessageOffset() const = 0; + + virtual void OnClientError(TStatus status) = 0; + + virtual void StartClientSession() = 0; + virtual void AddDataToClient(ui64 offset, ui64 rowSize) = 0; + virtual void UpdateClientOffset(ui64 offset) = 0; +}; + +class ITopicFormatHandler : public TNonCopyable { +private: + class TDestroy { + public: + static void Destroy(ITopicFormatHandler* handler); + }; + +public: + using TPtr = THolder; + + struct TSettings { + TString ParsingFormat = "raw"; + + std::strong_ordering operator<=>(const TSettings& other) const = default; + }; + +public: + virtual void ParseMessages(const TVector& messages) = 0; + + // vector of (messages batch, [offsets]) + virtual TQueue>> ExtractClientData(NActors::TActorId clientId) = 0; + + virtual TStatus AddClient(IClientDataConsumer::TPtr client) = 0; + virtual void RemoveClient(NActors::TActorId clientId) = 0; + virtual bool HasClients() const = 0; + + virtual TFormatHandlerStatistic GetStatistics() = 0; + virtual void ForceRefresh() = 0; + +protected: + virtual NActors::TActorId GetSelfId() const = 0; +}; + +// Static properties for all format handlers +struct TFormatHandlerConfig { + TJsonParserConfig JsonParserConfig; + TTopicFiltersConfig FiltersConfig; +}; + +ITopicFormatHandler::TPtr CreateTopicFormatHandler(const NActors::TActorContext& owner, const TFormatHandlerConfig& config, const ITopicFormatHandler::TSettings& settings, const TCountersDesc& counters); +TFormatHandlerConfig CreateFormatHandlerConfig(const NConfig::TRowDispatcherConfig& rowDispatcherConfig, NActors::TActorId compileServiceId); + +namespace NTests { + +ITopicFormatHandler::TPtr CreateTestFormatHandler(const TFormatHandlerConfig& config, const ITopicFormatHandler::TSettings& settings); + +} // namespace NTests + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/json_parser.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/json_parser.cpp new file mode 100644 index 000000000000..89472c928877 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/json_parser.cpp @@ -0,0 +1,514 @@ +#include "json_parser.h" + +#include "parser_base.h" + +#include + +#include + +#include + +#include +#include +#include +#include + +namespace NFq::NRowDispatcher { + +namespace { + +#define CHECK_JSON_ERROR(value) \ + const simdjson::error_code error = value; \ + if (Y_UNLIKELY(error)) \ + +struct TJsonParserBuffer { + size_t NumberValues = 0; + bool Finished = false; + TInstant CreationStartTime = TInstant::Now(); + TVector Offsets = {}; + + bool IsReady() const { + return !Finished && NumberValues > 0; + } + + size_t GetSize() const { + return Values.size(); + } + + void Reserve(size_t size, size_t numberValues) { + Values.reserve(size + simdjson::SIMDJSON_PADDING); + Offsets.reserve(numberValues); + } + + void AddMessage(const NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage& message) { + Y_ENSURE(!Finished, "Cannot add messages into finished buffer"); + + const auto offset = message.GetOffset(); + if (Y_UNLIKELY(Offsets && Offsets.back() > offset)) { + LOG_ROW_DISPATCHER_WARN("Got message with offset " << offset << " which is less than previous offset " << Offsets.back()); + } + + NumberValues++; + Values << message.GetData(); + Offsets.emplace_back(offset); + } + + std::pair Finish() { + Y_ENSURE(!Finished, "Cannot finish buffer twice"); + Finished = true; + Values << TString(simdjson::SIMDJSON_PADDING, ' '); + return {Values.data(), Values.size()}; + } + + void Clear() { + Y_ENSURE(Finished, "Cannot clear not finished buffer"); + NumberValues = 0; + Finished = false; + CreationStartTime = TInstant::Now(); + Values.clear(); + Offsets.clear(); + } + +private: + TStringBuilder Values = {}; + const TString LogPrefix = "TJsonParser: Buffer: "; +}; + +class TColumnParser { +public: + const std::string Name; // Used for column index by std::string_view + const TString TypeYson; + +public: + TColumnParser(const TString& name, const TString& typeYson, ui64 maxNumberRows) + : Name(name) + , TypeYson(typeYson) + , Status(TStatus::Success()) + { + ParsedRows.reserve(maxNumberRows); + } + + TStatus InitParser(const NKikimr::NMiniKQL::TType* typeMkql) { + return Status = ExtractDataSlot(typeMkql); + } + + const TVector& GetParsedRows() const { + return ParsedRows; + } + + TStatus GetStatus() const { + return Status; + } + + void ParseJsonValue(ui64 offset, ui64 rowId, simdjson::builtin::ondemand::value jsonValue, NYql::NUdf::TUnboxedValue& resultValue) { + if (Y_UNLIKELY(Status.IsFail())) { + return; + } + ParsedRows.emplace_back(rowId); + + if (DataSlot != NYql::NUdf::EDataSlot::Json) { + ParseDataType(std::move(jsonValue), resultValue, Status); + } else { + ParseJsonType(std::move(jsonValue), resultValue, Status); + } + + if (IsOptional && resultValue) { + resultValue = resultValue.MakeOptional(); + } + + if (Y_UNLIKELY(Status.IsFail())) { + Status.AddParentIssue(TStringBuilder() << "Failed to parse json string at offset " << offset << ", got parsing error for column '" << Name << "' with type " << TypeYson); + } + } + + void ValidateNumberValues(size_t expectedNumberValues, ui64 firstOffset) { + if (Status.IsFail()) { + return; + } + if (Y_UNLIKELY(!IsOptional && ParsedRows.size() < expectedNumberValues)) { + Status = TStatus::Fail(EStatusId::PRECONDITION_FAILED, TStringBuilder() << "Failed to parse json messages, found " << expectedNumberValues - ParsedRows.size() << " missing values from offset " << firstOffset << " in non optional column '" << Name << "' with type " << TypeYson); + } + } + + void ClearParsedRows() { + ParsedRows.clear(); + Status = TStatus::Success(); + } + +private: + TStatus ExtractDataSlot(const NKikimr::NMiniKQL::TType* type) { + switch (type->GetKind()) { + case NKikimr::NMiniKQL::TTypeBase::EKind::Data: { + auto slotStatus = GetDataSlot(type); + if (slotStatus.IsFail()) { + return slotStatus; + } + DataSlot = slotStatus.DetachResult(); + DataTypeName = NYql::NUdf::GetDataTypeInfo(DataSlot).Name; + return TStatus::Success(); + } + + case NKikimr::NMiniKQL::TTypeBase::EKind::Optional: { + if (IsOptional) { + return TStatus::Fail(EStatusId::UNSUPPORTED, TStringBuilder() << "Nested optionals is not supported as input type"); + } + IsOptional = true; + return ExtractDataSlot(AS_TYPE(NKikimr::NMiniKQL::TOptionalType, type)->GetItemType()); + } + + default: { + return TStatus::Fail(EStatusId::UNSUPPORTED, TStringBuilder() << "Unsupported type kind: " << type->GetKindAsStr()); + } + } + } + + Y_FORCE_INLINE void ParseDataType(simdjson::builtin::ondemand::value jsonValue, NYql::NUdf::TUnboxedValue& resultValue, TStatus& status) const { + simdjson::builtin::ondemand::json_type cellType; + CHECK_JSON_ERROR(jsonValue.type().get(cellType)) { + return GetParsingError(error, jsonValue, "determine json value type", status); + } + + switch (cellType) { + case simdjson::builtin::ondemand::json_type::number: { + switch (DataSlot) { + case NYql::NUdf::EDataSlot::Int8: + ParseJsonNumber(jsonValue.get_int64(), resultValue, status); + break; + case NYql::NUdf::EDataSlot::Int16: + ParseJsonNumber(jsonValue.get_int64(), resultValue, status); + break; + case NYql::NUdf::EDataSlot::Int32: + ParseJsonNumber(jsonValue.get_int64(), resultValue, status); + break; + case NYql::NUdf::EDataSlot::Int64: + ParseJsonNumber(jsonValue.get_int64(), resultValue, status); + break; + + case NYql::NUdf::EDataSlot::Uint8: + ParseJsonNumber(jsonValue.get_uint64(), resultValue, status); + break; + case NYql::NUdf::EDataSlot::Uint16: + ParseJsonNumber(jsonValue.get_uint64(), resultValue, status); + break; + case NYql::NUdf::EDataSlot::Uint32: + ParseJsonNumber(jsonValue.get_uint64(), resultValue, status); + break; + case NYql::NUdf::EDataSlot::Uint64: + ParseJsonNumber(jsonValue.get_uint64(), resultValue, status); + break; + + case NYql::NUdf::EDataSlot::Double: + ParseJsonDouble(jsonValue.get_double(), resultValue, status); + break; + case NYql::NUdf::EDataSlot::Float: + ParseJsonDouble(jsonValue.get_double(), resultValue, status); + break; + + default: + status = TStatus::Fail(EStatusId::PRECONDITION_FAILED, TStringBuilder() << "Number value is not expected for data type " << DataTypeName); + break; + } + if (Y_UNLIKELY(status.IsFail())) { + status.AddParentIssue(TStringBuilder() << "Failed to parse data type " << DataTypeName << " from json number (raw: '" << TruncateString(jsonValue.raw_json_token()) << "')"); + } + return; + } + + case simdjson::builtin::ondemand::json_type::string: { + std::string_view rawString; + CHECK_JSON_ERROR(jsonValue.get_string(rawString)) { + return GetParsingError(error, jsonValue, "extract json string", status); + } + + resultValue = LockObject(NKikimr::NMiniKQL::ValueFromString(DataSlot, rawString)); + if (Y_UNLIKELY(!resultValue)) { + status = TStatus::Fail(EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse data type " << DataTypeName << " from json string: '" << TruncateString(rawString) << "'"); + } + return; + } + + case simdjson::builtin::ondemand::json_type::array: + case simdjson::builtin::ondemand::json_type::object: { + std::string_view rawJson; + CHECK_JSON_ERROR(jsonValue.raw_json().get(rawJson)) { + return GetParsingError(error, jsonValue, "extract json value", status); + } + status = TStatus::Fail(EStatusId::PRECONDITION_FAILED, TStringBuilder() << "Found unexpected nested value (raw: '" << TruncateString(rawJson) << "'), expected data type " < + Y_FORCE_INLINE static void ParseJsonNumber(simdjson::simdjson_result jsonNumber, NYql::NUdf::TUnboxedValue& resultValue, TStatus& status) { + CHECK_JSON_ERROR(jsonNumber.error()) { + status = TStatus::Fail(EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to extract json integer number, error: " << simdjson::error_message(error)); + return; + } + + TJsonNumber number = jsonNumber.value(); + if (Y_UNLIKELY(number < std::numeric_limits::min() || std::numeric_limits::max() < number)) { + status = TStatus::Fail(EStatusId::BAD_REQUEST, TStringBuilder() << "Number is out of range [" << ToString(std::numeric_limits::min()) << ", " << ToString(std::numeric_limits::max()) << "]"); + return; + } + + resultValue = NYql::NUdf::TUnboxedValuePod(static_cast(number)); + } + + template + Y_FORCE_INLINE static void ParseJsonDouble(simdjson::simdjson_result jsonNumber, NYql::NUdf::TUnboxedValue& resultValue, TStatus& status) { + CHECK_JSON_ERROR(jsonNumber.error()) { + status = TStatus::Fail(EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to extract json float number, error: " << simdjson::error_message(error)); + return; + } + + resultValue = NYql::NUdf::TUnboxedValuePod(static_cast(jsonNumber.value())); + } + + static void GetParsingError(simdjson::error_code error, simdjson::builtin::ondemand::value jsonValue, const TString& description, TStatus& status) { + status = TStatus::Fail(EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to " << description << ", current token: '" << TruncateString(jsonValue.raw_json_token()) << "', error: " << simdjson::error_message(error)); + } + +private: + NYql::NUdf::EDataSlot DataSlot; + TString DataTypeName; + bool IsOptional = false; + + TVector ParsedRows; + TStatus Status; +}; + +class TJsonParser : public TTopicParserBase { +public: + using TBase = TTopicParserBase; + using TPtr = TIntrusivePtr; + +public: + TJsonParser(IParsedDataConsumer::TPtr consumer, const TJsonParserConfig& config, const TCountersDesc& counters) + : TBase(std::move(consumer), __LOCATION__, counters) + , Config(config) + , NumberColumns(Consumer->GetColumns().size()) + , MaxNumberRows((config.BufferCellCount - 1) / NumberColumns + 1) + , LogPrefix("TJsonParser: ") + , ParsedValues(NumberColumns) + { + Columns.reserve(NumberColumns); + for (const auto& column : Consumer->GetColumns()) { + Columns.emplace_back(column.Name, column.TypeYson, MaxNumberRows); + } + + ColumnsIndex.reserve(NumberColumns); + for (size_t i = 0; i < NumberColumns; i++) { + ColumnsIndex.emplace(std::string_view(Columns[i].Name), i); + } + + for (size_t i = 0; i < NumberColumns; i++) { + ParsedValues[i].resize(MaxNumberRows); + } + + Buffer.Reserve(Config.BatchSize, MaxNumberRows); + + LOG_ROW_DISPATCHER_INFO("Simdjson active implementation " << simdjson::get_active_implementation()->name()); + Parser.threaded = false; + } + + TStatus InitColumnsParsers() { + for (auto& column : Columns) { + auto typeStatus = ParseTypeYson(column.TypeYson); + if (typeStatus.IsFail()) { + return TStatus(typeStatus).AddParentIssue(TStringBuilder() << "Failed to parse column '" << column.Name << "' type " << column.TypeYson); + } + if (auto status = column.InitParser(typeStatus.DetachResult()); status.IsFail()) { + return status.AddParentIssue(TStringBuilder() << "Failed to create parser for column '" << column.Name << "' with type " << column.TypeYson); + } + } + return TStatus::Success(); + } + +public: + void ParseMessages(const TVector& messages) override { + LOG_ROW_DISPATCHER_TRACE("Add " << messages.size() << " messages to parse"); + + Y_ENSURE(!Buffer.Finished, "Cannot parse messages with finished buffer"); + for (const auto& message : messages) { + Buffer.AddMessage(message); + if (Buffer.IsReady() && (Buffer.NumberValues >= MaxNumberRows || Buffer.GetSize() >= Config.BatchSize)) { + ParseBuffer(); + } + } + + if (Buffer.IsReady()) { + if (!Config.LatencyLimit) { + ParseBuffer(); + } else { + LOG_ROW_DISPATCHER_TRACE("Collecting data to parse, skip parsing, current buffer size: " << Buffer.GetSize()); + } + } + } + + void Refresh(bool force) override { + TBase::Refresh(force); + + if (!Buffer.IsReady()) { + return; + } + + const auto creationDuration = TInstant::Now() - Buffer.CreationStartTime; + if (force || creationDuration > Config.LatencyLimit) { + ParseBuffer(); + } else { + LOG_ROW_DISPATCHER_TRACE("Refresh, skip parsing, buffer creation duration: " << creationDuration); + } + } + + const TVector& GetOffsets() const override { + return Buffer.Offsets; + } + + TValueStatus*> GetParsedColumn(ui64 columnId) const override { + if (auto status = Columns[columnId].GetStatus(); status.IsFail()) { + return status; + } + return &ParsedValues[columnId]; + } + +protected: + TStatus DoParsing() override { + Y_ENSURE(Buffer.IsReady(), "Nothing to parse"); + Y_ENSURE(Buffer.NumberValues <= MaxNumberRows, "Too many values to parse"); + + const auto [values, size] = Buffer.Finish(); + LOG_ROW_DISPATCHER_TRACE("Do parsing, first offset: " << Buffer.Offsets.front() << ", values:\n" << values); + + simdjson::ondemand::document_stream documents; + CHECK_JSON_ERROR(Parser.iterate_many(values, size, simdjson::ondemand::DEFAULT_BATCH_SIZE).get(documents)) { + return TStatus::Fail(EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse message batch from offset " << Buffer.Offsets.front() << ", json documents was corrupted: " << simdjson::error_message(error) << " Current data batch: " << TruncateString(std::string_view(values, size))); + } + + size_t rowId = 0; + for (auto document : documents) { + if (Y_UNLIKELY(rowId >= Buffer.NumberValues)) { + return TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Failed to parse json messages, expected " << Buffer.NumberValues << " json rows from offset " << Buffer.Offsets.front() << " but got " << rowId + 1 << " (expected one json row for each offset from topic API in json each row format, maybe initial data was corrupted or messages is not in json format), current data batch: " << TruncateString(std::string_view(values, size))); + } + + const ui64 offset = Buffer.Offsets[rowId]; + CHECK_JSON_ERROR(document.error()) { + return TStatus::Fail(EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse json message for offset " << offset << ", json document was corrupted: " << simdjson::error_message(error) << " Current data batch: " << TruncateString(std::string_view(values, size))); + } + + for (auto item : document.get_object()) { + CHECK_JSON_ERROR(item.error()) { + return TStatus::Fail(EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse json message for offset " << offset << ", json item was corrupted: " << simdjson::error_message(error) << " Current data batch: " << TruncateString(std::string_view(values, size))); + } + + const auto it = ColumnsIndex.find(item.escaped_key().value()); + if (it == ColumnsIndex.end()) { + continue; + } + + const size_t columnId = it->second; + Columns[columnId].ParseJsonValue(offset, rowId, item.value(), ParsedValues[columnId][rowId]); + } + rowId++; + } + + if (Y_UNLIKELY(rowId != Buffer.NumberValues)) { + return TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Failed to parse json messages, expected " << Buffer.NumberValues << " json rows from offset " << Buffer.Offsets.front() << " but got " << rowId << " (expected one json row for each offset from topic API in json each row format, maybe initial data was corrupted or messages is not in json format), current data batch: " << TruncateString(std::string_view(values, size))); + } + + const ui64 firstOffset = Buffer.Offsets.front(); + for (auto& column : Columns) { + column.ValidateNumberValues(rowId, firstOffset); + } + + return TStatus::Success(); + } + + void ClearBuffer() override { + for (size_t i = 0; i < Columns.size(); ++i) { + auto& parsedColumn = ParsedValues[i]; + for (size_t rowId : Columns[i].GetParsedRows()) { + ClearObject(parsedColumn[rowId]); + } + Columns[i].ClearParsedRows(); + } + Buffer.Clear(); + } + +private: + const TJsonParserConfig Config; + const ui64 NumberColumns; + const ui64 MaxNumberRows; + const TString LogPrefix; + + TVector Columns; + absl::flat_hash_map ColumnsIndex; + + TJsonParserBuffer Buffer; + simdjson::ondemand::parser Parser; + TVector> ParsedValues; +}; + +} // anonymous namespace + +TValueStatus CreateJsonParser(IParsedDataConsumer::TPtr consumer, const TJsonParserConfig& config, const TCountersDesc& counters) { + TJsonParser::TPtr parser = MakeIntrusive(consumer, config, counters); + if (auto status = parser->InitColumnsParsers(); status.IsFail()) { + return status; + } + + return ITopicParser::TPtr(parser); +} + +TJsonParserConfig CreateJsonParserConfig(const NConfig::TJsonParserConfig& parserConfig) { + TJsonParserConfig result; + if (const auto batchSize = parserConfig.GetBatchSizeBytes()) { + result.BatchSize = batchSize; + } + if (const auto bufferCellCount = parserConfig.GetBufferCellCount()) { + result.BufferCellCount = bufferCellCount; + } + result.LatencyLimit = TDuration::MilliSeconds(parserConfig.GetBatchCreationTimeoutMs()); + return result; +} + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/json_parser.h b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/json_parser.h new file mode 100644 index 000000000000..c8afd38464f9 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/json_parser.h @@ -0,0 +1,20 @@ +#pragma once + +#include "parser_abstract.h" + +#include + +#include + +namespace NFq::NRowDispatcher { + +struct TJsonParserConfig { + ui64 BatchSize = 1_MB; + TDuration LatencyLimit; + ui64 BufferCellCount = 1000000; // (number rows) * (number columns) limit +}; + +TValueStatus CreateJsonParser(IParsedDataConsumer::TPtr consumer, const TJsonParserConfig& config, const TCountersDesc& counters); +TJsonParserConfig CreateJsonParserConfig(const NConfig::TJsonParserConfig& parserConfig); + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_abstract.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_abstract.cpp new file mode 100644 index 000000000000..696eab2c06fb --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_abstract.cpp @@ -0,0 +1 @@ +#include "parser_abstract.h" diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_abstract.h b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_abstract.h new file mode 100644 index 000000000000..c952ba8e4b56 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_abstract.h @@ -0,0 +1,37 @@ +#pragma once + +#include +#include + +#include + +#include + +namespace NFq::NRowDispatcher { + +class IParsedDataConsumer : public TThrRefBase { +public: + using TPtr = TIntrusivePtr; + +public: + virtual const TVector& GetColumns() const = 0; + + virtual void OnParsingError(TStatus status) = 0; + virtual void OnParsedData(ui64 numberRows) = 0; +}; + +class ITopicParser : public TThrRefBase, public TNonCopyable { +public: + using TPtr = TIntrusivePtr; + +public: + virtual void ParseMessages(const TVector& messages) = 0; + virtual void Refresh(bool force = false) = 0; + + virtual const TVector& GetOffsets() const = 0; + virtual TValueStatus*> GetParsedColumn(ui64 columnId) const = 0; + + virtual void FillStatistics(TFormatHandlerStatistic& statistic) = 0; +}; + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_base.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_base.cpp new file mode 100644 index 000000000000..916795ea4ff5 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_base.cpp @@ -0,0 +1,140 @@ +#include "parser_base.h" + +#include + +#include +#include +#include + +namespace NFq::NRowDispatcher { + +//// TTypeParser + +TTypeParser::TTypeParser(const TSourceLocation& location, const TCountersDesc& counters) + : Alloc(location, NKikimr::TAlignedPagePoolCounters(counters.CountersRoot, counters.MkqlCountersName), true, false) + , FunctionRegistry(NKikimr::NMiniKQL::CreateFunctionRegistry(&PrintBackTrace, NKikimr::NMiniKQL::CreateBuiltinRegistry(), false, {})) + , TypeEnv(std::make_unique(Alloc)) + , ProgramBuilder(std::make_unique(*TypeEnv, *FunctionRegistry)) +{} + +TTypeParser::~TTypeParser() { + with_lock (Alloc) { + FunctionRegistry.Reset(); + TypeEnv.reset(); + ProgramBuilder.reset(); + } +} + +TValueStatus TTypeParser::ParseTypeYson(const TString& typeYson) const { + TString parseTypeError; + TStringOutput errorStream(parseTypeError); + NKikimr::NMiniKQL::TType* typeMkql = nullptr; + + with_lock (Alloc) { + typeMkql = NYql::NCommon::ParseTypeFromYson(TStringBuf(typeYson), *ProgramBuilder, errorStream); + } + + if (!typeMkql) { + return TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Failed to parse type from yson: " << parseTypeError); + } + return typeMkql; +} + +//// TTopicParserBase::TStats + +void TTopicParserBase::TStats::AddParserLatency(TDuration parserLatency) { + ParserLatency = std::max(ParserLatency, parserLatency); +} + +void TTopicParserBase::TStats::AddParseAndFilterLatency(TDuration parseAndFilterLatency) { + ParseAndFilterLatency = std::max(ParseAndFilterLatency, parseAndFilterLatency); +} + +void TTopicParserBase::TStats::Clear() { + ParserLatency = TDuration::Zero(); + ParseAndFilterLatency = TDuration::Zero(); +} + +//// TTopicParserBase + +TTopicParserBase::TTopicParserBase(IParsedDataConsumer::TPtr consumer, const TSourceLocation& location, const TCountersDesc& counters) + : TTypeParser(location, counters) + , Consumer(std::move(consumer)) +{} + +void TTopicParserBase::Refresh(bool force) { + Y_UNUSED(force); +} + +void TTopicParserBase::FillStatistics(TFormatHandlerStatistic& statistic) { + statistic.ParseAndFilterLatency = Stats.ParseAndFilterLatency; + statistic.ParserStats.ParserLatency = Stats.ParserLatency; + + Stats.Clear(); +} + +void TTopicParserBase::ParseBuffer() { + const TInstant startParseAndFilter = TInstant::Now(); + try { + Y_DEFER { + with_lock(Alloc) { + ClearBuffer(); + } + }; + + TStatus status = TStatus::Success(); + const TInstant startParse = TInstant::Now(); + with_lock(Alloc) { + status = DoParsing(); + } + Stats.AddParserLatency(TInstant::Now() - startParse); + + if (status.IsSuccess()) { + Consumer->OnParsedData(GetOffsets().size()); + } else { + Consumer->OnParsingError(status); + } + } catch (...) { + Consumer->OnParsingError(TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Failed to parse messages from offset " << GetOffsets().front() << ", got unexpected exception: " << CurrentExceptionMessage())); + } + Stats.AddParseAndFilterLatency(TInstant::Now() - startParseAndFilter); +} + +//// Functions + +NYql::NUdf::TUnboxedValue LockObject(NYql::NUdf::TUnboxedValue&& value) { + // All UnboxedValue's with type Boxed or String should be locked + // because after parsing they will be used under another MKQL allocator in purecalc filters + + const i32 numberRefs = value.LockRef(); + + // -1 - value is embbeded or empty, otherwise value should have exactly one ref + Y_ENSURE(numberRefs == -1 || numberRefs == 1); + + return value; +} + +void ClearObject(NYql::NUdf::TUnboxedValue& value) { + // Value should be unlocked with same number of refs + value.UnlockRef(1); + value.Clear(); +} + +TValueStatus GetDataSlot(const NKikimr::NMiniKQL::TType* type) { + Y_ENSURE(type->GetKind() == NKikimr::NMiniKQL::TTypeBase::EKind::Data, "Expected data type"); + + const auto* dataType = AS_TYPE(NKikimr::NMiniKQL::TDataType, type); + if (const auto dataSlot = dataType->GetDataSlot()) { + return *dataSlot; + } + return TStatus::Fail(EStatusId::UNSUPPORTED, TStringBuilder() << "Unsupported data type with id: " << dataType->GetSchemeType()); +} + +TString TruncateString(std::string_view rawString, size_t maxSize) { + if (rawString.size() <= maxSize) { + return TString(rawString); + } + return TStringBuilder() << rawString.substr(0, maxSize) << " truncated..."; +} + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_base.h b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_base.h new file mode 100644 index 000000000000..b28fadbd4c29 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/parser_base.h @@ -0,0 +1,68 @@ +#pragma once + +#include "parser_abstract.h" + +#include +#include + +namespace NFq::NRowDispatcher { + +class TTypeParser { +public: + explicit TTypeParser(const TSourceLocation& location, const TCountersDesc& counters); + virtual ~TTypeParser(); + + TValueStatus ParseTypeYson(const TString& typeYson) const; + +protected: + NKikimr::NMiniKQL::TScopedAlloc Alloc; + NKikimr::NMiniKQL::IFunctionRegistry::TPtr FunctionRegistry; + std::unique_ptr TypeEnv; + std::unique_ptr ProgramBuilder; +}; + +class TTopicParserBase : public ITopicParser, public TTypeParser { +private: + struct TStats { + void AddParserLatency(TDuration parserLatency); + void AddParseAndFilterLatency(TDuration parseAndFilterLatency); + void Clear(); + + TDuration ParserLatency; + TDuration ParseAndFilterLatency; + }; + +public: + using TPtr = TIntrusivePtr; + +public: + TTopicParserBase(IParsedDataConsumer::TPtr consumer, const TSourceLocation& location, const TCountersDesc& counters); + virtual ~TTopicParserBase() = default; + +public: + virtual void Refresh(bool force = false) override; + virtual void FillStatistics(TFormatHandlerStatistic& statistic) override; + +protected: + // Called with binded alloc + virtual TStatus DoParsing() = 0; + virtual void ClearBuffer() = 0; + +protected: + void ParseBuffer(); + +protected: + const IParsedDataConsumer::TPtr Consumer; + +private: + TStats Stats; +}; + +NYql::NUdf::TUnboxedValue LockObject(NYql::NUdf::TUnboxedValue&& value); +void ClearObject(NYql::NUdf::TUnboxedValue& value); + +TValueStatus GetDataSlot(const NKikimr::NMiniKQL::TType* type); + +TString TruncateString(std::string_view rawString, size_t maxSize = 1_KB); + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/raw_parser.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/raw_parser.cpp new file mode 100644 index 000000000000..32d4f4c4c17c --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/raw_parser.cpp @@ -0,0 +1,125 @@ +#include "raw_parser.h" + +#include "parser_base.h" + +#include + +#include +#include + +namespace NFq::NRowDispatcher { + +namespace { + +class TRawParser : public TTopicParserBase { +public: + using TBase = TTopicParserBase; + using TPtr = TIntrusivePtr; + +public: + TRawParser(IParsedDataConsumer::TPtr consumer, const TSchemaColumn& schema, const TCountersDesc& counters) + : TBase(std::move(consumer), __LOCATION__, counters) + , Schema(schema) + , LogPrefix("TRawParser: ") + {} + + TStatus InitColumnParser() { + auto typeStatus = ParseTypeYson(Schema.TypeYson); + if (typeStatus.IsFail()) { + return typeStatus; + } + + for (NKikimr::NMiniKQL::TType* type = typeStatus.DetachResult(); true; type = AS_TYPE(NKikimr::NMiniKQL::TOptionalType, type)->GetItemType()) { + if (type->GetKind() == NKikimr::NMiniKQL::TTypeBase::EKind::Data) { + auto slotStatus = GetDataSlot(type); + if (slotStatus.IsFail()) { + return slotStatus; + } + DataSlot = slotStatus.DetachResult(); + return TStatus::Success(); + } + + if (type->GetKind() != NKikimr::NMiniKQL::TTypeBase::EKind::Optional) { + return TStatus::Fail(EStatusId::UNSUPPORTED, TStringBuilder() << "Unsupported type kind for raw format: " << type->GetKindAsStr()); + } + + NumberOptionals++; + } + + return TStatus::Success(); + } + +public: + void ParseMessages(const TVector& messages) override { + LOG_ROW_DISPATCHER_TRACE("Add " << messages.size() << " messages to parse"); + + for (const auto& message : messages) { + CurrentMessage = message.GetData(); + Offsets.emplace_back(message.GetOffset()); + ParseBuffer(); + } + } + + const TVector& GetOffsets() const override { + return Offsets; + } + + TValueStatus*> GetParsedColumn(ui64 columnId) const override { + Y_ENSURE(columnId == 0, "Invalid column id for raw parser"); + return &ParsedColumn; + } + +protected: + TStatus DoParsing() override { + LOG_ROW_DISPATCHER_TRACE("Do parsing, first offset: " << Offsets.front() << ", value: " << CurrentMessage); + + NYql::NUdf::TUnboxedValue value = LockObject(NKikimr::NMiniKQL::ValueFromString(DataSlot, CurrentMessage)); + if (value) { + for (size_t i = 0; i < NumberOptionals; ++i) { + value = value.MakeOptional(); + } + } else if (!NumberOptionals) { + return TStatus::Fail(EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse massege at offset " << Offsets.back() << ", can't parse data type " << NYql::NUdf::GetDataTypeInfo(DataSlot).Name << " from string: '" << TruncateString(CurrentMessage) << "'"); + } + + ParsedColumn.emplace_back(std::move(value)); + return TStatus::Success(); + } + + void ClearBuffer() override { + for (auto& parsedValue : ParsedColumn) { + ClearObject(parsedValue); + } + ParsedColumn.clear(); + Offsets.clear(); + } + +private: + const TSchemaColumn Schema; + const TString LogPrefix; + + NYql::NUdf::EDataSlot DataSlot; + ui64 NumberOptionals = 0; + + TString CurrentMessage; + TVector Offsets; + TVector ParsedColumn; +}; + +} // anonymous namespace + +TValueStatus CreateRawParser(IParsedDataConsumer::TPtr consumer, const TCountersDesc& counters) { + const auto& columns = consumer->GetColumns(); + if (columns.size() != 1) { + return TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Expected only one column for raw format, but got " << columns.size()); + } + + TRawParser::TPtr parser = MakeIntrusive(consumer, columns[0], counters); + if (auto status = parser->InitColumnParser(); status.IsFail()) { + return status.AddParentIssue(TStringBuilder() << "Failed to create raw parser for column " << columns[0].ToString()); + } + + return ITopicParser::TPtr(parser); +} + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/raw_parser.h b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/raw_parser.h new file mode 100644 index 000000000000..6588844884a3 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/raw_parser.h @@ -0,0 +1,9 @@ +#pragma once + +#include "parser_abstract.h" + +namespace NFq::NRowDispatcher { + +TValueStatus CreateRawParser(IParsedDataConsumer::TPtr consumer, const TCountersDesc& counters); + +} // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/ya.make b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/ya.make new file mode 100644 index 000000000000..ec51348a980c --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/parsers/ya.make @@ -0,0 +1,33 @@ +LIBRARY() + +SRCS( + parser_abstract.cpp + parser_base.cpp + json_parser.cpp + raw_parser.cpp +) + +PEERDIR( + contrib/libs/simdjson + + library/cpp/containers/absl_flat_hash + + ydb/core/fq/libs/actors/logging + ydb/core/fq/libs/row_dispatcher/events + ydb/core/fq/libs/row_dispatcher/format_handler/common + + ydb/public/sdk/cpp/client/ydb_topic/include + + yql/essentials/minikql + yql/essentials/minikql/dom + yql/essentials/minikql/invoke_builtins + yql/essentials/providers/common/schema +) + +CFLAGS( + -Wno-assume +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/ut/common/ut_common.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/common/ut_common.cpp new file mode 100644 index 000000000000..0b0c8dc0cc94 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/common/ut_common.cpp @@ -0,0 +1,252 @@ +#include "ut_common.h" + +#include + +#include + +#include + +#include + +#include +#include + +namespace NFq::NRowDispatcher::NTests { + +namespace { + +class TPurecalcCompileServiceMock : public NActors::TActor { + using TBase = NActors::TActor; + +public: + TPurecalcCompileServiceMock(NActors::TActorId owner) + : TBase(&TPurecalcCompileServiceMock::StateFunc) + , Owner(owner) + , ProgramFactory(NYql::NPureCalc::MakeProgramFactory()) + {} + + STRICT_STFUNC(StateFunc, + hFunc(TEvRowDispatcher::TEvPurecalcCompileRequest, Handle); + ) + + void Handle(TEvRowDispatcher::TEvPurecalcCompileRequest::TPtr& ev) { + IProgramHolder::TPtr programHolder = std::move(ev->Get()->ProgramHolder); + + try { + programHolder->CreateProgram(ProgramFactory); + } catch (const NYql::NPureCalc::TCompileError& e) { + UNIT_FAIL("Failed to compile purecalc filter: sql: " << e.GetYql() << ", error: " << e.GetIssues()); + } + + Send(ev->Sender, new TEvRowDispatcher::TEvPurecalcCompileResponse(std::move(programHolder)), 0, ev->Cookie); + Send(Owner, new NActors::TEvents::TEvPing()); + } + +private: + const NActors::TActorId Owner; + const NYql::NPureCalc::IProgramFactoryPtr ProgramFactory; +}; + +void SegmentationFaultHandler(int) { + Cerr << "segmentation fault call stack:" << Endl; + FormatBackTrace(&Cerr); + abort(); +} + +//// TBaseFixture::ICell + +class TOptionalCell : public TBaseFixture::ICell { +public: + TOptionalCell(ICell::TPtr value) + : Value(value) + {} + +public: + TString GetType() const override { + return TStringBuilder() << "[OptionalType; " << Value->GetType() << "]"; + } + + void Validate(const NYql::NUdf::TUnboxedValue& parsedValue) const override { + if (!parsedValue) { + UNIT_FAIL("Unexpected NULL value for optional cell"); + return; + } + Value->Validate(parsedValue.GetOptionalValue()); + } + +private: + const ICell::TPtr Value; +}; + +class TStringSell : public TBaseFixture::ICell { +public: + TStringSell(const TString& value) + : Value(value) + {} + +public: + TString GetType() const override { + return "[DataType; String]"; + } + + void Validate(const NYql::NUdf::TUnboxedValue& parsedValue) const override { + UNIT_ASSERT_VALUES_EQUAL(Value, TString(parsedValue.AsStringRef())); + } + +private: + const TString Value; +}; + +class TUint64Sell : public TBaseFixture::ICell { +public: + TUint64Sell(ui64 value) + : Value(value) + {} + +public: + TString GetType() const override { + return "[DataType; Uint64]"; + } + + void Validate(const NYql::NUdf::TUnboxedValue& parsedValue) const override { + UNIT_ASSERT_VALUES_EQUAL(Value, parsedValue.Get()); + } + +private: + const ui64 Value; +}; + +} // anonymous namespace + +//// TBaseFixture::TRow + +TBaseFixture::TRow& TBaseFixture::TRow::AddCell(ICell::TPtr cell, bool optional) { + if (optional) { + cell = MakeIntrusive(cell); + } + + Cells.emplace_back(cell); + return *this; +} + +TBaseFixture::TRow& TBaseFixture::TRow::AddString(const TString& value, bool optional) { + return AddCell(MakeIntrusive(value), optional); +} + +TBaseFixture::TRow& TBaseFixture::TRow::AddUint64(ui64 value, bool optional) { + return AddCell(MakeIntrusive(value), optional); +} + +//// TBaseFixture::TBatch + +TBaseFixture::TBatch::TBatch(std::initializer_list rows) + : Rows(rows) +{} + +TBaseFixture::TBatch& TBaseFixture::TBatch::AddRow(TRow row) { + Rows.emplace_back(row); + return *this; +} + +//// TBaseFixture + +TBaseFixture::TBaseFixture() + : TTypeParser(__LOCATION__, {}) + , MemoryInfo("TBaseFixture alloc") + , HolderFactory(std::make_unique(Alloc.Ref(), MemoryInfo)) + , Runtime(1, true) +{ + NKikimr::EnableYDBBacktraceFormat(); + signal(SIGSEGV, &SegmentationFaultHandler); + + Alloc.Ref().UseRefLocking = true; +} + +void TBaseFixture::SetUp(NUnitTest::TTestContext&) { + // Init runtime + TAutoPtr app = new NKikimr::TAppPrepare(); + Runtime.SetLogBackend(NActors::CreateStderrBackend()); + Runtime.SetLogPriority(NKikimrServices::FQ_ROW_DISPATCHER, NActors::NLog::PRI_TRACE); + Runtime.SetDispatchTimeout(WAIT_TIMEOUT); + Runtime.Initialize(app->Unwrap()); + + // Init tls context + auto* actorSystem = Runtime.GetActorSystem(0); + Mailbox = std::make_unique(); + ExecutorThread = std::make_unique(0, actorSystem, nullptr, nullptr, "test thread"); + ActorCtx = std::make_unique(*Mailbox, *ExecutorThread, GetCycleCountFast(), NActors::TActorId()); + PrevActorCtx = NActors::TlsActivationContext; + NActors::TlsActivationContext = ActorCtx.get(); +} + +void TBaseFixture::TearDown(NUnitTest::TTestContext&) { + with_lock(Alloc) { + ProgramBuilder.reset(); + TypeEnv.reset(); + FunctionRegistry.Reset(); + HolderFactory.reset(); + } + + // Release tls context + NActors::TlsActivationContext = PrevActorCtx; + PrevActorCtx = nullptr; +} + +void TBaseFixture::CheckMessageBatch(TRope serializedBatch, const TBatch& expectedBatch) const { + const auto& expectedRows = expectedBatch.Rows; + UNIT_ASSERT_C(!expectedRows.empty(), "Expected batch should not be empty"); + + // Parse row type (take first row for infer) + const auto& expectedFirstRow = expectedRows.front().Cells; + TVector columnTypes; + columnTypes.reserve(expectedFirstRow.size()); + for (const auto& cell : expectedFirstRow) { + columnTypes.emplace_back(CheckSuccess(ParseTypeYson(cell->GetType()))); + } + + with_lock(Alloc) { + // Parse messages + const auto rowType = ProgramBuilder->NewMultiType(columnTypes); + const auto dataUnpacker = std::make_unique>(rowType); + + NKikimr::NMiniKQL::TUnboxedValueBatch parsedData(rowType); + dataUnpacker->UnpackBatch(NYql::MakeChunkedBuffer(std::move(serializedBatch)), *HolderFactory, parsedData); + + // Validate data + UNIT_ASSERT_VALUES_EQUAL(parsedData.RowCount(), expectedRows.size()); + UNIT_ASSERT_VALUES_EQUAL(parsedData.Width(), expectedFirstRow.size()); + + ui64 rowIndex = 0; + parsedData.ForEachRowWide([&](NYql::NUdf::TUnboxedValue* values, ui32 width){ + UNIT_ASSERT_GE(expectedRows.size(), rowIndex + 1); + const auto& expectedRow = expectedRows[rowIndex++].Cells; + + UNIT_ASSERT_VALUES_EQUAL(width, expectedRow.size()); + for (ui32 i = 0; i < width; ++i) { + expectedRow[i]->Validate(values[i]); + } + }); + } +} + +NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage TBaseFixture::GetMessage(ui64 offset, const TString& data) { + NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessageInformation info(offset, "", 0, TInstant::Zero(), TInstant::Zero(), nullptr, nullptr, 0, ""); + return NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage(data, nullptr, info, nullptr); +} + +//// Functions + +NActors::IActor* CreatePurecalcCompileServiceMock(NActors::TActorId owner) { + return new TPurecalcCompileServiceMock(owner); +} + +void CheckSuccess(const TStatus& status) { + UNIT_ASSERT_C(status.IsSuccess(), "Status is not success, " << status.GetErrorMessage()); +} + +void CheckError(const TStatus& status, TStatusCode expectedStatusCode, const TString& expectedMessage) { + UNIT_ASSERT_C(status.GetStatus() == expectedStatusCode, "Expected error status " << NYql::NDqProto::StatusIds_StatusCode_Name(expectedStatusCode) << ", but got: " << status.GetErrorMessage()); + UNIT_ASSERT_STRING_CONTAINS_C(status.GetErrorMessage(), expectedMessage, "Unexpected error message, Status: " << NYql::NDqProto::StatusIds_StatusCode_Name(status.GetStatus())); +} + +} // namespace NFq::NRowDispatcher::NTests diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/ut/common/ut_common.h b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/common/ut_common.h new file mode 100644 index 000000000000..ce54150b4608 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/common/ut_common.h @@ -0,0 +1,89 @@ +#pragma once + +#include + +#include + +#include + +#include + +namespace NFq::NRowDispatcher::NTests { + +static constexpr TDuration WAIT_TIMEOUT = TDuration::Seconds(20); + +class TBaseFixture : public NUnitTest::TBaseFixture, public TTypeParser { +public: + // Helper classes for checking serialized rows in multi type format + class ICell : public TThrRefBase { + public: + using TPtr = TIntrusivePtr; + + public: + virtual TString GetType() const = 0; + virtual void Validate(const NYql::NUdf::TUnboxedValue& parsedValue) const = 0; + }; + + class TRow { + public: + TRow() = default; + + TRow& AddCell(ICell::TPtr cell, bool optional); + + TRow& AddString(const TString& value, bool optional = false); + TRow& AddUint64(ui64 value, bool optional = false); + + public: + TVector Cells; + }; + + class TBatch { + public: + TBatch() = default; + TBatch(std::initializer_list rows); + + TBatch& AddRow(TRow row); + + public: + TVector Rows; + }; + +public: + TBaseFixture(); + +public: + virtual void SetUp(NUnitTest::TTestContext& ctx) override; + virtual void TearDown(NUnitTest::TTestContext& ctx) override; + +public: + void CheckMessageBatch(TRope serializedBatch, const TBatch& expectedBatch) const; + + static NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage GetMessage(ui64 offset, const TString& data); + +public: + NKikimr::NMiniKQL::TMemoryUsageInfo MemoryInfo; + std::unique_ptr HolderFactory; + NActors::TTestActorRuntime Runtime; + +private: + // Like NKikimr::TActorSystemStub but with Runtime as actor system in tls context + // it enables logging in unit test thread + // and using NActors::TActivationContext::ActorSystem() method + std::unique_ptr Mailbox; + std::unique_ptr ExecutorThread; + std::unique_ptr ActorCtx; + NActors::TActivationContext* PrevActorCtx; +}; + +NActors::IActor* CreatePurecalcCompileServiceMock(NActors::TActorId owner); + +void CheckSuccess(const TStatus& status); +void CheckError(const TStatus& status, TStatusCode expectedStatusCode, const TString& expectedMessage); + +template +TValue CheckSuccess(TValueStatus valueStatus) { + UNIT_ASSERT_C(valueStatus.IsSuccess(), "Value status is not success, " << valueStatus.GetErrorMessage()); + return valueStatus.DetachResult(); +} + +} // namespace NFq::NRowDispatcher::NTests diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/ut/common/ya.make b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/common/ya.make new file mode 100644 index 000000000000..11a0802e50b6 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/common/ya.make @@ -0,0 +1,29 @@ +LIBRARY() + +SRCS( + ut_common.cpp +) + +PEERDIR( + library/cpp/testing/unittest + + ydb/core/base + + ydb/core/fq/libs/row_dispatcher/format_handler/common + + ydb/core/testlib + ydb/core/testlib/actors + ydb/core/testlib/basics + + ydb/library/yql/dq/actors + ydb/library/yql/dq/common + + yql/essentials/minikql + yql/essentials/minikql/computation + yql/essentials/minikql/invoke_builtins + yql/essentials/providers/common/schema/mkql +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/ut/format_handler_ut.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/format_handler_ut.cpp new file mode 100644 index 000000000000..0e436f22fa27 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/format_handler_ut.cpp @@ -0,0 +1,389 @@ +#include +#include + +namespace NFq::NRowDispatcher::NTests { + +namespace { + +class TFormatHandlerFixture : public TBaseFixture { +public: + using TBase = TBaseFixture; + using TCallback = std::function>>&& data)>; + + class TClientDataConsumer : public IClientDataConsumer { + public: + using TPtr = TIntrusivePtr; + + public: + TClientDataConsumer(NActors::TActorId clientId, const TVector& columns, const TString& whereFilter, TCallback callback, ui64 expectedFilteredRows = 0) + : Callback(callback) + , ClientId(clientId) + , Columns(columns) + , WhereFilter(whereFilter) + , ExpectedFilteredRows(expectedFilteredRows) + {} + + void Freeze() { + Frozen = true; + } + + void ExpectOffsets(const TVector expectedOffsets) { + if (Frozen) { + return; + } + for (const ui64 offset : expectedOffsets) { + Offsets.emplace(offset); + } + HasData = false; + } + + void ExpectError(TStatusCode statusCode, const TString& message) { + UNIT_ASSERT_C(!ExpectedError, "Can not add existing error, client id: " << ClientId); + ExpectedError = {statusCode, message}; + } + + void Validate() const { + UNIT_ASSERT_C(Offsets.empty(), "Found " << Offsets.size() << " missing batches, client id: " << ClientId); + UNIT_ASSERT_VALUES_EQUAL_C(ExpectedFilteredRows, 0, "Found " << ExpectedFilteredRows << " not filtered rows, client id: " << ClientId); + UNIT_ASSERT_C(!ExpectedError, "Expected error: " << ExpectedError->second << ", client id: " << ClientId); + } + + bool IsStarted() const { + return Started; + } + + bool IsFinished() const { + return Frozen || !HasData; + } + + public: + TVector GetColumns() const override { + return Columns; + } + + const TString& GetWhereFilter() const override { + return WhereFilter; + } + + TPurecalcCompileSettings GetPurecalcSettings() const override { + return {.EnabledLLVM = false}; + } + + virtual NActors::TActorId GetClientId() const override { + return ClientId; + } + + TMaybe GetNextMessageOffset() const override { + return Nothing(); + } + + void OnClientError(TStatus status) override { + UNIT_ASSERT_C(!Offsets.empty(), "Unexpected message batch, status: " << status.GetErrorMessage() << ", client id: " << ClientId); + + if (ExpectedError) { + CheckError(status, ExpectedError->first, ExpectedError->second); + ExpectedError = std::nullopt; + } else { + CheckSuccess(status); + } + } + + void StartClientSession() override { + Started = true; + } + + void AddDataToClient(ui64 offset, ui64 rowSize) override { + UNIT_ASSERT_C(Started, "Unexpected data for not started session"); + UNIT_ASSERT_C(rowSize >= 0, "Expected non zero row size"); + UNIT_ASSERT_C(!ExpectedError, "Expected error: " << ExpectedError->second << ", client id: " << ClientId); + UNIT_ASSERT_C(ExpectedFilteredRows > 0, "Too many rows filtered, client id: " << ClientId); + UNIT_ASSERT_C(!Offsets.empty(), "Unexpected message batch, offset: " << offset << ", client id: " << ClientId); + ExpectedFilteredRows--; + HasData = true; + } + + void UpdateClientOffset(ui64 offset) override { + UNIT_ASSERT_C(Started, "Unexpected offset for not started session"); + UNIT_ASSERT_C(!ExpectedError, "Error is not handled: " << ExpectedError->second << ", client id: " << ClientId); + UNIT_ASSERT_C(!Offsets.empty(), "Unexpected message batch, offset: " << offset << ", client id: " << ClientId); + UNIT_ASSERT_VALUES_EQUAL_C(Offsets.front(), offset, "Unexpected commit offset, client id: " << ClientId); + Offsets.pop(); + } + + public: + const TCallback Callback; + + private: + const NActors::TActorId ClientId; + const TVector Columns; + const TString WhereFilter; + + bool Started = false; + bool HasData = false; + bool Frozen = false; + ui64 ExpectedFilteredRows = 0; + std::queue Offsets; + std::optional> ExpectedError; + }; + +public: + void SetUp(NUnitTest::TTestContext& ctx) override { + TBase::SetUp(ctx); + + CompileNotifier = Runtime.AllocateEdgeActor(); + CompileService = Runtime.Register(CreatePurecalcCompileServiceMock(CompileNotifier)); + + CreateFormatHandler({ + .JsonParserConfig = {}, + .FiltersConfig = {.CompileServiceId = CompileService} + }); + } + + void TearDown(NUnitTest::TTestContext& ctx) override { + for (const auto& client : Clients) { + client->Validate(); + } + ClientIds.clear(); + Clients.clear(); + FormatHandler.Reset(); + + TBase::TearDown(ctx); + } + +public: + void CreateFormatHandler(const TFormatHandlerConfig& config, ITopicFormatHandler::TSettings settings = {.ParsingFormat = "json_each_row"}) { + FormatHandler = CreateTestFormatHandler(config, settings); + } + + [[nodiscard]] TStatus MakeClient(const TVector& columns, const TString& whereFilter, TCallback callback, ui64 expectedFilteredRows = 1) { + ClientIds.emplace_back(ClientIds.size(), 0, 0, 0); + + auto client = MakeIntrusive(ClientIds.back(), columns, whereFilter, callback, expectedFilteredRows); + auto status = FormatHandler->AddClient(client); + if (status.IsFail()) { + return status; + } + + Clients.emplace_back(client); + if (!client->IsStarted()) { + // Wait filter compilation + const auto response = Runtime.GrabEdgeEvent(CompileNotifier, TDuration::Seconds(5)); + UNIT_ASSERT_C(response, "Compilation is not performed for filter: " << whereFilter); + } + + return TStatus::Success(); + } + + void ParseMessages(const TVector& messages, TVector expectedOffsets = {}) { + for (auto& client : Clients) { + client->ExpectOffsets(expectedOffsets ? expectedOffsets : TVector{messages.back().GetOffset()}); + } + FormatHandler->ParseMessages(messages); + ExtractClientsData(); + } + + void CheckClientError(const TVector& messages, NActors::TActorId clientId, TStatusCode statusCode, const TString& message) { + for (auto& client : Clients) { + client->ExpectOffsets({messages.back().GetOffset()}); + if (client->GetClientId() == clientId) { + client->ExpectError(statusCode, message); + } + } + FormatHandler->ParseMessages(messages); + ExtractClientsData(); + } + + void RemoveClient(NActors::TActorId clientId) { + for (auto& client : Clients) { + if (client->GetClientId() == clientId) { + client->Freeze(); + } + } + FormatHandler->RemoveClient(clientId); + } + +public: + static TCallback EmptyCheck() { + return [&](TQueue>>&& data) {}; + } + + static TCallback OneBatchCheck(std::function&& offsets)> callback) { + return [callback](TQueue>>&& data) { + UNIT_ASSERT_VALUES_EQUAL(data.size(), 1); + auto [messages, offsets] = data.front(); + + UNIT_ASSERT(!offsets.empty()); + callback(std::move(messages), std::move(offsets)); + }; + } + + TCallback OneRowCheck(ui64 offset, const TRow& row) const { + return OneBatchCheck([this, offset, row](TRope&& messages, TVector&& offsets) { + UNIT_ASSERT_VALUES_EQUAL(offsets.size(), 1); + UNIT_ASSERT_VALUES_EQUAL(offsets.front(), offset); + + CheckMessageBatch(messages, TBatch().AddRow(row)); + }); + } + +private: + void ExtractClientsData() { + for (auto& client : Clients) { + auto data = FormatHandler->ExtractClientData(client->GetClientId()); + if (client->IsFinished()) { + UNIT_ASSERT_VALUES_EQUAL_C(data.size(), 0, "Expected empty data for finished clients"); + } else { + client->Callback(std::move(data)); + } + } + } + +public: + NActors::TActorId CompileNotifier; + NActors::TActorId CompileService; + TVector ClientIds; + TVector Clients; + ITopicFormatHandler::TPtr FormatHandler; +}; + +} // anonymous namespace + + +Y_UNIT_TEST_SUITE(TestFormatHandler) { + Y_UNIT_TEST_F(ManyJsonClients, TFormatHandlerFixture) { + const ui64 firstOffset = 42; + const TSchemaColumn commonColumn = {"com_col", "[DataType; String]"}; + + CheckSuccess(MakeClient( + {commonColumn, {"col_first", "[DataType; String]"}}, + "WHERE col_first = \"str_first__large__\"", + OneRowCheck(firstOffset + 1, TRow().AddString("event2").AddString("str_first__large__")) + )); + + CheckSuccess(MakeClient( + {commonColumn, {"col_second", "[DataType; String]"}}, + "WHERE col_second = \"str_second\"", + OneRowCheck(firstOffset, TRow().AddString("event1").AddString("str_second")) + )); + + ParseMessages({ + GetMessage(firstOffset, R"({"com_col": "event1", "col_first": "some_str", "col_second": "str_second"})"), + GetMessage(firstOffset + 1, R"({"com_col": "event2", "col_second": "some_str", "col_first": "str_first__large__"})") + }); + + RemoveClient(ClientIds.back()); + + ParseMessages({ + GetMessage(firstOffset + 2, R"({"com_col": "event1", "col_first": "some_str", "col_second": "str_second"})") + }); + } + + Y_UNIT_TEST_F(ManyRawClients, TFormatHandlerFixture) { + CreateFormatHandler( + {.JsonParserConfig = {}, .FiltersConfig = {.CompileServiceId = CompileService}}, + {.ParsingFormat = "raw"} + ); + + const ui64 firstOffset = 42; + const TVector schema = {{"data", "[DataType; String]"}}; + const TVector testData = { + R"({"col_a": "str1__++___str2", "col_b": 12345})", + R"({"col_a": "str13__++___str23", "col_b": ["A", "B", "C"]})", + R"({"col_a": false, "col_b": {"X": "Y"}})" + }; + + CheckSuccess(MakeClient(schema, "WHERE FALSE", EmptyCheck(), 0)); + + const auto trueChacker = OneBatchCheck([&](TRope&& messages, TVector&& offsets) { + TBatch expectedBatch; + for (ui64 offset : offsets) { + UNIT_ASSERT(offset - firstOffset < testData.size()); + expectedBatch.AddRow( + TRow().AddString(testData[offset - firstOffset]) + ); + } + + CheckMessageBatch(messages, expectedBatch); + }); + CheckSuccess(MakeClient(schema, "WHERE TRUE", trueChacker, 3)); + CheckSuccess(MakeClient(schema, "", trueChacker, 2)); + + ParseMessages({ + GetMessage(firstOffset, testData[0]), + GetMessage(firstOffset + 1, testData[1]) + }, {firstOffset, firstOffset + 1}); + + RemoveClient(ClientIds.back()); + + ParseMessages({ + GetMessage(firstOffset + 2, testData[2]) + }); + } + + Y_UNIT_TEST_F(ClientValidation, TFormatHandlerFixture) { + const TVector schema = {{"data", "[DataType; String]"}}; + const TString filter = "WHERE FALSE"; + const auto callback = EmptyCheck(); + CheckSuccess(MakeClient(schema, filter, callback, 0)); + + CheckError( + FormatHandler->AddClient(MakeIntrusive(ClientIds.back(), schema, filter, callback)), + EStatusId::INTERNAL_ERROR, + "Failed to create new client, client with id [0:0:0] already exists" + ); + + CheckError( + MakeClient({{"data", "[OptionalType; [DataType; String]]"}}, filter, callback, 0), + EStatusId::SCHEME_ERROR, + "Failed to modify common parsing schema" + ); + + CheckError( + MakeClient({{"data_2", "[ListType; [DataType; String]]"}}, filter, callback, 0), + EStatusId::UNSUPPORTED, + "Failed to update parser with new client" + ); + } + + Y_UNIT_TEST_F(ClientError, TFormatHandlerFixture) { + const ui64 firstOffset = 42; + const TSchemaColumn commonColumn = {"com_col", "[DataType; String]"}; + + CheckSuccess(MakeClient({commonColumn, {"col_first", "[OptionalType; [DataType; Uint8]]"}}, "WHERE TRUE", EmptyCheck(), 0)); + + CheckSuccess(MakeClient( + {commonColumn, {"col_second", "[DataType; String]"}}, + "WHERE col_second = \"str_second\"", + OneRowCheck(firstOffset, TRow().AddString("event1").AddString("str_second")) + )); + + CheckClientError( + {GetMessage(firstOffset, R"({"com_col": "event1", "col_first": "some_str", "col_second": "str_second"})")}, + ClientIds[0], + EStatusId::BAD_REQUEST, + TStringBuilder() << "Failed to parse json string at offset " << firstOffset << ", got parsing error for column 'col_first' with type [OptionalType; [DataType; Uint8]]" + ); + } + + Y_UNIT_TEST_F(ClientErrorWithEmptyFilter, TFormatHandlerFixture) { + const ui64 firstOffset = 42; + const TSchemaColumn commonColumn = {"com_col", "[DataType; String]"}; + + CheckSuccess(MakeClient({commonColumn, {"col_first", "[DataType; String]"}}, "", EmptyCheck(), 0)); + + CheckSuccess(MakeClient( + {commonColumn, {"col_second", "[DataType; String]"}}, + "WHERE col_second = \"str_second\"", + OneRowCheck(firstOffset, TRow().AddString("event1").AddString("str_second")) + )); + + CheckClientError( + {GetMessage(firstOffset, R"({"com_col": "event1", "col_second": "str_second"})")}, + ClientIds[0], + EStatusId::PRECONDITION_FAILED, + TStringBuilder() << "Failed to parse json messages, found 1 missing values from offset " << firstOffset << " in non optional column 'col_first' with type [DataType; String]" + ); + } +} + +} // namespace NFq::NRowDispatcher::NTests diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/ut/topic_filter_ut.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/topic_filter_ut.cpp new file mode 100644 index 000000000000..55f6fd8f73d4 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/topic_filter_ut.cpp @@ -0,0 +1,434 @@ +#include +#include +#include + +#include + +namespace NFq::NRowDispatcher::NTests { + +namespace { + +class TFiterFixture : public TBaseFixture { +public: + using TBase = TBaseFixture; + using TCallback = std::function; + + class TFilterConsumer : public IFilteredDataConsumer { + public: + using TPtr = TIntrusivePtr; + + public: + TFilterConsumer(const TVector& columns, const TString& whereFilter, TCallback callback, std::optional> compileError) + : Columns(columns) + , WhereFilter(whereFilter) + , Callback(callback) + , CompileError(compileError) + {} + + public: + const TVector& GetColumns() const override { + return Columns; + } + + const TString& GetWhereFilter() const override { + return WhereFilter; + } + + TPurecalcCompileSettings GetPurecalcSettings() const override { + return {.EnabledLLVM = false}; + } + + NActors::TActorId GetFilterId() const override { + return FilterId; + } + + const TVector& GetColumnIds() const override { + return ColumnIds; + } + + TMaybe GetNextMessageOffset() const override { + return Nothing(); + } + + void OnFilterStarted() override { + Started = true; + UNIT_ASSERT_C(!CompileError, "Expected compile error: " << CompileError->second); + } + + void OnFilteringError(TStatus status) override { + if (CompileError) { + Started = true; + CheckError(status, CompileError->first, CompileError->second); + } else { + UNIT_FAIL("Filtering failed: " << status.GetErrorMessage()); + } + } + + void OnFilteredBatch(ui64 firstRow, ui64 lastRow) override { + UNIT_ASSERT_C(Started, "Unexpected data for not started filter"); + for (ui64 rowId = firstRow; rowId <= lastRow; ++rowId) { + Callback(rowId); + } + } + + void OnFilteredData(ui64 rowId) override { + UNIT_ASSERT_C(Started, "Unexpected data for not started filter"); + Callback(rowId); + } + + protected: + NActors::TActorId FilterId; + TVector ColumnIds; + bool Started = false; + + private: + const TVector Columns; + const TString WhereFilter; + const TCallback Callback; + const std::optional> CompileError; + }; + +public: + virtual void SetUp(NUnitTest::TTestContext& ctx) override { + TBase::SetUp(ctx); + + CompileServiceActorId = Runtime.Register(CreatePurecalcCompileService({}, MakeIntrusive())); + } + + virtual void TearDown(NUnitTest::TTestContext& ctx) override { + with_lock (Alloc) { + for (auto& holder : Holders) { + for (auto& value : holder) { + ClearObject(value); + } + } + Holders.clear(); + } + Filter.Reset(); + FilterHandler.Reset(); + + TBase::TearDown(ctx); + } + +public: + virtual TStatus MakeFilter(const TVector& columns, const TString& whereFilter, TCallback callback) { + FilterHandler = MakeIntrusive(columns, whereFilter, callback, CompileError); + + auto filterStatus = CreatePurecalcFilter(FilterHandler); + if (filterStatus.IsFail()) { + return filterStatus; + } + + Filter = filterStatus.DetachResult(); + CompileFilter(); + return TStatus::Success(); + } + + void Push(const TVector*>& values, ui64 numberRows = 0) { + Filter->FilterData(values, numberRows ? numberRows : values.front()->size()); + } + + const TVector* MakeVector(size_t size, std::function valueCreator) { + with_lock (Alloc) { + auto& holder = Holders.emplace_front(); + for (size_t i = 0; i < size; ++i) { + holder.emplace_back(LockObject(valueCreator(i))); + } + return &holder; + } + } + + template + const TVector* MakeVector(const TVector& values, bool optional = false) { + return MakeVector(values.size(), [&](size_t i) { + NYql::NUdf::TUnboxedValuePod unboxedValue = NYql::NUdf::TUnboxedValuePod(values[i]); + return optional ? unboxedValue.MakeOptional() : unboxedValue; + }); + } + + const TVector* MakeStringVector(const TVector& values, bool optional = false) { + return MakeVector(values.size(), [&](size_t i) { + NYql::NUdf::TUnboxedValuePod stringValue = NKikimr::NMiniKQL::MakeString(values[i]); + return optional ? stringValue.MakeOptional() : stringValue; + }); + } + + const TVector* MakeEmptyVector(size_t size) { + return MakeVector(size, [&](size_t) { + return NYql::NUdf::TUnboxedValuePod(); + }); + } + +private: + void CompileFilter() { + const auto edgeActor = Runtime.AllocateEdgeActor(); + Runtime.Send(CompileServiceActorId, edgeActor, Filter->GetCompileRequest().release()); + auto response = Runtime.GrabEdgeEvent(edgeActor, TDuration::Seconds(5)); + + UNIT_ASSERT_C(response, "Failed to get compile response"); + if (!CompileError) { + UNIT_ASSERT_C(response->Get()->ProgramHolder, "Failed to compile program, error: " << response->Get()->Issues.ToOneLineString()); + Filter->OnCompileResponse(std::move(response)); + FilterHandler->OnFilterStarted(); + } else { + CheckError(TStatus::Fail(response->Get()->Status, response->Get()->Issues), CompileError->first, CompileError->second); + } + } + +public: + NActors::TActorId CompileServiceActorId; + TFilterConsumer::TPtr FilterHandler; + IPurecalcFilter::TPtr Filter; + TList> Holders; + + std::optional> CompileError; +}; + +class TFilterSetFixture : public TFiterFixture { +public: + using TBase = TFiterFixture; + + class TFilterSetConsumer : public TFilterConsumer { + public: + using TBase = TFilterConsumer; + using TPtr = TIntrusivePtr; + + public: + TFilterSetConsumer(NActors::TActorId filterId, const TVector& columnIds, const TVector& columns, const TString& whereFilter, TCallback callback, std::optional> compileError) + : TBase(columns, whereFilter, callback, compileError) + { + FilterId = filterId; + ColumnIds = columnIds; + } + + bool IsStarted() const { + return Started; + } + }; + +public: + void SetUp(NUnitTest::TTestContext& ctx) override { + TBase::SetUp(ctx); + + CompileNotifier = Runtime.AllocateEdgeActor(); + FiltersSet = CreateTopicFilters(CompileNotifier, {.CompileServiceId = CompileServiceActorId}, MakeIntrusive()); + } + + void TearDown(NUnitTest::TTestContext& ctx) override { + FilterIds.clear(); + FiltersSet.Reset(); + + TBase::TearDown(ctx); + } + +public: + TStatus MakeFilter(const TVector& columns, const TString& whereFilter, TCallback callback) override { + TVector columnIds; + columnIds.reserve(columns.size()); + for (const auto& column : columns) { + if (const auto it = ColumnIndex.find(column.Name); it != ColumnIndex.end()) { + columnIds.emplace_back(it->second); + } else { + columnIds.emplace_back(ColumnIndex.size()); + ColumnIndex.insert({column.Name, ColumnIndex.size()}); + } + } + FilterIds.emplace_back(FilterIds.size(), 0, 0, 0); + + auto filterSetHandler = MakeIntrusive(FilterIds.back(), columnIds, columns, whereFilter, callback, CompileError); + if (auto status = FiltersSet->AddFilter(filterSetHandler); status.IsFail()) { + return status; + } + + if (!filterSetHandler->IsStarted()) { + // Wait filter compilation + auto response = Runtime.GrabEdgeEvent(CompileNotifier, TDuration::Seconds(5)); + UNIT_ASSERT_C(response, "Compilation is not performed for filter: " << whereFilter); + FiltersSet->OnCompileResponse(std::move(response)); + } + + return TStatus::Success(); + } + + void FilterData(const TVector& columnIndex, const TVector*>& values, ui64 numberRows = 0) { + numberRows = numberRows ? numberRows : values.front()->size(); + FiltersSet->FilterData(columnIndex, TVector(numberRows, std::numeric_limits::max()), values, numberRows); + } + +public: + TVector FilterIds; + std::unordered_map ColumnIndex; + + NActors::TActorId CompileNotifier; + ITopicFilters::TPtr FiltersSet; +}; + +} // anonymous namespace + +Y_UNIT_TEST_SUITE(TestPurecalcFilter) { + Y_UNIT_TEST_F(Simple1, TFiterFixture) { + std::unordered_set result; + CheckSuccess(MakeFilter( + {{"a1", "[DataType; String]"}, {"a2", "[DataType; Uint64]"}, {"a@3", "[OptionalType; [DataType; String]]"}}, + "where a2 > 100", + [&](ui64 offset) { + result.insert(offset); + } + )); + + Push({MakeStringVector({"hello1"}), MakeVector({99}), MakeStringVector({"zapuskaem"}, true)}); + UNIT_ASSERT_VALUES_EQUAL(0, result.size()); + + Push({MakeStringVector({"hello2"}), MakeVector({101}), MakeStringVector({"gusya"}, true)}); + UNIT_ASSERT_VALUES_EQUAL(1, result.size()); + UNIT_ASSERT_VALUES_EQUAL(*result.begin(), 0); + } + + Y_UNIT_TEST_F(Simple2, TFiterFixture) { + std::unordered_set result; + CheckSuccess(MakeFilter( + {{"a2", "[DataType; Uint64]"}, {"a1", "[DataType; String]"}}, + "where a2 > 100", + [&](ui64 offset) { + result.insert(offset); + } + )); + + Push({MakeVector({99}), MakeStringVector({"hello1"})}); + UNIT_ASSERT_VALUES_EQUAL(0, result.size()); + + Push({MakeVector({101}), MakeStringVector({"hello2"})}); + UNIT_ASSERT_VALUES_EQUAL(1, result.size()); + UNIT_ASSERT_VALUES_EQUAL(*result.begin(), 0); + } + + Y_UNIT_TEST_F(ManyValues, TFiterFixture) { + std::unordered_set result; + CheckSuccess(MakeFilter( + {{"a1", "[DataType; String]"}, {"a2", "[DataType; Uint64]"}, {"a3", "[DataType; String]"}}, + "where a2 > 100", + [&](ui64 offset) { + result.insert(offset); + } + )); + + const TString largeString = "abcdefghjkl1234567890+abcdefghjkl1234567890"; + for (ui64 i = 0; i < 5; ++i) { + result.clear(); + Push({MakeStringVector({"hello1", "hello2"}), MakeVector({99, 101}), MakeStringVector({largeString, largeString})}); + UNIT_ASSERT_VALUES_EQUAL_C(1, result.size(), i); + UNIT_ASSERT_VALUES_EQUAL_C(*result.begin(), 1, i); + } + } + + Y_UNIT_TEST_F(NullValues, TFiterFixture) { + std::unordered_set result; + CheckSuccess(MakeFilter( + {{"a1", "[OptionalType; [DataType; Uint64]]"}, {"a2", "[DataType; String]"}}, + "where a1 is null", + [&](ui64 offset) { + result.insert(offset); + } + )); + + Push({MakeEmptyVector(1), MakeStringVector({"str"})}); + UNIT_ASSERT_VALUES_EQUAL(1, result.size()); + UNIT_ASSERT_VALUES_EQUAL(*result.begin(), 0); + } + + Y_UNIT_TEST_F(PartialPush, TFiterFixture) { + std::unordered_set result; + CheckSuccess(MakeFilter( + {{"a1", "[DataType; String]"}, {"a2", "[DataType; Uint64]"}, {"a@3", "[OptionalType; [DataType; String]]"}}, + "where a2 > 50", + [&](ui64 offset) { + result.insert(offset); + } + )); + + Push({MakeStringVector({"hello1", "hello2"}), MakeVector({99, 101}), MakeStringVector({"zapuskaem", "gusya"}, true)}, 1); + UNIT_ASSERT_VALUES_EQUAL(1, result.size()); + UNIT_ASSERT_VALUES_EQUAL(*result.begin(), 0); + } + + Y_UNIT_TEST_F(CompilationValidation, TFiterFixture) { + CompileError = {EStatusId::INTERNAL_ERROR, "Failed to compile purecalc program subissue: {

: Error: Compile issues: generated.sql:2:36: Error: Unexpected token '.' : cannot match to any predicted input... } subissue: {
: Error: Final yql:"}; + MakeFilter( + {{"a1", "[DataType; String]"}}, + "where a2 ... 50", + [&](ui64 offset) {} + ); + } +} + +Y_UNIT_TEST_SUITE(TestFilterSet) { + Y_UNIT_TEST_F(FilterGroup, TFilterSetFixture) { + const TSchemaColumn commonColumn = {"common_col", "[DataType; String]"}; + const TVector whereFilters = { + "where col_0 == \"str1\"", + "where col_1 == \"str2\"", + "" // Empty filter <=> where true + }; + + TVector> results(whereFilters.size()); + for (size_t i = 0; i < results.size(); ++i) { + CheckSuccess(MakeFilter( + {commonColumn, {TStringBuilder() << "col_" << i, "[DataType; String]"}}, + whereFilters[i], + [&, index = i](ui64 offset) { + results[index].push_back(offset); + } + )); + } + + FilterData({0, 1, 2, 3}, { + MakeStringVector({"common_1", "common_2", "common_3"}), + MakeStringVector({"str1", "str2", "str3"}), + MakeStringVector({"str1", "str3", "str2"}), + MakeStringVector({"str2", "str3", "str1"}) + }); + + FiltersSet->RemoveFilter(FilterIds.back()); + + FilterData({0, 3, 1, 2}, { + MakeStringVector({"common_3"}), + MakeStringVector({"str2"}), + MakeStringVector({"str3"}), + MakeStringVector({"str1"}) + }); + + TVector> expectedResults = { + {0, 0}, + {2, 0}, + {0, 1, 2} + }; + for (size_t i = 0; i < results.size(); ++i) { + UNIT_ASSERT_VALUES_EQUAL_C(results[i], expectedResults[i], i); + } + } + + Y_UNIT_TEST_F(DuplicationValidation, TFilterSetFixture) { + CheckSuccess(MakeFilter( + {{"a1", "[DataType; String]"}}, + "where a1 = \"str1\"", + [&](ui64 offset) {} + )); + + CheckError( + FiltersSet->AddFilter(MakeIntrusive(FilterIds.back(), TVector(), TVector(), TString(), [&](ui64 offset) {}, CompileError)), + EStatusId::INTERNAL_ERROR, + "Failed to create new filter, filter with id [0:0:0] already exists" + ); + } + + Y_UNIT_TEST_F(CompilationValidation, TFilterSetFixture) { + CompileError = {EStatusId::INTERNAL_ERROR, "Failed to compile client filter subissue: {
: Error: Failed to compile purecalc program subissue: {
: Error: Compile issues: generated.sql:2:36: Error: Unexpected token '.' : cannot match to any predicted input... } subissue: {
: Error: Final yql:"}; + MakeFilter( + {{"a1", "[DataType; String]"}}, + "where a2 ... 50", + [&](ui64 offset) {} + ); + } +} + +} // namespace NFq::NRowDispatcher::NTests diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/ut/topic_parser_ut.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/topic_parser_ut.cpp new file mode 100644 index 000000000000..0895466e72f9 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/topic_parser_ut.cpp @@ -0,0 +1,477 @@ +#include +#include +#include + +namespace NFq::NRowDispatcher::NTests { + +namespace { + +class TBaseParserFixture : public TBaseFixture { +public: + static constexpr ui64 FIRST_OFFSET = 42; + + using TBase = TBaseFixture; + using TCallback = std::function*> result)>; + + class TParsedDataConsumer : public IParsedDataConsumer { + public: + using TPtr = TIntrusivePtr; + + public: + TParsedDataConsumer(const TBaseParserFixture& self, const TVector& columns, TCallback callback) + : Self(self) + , Columns(columns) + , Callback(callback) + {} + + void ExpectColumnError(ui64 columnId, TStatusCode statusCode, const TString& message) { + UNIT_ASSERT_C(ExpectedErrors.insert({columnId, {statusCode, message}}).second, "Can not add existing column error"); + } + + void ExpectCommonError(TStatusCode statusCode, const TString& message) { + UNIT_ASSERT_C(!ExpectedCommonError, "Can not add existing common error"); + ExpectedCommonError = {statusCode, message}; + } + + public: + const TVector& GetColumns() const override { + return Columns; + } + + void OnParsingError(TStatus status) override { + NumberBatches++; + CurrentOffset++; + if (ExpectedCommonError) { + CheckError(status, ExpectedCommonError->first, ExpectedCommonError->second); + ExpectedCommonError = std::nullopt; + } else { + CheckSuccess(status); + } + } + + void OnParsedData(ui64 numberRows) override { + NumberBatches++; + UNIT_ASSERT_C(!ExpectedCommonError, "Expected common error: " << ExpectedCommonError->second); + + const auto& offsets = Self.Parser->GetOffsets(); + UNIT_ASSERT_VALUES_EQUAL_C(offsets.size(), numberRows, "Unexpected offsets size"); + for (const ui64 offset : offsets) { + UNIT_ASSERT_VALUES_EQUAL_C(offset, CurrentOffset, "Unexpected offset"); + CurrentOffset++; + } + + TVector*> result(Columns.size(), nullptr); + for (ui64 i = 0; i < Columns.size(); ++i) { + if (const auto it = ExpectedErrors.find(i); it != ExpectedErrors.end()) { + CheckError(Self.Parser->GetParsedColumn(i), it->second.first, it->second.second); + ExpectedErrors.erase(i); + } else { + result[i] = CheckSuccess(Self.Parser->GetParsedColumn(i)); + } + } + Callback(numberRows, std::move(result)); + } + + public: + ui64 CurrentOffset = FIRST_OFFSET; + ui64 NumberBatches = 0; + + private: + const TBaseParserFixture& Self; + const TVector Columns; + const TCallback Callback; + + std::optional> ExpectedCommonError; + std::unordered_map> ExpectedErrors; + }; + +public: + void TearDown(NUnitTest::TTestContext& ctx) override { + if (ParserHandler) { + UNIT_ASSERT_VALUES_EQUAL_C(ExpectedBatches, ParserHandler->NumberBatches, "Unexpected number of batches"); + } + Parser.Reset(); + ParserHandler.Reset(); + + TBase::TearDown(ctx); + } + +public: + TStatus MakeParser(TVector columns, TCallback callback) { + ParserHandler = MakeIntrusive(*this, columns, callback); + + auto parserStatus = CreateParser(); + if (parserStatus.IsFail()) { + return parserStatus; + } + + Parser = parserStatus.DetachResult(); + return TStatus::Success(); + } + + TStatus MakeParser(TVector columnNames, TString columnType, TCallback callback) { + TVector columns; + for (const auto& columnName : columnNames) { + columns.push_back({.Name = columnName, .TypeYson = columnType}); + } + return MakeParser(columns, callback); + } + + TStatus MakeParser(TVector columnNames, TString columnType) { + return MakeParser(columnNames, columnType, [](ui64, TVector*>) {}); + } + + TStatus MakeParser(TVector columns) { + return MakeParser(columns, [](ui64, TVector*>) {}); + } + + void PushToParser(ui64 offset, const TString& data) { + ExpectedBatches++; + Parser->ParseMessages({GetMessage(offset, data)}); + } + + void CheckColumnError(const TString& data, ui64 columnId, TStatusCode statusCode, const TString& message) { + ExpectedBatches++; + ParserHandler->ExpectColumnError(columnId, statusCode, message); + Parser->ParseMessages({GetMessage(ParserHandler->CurrentOffset, data)}); + } + + void CheckBatchError(const TString& data, TStatusCode statusCode, const TString& message) { + ExpectedBatches++; + ParserHandler->ExpectCommonError(statusCode, message); + Parser->ParseMessages({GetMessage(ParserHandler->CurrentOffset, data)}); + } + +protected: + virtual TValueStatus CreateParser() = 0; + +public: + TParsedDataConsumer::TPtr ParserHandler; + ITopicParser::TPtr Parser; + ui64 ExpectedBatches = 0; +}; + +class TJsonParserFixture : public TBaseParserFixture { + using TBase = TBaseParserFixture; + +public: + TJsonParserFixture() + : TBase() + , Config({.BatchSize = 1_MB, .LatencyLimit = TDuration::Zero(), .BufferCellCount = 1000}) + {} + +protected: + TValueStatus CreateParser() override { + return CreateJsonParser(ParserHandler, Config, {}); + } + +public: + TJsonParserConfig Config; +}; + +class TRawParserFixture : public TBaseParserFixture { +protected: + TValueStatus CreateParser() override { + return CreateRawParser(ParserHandler, {}); + } +}; + +} // anonymous namespace + +Y_UNIT_TEST_SUITE(TestJsonParser) { + Y_UNIT_TEST_F(Simple1, TJsonParserFixture) { + CheckSuccess(MakeParser({{"a1", "[DataType; String]"}, {"a2", "[OptionalType; [DataType; Uint64]]"}}, [](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(1, numberRows); + UNIT_ASSERT_VALUES_EQUAL(2, result.size()); + UNIT_ASSERT_VALUES_EQUAL("hello1", TString(result[0]->at(0).AsStringRef())); + UNIT_ASSERT_VALUES_EQUAL(101, result[1]->at(0).GetOptionalValue().Get()); + })); + PushToParser(FIRST_OFFSET, R"({"a1": "hello1", "a2": 101, "event": "event1"})"); + } + + Y_UNIT_TEST_F(Simple2, TJsonParserFixture) { + CheckSuccess(MakeParser({"a2", "a1"}, "[DataType; String]", [](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(1, numberRows); + UNIT_ASSERT_VALUES_EQUAL(2, result.size()); + UNIT_ASSERT_VALUES_EQUAL("101", TString(result[0]->at(0).AsStringRef())); + UNIT_ASSERT_VALUES_EQUAL("hello1", TString(result[1]->at(0).AsStringRef())); + })); + PushToParser(FIRST_OFFSET, R"({"a1": "hello1", "a2": "101", "event": "event1"})"); + } + + Y_UNIT_TEST_F(Simple3, TJsonParserFixture) { + CheckSuccess(MakeParser({"a1", "a2"}, "[DataType; String]", [](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(1, numberRows); + UNIT_ASSERT_VALUES_EQUAL(2, result.size()); + UNIT_ASSERT_VALUES_EQUAL("101", TString(result[0]->at(0).AsStringRef())); + UNIT_ASSERT_VALUES_EQUAL("hello1", TString(result[1]->at(0).AsStringRef())); + })); + PushToParser(FIRST_OFFSET,R"({"a2": "hello1", "a1": "101", "event": "event1"})"); + } + + Y_UNIT_TEST_F(Simple4, TJsonParserFixture) { + CheckSuccess(MakeParser({"a2", "a1"}, "[DataType; String]", [](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(1, numberRows); + UNIT_ASSERT_VALUES_EQUAL(2, result.size()); + UNIT_ASSERT_VALUES_EQUAL("hello1", TString(result[0]->at(0).AsStringRef())); + UNIT_ASSERT_VALUES_EQUAL("101", TString(result[1]->at(0).AsStringRef())); + })); + PushToParser(FIRST_OFFSET, R"({"a2": "hello1", "a1": "101", "event": "event1"})"); + } + + Y_UNIT_TEST_F(LargeStrings, TJsonParserFixture) { + ExpectedBatches = 1; + + const TString largeString = "abcdefghjkl1234567890+abcdefghjkl1234567890"; + + CheckSuccess(MakeParser({"col"}, "[DataType; String]", [&](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(2, numberRows); + UNIT_ASSERT_VALUES_EQUAL(1, result.size()); + UNIT_ASSERT_VALUES_EQUAL(largeString, TString(result[0]->at(0).AsStringRef())); + UNIT_ASSERT_VALUES_EQUAL(largeString, TString(result[0]->at(1).AsStringRef())); + })); + + const TString jsonString = TStringBuilder() << "{\"col\": \"" << largeString << "\"}"; + Parser->ParseMessages({ + GetMessage(FIRST_OFFSET, jsonString), + GetMessage(FIRST_OFFSET + 1, jsonString) + }); + } + + Y_UNIT_TEST_F(ManyValues, TJsonParserFixture) { + ExpectedBatches = 1; + + CheckSuccess(MakeParser({"a1", "a2"}, "[DataType; String]", [&](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(3, numberRows); + UNIT_ASSERT_VALUES_EQUAL(2, result.size()); + for (size_t i = 0; i < numberRows; ++i) { + UNIT_ASSERT_VALUES_EQUAL_C("hello1", TString(result[0]->at(i).AsStringRef()), i); + UNIT_ASSERT_VALUES_EQUAL_C("101", TString(result[1]->at(i).AsStringRef()), i); + } + })); + + Parser->ParseMessages({ + GetMessage(FIRST_OFFSET, R"({"a1": "hello1", "a2": "101", "event": "event1"})"), + GetMessage(FIRST_OFFSET + 1, R"({"a1": "hello1", "a2": "101", "event": "event2"})"), + GetMessage(FIRST_OFFSET + 2, R"({"a2": "101", "a1": "hello1", "event": "event3"})") + }); + } + + Y_UNIT_TEST_F(MissingFields, TJsonParserFixture) { + ExpectedBatches = 1; + + CheckSuccess(MakeParser({{"a1", "[OptionalType; [DataType; String]]"}, {"a2", "[OptionalType; [DataType; Uint64]]"}}, [&](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(3, numberRows); + UNIT_ASSERT_VALUES_EQUAL(2, result.size()); + for (size_t i = 0; i < numberRows; ++i) { + if (i == 2) { + UNIT_ASSERT_C(!result[0]->at(i), i); + } else { + NYql::NUdf::TUnboxedValue value = result[0]->at(i).GetOptionalValue(); + UNIT_ASSERT_VALUES_EQUAL_C("hello1", TString(value.AsStringRef()), i); + } + if (i == 1) { + UNIT_ASSERT_C(!result[1]->at(i), i); + } else { + UNIT_ASSERT_VALUES_EQUAL_C(101, result[1]->at(i).GetOptionalValue().Get(), i); + } + } + })); + + Parser->ParseMessages({ + GetMessage(FIRST_OFFSET, R"({"a1": "hello1", "a2": 101 , "event": "event1"})"), + GetMessage(FIRST_OFFSET + 1, R"({"a1": "hello1", "event": "event2"})"), + GetMessage(FIRST_OFFSET + 2, R"({"a2": "101", "a1": null, "event": "event3"})") + }); + } + + Y_UNIT_TEST_F(NestedTypes, TJsonParserFixture) { + ExpectedBatches = 1; + + CheckSuccess(MakeParser({{"nested", "[OptionalType; [DataType; Json]]"}, {"a1", "[DataType; String]"}}, [&](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(4, numberRows); + + UNIT_ASSERT_VALUES_EQUAL(2, result.size()); + UNIT_ASSERT_VALUES_EQUAL("{\"key\": \"value\"}", TString(result[0]->at(0).AsStringRef())); + UNIT_ASSERT_VALUES_EQUAL("hello1", TString(result[1]->at(0).AsStringRef())); + + UNIT_ASSERT_VALUES_EQUAL("[\"key1\", \"key2\"]", TString(result[0]->at(1).AsStringRef())); + UNIT_ASSERT_VALUES_EQUAL("hello2", TString(result[1]->at(1).AsStringRef())); + + UNIT_ASSERT_VALUES_EQUAL("\"some string\"", TString(result[0]->at(2).AsStringRef())); + UNIT_ASSERT_VALUES_EQUAL("hello3", TString(result[1]->at(2).AsStringRef())); + + UNIT_ASSERT_VALUES_EQUAL("123456", TString(result[0]->at(3).AsStringRef())); + UNIT_ASSERT_VALUES_EQUAL("hello4", TString(result[1]->at(3).AsStringRef())); + })); + + Parser->ParseMessages({ + GetMessage(FIRST_OFFSET, R"({"a1": "hello1", "nested": {"key": "value"}})"), + GetMessage(FIRST_OFFSET + 1, R"({"a1": "hello2", "nested": ["key1", "key2"]})"), + GetMessage(FIRST_OFFSET + 2, R"({"a1": "hello3", "nested": "some string"})"), + GetMessage(FIRST_OFFSET + 3, R"({"a1": "hello4", "nested": 123456})") + }); + } + + Y_UNIT_TEST_F(SimpleBooleans, TJsonParserFixture) { + ExpectedBatches = 1; + + CheckSuccess(MakeParser({{"a", "[DataType; Bool]"}}, [&](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(2, numberRows); + + UNIT_ASSERT_VALUES_EQUAL(1, result.size()); + UNIT_ASSERT_VALUES_EQUAL(true, result[0]->at(0).Get()); + UNIT_ASSERT_VALUES_EQUAL(false, result[0]->at(1).Get()); + })); + + Parser->ParseMessages({ + GetMessage(FIRST_OFFSET, R"({"a": true})"), + GetMessage(FIRST_OFFSET + 1, R"({"a": false})") + }); + } + + Y_UNIT_TEST_F(ManyBatches, TJsonParserFixture) { + ExpectedBatches = 2; + Config.BufferCellCount = 1; + + const TString largeString = "abcdefghjkl1234567890+abcdefghjkl1234567890"; + CheckSuccess(MakeParser({"col"}, "[DataType; String]", [&](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(1, numberRows); + UNIT_ASSERT_VALUES_EQUAL(1, result.size()); + UNIT_ASSERT_VALUES_EQUAL(largeString, TString(result[0]->at(0).AsStringRef())); + })); + + const TString jsonString = TStringBuilder() << "{\"col\": \"" << largeString << "\"}"; + Parser->ParseMessages({ + GetMessage(FIRST_OFFSET, jsonString), + GetMessage(FIRST_OFFSET + 1, jsonString) + }); + } + + Y_UNIT_TEST_F(LittleBatches, TJsonParserFixture) { + ExpectedBatches = 2; + Config.BatchSize = 10; + + const TString largeString = "abcdefghjkl1234567890+abcdefghjkl1234567890"; + CheckSuccess(MakeParser({"col"}, "[DataType; String]", [&](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(1, numberRows); + UNIT_ASSERT_VALUES_EQUAL(1, result.size()); + UNIT_ASSERT_VALUES_EQUAL(largeString, TString(result[0]->at(0).AsStringRef())); + })); + + const TString jsonString = TStringBuilder() << "{\"col\": \"" << largeString << "\"}"; + Parser->ParseMessages({ + GetMessage(FIRST_OFFSET, jsonString), + GetMessage(FIRST_OFFSET + 1, jsonString) + }); + } + + Y_UNIT_TEST_F(MissingFieldsValidation, TJsonParserFixture) { + CheckSuccess(MakeParser({{"a1", "[DataType; String]"}, {"a2", "[DataType; Uint64]"}})); + CheckColumnError(R"({"a2": 105, "event": "event1"})", 0, EStatusId::PRECONDITION_FAILED, TStringBuilder() << "Failed to parse json messages, found 1 missing values from offset " << FIRST_OFFSET << " in non optional column 'a1' with type [DataType; String]"); + CheckColumnError(R"({"a1": "hello1", "a2": null, "event": "event1"})", 1, EStatusId::PRECONDITION_FAILED, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET + 1 << ", got parsing error for column 'a2' with type [DataType; Uint64] subissue: {
: Error: Found unexpected null value, expected non optional data type Uint64 }"); + } + + Y_UNIT_TEST_F(TypeKindsValidation, TJsonParserFixture) { + CheckError( + MakeParser({{"a1", "[[BAD TYPE]]"}}), + EStatusId::INTERNAL_ERROR, + "Failed to parse column 'a1' type [[BAD TYPE]] subissue: {
: Error: Failed to parse type from yson: Failed to parse scheme from YSON:" + ); + CheckError( + MakeParser({{"a2", "[OptionalType; [DataType; String]]"}, {"a1", "[ListType; [DataType; String]]"}}), + EStatusId::UNSUPPORTED, + "Failed to create parser for column 'a1' with type [ListType; [DataType; String]] subissue: {
: Error: Unsupported type kind: List }" + ); + } + + Y_UNIT_TEST_F(NumbersValidation, TJsonParserFixture) { + CheckSuccess(MakeParser({{"a1", "[OptionalType; [DataType; String]]"}, {"a2", "[DataType; Uint8]"}, {"a3", "[OptionalType; [DataType; Float]]"}})); + CheckColumnError(R"({"a1": 456, "a2": 42})", 0, EStatusId::PRECONDITION_FAILED, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET << ", got parsing error for column 'a1' with type [OptionalType; [DataType; String]] subissue: {
: Error: Failed to parse data type String from json number (raw: '456') subissue: {
: Error: Number value is not expected for data type String } }"); + CheckColumnError(R"({"a1": "456", "a2": -42})", 1, EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET + 1 << ", got parsing error for column 'a2' with type [DataType; Uint8] subissue: {
: Error: Failed to parse data type Uint8 from json number (raw: '-42') subissue: {
: Error: Failed to extract json integer number, error: INCORRECT_TYPE: The JSON element does not have the requested type. } }"); + CheckColumnError(R"({"a1": "str", "a2": 99999})", 1, EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET + 2 << ", got parsing error for column 'a2' with type [DataType; Uint8] subissue: {
: Error: Failed to parse data type Uint8 from json number (raw: '99999') subissue: {
: Error: Number is out of range [0, 255] } }"); + CheckColumnError(R"({"a1": "456", "a2": 42, "a3": 1.11.1})", 2, EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET + 3 << ", got parsing error for column 'a3' with type [OptionalType; [DataType; Float]] subissue: {
: Error: Failed to parse data type Float from json number (raw: '1.11.1') subissue: {
: Error: Failed to extract json float number, error: NUMBER_ERROR: Problem while parsing a number } }"); + } + + Y_UNIT_TEST_F(StringsValidation, TJsonParserFixture) { + CheckSuccess(MakeParser({{"a1", "[OptionalType; [DataType; Uint8]]"}})); + CheckColumnError(R"({"a1": "-456"})", 0, EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET << ", got parsing error for column 'a1' with type [OptionalType; [DataType; Uint8]] subissue: {
: Error: Failed to parse data type Uint8 from json string: '-456' }"); + } + + Y_UNIT_TEST_F(NestedJsonValidation, TJsonParserFixture) { + CheckSuccess(MakeParser({{"a1", "[OptionalType; [DataType; Json]]"}, {"a2", "[OptionalType; [DataType; String]]"}})); + CheckColumnError(R"({"a1": {"key": "value"}, "a2": {"key2": "value2"}})", 1, EStatusId::PRECONDITION_FAILED, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET << ", got parsing error for column 'a2' with type [OptionalType; [DataType; String]] subissue: {
: Error: Found unexpected nested value (raw: '{\"key2\": \"value2\"}'), expected data type String, please use Json type for nested values }"); + CheckColumnError(R"({"a1": {"key": "value", "nested": {"a": "b", "c":}}, "a2": "str"})", 0, EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET + 1 << ", got parsing error for column 'a1' with type [OptionalType; [DataType; Json]] subissue: {
: Error: Found bad json value: '{\"key\": \"value\", \"nested\": {\"a\": \"b\", \"c\":}}' }"); + CheckColumnError(R"({"a1": {"key" "value"}, "a2": "str"})", 0, EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET + 2 << ", got parsing error for column 'a1' with type [OptionalType; [DataType; Json]] subissue: {
: Error: Failed to extract json value, current token: '{', error: TAPE_ERROR: The JSON document has an improper structure: missing or superfluous commas, braces, missing keys, etc. }"); + } + + Y_UNIT_TEST_F(BoolsValidation, TJsonParserFixture) { + CheckSuccess(MakeParser({{"a1", "[OptionalType; [DataType; String]]"}, {"a2", "[DataType; Bool]"}})); + CheckColumnError(R"({"a1": true, "a2": false})", 0, EStatusId::PRECONDITION_FAILED, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET << ", got parsing error for column 'a1' with type [OptionalType; [DataType; String]] subissue: {
: Error: Found unexpected bool value, expected data type String }"); + CheckColumnError(R"({"a1": "true", "a2": falce})", 1, EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET + 1 << ", got parsing error for column 'a2' with type [DataType; Bool] subissue: {
: Error: Failed to extract json bool, current token: 'falce', error: INCORRECT_TYPE: The JSON element does not have the requested type. }"); + } + + Y_UNIT_TEST_F(JsonStructureValidation, TJsonParserFixture) { + CheckSuccess(MakeParser({{"a1", "[OptionalType; [DataType; String]]"}})); + CheckColumnError(R"({"a1": Yelse})", 0, EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse json string at offset " << FIRST_OFFSET << ", got parsing error for column 'a1' with type [OptionalType; [DataType; String]] subissue: {
: Error: Failed to determine json value type, current token: 'Yelse', error: TAPE_ERROR: The JSON document has an improper structure: missing or superfluous commas, braces, missing keys, etc. }"); + CheckBatchError(R"({"a1": "st""r"})", EStatusId::BAD_REQUEST, TStringBuilder() << "Failed to parse json message for offset " << FIRST_OFFSET + 1 << ", json item was corrupted: TAPE_ERROR: The JSON document has an improper structure: missing or superfluous commas, braces, missing keys, etc. Current data batch: {\"a1\": \"st\"\"r\"}"); + CheckBatchError(R"({"a1": "x"} {"a1": "y"})", EStatusId::INTERNAL_ERROR, TStringBuilder() << "Failed to parse json messages, expected 1 json rows from offset " << FIRST_OFFSET + 2 << " but got 2 (expected one json row for each offset from topic API in json each row format, maybe initial data was corrupted or messages is not in json format), current data batch: {\"a1\": \"x\"} {\"a1\": \"y\"}"); + CheckBatchError(R"({)", EStatusId::INTERNAL_ERROR, TStringBuilder() << "Failed to parse json messages, expected 1 json rows from offset " << FIRST_OFFSET + 3 << " but got 0 (expected one json row for each offset from topic API in json each row format, maybe initial data was corrupted or messages is not in json format), current data batch: {"); + } +} + +Y_UNIT_TEST_SUITE(TestRawParser) { + Y_UNIT_TEST_F(Simple, TRawParserFixture) { + CheckSuccess(MakeParser({{"data", "[OptionalType; [DataType; String]]"}}, [](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(1, numberRows); + UNIT_ASSERT_VALUES_EQUAL(1, result.size()); + + NYql::NUdf::TUnboxedValue value = result[0]->at(0).GetOptionalValue(); + UNIT_ASSERT_VALUES_EQUAL(R"({"a1": "hello1__large_str", "a2": 101, "event": "event1"})", TString(value.AsStringRef())); + })); + PushToParser(FIRST_OFFSET, R"({"a1": "hello1__large_str", "a2": 101, "event": "event1"})"); + } + + Y_UNIT_TEST_F(ManyValues, TRawParserFixture) { + TVector data = { + R"({"a1": "hello1", "a2": "101", "event": "event1"})", + R"({"a1": "hello1", "a2": "101", "event": "event2"})", + R"({"a2": "101", "a1": "hello1", "event": "event3"})" + }; + ExpectedBatches = data.size(); + + int i = 0; + CheckSuccess(MakeParser({"a1"}, "[DataType; String]", [&](ui64 numberRows, TVector*> result) { + UNIT_ASSERT_VALUES_EQUAL(1, numberRows); + UNIT_ASSERT_VALUES_EQUAL(1, result.size()); + UNIT_ASSERT_VALUES_EQUAL(data[i], TString(result[0]->at(0).AsStringRef())); + i++; + })); + + Parser->ParseMessages({ + GetMessage(FIRST_OFFSET, data[0]), + GetMessage(FIRST_OFFSET + 1, data[1]), + GetMessage(FIRST_OFFSET + 2, data[2]) + }); + } + + Y_UNIT_TEST_F(TypeKindsValidation, TRawParserFixture) { + CheckError( + MakeParser({{"a1", "[DataType; String]"}, {"a2", "[DataType; String]"}}), + EStatusId::INTERNAL_ERROR, + "Expected only one column for raw format, but got 2" + ); + CheckError( + MakeParser({{"a1", "[[BAD TYPE]]"}}), + EStatusId::INTERNAL_ERROR, + "Failed to create raw parser for column 'a1' : [[BAD TYPE]] subissue: {
: Error: Failed to parse type from yson: Failed to parse scheme from YSON:" + ); + CheckError( + MakeParser({{"a1", "[ListType; [DataType; String]]"}}), + EStatusId::UNSUPPORTED, + "Failed to create raw parser for column 'a1' : [ListType; [DataType; String]] subissue: {
: Error: Unsupported type kind for raw format: List }" + ); + } +} + +} // namespace NFq::NRowDispatcher::NTests diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/ut/ya.make b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/ya.make new file mode 100644 index 000000000000..54b30a2c7d48 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/ya.make @@ -0,0 +1,22 @@ +UNITTEST_FOR(ydb/core/fq/libs/row_dispatcher/format_handler) + +SRCS( + format_handler_ut.cpp + topic_filter_ut.cpp + topic_parser_ut.cpp +) + +PEERDIR( + ydb/core/fq/libs/row_dispatcher/format_handler + ydb/core/fq/libs/row_dispatcher/format_handler/filters + ydb/core/fq/libs/row_dispatcher/format_handler/parsers + ydb/core/fq/libs/row_dispatcher/format_handler/ut/common + + yql/essentials/sql/pg_dummy +) + +SIZE(MEDIUM) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/ya.make b/ydb/core/fq/libs/row_dispatcher/format_handler/ya.make new file mode 100644 index 000000000000..cb14937b89c2 --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/format_handler/ya.make @@ -0,0 +1,26 @@ +LIBRARY() + +SRCS( + format_handler.cpp +) + +PEERDIR( + ydb/core/fq/libs/actors/logging + ydb/core/fq/libs/row_dispatcher/events + ydb/core/fq/libs/row_dispatcher/format_handler/common + ydb/core/fq/libs/row_dispatcher/format_handler/filters + ydb/core/fq/libs/row_dispatcher/format_handler/parsers + + ydb/library/actors/core + ydb/library/actors/util + + ydb/library/yql/dq/common +) + +YQL_LAST_ABI_VERSION() + +END() + +RECURSE_FOR_TESTS( + ut +) diff --git a/ydb/core/fq/libs/row_dispatcher/json_filter.cpp b/ydb/core/fq/libs/row_dispatcher/json_filter.cpp deleted file mode 100644 index 0938e74da67c..000000000000 --- a/ydb/core/fq/libs/row_dispatcher/json_filter.cpp +++ /dev/null @@ -1,392 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include - -namespace { - -using TCallback = NFq::TJsonFilter::TCallback; -const char* OffsetFieldName = "_offset"; -TString LogPrefix = "JsonFilter: "; - -NYT::TNode CreateTypeNode(const TString& fieldType) { - return NYT::TNode::CreateList() - .Add("DataType") - .Add(fieldType); -} - -NYT::TNode CreateOptionalTypeNode(const TString& fieldType) { - return NYT::TNode::CreateList() - .Add("OptionalType") - .Add(CreateTypeNode(fieldType)); -} - -void AddField(NYT::TNode& node, const TString& fieldName, const TString& fieldType) { - node.Add( - NYT::TNode::CreateList() - .Add(fieldName) - .Add(CreateTypeNode(fieldType)) - ); -} - -void AddTypedField(NYT::TNode& node, const TString& fieldName, const TString& fieldTypeYson) { - NYT::TNode parsedType; - Y_ENSURE(NYql::NCommon::ParseYson(parsedType, fieldTypeYson, Cerr), "Invalid field type"); - - // TODO: remove this when the re-parsing is removed from pq read actor - if (parsedType == CreateTypeNode("Json")) { - parsedType = CreateTypeNode("String"); - } else if (parsedType == CreateOptionalTypeNode("Json")) { - parsedType = CreateOptionalTypeNode("String"); - } - - node.Add( - NYT::TNode::CreateList() - .Add(fieldName) - .Add(parsedType) - ); -} - -NYT::TNode MakeInputSchema(const TVector& columns, const TVector& types) { - auto structMembers = NYT::TNode::CreateList(); - AddField(structMembers, OffsetFieldName, "Uint64"); - for (size_t i = 0; i < columns.size(); ++i) { - AddTypedField(structMembers, columns[i], types[i]); - } - return NYT::TNode::CreateList().Add("StructType").Add(std::move(structMembers)); -} - -NYT::TNode MakeOutputSchema() { - auto structMembers = NYT::TNode::CreateList(); - AddField(structMembers, OffsetFieldName, "Uint64"); - AddField(structMembers, "data", "String"); - return NYT::TNode::CreateList().Add("StructType").Add(std::move(structMembers)); -} - -struct TInputType { - const TVector& Offsets; - const TVector*>& Values; - const ui64 RowsOffset; // offset of first value - const ui64 NumberRows; - - ui64 GetOffset(ui64 rowId) const { - return Offsets[rowId + RowsOffset]; - } -}; - -class TFilterInputSpec : public NYql::NPureCalc::TInputSpecBase { -public: - TFilterInputSpec(const NYT::TNode& schema) - : Schemas({schema}) { - } - - const TVector& GetSchemas() const override { - return Schemas; - } - -private: - TVector Schemas; -}; - -class TFilterInputConsumer : public NYql::NPureCalc::IConsumer { -public: - TFilterInputConsumer( - const TFilterInputSpec& spec, - NYql::NPureCalc::TWorkerHolder worker) - : Worker(std::move(worker)) { - const NKikimr::NMiniKQL::TStructType* structType = Worker->GetInputType(); - const auto count = structType->GetMembersCount(); - - THashMap schemaPositions; - for (ui32 i = 0; i < count; ++i) { - const auto name = structType->GetMemberName(i); - if (name == OffsetFieldName) { - OffsetPosition = i; - continue; - } - schemaPositions[name] = i; - } - - const NYT::TNode& schema = spec.GetSchemas()[0]; - const auto& fields = schema[1]; - Y_ENSURE(count == fields.Size()); - Y_ENSURE(fields.IsList()); - for (size_t i = 0; i < fields.Size(); ++i) { - auto name = fields[i][0].AsString(); - if (name == OffsetFieldName) { - continue; - } - FieldsPositions.push_back(schemaPositions[name]); - } - } - - ~TFilterInputConsumer() override { - with_lock(Worker->GetScopedAlloc()) { - Cache.Clear(); - } - } - - void OnObject(TInputType input) override { - Y_ENSURE(FieldsPositions.size() == input.Values.size()); - - NKikimr::NMiniKQL::TThrowingBindTerminator bind; - with_lock (Worker->GetScopedAlloc()) { - Y_DEFER { - // Clear cache after each object because - // values allocated on another allocator and should be released - Cache.Clear(); - Worker->GetGraph().Invalidate(); - }; - - auto& holderFactory = Worker->GetGraph().GetHolderFactory(); - - // TODO: use blocks here - for (size_t rowId = 0; rowId < input.NumberRows; ++rowId) { - NYql::NUdf::TUnboxedValue* items = nullptr; - - NYql::NUdf::TUnboxedValue result = Cache.NewArray( - holderFactory, - static_cast(input.Values.size() + 1), - items); - - items[OffsetPosition] = NYql::NUdf::TUnboxedValuePod(input.GetOffset(rowId)); - - size_t fieldId = 0; - for (const auto column : input.Values) { - items[FieldsPositions[fieldId++]] = column->at(rowId); - } - - Worker->Push(std::move(result)); - } - } - } - - void OnFinish() override { - NKikimr::NMiniKQL::TBindTerminator bind(Worker->GetGraph().GetTerminator()); - with_lock(Worker->GetScopedAlloc()) { - Worker->OnFinish(); - } - } - -private: - NYql::NPureCalc::TWorkerHolder Worker; - NKikimr::NMiniKQL::TPlainContainerCache Cache; - size_t OffsetPosition = 0; - TVector FieldsPositions; -}; - -class TFilterOutputConsumer: public NYql::NPureCalc::IConsumer> { -public: - TFilterOutputConsumer(TCallback callback) - : Callback(callback) { - } - - void OnObject(std::pair value) override { - Callback(value.first, value.second); - } - - void OnFinish() override { - Y_UNREACHABLE(); - } -private: - TCallback Callback; -}; - -class TFilterOutputSpec: public NYql::NPureCalc::TOutputSpecBase { -public: - explicit TFilterOutputSpec(const NYT::TNode& schema) - : Schema(schema) - {} - -public: - const NYT::TNode& GetSchema() const override { - return Schema; - } - -private: - NYT::TNode Schema; -}; - -class TFilterPushRelayImpl: public NYql::NPureCalc::IConsumer { -public: - TFilterPushRelayImpl(const TFilterOutputSpec& /*outputSpec*/, NYql::NPureCalc::IPushStreamWorker* worker, THolder>> underlying) - : Underlying(std::move(underlying)) - , Worker(worker) - {} -public: - void OnObject(const NYql::NUdf::TUnboxedValue* value) override { - auto unguard = Unguard(Worker->GetScopedAlloc()); - Y_ENSURE(value->GetListLength() == 2); - ui64 offset = value->GetElement(0).Get(); - const auto& cell = value->GetElement(1); - Y_ENSURE(cell); - TString str(cell.AsStringRef()); - Underlying->OnObject(std::make_pair(offset, str)); - } - - void OnFinish() override { - auto unguard = Unguard(Worker->GetScopedAlloc()); - Underlying->OnFinish(); - } - -private: - THolder>> Underlying; - NYql::NPureCalc::IWorker* Worker; -}; - -} - -template <> -struct NYql::NPureCalc::TInputSpecTraits { - static constexpr bool IsPartial = false; - static constexpr bool SupportPushStreamMode = true; - - using TConsumerType = THolder>; - - static TConsumerType MakeConsumer( - const TFilterInputSpec& spec, - NYql::NPureCalc::TWorkerHolder worker) - { - return MakeHolder(spec, std::move(worker)); - } -}; - -template <> -struct NYql::NPureCalc::TOutputSpecTraits { - static const constexpr bool IsPartial = false; - static const constexpr bool SupportPushStreamMode = true; - - static void SetConsumerToWorker(const TFilterOutputSpec& outputSpec, NYql::NPureCalc::IPushStreamWorker* worker, THolder>> consumer) { - worker->SetConsumer(MakeHolder(outputSpec, worker, std::move(consumer))); - } -}; - -namespace NFq { - -class TProgramHolder : public IProgramHolder { -public: - using TPtr = TIntrusivePtr; - -public: - TProgramHolder(const TVector& columns, const TVector& types, const TString& sql, TCallback callback) - : Columns(columns) - , Types(types) - , Sql(sql) - , Callback(callback) - { - Y_ENSURE(columns.size() == types.size(), "Number of columns and types should by equal"); - } - - NYql::NPureCalc::IConsumer& GetConsumer() { - Y_ENSURE(InputConsumer, "Program is not compiled"); - return *InputConsumer; - } - -public: - void CreateProgram(NYql::NPureCalc::IProgramFactoryPtr programFactory) override { - Program = programFactory->MakePushStreamProgram( - TFilterInputSpec(MakeInputSchema(Columns, Types)), - TFilterOutputSpec(MakeOutputSchema()), - Sql, - NYql::NPureCalc::ETranslationMode::SQL - ); - InputConsumer = Program->Apply(MakeHolder(Callback)); - } - -private: - const TVector Columns; - const TVector Types; - const TString Sql; - const TCallback Callback; - - THolder> Program; - THolder> InputConsumer; -}; - -class TJsonFilter::TImpl { -public: - TImpl(const TVector& columns, const TVector& types, const TString& whereFilter, TCallback callback, const TPurecalcCompileSettings& purecalcSettings) - : PurecalcSettings(purecalcSettings) - , Sql(GenerateSql(whereFilter)) - , ProgramHolder(MakeIntrusive(columns, types, Sql, callback)) - {} - - void Push(const TVector& offsets, const TVector*>& values, ui64 rowsOffset, ui64 numberRows) { - Y_ENSURE(ProgramHolder, "Program is not compiled"); - Y_ENSURE(values, "Expected non empty schema"); - ProgramHolder->GetConsumer().OnObject({.Offsets = offsets, .Values = values, .RowsOffset = rowsOffset, .NumberRows = numberRows}); - } - - TString GetSql() const { - return Sql; - } - - std::unique_ptr GetCompileRequest() { - Y_ENSURE(ProgramHolder, "Can not create compile request twice"); - auto result = std::make_unique(std::move(ProgramHolder), PurecalcSettings); - ProgramHolder = nullptr; - return result; - } - - void OnCompileResponse(TEvRowDispatcher::TEvPurecalcCompileResponse::TPtr ev) { - Y_ENSURE(!ProgramHolder, "Can not handle compile response twice"); - - auto result = static_cast(ev->Get()->ProgramHolder.Release()); - Y_ENSURE(result, "Unexpected compile response"); - - ProgramHolder = TIntrusivePtr(result); - } - -private: - TString GenerateSql(const TString& whereFilter) { - TStringStream str; - str << "PRAGMA config.flags(\"LLVM\", \"" << (PurecalcSettings.EnabledLLVM ? "ON" : "OFF") << "\");\n"; - str << "$filtered = SELECT * FROM Input " << whereFilter << ";\n"; - - str << "SELECT " << OffsetFieldName << ", Unwrap(Json::SerializeJson(Yson::From(RemoveMembers(TableRow(), [\"" << OffsetFieldName; - str << "\"])))) as data FROM $filtered"; - LOG_ROW_DISPATCHER_DEBUG("Generated sql: " << str.Str()); - return str.Str(); - } - -private: - const TPurecalcCompileSettings PurecalcSettings; - const TString Sql; - TProgramHolder::TPtr ProgramHolder; -}; - -TJsonFilter::TJsonFilter(const TVector& columns, const TVector& types, const TString& whereFilter, TCallback callback, const TPurecalcCompileSettings& purecalcSettings) - : Impl(std::make_unique(columns, types, whereFilter, callback, purecalcSettings)) -{} - -TJsonFilter::~TJsonFilter() { -} - -void TJsonFilter::Push(const TVector& offsets, const TVector*>& values, ui64 rowsOffset, ui64 numberRows) { - Impl->Push(offsets, values, rowsOffset, numberRows); -} - -TString TJsonFilter::GetSql() { - return Impl->GetSql(); -} - -std::unique_ptr TJsonFilter::GetCompileRequest() { - return Impl->GetCompileRequest(); -} - -void TJsonFilter::OnCompileResponse(TEvRowDispatcher::TEvPurecalcCompileResponse::TPtr ev) { - Impl->OnCompileResponse(std::move(ev)); -} - -std::unique_ptr NewJsonFilter(const TVector& columns, const TVector& types, const TString& whereFilter, TCallback callback, const TPurecalcCompileSettings& purecalcSettings) { - return std::unique_ptr(new TJsonFilter(columns, types, whereFilter, callback, purecalcSettings)); -} - -} // namespace NFq diff --git a/ydb/core/fq/libs/row_dispatcher/json_filter.h b/ydb/core/fq/libs/row_dispatcher/json_filter.h deleted file mode 100644 index dc3d51f5ecc8..000000000000 --- a/ydb/core/fq/libs/row_dispatcher/json_filter.h +++ /dev/null @@ -1,41 +0,0 @@ -#pragma once - -#include - -#include - -namespace NFq { - -class TJsonFilter { -public: - using TCallback = std::function; - -public: - TJsonFilter( - const TVector& columns, - const TVector& types, - const TString& whereFilter, - TCallback callback, - const TPurecalcCompileSettings& purecalcSettings); - - ~TJsonFilter(); - - void Push(const TVector& offsets, const TVector*>& values, ui64 rowsOffset, ui64 numberRows); - TString GetSql(); - - std::unique_ptr GetCompileRequest(); // Should be called exactly once - void OnCompileResponse(TEvRowDispatcher::TEvPurecalcCompileResponse::TPtr ev); - -private: - class TImpl; - const std::unique_ptr Impl; -}; - -std::unique_ptr NewJsonFilter( - const TVector& columns, - const TVector& types, - const TString& whereFilter, - TJsonFilter::TCallback callback, - const TPurecalcCompileSettings& purecalcSettings); - -} // namespace NFq diff --git a/ydb/core/fq/libs/row_dispatcher/json_parser.cpp b/ydb/core/fq/libs/row_dispatcher/json_parser.cpp deleted file mode 100644 index f5064ca0f07d..000000000000 --- a/ydb/core/fq/libs/row_dispatcher/json_parser.cpp +++ /dev/null @@ -1,486 +0,0 @@ -#include "json_parser.h" - -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -#include - -namespace { - -TString LogPrefix = "JsonParser: "; - -constexpr ui64 DEFAULT_BATCH_SIZE = 1_MB; -constexpr ui64 DEFAULT_BUFFER_CELL_COUNT = 1000000; - -struct TJsonParserBuffer { - size_t NumberValues = 0; - bool Finished = false; - TInstant CreationStartTime = TInstant::Now(); - TVector Offsets = {}; - - bool IsReady() const { - return !Finished && NumberValues > 0; - } - - size_t GetSize() const { - return Values.size(); - } - - void Reserve(size_t size, size_t numberValues) { - Values.reserve(size + simdjson::SIMDJSON_PADDING); - Offsets.reserve(numberValues); - } - - void AddMessage(const NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage& message) { - Y_ENSURE(!Finished, "Cannot add messages into finished buffer"); - NumberValues++; - Values << message.GetData(); - Offsets.emplace_back(message.GetOffset()); - } - - std::pair Finish() { - Y_ENSURE(!Finished, "Cannot finish buffer twice"); - Finished = true; - Values << TString(simdjson::SIMDJSON_PADDING, ' '); - return {Values.data(), Values.size()}; - } - - void Clear() { - Y_ENSURE(Finished, "Cannot clear not finished buffer"); - NumberValues = 0; - Finished = false; - CreationStartTime = TInstant::Now(); - Values.clear(); - Offsets.clear(); - } - -private: - TStringBuilder Values = {}; -}; - -class TColumnParser { - using TParser = std::function; - -public: - const std::string Name; - const TString TypeYson; - const NKikimr::NMiniKQL::TType* TypeMkql; - const bool IsOptional = false; - TVector ParsedRows; - -public: - TColumnParser(const TString& name, const TString& typeYson, ui64 maxNumberRows, NKikimr::NMiniKQL::TProgramBuilder& programBuilder) - : Name(name) - , TypeYson(typeYson) - , TypeMkql(NYql::NCommon::ParseTypeFromYson(TStringBuf(typeYson), programBuilder, Cerr)) - , IsOptional(TypeMkql->IsOptional()) - { - ParsedRows.reserve(maxNumberRows); - try { - Parser = CreateParser(TypeMkql); - } catch (...) { - throw NFq::TJsonParserError(Name) << "Failed to create parser for column '" << Name << "' with type " << TypeYson << ", description: " << CurrentExceptionMessage(); - } - } - - void ParseJsonValue(ui64 rowId, simdjson::builtin::ondemand::value jsonValue, NYql::NUdf::TUnboxedValue& resultValue) { - ParsedRows.emplace_back(rowId); - Parser(jsonValue, resultValue); - } - - void ValidateNumberValues(size_t expectedNumberValues, ui64 firstOffset) const { - if (Y_UNLIKELY(!IsOptional && ParsedRows.size() < expectedNumberValues)) { - throw NFq::TJsonParserError(Name) << "Failed to parse json messages, found " << expectedNumberValues - ParsedRows.size() << " missing values from offset " << firstOffset << " in non optional column '" << Name << "' with type " << TypeYson; - } - } - -private: - TParser CreateParser(const NKikimr::NMiniKQL::TType* type, bool optional = false) const { - switch (type->GetKind()) { - case NKikimr::NMiniKQL::TTypeBase::EKind::Data: { - const auto* dataType = AS_TYPE(NKikimr::NMiniKQL::TDataType, type); - if (const auto dataSlot = dataType->GetDataSlot()) { - return GetJsonValueParser(*dataSlot, optional); - } - throw NFq::TJsonParserError() << "unsupported data type with id " << dataType->GetSchemeType(); - } - - case NKikimr::NMiniKQL::TTypeBase::EKind::Optional: { - return AddOptional(CreateParser(AS_TYPE(NKikimr::NMiniKQL::TOptionalType, type)->GetItemType(), true)); - } - - default: { - throw NFq::TJsonParserError() << "unsupported type kind " << type->GetKindAsStr(); - } - } - } - - static TParser AddOptional(TParser parser) { - return [parser](simdjson::builtin::ondemand::value jsonValue, NYql::NUdf::TUnboxedValue& resultValue) { - parser(std::move(jsonValue), resultValue); - if (resultValue) { - resultValue = resultValue.MakeOptional(); - } - }; - } - - static TParser GetJsonValueParser(NYql::NUdf::EDataSlot dataSlot, bool optional) { - if (dataSlot == NYql::NUdf::EDataSlot::Json) { - return GetJsonValueExtractor(); - } - - const auto& typeInfo = NYql::NUdf::GetDataTypeInfo(dataSlot); - return [dataSlot, optional, &typeInfo](simdjson::builtin::ondemand::value jsonValue, NYql::NUdf::TUnboxedValue& resultValue) { - switch (jsonValue.type()) { - case simdjson::builtin::ondemand::json_type::number: { - try { - switch (dataSlot) { - case NYql::NUdf::EDataSlot::Int8: - resultValue = ParseJsonNumber(jsonValue.get_int64().value()); - break; - case NYql::NUdf::EDataSlot::Int16: - resultValue = ParseJsonNumber(jsonValue.get_int64().value()); - break; - case NYql::NUdf::EDataSlot::Int32: - resultValue = ParseJsonNumber(jsonValue.get_int64().value()); - break; - case NYql::NUdf::EDataSlot::Int64: - resultValue = ParseJsonNumber(jsonValue.get_int64().value()); - break; - - case NYql::NUdf::EDataSlot::Uint8: - resultValue = ParseJsonNumber(jsonValue.get_uint64().value()); - break; - case NYql::NUdf::EDataSlot::Uint16: - resultValue = ParseJsonNumber(jsonValue.get_uint64().value()); - break; - case NYql::NUdf::EDataSlot::Uint32: - resultValue = ParseJsonNumber(jsonValue.get_uint64().value()); - break; - case NYql::NUdf::EDataSlot::Uint64: - resultValue = ParseJsonNumber(jsonValue.get_uint64().value()); - break; - - case NYql::NUdf::EDataSlot::Double: - resultValue = NYql::NUdf::TUnboxedValuePod(jsonValue.get_double().value()); - break; - case NYql::NUdf::EDataSlot::Float: - resultValue = NYql::NUdf::TUnboxedValuePod(static_cast(jsonValue.get_double().value())); - break; - - default: - throw NFq::TJsonParserError() << "number value is not expected for data type " << typeInfo.Name; - } - } catch (...) { - throw NFq::TJsonParserError() << "failed to parse data type " << typeInfo.Name << " from json number (raw: '" << TruncateString(jsonValue.raw_json_token()) << "'), error: " << CurrentExceptionMessage(); - } - break; - } - - case simdjson::builtin::ondemand::json_type::string: { - const auto rawString = jsonValue.get_string().value(); - resultValue = NKikimr::NMiniKQL::ValueFromString(dataSlot, rawString); - if (Y_UNLIKELY(!resultValue)) { - throw NFq::TJsonParserError() << "failed to parse data type " << typeInfo.Name << " from json string: '" << TruncateString(rawString) << "'"; - } - LockObject(resultValue); - break; - } - - case simdjson::builtin::ondemand::json_type::array: - case simdjson::builtin::ondemand::json_type::object: { - throw NFq::TJsonParserError() << "found unexpected nested value (raw: '" << TruncateString(jsonValue.raw_json().value()) << "'), expected data type " < - static NYql::NUdf::TUnboxedValuePod ParseJsonNumber(TJsonNumber number) { - if (number < std::numeric_limits::min() || std::numeric_limits::max() < number) { - throw NFq::TJsonParserError() << "number is out of range"; - } - return NYql::NUdf::TUnboxedValuePod(static_cast(number)); - } - - static void LockObject(NYql::NUdf::TUnboxedValue& value) { - // All UnboxedValue's with type Boxed or String should be locked - // because after parsing they will be used under another MKQL allocator in purecalc filters - - const i32 numberRefs = value.LockRef(); - - // -1 - value is embbeded or empty, otherwise value should have exactly one ref - Y_ENSURE(numberRefs == -1 || numberRefs == 1); - } - - static TString TruncateString(std::string_view rawString, size_t maxSize = 1_KB) { - if (rawString.size() <= maxSize) { - return TString(rawString); - } - return TStringBuilder() << rawString.substr(0, maxSize) << " truncated..."; - } - -private: - TParser Parser; -}; - -} // anonymous namespace - -namespace NFq { - -//// TJsonParser - -class TJsonParser::TImpl { -public: - TImpl(const TVector& columns, const TVector& types, TCallback parseCallback, ui64 batchSize, TDuration batchCreationTimeout, ui64 bufferCellCount) - : Alloc(__LOCATION__, NKikimr::TAlignedPagePoolCounters(), true, false) - , TypeEnv(std::make_unique(Alloc)) - , BatchSize(batchSize ? batchSize : DEFAULT_BATCH_SIZE) - , MaxNumberRows(((bufferCellCount ? bufferCellCount : DEFAULT_BUFFER_CELL_COUNT) - 1) / columns.size() + 1) - , BatchCreationTimeout(batchCreationTimeout) - , ParseCallback(parseCallback) - , ParsedValues(columns.size()) - { - Y_ENSURE(columns.size() == types.size(), "Number of columns and types should by equal"); - - with_lock (Alloc) { - auto functonRegistry = NKikimr::NMiniKQL::CreateFunctionRegistry(&PrintBackTrace, NKikimr::NMiniKQL::CreateBuiltinRegistry(), false, {}); - NKikimr::NMiniKQL::TProgramBuilder programBuilder(*TypeEnv, *functonRegistry); - - Columns.reserve(columns.size()); - for (size_t i = 0; i < columns.size(); i++) { - Columns.emplace_back(columns[i], types[i], MaxNumberRows, programBuilder); - } - } - - ColumnsIndex.reserve(columns.size()); - for (size_t i = 0; i < columns.size(); i++) { - ColumnsIndex.emplace(std::string_view(Columns[i].Name), i); - } - - for (size_t i = 0; i < columns.size(); i++) { - ParsedValues[i].resize(MaxNumberRows); - } - - Buffer.Reserve(BatchSize, MaxNumberRows); - - LOG_ROW_DISPATCHER_INFO("Simdjson active implementation " << simdjson::get_active_implementation()->name()); - Parser.threaded = false; - } - - bool IsReady() const { - return Buffer.IsReady() && (Buffer.GetSize() >= BatchSize || TInstant::Now() - Buffer.CreationStartTime >= BatchCreationTimeout); - } - - TInstant GetCreationDeadline() const { - return Buffer.IsReady() ? Buffer.CreationStartTime + BatchCreationTimeout : TInstant::Zero(); - } - - size_t GetNumberValues() const { - return Buffer.IsReady() ? Buffer.NumberValues : 0; - } - - const TVector& GetOffsets() { - return Buffer.Offsets; - } - - void AddMessages(const TVector& messages) { - Y_ENSURE(!Buffer.Finished, "Cannot add messages into finished buffer"); - for (const auto& message : messages) { - Buffer.AddMessage(message); - if (Buffer.IsReady() && Buffer.GetSize() >= BatchSize) { - Parse(); - } - } - } - - void Parse() { - Y_ENSURE(Buffer.IsReady(), "Nothing to parse"); - - const auto [values, size] = Buffer.Finish(); - LOG_ROW_DISPATCHER_TRACE("Parse values:\n" << values); - - with_lock (Alloc) { - Y_DEFER { - // Clear all UV in case of exception - ClearColumns(); - Buffer.Clear(); - }; - - size_t rowId = 0; - size_t parsedRows = 0; - simdjson::ondemand::document_stream documents = Parser.iterate_many(values, size, simdjson::ondemand::DEFAULT_BATCH_SIZE); - for (auto document : documents) { - if (Y_UNLIKELY(parsedRows >= Buffer.NumberValues)) { - throw NFq::TJsonParserError() << "Failed to parse json messages, expected " << Buffer.NumberValues << " json rows from offset " << Buffer.Offsets.front() << " but got " << parsedRows + 1; - } - for (auto item : document.get_object()) { - const auto it = ColumnsIndex.find(item.escaped_key().value()); - if (it == ColumnsIndex.end()) { - continue; - } - - const size_t columnId = it->second; - auto& columnParser = Columns[columnId]; - try { - columnParser.ParseJsonValue(rowId, item.value(), ParsedValues[columnId][rowId]); - } catch (...) { - throw NFq::TJsonParserError(columnParser.Name) << "Failed to parse json string at offset " << Buffer.Offsets[rowId] << ", got parsing error for column '" << columnParser.Name << "' with type " << columnParser.TypeYson << ", description: " << CurrentExceptionMessage(); - } - } - - rowId++; - parsedRows++; - if (rowId == MaxNumberRows) { - FlushColumns(parsedRows, MaxNumberRows); - rowId = 0; - } - } - - if (Y_UNLIKELY(parsedRows != Buffer.NumberValues)) { - throw NFq::TJsonParserError() << "Failed to parse json messages, expected " << Buffer.NumberValues << " json rows from offset " << Buffer.Offsets.front() << " but got " << rowId; - } - if (rowId) { - FlushColumns(parsedRows, rowId); - } - } - } - - TString GetDescription() const { - TStringBuilder description = TStringBuilder() << "Columns: "; - for (const auto& column : Columns) { - description << "'" << column.Name << "':" << column.TypeYson << " "; - } - description << "\nNumber values in buffer: " << Buffer.NumberValues << ", buffer size: " << Buffer.GetSize() << ", finished: " << Buffer.Finished; - return description; - } - - ~TImpl() { - with_lock (Alloc) { - ParsedValues.clear(); - Columns.clear(); - TypeEnv.reset(); - } - } - -private: - void FlushColumns(size_t parsedRows, size_t savedRows) { - const ui64 firstOffset = Buffer.Offsets.front(); - for (const auto& column : Columns) { - column.ValidateNumberValues(savedRows, firstOffset); - } - - { - auto unguard = Unguard(Alloc); - ParseCallback(parsedRows - savedRows, savedRows, ParsedValues); - } - - ClearColumns(); - } - - void ClearColumns() { - for (size_t i = 0; i < Columns.size(); ++i) { - auto& parsedColumn = ParsedValues[i]; - for (size_t rowId : Columns[i].ParsedRows) { - auto& parsedRow = parsedColumn[rowId]; - parsedRow.UnlockRef(1); - parsedRow.Clear(); - } - Columns[i].ParsedRows.clear(); - } - } - -private: - NKikimr::NMiniKQL::TScopedAlloc Alloc; - std::unique_ptr TypeEnv; - - const ui64 BatchSize; - const ui64 MaxNumberRows; - const TDuration BatchCreationTimeout; - const TCallback ParseCallback; - TVector Columns; - absl::flat_hash_map ColumnsIndex; - - TJsonParserBuffer Buffer; - simdjson::ondemand::parser Parser; - - TVector> ParsedValues; -}; - -TJsonParser::TJsonParser(const TVector& columns, const TVector& types, TCallback parseCallback, ui64 batchSize, TDuration batchCreationTimeout, ui64 bufferCellCount) - : Impl(std::make_unique(columns, types, parseCallback, batchSize, batchCreationTimeout, bufferCellCount)) -{} - -TJsonParser::~TJsonParser() { -} - -void TJsonParser::AddMessages(const TVector& messages) { - Impl->AddMessages(messages); -} - -bool TJsonParser::IsReady() const { - return Impl->IsReady(); -} - -TInstant TJsonParser::GetCreationDeadline() const { - return Impl->GetCreationDeadline(); -} - -size_t TJsonParser::GetNumberValues() const { - return Impl->GetNumberValues(); -} - -const TVector& TJsonParser::GetOffsets() const { - return Impl->GetOffsets(); -} - -void TJsonParser::Parse() { - Impl->Parse(); -} - -TString TJsonParser::GetDescription() const { - return Impl->GetDescription(); -} - -std::unique_ptr NewJsonParser(const TVector& columns, const TVector& types, TJsonParser::TCallback parseCallback, ui64 batchSize, TDuration batchCreationTimeout, ui64 bufferCellCount) { - return std::unique_ptr(new TJsonParser(columns, types, parseCallback, batchSize, batchCreationTimeout, bufferCellCount)); -} - -} // namespace NFq diff --git a/ydb/core/fq/libs/row_dispatcher/json_parser.h b/ydb/core/fq/libs/row_dispatcher/json_parser.h deleted file mode 100644 index e96a8f752599..000000000000 --- a/ydb/core/fq/libs/row_dispatcher/json_parser.h +++ /dev/null @@ -1,50 +0,0 @@ -#pragma once - -#include - -#include - -namespace NFq { - -class TJsonParserError: public yexception { -public: - TJsonParserError() = default; - TJsonParserError(const std::string& fieldName) - : FieldName(fieldName) - {} - - TMaybe GetField() const noexcept { - return FieldName; - } - -private: - TMaybe FieldName; -}; - - -class TJsonParser { -public: - using TCallback = std::function>& parsedValues)>; - -public: - TJsonParser(const TVector& columns, const TVector& types, TCallback parseCallback, ui64 batchSize, TDuration batchCreationTimeout, ui64 bufferCellCount); - ~TJsonParser(); - - bool IsReady() const; - TInstant GetCreationDeadline() const; - size_t GetNumberValues() const; - const TVector& GetOffsets() const; - - void AddMessages(const TVector& messages); - void Parse(); - - TString GetDescription() const; - -private: - class TImpl; - const std::unique_ptr Impl; -}; - -std::unique_ptr NewJsonParser(const TVector& columns, const TVector& types, TJsonParser::TCallback parseCallback, ui64 batchSize, TDuration batchCreationTimeout, ui64 bufferCellCount); - -} // namespace NFq diff --git a/ydb/core/fq/libs/row_dispatcher/leader_election.cpp b/ydb/core/fq/libs/row_dispatcher/leader_election.cpp index 9d267183c0f3..2ba97f2cc45b 100644 --- a/ydb/core/fq/libs/row_dispatcher/leader_election.cpp +++ b/ydb/core/fq/libs/row_dispatcher/leader_election.cpp @@ -77,12 +77,12 @@ struct TLeaderElectionMetrics { explicit TLeaderElectionMetrics(const ::NMonitoring::TDynamicCounterPtr& counters) : Counters(counters) { Errors = Counters->GetCounter("LeaderElectionErrors", true); - LeaderChangedCount = Counters->GetCounter("LeaderElectionChangedCount"); + LeaderChanged = Counters->GetCounter("LeaderChanged", true); } ::NMonitoring::TDynamicCounterPtr Counters; ::NMonitoring::TDynamicCounters::TCounterPtr Errors; - ::NMonitoring::TDynamicCounters::TCounterPtr LeaderChangedCount; + ::NMonitoring::TDynamicCounters::TCounterPtr LeaderChanged; }; class TLeaderElection: public TActorBootstrapped { @@ -458,7 +458,7 @@ void TLeaderElection::Handle(TEvPrivate::TEvDescribeSemaphoreResult::TPtr& ev) { if (!LeaderActorId || (*LeaderActorId != id)) { LOG_ROW_DISPATCHER_INFO("Send TEvCoordinatorChanged to " << ParentId); TActivationContext::ActorSystem()->Send(ParentId, new NFq::TEvRowDispatcher::TEvCoordinatorChanged(id, generation)); - Metrics.LeaderChangedCount->Inc(); + Metrics.LeaderChanged->Inc(); } LeaderActorId = id; } diff --git a/ydb/core/fq/libs/row_dispatcher/probes.cpp b/ydb/core/fq/libs/row_dispatcher/probes.cpp new file mode 100644 index 000000000000..bbe9019ac5fd --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/probes.cpp @@ -0,0 +1,3 @@ +#include "probes.h" + +LWTRACE_DEFINE_PROVIDER(FQ_ROW_DISPATCHER_PROVIDER) \ No newline at end of file diff --git a/ydb/core/fq/libs/row_dispatcher/probes.h b/ydb/core/fq/libs/row_dispatcher/probes.h new file mode 100644 index 000000000000..c01ebd24cd3a --- /dev/null +++ b/ydb/core/fq/libs/row_dispatcher/probes.h @@ -0,0 +1,101 @@ +#pragma once + +#include + +#define FQ_ROW_DISPATCHER_PROVIDER(PROBE, EVENT, GROUPS, TYPES, NAMES) \ + PROBE(CoordinatorChanged, \ + GROUPS(), \ + TYPES(TString, ui64, TString, ui64, TString), \ + NAMES("sender", "newGeneration", "newCoordinator", "oldGeneration", "oldCoordinator")) \ + PROBE(NodeConnected, \ + GROUPS(), \ + TYPES(TString, ui32), \ + NAMES("sender", "nodeId")) \ + PROBE(NodeDisconnected, \ + GROUPS(), \ + TYPES(TString, ui32), \ + NAMES("sender", "nodeId")) \ + PROBE(UndeliveredStart, \ + GROUPS(), \ + TYPES(TString, ui64, ui64), \ + NAMES("sender", "reason", "generation")) \ + PROBE(UndeliveredSkipGeneration, \ + GROUPS(), \ + TYPES(TString, ui64, ui64, ui64), \ + NAMES("sender", "reason", "generation", "consumerGeneration")) \ + PROBE(UndeliveredDeleteConsumer, \ + GROUPS(), \ + TYPES(TString, ui64, ui64, TString), \ + NAMES("sender", "reason", "generation", "key")) \ + PROBE(CoordinatorPing, \ + GROUPS(), \ + TYPES(TString), \ + NAMES("coordinatorActor")) \ + PROBE(Pong, \ + GROUPS(), \ + TYPES(), \ + NAMES()) \ + PROBE(CoordinatorChangesSubscribe, \ + GROUPS(), \ + TYPES(TString, ui64, TString), \ + NAMES("sender", "coordinatorGeneration", "coordinatorActor")) \ + PROBE(StartSession, \ + GROUPS(), \ + TYPES(TString, TString, ui64), \ + NAMES("sender", "queryId", "size")) \ + PROBE(GetNextBatch, \ + GROUPS(), \ + TYPES(TString, ui32, TString, ui64), \ + NAMES("sender", "partitionId", "queryId", "size")) \ + PROBE(Heartbeat, \ + GROUPS(), \ + TYPES(TString, ui32, TString, ui64), \ + NAMES("sender", "partitionId", "queryId", "size")) \ + PROBE(StopSession, \ + GROUPS(), \ + TYPES(TString, TString, ui64), \ + NAMES("sender", "queryId", "size")) \ + PROBE(TryConnect, \ + GROUPS(), \ + TYPES(TString, ui32), \ + NAMES("sender", "nodeId")) \ + PROBE(PrivateHeartbeat, \ + GROUPS(), \ + TYPES(TString, TString, ui64), \ + NAMES("sender", "queryId", "generation")) \ + PROBE(NewDataArrived, \ + GROUPS(), \ + TYPES(TString, TString, TString, ui64, ui64), \ + NAMES("sender", "readActor", "queryId", "generation", "size")) \ + PROBE(MessageBatch, \ + GROUPS(), \ + TYPES(TString, TString, TString, ui64, ui64), \ + NAMES("sender", "readActor", "queryId", "generation", "size")) \ + PROBE(SessionError, \ + GROUPS(), \ + TYPES(TString, TString, TString, ui64, ui64), \ + NAMES("sender", "readActor", "queryId", "generation", "size")) \ + PROBE(Statistics, \ + GROUPS(), \ + TYPES(TString, TString, ui64, ui64), \ + NAMES("readActor", "queryId", "generation", "size")) \ + PROBE(UpdateMetrics, \ + GROUPS(), \ + TYPES(), \ + NAMES()) \ + PROBE(PrintStateToLog, \ + GROUPS(), \ + TYPES(ui64), \ + NAMES("printStateToLogPeriodSec")) \ + PROBE(SessionStatistic, \ + GROUPS(), \ + TYPES(TString, TString, TString, TString, TString, ui32, ui64, ui64, ui64, ui64, ui64), \ + NAMES("sender", "readGroup", "endpoint","database", "partitionId", "readBytes", "queuedBytes", "restartSessionByOffsets", "readEvents", "lastReadedOffset")) \ + PROBE(GetInternalState, \ + GROUPS(), \ + TYPES(TString, ui64), \ + NAMES("sender", "size")) \ + +// FQ_ROW_DISPATCHER_PROVIDER + +LWTRACE_DECLARE_PROVIDER(FQ_ROW_DISPATCHER_PROVIDER) diff --git a/ydb/core/fq/libs/row_dispatcher/protos/events.proto b/ydb/core/fq/libs/row_dispatcher/protos/events.proto index 8a981fd7c691..22f969ea0001 100644 --- a/ydb/core/fq/libs/row_dispatcher/protos/events.proto +++ b/ydb/core/fq/libs/row_dispatcher/protos/events.proto @@ -6,28 +6,39 @@ option cc_enable_arenas = true; import "ydb/library/actors/protos/actors.proto"; import "ydb/library/yql/providers/pq/proto/dq_io.proto"; import "ydb/library/yql/dq/actors/protos/dq_events.proto"; +import "ydb/library/yql/dq/actors/protos/dq_status_codes.proto"; +import "ydb/public/api/protos/ydb_issue_message.proto"; message TEvGetAddressRequest { NYql.NPq.NProto.TDqPqTopicSource Source = 1; - repeated uint32 PartitionId = 2; + repeated uint32 PartitionId = 2 [deprecated=true]; + repeated uint32 PartitionIds = 3; } message TEvPartitionAddress { - repeated uint32 PartitionId = 1; + repeated uint32 PartitionId = 1 [deprecated=true]; NActorsProto.TActorId ActorId = 2; + repeated uint32 PartitionIds = 3; } message TEvGetAddressResponse { repeated TEvPartitionAddress Partitions = 1; } +message TPartitionOffset { + uint32 PartitionId = 1; + uint64 Offset = 2; +} + message TEvStartSession { NYql.NPq.NProto.TDqPqTopicSource Source = 1; - uint32 PartitionId = 2; + uint32 PartitionId = 2 [deprecated=true]; string Token = 3; - optional uint64 Offset = 4; + optional uint64 Offset = 4 [deprecated=true]; uint64 StartingMessageTimestampMs = 5; string QueryId = 6; + repeated uint32 PartitionIds = 7; + repeated TPartitionOffset Offsets = 8; optional NYql.NDqProto.TMessageTransportMeta TransportMeta = 100; } @@ -48,13 +59,15 @@ message TEvNewDataArrived { message TEvStopSession { NYql.NPq.NProto.TDqPqTopicSource Source = 1; - uint32 PartitionId = 2; + uint32 PartitionId = 2 [deprecated=true]; optional NYql.NDqProto.TMessageTransportMeta TransportMeta = 100; } message TEvMessage { - string Json = 1; - uint64 Offset = 2; + reserved 1; + reserved 2; + uint32 PayloadId = 3; + repeated uint64 Offsets = 4; } message TEvMessageBatch { @@ -64,16 +77,29 @@ message TEvMessageBatch { optional NYql.NDqProto.TMessageTransportMeta TransportMeta = 100; } -message TEvStatistics { +message TPartitionStatistics { uint32 PartitionId = 1; uint64 NextMessageOffset = 2; +} + +message TEvStatistics { + uint32 PartitionId = 1; // deprecated + uint64 NextMessageOffset = 2; // deprecated uint64 ReadBytes = 3; + repeated TPartitionStatistics Partition = 4; + uint64 CpuMicrosec = 5; + uint64 FilteredBytes = 6; + uint64 FilteredRows = 7; + uint64 QueuedBytes = 8; + uint64 QueuedRows = 9; optional NYql.NDqProto.TMessageTransportMeta TransportMeta = 100; } message TEvSessionError { - string Message = 1; - uint32 PartitionId = 2; + reserved 1; + uint32 PartitionId = 2 [deprecated=true]; + NYql.NDqProto.StatusIds.StatusCode StatusCode = 3; + repeated Ydb.Issue.IssueMessage Issues = 4; optional NYql.NDqProto.TMessageTransportMeta TransportMeta = 100; } @@ -82,6 +108,10 @@ message TEvHeartbeat { optional NYql.NDqProto.TMessageTransportMeta TransportMeta = 100; } +message TEvNoSession { + uint32 PartitionId = 1; +} + message TEvGetInternalStateRequest { } diff --git a/ydb/core/fq/libs/row_dispatcher/protos/ya.make b/ydb/core/fq/libs/row_dispatcher/protos/ya.make index c2d06e232661..71a7967d22e8 100644 --- a/ydb/core/fq/libs/row_dispatcher/protos/ya.make +++ b/ydb/core/fq/libs/row_dispatcher/protos/ya.make @@ -8,6 +8,7 @@ PEERDIR( ydb/library/actors/protos ydb/library/yql/dq/actors/protos ydb/library/yql/providers/pq/proto + ydb/public/api/protos ) EXCLUDE_TAGS(GO_PROTO) diff --git a/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/compile_service.cpp b/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/compile_service.cpp index 02b29c005883..e79c26c17ace 100644 --- a/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/compile_service.cpp +++ b/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/compile_service.cpp @@ -1,7 +1,10 @@ #include "compile_service.h" +#include #include +#include +#include #include #include @@ -10,38 +13,174 @@ namespace NFq::NRowDispatcher { namespace { +struct TEvPrivate { + // Event ids + enum EEv : ui32 { + EvCompileFinished = EventSpaceBegin(NActors::TEvents::ES_PRIVATE), + EvEnd + }; + + static_assert(EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE)"); + + // Events + struct TEvCompileFinished : public NActors::TEventLocal { + TEvCompileFinished(NActors::TActorId requestActor, ui64 requestId) + : RequestActor(requestActor) + , RequestId(requestId) + {} + + const NActors::TActorId RequestActor; + const ui64 RequestId; + }; +}; + +class TPurecalcCompileActor : public NActors::TActorBootstrapped { +public: + TPurecalcCompileActor(NActors::TActorId owner, NYql::NPureCalc::IProgramFactoryPtr factory, TEvRowDispatcher::TEvPurecalcCompileRequest::TPtr request) + : Owner(owner) + , Factory(factory) + , LogPrefix(TStringBuilder() << "TPurecalcCompileActor " << request->Sender << " [id " << request->Cookie << "]: ") + , Request(std::move(request)) + {} + + static constexpr char ActorName[] = "FQ_ROW_DISPATCHER_COMPILE_ACTOR"; + + void Bootstrap() { + Y_DEFER { + Finish(); + }; + + LOG_ROW_DISPATCHER_TRACE("Started compile request"); + IProgramHolder::TPtr programHolder = std::move(Request->Get()->ProgramHolder); + + TStatus status = TStatus::Success(); + try { + programHolder->CreateProgram(Factory); + } catch (const NYql::NPureCalc::TCompileError& error) { + status = TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Compile issues: " << error.GetIssues()) + .AddIssue(TStringBuilder() << "Final yql: " << error.GetYql()) + .AddParentIssue(TStringBuilder() << "Failed to compile purecalc program"); + } catch (...) { + status = TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Failed to compile purecalc program, got unexpected exception: " << CurrentExceptionMessage()); + } + + if (status.IsFail()) { + LOG_ROW_DISPATCHER_ERROR("Compilation failed for request"); + Send(Request->Sender, new TEvRowDispatcher::TEvPurecalcCompileResponse(status.GetStatus(), status.GetErrorDescription()), 0, Request->Cookie); + } else { + LOG_ROW_DISPATCHER_TRACE("Compilation completed for request"); + Send(Request->Sender, new TEvRowDispatcher::TEvPurecalcCompileResponse(std::move(programHolder)), 0, Request->Cookie); + } + } + +private: + void Finish() { + Send(Owner, new TEvPrivate::TEvCompileFinished(Request->Sender, Request->Cookie)); + PassAway(); + } + +private: + const NActors::TActorId Owner; + const NYql::NPureCalc::IProgramFactoryPtr Factory; + const TString LogPrefix; + + TEvRowDispatcher::TEvPurecalcCompileRequest::TPtr Request; +}; + class TPurecalcCompileService : public NActors::TActor { using TBase = NActors::TActor; + struct TCounters { + const NMonitoring::TDynamicCounterPtr Counters; + + NMonitoring::TDynamicCounters::TCounterPtr ActiveCompileActors; + NMonitoring::TDynamicCounters::TCounterPtr CompileQueueSize; + + explicit TCounters(NMonitoring::TDynamicCounterPtr counters) + : Counters(counters) + { + Register(); + } + + private: + void Register() { + ActiveCompileActors = Counters->GetCounter("ActiveCompileActors", false); + CompileQueueSize = Counters->GetCounter("CompileQueueSize", false); + } + }; + public: - TPurecalcCompileService() + TPurecalcCompileService(const NConfig::TCompileServiceConfig& config, NMonitoring::TDynamicCounterPtr counters) : TBase(&TPurecalcCompileService::StateFunc) + , Config(config) + , InFlightLimit(Config.GetParallelCompilationLimit() ? Config.GetParallelCompilationLimit() : 1) + , LogPrefix("TPurecalcCompileService: ") + , Counters(counters) {} + static constexpr char ActorName[] = "FQ_ROW_DISPATCHER_COMPILE_SERVICE"; + STRICT_STFUNC(StateFunc, hFunc(TEvRowDispatcher::TEvPurecalcCompileRequest, Handle); + hFunc(TEvRowDispatcher::TEvPurecalcCompileAbort, Handle) + hFunc(TEvPrivate::TEvCompileFinished, Handle); ) void Handle(TEvRowDispatcher::TEvPurecalcCompileRequest::TPtr& ev) { - IProgramHolder::TPtr programHolder = std::move(ev->Get()->ProgramHolder); + const auto requestActor = ev->Sender; + const ui64 requestId = ev->Cookie; + LOG_ROW_DISPATCHER_TRACE("Add to compile queue request with id " << requestId << " from " << requestActor); - TString error; - try { - programHolder->CreateProgram(GetOrCreateFactory(ev->Get()->Settings)); - } catch (const NYql::NPureCalc::TCompileError& e) { - error = TStringBuilder() << "Failed to compile purecalc filter: sql: " << e.GetYql() << ", error: " << e.GetIssues(); - } catch (...) { - error = TStringBuilder() << "Failed to compile purecalc filter, unexpected exception: " << CurrentExceptionMessage(); + // Remove old compile request + RemoveRequest(requestActor, requestId); + + // Add new request + RequestsQueue.emplace_back(std::move(ev)); + Y_ENSURE(RequestsIndex.emplace(std::make_pair(requestActor, requestId), --RequestsQueue.end()).second); + Counters.CompileQueueSize->Inc(); + + StartCompilation(); + } + + void Handle(TEvRowDispatcher::TEvPurecalcCompileAbort::TPtr& ev) { + LOG_ROW_DISPATCHER_TRACE("Abort compile request with id " << ev->Cookie << " from " << ev->Sender); + + RemoveRequest(ev->Sender, ev->Cookie); + } + + void Handle(TEvPrivate::TEvCompileFinished::TPtr& ev) { + LOG_ROW_DISPATCHER_TRACE("Compile finished for request with id " << ev->Get()->RequestId << " from " << ev->Get()->RequestActor); + + InFlightCompilations.erase(ev->Sender); + Counters.ActiveCompileActors->Dec(); + + StartCompilation(); + } + +private: + void RemoveRequest(NActors::TActorId requestActor, ui64 requestId) { + const auto it = RequestsIndex.find(std::make_pair(requestActor, requestId)); + if (it == RequestsIndex.end()) { + return; } - if (error) { - Send(ev->Sender, new TEvRowDispatcher::TEvPurecalcCompileResponse(error), 0, ev->Cookie); - } else { - Send(ev->Sender, new TEvRowDispatcher::TEvPurecalcCompileResponse(std::move(programHolder)), 0, ev->Cookie); + RequestsQueue.erase(it->second); + RequestsIndex.erase(it); + Counters.CompileQueueSize->Dec(); + } + + void StartCompilation() { + while (!RequestsQueue.empty() && InFlightCompilations.size() < InFlightLimit) { + auto request = std::move(RequestsQueue.front()); + RemoveRequest(request->Sender, request->Cookie); + + const auto factory = GetOrCreateFactory(request->Get()->Settings); + const auto compileActor = Register(new TPurecalcCompileActor(SelfId(), factory, std::move(request))); + Y_ENSURE(InFlightCompilations.emplace(compileActor).second); + Counters.ActiveCompileActors->Inc(); } } -private: NYql::NPureCalc::IProgramFactoryPtr GetOrCreateFactory(const TPurecalcCompileSettings& settings) { const auto it = ProgramFactories.find(settings); if (it != ProgramFactories.end()) { @@ -54,13 +193,23 @@ class TPurecalcCompileService : public NActors::TActor } private: + const NConfig::TCompileServiceConfig Config; + const ui64 InFlightLimit; + const TString LogPrefix; + + std::list RequestsQueue; + THashMap, std::list::iterator> RequestsIndex; + std::unordered_set InFlightCompilations; + std::map ProgramFactories; + + const TCounters Counters; }; -} // namespace { +} // anonymous namespace -NActors::IActor* CreatePurecalcCompileService() { - return new TPurecalcCompileService(); +NActors::IActor* CreatePurecalcCompileService(const NConfig::TCompileServiceConfig& config, NMonitoring::TDynamicCounterPtr counters) { + return new TPurecalcCompileService(config, counters); } } // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/compile_service.h b/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/compile_service.h index 4675cf4d3be4..1ffd534ac99b 100644 --- a/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/compile_service.h +++ b/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/compile_service.h @@ -1,9 +1,11 @@ #pragma once +#include + #include namespace NFq::NRowDispatcher { -NActors::IActor* CreatePurecalcCompileService(); +NActors::IActor* CreatePurecalcCompileService(const NConfig::TCompileServiceConfig& config, NMonitoring::TDynamicCounterPtr counters); } // namespace NFq::NRowDispatcher diff --git a/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/ya.make b/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/ya.make index af6747855195..4f4b28718dfd 100644 --- a/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/ya.make +++ b/ydb/core/fq/libs/row_dispatcher/purecalc_compilation/ya.make @@ -5,7 +5,10 @@ SRCS( ) PEERDIR( + ydb/core/fq/libs/actors/logging + ydb/core/fq/libs/config/protos ydb/core/fq/libs/row_dispatcher/events + ydb/core/fq/libs/row_dispatcher/format_handler/common ydb/core/fq/libs/row_dispatcher/purecalc_no_pg_wrapper ydb/library/actors/core diff --git a/ydb/core/fq/libs/row_dispatcher/row_dispatcher.cpp b/ydb/core/fq/libs/row_dispatcher/row_dispatcher.cpp index 84870e1f81f9..caaba984513f 100644 --- a/ydb/core/fq/libs/row_dispatcher/row_dispatcher.cpp +++ b/ydb/core/fq/libs/row_dispatcher/row_dispatcher.cpp @@ -1,9 +1,9 @@ #include "row_dispatcher.h" #include "actors_factory.h" -#include "common.h" #include "coordinator.h" #include "leader_election.h" +#include "probes.h" #include #include @@ -15,6 +15,7 @@ #include #include +#include #include #include @@ -24,10 +25,14 @@ #include #include +#include + namespace NFq { using namespace NActors; +LWTRACE_USING(FQ_ROW_DISPATCHER_PROVIDER); + namespace { const ui64 CoordinatorPingPeriodSec = 2; @@ -50,6 +55,23 @@ struct TRowDispatcherMetrics { ::NMonitoring::TDynamicCounters::TCounterPtr NodesReconnect; }; +struct TUserPoolMetrics { + explicit TUserPoolMetrics(const ::NMonitoring::TDynamicCounterPtr& utilsCounters) { + auto execpoolGroup = utilsCounters->GetSubgroup("execpool", "User"); + auto microsecGroup = execpoolGroup->GetSubgroup("sensor", "ElapsedMicrosecByActivity"); + Session = microsecGroup->GetNamedCounter("activity", "FQ_ROW_DISPATCHER_SESSION", true); + RowDispatcher = microsecGroup->GetNamedCounter("activity", "FQ_ROW_DISPATCHER", true); + CompilerActor = microsecGroup->GetNamedCounter("activity", "FQ_ROW_DISPATCHER_COMPILE_ACTOR", true); + CompilerService = microsecGroup->GetNamedCounter("activity", "FQ_ROW_DISPATCHER_COMPILE_SERVICE", true); + FormatHandler = microsecGroup->GetNamedCounter("activity", "FQ_ROW_DISPATCHER_FORMAT_HANDLER", true); + } + ::NMonitoring::TDynamicCounters::TCounterPtr Session; + ::NMonitoring::TDynamicCounters::TCounterPtr RowDispatcher; + ::NMonitoring::TDynamicCounters::TCounterPtr CompilerActor; + ::NMonitoring::TDynamicCounters::TCounterPtr CompilerService; + ::NMonitoring::TDynamicCounters::TCounterPtr FormatHandler; +}; + struct TEvPrivate { // Event ids enum EEv : ui32 { @@ -58,6 +80,7 @@ struct TEvPrivate { EvUpdateMetrics, EvPrintStateToLog, EvTryConnect, + EvSendStatistic, EvEnd }; @@ -70,21 +93,40 @@ struct TEvPrivate { : NodeId(nodeId) {} ui32 NodeId = 0; }; + struct TEvSendStatistic : public NActors::TEventLocal {}; +}; + +struct TQueryStatKey { + TString QueryId; + TString ReadGroup; + + size_t Hash() const noexcept { + ui64 hash = std::hash()(QueryId); + hash = CombineHashes(hash, std::hash()(ReadGroup)); + return hash; + } + bool operator==(const TQueryStatKey& other) const { + return QueryId == other.QueryId && ReadGroup == other.ReadGroup; + } }; -using TQueryStatKey = std::pair; // QueryId / Topic +struct TQueryStatKeyHash { + size_t operator()(const TQueryStatKey& k) const { + return k.Hash(); + } +}; struct TAggQueryStat { - NYql::TCounters::TEntry ReadBytes; - NYql::TCounters::TEntry UnreadBytes; - NYql::TCounters::TEntry UnreadRows; + NYql::TCounters::TEntry FilteredBytes; + NYql::TCounters::TEntry QueuedBytes; + NYql::TCounters::TEntry QueuedRows; NYql::TCounters::TEntry ReadLagMessages; bool IsWaiting = false; - void Add(const TopicSessionClientStatistic& stat) { - ReadBytes.Add(NYql::TCounters::TEntry(stat.ReadBytes)); - UnreadBytes.Add(NYql::TCounters::TEntry(stat.UnreadBytes)); - UnreadRows.Add(NYql::TCounters::TEntry(stat.UnreadRows)); + void Add(const TTopicSessionClientStatistic& stat, ui64 filteredBytes) { + FilteredBytes.Add(NYql::TCounters::TEntry(filteredBytes)); + QueuedBytes.Add(NYql::TCounters::TEntry(stat.QueuedBytes)); + QueuedRows.Add(NYql::TCounters::TEntry(stat.QueuedRows)); ReadLagMessages.Add(NYql::TCounters::TEntry(stat.ReadLagMessages)); IsWaiting = IsWaiting || stat.IsWaiting; } @@ -92,52 +134,34 @@ struct TAggQueryStat { ui64 UpdateMetricsPeriodSec = 60; ui64 PrintStateToLogPeriodSec = 600; -ui64 PrintStateToLogSplitSize = 512000; +ui64 PrintStateToLogSplitSize = 64000; ui64 MaxSessionBufferSizeBytes = 16000000; class TRowDispatcher : public TActorBootstrapped { - struct ConsumerSessionKey { - TActorId ReadActorId; - ui32 PartitionId; - - size_t Hash() const noexcept { - ui64 hash = std::hash()(ReadActorId); - hash = CombineHashes(hash, std::hash()(PartitionId)); - return hash; - } - bool operator==(const ConsumerSessionKey& other) const { - return ReadActorId == other.ReadActorId && PartitionId == other.PartitionId; - } - }; - - struct ConsumerSessionKeyHash { - int operator()(const ConsumerSessionKey& k) const { - return k.Hash(); - } - }; - - struct TopicSessionKey { + struct TTopicSessionKey { + TString ReadGroup; TString Endpoint; TString Database; TString TopicPath; ui64 PartitionId; size_t Hash() const noexcept { - ui64 hash = std::hash()(Endpoint); + ui64 hash = std::hash()(ReadGroup); + hash = CombineHashes(hash, std::hash()(Endpoint)); hash = CombineHashes(hash, std::hash()(Database)); hash = CombineHashes(hash, std::hash()(TopicPath)); hash = CombineHashes(hash, std::hash()(PartitionId)); return hash; } - bool operator==(const TopicSessionKey& other) const { - return Endpoint == other.Endpoint && Database == other.Database + bool operator==(const TTopicSessionKey& other) const { + return ReadGroup == other.ReadGroup && Endpoint == other.Endpoint && Database == other.Database && TopicPath == other.TopicPath && PartitionId == other.PartitionId; } }; - struct TopicSessionKeyHash { - int operator()(const TopicSessionKey& k) const { + struct TTopicSessionKeyHash { + int operator()(const TTopicSessionKey& k) const { return k.Hash(); } }; @@ -242,8 +266,7 @@ class TRowDispatcher : public TActorBootstrapped { }; struct TAggregatedStats{ - NYql::TCounters::TEntry AllSessionsReadBytes; - TMap> LastQueryStats; + THashMap, TQueryStatKeyHash> LastQueryStats; TDuration LastUpdateMetricsPeriod; }; @@ -260,33 +283,43 @@ class TRowDispatcher : public TActorBootstrapped { TString Tenant; NFq::NRowDispatcher::IActorFactory::TPtr ActorFactory; const ::NMonitoring::TDynamicCounterPtr Counters; + const ::NMonitoring::TDynamicCounterPtr CountersRoot; TRowDispatcherMetrics Metrics; + TUserPoolMetrics UserPoolMetrics; NYql::IPqGateway::TPtr PqGateway; NActors::TMon* Monitoring; TNodesTracker NodesTracker; + NYql::TCounters::TEntry AllSessionsDateRate; TAggregatedStats AggrStats; + ui64 LastCpuTime = 0; - struct ConsumerCounters { + struct TConsumerCounters { ui64 NewDataArrived = 0; ui64 GetNextBatch = 0; ui64 MessageBatch = 0; }; - struct ConsumerInfo { - ConsumerInfo( + struct TConsumerPartition { + bool PendingGetNextBatch = false; + bool PendingNewDataArrived = false; + TActorId TopicSessionId; + TTopicSessionClientStatistic Stat; + ui64 FilteredBytes = 0; + bool StatisticsUpdated = false; + }; + + struct TConsumerInfo { + TConsumerInfo( NActors::TActorId readActorId, NActors::TActorId selfId, ui64 eventQueueId, NFq::NRowDispatcherProto::TEvStartSession& proto, - TActorId topicSessionId, bool alreadyConnected, ui64 generation) : ReadActorId(readActorId) , SourceParams(proto.GetSource()) - , PartitionId(proto.GetPartitionId()) , EventQueueId(eventQueueId) , Proto(proto) - , TopicSessionId(topicSessionId) , QueryId(proto.GetQueryId()) , Generation(generation) { EventsQueue.Init("txId", selfId, selfId, eventQueueId, /* KeepAlive */ true, /* UseConnect */ false); @@ -295,39 +328,36 @@ class TRowDispatcher : public TActorBootstrapped { NActors::TActorId ReadActorId; NYql::NPq::NProto::TDqPqTopicSource SourceParams; - ui64 PartitionId; NYql::NDq::TRetryEventsQueue EventsQueue; ui64 EventQueueId; NFq::NRowDispatcherProto::TEvStartSession Proto; - TActorId TopicSessionId; + THashMap Partitions; const TString QueryId; - ConsumerCounters Counters; - bool PendingGetNextBatch = false; - bool PendingNewDataArrived = false; - TopicSessionClientStatistic Stat; + TConsumerCounters Counters; + ui64 CpuMicrosec = 0; // Increment. ui64 Generation; }; - struct SessionInfo { - TMap> Consumers; // key - ReadActor actor id - TopicSessionCommonStatistic Stat; // Increments + struct TSessionInfo { + TMap> Consumers; // key - ReadActor actor id + TTopicSessionCommonStatistic Stat; // Increments NYql::TCounters::TEntry AggrReadBytes; }; - struct TopicSessionInfo { - TMap Sessions; // key - TopicSession actor id + struct TTopicSessionInfo { + TMap Sessions; // key - TopicSession actor id }; - struct ReadActorInfo { + struct TReadActorInfo { TString InternalState; TInstant RequestTime; TInstant ResponseTime; }; - THashMap, ConsumerSessionKeyHash> Consumers; - TMap> ConsumersByEventQueueId; - THashMap TopicSessions; - TMap ReadActorsInternalState; + THashMap> Consumers; // key - read actor id + TMap> ConsumersByEventQueueId; + THashMap TopicSessions; + TMap ReadActorsInternalState; public: explicit TRowDispatcher( @@ -338,6 +368,7 @@ class TRowDispatcher : public TActorBootstrapped { const TString& tenant, const NFq::NRowDispatcher::IActorFactory::TPtr& actorFactory, const ::NMonitoring::TDynamicCounterPtr& counters, + const ::NMonitoring::TDynamicCounterPtr& countersRoot, const NYql::IPqGateway::TPtr& pqGateway, NActors::TMon* monitoring = nullptr); @@ -359,26 +390,29 @@ class TRowDispatcher : public TActorBootstrapped { void Handle(NFq::TEvRowDispatcher::TEvNewDataArrived::TPtr& ev); void Handle(NFq::TEvRowDispatcher::TEvMessageBatch::TPtr& ev); void Handle(NFq::TEvRowDispatcher::TEvSessionError::TPtr& ev); - void Handle(NFq::TEvRowDispatcher::TEvStatistics::TPtr& ev); void Handle(NFq::TEvRowDispatcher::TEvSessionStatistic::TPtr& ev); void Handle(NFq::TEvRowDispatcher::TEvGetInternalStateResponse::TPtr& ev); void Handle(NFq::TEvRowDispatcher::TEvHeartbeat::TPtr& ev); + void Handle(NFq::TEvRowDispatcher::TEvNoSession::TPtr& ev); + void Handle(const TEvPrivate::TEvTryConnect::TPtr&); void Handle(const NYql::NDq::TEvRetryQueuePrivate::TEvEvHeartbeat::TPtr&); void Handle(NFq::TEvPrivate::TEvUpdateMetrics::TPtr&); void Handle(NFq::TEvPrivate::TEvPrintStateToLog::TPtr&); + void Handle(NFq::TEvPrivate::TEvSendStatistic::TPtr&); void Handle(const NMon::TEvHttpInfo::TPtr&); - void DeleteConsumer(const ConsumerSessionKey& key); + void DeleteConsumer(NActors::TActorId readActorId); void UpdateMetrics(); TString GetInternalState(); TString GetReadActorsInternalState(); void UpdateReadActorsInternalState(); template - bool CheckSession(TAtomicSharedPtr& consumer, const TEventPtr& ev); - void SetQueryMetrics(const TQueryStatKey& queryKey, ui64 unreadBytesMax, ui64 unreadBytesAvg, i64 readLagMessagesMax); + bool CheckSession(TAtomicSharedPtr& consumer, const TEventPtr& ev); + void SetQueryMetrics(const TQueryStatKey& queryKey, ui64 queuedBytesMax, ui64 queuedBytesAvg, i64 readLagMessagesMax); void PrintStateToLog(); + void UpdateCpuTime(); STRICT_STFUNC( StateFunc, { @@ -393,8 +427,8 @@ class TRowDispatcher : public TActorBootstrapped { hFunc(NFq::TEvRowDispatcher::TEvMessageBatch, Handle); hFunc(NFq::TEvRowDispatcher::TEvStartSession, Handle); hFunc(NFq::TEvRowDispatcher::TEvStopSession, Handle); + hFunc(NFq::TEvRowDispatcher::TEvNoSession, Handle); hFunc(NFq::TEvRowDispatcher::TEvSessionError, Handle); - hFunc(NFq::TEvRowDispatcher::TEvStatistics, Handle); hFunc(NFq::TEvRowDispatcher::TEvSessionStatistic, Handle); hFunc(NFq::TEvRowDispatcher::TEvGetInternalStateResponse, Handle); hFunc(TEvPrivate::TEvTryConnect, Handle); @@ -403,6 +437,7 @@ class TRowDispatcher : public TActorBootstrapped { hFunc(NFq::TEvRowDispatcher::TEvNewDataArrived, Handle); hFunc(NFq::TEvPrivate::TEvUpdateMetrics, Handle); hFunc(NFq::TEvPrivate::TEvPrintStateToLog, Handle); + hFunc(NFq::TEvPrivate::TEvSendStatistic, Handle); hFunc(NMon::TEvHttpInfo, Handle); }) }; @@ -415,6 +450,7 @@ TRowDispatcher::TRowDispatcher( const TString& tenant, const NFq::NRowDispatcher::IActorFactory::TPtr& actorFactory, const ::NMonitoring::TDynamicCounterPtr& counters, + const ::NMonitoring::TDynamicCounterPtr& countersRoot, const NYql::IPqGateway::TPtr& pqGateway, NActors::TMon* monitoring) : Config(config) @@ -425,7 +461,9 @@ TRowDispatcher::TRowDispatcher( , Tenant(tenant) , ActorFactory(actorFactory) , Counters(counters) + , CountersRoot(countersRoot) , Metrics(counters) + , UserPoolMetrics(countersRoot->GetSubgroup("counters", "utils")) , PqGateway(pqGateway) , Monitoring(monitoring) { @@ -439,13 +477,15 @@ void TRowDispatcher::Bootstrap() { auto coordinatorId = Register(NewCoordinator(SelfId(), config, YqSharedResources, Tenant, Counters).release()); Register(NewLeaderElection(SelfId(), coordinatorId, config, CredentialsProviderFactory, YqSharedResources, Tenant, Counters).release()); - CompileServiceActorId = Register(NRowDispatcher::CreatePurecalcCompileService()); + CompileServiceActorId = Register(NRowDispatcher::CreatePurecalcCompileService(Config.GetCompileService(), Counters)); Schedule(TDuration::Seconds(CoordinatorPingPeriodSec), new TEvPrivate::TEvCoordinatorPing()); Schedule(TDuration::Seconds(UpdateMetricsPeriodSec), new NFq::TEvPrivate::TEvUpdateMetrics()); Schedule(TDuration::Seconds(PrintStateToLogPeriodSec), new NFq::TEvPrivate::TEvPrintStateToLog()); + Schedule(TDuration::Seconds(Config.GetSendStatusPeriodSec()), new NFq::TEvPrivate::TEvSendStatistic()); if (Monitoring) { + NLwTraceMonPage::ProbeRegistry().AddProbesList(LWTRACE_GET_PROBES(FQ_ROW_DISPATCHER_PROVIDER)); ::NMonitoring::TIndexMonPage* actorsMonPage = Monitoring->RegisterIndexPage("actors", "Actors"); Monitoring->RegisterActorPage(actorsMonPage, "row_dispatcher", "Row Dispatcher", false, TlsActivationContext->ExecutorThread.ActorSystem, SelfId()); @@ -454,6 +494,7 @@ void TRowDispatcher::Bootstrap() { } void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvCoordinatorChanged::TPtr& ev) { + LWPROBE(CoordinatorChanged, ev->Sender.ToString(), ev->Get()->Generation, ev->Get()->CoordinatorActorId.ToString(), CoordinatorGeneration, CoordinatorActorId->ToString()); LOG_ROW_DISPATCHER_DEBUG("Coordinator changed, old leader " << CoordinatorActorId << ", new " << ev->Get()->CoordinatorActorId << " generation " << ev->Get()->Generation); if (ev->Get()->Generation < CoordinatorGeneration) { LOG_ROW_DISPATCHER_ERROR("New generation (" << ev->Get()->Generation << ") is less previous (" << CoordinatorGeneration << "), ignore updates"); @@ -471,6 +512,7 @@ void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvCoordinatorChanged::TPtr& } void TRowDispatcher::HandleConnected(TEvInterconnect::TEvNodeConnected::TPtr& ev) { + LWPROBE(NodeConnected, ev->Sender.ToString(), ev->Get()->NodeId); LOG_ROW_DISPATCHER_DEBUG("EvNodeConnected, node id " << ev->Get()->NodeId); Metrics.NodesReconnect->Inc(); NodesTracker.HandleNodeConnected(ev->Get()->NodeId); @@ -480,6 +522,7 @@ void TRowDispatcher::HandleConnected(TEvInterconnect::TEvNodeConnected::TPtr& ev } void TRowDispatcher::HandleDisconnected(TEvInterconnect::TEvNodeDisconnected::TPtr& ev) { + LWPROBE(NodeDisconnected, ev->Sender.ToString(), ev->Get()->NodeId); LOG_ROW_DISPATCHER_DEBUG("TEvNodeDisconnected, node id " << ev->Get()->NodeId); Metrics.NodesReconnect->Inc(); NodesTracker.HandleNodeDisconnected(ev->Get()->NodeId); @@ -489,13 +532,16 @@ void TRowDispatcher::HandleDisconnected(TEvInterconnect::TEvNodeDisconnected::TP } void TRowDispatcher::Handle(NActors::TEvents::TEvUndelivered::TPtr& ev) { - LOG_ROW_DISPATCHER_DEBUG("TEvUndelivered, from " << ev->Sender << ", reason " << ev->Get()->Reason); + LWPROBE(UndeliveredStart, ev->Sender.ToString(), ev->Get()->Reason, ev->Cookie); + LOG_ROW_DISPATCHER_TRACE("TEvUndelivered, from " << ev->Sender << ", reason " << ev->Get()->Reason); for (auto& [key, consumer] : Consumers) { if (ev->Cookie != consumer->Generation) { // Several partitions in one read_actor have different Generation. + LWPROBE(UndeliveredSkipGeneration, ev->Sender.ToString(), ev->Get()->Reason, ev->Cookie, consumer->Generation); continue; } if (consumer->EventsQueue.HandleUndelivered(ev) == NYql::NDq::TRetryEventsQueue::ESessionState::SessionClosed) { - DeleteConsumer(key); + LWPROBE(UndeliveredDeleteConsumer, ev->Sender.ToString(), ev->Get()->Reason, ev->Cookie, key.ToString()); + DeleteConsumer(ev->Sender); break; } } @@ -506,11 +552,13 @@ void TRowDispatcher::Handle(TEvPrivate::TEvCoordinatorPing::TPtr&) { if (!CoordinatorActorId) { return; } + LWPROBE(CoordinatorPing, CoordinatorActorId->ToString()); LOG_ROW_DISPATCHER_TRACE("Send ping to " << *CoordinatorActorId); Send(*CoordinatorActorId, new NActors::TEvents::TEvPing()); } void TRowDispatcher::Handle(NActors::TEvents::TEvPong::TPtr&) { + LWPROBE(Pong); LOG_ROW_DISPATCHER_TRACE("NActors::TEvents::TEvPong"); } @@ -521,6 +569,7 @@ void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvCoordinatorChangesSubscrib if (!CoordinatorActorId) { return; } + LWPROBE(CoordinatorChangesSubscribe, ev->Sender.ToString(), CoordinatorGeneration, CoordinatorActorId->ToString()); Send(ev->Sender, new NFq::TEvRowDispatcher::TEvCoordinatorChanged(*CoordinatorActorId, CoordinatorGeneration), IEventHandle::FlagTrackDelivery); } @@ -534,47 +583,54 @@ void TRowDispatcher::UpdateMetrics() { return; } - AggrStats.AllSessionsReadBytes = NYql::TCounters::TEntry(); + AllSessionsDateRate = NYql::TCounters::TEntry(); for (auto& [queryId, stat] : AggrStats.LastQueryStats) { stat = Nothing(); } for (auto& [key, sessionsInfo] : TopicSessions) { - const auto& topic = key.TopicPath; for (auto& [actorId, sessionInfo] : sessionsInfo.Sessions) { auto read = NYql::TCounters::TEntry(sessionInfo.Stat.ReadBytes); - AggrStats.AllSessionsReadBytes.Add(read); + AllSessionsDateRate.Add(read); sessionInfo.AggrReadBytes = read; sessionInfo.Stat.Clear(); for (auto& [readActorId, consumer] : sessionInfo.Consumers) { - auto& stat = AggrStats.LastQueryStats[TQueryStatKey{consumer->QueryId, topic}]; + const auto partionIt = consumer->Partitions.find(key.PartitionId); + if (partionIt == consumer->Partitions.end()) { + continue; + } + auto& partition = partionIt->second; + auto& stat = AggrStats.LastQueryStats[TQueryStatKey{consumer->QueryId, key.ReadGroup}]; if (!stat) { stat = TAggQueryStat(); } - stat->Add(consumer->Stat); - consumer->Stat.Clear(); + stat->Add(partition.Stat, partition.FilteredBytes); + partition.FilteredBytes = 0; } } } - for (auto it = AggrStats.LastQueryStats.begin(); it != AggrStats.LastQueryStats.end();) { - const auto& stats = it->second; + THashSet toDelete; + for (const auto& [key, stats] : AggrStats.LastQueryStats) { if (!stats) { - SetQueryMetrics(it->first, 0, 0, 0); - it = AggrStats.LastQueryStats.erase(it); + toDelete.insert(key); continue; } - SetQueryMetrics(it->first, stats->UnreadBytes.Max, stats->UnreadBytes.Avg, stats->ReadLagMessages.Max); - ++it; + SetQueryMetrics(key, stats->QueuedBytes.Max, stats->QueuedBytes.Avg, stats->ReadLagMessages.Max); + } + for (const auto& key : toDelete) { + SetQueryMetrics(key, 0, 0, 0); + Metrics.Counters->RemoveSubgroup("query_id", key.QueryId); + AggrStats.LastQueryStats.erase(key); } PrintStateToLog(); } -void TRowDispatcher::SetQueryMetrics(const TQueryStatKey& queryKey, ui64 unreadBytesMax, ui64 unreadBytesAvg, i64 readLagMessagesMax) { - auto queryGroup = Metrics.Counters->GetSubgroup("queryId", queryKey.first); - auto topicGroup = queryGroup->GetSubgroup("topic", CleanupCounterValueString(queryKey.second)); - topicGroup->GetCounter("MaxUnreadBytes")->Set(unreadBytesMax); - topicGroup->GetCounter("AvgUnreadBytes")->Set(unreadBytesAvg); +void TRowDispatcher::SetQueryMetrics(const TQueryStatKey& queryKey, ui64 queuedBytesMax, ui64 queuedBytesAvg, i64 readLagMessagesMax) { + auto queryGroup = Metrics.Counters->GetSubgroup("query_id", queryKey.QueryId); + auto topicGroup = queryGroup->GetSubgroup("read_group", SanitizeLabel(queryKey.ReadGroup)); + topicGroup->GetCounter("MaxQueuedBytes")->Set(queuedBytesMax); + topicGroup->GetCounter("AvgQueuedBytes")->Set(queuedBytesAvg); topicGroup->GetCounter("MaxReadLag")->Set(readLagMessagesMax); } @@ -601,75 +657,100 @@ TString TRowDispatcher::GetInternalState() { auto printDataRate = [&](NYql::TCounters::TEntry entry) { str << " (sum " << toHumanDR(entry.Sum) << " max " << toHumanDR(entry.Max) << " min " << toHumanDR(entry.Min) << ")"; }; + str << "SelfId: " << SelfId().ToString() << "\n"; str << "Consumers count: " << Consumers.size() << "\n"; str << "TopicSessions count: " << TopicSessions.size() << "\n"; str << "Max session buffer size: " << toHuman(MaxSessionBufferSizeBytes) << "\n"; + str << "CpuMicrosec: " << toHuman(LastCpuTime) << "\n"; str << "DataRate (all sessions): "; - printDataRate(AggrStats.AllSessionsReadBytes); + printDataRate(AllSessionsDateRate); str << "\n"; - TMap queryState; - TMap sessionCountByQuery; - ui64 unreadBytesSum = 0; + THashMap queryState; + THashMap sessionCountByQuery; + ui64 queuedBytesSum = 0; - for (auto& [key, sessionsInfo] : TopicSessions) { + for (auto& [sessionKey, sessionsInfo] : TopicSessions) { for (auto& [actorId, sessionInfo] : sessionsInfo.Sessions) { - const auto& topic = key.TopicPath; - unreadBytesSum += sessionInfo.Stat.UnreadBytes; + queuedBytesSum += sessionInfo.Stat.QueuedBytes; for (auto& [readActorId, consumer] : sessionInfo.Consumers) { - auto key = TQueryStatKey{consumer->QueryId, topic}; + const auto partionIt = consumer->Partitions.find(sessionKey.PartitionId); + if (partionIt == consumer->Partitions.end()) { + continue; + } + const auto& partitionStat = partionIt->second.Stat; + auto key = TQueryStatKey{consumer->QueryId, sessionKey.ReadGroup}; ++sessionCountByQuery[key]; - queryState[key].Add(consumer->Stat); + queryState[key].Add(partitionStat, 0); } } } if (TopicSessions.size()) { - str << "Buffer used: " << Prec(unreadBytesSum * 100.0 / (TopicSessions.size() * MaxSessionBufferSizeBytes), 4) << "% (" << toHuman(unreadBytesSum) << ")\n"; + str << "Buffer used: " << Prec(queuedBytesSum * 100.0 / (TopicSessions.size() * MaxSessionBufferSizeBytes), 4) << "% (" << toHuman(queuedBytesSum) << ")\n"; } str << "Queries:\n"; for (const auto& [queryStatKey, stat]: queryState) { - auto [queryId, topic] = queryStatKey; + auto [queryId, readGroup] = queryStatKey; const auto& aggStat = AggrStats.LastQueryStats[queryStatKey]; auto sessionsBufferSumSize = sessionCountByQuery[queryStatKey] * MaxSessionBufferSizeBytes; - auto used = sessionsBufferSumSize ? (stat.UnreadBytes.Sum * 100.0 / sessionsBufferSumSize) : 0.0; - str << " " << queryId << " / " << topic << ": buffer used (all partitions) " << LeftPad(Prec(used, 4), 10) << "% (" << toHuman(stat.UnreadBytes.Sum) << ") unread max (one partition) " << toHuman(stat.UnreadBytes.Max) << " data rate"; + auto used = sessionsBufferSumSize ? (stat.QueuedBytes.Sum * 100.0 / sessionsBufferSumSize) : 0.0; + str << " " << queryId << " / " << readGroup << ": buffer used (all partitions) " << LeftPad(Prec(used, 4), 10) << "% (" << toHuman(stat.QueuedBytes.Sum) << ") unread max (one partition) " << toHuman(stat.QueuedBytes.Max) << " data rate"; if (aggStat) { - printDataRate(aggStat->ReadBytes); + printDataRate(aggStat->FilteredBytes); } str << " waiting " << stat.IsWaiting << " max read lag " << stat.ReadLagMessages.Max; str << "\n"; } str << "TopicSessions:\n"; for (auto& [key, sessionsInfo] : TopicSessions) { - str << " " << key.TopicPath << " / " << key.PartitionId; + str << " " << key.TopicPath << " / " << key.PartitionId << " / " << key.ReadGroup; for (auto& [actorId, sessionInfo] : sessionsInfo.Sessions) { str << " / " << LeftPad(actorId, 32) - << " data rate " << toHumanDR(sessionInfo.AggrReadBytes.Sum) << " unread bytes " << toHuman(sessionInfo.Stat.UnreadBytes) - << " offset " << LeftPad(sessionInfo.Stat.LastReadedOffset, 12) << " restarts by offsets " << sessionInfo.Stat.RestartSessionByOffsets - << " parse and filter lantecy " << sessionInfo.Stat.ParseAndFilterLatency << "\n"; + << " data rate " << toHumanDR(sessionInfo.AggrReadBytes.Sum) << " unread bytes " << toHuman(sessionInfo.Stat.QueuedBytes) + << " offset " << LeftPad(sessionInfo.Stat.LastReadedOffset, 12) << " restarts by offsets " << sessionInfo.Stat.RestartSessionByOffsets << "\n"; ui64 maxInitialOffset = 0; ui64 minInitialOffset = std::numeric_limits::max(); + for (const auto& [formatName, formatStats] : sessionInfo.Stat.FormatHandlers) { + str << " " << formatName + << " parse and filter lantecy " << formatStats.ParseAndFilterLatency + << " (parse " << formatStats.ParserStats.ParserLatency << ", filter " << formatStats.FilterStats.FilterLatency << ")\n"; + } + for (auto& [readActorId, consumer] : sessionInfo.Consumers) { - str << " " << consumer->QueryId << " " << LeftPad(readActorId, 32) << " unread bytes " - << toHuman(consumer->Stat.UnreadBytes) << " (" << leftPad(consumer->Stat.UnreadRows) << " rows) " - << " offset " << leftPad(consumer->Stat.Offset) << " init offset " << leftPad(consumer->Stat.InitialOffset) + if (!consumer->Partitions.contains(key.PartitionId)) { + continue; + } + const auto& partition = consumer->Partitions[key.PartitionId]; + const auto& stat = partition.Stat; + str << " " << consumer->QueryId << " " << LeftPad(readActorId, 33) << " unread bytes " + << toHuman(stat.QueuedBytes) << " (" << leftPad(stat.QueuedRows) << " rows) " + << " offset " << leftPad(stat.Offset) << " init offset " << leftPad(stat.InitialOffset) << " get " << leftPad(consumer->Counters.GetNextBatch) << " arr " << leftPad(consumer->Counters.NewDataArrived) << " btc " << leftPad(consumer->Counters.MessageBatch) - << " pend get " << leftPad(consumer->PendingGetNextBatch) << " pend new " << leftPad(consumer->PendingNewDataArrived) - << " waiting " << consumer->Stat.IsWaiting << " read lag " << leftPad(consumer->Stat.ReadLagMessages) - << " conn id " << consumer->Generation << " "; - str << " retry queue: "; - consumer->EventsQueue.PrintInternalState(str); - - maxInitialOffset = std::max(maxInitialOffset, consumer->Stat.InitialOffset); - minInitialOffset = std::min(minInitialOffset, consumer->Stat.InitialOffset); + << " pend get " << leftPad(partition.PendingGetNextBatch) << " pend new " << leftPad(partition.PendingNewDataArrived) + << " waiting " << stat.IsWaiting << " read lag " << leftPad(stat.ReadLagMessages) + << " conn id " << consumer->Generation << "\n"; + maxInitialOffset = std::max(maxInitialOffset, stat.InitialOffset); + minInitialOffset = std::min(minInitialOffset, stat.InitialOffset); } str << " initial offset max " << leftPad(maxInitialOffset) << " min " << leftPad(minInitialOffset) << "\n";; } } + + str << "Consumers:\n"; + for (auto& [readActorId, consumer] : Consumers) { + str << " " << consumer->QueryId << " " << LeftPad(readActorId, 32) << " Generation " << consumer->Generation << "\n"; + str << " partitions: "; + for (const auto& [partitionId, info] : consumer->Partitions) { + str << partitionId << ","; + } + str << "\n retry queue: "; + consumer->EventsQueue.PrintInternalState(str); + } + return str.Str(); } @@ -683,8 +764,8 @@ TString TRowDispatcher::GetReadActorsInternalState() { void TRowDispatcher::UpdateReadActorsInternalState() { TSet ReadActors; - for (const auto& [key, _]: Consumers) { - ReadActors.insert(key.ReadActorId); + for (const auto& [readActorId, _]: Consumers) { + ReadActors.insert(readActorId); } for(auto it = ReadActorsInternalState.begin(); it != ReadActorsInternalState.end();) { @@ -707,20 +788,16 @@ void TRowDispatcher::UpdateReadActorsInternalState() { } void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev) { - LOG_ROW_DISPATCHER_DEBUG("Received TEvStartSession from " << ev->Sender << ", topicPath " << ev->Get()->Record.GetSource().GetTopicPath() << - " part id " << ev->Get()->Record.GetPartitionId() << " query id " << ev->Get()->Record.GetQueryId() << " cookie " << ev->Cookie); - auto queryGroup = Metrics.Counters->GetSubgroup("queryId", ev->Get()->Record.GetQueryId()); - auto topicGroup = queryGroup->GetSubgroup("topic", CleanupCounterValueString(ev->Get()->Record.GetSource().GetTopicPath())); + LOG_ROW_DISPATCHER_DEBUG("Received TEvStartSession from " << ev->Sender << ", read group " << ev->Get()->Record.GetSource().GetReadGroup() << ", topicPath " << ev->Get()->Record.GetSource().GetTopicPath() << + " part id " << JoinSeq(',', ev->Get()->Record.GetPartitionIds()) << " query id " << ev->Get()->Record.GetQueryId() << " cookie " << ev->Cookie); + auto queryGroup = Metrics.Counters->GetSubgroup("query_id", ev->Get()->Record.GetQueryId()); + auto topicGroup = queryGroup->GetSubgroup("read_group", SanitizeLabel(ev->Get()->Record.GetSource().GetReadGroup())); topicGroup->GetCounter("StartSession", true)->Inc(); - NodesTracker.AddNode(ev->Sender.NodeId()); - TMaybe readOffset; - if (ev->Get()->Record.HasOffset()) { - readOffset = ev->Get()->Record.GetOffset(); - } + LWPROBE(StartSession, ev->Sender.ToString(), ev->Get()->Record.GetQueryId(), ev->Get()->Record.ByteSizeLong()); - ConsumerSessionKey key{ev->Sender, ev->Get()->Record.GetPartitionId()}; - auto it = Consumers.find(key); + NodesTracker.AddNode(ev->Sender.NodeId()); + auto it = Consumers.find(ev->Sender); if (it != Consumers.end()) { if (ev->Cookie <= it->second->Generation) { LOG_ROW_DISPATCHER_WARN("Consumer already exists, ignore StartSession"); @@ -728,142 +805,173 @@ void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev) { } LOG_ROW_DISPATCHER_WARN("Consumer already exists, new consumer with new generation (" << ev->Cookie << ", current " << it->second->Generation << "), remove old consumer, sender " << ev->Sender << ", topicPath " - << ev->Get()->Record.GetSource().GetTopicPath() <<" partitionId " << ev->Get()->Record.GetPartitionId() << " cookie " << ev->Cookie); - DeleteConsumer(key); + << ev->Get()->Record.GetSource().GetTopicPath() << " cookie " << ev->Cookie); + DeleteConsumer(ev->Sender); } const auto& source = ev->Get()->Record.GetSource(); - TActorId sessionActorId; - TopicSessionKey topicKey{source.GetEndpoint(), source.GetDatabase(), source.GetTopicPath(), ev->Get()->Record.GetPartitionId()}; - TopicSessionInfo& topicSessionInfo = TopicSessions[topicKey]; - Y_ENSURE(topicSessionInfo.Sessions.size() <= 1); + auto consumerInfo = MakeAtomicShared(ev->Sender, SelfId(), NextEventQueueId++, ev->Get()->Record, + NodesTracker.GetNodeConnected(ev->Sender.NodeId()), ev->Cookie); - auto consumerInfo = MakeAtomicShared(ev->Sender, SelfId(), NextEventQueueId++, ev->Get()->Record, TActorId(), NodesTracker.GetNodeConnected(ev->Sender.NodeId()), ev->Cookie); - Consumers[key] = consumerInfo; + Consumers[ev->Sender] = consumerInfo; ConsumersByEventQueueId[consumerInfo->EventQueueId] = consumerInfo; if (!CheckSession(consumerInfo, ev)) { return; } - if (topicSessionInfo.Sessions.empty()) { - LOG_ROW_DISPATCHER_DEBUG("Create new session, offset " << readOffset); - sessionActorId = ActorFactory->RegisterTopicSession( - source.GetTopicPath(), - source.GetEndpoint(), - source.GetDatabase(), - Config, - SelfId(), - CompileServiceActorId, - ev->Get()->Record.GetPartitionId(), - YqSharedResources->UserSpaceYdbDriver, - CreateCredentialsProviderFactoryForStructuredToken( - CredentialsFactory, - ev->Get()->Record.GetToken(), - source.GetAddBearerToToken()), - Counters, - PqGateway, - MaxSessionBufferSizeBytes - ); - SessionInfo& sessionInfo = topicSessionInfo.Sessions[sessionActorId]; - sessionInfo.Consumers[ev->Sender] = consumerInfo; - } else { - auto sessionIt = topicSessionInfo.Sessions.begin(); - SessionInfo& sessionInfo = sessionIt->second; - sessionInfo.Consumers[ev->Sender] = consumerInfo; - sessionActorId = sessionIt->first; + for (auto partitionId : ev->Get()->Record.GetPartitionIds()) { + TActorId sessionActorId; + TTopicSessionKey topicKey{source.GetReadGroup(), source.GetEndpoint(), source.GetDatabase(), source.GetTopicPath(), partitionId}; + TTopicSessionInfo& topicSessionInfo = TopicSessions[topicKey]; + Y_ENSURE(topicSessionInfo.Sessions.size() <= 1); + + if (topicSessionInfo.Sessions.empty()) { + LOG_ROW_DISPATCHER_DEBUG("Create new session: read group " << source.GetReadGroup() << " topic " << source.GetTopicPath() + << " part id " << partitionId); + sessionActorId = ActorFactory->RegisterTopicSession( + source.GetReadGroup(), + source.GetTopicPath(), + source.GetEndpoint(), + source.GetDatabase(), + Config, + SelfId(), + CompileServiceActorId, + partitionId, + YqSharedResources->UserSpaceYdbDriver, + CreateCredentialsProviderFactoryForStructuredToken( + CredentialsFactory, + ev->Get()->Record.GetToken(), + source.GetAddBearerToToken()), + Counters, + CountersRoot, + PqGateway, + MaxSessionBufferSizeBytes + ); + TSessionInfo& sessionInfo = topicSessionInfo.Sessions[sessionActorId]; + sessionInfo.Consumers[ev->Sender] = consumerInfo; + } else { + auto sessionIt = topicSessionInfo.Sessions.begin(); + TSessionInfo& sessionInfo = sessionIt->second; + sessionInfo.Consumers[ev->Sender] = consumerInfo; + sessionActorId = sessionIt->first; + } + consumerInfo->Partitions[partitionId].TopicSessionId = sessionActorId; + + auto event = std::make_unique(); + event->Record.CopyFrom(ev->Get()->Record); + Send(new IEventHandle(sessionActorId, ev->Sender, event.release(), 0)); } - consumerInfo->TopicSessionId = sessionActorId; consumerInfo->EventsQueue.Send(new NFq::TEvRowDispatcher::TEvStartSessionAck(consumerInfo->Proto), consumerInfo->Generation); - - Forward(ev, sessionActorId); Metrics.ClientsCount->Set(Consumers.size()); } void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvGetNextBatch::TPtr& ev) { - ConsumerSessionKey key{ev->Sender, ev->Get()->Record.GetPartitionId()}; - auto it = Consumers.find(key); + auto it = Consumers.find(ev->Sender); if (it == Consumers.end()) { LOG_ROW_DISPATCHER_WARN("Ignore (no consumer) TEvGetNextBatch from " << ev->Sender << " part id " << ev->Get()->Record.GetPartitionId()); return; } + auto& session = it->second; + LWPROBE(GetNextBatch, ev->Sender.ToString(), ev->Get()->Record.GetPartitionId(), session->QueryId, ev->Get()->Record.ByteSizeLong()); LOG_ROW_DISPATCHER_TRACE("Received TEvGetNextBatch from " << ev->Sender << " part id " << ev->Get()->Record.GetPartitionId() << " query id " << it->second->QueryId); - if (!CheckSession(it->second, ev)) { + if (!CheckSession(session, ev)) { return; } - it->second->PendingNewDataArrived = false; - it->second->PendingGetNextBatch = true; - it->second->Counters.GetNextBatch++; - Forward(ev, it->second->TopicSessionId); + auto partitionIt = session->Partitions.find(ev->Get()->Record.GetPartitionId()); + if (partitionIt == session->Partitions.end()) { + LOG_ROW_DISPATCHER_ERROR("Ignore TEvGetNextBatch from " << ev->Sender << ", wrong partition id " << ev->Get()->Record.GetPartitionId()); + return; + } + partitionIt->second.PendingNewDataArrived = false; + partitionIt->second.PendingGetNextBatch = true; + session->Counters.GetNextBatch++; + Forward(ev, partitionIt->second.TopicSessionId); } void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvHeartbeat::TPtr& ev) { - ConsumerSessionKey key{ev->Sender, ev->Get()->Record.GetPartitionId()}; - auto it = Consumers.find(key); + auto it = Consumers.find(ev->Sender); if (it == Consumers.end()) { LOG_ROW_DISPATCHER_WARN("Wrong consumer, sender " << ev->Sender << ", part id " << ev->Get()->Record.GetPartitionId()); return; } + LWPROBE(Heartbeat, ev->Sender.ToString(), ev->Get()->Record.GetPartitionId(), it->second->QueryId, ev->Get()->Record.ByteSizeLong()); LOG_ROW_DISPATCHER_TRACE("Received TEvHeartbeat from " << ev->Sender << ", part id " << ev->Get()->Record.GetPartitionId() << " query id " << it->second->QueryId); CheckSession(it->second, ev); } +void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvNoSession::TPtr& ev) { + LOG_ROW_DISPATCHER_DEBUG("Received TEvNoSession from " << ev->Sender << ", generation " << ev->Cookie); + auto consumerIt = Consumers.find(ev->Sender); + if (consumerIt == Consumers.end()) { + return; + } + const auto& consumer = consumerIt->second; + if (consumer->Generation != ev->Cookie) { + return; + } + DeleteConsumer(ev->Sender); +} + template -bool TRowDispatcher::CheckSession(TAtomicSharedPtr& consumer, const TEventPtr& ev) { +bool TRowDispatcher::CheckSession(TAtomicSharedPtr& consumer, const TEventPtr& ev) { if (ev->Cookie != consumer->Generation) { LOG_ROW_DISPATCHER_WARN("Wrong message generation (" << typeid(TEventPtr).name() << "), sender " << ev->Sender << " cookie " << ev->Cookie << ", session generation " << consumer->Generation << ", query id " << consumer->QueryId); return false; } if (!consumer->EventsQueue.OnEventReceived(ev)) { const NYql::NDqProto::TMessageTransportMeta& meta = ev->Get()->Record.GetTransportMeta(); - LOG_ROW_DISPATCHER_WARN("Wrong seq num ignore message (" << typeid(TEventPtr).name() << ") seqNo " << meta.GetSeqNo() << " from " << ev->Sender.ToString() << ", query id " << consumer->QueryId); + LOG_ROW_DISPATCHER_WARN("Wrong seq num, ignore message (" << typeid(TEventPtr).name() << ") seqNo " << meta.GetSeqNo() << " from " << ev->Sender.ToString() << ", query id " << consumer->QueryId); return false; } return true; } void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvStopSession::TPtr& ev) { - ConsumerSessionKey key{ev->Sender, ev->Get()->Record.GetPartitionId()}; - auto it = Consumers.find(key); + auto it = Consumers.find(ev->Sender); if (it == Consumers.end()) { - LOG_ROW_DISPATCHER_WARN("Ignore TEvStopSession from " << ev->Sender << " part id " << ev->Get()->Record.GetPartitionId()); + LOG_ROW_DISPATCHER_WARN("Ignore TEvStopSession from " << ev->Sender); return; } - LOG_ROW_DISPATCHER_DEBUG("Received TEvStopSession, topicPath " << ev->Get()->Record.GetSource().GetTopicPath() << - " partitionId " << ev->Get()->Record.GetPartitionId() << " query id " << it->second->QueryId); + + LWPROBE(StopSession, ev->Sender.ToString(), it->second->QueryId, ev->Get()->Record.ByteSizeLong()); + LOG_ROW_DISPATCHER_DEBUG("Received TEvStopSession from " << ev->Sender << " topic " << ev->Get()->Record.GetSource().GetTopicPath() << " query id " << it->second->QueryId); if (!CheckSession(it->second, ev)) { return; } - DeleteConsumer(key); + DeleteConsumer(ev->Sender); } -void TRowDispatcher::DeleteConsumer(const ConsumerSessionKey& key) { - auto consumerIt = Consumers.find(key); +void TRowDispatcher::DeleteConsumer(NActors::TActorId readActorId) { + auto consumerIt = Consumers.find(readActorId); if (consumerIt == Consumers.end()) { - LOG_ROW_DISPATCHER_ERROR("Ignore (no consumer) DeleteConsumer, " << " read actor id " << key.ReadActorId << " part id " << key.PartitionId); + LOG_ROW_DISPATCHER_ERROR("Ignore (no consumer) DeleteConsumer, " << " read actor id " << readActorId); return; } const auto& consumer = consumerIt->second; - LOG_ROW_DISPATCHER_DEBUG("DeleteConsumer, readActorId " << key.ReadActorId << " partitionId " << key.PartitionId << " query id " << consumer->QueryId); - auto event = std::make_unique(); - *event->Record.MutableSource() = consumer->SourceParams; - event->Record.SetPartitionId(consumer->PartitionId); - Send(new IEventHandle(consumerIt->second->TopicSessionId, consumer->ReadActorId, event.release(), 0)); - - TopicSessionKey topicKey{ - consumer->SourceParams.GetEndpoint(), - consumer->SourceParams.GetDatabase(), - consumer->SourceParams.GetTopicPath(), - consumer->PartitionId}; - TopicSessionInfo& topicSessionInfo = TopicSessions[topicKey]; - SessionInfo& sessionInfo = topicSessionInfo.Sessions[consumerIt->second->TopicSessionId]; - Y_ENSURE(sessionInfo.Consumers.count(consumer->ReadActorId)); - sessionInfo.Consumers.erase(consumer->ReadActorId); - if (sessionInfo.Consumers.empty()) { - LOG_ROW_DISPATCHER_DEBUG("Session is not used, sent TEvPoisonPill to " << consumerIt->second->TopicSessionId); - topicSessionInfo.Sessions.erase(consumerIt->second->TopicSessionId); - Send(consumerIt->second->TopicSessionId, new NActors::TEvents::TEvPoisonPill()); - if (topicSessionInfo.Sessions.empty()) { - TopicSessions.erase(topicKey); + LOG_ROW_DISPATCHER_DEBUG("DeleteConsumer, readActorId " << readActorId << " query id " << consumer->QueryId); + for (auto& [partitionId, partition] : consumer->Partitions) { + auto event = std::make_unique(); + *event->Record.MutableSource() = consumer->SourceParams; + Send(new IEventHandle(partition.TopicSessionId, consumer->ReadActorId, event.release(), 0)); + + TTopicSessionKey topicKey{ + consumer->SourceParams.GetReadGroup(), + consumer->SourceParams.GetEndpoint(), + consumer->SourceParams.GetDatabase(), + consumer->SourceParams.GetTopicPath(), + partitionId}; + TTopicSessionInfo& topicSessionInfo = TopicSessions[topicKey]; + TSessionInfo& sessionInfo = topicSessionInfo.Sessions[partition.TopicSessionId]; + if (!sessionInfo.Consumers.erase(consumer->ReadActorId)) { + LOG_ROW_DISPATCHER_ERROR("Wrong readActorId " << consumer->ReadActorId << ", no such consumer"); + } + if (sessionInfo.Consumers.empty()) { + LOG_ROW_DISPATCHER_DEBUG("Session is not used, sent TEvPoisonPill to " << partition.TopicSessionId); + topicSessionInfo.Sessions.erase(partition.TopicSessionId); + Send(partition.TopicSessionId, new NActors::TEvents::TEvPoisonPill()); + if (topicSessionInfo.Sessions.empty()) { + TopicSessions.erase(topicKey); + } } } ConsumersByEventQueueId.erase(consumerIt->second->EventQueueId); @@ -872,6 +980,7 @@ void TRowDispatcher::DeleteConsumer(const ConsumerSessionKey& key) { } void TRowDispatcher::Handle(const TEvPrivate::TEvTryConnect::TPtr& ev) { + LWPROBE(TryConnect, ev->Sender.ToString(), ev->Get()->NodeId); LOG_ROW_DISPATCHER_TRACE("TEvTryConnect to node id " << ev->Get()->NodeId); NodesTracker.TryConnect(ev->Get()->NodeId); } @@ -879,76 +988,80 @@ void TRowDispatcher::Handle(const TEvPrivate::TEvTryConnect::TPtr& ev) { void TRowDispatcher::Handle(const NYql::NDq::TEvRetryQueuePrivate::TEvEvHeartbeat::TPtr& ev) { auto it = ConsumersByEventQueueId.find(ev->Get()->EventQueueId); if (it == ConsumersByEventQueueId.end()) { - LOG_ROW_DISPATCHER_WARN("No consumer with EventQueueId = " << ev->Get()->EventQueueId); + LOG_ROW_DISPATCHER_TRACE("No consumer with EventQueueId = " << ev->Get()->EventQueueId); return; } auto& sessionInfo = it->second; + LWPROBE(PrivateHeartbeat, ev->Sender.ToString(), sessionInfo->QueryId, sessionInfo->Generation); bool needSend = sessionInfo->EventsQueue.Heartbeat(); if (needSend) { LOG_ROW_DISPATCHER_TRACE("Send TEvHeartbeat to " << sessionInfo->ReadActorId << " query id " << sessionInfo->QueryId); - sessionInfo->EventsQueue.Send(new NFq::TEvRowDispatcher::TEvHeartbeat(sessionInfo->PartitionId), sessionInfo->Generation); + auto event = std::make_unique(); + sessionInfo->EventsQueue.Send(new NFq::TEvRowDispatcher::TEvHeartbeat(), sessionInfo->Generation); } } void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvNewDataArrived::TPtr& ev) { - ConsumerSessionKey key{ev->Get()->ReadActorId, ev->Get()->Record.GetPartitionId()}; - auto it = Consumers.find(key); + auto it = Consumers.find(ev->Get()->ReadActorId); if (it == Consumers.end()) { LOG_ROW_DISPATCHER_WARN("Ignore (no consumer) TEvNewDataArrived from " << ev->Sender << " part id " << ev->Get()->Record.GetPartitionId()); return; } - LOG_ROW_DISPATCHER_TRACE("Forward TEvNewDataArrived from " << ev->Sender << " to " << ev->Get()->ReadActorId << " query id " << it->second->QueryId); - it->second->PendingNewDataArrived = true; - it->second->Counters.NewDataArrived++; - it->second->EventsQueue.Send(ev->Release().Release(), it->second->Generation); + auto consumerInfoPtr = it->second; + LWPROBE(NewDataArrived, ev->Sender.ToString(), ev->Get()->ReadActorId.ToString(), consumerInfoPtr->QueryId, consumerInfoPtr->Generation, ev->Get()->Record.ByteSizeLong()); + LOG_ROW_DISPATCHER_TRACE("Forward TEvNewDataArrived from " << ev->Sender << " to " << ev->Get()->ReadActorId << " query id " << consumerInfoPtr->QueryId); + auto partitionIt = consumerInfoPtr->Partitions.find(ev->Get()->Record.GetPartitionId()); + if (partitionIt == consumerInfoPtr->Partitions.end()) { + // Ignore TEvNewDataArrived because read actor now read others partitions. + return; + } + partitionIt->second.PendingNewDataArrived = true; + consumerInfoPtr->Counters.NewDataArrived++; + consumerInfoPtr->EventsQueue.Send(ev->Release().Release(), it->second->Generation); } void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvMessageBatch::TPtr& ev) { - ConsumerSessionKey key{ev->Get()->ReadActorId, ev->Get()->Record.GetPartitionId()}; - auto it = Consumers.find(key); + auto it = Consumers.find(ev->Get()->ReadActorId); if (it == Consumers.end()) { - LOG_ROW_DISPATCHER_WARN("Ignore (no consumer) TEvMessageBatch from " << ev->Sender); + LOG_ROW_DISPATCHER_WARN("Ignore (no consumer) TEvMessageBatch from " << ev->Sender << " to " << ev->Get()->ReadActorId); return; } - LOG_ROW_DISPATCHER_TRACE("Forward TEvMessageBatch from " << ev->Sender << " to " << ev->Get()->ReadActorId << " query id " << it->second->QueryId); + auto consumerInfoPtr = it->second; + LWPROBE(MessageBatch, ev->Sender.ToString(), ev->Get()->ReadActorId.ToString(), consumerInfoPtr->QueryId, consumerInfoPtr->Generation, ev->Get()->Record.ByteSizeLong()); + LOG_ROW_DISPATCHER_TRACE("Forward TEvMessageBatch from " << ev->Sender << " to " << ev->Get()->ReadActorId << " query id " << consumerInfoPtr->QueryId); Metrics.RowsSent->Add(ev->Get()->Record.MessagesSize()); - it->second->PendingGetNextBatch = false; - it->second->Counters.MessageBatch++; - it->second->EventsQueue.Send(ev->Release().Release(), it->second->Generation); -} - -void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvSessionError::TPtr& ev) { - ConsumerSessionKey key{ev->Get()->ReadActorId, ev->Get()->Record.GetPartitionId()}; - auto it = Consumers.find(key); - if (it == Consumers.end()) { - LOG_ROW_DISPATCHER_WARN("Ignore (no consumer) TEvSessionError from " << ev->Sender << " to " << ev->Get()->ReadActorId << " part id " << ev->Get()->Record.GetPartitionId()); + auto partitionIt = consumerInfoPtr->Partitions.find(ev->Get()->Record.GetPartitionId()); + if (partitionIt == consumerInfoPtr->Partitions.end()) { + // Ignore TEvMessageBatch because read actor now read others partitions. return; } - ++*Metrics.ErrorsCount; - LOG_ROW_DISPATCHER_TRACE("Forward TEvSessionError from " << ev->Sender << " to " << ev->Get()->ReadActorId << " part id " << ev->Get()->Record.GetPartitionId() << " query id " << it->second->QueryId); - it->second->EventsQueue.Send(ev->Release().Release(), it->second->Generation); - DeleteConsumer(key); + partitionIt->second.PendingGetNextBatch = false; + consumerInfoPtr->Counters.MessageBatch++; + consumerInfoPtr->EventsQueue.Send(ev->Release().Release(), it->second->Generation); } -void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvStatistics::TPtr& ev) { - LOG_ROW_DISPATCHER_TRACE("TEvStatistics from " << ev->Sender); - ConsumerSessionKey key{ev->Get()->ReadActorId, ev->Get()->Record.GetPartitionId()}; - auto it = Consumers.find(key); +void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvSessionError::TPtr& ev) { + auto it = Consumers.find(ev->Get()->ReadActorId); if (it == Consumers.end()) { - LOG_ROW_DISPATCHER_WARN("Ignore (no consumer) TEvStatistics from " << ev->Sender << " to " << ev->Get()->ReadActorId << " part id " << ev->Get()->Record.GetPartitionId()); + LOG_ROW_DISPATCHER_WARN("Ignore (no consumer) TEvSessionError from " << ev->Sender << " to " << ev->Get()->ReadActorId); return; } - LOG_ROW_DISPATCHER_TRACE("Forward TEvStatus from " << ev->Sender << " to " << ev->Get()->ReadActorId << " part id " << ev->Get()->Record.GetPartitionId() << " query id " << it->second->QueryId); + LWPROBE(SessionError, ev->Sender.ToString(), ev->Get()->ReadActorId.ToString(), it->second->QueryId, it->second->Generation, ev->Get()->Record.ByteSizeLong()); + ++*Metrics.ErrorsCount; + LOG_ROW_DISPATCHER_TRACE("Forward TEvSessionError from " << ev->Sender << " to " << ev->Get()->ReadActorId << " query id " << it->second->QueryId); it->second->EventsQueue.Send(ev->Release().Release(), it->second->Generation); + DeleteConsumer(ev->Get()->ReadActorId); } void TRowDispatcher::Handle(NFq::TEvPrivate::TEvUpdateMetrics::TPtr&) { + LWPROBE(UpdateMetrics); Schedule(TDuration::Seconds(UpdateMetricsPeriodSec), new NFq::TEvPrivate::TEvUpdateMetrics()); UpdateMetrics(); } void TRowDispatcher::Handle(NFq::TEvPrivate::TEvPrintStateToLog::TPtr&) { + LWPROBE(PrintStateToLog, PrintStateToLogPeriodSec); PrintStateToLog(); Schedule(TDuration::Seconds(PrintStateToLogPeriodSec), new NFq::TEvPrivate::TEvPrintStateToLog()); } @@ -956,12 +1069,50 @@ void TRowDispatcher::Handle(NFq::TEvPrivate::TEvPrintStateToLog::TPtr&) { void TRowDispatcher::PrintStateToLog() { auto str = GetInternalState(); auto buf = TStringBuf(str); - i64 size = buf.size(); - ui64 offset = 0; - while (size > 0) { + for (ui64 offset = 0; offset < buf.size(); offset += PrintStateToLogSplitSize) { LOG_ROW_DISPATCHER_DEBUG(buf.SubString(offset, PrintStateToLogSplitSize)); - offset += PrintStateToLogSplitSize; - size -= PrintStateToLogSplitSize; + } +} + +void TRowDispatcher::Handle(NFq::TEvPrivate::TEvSendStatistic::TPtr&) { + LOG_ROW_DISPATCHER_TRACE("TEvPrivate::TEvSendStatistic"); + + UpdateCpuTime(); + Schedule(TDuration::Seconds(Config.GetSendStatusPeriodSec()), new NFq::TEvPrivate::TEvSendStatistic()); + for (auto& [actorId, consumer] : Consumers) { + if (!NodesTracker.GetNodeConnected(actorId.NodeId())) { + continue; // Wait Connected to prevent retry_queue increases. + } + auto event = std::make_unique(); + ui64 readBytes = 0; + ui64 filteredBytes = 0; + ui64 filteredRows = 0; + ui64 queuedBytes = 0; + ui64 queuedRows = 0; + for (auto& [partitionId, partition] : consumer->Partitions) { + if (!partition.StatisticsUpdated) { + continue; + } + auto* partitionsProto = event->Record.AddPartition(); + partitionsProto->SetPartitionId(partitionId); + partitionsProto->SetNextMessageOffset(partition.Stat.Offset); + readBytes += partition.Stat.ReadBytes; + filteredBytes += partition.Stat.FilteredBytes; + filteredRows += partition.Stat.FilteredRows; + queuedBytes += partition.Stat.QueuedBytes; + queuedRows += partition.Stat.QueuedRows; + partition.Stat.Clear(); + partition.StatisticsUpdated = false; + } + event->Record.SetReadBytes(readBytes); + event->Record.SetCpuMicrosec(consumer->CpuMicrosec); + consumer->CpuMicrosec = 0; + event->Record.SetFilteredBytes(filteredBytes); + event->Record.SetFilteredRows(filteredRows); + event->Record.SetQueuedBytes(queuedBytes); + event->Record.SetQueuedRows(queuedRows); + LWPROBE(Statistics, consumer->ReadActorId.ToString(), consumer->QueryId, consumer->Generation, event->Record.ByteSizeLong()); + consumer->EventsQueue.Send(event.release(), consumer->Generation); } } @@ -983,8 +1134,22 @@ void TRowDispatcher::Handle(const NMon::TEvHttpInfo::TPtr& ev) { void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvSessionStatistic::TPtr& ev) { LOG_ROW_DISPATCHER_TRACE("TEvSessionStatistic from " << ev->Sender); - const auto& key = ev->Get()->Stat.SessionKey; - TopicSessionKey sessionKey{key.Endpoint, key.Database, key.TopicPath, key.PartitionId}; + const auto& stat = ev->Get()->Stat; + const auto& key = stat.SessionKey; + + LWPROBE(SessionStatistic, + ev->Sender.ToString(), + key.ReadGroup, + key.Endpoint, + key.Database, + key.TopicPath, + key.PartitionId, + stat.Common.ReadBytes, + stat.Common.QueuedBytes, + stat.Common.RestartSessionByOffsets, + stat.Common.ReadEvents, + stat.Common.LastReadedOffset); + TTopicSessionKey sessionKey{key.ReadGroup, key.Endpoint, key.Database, key.TopicPath, key.PartitionId}; auto sessionsIt = TopicSessions.find(sessionKey); if (sessionsIt == TopicSessions.end()) { @@ -1004,16 +1169,39 @@ void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvSessionStatistic::TPtr& ev continue; } auto consumerInfoPtr = it->second; - consumerInfoPtr->Stat.Add(clientStat); + auto partitionIt = consumerInfoPtr->Partitions.find(key.PartitionId); + if (partitionIt == consumerInfoPtr->Partitions.end()) { + continue; + } + partitionIt->second.Stat.Add(clientStat); + partitionIt->second.FilteredBytes += clientStat.FilteredBytes; + partitionIt->second.StatisticsUpdated = true; } } void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvGetInternalStateResponse::TPtr& ev) { + LWPROBE(GetInternalState, ev->Sender.ToString(), ev->Get()->Record.ByteSizeLong()); auto& readActorInternalState = ReadActorsInternalState[ev->Sender]; readActorInternalState.InternalState = ev->Get()->Record.GetInternalState(); readActorInternalState.ResponseTime = TInstant::Now(); } +void TRowDispatcher::UpdateCpuTime() { + if (Consumers.empty()) { + return; + } + auto currentCpuTime = UserPoolMetrics.Session->Val() + + UserPoolMetrics.RowDispatcher->Val() + + UserPoolMetrics.CompilerActor->Val() + + UserPoolMetrics.CompilerService->Val() + + UserPoolMetrics.FormatHandler->Val(); + auto diff = (currentCpuTime - LastCpuTime) / Consumers.size(); + for (auto& [actorId, consumer] : Consumers) { + consumer->CpuMicrosec += diff; + } + LastCpuTime = currentCpuTime; +} + } // namespace //////////////////////////////////////////////////////////////////////////////// @@ -1026,6 +1214,7 @@ std::unique_ptr NewRowDispatcher( const TString& tenant, const NFq::NRowDispatcher::IActorFactory::TPtr& actorFactory, const ::NMonitoring::TDynamicCounterPtr& counters, + const ::NMonitoring::TDynamicCounterPtr& countersRoot, const NYql::IPqGateway::TPtr& pqGateway, NActors::TMon* monitoring) { @@ -1037,6 +1226,7 @@ std::unique_ptr NewRowDispatcher( tenant, actorFactory, counters, + countersRoot, pqGateway, monitoring)); } diff --git a/ydb/core/fq/libs/row_dispatcher/row_dispatcher.h b/ydb/core/fq/libs/row_dispatcher/row_dispatcher.h index f956ac3d8b1b..947afcbb798f 100644 --- a/ydb/core/fq/libs/row_dispatcher/row_dispatcher.h +++ b/ydb/core/fq/libs/row_dispatcher/row_dispatcher.h @@ -27,6 +27,7 @@ std::unique_ptr NewRowDispatcher( const TString& tenant, const NFq::NRowDispatcher::IActorFactory::TPtr& actorFactory, const ::NMonitoring::TDynamicCounterPtr& counters, + const ::NMonitoring::TDynamicCounterPtr& countersRoot, const NYql::IPqGateway::TPtr& pqGateway, NActors::TMon* monitoring = nullptr); diff --git a/ydb/core/fq/libs/row_dispatcher/row_dispatcher_service.cpp b/ydb/core/fq/libs/row_dispatcher/row_dispatcher_service.cpp index 65a63bdea2a8..bc406e6f6294 100644 --- a/ydb/core/fq/libs/row_dispatcher/row_dispatcher_service.cpp +++ b/ydb/core/fq/libs/row_dispatcher/row_dispatcher_service.cpp @@ -17,7 +17,8 @@ std::unique_ptr NewRowDispatcherService( const TString& tenant, const ::NMonitoring::TDynamicCounterPtr& counters, const NYql::IPqGateway::TPtr& pqGateway, - NActors::TMon* monitoring) + NActors::TMon* monitoring, + ::NMonitoring::TDynamicCounterPtr countersRoot) { return NewRowDispatcher( config, @@ -27,6 +28,7 @@ std::unique_ptr NewRowDispatcherService( tenant, NFq::NRowDispatcher::CreateActorFactory(), counters, + countersRoot, pqGateway, monitoring); } diff --git a/ydb/core/fq/libs/row_dispatcher/row_dispatcher_service.h b/ydb/core/fq/libs/row_dispatcher/row_dispatcher_service.h index 055240eb2ede..ebcce4b1e8b2 100644 --- a/ydb/core/fq/libs/row_dispatcher/row_dispatcher_service.h +++ b/ydb/core/fq/libs/row_dispatcher/row_dispatcher_service.h @@ -27,6 +27,7 @@ std::unique_ptr NewRowDispatcherService( const TString& tenant, const ::NMonitoring::TDynamicCounterPtr& counters, const NYql::IPqGateway::TPtr& pqGateway, - NActors::TMon* monitoring = nullptr); + NActors::TMon* monitoring = nullptr, + ::NMonitoring::TDynamicCounterPtr countersRoot = MakeIntrusive<::NMonitoring::TDynamicCounters>()); } // namespace NFq diff --git a/ydb/core/fq/libs/row_dispatcher/topic_session.cpp b/ydb/core/fq/libs/row_dispatcher/topic_session.cpp index e1138bf9de72..976937172a0d 100644 --- a/ydb/core/fq/libs/row_dispatcher/topic_session.cpp +++ b/ydb/core/fq/libs/row_dispatcher/topic_session.cpp @@ -1,115 +1,82 @@ #include "topic_session.h" -#include "common.h" - #include - +#include #include -#include +#include + #include #include -#include -#include -#include +#include + #include -#include #include -#include -#include -#include - namespace NFq { using namespace NActors; +using namespace NRowDispatcher; namespace { //////////////////////////////////////////////////////////////////////////////// struct TTopicSessionMetrics { - void Init(const ::NMonitoring::TDynamicCounterPtr& counters, const TString& topicPath, ui32 partitionId) { + void Init(const ::NMonitoring::TDynamicCounterPtr& counters, const TString& topicPath, const TString& readGroup, ui32 partitionId) { + TopicGroup = counters->GetSubgroup("topic", SanitizeLabel(topicPath)); + ReadGroup = TopicGroup->GetSubgroup("read_group", SanitizeLabel(readGroup)); + PartitionGroup = ReadGroup->GetSubgroup("partition", ToString(partitionId)); - TopicGroup = counters->GetSubgroup("topic", CleanupCounterValueString(topicPath)); - PartitionGroup = TopicGroup->GetSubgroup("partition", ToString(partitionId)); + AllSessionsDataRate = ReadGroup->GetCounter("AllSessionsDataRate", true); InFlyAsyncInputData = PartitionGroup->GetCounter("InFlyAsyncInputData"); InFlySubscribe = PartitionGroup->GetCounter("InFlySubscribe"); ReconnectRate = PartitionGroup->GetCounter("ReconnectRate", true); - RestartSessionByOffsets = counters->GetCounter("RestartSessionByOffsets", true); + RestartSessionByOffsets = PartitionGroup->GetCounter("RestartSessionByOffsets", true); SessionDataRate = PartitionGroup->GetCounter("SessionDataRate", true); - WaitEventTimeMs = PartitionGroup->GetHistogram("WaitEventTimeMs", NMonitoring::ExponentialHistogram(13, 2, 1)); // ~ 1ms -> ~ 8s - AllSessionsDataRate = counters->GetCounter("AllSessionsDataRate", true); - InFlightCompileRequests = PartitionGroup->GetCounter("InFlightCompileRequests"); + WaitEventTimeMs = PartitionGroup->GetHistogram("WaitEventTimeMs", NMonitoring::ExplicitHistogram({5, 20, 100, 500, 2000})); + QueuedBytes = PartitionGroup->GetCounter("QueuedBytes"); } ::NMonitoring::TDynamicCounterPtr TopicGroup; + ::NMonitoring::TDynamicCounterPtr ReadGroup; ::NMonitoring::TDynamicCounterPtr PartitionGroup; ::NMonitoring::TDynamicCounters::TCounterPtr InFlyAsyncInputData; ::NMonitoring::TDynamicCounters::TCounterPtr InFlySubscribe; ::NMonitoring::TDynamicCounters::TCounterPtr ReconnectRate; ::NMonitoring::TDynamicCounters::TCounterPtr RestartSessionByOffsets; ::NMonitoring::TDynamicCounters::TCounterPtr SessionDataRate; - ::NMonitoring::TDynamicCounters::TCounterPtr InFlightCompileRequests; ::NMonitoring::THistogramPtr WaitEventTimeMs; ::NMonitoring::TDynamicCounters::TCounterPtr AllSessionsDataRate; - + ::NMonitoring::TDynamicCounters::TCounterPtr QueuedBytes; }; struct TEvPrivate { // Event ids enum EEv : ui32 { - EvBegin = EventSpaceBegin(NActors::TEvents::ES_PRIVATE), + EvBegin = EventSpaceBegin(TEvents::ES_PRIVATE), EvPqEventsReady = EvBegin + 10, EvCreateSession, - EvSendStatisticToReadActor, - EvDataAfterFilteration, - EvDataFiltered, - EvSendStatisticToRowDispatcher, - EvStartParsing, + EvSendStatistic, EvReconnectSession, EvEnd }; - static_assert(EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE)"); + static_assert(EvEnd < EventSpaceEnd(TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(TEvents::ES_PRIVATE)"); // Events - struct TEvPqEventsReady : public NActors::TEventLocal {}; - struct TEvCreateSession : public NActors::TEventLocal {}; - struct TEvSendStatisticToRowDispatcher : public NActors::TEventLocal {}; - struct TEvSendStatisticToReadActor : public NActors::TEventLocal {}; - struct TEvStartParsing : public NActors::TEventLocal {}; - struct TEvReconnectSession : public NActors::TEventLocal {}; - - struct TEvDataFiltered : public NActors::TEventLocal { - explicit TEvDataFiltered(ui64 offset) - : Offset(offset) - {} - const ui64 Offset; - }; - - struct TEvDataAfterFilteration : public NActors::TEventLocal { - TEvDataAfterFilteration(ui64 offset, const TString& json, TActorId readActorId) - : Offset(offset) - , Json(json) - , ReadActorId(readActorId) { } - ui64 Offset; - TString Json; - TActorId ReadActorId; - }; + struct TEvPqEventsReady : public TEventLocal {}; + struct TEvCreateSession : public TEventLocal {}; + struct TEvSendStatistic : public TEventLocal {}; + struct TEvReconnectSession : public TEventLocal {}; }; -ui64 SendStatisticPeriodSec = 2; -ui64 MaxBatchSizeBytes = 10000000; -ui64 MaxHandledEventsCount = 1000; -ui64 MaxHandledEventsSize = 1000000; - -TVector GetVector(const google::protobuf::RepeatedPtrField& value) { - return {value.begin(), value.end()}; -} +constexpr ui64 SendStatisticPeriodSec = 2; +constexpr ui64 MaxHandledEventsCount = 1000; +constexpr ui64 MaxHandledEventsSize = 1000000; class TTopicSession : public TActorBootstrapped { private: - using TParserInputType = TSet>; + using TBase = TActorBootstrapped; struct TStats { void Add(ui64 dataSize, ui64 events) { @@ -119,43 +86,120 @@ class TTopicSession : public TActorBootstrapped { void Clear() { Bytes = 0; Events = 0; - ParseAndFilterLatency = TDuration::Zero(); } ui64 Bytes = 0; ui64 Events = 0; - TDuration ParseAndFilterLatency; }; - struct TClientsInfo { - TClientsInfo( - const NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev, - NMonitoring::TDynamicCounterPtr& counters) - : Settings(ev->Get()->Record) + struct TClientsInfo : public IClientDataConsumer { + using TPtr = TIntrusivePtr; + + TClientsInfo(TTopicSession& self, const TString& logPrefix, const ITopicFormatHandler::TSettings& handlerSettings, const NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev, const NMonitoring::TDynamicCounterPtr& counters, const TString& readGroup, TMaybe offset) + : Self(self) + , LogPrefix(logPrefix) + , HandlerSettings(handlerSettings) + , Settings(ev->Get()->Record) , ReadActorId(ev->Sender) - , FilteredDataRate(counters->GetCounter("FilteredDataRate", true)) - , RestartSessionByOffsetsByQuery(counters->GetCounter("RestartSessionByOffsetsByQuery", true)) + , Counters(counters) { - if (Settings.HasOffset()) { - NextMessageOffset = Settings.GetOffset(); - InitialOffset = Settings.GetOffset(); + if (offset) { + NextMessageOffset = *offset; + InitialOffset = *offset; } Y_UNUSED(TDuration::TryParse(Settings.GetSource().GetReconnectPeriod(), ReconnectPeriod)); + auto queryGroup = Counters->GetSubgroup("query_id", ev->Get()->Record.GetQueryId()); + auto readSubGroup = queryGroup->GetSubgroup("read_group", SanitizeLabel(readGroup)); + FilteredDataRate = readSubGroup->GetCounter("FilteredDataRate", true); + RestartSessionByOffsetsByQuery = readSubGroup->GetCounter("RestartSessionByOffsetsByQuery", true); } - NFq::NRowDispatcherProto::TEvStartSession Settings; - NActors::TActorId ReadActorId; - std::unique_ptr Filter; // empty if no predicate - ui64 InFlightCompilationId = 0; - TQueue> Buffer; - ui64 UnreadBytes = 0; - bool DataArrivedSent = false; - TMaybe NextMessageOffset; - ui64 LastSendedNextMessageOffset = 0; - TVector FieldsIds; + + ~TClientsInfo() { + Counters->RemoveSubgroup("query_id", Settings.GetQueryId()); + } + + TActorId GetClientId() const override { + return ReadActorId; + } + + TMaybe GetNextMessageOffset() const override { + return NextMessageOffset; + } + + TVector GetColumns() const override { + const auto& source = Settings.GetSource(); + Y_ENSURE(source.ColumnsSize() == source.ColumnTypesSize(), "Columns size and types size should be equal, but got " << source.ColumnsSize() << " columns and " << source.ColumnTypesSize() << " types"); + + TVector Columns; + Columns.reserve(source.ColumnsSize()); + for (ui64 i = 0; i < source.ColumnsSize(); ++i) { + Columns.emplace_back(TSchemaColumn{.Name = source.GetColumns().Get(i), .TypeYson = source.GetColumnTypes().Get(i)}); + } + + return Columns; + } + + const TString& GetWhereFilter() const override { + return Settings.GetSource().GetPredicate(); + } + + TPurecalcCompileSettings GetPurecalcSettings() const override { + return {.EnabledLLVM = Settings.GetSource().GetEnabledLLVM()}; + } + + void OnClientError(TStatus status) override { + Self.SendSessionError(ReadActorId, status); + } + + void StartClientSession() override { + Self.StartClientSession(*this); + } + + void AddDataToClient(ui64 offset, ui64 rowSize) override { + Y_ENSURE(!NextMessageOffset || offset >= *NextMessageOffset, "Unexpected historical offset"); + + LOG_ROW_DISPATCHER_TRACE("AddDataToClient to " << ReadActorId << ", offset: " << offset << ", serialized size: " << rowSize); + + NextMessageOffset = offset + 1; + QueuedRows++; + QueuedBytes += rowSize; + Self.QueuedBytes += rowSize; + Self.SendDataArrived(*this); + Self.Metrics.QueuedBytes->Add(rowSize); + } + + void UpdateClientOffset(ui64 offset) override { + LOG_ROW_DISPATCHER_TRACE("UpdateClientOffset for " << ReadActorId << ", new offset: " << offset); + if (!NextMessageOffset || *NextMessageOffset < offset + 1) { + NextMessageOffset = offset + 1; + } + if (!QueuedRows) { + if (!ProcessedNextMessageOffset || *ProcessedNextMessageOffset < offset + 1) { + ProcessedNextMessageOffset = offset + 1; + } + } + } + + // Settings + TTopicSession& Self; + const TString& LogPrefix; + const ITopicFormatHandler::TSettings HandlerSettings; + const NFq::NRowDispatcherProto::TEvStartSession Settings; + const TActorId ReadActorId; TDuration ReconnectPeriod; - TStats Stat; // Send (filtered) to read_actor + + // State + ui64 QueuedRows = 0; + ui64 QueuedBytes = 0; + bool DataArrivedSent = false; + TMaybe NextMessageOffset; // offset to restart topic session + TMaybe ProcessedNextMessageOffset; // offset of fully processed data (to save to checkpoint) + + // Metrics + ui64 InitialOffset = 0; + TStats FilteredStat; + const ::NMonitoring::TDynamicCounterPtr Counters; NMonitoring::TDynamicCounters::TCounterPtr FilteredDataRate; // filtered NMonitoring::TDynamicCounters::TCounterPtr RestartSessionByOffsetsByQuery; - ui64 InitialOffset = 0; }; struct TTopicEventProcessor { @@ -170,71 +214,61 @@ class TTopicSession : public TActorBootstrapped { TTopicSession& Self; const TString& LogPrefix; - ui64& dataReceivedEventSize; - }; - - struct TParserSchema { - TVector FieldsMap; // index - FieldId (from FieldsIndexes), value - parsing schema offset - TParserInputType InputType; - }; - - struct TFieldDescription { - ui64 IndexInParserSchema = 0; - TString Type; + ui64& DataReceivedEventSize; }; - struct TFiltersCompileState { - ui64 FreeId = 0; - std::unordered_map InFlightCompilations = {}; - }; - - bool InflightReconnect = false; - TDuration ReconnectPeriod; + // Settings + const TString ReadGroup; const TString TopicPath; const TString TopicPathPartition; const TString Endpoint; const TString Database; - NActors::TActorId RowDispatcherActorId; - NActors::TActorId CompileServiceActorId; - ui32 PartitionId; - NYdb::TDriver Driver; - std::shared_ptr CredentialsProviderFactory; - NYql::ITopicClient::TPtr TopicClient; - std::shared_ptr ReadSession; + const TActorId RowDispatcherActorId; + const ui32 PartitionId; + const NYdb::TDriver Driver; + const NYql::IPqGateway::TPtr PqGateway; + const std::shared_ptr CredentialsProviderFactory; + const NConfig::TRowDispatcherConfig Config; + const TFormatHandlerConfig FormatHandlerConfig; const i64 BufferSize; TString LogPrefix; - TStats SessionStats; - TStats ClientsStats; + + // State + bool InflightReconnect = false; + TDuration ReconnectPeriod; + + NYql::ITopicClient::TPtr TopicClient; + std::shared_ptr ReadSession; + std::map FormatHandlers; + std::unordered_map Clients; + ui64 LastMessageOffset = 0; bool IsWaitingEvents = false; - bool IsStartParsingScheduled = false; - THashMap Clients; - THashSet ClientsWithoutPredicate; - std::unique_ptr Parser; - TFiltersCompileState FiltersCompilation; - NConfig::TRowDispatcherConfig Config; - ui64 UnreadBytes = 0; - const ::NMonitoring::TDynamicCounterPtr Counters; - TTopicSessionMetrics Metrics; - TParserSchema ParserSchema; - THashMap FieldsIndexes; - NYql::IPqGateway::TPtr PqGateway; + ui64 QueuedBytes = 0; TMaybe ConsumerName; - ui64 RestartSessionByOffsets = 0; + + // Metrics TInstant WaitEventStartedAt; + ui64 RestartSessionByOffsets = 0; + TStats Statistics; + TTopicSessionMetrics Metrics; + const ::NMonitoring::TDynamicCounterPtr Counters; + const ::NMonitoring::TDynamicCounterPtr CountersRoot; public: explicit TTopicSession( + const TString& readGroup, const TString& topicPath, const TString& endpoint, const TString& database, const NConfig::TRowDispatcherConfig& config, - NActors::TActorId rowDispatcherActorId, - NActors::TActorId compileServiceActorId, + TActorId rowDispatcherActorId, + TActorId compileServiceActorId, ui32 partitionId, NYdb::TDriver driver, std::shared_ptr credentialsProviderFactory, const ::NMonitoring::TDynamicCounterPtr& counters, + const ::NMonitoring::TDynamicCounterPtr& countersRoot, const NYql::IPqGateway::TPtr& pqGateway, ui64 maxBufferSize); @@ -251,121 +285,107 @@ class TTopicSession : public TActorBootstrapped { void CloseTopicSession(); void SubscribeOnNextEvent(); void SendToParsing(const TVector& messages); - void DoParsing(bool force = false); - void DoFiltering(ui64 rowsOffset, ui64 numberRows, const TVector>& parsedValues); void SendData(TClientsInfo& info); - void UpdateParser(); - void FatalError(const TString& message, const std::unique_ptr* filter, bool addParserDescription, const TMaybe& fieldName); + void FatalError(const TStatus& status); + void ThrowFatalError(const TStatus& status); void SendDataArrived(TClientsInfo& client); void StopReadSession(); TString GetSessionId() const; void HandleNewEvents(); TInstant GetMinStartingMessageTimestamp() const; - void AddDataToClient(TClientsInfo& client, ui64 offset, const TString& json); void StartClientSession(TClientsInfo& info); - std::pair CreateItem(const NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage& message); - void Handle(NFq::TEvPrivate::TEvPqEventsReady::TPtr&); void Handle(NFq::TEvPrivate::TEvCreateSession::TPtr&); void Handle(NFq::TEvPrivate::TEvReconnectSession::TPtr&); - void Handle(NFq::TEvPrivate::TEvDataAfterFilteration::TPtr&); - void Handle(NFq::TEvPrivate::TEvSendStatisticToReadActor::TPtr&); - void Handle(NFq::TEvPrivate::TEvDataFiltered::TPtr&); - void Handle(NFq::TEvPrivate::TEvSendStatisticToRowDispatcher::TPtr&); + void Handle(NFq::TEvPrivate::TEvSendStatistic::TPtr&); void Handle(TEvRowDispatcher::TEvGetNextBatch::TPtr&); void Handle(NFq::TEvRowDispatcher::TEvStopSession::TPtr& ev); void Handle(NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev); - void Handle(TEvRowDispatcher::TEvPurecalcCompileResponse::TPtr& ev); void HandleException(const std::exception& err); - void SendStatisticToRowDispatcher(); - void SendSessionError(NActors::TActorId readActorId, const TString& message); - TVector*> RebuildJson(const TClientsInfo& info, const TVector>& parsedValues); - void UpdateParserSchema(const TParserInputType& inputType); - void UpdateFieldsIds(TClientsInfo& clientInfo); + void SendStatistics(); bool CheckNewClient(NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev); - TString GetAnyQueryIdByFieldName(const TString& fieldName); + TMaybe GetOffset(const NFq::NRowDispatcherProto::TEvStartSession& settings); + void SendSessionError(TActorId readActorId, TStatus status); + void RestartSessionIfOldestClient(const TClientsInfo& info); + void RefreshParsers(); private: STRICT_STFUNC_EXC(StateFunc, hFunc(NFq::TEvPrivate::TEvPqEventsReady, Handle); hFunc(NFq::TEvPrivate::TEvCreateSession, Handle); - hFunc(NFq::TEvPrivate::TEvDataAfterFilteration, Handle); - hFunc(NFq::TEvPrivate::TEvSendStatisticToReadActor, Handle); - hFunc(NFq::TEvPrivate::TEvDataFiltered, Handle); - hFunc(NFq::TEvPrivate::TEvSendStatisticToRowDispatcher, Handle); + hFunc(NFq::TEvPrivate::TEvSendStatistic, Handle); hFunc(NFq::TEvPrivate::TEvReconnectSession, Handle); hFunc(TEvRowDispatcher::TEvGetNextBatch, Handle); hFunc(NFq::TEvRowDispatcher::TEvStartSession, Handle); - hFunc(TEvRowDispatcher::TEvPurecalcCompileResponse, Handle); - sFunc(NFq::TEvPrivate::TEvStartParsing, DoParsing); - cFunc(NActors::TEvents::TEvPoisonPill::EventType, PassAway); + cFunc(TEvents::TEvPoisonPill::EventType, PassAway); hFunc(NFq::TEvRowDispatcher::TEvStopSession, Handle);, ExceptionFunc(std::exception, HandleException) ) - STRICT_STFUNC(ErrorState, { - cFunc(NActors::TEvents::TEvPoisonPill::EventType, PassAway); + STRICT_STFUNC_EXC(ErrorState, + cFunc(TEvents::TEvPoisonPill::EventType, PassAway); IgnoreFunc(NFq::TEvPrivate::TEvPqEventsReady); IgnoreFunc(NFq::TEvPrivate::TEvCreateSession); - IgnoreFunc(NFq::TEvPrivate::TEvDataAfterFilteration); - IgnoreFunc(NFq::TEvPrivate::TEvSendStatisticToReadActor); - IgnoreFunc(NFq::TEvPrivate::TEvDataFiltered); IgnoreFunc(TEvRowDispatcher::TEvGetNextBatch); IgnoreFunc(NFq::TEvRowDispatcher::TEvStartSession); IgnoreFunc(NFq::TEvRowDispatcher::TEvStopSession); - IgnoreFunc(NFq::TEvPrivate::TEvSendStatisticToRowDispatcher); - IgnoreFunc(TEvRowDispatcher::TEvPurecalcCompileResponse); - }) + IgnoreFunc(NFq::TEvPrivate::TEvSendStatistic); + IgnoreFunc(NFq::TEvPrivate::TEvReconnectSession);, + ExceptionFunc(std::exception, HandleException) + ) }; TTopicSession::TTopicSession( + const TString& readGroup, const TString& topicPath, const TString& endpoint, const TString& database, const NConfig::TRowDispatcherConfig& config, - NActors::TActorId rowDispatcherActorId, - NActors::TActorId compileServiceActorId, + TActorId rowDispatcherActorId, + TActorId compileServiceActorId, ui32 partitionId, NYdb::TDriver driver, std::shared_ptr credentialsProviderFactory, const ::NMonitoring::TDynamicCounterPtr& counters, + const ::NMonitoring::TDynamicCounterPtr& countersRoot, const NYql::IPqGateway::TPtr& pqGateway, ui64 maxBufferSize) - : TopicPath(topicPath) + : ReadGroup(readGroup) + , TopicPath(topicPath) , TopicPathPartition(TStringBuilder() << topicPath << "/" << partitionId) , Endpoint(endpoint) , Database(database) , RowDispatcherActorId(rowDispatcherActorId) - , CompileServiceActorId(compileServiceActorId) , PartitionId(partitionId) , Driver(std::move(driver)) + , PqGateway(pqGateway) , CredentialsProviderFactory(credentialsProviderFactory) + , Config(config) + , FormatHandlerConfig(CreateFormatHandlerConfig(config, compileServiceActorId)) , BufferSize(maxBufferSize) , LogPrefix("TopicSession") - , Config(config) , Counters(counters) - , PqGateway(pqGateway) -{ -} + , CountersRoot(countersRoot) +{} void TTopicSession::Bootstrap() { Become(&TTopicSession::StateFunc); - Metrics.Init(Counters, TopicPath, PartitionId); + Metrics.Init(Counters, TopicPath, ReadGroup, PartitionId); LogPrefix = LogPrefix + " " + SelfId().ToString() + " "; LOG_ROW_DISPATCHER_DEBUG("Bootstrap " << TopicPathPartition << ", Timeout " << Config.GetTimeoutBeforeStartSessionSec() << " sec, StatusPeriod " << Config.GetSendStatusPeriodSec() << " sec"); Y_ENSURE(Config.GetSendStatusPeriodSec() > 0); - Schedule(TDuration::Seconds(Config.GetSendStatusPeriodSec()), new NFq::TEvPrivate::TEvSendStatisticToReadActor()); - Schedule(TDuration::Seconds(SendStatisticPeriodSec), new NFq::TEvPrivate::TEvSendStatisticToRowDispatcher()); + Schedule(TDuration::Seconds(SendStatisticPeriodSec), new NFq::TEvPrivate::TEvSendStatistic()); } void TTopicSession::PassAway() { LOG_ROW_DISPATCHER_DEBUG("PassAway"); StopReadSession(); - NActors::TActorBootstrapped::PassAway(); + FormatHandlers.clear(); + TBase::PassAway(); } void TTopicSession::SubscribeOnNextEvent() { @@ -373,15 +393,15 @@ void TTopicSession::SubscribeOnNextEvent() { return; } - if (Config.GetMaxSessionUsedMemory() && UnreadBytes > Config.GetMaxSessionUsedMemory()) { - LOG_ROW_DISPATCHER_TRACE("Too much used memory (" << UnreadBytes << " bytes), skip subscribing to WaitEvent()"); + if (Config.GetMaxSessionUsedMemory() && QueuedBytes > Config.GetMaxSessionUsedMemory()) { + LOG_ROW_DISPATCHER_TRACE("Too much used memory (" << QueuedBytes << " bytes), skip subscribing to WaitEvent()"); return; } LOG_ROW_DISPATCHER_TRACE("SubscribeOnNextEvent"); IsWaitingEvents = true; Metrics.InFlySubscribe->Inc(); - NActors::TActorSystem* actorSystem = NActors::TActivationContext::ActorSystem(); + TActorSystem* actorSystem = TActivationContext::ActorSystem(); WaitEventStartedAt = TInstant::Now(); ReadSession->WaitEvent().Subscribe([actorSystem, selfId = SelfId()](const auto&){ actorSystem->Send(selfId, new NFq::TEvPrivate::TEvPqEventsReady()); @@ -389,12 +409,11 @@ void TTopicSession::SubscribeOnNextEvent() { } NYdb::NTopic::TTopicClientSettings TTopicSession::GetTopicClientSettings(const NYql::NPq::NProto::TDqPqTopicSource& sourceParams) const { - NYdb::NTopic::TTopicClientSettings opts; - opts.Database(Database) + return PqGateway->GetTopicClientSettings() + .Database(Database) .DiscoveryEndpoint(Endpoint) .SslCredentials(NYdb::TSslCredentials(sourceParams.GetUseSsl())) .CredentialsProviderFactory(CredentialsProviderFactory); - return opts; } NYql::ITopicClient& TTopicSession::GetTopicClient(const NYql::NPq::NProto::TDqPqTopicSource& sourceParams) { @@ -408,7 +427,7 @@ TInstant TTopicSession::GetMinStartingMessageTimestamp() const { auto result = TInstant::Max(); Y_ENSURE(!Clients.empty()); for (const auto& [actorId, info] : Clients) { - ui64 time = info.Settings.GetStartingMessageTimestampMs(); + ui64 time = info->Settings.GetStartingMessageTimestampMs(); result = std::min(result, TInstant::MilliSeconds(time)); } return result; @@ -442,17 +461,15 @@ void TTopicSession::CreateTopicSession() { } if (!ReadSession) { - UpdateParser(); - // Use any sourceParams. - const NYql::NPq::NProto::TDqPqTopicSource& sourceParams = Clients.begin()->second.Settings.GetSource(); + const NYql::NPq::NProto::TDqPqTopicSource& sourceParams = Clients.begin()->second->Settings.GetSource(); ReadSession = GetTopicClient(sourceParams).CreateReadSession(GetReadSessionSettings(sourceParams)); SubscribeOnNextEvent(); } - if (!InflightReconnect && Clients) { + if (!InflightReconnect && !Clients.empty()) { // Use any sourceParams. - ReconnectPeriod = Clients.begin()->second.ReconnectPeriod; + ReconnectPeriod = Clients.begin()->second->ReconnectPeriod; if (ReconnectPeriod != TDuration::Zero()) { LOG_ROW_DISPATCHER_INFO("ReconnectPeriod " << ReconnectPeriod.ToString()); Metrics.ReconnectRate->Inc(); @@ -476,70 +493,18 @@ void TTopicSession::Handle(NFq::TEvPrivate::TEvCreateSession::TPtr&) { CreateTopicSession(); } -TVector*> TTopicSession::RebuildJson(const TClientsInfo& info, const TVector>& parsedValues) { - TVector*> result; - const auto& offsets = ParserSchema.FieldsMap; - result.reserve(info.FieldsIds.size()); - for (auto fieldId : info.FieldsIds) { - Y_ENSURE(fieldId < offsets.size(), "fieldId " << fieldId << ", offsets.size() " << offsets.size()); - auto offset = offsets[fieldId]; - Y_ENSURE(offset < parsedValues.size(), "offset " << offset << ", jsonBatch.size() " << parsedValues.size()); - result.push_back(&parsedValues[offset]); - } - return result; -} - -void TTopicSession::Handle(NFq::TEvPrivate::TEvDataAfterFilteration::TPtr& ev) { - LOG_ROW_DISPATCHER_TRACE("TEvDataAfterFilteration, read actor id " << ev->Get()->ReadActorId.ToString() << ", " << ev->Get()->Json); - auto it = Clients.find(ev->Get()->ReadActorId); - if (it == Clients.end()) { - LOG_ROW_DISPATCHER_ERROR("Skip DataAfterFilteration, wrong read actor, id " << ev->Get()->ReadActorId.ToString()); - return; - } - AddDataToClient(it->second, ev->Get()->Offset, ev->Get()->Json); -} - -void TTopicSession::Handle(NFq::TEvPrivate::TEvSendStatisticToReadActor::TPtr&) { - LOG_ROW_DISPATCHER_TRACE("TEvSendStatisticToReadActor"); - Schedule(TDuration::Seconds(Config.GetSendStatusPeriodSec()), new NFq::TEvPrivate::TEvSendStatisticToReadActor()); - - auto readBytes = ClientsStats.Bytes; - for (auto& [actorId, info] : Clients) { - if (!info.NextMessageOffset) { - continue; - } - auto event = std::make_unique(); - event->Record.SetPartitionId(PartitionId); - event->Record.SetNextMessageOffset(*info.NextMessageOffset); - event->Record.SetReadBytes(readBytes); - info.LastSendedNextMessageOffset = *info.NextMessageOffset; - event->ReadActorId = info.ReadActorId; - LOG_ROW_DISPATCHER_TRACE("Send status to " << info.ReadActorId << ", offset " << *info.NextMessageOffset); - Send(RowDispatcherActorId, event.release()); - } - ClientsStats.Clear(); -} - void TTopicSession::Handle(NFq::TEvPrivate::TEvReconnectSession::TPtr&) { Metrics.ReconnectRate->Inc(); TInstant minTime = GetMinStartingMessageTimestamp(); LOG_ROW_DISPATCHER_DEBUG("Reconnect topic session, " << TopicPathPartition << ", StartingMessageTimestamp " << minTime << ", BufferSize " << BufferSize << ", WithoutConsumer " << Config.GetWithoutConsumer()); + RefreshParsers(); StopReadSession(); CreateTopicSession(); Schedule(ReconnectPeriod, new NFq::TEvPrivate::TEvReconnectSession()); } -void TTopicSession::Handle(NFq::TEvPrivate::TEvDataFiltered::TPtr& ev) { - LOG_ROW_DISPATCHER_TRACE("TEvDataFiltered, last offset " << ev->Get()->Offset); - for (auto& [actorId, info] : Clients) { - if (!info.NextMessageOffset || *info.NextMessageOffset < ev->Get()->Offset + 1) { - info.NextMessageOffset = ev->Get()->Offset + 1; - } - } -} - void TTopicSession::Handle(TEvRowDispatcher::TEvGetNextBatch::TPtr& ev) { LOG_ROW_DISPATCHER_TRACE("TEvGetNextBatch from " << ev->Sender.ToString()); Metrics.InFlyAsyncInputData->Set(0); @@ -548,7 +513,7 @@ void TTopicSession::Handle(TEvRowDispatcher::TEvGetNextBatch::TPtr& ev) { LOG_ROW_DISPATCHER_ERROR("Wrong client, sender " << ev->Sender); return; } - SendData(it->second); + SendData(*it->second); SubscribeOnNextEvent(); } @@ -559,8 +524,8 @@ void TTopicSession::HandleNewEvents() { if (!ReadSession) { return; } - if (Config.GetMaxSessionUsedMemory() && UnreadBytes > Config.GetMaxSessionUsedMemory()) { - LOG_ROW_DISPATCHER_TRACE("Too much used memory (" << UnreadBytes << " bytes), stop reading from yds"); + if (Config.GetMaxSessionUsedMemory() && QueuedBytes > Config.GetMaxSessionUsedMemory()) { + LOG_ROW_DISPATCHER_TRACE("Too much used memory (" << QueuedBytes << " bytes), stop reading from yds"); break; } TMaybe event = ReadSession->GetEvent(false); @@ -592,20 +557,21 @@ void TTopicSession::TTopicEventProcessor::operator()(NYdb::NTopic::TReadSessionE Self.LastMessageOffset = message.GetOffset(); } - Self.SessionStats.Add(dataSize, event.GetMessages().size()); - Self.ClientsStats.Add(dataSize, event.GetMessages().size()); + Self.Statistics.Add(dataSize, event.GetMessages().size()); Self.Metrics.SessionDataRate->Add(dataSize); Self.Metrics.AllSessionsDataRate->Add(dataSize); - dataReceivedEventSize += dataSize; + DataReceivedEventSize += dataSize; Self.SendToParsing(event.GetMessages()); } void TTopicSession::TTopicEventProcessor::operator()(NYdb::NTopic::TSessionClosedEvent& ev) { - TString message = TStringBuilder() << "Read session to topic \"" << Self.TopicPathPartition << "\" was closed: " << ev.DebugString(); - LOG_ROW_DISPATCHER_DEBUG(message); - NYql::TIssues issues; - issues.AddIssue(message); - Self.FatalError(issues.ToOneLineString(), nullptr, false, Nothing()); + const TString message = TStringBuilder() << "Read session to topic \"" << Self.TopicPathPartition << "\" was closed"; + LOG_ROW_DISPATCHER_DEBUG(message << ": " << ev.DebugString()); + + Self.ThrowFatalError(TStatus::Fail( + NYql::NDq::YdbStatusToDqStatus(static_cast(ev.GetStatus())), + ev.GetIssues() + ).AddParentIssue(message)); } void TTopicSession::TTopicEventProcessor::operator()(NYdb::NTopic::TReadSessionEvent::TStartPartitionSessionEvent& event) { @@ -613,10 +579,9 @@ void TTopicSession::TTopicEventProcessor::operator()(NYdb::NTopic::TReadSessionE TMaybe minOffset; for (const auto& [actorId, info] : Self.Clients) { - if (!minOffset - || (info.NextMessageOffset && (info.NextMessageOffset < *minOffset))) { - minOffset = info.NextMessageOffset; - } + if (!minOffset || (info->NextMessageOffset && *info->NextMessageOffset < *minOffset)) { + minOffset = info->NextMessageOffset; + } } LOG_ROW_DISPATCHER_DEBUG("Confirm StartPartitionSession with offset " << minOffset); event.Confirm(minOffset); @@ -635,109 +600,35 @@ void TTopicSession::TTopicEventProcessor::operator()(NYdb::NTopic::TReadSessionE LOG_ROW_DISPATCHER_WARN("TPartitionSessionClosedEvent"); } -std::pair TTopicSession::CreateItem(const NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage& message) { - const TString& data = message.GetData(); - i64 usedSpace = data.size(); - NYql::NUdf::TUnboxedValuePod item = NKikimr::NMiniKQL::MakeString(NYql::NUdf::TStringRef(data.data(), data.size())); - return std::make_pair(item, usedSpace); -} - TString TTopicSession::GetSessionId() const { return ReadSession ? ReadSession->GetSessionId() : TString{"empty"}; } void TTopicSession::SendToParsing(const TVector& messages) { - for (const auto& readActorId : ClientsWithoutPredicate) { - const auto it = Clients.find(readActorId); - Y_ENSURE(it != Clients.end(), "Internal error: unknown client"); - auto& info = it->second; - if (info.Filter) { - continue; - } - - for (const auto& message : messages) { - LOG_ROW_DISPATCHER_TRACE("Send message with offset " << message.GetOffset() << " to client " << info.ReadActorId <<" without parsing/filtering"); - AddDataToClient(info, message.GetOffset(), message.GetData()); + LOG_ROW_DISPATCHER_TRACE("SendToParsing, messages: " << messages.size()); + for (const auto& [_, formatHandler] : FormatHandlers) { + if (formatHandler->HasClients()) { + formatHandler->ParseMessages(messages); } } - - if (ClientsWithoutPredicate.size() == Clients.size()) { - return; - } - - Parser->AddMessages(messages); - DoParsing(); -} - -void TTopicSession::DoParsing(bool force) { - if (!Parser->IsReady() && !force) { - const TInstant batchCreationDeadline = Parser->GetCreationDeadline(); - LOG_ROW_DISPATCHER_TRACE("Collecting data to parse, skip parsing, creation deadline " << batchCreationDeadline); - if (!IsStartParsingScheduled && batchCreationDeadline) { - IsStartParsingScheduled = true; - Schedule(batchCreationDeadline, new TEvPrivate::TEvStartParsing()); - } - return; - } - - if (!Parser->GetNumberValues()) { - return; - } - - IsStartParsingScheduled = false; - LOG_ROW_DISPATCHER_TRACE("SendToParsing, first offset: " << Parser->GetOffsets().front() << ", number values in buffer " << Parser->GetOffsets().size()); - - TInstant startParseAndFilter = TInstant::Now(); - try { - Parser->Parse(); // this call processes parse and filter steps - } catch (const NFq::TJsonParserError& e) { - FatalError(e.what(), nullptr, true, e.GetField()); - } catch (const std::exception& e) { - FatalError(e.what(), nullptr, true, Nothing()); - } - auto parseAndFilterLatency = TInstant::Now() - startParseAndFilter; - SessionStats.ParseAndFilterLatency = Max(SessionStats.ParseAndFilterLatency, parseAndFilterLatency); - ClientsStats.ParseAndFilterLatency = Max(ClientsStats.ParseAndFilterLatency, parseAndFilterLatency); } -void TTopicSession::DoFiltering(ui64 rowsOffset, ui64 numberRows, const TVector>& parsedValues) { - const auto& offsets = Parser->GetOffsets(); - Y_ENSURE(rowsOffset < offsets.size(), "Invalid first row ofset"); - Y_ENSURE(numberRows, "Expected non empty parsed batch"); - Y_ENSURE(parsedValues, "Expected non empty schema"); - auto lastOffset = offsets[rowsOffset + numberRows - 1]; - LOG_ROW_DISPATCHER_TRACE("SendToFiltering, first offset: " << offsets[rowsOffset] << ", last offset: " << lastOffset); - - for (auto& [actorId, info] : Clients) { - if (info.InFlightCompilationId) { // filter compilation in flight - continue; - } - if (info.NextMessageOffset && lastOffset < info.NextMessageOffset) { // the batch has already been processed - continue; - } - try { - if (info.Filter) { - info.Filter->Push(offsets, RebuildJson(info, parsedValues), rowsOffset, numberRows); - } - } catch (const std::exception& e) { - FatalError(e.what(), &info.Filter, false, Nothing()); - } +void TTopicSession::SendData(TClientsInfo& info) { + TQueue>> buffer; + if (const auto formatIt = FormatHandlers.find(info.HandlerSettings); formatIt != FormatHandlers.end()) { + buffer = formatIt->second->ExtractClientData(info.GetClientId()); } - Send(SelfId(), new TEvPrivate::TEvDataFiltered(offsets.back())); -} - -void TTopicSession::SendData(TClientsInfo& info) { info.DataArrivedSent = false; - if (info.Buffer.empty()) { + if (buffer.empty()) { LOG_ROW_DISPATCHER_TRACE("Buffer empty"); } ui64 dataSize = 0; - ui64 eventsSize = info.Buffer.size(); + ui64 eventsSize = info.QueuedRows; if (!info.NextMessageOffset) { LOG_ROW_DISPATCHER_ERROR("Try SendData() without NextMessageOffset, " << info.ReadActorId - << " unread " << info.UnreadBytes << " DataArrivedSent " << info.DataArrivedSent); + << " unread " << info.QueuedBytes << " DataArrivedSent " << info.DataArrivedSent); return; } @@ -747,62 +638,46 @@ void TTopicSession::SendData(TClientsInfo& info) { event->ReadActorId = info.ReadActorId; ui64 batchSize = 0; - while (!info.Buffer.empty()) { - const auto& [offset, json] = info.Buffer.front(); - info.UnreadBytes -= json.size(); - UnreadBytes -= json.size(); - batchSize += json.size(); + while (!buffer.empty()) { + auto [serializedData, offsets] = std::move(buffer.front()); + Y_ENSURE(!offsets.empty(), "Expected non empty message batch"); + buffer.pop(); + + batchSize += serializedData.GetSize(); + NFq::NRowDispatcherProto::TEvMessage message; - message.SetJson(json); - message.SetOffset(offset); - event->Record.AddMessages()->CopyFrom(message); - event->Record.SetNextMessageOffset(offset + 1); - info.Buffer.pop(); + message.SetPayloadId(event->AddPayload(std::move(serializedData))); + message.MutableOffsets()->Assign(offsets.begin(), offsets.end()); + event->Record.AddMessages()->CopyFrom(std::move(message)); + event->Record.SetNextMessageOffset(offsets.back() + 1); - if (batchSize > MaxBatchSizeBytes) { + if (batchSize > MAX_BATCH_SIZE) { break; } } dataSize += batchSize; - if (info.Buffer.empty()) { + if (buffer.empty()) { event->Record.SetNextMessageOffset(*info.NextMessageOffset); } LOG_ROW_DISPATCHER_TRACE("SendData to " << info.ReadActorId << ", batch size " << event->Record.MessagesSize()); Send(RowDispatcherActorId, event.release()); - } while(!info.Buffer.empty()); - info.Stat.Add(dataSize, eventsSize); - info.FilteredDataRate->Add(dataSize); - info.LastSendedNextMessageOffset = *info.NextMessageOffset; -} + } while(!buffer.empty()); -void TTopicSession::UpdateFieldsIds(TClientsInfo& info) { - const auto& source = info.Settings.GetSource(); - for (size_t i = 0; i < source.ColumnsSize(); ++i) { - const auto& name = source.GetColumns().Get(i); - auto it = FieldsIndexes.find(name); - if (it == FieldsIndexes.end()) { - auto nextIndex = FieldsIndexes.size(); - info.FieldsIds.push_back(nextIndex); - FieldsIndexes[name] = {nextIndex, source.GetColumnTypes().Get(i)}; - } else { - info.FieldsIds.push_back(it->second.IndexInParserSchema); - } - } -} + QueuedBytes -= info.QueuedBytes; + Metrics.QueuedBytes->Sub(info.QueuedBytes); + info.QueuedRows = 0; + info.QueuedBytes = 0; -bool HasJsonColumns(const NYql::NPq::NProto::TDqPqTopicSource& sourceParams) { - for (const auto& type : sourceParams.GetColumnTypes()) { - if (type.Contains("Json")) { - return true; - } - } - return false; + info.FilteredStat.Add(dataSize, eventsSize); + info.FilteredDataRate->Add(dataSize); + info.ProcessedNextMessageOffset = *info.NextMessageOffset; } void TTopicSession::StartClientSession(TClientsInfo& info) { if (ReadSession) { - if (info.Settings.HasOffset() && info.Settings.GetOffset() <= LastMessageOffset) { - LOG_ROW_DISPATCHER_INFO("New client has less offset (" << info.Settings.GetOffset() << ") than the last message (" << LastMessageOffset << "), stop (restart) topic session"); + auto offset = GetOffset(info.Settings); + if (offset && offset <= LastMessageOffset) { + LOG_ROW_DISPATCHER_INFO("New client has less offset (" << offset << ") than the last message (" << LastMessageOffset << "), stop (restart) topic session"); Metrics.RestartSessionByOffsets->Inc(); ++RestartSessionByOffsets; info.RestartSessionByOffsetsByQuery->Inc(); @@ -810,196 +685,126 @@ void TTopicSession::StartClientSession(TClientsInfo& info) { } } - if (Parser) { - // Parse remains data before changing parsing schema - DoParsing(true); - } - UpdateParser(); - if (!ReadSession) { Schedule(TDuration::Seconds(Config.GetTimeoutBeforeStartSessionSec()), new NFq::TEvPrivate::TEvCreateSession()); } } void TTopicSession::Handle(NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev) { - LOG_ROW_DISPATCHER_INFO("New client: read actor id " << ev->Sender.ToString() << ", predicate: " - << ev->Get()->Record.GetSource().GetPredicate() << ", offset: " << ev->Get()->Record.GetOffset()); + auto offset = GetOffset(ev->Get()->Record); + const auto& source = ev->Get()->Record.GetSource(); + LOG_ROW_DISPATCHER_INFO("New client: read actor id " << ev->Sender.ToString() << ", predicate: " << source.GetPredicate() << ", offset: " << offset); if (!CheckNewClient(ev)) { return; } - auto columns = GetVector(ev->Get()->Record.GetSource().GetColumns()); - auto types = GetVector(ev->Get()->Record.GetSource().GetColumnTypes()); + const TString& format = source.GetFormat(); + ITopicFormatHandler::TSettings handlerSettings = {.ParsingFormat = format ? format : "raw"}; - try { - auto queryGroup = Counters->GetSubgroup("queryId", ev->Get()->Record.GetQueryId()); - auto topicGroup = queryGroup->GetSubgroup("topic", CleanupCounterValueString(TopicPath)); - auto& clientInfo = Clients.emplace( - std::piecewise_construct, - std::forward_as_tuple(ev->Sender), - std::forward_as_tuple(ev, topicGroup)).first->second; - UpdateFieldsIds(clientInfo); - - const auto& source = clientInfo.Settings.GetSource(); - TString predicate = source.GetPredicate(); - - // TODO: remove this when the re-parsing is removed from pq read actor - if (predicate.empty() && HasJsonColumns(source)) { - predicate = "WHERE TRUE"; - } - - if (!predicate.empty()) { - clientInfo.Filter = NewJsonFilter( - columns, - types, - predicate, - [&, actorId = clientInfo.ReadActorId](ui64 offset, const TString& json){ - Send(SelfId(), new NFq::TEvPrivate::TEvDataAfterFilteration(offset, json, actorId)); - }, - {.EnabledLLVM = source.GetEnabledLLVM()} - ); - - clientInfo.InFlightCompilationId = ++FiltersCompilation.FreeId; - Y_ENSURE(FiltersCompilation.InFlightCompilations.emplace(clientInfo.InFlightCompilationId, ev->Sender).second, "Got duplicated compilation event id"); - LOG_ROW_DISPATCHER_TRACE("Send compile request with id " << clientInfo.InFlightCompilationId); - - Send(CompileServiceActorId, clientInfo.Filter->GetCompileRequest().release(), 0, clientInfo.InFlightCompilationId); - Metrics.InFlightCompileRequests->Inc(); - } else { - ClientsWithoutPredicate.insert(ev->Sender); - - // In case of in flight compilation topic session will be checked after getting compile response - StartClientSession(clientInfo); - } - } catch (...) { - FatalError("Adding new client failed, got unexpected exception: " + CurrentExceptionMessage(), nullptr, true, Nothing()); + auto clientInfo = Clients.insert({ev->Sender, MakeIntrusive(*this, LogPrefix, handlerSettings, ev, Counters, ReadGroup, offset)}).first->second; + auto formatIt = FormatHandlers.find(handlerSettings); + if (formatIt == FormatHandlers.end()) { + formatIt = FormatHandlers.insert({handlerSettings, CreateTopicFormatHandler( + ActorContext(), + FormatHandlerConfig, + handlerSettings, + {.CountersRoot = CountersRoot, .CountersSubgroup = Metrics.PartitionGroup} + )}).first; } - ConsumerName = ev->Get()->Record.GetSource().GetConsumerName(); - SendStatisticToRowDispatcher(); -} -void TTopicSession::AddDataToClient(TClientsInfo& info, ui64 offset, const TString& json) { - if (info.NextMessageOffset && offset < info.NextMessageOffset) { + if (auto status = formatIt->second->AddClient(clientInfo); status.IsFail()) { + SendSessionError(clientInfo->ReadActorId, status); return; } - info.NextMessageOffset = offset + 1; - info.Buffer.push(std::make_pair(offset, json)); - info.UnreadBytes += json.size(); - UnreadBytes += json.size(); - SendDataArrived(info); + + ConsumerName = source.GetConsumerName(); + SendStatistics(); } void TTopicSession::Handle(NFq::TEvRowDispatcher::TEvStopSession::TPtr& ev) { - LOG_ROW_DISPATCHER_DEBUG("TEvStopSession from " << ev->Sender << " topicPath " << ev->Get()->Record.GetSource().GetTopicPath() << - " partitionId " << ev->Get()->Record.GetPartitionId() << " clients count " << Clients.size()); + LOG_ROW_DISPATCHER_DEBUG("TEvStopSession from " << ev->Sender << " topicPath " << ev->Get()->Record.GetSource().GetTopicPath() << " clients count " << Clients.size()); auto it = Clients.find(ev->Sender); if (it == Clients.end()) { - LOG_ROW_DISPATCHER_DEBUG("Wrong ClientSettings"); + LOG_ROW_DISPATCHER_WARN("Ignore TEvStopSession from " << ev->Sender << ", no client"); return; } - auto& info = it->second; - UnreadBytes -= info.UnreadBytes; + auto& info = *it->second; + RestartSessionIfOldestClient(info); + + QueuedBytes -= info.QueuedBytes; + Metrics.QueuedBytes->Sub(info.QueuedBytes); + if (const auto formatIt = FormatHandlers.find(info.HandlerSettings); formatIt != FormatHandlers.end()) { + formatIt->second->RemoveClient(info.GetClientId()); + if (!formatIt->second->HasClients()) { + FormatHandlers.erase(formatIt); + } + } Clients.erase(it); - ClientsWithoutPredicate.erase(ev->Sender); if (Clients.empty()) { StopReadSession(); } - UpdateParser(); SubscribeOnNextEvent(); } -void CollectColumns(const NYql::NPq::NProto::TDqPqTopicSource& sourceParams, TSet>& columns) { - auto size = sourceParams.GetColumns().size(); - Y_ENSURE(size == sourceParams.GetColumnTypes().size()); - - for (int i = 0; i < size; ++i) { - auto name = sourceParams.GetColumns().Get(i); - auto type = sourceParams.GetColumnTypes().Get(i); - columns.emplace(name, type); - } -} +void TTopicSession::RestartSessionIfOldestClient(const TClientsInfo& info) { + // if we read historical data (because of this client), then we restart the session. -void TTopicSession::UpdateParserSchema(const TParserInputType& inputType) { - ParserSchema.FieldsMap.clear(); - ParserSchema.FieldsMap.resize(FieldsIndexes.size()); - ui64 offset = 0; - for (const auto& [name, type]: inputType) { - Y_ENSURE(FieldsIndexes.contains(name)); - ui64 index = FieldsIndexes[name].IndexInParserSchema; - ParserSchema.FieldsMap[index] = offset++; + if (!ReadSession || !info.NextMessageOffset) { + return; } - ParserSchema.InputType = inputType; -} - -void TTopicSession::UpdateParser() { - TSet> namesWithTypes; - for (auto& [readActorId, info] : Clients) { - CollectColumns(info.Settings.GetSource(), namesWithTypes); + TMaybe minMessageOffset; + for (auto& [readActorId, clientPtr] : Clients) { + if (info.ReadActorId == readActorId || !clientPtr->NextMessageOffset) { + continue; + } + if (!minMessageOffset) { + minMessageOffset = clientPtr->NextMessageOffset; + continue; + } + minMessageOffset = std::min(minMessageOffset, clientPtr->NextMessageOffset); } - - if (namesWithTypes == ParserSchema.InputType) { + if (!minMessageOffset) { return; } - if (namesWithTypes.empty()) { - LOG_ROW_DISPATCHER_INFO("No columns to parse, reset parser"); - Parser.reset(); + + if (info.NextMessageOffset >= minMessageOffset) { return; } + LOG_ROW_DISPATCHER_INFO("Client (on StopSession) has less offset (" << info.NextMessageOffset << ") than others clients (" << minMessageOffset << "), stop (restart) topic session"); + Metrics.RestartSessionByOffsets->Inc(); + ++RestartSessionByOffsets; + info.RestartSessionByOffsetsByQuery->Inc(); + RefreshParsers(); + StopReadSession(); - try { - UpdateParserSchema(namesWithTypes); - - TVector names; - TVector types; - names.reserve(namesWithTypes.size()); - types.reserve(namesWithTypes.size()); - for (const auto& [name, type] : namesWithTypes) { - names.push_back(name); - types.push_back(type); - } - - LOG_ROW_DISPATCHER_TRACE("Init JsonParser with columns: " << JoinSeq(',', names)); - const auto& parserConfig = Config.GetJsonParser(); - Parser = NewJsonParser(names, types, [this](ui64 rowsOffset, ui64 numberRows, const TVector>& parsedValues) { - DoFiltering(rowsOffset, numberRows, parsedValues); - }, parserConfig.GetBatchSizeBytes(), TDuration::MilliSeconds(parserConfig.GetBatchCreationTimeoutMs()), parserConfig.GetBufferCellCount()); - } catch (const NFq::TJsonParserError& e) { - FatalError(e.what(), nullptr, true, e.GetField()); - } catch (const NYql::NPureCalc::TCompileError& e) { - FatalError(e.GetIssues(), nullptr, true, Nothing()); + if (!ReadSession) { + Schedule(TDuration::Seconds(Config.GetTimeoutBeforeStartSessionSec()), new NFq::TEvPrivate::TEvCreateSession()); } } -void TTopicSession::FatalError(const TString& message, const std::unique_ptr* filter, bool addParserDescription, const TMaybe& fieldName) { - TStringStream str; - str << message; - if (Parser && addParserDescription) { - str << ", parser description:\n" << Parser->GetDescription(); - } - if (filter) { - str << ", filter sql:\n" << (*filter)->GetSql(); - } - if (fieldName) { - auto queryId = GetAnyQueryIdByFieldName(*fieldName); - str << ", the field (" << *fieldName << ") has been added by query: " + queryId; - } - LOG_ROW_DISPATCHER_ERROR("FatalError: " << str.Str()); +void TTopicSession::FatalError(const TStatus& status) { + LOG_ROW_DISPATCHER_ERROR("FatalError: " << status.GetErrorMessage()); for (auto& [readActorId, info] : Clients) { LOG_ROW_DISPATCHER_DEBUG("Send TEvSessionError to " << readActorId); - SendSessionError(readActorId, str.Str()); + SendSessionError(readActorId, status); } StopReadSession(); Become(&TTopicSession::ErrorState); - ythrow yexception() << "FatalError: " << str.Str(); // To exit from current stack and call once PassAway() in HandleException(). } -void TTopicSession::SendSessionError(NActors::TActorId readActorId, const TString& message) { +void TTopicSession::ThrowFatalError(const TStatus& status) { + FatalError(status); + ythrow yexception() << "FatalError: " << status.GetErrorMessage(); +} + +void TTopicSession::SendSessionError(TActorId readActorId, TStatus status) { + LOG_ROW_DISPATCHER_WARN("SendSessionError to " << readActorId << ", status: " << status.GetErrorMessage()); auto event = std::make_unique(); - event->Record.SetMessage(message); - event->Record.SetPartitionId(PartitionId); + event->Record.SetStatusCode(status.GetStatus()); + NYql::IssuesToMessage(status.GetErrorDescription(), event->Record.MutableIssues()); event->ReadActorId = readActorId; Send(RowDispatcherActorId, event.release()); } @@ -1014,7 +819,7 @@ void TTopicSession::StopReadSession() { } void TTopicSession::SendDataArrived(TClientsInfo& info) { - if (info.Buffer.empty() || info.DataArrivedSent) { + if (!info.QueuedBytes || info.DataArrivedSent) { return; } info.DataArrivedSent = true; @@ -1030,144 +835,108 @@ void TTopicSession::HandleException(const std::exception& e) { if (CurrentStateFunc() == &TThis::ErrorState) { return; } - FatalError(TString("Internal error: exception: ") + e.what(), nullptr, false, Nothing()); + FatalError(TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Session error, got unexpected exception: " << e.what())); } -void TTopicSession::SendStatisticToRowDispatcher() { - TopicSessionStatistic sessionStatistic; +void TTopicSession::SendStatistics() { + LOG_ROW_DISPATCHER_TRACE("SendStatistics"); + TTopicSessionStatistic sessionStatistic; auto& commonStatistic = sessionStatistic.Common; - commonStatistic.UnreadBytes = UnreadBytes; + commonStatistic.QueuedBytes = QueuedBytes; commonStatistic.RestartSessionByOffsets = RestartSessionByOffsets; - commonStatistic.ReadBytes = SessionStats.Bytes; - commonStatistic.ReadEvents = SessionStats.Events; - commonStatistic.ParseAndFilterLatency = SessionStats.ParseAndFilterLatency; + commonStatistic.ReadBytes = Statistics.Bytes; + commonStatistic.ReadEvents = Statistics.Events; commonStatistic.LastReadedOffset = LastMessageOffset; - SessionStats.Clear(); - sessionStatistic.SessionKey = TopicSessionParams{Endpoint, Database, TopicPath, PartitionId}; + sessionStatistic.SessionKey = TTopicSessionParams{ReadGroup, Endpoint, Database, TopicPath, PartitionId}; sessionStatistic.Clients.reserve(Clients.size()); - for (auto& [readActorId, info] : Clients) { - TopicSessionClientStatistic clientStatistic; + for (const auto& [readActorId, infoPtr] : Clients) { + auto& info = *infoPtr; + TTopicSessionClientStatistic clientStatistic; clientStatistic.PartitionId = PartitionId; clientStatistic.ReadActorId = readActorId; - clientStatistic.UnreadRows = info.Buffer.size(); - clientStatistic.UnreadBytes = info.UnreadBytes; - clientStatistic.Offset = info.NextMessageOffset.GetOrElse(0); - clientStatistic.ReadBytes = info.Stat.Bytes; + clientStatistic.QueuedRows = info.QueuedRows; + clientStatistic.QueuedBytes = info.QueuedBytes; + clientStatistic.Offset = info.ProcessedNextMessageOffset.GetOrElse(0); + clientStatistic.FilteredBytes = info.FilteredStat.Bytes; + clientStatistic.FilteredRows = info.FilteredStat.Events; + clientStatistic.ReadBytes = Statistics.Bytes; clientStatistic.IsWaiting = LastMessageOffset + 1 < info.NextMessageOffset.GetOrElse(0); clientStatistic.ReadLagMessages = info.NextMessageOffset.GetOrElse(0) - LastMessageOffset - 1; clientStatistic.InitialOffset = info.InitialOffset; - info.Stat.Clear(); + info.FilteredStat.Clear(); sessionStatistic.Clients.emplace_back(std::move(clientStatistic)); } + + commonStatistic.FormatHandlers.reserve(FormatHandlers.size()); + for (const auto& [settings, handler] : FormatHandlers) { + commonStatistic.FormatHandlers.emplace(settings.ParsingFormat, handler->GetStatistics()); + } + auto event = std::make_unique(sessionStatistic); Send(RowDispatcherActorId, event.release()); + Statistics.Clear(); } -void TTopicSession::Handle(NFq::TEvPrivate::TEvSendStatisticToRowDispatcher::TPtr&) { - Schedule(TDuration::Seconds(SendStatisticPeriodSec), new NFq::TEvPrivate::TEvSendStatisticToRowDispatcher()); - SendStatisticToRowDispatcher(); +void TTopicSession::Handle(NFq::TEvPrivate::TEvSendStatistic::TPtr&) { + SendStatistics(); + Schedule(TDuration::Seconds(SendStatisticPeriodSec), new NFq::TEvPrivate::TEvSendStatistic()); } bool TTopicSession::CheckNewClient(NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev) { auto it = Clients.find(ev->Sender); if (it != Clients.end()) { LOG_ROW_DISPATCHER_ERROR("Such a client already exists"); - SendSessionError(ev->Sender, "Internal error: such a client already exists"); + SendSessionError(ev->Sender, TStatus::Fail(EStatusId::INTERNAL_ERROR, TStringBuilder() << "Client with id " << ev->Sender << " already exists")); return false; } const auto& source = ev->Get()->Record.GetSource(); if (!Config.GetWithoutConsumer() && ConsumerName && ConsumerName != source.GetConsumerName()) { LOG_ROW_DISPATCHER_INFO("Different consumer, expected " << ConsumerName << ", actual " << source.GetConsumerName() << ", send error"); - SendSessionError(ev->Sender, TStringBuilder() << "Use the same consumer in all queries via RD (current consumer " << ConsumerName << ")"); + SendSessionError(ev->Sender, TStatus::Fail(EStatusId::PRECONDITION_FAILED, TStringBuilder() << "Use the same consumer in all queries via RD (current consumer " << ConsumerName << ")")); return false; } - Y_ENSURE(source.ColumnsSize() == source.ColumnTypesSize()); - for (size_t i = 0; i < source.ColumnsSize(); ++i) { - const auto& name = source.GetColumns().Get(i); - const auto& type = source.GetColumnTypes().Get(i); - const auto it = FieldsIndexes.find(name); - if (it != FieldsIndexes.end() && it->second.Type != type) { - LOG_ROW_DISPATCHER_INFO("Different column `" << name << "` type, expected " << it->second.Type << ", actual " << type << ", send error"); - SendSessionError(ev->Sender, TStringBuilder() << "Use the same column type in all queries via RD, current type for column `" << name << "` is " << it->second.Type << " (requested type is " << type <<")"); - return false; - } - } - return true; } -void TTopicSession::Handle(TEvRowDispatcher::TEvPurecalcCompileResponse::TPtr& ev) { - LOG_ROW_DISPATCHER_TRACE("Got compile response for reauest with id " << ev->Cookie); - - const auto requestIt = FiltersCompilation.InFlightCompilations.find(ev->Cookie); - if (requestIt == FiltersCompilation.InFlightCompilations.end()) { - LOG_ROW_DISPATCHER_TRACE("Compile response ignored for id " << ev->Cookie); - return; - } - - const auto clientId = requestIt->second; - FiltersCompilation.InFlightCompilations.erase(requestIt); - Metrics.InFlightCompileRequests->Dec(); - - if (!ev->Get()->ProgramHolder) { - TString message = TStringBuilder() << "Filed to compile purecalc program, error: " << ev->Get()->Error; - LOG_ROW_DISPATCHER_ERROR(message); - FatalError(message, nullptr, false, Nothing()); - return; - } - - const auto clientIt = Clients.find(clientId); - if (clientIt == Clients.end()) { - LOG_ROW_DISPATCHER_TRACE("Compile response ignored for id " << ev->Cookie << ", client with id " << clientId << " not found"); - return; - } - - auto& clientInfo = clientIt->second; - if (ev->Cookie != clientInfo.InFlightCompilationId) { - LOG_ROW_DISPATCHER_TRACE("Outdated compiler response ignored for id " << ev->Cookie << ", client with id " << clientId << " changed"); - return; +TMaybe TTopicSession::GetOffset(const NFq::NRowDispatcherProto::TEvStartSession& settings) { + for (auto p : settings.GetOffsets()) { + if (p.GetPartitionId() != PartitionId) { + continue; + } + return p.GetOffset(); } - - Y_ENSURE(clientInfo.Filter, "Unexpected completion response for client without filter"); - clientInfo.Filter->OnCompileResponse(std::move(ev)); - clientInfo.InFlightCompilationId = 0; - StartClientSession(clientInfo); + return Nothing(); } -TString TTopicSession::GetAnyQueryIdByFieldName(const TString& fieldName) { - TSet> namesWithTypes; - for (auto& [readActorId, info] : Clients) { - for (const auto& name : info.Settings.GetSource().GetColumns()) { - if (name != fieldName) { - continue; - } - return info.Settings.GetQueryId(); - } +void TTopicSession::RefreshParsers() { + for (const auto& [_, formatHandler] : FormatHandlers) { + formatHandler->ForceRefresh(); } - return "Unknown"; } -} // namespace +} // anonymous namespace //////////////////////////////////////////////////////////////////////////////// - -std::unique_ptr NewTopicSession( + +std::unique_ptr NewTopicSession( + const TString& readGroup, const TString& topicPath, const TString& endpoint, const TString& database, const NConfig::TRowDispatcherConfig& config, - NActors::TActorId rowDispatcherActorId, - NActors::TActorId compileServiceActorId, + TActorId rowDispatcherActorId, + TActorId compileServiceActorId, ui32 partitionId, NYdb::TDriver driver, std::shared_ptr credentialsProviderFactory, const ::NMonitoring::TDynamicCounterPtr& counters, + const ::NMonitoring::TDynamicCounterPtr& countersRoot, const NYql::IPqGateway::TPtr& pqGateway, ui64 maxBufferSize) { - return std::unique_ptr(new TTopicSession(topicPath, endpoint, database, config, rowDispatcherActorId, compileServiceActorId, partitionId, std::move(driver), credentialsProviderFactory, counters, pqGateway, maxBufferSize)); + return std::unique_ptr(new TTopicSession(readGroup, topicPath, endpoint, database, config, rowDispatcherActorId, compileServiceActorId, partitionId, std::move(driver), credentialsProviderFactory, counters, countersRoot, pqGateway, maxBufferSize)); } -} // namespace NFq +} // namespace NFq diff --git a/ydb/core/fq/libs/row_dispatcher/topic_session.h b/ydb/core/fq/libs/row_dispatcher/topic_session.h index 3cf22cb0a904..d24201d0d601 100644 --- a/ydb/core/fq/libs/row_dispatcher/topic_session.h +++ b/ydb/core/fq/libs/row_dispatcher/topic_session.h @@ -16,6 +16,7 @@ namespace NFq { std::unique_ptr NewTopicSession( + const TString& readGroup, const TString& topicPath, const TString& endpoint, const TString& database, @@ -26,6 +27,7 @@ std::unique_ptr NewTopicSession( NYdb::TDriver driver, std::shared_ptr credentialsProviderFactory, const ::NMonitoring::TDynamicCounterPtr& counters, + const ::NMonitoring::TDynamicCounterPtr& countersRoot, const NYql::IPqGateway::TPtr& pqGateway, ui64 maxBufferSize); diff --git a/ydb/core/fq/libs/row_dispatcher/ut/coordinator_ut.cpp b/ydb/core/fq/libs/row_dispatcher/ut/coordinator_ut.cpp index 478326acf53c..152ed937f64d 100644 --- a/ydb/core/fq/libs/row_dispatcher/ut/coordinator_ut.cpp +++ b/ydb/core/fq/libs/row_dispatcher/ut/coordinator_ut.cpp @@ -35,6 +35,7 @@ class TFixture : public NUnitTest::TBaseFixture { NConfig::TRowDispatcherCoordinatorConfig config; config.SetCoordinationNodePath("RowDispatcher"); + config.SetTopicPartitionsLimitPerNode(1); auto& database = *config.MutableDatabase(); database.SetEndpoint("YDB_ENDPOINT"); database.SetDatabase("YDB_DATABASE"); @@ -59,12 +60,12 @@ class TFixture : public NUnitTest::TBaseFixture { } NYql::NPq::NProto::TDqPqTopicSource BuildPqTopicSourceSettings( - TString topic) + TString endpoint, TString topic) { NYql::NPq::NProto::TDqPqTopicSource settings; settings.SetTopicPath(topic); settings.SetConsumerName("PqConsumer"); - settings.SetEndpoint("Endpoint"); + settings.SetEndpoint(endpoint); settings.MutableToken()->SetName("token"); settings.SetDatabase("Database"); return settings; @@ -84,9 +85,9 @@ class TFixture : public NUnitTest::TBaseFixture { //UNIT_ASSERT(eventHolder.Get() != nullptr); } - void MockRequest(NActors::TActorId readActorId, TString topicName, const std::vector& partitionId) { + void MockRequest(NActors::TActorId readActorId, TString endpoint, TString topicName, const std::vector& partitionId) { auto event = new NFq::TEvRowDispatcher::TEvCoordinatorRequest( - BuildPqTopicSourceSettings(topicName), + BuildPqTopicSourceSettings(endpoint, topicName), partitionId); Runtime.Send(new NActors::IEventHandle(Coordinator, readActorId, event)); } @@ -107,8 +108,6 @@ class TFixture : public NUnitTest::TBaseFixture { NActors::TActorId RowDispatcher2Id; NActors::TActorId ReadActor1; NActors::TActorId ReadActor2; - - NYql::NPq::NProto::TDqPqTopicSource Source1 = BuildPqTopicSourceSettings("Source1"); }; Y_UNIT_TEST_SUITE(CoordinatorTests) { @@ -121,17 +120,17 @@ Y_UNIT_TEST_SUITE(CoordinatorTests) { Ping(id); } - MockRequest(ReadActor1, "topic1", {0}); + MockRequest(ReadActor1, "endpoint", "topic1", {0}); auto result1 = ExpectResult(ReadActor1); - MockRequest(ReadActor2, "topic1", {0}); + MockRequest(ReadActor2, "endpoint", "topic1", {0}); auto result2 = ExpectResult(ReadActor2); UNIT_ASSERT(result1.PartitionsSize() == 1); UNIT_ASSERT(result2.PartitionsSize() == 1); UNIT_ASSERT(google::protobuf::util::MessageDifferencer::Equals(result1, result2)); - MockRequest(ReadActor2, "topic1", {1}); + MockRequest(ReadActor2, "endpoint", "topic1", {1}); auto result3 = ExpectResult(ReadActor2); TActorId actualRowDispatcher1 = ActorIdFromProto(result1.GetPartitions(0).GetActorId()); @@ -151,15 +150,29 @@ Y_UNIT_TEST_SUITE(CoordinatorTests) { auto newDispatcher2Id = Runtime.AllocateEdgeActor(1); Ping(newDispatcher2Id); - MockRequest(ReadActor1, "topic1", {0}); + MockRequest(ReadActor1, "endpoint", "topic1", {0}); auto result4 = ExpectResult(ReadActor1); - MockRequest(ReadActor2, "topic1", {1}); + MockRequest(ReadActor2, "endpoint", "topic1", {1}); auto result5 = ExpectResult(ReadActor2); UNIT_ASSERT(!google::protobuf::util::MessageDifferencer::Equals(result1, result4) || !google::protobuf::util::MessageDifferencer::Equals(result3, result5)); } + + Y_UNIT_TEST_F(RouteTwoTopicWichSameName, TFixture) { + ExpectCoordinatorChangesSubscribe(); + TSet rowDispatcherIds{RowDispatcher1Id, RowDispatcher2Id, LocalRowDispatcherId}; + for (auto id : rowDispatcherIds) { + Ping(id); + } + + MockRequest(ReadActor1, "endpoint1", "topic1", {0, 1, 2}); + ExpectResult(ReadActor1); + + MockRequest(ReadActor2, "endpoint2", "topic1", {3}); + ExpectResult(ReadActor2); + } } } diff --git a/ydb/core/fq/libs/row_dispatcher/ut/json_filter_ut.cpp b/ydb/core/fq/libs/row_dispatcher/ut/json_filter_ut.cpp deleted file mode 100644 index 7f4e3ed1fc4a..000000000000 --- a/ydb/core/fq/libs/row_dispatcher/ut/json_filter_ut.cpp +++ /dev/null @@ -1,206 +0,0 @@ -#include - -#include -#include - -#include -#include -#include - -#include -#include -#include - -#include - -#include - -namespace { - -using namespace NKikimr; -using namespace NFq; - -class TFixture : public NUnitTest::TBaseFixture { - -public: - TFixture() - : Runtime(true) - , Alloc(__LOCATION__, NKikimr::TAlignedPagePoolCounters(), true, false) - { - Alloc.Ref().UseRefLocking = true; - } - - static void SegmentationFaultHandler(int) { - Cerr << "segmentation fault call stack:" << Endl; - FormatBackTrace(&Cerr); - abort(); - } - - void SetUp(NUnitTest::TTestContext&) override { - NKikimr::EnableYDBBacktraceFormat(); - signal(SIGSEGV, &SegmentationFaultHandler); - - TAutoPtr app = new TAppPrepare(); - Runtime.Initialize(app->Unwrap()); - Runtime.SetLogPriority(NKikimrServices::FQ_ROW_DISPATCHER, NLog::PRI_DEBUG); - Runtime.SetDispatchTimeout(TDuration::Seconds(5)); - - CompileServiceActorId = Runtime.Register(NRowDispatcher::CreatePurecalcCompileService()); - } - - void TearDown(NUnitTest::TTestContext& /* context */) override { - with_lock (Alloc) { - for (const auto& holder : Holders) { - for (const auto& value : holder) { - Alloc.Ref().UnlockObject(value); - } - } - Holders.clear(); - } - Filter.reset(); - } - - void MakeFilter( - const TVector& columns, - const TVector& types, - const TString& whereFilter, - NFq::TJsonFilter::TCallback callback) { - Filter = NFq::NewJsonFilter( - columns, - types, - whereFilter, - callback, - {.EnabledLLVM = false}); - - const auto edgeActor = Runtime.AllocateEdgeActor(); - Runtime.Send(CompileServiceActorId, edgeActor, Filter->GetCompileRequest().release()); - auto response = Runtime.GrabEdgeEvent(edgeActor, TDuration::Seconds(5)); - - UNIT_ASSERT_C(response, "Failed to get compile response"); - UNIT_ASSERT_C(response->Get()->ProgramHolder, "Failed to compile program, error: " << response->Get()->Error); - Filter->OnCompileResponse(std::move(response)); - } - - void Push(const TVector& offsets, const TVector*>& values) { - Filter->Push(offsets, values, 0, values.front()->size()); - } - - const TVector* MakeVector(size_t size, std::function valueCreator) { - with_lock (Alloc) { - Holders.emplace_front(); - for (size_t i = 0; i < size; ++i) { - Holders.front().emplace_back(valueCreator(i)); - Alloc.Ref().LockObject(Holders.front().back()); - } - return &Holders.front(); - } - } - - template - const TVector* MakeVector(const TVector& values, bool optional = false) { - return MakeVector(values.size(), [&](size_t i) { - NYql::NUdf::TUnboxedValuePod unboxedValue = NYql::NUdf::TUnboxedValuePod(values[i]); - return optional ? unboxedValue.MakeOptional() : unboxedValue; - }); - } - - const TVector* MakeStringVector(const TVector& values, bool optional = false) { - return MakeVector(values.size(), [&](size_t i) { - NYql::NUdf::TUnboxedValuePod stringValue = NKikimr::NMiniKQL::MakeString(values[i]); - return optional ? stringValue.MakeOptional() : stringValue; - }); - } - - const TVector* MakeEmptyVector(size_t size) { - return MakeVector(size, [&](size_t) { - return NYql::NUdf::TUnboxedValuePod(); - }); - } - - NActors::TTestActorRuntime Runtime; - TActorSystemStub ActorSystemStub; - TActorId CompileServiceActorId; - std::unique_ptr Filter; - - NKikimr::NMiniKQL::TScopedAlloc Alloc; - TList> Holders; -}; - -Y_UNIT_TEST_SUITE(TJsonFilterTests) { - Y_UNIT_TEST_F(Simple1, TFixture) { - TMap result; - MakeFilter( - {"a1", "a2", "a@3"}, - {"[DataType; String]", "[DataType; Uint64]", "[OptionalType; [DataType; String]]"}, - "where a2 > 100", - [&](ui64 offset, const TString& json) { - result[offset] = json; - }); - Push({5}, {MakeStringVector({"hello1"}), MakeVector({99}), MakeStringVector({"zapuskaem"}, true)}); - Push({6}, {MakeStringVector({"hello2"}), MakeVector({101}), MakeStringVector({"gusya"}, true)}); - UNIT_ASSERT_VALUES_EQUAL(1, result.size()); - UNIT_ASSERT_VALUES_EQUAL(R"({"a1":"hello2","a2":101,"a@3":"gusya"})", result[6]); - } - - Y_UNIT_TEST_F(Simple2, TFixture) { - TMap result; - MakeFilter( - {"a2", "a1"}, - {"[DataType; Uint64]", "[DataType; String]"}, - "where a2 > 100", - [&](ui64 offset, const TString& json) { - result[offset] = json; - }); - Push({5}, {MakeVector({99}), MakeStringVector({"hello1"})}); - Push({6}, {MakeVector({101}), MakeStringVector({"hello2"})}); - UNIT_ASSERT_VALUES_EQUAL(1, result.size()); - UNIT_ASSERT_VALUES_EQUAL(R"({"a1":"hello2","a2":101})", result[6]); - } - - Y_UNIT_TEST_F(ManyValues, TFixture) { - TMap result; - MakeFilter( - {"a1", "a2", "a3"}, - {"[DataType; String]", "[DataType; Uint64]", "[DataType; String]"}, - "where a2 > 100", - [&](ui64 offset, const TString& json) { - result[offset] = json; - }); - const TString largeString = "abcdefghjkl1234567890+abcdefghjkl1234567890"; - for (ui64 i = 0; i < 5; ++i) { - Push({2 * i, 2 * i + 1}, {MakeStringVector({"hello1", "hello2"}), MakeVector({99, 101}), MakeStringVector({largeString, largeString})}); - UNIT_ASSERT_VALUES_EQUAL_C(i + 1, result.size(), i); - UNIT_ASSERT_VALUES_EQUAL_C(TStringBuilder() << "{\"a1\":\"hello2\",\"a2\":101,\"a3\":\"" << largeString << "\"}", result[2 * i + 1], i); - } - } - - Y_UNIT_TEST_F(NullValues, TFixture) { - TMap result; - MakeFilter( - {"a1", "a2"}, - {"[OptionalType; [DataType; Uint64]]", "[DataType; String]"}, - "where a1 is null", - [&](ui64 offset, const TString& json) { - result[offset] = json; - }); - Push({5}, {MakeEmptyVector(1), MakeStringVector({"str"})}); - UNIT_ASSERT_VALUES_EQUAL(1, result.size()); - UNIT_ASSERT_VALUES_EQUAL(R"({"a1":null,"a2":"str"})", result[5]); - } - - Y_UNIT_TEST_F(PartialPush, TFixture) { - TMap result; - MakeFilter( - {"a1", "a2", "a@3"}, - {"[DataType; String]", "[DataType; Uint64]", "[OptionalType; [DataType; String]]"}, - "where a2 > 50", - [&](ui64 offset, const TString& json) { - result[offset] = json; - }); - Filter->Push({5, 6, 7}, {MakeStringVector({"hello1", "hello2"}), MakeVector({99, 101}), MakeStringVector({"zapuskaem", "gusya"}, true)}, 1, 1); - UNIT_ASSERT_VALUES_EQUAL(1, result.size()); - UNIT_ASSERT_VALUES_EQUAL(R"({"a1":"hello1","a2":99,"a@3":"zapuskaem"})", result[6]); - } -} - -} diff --git a/ydb/core/fq/libs/row_dispatcher/ut/json_parser_ut.cpp b/ydb/core/fq/libs/row_dispatcher/ut/json_parser_ut.cpp deleted file mode 100644 index 80a2b29322c6..000000000000 --- a/ydb/core/fq/libs/row_dispatcher/ut/json_parser_ut.cpp +++ /dev/null @@ -1,319 +0,0 @@ -#include - -#include -#include - -#include - -#include -#include -#include - -#include - -#include - -namespace { - -using namespace NKikimr; -using namespace NFq; - -class TFixture : public NUnitTest::TBaseFixture { - -public: - TFixture() - : Runtime(true) {} - - static void SegmentationFaultHandler(int) { - Cerr << "segmentation fault call stack:" << Endl; - FormatBackTrace(&Cerr); - abort(); - } - - void SetUp(NUnitTest::TTestContext&) override { - NKikimr::EnableYDBBacktraceFormat(); - signal(SIGSEGV, &SegmentationFaultHandler); - - TAutoPtr app = new TAppPrepare(); - Runtime.SetLogBackend(CreateStderrBackend()); - Runtime.SetLogPriority(NKikimrServices::FQ_ROW_DISPATCHER, NLog::PRI_TRACE); - Runtime.Initialize(app->Unwrap()); - } - - void TearDown(NUnitTest::TTestContext& /* context */) override { - if (Parser) { - Parser.reset(); - } - } - - void MakeParser(TVector columns, TVector types, TJsonParser::TCallback callback, ui64 batchSize = 1_MB, ui64 bufferCellCount = 1000) { - Parser = NFq::NewJsonParser(columns, types, callback, batchSize, TDuration::Hours(1), bufferCellCount); - } - - void MakeParser(TVector columns, TJsonParser::TCallback callback) { - MakeParser(columns, TVector(columns.size(), "[DataType; String]"), callback); - } - - void MakeParser(TVector columns, TVector types) { - MakeParser(columns, types, [](ui64, ui64, const TVector>&) {}); - } - - void MakeParser(TVector columns) { - MakeParser(columns, TVector(columns.size(), "[DataType; String]")); - } - - void PushToParser(ui64 offset, const TString& data) { - Parser->AddMessages({GetMessage(offset, data)}); - Parser->Parse(); - } - - static NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage GetMessage(ui64 offset, const TString& data) { - NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessageInformation info(offset, "", 0, TInstant::Zero(), TInstant::Zero(), nullptr, nullptr, 0, ""); - return NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage(data, nullptr, info, nullptr); - } - - TActorSystemStub actorSystemStub; - NActors::TTestActorRuntime Runtime; - std::unique_ptr Parser; -}; - -Y_UNIT_TEST_SUITE(TJsonParserTests) { - Y_UNIT_TEST_F(Simple1, TFixture) { - MakeParser({"a1", "a2"}, {"[DataType; String]", "[OptionalType; [DataType; Uint64]]"}, [](ui64 rowsOffset, ui64 numberRows, const TVector>& result) { - UNIT_ASSERT_VALUES_EQUAL(0, rowsOffset); - UNIT_ASSERT_VALUES_EQUAL(1, numberRows); - UNIT_ASSERT_VALUES_EQUAL(2, result.size()); - UNIT_ASSERT_VALUES_EQUAL("hello1", TString(result[0][0].AsStringRef())); - UNIT_ASSERT_VALUES_EQUAL(101, result[1][0].GetOptionalValue().Get()); - }); - PushToParser(42,R"({"a1": "hello1", "a2": 101, "event": "event1"})"); - } - - Y_UNIT_TEST_F(Simple2, TFixture) { - MakeParser({"a2", "a1"}, [](ui64 rowsOffset, ui64 numberRows, const TVector>& result) { - UNIT_ASSERT_VALUES_EQUAL(0, rowsOffset); - UNIT_ASSERT_VALUES_EQUAL(1, numberRows); - UNIT_ASSERT_VALUES_EQUAL(2, result.size()); - UNIT_ASSERT_VALUES_EQUAL("101", TString(result[0][0].AsStringRef())); - UNIT_ASSERT_VALUES_EQUAL("hello1", TString(result[1][0].AsStringRef())); - }); - PushToParser(42,R"({"a1": "hello1", "a2": "101", "event": "event1"})"); - } - - Y_UNIT_TEST_F(Simple3, TFixture) { - MakeParser({"a1", "a2"}, [](ui64 rowsOffset, ui64 numberRows, const TVector>& result) { - UNIT_ASSERT_VALUES_EQUAL(0, rowsOffset); - UNIT_ASSERT_VALUES_EQUAL(1, numberRows); - UNIT_ASSERT_VALUES_EQUAL(2, result.size()); - UNIT_ASSERT_VALUES_EQUAL("101", TString(result[0][0].AsStringRef())); - UNIT_ASSERT_VALUES_EQUAL("hello1", TString(result[1][0].AsStringRef())); - }); - PushToParser(42,R"({"a2": "hello1", "a1": "101", "event": "event1"})"); - } - - Y_UNIT_TEST_F(Simple4, TFixture) { - MakeParser({"a2", "a1"}, [](ui64 rowsOffset, ui64 numberRows, const TVector>& result) { - UNIT_ASSERT_VALUES_EQUAL(0, rowsOffset); - UNIT_ASSERT_VALUES_EQUAL(1, numberRows); - UNIT_ASSERT_VALUES_EQUAL(2, result.size()); - UNIT_ASSERT_VALUES_EQUAL("hello1", TString(result[0][0].AsStringRef())); - UNIT_ASSERT_VALUES_EQUAL("101", TString(result[1][0].AsStringRef())); - }); - PushToParser(42, R"({"a2": "hello1", "a1": "101", "event": "event1"})"); - } - - Y_UNIT_TEST_F(LargeStrings, TFixture) { - const TString largeString = "abcdefghjkl1234567890+abcdefghjkl1234567890"; - - MakeParser({"col"}, [&](ui64 rowsOffset, ui64 numberRows, const TVector>& result) { - UNIT_ASSERT_VALUES_EQUAL(0, rowsOffset); - UNIT_ASSERT_VALUES_EQUAL(2, numberRows); - UNIT_ASSERT_VALUES_EQUAL(1, result.size()); - UNIT_ASSERT_VALUES_EQUAL(largeString, TString(result[0][0].AsStringRef())); - UNIT_ASSERT_VALUES_EQUAL(largeString, TString(result[0][1].AsStringRef())); - }); - - const TString jsonString = TStringBuilder() << "{\"col\": \"" << largeString << "\"}"; - Parser->AddMessages({ - GetMessage(42, jsonString), - GetMessage(43, jsonString) - }); - Parser->Parse(); - } - - Y_UNIT_TEST_F(ManyValues, TFixture) { - MakeParser({"a1", "a2"}, [&](ui64 rowsOffset, ui64 numberRows, const TVector>& result) { - UNIT_ASSERT_VALUES_EQUAL(0, rowsOffset); - UNIT_ASSERT_VALUES_EQUAL(3, numberRows); - UNIT_ASSERT_VALUES_EQUAL(2, result.size()); - for (size_t i = 0; i < numberRows; ++i) { - UNIT_ASSERT_VALUES_EQUAL_C("hello1", TString(result[0][i].AsStringRef()), i); - UNIT_ASSERT_VALUES_EQUAL_C("101", TString(result[1][i].AsStringRef()), i); - } - }); - - Parser->AddMessages({ - GetMessage(42, R"({"a1": "hello1", "a2": "101", "event": "event1"})"), - GetMessage(43, R"({"a1": "hello1", "a2": "101", "event": "event2"})"), - GetMessage(44, R"({"a2": "101", "a1": "hello1", "event": "event3"})") - }); - Parser->Parse(); - } - - Y_UNIT_TEST_F(MissingFields, TFixture) { - MakeParser({"a1", "a2"}, {"[OptionalType; [DataType; String]]", "[OptionalType; [DataType; Uint64]]"}, [&](ui64 rowsOffset, ui64 numberRows, const TVector>& result) { - UNIT_ASSERT_VALUES_EQUAL(0, rowsOffset); - UNIT_ASSERT_VALUES_EQUAL(3, numberRows); - UNIT_ASSERT_VALUES_EQUAL(2, result.size()); - for (size_t i = 0; i < numberRows; ++i) { - if (i == 2) { - UNIT_ASSERT_C(!result[0][i], i); - } else { - NYql::NUdf::TUnboxedValue value = result[0][i].GetOptionalValue(); - UNIT_ASSERT_VALUES_EQUAL_C("hello1", TString(value.AsStringRef()), i); - } - if (i == 1) { - UNIT_ASSERT_C(!result[1][i], i); - } else { - UNIT_ASSERT_VALUES_EQUAL_C(101, result[1][i].GetOptionalValue().Get(), i); - } - } - }); - - Parser->AddMessages({ - GetMessage(42, R"({"a1": "hello1", "a2": 101 , "event": "event1"})"), - GetMessage(43, R"({"a1": "hello1", "event": "event2"})"), - GetMessage(44, R"({"a2": "101", "a1": null, "event": "event3"})") - }); - Parser->Parse(); - } - - Y_UNIT_TEST_F(NestedTypes, TFixture) { - MakeParser({"nested", "a1"}, {"[OptionalType; [DataType; Json]]", "[DataType; String]"}, [&](ui64 rowsOffset, ui64 numberRows, const TVector>& result) { - UNIT_ASSERT_VALUES_EQUAL(0, rowsOffset); - UNIT_ASSERT_VALUES_EQUAL(4, numberRows); - - UNIT_ASSERT_VALUES_EQUAL(2, result.size()); - UNIT_ASSERT_VALUES_EQUAL("{\"key\": \"value\"}", TString(result[0][0].AsStringRef())); - UNIT_ASSERT_VALUES_EQUAL("hello1", TString(result[1][0].AsStringRef())); - - UNIT_ASSERT_VALUES_EQUAL("[\"key1\", \"key2\"]", TString(result[0][1].AsStringRef())); - UNIT_ASSERT_VALUES_EQUAL("hello2", TString(result[1][1].AsStringRef())); - - UNIT_ASSERT_VALUES_EQUAL("\"some string\"", TString(result[0][2].AsStringRef())); - UNIT_ASSERT_VALUES_EQUAL("hello3", TString(result[1][2].AsStringRef())); - - UNIT_ASSERT_VALUES_EQUAL("123456", TString(result[0][3].AsStringRef())); - UNIT_ASSERT_VALUES_EQUAL("hello4", TString(result[1][3].AsStringRef())); - }); - - Parser->AddMessages({ - GetMessage(42, R"({"a1": "hello1", "nested": {"key": "value"}})"), - GetMessage(43, R"({"a1": "hello2", "nested": ["key1", "key2"]})"), - GetMessage(43, R"({"a1": "hello3", "nested": "some string"})"), - GetMessage(43, R"({"a1": "hello4", "nested": 123456})") - }); - Parser->Parse(); - } - - Y_UNIT_TEST_F(SimpleBooleans, TFixture) { - MakeParser({"a"}, {"[DataType; Bool]"}, [&](ui64 rowsOffset, ui64 numberRows, const TVector>& result) { - UNIT_ASSERT_VALUES_EQUAL(0, rowsOffset); - UNIT_ASSERT_VALUES_EQUAL(2, numberRows); - - UNIT_ASSERT_VALUES_EQUAL(1, result.size()); - UNIT_ASSERT_VALUES_EQUAL(true, result[0][0].Get()); - UNIT_ASSERT_VALUES_EQUAL(false, result[0][1].Get()); - }); - - Parser->AddMessages({ - GetMessage(42, R"({"a": true})"), - GetMessage(43, R"({"a": false})") - }); - Parser->Parse(); - } - - Y_UNIT_TEST_F(ManyBatches, TFixture) { - const TString largeString = "abcdefghjkl1234567890+abcdefghjkl1234567890"; - - ui64 currentOffset = 0; - MakeParser({"col"}, {"[DataType; String]"}, [&](ui64 rowsOffset, ui64 numberRows, const TVector>& result) { - UNIT_ASSERT_VALUES_EQUAL(currentOffset, rowsOffset); - currentOffset++; - - UNIT_ASSERT_VALUES_EQUAL(1, numberRows); - UNIT_ASSERT_VALUES_EQUAL(1, result.size()); - UNIT_ASSERT_VALUES_EQUAL(largeString, TString(result[0][0].AsStringRef())); - }, 1_MB, 1); - - const TString jsonString = TStringBuilder() << "{\"col\": \"" << largeString << "\"}"; - Parser->AddMessages({ - GetMessage(42, jsonString), - GetMessage(43, jsonString) - }); - Parser->Parse(); - } - - Y_UNIT_TEST_F(LittleBatches, TFixture) { - const TString largeString = "abcdefghjkl1234567890+abcdefghjkl1234567890"; - - ui64 currentOffset = 42; - MakeParser({"col"}, {"[DataType; String]"}, [&](ui64 rowsOffset, ui64 numberRows, const TVector>& result) { - UNIT_ASSERT_VALUES_EQUAL(Parser->GetOffsets().size(), 1); - UNIT_ASSERT_VALUES_EQUAL(Parser->GetOffsets().front(), currentOffset); - currentOffset++; - - UNIT_ASSERT_VALUES_EQUAL(0, rowsOffset); - UNIT_ASSERT_VALUES_EQUAL(1, numberRows); - UNIT_ASSERT_VALUES_EQUAL(1, result.size()); - UNIT_ASSERT_VALUES_EQUAL(largeString, TString(result[0][0].AsStringRef())); - }, 10); - - const TString jsonString = TStringBuilder() << "{\"col\": \"" << largeString << "\"}"; - Parser->AddMessages({ - GetMessage(42, jsonString), - GetMessage(43, jsonString) - }); - UNIT_ASSERT_VALUES_EQUAL(Parser->GetNumberValues(), 0); - } - - Y_UNIT_TEST_F(MissingFieldsValidation, TFixture) { - MakeParser({"a1", "a2"}, {"[DataType; String]", "[DataType; Uint64]"}); - UNIT_ASSERT_EXCEPTION_CONTAINS(PushToParser(42, R"({"a1": "hello1", "a2": null, "event": "event1"})"), TJsonParserError, "Failed to parse json string at offset 42, got parsing error for column 'a2' with type [DataType; Uint64], description: (NFq::TJsonParserError) found unexpected null value, expected non optional data type Uint64"); - UNIT_ASSERT_EXCEPTION_CONTAINS(PushToParser(42, R"({"a2": 105, "event": "event1"})"), TJsonParserError, "Failed to parse json messages, found 1 missing values from offset 42 in non optional column 'a1' with type [DataType; String]"); - } - - Y_UNIT_TEST_F(TypeKindsValidation, TFixture) { - UNIT_ASSERT_EXCEPTION_CONTAINS( - MakeParser({"a2", "a1"}, {"[OptionalType; [DataType; String]]", "[ListType; [DataType; String]]"}), - NFq::TJsonParserError, - "Failed to create parser for column 'a1' with type [ListType; [DataType; String]], description: (NFq::TJsonParserError) unsupported type kind List" - ); - } - - Y_UNIT_TEST_F(NumbersValidation, TFixture) { - MakeParser({"a1", "a2"}, {"[OptionalType; [DataType; String]]", "[DataType; Uint8]"}); - UNIT_ASSERT_EXCEPTION_CONTAINS(PushToParser(42, R"({"a1": 456, "a2": 42})"), NFq::TJsonParserError, "Failed to parse json string at offset 42, got parsing error for column 'a1' with type [OptionalType; [DataType; String]], description: (NFq::TJsonParserError) failed to parse data type String from json number (raw: '456'), error: (NFq::TJsonParserError) number value is not expected for data type String"); - UNIT_ASSERT_EXCEPTION_CONTAINS(PushToParser(42, R"({"a1": "456", "a2": -42})"), NFq::TJsonParserError, "Failed to parse json string at offset 42, got parsing error for column 'a2' with type [DataType; Uint8], description: (NFq::TJsonParserError) failed to parse data type Uint8 from json number (raw: '-42'), error: (simdjson::simdjson_error) INCORRECT_TYPE: The JSON element does not have the requested type."); - UNIT_ASSERT_EXCEPTION_CONTAINS(PushToParser(42, R"({"a1": "str", "a2": 99999})"), NFq::TJsonParserError, "Failed to parse json string at offset 42, got parsing error for column 'a2' with type [DataType; Uint8], description: (NFq::TJsonParserError) failed to parse data type Uint8 from json number (raw: '99999'), error: (NFq::TJsonParserError) number is out of range"); - } - - Y_UNIT_TEST_F(NestedJsonValidation, TFixture) { - MakeParser({"a1", "a2"}, {"[OptionalType; [DataType; Json]]", "[OptionalType; [DataType; String]]"}); - UNIT_ASSERT_EXCEPTION_CONTAINS(PushToParser(42, R"({"a1": {"key": "value"}, "a2": {"key2": "value2"}})"), NFq::TJsonParserError, "Failed to parse json string at offset 42, got parsing error for column 'a2' with type [OptionalType; [DataType; String]], description: (NFq::TJsonParserError) found unexpected nested value (raw: '{\"key2\": \"value2\"}'), expected data type String, please use Json type for nested values"); - UNIT_ASSERT_EXCEPTION_CONTAINS(PushToParser(42, R"({"a1": {"key" "value"}, "a2": "str"})"), NFq::TJsonParserError, "Failed to parse json string at offset 42, got parsing error for column 'a1' with type [OptionalType; [DataType; Json]], description: (simdjson::simdjson_error) TAPE_ERROR: The JSON document has an improper structure: missing or superfluous commas, braces, missing keys, etc."); - } - - Y_UNIT_TEST_F(BoolsValidation, TFixture) { - MakeParser({"a1", "a2"}, {"[OptionalType; [DataType; String]]", "[DataType; Bool]"}); - UNIT_ASSERT_EXCEPTION_CONTAINS(PushToParser(42, R"({"a1": true, "a2": false})"), NFq::TJsonParserError, "Failed to parse json string at offset 42, got parsing error for column 'a1' with type [OptionalType; [DataType; String]], description: (NFq::TJsonParserError) found unexpected bool value, expected data type String"); - } - - Y_UNIT_TEST_F(ThrowExceptionByError, TFixture) { - MakeParser({"a"}); - UNIT_ASSERT_EXCEPTION_CONTAINS(PushToParser(42, R"(ydb)"), simdjson::simdjson_error, "INCORRECT_TYPE: The JSON element does not have the requested type."); - UNIT_ASSERT_EXCEPTION_CONTAINS(PushToParser(42, R"({"a": "value1"} {"a": "value2"})"), NFq::TJsonParserError, "Failed to parse json messages, expected 1 json rows from offset 42 but got 2"); - } -} - -} diff --git a/ydb/core/fq/libs/row_dispatcher/ut/row_dispatcher_ut.cpp b/ydb/core/fq/libs/row_dispatcher/ut/row_dispatcher_ut.cpp index 912257967d47..cb0f1caa1868 100644 --- a/ydb/core/fq/libs/row_dispatcher/ut/row_dispatcher_ut.cpp +++ b/ydb/core/fq/libs/row_dispatcher/ut/row_dispatcher_ut.cpp @@ -27,6 +27,7 @@ struct TTestActorFactory : public NFq::NRowDispatcher::IActorFactory { } NActors::TActorId RegisterTopicSession( + const TString& /*readGroup*/, const TString& /*topicPath*/, const TString& /*endpoint*/, const TString& /*database*/, @@ -37,6 +38,7 @@ struct TTestActorFactory : public NFq::NRowDispatcher::IActorFactory { NYdb::TDriver /*driver*/, std::shared_ptr /*credentialsProviderFactory*/, const ::NMonitoring::TDynamicCounterPtr& /*counters*/, + const ::NMonitoring::TDynamicCounterPtr& /*counters*/, const NYql::IPqGateway::TPtr& /*pqGateway*/, ui64 /*maxBufferSize*/) const override { auto actorId = Runtime.AllocateEdgeActor(); @@ -49,17 +51,29 @@ struct TTestActorFactory : public NFq::NRowDispatcher::IActorFactory { }; class TFixture : public NUnitTest::TBaseFixture { - + const ui64 NodesCount = 2; public: TFixture() - : Runtime(1) {} + : Runtime(NodesCount) {} void SetUp(NUnitTest::TTestContext&) override { + TIntrusivePtr nameserverTable(new TTableNameserverSetup()); + TPortManager pm; + for (ui32 i = 0; i < NodesCount; ++i) { + nameserverTable->StaticNodeTable[Runtime.GetNodeId(i)] = std::pair("127.0.0." + std::to_string(i + 1), pm.GetPort(12001 + i)); + } + const TActorId nameserviceId = GetNameserviceActorId(); + for (ui32 i = 0; i < NodesCount; ++i) { + TActorSetupCmd nameserviceSetup(CreateNameserverTable(nameserverTable), TMailboxType::Simple, 0); + Runtime.AddLocalService(nameserviceId, std::move(nameserviceSetup), i); + } + TAutoPtr app = new TAppPrepare(); Runtime.Initialize(app->Unwrap()); Runtime.SetLogPriority(NKikimrServices::FQ_ROW_DISPATCHER, NLog::PRI_TRACE); NConfig::TRowDispatcherConfig config; config.SetEnabled(true); + config.SetSendStatusPeriodSec(1); NConfig::TRowDispatcherCoordinatorConfig& coordinatorConfig = *config.MutableCoordinator(); coordinatorConfig.SetCoordinationNodePath("RowDispatcher"); auto& database = *coordinatorConfig.MutableDatabase(); @@ -76,6 +90,7 @@ class TFixture : public NUnitTest::TBaseFixture { EdgeActor = Runtime.AllocateEdgeActor(); ReadActorId1 = Runtime.AllocateEdgeActor(); ReadActorId2 = Runtime.AllocateEdgeActor(); + ReadActorId3 = Runtime.AllocateEdgeActor(1); TestActorFactory = MakeIntrusive(Runtime); NYql::TPqGatewayServices pqServices( @@ -93,6 +108,7 @@ class TFixture : public NUnitTest::TBaseFixture { "Tenant", TestActorFactory, MakeIntrusive(), + MakeIntrusive(), CreatePqNativeGateway(pqServices) ).release()); @@ -109,7 +125,8 @@ class TFixture : public NUnitTest::TBaseFixture { NYql::NPq::NProto::TDqPqTopicSource BuildPqTopicSourceSettings( TString endpoint, TString database, - TString topic) + TString topic, + TString readGroup) { NYql::NPq::NProto::TDqPqTopicSource settings; settings.SetTopicPath(topic); @@ -117,27 +134,34 @@ class TFixture : public NUnitTest::TBaseFixture { settings.SetEndpoint(endpoint); settings.MutableToken()->SetName("token"); settings.SetDatabase(database); + settings.SetReadGroup(readGroup); return settings; } - void MockAddSession(const NYql::NPq::NProto::TDqPqTopicSource& source, ui64 partitionId, TActorId readActorId, ui64 generation = 1) { + void MockAddSession(const NYql::NPq::NProto::TDqPqTopicSource& source, const std::set& partitionIds, TActorId readActorId, ui64 generation = 1) { auto event = new NFq::TEvRowDispatcher::TEvStartSession( source, - partitionId, // partitionId + partitionIds, "Token", - Nothing(), // readOffset, + {}, // readOffset, 0, // StartingMessageTimestamp; "QueryId"); + event->Record.MutableTransportMeta()->SetSeqNo(1); Runtime.Send(new IEventHandle(RowDispatcher, readActorId, event, 0, generation)); } - void MockStopSession(const NYql::NPq::NProto::TDqPqTopicSource& source, ui64 partitionId, TActorId readActorId) { + void MockStopSession(const NYql::NPq::NProto::TDqPqTopicSource& source, TActorId readActorId) { auto event = std::make_unique(); event->Record.MutableSource()->CopyFrom(source); - event->Record.SetPartitionId(partitionId); + event->Record.MutableTransportMeta()->SetSeqNo(1); Runtime.Send(new IEventHandle(RowDispatcher, readActorId, event.release(), 0, 1)); } + void MockNoSession(TActorId readActorId, ui64 generation) { + auto event = std::make_unique(); + Runtime.Send(new IEventHandle(RowDispatcher, readActorId, event.release(), 0, generation)); + } + void MockNewDataArrived(ui64 partitionId, TActorId topicSessionId, TActorId readActorId) { auto event = std::make_unique(); event->Record.SetPartitionId(partitionId); @@ -152,16 +176,16 @@ class TFixture : public NUnitTest::TBaseFixture { Runtime.Send(new IEventHandle(RowDispatcher, topicSessionId, event.release(), 0, generation)); } - void MockSessionError(ui64 partitionId, TActorId topicSessionId, TActorId readActorId) { + void MockSessionError(TActorId topicSessionId, TActorId readActorId) { auto event = std::make_unique(); - event->Record.SetPartitionId(partitionId); event->ReadActorId = readActorId; Runtime.Send(new IEventHandle(RowDispatcher, topicSessionId, event.release())); } - void MockGetNextBatch(ui64 partitionId, TActorId readActorId, ui64 generation) { + void MockGetNextBatch(ui64 partitionId, TActorId readActorId, ui64 generation, ui64 seqNo = 2) { auto event = std::make_unique(); event->Record.SetPartitionId(partitionId); + event->Record.MutableTransportMeta()->SetSeqNo(seqNo); Runtime.Send(new IEventHandle(RowDispatcher, readActorId, event.release(), 0, generation)); } @@ -175,10 +199,9 @@ class TFixture : public NUnitTest::TBaseFixture { UNIT_ASSERT(eventHolder.Get() != nullptr); } - void ExpectStopSession(NActors::TActorId actorId, ui64 partitionId) { + void ExpectStopSession(NActors::TActorId actorId) { auto eventHolder = Runtime.GrabEdgeEvent(actorId); UNIT_ASSERT(eventHolder.Get() != nullptr); - UNIT_ASSERT(eventHolder->Get()->Record.GetPartitionId() == partitionId); } void ExpectGetNextBatch(NActors::TActorId topicSessionId, ui64 partitionId) { @@ -204,10 +227,9 @@ class TFixture : public NUnitTest::TBaseFixture { UNIT_ASSERT(eventHolder.Get() != nullptr); } - void ExpectSessionError(NActors::TActorId readActorId, ui64 partitionId) { + void ExpectSessionError(NActors::TActorId readActorId) { auto eventHolder = Runtime.GrabEdgeEvent(readActorId); UNIT_ASSERT(eventHolder.Get() != nullptr); - UNIT_ASSERT(eventHolder->Get()->Record.GetPartitionId() == partitionId); } NActors::TActorId ExpectRegisterTopicSession() { @@ -215,11 +237,11 @@ class TFixture : public NUnitTest::TBaseFixture { return actorId; } - void ProcessData(NActors::TActorId readActorId, ui64 partId, NActors::TActorId topicSessionActorId, ui64 generation = 1) { + void ProcessData(NActors::TActorId readActorId, ui64 partId, NActors::TActorId topicSessionActorId, ui64 generation = 1, ui64 seqNo = 1) { MockNewDataArrived(partId, topicSessionActorId, readActorId); ExpectNewDataArrived(readActorId, partId); - MockGetNextBatch(partId, readActorId, generation); + MockGetNextBatch(partId, readActorId, generation, seqNo); ExpectGetNextBatch(topicSessionActorId, partId); MockMessageBatch(partId, topicSessionActorId, readActorId, generation); @@ -234,56 +256,58 @@ class TFixture : public NUnitTest::TBaseFixture { NActors::TActorId EdgeActor; NActors::TActorId ReadActorId1; NActors::TActorId ReadActorId2; + NActors::TActorId ReadActorId3; TIntrusivePtr TestActorFactory; - NYql::NPq::NProto::TDqPqTopicSource Source1 = BuildPqTopicSourceSettings("Endpoint1", "Database1", "topic"); - NYql::NPq::NProto::TDqPqTopicSource Source2 = BuildPqTopicSourceSettings("Endpoint2", "Database1", "topic"); + NYql::NPq::NProto::TDqPqTopicSource Source1 = BuildPqTopicSourceSettings("Endpoint1", "Database1", "topic", "connection_id1"); + NYql::NPq::NProto::TDqPqTopicSource Source2 = BuildPqTopicSourceSettings("Endpoint2", "Database1", "topic", "connection_id1"); + NYql::NPq::NProto::TDqPqTopicSource Source1Connection2 = BuildPqTopicSourceSettings("Endpoint1", "Database1", "topic", "connection_id2"); - ui64 PartitionId0 = 0; - ui64 PartitionId1 = 1; + ui32 PartitionId0 = 0; + ui32 PartitionId1 = 1; }; Y_UNIT_TEST_SUITE(RowDispatcherTests) { Y_UNIT_TEST_F(OneClientOneSession, TFixture) { - MockAddSession(Source1, PartitionId0, ReadActorId1); + MockAddSession(Source1, {PartitionId0}, ReadActorId1); auto topicSessionId = ExpectRegisterTopicSession(); ExpectStartSessionAck(ReadActorId1); ExpectStartSession(topicSessionId); ProcessData(ReadActorId1, PartitionId0, topicSessionId); - MockStopSession(Source1, PartitionId0, ReadActorId1); - ExpectStopSession(topicSessionId, PartitionId0); + MockStopSession(Source1, ReadActorId1); + ExpectStopSession(topicSessionId); } Y_UNIT_TEST_F(TwoClientOneSession, TFixture) { - MockAddSession(Source1, PartitionId0, ReadActorId1); + MockAddSession(Source1, {PartitionId0}, ReadActorId1); auto topicSessionId = ExpectRegisterTopicSession(); ExpectStartSessionAck(ReadActorId1); ExpectStartSession(topicSessionId); - MockAddSession(Source1, PartitionId0, ReadActorId2); + MockAddSession(Source1, {PartitionId0}, ReadActorId2); ExpectStartSessionAck(ReadActorId2); ExpectStartSession(topicSessionId); ProcessData(ReadActorId1, PartitionId0, topicSessionId); ProcessData(ReadActorId2, PartitionId0, topicSessionId); - MockSessionError(PartitionId0, topicSessionId, ReadActorId1); - ExpectSessionError(ReadActorId1, PartitionId0); + MockSessionError(topicSessionId, ReadActorId1); + ExpectSessionError(ReadActorId1); - MockSessionError(PartitionId0, topicSessionId, ReadActorId2); - ExpectSessionError(ReadActorId2, PartitionId0); + MockSessionError(topicSessionId, ReadActorId2); + ExpectSessionError(ReadActorId2); } Y_UNIT_TEST_F(SessionError, TFixture) { - MockAddSession(Source1, PartitionId0, ReadActorId1); + MockAddSession(Source1, {PartitionId0}, ReadActorId1); auto topicSessionId = ExpectRegisterTopicSession(); ExpectStartSessionAck(ReadActorId1); ExpectStartSession(topicSessionId); - MockSessionError(PartitionId0, topicSessionId, ReadActorId1); - ExpectSessionError(ReadActorId1, PartitionId0); + MockSessionError(topicSessionId, ReadActorId1); + ExpectSessionError(ReadActorId1); } Y_UNIT_TEST_F(CoordinatorSubscribe, TFixture) { @@ -314,24 +338,18 @@ Y_UNIT_TEST_SUITE(RowDispatcherTests) { Y_UNIT_TEST_F(TwoClients4Sessions, TFixture) { - MockAddSession(Source1, PartitionId0, ReadActorId1); + MockAddSession(Source1, {PartitionId0, PartitionId1}, ReadActorId1); auto topicSession1 = ExpectRegisterTopicSession(); - ExpectStartSessionAck(ReadActorId1); - ExpectStartSession(topicSession1); - - MockAddSession(Source1, PartitionId1, ReadActorId1); auto topicSession2 = ExpectRegisterTopicSession(); ExpectStartSessionAck(ReadActorId1); + ExpectStartSession(topicSession1); ExpectStartSession(topicSession2); - MockAddSession(Source2, PartitionId0, ReadActorId2); + MockAddSession(Source2, {PartitionId0, PartitionId1}, ReadActorId2); auto topicSession3 = ExpectRegisterTopicSession(); - ExpectStartSessionAck(ReadActorId2); - ExpectStartSession(topicSession3); - - MockAddSession(Source2, PartitionId1, ReadActorId2); auto topicSession4 = ExpectRegisterTopicSession(); ExpectStartSessionAck(ReadActorId2); + ExpectStartSession(topicSession3); ExpectStartSession(topicSession4); ProcessData(ReadActorId1, PartitionId0, topicSession1); @@ -339,60 +357,114 @@ Y_UNIT_TEST_SUITE(RowDispatcherTests) { ProcessData(ReadActorId2, PartitionId0, topicSession3); ProcessData(ReadActorId2, PartitionId1, topicSession4); - MockSessionError(PartitionId0, topicSession1, ReadActorId1); - ExpectSessionError(ReadActorId1, PartitionId0); + MockSessionError(topicSession1, ReadActorId1); + ExpectSessionError(ReadActorId1); ProcessData(ReadActorId1, PartitionId1, topicSession2); ProcessData(ReadActorId2, PartitionId0, topicSession3); ProcessData(ReadActorId2, PartitionId1, topicSession4); - MockStopSession(Source1, PartitionId1, ReadActorId1); - ExpectStopSession(topicSession2, PartitionId1); + MockStopSession(Source1, ReadActorId1); + ExpectStopSession(topicSession2); - MockStopSession(Source2, PartitionId0, ReadActorId2); - ExpectStopSession(topicSession3, PartitionId0); + MockStopSession(Source2, ReadActorId2); + ExpectStopSession(topicSession3); - MockStopSession(Source2, PartitionId1, ReadActorId2); - ExpectStopSession(topicSession4, PartitionId1); + MockStopSession(Source2, ReadActorId2); + ExpectStopSession(topicSession4); // Ignore data after StopSession MockMessageBatch(PartitionId1, topicSession4, ReadActorId2, 1); } Y_UNIT_TEST_F(ReinitConsumerIfNewGeneration, TFixture) { - MockAddSession(Source1, PartitionId0, ReadActorId1, 1); + MockAddSession(Source1, {PartitionId0}, ReadActorId1, 1); auto topicSessionId = ExpectRegisterTopicSession(); ExpectStartSessionAck(ReadActorId1); ExpectStartSession(topicSessionId); ProcessData(ReadActorId1, PartitionId0, topicSessionId); // ignore StartSession with same generation - MockAddSession(Source1, PartitionId0, ReadActorId1, 1); + MockAddSession(Source1, {PartitionId0}, ReadActorId1, 1); // reinit consumer - MockAddSession(Source1, PartitionId0, ReadActorId1, 2); + MockAddSession(Source1, {PartitionId0}, ReadActorId1, 2); ExpectStartSessionAck(ReadActorId1, 2); } Y_UNIT_TEST_F(HandleTEvUndelivered, TFixture) { - MockAddSession(Source1, PartitionId0, ReadActorId1, 1); + MockAddSession(Source1, {PartitionId0, PartitionId1}, ReadActorId1, 1); auto topicSession1 = ExpectRegisterTopicSession(); + auto topicSession2 = ExpectRegisterTopicSession(); ExpectStartSessionAck(ReadActorId1, 1); ExpectStartSession(topicSession1); + ExpectStartSession(topicSession2); - MockAddSession(Source1, PartitionId1, ReadActorId1, 2); - auto topicSession2 = ExpectRegisterTopicSession(); - ExpectStartSessionAck(ReadActorId1, 2); + MockAddSession(Source1, {PartitionId0, PartitionId1}, ReadActorId2, 1); + ExpectStartSessionAck(ReadActorId2, 1); + ExpectStartSession(topicSession1); ExpectStartSession(topicSession2); ProcessData(ReadActorId1, PartitionId0, topicSession1, 1); - ProcessData(ReadActorId1, PartitionId1, topicSession2, 2); - - MockUndelivered(ReadActorId1, 2); - ExpectStopSession(topicSession2, PartitionId1); + ProcessData(ReadActorId1, PartitionId1, topicSession2, 1); + ProcessData(ReadActorId2, PartitionId0, topicSession1, 1); + ProcessData(ReadActorId2, PartitionId1, topicSession2, 1); MockUndelivered(ReadActorId1, 1); - ExpectStopSession(topicSession1, PartitionId0); + ExpectStopSession(topicSession1); + ExpectStopSession(topicSession2); + + MockUndelivered(ReadActorId2, 1); + ExpectStopSession(topicSession1); + ExpectStopSession(topicSession2); + } + + Y_UNIT_TEST_F(TwoClientTwoConnection, TFixture) { + MockAddSession(Source1, {PartitionId0}, ReadActorId1); + auto session1 = ExpectRegisterTopicSession(); + ExpectStartSessionAck(ReadActorId1); + ExpectStartSession(session1); + + MockAddSession(Source1Connection2, {PartitionId0}, ReadActorId2); + auto session2 = ExpectRegisterTopicSession(); + ExpectStartSessionAck(ReadActorId2); + ExpectStartSession(session2); + + ProcessData(ReadActorId1, PartitionId0, session1); + ProcessData(ReadActorId2, PartitionId0, session2); + + MockStopSession(Source1, ReadActorId1); + ExpectStopSession(session1); + + MockStopSession(Source1Connection2, ReadActorId2); + ExpectStopSession(session2); + } + + Y_UNIT_TEST_F(ProcessNoSession, TFixture) { + ui64 generation = 42; + MockAddSession(Source1, {PartitionId0}, ReadActorId3, generation); + auto topicSessionId = ExpectRegisterTopicSession(); + ExpectStartSessionAck(ReadActorId3, generation); + ExpectStartSession(topicSessionId); + ProcessData(ReadActorId3, PartitionId0, topicSessionId, generation, 2); + + MockNoSession(ReadActorId3, generation - 1); // Ignore NoSession with wrong generation. + ProcessData(ReadActorId3, PartitionId0, topicSessionId, generation, 3); + + MockNoSession(ReadActorId3, generation); + ExpectStopSession(topicSessionId); + } + + Y_UNIT_TEST_F(IgnoreWrongPartitionId, TFixture) { + MockAddSession(Source1, {PartitionId0}, ReadActorId1); + auto topicSessionId = ExpectRegisterTopicSession(); + ExpectStartSessionAck(ReadActorId1); + ExpectStartSession(topicSessionId); + + MockNewDataArrived(PartitionId1, topicSessionId, ReadActorId1); + + MockStopSession(Source1, ReadActorId1); + ExpectStopSession(topicSessionId); } } diff --git a/ydb/core/fq/libs/row_dispatcher/ut/topic_session_ut.cpp b/ydb/core/fq/libs/row_dispatcher/ut/topic_session_ut.cpp index 382f32957ff2..75af42af21bd 100644 --- a/ydb/core/fq/libs/row_dispatcher/ut/topic_session_ut.cpp +++ b/ydb/core/fq/libs/row_dispatcher/ut/topic_session_ut.cpp @@ -1,76 +1,41 @@ -#include - #include #include #include #include -#include +#include #include #include #include #include +#include #include #include #include +#include + +namespace NFq::NRowDispatcher::NTests { namespace { using namespace NKikimr; -using namespace NFq; using namespace NYql::NDq; -const ui64 TimeoutBeforeStartSessionSec = 3; -const ui64 GrabTimeoutSec = 4 * TimeoutBeforeStartSessionSec; - -class TPurecalcCompileServiceMock : public NActors::TActor { - using TBase = NActors::TActor; +constexpr ui64 TimeoutBeforeStartSessionSec = 3; +constexpr ui64 GrabTimeoutSec = 4 * TimeoutBeforeStartSessionSec; +static_assert(GrabTimeoutSec <= WAIT_TIMEOUT.Seconds()); +template +class TFixture : public NTests::TBaseFixture { public: - TPurecalcCompileServiceMock(TActorId owner) - : TBase(&TPurecalcCompileServiceMock::StateFunc) - , Owner(owner) - , ProgramFactory(NYql::NPureCalc::MakeProgramFactory()) - {} - - STRICT_STFUNC(StateFunc, - hFunc(TEvRowDispatcher::TEvPurecalcCompileRequest, Handle); - ) - - void Handle(TEvRowDispatcher::TEvPurecalcCompileRequest::TPtr& ev) { - IProgramHolder::TPtr programHolder = std::move(ev->Get()->ProgramHolder); - - try { - programHolder->CreateProgram(ProgramFactory); - } catch (const NYql::NPureCalc::TCompileError& e) { - UNIT_ASSERT_C(false, "Failed to compile purecalc filter: sql: " << e.GetYql() << ", error: " << e.GetIssues()); - } - - Send(ev->Sender, new TEvRowDispatcher::TEvPurecalcCompileResponse(std::move(programHolder)), 0, ev->Cookie); - Send(Owner, new NActors::TEvents::TEvPing()); - } - -private: - const TActorId Owner; - const NYql::NPureCalc::IProgramFactoryPtr ProgramFactory; -}; + using TBase = NTests::TBaseFixture; -class TFixture : public NUnitTest::TBaseFixture { public: - TFixture() - : Runtime(true) - {} - - void SetUp(NUnitTest::TTestContext&) override { - TAutoPtr app = new TAppPrepare(); - Runtime.Initialize(app->Unwrap()); - Runtime.SetLogPriority(NKikimrServices::FQ_ROW_DISPATCHER, NLog::PRI_TRACE); - Runtime.SetDispatchTimeout(TDuration::Seconds(5)); - - NKikimr::EnableYDBBacktraceFormat(); + void SetUp(NUnitTest::TTestContext& ctx) override { + TBase::SetUp(ctx); ReadActorId1 = Runtime.AllocateEdgeActor(); ReadActorId2 = Runtime.AllocateEdgeActor(); @@ -79,6 +44,7 @@ class TFixture : public NUnitTest::TBaseFixture { } void Init(const TString& topicPath, ui64 maxSessionUsedMemory = std::numeric_limits::max()) { + TopicPath = topicPath; Config.SetTimeoutBeforeStartSessionSec(TimeoutBeforeStartSessionSec); Config.SetMaxSessionUsedMemory(maxSessionUsedMemory); Config.SetSendStatusPeriodSec(2); @@ -95,9 +61,15 @@ class TFixture : public NUnitTest::TBaseFixture { nullptr); CompileNotifier = Runtime.AllocateEdgeActor(); - const auto compileServiceActorId = Runtime.Register(new TPurecalcCompileServiceMock(CompileNotifier)); + const auto compileServiceActorId = Runtime.Register(CreatePurecalcCompileServiceMock(CompileNotifier)); + + if (MockTopicSession) { + PqGatewayNotifier = Runtime.AllocateEdgeActor(); + MockPqGateway = CreateMockPqGateway(Runtime, PqGatewayNotifier); + } TopicSession = Runtime.Register(NewTopicSession( + "read_group", topicPath, GetDefaultPqEndpoint(), GetDefaultPqDatabase(), @@ -108,25 +80,23 @@ class TFixture : public NUnitTest::TBaseFixture { Driver, CredentialsProviderFactory, MakeIntrusive(), - CreatePqNativeGateway(pqServices), + MakeIntrusive(), + !MockTopicSession ? CreatePqNativeGateway(pqServices) : MockPqGateway, 16000000 ).release()); Runtime.EnableScheduleForActor(TopicSession); - - TDispatchOptions options; - options.FinalEvents.emplace_back(NActors::TEvents::TSystem::Bootstrap, 1); - UNIT_ASSERT(Runtime.DispatchEvents(options)); - } - - void TearDown(NUnitTest::TTestContext& /* context */) override { } void StartSession(TActorId readActorId, const NYql::NPq::NProto::TDqPqTopicSource& source, TMaybe readOffset = Nothing(), bool expectedError = false) { + std::map readOffsets; + if (readOffset) { + readOffsets[PartitionId] = *readOffset; + } auto event = new NFq::TEvRowDispatcher::TEvStartSession( source, - PartitionId, + {PartitionId}, "Token", - readOffset, // readOffset, + readOffsets, 0, // StartingMessageTimestamp; "QueryId"); Runtime.Send(new IEventHandle(TopicSession, readActorId, event)); @@ -137,13 +107,19 @@ class TFixture : public NUnitTest::TBaseFixture { const auto ping = Runtime.GrabEdgeEvent(CompileNotifier); UNIT_ASSERT_C(ping, "Compilation is not performed for predicate: " << predicate); } + + if (MockTopicSession) { + Runtime.GrabEdgeEvent(PqGatewayNotifier, TDuration::Seconds(GrabTimeoutSec)); + MockPqGateway->AddEvent(TopicPath, NYdb::NTopic::TReadSessionEvent::TStartPartitionSessionEvent(nullptr, 0, 0), 0); + } } - NYql::NPq::NProto::TDqPqTopicSource BuildSource(TString topic, bool emptyPredicate = false, const TString& consumer = DefaultPqConsumer) { + NYql::NPq::NProto::TDqPqTopicSource BuildSource(bool emptyPredicate = false, const TString& consumer = DefaultPqConsumer) { NYql::NPq::NProto::TDqPqTopicSource settings; settings.SetEndpoint(GetDefaultPqEndpoint()); - settings.SetTopicPath(topic); + settings.SetTopicPath(TopicPath); settings.SetConsumerName(consumer); + settings.SetFormat("json_each_row"); settings.MutableToken()->SetName("token"); settings.SetDatabase(GetDefaultPqDatabase()); settings.AddColumns("dt"); @@ -159,32 +135,31 @@ class TFixture : public NUnitTest::TBaseFixture { void StopSession(NActors::TActorId readActorId, const NYql::NPq::NProto::TDqPqTopicSource& source) { auto event = std::make_unique(); *event->Record.MutableSource() = source; - event->Record.SetPartitionId(PartitionId); Runtime.Send(new IEventHandle(TopicSession, readActorId, event.release())); } - void ExpectMessageBatch(NActors::TActorId readActorId, const std::vector& expected) { + void ExpectMessageBatch(NActors::TActorId readActorId, const TBatch& expected) { Runtime.Send(new IEventHandle(TopicSession, readActorId, new TEvRowDispatcher::TEvGetNextBatch())); auto eventHolder = Runtime.GrabEdgeEvent(RowDispatcherActorId, TDuration::Seconds(GrabTimeoutSec)); UNIT_ASSERT(eventHolder.Get() != nullptr); UNIT_ASSERT_VALUES_EQUAL(eventHolder->Get()->ReadActorId, readActorId); - UNIT_ASSERT_VALUES_EQUAL(expected.size(), eventHolder->Get()->Record.MessagesSize()); - for (size_t i = 0; i < expected.size(); ++i) { - NFq::NRowDispatcherProto::TEvMessage message = eventHolder->Get()->Record.GetMessages(i); - std::cerr << "message.GetJson() " << message.GetJson() << std::endl; - UNIT_ASSERT_VALUES_EQUAL(expected[i], message.GetJson()); - } + UNIT_ASSERT_VALUES_EQUAL(1, eventHolder->Get()->Record.MessagesSize()); + + NFq::NRowDispatcherProto::TEvMessage message = eventHolder->Get()->Record.GetMessages(0); + UNIT_ASSERT_VALUES_EQUAL(message.OffsetsSize(), expected.Rows.size()); + CheckMessageBatch(eventHolder->Get()->GetPayload(message.GetPayloadId()), expected); } - TString ExpectSessionError(NActors::TActorId readActorId, TMaybe message = Nothing()) { + void ExpectSessionError(NActors::TActorId readActorId, TStatusCode statusCode, TString message = "") { auto eventHolder = Runtime.GrabEdgeEvent(RowDispatcherActorId, TDuration::Seconds(GrabTimeoutSec)); UNIT_ASSERT(eventHolder.Get() != nullptr); UNIT_ASSERT_VALUES_EQUAL(eventHolder->Get()->ReadActorId, readActorId); - if (message) { - UNIT_ASSERT_STRING_CONTAINS(TString(eventHolder->Get()->Record.GetMessage()), *message); - } - return eventHolder->Get()->Record.GetMessage(); + + const auto& record = eventHolder->Get()->Record; + NYql::TIssues issues; + NYql::IssuesFromMessage(record.GetIssues(), issues); + NTests::CheckError(TStatus::Fail(record.GetStatusCode(), std::move(issues)), statusCode, message); } void ExpectNewDataArrived(TSet readActorIds) { @@ -202,31 +177,100 @@ class TFixture : public NUnitTest::TBaseFixture { auto eventHolder = Runtime.GrabEdgeEvent(RowDispatcherActorId, TDuration::Seconds(GrabTimeoutSec)); UNIT_ASSERT(eventHolder.Get() != nullptr); UNIT_ASSERT_VALUES_EQUAL(eventHolder->Get()->ReadActorId, readActorId); - return eventHolder->Get()->Record.MessagesSize(); + + size_t numberMessages = 0; + for (const auto& message : eventHolder->Get()->Record.GetMessages()) { + numberMessages += message.OffsetsSize(); + } + + return numberMessages; } - void ExpectStatisticToReadActor(TSet readActorIds) { - size_t count = readActorIds.size(); - for (size_t i = 0; i < count; ++i) { - auto eventHolder = Runtime.GrabEdgeEvent(RowDispatcherActorId, TDuration::Seconds(GrabTimeoutSec)); + void ExpectStatistics(TMap clients) { + auto check = [&]() -> bool { + auto eventHolder = Runtime.GrabEdgeEvent(RowDispatcherActorId, TDuration::Seconds(GrabTimeoutSec)); UNIT_ASSERT(eventHolder.Get() != nullptr); - UNIT_ASSERT(readActorIds.contains(eventHolder->Get()->ReadActorId)); - readActorIds.erase(eventHolder->Get()->ReadActorId); + if (clients.size() != eventHolder->Get()->Stat.Clients.size()) { + return false; + } + for (const auto& client : eventHolder->Get()->Stat.Clients) { + if (!clients.contains(client.ReadActorId)) { + return false; + } + if (clients[client.ReadActorId] != client.Offset) { + return false; + } + } + return true; + }; + auto start = TInstant::Now(); + while (TInstant::Now() - start < TDuration::Seconds(5)) { + if (check()) { + return; + } } + UNIT_FAIL("ExpectStatistics timeout"); + } + + static TRow JsonMessage(ui64 index) { + return TRow().AddUint64(100 * index).AddString(TStringBuilder() << "value" << index); + } + + using TMessageInformation = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessageInformation; + using TMessage = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage; + + TMessageInformation MakeNextMessageInformation(size_t offset, size_t uncompressedSize) { + auto now = TInstant::Now(); + TMessageInformation msgInfo( + offset, + "ProducerId", + 0, + now, + now, + MakeIntrusive(), + MakeIntrusive(), + uncompressedSize, + "messageGroupId" + ); + return msgInfo; } - NActors::TTestActorRuntime Runtime; - TActorSystemStub ActorSystemStub; + void PQWrite( + const std::vector& sequence, + ui64 firstMessageOffset = 0) { + if (!MockTopicSession) { + NYql::NDq::PQWrite(sequence, TopicPath, GetDefaultPqEndpoint()); + } else { + ui64 offset = firstMessageOffset; + TVector msgs; + size_t size = 0; + for (const auto& s : sequence) { + TMessage msg(s, nullptr, MakeNextMessageInformation(offset++, s.size()), CreatePartitionSession()); + msgs.emplace_back(msg); + size += s.size(); + } + MockPqGateway->AddEvent(TopicPath, NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent(msgs, {}, CreatePartitionSession()), size); + } + } + + void PassAway() { + Runtime.Send(new IEventHandle(TopicSession, RowDispatcherActorId, new NActors::TEvents::TEvPoisonPill)); + } + +public: + TString TopicPath; NActors::TActorId TopicSession; NActors::TActorId RowDispatcherActorId; NActors::TActorId CompileNotifier; + NActors::TActorId PqGatewayNotifier; NYdb::TDriver Driver = NYdb::TDriver(NYdb::TDriverConfig().SetLog(CreateLogBackend("cerr"))); std::shared_ptr CredentialsProviderFactory; NActors::TActorId ReadActorId1; NActors::TActorId ReadActorId2; NActors::TActorId ReadActorId3; - ui64 PartitionId = 0; + ui32 PartitionId = 0; NConfig::TRowDispatcherConfig Config; + TIntrusivePtr MockPqGateway; const TString Json1 = "{\"dt\":100,\"value\":\"value1\"}"; const TString Json2 = "{\"dt\":200,\"value\":\"value2\"}"; @@ -234,192 +278,221 @@ class TFixture : public NUnitTest::TBaseFixture { const TString Json4 = "{\"dt\":400,\"value\":\"value4\"}"; }; +using TRealTopicFixture = TFixture; +using TMockTopicFixture = TFixture; + +} // anonymous namespace + Y_UNIT_TEST_SUITE(TopicSessionTests) { - Y_UNIT_TEST_F(TwoSessionsWithoutOffsets, TFixture) { + + Y_UNIT_TEST_F(TwoSessionsWithoutOffsets, TRealTopicFixture) { const TString topicName = "topic1"; PQCreateStream(topicName); Init(topicName); - auto source = BuildSource(topicName); + auto source = BuildSource(); StartSession(ReadActorId1, source); + ExpectStatistics({{ReadActorId1, 0}}); StartSession(ReadActorId2, source); + ExpectStatistics({{ReadActorId1, 0}, {ReadActorId2, 0}}); - const std::vector data = { Json1 }; - PQWrite(data, topicName); + std::vector data = { Json1 }; + PQWrite(data); ExpectNewDataArrived({ReadActorId1, ReadActorId2}); - ExpectMessageBatch(ReadActorId1, { Json1 }); - ExpectMessageBatch(ReadActorId2, { Json1 }); - ExpectStatisticToReadActor({ReadActorId1, ReadActorId2}); - auto source2 = BuildSource(topicName, false, "OtherConsumer"); + ExpectMessageBatch(ReadActorId1, { JsonMessage(1) }); + ExpectMessageBatch(ReadActorId2, { JsonMessage(1) }); + ExpectStatistics({{ReadActorId1, 1}, {ReadActorId2, 1}}); + + data = { Json2 }; + PQWrite(data); + ExpectNewDataArrived({ReadActorId1, ReadActorId2}); + + ExpectStatistics({{ReadActorId1, 1}, {ReadActorId2, 1}}); + ExpectMessageBatch(ReadActorId1, { JsonMessage(2) }); + ExpectMessageBatch(ReadActorId2, { JsonMessage(2) }); + ExpectStatistics({{ReadActorId1, 2}, {ReadActorId2, 2}}); + + auto source2 = BuildSource(false, "OtherConsumer"); StartSession(ReadActorId3, source2, Nothing(), true); - ExpectSessionError(ReadActorId3, "Use the same consumer"); + ExpectSessionError(ReadActorId3, EStatusId::PRECONDITION_FAILED, "Use the same consumer"); StopSession(ReadActorId1, source); StopSession(ReadActorId2, source); } - Y_UNIT_TEST_F(TwoSessionWithoutPredicate, TFixture) { + Y_UNIT_TEST_F(TwoSessionWithoutPredicate, TRealTopicFixture) { const TString topicName = "twowithoutpredicate"; PQCreateStream(topicName); Init(topicName); - auto source1 = BuildSource(topicName, true); - auto source2 = BuildSource(topicName, true); + auto source1 = BuildSource(true); + auto source2 = BuildSource(true); StartSession(ReadActorId1, source1); StartSession(ReadActorId2, source2); const std::vector data = { Json1 }; - PQWrite(data, topicName); + PQWrite(data); ExpectNewDataArrived({ReadActorId1, ReadActorId2}); Runtime.Send(new IEventHandle(TopicSession, ReadActorId1, new TEvRowDispatcher::TEvGetNextBatch())); Runtime.Send(new IEventHandle(TopicSession, ReadActorId2, new TEvRowDispatcher::TEvGetNextBatch())); - ExpectMessageBatch(ReadActorId1, { Json1 }); - ExpectMessageBatch(ReadActorId2, { Json1 }); + ExpectMessageBatch(ReadActorId1, { JsonMessage(1) }); + ExpectMessageBatch(ReadActorId2, { JsonMessage(1) }); StopSession(ReadActorId1, source1); StopSession(ReadActorId2, source2); } - Y_UNIT_TEST_F(SessionWithPredicateAndSessionWithoutPredicate, TFixture) { + Y_UNIT_TEST_F(SessionWithPredicateAndSessionWithoutPredicate, TRealTopicFixture) { const TString topicName = "topic2"; PQCreateStream(topicName); Init(topicName); - auto source1 = BuildSource(topicName, false); - auto source2 = BuildSource(topicName, true); + auto source1 = BuildSource(false); + auto source2 = BuildSource(true); StartSession(ReadActorId1, source1); StartSession(ReadActorId2, source2); const std::vector data = { Json1 }; - PQWrite(data, topicName); + PQWrite(data); ExpectNewDataArrived({ReadActorId1, ReadActorId2}); - ExpectMessageBatch(ReadActorId1, { Json1 }); - ExpectMessageBatch(ReadActorId2, { Json1 }); + ExpectMessageBatch(ReadActorId1, { JsonMessage(1) }); + ExpectMessageBatch(ReadActorId2, { JsonMessage(1) }); StopSession(ReadActorId1, source1); StopSession(ReadActorId2, source2); } - Y_UNIT_TEST_F(SecondSessionWithoutOffsetsAfterSessionConnected, TFixture) { + Y_UNIT_TEST_F(SecondSessionWithoutOffsetsAfterSessionConnected, TRealTopicFixture) { const TString topicName = "topic3"; PQCreateStream(topicName); Init(topicName); - auto source = BuildSource(topicName); + auto source = BuildSource(); StartSession(ReadActorId1, source); const std::vector data = { Json1 }; - PQWrite(data, topicName); + PQWrite(data); ExpectNewDataArrived({ReadActorId1}); - ExpectMessageBatch(ReadActorId1, data); + ExpectMessageBatch(ReadActorId1, { JsonMessage(1) }); StartSession(ReadActorId2, source); const std::vector data2 = { Json2 }; - PQWrite(data2, topicName); + PQWrite(data2); ExpectNewDataArrived({ReadActorId1, ReadActorId2}); - - ExpectMessageBatch(ReadActorId1, data2); - ExpectMessageBatch(ReadActorId2, data2); + + ExpectMessageBatch(ReadActorId1, { JsonMessage(2) }); + ExpectMessageBatch(ReadActorId2, { JsonMessage(2) }); StopSession(ReadActorId1, source); StopSession(ReadActorId2, source); } - Y_UNIT_TEST_F(TwoSessionsWithOffsets, TFixture) { + Y_UNIT_TEST_F(TwoSessionsWithOffsets, TRealTopicFixture) { const TString topicName = "topic4"; PQCreateStream(topicName); Init(topicName); - auto source = BuildSource(topicName); + auto source = BuildSource(); const std::vector data = { Json1, Json2, Json3}; - PQWrite(data, topicName); + PQWrite(data); StartSession(ReadActorId1, source, 1); StartSession(ReadActorId2, source, 2); ExpectNewDataArrived({ReadActorId1, ReadActorId2}); - std::vector expected1 = { Json2, Json3}; + TBatch expected1 = { JsonMessage(2), JsonMessage(3) }; ExpectMessageBatch(ReadActorId1, expected1); - std::vector expected2 = { Json3 }; + TBatch expected2 = { JsonMessage(3) }; ExpectMessageBatch(ReadActorId2, expected2); const std::vector data2 = { Json4 }; - PQWrite(data2, topicName); + PQWrite(data2); ExpectNewDataArrived({ReadActorId1, ReadActorId2}); - ExpectMessageBatch(ReadActorId1, data2); - ExpectMessageBatch(ReadActorId2, data2); + ExpectMessageBatch(ReadActorId1, { JsonMessage(4) }); + ExpectMessageBatch(ReadActorId2, { JsonMessage(4) }); StopSession(ReadActorId1, source); StopSession(ReadActorId2, source); } - Y_UNIT_TEST_F(BadDataSessionError, TFixture) { + Y_UNIT_TEST_F(BadDataSessionError, TRealTopicFixture) { const TString topicName = "topic5"; PQCreateStream(topicName); Init(topicName); - auto source = BuildSource(topicName); + auto source = BuildSource(); StartSession(ReadActorId1, source); const std::vector data = { "not json", "noch einmal / nicht json" }; - PQWrite(data, topicName); + PQWrite(data); - ExpectSessionError(ReadActorId1, "INCORRECT_TYPE: The JSON element does not have the requested type."); + ExpectSessionError(ReadActorId1, EStatusId::BAD_REQUEST, "INCORRECT_TYPE: The JSON element does not have the requested type."); StopSession(ReadActorId1, source); } - Y_UNIT_TEST_F(WrongFieldType, TFixture) { + Y_UNIT_TEST_F(WrongFieldType, TRealTopicFixture) { const TString topicName = "wrong_field"; PQCreateStream(topicName); Init(topicName); - auto source = BuildSource(topicName); + + auto source = BuildSource(); StartSession(ReadActorId1, source); - const std::vector data = {"{\"dt\":100}"}; - PQWrite(data, topicName); - auto error = ExpectSessionError(ReadActorId1); - UNIT_ASSERT_STRING_CONTAINS(error, "Failed to parse json messages, found 1 missing values"); - UNIT_ASSERT_STRING_CONTAINS(error, "the field (value) has been added by query"); + source.AddColumns("field1"); + source.AddColumnTypes("[DataType; String]"); + StartSession(ReadActorId2, source); + + PQWrite({ Json1 }); + ExpectNewDataArrived({ReadActorId1}); + ExpectMessageBatch(ReadActorId1, { JsonMessage(1) }); + ExpectSessionError(ReadActorId2, EStatusId::PRECONDITION_FAILED, "Failed to parse json messages, found 1 missing values"); + + PQWrite({ Json2 }); + ExpectNewDataArrived({ReadActorId1}); + ExpectMessageBatch(ReadActorId1, { JsonMessage(2) }); + ExpectSessionError(ReadActorId2, EStatusId::PRECONDITION_FAILED, "Failed to parse json messages, found 1 missing values"); + StopSession(ReadActorId1, source); + StopSession(ReadActorId2, source); } - Y_UNIT_TEST_F(RestartSessionIfNewClientWithOffset, TFixture) { + Y_UNIT_TEST_F(RestartSessionIfNewClientWithOffset, TRealTopicFixture) { const TString topicName = "topic6"; PQCreateStream(topicName); Init(topicName); - auto source = BuildSource(topicName); + auto source = BuildSource(); StartSession(ReadActorId1, source); const std::vector data = { Json1, Json2, Json3 }; // offset 0, 1, 2 - PQWrite(data, topicName); + PQWrite(data); ExpectNewDataArrived({ReadActorId1}); - ExpectMessageBatch(ReadActorId1, data); + ExpectMessageBatch(ReadActorId1, { JsonMessage(1), JsonMessage(2), JsonMessage(3) }); // Restart topic session. StartSession(ReadActorId2, source, 1); ExpectNewDataArrived({ReadActorId2}); - PQWrite({ Json4 }, topicName); + PQWrite({ Json4 }); ExpectNewDataArrived({ReadActorId1}); - ExpectMessageBatch(ReadActorId1, { Json4 }); - ExpectMessageBatch(ReadActorId2, { Json2, Json3, Json4 }); + ExpectMessageBatch(ReadActorId1, { JsonMessage(4) }); + ExpectMessageBatch(ReadActorId2, { JsonMessage(2), JsonMessage(3), JsonMessage(4) }); StopSession(ReadActorId1, source); StopSession(ReadActorId2, source); } - Y_UNIT_TEST_F(ReadNonExistentTopic, TFixture) { + Y_UNIT_TEST_F(ReadNonExistentTopic, TRealTopicFixture) { const TString topicName = "topic7"; Init(topicName); - auto source = BuildSource(topicName); + auto source = BuildSource(); StartSession(ReadActorId1, source); - ExpectSessionError(ReadActorId1, "no path"); + ExpectSessionError(ReadActorId1, EStatusId::SCHEME_ERROR, "no path"); StopSession(ReadActorId1, source); } - Y_UNIT_TEST_F(SlowSession, TFixture) { + Y_UNIT_TEST_F(SlowSession, TRealTopicFixture) { const TString topicName = "topic8"; PQCreateStream(topicName); - Init(topicName, 50); - auto source = BuildSource(topicName); + Init(topicName, 40); + auto source = BuildSource(); StartSession(ReadActorId1, source); StartSession(ReadActorId2, source); // slow session @@ -427,110 +500,134 @@ Y_UNIT_TEST_SUITE(TopicSessionTests) { auto writeMessages = [&]() { for (size_t i = 0; i < messagesSize; ++i) { const std::vector data = { Json1 }; - PQWrite(data, topicName); + PQWrite(data); } Sleep(TDuration::MilliSeconds(100)); - Runtime.DispatchEvents({}, Runtime.GetCurrentTime() - TDuration::MilliSeconds(1)); }; - + writeMessages(); ExpectNewDataArrived({ReadActorId1, ReadActorId2}); auto readMessages = ReadMessages(ReadActorId1); - UNIT_ASSERT(readMessages == messagesSize); + UNIT_ASSERT_VALUES_EQUAL(readMessages, messagesSize); // Reading from yds is stopped. writeMessages(); readMessages = ReadMessages(ReadActorId1); - UNIT_ASSERT(readMessages == 0); + UNIT_ASSERT_VALUES_EQUAL(readMessages, 0); readMessages = ReadMessages(ReadActorId2); - UNIT_ASSERT(readMessages == messagesSize); + UNIT_ASSERT_VALUES_EQUAL(readMessages, messagesSize); Sleep(TDuration::MilliSeconds(100)); - Runtime.DispatchEvents({}, Runtime.GetCurrentTime() - TDuration::MilliSeconds(1)); readMessages = ReadMessages(ReadActorId1); - UNIT_ASSERT(readMessages == messagesSize); + UNIT_ASSERT_VALUES_EQUAL(readMessages, messagesSize); writeMessages(); StopSession(ReadActorId2, source); // delete slow client, clear unread buffer Sleep(TDuration::MilliSeconds(100)); - Runtime.DispatchEvents({}, Runtime.GetCurrentTime() - TDuration::MilliSeconds(1)); readMessages = ReadMessages(ReadActorId1); - UNIT_ASSERT(readMessages == messagesSize); + UNIT_ASSERT_VALUES_EQUAL(readMessages, messagesSize); StopSession(ReadActorId1, source); } - Y_UNIT_TEST_F(TwoSessionsWithDifferentSchemes, TFixture) { + Y_UNIT_TEST_F(TwoSessionsWithDifferentSchemes, TRealTopicFixture) { const TString topicName = "dif_schemes"; PQCreateStream(topicName); Init(topicName); - auto source1 = BuildSource(topicName); - auto source2 = BuildSource(topicName); + auto source1 = BuildSource(); + auto source2 = BuildSource(); source2.AddColumns("field1"); source2.AddColumnTypes("[DataType; String]"); StartSession(ReadActorId1, source1); StartSession(ReadActorId2, source2); - TString json1 = "{\"dt\":101,\"value\":\"value1\", \"field1\":\"field1\"}"; - TString json2 = "{\"dt\":102,\"value\":\"value2\", \"field1\":\"field2\"}"; + TString json1 = "{\"dt\":100,\"value\":\"value1\", \"field1\":\"field1\"}"; + TString json2 = "{\"dt\":200,\"value\":\"value2\", \"field1\":\"field2\"}"; - PQWrite({ json1, json2 }, topicName); + PQWrite({ json1, json2 }); ExpectNewDataArrived({ReadActorId1, ReadActorId2}); - ExpectMessageBatch(ReadActorId1, { "{\"dt\":101,\"value\":\"value1\"}", "{\"dt\":102,\"value\":\"value2\"}" }); - ExpectMessageBatch(ReadActorId2, { "{\"dt\":101,\"field1\":\"field1\",\"value\":\"value1\"}", "{\"dt\":102,\"field1\":\"field2\",\"value\":\"value2\"}" }); + ExpectMessageBatch(ReadActorId1, { JsonMessage(1), JsonMessage(2) }); + ExpectMessageBatch(ReadActorId2, { JsonMessage(1).AddString("field1"), JsonMessage(2).AddString("field2") }); - auto source3 = BuildSource(topicName); + auto source3 = BuildSource(); source3.AddColumns("field2"); source3.AddColumnTypes("[DataType; String]"); auto readActorId3 = Runtime.AllocateEdgeActor(); StartSession(readActorId3, source3); - TString json3 = "{\"dt\":103,\"value\":\"value3\", \"field1\":\"value1_field1\", \"field2\":\"value1_field2\"}"; - PQWrite({ json3 }, topicName); + TString json3 = "{\"dt\":300,\"value\":\"value3\", \"field1\":\"value1_field1\", \"field2\":\"value1_field2\"}"; + PQWrite({ json3 }); ExpectNewDataArrived({ReadActorId1, ReadActorId2, readActorId3}); - ExpectMessageBatch(ReadActorId1, { "{\"dt\":103,\"value\":\"value3\"}" }); - ExpectMessageBatch(ReadActorId2, { "{\"dt\":103,\"field1\":\"value1_field1\",\"value\":\"value3\"}" }); - ExpectMessageBatch(readActorId3, { "{\"dt\":103,\"field2\":\"value1_field2\",\"value\":\"value3\"}" }); + ExpectMessageBatch(ReadActorId1, { JsonMessage(3) }); + ExpectMessageBatch(ReadActorId2, { JsonMessage(3).AddString("value1_field1") }); + ExpectMessageBatch(readActorId3, { JsonMessage(3).AddString("value1_field2") }); StopSession(ReadActorId1, source3); StopSession(readActorId3, source3); - TString json4 = "{\"dt\":104,\"value\":\"value4\", \"field1\":\"value2_field1\", \"field2\":\"value2_field2\"}"; - TString json5 = "{\"dt\":105,\"value\":\"value5\", \"field1\":\"value2_field1\", \"field2\":\"value2_field2\"}"; - PQWrite({ json4, json5 }, topicName); + TString json4 = "{\"dt\":400,\"value\":\"value4\", \"field1\":\"value2_field1\", \"field2\":\"value2_field2\"}"; + TString json5 = "{\"dt\":500,\"value\":\"value5\", \"field1\":\"value3_field1\", \"field2\":\"value3_field2\"}"; + PQWrite({ json4, json5 }); ExpectNewDataArrived({ReadActorId2}); - ExpectMessageBatch(ReadActorId2, { "{\"dt\":104,\"field1\":\"value2_field1\",\"value\":\"value4\"}", "{\"dt\":105,\"field1\":\"value2_field1\",\"value\":\"value5\"}" }); + ExpectMessageBatch(ReadActorId2, { JsonMessage(4).AddString("value2_field1"), JsonMessage(5).AddString("value3_field1") }); StopSession(ReadActorId1, source1); StopSession(ReadActorId2, source2); } - Y_UNIT_TEST_F(TwoSessionsWithDifferentColumnTypes, TFixture) { + Y_UNIT_TEST_F(TwoSessionsWithDifferentColumnTypes, TRealTopicFixture) { const TString topicName = "dif_types"; PQCreateStream(topicName); Init(topicName); - auto source1 = BuildSource(topicName); + auto source1 = BuildSource(); source1.AddColumns("field1"); source1.AddColumnTypes("[OptionalType; [DataType; String]]"); StartSession(ReadActorId1, source1); - TString json1 = "{\"dt\":101,\"field1\":null,\"value\":\"value1\"}"; - PQWrite({ json1 }, topicName); + TString json1 = "{\"dt\":100,\"field1\":\"str\",\"value\":\"value1\"}"; + PQWrite({ json1 }); ExpectNewDataArrived({ReadActorId1}); - ExpectMessageBatch(ReadActorId1, { json1 }); + ExpectMessageBatch(ReadActorId1, { JsonMessage(1).AddString("str", true) }); - auto source2 = BuildSource(topicName); + auto source2 = BuildSource(); source2.AddColumns("field1"); source2.AddColumnTypes("[DataType; String]"); StartSession(ReadActorId2, source2, Nothing(), true); - ExpectSessionError(ReadActorId2, "Use the same column type in all queries via RD, current type for column `field1` is [OptionalType; [DataType; String]] (requested type is [DataType; String])"); - } -} + ExpectSessionError(ReadActorId2, EStatusId::SCHEME_ERROR, "Use the same column type in all queries via RD, current type for column `field1` is [OptionalType; [DataType; String]] (requested type is [DataType; String])"); + } + + Y_UNIT_TEST_F(RestartSessionIfQueryStopped, TMockTopicFixture) { + Init("fake_topic", 1000); + auto source = BuildSource(); + StartSession(ReadActorId1, source); + std::vector data = { Json1, Json2, Json3 }; + PQWrite(data, 1); + ExpectNewDataArrived({ReadActorId1}); + ExpectMessageBatch(ReadActorId1, { JsonMessage(1), JsonMessage(2), JsonMessage(3) }); + + StartSession(ReadActorId2, source, 1); + std::vector data2 = { Json1 }; + PQWrite(data2, 1); + ExpectNewDataArrived({ReadActorId2}); + ExpectMessageBatch(ReadActorId2, { JsonMessage(1)}); + + StopSession(ReadActorId2, source); + Runtime.GrabEdgeEvent(PqGatewayNotifier, TDuration::Seconds(GrabTimeoutSec)); + MockPqGateway->AddEvent(TopicPath, NYdb::NTopic::TReadSessionEvent::TStartPartitionSessionEvent(nullptr, 0, 0), 0); + + std::vector data3 = { Json4 }; + PQWrite(data3, 4); + ExpectNewDataArrived({ReadActorId1}); + ExpectMessageBatch(ReadActorId1, { JsonMessage(4) }); + + PassAway(); + } } +} // namespace NFq::NRowDispatcher::NTests diff --git a/ydb/core/fq/libs/row_dispatcher/ut/ya.make b/ydb/core/fq/libs/row_dispatcher/ut/ya.make index 7049d37ee118..782b4498a291 100644 --- a/ydb/core/fq/libs/row_dispatcher/ut/ya.make +++ b/ydb/core/fq/libs/row_dispatcher/ut/ya.make @@ -4,8 +4,6 @@ INCLUDE(${ARCADIA_ROOT}/ydb/tests/tools/fq_runner/ydb_runner_with_datastreams.in SRCS( coordinator_ut.cpp - json_filter_ut.cpp - json_parser_ut.cpp leader_election_ut.cpp row_dispatcher_ut.cpp topic_session_ut.cpp @@ -13,13 +11,13 @@ SRCS( PEERDIR( library/cpp/testing/unittest + ydb/core/fq/libs/row_dispatcher/format_handler/ut/common ydb/core/fq/libs/row_dispatcher ydb/core/testlib ydb/core/testlib/actors - yql/essentials/udfs/common/json2 - yql/essentials/udfs/common/yson2 ydb/tests/fq/pq_async_io yql/essentials/sql/pg_dummy + ydb/library/yql/providers/pq/gateway/dummy ) SIZE(MEDIUM) diff --git a/ydb/core/fq/libs/row_dispatcher/ya.make b/ydb/core/fq/libs/row_dispatcher/ya.make index 0b6f68bfc03a..d1039a939529 100644 --- a/ydb/core/fq/libs/row_dispatcher/ya.make +++ b/ydb/core/fq/libs/row_dispatcher/ya.make @@ -2,46 +2,45 @@ LIBRARY() SRCS( actors_factory.cpp - common.cpp coordinator.cpp - json_filter.cpp - json_parser.cpp leader_election.cpp - row_dispatcher_service.cpp + probes.cpp row_dispatcher.cpp + row_dispatcher_service.cpp topic_session.cpp ) PEERDIR( - contrib/libs/fmt - contrib/libs/simdjson ydb/core/fq/libs/actors/logging ydb/core/fq/libs/config/protos + ydb/core/fq/libs/metrics ydb/core/fq/libs/row_dispatcher/events + ydb/core/fq/libs/row_dispatcher/format_handler ydb/core/fq/libs/row_dispatcher/purecalc_compilation ydb/core/fq/libs/shared_resources ydb/core/fq/libs/ydb ydb/core/mon + ydb/library/actors/core ydb/library/security + ydb/library/yql/dq/actors ydb/library/yql/dq/actors/common ydb/library/yql/dq/actors/compute ydb/library/yql/dq/proto ydb/library/yql/providers/pq/provider - ydb/core/fq/libs/row_dispatcher/purecalc_no_pg_wrapper + ydb/public/sdk/cpp/client/ydb_scheme ydb/public/sdk/cpp/client/ydb_table ) -CFLAGS( - -Wno-assume -) - YQL_LAST_ABI_VERSION() END() -RECURSE(purecalc_no_pg_wrapper) +RECURSE( + purecalc_no_pg_wrapper + format_handler +) IF(NOT EXPORT_CMAKE) RECURSE_FOR_TESTS( diff --git a/ydb/core/fq/libs/test_connection/counters.cpp b/ydb/core/fq/libs/test_connection/counters.cpp index 0d708ebacad9..ef656b44b998 100644 --- a/ydb/core/fq/libs/test_connection/counters.cpp +++ b/ydb/core/fq/libs/test_connection/counters.cpp @@ -15,7 +15,7 @@ void TTestConnectionRequestCounters::Register(const ::NMonitoring::TDynamicCount } NMonitoring::IHistogramCollectorPtr TTestConnectionRequestCounters::GetLatencyHistogramBuckets() { - return NMonitoring::ExplicitHistogram({0, 1, 2, 5, 10, 20, 50, 100, 500, 1000, 2000, 5000, 10000, 30000, 50000, 500000}); + return NMonitoring::ExplicitHistogram({0, 10, 100, 1000, 10000}); } } // NFq diff --git a/ydb/core/fq/libs/ydb/ut/ydb_ut.cpp b/ydb/core/fq/libs/ydb/ut/ydb_ut.cpp index a0cc5edfdee6..958376ecb34f 100644 --- a/ydb/core/fq/libs/ydb/ut/ydb_ut.cpp +++ b/ydb/core/fq/libs/ydb/ut/ydb_ut.cpp @@ -466,6 +466,18 @@ Y_UNIT_TEST_SUITE(TFqYdbTest) { UNIT_ASSERT(issues.Size() == 1); UNIT_ASSERT(issues.ToString().Contains(text)); } + + Y_UNIT_TEST(ShouldStatusToIssuesProcessEmptyIssues) + { + auto promise = NThreading::NewPromise(); + auto future = promise.GetFuture(); + promise.SetValue(TStatus(EStatus::BAD_REQUEST, NYql::TIssues{})); + NThreading::TFuture future2 = NFq::StatusToIssues(future); + + NYql::TIssues issues = future2.GetValueSync(); + UNIT_ASSERT_C(issues.Size() == 1, issues.ToString()); + UNIT_ASSERT(issues.ToString().Contains("empty issues")); + } } } // namespace NFq diff --git a/ydb/core/fq/libs/ydb/ydb.cpp b/ydb/core/fq/libs/ydb/ydb.cpp index 3bcdc01f5fbe..4b3958a726be 100644 --- a/ydb/core/fq/libs/ydb/ydb.cpp +++ b/ydb/core/fq/libs/ydb/ydb.cpp @@ -203,6 +203,11 @@ NYql::TIssues StatusToIssues(const NYdb::TStatus& status) { TIssues issues; if (!status.IsSuccess()) { issues = status.GetIssues(); + if (!issues) { + TStringStream str; + str << "Internal error: empty issues with failed status (" << status.GetStatus() << ")"; + issues.AddIssue(str.Str()); + } } return issues; } diff --git a/ydb/core/grpc_services/grpc_request_proxy_simple.cpp b/ydb/core/grpc_services/grpc_request_proxy_simple.cpp index c2274c3d7be3..5f4168ff4575 100644 --- a/ydb/core/grpc_services/grpc_request_proxy_simple.cpp +++ b/ydb/core/grpc_services/grpc_request_proxy_simple.cpp @@ -86,6 +86,11 @@ class TGRpcRequestProxySimple request->SendResult(*result, Ydb::StatusIds::SUCCESS); } + void Handle(TEvRequestAuthAndCheck::TPtr& ev, const TActorContext&) { + ev->Get()->FinishSpan(); + ev->Get()->ReplyWithYdbStatus(Ydb::StatusIds::SUCCESS); + } + void Handle(TEvProxyRuntimeEvent::TPtr& event, const TActorContext&) { IRequestProxyCtx* requestBaseCtx = event->Get(); TString validationError; @@ -221,6 +226,7 @@ void TGRpcRequestProxySimple::StateFunc(TAutoPtr& ev) { hFunc(TEvents::TEvUndelivered, HandleUndelivery); HFunc(TEvListEndpointsRequest, PreHandle); HFunc(TEvProxyRuntimeEvent, PreHandle); + HFunc(TEvRequestAuthAndCheck, PreHandle); default: Y_ABORT("Unknown request: %u\n", ev->GetTypeRewrite()); break; diff --git a/ydb/core/kqp/ut/federated_query/common/common.cpp b/ydb/core/kqp/ut/federated_query/common/common.cpp index bd95dd3a6246..508862dd7313 100644 --- a/ydb/core/kqp/ut/federated_query/common/common.cpp +++ b/ydb/core/kqp/ut/federated_query/common/common.cpp @@ -3,6 +3,16 @@ #include namespace NKikimr::NKqp::NFederatedQueryTest { + TString GetSymbolsString(char start, char end, const TString& skip) { + TStringBuilder result; + for (char symbol = start; symbol <= end; ++symbol) { + if (skip.Contains(symbol)) { + continue; + } + result << symbol; + } + return result; + } NYdb::NQuery::TScriptExecutionOperation WaitScriptExecutionOperation(const NYdb::TOperation::TOperationId& operationId, const NYdb::TDriver& ydbDriver) { NYdb::NOperation::TOperationClient client(ydbDriver); diff --git a/ydb/core/kqp/ut/federated_query/common/common.h b/ydb/core/kqp/ut/federated_query/common/common.h index a06d044a3e6d..c539e9a525c2 100644 --- a/ydb/core/kqp/ut/federated_query/common/common.h +++ b/ydb/core/kqp/ut/federated_query/common/common.h @@ -8,6 +8,8 @@ namespace NKikimr::NKqp::NFederatedQueryTest { using namespace NKikimr::NKqp; + TString GetSymbolsString(char start, char end, const TString& skip = ""); + NYdb::NQuery::TScriptExecutionOperation WaitScriptExecutionOperation( const NYdb::TOperation::TOperationId& operationId, const NYdb::TDriver& ydbDriver); diff --git a/ydb/core/kqp/ut/federated_query/generic_ut/kqp_generic_provider_ut.cpp b/ydb/core/kqp/ut/federated_query/generic_ut/kqp_generic_provider_ut.cpp index 0a5939a82c98..480ad5f29912 100644 --- a/ydb/core/kqp/ut/federated_query/generic_ut/kqp_generic_provider_ut.cpp +++ b/ydb/core/kqp/ut/federated_query/generic_ut/kqp_generic_provider_ut.cpp @@ -35,7 +35,7 @@ namespace NKikimr::NKqp { Ydb, }; - NApi::TDataSourceInstance MakeDataSourceInstance(EProviderType providerType) { + NYql::TGenericDataSourceInstance MakeDataSourceInstance(EProviderType providerType) { switch (providerType) { case EProviderType::PostgreSQL: return TConnectorClientMock::TPostgreSQLDataSourceInstanceBuilder<>().GetResult(); @@ -96,7 +96,7 @@ namespace NKikimr::NKqp { // prepare mock auto clientMock = std::make_shared(); - const NApi::TDataSourceInstance dataSourceInstance = MakeDataSourceInstance(providerType); + const NYql::TGenericDataSourceInstance dataSourceInstance = MakeDataSourceInstance(providerType); // step 1: DescribeTable // clang-format off @@ -192,7 +192,7 @@ namespace NKikimr::NKqp { // prepare mock auto clientMock = std::make_shared(); - const NApi::TDataSourceInstance dataSourceInstance = MakeDataSourceInstance(providerType); + const NYql::TGenericDataSourceInstance dataSourceInstance = MakeDataSourceInstance(providerType); constexpr size_t ROWS_COUNT = 5; @@ -285,7 +285,7 @@ namespace NKikimr::NKqp { // prepare mock auto clientMock = std::make_shared(); - const NApi::TDataSourceInstance dataSourceInstance = MakeDataSourceInstance(providerType); + const NYql::TGenericDataSourceInstance dataSourceInstance = MakeDataSourceInstance(providerType); constexpr size_t ROWS_COUNT = 5; @@ -374,7 +374,7 @@ namespace NKikimr::NKqp { // prepare mock auto clientMock = std::make_shared(); - const NApi::TDataSourceInstance dataSourceInstance = MakeDataSourceInstance(providerType); + const NYql::TGenericDataSourceInstance dataSourceInstance = MakeDataSourceInstance(providerType); // clang-format off const NApi::TSelect select = TConnectorClientMock::TSelectBuilder<>() .DataSourceInstance(dataSourceInstance) diff --git a/ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp b/ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp index 2a980bae8c50..c66575f92176 100644 --- a/ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp +++ b/ydb/core/kqp/ut/federated_query/s3/kqp_federated_query_ut.cpp @@ -20,17 +20,6 @@ using namespace NTestUtils; using namespace fmt::literals; Y_UNIT_TEST_SUITE(KqpFederatedQuery) { - TString GetSymbolsString(char start, char end, const TString& skip = "") { - TStringBuilder result; - for (char symbol = start; symbol <= end; ++symbol) { - if (skip.Contains(symbol)) { - continue; - } - result << symbol; - } - return result; - } - Y_UNIT_TEST(ExecuteScriptWithExternalTableResolve) { const TString externalDataSourceName = "/Root/external_data_source"; const TString externalTableName = "/Root/test_binding_resolve"; diff --git a/ydb/core/kqp/ut/federated_query/s3/kqp_s3_plan_ut.cpp b/ydb/core/kqp/ut/federated_query/s3/kqp_s3_plan_ut.cpp index bc5a6ce27256..c81295685d5a 100644 --- a/ydb/core/kqp/ut/federated_query/s3/kqp_s3_plan_ut.cpp +++ b/ydb/core/kqp/ut/federated_query/s3/kqp_s3_plan_ut.cpp @@ -143,14 +143,15 @@ Y_UNIT_TEST_SUITE(KqpS3PlanTest) { UNIT_ASSERT(NJson::ReadJsonTree(*queryResult.GetStats()->GetPlan(), &plan)); const auto& writeStagePlan = plan["Plan"]["Plans"][0]["Plans"][0]; - UNIT_ASSERT_VALUES_EQUAL(writeStagePlan["Node Type"].GetStringSafe(), "Limit-Sink"); - UNIT_ASSERT(writeStagePlan["Operators"].GetArraySafe().size() >= 2); - const auto& sinkOp = writeStagePlan["Operators"].GetArraySafe()[1]; + UNIT_ASSERT_VALUES_EQUAL(writeStagePlan["Node Type"].GetStringSafe(), "Stage-Sink"); + UNIT_ASSERT(writeStagePlan["Operators"].GetArraySafe().size() >= 1); + const auto& sinkOp = writeStagePlan["Operators"].GetArraySafe()[0]; UNIT_ASSERT_VALUES_EQUAL(sinkOp["ExternalDataSource"].GetStringSafe(), "write_data_source"); UNIT_ASSERT_VALUES_EQUAL(sinkOp["Compression"].GetStringSafe(), "gzip"); - const auto& readStagePlan = plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]; + const auto& readStagePlan = plan["Plan"]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]["Plans"][0]; UNIT_ASSERT_VALUES_EQUAL(readStagePlan["Node Type"].GetStringSafe(), "Source"); + UNIT_ASSERT(readStagePlan["Operators"].GetArraySafe().size() >= 1); const auto& sourceOp = readStagePlan["Operators"].GetArraySafe()[0]; UNIT_ASSERT_VALUES_EQUAL(sourceOp["ExternalDataSource"].GetStringSafe(), "read_data_source"); diff --git a/ydb/core/public_http/fq_handlers.h b/ydb/core/public_http/fq_handlers.h index 6fd53b9701d3..1a30ba6b409f 100644 --- a/ydb/core/public_http/fq_handlers.h +++ b/ydb/core/public_http/fq_handlers.h @@ -10,6 +10,8 @@ #include #include +#include + namespace NKikimr::NPublicHttp { using namespace NActors; @@ -249,6 +251,7 @@ void SetIdempotencyKey(T& dst, const TString& key) { template class TGrpcCallWrapper : public TActorBootstrapped> { +protected: THttpRequestContext RequestContext; typedef std::function(TIntrusivePtr ctx)> TGrpcProxyEventFactory; @@ -278,6 +281,10 @@ class TGrpcCallWrapper : public TActorBootstrapped(); if (Parse(*grpcRequest)) { TIntrusivePtr requestContext = new TGrpcRequestContextWrapper(RequestContext, std::move(grpcRequest), &SendReply); @@ -354,7 +361,6 @@ class TGrpcCallWrapper : public TActorBootstrapped(resp->GetArena()); FqConvert(typedResponse->operation(), *httpResult); FqPackToJson(json, *httpResult, jsonSettings); - requestContext.ResponseBadRequestJson(typedResponse->operation().status(), json.Str()); return; } @@ -396,4 +402,72 @@ DECLARE_YQ_GRPC_ACTOR(GetQueryStatus, GetQueryStatus); DECLARE_YQ_GRPC_ACTOR_WIHT_EMPTY_RESULT(StopQuery, ControlQuery); DECLARE_YQ_GRPC_ACTOR(GetResultData, GetResultData); +class TJsonStartQuery : public TGrpcCallWrapper { +public: + typedef TGrpcCallWrapper TGrpcCallWrapperBase; + + TJsonStartQuery(const THttpRequestContext& ctx) + : TGrpcCallWrapperBase(ctx, &NGRpcService::CreateFederatedQueryDescribeQueryRequestOperationCall) + {} + + void BootstrapWrapper(const TActorContext& ctx) override { + + auto describeRequest = std::make_unique(); + if (!Parse(*describeRequest)) { + this->Die(ctx); + return; + } + + TProtoStringType queryId = describeRequest->Getquery_id(); + TIntrusivePtr requestContext = MakeIntrusive( + RequestContext, + std::move(describeRequest), + [query_id = std::move(queryId), actorSystem = TActivationContext::ActorSystem()](const THttpRequestContext& requestContext, const TJsonSettings& jsonSettings, NProtoBuf::Message* resp, ui32 status) { + + Y_ABORT_UNLESS(resp); + Y_ABORT_UNLESS(resp->GetArena()); + + auto* typedResponse = static_cast(resp); + if (!typedResponse->operation().result().template Is()) { + TStringStream json; + auto httpResult = std::unique_ptr(new FQHttp::Error()); + FqConvert(typedResponse->operation(), *httpResult); + FqPackToJson(json, *httpResult, jsonSettings); + requestContext.ResponseBadRequestJson(typedResponse->operation().status(), json.Str()); + return; + } + + std::unique_ptr describeResult = std::unique_ptr(new FederatedQuery::DescribeQueryResult()); + if (!typedResponse->operation().result().UnpackTo(&*describeResult)) { + requestContext.ResponseBadRequest(Ydb::StatusIds::INTERNAL_ERROR, "Error in response unpack"); + return; + } + + // modify + auto modifyRequest = std::unique_ptr(new FederatedQuery::ModifyQueryRequest()); + + modifyRequest->set_query_id(query_id); + *modifyRequest->mutable_content() = describeResult->query().content(); + modifyRequest->set_execute_mode(::FederatedQuery::ExecuteMode::RUN); + modifyRequest->set_state_load_mode(::FederatedQuery::StateLoadMode::STATE_LOAD_MODE_UNSPECIFIED); + modifyRequest->set_previous_revision(describeResult->query().meta().Getlast_job_query_revision()); + modifyRequest->set_idempotency_key(requestContext.GetIdempotencyKey()); + + TIntrusivePtr requestContextModify = new TGrpcRequestContextWrapper( + requestContext, + std::move(modifyRequest), + TGrpcCallWrapper::SendReply + ); + + // new event -> new EventFactory + actorSystem->Send(NGRpcService::CreateGRpcRequestProxyId(), NGRpcService::CreateFederatedQueryModifyQueryRequestOperationCall(std::move(requestContextModify)).release()); + }); + + ctx.Send(NGRpcService::CreateGRpcRequestProxyId(), EventFactory(std::move(requestContext)).release()); + this->Die(ctx); + } +}; + +#undef TGrpcCallWrapperBase + } // namespace NKikimr::NPublicHttp diff --git a/ydb/core/public_http/http_service.cpp b/ydb/core/public_http/http_service.cpp index ae6d72209d46..81fa9c66ebfb 100644 --- a/ydb/core/public_http/http_service.cpp +++ b/ydb/core/public_http/http_service.cpp @@ -57,6 +57,7 @@ namespace { Router.RegisterHandler(HTTP_METHOD_GET, "/api/fq/v1/queries/{query_id}/status", CreateHttpHandler()); Router.RegisterHandler(HTTP_METHOD_GET, "/api/fq/v1/queries/{query_id}/results/{result_set_index}", CreateHttpHandler()); Router.RegisterHandler(HTTP_METHOD_POST, "/api/fq/v1/queries/{query_id}/stop", CreateHttpHandler()); + Router.RegisterHandler(HTTP_METHOD_POST, "/api/fq/v1/queries/{query_id}/start", CreateHttpHandler()); } void Bootstrap(const TActorContext& ctx) { diff --git a/ydb/core/public_http/openapi/openapi.yaml b/ydb/core/public_http/openapi/openapi.yaml index 6a490bebd3d4..358dfc1bbd17 100644 --- a/ydb/core/public_http/openapi/openapi.yaml +++ b/ydb/core/public_http/openapi/openapi.yaml @@ -185,6 +185,31 @@ paths: required: true schema: type: string + '/queries/{query_id}/start': + post: + responses: + '204': + description: No Content + '400': + description: Bad Request + content: + application/json: + schema: + $ref: '#/components/schemas/GenericError' + parameters: + - $ref: '#/components/parameters/Idempotency-Key' + - $ref: '#/components/parameters/Authorization' + - $ref: '#/components/parameters/x-request-id' + - $ref: '#/components/parameters/db' + - $ref: '#/components/parameters/project' + summary: start stopped query + operationId: start-query + parameters: + - name: query_id + in: path + required: true + schema: + type: string '/queries/{query_id}/results/{result_set_index}': parameters: - name: query_id diff --git a/ydb/core/testlib/basics/helpers.h b/ydb/core/testlib/basics/helpers.h index e6be7bca98ef..24126aca2cc4 100644 --- a/ydb/core/testlib/basics/helpers.h +++ b/ydb/core/testlib/basics/helpers.h @@ -18,6 +18,8 @@ namespace NFake { ui64 SectorSize = 0; ui64 ChunkSize = 0; ui64 DiskSize = 0; + bool FormatDisk = true; + TString DiskPath; }; struct INode { diff --git a/ydb/core/testlib/basics/storage.h b/ydb/core/testlib/basics/storage.h index 65b58a076fb1..e1aac88071d2 100644 --- a/ydb/core/testlib/basics/storage.h +++ b/ydb/core/testlib/basics/storage.h @@ -54,7 +54,7 @@ namespace NKikimr { static ui64 keySalt = 0; ui64 salt = ++keySalt; - TString baseDir = Runtime.GetTempDir(); + TString baseDir = conf.DiskPath ? conf.DiskPath : Runtime.GetTempDir(); if (Conf.UseDisk) { MakeDirIfNotExist(baseDir.c_str()); @@ -62,7 +62,7 @@ namespace NKikimr { PDiskPath = TStringBuilder() << baseDir << "pdisk_1.dat"; - if (!Mock) { + if (!Mock && conf.FormatDisk) { FormatPDisk(PDiskPath, Conf.DiskSize, Conf.SectorSize, Conf.ChunkSize, PDiskGuid, 0x123 + salt, 0x456 + salt, 0x789 + salt, mainKey, diff --git a/ydb/core/testlib/test_client.cpp b/ydb/core/testlib/test_client.cpp index e9113ad903d5..e67daa9777c6 100644 --- a/ydb/core/testlib/test_client.cpp +++ b/ydb/core/testlib/test_client.cpp @@ -58,6 +58,7 @@ #include #include #include +#include #include #include #include @@ -93,6 +94,7 @@ #include #include #include +#include #include #include #include @@ -439,7 +441,7 @@ namespace Tests { GRpcServer->AddService(new NGRpcService::TGRpcMonitoringService(system, counters, grpcRequestProxies[0], true)); GRpcServer->AddService(new NGRpcService::TGRpcYdbQueryService(system, counters, grpcRequestProxies, true, 1)); GRpcServer->AddService(new NGRpcService::TGRpcYdbTabletService(system, counters, grpcRequestProxies, true, 1)); - if (Settings->EnableYq) { + if (Settings->EnableYq || Settings->EnableYqGrpc) { GRpcServer->AddService(new NGRpcService::TGRpcFederatedQueryService(system, counters, grpcRequestProxies[0])); GRpcServer->AddService(new NGRpcService::TGRpcFqPrivateTaskService(system, counters, grpcRequestProxies[0])); } @@ -610,6 +612,7 @@ namespace Tests { NKikimrBlobStorage::TDefineBox boxConfig; boxConfig.SetBoxId(Settings->BOX_ID); + boxConfig.SetItemConfigGeneration(Settings->StorageGeneration); ui32 nodeId = Runtime->GetNodeId(0); Y_ABORT_UNLESS(nodesInfo->Nodes[0].NodeId == nodeId); @@ -617,11 +620,13 @@ namespace Tests { NKikimrBlobStorage::TDefineHostConfig hostConfig; hostConfig.SetHostConfigId(nodeId); + hostConfig.SetItemConfigGeneration(Settings->StorageGeneration); TString path; if (Settings->UseSectorMap) { path ="SectorMap:test-client[:2000]"; } else { - path = TStringBuilder() << Runtime->GetTempDir() << "pdisk_1.dat"; + TString diskPath = Settings->CustomDiskParams.DiskPath; + path = TStringBuilder() << (diskPath ? diskPath : Runtime->GetTempDir()) << "pdisk_1.dat"; } hostConfig.AddDrive()->SetPath(path); if (Settings->Verbose) { @@ -637,7 +642,9 @@ namespace Tests { for (const auto& [poolKind, storagePool] : Settings->StoragePoolTypes) { if (storagePool.GetNumGroups() > 0) { - bsConfigureRequest->Record.MutableRequest()->AddCommand()->MutableDefineStoragePool()->CopyFrom(storagePool); + auto* command = bsConfigureRequest->Record.MutableRequest()->AddCommand()->MutableDefineStoragePool(); + command->CopyFrom(storagePool); + command->SetItemConfigGeneration(Settings->StorageGeneration); } } @@ -1071,6 +1078,34 @@ namespace Tests { } } + { + if (Settings->NeedStatsCollectors) { + TString filePathPrefix; + if (Settings->AppConfig->HasMonitoringConfig()) { + filePathPrefix = Settings->AppConfig->GetMonitoringConfig().GetMemAllocDumpPathPrefix(); + } + + const TIntrusivePtr processMemoryInfoProvider(MakeIntrusive()); + + IActor* monitorActor = CreateMemProfMonitor(TDuration::Seconds(1), processMemoryInfoProvider, + Runtime->GetAppData(nodeIdx).Counters, filePathPrefix); + const TActorId monitorActorId = Runtime->Register(monitorActor, nodeIdx, Runtime->GetAppData(nodeIdx).BatchPoolId); + Runtime->RegisterService(MakeMemProfMonitorID(Runtime->GetNodeId(nodeIdx)), monitorActorId, nodeIdx); + + IActor* controllerActor = NMemory::CreateMemoryController(TDuration::Seconds(1), processMemoryInfoProvider, + Settings->AppConfig->GetMemoryControllerConfig(), NKikimrConfigHelpers::CreateMemoryControllerResourceBrokerConfig(*Settings->AppConfig), + Runtime->GetAppData(nodeIdx).Counters); + const TActorId controllerActorId = Runtime->Register(controllerActor, nodeIdx, Runtime->GetAppData(nodeIdx).BatchPoolId); + Runtime->RegisterService(NMemory::MakeMemoryControllerId(0), controllerActorId, nodeIdx); + } + } + + { + auto statActor = NStat::CreateStatService(); + const TActorId statActorId = Runtime->Register(statActor.Release(), nodeIdx, Runtime->GetAppData(nodeIdx).UserPoolId); + Runtime->RegisterService(NStat::MakeStatServiceID(Runtime->GetNodeId(nodeIdx)), statActorId, nodeIdx); + } + { IActor* kesusService = NKesus::CreateKesusProxyService(); TActorId kesusServiceId = Runtime->Register(kesusService, nodeIdx, userPoolId); diff --git a/ydb/core/testlib/test_client.h b/ydb/core/testlib/test_client.h index ce556d86c323..157ae243065b 100644 --- a/ydb/core/testlib/test_client.h +++ b/ydb/core/testlib/test_client.h @@ -122,6 +122,7 @@ namespace Tests { TString DomainName = TestDomainName; ui32 NodeCount = 1; ui32 DynamicNodeCount = 0; + ui64 StorageGeneration = 0; NFake::TStorage CustomDiskParams; TControls Controls; TAppPrepare::TFnReg FrFactory = &DefaultFrFactory; @@ -139,6 +140,7 @@ namespace Tests { bool UseRealThreads = true; bool EnableKqpSpilling = false; bool EnableYq = false; + bool EnableYqGrpc = false; TDuration KeepSnapshotTimeout = TDuration::Zero(); ui64 ChangesQueueItemsLimit = 0; ui64 ChangesQueueBytesLimit = 0; @@ -177,6 +179,7 @@ namespace Tests { TServerSettings& SetDomainName(const TString& value); TServerSettings& SetNodeCount(ui32 value) { NodeCount = value; return *this; } TServerSettings& SetDynamicNodeCount(ui32 value) { DynamicNodeCount = value; return *this; } + TServerSettings& SetStorageGeneration(ui64 value) { StorageGeneration = value; return *this; } TServerSettings& SetCustomDiskParams(const NFake::TStorage& value) { CustomDiskParams = value; return *this; } TServerSettings& SetControls(const TControls& value) { Controls = value; return *this; } TServerSettings& SetFrFactory(const TAppPrepare::TFnReg& value) { FrFactory = value; return *this; } @@ -203,6 +206,7 @@ namespace Tests { TServerSettings& SetEnableDbCounters(bool value) { FeatureFlags.SetEnableDbCounters(value); return *this; } TServerSettings& SetEnablePersistentQueryStats(bool value) { FeatureFlags.SetEnablePersistentQueryStats(value); return *this; } TServerSettings& SetEnableYq(bool value) { EnableYq = value; return *this; } + TServerSettings& SetEnableYqGrpc(bool value) { EnableYqGrpc = value; return *this; } TServerSettings& SetKeepSnapshotTimeout(TDuration value) { KeepSnapshotTimeout = value; return *this; } TServerSettings& SetChangesQueueItemsLimit(ui64 value) { ChangesQueueItemsLimit = value; return *this; } TServerSettings& SetChangesQueueBytesLimit(ui64 value) { ChangesQueueBytesLimit = value; return *this; } diff --git a/ydb/library/conclusion/generic/generic_status.h b/ydb/library/conclusion/generic/generic_status.h new file mode 100644 index 000000000000..e0a66bee606d --- /dev/null +++ b/ydb/library/conclusion/generic/generic_status.h @@ -0,0 +1,74 @@ +#pragma once + +#include +#include + +#include + +namespace NKikimr { + +template +class TConclusionStatusGenericImpl { +protected: + std::optional ErrorDescription; + TStatus Status = StatusOk; + + TConclusionStatusGenericImpl() = default; + + TConclusionStatusGenericImpl(const TErrorDescription& error, TStatus status = DefaultError) + : ErrorDescription(error) + , Status(status) { + Y_ABORT_UNLESS(!!ErrorDescription); + } + + TConclusionStatusGenericImpl(TErrorDescription&& error, TStatus status = DefaultError) + : ErrorDescription(std::move(error)) + , Status(status) { + Y_ABORT_UNLESS(!!ErrorDescription); + } + +public: + virtual ~TConclusionStatusGenericImpl() = default; + +public: + [[nodiscard]] const TErrorDescription& GetErrorDescription() const { + return ErrorDescription ? *ErrorDescription : Default(); + } + + [[nodiscard]] TStatus GetStatus() const { + return Status; + } + + template + [[nodiscard]] static TDerived Fail(const TErrorMessage& errorMessage) { + return TDerived(errorMessage); + } + + template + [[nodiscard]] static TDerived Fail(const TStatus& status, const TErrorMessage& errorMessage) { + Y_ABORT_UNLESS(DefaultError == StatusOk || status != StatusOk); + return TDerived(errorMessage, status); + } + + [[nodiscard]] bool IsFail() const { + return !Ok(); + } + + [[nodiscard]] bool IsSuccess() const { + return Ok(); + } + + [[nodiscard]] bool Ok() const { + return !ErrorDescription; + } + + [[nodiscard]] bool operator!() const { + return !!ErrorDescription; + } + + [[nodiscard]] static TDerived Success() { + return TDerived(); + } +}; + +} // namespace NKikimr diff --git a/ydb/library/conclusion/generic/result.h b/ydb/library/conclusion/generic/result.h index b0d93d3a404d..768ef2bceef3 100644 --- a/ydb/library/conclusion/generic/result.h +++ b/ydb/library/conclusion/generic/result.h @@ -1,5 +1,6 @@ #pragma once #include +#include #include #include @@ -55,13 +56,13 @@ class TConclusionImpl { const TResult& GetResult() const { auto result = std::get_if(&Result); - Y_ABORT_UNLESS(result, "incorrect object for result request"); + Y_ABORT_UNLESS(result, "incorrect object for result request: %s", GetErrorMessage().data()); return *result; } TResult& MutableResult() { auto result = std::get_if(&Result); - Y_ABORT_UNLESS(result, "incorrect object for result request"); + Y_ABORT_UNLESS(result, "incorrect object for result request: %s", GetErrorMessage().data()); return *result; } @@ -91,13 +92,10 @@ class TConclusionImpl { return GetError(); } - const TString& GetErrorMessage() const { + TString GetErrorMessage() const { auto* status = std::get_if(&Result); - if (!status) { - return Default(); - } else { - return status->GetErrorMessage(); - } + Y_ABORT_UNLESS(status, "incorrect object for extracting error message"); + return status->GetErrorMessage(); } auto GetStatus() const { diff --git a/ydb/library/conclusion/generic/status.h b/ydb/library/conclusion/generic/status.h deleted file mode 100644 index 26be88712b50..000000000000 --- a/ydb/library/conclusion/generic/status.h +++ /dev/null @@ -1,94 +0,0 @@ -#pragma once - -#include -#include - -#include - -namespace NKikimr { - -template -class TConclusionStatusImpl { -private: - std::optional ErrorMessage; - TStatus Status = StatusOk; - TConclusionStatusImpl() = default; - TConclusionStatusImpl(const TString& errorMessage, TStatus status = DefaultError) - : ErrorMessage(errorMessage) - , Status(status) { - Y_ABORT_UNLESS(!!ErrorMessage); - } - - TConclusionStatusImpl(const char* errorMessage, TStatus status = DefaultError) - : ErrorMessage(errorMessage) - , Status(status) { - Y_ABORT_UNLESS(!!ErrorMessage); - } - - TConclusionStatusImpl(const std::string& errorMessage, TStatus status = DefaultError) - : ErrorMessage(TString(errorMessage.data(), errorMessage.size())) - , Status(status) { - Y_ABORT_UNLESS(!!ErrorMessage); - } - -public: - void Validate(const TString& processInfo = Default()) const { - if (processInfo) { - Y_ABORT_UNLESS(Ok(), "error=%s, processInfo=%s", GetErrorMessage().c_str(), processInfo.c_str()); - } else { - Y_ABORT_UNLESS(Ok(), "error=%s", GetErrorMessage().c_str()); - } - } - - [[nodiscard]] const TString& GetErrorMessage() const { - return ErrorMessage ? *ErrorMessage : Default(); - } - - [[nodiscard]] TStatus GetStatus() const { - return Status; - } - - [[nodiscard]] static TConclusionStatusImpl Fail(const char* errorMessage) { - return TConclusionStatusImpl(errorMessage); - } - - [[nodiscard]] static TConclusionStatusImpl Fail(const TString& errorMessage) { - return TConclusionStatusImpl(errorMessage); - } - - [[nodiscard]] static TConclusionStatusImpl Fail(const std::string& errorMessage) { - return TConclusionStatusImpl(errorMessage); - } - - [[nodiscard]] static TConclusionStatusImpl Fail(const TStatus& status, const char* errorMessage) { - Y_ABORT_UNLESS(DefaultError == StatusOk || status != StatusOk); - return TConclusionStatusImpl(errorMessage, status); - } - - [[nodiscard]] static TConclusionStatusImpl Fail(const TStatus& status, const TString& errorMessage) { - Y_ABORT_UNLESS(DefaultError == StatusOk || status != StatusOk); - return TConclusionStatusImpl(errorMessage, status); - } - - [[nodiscard]] bool IsFail() const { - return !Ok(); - } - - [[nodiscard]] bool IsSuccess() const { - return Ok(); - } - - [[nodiscard]] bool Ok() const { - return !ErrorMessage; - } - - [[nodiscard]] bool operator!() const { - return !!ErrorMessage; - } - - [[nodiscard]] static TConclusionStatusImpl Success() { - return TConclusionStatusImpl(); - } -}; - -} // namespace NKikimr diff --git a/ydb/library/conclusion/generic/string_status.h b/ydb/library/conclusion/generic/string_status.h new file mode 100644 index 000000000000..ccb8ff112146 --- /dev/null +++ b/ydb/library/conclusion/generic/string_status.h @@ -0,0 +1,42 @@ +#pragma once + +#include "generic_status.h" + +#include + +namespace NKikimr { + +template +class TConclusionStatusImpl : public TConclusionStatusGenericImpl, TString, TStatus, StatusOk, DefaultError> { +protected: + using TSelf = TConclusionStatusImpl; + using TBase = TConclusionStatusGenericImpl; + using TBase::TBase; + + friend class TConclusionStatusGenericImpl; + + TConclusionStatusImpl() = default; + + TConclusionStatusImpl(const char* errorMessage, TStatus status = DefaultError) + : TBase(TString(errorMessage), status) { + } + + TConclusionStatusImpl(const std::string& errorMessage, TStatus status = DefaultError) + : TBase(TString(errorMessage), status) { + } + +public: + void Validate(const TString& processInfo = Default()) const { + if (processInfo) { + Y_ABORT_UNLESS(TBase::Ok(), "error=%s, processInfo=%s", GetErrorMessage().c_str(), processInfo.c_str()); + } else { + Y_ABORT_UNLESS(TBase::Ok(), "error=%s", GetErrorMessage().c_str()); + } + } + + [[nodiscard]] TString GetErrorMessage() const { + return TBase::GetErrorDescription(); + } +}; + +} // namespace NKikimr diff --git a/ydb/library/conclusion/generic/ya.make b/ydb/library/conclusion/generic/ya.make index 1e614b2bfb5c..5232905ff9ff 100644 --- a/ydb/library/conclusion/generic/ya.make +++ b/ydb/library/conclusion/generic/ya.make @@ -4,6 +4,8 @@ SRCS() PEERDIR( util + + yql/essentials/public/issue ) END() diff --git a/ydb/library/conclusion/generic/yql_status.h b/ydb/library/conclusion/generic/yql_status.h new file mode 100644 index 000000000000..8b1d4f10107d --- /dev/null +++ b/ydb/library/conclusion/generic/yql_status.h @@ -0,0 +1,55 @@ +#pragma once + +#include "generic_status.h" + +#include + +namespace NKikimr { + +template +class TYQLConclusionStatusImpl : public TConclusionStatusGenericImpl, NYql::TIssues, TStatus, StatusOk, DefaultError> { +protected: + using TSelf = TYQLConclusionStatusImpl; + using TBase = TConclusionStatusGenericImpl; + using TBase::TBase; + + friend class TConclusionStatusGenericImpl; + + TYQLConclusionStatusImpl() = default; + + TYQLConclusionStatusImpl(const TString& errorMessage, TStatus status = DefaultError) + : TBase({NYql::TIssue(errorMessage)}, status) { + } + +public: + TYQLConclusionStatusImpl& AddParentIssue(NYql::TIssue issue) { + Y_ABORT_UNLESS(!!TBase::ErrorDescription); + for (const auto& childIssue : *TBase::ErrorDescription) { + issue.AddSubIssue(MakeIntrusive(childIssue)); + } + TBase::ErrorDescription = {std::move(issue)}; + return *this; + } + + TYQLConclusionStatusImpl& AddParentIssue(const TString& message) { + AddParentIssue(NYql::TIssue(message)); + return *this; + } + + TYQLConclusionStatusImpl& AddIssue(NYql::TIssue issue) { + Y_ABORT_UNLESS(!!TBase::ErrorDescription); + TBase::ErrorDescription->AddIssue(std::move(issue)); + return *this; + } + + TYQLConclusionStatusImpl& AddIssue(const TString& message) { + AddIssue(NYql::TIssue(message)); + return *this; + } + + [[nodiscard]] TString GetErrorMessage() const { + return TBase::GetErrorDescription().ToOneLineString(); + } +}; + +} // namespace NKikimr diff --git a/ydb/library/conclusion/status.h b/ydb/library/conclusion/status.h index b6a0830cf7a6..3a56012ca36e 100644 --- a/ydb/library/conclusion/status.h +++ b/ydb/library/conclusion/status.h @@ -1,5 +1,7 @@ #pragma once -#include + +#include +#include namespace NKikimr { @@ -8,4 +10,7 @@ using TConclusionStatus = TConclusionStatusImpl<::TNull, ::TNull{}, ::TNull{}>; template using TConclusionSpecialStatus = TConclusionStatusImpl; +template +using TYQLConclusionSpecialStatus = TYQLConclusionStatusImpl; + } diff --git a/ydb/library/grpc/client/grpc_client_low.cpp b/ydb/library/grpc/client/grpc_client_low.cpp index bf7af6cc91dc..687b4b6e45a9 100644 --- a/ydb/library/grpc/client/grpc_client_low.cpp +++ b/ydb/library/grpc/client/grpc_client_low.cpp @@ -141,15 +141,8 @@ void TChannelPool::GetStubsHolderLocked( } } } - TGRpcKeepAliveSocketMutator* mutator = nullptr; + auto mutator = NImpl::CreateGRpcKeepAliveSocketMutator(TcpKeepAliveSettings_); // will be destroyed inside grpc - if (TcpKeepAliveSettings_.Enabled) { - mutator = new TGRpcKeepAliveSocketMutator( - TcpKeepAliveSettings_.Idle, - TcpKeepAliveSettings_.Count, - TcpKeepAliveSettings_.Interval - ); - } cb(Pool_.emplace(channelId, CreateChannelInterface(config, mutator)).first->second); LastUsedQueue_.emplace(Pool_.at(channelId).GetLastUseTime(), channelId); } @@ -588,4 +581,16 @@ void TGRpcClientLow::ForgetContext(TContextImpl* context) { } } +grpc_socket_mutator* NImpl::CreateGRpcKeepAliveSocketMutator(const TTcpKeepAliveSettings& TcpKeepAliveSettings_) { + TGRpcKeepAliveSocketMutator* mutator = nullptr; + if (TcpKeepAliveSettings_.Enabled) { + mutator = new TGRpcKeepAliveSocketMutator( + TcpKeepAliveSettings_.Idle, + TcpKeepAliveSettings_.Count, + TcpKeepAliveSettings_.Interval + ); + } + return mutator; +} + } // namespace NGRpc diff --git a/ydb/library/grpc/client/grpc_client_low.h b/ydb/library/grpc/client/grpc_client_low.h index 4f944d7c559b..233f9fbab39b 100644 --- a/ydb/library/grpc/client/grpc_client_low.h +++ b/ydb/library/grpc/client/grpc_client_low.h @@ -451,6 +451,10 @@ class IStreamRequestReadWriteProcessor : public IStreamRequestReadProcessor>(new TServiceConnection(CreateChannelInterface(config), this)); } + template + std::unique_ptr> CreateGRpcServiceConnection(const TGRpcClientConfig& config, const TTcpKeepAliveSettings& keepAlive) { + auto mutator = NImpl::CreateGRpcKeepAliveSocketMutator(keepAlive); + // will be destroyed inside grpc + return std::unique_ptr>(new TServiceConnection(CreateChannelInterface(config, mutator), this)); + } + template std::unique_ptr> CreateGRpcServiceConnection(TStubsHolder& holder) { return std::unique_ptr>(new TServiceConnection(holder, this)); diff --git a/ydb/library/yql/dq/actors/common/retry_queue.h b/ydb/library/yql/dq/actors/common/retry_queue.h index ab4e0f84c498..2124a0e9f6e2 100644 --- a/ydb/library/yql/dq/actors/common/retry_queue.h +++ b/ydb/library/yql/dq/actors/common/retry_queue.h @@ -186,6 +186,9 @@ class TRetryEventsQueue { THolder ev = MakeHolder(); ev->Record = Event->Record; ev->Record.MutableTransportMeta()->SetConfirmedSeqNo(confirmedSeqNo); + for (ui32 i = 0; i < Event->GetPayloadCount(); ++i) { + ev->AddPayload(TRope(Event->GetPayload(i))); + } return MakeHolder(Recipient, Sender, ev.Release(), NActors::IEventHandle::FlagTrackDelivery, Cookie); } diff --git a/ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp b/ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp index ed045434a9a9..167a9490323d 100644 --- a/ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp +++ b/ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp @@ -173,13 +173,15 @@ class TDqAsyncComputeActor : public TDqComputeActorBaseState"; html << "
" << ComputeActorState.DebugString() << "
"; -#define DUMP(P, X) html << #X ": " << P.X << "
" +#define DUMP(P, X,...) html << #X ": " << P.X __VA_ARGS__ << "
" +#define DUMP_PREFIXED(TITLE, S, FIELD,...) html << TITLE << #FIELD ": " << S . FIELD __VA_ARGS__ << "
" html << "

ProcessSourcesState

"; DUMP(ProcessSourcesState, Inflight); html << "

ProcessOutputsState

"; DUMP(ProcessOutputsState, Inflight); DUMP(ProcessOutputsState, ChannelsReady); DUMP(ProcessOutputsState, HasDataToSend); + DUMP(ProcessOutputsState, DataWasSent); DUMP(ProcessOutputsState, AllOutputsFinished); DUMP(ProcessOutputsState, LastRunStatus); DUMP(ProcessOutputsState, LastPopReturnedNoData); @@ -192,12 +194,11 @@ class TDqAsyncComputeActor : public TDqComputeActorBaseCPU Quota"; html << "QuoterServiceActorId: " << QuoterServiceActorId.ToString() << "
"; if (ContinueRunEvent) { - html << "ContinueRunEvent.AskFreeSpace: " << ContinueRunEvent->AskFreeSpace << "
"; - html << "ContinueRunEvent.AskFreeSpace: " << ContinueRunEvent->CheckpointOnly << "
"; - html << "ContinueRunEvent.AskFreeSpace: " << ContinueRunEvent->CheckpointRequest.Defined() << "
"; - html << "ContinueRunEvent.AskFreeSpace: " << ContinueRunEvent->WatermarkRequest.Defined() << "
"; - html << "ContinueRunEvent.AskFreeSpace: " << ContinueRunEvent->CheckpointOnly << "
"; - html << "ContinueRunEvent.AskFreeSpace: " << ContinueRunEvent->MemLimit << "
"; + DUMP_PREFIXED("ContinueRunEvent.", (*ContinueRunEvent), AskFreeSpace); + DUMP_PREFIXED("ContinueRunEvent.", (*ContinueRunEvent), CheckpointOnly); + DUMP_PREFIXED("ContinueRunEvent.", (*ContinueRunEvent), CheckpointRequest, .Defined()); + DUMP_PREFIXED("ContinueRunEvent.", (*ContinueRunEvent), WatermarkRequest, .Defined()); + DUMP_PREFIXED("ContinueRunEvent.", (*ContinueRunEvent), MemLimit); for (const auto& sinkId: ContinueRunEvent->SinkIds) { html << "ContinueRunEvent.SinkIds: " << sinkId << "
"; } @@ -207,46 +208,49 @@ class TDqAsyncComputeActor : public TDqComputeActorBase"; - html << "ContinueRunInflight: " << ContinueRunInflight << "
"; - html << "CpuTimeSpent: " << CpuTimeSpent.ToString() << "
"; - html << "CpuTimeQuotaAsked: " << CpuTimeQuotaAsked.ToString() << "
"; - html << "UseCpuQuota: " << UseCpuQuota() << "
"; - + DUMP((*this), ContinueRunStartWaitTime, .ToString()); + DUMP((*this), ContinueRunInflight); + DUMP((*this), CpuTimeSpent, .ToString()); + DUMP((*this), CpuTimeQuotaAsked, .ToString()); + DUMP((*this), UseCpuQuota, ()); html << "

Checkpoints

"; - html << "ReadyToCheckpoint: " << ReadyToCheckpoint() << "
"; - html << "CheckpointRequestedFromTaskRunner: " << CheckpointRequestedFromTaskRunner << "
"; + DUMP((*this), ReadyToCheckpoint, ()); + DUMP((*this), CheckpointRequestedFromTaskRunner); auto dumpAsyncStats = [&](auto prefix, auto& asyncStats) { html << prefix << "Level: " << static_cast(asyncStats.Level) << "
"; - html << prefix << "MinWaitDuration: " << asyncStats.MinWaitDuration.ToString() << "
"; + DUMP_PREFIXED(prefix, asyncStats, MinWaitDuration, .ToString()); html << prefix << "CurrentPauseTs: " << (asyncStats.CurrentPauseTs ? asyncStats.CurrentPauseTs->ToString() : TString{}) << "
"; - html << prefix << "MergeWaitPeriod: " << asyncStats.MergeWaitPeriod << "
"; - html << prefix << "Bytes: " << asyncStats.Bytes << "
"; - html << prefix << "DecompressedBytes: " << asyncStats.DecompressedBytes << "
"; - html << prefix << "Rows: " << asyncStats.Rows << "
"; - html << prefix << "Chunks: " << asyncStats.Chunks << "
"; - html << prefix << "Splits: " << asyncStats.Splits << "
"; - html << prefix << "FirstMessageTs: " << asyncStats.FirstMessageTs.ToString() << "
"; - html << prefix << "PauseMessageTs: " << asyncStats.PauseMessageTs.ToString() << "
"; - html << prefix << "ResumeMessageTs: " << asyncStats.ResumeMessageTs.ToString() << "
"; - html << prefix << "LastMessageTs: " << asyncStats.LastMessageTs.ToString() << "
"; - html << prefix << "WaitTime: " << asyncStats.WaitTime.ToString() << "
"; + DUMP_PREFIXED(prefix, asyncStats, MergeWaitPeriod); + DUMP_PREFIXED(prefix, asyncStats, Bytes); + DUMP_PREFIXED(prefix, asyncStats, DecompressedBytes); + DUMP_PREFIXED(prefix, asyncStats, Rows); + DUMP_PREFIXED(prefix, asyncStats, Chunks); + DUMP_PREFIXED(prefix, asyncStats, Splits); + DUMP_PREFIXED(prefix, asyncStats, FilteredBytes); + DUMP_PREFIXED(prefix, asyncStats, FilteredRows); + DUMP_PREFIXED(prefix, asyncStats, QueuedBytes); + DUMP_PREFIXED(prefix, asyncStats, QueuedRows); + DUMP_PREFIXED(prefix, asyncStats, FirstMessageTs, .ToString()); + DUMP_PREFIXED(prefix, asyncStats, PauseMessageTs, .ToString()); + DUMP_PREFIXED(prefix, asyncStats, ResumeMessageTs, .ToString()); + DUMP_PREFIXED(prefix, asyncStats, LastMessageTs, .ToString()); + DUMP_PREFIXED(prefix, asyncStats, WaitTime, .ToString()); }; auto dumpOutputStats = [&](auto prefix, auto& outputStats) { - html << prefix << "MaxMemoryUsage: " << outputStats.MaxMemoryUsage << "
"; - html << prefix << "MaxRowsInMemory: " << outputStats.MaxRowsInMemory << "
"; + DUMP_PREFIXED(prefix, outputStats, MaxMemoryUsage); + DUMP_PREFIXED(prefix, outputStats, MaxRowsInMemory); dumpAsyncStats(prefix, outputStats); }; auto dumpInputChannelStats = [&](auto prefix, auto& pushStats) { - html << prefix << "ChannelId: " << pushStats.ChannelId << "
"; - html << prefix << "SrcStageId: " << pushStats.SrcStageId << "
"; - html << prefix << "RowsInMemory: " << pushStats.RowsInMemory << "
"; - html << prefix << "MaxMemoryUsage: " << pushStats.MaxMemoryUsage << "
"; - html << prefix << "DeserializationTime: " << pushStats.DeserializationTime.ToString() << "
"; + DUMP_PREFIXED(prefix, pushStats, ChannelId); + DUMP_PREFIXED(prefix, pushStats, SrcStageId); + DUMP_PREFIXED(prefix, pushStats, RowsInMemory); + DUMP_PREFIXED(prefix, pushStats, MaxMemoryUsage); + DUMP_PREFIXED(prefix, pushStats, DeserializationTime, .ToString()); dumpAsyncStats(prefix, pushStats); }; @@ -265,19 +269,32 @@ class TDqAsyncComputeActor : public TDqComputeActorBase"; DUMP(info, FreeSpace); html << "IsPaused: " << info.IsPaused() << "
"; - if (info.Channel) { - html << "DqInputChannel.ChannelId: " << info.Channel->GetChannelId() << "
"; - html << "DqInputChannel.FreeSpace: " << info.Channel->GetFreeSpace() << "
"; - html << "DqInputChannel.StoredBytes: " << info.Channel->GetStoredBytes() << "
"; - html << "DqInputChannel.Empty: " << info.Channel->Empty() << "
"; - html << "DqInputChannel.InputType: " << (info.Channel->GetInputType() ? info.Channel->GetInputType()->GetKindAsStr() : TString{"unknown"}) << "
"; - html << "DqInputChannel.InputWidth: " << (info.Channel->GetInputWidth() ? ToString(*info.Channel->GetInputWidth()) : TString{"unknown"}) << "
"; - html << "DqInputChannel.IsFinished: " << info.Channel->IsFinished() << "
"; - - const auto& pushStats = info.Channel->GetPushStats(); + auto channel = info.Channel; + if (!channel) { + auto stats = GetTaskRunnerStats(); + if (stats) { + auto stageIt = stats->InputChannels.find(info.SrcStageId); + if (stageIt != stats->InputChannels.end()) { + auto channelIt = stageIt->second.find(info.ChannelId); + if (channelIt != stageIt->second.end()) { + channel = channelIt->second; + } + } + } + } + if (channel) { + html << "DqInputChannel.ChannelId: " << channel->GetChannelId() << "
"; + html << "DqInputChannel.FreeSpace: " << channel->GetFreeSpace() << "
"; + html << "DqInputChannel.StoredBytes: " << channel->GetStoredBytes() << "
"; + html << "DqInputChannel.Empty: " << channel->Empty() << "
"; + html << "DqInputChannel.InputType: " << (channel->GetInputType() ? channel->GetInputType()->GetKindAsStr() : TString{"unknown"}) << "
"; + html << "DqInputChannel.InputWidth: " << (channel->GetInputWidth() ? ToString(*channel->GetInputWidth()) : TString{"unknown"}) << "
"; + html << "DqInputChannel.IsFinished: " << channel->IsFinished() << "
"; + + const auto& pushStats = channel->GetPushStats(); dumpInputChannelStats("DqInputChannel.PushStats.", pushStats); - const auto& popStats = info.Channel->GetPopStats(); + const auto& popStats = channel->GetPopStats(); dumpInputStats("DqInputChannel.PopStats."sv, popStats); } } @@ -290,18 +307,19 @@ class TDqAsyncComputeActor : public TDqComputeActorBaseToString()) << "
"; html << "WatermarksMode: " << NDqProto::EWatermarksMode_Name(info.WatermarksMode) << "
"; html << "FreeSpace: " << info.GetFreeSpace() << "
"; - if (info.Buffer) { - html << "DqInputBuffer.InputIndex: " << info.Buffer->GetInputIndex() << "
"; - html << "DqInputBuffer.FreeSpace: " << info.Buffer->GetFreeSpace() << "
"; - html << "DqInputBuffer.StoredBytes: " << info.Buffer->GetStoredBytes() << "
"; - html << "DqInputBuffer.Empty: " << info.Buffer->Empty() << "
"; - html << "DqInputBuffer.InputType: " << (info.Buffer->GetInputType() ? info.Buffer->GetInputType()->GetKindAsStr() : TString{"unknown"}) << "
"; - html << "DqInputBuffer.InputWidth: " << (info.Buffer->GetInputWidth() ? ToString(*info.Buffer->GetInputWidth()) : TString{"unknown"}) << "
"; - html << "DqInputBuffer.IsFinished: " << info.Buffer->IsFinished() << "
"; - html << "DqInputBuffer.IsPaused: " << info.Buffer->IsPaused() << "
"; - html << "DqInputBuffer.IsPending: " << info.Buffer->IsPending() << "
"; - - const auto& popStats = info.Buffer->GetPopStats(); + auto buffer = info.Buffer; + if (buffer) { + html << "DqInputBuffer.InputIndex: " << buffer->GetInputIndex() << "
"; + html << "DqInputBuffer.FreeSpace: " << buffer->GetFreeSpace() << "
"; + html << "DqInputBuffer.StoredBytes: " << buffer->GetStoredBytes() << "
"; + html << "DqInputBuffer.Empty: " << buffer->Empty() << "
"; + html << "DqInputBuffer.InputType: " << (buffer->GetInputType() ? buffer->GetInputType()->GetKindAsStr() : TString{"unknown"}) << "
"; + html << "DqInputBuffer.InputWidth: " << (buffer->GetInputWidth() ? ToString(*buffer->GetInputWidth()) : TString{"unknown"}) << "
"; + html << "DqInputBuffer.IsFinished: " << buffer->IsFinished() << "
"; + html << "DqInputBuffer.IsPaused: " << buffer->IsPaused() << "
"; + html << "DqInputBuffer.IsPending: " << buffer->IsPending() << "
"; + + const auto& popStats = buffer->GetPopStats(); dumpInputStats("DqInputBuffer."sv, popStats); } if (info.AsyncInput) { @@ -332,18 +350,32 @@ class TDqAsyncComputeActor : public TDqComputeActorBaseWatermark << "
"; } - if (info.Channel) { - html << "DqOutputChannel.ChannelId: " << info.Channel->GetChannelId() << "
"; - html << "DqOutputChannel.ValuesCount: " << info.Channel->GetValuesCount() << "
"; - html << "DqOutputChannel.IsFull: " << info.Channel->IsFull() << "
"; - html << "DqOutputChannel.HasData: " << info.Channel->HasData() << "
"; - html << "DqOutputChannel.IsFinished: " << info.Channel->IsFinished() << "
"; - html << "DqInputChannel.OutputType: " << (info.Channel->GetOutputType() ? info.Channel->GetOutputType()->GetKindAsStr() : TString{"unknown"}) << "
"; + auto channel = info.Channel; + if (!channel) { + auto stats = GetTaskRunnerStats(); + if (stats) { + auto stageIt = stats->OutputChannels.find(info.DstStageId); + if (stageIt != stats->OutputChannels.end()) { + auto channelIt = stageIt->second.find(info.ChannelId); + if (channelIt != stageIt->second.end()) { + channel = channelIt->second; + } + } + } + } + + if (channel) { + html << "DqOutputChannel.ChannelId: " << channel->GetChannelId() << "
"; + html << "DqOutputChannel.ValuesCount: " << channel->GetValuesCount() << "
"; + html << "DqOutputChannel.IsFull: " << channel->IsFull() << "
"; + html << "DqOutputChannel.HasData: " << channel->HasData() << "
"; + html << "DqOutputChannel.IsFinished: " << channel->IsFinished() << "
"; + html << "DqInputChannel.OutputType: " << (channel->GetOutputType() ? channel->GetOutputType()->GetKindAsStr() : TString{"unknown"}) << "
"; - const auto& pushStats = info.Channel->GetPushStats(); + const auto& pushStats = channel->GetPushStats(); dumpOutputStats("DqOutputChannel.PushStats."sv, pushStats); - const auto& popStats = info.Channel->GetPopStats(); + const auto& popStats = channel->GetPopStats(); html << "DqOutputChannel.PopStats.ChannelId: " << popStats.ChannelId << "
"; html << "DqOutputChannel.PopStats.DstStageId: " << popStats.DstStageId << "
"; html << "DqOutputChannel.PopStats.MaxMemoryUsage: " << popStats.MaxMemoryUsage << "
"; @@ -364,8 +396,8 @@ class TDqAsyncComputeActor : public TDqComputeActorBase"; html << "DqOutputBuffer.IsFull: " << buffer.IsFull() << "
"; html << "DqOutputBuffer.OutputType: " << (buffer.GetOutputType() ? buffer.GetOutputType()->GetKindAsStr() : TString{"unknown"}) << "
"; @@ -404,6 +436,7 @@ class TDqAsyncComputeActor : public TDqComputeActorBaseSender, new NActors::NMon::TEvHttpInfoRes(html.Str())); } @@ -451,6 +484,7 @@ class TDqAsyncComputeActor : public TDqComputeActorBaseGet()->Stats) { CA_LOG_T("update task runner stats"); TaskRunnerStats = std::move(ev->Get()->Stats); + TaskRunnerActorElapsedTicks = TaskRunnerStats.GetActorElapsedTicks(); } ComputeActorState = NDqProto::TEvComputeActorState(); ComputeActorState.SetState(NDqProto::COMPUTE_STATE_EXECUTING); @@ -765,12 +799,19 @@ class TDqAsyncComputeActor : public TDqComputeActorBaseGet()->ComputeTime; + CpuTimeSpent += TakeCpuTimeDelta(); AskCpuQuota(); ProcessContinueRun(); } } + TDuration TakeCpuTimeDelta() { + auto newTicks = ComputeActorElapsedTicks + TaskRunnerActorElapsedTicks; + auto result = newTicks - LastQuotaElapsedTicks; + LastQuotaElapsedTicks = newTicks; + return TDuration::MicroSeconds(NHPTimer::GetSeconds(result) * 1'000'000ull); + } + void SaveState(const NDqProto::TCheckpoint& checkpoint, TComputeActorState& state) const override { CA_LOG_D("Save state"); Y_ABORT_UNLESS(ProgramState); @@ -1184,6 +1225,7 @@ class TDqAsyncComputeActor : public TDqComputeActorBase ContinueRunEvent; TInstant ContinueRunStartWaitTime; bool ContinueRunInflight = false; diff --git a/ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h b/ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h index def8b094db14..485f12c40224 100644 --- a/ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h +++ b/ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h @@ -98,6 +98,26 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped , public TSinkCallbacks , public TOutputTransformCallbacks { +private: + struct TEvPrivate { + enum EEv : ui32 { + EvAsyncOutputError = EventSpaceBegin(NActors::TEvents::ES_PRIVATE), + EvEnd + }; + + static_assert(EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(TEvents::ES_PRIVATE)"); + + struct TEvAsyncOutputError : public NActors::TEventLocal { + TEvAsyncOutputError(NYql::NDqProto::StatusIds::StatusCode statusCode, const TIssues& issues) + : StatusCode(statusCode) + , Issues(issues) + {} + + NYql::NDqProto::StatusIds::StatusCode StatusCode; + NYql::TIssues Issues; + }; + }; + protected: enum EEvWakeupTag : ui64 { TimeoutTag = 1, @@ -248,8 +268,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped void ReportEventElapsedTime() { if (RuntimeSettings.CollectBasic()) { - ui64 elapsedMicros = NActors::TlsActivationContext->GetCurrentEventTicksAsSeconds() * 1'000'000ull; - CpuTime += TDuration::MicroSeconds(elapsedMicros); + ComputeActorElapsedTicks += NActors::TlsActivationContext->GetCurrentEventTicks(); } } @@ -297,6 +316,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped hFunc(NActors::TEvInterconnect::TEvNodeConnected, HandleExecuteBase); hFunc(IDqComputeActorAsyncInput::TEvNewAsyncInputDataArrived, OnNewAsyncInputDataArrived); hFunc(IDqComputeActorAsyncInput::TEvAsyncInputError, OnAsyncInputError); + hFunc(TEvPrivate::TEvAsyncOutputError, HandleAsyncOutputError); default: { CA_LOG_C("TDqComputeActorBase, unexpected event: " << ev->GetTypeRewrite() << " (" << GetEventTypeString(ev) << ")"); InternalError(NYql::NDqProto::StatusIds::INTERNAL_ERROR, TIssuesIds::DEFAULT_ERROR, TStringBuilder() << "Unexpected event: " << ev->GetTypeRewrite() << " (" << GetEventTypeString(ev) << ")"); @@ -368,13 +388,13 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped } void ProcessOutputsImpl(ERunStatus status) { - ProcessOutputsState.LastRunStatus = status; - CA_LOG_T("ProcessOutputsState.Inflight: " << ProcessOutputsState.Inflight); if (ProcessOutputsState.Inflight == 0) { ProcessOutputsState = TProcessOutputsState(); } + ProcessOutputsState.LastRunStatus = status; + for (auto& entry : OutputChannelsMap) { const ui64 channelId = entry.first; TOutputChannelInfo& outputChannel = entry.second; @@ -439,6 +459,14 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped return; } + if (status != ERunStatus::Finished) { + for (auto& [id, inputTransform] : InputTransformsMap) { + if (!inputTransform.Buffer->Empty()) { + ContinueExecute(EResumeSource::CAPendingInput); + } + } + } + if (status != ERunStatus::Finished) { // If the incoming channel's buffer was full at the moment when last ChannelDataAck event had been sent, // there will be no attempts to send a new piece of data from the other side of this channel. @@ -1517,7 +1545,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped } CA_LOG_E("Sink[" << outputIndex << "] fatal error: " << issues.ToOneLineString()); - InternalError(fatalCode, issues); + this->Send(this->SelfId(), new TEvPrivate::TEvAsyncOutputError(fatalCode, issues)); } void OnOutputTransformError(ui64 outputIndex, const TIssues& issues, NYql::NDqProto::StatusIds::StatusCode fatalCode) override final { @@ -1527,7 +1555,11 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped } CA_LOG_E("OutputTransform[" << outputIndex << "] fatal error: " << issues.ToOneLineString()); - InternalError(fatalCode, issues); + this->Send(this->SelfId(), new TEvPrivate::TEvAsyncOutputError(fatalCode, issues)); + } + + void HandleAsyncOutputError(const TEvPrivate::TEvAsyncOutputError::TPtr& ev) { + InternalError(ev->Get()->StatusCode, ev->Get()->Issues); } bool AllAsyncOutputsFinished() const { @@ -1696,7 +1728,8 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped ReportEventElapsedTime(); } - dst->SetCpuTimeUs(CpuTime.MicroSeconds() + SourceCpuTime.MicroSeconds() + InputTransformCpuTime.MicroSeconds()); + ui64 computeActorElapsedUs = NHPTimer::GetSeconds(ComputeActorElapsedTicks) * 1'000'000ull; + dst->SetCpuTimeUs(computeActorElapsedUs + SourceCpuTime.MicroSeconds() + InputTransformCpuTime.MicroSeconds()); dst->SetMaxMemoryUsage(MemoryLimits.MemoryQuotaManager->GetMaxMemorySize()); if (auto memProfileStats = GetMemoryProfileStats(); memProfileStats) { @@ -1732,14 +1765,18 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped auto cpuTimeUs = taskStats->ComputeCpuTime.MicroSeconds() + taskStats->BuildCpuTime.MicroSeconds(); if (TDerived::HasAsyncTaskRunner) { // Async TR is another actor, summarize CPU usage - cpuTimeUs += CpuTime.MicroSeconds(); + cpuTimeUs = NHPTimer::GetSeconds(ComputeActorElapsedTicks + TaskRunnerActorElapsedTicks) * 1'000'000ull; } - // CpuTimeUs does include SourceCpuTime + // cpuTimeUs does include SourceCpuTime protoTask->SetCpuTimeUs(cpuTimeUs + SourceCpuTime.MicroSeconds() + InputTransformCpuTime.MicroSeconds()); protoTask->SetSourceCpuTimeUs(SourceCpuTime.MicroSeconds()); ui64 ingressBytes = 0; ui64 ingressRows = 0; + ui64 filteredBytes = 0; + ui64 filteredRows = 0; + ui64 queuedBytes = 0; + ui64 queuedRows = 0; ui64 ingressDecompressedBytes = 0; auto startTimeMs = protoTask->GetStartTimeMs(); @@ -1749,6 +1786,9 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped auto inputIndex = protoSource.GetInputIndex(); if (auto* sourceInfoPtr = SourcesMap.FindPtr(inputIndex)) { auto& sourceInfo = *sourceInfoPtr; + if (!sourceInfo.AsyncInput) + continue; + protoSource.SetIngressName(sourceInfo.Type); const auto& ingressStats = sourceInfo.AsyncInput->GetIngressStats(); FillAsyncStats(*protoSource.MutableIngress(), ingressStats); @@ -1762,6 +1802,10 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped startTimeMs = firstMessageMs; } } + filteredBytes += ingressStats.FilteredBytes; + filteredRows += ingressStats.FilteredRows; + queuedBytes += ingressStats.QueuedBytes; + queuedRows += ingressStats.QueuedRows; } } } else { @@ -1775,6 +1819,10 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped // ingress rows are usually not reported, so we count rows in task runner input ingressRows += ingressStats.Rows ? ingressStats.Rows : taskStats->Sources.at(inputIndex)->GetPopStats().Rows; ingressDecompressedBytes += ingressStats.DecompressedBytes; + filteredBytes += ingressStats.FilteredBytes; + filteredRows += ingressStats.FilteredRows; + queuedBytes += ingressStats.QueuedBytes; + queuedRows += ingressStats.QueuedRows; } } @@ -1785,6 +1833,10 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped protoTask->SetIngressBytes(ingressBytes); protoTask->SetIngressRows(ingressRows); protoTask->SetIngressDecompressedBytes(ingressDecompressedBytes); + protoTask->SetIngressFilteredBytes(filteredBytes); + protoTask->SetIngressFilteredRows(filteredRows); + protoTask->SetIngressQueuedBytes(queuedBytes); + protoTask->SetIngressQueuedRows(queuedRows); ui64 egressBytes = 0; ui64 egressRows = 0; @@ -1984,7 +2036,8 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped bool ResumeEventScheduled = false; NDqProto::EComputeState State; TIntrusivePtr RequestContext; - TDuration CpuTime; + ui64 ComputeActorElapsedTicks = 0; + ui64 TaskRunnerActorElapsedTicks = 0; struct TProcessOutputsState { int Inflight = 0; diff --git a/ydb/library/yql/dq/actors/compute/dq_compute_actor_stats.cpp b/ydb/library/yql/dq/actors/compute/dq_compute_actor_stats.cpp index 8a35905543cd..cba8011bc33e 100644 --- a/ydb/library/yql/dq/actors/compute/dq_compute_actor_stats.cpp +++ b/ydb/library/yql/dq/actors/compute/dq_compute_actor_stats.cpp @@ -13,6 +13,10 @@ void FillAsyncStats(NDqProto::TDqAsyncBufferStats& proto, TDqAsyncStats stats) { proto.SetRows(stats.Rows); proto.SetChunks(stats.Chunks); proto.SetSplits(stats.Splits); + proto.SetFilteredBytes(stats.FilteredBytes); + proto.SetFilteredRows(stats.FilteredRows); + proto.SetQueuedBytes(stats.QueuedBytes); + proto.SetQueuedRows(stats.QueuedRows); if (stats.CollectFull()) { proto.SetFirstMessageMs(stats.FirstMessageTs.MilliSeconds()); proto.SetPauseMessageMs(stats.PauseMessageTs.MilliSeconds()); diff --git a/ydb/library/yql/dq/actors/input_transforms/dq_input_transform_lookup.cpp b/ydb/library/yql/dq/actors/input_transforms/dq_input_transform_lookup.cpp index a3b683f54157..73d7fec1050a 100644 --- a/ydb/library/yql/dq/actors/input_transforms/dq_input_transform_lookup.cpp +++ b/ydb/library/yql/dq/actors/input_transforms/dq_input_transform_lookup.cpp @@ -22,9 +22,10 @@ using TOutputRowColumnOrder = std::vector> //Design note: Base implementation is optimized for wide channels class TInputTransformStreamLookupBase - : public NActors::TActorBootstrapped + : public NActors::TActor , public NYql::NDq::IDqComputeActorAsyncInput { + using TActor = NActors::TActor; public: TInputTransformStreamLookupBase( std::shared_ptr alloc, @@ -47,7 +48,8 @@ class TInputTransformStreamLookupBase size_t cacheLimit, std::chrono::seconds cacheTtl ) - : Alloc(alloc) + : TActor(&TInputTransformStreamLookupBase::StateFunc) + , Alloc(alloc) , HolderFactory(holderFactory) , TypeEnv(typeEnv) , InputIndex(inputIndex) @@ -68,7 +70,10 @@ class TInputTransformStreamLookupBase , LruCache(std::make_unique(cacheLimit, lookupKeyType)) , MaxDelayedRows(maxDelayedRows) , CacheTtl(cacheTtl) + , MinimumRowSize(OutputRowColumnOrder.size()*sizeof(NUdf::TUnboxedValuePod)) + , PayloadExtraSize(0) , ReadyQueue(OutputRowType) + , LastLruSize(0) { Y_ABORT_UNLESS(Alloc); for (size_t i = 0; i != LookupInputIndexes.size(); ++i) { @@ -85,26 +90,6 @@ class TInputTransformStreamLookupBase Free(); } - void Bootstrap() { - Become(&TInputTransformStreamLookupBase::StateFunc); - NDq::IDqAsyncIoFactory::TLookupSourceArguments lookupSourceArgs { - .Alloc = Alloc, - .KeyTypeHelper = KeyTypeHelper, - .ParentId = SelfId(), - .TaskCounters = TaskCounters, - .LookupSource = Settings.GetRightSource().GetLookupSource(), - .KeyType = LookupKeyType, - .PayloadType = LookupPayloadType, - .TypeEnv = TypeEnv, - .HolderFactory = HolderFactory, - .MaxKeysInRequest = 1000 // TODO configure me - }; - auto guard = Guard(*Alloc); - auto [lookupSource, lookupSourceActor] = Factory->CreateDqLookupSource(Settings.GetRightSource().GetProviderName(), std::move(lookupSourceArgs)); - MaxKeysInRequest = lookupSource->GetMaxSupportedKeysInRequest(); - LookupSourceId = RegisterWithSameMailbox(lookupSourceActor); - KeysForLookup = std::make_shared(MaxKeysInRequest, KeyTypeHelper->GetValueHash(), KeyTypeHelper->GetValueEqual()); - } protected: virtual NUdf::EFetchStatus FetchWideInputValue(NUdf::TUnboxedValue* inputRowItems) = 0; virtual void PushOutputValue(NKikimr::NMiniKQL::TUnboxedValueBatch& batch, NUdf::TUnboxedValue* outputRowItems) = 0; @@ -147,6 +132,9 @@ class TInputTransformStreamLookupBase Y_ABORT(); break; } + if (outputRowItems[i].IsString()) { + PayloadExtraSize += outputRowItems[i].AsStringRef().size(); + } } ReadyQueue.PushRow(outputRowItems, OutputRowType->GetElementsCount()); } @@ -172,11 +160,14 @@ class TInputTransformStreamLookupBase LruCache->Update(NUdf::TUnboxedValue(const_cast(k)), std::move(v), now + CacheTtl); } KeysForLookup->clear(); + auto deltaLruSize = (i64)LruCache->Size() - LastLruSize; auto deltaTime = GetCpuTimeDelta(startCycleCount); CpuTime += deltaTime; if (CpuTimeUs) { + LruSize->Add(deltaLruSize); // Note: there can be several streamlookup tied to same counter, so Add instead of Set CpuTimeUs->Add(deltaTime.MicroSeconds()); } + LastLruSize += deltaLruSize; Send(ComputeActorId, new TEvNewAsyncInputDataArrived{InputIndex}); } @@ -195,11 +186,16 @@ class TInputTransformStreamLookupBase } void PassAway() final { + InputFlowFetchStatus = NUdf::EFetchStatus::Finish; Send(LookupSourceId, new NActors::TEvents::TEvPoison{}); Free(); } void Free() { + if (LruSize && LastLruSize) { + LruSize->Add(-LastLruSize); + LastLruSize = 0; + } auto guard = BindAllocator(); //All resources, held by this class, that have been created with mkql allocator, must be deallocated here KeysForLookup.reset(); @@ -217,6 +213,30 @@ class TInputTransformStreamLookupBase } } + std::shared_ptr GetKeysForLookup() { // must be called with mkql allocator + if (!KeysForLookup) { + Y_ENSURE(SelfId()); + Y_ENSURE(!LookupSourceId); + NDq::IDqAsyncIoFactory::TLookupSourceArguments lookupSourceArgs { + .Alloc = Alloc, + .KeyTypeHelper = KeyTypeHelper, + .ParentId = SelfId(), + .TaskCounters = TaskCounters, + .LookupSource = Settings.GetRightSource().GetLookupSource(), + .KeyType = LookupKeyType, + .PayloadType = LookupPayloadType, + .TypeEnv = TypeEnv, + .HolderFactory = HolderFactory, + .MaxKeysInRequest = 1000 // TODO configure me + }; + auto [lookupSource, lookupSourceActor] = Factory->CreateDqLookupSource(Settings.GetRightSource().GetProviderName(), std::move(lookupSourceArgs)); + MaxKeysInRequest = lookupSource->GetMaxSupportedKeysInRequest(); + LookupSourceId = RegisterWithSameMailbox(lookupSourceActor); + KeysForLookup = std::make_shared(MaxKeysInRequest, KeyTypeHelper->GetValueHash(), KeyTypeHelper->GetValueEqual()); + } + return KeysForLookup; + } + i64 GetAsyncInputData(NKikimr::NMiniKQL::TUnboxedValueBatch& batch, TMaybe&, bool& finished, i64 freeSpace) final { Y_UNUSED(freeSpace); auto startCycleCount = GetCycleCountFast(); @@ -224,7 +244,7 @@ class TInputTransformStreamLookupBase DrainReadyQueue(batch); - if (InputFlowFetchStatus != NUdf::EFetchStatus::Finish && KeysForLookup->empty()) { + if (InputFlowFetchStatus != NUdf::EFetchStatus::Finish && GetKeysForLookup()->empty()) { Y_DEBUG_ABORT_UNLESS(AwaitingQueue.empty()); NUdf::TUnboxedValue* inputRowItems; NUdf::TUnboxedValue inputRow = HolderFactory.CreateDirectArrayHolder(InputRowType->GetElementsCount(), inputRowItems); @@ -280,7 +300,18 @@ class TInputTransformStreamLookupBase CpuTimeUs->Add(deltaTime.MicroSeconds()); } finished = IsFinished(); - return AwaitingQueue.size(); + if (batch.empty()) { + // Use non-zero value to signal presence of pending request + // (value is NOT used for used space accounting) + return !AwaitingQueue.empty(); + } else { + // Attempt to estimate actual byte size; + // May be over-estimated for shared strings; + // May be under-estimated for complex types; + auto usedSpace = batch.RowCount() * MinimumRowSize + PayloadExtraSize; + PayloadExtraSize = 0; + return usedSpace; + } } void InitMonCounters(const ::NMonitoring::TDynamicCounterPtr& taskCounters) { @@ -290,6 +321,7 @@ class TInputTransformStreamLookupBase auto component = taskCounters->GetSubgroup("component", "Lookup"); LruHits = component->GetCounter("Hits"); LruMiss = component->GetCounter("Miss"); + LruSize = component->GetCounter("Size"); CpuTimeUs = component->GetCounter("CpuUs"); Batches = component->GetCounter("Batches"); } @@ -355,12 +387,16 @@ class TInputTransformStreamLookupBase using TInputKeyOtherPair = std::pair; using TAwaitingQueue = std::deque>; //input row split in two parts: key columns and other columns TAwaitingQueue AwaitingQueue; + size_t MinimumRowSize; // only account for unboxed parts + size_t PayloadExtraSize; // non-embedded part of strings in ReadyQueue NKikimr::NMiniKQL::TUnboxedValueBatch ReadyQueue; NYql::NDq::TDqAsyncStats IngressStats; std::shared_ptr KeysForLookup; + i64 LastLruSize; ::NMonitoring::TDynamicCounters::TCounterPtr LruHits; ::NMonitoring::TDynamicCounters::TCounterPtr LruMiss; + ::NMonitoring::TDynamicCounters::TCounterPtr LruSize; ::NMonitoring::TDynamicCounters::TCounterPtr CpuTimeUs; ::NMonitoring::TDynamicCounters::TCounterPtr Batches; TDuration CpuTime; @@ -504,6 +540,16 @@ std::pair< } else { result[i] = { EOutputRowItemSource::LookupOther, lookupPayloadColumns.at(name) }; } + } else if (leftLabel.empty()) { + const auto name = prefixedName; + if (auto j = leftJoinColumns.FindPtr(name)) { + result[i] = { EOutputRowItemSource::InputKey, lookupKeyColumns.at(rightNames[*j]) }; + } else if (auto k = inputColumns.FindPtr(name)) { + result[i] = { EOutputRowItemSource::InputOther, otherInputIndexes.size() }; + otherInputIndexes.push_back(*k); + } else { + Y_ABORT(); + } } else { Y_ABORT(); } diff --git a/ydb/library/yql/dq/actors/protos/dq_stats.proto b/ydb/library/yql/dq/actors/protos/dq_stats.proto index 733bf969dc2c..e261c967e4ba 100644 --- a/ydb/library/yql/dq/actors/protos/dq_stats.proto +++ b/ydb/library/yql/dq/actors/protos/dq_stats.proto @@ -30,6 +30,11 @@ message TDqAsyncBufferStats { uint64 LastMessageMs = 8; // last message processed uint64 WaitTimeUs = 9; // SUM(Resume_i - Pause_i) in us uint64 WaitPeriods = 10; // COUNT(Resume_i - Pause_i) + + uint64 FilteredBytes = 12; + uint64 FilteredRows = 13; + uint64 QueuedBytes = 14; + uint64 QueuedRows = 15; } message TDqAsyncInputBufferStats { @@ -187,6 +192,10 @@ message TDqTaskStats { uint64 IngressRows = 17; uint64 EgressBytes = 18; uint64 EgressRows = 19; + uint64 IngressFilteredBytes = 23; + uint64 IngressFilteredRows = 24; + uint64 IngressQueuedBytes = 25; + uint64 IngressQueuedRows = 26; // full stats repeated TDqAsyncInputBufferStats Sources = 150; @@ -286,6 +295,11 @@ message TDqAsyncStatsAggr { TDqStatsAggr WaitTimeUs = 9; TDqStatsAggr WaitPeriods = 10; TDqStatsAggr ActiveTimeUs = 11; + + TDqStatsAggr FilteredBytes = 13; + TDqStatsAggr FilteredRows = 14; + TDqStatsAggr QueuedBytes = 15; + TDqStatsAggr QueuedRows = 16; } message TDqAsyncBufferStatsAggr { @@ -331,6 +345,10 @@ message TDqStageStats { TDqStatsAggr IngressBytes = 28; TDqStatsAggr IngressDecompressedBytes = 37; + TDqStatsAggr IngressFilteredBytes = 46; + TDqStatsAggr IngressFilteredRows = 47; + TDqStatsAggr IngressQueuedBytes = 48; + TDqStatsAggr IngressQueuedRows = 49; TDqStatsAggr IngressRows = 29; TDqStatsAggr EgressBytes = 30; TDqStatsAggr EgressRows = 31; diff --git a/ydb/library/yql/dq/actors/task_runner/task_runner_actor_local.cpp b/ydb/library/yql/dq/actors/task_runner/task_runner_actor_local.cpp index 32f411181672..85204a16ba0f 100644 --- a/ydb/library/yql/dq/actors/task_runner/task_runner_actor_local.cpp +++ b/ydb/library/yql/dq/actors/task_runner/task_runner_actor_local.cpp @@ -83,6 +83,7 @@ class TLocalTaskRunnerActor /*flags=*/0, ev->Cookie); } + ActorElapsedTicks += NActors::TlsActivationContext->GetCurrentEventTicks(); } private: @@ -99,7 +100,7 @@ class TLocalTaskRunnerActor inputTransforms[inputTransformId] = TaskRunner->GetInputTransform(inputTransformId)->second.Get(); } - ev->Get()->Stats = TDqTaskRunnerStatsView(TaskRunner->GetStats(), std::move(sinks), std::move(inputTransforms)); + ev->Get()->Stats = TDqTaskRunnerStatsView(TaskRunner->GetStats(), std::move(sinks), std::move(inputTransforms), ActorElapsedTicks); Send( ParentId, ev->Release().Release(), @@ -247,7 +248,7 @@ class TLocalTaskRunnerActor inputTransforms[inputTransformId] = TaskRunner->GetInputTransform(inputTransformId)->second.Get(); } - st->Stats = TDqTaskRunnerStatsView(TaskRunner->GetStats(), std::move(sinks), std::move(inputTransforms)); + st->Stats = TDqTaskRunnerStatsView(TaskRunner->GetStats(), std::move(sinks), std::move(inputTransforms), ActorElapsedTicks); Send(ParentId, st.Release()); } @@ -508,6 +509,7 @@ class TLocalTaskRunnerActor TIntrusivePtr TaskRunner; THashSet InputChannelsWithDisabledCheckpoints; THolder MemoryQuota; + ui64 ActorElapsedTicks = 0; }; struct TLocalTaskRunnerActorFactory: public ITaskRunnerActorFactory { diff --git a/ydb/library/yql/dq/expr_nodes/dq_expr_nodes.json b/ydb/library/yql/dq/expr_nodes/dq_expr_nodes.json index e9efb8c20423..7a5766428f7d 100644 --- a/ydb/library/yql/dq/expr_nodes/dq_expr_nodes.json +++ b/ydb/library/yql/dq/expr_nodes/dq_expr_nodes.json @@ -44,7 +44,8 @@ "Match": {"Type": "Callable", "Name": "DqJoin"}, "Children": [ {"Index": 8, "Name": "JoinAlgo", "Type": "TCoAtom"}, - {"Index": 9, "Name": "Flags", "Type": "TCoAtomList", "Optional": true} + {"Index": 9, "Name": "Flags", "Type": "TCoAtomList", "Optional": true}, + {"Index": 10, "Name": "JoinAlgoOptions", "Type": "TCoNameValueTupleList", "Optional": true} ] }, { diff --git a/ydb/library/yql/dq/opt/dq_opt_join.cpp b/ydb/library/yql/dq/opt/dq_opt_join.cpp index 9ffc9d365689..c7f6bd7b3126 100644 --- a/ydb/library/yql/dq/opt/dq_opt_join.cpp +++ b/ydb/library/yql/dq/opt/dq_opt_join.cpp @@ -242,7 +242,7 @@ TMaybe BuildDqJoin( rightJoinKeyNames.emplace_back(rightColumnName); } - if (EHashJoinMode::Off == mode || EHashJoinMode::Map == mode || !(leftAny || rightAny)) { + if ((linkSettings.JoinAlgo != EJoinAlgoType::StreamLookupJoin && (EHashJoinMode::Off == mode || EHashJoinMode::Map == mode)) || !(leftAny || rightAny || !linkSettings.JoinAlgoOptions.empty())) { auto dqJoin = Build(ctx, joinTuple.Pos()) .LeftInput(BuildDqJoinInput(ctx, joinTuple.Pos(), left->Input, leftJoinKeys, leftAny)) .LeftLabel(leftTableLabel) @@ -266,6 +266,15 @@ TMaybe BuildDqJoin( if (rightAny) flags.emplace_back(ctx.NewAtom(joinTuple.Pos(), "RightAny", TNodeFlags::Default)); + TVector joinAlgoOptions; + for (ui32 i = 0; i + 1 < linkSettings.JoinAlgoOptions.size(); i += 2) { + joinAlgoOptions.push_back( + Build(ctx, joinTuple.Pos()) + .Name().Build(linkSettings.JoinAlgoOptions[i]) + .Value().Build(linkSettings.JoinAlgoOptions[i + 1]) + .Done()); + } + auto dqJoin = Build(ctx, joinTuple.Pos()) .LeftInput(BuildDqJoinInput(ctx, joinTuple.Pos(), left->Input, leftJoinKeys, false)) .LeftLabel(leftTableLabel) @@ -280,9 +289,11 @@ TMaybe BuildDqJoin( .Add(rightJoinKeyNames) .Build() .JoinAlgo(joinAlgo) - .Flags().Add(std::move(flags)).Build() - .Done(); - return TJoinInputDesc(Nothing(), dqJoin, std::move(resultKeys)); + .Flags().Add(std::move(flags)).Build(); + if (!joinAlgoOptions.empty()) { + dqJoin.JoinAlgoOptions().Add(std::move(joinAlgoOptions)).Build(); + } + return TJoinInputDesc(Nothing(), dqJoin.Done(), std::move(resultKeys)); } } diff --git a/ydb/library/yql/dq/runtime/dq_async_stats.h b/ydb/library/yql/dq/runtime/dq_async_stats.h index 6fb51809b840..d1e7b7111161 100644 --- a/ydb/library/yql/dq/runtime/dq_async_stats.h +++ b/ydb/library/yql/dq/runtime/dq_async_stats.h @@ -68,12 +68,21 @@ struct TDqAsyncStats { TDuration WaitTime; ui64 WaitPeriods = 0; + ui64 FilteredBytes = 0; + ui64 FilteredRows = 0; + ui64 QueuedBytes = 0; + ui64 QueuedRows = 0; + void MergeData(const TDqAsyncStats& other) { Bytes += other.Bytes; DecompressedBytes += other.DecompressedBytes; Rows += other.Rows; Chunks += other.Chunks; Splits += other.Splits; + FilteredBytes += other.FilteredBytes; + FilteredRows += other.FilteredRows; + QueuedBytes += other.QueuedBytes; + QueuedRows += other.QueuedRows; } void MergeTime(const TDqAsyncStats& other) { diff --git a/ydb/library/yql/dq/runtime/dq_tasks_runner.h b/ydb/library/yql/dq/runtime/dq_tasks_runner.h index 636e8600346f..4910e05d889d 100644 --- a/ydb/library/yql/dq/runtime/dq_tasks_runner.h +++ b/ydb/library/yql/dq/runtime/dq_tasks_runner.h @@ -89,11 +89,12 @@ class TDqTaskRunnerStatsView { } TDqTaskRunnerStatsView(const TDqTaskRunnerStats* stats, THashMap&& sinks, - THashMap&& inputTransforms) + THashMap&& inputTransforms, ui64 actorElapsedTicks) : StatsPtr(stats) , IsDefined(true) , Sinks(std::move(sinks)) - , InputTransforms(std::move(inputTransforms)) { + , InputTransforms(std::move(inputTransforms)) + , ActorElapsedTicks(actorElapsedTicks) { } const TTaskRunnerStatsBase* Get() { @@ -115,11 +116,16 @@ class TDqTaskRunnerStatsView { return InputTransforms.at(inputTransformId); } + ui64 GetActorElapsedTicks() { + return ActorElapsedTicks; + } + private: const TDqTaskRunnerStats* StatsPtr; bool IsDefined; THashMap Sinks; THashMap InputTransforms; + ui64 ActorElapsedTicks = 0; }; struct TDqTaskRunnerContext { diff --git a/ydb/library/yql/dq/type_ann/dq_type_ann.cpp b/ydb/library/yql/dq/type_ann/dq_type_ann.cpp index 90b36e1f8e3d..63b745a8ae27 100644 --- a/ydb/library/yql/dq/type_ann/dq_type_ann.cpp +++ b/ydb/library/yql/dq/type_ann/dq_type_ann.cpp @@ -1,5 +1,4 @@ #include "dq_type_ann.h" - #include #include #include @@ -85,6 +84,16 @@ const TTypeAnnotationNode* GetColumnType(const TDqConnection& node, const TStruc return result; } +template +bool EnsureConvertibleTo(const TExprNode& value, const TStringBuf name, TExprContext& ctx) { + auto&& stringValue = value.Content(); + if (!TryFromString(stringValue)) { + ctx.AddError(TIssue(ctx.GetPosition(value.Pos()), TStringBuilder() << "Unsupported " << name << " value: " << stringValue)); + return false; + } + return true; +} + template TStatus AnnotateStage(const TExprNode::TPtr& stage, TExprContext& ctx) { if (!EnsureMinMaxArgsCount(*stage, 3, 4, ctx)) { @@ -428,7 +437,7 @@ const TStructExprType* GetDqJoinResultType(TPositionHandle pos, const TStructExp template const TStructExprType* GetDqJoinResultType(const TExprNode::TPtr& input, bool stream, TExprContext& ctx) { - if (!EnsureMinMaxArgsCount(*input, 8, 10, ctx)) { + if (!EnsureMinMaxArgsCount(*input, 8, 11, ctx)) { return nullptr; } @@ -444,7 +453,8 @@ const TStructExprType* GetDqJoinResultType(const TExprNode::TPtr& input, bool st } } - if (!EnsureAtom(*input->Child(TDqJoin::idx_JoinType), ctx)) { + const auto& joinType = *input->Child(TDqJoin::idx_JoinType); + if (!EnsureAtom(joinType, ctx)) { return nullptr; } @@ -502,9 +512,39 @@ const TStructExprType* GetDqJoinResultType(const TExprNode::TPtr& input, bool st ? join.RightLabel().Cast().Value() : TStringBuf(""); - if (input->ChildrenSize() > 9U) { - for (auto i = 0U; i < input->Tail().ChildrenSize(); ++i) { - if (const auto& flag = *input->Tail().Child(i); !flag.IsAtom({"LeftAny", "RightAny"})) { + if (input->ChildrenSize() > TDqJoin::idx_JoinAlgoOptions) { + const auto& joinAlgo = *input->Child(TDqJoin::idx_JoinAlgo); + if (!EnsureAtom(joinAlgo, ctx)) { + return nullptr; + } + auto& joinAlgoOptions = *input->Child(TDqJoin::idx_JoinAlgoOptions); + for (ui32 i = 0; i < joinAlgoOptions.ChildrenSize(); ++i) { + auto& joinAlgoOption = *joinAlgoOptions.Child(i); + if (!EnsureTupleOfAtoms(joinAlgoOption, ctx) || !EnsureTupleMinSize(joinAlgoOption, 1, ctx)) { + return nullptr; + } + auto& name = *joinAlgoOption.Child(TCoNameValueTuple::idx_Name); + if (joinAlgo.IsAtom("StreamLookupJoin")) { + if (name.IsAtom({"TTL", "MaxCachedRows", "MaxDelayedRows"})) { + if (!EnsureTupleSize(joinAlgoOption, 2, ctx)) { + return nullptr; + } + auto& value = *joinAlgoOption.Child(TCoNameValueTuple::idx_Value); + if (!EnsureConvertibleTo(value, name.Content(), ctx)) { + return nullptr; + } + continue; + } + } + ctx.AddError(TIssue(ctx.GetPosition(joinAlgoOption.Pos()), TStringBuilder() << "DqJoin: Unsupported DQ join option: " << name.Content())); + return nullptr; + } + } + + if (input->ChildrenSize() > TDqJoin::idx_Flags) { + auto& flags = *input->Child(TDqJoin::idx_Flags); + for (auto i = 0U; i < flags.ChildrenSize(); ++i) { + if (const auto& flag = *flags.Child(i); !flag.IsAtom({"LeftAny", "RightAny"})) { ctx.AddError(TIssue(ctx.GetPosition(flag.Pos()), TStringBuilder() << "Unsupported DQ join option: " << flag.Content())); return nullptr; } @@ -576,24 +616,97 @@ TStatus AnnotateDqConnection(const TExprNode::TPtr& input, TExprContext& ctx) { } TStatus AnnotateDqCnStreamLookup(const TExprNode::TPtr& input, TExprContext& ctx) { + if (!EnsureArgsCount(*input, 11, ctx)) { + return TStatus::Error; + } + if (!EnsureCallable(*input->Child(TDqCnStreamLookup::idx_Output), ctx)) { + return TStatus::Error; + } + if (!TDqOutput::Match(input->Child(TDqCnStreamLookup::idx_Output))) { + ctx.AddError(TIssue(ctx.GetPosition(input->Child(TDqCnStreamLookup::idx_Output)->Pos()), TStringBuilder() << "Expected " << TDqOutput::CallableName())); + return TStatus::Error; + } + if (!EnsureAtom(*input->Child(TDqCnStreamLookup::idx_LeftLabel), ctx)) { + return TStatus::Error; + } + if (!EnsureCallable(*input->Child(TDqCnStreamLookup::idx_RightInput), ctx)) { + return TStatus::Error; + } + if (!EnsureAtom(*input->Child(TDqCnStreamLookup::idx_RightLabel), ctx)) { + return TStatus::Error; + } + if (!EnsureAtom(*input->Child(TDqCnStreamLookup::idx_JoinType), ctx)) { + return TStatus::Error; + } + if (!EnsureTuple(*input->Child(TDqCnStreamLookup::idx_JoinKeys), ctx)) { + return TStatus::Error; + } + for (auto& child: input->Child(TDqCnStreamLookup::idx_JoinKeys)->Children()) { + if (!EnsureTupleSize(*child, 4, ctx)) { + return TStatus::Error; + } + for (auto& subChild: child->Children()) { + if (!EnsureAtom(*subChild, ctx)) { + return TStatus::Error; + } + } + } + if (!EnsureTupleOfAtoms(*input->Child(TDqCnStreamLookup::idx_LeftJoinKeyNames), ctx)) { + return TStatus::Error; + } + if (!EnsureTupleOfAtoms(*input->Child(TDqCnStreamLookup::idx_RightJoinKeyNames), ctx)) { + return TStatus::Error; + } + if (!EnsureAtom(*input->Child(TDqCnStreamLookup::idx_TTL), ctx)) { + return TStatus::Error; + } + if (!EnsureAtom(*input->Child(TDqCnStreamLookup::idx_MaxDelayedRows), ctx)) { + return TStatus::Error; + } + if (!EnsureAtom(*input->Child(TDqCnStreamLookup::idx_MaxCachedRows), ctx)) { + return TStatus::Error; + } auto cnStreamLookup = TDqCnStreamLookup(input); auto leftInputType = GetDqConnectionType(TDqConnection(input), ctx); if (!leftInputType) { return TStatus::Error; } - const auto leftRowType = GetSeqItemType(leftInputType); - const auto rightRowType = GetSeqItemType(cnStreamLookup.RightInput().Raw()->GetTypeAnn()); + if (auto joinType = cnStreamLookup.JoinType(); joinType != TStringBuf("Left")) { + ctx.AddError(TIssue(ctx.GetPosition(joinType.Pos()), "Streamlookup supports only LEFT JOIN ... ANY")); + return TStatus::Error; + } + auto rightInput = cnStreamLookup.RightInput(); + if (!rightInput.Raw()->IsCallable("TDqLookupSourceWrap")) { + ctx.AddError(TIssue(ctx.GetPosition(rightInput.Pos()), TStringBuilder() << "DqCnStreamLookup: RightInput: Expected TDqLookupSourceWrap, but got " << rightInput.Raw()->Content())); + return TStatus::Error; + } + const auto& leftRowType = GetSeqItemType(*leftInputType); + if (!EnsureStructType(input->Pos(), leftRowType, ctx)) { + return TStatus::Error; + } + const auto rightInputType = rightInput.Raw()->GetTypeAnn(); + const auto& rightRowType = GetSeqItemType(*rightInputType); + if (!EnsureStructType(input->Pos(), rightRowType, ctx)) { + return TStatus::Error; + } const auto outputRowType = GetDqJoinResultType( input->Pos(), - *leftRowType->Cast(), + *leftRowType.Cast(), cnStreamLookup.LeftLabel().Cast().StringValue(), - *rightRowType->Cast(), + *rightRowType.Cast(), cnStreamLookup.RightLabel().StringValue(), cnStreamLookup.JoinType().StringValue(), cnStreamLookup.JoinKeys(), ctx ); - //TODO (YQ-2068) verify lookup parameters + if (!outputRowType) { + return TStatus::Error; + } + if (!EnsureConvertibleTo(cnStreamLookup.MaxCachedRows().Ref(), "MaxCachedRows", ctx) || + !EnsureConvertibleTo(cnStreamLookup.TTL().Ref(), "TTL", ctx) || + !EnsureConvertibleTo(cnStreamLookup.MaxDelayedRows().Ref(), "MaxDelayedRows", ctx)) { + return TStatus::Error; + } input->SetTypeAnn(ctx.MakeType(outputRowType)); return TStatus::Ok; } @@ -1212,16 +1325,13 @@ bool TDqStageSettings::Validate(const TExprNode& stage, TExprContext& ctx) { return false; } - if (name == LogicalIdSettingName && !TryFromString(value->Content())) { - ctx.AddError(TIssue(ctx.GetPosition(setting->Pos()), TStringBuilder() << "Setting " << name << " should contain ui64 value, but got: " << value->Content())); + if (name == LogicalIdSettingName && !EnsureConvertibleTo(*value, name, ctx)) { return false; } - if (name == BlockStatusSettingName && !TryFromString(value->Content())) { - ctx.AddError(TIssue(ctx.GetPosition(setting->Pos()), TStringBuilder() << "Unsupported " << name << " value: " << value->Content())); + if (name == BlockStatusSettingName && !EnsureConvertibleTo(*value, name, ctx)) { return false; } - if (name == PartitionModeSettingName && !TryFromString(value->Content())) { - ctx.AddError(TIssue(ctx.GetPosition(setting->Pos()), TStringBuilder() << "Unsupported " << name << " value: " << value->Content())); + if (name == PartitionModeSettingName && !EnsureConvertibleTo(*value, name, ctx)) { return false; } } else if (name == WideChannelsSettingName) { diff --git a/ydb/library/yql/providers/common/db_id_async_resolver/db_async_resolver.h b/ydb/library/yql/providers/common/db_id_async_resolver/db_async_resolver.h index 94237dd434d2..06ce03a48e9e 100644 --- a/ydb/library/yql/providers/common/db_id_async_resolver/db_async_resolver.h +++ b/ydb/library/yql/providers/common/db_id_async_resolver/db_async_resolver.h @@ -3,7 +3,7 @@ #include #include #include -#include +#include #include namespace NYql { @@ -22,47 +22,47 @@ enum class EDatabaseType { Logging }; -inline EDatabaseType DatabaseTypeFromDataSourceKind(NConnector::NApi::EDataSourceKind dataSourceKind) { +inline EDatabaseType DatabaseTypeFromDataSourceKind(NYql::EGenericDataSourceKind dataSourceKind) { switch (dataSourceKind) { - case NConnector::NApi::EDataSourceKind::POSTGRESQL: + case NYql::EGenericDataSourceKind::POSTGRESQL: return EDatabaseType::PostgreSQL; - case NConnector::NApi::EDataSourceKind::CLICKHOUSE: + case NYql::EGenericDataSourceKind::CLICKHOUSE: return EDatabaseType::ClickHouse; - case NConnector::NApi::EDataSourceKind::YDB: + case NYql::EGenericDataSourceKind::YDB: return EDatabaseType::Ydb; - case NConnector::NApi::EDataSourceKind::MYSQL: + case NYql::EGenericDataSourceKind::MYSQL: return EDatabaseType::MySQL; - case NConnector::NApi::EDataSourceKind::GREENPLUM: + case NYql::EGenericDataSourceKind::GREENPLUM: return EDatabaseType::Greenplum; - case NConnector::NApi::EDataSourceKind::MS_SQL_SERVER: + case NYql::EGenericDataSourceKind::MS_SQL_SERVER: return EDatabaseType::MsSQLServer; - case NConnector::NApi::EDataSourceKind::ORACLE: + case NYql::EGenericDataSourceKind::ORACLE: return EDatabaseType::Oracle; - case NConnector::NApi::EDataSourceKind::LOGGING: + case NYql::EGenericDataSourceKind::LOGGING: return EDatabaseType::Logging; default: - ythrow yexception() << "Unknown data source kind: " << NConnector::NApi::EDataSourceKind_Name(dataSourceKind); + ythrow yexception() << "Unknown data source kind: " << NYql::EGenericDataSourceKind_Name(dataSourceKind); } } -inline NConnector::NApi::EDataSourceKind DatabaseTypeToDataSourceKind(EDatabaseType databaseType) { +inline NYql::EGenericDataSourceKind DatabaseTypeToDataSourceKind(EDatabaseType databaseType) { switch (databaseType) { case EDatabaseType::PostgreSQL: - return NConnector::NApi::EDataSourceKind::POSTGRESQL; + return NYql::EGenericDataSourceKind::POSTGRESQL; case EDatabaseType::ClickHouse: - return NConnector::NApi::EDataSourceKind::CLICKHOUSE; + return NYql::EGenericDataSourceKind::CLICKHOUSE; case EDatabaseType::Ydb: - return NConnector::NApi::EDataSourceKind::YDB; + return NYql::EGenericDataSourceKind::YDB; case EDatabaseType::MySQL: - return NConnector::NApi::EDataSourceKind::MYSQL; + return NYql::EGenericDataSourceKind::MYSQL; case EDatabaseType::Greenplum: - return NConnector::NApi::EDataSourceKind::GREENPLUM; + return NYql::EGenericDataSourceKind::GREENPLUM; case EDatabaseType::MsSQLServer: - return NConnector::NApi::EDataSourceKind::MS_SQL_SERVER; + return NYql::EGenericDataSourceKind::MS_SQL_SERVER; case EDatabaseType::Oracle: - return NConnector::NApi::EDataSourceKind::ORACLE; + return NYql::EGenericDataSourceKind::ORACLE; case EDatabaseType::Logging: - return NConnector::NApi::EDataSourceKind::LOGGING; + return NYql::EGenericDataSourceKind::LOGGING; default: ythrow yexception() << "Unknown database type: " << ToString(databaseType); } @@ -93,7 +93,7 @@ struct TDatabaseAuth { // For some of the data sources accessible via generic provider it's possible to specify the connection protocol. // This setting may impact the throughput. - NConnector::NApi::EProtocol Protocol = NConnector::NApi::EProtocol::PROTOCOL_UNSPECIFIED; + NYql::EGenericProtocol Protocol = NYql::EGenericProtocol::PROTOCOL_UNSPECIFIED; bool operator==(const TDatabaseAuth& other) const { return std::tie(StructuredToken, AddBearerToToken, UseTls, Protocol) == std::tie(other.StructuredToken, other.AddBearerToToken, other.UseTls, Protocol); diff --git a/ydb/library/yql/providers/common/db_id_async_resolver/mdb_endpoint_generator.h b/ydb/library/yql/providers/common/db_id_async_resolver/mdb_endpoint_generator.h index abf38158aa77..6eabe6df83d3 100644 --- a/ydb/library/yql/providers/common/db_id_async_resolver/mdb_endpoint_generator.h +++ b/ydb/library/yql/providers/common/db_id_async_resolver/mdb_endpoint_generator.h @@ -13,7 +13,7 @@ namespace NYql { NYql::EDatabaseType DatabaseType; TString MdbHost; bool UseTls; - NConnector::NApi::EProtocol Protocol; + NYql::EGenericProtocol Protocol; }; using TPtr = std::shared_ptr; diff --git a/ydb/library/yql/providers/common/db_id_async_resolver/ya.make b/ydb/library/yql/providers/common/db_id_async_resolver/ya.make index e2a48a290cbd..48142fe1e922 100644 --- a/ydb/library/yql/providers/common/db_id_async_resolver/ya.make +++ b/ydb/library/yql/providers/common/db_id_async_resolver/ya.make @@ -7,7 +7,7 @@ SRCS( PEERDIR( library/cpp/threading/future - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto yql/essentials/public/issue ) diff --git a/ydb/library/yql/providers/common/http_gateway/yql_aws_signature.cpp b/ydb/library/yql/providers/common/http_gateway/yql_aws_signature.cpp index 699d43f4f119..eea7e35507d4 100644 --- a/ydb/library/yql/providers/common/http_gateway/yql_aws_signature.cpp +++ b/ydb/library/yql/providers/common/http_gateway/yql_aws_signature.cpp @@ -138,11 +138,11 @@ TString TAwsSignature::HashSHA256(TStringBuf data) { return to_lower(HexEncode(hash, SHA256_DIGEST_LENGTH)); } -TString TAwsSignature::UriEncode(const TStringBuf input, bool encodeSlash) { +TString TAwsSignature::UriEncode(const TStringBuf input, bool encodeSlash, bool encodePercent) { TStringStream result; for (const char ch : input) { if ((ch >= 'A' && ch <= 'Z') || (ch >= 'a' && ch <= 'z') || (ch >= '0' && ch <= '9') || ch == '_' || - ch == '-' || ch == '~' || ch == '.') { + ch == '-' || ch == '~' || ch == '.' || (ch == '%' && !encodePercent)) { result << ch; } else if (ch == '/') { if (encodeSlash) { @@ -175,11 +175,10 @@ void TAwsSignature::PrepareCgiParameters() { auto printSingleParam = [&canonicalCgi](const TString& key, const TVector& values) { auto it = values.begin(); - canonicalCgi << UriEncode(key, true) << "=" << UriEncode(*it, true); + canonicalCgi << UriEncode(key, true, true) << "=" << UriEncode(*it, true, true); while (++it != values.end()) { - canonicalCgi << "&" << UriEncode(key, true) << "=" << UriEncode(*it, true); + canonicalCgi << "&" << UriEncode(key, true, true) << "=" << UriEncode(*it, true, true); } - }; auto it = sortedCgi.begin(); diff --git a/ydb/library/yql/providers/common/http_gateway/yql_aws_signature.h b/ydb/library/yql/providers/common/http_gateway/yql_aws_signature.h index 7a5032814b5b..469c03e10431 100644 --- a/ydb/library/yql/providers/common/http_gateway/yql_aws_signature.h +++ b/ydb/library/yql/providers/common/http_gateway/yql_aws_signature.h @@ -32,7 +32,7 @@ struct TAwsSignature { static TString HashSHA256(TStringBuf data); - static TString UriEncode(const TStringBuf input, bool encodeSlash = false); + static TString UriEncode(const TStringBuf input, bool encodeSlash = false, bool encodePercent = false); void PrepareCgiParameters(); diff --git a/ydb/library/yql/providers/common/http_gateway/yql_aws_signature_ut.cpp b/ydb/library/yql/providers/common/http_gateway/yql_aws_signature_ut.cpp index b3f053bf216f..bf6e3e01a8bb 100644 --- a/ydb/library/yql/providers/common/http_gateway/yql_aws_signature_ut.cpp +++ b/ydb/library/yql/providers/common/http_gateway/yql_aws_signature_ut.cpp @@ -2,6 +2,7 @@ #include #include +#include namespace NYql { @@ -68,5 +69,15 @@ Y_UNIT_TEST_SUITE(TAwsSignature) { UNIT_ASSERT_VALUES_EQUAL(signature1.GetAmzDate(), signature2.GetAmzDate()); UNIT_ASSERT_VALUES_EQUAL(signature1.GetAuthorization(), signature2.GetAuthorization()); } + + Y_UNIT_TEST(SignWithEscaping) { + auto time = TInstant::FromValue(30); + NYql::TAwsSignature signature("GET", UrlEscapeRet("http://os.com/my-bucket/ !\"#$%&'()+,-./0123456789:;<=>@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz|~/", true), "application/json", {}, "key", "pwd", time); + UNIT_ASSERT_VALUES_EQUAL(signature.GetContentType(), "application/json"); + UNIT_ASSERT_VALUES_EQUAL(signature.GetXAmzContentSha256(), "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"); + UNIT_ASSERT_VALUES_EQUAL(signature.GetAuthorization(), "AWS4-HMAC-SHA256 Credential=/19700101///aws4_request, SignedHeaders=content-type;host;x-amz-content-sha256;x-amz-date, Signature=21470c8f999941fdc785c508f0c55afa1a12735eddd868aa7276e532d687c436"); + UNIT_ASSERT_VALUES_UNEQUAL(signature.GetAmzDate(), ""); + } } // Y_UNIT_TEST_SUITE(TAwsSignature) + } // namespace NYql diff --git a/ydb/library/yql/providers/dq/actors/task_controller_impl.h b/ydb/library/yql/providers/dq/actors/task_controller_impl.h index 698946386504..a0ecbad30fdf 100644 --- a/ydb/library/yql/providers/dq/actors/task_controller_impl.h +++ b/ydb/library/yql/providers/dq/actors/task_controller_impl.h @@ -295,6 +295,16 @@ class TTaskControllerImpl: public NActors::TActor { } else if (name == "EgressRows" && taskLevelCounter) { publicCounterName = "query.sink_output_records"; isDeriv = true; + } else if (name == "IngressFilteredBytes" && taskLevelCounter) { + publicCounterName = "query.input_filtered_bytes"; + isDeriv = true; + } else if (name == "IngressFilteredRows" && taskLevelCounter) { + publicCounterName = "query.source_input_filtered_records"; + isDeriv = true; + } else if (name == "IngressQueuedBytes" && taskLevelCounter) { + publicCounterName = "query.input_queued_bytes"; + } else if (name == "IngressQueuedRows" && taskLevelCounter) { + publicCounterName = "query.source_input_queued_records"; } else if (name == "Tasks") { publicCounterName = "query.running_tasks"; isDeriv = false; @@ -351,9 +361,12 @@ class TTaskControllerImpl: public NActors::TActor { ui64 taskId = s.GetTaskId(); ui64 stageId = Stages.Value(taskId, s.GetStageId()); +#define SET_COUNTER(name) \ + TaskStat.SetCounter(TaskStat.GetCounterName("TaskRunner", labels, #name), stats.Get ## name ()); \ + #define ADD_COUNTER(name) \ if (stats.Get ## name()) { \ - TaskStat.SetCounter(TaskStat.GetCounterName("TaskRunner", labels, #name), stats.Get ## name ()); \ + SET_COUNTER(name); \ } std::map commonLabels = { @@ -380,6 +393,11 @@ class TTaskControllerImpl: public NActors::TActor { ADD_COUNTER(ResultRows) ADD_COUNTER(ResultBytes) + ADD_COUNTER(IngressFilteredBytes) + ADD_COUNTER(IngressFilteredRows) + SET_COUNTER(IngressQueuedBytes) + SET_COUNTER(IngressQueuedRows) + ADD_COUNTER(StartTimeMs) ADD_COUNTER(FinishTimeMs) ADD_COUNTER(WaitInputTimeUs) diff --git a/ydb/library/yql/providers/dq/counters/task_counters.h b/ydb/library/yql/providers/dq/counters/task_counters.h index 65a393534338..3000787f0b93 100644 --- a/ydb/library/yql/providers/dq/counters/task_counters.h +++ b/ydb/library/yql/providers/dq/counters/task_counters.h @@ -33,6 +33,10 @@ struct TTaskCounters : public TCounters { (TInstant::MilliSeconds(lastMessageMs) - TInstant::MilliSeconds(firstMessageMs)).MicroSeconds() ); } + if (auto v = stats.GetFilteredBytes(); v) SetCounter(GetCounterName("TaskRunner", l, p + "FilteredBytes"), v); + if (auto v = stats.GetFilteredRows(); v) SetCounter(GetCounterName("TaskRunner", l, p + "FilteredRows"), v); + SetCounter(GetCounterName("TaskRunner", l, p + "QueuedBytes"), stats.GetQueuedBytes()); + SetCounter(GetCounterName("TaskRunner", l, p + "QueuedRows"), stats.GetQueuedRows()); } void AddAsyncStats(const NDq::TDqAsyncStats stats, const std::map& l, const TString& p) { @@ -53,6 +57,10 @@ struct TTaskCounters : public TCounters { if (activeTime) { SetCounter(GetCounterName("TaskRunner", l, p + "ActiveTimeUs"), activeTime.MicroSeconds()); } + if (stats.FilteredBytes) SetCounter(GetCounterName("TaskRunner", l, p + "FilteredBytes"), stats.FilteredBytes); + if (stats.FilteredRows) SetCounter(GetCounterName("TaskRunner", l, p + "FilteredRows"), stats.FilteredRows); + SetCounter(GetCounterName("TaskRunner", l, p + "QueuedBytes"), stats.QueuedBytes); + SetCounter(GetCounterName("TaskRunner", l, p + "QueuedRows"), stats.QueuedRows); } void AddInputChannelStats( diff --git a/ydb/library/yql/providers/dq/opt/logical_optimize.cpp b/ydb/library/yql/providers/dq/opt/logical_optimize.cpp index c9f72f182955..cf809f3ab6f9 100644 --- a/ydb/library/yql/providers/dq/opt/logical_optimize.cpp +++ b/ydb/library/yql/providers/dq/opt/logical_optimize.cpp @@ -225,36 +225,64 @@ class TDqsLogicalOptProposalTransformer : public TOptimizeTransformerBase { return TDqLookupSourceWrap(lookupSourceWrap); } - TMaybeNode RewriteStreamEquiJoinWithLookup(TExprBase node, TExprContext& ctx) { - Y_UNUSED(ctx); - const auto equiJoin = node.Cast(); - if (equiJoin.ArgCount() != 4) { // 2 parties join - return node; + // Recursively walk join tree and replace right-side of StreamLookupJoin + ui32 RewriteStreamJoinTuple(ui32 idx, const TCoEquiJoin& equiJoin, const TCoEquiJoinTuple& joinTuple, std::vector& args, TExprContext& ctx, bool& changed) { + // recursion depth O(args.size()) + Y_ENSURE(idx < args.size()); + // handle left side + if (!joinTuple.LeftScope().Maybe()) { + idx = RewriteStreamJoinTuple(idx, equiJoin, joinTuple.LeftScope().Cast(), args, ctx, changed); + } else { + ++idx; + } + // handle right side + if (!joinTuple.RightScope().Maybe()) { + return RewriteStreamJoinTuple(idx, equiJoin, joinTuple.RightScope().Cast(), args, ctx, changed); } - const auto left = equiJoin.Arg(0).Cast().List(); - const auto right = equiJoin.Arg(1).Cast().List(); - const auto joinTuple = equiJoin.Arg(equiJoin.ArgCount() - 2).Cast(); + Y_ENSURE(idx < args.size()); if (!IsStreamLookup(joinTuple)) { - return node; + return idx + 1; } - if (!right.Maybe() && !right.Maybe()) { - return node; + auto right = equiJoin.Arg(idx).Cast(); + auto rightList = right.List(); + if (auto maybeExtractMembers = rightList.Maybe()) { + rightList = maybeExtractMembers.Cast().Input(); + } + TExprNode::TPtr lookupSourceWrap; + if (auto maybeSource = rightList.Maybe()) { + lookupSourceWrap = LookupSourceFromSource(maybeSource.Cast(), ctx).Ptr(); + } else if (auto maybeRead = rightList.Maybe()) { + lookupSourceWrap = LookupSourceFromRead(maybeRead.Cast(), ctx).Ptr(); + } else { + return idx + 1; } + changed = true; + args[idx] = + Build(ctx, joinTuple.Pos()) + .List(lookupSourceWrap) + .Scope(right.Scope()) + .Done().Ptr(); + return idx + 1; + } - TDqLookupSourceWrap lookupSourceWrap = right.Maybe() - ? LookupSourceFromSource(right.Cast(), ctx) - : LookupSourceFromRead(right.Cast(), ctx) - ; - - return Build(ctx, node.Pos()) - .Add(equiJoin.Arg(0)) - .Add() - .List(lookupSourceWrap) - .Scope(equiJoin.Arg(1).Cast().Scope()) - .Build() - .Add(equiJoin.Arg(2)) - .Add(equiJoin.Arg(3)) - .Done(); + TMaybeNode RewriteStreamEquiJoinWithLookup(TExprBase node, TExprContext& ctx) { + const auto equiJoin = node.Cast(); + auto argCount = equiJoin.ArgCount(); + const auto joinTuple = equiJoin.Arg(argCount - 2).Cast(); + std::vector args(argCount); + bool changed = false; + auto rightIdx = RewriteStreamJoinTuple(0u, equiJoin, joinTuple, args, ctx, changed); + Y_ENSURE(rightIdx + 2 == argCount); + if (!changed) { + return node; + } + // fill copies of remaining args + for (ui32 i = 0; i < argCount; ++i) { + if (!args[i]) { + args[i] = equiJoin.Arg(i).Ptr(); + } + } + return Build(ctx, node.Pos()).Add(std::move(args)).Done(); } TMaybeNode OptimizeEquiJoinWithCosts(TExprBase node, TExprContext& ctx) { diff --git a/ydb/library/yql/providers/dq/opt/physical_optimize.cpp b/ydb/library/yql/providers/dq/opt/physical_optimize.cpp index c343cb12b09d..98691d51131d 100644 --- a/ydb/library/yql/providers/dq/opt/physical_optimize.cpp +++ b/ydb/library/yql/providers/dq/opt/physical_optimize.cpp @@ -241,6 +241,36 @@ class TDqsPhysicalOptProposalTransformer : public TOptimizeTransformerBase { return DqRewriteLeftPureJoin(node, ctx, *getParents(), IsGlobal); } + bool ValidateStreamLookupJoinFlags(const TDqJoin& join, TExprContext& ctx) { + bool leftAny = false; + bool rightAny = false; + if (const auto maybeFlags = join.Flags()) { + for (auto&& flag: maybeFlags.Cast()) { + auto&& name = flag.StringValue(); + if (name == "LeftAny"sv) { + leftAny = true; + continue; + } else if (name == "RightAny"sv) { + rightAny = true; + continue; + } + } + if (leftAny) { + ctx.AddError(TIssue(ctx.GetPosition(maybeFlags.Cast().Pos()), "Streamlookup ANY LEFT join is not implemented")); + return false; + } + } + if (!rightAny) { + if (false) { // Tempoarily change to waring to allow for smooth transition + ctx.AddError(TIssue(ctx.GetPosition(join.Pos()), "Streamlookup: must be LEFT JOIN /*+streamlookup(...)*/ ANY")); + return false; + } else { + ctx.AddWarning(TIssue(ctx.GetPosition(join.Pos()), "(Deprecation) Streamlookup: must be LEFT JOIN /*+streamlookup(...)*/ ANY")); + } + } + return true; + } + TMaybeNode RewriteStreamLookupJoin(TExprBase node, TExprContext& ctx) { const auto join = node.Cast(); if (join.JoinAlgo().StringValue() != "StreamLookupJoin") { @@ -252,18 +282,54 @@ class TDqsPhysicalOptProposalTransformer : public TOptimizeTransformerBase { if (!left) { return node; } + + if (!ValidateStreamLookupJoinFlags(join, ctx)) { + return {}; + } + + TExprNode::TPtr ttl; + TExprNode::TPtr maxCachedRows; + TExprNode::TPtr maxDelayedRows; + if (const auto maybeOptions = join.JoinAlgoOptions()) { + for (auto&& option: maybeOptions.Cast()) { + auto&& name = option.Name().Value(); + if (name == "TTL"sv) { + ttl = option.Value().Cast().Ptr(); + } else if (name == "MaxCachedRows"sv) { + maxCachedRows = option.Value().Cast().Ptr(); + } else if (name == "MaxDelayedRows"sv) { + maxDelayedRows = option.Value().Cast().Ptr(); + } + } + } + + if (!ttl) { + ttl = ctx.NewAtom(pos, 300); + } + if (!maxCachedRows) { + maxCachedRows = ctx.NewAtom(pos, 1'000'000); + } + if (!maxDelayedRows) { + maxDelayedRows = ctx.NewAtom(pos, 1'000'000); + } + auto rightInput = join.RightInput().Ptr(); + if (auto maybe = TExprBase(rightInput).Maybe()) { + rightInput = maybe.Cast().Input().Ptr(); + } + auto leftLabel = join.LeftLabel().Maybe() ? join.LeftLabel().Cast().Ptr() : ctx.NewAtom(pos, ""); + Y_ENSURE(join.RightLabel().Maybe()); auto cn = Build(ctx, pos) .Output(left.Output().Cast()) - .LeftLabel(join.LeftLabel().Cast()) - .RightInput(join.RightInput()) + .LeftLabel(leftLabel) + .RightInput(rightInput) .RightLabel(join.RightLabel().Cast()) .JoinKeys(join.JoinKeys()) .JoinType(join.JoinType()) .LeftJoinKeyNames(join.LeftJoinKeyNames()) .RightJoinKeyNames(join.RightJoinKeyNames()) - .TTL(ctx.NewAtom(pos, 300)) //TODO configure me - .MaxCachedRows(ctx.NewAtom(pos, 1'000'000)) //TODO configure me - .MaxDelayedRows(ctx.NewAtom(pos, 1'000'000)) //Configure me + .TTL(ttl) + .MaxCachedRows(maxCachedRows) + .MaxDelayedRows(maxDelayedRows) .Done(); auto lambda = Build(ctx, pos) diff --git a/ydb/library/yql/providers/dq/planner/execution_planner.cpp b/ydb/library/yql/providers/dq/planner/execution_planner.cpp index bf0a0cba842a..90bd08f2a91e 100644 --- a/ydb/library/yql/providers/dq/planner/execution_planner.cpp +++ b/ydb/library/yql/providers/dq/planner/execution_planner.cpp @@ -608,7 +608,7 @@ namespace NYql::NDqs { settings.SetRightLabel(streamLookup.RightLabel().StringValue()); settings.SetJoinType(streamLookup.JoinType().StringValue()); for (const auto& k: streamLookup.LeftJoinKeyNames()) { - *settings.AddLeftJoinKeyNames() = RemoveAliases(k.StringValue()); + *settings.AddLeftJoinKeyNames() = streamLookup.LeftLabel().StringValue().empty() ? k.StringValue() : RemoveAliases(k.StringValue()); } for (const auto& k: streamLookup.RightJoinKeyNames()) { *settings.AddRightJoinKeyNames() = RemoveAliases(k.StringValue()); @@ -619,9 +619,9 @@ namespace NYql::NDqs { const auto narrowOutputRowType = GetSeqItemType(streamLookup.Ptr()->GetTypeAnn()); Y_ABORT_UNLESS(narrowOutputRowType->GetKind() == ETypeAnnotationKind::Struct); settings.SetNarrowOutputRowType(NYql::NCommon::GetSerializedTypeAnnotation(narrowOutputRowType)); - settings.SetMaxDelayedRows(1'000'000); //TODO configure me - settings.SetCacheLimit(1'000'000); //TODO configure me - settings.SetCacheTtlSeconds(60); //TODO configure me + settings.SetCacheLimit(FromString(streamLookup.MaxCachedRows().StringValue())); + settings.SetCacheTtlSeconds(FromString(streamLookup.TTL().StringValue())); + settings.SetMaxDelayedRows(FromString(streamLookup.MaxDelayedRows().StringValue())); const auto inputRowType = GetSeqItemType(streamLookup.Output().Stage().Program().Ref().GetTypeAnn()); const auto outputRowType = GetSeqItemType(stage.Program().Args().Arg(inputIndex).Ref().GetTypeAnn()); diff --git a/ydb/library/yql/providers/dq/runtime/task_command_executor.cpp b/ydb/library/yql/providers/dq/runtime/task_command_executor.cpp index dc7c6a16e002..e07d1a98a7e6 100644 --- a/ydb/library/yql/providers/dq/runtime/task_command_executor.cpp +++ b/ydb/library/yql/providers/dq/runtime/task_command_executor.cpp @@ -56,7 +56,10 @@ void ToProto(T& proto, const NDq::TDqAsyncStats& stats) proto.SetRows(stats.Rows); proto.SetChunks(stats.Chunks); proto.SetSplits(stats.Splits); - + proto.SetFilteredBytes(stats.FilteredBytes); + proto.SetFilteredRows(stats.FilteredRows); + proto.SetQueuedBytes(stats.QueuedBytes); + proto.SetQueuedRows(stats.QueuedRows); proto.SetFirstMessageMs(stats.FirstMessageTs.MilliSeconds()); proto.SetPauseMessageMs(stats.PauseMessageTs.MilliSeconds()); proto.SetResumeMessageMs(stats.ResumeMessageTs.MilliSeconds()); diff --git a/ydb/library/yql/providers/dq/task_runner/tasks_runner_pipe.cpp b/ydb/library/yql/providers/dq/task_runner/tasks_runner_pipe.cpp index 1c08a50b41ac..e8f71d108166 100644 --- a/ydb/library/yql/providers/dq/task_runner/tasks_runner_pipe.cpp +++ b/ydb/library/yql/providers/dq/task_runner/tasks_runner_pipe.cpp @@ -549,6 +549,11 @@ void LoadFromProto(TDqAsyncStats& stats, const NYql::NDqProto::TDqAsyncBufferSta stats.LastMessageTs = TInstant::MilliSeconds(f.GetLastMessageMs()); stats.WaitTime = TDuration::MicroSeconds(f.GetWaitTimeUs()); stats.WaitPeriods = f.GetWaitPeriods(); + + stats.FilteredBytes = f.GetFilteredBytes(); + stats.FilteredRows = f.GetFilteredRows(); + stats.QueuedBytes = f.GetQueuedBytes(); + stats.QueuedRows = f.GetQueuedRows(); } /*______________________________________________________________________________________________*/ diff --git a/ydb/library/yql/providers/generic/actors/ut/yql_generic_lookup_actor_ut.cpp b/ydb/library/yql/providers/generic/actors/ut/yql_generic_lookup_actor_ut.cpp index c940fdbdbbc0..95b057e7d304 100644 --- a/ydb/library/yql/providers/generic/actors/ut/yql_generic_lookup_actor_ut.cpp +++ b/ydb/library/yql/providers/generic/actors/ut/yql_generic_lookup_actor_ut.cpp @@ -81,13 +81,13 @@ Y_UNIT_TEST_SUITE(GenericProviderLookupActor) { runtime.Initialize(); auto edge = runtime.AllocateEdgeActor(); - NYql::NConnector::NApi::TDataSourceInstance dsi; - dsi.Setkind(NYql::NConnector::NApi::EDataSourceKind::YDB); + NYql::TGenericDataSourceInstance dsi; + dsi.Setkind(NYql::EGenericDataSourceKind::YDB); dsi.mutable_endpoint()->Sethost("some_host"); dsi.mutable_endpoint()->Setport(2135); dsi.Setdatabase("some_db"); dsi.Setuse_tls(true); - dsi.set_protocol(::NYql::NConnector::NApi::EProtocol::NATIVE); + dsi.set_protocol(::NYql::EGenericProtocol::NATIVE); auto token = dsi.mutable_credentials()->mutable_token(); token->Settype("IAM"); token->Setvalue("TEST_TOKEN"); @@ -126,6 +126,12 @@ Y_UNIT_TEST_SUITE(GenericProviderLookupActor) { .Operand().Equal().Column("optional_id").OptionalValue(100).Done().Done() .Done() .Done() + .Operand() + .Conjunction() + .Operand().Equal().Column("id").Value(2).Done().Done() + .Operand().Equal().Column("optional_id").OptionalValue(102).Done().Done() + .Done() + .Done() .Done() .Done() .Done() @@ -226,4 +232,195 @@ Y_UNIT_TEST_SUITE(GenericProviderLookupActor) { } } + Y_UNIT_TEST(LookupWithErrors) { + auto alloc = std::make_shared(__LOCATION__, NKikimr::TAlignedPagePoolCounters(), true, false); + NKikimr::NMiniKQL::TMemoryUsageInfo memUsage("TestMemUsage"); + NKikimr::NMiniKQL::THolderFactory holderFactory(alloc->Ref(), memUsage); + NKikimr::NMiniKQL::TTypeEnvironment typeEnv(*alloc); + NKikimr::NMiniKQL::TTypeBuilder typeBuilder(typeEnv); + + auto loggerConfig = NYql::NProto::TLoggingConfig(); + loggerConfig.set_allcomponentslevel(::NYql::NProto::TLoggingConfig_ELevel::TLoggingConfig_ELevel_TRACE); + NYql::NLog::InitLogger(loggerConfig, false); + + TTestActorRuntimeBase runtime(1, 1, true); + runtime.Initialize(); + auto edge = runtime.AllocateEdgeActor(); + + NYql::TGenericDataSourceInstance dsi; + dsi.Setkind(NYql::EGenericDataSourceKind::YDB); + dsi.mutable_endpoint()->Sethost("some_host"); + dsi.mutable_endpoint()->Setport(2135); + dsi.Setdatabase("some_db"); + dsi.Setuse_tls(true); + dsi.set_protocol(::NYql::EGenericProtocol::NATIVE); + auto token = dsi.mutable_credentials()->mutable_token(); + token->Settype("IAM"); + token->Setvalue("TEST_TOKEN"); + + auto connectorMock = std::make_shared(); + + // clang-format off + // step 1: ListSplits + { + ::testing::InSequence seq; + for (grpc::StatusCode readStatus : { grpc::StatusCode::UNAVAILABLE, grpc::StatusCode::OK }) { + for (grpc::StatusCode listStatus : { grpc::StatusCode::UNAVAILABLE, readStatus == grpc::StatusCode::OK ? grpc::StatusCode::DEADLINE_EXCEEDED : grpc::StatusCode::UNAVAILABLE, grpc::StatusCode::OK }) { + auto listBuilder = connectorMock->ExpectListSplits(); + listBuilder + .Select() + .DataSourceInstance(dsi) + .What() + .Column("id", Ydb::Type::UINT64) + .NullableColumn("optional_id", Ydb::Type::UINT64) + .NullableColumn("string_value", Ydb::Type::STRING) + .Done() + .Table("lookup_test") + .Where() + .Filter() + .Disjunction() + .Operand() + .Conjunction() + .Operand().Equal().Column("id").Value(2).Done().Done() + .Operand().Equal().Column("optional_id").OptionalValue(102).Done().Done() + .Done() + .Done() + .Operand() + .Conjunction() + .Operand().Equal().Column("id").Value(1).Done().Done() + .Operand().Equal().Column("optional_id").OptionalValue(101).Done().Done() + .Done() + .Done() + .Operand() + .Conjunction() + .Operand().Equal().Column("id").Value(0).Done().Done() + .Operand().Equal().Column("optional_id").OptionalValue(100).Done().Done() + .Done() + .Done() + .Operand() + .Conjunction() + .Operand().Equal().Column("id").Value(2).Done().Done() + .Operand().Equal().Column("optional_id").OptionalValue(102).Done().Done() + .Done() + .Done() + .Done() + .Done() + .Done() + .Done() + .MaxSplitCount(1) + ; + if (listStatus != grpc::StatusCode::OK) { + listBuilder + .Status(NYdbGrpc::TGrpcStatus(listStatus, "Mocked Error")) + ; + continue; + } + listBuilder + .Result() + .AddResponse(NewSuccess()) + .Description("Actual split info is not important") + ; + } + + auto readBuilder = connectorMock->ExpectReadSplits(); + readBuilder + .DataSourceInstance(dsi) + .Filtering(NYql::NConnector::NApi::TReadSplitsRequest::FILTERING_MANDATORY) + .Split() + .Description("Actual split info is not important") + .Done() + ; + if (readStatus != grpc::StatusCode::OK) { + readBuilder + .Status(NYdbGrpc::TGrpcStatus(readStatus, "Mocked Error")) + ; + continue; + } + readBuilder + .Result() + .AddResponse( + MakeRecordBatch( + MakeArray("id", {0, 1, 2}, arrow::uint64()), + MakeArray("optional_id", {100, 101, 103}, arrow::uint64()), // the last value is intentially wrong + MakeArray("string_value", {"a", "b", "c"}, arrow::utf8()) + ), + NewSuccess() + ) + ; + } + } + // clang-format on + + NYql::Generic::TLookupSource lookupSourceSettings; + *lookupSourceSettings.mutable_data_source_instance() = dsi; + lookupSourceSettings.Settable("lookup_test"); + lookupSourceSettings.SetServiceAccountId("testsaid"); + lookupSourceSettings.SetServiceAccountIdSignature("fake_signature"); + + google::protobuf::Any packedLookupSource; + Y_ABORT_UNLESS(packedLookupSource.PackFrom(lookupSourceSettings)); + + NKikimr::NMiniKQL::TStructTypeBuilder keyTypeBuilder{typeEnv}; + keyTypeBuilder.Add("id", typeBuilder.NewDataType(NYql::NUdf::EDataSlot::Uint64, false)); + keyTypeBuilder.Add("optional_id", typeBuilder.NewDataType(NYql::NUdf::EDataSlot::Uint64, true)); + NKikimr::NMiniKQL::TStructTypeBuilder outputypeBuilder{typeEnv}; + outputypeBuilder.Add("string_value", typeBuilder.NewDataType(NYql::NUdf::EDataSlot::String, true)); + + auto guard = Guard(*alloc.get()); + auto keyTypeHelper = std::make_shared(keyTypeBuilder.Build()); + + auto [lookupSource, actor] = NYql::NDq::CreateGenericLookupActor( + connectorMock, + std::make_shared(), + edge, + nullptr, + alloc, + keyTypeHelper, + std::move(lookupSourceSettings), + keyTypeBuilder.Build(), + outputypeBuilder.Build(), + typeEnv, + holderFactory, + 1'000'000); + auto lookupActor = runtime.Register(actor); + + auto request = std::make_shared(3, keyTypeHelper->GetValueHash(), keyTypeHelper->GetValueEqual()); + for (size_t i = 0; i != 3; ++i) { + NYql::NUdf::TUnboxedValue* keyItems; + auto key = holderFactory.CreateDirectArrayHolder(2, keyItems); + keyItems[0] = NYql::NUdf::TUnboxedValuePod(ui64(i)); + keyItems[1] = NYql::NUdf::TUnboxedValuePod(ui64(100 + i)); + request->emplace(std::move(key), NYql::NUdf::TUnboxedValue{}); + } + + guard.Release(); // let actors use alloc + + auto callLookupActor = new TCallLookupActor(alloc, lookupActor, request); + runtime.Register(callLookupActor); + + auto ev = runtime.GrabEdgeEventRethrow(edge); + auto guard2 = Guard(*alloc.get()); + auto lookupResult = ev->Get()->Result.lock(); + UNIT_ASSERT(lookupResult); + + UNIT_ASSERT_EQUAL(3, lookupResult->size()); + { + const auto* v = lookupResult->FindPtr(CreateStructValue(holderFactory, {0, 100})); + UNIT_ASSERT(v); + NYql::NUdf::TUnboxedValue val = v->GetElement(0); + UNIT_ASSERT(val.AsStringRef() == TStringBuf("a")); + } + { + const auto* v = lookupResult->FindPtr(CreateStructValue(holderFactory, {1, 101})); + UNIT_ASSERT(v); + NYql::NUdf::TUnboxedValue val = v->GetElement(0); + UNIT_ASSERT(val.AsStringRef() == TStringBuf("b")); + } + { + const auto* v = lookupResult->FindPtr(CreateStructValue(holderFactory, {2, 102})); + UNIT_ASSERT(v); + UNIT_ASSERT(!*v); + } + } + } // Y_UNIT_TEST_SUITE(GenericProviderLookupActor) diff --git a/ydb/library/yql/providers/generic/actors/yql_generic_base_actor.h b/ydb/library/yql/providers/generic/actors/yql_generic_base_actor.h index d3a07148ccea..49d769de4360 100644 --- a/ydb/library/yql/providers/generic/actors/yql_generic_base_actor.h +++ b/ydb/library/yql/providers/generic/actors/yql_generic_base_actor.h @@ -21,6 +21,7 @@ namespace NYql::NDq { EvReadSplitsPart, EvReadSplitsFinished, EvError, + EvRetry, EvEnd }; diff --git a/ydb/library/yql/providers/generic/actors/yql_generic_lookup_actor.cpp b/ydb/library/yql/providers/generic/actors/yql_generic_lookup_actor.cpp index e1f3854b79b6..8afcc2cc2be8 100644 --- a/ydb/library/yql/providers/generic/actors/yql_generic_lookup_actor.cpp +++ b/ydb/library/yql/providers/generic/actors/yql_generic_lookup_actor.cpp @@ -23,11 +23,15 @@ #include #include +#include + namespace NYql::NDq { using namespace NActors; namespace { + constexpr ui32 RequestRetriesLimit = 10; // TODO lookup parameters or PRAGMA? + constexpr TDuration RequestTimeout = TDuration::Minutes(3); // TODO lookup parameters or PRAGMA? const NKikimr::NMiniKQL::TStructType* MergeStructTypes(const NKikimr::NMiniKQL::TTypeEnvironment& env, const NKikimr::NMiniKQL::TStructType* t1, const NKikimr::NMiniKQL::TStructType* t2) { Y_ABORT_UNLESS(t1); @@ -45,7 +49,7 @@ namespace NYql::NDq { template T ExtractFromConstFuture(const NThreading::TFuture& f) { // We want to avoid making a copy of data stored in a future. - // But there is no direct way to extract data from a const future5 + // But there is no direct way to extract data from a const future // So, we make a copy of the future, that is cheap. Then, extract the value from this copy. // It destructs the value in the original future, but this trick is legal and documented here: // https://docs.yandex-team.ru/arcadia-cpp/cookbook/concurrency @@ -59,6 +63,12 @@ namespace NYql::NDq { public TGenericBaseActor { using TBase = TGenericBaseActor; + using ILookupRetryPolicy = IRetryPolicy; + using ILookupRetryState = ILookupRetryPolicy::IRetryState; + + struct TEvLookupRetry : NActors::TEventLocal { + }; + public: TGenericLookupActor( NConnector::IClient::TPtr connectorClient, @@ -85,6 +95,24 @@ namespace NYql::NDq { , HolderFactory(holderFactory) , ColumnDestinations(CreateColumnDestination()) , MaxKeysInRequest(maxKeysInRequest) + , RetryPolicy( + ILookupRetryPolicy::GetExponentialBackoffPolicy( + /* retryClassFunction */ + [](const NYdbGrpc::TGrpcStatus& status) { + if (NConnector::GrpcStatusNeedsRetry(status)) { + return ERetryErrorClass::ShortRetry; + } + if (status.GRpcStatusCode == grpc::DEADLINE_EXCEEDED) { + return ERetryErrorClass::ShortRetry; // TODO LongRetry? + } + return ERetryErrorClass::NoRetry; + }, + /* minDelay */ TDuration::MilliSeconds(1), + /* minLongRetryDelay */ TDuration::MilliSeconds(500), + /* maxDelay */ TDuration::Seconds(1), + /* maxRetries */ RequestRetriesLimit, + /* maxTime */ TDuration::Minutes(5), + /* scaleFactor */ 2)) { InitMonCounters(taskCounters); } @@ -96,6 +124,10 @@ namespace NYql::NDq { private: void Free() { auto guard = Guard(*Alloc); + if (Request && InFlight) { + // If request fails on (unrecoverable) error or cancelled, we may end up with non-zero InFlight (when request successfully completed, @Request is nullptr) + InFlight->Dec(); + } Request.reset(); KeyTypeHelper.reset(); } @@ -111,17 +143,18 @@ namespace NYql::NDq { ResultBytes = component->GetCounter("Bytes"); AnswerTime = component->GetCounter("AnswerMs"); CpuTime = component->GetCounter("CpuUs"); + InFlight = component->GetCounter("InFlight"); } public: void Bootstrap() { auto dsi = LookupSource.data_source_instance(); YQL_CLOG(INFO, ProviderGeneric) << "New generic proivider lookup source actor(ActorId=" << SelfId() << ") for" - << " kind=" << NYql::NConnector::NApi::EDataSourceKind_Name(dsi.kind()) + << " kind=" << NYql::EGenericDataSourceKind_Name(dsi.kind()) << ", endpoint=" << dsi.endpoint().ShortDebugString() << ", database=" << dsi.database() << ", use_tls=" << ToString(dsi.use_tls()) - << ", protocol=" << NYql::NConnector::NApi::EProtocol_Name(dsi.protocol()) + << ", protocol=" << NYql::EGenericProtocol_Name(dsi.protocol()) << ", table=" << LookupSource.table(); Become(&TGenericLookupActor::StateFunc); } @@ -150,12 +183,17 @@ namespace NYql::NDq { hFunc(TEvReadSplitsPart, Handle); hFunc(TEvReadSplitsFinished, Handle); hFunc(TEvError, Handle); + hFunc(TEvLookupRetry, Handle); hFunc(NActors::TEvents::TEvPoison, Handle);) void Handle(TEvListSplitsIterator::TPtr ev) { auto& iterator = ev->Get()->Iterator; iterator->ReadNext().Subscribe( - [actorSystem = TActivationContext::ActorSystem(), selfId = SelfId()](const NConnector::TAsyncResult& asyncResult) { + [ + actorSystem = TActivationContext::ActorSystem(), + selfId = SelfId(), + retryState = RetryState + ](const NConnector::TAsyncResult& asyncResult) { YQL_CLOG(DEBUG, ProviderGeneric) << "ActorId=" << selfId << " Got TListSplitsResponse from Connector"; auto result = ExtractFromConstFuture(asyncResult); if (result.Status.Ok()) { @@ -163,7 +201,7 @@ namespace NYql::NDq { auto ev = new TEvListSplitsPart(std::move(*result.Response)); actorSystem->Send(new NActors::IEventHandle(selfId, selfId, ev)); } else { - SendError(actorSystem, selfId, result.Status); + SendRetryOrError(actorSystem, selfId, result.Status, retryState); } }); } @@ -184,14 +222,18 @@ namespace NYql::NDq { *readRequest.add_splits() = split; readRequest.Setformat(NConnector::NApi::TReadSplitsRequest_EFormat::TReadSplitsRequest_EFormat_ARROW_IPC_STREAMING); readRequest.set_filtering(NConnector::NApi::TReadSplitsRequest::FILTERING_MANDATORY); - Connector->ReadSplits(readRequest).Subscribe([actorSystem = TActivationContext::ActorSystem(), selfId = SelfId()](const NConnector::TReadSplitsStreamIteratorAsyncResult& asyncResult) { + Connector->ReadSplits(readRequest, RequestTimeout).Subscribe([ + actorSystem = TActivationContext::ActorSystem(), + selfId = SelfId(), + retryState = RetryState + ](const NConnector::TReadSplitsStreamIteratorAsyncResult& asyncResult) { YQL_CLOG(DEBUG, ProviderGeneric) << "ActorId=" << selfId << " Got ReadSplitsStreamIterator from Connector"; auto result = ExtractFromConstFuture(asyncResult); if (result.Status.Ok()) { auto ev = new TEvReadSplitsIterator(std::move(result.Iterator)); actorSystem->Send(new NActors::IEventHandle(selfId, selfId, ev)); } else { - SendError(actorSystem, selfId, result.Status); + SendRetryOrError(actorSystem, selfId, result.Status, retryState); } }); } @@ -220,6 +262,11 @@ namespace NYql::NDq { actorSystem->Send(new NActors::IEventHandle(ParentId, SelfId(), errEv.release())); } + void Handle(TEvLookupRetry::TPtr) { + auto guard = Guard(*Alloc); + SendRequest(); + } + void Handle(NActors::TEvents::TEvPoison::TPtr) { PassAway(); } @@ -238,17 +285,22 @@ namespace NYql::NDq { if (!request) { return; } - auto startCycleCount = GetCycleCountFast(); SentTime = TInstant::Now(); YQL_CLOG(DEBUG, ProviderGeneric) << "ActorId=" << SelfId() << " Got LookupRequest for " << request->size() << " keys"; Y_ABORT_IF(request->size() == 0 || request->size() > MaxKeysInRequest); - if (Count) { Count->Inc(); + InFlight->Inc(); Keys->Add(request->size()); } Request = std::move(request); + RetryState = std::shared_ptr(RetryPolicy->CreateRetryState()); + SendRequest(); + } + + void SendRequest() { + auto startCycleCount = GetCycleCountFast(); NConnector::NApi::TListSplitsRequest splitRequest; auto error = FillSelect(*splitRequest.add_selects()); @@ -258,7 +310,11 @@ namespace NYql::NDq { }; splitRequest.Setmax_split_count(1); - Connector->ListSplits(splitRequest).Subscribe([actorSystem = TActivationContext::ActorSystem(), selfId = SelfId()](const NConnector::TListSplitsStreamIteratorAsyncResult& asyncResult) { + Connector->ListSplits(splitRequest, RequestTimeout).Subscribe([ + actorSystem = TActivationContext::ActorSystem(), + selfId = SelfId(), + retryState = RetryState + ](const NConnector::TListSplitsStreamIteratorAsyncResult& asyncResult) { auto result = ExtractFromConstFuture(asyncResult); if (result.Status.Ok()) { YQL_CLOG(DEBUG, ProviderGeneric) << "ActorId=" << selfId << " Got TListSplitsStreamIterator"; @@ -266,7 +322,7 @@ namespace NYql::NDq { auto ev = new TEvListSplitsIterator(std::move(result.Iterator)); actorSystem->Send(new NActors::IEventHandle(selfId, selfId, ev)); } else { - SendError(actorSystem, selfId, result.Status); + SendRetryOrError(actorSystem, selfId, result.Status, retryState); } }); if (CpuTime) { @@ -276,12 +332,17 @@ namespace NYql::NDq { void ReadNextData() { ReadSplitsIterator->ReadNext().Subscribe( - [actorSystem = TActivationContext::ActorSystem(), selfId = SelfId()](const NConnector::TAsyncResult& asyncResult) { + [ + actorSystem = TActivationContext::ActorSystem(), + selfId = SelfId(), + retryState = RetryState + ](const NConnector::TAsyncResult& asyncResult) { auto result = ExtractFromConstFuture(asyncResult); if (result.Status.Ok()) { YQL_CLOG(DEBUG, ProviderGeneric) << "ActorId=" << selfId << " Got DataChunk"; Y_ABORT_UNLESS(result.Response); auto& response = *result.Response; + // TODO: retry on some YDB errors if (NConnector::IsSuccess(response)) { auto ev = new TEvReadSplitsPart(std::move(response)); actorSystem->Send(new NActors::IEventHandle(selfId, selfId, ev)); @@ -293,7 +354,7 @@ namespace NYql::NDq { auto ev = new TEvReadSplitsFinished(std::move(result.Status)); actorSystem->Send(new NActors::IEventHandle(selfId, selfId, ev)); } else { - SendError(actorSystem, selfId, result.Status); + SendRetryOrError(actorSystem, selfId, result.Status, retryState); } }); } @@ -344,6 +405,7 @@ namespace NYql::NDq { auto ev = new IDqAsyncLookupSource::TEvLookupResult(Request); if (AnswerTime) { AnswerTime->Add((TInstant::Now() - SentTime).MilliSeconds()); + InFlight->Dec(); } Request.reset(); TActivationContext::ActorSystem()->Send(new NActors::IEventHandle(ParentId, SelfId(), ev)); @@ -358,7 +420,13 @@ namespace NYql::NDq { new TEvError(std::move(error))); } - static void SendError(NActors::TActorSystem* actorSystem, const NActors::TActorId& selfId, const NYdbGrpc::TGrpcStatus& status) { + static void SendRetryOrError(NActors::TActorSystem* actorSystem, const NActors::TActorId& selfId, const NYdbGrpc::TGrpcStatus& status, std::shared_ptr retryState) { + auto nextRetry = retryState->GetNextRetryDelay(status); + if (nextRetry) { + YQL_CLOG(WARN, ProviderGeneric) << "ActorId=" << selfId << " Got retrievable GRPC Error from Connector: " << status.ToDebugString() << ", retry scheduled in " << *nextRetry; + actorSystem->Schedule(*nextRetry, new IEventHandle(selfId, selfId, new TEvLookupRetry())); + return; + } SendError(actorSystem, selfId, NConnector::ErrorFromGRPCStatus(status)); } @@ -397,6 +465,19 @@ namespace NYql::NDq { return result; } + void AddClause(NConnector::NApi::TPredicate::TDisjunction &disjunction, + ui32 columnsCount, const NUdf::TUnboxedValue& keys) { + NConnector::NApi::TPredicate::TConjunction& conjunction = *disjunction.mutable_operands()->Add()->mutable_conjunction(); + for (ui32 c = 0; c != columnsCount; ++c) { + NConnector::NApi::TPredicate::TComparison& eq = *conjunction.mutable_operands()->Add()->mutable_comparison(); + eq.set_operation(NConnector::NApi::TPredicate::TComparison::EOperation::TPredicate_TComparison_EOperation_EQ); + eq.mutable_left_value()->set_column(TString(KeyType->GetMemberName(c))); + auto rightTypedValue = eq.mutable_right_value()->mutable_typed_value(); + ExportTypeToProto(KeyType->GetMemberType(c), *rightTypedValue->mutable_type()); + ExportValueToProto(KeyType->GetMemberType(c), keys.GetElement(c), *rightTypedValue->mutable_value()); + } + } + TString FillSelect(NConnector::NApi::TSelect& select) { auto dsi = LookupSource.data_source_instance(); auto error = TokenProvider->MaybeFillToken(dsi); @@ -413,19 +494,16 @@ namespace NYql::NDq { select.mutable_from()->Settable(LookupSource.table()); - NConnector::NApi::TPredicate_TDisjunction disjunction; - for (const auto& [k, _] : *Request) { - NConnector::NApi::TPredicate_TConjunction conjunction; - for (ui32 c = 0; c != KeyType->GetMembersCount(); ++c) { - NConnector::NApi::TPredicate_TComparison eq; - eq.Setoperation(NConnector::NApi::TPredicate_TComparison_EOperation::TPredicate_TComparison_EOperation_EQ); - eq.mutable_left_value()->Setcolumn(TString(KeyType->GetMemberName(c))); - auto rightTypedValue = eq.mutable_right_value()->mutable_typed_value(); - ExportTypeToProto(KeyType->GetMemberType(c), *rightTypedValue->mutable_type()); - ExportValueToProto(KeyType->GetMemberType(c), k.GetElement(c), *rightTypedValue->mutable_value()); - *conjunction.mutable_operands()->Add()->mutable_comparison() = eq; - } - *disjunction.mutable_operands()->Add()->mutable_conjunction() = conjunction; + NConnector::NApi::TPredicate::TDisjunction disjunction; + for (const auto& [keys, _] : *Request) { + // TODO consider skipping already retrieved keys + // ... but careful, can we end up with zero? TODO + AddClause(disjunction, KeyType->GetMembersCount(), keys); + } + auto& keys = Request->begin()->first; // Request is never empty + // Pad query with dummy clauses to improve caching + for (ui32 nRequests = Request->size(); !IsPowerOf2(nRequests) && nRequests < MaxKeysInRequest; ++nRequests) { + AddClause(disjunction, KeyType->GetMembersCount(), keys); } *select.mutable_where()->mutable_filter_typed()->mutable_disjunction() = disjunction; return {}; @@ -447,6 +525,8 @@ namespace NYql::NDq { std::shared_ptr Request; NConnector::IReadSplitsStreamIterator::TPtr ReadSplitsIterator; // TODO move me to TEvReadSplitsPart NKikimr::NMiniKQL::TKeyPayloadPairVector LookupResult; + ILookupRetryPolicy::TPtr RetryPolicy; + std::shared_ptr RetryState; ::NMonitoring::TDynamicCounters::TCounterPtr Count; ::NMonitoring::TDynamicCounters::TCounterPtr Keys; ::NMonitoring::TDynamicCounters::TCounterPtr ResultRows; @@ -454,6 +534,7 @@ namespace NYql::NDq { ::NMonitoring::TDynamicCounters::TCounterPtr ResultChunks; ::NMonitoring::TDynamicCounters::TCounterPtr AnswerTime; ::NMonitoring::TDynamicCounters::TCounterPtr CpuTime; + ::NMonitoring::TDynamicCounters::TCounterPtr InFlight; TInstant SentTime; }; diff --git a/ydb/library/yql/providers/generic/actors/yql_generic_read_actor.cpp b/ydb/library/yql/providers/generic/actors/yql_generic_read_actor.cpp index 3010e3e04018..28f59913cb2d 100644 --- a/ydb/library/yql/providers/generic/actors/yql_generic_read_actor.cpp +++ b/ydb/library/yql/providers/generic/actors/yql_generic_read_actor.cpp @@ -495,11 +495,11 @@ namespace NYql::NDq { { const auto dsi = source.select().data_source_instance(); YQL_CLOG(INFO, ProviderGeneric) << "Creating read actor with params:" - << " kind=" << NYql::NConnector::NApi::EDataSourceKind_Name(dsi.kind()) + << " kind=" << NYql::EGenericDataSourceKind_Name(dsi.kind()) << ", endpoint=" << dsi.endpoint().ShortDebugString() << ", database=" << dsi.database() << ", use_tls=" << ToString(dsi.use_tls()) - << ", protocol=" << NYql::NConnector::NApi::EProtocol_Name(dsi.protocol()); + << ", protocol=" << NYql::EGenericProtocol_Name(dsi.protocol()); // FIXME: strange piece of logic - authToken is created but not used: // https://a.yandex-team.ru/arcadia/ydb/library/yql/providers/clickhouse/actors/yql_ch_read_actor.cpp?rev=r11550199#L140 @@ -528,7 +528,8 @@ namespace NYql::NDq { auto tokenProvider = CreateGenericTokenProvider( source.GetToken(), - source.GetServiceAccountId(), source.GetServiceAccountIdSignature(), + source.GetServiceAccountId(), + source.GetServiceAccountIdSignature(), credentialsFactory); const auto actor = new TGenericReadActor( diff --git a/ydb/library/yql/providers/generic/actors/yql_generic_token_provider.cpp b/ydb/library/yql/providers/generic/actors/yql_generic_token_provider.cpp index 06673fbe16c0..753bb4c94027 100644 --- a/ydb/library/yql/providers/generic/actors/yql_generic_token_provider.cpp +++ b/ydb/library/yql/providers/generic/actors/yql_generic_token_provider.cpp @@ -10,14 +10,15 @@ namespace NYql::NDq { } TGenericTokenProvider::TGenericTokenProvider( - const TString& serviceAccountId, const TString& ServiceAccountIdSignature, + const TString& serviceAccountId, + const TString& serviceAccountIdSignature, const ISecuredServiceAccountCredentialsFactory::TPtr& credentialsFactory) { Y_ENSURE(!serviceAccountId.empty(), "No service account provided"); - Y_ENSURE(!ServiceAccountIdSignature.empty(), "No service account signature provided"); + Y_ENSURE(!serviceAccountIdSignature.empty(), "No service account signature provided"); Y_ENSURE(credentialsFactory, "CredentialsFactory is not initialized"); auto structuredTokenJSON = - TStructuredTokenBuilder().SetServiceAccountIdAuth(serviceAccountId, ServiceAccountIdSignature).ToJson(); + TStructuredTokenBuilder().SetServiceAccountIdAuth(serviceAccountId, serviceAccountIdSignature).ToJson(); Y_ENSURE(structuredTokenJSON, "empty structured token"); @@ -26,7 +27,7 @@ namespace NYql::NDq { CredentialsProvider_ = credentialsProviderFactory->CreateProvider(); } - TString TGenericTokenProvider::MaybeFillToken(NConnector::NApi::TDataSourceInstance& dsi) const { + TString TGenericTokenProvider::MaybeFillToken(NYql::TGenericDataSourceInstance& dsi) const { // 1. Don't need tokens if basic auth is set if (dsi.credentials().has_basic()) { return {}; @@ -58,7 +59,8 @@ namespace NYql::NDq { } TGenericTokenProvider::TPtr - CreateGenericTokenProvider(const TString& staticIamToken, const TString& serviceAccountId, + CreateGenericTokenProvider(const TString& staticIamToken, + const TString& serviceAccountId, const TString& serviceAccountIdSignature, const ISecuredServiceAccountCredentialsFactory::TPtr& credentialsFactory) { if (!staticIamToken.empty()) { diff --git a/ydb/library/yql/providers/generic/actors/yql_generic_token_provider.h b/ydb/library/yql/providers/generic/actors/yql_generic_token_provider.h index f780ecc60734..304398d4f41c 100644 --- a/ydb/library/yql/providers/generic/actors/yql_generic_token_provider.h +++ b/ydb/library/yql/providers/generic/actors/yql_generic_token_provider.h @@ -22,7 +22,7 @@ namespace NYql::NDq { // MaybeFillToken sets IAM-token within DataSourceInstance. // Returns string containing error, if it happened. - TString MaybeFillToken(NConnector::NApi::TDataSourceInstance& dsi) const; + TString MaybeFillToken(NYql::TGenericDataSourceInstance& dsi) const; private: TString StaticIAMToken_; @@ -32,6 +32,7 @@ namespace NYql::NDq { TGenericTokenProvider::TPtr CreateGenericTokenProvider( const TString& staticIamToken, - const TString& serviceAccountId, const TString& ServiceAccountIdSignature, + const TString& serviceAccountId, + const TString& serviceAccountIdSignature, const ISecuredServiceAccountCredentialsFactory::TPtr& credentialsFactory); } // namespace NYql::NDq diff --git a/ydb/library/yql/providers/generic/connector/api/common/data_source.proto b/ydb/library/yql/providers/generic/connector/api/common/data_source.proto deleted file mode 100644 index 68d5ca4c264e..000000000000 --- a/ydb/library/yql/providers/generic/connector/api/common/data_source.proto +++ /dev/null @@ -1,105 +0,0 @@ -syntax = "proto3"; - -package NYql.NConnector.NApi; - -import "ydb/library/yql/providers/generic/connector/api/common/endpoint.proto"; - -option go_package = "github.com/ydb-platform/ydb/ydb/library/yql/providers/generic/connector/api/common"; - -// TCredentials represents various ways of user authentication in the data source instance -message TCredentials { - message TBasic { - string username = 1; - string password = 2; - } - - message TToken { - string type = 1; - string value = 2; - } - - oneof payload { - TBasic basic = 1; - TToken token = 2; - } -} - -// EDataSourceKind enumerates the external data sources -// supported by the federated query system -enum EDataSourceKind { - DATA_SOURCE_KIND_UNSPECIFIED = 0; - CLICKHOUSE = 1; - POSTGRESQL = 2; - S3 = 3; - YDB = 4; - MYSQL = 5; - MS_SQL_SERVER = 6; - GREENPLUM = 7; - ORACLE = 8; - LOGGING = 9; -} - -// EProtocol generalizes various kinds of network protocols supported by different databases. -enum EProtocol { - PROTOCOL_UNSPECIFIED = 0; - NATIVE = 1; // CLICKHOUSE, POSTGRESQL - HTTP = 2; // CLICKHOUSE, S3 -} - -// TPostgreSQLDataSourceOptions represents settings specific to PostgreSQL -message TPostgreSQLDataSourceOptions { - // PostgreSQL schema - string schema = 1; -} - -// TClickhouseDataSourceOptions represents settings specific to Clickhouse -message TClickhouseDataSourceOptions { -} - -// TS3DataSourceOptions represents settings specific to S3 (Simple Storage Service) -message TS3DataSourceOptions { - // the region where data is stored - string region = 1; - // the bucket the object belongs to - string bucket = 2; -} - -// TGreenplumDataSourceOptions represents settings specific to Greenplum -message TGreenplumDataSourceOptions { - // Greenplum schema - string schema = 1; -} - -// TOracleDataSourceOptions represents settings specific to Oracle -message TOracleDataSourceOptions { - // Oracle service_name - alias to SID of oracle INSTANCE, or SID, or PDB. - // More about connection options in Oracle docs: - // https://docs.oracle.com/en/database/other-databases/essbase/21/essoa/connection-string-formats.html - string service_name = 1; -} - -// TDataSourceInstance helps to identify the instance of a data source to redirect request to. -message TDataSourceInstance { - // Data source kind - EDataSourceKind kind = 1; - // Network address - TEndpoint endpoint = 2; - // Database name - string database = 3; - // Credentials to access database - TCredentials credentials = 4; - // If true, Connector server will use secure connections to access remote data sources. - // Certificates will be obtained from the standard system paths. - bool use_tls = 5; - // Allows to specify network protocol that should be used between - // during the connection between Connector and the remote data source - EProtocol protocol = 6; - // Options specific to various data sources - oneof options { - TPostgreSQLDataSourceOptions pg_options = 7; - TClickhouseDataSourceOptions ch_options = 8; - TS3DataSourceOptions s3_options = 9; - TGreenplumDataSourceOptions gp_options = 10; - TOracleDataSourceOptions oracle_options = 11; - } -} diff --git a/ydb/library/yql/providers/generic/connector/api/common/endpoint.proto b/ydb/library/yql/providers/generic/connector/api/common/endpoint.proto deleted file mode 100644 index 1e4f19eff4f8..000000000000 --- a/ydb/library/yql/providers/generic/connector/api/common/endpoint.proto +++ /dev/null @@ -1,11 +0,0 @@ -syntax = "proto3"; - -package NYql.NConnector.NApi; - -option go_package = "github.com/ydb-platform/ydb/ydb/library/yql/providers/generic/connector/api/common"; - -// TEndpoint represents the network address of a data source instance -message TEndpoint { - string host = 1; - uint32 port = 2; -} diff --git a/ydb/library/yql/providers/generic/connector/api/common/ya.make b/ydb/library/yql/providers/generic/connector/api/common/ya.make deleted file mode 100644 index 0dd3d856cb78..000000000000 --- a/ydb/library/yql/providers/generic/connector/api/common/ya.make +++ /dev/null @@ -1,11 +0,0 @@ -# WARNING: never add dependency to YDB protofiles here, -# because it would break Go code generation. - -PROTO_LIBRARY() - -SRCS( - data_source.proto - endpoint.proto -) - -END() diff --git a/ydb/library/yql/providers/generic/connector/api/service/connector.proto b/ydb/library/yql/providers/generic/connector/api/service/connector.proto index ec1b855ff3c8..e70c8ab2270a 100644 --- a/ydb/library/yql/providers/generic/connector/api/service/connector.proto +++ b/ydb/library/yql/providers/generic/connector/api/service/connector.proto @@ -4,7 +4,7 @@ package NYql.NConnector.NApi; import "ydb/library/yql/providers/generic/connector/api/service/protos/connector.proto"; -option go_package = "github.com/ydb-platform/ydb/ydb/library/yql/providers/generic/connector/libgo/service"; +option go_package = "a.yandex-team.ru/contrib/ydb/library/yql/providers/generic/connector/libgo/service"; // Connector provides unified interface for various data sources that can be used to extend capabilities // of YQ and YQL services. diff --git a/ydb/library/yql/providers/generic/connector/api/service/protos/connector.proto b/ydb/library/yql/providers/generic/connector/api/service/protos/connector.proto index 1d638fab1238..a91d3eb21733 100644 --- a/ydb/library/yql/providers/generic/connector/api/service/protos/connector.proto +++ b/ydb/library/yql/providers/generic/connector/api/service/protos/connector.proto @@ -5,16 +5,16 @@ package NYql.NConnector.NApi; import "ydb/public/api/protos/ydb_value.proto"; import "ydb/public/api/protos/ydb_status_codes.proto"; import "ydb/public/api/protos/ydb_issue_message.proto"; -import "ydb/library/yql/providers/generic/connector/api/common/data_source.proto"; +import "yql/essentials/providers/common/proto/gateways_config.proto"; -option go_package = "github.com/ydb-platform/ydb/ydb/library/yql/providers/generic/connector/libgo/service/protos"; +option go_package = "a.yandex-team.ru/contrib/ydb/library/yql/providers/generic/connector/libgo/service/protos"; // ---------- API Requests ---------- // TListTablesRequest requests the list of tables in a particular database of the data source message TListTablesRequest { // Data source instance to connect - TDataSourceInstance data_source_instance = 1; + NYql.TGenericDataSourceInstance data_source_instance = 1; // There may be a huge number of tables in the data source, // and here are ways to extract only necessary ones: @@ -36,7 +36,7 @@ message TListTablesResponse { // TDescribeTableRequest requests table metadata message TDescribeTableRequest { // Data source instance to connect - TDataSourceInstance data_source_instance = 1; + NYql.TGenericDataSourceInstance data_source_instance = 1; // Table to describe string table = 2; // Rules for type mapping @@ -115,7 +115,7 @@ message TListSplitsResponse { // TODO: support JOIN, ORDER BY, GROUP BY message TSelect { // Data source instance to connect - TDataSourceInstance data_source_instance = 1; + NYql.TGenericDataSourceInstance data_source_instance = 1; // Describes what particularly to get from the data source message TWhat { @@ -200,7 +200,7 @@ message TReadSplitsRequest { // Data source instance to connect. // Deprecated field: server implementations must rely on // TDataSourceInstance provided in each TSelect. - TDataSourceInstance data_source_instance = 1 [deprecated = true]; + NYql.TGenericDataSourceInstance data_source_instance = 1 [deprecated = true]; // Splits that YQ engine would like to read. repeated TSplit splits = 2; diff --git a/ydb/library/yql/providers/generic/connector/api/service/protos/ya.make b/ydb/library/yql/providers/generic/connector/api/service/protos/ya.make index f9fe5ea3cadc..c05760b5e65c 100644 --- a/ydb/library/yql/providers/generic/connector/api/service/protos/ya.make +++ b/ydb/library/yql/providers/generic/connector/api/service/protos/ya.make @@ -1,7 +1,7 @@ PROTO_LIBRARY() PEERDIR( - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/public/api/protos ) diff --git a/ydb/library/yql/providers/generic/connector/api/ya.make b/ydb/library/yql/providers/generic/connector/api/ya.make index 254cdc978c14..b21d3fa468a5 100644 --- a/ydb/library/yql/providers/generic/connector/api/ya.make +++ b/ydb/library/yql/providers/generic/connector/api/ya.make @@ -1,4 +1,3 @@ RECURSE( - common # Keep this library pure of dependencies on YDB protofiles service # Use YDB protofiles only here ) diff --git a/ydb/library/yql/providers/generic/connector/libcpp/client.cpp b/ydb/library/yql/providers/generic/connector/libcpp/client.cpp index 91b10e58d15b..88d4e2636f6a 100644 --- a/ydb/library/yql/providers/generic/connector/libcpp/client.cpp +++ b/ydb/library/yql/providers/generic/connector/libcpp/client.cpp @@ -44,19 +44,25 @@ namespace NYql::NConnector { GrpcClient_ = std::make_unique(); // FIXME: is it OK to use single connection during the client lifetime? - GrpcConnection_ = GrpcClient_->CreateGRpcServiceConnection(GrpcConfig_); + GrpcConnection_ = GrpcClient_->CreateGRpcServiceConnection(GrpcConfig_, NYdbGrpc::TTcpKeepAliveSettings { + // TODO configure hardcoded values + .Enabled = true, + .Idle = 30, + .Count = 5, + .Interval = 10 + }); } - virtual TDescribeTableAsyncResult DescribeTable(const NApi::TDescribeTableRequest& request) override { - return UnaryCall(request, &NApi::Connector::Stub::AsyncDescribeTable); + virtual TDescribeTableAsyncResult DescribeTable(const NApi::TDescribeTableRequest& request, TDuration timeout = {}) override { + return UnaryCall(request, &NApi::Connector::Stub::AsyncDescribeTable, timeout); } - virtual TListSplitsStreamIteratorAsyncResult ListSplits(const NApi::TListSplitsRequest& request) override { - return ServerSideStreamingCall(request, &NApi::Connector::Stub::AsyncListSplits); + virtual TListSplitsStreamIteratorAsyncResult ListSplits(const NApi::TListSplitsRequest& request, TDuration timeout = {}) override { + return ServerSideStreamingCall(request, &NApi::Connector::Stub::AsyncListSplits, timeout); } - virtual TReadSplitsStreamIteratorAsyncResult ReadSplits(const NApi::TReadSplitsRequest& request) override { - return ServerSideStreamingCall(request, &NApi::Connector::Stub::AsyncReadSplits); + virtual TReadSplitsStreamIteratorAsyncResult ReadSplits(const NApi::TReadSplitsRequest& request, TDuration timeout = {}) override { + return ServerSideStreamingCall(request, &NApi::Connector::Stub::AsyncReadSplits, timeout); } ~TClientGRPC() { @@ -74,7 +80,7 @@ namespace NYql::NConnector { template TAsyncResult UnaryCall( const TRequest& request, - typename NYdbGrpc::TSimpleRequestProcessor::TAsyncRequest rpc) { + typename NYdbGrpc::TSimpleRequestProcessor::TAsyncRequest rpc, TDuration timeout = {}) { auto context = GrpcClient_->CreateContext(); if (!context) { throw yexception() << "Client is being shutted down"; @@ -89,7 +95,7 @@ namespace NYql::NConnector { std::move(request), std::move(callback), rpc, - {}, + { .Timeout = timeout }, context.get()); return promise.GetFuture(); @@ -98,7 +104,8 @@ namespace NYql::NConnector { template TIteratorAsyncResult> ServerSideStreamingCall( const TRequest& request, - TStreamRpc rpc) { + TStreamRpc rpc, + TDuration timeout = {}) { using TStreamProcessorPtr = typename NYdbGrpc::IStreamRequestReadProcessor::TPtr; using TStreamInitResult = std::pair; @@ -115,7 +122,7 @@ namespace NYql::NConnector { promise.SetValue({std::move(status), streamProcessor}); }, rpc, - {}, + { .Timeout = timeout }, context.get()); // TODO: async handling YQ-2513 diff --git a/ydb/library/yql/providers/generic/connector/libcpp/client.h b/ydb/library/yql/providers/generic/connector/libcpp/client.h index 448c9d6c5267..a06c18c1e907 100644 --- a/ydb/library/yql/providers/generic/connector/libcpp/client.h +++ b/ydb/library/yql/providers/generic/connector/libcpp/client.h @@ -89,9 +89,9 @@ namespace NYql::NConnector { public: using TPtr = std::shared_ptr; - virtual TDescribeTableAsyncResult DescribeTable(const NApi::TDescribeTableRequest& request) = 0; - virtual TListSplitsStreamIteratorAsyncResult ListSplits(const NApi::TListSplitsRequest& request) = 0; - virtual TReadSplitsStreamIteratorAsyncResult ReadSplits(const NApi::TReadSplitsRequest& request) = 0; + virtual TDescribeTableAsyncResult DescribeTable(const NApi::TDescribeTableRequest& request, TDuration timeout = {}) = 0; + virtual TListSplitsStreamIteratorAsyncResult ListSplits(const NApi::TListSplitsRequest& request, TDuration timeout = {}) = 0; + virtual TReadSplitsStreamIteratorAsyncResult ReadSplits(const NApi::TReadSplitsRequest& request, TDuration timeout = {}) = 0; virtual ~IClient() = default; }; diff --git a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.cpp b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.cpp index 8227b03b3535..ed92d42b65a7 100644 --- a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.cpp +++ b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.cpp @@ -32,7 +32,7 @@ namespace NYql::NConnector::NTest { void CreatePostgreSQLExternalDataSource( const std::shared_ptr& kikimr, const TString& dataSourceName, - NApi::EProtocol protocol, + NYql::EGenericProtocol protocol, const TString& host, int port, const TString& login, @@ -65,7 +65,7 @@ namespace NYql::NConnector::NTest { "login"_a = login, "password"_a = password, "use_tls"_a = useTls ? "TRUE" : "FALSE", - "protocol"_a = NApi::EProtocol_Name(protocol), + "protocol"_a = NYql::EGenericProtocol_Name(protocol), "source_type"_a = PG_SOURCE_TYPE, "database"_a = databaseName, "schema"_a = schema); @@ -76,7 +76,7 @@ namespace NYql::NConnector::NTest { void CreateClickHouseExternalDataSource( const std::shared_ptr& kikimr, const TString& dataSourceName, - NApi::EProtocol protocol, + NYql::EGenericProtocol protocol, const TString& clickHouseClusterId, const TString& login, const TString& password, @@ -110,7 +110,7 @@ namespace NYql::NConnector::NTest { "login"_a = login, "password"_a = password, "use_tls"_a = useTls ? "TRUE" : "FALSE", - "protocol"_a = NYql::NConnector::NApi::EProtocol_Name(protocol), + "protocol"_a = NYql::EGenericProtocol_Name(protocol), "service_account_id"_a = serviceAccountId, "service_account_id_signature"_a = serviceAccountIdSignature, "source_type"_a = ToString(NYql::EDatabaseType::ClickHouse), diff --git a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.h b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.h index 7f764611593a..cee03841f9ca 100644 --- a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.h +++ b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.h @@ -2,7 +2,7 @@ #include #include -#include +#include #include #include #include @@ -41,7 +41,7 @@ namespace NYql::NConnector::NTest { } #define DATA_SOURCE_INSTANCE_SUBBUILDER() \ - TBuilder& DataSourceInstance(const NApi::TDataSourceInstance& proto) { \ + TBuilder& DataSourceInstance(const NYql::TGenericDataSourceInstance& proto) { \ this->Result_->mutable_data_source_instance()->CopyFrom(proto); \ return static_cast(*this); \ } \ @@ -184,7 +184,7 @@ namespace NYql::NConnector::NTest { void CreatePostgreSQLExternalDataSource( const std::shared_ptr& kikimr, const TString& dataSourceName = DEFAULT_DATA_SOURCE_NAME, - NApi::EProtocol protocol = DEFAULT_PG_PROTOCOL, + NYql::EGenericProtocol protocol = DEFAULT_PG_PROTOCOL, const TString& host = DEFAULT_PG_HOST, int port = DEFAULT_PG_PORT, const TString& login = DEFAULT_LOGIN, @@ -196,7 +196,7 @@ namespace NYql::NConnector::NTest { void CreateClickHouseExternalDataSource( const std::shared_ptr& kikimr, const TString& dataSourceName = DEFAULT_DATA_SOURCE_NAME, - NApi::EProtocol protocol = DEFAULT_CH_PROTOCOL, + NYql::EGenericProtocol protocol = DEFAULT_CH_PROTOCOL, const TString& clickHouseClusterId = DEFAULT_CH_CLUSTER_ID, const TString& login = DEFAULT_LOGIN, const TString& password = DEFAULT_PASSWORD, @@ -225,11 +225,11 @@ namespace NYql::NConnector::NTest { // template - struct TBaseDataSourceInstanceBuilder: public TProtoBuilder { + struct TBaseDataSourceInstanceBuilder: public TProtoBuilder { using TBuilder = TDerived; - explicit TBaseDataSourceInstanceBuilder(NApi::TDataSourceInstance* result = nullptr, TParent* parent = nullptr) - : TProtoBuilder(result, parent) + explicit TBaseDataSourceInstanceBuilder(NYql::TGenericDataSourceInstance* result = nullptr, TParent* parent = nullptr) + : TProtoBuilder(result, parent) { } @@ -256,7 +256,7 @@ namespace NYql::NConnector::NTest { struct TPostgreSQLDataSourceInstanceBuilder: public TBaseDataSourceInstanceBuilder, TParent> { using TBase = TBaseDataSourceInstanceBuilder, TParent>; - explicit TPostgreSQLDataSourceInstanceBuilder(NApi::TDataSourceInstance* result = nullptr, TParent* parent = nullptr) + explicit TPostgreSQLDataSourceInstanceBuilder(NYql::TGenericDataSourceInstance* result = nullptr, TParent* parent = nullptr) : TBase(result, parent) { FillWithDefaults(); @@ -266,7 +266,7 @@ namespace NYql::NConnector::NTest { TBase::FillWithDefaults(); this->Host(DEFAULT_PG_HOST); this->Port(DEFAULT_PG_PORT); - this->Kind(NApi::EDataSourceKind::POSTGRESQL); + this->Kind(NYql::EGenericDataSourceKind::POSTGRESQL); this->Protocol(DEFAULT_PG_PROTOCOL); this->Schema(DEFAULT_PG_SCHEMA); } @@ -276,7 +276,7 @@ namespace NYql::NConnector::NTest { struct TClickHouseDataSourceInstanceBuilder: public TBaseDataSourceInstanceBuilder, TParent> { using TBase = TBaseDataSourceInstanceBuilder, TParent>; - explicit TClickHouseDataSourceInstanceBuilder(NApi::TDataSourceInstance* result = nullptr, TParent* parent = nullptr) + explicit TClickHouseDataSourceInstanceBuilder(NYql::TGenericDataSourceInstance* result = nullptr, TParent* parent = nullptr) : TBase(result, parent) { FillWithDefaults(); @@ -286,7 +286,7 @@ namespace NYql::NConnector::NTest { TBase::FillWithDefaults(); this->Host(DEFAULT_CH_HOST); this->Port(DEFAULT_CH_PORT); - this->Kind(NApi::EDataSourceKind::CLICKHOUSE); + this->Kind(NYql::EGenericDataSourceKind::CLICKHOUSE); this->Protocol(DEFAULT_CH_PROTOCOL); } }; @@ -295,7 +295,7 @@ namespace NYql::NConnector::NTest { struct TYdbDataSourceInstanceBuilder: public TBaseDataSourceInstanceBuilder, TParent> { using TBase = TBaseDataSourceInstanceBuilder, TParent>; - explicit TYdbDataSourceInstanceBuilder(NApi::TDataSourceInstance* result = nullptr, TParent* parent = nullptr) + explicit TYdbDataSourceInstanceBuilder(NYql::TGenericDataSourceInstance* result = nullptr, TParent* parent = nullptr) : TBase(result, parent) { FillWithDefaults(); @@ -305,7 +305,7 @@ namespace NYql::NConnector::NTest { TBase::FillWithDefaults(); this->Host(DEFAULT_YDB_HOST); this->Port(DEFAULT_YDB_PORT); - this->Kind(NApi::EDataSourceKind::YDB); + this->Kind(NYql::EGenericDataSourceKind::YDB); this->Protocol(DEFAULT_YDB_PROTOCOL); } }; @@ -681,6 +681,11 @@ namespace NYql::NConnector::NTest { return TListSplitsResultBuilder(ResponseResult_, this); } + auto& Status(const NYdbGrpc::TGrpcStatus& status) { + ResponseStatus_ = status; + return *this; + } + void FillWithDefaults() { Result(); } @@ -688,12 +693,13 @@ namespace NYql::NConnector::NTest { private: void SetExpectation() { EXPECT_CALL(*Mock_, ListSplitsImpl(ProtobufRequestMatcher(*Result_))) - .WillOnce(Return(TIteratorResult{NYdbGrpc::TGrpcStatus(), ResponseResult_})); + .WillOnce(Return(TIteratorResult{ResponseStatus_, ResponseResult_})); } private: TConnectorClientMock* Mock_ = nullptr; TListSplitsStreamIteratorMock::TPtr ResponseResult_ = std::make_shared(); + NYdbGrpc::TGrpcStatus ResponseStatus_ {}; }; template @@ -751,6 +757,11 @@ namespace NYql::NConnector::NTest { return TReadSplitsResultBuilder(ResponseResult_, this); } + auto& Status(const NYdbGrpc::TGrpcStatus& status) { + ResponseStatus_ = status; + return *this; + } + void FillWithDefaults() { Format(NApi::TReadSplitsRequest::ARROW_IPC_STREAMING); } @@ -758,12 +769,13 @@ namespace NYql::NConnector::NTest { private: void SetExpectation() { EXPECT_CALL(*Mock_, ReadSplitsImpl(ProtobufRequestMatcher(*Result_))) - .WillOnce(Return(TIteratorResult{NYdbGrpc::TGrpcStatus(), ResponseResult_})); + .WillOnce(Return(TIteratorResult{ResponseStatus_, ResponseResult_})); } private: TConnectorClientMock* Mock_ = nullptr; TReadSplitsStreamIteratorMock::TPtr ResponseResult_ = std::make_shared(); + NYdbGrpc::TGrpcStatus ResponseStatus_ {}; }; TDescribeTableExpectationBuilder ExpectDescribeTable() { @@ -778,7 +790,7 @@ namespace NYql::NConnector::NTest { return TReadSplitsExpectationBuilder(this); } - TDescribeTableAsyncResult DescribeTable(const NApi::TDescribeTableRequest& request) override { + TDescribeTableAsyncResult DescribeTable(const NApi::TDescribeTableRequest& request, TDuration = {}) override { Cerr << "Call DescribeTable.\n" << request.Utf8DebugString() << Endl; auto result = DescribeTableImpl(request); @@ -792,7 +804,7 @@ namespace NYql::NConnector::NTest { return NThreading::MakeFuture(std::move(result)); } - TListSplitsStreamIteratorAsyncResult ListSplits(const NApi::TListSplitsRequest& request) override { + TListSplitsStreamIteratorAsyncResult ListSplits(const NApi::TListSplitsRequest& request, TDuration = {}) override { Cerr << "Call ListSplits.\n" << request.Utf8DebugString() << Endl; auto result = ListSplitsImpl(request); @@ -801,7 +813,7 @@ namespace NYql::NConnector::NTest { return NThreading::MakeFuture(std::move(result)); } - TReadSplitsStreamIteratorAsyncResult ReadSplits(const NApi::TReadSplitsRequest& request) override { + TReadSplitsStreamIteratorAsyncResult ReadSplits(const NApi::TReadSplitsRequest& request, TDuration = {}) override { Cerr << "Call ReadSplits.\n" << request.Utf8DebugString() << Endl; auto result = ReadSplitsImpl(request); diff --git a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/defaults.h b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/defaults.h index bbca9127a4bd..3cc9bd395023 100644 --- a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/defaults.h +++ b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/defaults.h @@ -1,5 +1,5 @@ #pragma once -#include +#include #include #include @@ -15,19 +15,19 @@ namespace NYql::NConnector::NTest { constexpr int DEFAULT_PG_PORT = 5432; constexpr bool DEFAULT_USE_TLS = true; extern const TString PG_SOURCE_TYPE; - constexpr NApi::EProtocol DEFAULT_PG_PROTOCOL = NApi::EProtocol::NATIVE; + constexpr NYql::EGenericProtocol DEFAULT_PG_PROTOCOL = NYql::EGenericProtocol::NATIVE; extern const TString DEFAULT_PG_SCHEMA; extern const TString DEFAULT_CH_HOST; constexpr int DEFAULT_CH_PORT = 8443; extern const TString DEFAULT_CH_ENDPOINT; extern const TString DEFAULT_CH_CLUSTER_ID; - constexpr NApi::EProtocol DEFAULT_CH_PROTOCOL = NApi::EProtocol::HTTP; + constexpr NYql::EGenericProtocol DEFAULT_CH_PROTOCOL = NYql::EGenericProtocol::HTTP; extern const TString DEFAULT_CH_SERVICE_ACCOUNT_ID; extern const TString DEFAULT_CH_SERVICE_ACCOUNT_ID_SIGNATURE; extern const TString DEFAULT_YDB_HOST; constexpr int DEFAULT_YDB_PORT = 2136; extern const TString DEFAULT_YDB_ENDPOINT; - constexpr NApi::EProtocol DEFAULT_YDB_PROTOCOL = NApi::EProtocol::NATIVE; + constexpr NYql::EGenericProtocol DEFAULT_YDB_PROTOCOL = NYql::EGenericProtocol::NATIVE; } // namespace NYql::NConnector::NTest diff --git a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/ya.make b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/ya.make index 9bfc3b305e43..ef40fa52d639 100644 --- a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/ya.make +++ b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/ya.make @@ -14,7 +14,7 @@ PEERDIR( ydb/core/kqp/ut/common ydb/library/yql/providers/common/db_id_async_resolver yql/essentials/providers/common/structured_token - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/libcpp ) diff --git a/ydb/library/yql/providers/generic/connector/libcpp/ya.make b/ydb/library/yql/providers/generic/connector/libcpp/ya.make index 5b88a539d991..d31e47403a2e 100644 --- a/ydb/library/yql/providers/generic/connector/libcpp/ya.make +++ b/ydb/library/yql/providers/generic/connector/libcpp/ya.make @@ -14,7 +14,7 @@ PEERDIR( yql/essentials/ast ydb/library/yql/dq/actors/protos yql/essentials/providers/common/proto - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/api/service ydb/library/yql/providers/generic/connector/api/service/protos yql/essentials/public/issue diff --git a/ydb/library/yql/providers/generic/connector/tests/common_test_cases/base.py b/ydb/library/yql/providers/generic/connector/tests/common_test_cases/base.py index 6292077391c4..ce3f9e37306d 100644 --- a/ydb/library/yql/providers/generic/connector/tests/common_test_cases/base.py +++ b/ydb/library/yql/providers/generic/connector/tests/common_test_cases/base.py @@ -4,7 +4,7 @@ from typing import Dict import functools -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.library.yql.providers.generic.connector.api.service.protos.connector_pb2 import EDateTimeFormat from ydb.library.yql.providers.generic.connector.tests.utils.database import Database from ydb.library.yql.providers.generic.connector.tests.utils.settings import GenericSettings @@ -13,26 +13,26 @@ @dataclass class BaseTestCase: name_: str - data_source_kind: EDataSourceKind.ValueType + data_source_kind: EGenericDataSourceKind.ValueType pragmas: Dict[str, str] - protocol: EProtocol + protocol: EGenericProtocol @property def name(self) -> str: match self.data_source_kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: # ClickHouse has two kinds of network protocols: NATIVE and HTTP, # so we append protocol name to the test case name - return f'{self.name_}_{EProtocol.Name(self.protocol)}' - case EDataSourceKind.MS_SQL_SERVER: + return f'{self.name_}_{EGenericProtocol.Name(self.protocol)}' + case EGenericDataSourceKind.MS_SQL_SERVER: return self.name_ - case EDataSourceKind.MYSQL: + case EGenericDataSourceKind.MYSQL: return self.name_ - case EDataSourceKind.ORACLE: + case EGenericDataSourceKind.ORACLE: return self.name_ - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: return self.name_ - case EDataSourceKind.YDB: + case EGenericDataSourceKind.YDB: return self.name_ case _: raise Exception(f'invalid data source: {self.data_source_kind}') @@ -45,17 +45,17 @@ def database(self) -> Database: ''' # FIXME: do not hardcode databases here match self.data_source_kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: return Database(self.name, self.data_source_kind) - case EDataSourceKind.MS_SQL_SERVER: + case EGenericDataSourceKind.MS_SQL_SERVER: return Database("master", self.data_source_kind) - case EDataSourceKind.MYSQL: + case EGenericDataSourceKind.MYSQL: return Database("db", self.data_source_kind) - case EDataSourceKind.ORACLE: + case EGenericDataSourceKind.ORACLE: return Database(self.name, self.data_source_kind) - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: return Database(self.name, self.data_source_kind) - case EDataSourceKind.YDB: + case EGenericDataSourceKind.YDB: return Database("local", self.data_source_kind) @functools.cached_property @@ -65,17 +65,17 @@ def table_name(self) -> str: so we provide a random table name instead where necessary. ''' match self.data_source_kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: return self.name_ # without protocol - case EDataSourceKind.MS_SQL_SERVER: + case EGenericDataSourceKind.MS_SQL_SERVER: return self.name - case EDataSourceKind.MYSQL: + case EGenericDataSourceKind.MYSQL: return self.name - case EDataSourceKind.ORACLE: + case EGenericDataSourceKind.ORACLE: return self.name - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: return 't' + make_random_string(8) - case EDataSourceKind.YDB: + case EGenericDataSourceKind.YDB: return self.name case _: raise Exception(f'invalid data source: {self.data_source_kind}') @@ -90,34 +90,34 @@ def pragmas_sql_string(self) -> str: @property def generic_settings(self) -> GenericSettings: match self.data_source_kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: return GenericSettings( date_time_format=EDateTimeFormat.YQL_FORMAT, clickhouse_clusters=[ - GenericSettings.ClickHouseCluster(database=self.database.name, protocol=EProtocol.NATIVE) + GenericSettings.ClickHouseCluster(database=self.database.name, protocol=EGenericProtocol.NATIVE) ], ) - case EDataSourceKind.MS_SQL_SERVER: + case EGenericDataSourceKind.MS_SQL_SERVER: return GenericSettings( date_time_format=EDateTimeFormat.YQL_FORMAT, ms_sql_server_clusters=[GenericSettings.MsSQLServerCluster(database=self.database.name)], ) - case EDataSourceKind.MYSQL: + case EGenericDataSourceKind.MYSQL: return GenericSettings( date_time_format=EDateTimeFormat.YQL_FORMAT, mysql_clusters=[GenericSettings.MySQLCluster(database=self.database.name)], ) - case EDataSourceKind.ORACLE: + case EGenericDataSourceKind.ORACLE: return GenericSettings( date_time_format=EDateTimeFormat.YQL_FORMAT, oracle_clusters=[GenericSettings.OracleCluster(database=self.database.name, service_name=None)], ) - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: return GenericSettings( date_time_format=EDateTimeFormat.YQL_FORMAT, postgresql_clusters=[GenericSettings.PostgreSQLCluster(database=self.database.name, schema=None)], ) - case EDataSourceKind.YDB: + case EGenericDataSourceKind.YDB: return GenericSettings( date_time_format=EDateTimeFormat.YQL_FORMAT, ydb_clusters=[GenericSettings.YdbCluster(database=self.database.name)], diff --git a/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_missing_database.py b/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_missing_database.py index da33e6401036..c94735e59db4 100644 --- a/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_missing_database.py +++ b/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_missing_database.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.library.yql.providers.generic.connector.tests.common_test_cases.base import BaseTestCase from ydb.library.yql.providers.generic.connector.tests.utils.settings import GenericSettings @@ -16,14 +16,14 @@ def generic_settings(self) -> GenericSettings: gs = super().generic_settings # Overload setting for MySQL database - if self.data_source_kind == EDataSourceKind.MYSQL: + if self.data_source_kind == EGenericDataSourceKind.MYSQL: for cluster in gs.mysql_clusters: cluster.database = "missing_database" for cluster in gs.oracle_clusters: if self.service_name is not None: cluster.service_name = self.service_name - if self.data_source_kind == EDataSourceKind.MS_SQL_SERVER: + if self.data_source_kind == EGenericDataSourceKind.MS_SQL_SERVER: for cluster in gs.ms_sql_server_clusters: cluster.database = "missing_database" @@ -36,12 +36,12 @@ class Factory: def __init__(self, ss: Settings): self.ss = ss - def make_test_cases(self, data_source_kind: EDataSourceKind) -> List[TestCase]: + def make_test_cases(self, data_source_kind: EGenericDataSourceKind) -> List[TestCase]: return [ TestCase( name_="missing_database", data_source_kind=data_source_kind, - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), service_name=self.ss.oracle.service_name if self.ss.oracle is not None else None, ) diff --git a/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_missing_table.py b/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_missing_table.py index 3d24f0e3f9ce..9426d0d27dee 100644 --- a/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_missing_table.py +++ b/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_missing_table.py @@ -3,7 +3,7 @@ from typing import Sequence from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.library.yql.providers.generic.connector.tests.common_test_cases.base import BaseTestCase from ydb.library.yql.providers.generic.connector.tests.utils.settings import GenericSettings @@ -31,7 +31,7 @@ class Factory: def __init__(self, ss: Settings): self.ss = ss - def make_test_cases(self, data_source_kind: EDataSourceKind) -> List[TestCase]: + def make_test_cases(self, data_source_kind: EGenericDataSourceKind) -> List[TestCase]: test_cases = [] test_case_name = 'missing_table' @@ -39,7 +39,7 @@ def make_test_cases(self, data_source_kind: EDataSourceKind) -> List[TestCase]: test_case = TestCase( name_=test_case_name, data_source_kind=data_source_kind, - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), service_name=self.ss.oracle.service_name if self.ss.oracle is not None else None, ) diff --git a/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_positive_common.py b/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_positive_common.py index 99da19b19aae..f9cc66c2b797 100644 --- a/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_positive_common.py +++ b/ydb/library/yql/providers/generic/connector/tests/common_test_cases/select_positive_common.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, replace from typing import Sequence, Optional -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.public.api.protos.ydb_value_pb2 import Type from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings @@ -95,11 +95,11 @@ def _column_selection(self) -> Sequence[TestCase]: SelectWhat.asterisk(column_list=schema.columns), data_in, ( - EDataSourceKind.CLICKHOUSE, - EDataSourceKind.POSTGRESQL, - EDataSourceKind.YDB, - EDataSourceKind.MYSQL, - EDataSourceKind.MS_SQL_SERVER, + EGenericDataSourceKind.CLICKHOUSE, + EGenericDataSourceKind.POSTGRESQL, + EGenericDataSourceKind.YDB, + EGenericDataSourceKind.MYSQL, + EGenericDataSourceKind.MS_SQL_SERVER, ), ), # SELECT COL1 FROM table @@ -110,11 +110,11 @@ def _column_selection(self) -> Sequence[TestCase]: [10], ], ( - EDataSourceKind.CLICKHOUSE, + EGenericDataSourceKind.CLICKHOUSE, # NOTE: YQ-2264: doesn't work for PostgreSQL because of implicit cast to lowercase (COL1 -> col1) - EDataSourceKind.YDB, - EDataSourceKind.MYSQL, - EDataSourceKind.MS_SQL_SERVER, + EGenericDataSourceKind.YDB, + EGenericDataSourceKind.MYSQL, + EGenericDataSourceKind.MS_SQL_SERVER, ), ), # SELECT col1 FROM table @@ -124,7 +124,7 @@ def _column_selection(self) -> Sequence[TestCase]: [1], [10], ], - (EDataSourceKind.POSTGRESQL,), # works because of implicit cast to lowercase (COL1 -> col1) + (EGenericDataSourceKind.POSTGRESQL,), # works because of implicit cast to lowercase (COL1 -> col1) ), # SELECT col2 FROM table ( @@ -134,11 +134,11 @@ def _column_selection(self) -> Sequence[TestCase]: [20], ], ( - EDataSourceKind.CLICKHOUSE, - EDataSourceKind.POSTGRESQL, - EDataSourceKind.YDB, - EDataSourceKind.MYSQL, - EDataSourceKind.MS_SQL_SERVER, + EGenericDataSourceKind.CLICKHOUSE, + EGenericDataSourceKind.POSTGRESQL, + EGenericDataSourceKind.YDB, + EGenericDataSourceKind.MYSQL, + EGenericDataSourceKind.MS_SQL_SERVER, ), ), # SELECT col2, COL1 FROM table @@ -149,11 +149,11 @@ def _column_selection(self) -> Sequence[TestCase]: [20, 10], ], ( - EDataSourceKind.CLICKHOUSE, + EGenericDataSourceKind.CLICKHOUSE, # NOTE: YQ-2264: doesn't work for PostgreSQL because of implicit cast to lowercase (COL1 -> col1) - EDataSourceKind.YDB, - EDataSourceKind.MYSQL, - EDataSourceKind.MS_SQL_SERVER, + EGenericDataSourceKind.YDB, + EGenericDataSourceKind.MYSQL, + EGenericDataSourceKind.MS_SQL_SERVER, ), ), # SELECT col2, col1 FROM table @@ -163,7 +163,7 @@ def _column_selection(self) -> Sequence[TestCase]: [2, 1], [20, 10], ], - (EDataSourceKind.POSTGRESQL,), # works because of implicit cast to lowercase (COL1 -> col1) + (EGenericDataSourceKind.POSTGRESQL,), # works because of implicit cast to lowercase (COL1 -> col1) ), # Simple math computation: # SELECT COL1 + col2 AS col3 FROM table @@ -174,11 +174,11 @@ def _column_selection(self) -> Sequence[TestCase]: [30], ], ( - EDataSourceKind.CLICKHOUSE, + EGenericDataSourceKind.CLICKHOUSE, # NOTE: YQ-2264: doesn't work for PostgreSQL because of implicit cast to lowercase (COL1 -> col1) - EDataSourceKind.YDB, - EDataSourceKind.MYSQL, - EDataSourceKind.MS_SQL_SERVER, + EGenericDataSourceKind.YDB, + EGenericDataSourceKind.MYSQL, + EGenericDataSourceKind.MS_SQL_SERVER, ), ), # Select the same column multiple times with different aliases @@ -196,11 +196,11 @@ def _column_selection(self) -> Sequence[TestCase]: [20, 20, 20, 20, 20], ], ( - EDataSourceKind.CLICKHOUSE, - EDataSourceKind.POSTGRESQL, - EDataSourceKind.YDB, - EDataSourceKind.MYSQL, - EDataSourceKind.MS_SQL_SERVER, + EGenericDataSourceKind.CLICKHOUSE, + EGenericDataSourceKind.POSTGRESQL, + EGenericDataSourceKind.YDB, + EGenericDataSourceKind.MYSQL, + EGenericDataSourceKind.MS_SQL_SERVER, ), ), ) @@ -215,7 +215,7 @@ def _column_selection(self) -> Sequence[TestCase]: test_case = TestCase( data_in=data_in, data_source_kind=data_source_kind, - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, select_what=select_what, select_where=None, schema=schema, @@ -260,8 +260,8 @@ def _large_table(self) -> Sequence[TestCase]: ) data_source_kinds = ( - EDataSourceKind.CLICKHOUSE, - EDataSourceKind.POSTGRESQL, + EGenericDataSourceKind.CLICKHOUSE, + EGenericDataSourceKind.POSTGRESQL, ) test_case_name = 'large' @@ -269,11 +269,11 @@ def _large_table(self) -> Sequence[TestCase]: for data_source_kind in data_source_kinds: match data_source_kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: tc = TestCase( name_=test_case_name, data_source_kind=data_source_kind, - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, data_in=None, data_out_=[[999999]], # We put 1M of rows in the large table select_what=SelectWhat(SelectWhat.Item(name='MAX(col_00_int32)', kind='expr')), @@ -282,18 +282,18 @@ def _large_table(self) -> Sequence[TestCase]: pragmas=dict(), ) - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: # Assuming that request will look something like: # `SELECT * FROM table WHERE id = (SELECT MAX(id) FROM table)` # We expect last line to be the answer data_in = generate_table_data(schema=schema, bytes_soft_limit=table_size) data_out = [data_in[-1]] - data_source_kinds = [EDataSourceKind.CLICKHOUSE, EDataSourceKind.POSTGRESQL] + data_source_kinds = [EGenericDataSourceKind.CLICKHOUSE, EGenericDataSourceKind.POSTGRESQL] tc = TestCase( name_=test_case_name, data_source_kind=data_source_kind, - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, data_in=data_in, data_out_=data_out, select_what=SelectWhat.asterisk(schema.columns), @@ -311,31 +311,31 @@ def _large_table(self) -> Sequence[TestCase]: return test_cases - def make_test_cases(self, data_source_kind: EDataSourceKind) -> Sequence[TestCase]: + def make_test_cases(self, data_source_kind: EGenericDataSourceKind) -> Sequence[TestCase]: protocols = { - EDataSourceKind.CLICKHOUSE: [EProtocol.NATIVE, EProtocol.HTTP], - EDataSourceKind.POSTGRESQL: [EProtocol.NATIVE], - EDataSourceKind.YDB: [EProtocol.NATIVE], - EDataSourceKind.MYSQL: [EProtocol.NATIVE], - EDataSourceKind.MS_SQL_SERVER: [EProtocol.NATIVE], + EGenericDataSourceKind.CLICKHOUSE: [EGenericProtocol.NATIVE, EGenericProtocol.HTTP], + EGenericDataSourceKind.POSTGRESQL: [EGenericProtocol.NATIVE], + EGenericDataSourceKind.YDB: [EGenericProtocol.NATIVE], + EGenericDataSourceKind.MYSQL: [EGenericProtocol.NATIVE], + EGenericDataSourceKind.MS_SQL_SERVER: [EGenericProtocol.NATIVE], } base_test_cases = None if data_source_kind in [ - EDataSourceKind.YDB, - EDataSourceKind.MYSQL, - EDataSourceKind.MS_SQL_SERVER, + EGenericDataSourceKind.YDB, + EGenericDataSourceKind.MYSQL, + EGenericDataSourceKind.MS_SQL_SERVER, ]: base_test_cases = self._column_selection() - elif data_source_kind in [EDataSourceKind.CLICKHOUSE, EDataSourceKind.POSTGRESQL]: + elif data_source_kind in [EGenericDataSourceKind.CLICKHOUSE, EGenericDataSourceKind.POSTGRESQL]: base_test_cases = list( itertools.chain( self._column_selection(), self._large_table(), ) ) - elif data_source_kind == EDataSourceKind.ORACLE: + elif data_source_kind == EGenericDataSourceKind.ORACLE: raise 'Common test cases are not supported by Oracle due to the lack of Int32 columns' else: raise f'Unexpected data source kind: {data_source_kind}' diff --git a/ydb/library/yql/providers/generic/connector/tests/common_test_cases/ya.make b/ydb/library/yql/providers/generic/connector/tests/common_test_cases/ya.make index 143966c5cd4c..db6c24486d69 100644 --- a/ydb/library/yql/providers/generic/connector/tests/common_test_cases/ya.make +++ b/ydb/library/yql/providers/generic/connector/tests/common_test_cases/ya.make @@ -13,7 +13,7 @@ PY_SRCS( ) PEERDIR( - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/tests/utils ydb/library/yql/providers/generic/connector/api/service/protos ydb/public/api/protos diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/collection.py b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/collection.py index b317efd3b5ef..48bc1c1af596 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/collection.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/collection.py @@ -1,6 +1,6 @@ from typing import Sequence, Mapping -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind # test cases import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_database as select_missing_database @@ -17,10 +17,12 @@ class Collection(object): def __init__(self, ss: Settings): self._test_cases = { - 'select_missing_database': select_missing_database.Factory(ss).make_test_cases(EDataSourceKind.CLICKHOUSE), - 'select_missing_table': select_missing_table.Factory(ss).make_test_cases(EDataSourceKind.CLICKHOUSE), + 'select_missing_database': select_missing_database.Factory(ss).make_test_cases( + EGenericDataSourceKind.CLICKHOUSE + ), + 'select_missing_table': select_missing_table.Factory(ss).make_test_cases(EGenericDataSourceKind.CLICKHOUSE), 'select_positive': select_positive.Factory().make_test_cases() - + select_positive_common.Factory(ss).make_test_cases(EDataSourceKind.CLICKHOUSE), + + select_positive_common.Factory(ss).make_test_cases(EGenericDataSourceKind.CLICKHOUSE), 'select_datetime': select_datetime.Factory().make_test_cases(), } diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/conftest.py b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/conftest.py index b2869f8d1f5f..2788b5c4017b 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/conftest.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/conftest.py @@ -3,7 +3,7 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings docker_compose_dir: Final = pathlib.Path("ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse") @@ -11,4 +11,6 @@ @pytest.fixture def settings() -> Settings: - return Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EDataSourceKind.CLICKHOUSE]) + return Settings.from_env( + docker_compose_dir=docker_compose_dir, data_source_kinds=[EGenericDataSourceKind.CLICKHOUSE] + ) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/docker-compose.yml b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/docker-compose.yml index 5a0c45f76f75..cd94b3249893 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/docker-compose.yml +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/docker-compose.yml @@ -15,7 +15,7 @@ services: - ./init:/docker-entrypoint-initdb.d fq-connector-go: container_name: fq-tests-ch-fq-connector-go - image: ghcr.io/ydb-platform/fq-connector-go:v0.5.12-rc.2@sha256:84bb0b19f16f354b8a9ef7a020ee80f3ba7dc28db92f7007f235241153025b8a + image: ghcr.io/ydb-platform/fq-connector-go:v0.6.0-rc.1@sha256:4f74bd11e696218053f48622d1d065567d103906c187b3a24ea4e56f886c6c60 ports: - 2130 volumes: diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/select_datetime.py b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/select_datetime.py index 1057054a4b92..5e9dcbcc5e52 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/select_datetime.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/select_datetime.py @@ -2,7 +2,7 @@ import datetime from typing import Sequence -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.library.yql.providers.generic.connector.api.service.protos.connector_pb2 import EDateTimeFormat from ydb.public.api.protos.ydb_value_pb2 import Type @@ -120,8 +120,8 @@ def _make_test_yql(self) -> TestCase: data_out_=data_out, select_what=SelectWhat.asterisk(schema.columns), select_where=None, - data_source_kind=EDataSourceKind.CLICKHOUSE, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.CLICKHOUSE, + protocol=EGenericProtocol.NATIVE, schema=schema, pragmas=dict(), check_output_schema=True, @@ -177,12 +177,12 @@ def _make_test_string(self) -> TestCase: return TestCase( name_=test_case_name, date_time_format=EDateTimeFormat.STRING_FORMAT, - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, data_in=None, data_out_=data_out, select_what=SelectWhat.asterisk(schema.columns), select_where=None, - data_source_kind=EDataSourceKind.CLICKHOUSE, + data_source_kind=EGenericDataSourceKind.CLICKHOUSE, schema=schema, pragmas=dict(), check_output_schema=True, diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/select_positive.py b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/select_positive.py index da75133e3432..244a4e93add7 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/select_positive.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/select_positive.py @@ -3,7 +3,7 @@ from dataclasses import replace from typing import Sequence, Final -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.public.api.protos.ydb_value_pb2 import Type import ydb.library.yql.providers.generic.connector.tests.utils.types.clickhouse as clickhouse @@ -169,8 +169,8 @@ def _primitive_types_non_nullable(self) -> Sequence[TestCase]: datetime.datetime(2023, 3, 21, 11, 21, 31, 456000), ], ], - data_source_kind=EDataSourceKind.CLICKHOUSE, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.CLICKHOUSE, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), check_output_schema=True, ) @@ -271,8 +271,8 @@ def _primitive_types_nullable(self) -> Sequence[TestCase]: datetime.datetime(2023, 3, 21, 11, 21, 31, 456000), ], ], - data_source_kind=EDataSourceKind.CLICKHOUSE, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.CLICKHOUSE, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), check_output_schema=True, ) @@ -315,8 +315,8 @@ def _constant(self) -> Sequence[TestCase]: 42, ], ], - data_source_kind=EDataSourceKind.CLICKHOUSE, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.CLICKHOUSE, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), ) @@ -350,8 +350,8 @@ def _counts(self) -> Sequence[TestCase]: 4, ], ], - protocol=EProtocol.NATIVE, - data_source_kind=EDataSourceKind.CLICKHOUSE, + protocol=EGenericProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.CLICKHOUSE, pragmas=dict(), check_output_schema=False, # because the aggregate's value has other type ) @@ -378,7 +378,7 @@ def _pushdown(self) -> TestCase: ['one'], ] - data_source_kind = EDataSourceKind.CLICKHOUSE + data_source_kind = EGenericDataSourceKind.CLICKHOUSE test_case_name = 'pushdown' return [ @@ -390,7 +390,7 @@ def _pushdown(self) -> TestCase: select_what=SelectWhat(SelectWhat.Item(name='col_01_string')), select_where=SelectWhere('col_00_int32 = 1'), data_source_kind=data_source_kind, - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, schema=schema, # TODO: implement schema checkswhen selecting only one column check_output_schema=False, @@ -398,7 +398,7 @@ def _pushdown(self) -> TestCase: ] def make_test_cases(self) -> Sequence[TestCase]: - protocols = [EProtocol.NATIVE, EProtocol.HTTP] + protocols = [EGenericProtocol.NATIVE, EGenericProtocol.HTTP] base_test_cases = list( itertools.chain( diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/test.py b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/test.py index ab690a43f2d7..4380bf270bf8 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/test.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/test.py @@ -1,6 +1,6 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings from ydb.library.yql.providers.generic.connector.tests.utils.run.runners import runner_types, configure_runner import ydb.library.yql.providers.generic.connector.tests.utils.scenario.clickhouse as scenario @@ -14,7 +14,7 @@ import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_positive_common as select_positive_common one_time_waiter = OneTimeWaiter( - data_source_kind=EDataSourceKind.CLICKHOUSE, + data_source_kind=EGenericDataSourceKind.CLICKHOUSE, docker_compose_file_path=str(docker_compose_dir / 'docker-compose.yml'), expected_tables=[ "column_selection_A_b_C_d_E", @@ -36,7 +36,7 @@ # Global collection of test cases dependent on environment tc_collection = Collection( - Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EDataSourceKind.CLICKHOUSE]) + Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EGenericDataSourceKind.CLICKHOUSE]) ) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/ya.make b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/ya.make index dde9f62942d9..3a8f114c35f5 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/ya.make +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/clickhouse/ya.make @@ -65,7 +65,7 @@ TEST_SRCS( PEERDIR( contrib/python/pytest - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/api/service/protos ydb/library/yql/providers/generic/connector/tests/common_test_cases ydb/library/yql/providers/generic/connector/tests/utils diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/collection.py b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/collection.py index 19868608b43f..1643f2355483 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/collection.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/collection.py @@ -1,6 +1,6 @@ from typing import Sequence, Mapping -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_database as select_missing_database import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_table as select_missing_table import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_positive_common as select_positive_common @@ -18,11 +18,13 @@ class Collection(object): def __init__(self, ss: Settings): self._test_cases = { 'select_missing_database': select_missing_database.Factory(ss).make_test_cases( - EDataSourceKind.MS_SQL_SERVER + EGenericDataSourceKind.MS_SQL_SERVER + ), + 'select_missing_table': select_missing_table.Factory(ss).make_test_cases( + EGenericDataSourceKind.MS_SQL_SERVER ), - 'select_missing_table': select_missing_table.Factory(ss).make_test_cases(EDataSourceKind.MS_SQL_SERVER), 'select_positive': select_positive.Factory().make_test_cases() - + select_positive_common.Factory(ss).make_test_cases(EDataSourceKind.MS_SQL_SERVER), + + select_positive_common.Factory(ss).make_test_cases(EGenericDataSourceKind.MS_SQL_SERVER), 'select_datetime': select_datetime.Factory().make_test_cases(), } diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/conftest.py b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/conftest.py index c229347d688b..90b079354d5f 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/conftest.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/conftest.py @@ -4,7 +4,7 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings @@ -13,4 +13,6 @@ @pytest.fixture def settings() -> Settings: - return Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EDataSourceKind.MS_SQL_SERVER]) + return Settings.from_env( + docker_compose_dir=docker_compose_dir, data_source_kinds=[EGenericDataSourceKind.MS_SQL_SERVER] + ) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/docker-compose.yml b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/docker-compose.yml index aa7213e1f539..9a78ab8851c9 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/docker-compose.yml +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/docker-compose.yml @@ -1,7 +1,7 @@ services: fq-connector-go: container_name: fq-tests-mssql-fq-connector-go - image: ghcr.io/ydb-platform/fq-connector-go:v0.5.12-rc.2@sha256:84bb0b19f16f354b8a9ef7a020ee80f3ba7dc28db92f7007f235241153025b8a + image: ghcr.io/ydb-platform/fq-connector-go:v0.6.0-rc.1@sha256:4f74bd11e696218053f48622d1d065567d103906c187b3a24ea4e56f886c6c60 ports: - 2130 volumes: diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/select_datetime.py b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/select_datetime.py index 21db4ef79157..125f4ebad100 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/select_datetime.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/select_datetime.py @@ -2,7 +2,7 @@ import datetime from typing import Sequence -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.library.yql.providers.generic.connector.api.service.protos.connector_pb2 import EDateTimeFormat from ydb.public.api.protos.ydb_value_pb2 import Type @@ -96,8 +96,8 @@ def _make_test_yql(self) -> TestCase: data_out_=data_out, select_what=SelectWhat.asterisk(self._schema.columns), select_where=None, - data_source_kind=EDataSourceKind.MS_SQL_SERVER, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.MS_SQL_SERVER, + protocol=EGenericProtocol.NATIVE, schema=self._schema, pragmas=dict(), ) @@ -134,8 +134,8 @@ def _make_test_string(self) -> TestCase: data_out_=data_out, select_what=SelectWhat.asterisk(self._schema.columns), select_where=None, - data_source_kind=EDataSourceKind.MS_SQL_SERVER, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.MS_SQL_SERVER, + protocol=EGenericProtocol.NATIVE, schema=self._schema, pragmas=dict(), ) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/select_positive.py b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/select_positive.py index ce729a234ab7..d00811dfdfd7 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/select_positive.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/select_positive.py @@ -2,7 +2,7 @@ import itertools from typing import Sequence -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.public.api.protos.ydb_value_pb2 import Type import ydb.library.yql.providers.generic.connector.tests.utils.types.ms_sql_server as ms_sql_server @@ -212,8 +212,8 @@ def _primitive_types(self) -> Sequence[TestCase]: datetime.datetime(2023, 3, 21, 11, 21, 31, 0), ], ], - data_source_kind=EDataSourceKind.MS_SQL_SERVER, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.MS_SQL_SERVER, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), check_output_schema=True, ) @@ -254,8 +254,8 @@ def _constant(self) -> Sequence[TestCase]: 42, ], ], - data_source_kind=EDataSourceKind.MS_SQL_SERVER, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.MS_SQL_SERVER, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), ) @@ -289,8 +289,8 @@ def _count_rows(self) -> Sequence[TestCase]: 3, ], ], - data_source_kind=EDataSourceKind.MS_SQL_SERVER, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.MS_SQL_SERVER, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), ) @@ -324,11 +324,11 @@ def _pushdown(self) -> TestCase: name_=test_case_name, data_in=None, data_out_=[[4, None, None]], - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, pragmas=dict({'generic.UsePredicatePushdown': 'true'}), select_what=SelectWhat.asterisk(schema.columns), select_where=SelectWhere('col_00_id = 4'), - data_source_kind=EDataSourceKind.MS_SQL_SERVER, + data_source_kind=EGenericDataSourceKind.MS_SQL_SERVER, schema=schema, ), TestCase( @@ -337,11 +337,11 @@ def _pushdown(self) -> TestCase: data_out_=[ ['b'], ], - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, pragmas=dict({'generic.UsePredicatePushdown': 'true'}), select_what=SelectWhat(SelectWhat.Item(name='col_02_text')), select_where=SelectWhere('col_00_id = col_01_integer'), - data_source_kind=EDataSourceKind.MS_SQL_SERVER, + data_source_kind=EGenericDataSourceKind.MS_SQL_SERVER, schema=schema, ), ] diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/test.py b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/test.py index 57ac2806b79b..d313915cf4a2 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/test.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/test.py @@ -1,6 +1,6 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.one_time_waiter import OneTimeWaiter from ydb.library.yql.providers.generic.connector.tests.utils.log import make_logger from ydb.library.yql.providers.generic.connector.tests.utils.run.runners import runner_types, configure_runner @@ -19,11 +19,11 @@ # Global collection of test cases dependent on environment tc_collection = Collection( - Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EDataSourceKind.MS_SQL_SERVER]) + Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EGenericDataSourceKind.MS_SQL_SERVER]) ) one_time_waiter = OneTimeWaiter( - data_source_kind=EDataSourceKind.MS_SQL_SERVER, + data_source_kind=EGenericDataSourceKind.MS_SQL_SERVER, docker_compose_file_path=str(docker_compose_dir / 'docker-compose.yml'), expected_tables=[ 'column_selection_A_b_C_d_E', diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/ya.make b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/ya.make index b826eab12718..18f425f9f31a 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/ya.make +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ms_sql_server/ya.make @@ -65,7 +65,7 @@ TEST_SRCS( PEERDIR( contrib/python/pytest - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/tests/common_test_cases ydb/library/yql/providers/generic/connector/tests/utils ydb/library/yql/providers/generic/connector/tests/utils/run diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/collection.py b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/collection.py index 2f6d5334368f..97c4381dce20 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/collection.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/collection.py @@ -1,6 +1,6 @@ from typing import Sequence, Mapping -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_database as select_missing_database import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_table as select_missing_table import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_positive_common as select_positive_common @@ -17,10 +17,12 @@ class Collection(object): def __init__(self, ss: Settings): self._test_cases = { - 'select_missing_database': select_missing_database.Factory(ss).make_test_cases(EDataSourceKind.MYSQL), - 'select_missing_table': select_missing_table.Factory(ss).make_test_cases(EDataSourceKind.MYSQL), + 'select_missing_database': select_missing_database.Factory(ss).make_test_cases( + EGenericDataSourceKind.MYSQL + ), + 'select_missing_table': select_missing_table.Factory(ss).make_test_cases(EGenericDataSourceKind.MYSQL), 'select_positive': select_positive.Factory().make_test_cases() - + select_positive_common.Factory(ss).make_test_cases(EDataSourceKind.MYSQL), + + select_positive_common.Factory(ss).make_test_cases(EGenericDataSourceKind.MYSQL), 'select_datetime': select_datetime.Factory().make_test_cases(), } diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/conftest.py b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/conftest.py index 0724b709dbde..4fed6c744d2b 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/conftest.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/conftest.py @@ -4,7 +4,7 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings @@ -13,4 +13,4 @@ @pytest.fixture def settings() -> Settings: - return Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EDataSourceKind.MYSQL]) + return Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EGenericDataSourceKind.MYSQL]) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/docker-compose.yml b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/docker-compose.yml index 55e01b2cd08c..aecb80c7b28b 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/docker-compose.yml +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/docker-compose.yml @@ -1,7 +1,7 @@ services: fq-connector-go: container_name: fq-tests-mysql-fq-connector-go - image: ghcr.io/ydb-platform/fq-connector-go:v0.5.12-rc.2@sha256:84bb0b19f16f354b8a9ef7a020ee80f3ba7dc28db92f7007f235241153025b8a + image: ghcr.io/ydb-platform/fq-connector-go:v0.6.0-rc.1@sha256:4f74bd11e696218053f48622d1d065567d103906c187b3a24ea4e56f886c6c60 ports: - 2130 volumes: diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/select_datetime.py b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/select_datetime.py index 8cad456b187a..1c210b276c40 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/select_datetime.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/select_datetime.py @@ -2,7 +2,7 @@ import datetime from typing import Sequence -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.library.yql.providers.generic.connector.api.service.protos.connector_pb2 import EDateTimeFormat from ydb.public.api.protos.ydb_value_pb2 import Type @@ -90,8 +90,8 @@ def _make_test_yql(self) -> TestCase: data_out_=data_out, select_what=SelectWhat.asterisk(self._schema.columns), select_where=None, - data_source_kind=EDataSourceKind.MYSQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.MYSQL, + protocol=EGenericProtocol.NATIVE, schema=self._schema, pragmas=dict(), ) @@ -121,8 +121,8 @@ def _make_test_string(self) -> TestCase: data_out_=data_out, select_what=SelectWhat.asterisk(self._schema.columns), select_where=None, - data_source_kind=EDataSourceKind.MYSQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.MYSQL, + protocol=EGenericProtocol.NATIVE, schema=self._schema, pragmas=dict(), ) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/select_positive.py b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/select_positive.py index 96278c1ae29e..a08c823927d1 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/select_positive.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/select_positive.py @@ -2,7 +2,7 @@ import itertools from typing import Sequence -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.public.api.protos.ydb_value_pb2 import Type import ydb.library.yql.providers.generic.connector.tests.utils.types.mysql as mysql @@ -285,8 +285,8 @@ def _primitive_types(self) -> Sequence[TestCase]: '{ "TODO" : "unicode" }', ], ], - data_source_kind=EDataSourceKind.MYSQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.MYSQL, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), check_output_schema=True, ) @@ -327,8 +327,8 @@ def _constant(self) -> Sequence[TestCase]: 42, ], ], - data_source_kind=EDataSourceKind.MYSQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.MYSQL, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), ) @@ -362,8 +362,8 @@ def _count_rows(self) -> Sequence[TestCase]: 3, ], ], - data_source_kind=EDataSourceKind.MYSQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.MYSQL, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), ) @@ -397,11 +397,11 @@ def _pushdown(self) -> TestCase: name_=test_case_name, data_in=None, data_out_=[[4, None, None]], - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, pragmas=dict({'generic.UsePredicatePushdown': 'true'}), select_what=SelectWhat.asterisk(schema.columns), select_where=SelectWhere('col_00_id = 4'), - data_source_kind=EDataSourceKind.MYSQL, + data_source_kind=EGenericDataSourceKind.MYSQL, schema=schema, ), TestCase( @@ -410,11 +410,11 @@ def _pushdown(self) -> TestCase: data_out_=[ ['b'], ], - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, pragmas=dict({'generic.UsePredicatePushdown': 'true'}), select_what=SelectWhat(SelectWhat.Item(name='col_02_text')), select_where=SelectWhere('col_00_id = col_01_integer'), - data_source_kind=EDataSourceKind.MYSQL, + data_source_kind=EGenericDataSourceKind.MYSQL, schema=schema, ), ] @@ -444,10 +444,10 @@ def _json(self) -> TestCase: [None], [None], ], - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, select_what=SelectWhat(SelectWhat.Item(name='JSON_QUERY(col_01_json, "$.friends[0]")', kind='expr')), select_where=None, - data_source_kind=EDataSourceKind.MYSQL, + data_source_kind=EGenericDataSourceKind.MYSQL, pragmas=dict(), schema=schema, ), diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/test.py b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/test.py index 78203acf53fe..b1340face7bb 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/test.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/test.py @@ -1,6 +1,6 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.one_time_waiter import OneTimeWaiter from ydb.library.yql.providers.generic.connector.tests.utils.log import make_logger from ydb.library.yql.providers.generic.connector.tests.utils.run.runners import runner_types, configure_runner @@ -17,7 +17,7 @@ LOGGER = make_logger(__name__) one_time_waiter = OneTimeWaiter( - data_source_kind=EDataSourceKind.MYSQL, + data_source_kind=EGenericDataSourceKind.MYSQL, docker_compose_file_path=str(docker_compose_dir / 'docker-compose.yml'), expected_tables=[ 'column_selection_A_b_C_d_E', @@ -39,7 +39,7 @@ # Global collection of test cases dependent on environment tc_collection = Collection( - Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EDataSourceKind.MYSQL]) + Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EGenericDataSourceKind.MYSQL]) ) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/ya.make b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/ya.make index 9fc3764f7d36..70acbf909e38 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/ya.make +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/mysql/ya.make @@ -65,7 +65,7 @@ TEST_SRCS( PEERDIR( contrib/python/pytest - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/tests/common_test_cases ydb/library/yql/providers/generic/connector/tests/utils ydb/library/yql/providers/generic/connector/tests/utils/run diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/collection.py b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/collection.py index cf014c90af84..6f59eb08ed29 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/collection.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/collection.py @@ -1,6 +1,6 @@ from typing import Sequence, Mapping -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind # import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_database as select_missing_database import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_table as select_missing_table @@ -19,10 +19,10 @@ class Collection(object): def __init__(self, ss: Settings): self._test_cases = { - # 'select_missing_database': select_missing_database.Factory(ss).make_test_cases(EDataSourceKind.ORACLE), # TODO YQ-3413 - 'select_missing_table': select_missing_table.Factory(ss).make_test_cases(EDataSourceKind.ORACLE), + # 'select_missing_database': select_missing_database.Factory(ss).make_test_cases(EGenericDataSourceKind.ORACLE), # TODO YQ-3413 + 'select_missing_table': select_missing_table.Factory(ss).make_test_cases(EGenericDataSourceKind.ORACLE), 'select_positive': select_positive_with_service_name.Factory(ss).make_test_cases(), - # + select_positive_common.Factory(ss).make_test_cases(EDataSourceKind.ORACLE), # TODO does not work because of cast to uppercase and lack of Int32 column type + # + select_positive_common.Factory(ss).make_test_cases(EGenericDataSourceKind.ORACLE), # TODO does not work because of cast to uppercase and lack of Int32 column type 'select_datetime': select_datetime_with_service_name.Factory(ss).make_test_cases(), } diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/conftest.py b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/conftest.py index be427e6e3003..405a314b4fb9 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/conftest.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/conftest.py @@ -4,7 +4,7 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings @@ -13,4 +13,4 @@ @pytest.fixture def settings() -> Settings: - return Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EDataSourceKind.ORACLE]) + return Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EGenericDataSourceKind.ORACLE]) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/docker-compose.yml b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/docker-compose.yml index 0ad174526862..a663e34d6d21 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/docker-compose.yml +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/docker-compose.yml @@ -1,7 +1,7 @@ services: fq-connector-go: container_name: fq-tests-oracle-fq-connector-go - image: ghcr.io/ydb-platform/fq-connector-go:v0.5.12-rc.2@sha256:84bb0b19f16f354b8a9ef7a020ee80f3ba7dc28db92f7007f235241153025b8a + image: ghcr.io/ydb-platform/fq-connector-go:v0.6.0-rc.1@sha256:4f74bd11e696218053f48622d1d065567d103906c187b3a24ea4e56f886c6c60 ports: - 2130 volumes: diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/select_datetime_with_service_name.py b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/select_datetime_with_service_name.py index cf0cca4ec14f..084bc5462aaa 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/select_datetime_with_service_name.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/select_datetime_with_service_name.py @@ -3,7 +3,7 @@ from typing import Sequence from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.library.yql.providers.generic.connector.api.service.protos.connector_pb2 import EDateTimeFormat from ydb.public.api.protos.ydb_value_pb2 import Type @@ -91,8 +91,8 @@ def _make_test_yql(self) -> TestCase: data_out_=data_out, select_what=SelectWhat.asterisk(self._schema.columns), select_where=None, - data_source_kind=EDataSourceKind.ORACLE, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.ORACLE, + protocol=EGenericProtocol.NATIVE, schema=self._schema, pragmas=dict(), service_name=self.ss.oracle.service_name, @@ -121,8 +121,8 @@ def _make_test_string(self) -> TestCase: data_out_=data_out, select_what=SelectWhat.asterisk(self._schema.columns), select_where=None, - data_source_kind=EDataSourceKind.ORACLE, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.ORACLE, + protocol=EGenericProtocol.NATIVE, schema=self._schema, pragmas=dict(), service_name=self.ss.oracle.service_name, diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/select_positive_with_service_name.py b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/select_positive_with_service_name.py index 54b4a9be3383..dfdfca783946 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/select_positive_with_service_name.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/select_positive_with_service_name.py @@ -4,7 +4,7 @@ from typing import Sequence, Optional from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.public.api.protos.ydb_value_pb2 import Type from ydb.library.yql.providers.generic.connector.tests.utils.settings import GenericSettings @@ -250,8 +250,8 @@ def _primitive_types(self) -> Sequence[TestCase]: '{ "TODO" : "unicode" }', ], ], - data_source_kind=EDataSourceKind.ORACLE, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.ORACLE, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), check_output_schema=True, service_name=self.ss.oracle.service_name, @@ -296,8 +296,8 @@ def _longraw(self) -> Sequence[TestCase]: None, ], ], - data_source_kind=EDataSourceKind.ORACLE, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.ORACLE, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), check_output_schema=True, service_name=self.ss.oracle.service_name, @@ -342,8 +342,8 @@ def _long_table(self) -> Sequence[TestCase]: None, ], ], - data_source_kind=EDataSourceKind.ORACLE, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.ORACLE, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), check_output_schema=True, service_name=self.ss.oracle.service_name, @@ -385,8 +385,8 @@ def _constant(self) -> Sequence[TestCase]: 42, ], ], - data_source_kind=EDataSourceKind.ORACLE, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.ORACLE, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), service_name=self.ss.oracle.service_name, ) @@ -421,8 +421,8 @@ def _count_rows(self) -> Sequence[TestCase]: 3, ], ], - data_source_kind=EDataSourceKind.ORACLE, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.ORACLE, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), service_name=self.ss.oracle.service_name, ) @@ -457,11 +457,11 @@ def _pushdown(self) -> TestCase: name_=test_case_name, data_in=None, data_out_=[[4, None, None]], - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, pragmas=dict({'generic.UsePredicatePushdown': 'true'}), select_what=SelectWhat.asterisk(schema.columns), select_where=SelectWhere('COL_00_ID = 4'), - data_source_kind=EDataSourceKind.ORACLE, + data_source_kind=EGenericDataSourceKind.ORACLE, schema=schema, service_name=self.ss.oracle.service_name, ), @@ -471,11 +471,11 @@ def _pushdown(self) -> TestCase: data_out_=[ ['b'], ], - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, pragmas=dict({'generic.UsePredicatePushdown': 'true'}), select_what=SelectWhat(SelectWhat.Item(name='COL_02_TEXT')), select_where=SelectWhere('COL_00_ID = COL_01_INTEGER'), - data_source_kind=EDataSourceKind.ORACLE, + data_source_kind=EGenericDataSourceKind.ORACLE, schema=schema, service_name=self.ss.oracle.service_name, ), @@ -506,10 +506,10 @@ def _json(self) -> TestCase: [None], [None], ], - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, select_what=SelectWhat(SelectWhat.Item(name='JSON_QUERY(COL_01_JSON, "$.friends[0]")', kind='expr')), select_where=None, - data_source_kind=EDataSourceKind.ORACLE, + data_source_kind=EGenericDataSourceKind.ORACLE, pragmas=dict(), schema=schema, service_name=self.ss.oracle.service_name, diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/test.py b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/test.py index 8c57b8bce1a9..9ee77d1dc4a5 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/test.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/test.py @@ -1,6 +1,6 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.one_time_waiter import OneTimeWaiter from ydb.library.yql.providers.generic.connector.tests.utils.log import make_logger from ydb.library.yql.providers.generic.connector.tests.utils.run.runners import runner_types, configure_runner @@ -20,11 +20,11 @@ # Global collection of test cases dependent on environment tc_collection = Collection( - Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EDataSourceKind.ORACLE]) + Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EGenericDataSourceKind.ORACLE]) ) one_time_waiter = OneTimeWaiter( - data_source_kind=EDataSourceKind.ORACLE, + data_source_kind=EGenericDataSourceKind.ORACLE, docker_compose_file_path=str(docker_compose_dir / 'docker-compose.yml'), expected_tables=[ 'column_selection_A_b_C_d_E', diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/ya.make b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/ya.make index e34e0d53220a..11a157241fdf 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/ya.make +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/oracle/ya.make @@ -65,7 +65,7 @@ TEST_SRCS( PEERDIR( contrib/python/pytest - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/tests/common_test_cases ydb/library/yql/providers/generic/connector/tests/utils ydb/library/yql/providers/generic/connector/tests/utils/run diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/collection.py b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/collection.py index 3f077ca089ff..353871b7d3c7 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/collection.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/collection.py @@ -1,6 +1,6 @@ from typing import Sequence, Mapping -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_database as select_missing_database import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_table as select_missing_table import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_positive_common as select_positive_common @@ -16,10 +16,12 @@ class Collection(object): def __init__(self, ss: Settings): self._test_cases = { - 'select_missing_database': select_missing_database.Factory(ss).make_test_cases(EDataSourceKind.POSTGRESQL), - 'select_missing_table': select_missing_table.Factory(ss).make_test_cases(EDataSourceKind.POSTGRESQL), + 'select_missing_database': select_missing_database.Factory(ss).make_test_cases( + EGenericDataSourceKind.POSTGRESQL + ), + 'select_missing_table': select_missing_table.Factory(ss).make_test_cases(EGenericDataSourceKind.POSTGRESQL), 'select_positive': select_positive.Factory().make_test_cases() - + select_positive_common.Factory(ss).make_test_cases(EDataSourceKind.POSTGRESQL), + + select_positive_common.Factory(ss).make_test_cases(EGenericDataSourceKind.POSTGRESQL), 'select_positive_with_schema': select_positive_with_schema.Factory().make_test_cases(), 'select_datetime': select_datetime.Factory().make_test_cases(), } diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/conftest.py b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/conftest.py index e2e36087059d..7d3f0c096603 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/conftest.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/conftest.py @@ -4,7 +4,7 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings from ydb.library.yql.providers.generic.connector.tests.utils.clients.postgresql import Client @@ -14,7 +14,9 @@ @pytest.fixture def settings() -> Settings: - return Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EDataSourceKind.POSTGRESQL]) + return Settings.from_env( + docker_compose_dir=docker_compose_dir, data_source_kinds=[EGenericDataSourceKind.POSTGRESQL] + ) @pytest.fixture diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/docker-compose.yml b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/docker-compose.yml index 8a17be178649..42e8a7908c50 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/docker-compose.yml +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/docker-compose.yml @@ -1,7 +1,7 @@ services: fq-connector-go: container_name: fq-tests-pg-fq-connector-go - image: ghcr.io/ydb-platform/fq-connector-go:v0.5.12-rc.2@sha256:84bb0b19f16f354b8a9ef7a020ee80f3ba7dc28db92f7007f235241153025b8a + image: ghcr.io/ydb-platform/fq-connector-go:v0.6.0-rc.1@sha256:4f74bd11e696218053f48622d1d065567d103906c187b3a24ea4e56f886c6c60 ports: - 2130 volumes: diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_datetime.py b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_datetime.py index 7c396bd5013a..522bf9cecc92 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_datetime.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_datetime.py @@ -2,7 +2,7 @@ import datetime from typing import Sequence -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.library.yql.providers.generic.connector.api.service.protos.connector_pb2 import EDateTimeFormat from ydb.public.api.protos.ydb_value_pb2 import Type @@ -92,8 +92,8 @@ def _make_test_yql_postgresql(self) -> TestCase: data_out_=data_out, select_what=SelectWhat.asterisk(schema.columns), select_where=None, - data_source_kind=EDataSourceKind.POSTGRESQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.POSTGRESQL, + protocol=EGenericProtocol.NATIVE, schema=schema, pragmas=dict(), ) @@ -150,8 +150,8 @@ def _make_test_string_postgresql(self) -> TestCase: data_out_=data_out, select_what=SelectWhat.asterisk(schema.columns), select_where=None, - data_source_kind=EDataSourceKind.POSTGRESQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.POSTGRESQL, + protocol=EGenericProtocol.NATIVE, schema=schema, pragmas=dict(), ) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_positive.py b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_positive.py index c42ad0152b43..8db8a254cd29 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_positive.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_positive.py @@ -2,7 +2,7 @@ import itertools from typing import Sequence -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.public.api.protos.ydb_value_pb2 import Type import ydb.library.yql.providers.generic.connector.tests.utils.types.postgresql as postgresql @@ -254,8 +254,8 @@ def _primitive_types(self) -> Sequence[TestCase]: ], ], data_out_=None, - data_source_kind=EDataSourceKind.POSTGRESQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.POSTGRESQL, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), check_output_schema=True, ) @@ -295,8 +295,8 @@ def _upper_case_column(self) -> Sequence[TestCase]: 3, ], ], - data_source_kind=EDataSourceKind.POSTGRESQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.POSTGRESQL, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), ) @@ -346,8 +346,8 @@ def _constant(self) -> Sequence[TestCase]: 42, ], ], - data_source_kind=EDataSourceKind.POSTGRESQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.POSTGRESQL, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), ) @@ -391,8 +391,8 @@ def _count(self) -> Sequence[TestCase]: 3, ], ], - data_source_kind=EDataSourceKind.POSTGRESQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.POSTGRESQL, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), ) @@ -438,7 +438,7 @@ def _pushdown(self) -> TestCase: ['two'], ] - data_source_kind = EDataSourceKind.POSTGRESQL + data_source_kind = EGenericDataSourceKind.POSTGRESQL test_case_name = 'pushdown' @@ -447,7 +447,7 @@ def _pushdown(self) -> TestCase: name_=test_case_name, data_in=data_in, data_out_=data_out_1, - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, pragmas=dict({'generic.UsePredicatePushdown': 'true'}), select_what=SelectWhat(SelectWhat.Item(name='col_string')), select_where=SelectWhere('col_int32 = 1'), @@ -458,7 +458,7 @@ def _pushdown(self) -> TestCase: name_=test_case_name, data_in=data_in, data_out_=data_out_2, - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, pragmas=dict({'generic.UsePredicatePushdown': 'true'}), select_what=SelectWhat(SelectWhat.Item(name='col_string')), select_where=SelectWhere('col_int32 = col_int64'), @@ -490,7 +490,7 @@ def _json(self) -> TestCase: [None], ] - data_source_kind = EDataSourceKind.POSTGRESQL + data_source_kind = EGenericDataSourceKind.POSTGRESQL test_case_name = 'json' @@ -499,7 +499,7 @@ def _json(self) -> TestCase: name_=test_case_name, data_in=data_in, data_out_=data_out_1, - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, select_what=SelectWhat(SelectWhat.Item(name='JSON_QUERY(col_json, "$.friends[0]")', kind='expr')), select_where=None, data_source_kind=data_source_kind, diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_positive_with_schema.py b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_positive_with_schema.py index b2cb9bf457cf..338dc14159ff 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_positive_with_schema.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/select_positive_with_schema.py @@ -5,7 +5,7 @@ from string import ascii_lowercase, digits -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.public.api.protos.ydb_value_pb2 import Type from ydb.library.yql.providers.generic.connector.tests.utils.settings import GenericSettings @@ -67,8 +67,8 @@ def _select_with_pg_schema(self) -> Sequence[TestCase]: test_case = TestCase( name_=test_case_name, data_in=[[1, 2], [10, 20]], - data_source_kind=EDataSourceKind.POSTGRESQL, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.POSTGRESQL, + protocol=EGenericProtocol.NATIVE, select_what=select_what, schema=schema, data_out_=[[1, 2], [10, 20]], diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/test.py b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/test.py index a46e79609c34..c94f36b5797e 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/test.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/test.py @@ -1,6 +1,6 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings from ydb.library.yql.providers.generic.connector.tests.utils.run.runners import runner_types, configure_runner from ydb.library.yql.providers.generic.connector.tests.utils.clients.postgresql import Client @@ -16,7 +16,7 @@ # Global collection of test cases dependent on environment tc_collection = Collection( - Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EDataSourceKind.POSTGRESQL]) + Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EGenericDataSourceKind.POSTGRESQL]) ) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/ya.make b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/ya.make index 0a9ed43f265c..679a0dcd0518 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/ya.make +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/postgresql/ya.make @@ -66,7 +66,7 @@ TEST_SRCS( PEERDIR( contrib/python/pytest - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/tests/common_test_cases ydb/library/yql/providers/generic/connector/tests/utils ydb/library/yql/providers/generic/connector/tests/utils/run diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/collection.py b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/collection.py index cb0da9fe7b23..23a8e64c1745 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/collection.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/collection.py @@ -1,6 +1,6 @@ from typing import Sequence, Mapping -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_database as select_missing_database import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_table as select_missing_table @@ -15,9 +15,9 @@ class Collection(object): def __init__(self, ss: Settings): self._test_cases = { - 'select_missing_database': select_missing_database.Factory(ss).make_test_cases(EDataSourceKind.YDB), - 'select_missing_table': select_missing_table.Factory(ss).make_test_cases(EDataSourceKind.YDB), - 'select_positive': select_positive_common.Factory(ss).make_test_cases(EDataSourceKind.YDB) + 'select_missing_database': select_missing_database.Factory(ss).make_test_cases(EGenericDataSourceKind.YDB), + 'select_missing_table': select_missing_table.Factory(ss).make_test_cases(EGenericDataSourceKind.YDB), + 'select_positive': select_positive_common.Factory(ss).make_test_cases(EGenericDataSourceKind.YDB) + select_positive.Factory().make_test_cases(), } diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/docker-compose.yml b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/docker-compose.yml index 72680763e599..a2091af3ed92 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/docker-compose.yml +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/docker-compose.yml @@ -8,7 +8,7 @@ services: dnsmasq; /opt/ydb/bin/fq-connector-go server -c /opt/ydb/cfg/fq-connector-go.yaml" container_name: fq-tests-ydb-fq-connector-go - image: ghcr.io/ydb-platform/fq-connector-go:v0.5.12-rc.2@sha256:84bb0b19f16f354b8a9ef7a020ee80f3ba7dc28db92f7007f235241153025b8a + image: ghcr.io/ydb-platform/fq-connector-go:v0.6.0-rc.1@sha256:4f74bd11e696218053f48622d1d065567d103906c187b3a24ea4e56f886c6c60 ports: - 2130 volumes: diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/select_positive.py b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/select_positive.py index 7d23cecce8d4..af096e1de4a8 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/select_positive.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/select_positive.py @@ -2,7 +2,7 @@ import itertools from typing import Sequence -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.public.api.protos.ydb_value_pb2 import Type import ydb.library.yql.providers.generic.connector.tests.utils.types.ydb as types_ydb @@ -167,8 +167,8 @@ def _primitive_types(self) -> Sequence[TestCase]: '{ "TODO" : "unicode" }', ], ], - data_source_kind=EDataSourceKind.YDB, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.YDB, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), check_output_schema=True, ) @@ -341,8 +341,8 @@ def _optional_types(self) -> Sequence[TestCase]: None, ], ], - data_source_kind=EDataSourceKind.YDB, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.YDB, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), check_output_schema=True, ) @@ -383,8 +383,8 @@ def _constant(self) -> Sequence[TestCase]: 42, ], ], - data_source_kind=EDataSourceKind.YDB, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.YDB, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), ) @@ -418,8 +418,8 @@ def _count(self) -> Sequence[TestCase]: 4, ], ], - protocol=EProtocol.NATIVE, - data_source_kind=EDataSourceKind.YDB, + protocol=EGenericProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.YDB, pragmas=dict(), ) @@ -453,8 +453,8 @@ def _pushdown(self) -> TestCase: pragmas=dict({'generic.UsePredicatePushdown': 'true'}), select_what=SelectWhat(SelectWhat.Item(name='col_01_string')), select_where=SelectWhere('col_00_id = 1'), - data_source_kind=EDataSourceKind.YDB, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.YDB, + protocol=EGenericProtocol.NATIVE, schema=schema, ) ] @@ -492,8 +492,8 @@ def _unsupported_types(self) -> Sequence[TestCase]: 2, ], ], - data_source_kind=EDataSourceKind.YDB, - protocol=EProtocol.NATIVE, + data_source_kind=EGenericDataSourceKind.YDB, + protocol=EGenericProtocol.NATIVE, pragmas=dict(), check_output_schema=True, ) @@ -539,7 +539,7 @@ def _json(self) -> Sequence[TestCase]: [None], ] - data_source_kind = EDataSourceKind.YDB + data_source_kind = EGenericDataSourceKind.YDB test_case_name = 'json' @@ -548,7 +548,7 @@ def _json(self) -> Sequence[TestCase]: name_=test_case_name, data_in=data_in, data_out_=data_out_1, - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, select_what=SelectWhat(SelectWhat.Item(name='JSON_QUERY(col_01_json, "$.friends[0]")', kind='expr')), select_where=None, data_source_kind=data_source_kind, @@ -585,9 +585,9 @@ def _json_document(self) -> TestCase: data_out_=data_out, select_where=None, select_what=SelectWhat(SelectWhat.Item(name='col_01_data')), - data_source_kind=EDataSourceKind.YDB, + data_source_kind=EGenericDataSourceKind.YDB, pragmas=dict(), - protocol=EProtocol.NATIVE, + protocol=EGenericProtocol.NATIVE, schema=schema, # check_output_schema=True, ) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/test.py b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/test.py index 843b885b6906..05aea2d55286 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/test.py +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/test.py @@ -1,6 +1,6 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings from ydb.library.yql.providers.generic.connector.tests.utils.run.runners import runner_types, configure_runner from ydb.library.yql.providers.generic.connector.tests.utils.one_time_waiter import OneTimeWaiter @@ -14,7 +14,7 @@ from collection import Collection one_time_waiter = OneTimeWaiter( - data_source_kind=EDataSourceKind.YDB, + data_source_kind=EGenericDataSourceKind.YDB, docker_compose_file_path=str(docker_compose_dir / 'docker-compose.yml'), expected_tables=[ "column_selection_A_b_C_d_E", @@ -35,7 +35,7 @@ ], ) -settings = Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EDataSourceKind.YDB]) +settings = Settings.from_env(docker_compose_dir=docker_compose_dir, data_source_kinds=[EGenericDataSourceKind.YDB]) tc_collection = Collection(settings) diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/ya.make b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/ya.make index e2e749f8a532..d9fbc295a269 100644 --- a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/ya.make +++ b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/ya.make @@ -64,7 +64,7 @@ TEST_SRCS( PEERDIR( contrib/python/pytest - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/tests/common_test_cases ydb/library/yql/providers/generic/connector/tests/utils ydb/library/yql/providers/generic/connector/tests/utils/run diff --git a/ydb/library/yql/providers/generic/connector/tests/join/conftest.py b/ydb/library/yql/providers/generic/connector/tests/join/conftest.py index a2d90087ea6e..a896d70f9660 100644 --- a/ydb/library/yql/providers/generic/connector/tests/join/conftest.py +++ b/ydb/library/yql/providers/generic/connector/tests/join/conftest.py @@ -4,7 +4,7 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.clients.postgresql import Client as PostgreSQLClient from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings @@ -16,7 +16,7 @@ def settings() -> Settings: return Settings.from_env( docker_compose_dir=docker_compose_dir, - data_source_kinds=[EDataSourceKind.POSTGRESQL, EDataSourceKind.CLICKHOUSE], + data_source_kinds=[EGenericDataSourceKind.POSTGRESQL, EGenericDataSourceKind.CLICKHOUSE], ) diff --git a/ydb/library/yql/providers/generic/connector/tests/join/docker-compose.yml b/ydb/library/yql/providers/generic/connector/tests/join/docker-compose.yml index 7ab666f7df14..c9c920712828 100644 --- a/ydb/library/yql/providers/generic/connector/tests/join/docker-compose.yml +++ b/ydb/library/yql/providers/generic/connector/tests/join/docker-compose.yml @@ -15,7 +15,7 @@ services: - ./init/clickhouse:/docker-entrypoint-initdb.d fq-connector-go: container_name: fq-tests-join-fq-connector-go - image: ghcr.io/ydb-platform/fq-connector-go:v0.5.12-rc.2@sha256:84bb0b19f16f354b8a9ef7a020ee80f3ba7dc28db92f7007f235241153025b8a + image: ghcr.io/ydb-platform/fq-connector-go:v0.6.0-rc.1@sha256:4f74bd11e696218053f48622d1d065567d103906c187b3a24ea4e56f886c6c60 ports: - 2130 volumes: diff --git a/ydb/library/yql/providers/generic/connector/tests/join/scenario.py b/ydb/library/yql/providers/generic/connector/tests/join/scenario.py index 156ceb8cabd6..32929e2e997f 100644 --- a/ydb/library/yql/providers/generic/connector/tests/join/scenario.py +++ b/ydb/library/yql/providers/generic/connector/tests/join/scenario.py @@ -1,4 +1,4 @@ -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.comparator import assert_data_outs_equal from ydb.library.yql.providers.generic.connector.tests.utils.log import make_logger from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings @@ -22,10 +22,10 @@ def join( # prepare tables for data_source in test_case.data_sources: match data_source.kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: # do nothing as tables are initialized via init scripts continue - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: postgresql_scenario.prepare_table( test_name=test_name, client=postgresql_client, diff --git a/ydb/library/yql/providers/generic/connector/tests/join/test.py b/ydb/library/yql/providers/generic/connector/tests/join/test.py index 7623c10922ff..c163982860a4 100644 --- a/ydb/library/yql/providers/generic/connector/tests/join/test.py +++ b/ydb/library/yql/providers/generic/connector/tests/join/test.py @@ -1,6 +1,6 @@ import pytest -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings from ydb.library.yql.providers.generic.connector.tests.utils.run.runners import runner_types, configure_runner @@ -13,7 +13,7 @@ tc_collection = Collection( Settings.from_env( docker_compose_dir=conftest.docker_compose_dir, - data_source_kinds=[EDataSourceKind.CLICKHOUSE, EDataSourceKind.POSTGRESQL], + data_source_kinds=[EGenericDataSourceKind.CLICKHOUSE, EGenericDataSourceKind.POSTGRESQL], ) ) diff --git a/ydb/library/yql/providers/generic/connector/tests/join/test_case.py b/ydb/library/yql/providers/generic/connector/tests/join/test_case.py index ca0d0a0474ed..038b62665908 100644 --- a/ydb/library/yql/providers/generic/connector/tests/join/test_case.py +++ b/ydb/library/yql/providers/generic/connector/tests/join/test_case.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from typing import Sequence, Final -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.library.yql.providers.generic.connector.api.service.protos.connector_pb2 import EDateTimeFormat from ydb.public.api.protos.ydb_value_pb2 import Type @@ -36,7 +36,7 @@ def select_what(self) -> SelectWhat: class DataSource: database: Database table: Table - kind: EDataSourceKind.ValueType = EDataSourceKind.DATA_SOURCE_KIND_UNSPECIFIED + kind: EGenericDataSourceKind.ValueType = EGenericDataSourceKind.DATA_SOURCE_KIND_UNSPECIFIED @property def alias(self) -> str: @@ -70,12 +70,14 @@ def generic_settings(self) -> GenericSettings: for data_source in self.data_sources: match data_source.kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: clickhouse_clusters.add( - GenericSettings.ClickHouseCluster(database=data_source.database.name, protocol=EProtocol.NATIVE) + GenericSettings.ClickHouseCluster( + database=data_source.database.name, protocol=EGenericProtocol.NATIVE + ) ) - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: postgresql_clusters.add( GenericSettings.PostgreSQLCluster(database=data_source.database.name, schema=None) ) @@ -170,9 +172,9 @@ def __make_simple_test_cases(self) -> Sequence[TestCase]: data_out = list(map(lambda x: list(itertools.chain(*x)), zip(*(t.data_in for t in tables)))) - data_sources: Sequence[EDataSourceKind] = ( - EDataSourceKind.CLICKHOUSE, - EDataSourceKind.POSTGRESQL, + data_sources: Sequence[EGenericDataSourceKind] = ( + EGenericDataSourceKind.CLICKHOUSE, + EGenericDataSourceKind.POSTGRESQL, ) # For each test case we create a unique set of datasources; @@ -247,13 +249,13 @@ def __make_inner_join_test_case(self) -> Sequence[TestCase]: test_case_data_sources = [ DataSource( - kind=EDataSourceKind.CLICKHOUSE, - database=Database(kind=EDataSourceKind.CLICKHOUSE, name=test_case_name), + kind=EGenericDataSourceKind.CLICKHOUSE, + database=Database(kind=EGenericDataSourceKind.CLICKHOUSE, name=test_case_name), table=ch_table, ), DataSource( - kind=EDataSourceKind.POSTGRESQL, - database=Database(kind=EDataSourceKind.POSTGRESQL, name=test_case_name), + kind=EGenericDataSourceKind.POSTGRESQL, + database=Database(kind=EGenericDataSourceKind.POSTGRESQL, name=test_case_name), table=pg_table, ), ] diff --git a/ydb/library/yql/providers/generic/connector/tests/join/ya.make b/ydb/library/yql/providers/generic/connector/tests/join/ya.make index 1bb206bc3cab..cb8b8a62514e 100644 --- a/ydb/library/yql/providers/generic/connector/tests/join/ya.make +++ b/ydb/library/yql/providers/generic/connector/tests/join/ya.make @@ -65,7 +65,7 @@ TEST_SRCS( PEERDIR( contrib/python/pytest - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/api/service/protos ydb/library/yql/providers/generic/connector/tests/common_test_cases ydb/library/yql/providers/generic/connector/tests/utils diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/data_source_kind.py b/ydb/library/yql/providers/generic/connector/tests/utils/data_source_kind.py index a186fdd87ab0..2098dbb4f113 100644 --- a/ydb/library/yql/providers/generic/connector/tests/utils/data_source_kind.py +++ b/ydb/library/yql/providers/generic/connector/tests/utils/data_source_kind.py @@ -1,11 +1,11 @@ -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind -def data_source_kind_alias(kind: EDataSourceKind.ValueType) -> str: +def data_source_kind_alias(kind: EGenericDataSourceKind.ValueType) -> str: match (kind): - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: return "ch" - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: return "pg" case _: raise Exception(f'invalid data source: {kind}') diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/database.py b/ydb/library/yql/providers/generic/connector/tests/utils/database.py index 82b2f50139a7..ccbea9570f50 100644 --- a/ydb/library/yql/providers/generic/connector/tests/utils/database.py +++ b/ydb/library/yql/providers/generic/connector/tests/utils/database.py @@ -1,34 +1,34 @@ from dataclasses import dataclass -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind @dataclass class Database: name: str - def __init__(self, name: str, kind: EDataSourceKind.ValueType): + def __init__(self, name: str, kind: EGenericDataSourceKind.ValueType): self.kind = kind match kind: - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: # PostgreSQL implicitly converts all identifiers to lowercase, # so we'd better make it first on our own self.name = name[:63].lower() - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: # We use preinitialized database when working with ClickHouse. self.name = "db" - case EDataSourceKind.MS_SQL_SERVER: + case EGenericDataSourceKind.MS_SQL_SERVER: # For this kind of database this name is provided by the external logic self.name = name - case EDataSourceKind.MYSQL: + case EGenericDataSourceKind.MYSQL: # For this kind of database this name is provided by the external logic self.name = name - case EDataSourceKind.ORACLE: + case EGenericDataSourceKind.ORACLE: # Oracle is not sensitive for identifiers until they are inclosed in quota marks, # therefore, we'd better use uppercase for ease of testing self.name = name[:127].upper() # TODO: is it needed? max length of Oracle table name is 128 bytes/chars - case EDataSourceKind.YDB: + case EGenericDataSourceKind.YDB: # We use preinitialized database when working with YDB. self.name = "local" case _: @@ -36,14 +36,14 @@ def __init__(self, name: str, kind: EDataSourceKind.ValueType): def exists(self) -> str: match self.kind: - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: return f"SELECT 1 FROM pg_database WHERE datname = '{self.name}'" case _: raise Exception(f'invalid data source: {self.kind}') def create(self) -> str: match self.kind: - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: return f"CREATE DATABASE {self.name}" case _: raise Exception(f'invalid data source: {self.kind}') @@ -53,34 +53,34 @@ def sql_table_name(self, table_name: str) -> str: def missing_database_msg(self) -> str: match self.kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: return f"Database {self.name} doesn't exist" - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: return f'database "{self.name}" does not exist' - case EDataSourceKind.YDB: + case EGenericDataSourceKind.YDB: raise Exception("Fix me first in YQ-3315") - case EDataSourceKind.MS_SQL_SERVER: + case EGenericDataSourceKind.MS_SQL_SERVER: return 'Cannot open database' - case EDataSourceKind.MYSQL: + case EGenericDataSourceKind.MYSQL: return 'Unknown database' - case EDataSourceKind.ORACLE: + case EGenericDataSourceKind.ORACLE: raise Exception("Fix me first in YQ-3413") case _: raise Exception(f'invalid data source: {self.kind}') def missing_table_msg(self) -> str: match self.kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: return 'table does not exist' - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: return 'table does not exist' - case EDataSourceKind.YDB: + case EGenericDataSourceKind.YDB: return 'issues = [{\'Path not found\'}])' - case EDataSourceKind.MS_SQL_SERVER: + case EGenericDataSourceKind.MS_SQL_SERVER: return 'table does not exist' - case EDataSourceKind.MYSQL: + case EGenericDataSourceKind.MYSQL: return 'table does not exist' - case EDataSourceKind.ORACLE: + case EGenericDataSourceKind.ORACLE: return 'table does not exist' case _: raise Exception(f'invalid data source: {self.kind}') diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/docker_compose.py b/ydb/library/yql/providers/generic/connector/tests/utils/docker_compose.py index 19d308744e20..dd2acea25897 100644 --- a/ydb/library/yql/providers/generic/connector/tests/utils/docker_compose.py +++ b/ydb/library/yql/providers/generic/connector/tests/utils/docker_compose.py @@ -9,7 +9,7 @@ import yatest.common -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.log import make_logger LOGGER = make_logger(__name__) @@ -111,17 +111,17 @@ def get_internal_ip(self, service_name: str) -> str: def get_container_name(self, service_name: str) -> str: return self.docker_compose_yml_data['services'][service_name]['container_name'] - def list_tables(self, dataSourceKind: EDataSourceKind) -> Sequence[str]: + def list_tables(self, dataSourceKind: EGenericDataSourceKind) -> Sequence[str]: match dataSourceKind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: return self.list_clickhouse_tables() - case EDataSourceKind.YDB: + case EGenericDataSourceKind.YDB: return self._list_ydb_tables() - case EDataSourceKind.MYSQL: + case EGenericDataSourceKind.MYSQL: return self._list_mysql_tables() - case EDataSourceKind.MS_SQL_SERVER: + case EGenericDataSourceKind.MS_SQL_SERVER: return self._list_ms_sql_server_tables() - case EDataSourceKind.ORACLE: + case EGenericDataSourceKind.ORACLE: return self._list_oracle_tables() case _: raise ValueError("invalid data source kind: {dataSourceKind}") diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/one_time_waiter.py b/ydb/library/yql/providers/generic/connector/tests/utils/one_time_waiter.py index 3659b095beca..082ecba1325a 100644 --- a/ydb/library/yql/providers/generic/connector/tests/utils/one_time_waiter.py +++ b/ydb/library/yql/providers/generic/connector/tests/utils/one_time_waiter.py @@ -4,7 +4,7 @@ import yatest.common -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.library.yql.providers.generic.connector.tests.utils.log import make_logger from ydb.library.yql.providers.generic.connector.tests.utils.docker_compose import DockerComposeHelper @@ -17,7 +17,7 @@ class OneTimeWaiter: def __init__( self, docker_compose_file_path: str, - data_source_kind: EDataSourceKind, + data_source_kind: EGenericDataSourceKind, expected_tables: Sequence[str], ): docker_compose_file_abs_path = yatest.common.source_path(docker_compose_file_path) diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/run/dqrun.py b/ydb/library/yql/providers/generic/connector/tests/utils/run/dqrun.py index 9bdea4d763ec..5c95f875e3b2 100644 --- a/ydb/library/yql/providers/generic/connector/tests/utils/run/dqrun.py +++ b/ydb/library/yql/providers/generic/connector/tests/utils/run/dqrun.py @@ -6,7 +6,7 @@ from yt import yson -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericProtocol from ydb.library.yql.providers.generic.connector.api.service.protos.connector_pb2 import EDateTimeFormat import ydb.library.yql.providers.generic.connector.tests.utils.artifacts as artifacts @@ -76,10 +76,10 @@ class GatewaysConfRenderer: {% for cluster in generic_settings.clickhouse_clusters %} -{% if cluster.protocol == EProtocol.NATIVE %} +{% if cluster.protocol == EGenericProtocol.NATIVE %} {% set CLICKHOUSE_PORT = settings.clickhouse.native_port_internal %} {% set CLICKHOUSE_PROTOCOL = NATIVE %} -{% elif cluster.protocol == EProtocol.HTTP %} +{% elif cluster.protocol == EGenericProtocol.HTTP %} {% set CLICKHOUSE_PORT = settings.clickhouse.http_port_internal %} {% set CLICKHOUSE_PROTOCOL = HTTP %} {% endif %} @@ -244,7 +244,7 @@ def __init__(self): self.template = jinja2.Environment(loader=jinja2.BaseLoader, undefined=jinja2.DebugUndefined).from_string( self._template ) - self.template.globals['EProtocol'] = EProtocol + self.template.globals['EGenericProtocol'] = EGenericProtocol self.template.globals['EDateTimeFormat'] = EDateTimeFormat def render(self, file_path: Path, settings: Settings, generic_settings: GenericSettings) -> None: 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 index 36ee776144df..f1c5f3bb2464 100644 --- a/ydb/library/yql/providers/generic/connector/tests/utils/run/kqprun.py +++ b/ydb/library/yql/providers/generic/connector/tests/utils/run/kqprun.py @@ -5,7 +5,7 @@ import jinja2 -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericProtocol from ydb.library.yql.providers.generic.connector.api.service.protos.connector_pb2 import EDateTimeFormat import ydb.library.yql.providers.generic.connector.tests.utils.artifacts as artifacts @@ -63,10 +63,10 @@ class SchemeRenderer: {% for cluster in generic_settings.clickhouse_clusters %} -{% if cluster.protocol == EProtocol.NATIVE %} +{% if cluster.protocol == EGenericProtocol.NATIVE %} {% set CLICKHOUSE_PORT = settings.clickhouse.native_port_internal %} {% set CLICKHOUSE_PROTOCOL = NATIVE %} -{% elif cluster.protocol == EProtocol.HTTP %} +{% elif cluster.protocol == EGenericProtocol.HTTP %} {% set CLICKHOUSE_PORT = settings.clickhouse.http_port_internal %} {% set CLICKHOUSE_PROTOCOL = HTTP %} {% endif %} @@ -165,7 +165,7 @@ def __init__(self): self.template = jinja2.Environment(loader=jinja2.BaseLoader, undefined=jinja2.DebugUndefined).from_string( self.template_ ) - self.template.globals['EProtocol'] = EProtocol + self.template.globals['EGenericProtocol'] = EGenericProtocol def render(self, file_path: Path, settings: Settings, generic_settings: GenericSettings) -> None: content = self.template.render(dict(settings=settings, generic_settings=generic_settings)) diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/run/ya.make b/ydb/library/yql/providers/generic/connector/tests/utils/run/ya.make index 362ed82c97b2..8cca116ec013 100644 --- a/ydb/library/yql/providers/generic/connector/tests/utils/run/ya.make +++ b/ydb/library/yql/providers/generic/connector/tests/utils/run/ya.make @@ -17,7 +17,7 @@ PY_SRCS( PEERDIR( contrib/python/Jinja2 contrib/python/PyYAML - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/api/service/protos ydb/library/yql/providers/generic/connector/tests/utils ydb/public/api/protos diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/scenario/postgresql.py b/ydb/library/yql/providers/generic/connector/tests/utils/scenario/postgresql.py index 91e0522bdcfc..dbc555ba1800 100644 --- a/ydb/library/yql/providers/generic/connector/tests/utils/scenario/postgresql.py +++ b/ydb/library/yql/providers/generic/connector/tests/utils/scenario/postgresql.py @@ -1,6 +1,6 @@ from typing import Sequence -import ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 as data_source_pb2 +import yql.essentials.providers.common.proto.gateways_config_pb2 as data_source_pb2 import ydb.library.yql.providers.generic.connector.tests.utils.artifacts as artifacts from ydb.library.yql.providers.generic.connector.tests.utils.comparator import assert_data_outs_equal diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/schema.py b/ydb/library/yql/providers/generic/connector/tests/utils/schema.py index d931f31730a1..cbbb7abb4e7a 100644 --- a/ydb/library/yql/providers/generic/connector/tests/utils/schema.py +++ b/ydb/library/yql/providers/generic/connector/tests/utils/schema.py @@ -5,7 +5,7 @@ from yt import yson from yt.yson.yson_types import YsonEntity import ydb.public.api.protos.ydb_value_pb2 as ydb_value -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind from ydb.public.api.protos.ydb_value_pb2 import Type, OptionalType import ydb.library.yql.providers.generic.connector.tests.utils.types.clickhouse as clickhouse @@ -27,20 +27,20 @@ class DataSourceType: pg: postgresql.Type = None ydb: Ydb.Type = None - def pick(self, kind: EDataSourceKind.ValueType) -> str: + def pick(self, kind: EGenericDataSourceKind.ValueType) -> str: target = None match kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: target = self.ch - case EDataSourceKind.MS_SQL_SERVER: + case EGenericDataSourceKind.MS_SQL_SERVER: target = self.ms - case EDataSourceKind.MYSQL: + case EGenericDataSourceKind.MYSQL: target = self.my - case EDataSourceKind.ORACLE: + case EGenericDataSourceKind.ORACLE: target = self.ora - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: target = self.pg - case EDataSourceKind.YDB: + case EGenericDataSourceKind.YDB: target = self.ydb case _: raise Exception(f'invalid data source: {kind}') @@ -192,7 +192,7 @@ def __cast_primitive_type(self, primitive_type_id: ydb_value.Type.PrimitiveTypeI case _: raise Exception(f"invalid type '{primitive_type_id}' for value '{value}'") - def format_for_data_source(self, kind: EDataSourceKind.ValueType) -> str: + def format_for_data_source(self, kind: EGenericDataSourceKind.ValueType) -> str: return f'{self.name} {self.data_source_type.pick(kind)}' @@ -349,7 +349,7 @@ def from_yson(cls, src: YsonList): def from_json(cls, src: Dict): return cls(columns=ColumnList(*map(Column.from_json, src))) - def yql_column_list(self, kind: EDataSourceKind.ValueType) -> str: + def yql_column_list(self, kind: EGenericDataSourceKind.ValueType) -> str: return ", ".join(map(lambda col: col.format_for_data_source(kind), self.columns)) def select_every_column(self) -> SelectWhat: diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/settings.py b/ydb/library/yql/providers/generic/connector/tests/utils/settings.py index db76d0927539..7d639bfe5e84 100644 --- a/ydb/library/yql/providers/generic/connector/tests/utils/settings.py +++ b/ydb/library/yql/providers/generic/connector/tests/utils/settings.py @@ -4,7 +4,7 @@ import yatest.common -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind, EGenericProtocol from ydb.library.yql.providers.generic.connector.api.service.protos.connector_pb2 import EDateTimeFormat from ydb.library.yql.providers.generic.connector.tests.utils.docker_compose import DockerComposeHelper @@ -100,7 +100,9 @@ class Ydb: ydb: Ydb @classmethod - def from_env(cls, docker_compose_dir: pathlib.Path, data_source_kinds: Sequence[EDataSourceKind]) -> 'Settings': + def from_env( + cls, docker_compose_dir: pathlib.Path, data_source_kinds: Sequence[EGenericDataSourceKind] + ) -> 'Settings': docker_compose_file_relative_path = str(docker_compose_dir / 'docker-compose.yml') docker_compose_file_abs_path = yatest.common.source_path(docker_compose_file_relative_path) endpoint_determiner = DockerComposeHelper(docker_compose_file_abs_path) @@ -109,7 +111,7 @@ def from_env(cls, docker_compose_dir: pathlib.Path, data_source_kinds: Sequence[ for data_source_kind in data_source_kinds: match data_source_kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: data_sources[data_source_kind] = cls.ClickHouse( cluster_name='clickhouse_integration_test', host_external='0.0.0.0', @@ -125,7 +127,7 @@ def from_env(cls, docker_compose_dir: pathlib.Path, data_source_kinds: Sequence[ password='password', protocol='native', ) - case EDataSourceKind.MS_SQL_SERVER: + case EGenericDataSourceKind.MS_SQL_SERVER: data_sources[data_source_kind] = cls.MsSQLServer( cluster_name='ms_sql_server_integration_test', host_external='0.0.0.0', @@ -139,7 +141,7 @@ def from_env(cls, docker_compose_dir: pathlib.Path, data_source_kinds: Sequence[ username='sa', password='Qwerty12345!', ) - case EDataSourceKind.MYSQL: + case EGenericDataSourceKind.MYSQL: data_sources[data_source_kind] = cls.MySQL( cluster_name='mysql_integration_test', host_external='0.0.0.0', @@ -153,7 +155,7 @@ def from_env(cls, docker_compose_dir: pathlib.Path, data_source_kinds: Sequence[ username='root', password='password', ) - case EDataSourceKind.ORACLE: + case EGenericDataSourceKind.ORACLE: data_sources[data_source_kind] = cls.Oracle( cluster_name='oracle_integration_test', host_external='0.0.0.0', @@ -168,7 +170,7 @@ def from_env(cls, docker_compose_dir: pathlib.Path, data_source_kinds: Sequence[ password='password', service_name="FREE", ) - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: data_sources[data_source_kind] = cls.PostgreSQL( cluster_name='postgresql_integration_test', host_external='0.0.0.0', @@ -182,7 +184,7 @@ def from_env(cls, docker_compose_dir: pathlib.Path, data_source_kinds: Sequence[ username='user', password='password', ) - case EDataSourceKind.YDB: + case EGenericDataSourceKind.YDB: data_sources[data_source_kind] = cls.Ydb( cluster_name='ydb_integration_test', host_internal=endpoint_determiner.get_container_name('ydb'), @@ -195,36 +197,36 @@ def from_env(cls, docker_compose_dir: pathlib.Path, data_source_kinds: Sequence[ raise Exception(f'invalid data source: {data_source_kind}') return cls( - clickhouse=data_sources.get(EDataSourceKind.CLICKHOUSE), + clickhouse=data_sources.get(EGenericDataSourceKind.CLICKHOUSE), connector=cls.Connector( grpc_host='localhost', grpc_port=endpoint_determiner.get_external_port('fq-connector-go', 2130), paging_bytes_per_page=4 * 1024 * 1024, paging_prefetch_queue_capacity=2, ), - ms_sql_server=data_sources.get(EDataSourceKind.MS_SQL_SERVER), - mysql=data_sources.get(EDataSourceKind.MYSQL), - oracle=data_sources.get(EDataSourceKind.ORACLE), - postgresql=data_sources.get(EDataSourceKind.POSTGRESQL), - ydb=data_sources.get(EDataSourceKind.YDB), + ms_sql_server=data_sources.get(EGenericDataSourceKind.MS_SQL_SERVER), + mysql=data_sources.get(EGenericDataSourceKind.MYSQL), + oracle=data_sources.get(EGenericDataSourceKind.ORACLE), + postgresql=data_sources.get(EGenericDataSourceKind.POSTGRESQL), + ydb=data_sources.get(EGenericDataSourceKind.YDB), ) - def get_cluster_name(self, data_source_kind: EDataSourceKind) -> str: + def get_cluster_name(self, data_source_kind: EGenericDataSourceKind) -> str: match data_source_kind: - case EDataSourceKind.CLICKHOUSE: + case EGenericDataSourceKind.CLICKHOUSE: return self.clickhouse.cluster_name - case EDataSourceKind.MYSQL: + case EGenericDataSourceKind.MYSQL: return self.mysql.cluster_name - case EDataSourceKind.ORACLE: + case EGenericDataSourceKind.ORACLE: return self.oracle.cluster_name - case EDataSourceKind.MS_SQL_SERVER: + case EGenericDataSourceKind.MS_SQL_SERVER: return self.ms_sql_server.cluster_name - case EDataSourceKind.POSTGRESQL: + case EGenericDataSourceKind.POSTGRESQL: return self.postgresql.cluster_name - case EDataSourceKind.YDB: + case EGenericDataSourceKind.YDB: return self.ydb.cluster_name case _: - raise Exception(f'invalid data source: {EDataSourceKind.Name(data_source_kind)}') + raise Exception(f'invalid data source: {EGenericDataSourceKind.Name(data_source_kind)}') @dataclass @@ -237,7 +239,7 @@ def __hash__(self) -> int: return hash(self.database) + hash(self.protocol) database: str - protocol: EProtocol + protocol: EGenericProtocol clickhouse_clusters: Sequence[ClickHouseCluster] = field(default_factory=list) diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/ya.make b/ydb/library/yql/providers/generic/connector/tests/utils/ya.make index d0dea697ee67..6cb533c60941 100644 --- a/ydb/library/yql/providers/generic/connector/tests/utils/ya.make +++ b/ydb/library/yql/providers/generic/connector/tests/utils/ya.make @@ -21,7 +21,7 @@ ENDIF() PEERDIR( contrib/python/PyYAML - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/tests/utils/types ydb/public/api/protos yt/python/yt/yson diff --git a/ydb/library/yql/providers/generic/proto/source.proto b/ydb/library/yql/providers/generic/proto/source.proto index a78d4ddc27b3..a9a92d92e36b 100644 --- a/ydb/library/yql/providers/generic/proto/source.proto +++ b/ydb/library/yql/providers/generic/proto/source.proto @@ -5,7 +5,7 @@ option cc_enable_arenas = true; package NYql.Generic; import "ydb/library/yql/providers/generic/connector/api/service/protos/connector.proto"; -import "ydb/library/yql/providers/generic/connector/api/common/data_source.proto"; +import "yql/essentials/providers/common/proto/gateways_config.proto"; message TSource { // Prepared Select expression @@ -25,7 +25,7 @@ message TSource { //TODO(YQ-3093) move SA creds to some common message message TLookupSource { - NYql.NConnector.NApi.TDataSourceInstance data_source_instance = 1; + NYql.TGenericDataSourceInstance data_source_instance = 1; string table = 2; // Credentials used to access managed databases APIs. // When working with external data source instances deployed in clouds, diff --git a/ydb/library/yql/providers/generic/proto/ya.make b/ydb/library/yql/providers/generic/proto/ya.make index 39a31b04067d..31880aadde41 100644 --- a/ydb/library/yql/providers/generic/proto/ya.make +++ b/ydb/library/yql/providers/generic/proto/ya.make @@ -4,7 +4,7 @@ ONLY_TAGS(CPP_PROTO) PEERDIR( ydb/library/yql/providers/generic/connector/api/service/protos - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ) SRCS( diff --git a/ydb/library/yql/providers/generic/provider/ut/pushdown/pushdown_ut.cpp b/ydb/library/yql/providers/generic/provider/ut/pushdown/pushdown_ut.cpp index 73e538e024a8..841de470544e 100644 --- a/ydb/library/yql/providers/generic/provider/ut/pushdown/pushdown_ut.cpp +++ b/ydb/library/yql/providers/generic/provider/ut/pushdown/pushdown_ut.cpp @@ -74,7 +74,7 @@ struct TFakeDatabaseResolver: public IDatabaseAsyncResolver { }; struct TFakeGenericClient: public NConnector::IClient { - NConnector::TDescribeTableAsyncResult DescribeTable(const NConnector::NApi::TDescribeTableRequest& request) { + NConnector::TDescribeTableAsyncResult DescribeTable(const NConnector::NApi::TDescribeTableRequest& request, TDuration) override { UNIT_ASSERT_VALUES_EQUAL(request.table(), "test_table"); NConnector::TResult result; auto& resp = result.Response.emplace(); @@ -123,7 +123,7 @@ struct TFakeGenericClient: public NConnector::IClient { return NThreading::MakeFuture(std::move(result)); } - NConnector::TListSplitsStreamIteratorAsyncResult ListSplits(const NConnector::NApi::TListSplitsRequest& request) { + NConnector::TListSplitsStreamIteratorAsyncResult ListSplits(const NConnector::NApi::TListSplitsRequest& request, TDuration) override { Y_UNUSED(request); try { throw std::runtime_error("ListSplits unimplemented"); @@ -132,7 +132,7 @@ struct TFakeGenericClient: public NConnector::IClient { } } - NConnector::TReadSplitsStreamIteratorAsyncResult ReadSplits(const NConnector::NApi::TReadSplitsRequest& request) { + NConnector::TReadSplitsStreamIteratorAsyncResult ReadSplits(const NConnector::NApi::TReadSplitsRequest& request, TDuration) override { Y_UNUSED(request); try { throw std::runtime_error("ReadSplits unimplemented"); @@ -230,13 +230,13 @@ struct TPushdownFixture: public NUnitTest::TBaseFixture { auto* cluster = GatewaysCfg.MutableGeneric()->AddClusterMapping(); cluster->SetName("test_cluster"); - cluster->SetKind(NConnector::NApi::POSTGRESQL); + cluster->SetKind(NYql::EGenericDataSourceKind::POSTGRESQL); cluster->MutableEndpoint()->set_host("host"); cluster->MutableEndpoint()->set_port(42); cluster->MutableCredentials()->mutable_basic()->set_username("user"); cluster->MutableCredentials()->mutable_basic()->set_password("password"); cluster->SetDatabaseName("database"); - cluster->SetProtocol(NConnector::NApi::NATIVE); + cluster->SetProtocol(NYql::EGenericProtocol::NATIVE); } GenericState = MakeIntrusive( diff --git a/ydb/library/yql/providers/generic/provider/ya.make b/ydb/library/yql/providers/generic/provider/ya.make index b87b656e56bf..5ee4ed066a72 100644 --- a/ydb/library/yql/providers/generic/provider/ya.make +++ b/ydb/library/yql/providers/generic/provider/ya.make @@ -54,7 +54,7 @@ PEERDIR( ydb/library/yql/providers/dq/expr_nodes ydb/library/yql/providers/generic/expr_nodes ydb/library/yql/providers/generic/proto - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/libcpp yql/essentials/providers/result/expr_nodes ydb/library/yql/utils/plan diff --git a/ydb/library/yql/providers/generic/provider/yql_generic_cluster_config.cpp b/ydb/library/yql/providers/generic/provider/yql_generic_cluster_config.cpp index 9745f2c59461..ce35300b49e8 100644 --- a/ydb/library/yql/providers/generic/provider/yql_generic_cluster_config.cpp +++ b/ydb/library/yql/providers/generic/provider/yql_generic_cluster_config.cpp @@ -6,7 +6,7 @@ #include #include -#include +#include #include #include "yql_generic_cluster_config.h" @@ -200,25 +200,23 @@ namespace NYql { void ParseProtocol(const THashMap& properties, NYql::TGenericClusterConfig& clusterConfig) { - using namespace NConnector::NApi; - // For some datasources the PROTOCOL is not required - if (clusterConfig.GetKind() == EDataSourceKind::LOGGING) { - clusterConfig.SetProtocol(EProtocol::PROTOCOL_UNSPECIFIED); + if (clusterConfig.GetKind() == EGenericDataSourceKind::LOGGING) { + clusterConfig.SetProtocol(EGenericProtocol::PROTOCOL_UNSPECIFIED); return; } // For the most of transactional databases the PROTOCOL is always NATIVE if (IsIn({ - EDataSourceKind::GREENPLUM, - EDataSourceKind::YDB, - EDataSourceKind::MYSQL, - EDataSourceKind::MS_SQL_SERVER, - EDataSourceKind::ORACLE, + EGenericDataSourceKind::GREENPLUM, + EGenericDataSourceKind::YDB, + EGenericDataSourceKind::MYSQL, + EGenericDataSourceKind::MS_SQL_SERVER, + EGenericDataSourceKind::ORACLE, }, clusterConfig.GetKind() )) { - clusterConfig.SetProtocol(EProtocol::NATIVE); + clusterConfig.SetProtocol(EGenericProtocol::NATIVE); return; } @@ -231,13 +229,13 @@ namespace NYql { auto input = it->second; input.to_upper(); - EProtocol protocol; - if (!EProtocol_Parse(input, &protocol)) { + EGenericProtocol protocol; + if (!NYql::EGenericProtocol_Parse(input, &protocol)) { TStringBuilder b; b << "invalid 'PROTOCOL' value: '" << it->second << "', valid types are: "; - for (auto i = EProtocol_MIN + 1; i < EProtocol_MAX; i++) { - b << EProtocol_Name(i); - if (i != EProtocol_MAX - 1) { + for (auto i = EGenericProtocol_MIN + 1; i < EGenericProtocol_MAX; i++) { + b << NYql::EGenericProtocol_Name(i); + if (i != EGenericProtocol_MAX - 1) { b << ", "; } } @@ -351,7 +349,7 @@ namespace NYql { "context"_a = context, "msg"_a = msg, "name"_a = clusterConfig.GetName(), - "kind"_a = NConnector::NApi::EDataSourceKind_Name(clusterConfig.GetKind()), + "kind"_a = NYql::EGenericDataSourceKind_Name(clusterConfig.GetKind()), "host"_a = clusterConfig.GetEndpoint().Gethost(), "port"_a = clusterConfig.GetEndpoint().Getport(), "database_id"_a = clusterConfig.GetDatabaseId(), @@ -362,31 +360,31 @@ namespace NYql { "token"_a = ToString(clusterConfig.GetToken().size()), "use_ssl"_a = clusterConfig.GetUseSsl() ? "TRUE" : "FALSE", "database_name"_a = clusterConfig.GetDatabaseName(), - "protocol"_a = NConnector::NApi::EProtocol_Name(clusterConfig.GetProtocol()), + "protocol"_a = NYql::EGenericProtocol_Name(clusterConfig.GetProtocol()), "schema"_a = GetPropertyWithDefault(clusterConfig.datasourceoptions(), "schema"), "folder_id"_a = GetPropertyWithDefault(clusterConfig.datasourceoptions(), "folder_id") ); } - static const TSet managedDatabaseKinds{ - NConnector::NApi::EDataSourceKind::CLICKHOUSE, - NConnector::NApi::EDataSourceKind::GREENPLUM, - NConnector::NApi::EDataSourceKind::MYSQL, - NConnector::NApi::EDataSourceKind::POSTGRESQL, - NConnector::NApi::EDataSourceKind::YDB, + static const TSet managedDatabaseKinds{ + NYql::EGenericDataSourceKind::CLICKHOUSE, + NYql::EGenericDataSourceKind::GREENPLUM, + NYql::EGenericDataSourceKind::MYSQL, + NYql::EGenericDataSourceKind::POSTGRESQL, + NYql::EGenericDataSourceKind::YDB, }; - static const TSet traditionalRelationalDatabaseKinds{ - NConnector::NApi::EDataSourceKind::CLICKHOUSE, - NConnector::NApi::EDataSourceKind::GREENPLUM, - NConnector::NApi::EDataSourceKind::MS_SQL_SERVER, - NConnector::NApi::EDataSourceKind::MYSQL, - NConnector::NApi::EDataSourceKind::ORACLE, - NConnector::NApi::EDataSourceKind::POSTGRESQL, + static const TSet traditionalRelationalDatabaseKinds{ + NYql::EGenericDataSourceKind::CLICKHOUSE, + NYql::EGenericDataSourceKind::GREENPLUM, + NYql::EGenericDataSourceKind::MS_SQL_SERVER, + NYql::EGenericDataSourceKind::MYSQL, + NYql::EGenericDataSourceKind::ORACLE, + NYql::EGenericDataSourceKind::POSTGRESQL, }; - bool DataSourceMustHaveDataBaseName(const NConnector::NApi::EDataSourceKind& sourceKind) { - return traditionalRelationalDatabaseKinds.contains(sourceKind) && sourceKind != NConnector::NApi::ORACLE; + bool DataSourceMustHaveDataBaseName(const NYql::EGenericDataSourceKind& sourceKind) { + return traditionalRelationalDatabaseKinds.contains(sourceKind) && sourceKind != NYql::EGenericDataSourceKind::ORACLE; } void ValidateGenericClusterConfig( @@ -447,7 +445,7 @@ namespace NYql { // YDB: // * set database name when working with on-prem YDB instance; // * but set database ID when working with managed YDB. - if (clusterConfig.GetKind() == NConnector::NApi::YDB) { + if (clusterConfig.GetKind() == NYql::EGenericDataSourceKind::YDB) { if (clusterConfig.HasDatabaseName() && clusterConfig.HasDatabaseId()) { return ValidationError( clusterConfig, @@ -465,7 +463,7 @@ namespace NYql { // Oracle: // * always set service_name for oracle; - if (clusterConfig.GetKind() == NConnector::NApi::ORACLE) { + if (clusterConfig.GetKind() == NYql::EGenericDataSourceKind::ORACLE) { if (!clusterConfig.GetDataSourceOptions().contains("service_name")) { return ValidationError( clusterConfig, @@ -490,7 +488,7 @@ namespace NYql { return ValidationError(clusterConfig, context, "empty field 'Name'"); } - if (clusterConfig.GetKind() == NConnector::NApi::EDataSourceKind::DATA_SOURCE_KIND_UNSPECIFIED) { + if (clusterConfig.GetKind() == NYql::EGenericDataSourceKind::DATA_SOURCE_KIND_UNSPECIFIED) { return ValidationError(clusterConfig, context, "empty field 'Kind'"); } diff --git a/ydb/library/yql/providers/generic/provider/yql_generic_dq_integration.cpp b/ydb/library/yql/providers/generic/provider/yql_generic_dq_integration.cpp index c7a92f783037..c36536b6fd17 100644 --- a/ydb/library/yql/providers/generic/provider/yql_generic_dq_integration.cpp +++ b/ydb/library/yql/providers/generic/provider/yql_generic_dq_integration.cpp @@ -21,23 +21,23 @@ namespace NYql { namespace { - TString GetSourceType(NYql::NConnector::NApi::TDataSourceInstance dsi) { + TString GetSourceType(NYql::TGenericDataSourceInstance dsi) { switch (dsi.kind()) { - case NYql::NConnector::NApi::CLICKHOUSE: + case NYql::EGenericDataSourceKind::CLICKHOUSE: return "ClickHouseGeneric"; - case NYql::NConnector::NApi::POSTGRESQL: + case NYql::EGenericDataSourceKind::POSTGRESQL: return "PostgreSqlGeneric"; - case NYql::NConnector::NApi::MYSQL: + case NYql::EGenericDataSourceKind::MYSQL: return "MySqlGeneric"; - case NYql::NConnector::NApi::YDB: + case NYql::EGenericDataSourceKind::YDB: return "YdbGeneric"; - case NYql::NConnector::NApi::GREENPLUM: + case NYql::EGenericDataSourceKind::GREENPLUM: return "GreenplumGeneric"; - case NYql::NConnector::NApi::MS_SQL_SERVER: + case NYql::EGenericDataSourceKind::MS_SQL_SERVER: return "MsSQLServerGeneric"; - case NYql::NConnector::NApi::ORACLE: + case NYql::EGenericDataSourceKind::ORACLE: return "OracleGeneric"; - case NYql::NConnector::NApi::LOGGING: + case NYql::EGenericDataSourceKind::LOGGING: return "LoggingGeneric"; default: ythrow yexception() << "Data source kind is unknown or not specified"; @@ -164,10 +164,10 @@ namespace NYql { } } - // Managed YDB supports access via IAM token. + // Managed YDB (including YDB underlying Logging) supports access via IAM token. // If exist, copy service account creds to obtain tokens during request execution phase. // If exists, copy previously created token. - if (clusterConfig.kind() == NConnector::NApi::EDataSourceKind::YDB) { + if (IsIn({NYql::EGenericDataSourceKind::YDB, NYql::EGenericDataSourceKind::LOGGING}, clusterConfig.kind())) { source.SetServiceAccountId(clusterConfig.GetServiceAccountId()); source.SetServiceAccountIdSignature(clusterConfig.GetServiceAccountIdSignature()); source.SetToken(State_->Types->Credentials->FindCredentialContent( @@ -198,36 +198,36 @@ namespace NYql { properties["Table"] = table; auto [tableMeta, issue] = State_->GetTable(clusterName, table); if (!issue) { - const NConnector::NApi::TDataSourceInstance& dataSourceInstance = tableMeta.value()->DataSourceInstance; + const NYql::TGenericDataSourceInstance& dataSourceInstance = tableMeta.value()->DataSourceInstance; switch (dataSourceInstance.kind()) { - case NConnector::NApi::CLICKHOUSE: + case NYql::EGenericDataSourceKind::CLICKHOUSE: properties["SourceType"] = "ClickHouse"; break; - case NConnector::NApi::POSTGRESQL: + case NYql::EGenericDataSourceKind::POSTGRESQL: properties["SourceType"] = "PostgreSql"; break; - case NConnector::NApi::MYSQL: + case NYql::EGenericDataSourceKind::MYSQL: properties["SourceType"] = "MySql"; break; - case NConnector::NApi::YDB: + case NYql::EGenericDataSourceKind::YDB: properties["SourceType"] = "Ydb"; break; - case NConnector::NApi::GREENPLUM: + case NYql::EGenericDataSourceKind::GREENPLUM: properties["SourceType"] = "Greenplum"; break; - case NConnector::NApi::MS_SQL_SERVER: + case NYql::EGenericDataSourceKind::MS_SQL_SERVER: properties["SourceType"] = "MsSQLServer"; break; - case NConnector::NApi::ORACLE: + case NYql::EGenericDataSourceKind::ORACLE: properties["SourceType"] = "Oracle"; break; - case NConnector::NApi::LOGGING: + case NYql::EGenericDataSourceKind::LOGGING: properties["SourceType"] = "Logging"; break; - case NConnector::NApi::DATA_SOURCE_KIND_UNSPECIFIED: + case NYql::EGenericDataSourceKind::DATA_SOURCE_KIND_UNSPECIFIED: break; default: - properties["SourceType"] = NConnector::NApi::EDataSourceKind_Name(dataSourceInstance.kind()); + properties["SourceType"] = NYql::EGenericDataSourceKind_Name(dataSourceInstance.kind()); break; } @@ -236,15 +236,15 @@ namespace NYql { } switch (dataSourceInstance.protocol()) { - case NConnector::NApi::NATIVE: + case NYql::EGenericProtocol::NATIVE: properties["Protocol"] = "Native"; break; - case NConnector::NApi::HTTP: + case NYql::EGenericProtocol::HTTP: properties["Protocol"] = "Http"; break; - case NConnector::NApi::PROTOCOL_UNSPECIFIED: + case NYql::EGenericProtocol::PROTOCOL_UNSPECIFIED: default: - properties["Protocol"] = NConnector::NApi::EProtocol_Name(dataSourceInstance.protocol()); + properties["Protocol"] = NYql::EGenericProtocol_Name(dataSourceInstance.protocol()); break; } } @@ -291,7 +291,7 @@ namespace NYql { // Managed YDB supports access via IAM token. // If exist, copy service account creds to obtain tokens during request execution phase. // If exists, copy previously created token. - if (clusterConfig.kind() == NConnector::NApi::EDataSourceKind::YDB) { + if (clusterConfig.kind() == NYql::EGenericDataSourceKind::YDB) { source.SetServiceAccountId(clusterConfig.GetServiceAccountId()); source.SetServiceAccountIdSignature(clusterConfig.GetServiceAccountIdSignature()); source.SetToken(State_->Types->Credentials->FindCredentialContent( diff --git a/ydb/library/yql/providers/generic/provider/yql_generic_io_discovery.cpp b/ydb/library/yql/providers/generic/provider/yql_generic_io_discovery.cpp index 95f5b3f7fa97..0a3f8bd00deb 100644 --- a/ydb/library/yql/providers/generic/provider/yql_generic_io_discovery.cpp +++ b/ydb/library/yql/providers/generic/provider/yql_generic_io_discovery.cpp @@ -174,7 +174,7 @@ namespace NYql { // If we work with managed YDB, we find out database name // only after database id (== cluster id) resolving. - if (clusterConfigIter->second.kind() == NConnector::NApi::EDataSourceKind::YDB) { + if (clusterConfigIter->second.kind() == NYql::EGenericDataSourceKind::YDB) { clusterConfigIter->second.set_databasename(databaseDescription.Database); } diff --git a/ydb/library/yql/providers/generic/provider/yql_generic_load_meta.cpp b/ydb/library/yql/providers/generic/provider/yql_generic_load_meta.cpp index 2ca4fe6af541..71d3e977362a 100644 --- a/ydb/library/yql/providers/generic/provider/yql_generic_load_meta.cpp +++ b/ydb/library/yql/providers/generic/provider/yql_generic_load_meta.cpp @@ -28,7 +28,7 @@ namespace NYql { struct TGenericTableDescription { using TPtr = std::shared_ptr; - NConnector::NApi::TDataSourceInstance DataSourceInstance; + NYql::TGenericDataSourceInstance DataSourceInstance; std::optional Response; }; @@ -328,40 +328,49 @@ namespace NYql { request.set_schema(schema); } - void GetOracleServiceName(NYql::NConnector::NApi::TOracleDataSourceOptions& request, const TGenericClusterConfig& clusterConfig) { + void SetOracleServiceName(NYql::TOracleDataSourceOptions& options, const TGenericClusterConfig& clusterConfig) { const auto it = clusterConfig.GetDataSourceOptions().find("service_name"); if (it != clusterConfig.GetDataSourceOptions().end()) { - request.set_service_name(it->second); + options.set_service_name(it->second); + } + } + + void SetLoggingFolderId(NYql::TLoggingDataSourceOptions& options, const TGenericClusterConfig& clusterConfig) { + const auto it = clusterConfig.GetDataSourceOptions().find("folder_id"); + if (it != clusterConfig.GetDataSourceOptions().end()) { + options.set_folder_id(it->second); } } void FillDataSourceOptions(NConnector::NApi::TDescribeTableRequest& request, const TGenericClusterConfig& clusterConfig) { const auto dataSourceKind = clusterConfig.GetKind(); switch (dataSourceKind) { - case NYql::NConnector::NApi::CLICKHOUSE: + case NYql::EGenericDataSourceKind::CLICKHOUSE: break; - case NYql::NConnector::NApi::YDB: + case NYql::EGenericDataSourceKind::YDB: break; - case NYql::NConnector::NApi::MYSQL: + case NYql::EGenericDataSourceKind::MYSQL: break; - case NYql::NConnector::NApi::GREENPLUM: { + case NYql::EGenericDataSourceKind::GREENPLUM: { auto* options = request.mutable_data_source_instance()->mutable_gp_options(); SetSchema(*options, clusterConfig); } break; - case NYql::NConnector::NApi::MS_SQL_SERVER: + case NYql::EGenericDataSourceKind::MS_SQL_SERVER: break; - case NYql::NConnector::NApi::POSTGRESQL: { + case NYql::EGenericDataSourceKind::POSTGRESQL: { auto* options = request.mutable_data_source_instance()->mutable_pg_options(); SetSchema(*options, clusterConfig); } break; - case NYql::NConnector::NApi::ORACLE: { + case NYql::EGenericDataSourceKind::ORACLE: { auto* options = request.mutable_data_source_instance()->mutable_oracle_options(); - GetOracleServiceName(*options, clusterConfig); + SetOracleServiceName(*options, clusterConfig); + } break; + case NYql::EGenericDataSourceKind::LOGGING: { + auto* options = request.mutable_data_source_instance()->mutable_logging_options(); + SetLoggingFolderId(*options, clusterConfig); } break; - case NYql::NConnector::NApi::LOGGING: - break; default: - ythrow yexception() << "Unexpected data source kind: '" << NYql::NConnector::NApi::EDataSourceKind_Name(dataSourceKind) + ythrow yexception() << "Unexpected data source kind: '" << NYql::EGenericDataSourceKind_Name(dataSourceKind) << "'"; } } diff --git a/ydb/library/yql/providers/generic/provider/yql_generic_state.h b/ydb/library/yql/providers/generic/provider/yql_generic_state.h index c217db35eff0..9a4a003e048c 100644 --- a/ydb/library/yql/providers/generic/provider/yql_generic_state.h +++ b/ydb/library/yql/providers/generic/provider/yql_generic_state.h @@ -21,7 +21,7 @@ namespace NYql { const TStructExprType* ItemType = nullptr; TVector ColumnOrder; NYql::NConnector::NApi::TSchema Schema; - NYql::NConnector::NApi::TDataSourceInstance DataSourceInstance; + NYql::TGenericDataSourceInstance DataSourceInstance; }; using TGetTableResult = std::pair, std::optional>; diff --git a/ydb/library/yql/providers/generic/provider/yql_generic_utils.cpp b/ydb/library/yql/providers/generic/provider/yql_generic_utils.cpp index b2ec216b3b0c..9d88896b8896 100644 --- a/ydb/library/yql/providers/generic/provider/yql_generic_utils.cpp +++ b/ydb/library/yql/providers/generic/provider/yql_generic_utils.cpp @@ -6,12 +6,12 @@ namespace NYql { TString DumpGenericClusterConfig(const TGenericClusterConfig& clusterConfig) { TStringBuilder sb; sb << "name = " << clusterConfig.GetName() - << ", kind = " << NConnector::NApi::EDataSourceKind_Name(clusterConfig.GetKind()) + << ", kind = " << NYql::EGenericDataSourceKind_Name(clusterConfig.GetKind()) << ", database name = " << clusterConfig.GetDatabaseName() << ", database id = " << clusterConfig.GetDatabaseId() << ", endpoint = " << clusterConfig.GetEndpoint() << ", use tls = " << clusterConfig.GetUseSsl() - << ", protocol = " << NConnector::NApi::EProtocol_Name(clusterConfig.GetProtocol()); + << ", protocol = " << NYql::EGenericProtocol_Name(clusterConfig.GetProtocol()); for (const auto& [key, value] : clusterConfig.GetDataSourceOptions()) { sb << ", " << key << " = " << value; diff --git a/ydb/library/yql/providers/generic/provider/yql_generic_utils.h b/ydb/library/yql/providers/generic/provider/yql_generic_utils.h index 9bf173b1bf6a..4370e4f5ece1 100644 --- a/ydb/library/yql/providers/generic/provider/yql_generic_utils.h +++ b/ydb/library/yql/providers/generic/provider/yql_generic_utils.h @@ -1,7 +1,7 @@ #pragma once #include -#include +#include namespace NYql { TString DumpGenericClusterConfig(const TGenericClusterConfig& clusterConfig); diff --git a/ydb/library/yql/providers/pq/async_io/dq_pq_rd_read_actor.cpp b/ydb/library/yql/providers/pq/async_io/dq_pq_rd_read_actor.cpp index 426a4ab92f79..6d01d6d90c71 100644 --- a/ydb/library/yql/providers/pq/async_io/dq_pq_rd_read_actor.cpp +++ b/ydb/library/yql/providers/pq/async_io/dq_pq_rd_read_actor.cpp @@ -2,6 +2,7 @@ #include "probes.h" #include +#include #include #include #include @@ -11,11 +12,14 @@ #include #include +#include #include +#include #include #include #include #include +#include #include #include #include @@ -74,15 +78,16 @@ struct TRowDispatcherReadActorMetrics { explicit TRowDispatcherReadActorMetrics(const TTxId& txId, ui64 taskId, const ::NMonitoring::TDynamicCounterPtr& counters) : TxId(std::visit([](auto arg) { return ToString(arg); }, txId)) , Counters(counters) { - SubGroup = Counters->GetSubgroup("sink", "RdPqRead"); - auto sink = SubGroup->GetSubgroup("tx_id", TxId); - auto task = sink->GetSubgroup("task_id", ToString(taskId)); + SubGroup = Counters->GetSubgroup("source", "RdPqRead"); + auto source = SubGroup->GetSubgroup("tx_id", TxId); + auto task = source->GetSubgroup("task_id", ToString(taskId)); InFlyGetNextBatch = task->GetCounter("InFlyGetNextBatch"); InFlyAsyncInputData = task->GetCounter("InFlyAsyncInputData"); + ReInit = task->GetCounter("ReInit", true); } ~TRowDispatcherReadActorMetrics() { - SubGroup->RemoveSubgroup("id", TxId); + SubGroup->RemoveSubgroup("tx_id", TxId); } TString TxId; @@ -90,6 +95,7 @@ struct TRowDispatcherReadActorMetrics { ::NMonitoring::TDynamicCounterPtr SubGroup; ::NMonitoring::TDynamicCounters::TCounterPtr InFlyGetNextBatch; ::NMonitoring::TDynamicCounters::TCounterPtr InFlyAsyncInputData; + ::NMonitoring::TDynamicCounters::TCounterPtr ReInit; }; struct TEvPrivate { @@ -97,19 +103,21 @@ struct TEvPrivate { EvBegin = EventSpaceBegin(NActors::TEvents::ES_PRIVATE), EvPrintState = EvBegin + 20, EvProcessState = EvBegin + 21, + EvNotifyCA = EvBegin + 22, EvEnd }; static_assert(EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE)"); struct TEvPrintState : public NActors::TEventLocal {}; struct TEvProcessState : public NActors::TEventLocal {}; + struct TEvNotifyCA : public NActors::TEventLocal {}; }; class TDqPqRdReadActor : public NActors::TActor, public NYql::NDq::NInternal::TDqPqReadActorBase { const ui64 PrintStatePeriodSec = 300; - const ui64 ProcessStatePeriodSec = 2; - - using TDebugOffsets = TMaybe>; + const ui64 ProcessStatePeriodSec = 1; + const ui64 PrintStateToLogSplitSize = 64000; + const ui64 NotifyCAPeriodSec = 10; struct TReadyBatch { public: @@ -119,7 +127,7 @@ class TDqPqRdReadActor : public NActors::TActor, public NYql:: } public: - TVector Data; + TVector Data; i64 UsedSpace = 0; ui64 NextOffset = 0; ui64 PartitionId; @@ -155,54 +163,76 @@ class TDqPqRdReadActor : public NActors::TActor, public NYql:: }; private: - std::vector> MetadataFields; const TString Token; TMaybe CoordinatorActorId; NActors::TActorId LocalRowDispatcherActorId; std::queue ReadyBuffer; EState State = EState::INIT; + bool Inited = false; ui64 CoordinatorRequestCookie = 0; TRowDispatcherReadActorMetrics Metrics; - bool SchedulePrintStatePeriod = false; bool ProcessStateScheduled = false; bool InFlyAsyncInputData = false; - TCounters Counters; + TCounters Counters; + // Parsing info + std::vector> ColumnIndexes; // Output column index in schema passed into RowDispatcher + const TType* InputDataType = nullptr; // Multi type (comes from Row Dispatcher) + std::unique_ptr> DataUnpacker; + ui64 CpuMicrosec = 0; + + THashMap> NextOffsetFromRD; + + struct TPartition { + bool HasPendingData = false; + bool IsWaitingMessageBatch = false; + }; - struct SessionInfo { + struct TSession { enum class ESessionStatus { - NoSession, - Started, + INIT, + WAIT_START_SESSION_ACK, + STARTED, }; - SessionInfo( + + TSession( const TTxId& txId, const NActors::TActorId selfId, TActorId rowDispatcherActorId, - ui64 partitionId, ui64 eventQueueId, ui64 generation) - : RowDispatcherActorId(rowDispatcherActorId) - , PartitionId(partitionId) - , Generation(generation) { - EventsQueue.Init(txId, selfId, selfId, eventQueueId, /* KeepAlive */ true); + : TxId(txId) + , SelfId(selfId) + , EventQueueId(eventQueueId) + , RowDispatcherActorId(rowDispatcherActorId) + , Generation(generation) + { + EventsQueue.Init(TxId, SelfId, SelfId, EventQueueId, /* KeepAlive */ true); EventsQueue.OnNewRecipientId(rowDispatcherActorId); } - ESessionStatus Status = ESessionStatus::NoSession; - ui64 NextOffset = 0; - bool IsWaitingStartSessionAck = false; + ESessionStatus Status = ESessionStatus::INIT; + const TTxId TxId; + const NActors::TActorId SelfId; + const ui64 EventQueueId; NYql::NDq::TRetryEventsQueue EventsQueue; - bool HasPendingData = false; - bool IsWaitingMessageBatch = false; TActorId RowDispatcherActorId; - ui64 PartitionId; - ui64 Generation; + ui64 Generation = std::numeric_limits::max(); + THashMap Partitions; + bool IsWaitingStartSessionAck = false; + ui64 QueuedBytes = 0; + ui64 QueuedRows = 0; }; - TMap Sessions; + TMap Sessions; + THashMap ReadActorByEventQueueId; const THolderFactory& HolderFactory; const i64 MaxBufferSize; i64 ReadyBufferSizeBytes = 0; ui64 NextGeneration = 0; + ui64 NextEventQueueId = 0; + + TMap> LastUsedPartitionDistribution; + TMap> LastReceivedPartitionDistribution; public: TDqPqRdReadActor( @@ -211,6 +241,7 @@ class TDqPqRdReadActor : public NActors::TActor, public NYql:: const TTxId& txId, ui64 taskId, const THolderFactory& holderFactory, + const TTypeEnvironment& typeEnv, NPq::NProto::TDqPqTopicSource&& sourceParams, NPq::NProto::TDqReadTaskParams&& readParams, const NActors::TActorId& computeActorId, @@ -237,6 +268,7 @@ class TDqPqRdReadActor : public NActors::TActor, public NYql:: void Handle(const NFq::TEvRowDispatcher::TEvHeartbeat::TPtr&); void Handle(TEvPrivate::TEvPrintState::TPtr&); void Handle(TEvPrivate::TEvProcessState::TPtr&); + void Handle(TEvPrivate::TEvNotifyCA::TPtr&); STRICT_STFUNC(StateFunc, { hFunc(NFq::TEvRowDispatcher::TEvCoordinatorChanged, Handle); @@ -257,6 +289,7 @@ class TDqPqRdReadActor : public NActors::TActor, public NYql:: hFunc(NFq::TEvRowDispatcher::TEvHeartbeat, Handle); hFunc(TEvPrivate::TEvPrintState, Handle); hFunc(TEvPrivate::TEvProcessState, Handle); + hFunc(TEvPrivate::TEvNotifyCA, Handle); }) static constexpr char ActorName[] = "DQ_PQ_READ_ACTOR"; @@ -264,19 +297,27 @@ class TDqPqRdReadActor : public NActors::TActor, public NYql:: void CommitState(const NDqProto::TCheckpoint& checkpoint) override; void PassAway() override; i64 GetAsyncInputData(NKikimr::NMiniKQL::TUnboxedValueBatch& buffer, TMaybe& watermark, bool&, i64 freeSpace) override; + TDuration GetCpuTime() override; std::vector GetPartitionsToRead() const; - std::pair CreateItem(const TString& data); + void AddMessageBatch(TRope&& serializedBatch, NKikimr::NMiniKQL::TUnboxedValueBatch& buffer); void ProcessState(); - void Stop(const TString& message); - void StopSessions(); + void StopSession(TSession& sessionInfo); + void Stop(NDqProto::StatusIds::StatusCode status, TIssues issues); void ReInit(const TString& reason); void PrintInternalState(); + void TrySendGetNextBatch(TSession& sessionInfo); TString GetInternalState(); - void TrySendGetNextBatch(SessionInfo& sessionInfo); template - bool CheckSession(SessionInfo& session, const TEventPtr& ev, ui64 partitionId); - void SendStopSession(const NActors::TActorId& recipient, ui64 partitionId, ui64 cookie); + TSession* FindAndUpdateSession(const TEventPtr& ev); + void SendNoSession(const NActors::TActorId& recipient, ui64 cookie); void NotifyCA(); + void SendStartSession(TSession& sessionInfo); + void Init(); + void ScheduleProcessState(); + void ProcessGlobalState(); + void ProcessSessionsState(); + void UpdateSessions(); + void UpdateQueuedSize(); }; TDqPqRdReadActor::TDqPqRdReadActor( @@ -285,6 +326,7 @@ TDqPqRdReadActor::TDqPqRdReadActor( const TTxId& txId, ui64 taskId, const THolderFactory& holderFactory, + const TTypeEnvironment& typeEnv, NPq::NProto::TDqPqTopicSource&& sourceParams, NPq::NProto::TDqReadTaskParams&& readParams, const NActors::TActorId& computeActorId, @@ -300,43 +342,71 @@ TDqPqRdReadActor::TDqPqRdReadActor( , HolderFactory(holderFactory) , MaxBufferSize(bufferSize) { - MetadataFields.reserve(SourceParams.MetadataFieldsSize()); - TPqMetaExtractor fieldsExtractor; - for (const auto& fieldName : SourceParams.GetMetadataFields()) { - MetadataFields.emplace_back(fieldName, fieldsExtractor.FindExtractorLambda(fieldName)); + const auto programBuilder = std::make_unique(typeEnv, *holderFactory.GetFunctionRegistry()); + + // Parse output schema (expected struct output type) + const auto& outputTypeYson = SourceParams.GetRowType(); + const auto outputItemType = NCommon::ParseTypeFromYson(TStringBuf(outputTypeYson), *programBuilder, Cerr); + YQL_ENSURE(outputItemType, "Failed to parse output type: " << outputTypeYson); + YQL_ENSURE(outputItemType->IsStruct(), "Output type " << outputTypeYson << " is not struct"); + const auto structType = static_cast(outputItemType); + + // Build input schema and unpacker (for data comes from RowDispatcher) + TVector inputTypeParts; + inputTypeParts.reserve(SourceParams.ColumnsSize()); + ColumnIndexes.resize(structType->GetMembersCount()); + for (size_t i = 0; i < SourceParams.ColumnsSize(); ++i) { + const auto index = structType->GetMemberIndex(SourceParams.GetColumns().Get(i)); + inputTypeParts.emplace_back(structType->GetMemberType(index)); + ColumnIndexes[index] = i; } + InputDataType = programBuilder->NewMultiType(inputTypeParts); + DataUnpacker = std::make_unique>(InputDataType); IngressStats.Level = statsLevel; - SRC_LOG_I("Start read actor, local row dispatcher " << LocalRowDispatcherActorId.ToString() << ", metadatafields: " << JoinSeq(',', SourceParams.GetMetadataFields())); + SRC_LOG_I("Start read actor, local row dispatcher " << LocalRowDispatcherActorId.ToString() << ", metadatafields: " << JoinSeq(',', SourceParams.GetMetadataFields()) + << ", partitions: " << JoinSeq(',', GetPartitionsToRead())); } -void TDqPqRdReadActor::ProcessState() { +void TDqPqRdReadActor::Init() { + if (Inited) { + return; + } + LogPrefix = (TStringBuilder() << "SelfId: " << SelfId() << ", TxId: " << TxId << ", task: " << TaskId << ". PQ source. "); + + auto partitionToRead = GetPartitionsToRead(); + for (auto partitionId : partitionToRead) { + TPartitionKey partitionKey{TString{}, partitionId}; + const auto offsetIt = PartitionToOffset.find(partitionKey); + auto& nextOffset = NextOffsetFromRD[partitionId]; + if (offsetIt != PartitionToOffset.end()) { + nextOffset = offsetIt->second; + } + } + SRC_LOG_I("Send TEvCoordinatorChangesSubscribe to local RD (" << LocalRowDispatcherActorId << ")"); + Send(LocalRowDispatcherActorId, new NFq::TEvRowDispatcher::TEvCoordinatorChangesSubscribe()); + + Schedule(TDuration::Seconds(PrintStatePeriodSec), new TEvPrivate::TEvPrintState()); + Schedule(TDuration::Seconds(NotifyCAPeriodSec), new TEvPrivate::TEvNotifyCA()); + Inited = true; +} + +void TDqPqRdReadActor::ProcessGlobalState() { switch (State) { case EState::INIT: - LogPrefix = (TStringBuilder() << "SelfId: " << SelfId() << ", TxId: " << TxId << ", task: " << TaskId << ". PQ source. "); - if (!ReadyBuffer.empty()) { return; } - if (!ProcessStateScheduled) { - ProcessStateScheduled = true; - Schedule(TDuration::Seconds(ProcessStatePeriodSec), new TEvPrivate::TEvProcessState()); - } if (!CoordinatorActorId) { SRC_LOG_I("Send TEvCoordinatorChangesSubscribe to local row dispatcher, self id " << SelfId()); Send(LocalRowDispatcherActorId, new NFq::TEvRowDispatcher::TEvCoordinatorChangesSubscribe()); - if (!SchedulePrintStatePeriod) { - SchedulePrintStatePeriod = true; - Schedule(TDuration::Seconds(PrintStatePeriodSec), new TEvPrivate::TEvPrintState()); - } + State = EState::WAIT_COORDINATOR_ID; } - State = EState::WAIT_COORDINATOR_ID; [[fallthrough]]; case EState::WAIT_COORDINATOR_ID: { if (!CoordinatorActorId) { return; } - State = EState::WAIT_PARTITIONS_ADDRES; auto partitionToRead = GetPartitionsToRead(); auto cookie = ++CoordinatorRequestCookie; SRC_LOG_I("Send TEvCoordinatorRequest to coordinator " << CoordinatorActorId->ToString() << ", partIds: " @@ -346,99 +416,123 @@ void TDqPqRdReadActor::ProcessState() { new NFq::TEvRowDispatcher::TEvCoordinatorRequest(SourceParams, partitionToRead), IEventHandle::FlagTrackDelivery | IEventHandle::FlagSubscribeOnSession, cookie); - return; + LastReceivedPartitionDistribution.clear(); + State = EState::WAIT_PARTITIONS_ADDRES; + [[fallthrough]]; } case EState::WAIT_PARTITIONS_ADDRES: - if (Sessions.empty()) { - return; - } - - for (auto& [partitionId, sessionInfo] : Sessions) { - if (sessionInfo.Status == SessionInfo::ESessionStatus::NoSession) { - TMaybe readOffset; - TPartitionKey partitionKey{TString{}, partitionId}; - const auto offsetIt = PartitionToOffset.find(partitionKey); - if (offsetIt != PartitionToOffset.end()) { - SRC_LOG_D("ReadOffset found" ); - readOffset = offsetIt->second; - } - - SRC_LOG_I("Send TEvStartSession to " << sessionInfo.RowDispatcherActorId - << ", offset " << readOffset - << ", partitionId " << partitionId - << ", connection id " << sessionInfo.Generation); - - auto event = new NFq::TEvRowDispatcher::TEvStartSession( - SourceParams, - partitionId, - Token, - readOffset, - StartingMessageTimestamp.MilliSeconds(), - std::visit([](auto arg) { return ToString(arg); }, TxId)); - sessionInfo.EventsQueue.Send(event, sessionInfo.Generation); - sessionInfo.IsWaitingStartSessionAck = true; - sessionInfo.Status = SessionInfo::ESessionStatus::Started; - } + if (LastReceivedPartitionDistribution.empty()) { + break; } + UpdateSessions(); State = EState::STARTED; - return; + [[fallthrough]]; case EState::STARTED: - return; + break; } } +void TDqPqRdReadActor::ProcessSessionsState() { + for (auto& [rowDispatcherActorId, sessionInfo] : Sessions) { + switch (sessionInfo.Status) { + case TSession::ESessionStatus::INIT: + SendStartSession(sessionInfo); + sessionInfo.IsWaitingStartSessionAck = true; + sessionInfo.Status = TSession::ESessionStatus::WAIT_START_SESSION_ACK; + [[fallthrough]]; + case TSession::ESessionStatus::WAIT_START_SESSION_ACK: + if (sessionInfo.IsWaitingStartSessionAck) { + break; + } + sessionInfo.Status = TSession::ESessionStatus::STARTED; + [[fallthrough]]; + case TSession::ESessionStatus::STARTED: + break; + } + } +} -void TDqPqRdReadActor::CommitState(const NDqProto::TCheckpoint& /*checkpoint*/) { +void TDqPqRdReadActor::ProcessState() { + ProcessGlobalState(); + ProcessSessionsState(); } -void TDqPqRdReadActor::StopSessions() { - SRC_LOG_I("Stop all session"); - for (auto& [partitionId, sessionInfo] : Sessions) { - if (sessionInfo.Status == SessionInfo::ESessionStatus::NoSession) { +void TDqPqRdReadActor::SendStartSession(TSession& sessionInfo) { + auto str = TStringBuilder() << "Send TEvStartSession to " << sessionInfo.RowDispatcherActorId + << ", connection id " << sessionInfo.Generation << " partitions offsets "; + + std::set partitions; + std::map partitionOffsets; + for (auto& [partitionId, partition] : sessionInfo.Partitions) { + partitions.insert(partitionId); + auto nextOffset = NextOffsetFromRD[partitionId]; + str << "(" << partitionId << " / "; + if (!nextOffset) { + str << "),"; continue; } - auto event = std::make_unique(); - *event->Record.MutableSource() = SourceParams; - event->Record.SetPartitionId(partitionId); - SRC_LOG_I("Send StopSession to " << sessionInfo.RowDispatcherActorId); - sessionInfo.EventsQueue.Send(event.release(), sessionInfo.Generation); + partitionOffsets[partitionId] = *nextOffset; + str << nextOffset << "),"; } + SRC_LOG_I(str); + + auto event = new NFq::TEvRowDispatcher::TEvStartSession( + SourceParams, + partitions, + Token, + partitionOffsets, + StartingMessageTimestamp.MilliSeconds(), + std::visit([](auto arg) { return ToString(arg); }, TxId)); + sessionInfo.EventsQueue.Send(event, sessionInfo.Generation); +} + +void TDqPqRdReadActor::CommitState(const NDqProto::TCheckpoint& /*checkpoint*/) { +} + +void TDqPqRdReadActor::StopSession(TSession& sessionInfo) { + SRC_LOG_I("Send StopSession to " << sessionInfo.RowDispatcherActorId + << " generation " << sessionInfo.Generation); + auto event = std::make_unique(); + *event->Record.MutableSource() = SourceParams; + sessionInfo.EventsQueue.Send(event.release(), sessionInfo.Generation); } // IActor & IDqComputeActorAsyncInput void TDqPqRdReadActor::PassAway() { // Is called from Compute Actor SRC_LOG_I("PassAway"); PrintInternalState(); - StopSessions(); + for (auto& [rowDispatcherActorId, sessionInfo] : Sessions) { + StopSession(sessionInfo); + } TActor::PassAway(); - + // TODO: RetryQueue::Unsubscribe() } i64 TDqPqRdReadActor::GetAsyncInputData(NKikimr::NMiniKQL::TUnboxedValueBatch& buffer, TMaybe& /*watermark*/, bool&, i64 freeSpace) { SRC_LOG_T("GetAsyncInputData freeSpace = " << freeSpace); + Init(); Metrics.InFlyAsyncInputData->Set(0); InFlyAsyncInputData = false; - ProcessState(); if (ReadyBuffer.empty() || !freeSpace) { return 0; } i64 usedSpace = 0; buffer.clear(); do { - auto& readyBatch = ReadyBuffer.front(); + auto readyBatch = std::move(ReadyBuffer.front()); + ReadyBuffer.pop(); - for (const auto& message : readyBatch.Data) { - auto [item, size] = CreateItem(message); - buffer.push_back(std::move(item)); + for (auto& messageBatch : readyBatch.Data) { + AddMessageBatch(std::move(messageBatch), buffer); } + usedSpace += readyBatch.UsedSpace; freeSpace -= readyBatch.UsedSpace; TPartitionKey partitionKey{TString{}, readyBatch.PartitionId}; PartitionToOffset[partitionKey] = readyBatch.NextOffset; SRC_LOG_T("NextOffset " << readyBatch.NextOffset); - ReadyBuffer.pop(); } while (freeSpace > 0 && !ReadyBuffer.empty()); ReadyBufferSizeBytes -= usedSpace; @@ -447,16 +541,18 @@ i64 TDqPqRdReadActor::GetAsyncInputData(NKikimr::NMiniKQL::TUnboxedValueBatch& b if (!ReadyBuffer.empty()) { NotifyCA(); } - for (auto& [partitionId, sessionInfo] : Sessions) { + for (auto& [rowDispatcherActorId, sessionInfo] : Sessions) { TrySendGetNextBatch(sessionInfo); } - ProcessState(); return usedSpace; } +TDuration TDqPqRdReadActor::GetCpuTime() { + return TDuration::MicroSeconds(CpuMicrosec); +} + std::vector TDqPqRdReadActor::GetPartitionsToRead() const { std::vector res; - ui64 currentPartition = ReadParams.GetPartitioningParams().GetEachTopicPartitionGroupId(); do { res.emplace_back(currentPartition); // 0-based in topic API @@ -467,71 +563,59 @@ std::vector TDqPqRdReadActor::GetPartitionsToRead() const { void TDqPqRdReadActor::Handle(NFq::TEvRowDispatcher::TEvStartSessionAck::TPtr& ev) { const NYql::NDqProto::TMessageTransportMeta& meta = ev->Get()->Record.GetTransportMeta(); - SRC_LOG_I("TEvStartSessionAck from " << ev->Sender << ", seqNo " << meta.GetSeqNo() << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo()); + SRC_LOG_I("Received TEvStartSessionAck from " << ev->Sender << ", seqNo " << meta.GetSeqNo() << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo() << ", generation " << ev->Cookie); Counters.StartSessionAck++; - - ui64 partitionId = ev->Get()->Record.GetConsumer().GetPartitionId(); - auto sessionIt = Sessions.find(partitionId); - if (sessionIt == Sessions.end()) { - SRC_LOG_W("Ignore TEvStartSessionAck from " << ev->Sender << ", seqNo " << meta.GetSeqNo() - << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo() << ", PartitionId " << partitionId << ", cookie " << ev->Cookie); - YQL_ENSURE(State != EState::STARTED); - SendStopSession(ev->Sender, partitionId, ev->Cookie); + auto* session = FindAndUpdateSession(ev); + if (!session) { return; } - auto& sessionInfo = sessionIt->second; - if (!CheckSession(sessionInfo, ev, partitionId)) { - return; - } - sessionInfo.IsWaitingStartSessionAck = false; + session->IsWaitingStartSessionAck = false; + ProcessState(); } void TDqPqRdReadActor::Handle(NFq::TEvRowDispatcher::TEvSessionError::TPtr& ev) { const NYql::NDqProto::TMessageTransportMeta& meta = ev->Get()->Record.GetTransportMeta(); - SRC_LOG_I("TEvSessionError from " << ev->Sender << ", seqNo " << meta.GetSeqNo() << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo()); + SRC_LOG_I("Received TEvSessionError from " << ev->Sender << ", seqNo " << meta.GetSeqNo() << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo()); Counters.SessionError++; - ui64 partitionId = ev->Get()->Record.GetPartitionId(); - auto sessionIt = Sessions.find(partitionId); - if (sessionIt == Sessions.end()) { - SRC_LOG_W("Ignore TEvSessionError from " << ev->Sender << ", seqNo " << meta.GetSeqNo() - << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo() << ", PartitionId " << partitionId << ", cookie " << ev->Cookie); - YQL_ENSURE(State != EState::STARTED); - SendStopSession(ev->Sender, partitionId, ev->Cookie); + auto* session = FindAndUpdateSession(ev); + if (!session) { return; } - - auto& sessionInfo = sessionIt->second; - if (!CheckSession(sessionInfo, ev, partitionId)) { - return; - } - Stop(ev->Get()->Record.GetMessage()); + NYql::TIssues issues; + IssuesFromMessage(ev->Get()->Record.GetIssues(), issues); + Stop(ev->Get()->Record.GetStatusCode(), issues); } void TDqPqRdReadActor::Handle(NFq::TEvRowDispatcher::TEvStatistics::TPtr& ev) { const NYql::NDqProto::TMessageTransportMeta& meta = ev->Get()->Record.GetTransportMeta(); - SRC_LOG_T("TEvStatistics from " << ev->Sender << ", offset " << ev->Get()->Record.GetNextMessageOffset() << ", seqNo " << meta.GetSeqNo() << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo()); + SRC_LOG_T("Received TEvStatistics from " << ev->Sender << ", seqNo " << meta.GetSeqNo() << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo() << " generation " << ev->Cookie); Counters.Statistics++; - - ui64 partitionId = ev->Get()->Record.GetPartitionId(); - auto sessionIt = Sessions.find(partitionId); - if (sessionIt == Sessions.end()) { - SRC_LOG_W("Ignore TEvStatistics from " << ev->Sender << ", seqNo " << meta.GetSeqNo() - << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo() << ", PartitionId " << partitionId); - YQL_ENSURE(State != EState::STARTED); - SendStopSession(ev->Sender, partitionId, ev->Cookie); + CpuMicrosec += ev->Get()->Record.GetCpuMicrosec(); + auto* session = FindAndUpdateSession(ev); + if (!session) { return; } - auto& sessionInfo = sessionIt->second; IngressStats.Bytes += ev->Get()->Record.GetReadBytes(); - - if (!CheckSession(sessionInfo, ev, partitionId)) { - return; - } - - if (ReadyBuffer.empty()) { - TPartitionKey partitionKey{TString{}, partitionId}; - PartitionToOffset[partitionKey] = ev->Get()->Record.GetNextMessageOffset(); + IngressStats.FilteredBytes += ev->Get()->Record.GetFilteredBytes(); + IngressStats.FilteredRows += ev->Get()->Record.GetFilteredRows(); + session->QueuedBytes = ev->Get()->Record.GetQueuedBytes(); + session->QueuedRows = ev->Get()->Record.GetQueuedRows(); + UpdateQueuedSize(); + + for (auto partition : ev->Get()->Record.GetPartition()) { + ui64 partitionId = partition.GetPartitionId(); + auto& nextOffset = NextOffsetFromRD[partitionId]; + if (!nextOffset) { + nextOffset = partition.GetNextMessageOffset(); + } else { + nextOffset = std::max(*nextOffset, partition.GetNextMessageOffset()); + } + SRC_LOG_T("NextOffsetFromRD [" << partitionId << "]= " << nextOffset); + if (ReadyBuffer.empty()) { + TPartitionKey partitionKey{TString{}, partitionId}; + PartitionToOffset[partitionKey] = *nextOffset; + } } } @@ -543,35 +627,36 @@ void TDqPqRdReadActor::Handle(NFq::TEvRowDispatcher::TEvGetInternalStateRequest: void TDqPqRdReadActor::Handle(NFq::TEvRowDispatcher::TEvNewDataArrived::TPtr& ev) { const NYql::NDqProto::TMessageTransportMeta& meta = ev->Get()->Record.GetTransportMeta(); - SRC_LOG_T("TEvNewDataArrived from " << ev->Sender << ", part id " << ev->Get()->Record.GetPartitionId() << ", seqNo " << meta.GetSeqNo() << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo()); + SRC_LOG_T("Received TEvNewDataArrived from " << ev->Sender << ", partition " << ev->Get()->Record.GetPartitionId() << ", seqNo " << meta.GetSeqNo() << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo() << " generation " << ev->Cookie); Counters.NewDataArrived++; - - ui64 partitionId = ev->Get()->Record.GetPartitionId(); - auto sessionIt = Sessions.find(partitionId); - if (sessionIt == Sessions.end()) { - SRC_LOG_W("Ignore TEvNewDataArrived from " << ev->Sender << ", seqNo " << meta.GetSeqNo() - << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo() << ", PartitionId " << partitionId); - YQL_ENSURE(State != EState::STARTED); - SendStopSession(ev->Sender, partitionId, ev->Cookie); + + auto* session = FindAndUpdateSession(ev); + if (!session) { return; } - - auto& sessionInfo = sessionIt->second; - if (!CheckSession(sessionInfo, ev, partitionId)) { + auto partitionIt = session->Partitions.find(ev->Get()->Record.GetPartitionId()); + if (partitionIt == session->Partitions.end()) { + SRC_LOG_E("Received TEvNewDataArrived from " << ev->Sender << " with wrong partition id " << ev->Get()->Record.GetPartitionId()); + Stop(NDqProto::StatusIds::INTERNAL_ERROR, {TIssue(TStringBuilder() << LogPrefix << "No partition with id " << ev->Get()->Record.GetPartitionId())}); return; } - sessionInfo.HasPendingData = true; - TrySendGetNextBatch(sessionInfo); + partitionIt->second.HasPendingData = true; + TrySendGetNextBatch(*session); } void TDqPqRdReadActor::Handle(const NYql::NDq::TEvRetryQueuePrivate::TEvRetry::TPtr& ev) { - SRC_LOG_D("TEvRetry"); + SRC_LOG_T("Received TEvRetry, EventQueueId " << ev->Get()->EventQueueId); Counters.Retry++; - ui64 partitionId = ev->Get()->EventQueueId; - auto sessionIt = Sessions.find(partitionId); + auto readActorIt = ReadActorByEventQueueId.find(ev->Get()->EventQueueId); + if (readActorIt == ReadActorByEventQueueId.end()) { + SRC_LOG_D("Ignore TEvRetry, wrong EventQueueId " << ev->Get()->EventQueueId); + return; + } + + auto sessionIt = Sessions.find(readActorIt->second); if (sessionIt == Sessions.end()) { - SRC_LOG_W("Unknown partition id " << partitionId << ", skip TEvRetry"); + SRC_LOG_D("Ignore TEvRetry, wrong read actor id " << readActorIt->second); return; } sessionIt->second.EventsQueue.Retry(); @@ -580,35 +665,29 @@ void TDqPqRdReadActor::Handle(const NYql::NDq::TEvRetryQueuePrivate::TEvRetry::T void TDqPqRdReadActor::Handle(const NYql::NDq::TEvRetryQueuePrivate::TEvEvHeartbeat::TPtr& ev) { SRC_LOG_T("TEvRetryQueuePrivate::TEvEvHeartbeat"); Counters.PrivateHeartbeat++; - ui64 partitionId = ev->Get()->EventQueueId; + auto readActorIt = ReadActorByEventQueueId.find(ev->Get()->EventQueueId); + if (readActorIt == ReadActorByEventQueueId.end()) { + SRC_LOG_D("Ignore TEvEvHeartbeat, wrong EventQueueId " << ev->Get()->EventQueueId); + return; + } - auto sessionIt = Sessions.find(partitionId); + auto sessionIt = Sessions.find(readActorIt->second); if (sessionIt == Sessions.end()) { - SRC_LOG_W("Unknown partition id " << partitionId << ", skip TEvPing"); + SRC_LOG_D("Ignore TEvEvHeartbeat, wrong read actor id " << readActorIt->second); return; } auto& sessionInfo = sessionIt->second; bool needSend = sessionInfo.EventsQueue.Heartbeat(); if (needSend) { SRC_LOG_T("Send TEvEvHeartbeat"); - sessionInfo.EventsQueue.Send(new NFq::TEvRowDispatcher::TEvHeartbeat(sessionInfo.PartitionId), sessionInfo.Generation); + sessionInfo.EventsQueue.Send(new NFq::TEvRowDispatcher::TEvHeartbeat(), sessionInfo.Generation); } } void TDqPqRdReadActor::Handle(const NFq::TEvRowDispatcher::TEvHeartbeat::TPtr& ev) { - SRC_LOG_T("Received TEvHeartbeat from " << ev->Sender); + SRC_LOG_T("Received TEvHeartbeat from " << ev->Sender << ", generation " << ev->Cookie); Counters.Heartbeat++; - const NYql::NDqProto::TMessageTransportMeta& meta = ev->Get()->Record.GetTransportMeta(); - - ui64 partitionId = ev->Get()->Record.GetPartitionId(); - auto sessionIt = Sessions.find(partitionId); - if (sessionIt == Sessions.end()) { - SRC_LOG_W("Ignore TEvHeartbeat from " << ev->Sender << ", seqNo " << meta.GetSeqNo() - << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo() << ", PartitionId " << partitionId << ", cookie " << ev->Cookie); - SendStopSession(ev->Sender, partitionId, ev->Cookie); - return; - } - CheckSession(sessionIt->second, ev, partitionId); + FindAndUpdateSession(ev); } void TDqPqRdReadActor::Handle(NFq::TEvRowDispatcher::TEvCoordinatorChanged::TPtr& ev) { @@ -628,50 +707,45 @@ void TDqPqRdReadActor::Handle(NFq::TEvRowDispatcher::TEvCoordinatorChanged::TPtr CoordinatorActorId = ev->Get()->CoordinatorActorId; ReInit("Coordinator is changed"); - ProcessState(); + ScheduleProcessState(); +} + +void TDqPqRdReadActor::ScheduleProcessState() { + if (ProcessStateScheduled) { + return; + } + ProcessStateScheduled = true; + Schedule(TDuration::Seconds(ProcessStatePeriodSec), new TEvPrivate::TEvProcessState()); } void TDqPqRdReadActor::ReInit(const TString& reason) { SRC_LOG_I("ReInit state, reason " << reason); - StopSessions(); - Sessions.clear(); - State = EState::INIT; + Metrics.ReInit->Inc(); + + State = EState::WAIT_COORDINATOR_ID; if (!ReadyBuffer.empty()) { NotifyCA(); } - PrintInternalState(); } -void TDqPqRdReadActor::Stop(const TString& message) { - NYql::TIssues issues; - issues.AddIssue(NYql::TIssue{message}); - SRC_LOG_E("Stop read actor, error: " << message); - Send(ComputeActorId, new TEvAsyncInputError(InputIndex, issues, NYql::NDqProto::StatusIds::BAD_REQUEST)); // TODO: use UNAVAILABLE ? +void TDqPqRdReadActor::Stop(NDqProto::StatusIds::StatusCode status, TIssues issues) { + SRC_LOG_E("Stop read actor, status: " << NYql::NDqProto::StatusIds_StatusCode_Name(status) << ", issues: " << issues.ToOneLineString()); + Send(ComputeActorId, new TEvAsyncInputError(InputIndex, std::move(issues), status)); } void TDqPqRdReadActor::Handle(NFq::TEvRowDispatcher::TEvCoordinatorResult::TPtr& ev) { - SRC_LOG_I("TEvCoordinatorResult from " << ev->Sender.ToString() << ", cookie " << ev->Cookie); + SRC_LOG_I("Received TEvCoordinatorResult from " << ev->Sender.ToString() << ", cookie " << ev->Cookie); Counters.CoordinatorChanged++; if (ev->Cookie != CoordinatorRequestCookie) { SRC_LOG_W("Ignore TEvCoordinatorResult. wrong cookie"); return; } - if (State != EState::WAIT_PARTITIONS_ADDRES) { - SRC_LOG_W("Ignore TEvCoordinatorResult. wrong state " << static_cast(EState::WAIT_PARTITIONS_ADDRES)); - return; - } + LastReceivedPartitionDistribution.clear(); + TMap> distribution; for (auto& p : ev->Get()->Record.GetPartitions()) { TActorId rowDispatcherActorId = ActorIdFromProto(p.GetActorId()); - for (auto partitionId : p.GetPartitionId()) { - if (Sessions.contains(partitionId)) { - Stop("Internal error: session already exists"); - return; - } - SRC_LOG_I("Create session to RD (" << rowDispatcherActorId << "), partitionId " << partitionId); - Sessions.emplace( - std::piecewise_construct, - std::forward_as_tuple(partitionId), - std::forward_as_tuple(TxId, SelfId(), rowDispatcherActorId, partitionId, partitionId, ++NextGeneration)); + for (auto partitionId : p.GetPartitionIds()) { + LastReceivedPartitionDistribution[rowDispatcherActorId].insert(partitionId); } } ProcessState(); @@ -680,7 +754,7 @@ void TDqPqRdReadActor::Handle(NFq::TEvRowDispatcher::TEvCoordinatorResult::TPtr& void TDqPqRdReadActor::HandleConnected(TEvInterconnect::TEvNodeConnected::TPtr& ev) { SRC_LOG_D("EvNodeConnected " << ev->Get()->NodeId); Counters.NodeConnected++; - for (auto& [partitionId, sessionInfo] : Sessions) { + for (auto& [rowDispatcherActorId, sessionInfo] : Sessions) { sessionInfo.EventsQueue.HandleNodeConnected(ev->Get()->NodeId); } } @@ -688,7 +762,7 @@ void TDqPqRdReadActor::HandleConnected(TEvInterconnect::TEvNodeConnected::TPtr& void TDqPqRdReadActor::HandleDisconnected(TEvInterconnect::TEvNodeDisconnected::TPtr& ev) { SRC_LOG_D("TEvNodeDisconnected, node id " << ev->Get()->NodeId); Counters.NodeDisconnected++; - for (auto& [partitionId, sessionInfo] : Sessions) { + for (auto& [rowDispatcherActorId, sessionInfo] : Sessions) { sessionInfo.EventsQueue.HandleNodeDisconnected(ev->Get()->NodeId); } // In case of row dispatcher disconnection: wait connected or SessionClosed(). TODO: Stop actor after timeout. @@ -697,76 +771,97 @@ void TDqPqRdReadActor::HandleDisconnected(TEvInterconnect::TEvNodeDisconnected:: } void TDqPqRdReadActor::Handle(NActors::TEvents::TEvUndelivered::TPtr& ev) { - SRC_LOG_D("TEvUndelivered, " << ev->Get()->ToString() << " from " << ev->Sender.ToString()); + SRC_LOG_D("Received TEvUndelivered, " << ev->Get()->ToString() << " from " << ev->Sender.ToString() << ", reason " << ev->Get()->Reason << ", cookie " << ev->Cookie); Counters.Undelivered++; - for (auto& [partitionId, sessionInfo] : Sessions) { + + auto sessionIt = Sessions.find(ev->Sender); + if (sessionIt != Sessions.end()) { + auto& sessionInfo = sessionIt->second; if (sessionInfo.EventsQueue.HandleUndelivered(ev) == NYql::NDq::TRetryEventsQueue::ESessionState::SessionClosed) { - ReInit(TStringBuilder() << "Session closed, partition id " << sessionInfo.PartitionId); - break; + if (sessionInfo.Generation == ev->Cookie) { + SRC_LOG_D("Erase session to " << ev->Sender.ToString()); + ReadActorByEventQueueId.erase(sessionInfo.EventQueueId); + Sessions.erase(sessionIt); + ReInit("Reset session state (by TEvUndelivered)"); + ScheduleProcessState(); + } } } if (CoordinatorActorId && *CoordinatorActorId == ev->Sender) { ReInit("TEvUndelivered to coordinator"); + ScheduleProcessState(); + return; } } void TDqPqRdReadActor::Handle(NFq::TEvRowDispatcher::TEvMessageBatch::TPtr& ev) { const NYql::NDqProto::TMessageTransportMeta& meta = ev->Get()->Record.GetTransportMeta(); - SRC_LOG_T("TEvMessageBatch from " << ev->Sender << ", seqNo " << meta.GetSeqNo() << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo()); + SRC_LOG_T("Received TEvMessageBatch from " << ev->Sender << ", seqNo " << meta.GetSeqNo() << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo() << " generation " << ev->Cookie); Counters.MessageBatch++; - ui64 partitionId = ev->Get()->Record.GetPartitionId(); - auto sessionIt = Sessions.find(partitionId); - if (sessionIt == Sessions.end()) { - SRC_LOG_W("Ignore TEvMessageBatch from " << ev->Sender << ", seqNo " << meta.GetSeqNo() - << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo() << ", PartitionId " << partitionId << ", cookie " << ev->Cookie); - YQL_ENSURE(State != EState::STARTED); - SendStopSession(ev->Sender, partitionId, ev->Cookie); + auto* session = FindAndUpdateSession(ev); + if (!session) { return; } - - Metrics.InFlyGetNextBatch->Set(0); - auto& sessionInfo = sessionIt->second; - if (!CheckSession(sessionInfo, ev, partitionId)) { + auto partitionId = ev->Get()->Record.GetPartitionId(); + auto partitionIt = session->Partitions.find(partitionId); + if (partitionIt == session->Partitions.end()) { + SRC_LOG_E("TEvMessageBatch: wrong partition id " << partitionId); + Stop(NDqProto::StatusIds::INTERNAL_ERROR, {TIssue(TStringBuilder() << LogPrefix << "No partition with id " << partitionId)}); return; } + auto& partirtion = partitionIt->second; + Metrics.InFlyGetNextBatch->Set(0); ReadyBuffer.emplace(partitionId, ev->Get()->Record.MessagesSize()); TReadyBatch& activeBatch = ReadyBuffer.back(); + auto& nextOffset = NextOffsetFromRD[partitionId]; + ui64 bytes = 0; for (const auto& message : ev->Get()->Record.GetMessages()) { - SRC_LOG_T("Json: " << message.GetJson()); - activeBatch.Data.emplace_back(message.GetJson()); - sessionInfo.NextOffset = message.GetOffset() + 1; - bytes += message.GetJson().size(); - SRC_LOG_T("TEvMessageBatch NextOffset " << sessionInfo.NextOffset); + const auto& offsets = message.GetOffsets(); + if (offsets.empty()) { + Stop(NDqProto::StatusIds::INTERNAL_ERROR, {TIssue(TStringBuilder() << LogPrefix << "Got unexpected empty batch from row dispatcher")}); + return; + } + + activeBatch.Data.emplace_back(ev->Get()->GetPayload(message.GetPayloadId())); + bytes += activeBatch.Data.back().GetSize(); + + nextOffset = *offsets.rbegin() + 1; + SRC_LOG_T("TEvMessageBatch NextOffset " << nextOffset); } activeBatch.UsedSpace = bytes; ReadyBufferSizeBytes += bytes; activeBatch.NextOffset = ev->Get()->Record.GetNextMessageOffset(); - sessionInfo.IsWaitingMessageBatch = false; + partirtion.IsWaitingMessageBatch = false; NotifyCA(); } -std::pair TDqPqRdReadActor::CreateItem(const TString& data) { - i64 usedSpace = 0; - NUdf::TUnboxedValuePod item; - if (MetadataFields.empty()) { - item = NKikimr::NMiniKQL::MakeString(NUdf::TStringRef(data.data(), data.size())); - usedSpace += data.size(); - return std::make_pair(item, usedSpace); - } +void TDqPqRdReadActor::AddMessageBatch(TRope&& messageBatch, NKikimr::NMiniKQL::TUnboxedValueBatch& buffer) { + // TDOD: pass multi type directly to CA, without transforming it into struct + + NKikimr::NMiniKQL::TUnboxedValueBatch parsedData(InputDataType); + DataUnpacker->UnpackBatch(MakeChunkedBuffer(std::move(messageBatch)), HolderFactory, parsedData); - NUdf::TUnboxedValue* itemPtr; - item = HolderFactory.CreateDirectArrayHolder(MetadataFields.size() + 1, itemPtr); - *(itemPtr++) = NKikimr::NMiniKQL::MakeString(NUdf::TStringRef(data.data(), data.size())); - usedSpace += data.size(); + while (!parsedData.empty()) { + const auto* parsedRow = parsedData.Head(); + + NUdf::TUnboxedValue* itemPtr; + NUdf::TUnboxedValuePod item = HolderFactory.CreateDirectArrayHolder(ColumnIndexes.size(), itemPtr); + for (const auto index : ColumnIndexes) { + if (index) { + YQL_ENSURE(*index < parsedData.Width(), "Unexpected data width " << parsedData.Width() << ", failed to extract column by index " << index); + *(itemPtr++) = parsedRow[*index]; + } else { + // TODO: support metadata fields here + *(itemPtr++) = NUdf::TUnboxedValuePod::Zero(); + } + } - for ([[maybe_unused]] const auto& [name, extractor] : MetadataFields) { - auto ub = NYql::NUdf::TUnboxedValuePod(0); // TODO: use real values - *(itemPtr++) = std::move(ub); + buffer.emplace_back(std::move(item)); + parsedData.Pop(); } - return std::make_pair(item, usedSpace); } void TDqPqRdReadActor::Handle(NActors::TEvents::TEvPong::TPtr& ev) { @@ -781,12 +876,16 @@ void TDqPqRdReadActor::Handle(TEvPrivate::TEvPrintState::TPtr&) { } void TDqPqRdReadActor::PrintInternalState() { - SRC_LOG_I(GetInternalState()); + auto str = GetInternalState(); + auto buf = TStringBuf(str); + for (ui64 offset = 0; offset < buf.size(); offset += PrintStateToLogSplitSize) { + SRC_LOG_I(buf.SubString(offset, PrintStateToLogSplitSize)); + } } TString TDqPqRdReadActor::GetInternalState() { TStringStream str; - str << "State: used buffer size " << ReadyBufferSizeBytes << " ready buffer event size " << ReadyBuffer.size() << " state " << static_cast(State) << " InFlyAsyncInputData " << InFlyAsyncInputData << "\n"; + str << LogPrefix << "State: used buffer size " << ReadyBufferSizeBytes << " ready buffer event size " << ReadyBuffer.size() << " state " << static_cast(State) << " InFlyAsyncInputData " << InFlyAsyncInputData << "\n"; str << "Counters: GetAsyncInputData " << Counters.GetAsyncInputData << " CoordinatorChanged " << Counters.CoordinatorChanged << " CoordinatorResult " << Counters.CoordinatorResult << " MessageBatch " << Counters.MessageBatch << " StartSessionAck " << Counters.StartSessionAck << " NewDataArrived " << Counters.NewDataArrived << " SessionError " << Counters.SessionError << " Statistics " << Counters.Statistics << " NodeDisconnected " << Counters.NodeDisconnected @@ -795,56 +894,80 @@ TString TDqPqRdReadActor::GetInternalState() { << " Heartbeat " << Counters.Heartbeat << " PrintState " << Counters.PrintState << " ProcessState " << Counters.ProcessState << " NotifyCA " << Counters.NotifyCA << "\n"; - for (auto& [partitionId, sessionInfo] : Sessions) { - str << " partId " << partitionId << " status " << static_cast(sessionInfo.Status) - << " next offset " << sessionInfo.NextOffset - << " is waiting ack " << sessionInfo.IsWaitingStartSessionAck << " is waiting batch " << sessionInfo.IsWaitingMessageBatch - << " has pending data " << sessionInfo.HasPendingData << " connection id " << sessionInfo.Generation << " "; + for (auto& [rowDispatcherActorId, sessionInfo] : Sessions) { + str << " " << rowDispatcherActorId << " status " << static_cast(sessionInfo.Status) + << " is waiting ack " << sessionInfo.IsWaitingStartSessionAck << " connection id " << sessionInfo.Generation << " "; sessionInfo.EventsQueue.PrintInternalState(str); + for (const auto& [partitionId, partition] : sessionInfo.Partitions) { + const auto offsetIt = NextOffsetFromRD.find(partitionId); + str << " partId " << partitionId + << " next offset " << ((offsetIt != NextOffsetFromRD.end()) ? ToString(offsetIt->second) : TString("")) + << " is waiting batch " << partition.IsWaitingMessageBatch + << " has pending data " << partition.HasPendingData << "\n"; + } + str << "\n"; } return str.Str(); } void TDqPqRdReadActor::Handle(TEvPrivate::TEvProcessState::TPtr&) { + ProcessStateScheduled = false; Counters.ProcessState++; - Schedule(TDuration::Seconds(ProcessStatePeriodSec), new TEvPrivate::TEvProcessState()); ProcessState(); } -void TDqPqRdReadActor::TrySendGetNextBatch(SessionInfo& sessionInfo) { - if (!sessionInfo.HasPendingData) { - return; - } +void TDqPqRdReadActor::TrySendGetNextBatch(TSession& sessionInfo) { if (ReadyBufferSizeBytes > MaxBufferSize) { return; } - Metrics.InFlyGetNextBatch->Inc(); - auto event = std::make_unique(); - sessionInfo.HasPendingData = false; - sessionInfo.IsWaitingMessageBatch = true; - event->Record.SetPartitionId(sessionInfo.PartitionId); - sessionInfo.EventsQueue.Send(event.release(), sessionInfo.Generation); + for (auto& [partitionId, partition] : sessionInfo.Partitions) { + if (!partition.HasPendingData) { + continue; + } + Metrics.InFlyGetNextBatch->Inc(); + auto event = std::make_unique(); + partition.HasPendingData = false; + partition.IsWaitingMessageBatch = true; + event->Record.SetPartitionId(partitionId); + sessionInfo.EventsQueue.Send(event.release(), sessionInfo.Generation); + } } template -bool TDqPqRdReadActor::CheckSession(SessionInfo& session, const TEventPtr& ev, ui64 partitionId) { +TDqPqRdReadActor::TSession* TDqPqRdReadActor::FindAndUpdateSession(const TEventPtr& ev) { + auto sessionIt = Sessions.find(ev->Sender); + if (sessionIt == Sessions.end()) { + SRC_LOG_W("Ignore " << typeid(TEventPtr).name() << " from " << ev->Sender); + SendNoSession(ev->Sender, ev->Cookie); + return nullptr; + } + auto& session = sessionIt->second; + if (ev->Cookie != session.Generation) { - SRC_LOG_W("Wrong message generation (" << typeid(TEventPtr).name() << "), sender " << ev->Sender << " cookie " << ev->Cookie << ", session generation " << session.Generation << ", send TEvStopSession"); - SendStopSession(ev->Sender, partitionId, ev->Cookie); - return false; + SRC_LOG_W("Wrong message generation (" << typeid(TEventPtr).name() << "), sender " << ev->Sender << " cookie " << ev->Cookie << ", session generation " << session.Generation << ", send TEvNoSession"); + SendNoSession(ev->Sender, ev->Cookie); + return nullptr; } if (!session.EventsQueue.OnEventReceived(ev)) { const NYql::NDqProto::TMessageTransportMeta& meta = ev->Get()->Record.GetTransportMeta(); - SRC_LOG_W("Wrong seq num ignore message (" << typeid(TEventPtr).name() << ") seqNo " << meta.GetSeqNo() << " from " << ev->Sender.ToString()); - return false; + SRC_LOG_W("Ignore " << typeid(TEventPtr).name() << " from " << ev->Sender << ", wrong seq num, seqNo " << meta.GetSeqNo()); + return nullptr; + } + + auto expectedStatus = TSession::ESessionStatus::STARTED; + if constexpr (std::is_same_v) { + expectedStatus = TSession::ESessionStatus::WAIT_START_SESSION_ACK; } - return true; + + if (session.Status != expectedStatus) { + SRC_LOG_E("Wrong " << typeid(TEventPtr).name() << " from " << ev->Sender << " session status " << static_cast(session.Status) << " expected " << static_cast(expectedStatus)); + return nullptr; + } + return &session; } -void TDqPqRdReadActor::SendStopSession(const NActors::TActorId& recipient, ui64 partitionId, ui64 cookie) { - auto event = std::make_unique(); - *event->Record.MutableSource() = SourceParams; - event->Record.SetPartitionId(partitionId); +void TDqPqRdReadActor::SendNoSession(const NActors::TActorId& recipient, ui64 cookie) { + auto event = std::make_unique(); Send(recipient, event.release(), 0, cookie); } @@ -855,7 +978,56 @@ void TDqPqRdReadActor::NotifyCA() { Send(ComputeActorId, new TEvNewAsyncInputDataArrived(InputIndex)); } +void TDqPqRdReadActor::UpdateSessions() { + SRC_LOG_I("UpdateSessions, Sessions size " << Sessions.size()); + + if (LastUsedPartitionDistribution != LastReceivedPartitionDistribution) { + SRC_LOG_I("Distribution is changed, remove sessions"); + for (auto& [rowDispatcherActorId, sessionInfo] : Sessions) { + StopSession(sessionInfo); + ReadActorByEventQueueId.erase(sessionInfo.EventQueueId); + } + Sessions.clear(); + } + + for (const auto& [rowDispatcherActorId, partitions] : LastReceivedPartitionDistribution) { + if (Sessions.contains(rowDispatcherActorId)) { + continue; + } + + auto queueId = ++NextEventQueueId; + Sessions.emplace( + std::piecewise_construct, + std::forward_as_tuple(rowDispatcherActorId), + std::forward_as_tuple(TxId, SelfId(), rowDispatcherActorId, queueId, ++NextGeneration)); + auto& session = Sessions.at(rowDispatcherActorId); + SRC_LOG_I("Create session to " << rowDispatcherActorId << ", generation " << session.Generation); + for (auto partitionId : partitions) { + session.Partitions[partitionId]; + } + ReadActorByEventQueueId[queueId] = rowDispatcherActorId; + } + LastUsedPartitionDistribution = LastReceivedPartitionDistribution; +} + +void TDqPqRdReadActor::UpdateQueuedSize() { + ui64 queuedBytes = 0; + ui64 queuedRows = 0; + for (auto& [_, sessionInfo] : Sessions) { + queuedBytes += sessionInfo.QueuedBytes; + queuedRows += sessionInfo.QueuedRows; + } + IngressStats.QueuedBytes = queuedBytes; + IngressStats.QueuedRows = queuedRows; +} + +void TDqPqRdReadActor::Handle(TEvPrivate::TEvNotifyCA::TPtr&) { + Schedule(TDuration::Seconds(NotifyCAPeriodSec), new TEvPrivate::TEvNotifyCA()); + NotifyCA(); +} + std::pair CreateDqPqRdReadActor( + const TTypeEnvironment& typeEnv, NPq::NProto::TDqPqTopicSource&& settings, ui64 inputIndex, TCollectStatsLevel statsLevel, @@ -884,6 +1056,7 @@ std::pair CreateDqPqRdReadActor( txId, taskId, holderFactory, + typeEnv, std::move(settings), std::move(readTaskParamsMsg), computeActorId, diff --git a/ydb/library/yql/providers/pq/async_io/dq_pq_rd_read_actor.h b/ydb/library/yql/providers/pq/async_io/dq_pq_rd_read_actor.h index a218ae18c587..6e09d4637eba 100644 --- a/ydb/library/yql/providers/pq/async_io/dq_pq_rd_read_actor.h +++ b/ydb/library/yql/providers/pq/async_io/dq_pq_rd_read_actor.h @@ -23,6 +23,7 @@ class TDqAsyncIoFactory; const i64 PQRdReadDefaultFreeSpace = 256_MB; std::pair CreateDqPqRdReadActor( + const NKikimr::NMiniKQL::TTypeEnvironment& typeEnv, NPq::NProto::TDqPqTopicSource&& settings, ui64 inputIndex, TCollectStatsLevel statsLevel, diff --git a/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor.cpp b/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor.cpp index efae17d4d68e..f9d5e968b290 100644 --- a/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor.cpp +++ b/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor.cpp @@ -94,19 +94,19 @@ class TDqPqReadActor : public NActors::TActor, public NYql::NDq: TMetrics(const TTxId& txId, ui64 taskId, const ::NMonitoring::TDynamicCounterPtr& counters) : TxId(std::visit([](auto arg) { return ToString(arg); }, txId)) , Counters(counters) { - SubGroup = Counters->GetSubgroup("sink", "PqRead"); - auto sink = SubGroup->GetSubgroup("tx_id", TxId); - auto task = sink->GetSubgroup("task_id", ToString(taskId)); + SubGroup = Counters->GetSubgroup("source", "PqRead"); + auto source = SubGroup->GetSubgroup("tx_id", TxId); + auto task = source->GetSubgroup("task_id", ToString(taskId)); InFlyAsyncInputData = task->GetCounter("InFlyAsyncInputData"); InFlySubscribe = task->GetCounter("InFlySubscribe"); AsyncInputDataRate = task->GetCounter("AsyncInputDataRate", true); ReconnectRate = task->GetCounter("ReconnectRate", true); DataRate = task->GetCounter("DataRate", true); - WaitEventTimeMs = sink->GetHistogram("WaitEventTimeMs", NMonitoring::ExponentialHistogram(13, 2, 1)); // ~ 1ms -> ~ 8s + WaitEventTimeMs = source->GetHistogram("WaitEventTimeMs", NMonitoring::ExplicitHistogram({5, 20, 100, 500, 2000})); } ~TMetrics() { - SubGroup->RemoveSubgroup("id", TxId); + SubGroup->RemoveSubgroup("tx_id", TxId); } TString TxId; @@ -159,7 +159,7 @@ class TDqPqReadActor : public NActors::TActor, public NYql::NDq: } NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() const { - NYdb::NTopic::TTopicClientSettings opts; + NYdb::NTopic::TTopicClientSettings opts = PqGateway->GetTopicClientSettings(); opts.Database(SourceParams.GetDatabase()) .DiscoveryEndpoint(SourceParams.GetEndpoint()) .SslCredentials(NYdb::TSslCredentials(SourceParams.GetUseSsl())) @@ -288,7 +288,7 @@ class TDqPqReadActor : public NActors::TActor, public NYql::NDq: if (!InflightReconnect && ReconnectPeriod != TDuration::Zero()) { Metrics.ReconnectRate->Inc(); - Schedule(ReconnectPeriod, new TEvPrivate::TEvSourceDataReady()); + Schedule(ReconnectPeriod, new TEvPrivate::TEvReconnectSession()); InflightReconnect = true; } @@ -698,6 +698,7 @@ void RegisterDqPqReadActorFactory(TDqAsyncIoFactory& factory, NYdb::TDriver driv } return CreateDqPqRdReadActor( + args.TypeEnv, std::move(settings), args.InputIndex, args.StatsLevel, diff --git a/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor_base.cpp b/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor_base.cpp index 4bb5655013db..f4c4162a9f27 100644 --- a/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor_base.cpp +++ b/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor_base.cpp @@ -29,6 +29,8 @@ void TDqPqReadActorBase::SaveState(const NDqProto::TCheckpoint& /*checkpoint*/, topic->SetDatabase(SourceParams.GetDatabase()); topic->SetTopicPath(SourceParams.GetTopicPath()); + TStringStream str; + str << "SessionId: " << GetSessionId() << " SaveState, offsets: "; for (const auto& [clusterAndPartition, offset] : PartitionToOffset) { const auto& [cluster, partition] = clusterAndPartition; NPq::NProto::TDqPqTopicSourceState::TPartitionReadState* partitionState = stateProto.AddPartitions(); @@ -36,8 +38,9 @@ void TDqPqReadActorBase::SaveState(const NDqProto::TCheckpoint& /*checkpoint*/, partitionState->SetCluster(cluster); partitionState->SetPartition(partition); partitionState->SetOffset(offset); - SRC_LOG_D("SessionId: " << GetSessionId() << " SaveState: partition " << partition << ", offset: " << offset); + str << "{" << partition << "," << offset << "},"; } + SRC_LOG_D(str.Str()); stateProto.SetStartingMessageTimestampMs(StartingMessageTimestamp.MilliSeconds()); stateProto.SetIngressBytes(IngressStats.Bytes); @@ -70,9 +73,12 @@ void TDqPqReadActorBase::LoadState(const TSourceState& state) { minStartingMessageTs = Min(minStartingMessageTs, TInstant::MilliSeconds(stateProto.GetStartingMessageTimestampMs())); ingressBytes += stateProto.GetIngressBytes(); } + TStringStream str; + str << "SessionId: " << GetSessionId() << " Restoring offset: "; for (const auto& [key, value] : PartitionToOffset) { - SRC_LOG_D("SessionId: " << GetSessionId() << " Restoring offset: cluster " << key.first << ", partition id " << key.second << ", offset: " << value); + str << "{" << key.first << "," << key.second << "," << value << "},"; } + SRC_LOG_D(str.Str()); StartingMessageTimestamp = minStartingMessageTs; IngressStats.Bytes += ingressBytes; IngressStats.Chunks++; diff --git a/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor_base.h b/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor_base.h index 5028a93078c2..7b965b7f94b4 100644 --- a/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor_base.h +++ b/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor_base.h @@ -1,6 +1,8 @@ #pragma once #include +#include +#include namespace NYql::NDq::NInternal { diff --git a/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp b/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp index 18bb17683c4f..8d0085b9587e 100644 --- a/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp +++ b/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp @@ -103,10 +103,11 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu InFlyData = task->GetCounter("InFlyData"); AlreadyWritten = task->GetCounter("AlreadyWritten"); FirstContinuationTokenMs = task->GetCounter("FirstContinuationTokenMs"); + EgressDataRate = task->GetCounter("EgressDataRate", true); } ~TMetrics() { - SubGroup->RemoveSubgroup("id", TxId); + SubGroup->RemoveSubgroup("tx_id", TxId); } TString TxId; @@ -117,6 +118,7 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu ::NMonitoring::TDynamicCounters::TCounterPtr InFlyData; ::NMonitoring::TDynamicCounters::TCounterPtr AlreadyWritten; ::NMonitoring::TDynamicCounters::TCounterPtr FirstContinuationTokenMs; + ::NMonitoring::TDynamicCounters::TCounterPtr EgressDataRate; }; struct TAckInfo { @@ -140,7 +142,8 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu std::shared_ptr credentialsProviderFactory, IDqComputeActorAsyncOutput::ICallbacks* callbacks, const ::NMonitoring::TDynamicCounterPtr& counters, - i64 freeSpace) + i64 freeSpace, + const IPqGateway::TPtr& pqGateway) : TActor(&TDqPqWriteActor::StateFunc) , OutputIndex(outputIndex) , TxId(txId) @@ -151,7 +154,7 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu , Callbacks(callbacks) , LogPrefix(TStringBuilder() << "SelfId: " << this->SelfId() << ", TxId: " << TxId << ", TaskId: " << taskId << ", PQ sink. ") , FreeSpace(freeSpace) - , TopicClient(Driver, GetTopicClientSettings()) + , PqGateway(pqGateway) { EgressStats.Level = statsLevel; } @@ -294,14 +297,22 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu NYdb::NTopic::TWriteSessionSettings GetWriteSessionSettings() { return NYdb::NTopic::TWriteSessionSettings(SinkParams.GetTopicPath(), GetSourceId(), GetSourceId()) + .TraceId(LogPrefix) .MaxMemoryUsage(FreeSpace) .Codec(SinkParams.GetClusterType() == NPq::NProto::DataStreams ? NYdb::NTopic::ECodec::RAW : NYdb::NTopic::ECodec::GZIP); } + ITopicClient& GetTopicClient() { + if (!TopicClient) { + TopicClient = PqGateway->GetTopicClient(Driver, GetTopicClientSettings()); + } + return *TopicClient; + } + NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() { - return NYdb::NTopic::TTopicClientSettings() + return PqGateway->GetTopicClientSettings() .Database(SinkParams.GetDatabase()) .DiscoveryEndpoint(SinkParams.GetEndpoint()) .SslCredentials(NYdb::TSslCredentials(SinkParams.GetUseSsl())) @@ -314,7 +325,7 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu void CreateSessionIfNotExists() { if (!WriteSession) { - WriteSession = TopicClient.CreateWriteSession(GetWriteSessionSettings()); + WriteSession = GetTopicClient().CreateWriteSession(GetWriteSessionSettings()); SubscribeOnNextEvent(); } } @@ -375,6 +386,7 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu auto itemSize = GetItemSize(Buffer.front()); WaitingAcks.emplace(itemSize, TInstant::Now()); EgressStats.Bytes += itemSize; + Metrics.EgressDataRate->Add(itemSize); Buffer.pop(); } @@ -471,7 +483,7 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu i64 FreeSpace = 0; bool Finished = false; - NYdb::NTopic::TTopicClient TopicClient; + ITopicClient::TPtr TopicClient; std::shared_ptr WriteSession; TString SourceId; ui64 NextSeqNo = 1; @@ -482,6 +494,7 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu std::queue Buffer; std::queue WaitingAcks; // Size of items which are waiting for acks (used to update free space) std::queue> DeferredCheckpoints; + IPqGateway::TPtr PqGateway; }; std::pair CreateDqPqWriteActor( @@ -495,6 +508,7 @@ std::pair CreateDqPqWriteActor( ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, IDqComputeActorAsyncOutput::ICallbacks* callbacks, const ::NMonitoring::TDynamicCounterPtr& counters, + IPqGateway::TPtr pqGateway, i64 freeSpace) { const TString& tokenName = settings.GetToken().GetName(); @@ -511,13 +525,14 @@ std::pair CreateDqPqWriteActor( CreateCredentialsProviderFactoryForStructuredToken(credentialsFactory, token, addBearerToToken), callbacks, counters, - freeSpace); + freeSpace, + pqGateway); return {actor, actor}; } -void RegisterDqPqWriteActorFactory(TDqAsyncIoFactory& factory, NYdb::TDriver driver, ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, const ::NMonitoring::TDynamicCounterPtr& counters) { +void RegisterDqPqWriteActorFactory(TDqAsyncIoFactory& factory, NYdb::TDriver driver, ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, const IPqGateway::TPtr& pqGateway, const ::NMonitoring::TDynamicCounterPtr& counters) { factory.RegisterSink("PqSink", - [driver = std::move(driver), credentialsFactory = std::move(credentialsFactory), counters]( + [driver = std::move(driver), credentialsFactory = std::move(credentialsFactory), counters, pqGateway]( NPq::NProto::TDqPqTopicSink&& settings, IDqAsyncIoFactory::TSinkArguments&& args) { @@ -532,7 +547,8 @@ void RegisterDqPqWriteActorFactory(TDqAsyncIoFactory& factory, NYdb::TDriver dri driver, credentialsFactory, args.Callback, - counters + counters, + pqGateway ); }); } diff --git a/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.h b/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.h index 3af8c6e63ec8..39f388a79dcc 100644 --- a/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.h +++ b/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.h @@ -4,6 +4,7 @@ #include #include +#include #include #include @@ -31,8 +32,9 @@ std::pair CreateDqPqWriteActor( ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, IDqComputeActorAsyncOutput::ICallbacks* callbacks, const ::NMonitoring::TDynamicCounterPtr& counters, + IPqGateway::TPtr pqGateway, i64 freeSpace = DqPqDefaultFreeSpace); -void RegisterDqPqWriteActorFactory(TDqAsyncIoFactory& factory, NYdb::TDriver driver, ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, const ::NMonitoring::TDynamicCounterPtr& counters = MakeIntrusive<::NMonitoring::TDynamicCounters>()); +void RegisterDqPqWriteActorFactory(TDqAsyncIoFactory& factory, NYdb::TDriver driver, ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, const IPqGateway::TPtr& pqGateway, const ::NMonitoring::TDynamicCounterPtr& counters = MakeIntrusive<::NMonitoring::TDynamicCounters>()); } // namespace NYql::NDq diff --git a/ydb/library/yql/providers/pq/expr_nodes/yql_pq_expr_nodes.json b/ydb/library/yql/providers/pq/expr_nodes/yql_pq_expr_nodes.json index 1fe5bf0b57dd..9a6603bd6631 100644 --- a/ydb/library/yql/providers/pq/expr_nodes/yql_pq_expr_nodes.json +++ b/ydb/library/yql/providers/pq/expr_nodes/yql_pq_expr_nodes.json @@ -72,7 +72,8 @@ {"Index": 2, "Name": "Columns", "Type": "TExprBase"}, {"Index": 3, "Name": "Settings", "Type": "TCoNameValueTupleList"}, {"Index": 4, "Name": "Token", "Type": "TCoSecureParam"}, - {"Index": 5, "Name": "FilterPredicate", "Type": "TCoAtom"} + {"Index": 5, "Name": "FilterPredicate", "Type": "TCoAtom"}, + {"Index": 6, "Name": "RowType", "Type": "TExprBase"} ] }, { diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_blocking_queue.h b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_blocking_queue.h new file mode 100644 index 000000000000..7a8db40dc554 --- /dev/null +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_blocking_queue.h @@ -0,0 +1,85 @@ +#include + +#include +#include + +namespace NYql { + +template +class TBlockingEQueue { +public: + explicit TBlockingEQueue(size_t maxSize) + : MaxSize_(maxSize) + { + } + void Push(TEvent&& e, size_t size = 0) { + with_lock(Mutex_) { + CanPush_.WaitI(Mutex_, [this] () {return CanPushPredicate();}); + Events_.emplace_back(std::move(e), size ); + Size_ += size; + } + CanPop_.BroadCast(); + } + + void BlockUntilEvent() { + with_lock(Mutex_) { + CanPop_.WaitI(Mutex_, [this] () {return CanPopPredicate();}); + } + } + + TMaybe Pop(bool block) { + with_lock(Mutex_) { + if (block) { + CanPop_.WaitI(Mutex_, [this] () {return CanPopPredicate();}); + } else { + if (!CanPopPredicate()) { + return {}; + } + } + if (Events_.empty()) { + return {}; + } + + auto [front, size] = std::move(Events_.front()); + Events_.pop_front(); + Size_ -= size; + if (Size_ < MaxSize_) { + CanPush_.BroadCast(); + } + return std::move(front); // cast to TMaybe<> + } + } + + void Stop() { + with_lock(Mutex_) { + Stopped_ = true; + CanPop_.BroadCast(); + CanPush_.BroadCast(); + } + } + + bool IsStopped() { + with_lock(Mutex_) { + return Stopped_; + } + } + +private: + bool CanPopPredicate() const { + return !Events_.empty() || Stopped_; + } + + bool CanPushPredicate() const { + return Size_ < MaxSize_ || Stopped_; + } + + size_t MaxSize_; + size_t Size_ = 0; + TDeque> Events_; + bool Stopped_ = false; + TMutex Mutex_; + TCondVar CanPop_; + TCondVar CanPush_; +}; + +} diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp index f1575567e468..f4b56a3ea280 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp @@ -11,7 +11,7 @@ namespace NYql { NThreading::TFuture TDummyPqGateway::OpenSession(const TString& sessionId, const TString& username) { with_lock (Mutex) { Y_ENSURE(sessionId); - Y_ENSURE(username); + Y_UNUSED(username); Y_ENSURE(!IsIn(OpenedSessions, sessionId), "Session " << sessionId << " is already opened in pq gateway"); OpenedSessions.insert(sessionId); @@ -53,8 +53,8 @@ NThreading::TFuture TDummyPqGateway::ListStreams(const TDummyPqGateway& TDummyPqGateway::AddDummyTopic(const TDummyTopic& topic) { with_lock (Mutex) { Y_ENSURE(topic.Cluster); - Y_ENSURE(topic.Path); - const auto key = std::make_pair(topic.Cluster, topic.Path); + Y_ENSURE(topic.TopicName); + const auto key = std::make_pair(topic.Cluster, topic.TopicName); Y_ENSURE(Topics.emplace(key, topic).second, "Already inserted dummy topic {" << topic.Cluster << ", " << topic.Path << "}"); return *this; } @@ -80,4 +80,28 @@ void TDummyPqGateway::UpdateClusterConfigs( Y_UNUSED(secure); } +void TDummyPqGateway::UpdateClusterConfigs(const TPqGatewayConfigPtr& config) { + Y_UNUSED(config); +} + +NYdb::NTopic::TTopicClientSettings TDummyPqGateway::GetTopicClientSettings() const { + return NYdb::NTopic::TTopicClientSettings(); +} + +class TPqFileGatewayFactory : public IPqGatewayFactory { +public: + TPqFileGatewayFactory(const TDummyPqGateway::TPtr pqFileGateway) + : PqFileGateway(pqFileGateway) {} + + IPqGateway::TPtr CreatePqGateway() override { + return PqFileGateway; + } +private: + const TDummyPqGateway::TPtr PqFileGateway; +}; + +IPqGatewayFactory::TPtr CreatePqFileGatewayFactory(const TDummyPqGateway::TPtr pqFileGateway) { + return MakeIntrusive(pqFileGateway); +} + } // namespace NYql diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h index 84a394531ee0..084dd07f572c 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h @@ -9,10 +9,11 @@ namespace NYql { struct TDummyTopic { - TDummyTopic(const TString& cluster, const TString& path, const TMaybe& filePath = {}) + TDummyTopic(const TString& cluster, const TString& topicName, const TMaybe& path = {}, size_t partitionCount = 1) : Cluster(cluster) + , TopicName(topicName) , Path(path) - , FilePath(filePath) + , PartitionsCount(partitionCount) { } @@ -22,9 +23,10 @@ struct TDummyTopic { } TString Cluster; - TString Path; - TMaybe FilePath; - size_t PartitionsCount = 1; + TString TopicName; + TMaybe Path; + size_t PartitionsCount; + bool CancelOnFileFinish = false; }; // Dummy Pq gateway for tests. @@ -56,8 +58,11 @@ class TDummyPqGateway : public IPqGateway { const TString& endpoint, const TString& database, bool secure) override; - + + void UpdateClusterConfigs(const TPqGatewayConfigPtr& config) override; + ITopicClient::TPtr GetTopicClient(const NYdb::TDriver& driver, const NYdb::NTopic::TTopicClientSettings& settings) override; + NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() const override; using TClusterNPath = std::pair; private: @@ -69,4 +74,6 @@ class TDummyPqGateway : public IPqGateway { IPqGateway::TPtr CreatePqFileGateway(); +IPqGatewayFactory::TPtr CreatePqFileGatewayFactory(const TDummyPqGateway::TPtr pqFileGateway); + } // namespace NYql diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp index 55f284b8f865..e85243df7e58 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp @@ -3,89 +3,28 @@ #include -#include #include +#include #include +#include "yql_pq_blocking_queue.h" namespace NYql { -class TBlockingEQueue { -public: - TBlockingEQueue(size_t maxSize):MaxSize_(maxSize) { - } - void Push(NYdb::NTopic::TReadSessionEvent::TEvent&& e, size_t size) { - with_lock(Mutex_) { - CanPush_.WaitI(Mutex_, [this] () {return Stopped_ || Size_ < MaxSize_;}); - Events_.emplace_back(std::move(e), size ); - Size_ += size; - } - CanPop_.BroadCast(); - } - - void BlockUntilEvent() { - with_lock(Mutex_) { - CanPop_.WaitI(Mutex_, [this] () {return Stopped_ || !Events_.empty();}); - } - } - - TMaybe Pop(bool block) { - with_lock(Mutex_) { - if (block) { - CanPop_.WaitI(Mutex_, [this] () {return CanPopPredicate();}); - } else { - if (!CanPopPredicate()) { - return {}; - } - } - auto [front, size] = std::move(Events_.front()); - Events_.pop_front(); - Size_ -= size; - if (Size_ < MaxSize_) { - CanPush_.BroadCast(); - } - return front; - } - } - - void Stop() { - with_lock(Mutex_) { - Stopped_ = true; - CanPop_.BroadCast(); - CanPush_.BroadCast(); - } - } - - bool IsStopped() { - with_lock(Mutex_) { - return Stopped_; - } - } - -private: - bool CanPopPredicate() { - return !Events_.empty() && !Stopped_; - } - - size_t MaxSize_; - size_t Size_ = 0; - TDeque> Events_; - bool Stopped_ = false; - TMutex Mutex_; - TCondVar CanPop_; - TCondVar CanPush_; -}; - class TFileTopicReadSession : public NYdb::NTopic::IReadSession { -constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); +constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); public: - TFileTopicReadSession(TFile file, NYdb::NTopic::TPartitionSession::TPtr session, const TString& producerId = ""): - File_(std::move(file)), Session_(std::move(session)), ProducerId_(producerId), - FilePoller_([this] () { - PollFileForChanges(); - }), Counters_() + TFileTopicReadSession(TFile file, NYdb::NTopic::TPartitionSession::TPtr session, const TString& producerId = "", bool cancelOnFileFinish = false) + : File_(std::move(file)) + , Session_(std::move(session)) + , ProducerId_(producerId) + , FilePoller_([this] () { + PollFileForChanges(); + }) + , Counters_() + , CancelOnFileFinish_(cancelOnFileFinish) { Pool_.Start(1); } @@ -102,8 +41,8 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); Y_UNUSED(maxByteSize); TVector res; - for (auto event = EventsQ_.Pop(block); !event.Empty() && res.size() <= maxEventsCount.GetOrElse(std::numeric_limits::max()); event = EventsQ_.Pop(/*block=*/ false)) { - res.push_back(*event); + for (auto event = EventsQ_.Pop(block); !event.Empty() && res.size() < maxEventsCount.GetOrElse(std::numeric_limits::max()); event = EventsQ_.Pop(/*block=*/ false)) { + res.push_back(std::move(*event)); } return res; } @@ -125,14 +64,15 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); bool Close(TDuration timeout = TDuration::Max()) override { Y_UNUSED(timeout); - // TOOD send TSessionClosedEvent + // TODO send TSessionClosedEvent + // XXX (... but if we stop queues, nobody will receive it, needs rethinking) EventsQ_.Stop(); Pool_.Stop(); if (FilePoller_.joinable()) { FilePoller_.join(); } - return true; + return true; // TODO incorrect if EventQ_ was non-empty } NYdb::NTopic::TReaderCounters::TPtr GetCounters() const override { @@ -152,11 +92,11 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); } private: - using TMessageInformation = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessageInformation; - using TMessage = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage; + using TMessageInformation = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessageInformation; + using TMessage = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage; - TMessageInformation MakeNextMessageInformation(size_t offset, size_t uncompressedSize, const TString& messageGroupId = "") { - auto now = TInstant::Now(); + TMessageInformation MakeNextMessageInformation(size_t offset, size_t uncompressedSize, const TString& messageGroupId = "") { + auto now = TInstant::Now(); TMessageInformation msgInfo( offset, ProducerId_, @@ -170,7 +110,7 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); ); return msgInfo; } - + TMessage MakeNextMessage(const TString& msgBuff) { TMessage msg(msgBuff, nullptr, MakeNextMessageInformation(MsgOffset_, msgBuff.size()), Session_); return msg; @@ -184,7 +124,7 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); size_t size = 0; ui64 maxBatchRowSize = 100; - while (size_t read = fi.ReadLine(rawMsg)) { + while (fi.ReadLine(rawMsg)) { msgs.emplace_back(MakeNextMessage(rawMsg)); MsgOffset_++; if (!maxBatchRowSize--) { @@ -194,24 +134,200 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); } if (!msgs.empty()) { EventsQ_.Push(NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent(msgs, {}, Session_), size); + } else if (CancelOnFileFinish_) { + EventsQ_.Push(NYdb::NTopic::TSessionClosedEvent(NYdb::EStatus::CANCELLED, {NYql::TIssue("PQ file topic was finished")}), size); } Sleep(FILE_POLL_PERIOD); } } - + TFile File_; - TBlockingEQueue EventsQ_ {4_MB}; + TBlockingEQueue EventsQ_ {4_MB}; NYdb::NTopic::TPartitionSession::TPtr Session_; TString ProducerId_; std::thread FilePoller_; NYdb::NTopic::TReaderCounters::TPtr Counters_; + bool CancelOnFileFinish_ = false; TThreadPool Pool_; size_t MsgOffset_ = 0; ui64 SeqNo_ = 0; }; +class TFileTopicWriteSession : public NYdb::NTopic::IWriteSession, private NYdb::NTopic::TContinuationTokenIssuer { +public: + explicit TFileTopicWriteSession(TFile file) + : File_(std::move(file)) + , FileWriter_([this] () { + PushToFile(); + }) + , Counters_() + { + Pool_.Start(1); + EventsQ_.Push(NYdb::NTopic::TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()}); + } + + NThreading::TFuture WaitEvent() override { + return NThreading::Async([this] () { + EventsQ_.BlockUntilEvent(); + return NThreading::MakeFuture(); + }, Pool_); + } + + TMaybe GetEvent(bool block) override { + return EventsQ_.Pop(block); + } + + TVector GetEvents(bool block, TMaybe maxEventsCount) override { + TVector res; + for (auto event = EventsQ_.Pop(block); !event.Empty() && res.size() < maxEventsCount.GetOrElse(std::numeric_limits::max()); event = EventsQ_.Pop(/*block=*/ false)) { + res.push_back(std::move(*event)); + } + return res; + } + + NThreading::TFuture GetInitSeqNo() override { + return NThreading::MakeFuture(SeqNo_); + } + + void Write(NYdb::NTopic::TContinuationToken&&, NYdb::NTopic::TWriteMessage&& message, + NYdb::NTable::TTransaction* tx) override { + Y_UNUSED(tx); + + auto size = message.Data.size(); + EventsMsgQ_.Push(TOwningWriteMessage(std::move(message)), size); + } + + void Write(NYdb::NTopic::TContinuationToken&& token, TStringBuf data, TMaybe seqNo, + TMaybe createTimestamp) override { + NYdb::NTopic::TWriteMessage message(data); + if (seqNo.Defined()) { + message.SeqNo(*seqNo); + } + if (createTimestamp.Defined()) { + message.CreateTimestamp(*createTimestamp); + } + + Write(std::move(token), std::move(message), nullptr); + } + + // Ignores codec in message and always writes raw for debugging purposes + void WriteEncoded(NYdb::NTopic::TContinuationToken&& token, NYdb::NTopic::TWriteMessage&& params, + NYdb::NTable::TTransaction* tx) override { + Y_UNUSED(tx); + + NYdb::NTopic::TWriteMessage message(params.Data); + + if (params.CreateTimestamp_.Defined()) { + message.CreateTimestamp(*params.CreateTimestamp_); + } + if (params.SeqNo_) { + message.SeqNo(*params.SeqNo_); + } + message.MessageMeta(params.MessageMeta_); + + Write(std::move(token), std::move(message), nullptr); + } + + // Ignores codec in message and always writes raw for debugging purposes + void WriteEncoded(NYdb::NTopic::TContinuationToken&& token, TStringBuf data, NYdb::NTopic::ECodec codec, ui32 originalSize, + TMaybe seqNo, TMaybe createTimestamp) override { + Y_UNUSED(codec); + Y_UNUSED(originalSize); + + NYdb::NTopic::TWriteMessage message(data); + if (seqNo.Defined()) { + message.SeqNo(*seqNo); + } + if (createTimestamp.Defined()) { + message.CreateTimestamp(*createTimestamp); + } + + Write(std::move(token), std::move(message), nullptr); + } + + bool Close(TDuration timeout = TDuration::Max()) override { + Y_UNUSED(timeout); + // TODO send TSessionClosedEvent + // XXX (... but if we stop queues, nobody will receive it, needs rethinking) + EventsQ_.Stop(); + EventsMsgQ_.Stop(); + Pool_.Stop(); + + if (FileWriter_.joinable()) { + FileWriter_.join(); + } + return true; // TODO incorrect if Event*Q_ was non-empty + } + + NYdb::NTopic::TWriterCounters::TPtr GetCounters() override { + return Counters_; + } + + ~TFileTopicWriteSession() override { + EventsQ_.Stop(); + EventsMsgQ_.Stop(); + Pool_.Stop(); + if (FileWriter_.joinable()) { + FileWriter_.join(); + } + } + +private: + void PushToFile() { + TFileOutput fo(File_); + ui64 offset = 0; // FIXME dummy + ui64 partitionId = 0; // FIXME dummy + while (auto maybeMsg = EventsMsgQ_.Pop(true)) { + NYdb::NTopic::TWriteSessionEvent::TAcksEvent acks; + do { + auto& [content, msg] = *maybeMsg; + NYdb::NTopic::TWriteSessionEvent::TWriteAck ack; + if (msg.SeqNo_.Defined()) { // FIXME should be auto generated otherwise + ack.SeqNo = *msg.SeqNo_; + } + ack.State = NYdb::NTopic::TWriteSessionEvent::TWriteAck::EES_WRITTEN; + ack.Details.ConstructInPlace(offset, partitionId); + acks.Acks.emplace_back(std::move(ack)); + offset += content.size() + 1; + fo.Write(content); + fo.Write('\n'); + } while ((maybeMsg = EventsMsgQ_.Pop(false))); + fo.Flush(); + EventsQ_.Push(std::move(acks), 1 + acks.Acks.size()); + EventsQ_.Push(NYdb::NTopic::TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()}, 1); + if (EventsQ_.IsStopped()) { + break; + } + } + } + + TFile File_; + + // We acquire ownership of messages immediately + // TODO: remove extra message copying to and from queue + struct TOwningWriteMessage { + TString content; + NYdb::NTopic::TWriteMessage msg; + + explicit TOwningWriteMessage(NYdb::NTopic::TWriteMessage&& msg) + : content(msg.Data) + , msg(std::move(msg)) + { + msg.Data = content; + } + }; + TBlockingEQueue EventsMsgQ_ {4_MB}; + + TBlockingEQueue EventsQ_ {128_KB}; + std::thread FileWriter_; + + TThreadPool Pool_; + NYdb::NTopic::TWriterCounters::TPtr Counters_; + ui64 SeqNo_ = 0; +}; + struct TDummyPartitionSession: public NYdb::NTopic::TPartitionSession { TDummyPartitionSession(ui64 sessionId, const TString& topicPath, ui64 partId) { PartitionSessionId = sessionId; @@ -224,22 +340,34 @@ struct TDummyPartitionSession: public NYdb::NTopic::TPartitionSession { } }; +TFileTopicClient::TFileTopicClient(THashMap topics) + : Topics_(std::move(topics)) +{} + std::shared_ptr TFileTopicClient::CreateReadSession(const NYdb::NTopic::TReadSessionSettings& settings) { Y_ENSURE(!settings.Topics_.empty()); - TString topicPath = settings.Topics_.front().Path_; - + const auto& topic = settings.Topics_.front(); + auto topicPath = topic.Path_; + Y_ENSURE(topic.PartitionIds_.size() >= 1); + ui64 partitionId = topic.PartitionIds_.front(); auto topicsIt = Topics_.find(make_pair("pq", topicPath)); Y_ENSURE(topicsIt != Topics_.end()); - auto filePath = topicsIt->second.FilePath; + auto filePath = topicsIt->second.Path; Y_ENSURE(filePath); - + + TFsPath fsPath(*filePath); + if (fsPath.IsDirectory()) { + filePath = TStringBuilder() << *filePath << "/" << ToString(partitionId); + } else if (!fsPath.Exists()) { + filePath = TStringBuilder() << *filePath << "_" << partitionId; + } + // TODO ui64 sessionId = 0; - ui64 partitionId = 0; - return std::make_shared( TFile(*filePath, EOpenMode::TEnum::RdOnly), - MakeIntrusive(sessionId, topicPath, partitionId) + MakeIntrusive(sessionId, TString{topicPath}, partitionId), + "", topicsIt->second.CancelOnFileFinish ); } @@ -261,7 +389,7 @@ NYdb::TAsyncStatus TFileTopicClient::DropTopic(const TString& path, const NYdb:: return NThreading::MakeFuture(NYdb::TStatus(NYdb::EStatus::SUCCESS, {})); } -NYdb::NTopic::TAsyncDescribeTopicResult TFileTopicClient::DescribeTopic(const TString& path, +NYdb::NTopic::TAsyncDescribeTopicResult TFileTopicClient::DescribeTopic(const TString& path, const NYdb::NTopic::TDescribeTopicSettings& settings) { Y_UNUSED(path); Y_UNUSED(settings); @@ -270,7 +398,7 @@ NYdb::NTopic::TAsyncDescribeTopicResult TFileTopicClient::DescribeTopic(const TS return NThreading::MakeFuture(NYdb::NTopic::TDescribeTopicResult(std::move(success), {})); } -NYdb::NTopic::TAsyncDescribeConsumerResult TFileTopicClient::DescribeConsumer(const TString& path, const TString& consumer, +NYdb::NTopic::TAsyncDescribeConsumerResult TFileTopicClient::DescribeConsumer(const TString& path, const TString& consumer, const NYdb::NTopic::TDescribeConsumerSettings& settings) { Y_UNUSED(path); Y_UNUSED(consumer); @@ -280,7 +408,7 @@ NYdb::NTopic::TAsyncDescribeConsumerResult TFileTopicClient::DescribeConsumer(co return NThreading::MakeFuture(NYdb::NTopic::TDescribeConsumerResult(std::move(success), {})); } -NYdb::NTopic::TAsyncDescribePartitionResult TFileTopicClient::DescribePartition(const TString& path, i64 partitionId, +NYdb::NTopic::TAsyncDescribePartitionResult TFileTopicClient::DescribePartition(const TString& path, i64 partitionId, const NYdb::NTopic::TDescribePartitionSettings& settings) { Y_UNUSED(path); Y_UNUSED(partitionId); @@ -297,8 +425,13 @@ std::shared_ptr TFileTopicClient::Cre } std::shared_ptr TFileTopicClient::CreateWriteSession(const NYdb::NTopic::TWriteSessionSettings& settings) { - Y_UNUSED(settings); - return nullptr; + TString topicPath = settings.Path_; + auto topicsIt = Topics_.find(make_pair("pq", topicPath)); + Y_ENSURE(topicsIt != Topics_.end()); + auto filePath = topicsIt->second.Path; + Y_ENSURE(filePath); + + return std::make_shared(TFile(*filePath, EOpenMode::TEnum::RdWr)); } NYdb::TAsyncStatus TFileTopicClient::CommitOffset(const TString& path, ui64 partitionId, const TString& consumerName, ui64 offset, diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.h b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.h index f80426726159..7b4ccae0d85c 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.h +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.h @@ -4,7 +4,7 @@ namespace NYql { struct TFileTopicClient : public ITopicClient { - TFileTopicClient(THashMap topics): Topics_(topics) {} + explicit TFileTopicClient(THashMap topics); NYdb::TAsyncStatus CreateTopic(const TString& path, const NYdb::NTopic::TCreateTopicSettings& settings = {}) override; @@ -32,5 +32,7 @@ struct TFileTopicClient : public ITopicClient { private: THashMap Topics_; + bool CancelOnFileFinish_ = false; }; -} \ No newline at end of file + +} diff --git a/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.cpp b/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.cpp index 8c94f126ecda..063557322848 100644 --- a/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.cpp +++ b/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.cpp @@ -40,10 +40,12 @@ class TPqNativeGateway : public IPqGateway { const TString& database, bool secure) override; + void UpdateClusterConfigs(const TPqGatewayConfigPtr& config) override; + ITopicClient::TPtr GetTopicClient(const NYdb::TDriver& driver, const NYdb::NTopic::TTopicClientSettings& settings) override; + NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() const override; private: - void InitClusterConfigs(); TPqSession::TPtr GetExistingSession(const TString& sessionId) const; private: @@ -56,6 +58,7 @@ class TPqNativeGateway : public IPqGateway { NYdb::TDriver YdbDriver; TPqClusterConfigsMapPtr ClusterConfigs; THashMap Sessions; + TMaybe CommonTopicClientSettings; }; TPqNativeGateway::TPqNativeGateway(const TPqGatewayServices& services) @@ -65,14 +68,15 @@ TPqNativeGateway::TPqNativeGateway(const TPqGatewayServices& services) , CredentialsFactory(services.CredentialsFactory) , CmConnections(services.CmConnections) , YdbDriver(services.YdbDriver) + , CommonTopicClientSettings(services.CommonTopicClientSettings) { Y_UNUSED(FunctionRegistry); - InitClusterConfigs(); + UpdateClusterConfigs(Config); } -void TPqNativeGateway::InitClusterConfigs() { +void TPqNativeGateway::UpdateClusterConfigs(const TPqGatewayConfigPtr& config) { ClusterConfigs = std::make_shared(); - for (const auto& cfg : Config->GetClusterMapping()) { + for (const auto& cfg : config->GetClusterMapping()) { auto& config = (*ClusterConfigs)[cfg.GetName()]; config = cfg; } @@ -144,8 +148,28 @@ ITopicClient::TPtr TPqNativeGateway::GetTopicClient(const NYdb::TDriver& driver, return MakeIntrusive(driver, settings); } +NYdb::NTopic::TTopicClientSettings TPqNativeGateway::GetTopicClientSettings() const { + return CommonTopicClientSettings ? *CommonTopicClientSettings : NYdb::NTopic::TTopicClientSettings(); +} + TPqNativeGateway::~TPqNativeGateway() { Sessions.clear(); } +class TPqNativeGatewayFactory : public IPqGatewayFactory { +public: + TPqNativeGatewayFactory(const NYql::TPqGatewayServices& services) + : Services(services) {} + + IPqGateway::TPtr CreatePqGateway() override { + return CreatePqNativeGateway(Services); + } + const NYql::TPqGatewayServices Services; +}; + +IPqGatewayFactory::TPtr CreatePqNativeGatewayFactory(const NYql::TPqGatewayServices& services) { + return MakeIntrusive(services); +} + + } // namespace NYql diff --git a/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.h b/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.h index 1294abfe6a34..32d28a2549a6 100644 --- a/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.h +++ b/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.h @@ -26,6 +26,7 @@ struct TPqGatewayServices { ISecuredServiceAccountCredentialsFactory::TPtr CredentialsFactory; ::NPq::NConfigurationManager::IConnections::TPtr CmConnections; NYdb::TDriver YdbDriver; + TMaybe CommonTopicClientSettings; TPqGatewayServices( NYdb::TDriver driver, @@ -33,17 +34,21 @@ struct TPqGatewayServices { ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, TPqGatewayConfigPtr config, const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry, - IMetricsRegistryPtr metrics = nullptr) + IMetricsRegistryPtr metrics = nullptr, + TMaybe commonTopicClientSettings = Nothing()) : FunctionRegistry(functionRegistry) , Config(std::move(config)) , Metrics(std::move(metrics)) , CredentialsFactory(std::move(credentialsFactory)) , CmConnections(std::move(cmConnections)) , YdbDriver(std::move(driver)) + , CommonTopicClientSettings(commonTopicClientSettings) { } }; IPqGateway::TPtr CreatePqNativeGateway(const TPqGatewayServices& services); +IPqGatewayFactory::TPtr CreatePqNativeGatewayFactory(const NYql::TPqGatewayServices& services); + } // namespace NYql diff --git a/ydb/library/yql/providers/pq/proto/dq_io.proto b/ydb/library/yql/providers/pq/proto/dq_io.proto index 5cda0945d7a9..93b5d8a06026 100644 --- a/ydb/library/yql/providers/pq/proto/dq_io.proto +++ b/ydb/library/yql/providers/pq/proto/dq_io.proto @@ -40,6 +40,8 @@ message TDqPqTopicSource { string ReconnectPeriod = 16; // disabled by default, example of a parameter: 5m bool EnabledLLVM = 17; string ReadGroup = 18; + string Format = 19; + string RowType = 20; // Final row type with metadata columns } message TDqPqTopicSink { diff --git a/ydb/library/yql/providers/pq/provider/ut/yql_pq_ut.cpp b/ydb/library/yql/providers/pq/provider/ut/yql_pq_ut.cpp index d6f1bd5a46d2..3cd2c1538d4d 100644 --- a/ydb/library/yql/providers/pq/provider/ut/yql_pq_ut.cpp +++ b/ydb/library/yql/providers/pq/provider/ut/yql_pq_ut.cpp @@ -45,7 +45,7 @@ NDq::IDqAsyncIoFactory::TPtr CreateAsyncIoFactory(const NYdb::TDriver& driver, c auto factory = MakeIntrusive(); RegisterDqPqReadActorFactory(*factory, driver, nullptr, pqGateway); - RegisterDqPqWriteActorFactory(*factory, driver, nullptr); + RegisterDqPqWriteActorFactory(*factory, driver, nullptr, pqGateway); return factory; } diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_datasource_type_ann.cpp b/ydb/library/yql/providers/pq/provider/yql_pq_datasource_type_ann.cpp index 988cc73bea76..82cebc19e0cc 100644 --- a/ydb/library/yql/providers/pq/provider/yql_pq_datasource_type_ann.cpp +++ b/ydb/library/yql/providers/pq/provider/yql_pq_datasource_type_ann.cpp @@ -1,4 +1,5 @@ #include "yql_pq_provider_impl.h" +#include "yql_pq_helpers.h" #include #include @@ -7,6 +8,7 @@ #include #include #include +#include #include #include @@ -132,7 +134,7 @@ class TPqDataSourceTypeAnnotationTransformer : public TVisitorTransformerBase { } TStatus HandleDqTopicSource(TExprBase input, TExprContext& ctx) { - if (!EnsureArgsCount(input.Ref(), 6, ctx)) { + if (!EnsureArgsCount(input.Ref(), 7, ctx)) { return TStatus::Error; } @@ -156,6 +158,13 @@ class TPqDataSourceTypeAnnotationTransformer : public TVisitorTransformerBase { return TStatus::Error; } + if (const auto maybeSharedReadingSetting = FindSetting(topicSource.Settings().Ptr(), SharedReading)) { + const TExprNode& value = maybeSharedReadingSetting.Cast().Ref(); + if (value.IsAtom() && FromString(value.Content())) { + input.Ptr()->SetTypeAnn(ctx.MakeType(topicSource.RowType().Ref().GetTypeAnn())); + return TStatus::Ok; + } + } if (topic.Metadata().Empty()) { input.Ptr()->SetTypeAnn(ctx.MakeType(ctx.MakeType(EDataSlot::String))); diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_dq_integration.cpp b/ydb/library/yql/providers/pq/provider/yql_pq_dq_integration.cpp index 1932fbe9a9a2..a91029b6566f 100644 --- a/ydb/library/yql/providers/pq/provider/yql_pq_dq_integration.cpp +++ b/ydb/library/yql/providers/pq/provider/yql_pq_dq_integration.cpp @@ -81,48 +81,7 @@ class TPqDqIntegration: public TDqIntegrationBase { ->Cast()->GetItems().back()->Cast() ->GetItemType()->Cast(); const auto& clusterName = pqReadTopic.DataSource().Cluster().StringValue(); - - TVector settings; - settings.push_back(Build(ctx, pqReadTopic.Pos()) - .Name().Build("format") - .Value(pqReadTopic.Format()) - .Done()); - - auto format = pqReadTopic.Format().Ref().Content(); - - TVector innerSettings; - if (pqReadTopic.Compression() != "") { - innerSettings.push_back(Build(ctx, pqReadTopic.Pos()) - .Name().Build("compression") - .Value(pqReadTopic.Compression()) - .Done()); - } - - if (!innerSettings.empty()) { - settings.push_back(Build(ctx, pqReadTopic.Pos()) - .Name().Build("settings") - .Value() - .Add(innerSettings) - .Build() - .Done()); - } - - TExprNode::TListType metadataFieldsList; - for (auto sysColumn : AllowedPqMetaSysColumns()) { - metadataFieldsList.push_back(ctx.NewAtom(pqReadTopic.Pos(), sysColumn)); - } - - settings.push_back(Build(ctx, pqReadTopic.Pos()) - .Name().Build("metadataColumns") - .Value(ctx.NewList(pqReadTopic.Pos(), std::move(metadataFieldsList))) - .Done()); - - - settings.push_back(Build(ctx, pqReadTopic.Pos()) - .Name().Build("formatSettings") - .Value(std::move(pqReadTopic.Settings())) - .Done()); - + const auto format = pqReadTopic.Format().Ref().Content(); const auto token = "cluster:default_" + clusterName; const auto& typeItems = pqReadTopic.Topic().RowSpec().Ref().GetTypeAnn()->Cast()->GetType()->Cast()->GetItems(); @@ -136,25 +95,21 @@ class TPqDqIntegration: public TDqIntegrationBase { }); auto columnNames = ctx.NewList(pos, std::move(colNames)); - auto row = Build(ctx, read->Pos()) - .Name("row") - .Done(); - TString emptyPredicate; - - return Build(ctx, read->Pos()) + return Build(ctx, pos) .Input() .World(pqReadTopic.World()) .Topic(pqReadTopic.Topic()) .Columns(std::move(columnNames)) - .Settings(BuildTopicReadSettings(clusterName, wrSettings, read->Pos(), format, ctx)) + .Settings(BuildTopicReadSettings(clusterName, wrSettings, pos, format, ctx)) .Token() .Name().Build(token) .Build() - .FilterPredicate().Value(emptyPredicate).Build() + .FilterPredicate().Value(TString()).Build() // Empty predicate by default <=> WHERE TRUE + .RowType(ExpandType(pqReadTopic.Pos(), *rowType, ctx)) .Build() .RowType(ExpandType(pqReadTopic.Pos(), *rowType, ctx)) .DataSource(pqReadTopic.DataSource().Cast()) - .Settings(Build(ctx, read->Pos()).Add(settings).Done()) + .Settings(BuildDqSourceWrapSettings(pqReadTopic, pos, ctx)) .Done().Ptr(); } return read; @@ -209,6 +164,9 @@ class TPqDqIntegration: public TDqIntegrationBase { srcDesc.SetClusterType(ToClusterType(clusterDesc->ClusterType)); srcDesc.SetDatabaseId(clusterDesc->DatabaseId); + const TStructExprType* fullRowType = topicSource.RowType().Ref().GetTypeAnn()->Cast()->GetType()->Cast(); + srcDesc.SetRowType(NCommon::WriteTypeToYson(fullRowType, NYT::NYson::EYsonFormat::Text)); + if (const auto& types = State_->Types) { if (const auto& optLLVM = types->OptLLVM) { srcDesc.SetEnabledLLVM(!optLLVM->empty() && *optLLVM != "OFF"); @@ -247,6 +205,7 @@ class TPqDqIntegration: public TDqIntegrationBase { srcDesc.MutableWatermarks()->SetIdlePartitionsEnabled(true); } } + srcDesc.SetFormat(format); if (auto maybeToken = TMaybeNode(topicSource.Token().Raw())) { srcDesc.MutableToken()->SetName(TString(maybeToken.Cast().Name().Value())); @@ -270,12 +229,9 @@ class TPqDqIntegration: public TDqIntegrationBase { auto serializedProto = topicSource.FilterPredicate().Ref().Content(); YQL_ENSURE (predicateProto.ParseFromString(serializedProto)); - sharedReading = sharedReading && (format == "json_each_row" || format == "raw"); TString predicateSql = NYql::FormatWhere(predicateProto); if (sharedReading) { - if (format == "json_each_row") { - srcDesc.SetPredicate(predicateSql); - } + srcDesc.SetPredicate(predicateSql); srcDesc.SetSharedReading(true); } protoSettings.PackFrom(srcDesc); @@ -341,13 +297,10 @@ class TPqDqIntegration: public TDqIntegrationBase { } } - auto clusterConfiguration = State_->Configuration->ClustersConfigurationSettings.FindPtr(cluster); - if (!clusterConfiguration) { - ythrow yexception() << "Unknown pq cluster \"" << cluster << "\""; - } + auto clusterConfiguration = GetClusterConfiguration(cluster); Add(props, EndpointSetting, clusterConfiguration->Endpoint, pos, ctx); - Add(props, SharedReading, ToString(clusterConfiguration->SharedReading), pos, ctx); + Add(props, SharedReading, ToString(UseSharedReading(clusterConfiguration, format)), pos, ctx); Add(props, ReconnectPeriod, ToString(clusterConfiguration->ReconnectPeriod), pos, ctx); Add(props, Format, format, pos, ctx); Add(props, ReadGroup, clusterConfiguration->ReadGroup, pos, ctx); @@ -383,6 +336,65 @@ class TPqDqIntegration: public TDqIntegrationBase { .Done(); } + NNodes::TCoNameValueTupleList BuildDqSourceWrapSettings(const TPqReadTopic& pqReadTopic, TPositionHandle pos, TExprContext& ctx) const { + TVector settings; + settings.push_back(Build(ctx, pos) + .Name().Build("format") + .Value(pqReadTopic.Format()) + .Done()); + + TExprNode::TListType metadataFieldsList; + for (const auto& sysColumn : AllowedPqMetaSysColumns()) { + metadataFieldsList.push_back(ctx.NewAtom(pos, sysColumn)); + } + + settings.push_back(Build(ctx, pos) + .Name().Build("metadataColumns") + .Value(ctx.NewList(pos, std::move(metadataFieldsList))) + .Done()); + + settings.push_back(Build(ctx, pos) + .Name().Build("formatSettings") + .Value(std::move(pqReadTopic.Settings())) + .Done()); + + TVector innerSettings; + if (pqReadTopic.Compression() != "") { + innerSettings.push_back(Build(ctx, pos) + .Name().Build("compression") + .Value(pqReadTopic.Compression()) + .Done()); + } + + const auto clusterConfiguration = GetClusterConfiguration(pqReadTopic.DataSource().Cluster().StringValue()); + Add(innerSettings, SharedReading, ToString(UseSharedReading(clusterConfiguration, pqReadTopic.Format().Ref().Content())), pos, ctx); + + if (!innerSettings.empty()) { + settings.push_back(Build(ctx, pos) + .Name().Build("settings") + .Value() + .Add(innerSettings) + .Build() + .Done()); + } + + return Build(ctx, pos) + .Add(settings) + .Done(); + } + + const TPqClusterConfigurationSettings* GetClusterConfiguration(const TString& cluster) const { + const auto clusterConfiguration = State_->Configuration->ClustersConfigurationSettings.FindPtr(cluster); + if (!clusterConfiguration) { + ythrow yexception() << "Unknown pq cluster \"" << cluster << "\""; + } + return clusterConfiguration; + } + + static bool UseSharedReading(const TPqClusterConfigurationSettings* clusterConfiguration, std::string_view format) { + return clusterConfiguration->SharedReading && (format == "json_each_row" || format == "raw"); + } + private: TPqState* State_; // State owns dq integration, so back reference must be not smart. }; diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h b/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h index e9dd09de6009..cbe70bb3b3c4 100644 --- a/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h +++ b/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h @@ -6,13 +6,23 @@ #include #include +#include #include #include #include +namespace NKikimr { +namespace NMiniKQL { +class IFunctionRegistry; +} // namespace NMiniKQL +} + namespace NYql { +class TPqGatewayConfig; +using TPqGatewayConfigPtr = std::shared_ptr; + struct IPqGateway : public TThrRefBase { using TPtr = TIntrusivePtr; @@ -36,6 +46,17 @@ struct IPqGateway : public TThrRefBase { const TString& endpoint, const TString& database, bool secure) = 0; + + virtual void UpdateClusterConfigs(const TPqGatewayConfigPtr& config) = 0; + + virtual NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() const = 0; +}; + +struct IPqGatewayFactory : public TThrRefBase { + using TPtr = TIntrusivePtr; + + virtual ~IPqGatewayFactory() = default; + virtual IPqGateway::TPtr CreatePqGateway() = 0; }; } // namespace NYql diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_helpers.cpp b/ydb/library/yql/providers/pq/provider/yql_pq_helpers.cpp index c4fc523825f1..de525c933685 100644 --- a/ydb/library/yql/providers/pq/provider/yql_pq_helpers.cpp +++ b/ydb/library/yql/providers/pq/provider/yql_pq_helpers.cpp @@ -101,5 +101,20 @@ void FillSettingsWithResolvedYdsIds( } } +TMaybeNode FindSetting(TExprNode::TPtr settings, TStringBuf name) { + const auto maybeSettingsList = TMaybeNode(settings); + if (!maybeSettingsList) { + return nullptr; + } + const auto settingsList = maybeSettingsList.Cast(); + + for (size_t i = 0; i < settingsList.Size(); ++i) { + TCoNameValueTuple setting = settingsList.Item(i); + if (setting.Name().Value() == name) { + return setting.Value(); + } + } + return nullptr; +} } // namespace NYql diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_helpers.h b/ydb/library/yql/providers/pq/provider/yql_pq_helpers.h index 0a64879f0352..bf328554a817 100644 --- a/ydb/library/yql/providers/pq/provider/yql_pq_helpers.h +++ b/ydb/library/yql/providers/pq/provider/yql_pq_helpers.h @@ -21,4 +21,6 @@ void FillSettingsWithResolvedYdsIds( const TPqState::TPtr& state, const TDatabaseResolverResponse::TDatabaseDescriptionMap& fullResolvedIds); +NNodes::TMaybeNode FindSetting(TExprNode::TPtr settings, TStringBuf name); + } // namespace NYql diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_logical_opt.cpp b/ydb/library/yql/providers/pq/provider/yql_pq_logical_opt.cpp index 60a541aa94c9..fc667dbe6030 100644 --- a/ydb/library/yql/providers/pq/provider/yql_pq_logical_opt.cpp +++ b/ydb/library/yql/providers/pq/provider/yql_pq_logical_opt.cpp @@ -216,6 +216,7 @@ class TPqLogicalOptProposalTransformer : public TOptimizeTransformerBase { .RowSpec(DropUnusedRowItems(pqTopic.RowSpec().Pos(), inputRowType, usedColumnNames, ctx)) .Build() .Columns(DropUnusedColumns(dqPqTopicSource.Columns(), usedColumnNames, ctx)) + .RowType(DropUnusedRowItems(dqPqTopicSource.RowType().Pos(), oldRowType, usedColumnNames, ctx)) .Done() .Ptr(); diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_mkql_compiler.cpp b/ydb/library/yql/providers/pq/provider/yql_pq_mkql_compiler.cpp index cd708f07c0ef..483393d9b3a8 100644 --- a/ydb/library/yql/providers/pq/provider/yql_pq_mkql_compiler.cpp +++ b/ydb/library/yql/providers/pq/provider/yql_pq_mkql_compiler.cpp @@ -1,5 +1,7 @@ #include "yql_pq_mkql_compiler.h" +#include "yql_pq_helpers.h" +#include #include #include #include @@ -9,10 +11,52 @@ namespace NYql { using namespace NKikimr::NMiniKQL; using namespace NNodes; +namespace { + +bool UseSharedReading(TExprNode::TPtr settings) { + const auto maybeInnerSettings = FindSetting(settings, "settings"); + if (!maybeInnerSettings) { + return false; + } + + const auto maybeSharedReadingSetting = FindSetting(maybeInnerSettings.Cast().Ptr(), SharedReading); + if (!maybeSharedReadingSetting) { + return false; + } + + const TExprNode& value = maybeSharedReadingSetting.Cast().Ref(); + return value.IsAtom() && FromString(value.Content()); +} + +TRuntimeNode WrapSharedReading(const TDqSourceWrapBase &wrapper, NCommon::TMkqlBuildContext& ctx) { + const auto input = MkqlBuildExpr(wrapper.Input().Ref(), ctx); + const auto flow = ctx.ProgramBuilder.ToFlow(input); + + const TStructExprType* rowType = wrapper.RowType().Ref().GetTypeAnn()->Cast()->GetType()->Cast(); + const auto* finalItemStructType = static_cast(NCommon::BuildType(wrapper.RowType().Ref(), *rowType, ctx.ProgramBuilder)); + + return ctx.ProgramBuilder.ExpandMap(flow, [&](TRuntimeNode item) -> TRuntimeNode::TList { + TRuntimeNode::TList fields; + fields.reserve(finalItemStructType->GetMembersCount()); + for (ui32 i = 0; i < finalItemStructType->GetMembersCount(); ++i) { + fields.push_back(ctx.ProgramBuilder.Member(item, finalItemStructType->GetMemberName(i))); + } + return fields; + }); +} + +} + void RegisterDqPqMkqlCompilers(NCommon::TMkqlCallableCompilerBase& compiler) { compiler.ChainCallable(TDqSourceWideWrap::CallableName(), [](const TExprNode& node, NCommon::TMkqlBuildContext& ctx) { if (const auto wrapper = TDqSourceWideWrap(&node); wrapper.DataSource().Category().Value() == PqProviderName) { + if (const auto maybeSettings = wrapper.Settings()) { + if (UseSharedReading(maybeSettings.Cast().Ptr())) { + return WrapSharedReading(wrapper, ctx); + } + } + const auto wrapped = TryWrapWithParser(wrapper, ctx); if (wrapped) { return *wrapped; diff --git a/ydb/library/yql/providers/s3/actors/yql_s3_applicator_actor.cpp b/ydb/library/yql/providers/s3/actors/yql_s3_applicator_actor.cpp index 4cf07f3f8e16..d569df2bd037 100644 --- a/ydb/library/yql/providers/s3/actors/yql_s3_applicator_actor.cpp +++ b/ydb/library/yql/providers/s3/actors/yql_s3_applicator_actor.cpp @@ -49,7 +49,7 @@ struct TCompleteMultipartUpload { } TString BuildUrl() const { - TUrlBuilder urlBuilder(Url); + NS3Util::TUrlBuilder urlBuilder(NS3Util::UrlEscapeRet(Url)); urlBuilder.AddUrlParam("uploadId", UploadId); return urlBuilder.Build(); } @@ -87,7 +87,7 @@ struct TListMultipartUploads { // This requirement will be fixed in the curl library // https://github.com/curl/curl/commit/fc76a24c53b08cdf6eec8ba787d8eac64651d56e // https://github.com/curl/curl/commit/c87920353883ef9d5aa952e724a8e2589d76add5 - TUrlBuilder urlBuilder(Url); + NS3Util::TUrlBuilder urlBuilder(NS3Util::UrlEscapeRet(Url)); if (KeyMarker) { urlBuilder.AddUrlParam("key-marker", KeyMarker); } @@ -114,7 +114,7 @@ struct TAbortMultipartUpload { } TString BuildUrl() const { - TUrlBuilder urlBuilder(Url); + NS3Util::TUrlBuilder urlBuilder(NS3Util::UrlEscapeRet(Url)); urlBuilder.AddUrlParam("uploadId", UploadId); return urlBuilder.Build(); } @@ -141,7 +141,7 @@ struct TListParts { // This requirement will be fixed in the curl library // https://github.com/curl/curl/commit/fc76a24c53b08cdf6eec8ba787d8eac64651d56e // https://github.com/curl/curl/commit/c87920353883ef9d5aa952e724a8e2589d76add5 - TUrlBuilder urlBuilder(Url); + NS3Util::TUrlBuilder urlBuilder(NS3Util::UrlEscapeRet(Url)); if (PartNumberMarker) { urlBuilder.AddUrlParam("part-number-marker", PartNumberMarker); } @@ -682,4 +682,4 @@ THolder MakeS3ApplicatorActor( ); } -} // namespace NYql::NDq \ No newline at end of file +} // namespace NYql::NDq diff --git a/ydb/library/yql/providers/s3/actors/yql_s3_write_actor.cpp b/ydb/library/yql/providers/s3/actors/yql_s3_write_actor.cpp index c737921bb320..da2ba03a2c97 100644 --- a/ydb/library/yql/providers/s3/actors/yql_s3_write_actor.cpp +++ b/ydb/library/yql/providers/s3/actors/yql_s3_write_actor.cpp @@ -573,15 +573,15 @@ class TS3WriteActor : public TActorBootstrapped, public IDqComput const auto& key = MakePartitionKey(row); const auto [keyIt, insertedNew] = FileWriteActors.emplace(key, std::vector()); if (insertedNew || keyIt->second.empty() || keyIt->second.back()->IsFinishing()) { - auto fileWrite = std::make_unique( - TxId, - Gateway, - Credentials, - key, - NS3Util::UrlEscapeRet(Url + Path + key + MakeOutputName() + Extension), - Compression, - RetryPolicy, DirtyWrite, Token); - keyIt->second.emplace_back(fileWrite.get()); + auto fileWrite = std::make_unique( + TxId, + Gateway, + Credentials, + key, + NS3Util::UrlEscapeRet(Url + Path + key + MakeOutputName() + Extension), + Compression, + RetryPolicy, DirtyWrite, Token); + keyIt->second.emplace_back(fileWrite.get()); RegisterWithSameMailbox(fileWrite.release()); } @@ -619,6 +619,10 @@ class TS3WriteActor : public TActorBootstrapped, public IDqComput NDqProto::StatusIds::StatusCode statusCode = result->Get()->StatusCode; if (statusCode == NDqProto::StatusIds::UNSPECIFIED) { statusCode = StatusFromS3ErrorCode(result->Get()->S3ErrorCode); + if (statusCode == NDqProto::StatusIds::UNSPECIFIED) { + statusCode = NDqProto::StatusIds::INTERNAL_ERROR; + result->Get()->Issues.AddIssue("Got upload error with unspecified error code."); + } } Callbacks->OnAsyncOutputError(OutputIndex, result->Get()->Issues, statusCode); @@ -640,10 +644,15 @@ class TS3WriteActor : public TActorBootstrapped, public IDqComput if (const auto ft = std::find_if(it->second.cbegin(), it->second.cend(), [&](TS3FileWriteActor* actor){ return result->Get()->Url == actor->GetUrl(); }); it->second.cend() != ft) { (*ft)->PassAway(); it->second.erase(ft); - if (it->second.empty()) + if (it->second.empty()) { FileWriteActors.erase(it); + } } } + if (!Finished && GetFreeSpace() > 0) { + LOG_D("TS3WriteActor", "Has free space, notify owner"); + Callbacks->ResumeExecution(); + } FinishIfNeeded(); } diff --git a/ydb/library/yql/providers/s3/common/util.cpp b/ydb/library/yql/providers/s3/common/util.cpp index 074b116cc8c2..54c99fbad3fe 100644 --- a/ydb/library/yql/providers/s3/common/util.cpp +++ b/ydb/library/yql/providers/s3/common/util.cpp @@ -64,4 +64,34 @@ bool ValidateS3ReadWriteSchema(const TStructExprType* schemaStructRowType, TExpr return true; } +TUrlBuilder::TUrlBuilder(const TString& uri) + : MainUri(uri) +{} + +TUrlBuilder& TUrlBuilder::AddUrlParam(const TString& name, const TString& value) { + Params.emplace_back(name, value); + return *this; +} + +TString TUrlBuilder::Build() const { + if (Params.empty()) { + return MainUri; + } + + TStringBuilder result; + result << MainUri << "?"; + + TStringBuf separator = ""sv; + for (const auto& p : Params) { + result << separator << p.Name; + if (auto value = p.Value) { + Quote(value, ""); + result << "=" << value; + } + separator = "&"sv; + } + + return std::move(result); +} + } diff --git a/ydb/library/yql/providers/s3/common/util.h b/ydb/library/yql/providers/s3/common/util.h index d364e971078f..00df60e48f89 100644 --- a/ydb/library/yql/providers/s3/common/util.h +++ b/ydb/library/yql/providers/s3/common/util.h @@ -15,4 +15,22 @@ TString UrlEscapeRet(const TStringBuf from); bool ValidateS3ReadWriteSchema(const TStructExprType* schemaStructRowType, TExprContext& ctx); +class TUrlBuilder { + struct TParam { + TString Name; + TString Value; + }; + +public: + explicit TUrlBuilder(const TString& uri); + + TUrlBuilder& AddUrlParam(const TString& name, const TString& value = ""); + + TString Build() const; + +private: + std::vector Params; + TString MainUri; +}; + } diff --git a/ydb/library/yql/providers/s3/common/util_ut.cpp b/ydb/library/yql/providers/s3/common/util_ut.cpp index 2dcbf47ceef3..6e6a02e936e7 100644 --- a/ydb/library/yql/providers/s3/common/util_ut.cpp +++ b/ydb/library/yql/providers/s3/common/util_ut.cpp @@ -30,4 +30,37 @@ Y_UNIT_TEST_SUITE(TestS3UrlEscape) { } } +Y_UNIT_TEST_SUITE(TestUrlBuilder) { + Y_UNIT_TEST(UriOnly) { + TUrlBuilder builder("https://localhost/abc"); + UNIT_ASSERT_VALUES_EQUAL(builder.Build(), "https://localhost/abc"); + } + + Y_UNIT_TEST(Basic) { + TUrlBuilder builder("https://localhost/abc"); + builder.AddUrlParam("param1", "val1"); + builder.AddUrlParam("param2", "val2"); + + UNIT_ASSERT_VALUES_EQUAL(builder.Build(), "https://localhost/abc?param1=val1¶m2=val2"); + } + + Y_UNIT_TEST(BasicWithEncoding) { + auto url = TUrlBuilder("https://localhost/abc") + .AddUrlParam("param1", "=!@#$%^&*(){}[]\" ") + .AddUrlParam("param2", "val2") + .Build(); + + UNIT_ASSERT_VALUES_EQUAL(url, "https://localhost/abc?param1=%3D%21%40%23%24%25%5E%26%2A%28%29%7B%7D%5B%5D%22+¶m2=val2"); + } + + Y_UNIT_TEST(BasicWithAdditionalEncoding) { + auto url = TUrlBuilder("https://localhost/abc") + .AddUrlParam("param1", ":/?#[]@!$&\'()*+,;=") + .AddUrlParam("param2", "val2") + .Build(); + + UNIT_ASSERT_VALUES_EQUAL(url, "https://localhost/abc?param1=%3A%2F%3F%23%5B%5D%40%21%24%26%27%28%29%2A%2B%2C%3B%3D¶m2=val2"); + } +} + } // namespace NYql::NS3Util diff --git a/ydb/library/yql/providers/s3/object_listers/yql_s3_list.cpp b/ydb/library/yql/providers/s3/object_listers/yql_s3_list.cpp index c12bd71f0613..ea05c077a71e 100644 --- a/ydb/library/yql/providers/s3/object_listers/yql_s3_list.cpp +++ b/ydb/library/yql/providers/s3/object_listers/yql_s3_list.cpp @@ -304,7 +304,7 @@ class TS3Lister : public IS3Lister { // This requirement will be fixed in the curl library // https://github.com/curl/curl/commit/fc76a24c53b08cdf6eec8ba787d8eac64651d56e // https://github.com/curl/curl/commit/c87920353883ef9d5aa952e724a8e2589d76add5 - TUrlBuilder urlBuilder(ctx.ListingRequest.Url); + NS3Util::TUrlBuilder urlBuilder(ctx.ListingRequest.Url); if (ctx.ContinuationToken.Defined()) { urlBuilder.AddUrlParam("continuation-token", *ctx.ContinuationToken); } diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_phy_opt.cpp b/ydb/library/yql/providers/s3/provider/yql_s3_phy_opt.cpp index fe533331cdf5..3feae46c3495 100644 --- a/ydb/library/yql/providers/s3/provider/yql_s3_phy_opt.cpp +++ b/ydb/library/yql/providers/s3/provider/yql_s3_phy_opt.cpp @@ -124,13 +124,15 @@ class TS3PhysicalOptProposalTransformer : public TOptimizeTransformerBase { auto keys = GetPartitionKeys(partBy); auto sinkSettingsBuilder = Build(ctx, target.Pos()); - if (partBy) + if (partBy) { sinkSettingsBuilder.Add(std::move(partBy)); + } auto compression = GetCompression(settings); const auto& extension = GetExtension(target.Format().Value(), compression ? compression->Tail().Content() : ""sv); - if (compression) + if (compression) { sinkSettingsBuilder.Add(std::move(compression)); + } auto sinkOutputSettingsBuilder = Build(ctx, target.Pos()); if (auto csvDelimiter = GetCsvDelimiter(settings)) { @@ -199,7 +201,7 @@ class TS3PhysicalOptProposalTransformer : public TOptimizeTransformerBase { } } - if (!FindNode(input.Ptr(), [] (const TExprNode::TPtr& node) { return node->IsCallable(TCoDataSource::CallableName()); })) { + if (IsDqPureExpr(input)) { YQL_CLOG(INFO, ProviderS3) << "Rewrite pure S3WriteObject `" << cluster << "`.`" << target.Path().StringValue() << "` as stage with sink."; return keys.empty() ? Build(ctx, writePos) @@ -305,23 +307,26 @@ class TS3PhysicalOptProposalTransformer : public TOptimizeTransformerBase { .Build() .Done(); - auto outputsBuilder = Build(ctx, target.Pos()); - if (inputStage.Outputs() && keys.empty()) { - outputsBuilder.InitFrom(inputStage.Outputs().Cast()); - } - outputsBuilder.Add(sink); + auto outputsBuilder = Build(ctx, target.Pos()) + .Add(sink); if (keys.empty()) { - const auto outputBuilder = Build(ctx, target.Pos()) - .Input(inputStage.Program().Body().Ptr()) - .Format(target.Format()) - .KeyColumns().Add(std::move(keys)).Build() - .Settings(sinkOutputSettingsBuilder.Done()) - .Done(); - return Build(ctx, writePos) - .InitFrom(inputStage) - .Program(ctx.DeepCopyLambda(inputStage.Program().Ref(), outputBuilder.Ptr())) + .Inputs() + .Add() + .Output(dqUnion.Output()) + .Build() + .Build() + .Program() + .Args({"in"}) + .Body() + .Input("in") + .Format(target.Format()) + .KeyColumns().Add(std::move(keys)).Build() + .Settings(sinkOutputSettingsBuilder.Done()) + .Build() + .Build() + .Settings().Build() .Outputs(outputsBuilder.Done()) .Done(); } else { diff --git a/ydb/library/yql/tools/dq/worker_node/main.cpp b/ydb/library/yql/tools/dq/worker_node/main.cpp index f43e31656864..9a349b0f493c 100644 --- a/ydb/library/yql/tools/dq/worker_node/main.cpp +++ b/ydb/library/yql/tools/dq/worker_node/main.cpp @@ -115,11 +115,12 @@ NDq::IDqAsyncIoFactory::TPtr CreateAsyncIoFactory(const NYdb::TDriver& driver, I std::make_shared(), nullptr ); - RegisterDqPqReadActorFactory(*factory, driver, nullptr, CreatePqNativeGateway(std::move(pqServices))); + auto pqGateway = CreatePqNativeGateway(std::move(pqServices)); + RegisterDqPqReadActorFactory(*factory, driver, nullptr, pqGateway); RegisterYdbReadActorFactory(*factory, driver, nullptr); RegisterClickHouseReadActorFactory(*factory, nullptr, httpGateway); - RegisterDqPqWriteActorFactory(*factory, driver, nullptr); + RegisterDqPqWriteActorFactory(*factory, driver, nullptr, pqGateway); auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory(); auto retryPolicy = GetHTTPDefaultRetryPolicy(); diff --git a/ydb/library/yql/tools/dqrun/dqrun.cpp b/ydb/library/yql/tools/dqrun/dqrun.cpp index 64af597ec8b2..0b109a3af273 100644 --- a/ydb/library/yql/tools/dqrun/dqrun.cpp +++ b/ydb/library/yql/tools/dqrun/dqrun.cpp @@ -311,7 +311,7 @@ NDq::IDqAsyncIoFactory::TPtr CreateAsyncIoFactory( RegisterDQSolomonReadActorFactory(*factory, nullptr); RegisterClickHouseReadActorFactory(*factory, nullptr, httpGateway); RegisterGenericProviderFactories(*factory, credentialsFactory, genericClient); - RegisterDqPqWriteActorFactory(*factory, driver, nullptr); + RegisterDqPqWriteActorFactory(*factory, driver, nullptr, pqGateway); auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory(); s3ActorsFactory->RegisterS3WriteActorFactory(*factory, nullptr, httpGateway, GetHTTPDefaultRetryPolicy()); diff --git a/ydb/library/yql/tools/mrrun/mrrun.cpp b/ydb/library/yql/tools/mrrun/mrrun.cpp index 8e1a09477269..37fdd07c5802 100644 --- a/ydb/library/yql/tools/mrrun/mrrun.cpp +++ b/ydb/library/yql/tools/mrrun/mrrun.cpp @@ -245,10 +245,11 @@ NDq::IDqAsyncIoFactory::TPtr CreateAsyncIoFactory(const NYdb::TDriver& driver, I std::make_shared(), nullptr ); - RegisterDqPqReadActorFactory(*factory, driver, nullptr, CreatePqNativeGateway(std::move(pqServices))); + auto pqGateway = CreatePqNativeGateway(std::move(pqServices)); + RegisterDqPqReadActorFactory(*factory, driver, nullptr, pqGateway); RegisterYdbReadActorFactory(*factory, driver, nullptr); RegisterClickHouseReadActorFactory(*factory, nullptr, httpGateway); - RegisterDqPqWriteActorFactory(*factory, driver, nullptr); + RegisterDqPqWriteActorFactory(*factory, driver, nullptr, pqGateway); auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory(); auto retryPolicy = GetHTTPDefaultRetryPolicy(); diff --git a/ydb/public/sdk/cpp/client/ydb_topic/impl/write_session_impl.cpp b/ydb/public/sdk/cpp/client/ydb_topic/impl/write_session_impl.cpp index 46b774419a6b..683b1d6c0b91 100644 --- a/ydb/public/sdk/cpp/client/ydb_topic/impl/write_session_impl.cpp +++ b/ydb/public/sdk/cpp/client/ydb_topic/impl/write_session_impl.cpp @@ -965,6 +965,7 @@ void TWriteSessionImpl::OnReadDone(NYdbGrpc::TGrpcStatus&& grpcStatus, size_t co TStringBuilder TWriteSessionImpl::LogPrefix() const { TStringBuilder ret; + ret << " TraceId [" << Settings.TraceId_ << "] "; ret << " SessionId [" << SessionId << "] "; if (Settings.PartitionId_.Defined() || DirectWriteToPartitionId.Defined()) { diff --git a/ydb/tests/fq/control_plane_storage/in_memory_control_plane_storage_ut.cpp b/ydb/tests/fq/control_plane_storage/in_memory_control_plane_storage_ut.cpp index c7266237bd7f..de0bdb869de4 100644 --- a/ydb/tests/fq/control_plane_storage/in_memory_control_plane_storage_ut.cpp +++ b/ydb/tests/fq/control_plane_storage/in_memory_control_plane_storage_ut.cpp @@ -1,3 +1,5 @@ +#include + #include #include @@ -12,74 +14,111 @@ #include #include +#include + namespace NFq { using namespace NActors; using namespace NKikimr; +using namespace NFqRun; +using namespace NKikimrRun; namespace { -//////////////////////////////////////////////////////////////////////////////// +class TTestFuxture : public NUnitTest::TBaseFixture { +protected: + using TBase = NUnitTest::TBaseFixture; -using TRuntimePtr = std::unique_ptr; + static constexpr TDuration WAIT_TIMEOUT = TDuration::Seconds(15); -struct TTestBootstrap { - NConfig::TControlPlaneStorageConfig Config; - TRuntimePtr Runtime; +public: + void SetUp(NUnitTest::TTestContext& ctx) override { + TBase::SetUp(ctx); - const TDuration RequestTimeout = TDuration::Seconds(10); + SetupSignalActions(); + EnableYDBBacktraceFormat(); - TTestBootstrap() - { - Runtime = PrepareTestActorRuntime(); - } + TFqSetupSettings settings; + settings.VerboseLevel = TFqSetupSettings::EVerbose::Max; - std::pair CreateQuery() - { - TActorId sender = Runtime->AllocateEdgeActor(); - FederatedQuery::CreateQueryRequest proto; + settings.LogConfig.SetDefaultLevel(NLog::EPriority::PRI_WARN); + ModifyLogPriorities({{NKikimrServices::EServiceKikimr::YQ_CONTROL_PLANE_STORAGE, NLog::EPriority::PRI_DEBUG}}, settings.LogConfig); - FederatedQuery::QueryContent& content = *proto.mutable_content(); - content.set_name("my_query_1"); - content.set_text("SELECT 1;"); - content.set_type(FederatedQuery::QueryContent::ANALYTICS); - content.mutable_acl()->set_visibility(FederatedQuery::Acl::SCOPE); + auto& cpStorageConfig = *settings.FqConfig.MutableControlPlaneStorage(); + cpStorageConfig.SetEnabled(true); + cpStorageConfig.SetUseInMemory(true); - auto request = std::make_unique("yandexcloud://my_cloud_1/my_folder_1", proto, "user@staff", "", "mock_cloud", TPermissions{}, TQuotaMap{}, std::make_shared(), FederatedQuery::Internal::ComputeDatabaseInternal{}); - Runtime->Send(new IEventHandle(ControlPlaneStorageServiceActorId(), sender, request.release())); + auto& privateApiConfig = *settings.FqConfig.MutablePrivateApi(); + privateApiConfig.SetEnabled(true); + privateApiConfig.SetLoopback(true); + + settings.FqConfig.SetEnableDynamicNameservice(true); + settings.FqConfig.MutableControlPlaneProxy()->SetEnabled(true); + settings.FqConfig.MutableDbPool()->SetEnabled(true); + settings.FqConfig.MutableNodesManager()->SetEnabled(true); + settings.FqConfig.MutablePendingFetcher()->SetEnabled(true); + settings.FqConfig.MutablePrivateProxy()->SetEnabled(true); + settings.FqConfig.MutableGateways()->SetEnabled(true); + settings.FqConfig.MutableResourceManager()->SetEnabled(true); + settings.FqConfig.MutableCommon()->SetIdsPrefix("ut"); + + FqSetup = std::make_unique(settings); + } - TAutoPtr handle; - TEvControlPlaneStorage::TEvCreateQueryResponse* event = Runtime->GrabEdgeEvent(handle); - return {event->Result, event->Issues}; +protected: + static void CheckSuccess(const TRequestResult& result) { + UNIT_ASSERT_VALUES_EQUAL_C(result.Status, Ydb::StatusIds::SUCCESS, result.Issues.ToOneLineString()); } -private: - TRuntimePtr PrepareTestActorRuntime() { - TRuntimePtr runtime(new TTestBasicRuntime(1)); - runtime->SetLogPriority(NKikimrServices::YQ_CONTROL_PLANE_STORAGE, NLog::PRI_DEBUG); + TExecutionMeta WaitQueryExecution(const TString& queryId, TDuration timeout = WAIT_TIMEOUT) const { + using EStatus = FederatedQuery::QueryMeta; - auto controlPlaneProxy = CreateInMemoryControlPlaneStorageServiceActor(Config); + const TInstant start = TInstant::Now(); + while (TInstant::Now() - start <= timeout) { + TExecutionMeta meta; + CheckSuccess(FqSetup->DescribeQuery(queryId, meta)); - runtime->AddLocalService( - ControlPlaneStorageServiceActorId(), - TActorSetupCmd(controlPlaneProxy, TMailboxType::Simple, 0)); + if (!IsIn({EStatus::FAILED, EStatus::COMPLETED, EStatus::ABORTED_BY_USER, EStatus::ABORTED_BY_SYSTEM}, meta.Status)) { + Cerr << "Wait query execution " << TInstant::Now() - start << ": " << EStatus::ComputeStatus_Name(meta.Status) << "\n"; + Sleep(TDuration::Seconds(1)); + continue; + } - SetupTabletServices(*runtime); + UNIT_ASSERT_C(meta.Status == EStatus::COMPLETED, "issues: " << meta.Issues.ToOneLineString() << ", transient issues: " << meta.TransientIssues.ToOneLineString()); + return meta; + } - return runtime; + UNIT_ASSERT_C(false, "Waiting query execution timeout. Spent time " << TInstant::Now() - start << " exceeds limit " << timeout); + return {}; } + +protected: + std::unique_ptr FqSetup; }; -} +} // anonymous namespace + +Y_UNIT_TEST_SUITE(InMemoryControlPlaneStorage) { + Y_UNIT_TEST_F(ExecuteSimpleQuery, TTestFuxture) { + TString queryId; + CheckSuccess(FqSetup->StreamRequest({ + .Query = "SELECT 42 AS result_value" + }, queryId)); + UNIT_ASSERT(queryId); -Y_UNIT_TEST_SUITE(CreateQueryRequest) { - Y_UNIT_TEST(ShouldCreateSimpleQuery) - { - TTestBootstrap bootstrap; - const auto [result, issues] = bootstrap.CreateQuery(); - UNIT_ASSERT(!issues); + const auto meta = WaitQueryExecution(queryId); + UNIT_ASSERT_VALUES_EQUAL(meta.ResultSetSizes.size(), 1); + UNIT_ASSERT_VALUES_EQUAL(meta.ResultSetSizes[0], 1); + + Ydb::ResultSet resultSet; + CheckSuccess(FqSetup->FetchQueryResults(queryId, 0, resultSet)); + + NYdb::TResultSetParser parser(resultSet); + UNIT_ASSERT_VALUES_EQUAL(parser.ColumnsCount(), 1); + UNIT_ASSERT_VALUES_EQUAL(parser.RowsCount(), 1); + UNIT_ASSERT(parser.TryNextRow()); + UNIT_ASSERT_VALUES_EQUAL(parser.ColumnParser("result_value").GetInt32(), 42); } } - } // NFq diff --git a/ydb/tests/fq/control_plane_storage/ya.make b/ydb/tests/fq/control_plane_storage/ya.make index 889ba53f9d36..b8ee3698c58d 100644 --- a/ydb/tests/fq/control_plane_storage/ya.make +++ b/ydb/tests/fq/control_plane_storage/ya.make @@ -14,6 +14,7 @@ ENDIF() PEERDIR( library/cpp/retry library/cpp/testing/unittest + ydb/core/base ydb/core/external_sources ydb/core/fq/libs/actors/logging ydb/core/fq/libs/init @@ -22,6 +23,8 @@ PEERDIR( ydb/core/fq/libs/rate_limiter/events ydb/core/testlib/default ydb/library/security + ydb/tests/tools/fqrun/src + ydb/tests/tools/kqprun/runlib ) INCLUDE(${ARCADIA_ROOT}/ydb/public/tools/ydb_recipe/recipe.inc) diff --git a/ydb/tests/fq/control_plane_storage/ydb_test_bootstrap.h b/ydb/tests/fq/control_plane_storage/ydb_test_bootstrap.h index 7620c35a1578..a80b7cef0084 100644 --- a/ydb/tests/fq/control_plane_storage/ydb_test_bootstrap.h +++ b/ydb/tests/fq/control_plane_storage/ydb_test_bootstrap.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -31,6 +32,9 @@ #include #include #include + +#include + #include #include @@ -41,6 +45,7 @@ namespace NFq { using namespace NActors; using namespace NKikimr; +using namespace NKikimrRun; ////////////////////////////////////////////////////// @@ -96,6 +101,9 @@ struct TTestBootstrap { TTestBootstrap(std::string tablePrefix, const NConfig::TControlPlaneStorageConfig& config = {}) : Config(config) { + SetupSignalActions(); + EnableYDBBacktraceFormat(); + Cerr << "Netstat: " << Exec("netstat --all --program") << Endl; Cerr << "Process stat: " << Exec("ps aux") << Endl; Cerr << "YDB receipt endpoint: " << GetEnv("YDB_ENDPOINT") << ", database: " << GetEnv("YDB_DATABASE") << Endl; diff --git a/ydb/tests/fq/generic/analytics/docker-compose.yml b/ydb/tests/fq/generic/analytics/docker-compose.yml index b71cb5fee9c9..6288502d252e 100644 --- a/ydb/tests/fq/generic/analytics/docker-compose.yml +++ b/ydb/tests/fq/generic/analytics/docker-compose.yml @@ -19,7 +19,7 @@ services: echo \"$$(dig tests-fq-generic-analytics-ydb +short) tests-fq-generic-analytics-ydb\" >> /etc/hosts; cat /etc/hosts; /opt/ydb/bin/fq-connector-go server -c /opt/ydb/cfg/fq-connector-go.yaml" container_name: tests-fq-generic-analytics-fq-connector-go - image: ghcr.io/ydb-platform/fq-connector-go:v0.5.12-rc.2@sha256:84bb0b19f16f354b8a9ef7a020ee80f3ba7dc28db92f7007f235241153025b8a + image: ghcr.io/ydb-platform/fq-connector-go:v0.6.0-rc.1@sha256:4f74bd11e696218053f48622d1d065567d103906c187b3a24ea4e56f886c6c60 ports: - "2130" volumes: diff --git a/ydb/tests/fq/generic/analytics/test_join.py b/ydb/tests/fq/generic/analytics/test_join.py index 89b1cdc93aec..336a6cd673fd 100644 --- a/ydb/tests/fq/generic/analytics/test_join.py +++ b/ydb/tests/fq/generic/analytics/test_join.py @@ -7,12 +7,12 @@ from ydb.tests.tools.fq_runner.fq_client import FederatedQueryClient from ydb.tests.fq.generic.utils.settings import Settings from ydb.library.yql.providers.generic.connector.tests.utils.one_time_waiter import OneTimeWaiter -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind import conftest one_time_waiter = OneTimeWaiter( - data_source_kind=EDataSourceKind.YDB, + data_source_kind=EGenericDataSourceKind.YDB, docker_compose_file_path=conftest.docker_compose_file_path, expected_tables=["join_table", "dummy_table"], ) diff --git a/ydb/tests/fq/generic/analytics/test_ydb.py b/ydb/tests/fq/generic/analytics/test_ydb.py index 0bddcc30fe66..3ee8911ee033 100644 --- a/ydb/tests/fq/generic/analytics/test_ydb.py +++ b/ydb/tests/fq/generic/analytics/test_ydb.py @@ -8,12 +8,12 @@ from ydb.tests.tools.fq_runner.fq_client import FederatedQueryClient from ydb.tests.fq.generic.utils.settings import Settings from ydb.library.yql.providers.generic.connector.tests.utils.one_time_waiter import OneTimeWaiter -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind import conftest one_time_waiter = OneTimeWaiter( - data_source_kind=EDataSourceKind.YDB, + data_source_kind=EGenericDataSourceKind.YDB, docker_compose_file_path=conftest.docker_compose_file_path, expected_tables=["simple_table", "dummy_table"], ) diff --git a/ydb/tests/fq/generic/analytics/ya.make b/ydb/tests/fq/generic/analytics/ya.make index a17d05b8ee48..33f675520edb 100644 --- a/ydb/tests/fq/generic/analytics/ya.make +++ b/ydb/tests/fq/generic/analytics/ya.make @@ -54,7 +54,7 @@ IF (OPENSOURCE) ENDIF() PEERDIR( - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/tests/utils ydb/tests/fq/generic/utils library/python/testing/recipe diff --git a/ydb/tests/fq/generic/streaming/docker-compose.yml b/ydb/tests/fq/generic/streaming/docker-compose.yml index d6ea9aa7b69b..4bdfb1219436 100644 --- a/ydb/tests/fq/generic/streaming/docker-compose.yml +++ b/ydb/tests/fq/generic/streaming/docker-compose.yml @@ -5,7 +5,7 @@ services: echo \"$$(dig tests-fq-generic-streaming-ydb +short) tests-fq-generic-streaming-ydb\" >> /etc/hosts; cat /etc/hosts; /opt/ydb/bin/fq-connector-go server -c /opt/ydb/cfg/fq-connector-go.yaml" container_name: tests-fq-generic-streaming-fq-connector-go - image: ghcr.io/ydb-platform/fq-connector-go:v0.5.12-rc.2@sha256:84bb0b19f16f354b8a9ef7a020ee80f3ba7dc28db92f7007f235241153025b8a + image: ghcr.io/ydb-platform/fq-connector-go:v0.6.0-rc.1@sha256:4f74bd11e696218053f48622d1d065567d103906c187b3a24ea4e56f886c6c60 ports: - "2130" volumes: diff --git a/ydb/tests/fq/generic/streaming/test_join.py b/ydb/tests/fq/generic/streaming/test_join.py index b61e6a86f572..d982410e22d9 100644 --- a/ydb/tests/fq/generic/streaming/test_join.py +++ b/ydb/tests/fq/generic/streaming/test_join.py @@ -12,7 +12,7 @@ from ydb.tests.tools.datastreams_helpers.test_yds_base import TestYdsBase from ydb.library.yql.providers.generic.connector.tests.utils.one_time_waiter import OneTimeWaiter -from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind +from yql.essentials.providers.common.proto.gateways_config_pb2 import EGenericDataSourceKind import conftest @@ -53,7 +53,7 @@ def freeze(json): e.Data as data, u.id as lookup from $input as e - left join {streamlookup} ydb_conn_{table_name}.{table_name} as u + left join {streamlookup} any ydb_conn_{table_name}.{table_name} as u on(e.Data = u.data) ; @@ -83,7 +83,7 @@ def freeze(json): e.Data as data, CAST(e.Data AS Int32) as id, u.data as lookup from $input as e - left join {streamlookup} ydb_conn_{table_name}.{table_name} as u + left join {streamlookup} any ydb_conn_{table_name}.{table_name} as u on(CAST(e.Data AS Int32) = u.id) ; @@ -121,7 +121,7 @@ def freeze(json): u.data as lookup from $input as e - left join {streamlookup} ydb_conn_{table_name}.{table_name} as u + left join {streamlookup} any ydb_conn_{table_name}.{table_name} as u on(e.user = u.id) ; @@ -165,7 +165,7 @@ def freeze(json): u.data as lookup from $input as e - left join {streamlookup} ydb_conn_{table_name}.{table_name} as u + left join {streamlookup} any ydb_conn_{table_name}.{table_name} as u on(e.user = u.id) ; @@ -230,7 +230,7 @@ def freeze(json): u.age as age from $input as e - left join {streamlookup} ydb_conn_{table_name}.`users` as u + left join {streamlookup} any ydb_conn_{table_name}.`users` as u on(e.user = u.id) ; @@ -270,6 +270,12 @@ def freeze(json): ] * 1000 ), + "TTL", + "10", + "MaxCachedRows", + "5", + "MaxDelayedRows", + "100", ), # 5 ( @@ -290,7 +296,7 @@ def freeze(json): eu.id as uid from $input as e - left join {streamlookup} ydb_conn_{table_name}.`users` as eu + left join {streamlookup} any ydb_conn_{table_name}.`users` as eu on(e.user = eu.id) ; @@ -333,7 +339,7 @@ def freeze(json): $enriched = select a, b, c, d, e, f, za, yb, yc, zd from $input as e - left join {streamlookup} ydb_conn_{table_name}.db as u + left join {streamlookup} any ydb_conn_{table_name}.db as u on(e.yb = u.b AND e.za = u.a ) ; @@ -378,7 +384,7 @@ def freeze(json): $enriched = select a, b, c, d, e, f, za, yb, yc, zd from $input as e - left join {streamlookup} ydb_conn_{table_name}.db as u + left join {streamlookup} any ydb_conn_{table_name}.db as u on(e.za = u.a AND e.yb = u.b) ; @@ -406,11 +412,120 @@ def freeze(json): ] ), ), + # 8 + ( + R''' + $input = SELECT * FROM myyds.`{input_topic}` + WITH ( + FORMAT=json_each_row, + SCHEMA ( + za Int32, + yb STRING, + yc Int32, + zd Int32, + ) + ) ; + + $enriched1 = select a, b, c, d, e, f, za, yb, yc, zd + from + $input as e + left join {streamlookup} any ydb_conn_{table_name}.db as u + on(e.za = u.a AND e.yb = u.b) + ; + + $enriched2 = SELECT e.a AS a, e.b AS b, e.c AS c, e.d AS d, e.e AS e, e.f AS f, za, yb, yc, zd, u.c AS c2, u.d AS d2 + from + $enriched1 as e + left join {streamlookup} any ydb_conn_{table_name}.db as u + on(e.za = u.a AND e.yb = u.b) + ; + + $enriched = select a, b, c, d, e, f, za, yb, yc, zd, (c2 IS NOT DISTINCT FROM c) as eq1, (d2 IS NOT DISTINCT FROM d) as eq2 + from + $enriched2 as e + ; + + insert into myyds.`{output_topic}` + select Unwrap(Yson::SerializeJson(Yson::From(TableRow()))) from $enriched; + ''', + ResequenceId( + [ + ( + '{"id":1,"za":1,"yb":"2","yc":100,"zd":101}', + '{"a":1,"b":"2","c":3,"d":4,"e":5,"f":6,"za":1,"yb":"2","yc":100,"zd":101,"eq1":true,"eq2":true}', + ), + ( + '{"id":2,"za":7,"yb":"8","yc":106,"zd":107}', + '{"a":7,"b":"8","c":9,"d":10,"e":11,"f":12,"za":7,"yb":"8","yc":106,"zd":107,"eq1":true,"eq2":true}', + ), + ( + '{"id":3,"za":2,"yb":"1","yc":114,"zd":115}', + '{"a":null,"b":null,"c":null,"d":null,"e":null,"f":null,"za":2,"yb":"1","yc":114,"zd":115,"eq1":true,"eq2":true}', + ), + ( + '{"id":3,"za":null,"yb":"1","yc":114,"zd":115}', + '{"a":null,"b":null,"c":null,"d":null,"e":null,"f":null,"za":null,"yb":"1","yc":114,"zd":115,"eq1":true,"eq2":true}', + ), + ] + ), + ), + # 9 + ( + R''' + $input = SELECT * FROM myyds.`{input_topic}` + WITH ( + FORMAT=json_each_row, + SCHEMA ( + a Int32, + b STRING, + c Int32, + d Int32, + ) + ) ; + + $enriched12 = select u.a as a, u.b as b, u.c as c, u.d as d, u.e as e, u.f as f, e.a as za, e.b as yb, e.c as yc, e.d as zd, u2.c as c2, u2.d as d2 + from + $input as e + left join {streamlookup} any ydb_conn_{table_name}.db as u + on(e.a = u.a AND e.b = u.b) + left join {streamlookup} any ydb_conn_{table_name}.db as u2 + on(e.b = u2.b AND e.a = u2.a) + ; + + $enriched = select a, b, c, d, e, f, za, yb, yc, zd, (c2 IS NOT DISTINCT FROM c) as eq1, (d2 IS NOT DISTINCT FROM d) as eq2 + from + $enriched12 as e + ; + + insert into myyds.`{output_topic}` + select Unwrap(Yson::SerializeJson(Yson::From(TableRow()))) from $enriched; + ''', + ResequenceId( + [ + ( + '{"id":1,"a":1,"b":"2","c":100,"d":101}', + '{"a":1,"b":"2","c":3,"d":4,"e":5,"f":6,"za":1,"yb":"2","yc":100,"zd":101,"eq1":true,"eq2":true}', + ), + ( + '{"id":2,"a":7,"b":"8","c":106,"d":107}', + '{"a":7,"b":"8","c":9,"d":10,"e":11,"f":12,"za":7,"yb":"8","yc":106,"zd":107,"eq1":true,"eq2":true}', + ), + ( + '{"id":3,"a":2,"b":"1","c":114,"d":115}', + '{"a":null,"b":null,"c":null,"d":null,"e":null,"f":null,"za":2,"yb":"1","yc":114,"zd":115,"eq1":true,"eq2":true}', + ), + ( + '{"id":3,"a":null,"b":"1","c":114,"d":115}', + '{"a":null,"b":null,"c":null,"d":null,"e":null,"f":null,"za":null,"yb":"1","yc":114,"zd":115,"eq1":true,"eq2":true}', + ), + ] + ), + ), ] one_time_waiter = OneTimeWaiter( - data_source_kind=EDataSourceKind.YDB, + data_source_kind=EGenericDataSourceKind.YDB, docker_compose_file_path=conftest.docker_compose_file_path, expected_tables=["simple_table", "join_table", "dummy_table"], ) @@ -501,12 +616,12 @@ def test_streamlookup( database_id='local', ) - sql, messages = TESTCASES[testcase] + sql, messages, *options = TESTCASES[testcase] sql = sql.format( input_topic=self.input_topic, output_topic=self.output_topic, table_name=table_name, - streamlookup=R'/*+ streamlookup() */' if streamlookup else '', + streamlookup=Rf'/*+ streamlookup({" ".join(options)}) */' if streamlookup else '', ) one_time_waiter.wait() diff --git a/ydb/tests/fq/generic/streaming/ya.make b/ydb/tests/fq/generic/streaming/ya.make index 2183f7213e63..fddcf60053a1 100644 --- a/ydb/tests/fq/generic/streaming/ya.make +++ b/ydb/tests/fq/generic/streaming/ya.make @@ -56,7 +56,7 @@ DEPENDS( ) PEERDIR( - ydb/library/yql/providers/generic/connector/api/common + yql/essentials/providers/common/proto ydb/library/yql/providers/generic/connector/tests/utils ydb/tests/fq/generic/utils ydb/tests/tools/datastreams_helpers diff --git a/ydb/tests/fq/generic/streaming/ydb/01_basic.sh b/ydb/tests/fq/generic/streaming/ydb/01_basic.sh index e53213f2e6ad..d8a73368b075 100755 --- a/ydb/tests/fq/generic/streaming/ydb/01_basic.sh +++ b/ydb/tests/fq/generic/streaming/ydb/01_basic.sh @@ -35,7 +35,7 @@ set -ex (56, 12, "2a02:1812:1713:4f00:517e:1d79:c88b:704", "Elena", 2), (18, 17, "ivalid ip", "newUser", 12); COMMIT; - CREATE TABLE db (b STRING NOT NULL, c Int32, a Int32 NOT NULL, d Int32, f Int32, e Int32, PRIMARY KEY(b, a)); + CREATE TABLE db (b STRING NOT NULL, c Int32, a Int32 NOT NULL, d Int32, f Int32, e Int32, g Int32, h Int32, PRIMARY KEY(b, a)); COMMIT; INSERT INTO db (a, b, c, d, e, f) VALUES (1, "2", 3, 4, 5, 6), diff --git a/ydb/tests/fq/http_api/test_http_api.py b/ydb/tests/fq/http_api/test_http_api.py index 74e454afe53e..08130fb11c18 100644 --- a/ydb/tests/fq/http_api/test_http_api.py +++ b/ydb/tests/fq/http_api/test_http_api.py @@ -76,7 +76,7 @@ def test_simple_analytics_query(self): assert len(query_id) == 20 status = client.get_query_status(query_id) - assert status in ["FAILED", "RUNNING", "COMPLETED"] + assert status in ["STARTING", "RUNNING", "COMPLETED", "COMPLETING"] wait_for_query_status(client, query_id, ["COMPLETED"]) query_json = client.get_query(query_id) @@ -98,6 +98,14 @@ def test_simple_analytics_query(self): response = client.stop_query(query_id) assert response.status_code == 204 + response = client.start_query(query_id) + assert response.status_code == 204 + + assert client.get_query_status(query_id) in ["STARTING", "RUNNING", "COMPLETED", "COMPLETING"] + + response = client.stop_query(query_id) + assert response.status_code == 204 + def test_empty_query(self): with self.create_client() as client: with pytest.raises( @@ -228,6 +236,28 @@ def test_stop_idempotency(self): self.streaming_over_kikimr.compute_plane.start() c.wait_query_status(query_id, fq.QueryMeta.ABORTED_BY_USER) + def test_restart_idempotency(self): + c = FederatedQueryClient("my_folder", streaming_over_kikimr=self.streaming_over_kikimr) + self.streaming_over_kikimr.compute_plane.stop() + query_id = c.create_query("select1", "select 1").result.query_id + c.wait_query_status(query_id, fq.QueryMeta.STARTING) + + with self.create_client() as client: + response1 = client.stop_query(query_id, idempotency_key="Z") + assert response1.status_code == 204 + + response2 = client.start_query(query_id, idempotency_key="Z") + assert response2.status_code == 204 + + response2 = client.start_query(query_id, idempotency_key="Z") + assert response2.status_code == 204 + + response1 = client.stop_query(query_id, idempotency_key="Z") + assert response1.status_code == 204 + + self.streaming_over_kikimr.compute_plane.start() + c.wait_query_status(query_id, fq.QueryMeta.COMPLETED) + def test_simple_streaming_query(self): self.init_topics("simple_streaming_query", create_output=False) c = FederatedQueryClient("my_folder", streaming_over_kikimr=self.streaming_over_kikimr) diff --git a/ydb/tests/fq/pq_async_io/mock_pq_gateway.cpp b/ydb/tests/fq/pq_async_io/mock_pq_gateway.cpp new file mode 100644 index 000000000000..5f41cf2b9ac3 --- /dev/null +++ b/ydb/tests/fq/pq_async_io/mock_pq_gateway.cpp @@ -0,0 +1,190 @@ +#include "mock_pq_gateway.h" + +#include +#include + +namespace NYql::NDq { + +namespace { + +using TQueue = NYql::TBlockingEQueue; + +class TMockTopicReadSession : public NYdb::NTopic::IReadSession { +public: + TMockTopicReadSession(NYdb::NTopic::TPartitionSession::TPtr session, std::shared_ptr queue) + : Session(std::move(session)) + , Queue(queue) { + if (Queue->IsStopped()) { + Queue->~TBlockingEQueue(); + new (Queue.get()) TQueue(4_MB); + } + ThreadPool.Start(1); + } + + NThreading::TFuture WaitEvent() override { + return NThreading::Async([&] () { + Queue->BlockUntilEvent(); + return NThreading::MakeFuture(); + }, ThreadPool); + } + + TVector GetEvents(bool block, TMaybe maxEventsCount, size_t /*maxByteSize*/) override { + TVector res; + for (auto event = Queue->Pop(block); !event.Empty() && res.size() <= maxEventsCount.GetOrElse(std::numeric_limits::max()); event = Queue->Pop(/*block=*/ false)) { + res.push_back(*event); + } + return res; + } + + TVector GetEvents(const NYdb::NTopic::TReadSessionGetEventSettings& settings) override { + return GetEvents(settings.Block_, settings.MaxEventsCount_, settings.MaxByteSize_); + } + + TMaybe GetEvent(bool block, size_t /*maxByteSize*/) override { + return Queue->Pop(block); + } + + TMaybe GetEvent(const NYdb::NTopic::TReadSessionGetEventSettings& settings) override { + return GetEvent(settings.Block_, settings.MaxByteSize_); + } + + bool Close(TDuration timeout = TDuration::Max()) override { + Y_UNUSED(timeout); + Queue->Stop(); + ThreadPool.Stop(); + return true; + } + + NYdb::NTopic::TReaderCounters::TPtr GetCounters() const override {return nullptr;} + TString GetSessionId() const override {return "fake";} +private: + TThreadPool ThreadPool; + NYdb::NTopic::TPartitionSession::TPtr Session; + std::shared_ptr Queue; +}; + +struct TDummyPartitionSession: public NYdb::NTopic::TPartitionSession { + TDummyPartitionSession() {} + void RequestStatus() override {} +}; + +class TMockPqGateway : public IMockPqGateway { + + struct TMockTopicClient : public NYql::ITopicClient { + + TMockTopicClient(TMockPqGateway* self): Self(self) { } + + NYdb::TAsyncStatus CreateTopic(const TString& /*path*/, const NYdb::NTopic::TCreateTopicSettings& /*settings*/ = {}) override {return NYdb::TAsyncStatus{};} + NYdb::TAsyncStatus AlterTopic(const TString& /*path*/, const NYdb::NTopic::TAlterTopicSettings& /*settings*/ = {}) override {return NYdb::TAsyncStatus{};} + NYdb::TAsyncStatus DropTopic(const TString& /*path*/, const NYdb::NTopic::TDropTopicSettings& /*settings*/ = {}) override {return NYdb::TAsyncStatus{};} + NYdb::NTopic::TAsyncDescribeTopicResult DescribeTopic(const TString& /*path*/, + const NYdb::NTopic::TDescribeTopicSettings& /*settings*/ = {}) override {return NYdb::NTopic::TAsyncDescribeTopicResult{};} + + NYdb::NTopic::TAsyncDescribeConsumerResult DescribeConsumer(const TString& /*path*/, const TString& /*consumer*/, + const NYdb::NTopic::TDescribeConsumerSettings& /*settings*/ = {}) override {return NYdb::NTopic::TAsyncDescribeConsumerResult{};} + + NYdb::NTopic::TAsyncDescribePartitionResult DescribePartition(const TString& /*path*/, i64 /*partitionId*/, + const NYdb::NTopic::TDescribePartitionSettings& /*settings*/ = {}) override {return NYdb::NTopic::TAsyncDescribePartitionResult{};} + + std::shared_ptr CreateReadSession(const NYdb::NTopic::TReadSessionSettings& settings) override { + Y_ENSURE(!settings.Topics_.empty()); + TString topic = settings.Topics_.front().Path_; + Self->Runtime.Send(new NActors::IEventHandle(Self->Notifier, NActors::TActorId(), new NYql::NDq::TEvMockPqEvents::TEvCreateSession())); + return std::make_shared(MakeIntrusive(), Self->GetEventQueue(topic)); + } + + std::shared_ptr CreateSimpleBlockingWriteSession( + const NYdb::NTopic::TWriteSessionSettings& /*settings*/) override { + return nullptr; + } + std::shared_ptr CreateWriteSession(const NYdb::NTopic::TWriteSessionSettings& /*settings*/) override { + return nullptr; + } + + NYdb::TAsyncStatus CommitOffset(const TString& /*path*/, ui64 /*partitionId*/, const TString& /*consumerName*/, ui64 /*offset*/, + const NYdb::NTopic::TCommitOffsetSettings& /*settings*/ = {}) override {return NYdb::TAsyncStatus{};} + + TMockPqGateway* Self; + }; +public: + + TMockPqGateway( + NActors::TTestActorRuntime& runtime, + NActors::TActorId notifier) + : Runtime(runtime) + , Notifier(notifier) {} + + ~TMockPqGateway() {} + + NThreading::TFuture OpenSession(const TString& /*sessionId*/, const TString& /*username*/) override { + return NThreading::MakeFuture(); + } + NThreading::TFuture CloseSession(const TString& /*sessionId*/) override { + return NThreading::MakeFuture(); + } + + ::NPq::NConfigurationManager::TAsyncDescribePathResult DescribePath( + const TString& /*sessionId*/, + const TString& /*cluster*/, + const TString& /*database*/, + const TString& /*path*/, + const TString& /*token*/) override { + return ::NPq::NConfigurationManager::TAsyncDescribePathResult{}; + } + + NThreading::TFuture ListStreams( + const TString& /*sessionId*/, + const TString& /*cluster*/, + const TString& /*database*/, + const TString& /*token*/, + ui32 /*limit*/, + const TString& /*exclusiveStartStreamName*/ = {}) override { + return NThreading::TFuture{}; + } + + void UpdateClusterConfigs( + const TString& /*clusterName*/, + const TString& /*endpoint*/, + const TString& /*database*/, + bool /*secure*/) override {} + + void UpdateClusterConfigs(const TPqGatewayConfigPtr& /*config*/) override {}; + + NYql::ITopicClient::TPtr GetTopicClient(const NYdb::TDriver& /*driver*/, const NYdb::NTopic::TTopicClientSettings& /*settings*/) override { + return MakeIntrusive(this); + } + + std::shared_ptr GetEventQueue(const TString& topic) { + if (!Queues.contains(topic)) { + Queues[topic] = std::make_shared(4_MB); + } + return Queues[topic]; + } + + void AddEvent(const TString& topic, NYdb::NTopic::TReadSessionEvent::TEvent&& e, size_t size) override { + GetEventQueue(topic)->Push(std::move(e), size); + } + + NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() const override { + return NYdb::NTopic::TTopicClientSettings(); + } + +private: + std::unordered_map> Queues; + NActors::TTestActorRuntime& Runtime; + NActors::TActorId Notifier; +}; + +} + +NYdb::NTopic::TPartitionSession::TPtr CreatePartitionSession() { + return MakeIntrusive(); +} + +TIntrusivePtr CreateMockPqGateway( + NActors::TTestActorRuntime& runtime, + NActors::TActorId notifier) { + return MakeIntrusive(runtime, notifier); +} + +} diff --git a/ydb/tests/fq/pq_async_io/mock_pq_gateway.h b/ydb/tests/fq/pq_async_io/mock_pq_gateway.h new file mode 100644 index 000000000000..b6a05e170b7e --- /dev/null +++ b/ydb/tests/fq/pq_async_io/mock_pq_gateway.h @@ -0,0 +1,42 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include + +#include + +namespace NYql::NDq { + +struct TEvMockPqEvents { + enum EEv : ui32 { + EvBegin = EventSpaceBegin(NActors::TEvents::ES_PRIVATE), + EvCreateSession = EvBegin, + EvEnd + }; + static_assert(EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(TEvents::ES_PRIVATE)"); + struct TEvCreateSession : public NActors::TEventLocal {}; +}; + +class IMockPqGateway : public NYql::IPqGateway { +public: + virtual void AddEvent(const TString& topic, NYdb::NTopic::TReadSessionEvent::TEvent&& e, size_t size) = 0; +}; + +NYdb::NTopic::TPartitionSession::TPtr CreatePartitionSession(); + +TIntrusivePtr CreateMockPqGateway( + NActors::TTestActorRuntime& runtime, + NActors::TActorId notifier); + +} diff --git a/ydb/tests/fq/pq_async_io/ut/dq_pq_rd_read_actor_ut.cpp b/ydb/tests/fq/pq_async_io/ut/dq_pq_rd_read_actor_ut.cpp index 83fc68b55d0b..34ef1eacb7a8 100644 --- a/ydb/tests/fq/pq_async_io/ut/dq_pq_rd_read_actor_ut.cpp +++ b/ydb/tests/fq/pq_async_io/ut/dq_pq_rd_read_actor_ut.cpp @@ -1,21 +1,25 @@ #include +#include +#include +#include #include -#include #include #include #include +#include #include +#include #include namespace NYql::NDq { -const ui64 PartitionId = 666; +const ui64 PartitionId1 = 666; +const ui64 PartitionId2 = 667; struct TFixture : public TPqIoTestFixture { - TFixture() { LocalRowDispatcherId = CaSetup->Runtime->AllocateEdgeActor(); Coordinator1Id = CaSetup->Runtime->AllocateEdgeActor(); @@ -26,13 +30,14 @@ struct TFixture : public TPqIoTestFixture { void InitRdSource( const NYql::NPq::NProto::TDqPqTopicSource& settings, - i64 freeSpace = 1_MB) + i64 freeSpace = 1_MB, + ui64 partitionCount = 1) { CaSetup->Execute([&](TFakeActor& actor) { NPq::NProto::TDqReadTaskParams params; auto* partitioninigParams = params.MutablePartitioningParams(); - partitioninigParams->SetTopicPartitionsCount(1); - partitioninigParams->SetEachTopicPartitionGroupId(PartitionId); + partitioninigParams->SetTopicPartitionsCount(partitionCount); + partitioninigParams->SetEachTopicPartitionGroupId(PartitionId1); partitioninigParams->SetDqPartitionsCount(1); TString serializedParams; @@ -43,6 +48,7 @@ struct TFixture : public TPqIoTestFixture { NYql::NPq::NProto::TDqPqTopicSource copySettings = settings; auto [dqSource, dqSourceAsActor] = CreateDqPqRdReadActor( + actor.TypeEnv, std::move(copySettings), 0, NYql::NDq::TCollectStatsLevel::None, @@ -72,23 +78,33 @@ struct TFixture : public TPqIoTestFixture { return eventHolder; } - void ExpectStartSession(ui64 expectedOffset, NActors::TActorId rowDispatcherId, ui64 expectedGeneration = 1) { + void ExpectStartSession(const TMap& expectedOffsets, NActors::TActorId rowDispatcherId, ui64 expectedGeneration = 1) { auto eventHolder = CaSetup->Runtime->GrabEdgeEvent(rowDispatcherId, TDuration::Seconds(5)); UNIT_ASSERT(eventHolder.Get() != nullptr); - UNIT_ASSERT(eventHolder->Get()->Record.GetOffset() == expectedOffset); - UNIT_ASSERT(eventHolder->Cookie == expectedGeneration); + TMap offsets; + for (auto p : eventHolder->Get()->Record.GetOffsets()) { + offsets[p.GetPartitionId()] = p.GetOffset(); + } + UNIT_ASSERT_EQUAL(offsets, expectedOffsets); + UNIT_ASSERT_EQUAL(eventHolder->Cookie, expectedGeneration); } void ExpectStopSession(NActors::TActorId rowDispatcherId, ui64 expectedGeneration = 1) { auto eventHolder = CaSetup->Runtime->GrabEdgeEvent(rowDispatcherId, TDuration::Seconds(5)); UNIT_ASSERT(eventHolder.Get() != nullptr); + UNIT_ASSERT_EQUAL(eventHolder->Cookie, expectedGeneration); + } + + void ExpectNoSession(NActors::TActorId rowDispatcherId, ui64 expectedGeneration = 1) { + auto eventHolder = CaSetup->Runtime->GrabEdgeEvent(rowDispatcherId, TDuration::Seconds(5)); + UNIT_ASSERT(eventHolder.Get() != nullptr); UNIT_ASSERT(eventHolder->Cookie == expectedGeneration); } - void ExpectGetNextBatch(NActors::TActorId rowDispatcherId) { + void ExpectGetNextBatch(NActors::TActorId rowDispatcherId, ui64 partitionId = PartitionId1) { auto eventHolder = CaSetup->Runtime->GrabEdgeEvent(rowDispatcherId, TDuration::Seconds(5)); UNIT_ASSERT(eventHolder.Get() != nullptr); - UNIT_ASSERT(eventHolder->Get()->Record.GetPartitionId() == PartitionId); + UNIT_ASSERT_EQUAL(eventHolder->Get()->Record.GetPartitionId(), partitionId); } void MockCoordinatorChanged(NActors::TActorId coordinatorId) { @@ -98,20 +114,23 @@ struct TFixture : public TPqIoTestFixture { }); } - void MockCoordinatorResult(NActors::TActorId rowDispatcherId, ui64 cookie = 0) { + void MockCoordinatorResult(const TMap result, ui64 cookie = 0) { CaSetup->Execute([&](TFakeActor& actor) { auto event = new NFq::TEvRowDispatcher::TEvCoordinatorResult(); - auto* partitions = event->Record.AddPartitions(); - partitions->AddPartitionId(PartitionId); - ActorIdToProto(rowDispatcherId, partitions->MutableActorId()); + + for (const auto& [rowDispatcherId, partitionId] : result) { + auto* partitions = event->Record.AddPartitions(); + partitions->AddPartitionIds(partitionId); + ActorIdToProto(rowDispatcherId, partitions->MutableActorId()); + } CaSetup->Runtime->Send(new NActors::IEventHandle(*actor.DqAsyncInputActorId, Coordinator1Id, event, 0, cookie)); }); } - void MockAck(NActors::TActorId rowDispatcherId, ui64 generation = 1) { + void MockAck(NActors::TActorId rowDispatcherId, ui64 generation = 1, ui64 partitionId = PartitionId1) { CaSetup->Execute([&](TFakeActor& actor) { NFq::NRowDispatcherProto::TEvStartSession proto; - proto.SetPartitionId(PartitionId); + proto.AddPartitionIds(partitionId); auto event = new NFq::TEvRowDispatcher::TEvStartSessionAck(proto); CaSetup->Runtime->Send(new NActors::IEventHandle(*actor.DqAsyncInputActorId, rowDispatcherId, event, 0, generation)); }); @@ -119,29 +138,48 @@ struct TFixture : public TPqIoTestFixture { void MockHeartbeat(NActors::TActorId rowDispatcherId, ui64 generation = 1) { CaSetup->Execute([&](TFakeActor& actor) { - auto event = new NFq::TEvRowDispatcher::TEvHeartbeat(PartitionId); + auto event = new NFq::TEvRowDispatcher::TEvHeartbeat(); CaSetup->Runtime->Send(new NActors::IEventHandle(*actor.DqAsyncInputActorId, rowDispatcherId, event, 0, generation)); }); } - void MockNewDataArrived(NActors::TActorId rowDispatcherId, ui64 generation = 1) { + void MockNewDataArrived(NActors::TActorId rowDispatcherId, ui64 generation = 1, ui64 partitionId = PartitionId1) { CaSetup->Execute([&](TFakeActor& actor) { auto event = new NFq::TEvRowDispatcher::TEvNewDataArrived(); - event->Record.SetPartitionId(PartitionId); + event->Record.SetPartitionId(partitionId); CaSetup->Runtime->Send(new NActors::IEventHandle(*actor.DqAsyncInputActorId, rowDispatcherId, event, 0, generation)); }); } - void MockMessageBatch(ui64 offset, const std::vector& jsons, NActors::TActorId rowDispatcherId, ui64 generation = 1) { + TRope SerializeItem(TFakeActor& actor, ui64 intValue, const TString& strValue) { + NKikimr::NMiniKQL::TType* typeMkql = actor.ProgramBuilder.NewMultiType({ + NYql::NCommon::ParseTypeFromYson(TStringBuf("[DataType; Uint64]"), actor.ProgramBuilder, Cerr), + NYql::NCommon::ParseTypeFromYson(TStringBuf("[DataType; String]"), actor.ProgramBuilder, Cerr) + }); + UNIT_ASSERT_C(typeMkql, "Failed to create multi type"); + + NKikimr::NMiniKQL::TValuePackerTransport packer(typeMkql); + + TVector values = { + NUdf::TUnboxedValuePod(intValue), + NKikimr::NMiniKQL::MakeString(strValue) + }; + packer.AddWideItem(values.data(), 2); + + return NYql::MakeReadOnlyRope(packer.Finish()); + } + + // Supported schema (Uint64, String) + void MockMessageBatch(ui64 offset, const std::vector>& messages, NActors::TActorId rowDispatcherId, ui64 generation = 1, ui64 partitionId = PartitionId1) { CaSetup->Execute([&](TFakeActor& actor) { auto event = new NFq::TEvRowDispatcher::TEvMessageBatch(); - for (const auto& json :jsons) { + for (const auto& item : messages) { NFq::NRowDispatcherProto::TEvMessage message; - message.SetJson(json); - message.SetOffset(offset++); + message.SetPayloadId(event->AddPayload(SerializeItem(actor, item.first, item.second))); + message.AddOffsets(offset++); *event->Record.AddMessages() = message; } - event->Record.SetPartitionId(PartitionId); + event->Record.SetPartitionId(partitionId); event->Record.SetNextMessageOffset(offset); CaSetup->Runtime->Send(new NActors::IEventHandle(*actor.DqAsyncInputActorId, rowDispatcherId, event, 0, generation)); }); @@ -150,12 +188,22 @@ struct TFixture : public TPqIoTestFixture { void MockSessionError() { CaSetup->Execute([&](TFakeActor& actor) { auto event = new NFq::TEvRowDispatcher::TEvSessionError(); - event->Record.SetMessage("A problem has been detected and session has been shut down to prevent damage your life"); - event->Record.SetPartitionId(PartitionId); + event->Record.SetStatusCode(::NYql::NDqProto::StatusIds::BAD_REQUEST); + IssueToMessage(TIssue("A problem has been detected and session has been shut down to prevent damage your life"), event->Record.AddIssues()); CaSetup->Runtime->Send(new NActors::IEventHandle(*actor.DqAsyncInputActorId, RowDispatcher1, event, 0, 1)); }); } + void MockStatistics(NActors::TActorId rowDispatcherId, ui64 nextOffset, ui64 generation, ui64 partitionId) { + CaSetup->Execute([&](TFakeActor& actor) { + auto event = new NFq::TEvRowDispatcher::TEvStatistics(); + auto* partitionsProto = event->Record.AddPartition(); + partitionsProto->SetPartitionId(partitionId); + partitionsProto->SetNextMessageOffset(nextOffset); + CaSetup->Runtime->Send(new NActors::IEventHandle(*actor.DqAsyncInputActorId, rowDispatcherId, event, 0, generation)); + }); + } + template void AssertDataWithWatermarks( const std::vector>& actual, @@ -177,7 +225,7 @@ struct TFixture : public TPqIoTestFixture { watermarksBeforeIter == watermarkBeforePositions.end() || *watermarksBeforeIter > expectedPos, "Watermark before item on position " << expectedPos << " was expected"); - UNIT_ASSERT_EQUAL(std::get(item), expected.at(expectedPos)); + UNIT_ASSERT_VALUES_EQUAL(std::get(item), expected.at(expectedPos)); expectedPos++; } } @@ -197,41 +245,43 @@ struct TFixture : public TPqIoTestFixture { }); } - void MockUndelivered() { + void MockUndelivered(NActors::TActorId rowDispatcherId, ui64 generation = 0, NActors::TEvents::TEvUndelivered::EReason reason = NActors::TEvents::TEvUndelivered::ReasonActorUnknown) { CaSetup->Execute([&](TFakeActor& actor) { - auto event = new NActors::TEvents::TEvUndelivered(0, NActors::TEvents::TEvUndelivered::ReasonActorUnknown); - CaSetup->Runtime->Send(new NActors::IEventHandle(*actor.DqAsyncInputActorId, RowDispatcher1, event)); + auto event = new NActors::TEvents::TEvUndelivered(0, reason); + CaSetup->Runtime->Send(new NActors::IEventHandle(*actor.DqAsyncInputActorId, rowDispatcherId, event, 0, generation)); }); } - void StartSession(NYql::NPq::NProto::TDqPqTopicSource& settings, i64 freeSpace = 1_MB) { - InitRdSource(settings, freeSpace); - SourceRead(UVParser); + void StartSession(NYql::NPq::NProto::TDqPqTopicSource& settings, i64 freeSpace = 1_MB, ui64 partitionCount = 1) { + InitRdSource(settings, freeSpace, partitionCount); + SourceRead>(UVPairParser); ExpectCoordinatorChangesSubscribe(); MockCoordinatorChanged(Coordinator1Id); - auto req =ExpectCoordinatorRequest(Coordinator1Id); + auto req = ExpectCoordinatorRequest(Coordinator1Id); - MockCoordinatorResult(RowDispatcher1, req->Cookie); - ExpectStartSession(0, RowDispatcher1); + MockCoordinatorResult({{RowDispatcher1, PartitionId1}}, req->Cookie); + ExpectStartSession({}, RowDispatcher1); MockAck(RowDispatcher1); } - void ProcessSomeJsons(ui64 offset, const std::vector& jsons, NActors::TActorId rowDispatcherId, - std::function(const NUdf::TUnboxedValue&)> uvParser = UVParser, ui64 generation = 1) { - MockNewDataArrived(rowDispatcherId, generation); - ExpectGetNextBatch(rowDispatcherId); - - MockMessageBatch(offset, jsons, rowDispatcherId, generation); + void ProcessSomeMessages(ui64 offset, const std::vector>& messages, NActors::TActorId rowDispatcherId, + std::function>(const NUdf::TUnboxedValue&)> uvParser = UVPairParser, ui64 generation = 1, + ui64 partitionId = PartitionId1, bool readedByCA = true) { + MockNewDataArrived(rowDispatcherId, generation, partitionId); + ExpectGetNextBatch(rowDispatcherId, partitionId); - auto result = SourceReadDataUntil(uvParser, jsons.size()); - AssertDataWithWatermarks(result, jsons, {}); + MockMessageBatch(offset, messages, rowDispatcherId, generation, partitionId); + if (readedByCA) { + auto result = SourceReadDataUntil>(uvParser, messages.size()); + AssertDataWithWatermarks(result, messages, {}); + } } - const TString Json1 = "{\"dt\":100,\"value\":\"value1\"}"; - const TString Json2 = "{\"dt\":200,\"value\":\"value2\"}"; - const TString Json3 = "{\"dt\":300,\"value\":\"value3\"}"; - const TString Json4 = "{\"dt\":400,\"value\":\"value4\"}"; + const std::pair Message1 = {100, "value1"}; + const std::pair Message2 = {200, "value2"}; + const std::pair Message3 = {300, "value3"}; + const std::pair Message4 = {400, "value4"}; NYql::NPq::NProto::TDqPqTopicSource Source1 = BuildPqTopicSourceSettings("topicName"); @@ -243,9 +293,16 @@ struct TFixture : public TPqIoTestFixture { }; Y_UNIT_TEST_SUITE(TDqPqRdReadActorTests) { - Y_UNIT_TEST_F(TestReadFromTopic, TFixture) { + Y_UNIT_TEST_F(TestReadFromTopic2, TFixture) { StartSession(Source1); - ProcessSomeJsons(0, {Json1, Json2}, RowDispatcher1); + ProcessSomeMessages(0, {Message1, Message2}, RowDispatcher1); + } + + Y_UNIT_TEST_F(IgnoreUndeliveredWithWrongGeneration, TFixture) { + StartSession(Source1); + ProcessSomeMessages(0, {Message1, Message2}, RowDispatcher1); + MockUndelivered(RowDispatcher1, 999, NActors::TEvents::TEvUndelivered::Disconnected); + ProcessSomeMessages(2, {Message3}, RowDispatcher1); } Y_UNIT_TEST_F(SessionError, TFixture) { @@ -257,7 +314,7 @@ Y_UNIT_TEST_SUITE(TDqPqRdReadActorTests) { bool failured = false; while (Now() < deadline) { - SourceRead(UVParser); + SourceRead>(UVPairParser); if (future.HasValue()) { UNIT_ASSERT_STRING_CONTAINS(future.GetValue().ToOneLineString(), "damage your life"); failured = true; @@ -273,17 +330,18 @@ Y_UNIT_TEST_SUITE(TDqPqRdReadActorTests) { MockNewDataArrived(RowDispatcher1); ExpectGetNextBatch(RowDispatcher1); - const std::vector data1 = {Json1, Json2}; + const std::vector> data1 = {Message1, Message2}; MockMessageBatch(0, data1, RowDispatcher1); - const std::vector data2 = {Json3, Json4}; + const std::vector> data2 = {Message3, Message4}; MockMessageBatch(2, data2, RowDispatcher1); - auto result = SourceReadDataUntil(UVParser, 1, 1); - std::vector expected{data1}; + auto result = SourceReadDataUntil>(UVPairParser, 1, 1); + std::vector> expected{data1}; AssertDataWithWatermarks(result, expected, {}); - UNIT_ASSERT_EQUAL(SourceRead(UVParser, 0).size(), 0); + auto readSize = SourceRead>(UVPairParser, 0).size(); + UNIT_ASSERT_VALUES_EQUAL(readSize, 0); } Y_UNIT_TEST(TestSaveLoadPqRdRead) { @@ -292,7 +350,7 @@ Y_UNIT_TEST_SUITE(TDqPqRdReadActorTests) { { TFixture f; f.StartSession(f.Source1); - f.ProcessSomeJsons(0, {f.Json1, f.Json2}, f.RowDispatcher1); // offsets: 0, 1 + f.ProcessSomeMessages(0, {f.Message1, f.Message2}, f.RowDispatcher1); // offsets: 0, 1 f.SaveSourceState(CreateCheckpoint(), state); Cerr << "State saved" << Endl; @@ -300,19 +358,18 @@ Y_UNIT_TEST_SUITE(TDqPqRdReadActorTests) { { TFixture f; f.InitRdSource(f.Source1); - f.SourceRead(UVParser); f.LoadSource(state); - f.SourceRead(UVParser); + f.SourceRead>(UVPairParser); f.ExpectCoordinatorChangesSubscribe(); f.MockCoordinatorChanged(f.Coordinator1Id); auto req = f.ExpectCoordinatorRequest(f.Coordinator1Id); - f.MockCoordinatorResult(f.RowDispatcher1, req->Cookie); - f.ExpectStartSession(2, f.RowDispatcher1); + f.MockCoordinatorResult({{f.RowDispatcher1, PartitionId1}}, req->Cookie); + f.ExpectStartSession({{PartitionId1, 2}}, f.RowDispatcher1); f.MockAck(f.RowDispatcher1); - f.ProcessSomeJsons(2, {f.Json3}, f.RowDispatcher1); // offsets: 2 + f.ProcessSomeMessages(2, {f.Message3}, f.RowDispatcher1); // offsets: 2 state.Data.clear(); f.SaveSourceState(CreateCheckpoint(), state); Cerr << "State saved" << Endl; @@ -320,110 +377,165 @@ Y_UNIT_TEST_SUITE(TDqPqRdReadActorTests) { { TFixture f; f.InitRdSource(f.Source1); - f.SourceRead(UVParser); f.LoadSource(state); - f.SourceRead(UVParser); + f.SourceRead>(UVPairParser); f.ExpectCoordinatorChangesSubscribe(); f.MockCoordinatorChanged(f.Coordinator1Id); auto req = f.ExpectCoordinatorRequest(f.Coordinator1Id); - f.MockCoordinatorResult(f.RowDispatcher1, req->Cookie); - f.ExpectStartSession(3, f.RowDispatcher1); + f.MockCoordinatorResult({{f.RowDispatcher1, PartitionId1}}, req->Cookie); + f.ExpectStartSession({{PartitionId1, 3}}, f.RowDispatcher1); f.MockAck(f.RowDispatcher1); - f.ProcessSomeJsons(3, {f.Json4}, f.RowDispatcher1); // offsets: 3 + f.ProcessSomeMessages(3, {f.Message4}, f.RowDispatcher1); // offsets: 3 } } Y_UNIT_TEST_F(CoordinatorChanged, TFixture) { StartSession(Source1); - ProcessSomeJsons(0, {Json1, Json2}, RowDispatcher1); - MockMessageBatch(2, {Json3}, RowDispatcher1); + ProcessSomeMessages(0, {Message1, Message2}, RowDispatcher1); + MockMessageBatch(2, {Message3}, RowDispatcher1); // change active Coordinator MockCoordinatorChanged(Coordinator2Id); - ExpectStopSession(RowDispatcher1); - - auto result = SourceReadDataUntil(UVParser, 1); - AssertDataWithWatermarks(result, {Json3}, {}); + // continue use old sessions + MockMessageBatch(3, {Message4}, RowDispatcher1); auto req = ExpectCoordinatorRequest(Coordinator2Id); - MockCoordinatorResult(RowDispatcher2, req->Cookie); - ExpectStartSession(3, RowDispatcher2, 2); + auto result = SourceReadDataUntil>(UVPairParser, 2); + AssertDataWithWatermarks(result, {Message3, Message4}, {}); + + MockCoordinatorResult({{RowDispatcher2, PartitionId1}}, req->Cookie); // change distribution + ExpectStopSession(RowDispatcher1); + + ExpectStartSession({{PartitionId1, 4}}, RowDispatcher2, 2); MockAck(RowDispatcher2, 2); - ProcessSomeJsons(3, {Json4}, RowDispatcher2, UVParser, 2); + ProcessSomeMessages(4, {Message4}, RowDispatcher2, UVPairParser, 2); MockHeartbeat(RowDispatcher1, 1); // old generation - ExpectStopSession(RowDispatcher1); + ExpectNoSession(RowDispatcher1, 1); + + // change active Coordinator + MockCoordinatorChanged(Coordinator1Id); + req = ExpectCoordinatorRequest(Coordinator1Id); + MockCoordinatorResult({{RowDispatcher2, PartitionId1}}, req->Cookie); // distribution is not changed + ProcessSomeMessages(5, {Message2}, RowDispatcher2, UVPairParser, 2); } Y_UNIT_TEST_F(Backpressure, TFixture) { StartSession(Source1, 2_KB); - TString json(900, 'c'); - ProcessSomeJsons(0, {json}, RowDispatcher1); + std::pair message = {100500, TString(900, 'c')}; + ProcessSomeMessages(0, {message}, RowDispatcher1); MockNewDataArrived(RowDispatcher1); ExpectGetNextBatch(RowDispatcher1); - MockMessageBatch(0, {json, json, json}, RowDispatcher1); + MockMessageBatch(1, {message, message, message}, RowDispatcher1); MockNewDataArrived(RowDispatcher1); ASSERT_THROW( CaSetup->Runtime->GrabEdgeEvent(RowDispatcher1, TDuration::Seconds(0)), NActors::TEmptyEventQueueException); - auto result = SourceReadDataUntil(UVParser, 3); - AssertDataWithWatermarks(result, {json, json, json}, {}); + auto result = SourceReadDataUntil>(UVPairParser, 3); + AssertDataWithWatermarks(result, {message, message, message}, {}); ExpectGetNextBatch(RowDispatcher1); - MockMessageBatch(3, {Json1}, RowDispatcher1); - result = SourceReadDataUntil(UVParser, 1); - AssertDataWithWatermarks(result, {Json1}, {}); + MockMessageBatch(4, {Message1}, RowDispatcher1); + result = SourceReadDataUntil>(UVPairParser, 1); + AssertDataWithWatermarks(result, {Message1}, {}); } - Y_UNIT_TEST_F(RowDispatcherIsRestarted, TFixture) { + Y_UNIT_TEST_F(RowDispatcherIsRestarted2, TFixture) { StartSession(Source1); - ProcessSomeJsons(0, {Json1, Json2}, RowDispatcher1); + ProcessSomeMessages(0, {Message1, Message2}, RowDispatcher1); MockDisconnected(); MockConnected(); - MockUndelivered(); + MockUndelivered(RowDispatcher1, 1); auto req = ExpectCoordinatorRequest(Coordinator1Id); - MockCoordinatorResult(RowDispatcher1, req->Cookie); - ExpectStartSession(2, RowDispatcher1, 2); + MockCoordinatorResult({{RowDispatcher1, PartitionId1}}, req->Cookie); + ExpectStartSession({{PartitionId1, 2}}, RowDispatcher1, 2); MockAck(RowDispatcher1, 2); - ProcessSomeJsons(2, {Json3}, RowDispatcher1, UVParser, 2); + ProcessSomeMessages(2, {Message3}, RowDispatcher1, UVPairParser, 2); + } + + Y_UNIT_TEST_F(TwoPartitionsRowDispatcherIsRestarted, TFixture) { + InitRdSource(Source1, 1_MB, PartitionId2 + 1); + SourceRead(UVParser); + ExpectCoordinatorChangesSubscribe(); + MockCoordinatorChanged(Coordinator1Id); + auto req = ExpectCoordinatorRequest(Coordinator1Id); + MockCoordinatorResult({{RowDispatcher1, PartitionId1}, {RowDispatcher2, PartitionId2}}, req->Cookie); + ExpectStartSession({}, RowDispatcher1, 1); + ExpectStartSession({}, RowDispatcher2, 2); + MockAck(RowDispatcher1, 1, PartitionId1); + MockAck(RowDispatcher2, 2, PartitionId2); + + ProcessSomeMessages(0, {Message1, Message2}, RowDispatcher1, UVPairParser, 1, PartitionId1); + ProcessSomeMessages(0, {Message3}, RowDispatcher2, UVPairParser, 2, PartitionId2, false); // not read by CA + MockStatistics(RowDispatcher2, 10, 2, PartitionId2); + + // Restart node 2 (RowDispatcher2) + MockDisconnected(); + MockConnected(); + MockUndelivered(RowDispatcher2, 2); + + // session1 is still working + ProcessSomeMessages(2, {Message4}, RowDispatcher1, UVPairParser, 1, PartitionId1, false); + + // Reinit session to RowDispatcher2 + auto req2 = ExpectCoordinatorRequest(Coordinator1Id); + MockCoordinatorResult({{RowDispatcher1, PartitionId1}, {RowDispatcher2, PartitionId2}}, req2->Cookie); + ExpectStartSession({{PartitionId2, 10}}, RowDispatcher2, 3); + MockAck(RowDispatcher2, 3, PartitionId2); + + ProcessSomeMessages(3, {Message4}, RowDispatcher1, UVPairParser, 1, PartitionId1, false); + ProcessSomeMessages(10, {Message4}, RowDispatcher2, UVPairParser, 3, PartitionId2, false); + + std::vector> expected{Message3, Message4, Message4, Message4}; + auto result = SourceReadDataUntil>(UVPairParser, expected.size()); + AssertDataWithWatermarks(result, expected, {}); } Y_UNIT_TEST_F(IgnoreMessageIfNoSessions, TFixture) { StartSession(Source1); MockCoordinatorChanged(Coordinator2Id); + MockUndelivered(RowDispatcher1, 1); MockSessionError(); } Y_UNIT_TEST_F(MetadataFields, TFixture) { + auto metadataUVParser = [](const NUdf::TUnboxedValue& item) -> std::vector> { + UNIT_ASSERT_VALUES_EQUAL(item.GetListLength(), 3); + auto stringElement = item.GetElement(2); + return { {item.GetElement(1).Get(), TString(stringElement.AsStringRef())} }; + }; + auto source = BuildPqTopicSourceSettings("topicName"); source.AddMetadataFields("_yql_sys_create_time"); + source.SetRowType("[StructType; [[_yql_sys_create_time; [DataType; Uint32]]; [dt; [DataType; Uint64]]; [value; [DataType; String]]]]"); StartSession(source); - ProcessSomeJsons(0, {Json1}, RowDispatcher1, UVParserWithMetadatafields); + ProcessSomeMessages(0, {Message1}, RowDispatcher1, metadataUVParser); } Y_UNIT_TEST_F(IgnoreCoordinatorResultIfWrongState, TFixture) { StartSession(Source1); - ProcessSomeJsons(0, {Json1, Json2}, RowDispatcher1); + ProcessSomeMessages(0, {Message1, Message2}, RowDispatcher1); MockCoordinatorChanged(Coordinator2Id); auto req = ExpectCoordinatorRequest(Coordinator2Id); + MockUndelivered(RowDispatcher1, 1); + auto req2 = ExpectCoordinatorRequest(Coordinator2Id); - MockUndelivered(); + MockCoordinatorResult({{RowDispatcher1, PartitionId1}}, req->Cookie); + MockCoordinatorResult({{RowDispatcher1, PartitionId1}}, req2->Cookie); + ExpectStartSession({{PartitionId1, 2}}, RowDispatcher1, 2); - MockCoordinatorResult(RowDispatcher1, req->Cookie); - MockCoordinatorResult(RowDispatcher1, req->Cookie); - ExpectStartSession(2, RowDispatcher1, 2); MockAck(RowDispatcher1); } } diff --git a/ydb/tests/fq/pq_async_io/ut_helpers.cpp b/ydb/tests/fq/pq_async_io/ut_helpers.cpp index a9d716d730c1..e011029479ca 100644 --- a/ydb/tests/fq/pq_async_io/ut_helpers.cpp +++ b/ydb/tests/fq/pq_async_io/ut_helpers.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -12,6 +13,16 @@ namespace NYql::NDq { +namespace { + +void SegmentationFaultHandler(int) { + Cerr << "segmentation fault call stack:" << Endl; + FormatBackTrace(&Cerr); + abort(); +} + +} + using namespace NActors; NYql::NPq::NProto::TDqPqTopicSource BuildPqTopicSourceSettings( @@ -26,6 +37,11 @@ NYql::NPq::NProto::TDqPqTopicSource BuildPqTopicSourceSettings( settings.SetEndpoint(GetDefaultPqEndpoint()); settings.MutableToken()->SetName("token"); settings.SetDatabase(GetDefaultPqDatabase()); + settings.SetRowType("[StructType; [[dt; [DataType; Uint64]]; [value; [DataType; String]]]]"); + settings.AddColumns("dt"); + settings.AddColumns("value"); + settings.AddColumnTypes("[DataType; Uint64]"); + settings.AddColumnTypes("[DataType; String]"); if (watermarksPeriod) { settings.MutableWatermarks()->SetEnabled(true); settings.MutableWatermarks()->SetGranularityUs(watermarksPeriod->MicroSeconds()); @@ -48,6 +64,8 @@ NYql::NPq::NProto::TDqPqTopicSink BuildPqTopicSinkSettings(TString topic) { } TPqIoTestFixture::TPqIoTestFixture() { + NKikimr::EnableYDBBacktraceFormat(); + signal(SIGSEGV, &SegmentationFaultHandler); } TPqIoTestFixture::~TPqIoTestFixture() { @@ -106,6 +124,14 @@ void TPqIoTestFixture::InitAsyncOutput( { const THashMap secureParams; + TPqGatewayServices pqServices( + Driver, + nullptr, + nullptr, + std::make_shared(), + nullptr + ); + CaSetup->Execute([&](TFakeActor& actor) { auto [dqAsyncOutput, dqAsyncOutputAsActor] = CreateDqPqWriteActor( std::move(settings), @@ -118,6 +144,7 @@ void TPqIoTestFixture::InitAsyncOutput( nullptr, &actor.GetAsyncOutputCallbacks(), MakeIntrusive(), + CreatePqNativeGateway(std::move(pqServices)), freeSpace); actor.InitAsyncOutput(dqAsyncOutput, dqAsyncOutputAsActor); @@ -239,14 +266,14 @@ void AddReadRule(NYdb::TDriver& driver, const TString& streamName) { UNIT_ASSERT_VALUES_EQUAL(result.IsTransportError(), false); } -std::vector UVParser(const NUdf::TUnboxedValue& item) { - return { TString(item.AsStringRef()) }; +std::vector> UVPairParser(const NUdf::TUnboxedValue& item) { + UNIT_ASSERT_VALUES_EQUAL(item.GetListLength(), 2); + auto stringElement = item.GetElement(1); + return { {item.GetElement(0).Get(), TString(stringElement.AsStringRef())} }; } -std::vector UVParserWithMetadatafields(const NUdf::TUnboxedValue& item) { - const auto& cell = item.GetElement(0); - TString str(cell.AsStringRef()); - return {str}; +std::vector UVParser(const NUdf::TUnboxedValue& item) { + return { TString(item.AsStringRef()) }; } void TPqIoTestFixture::AsyncOutputWrite(std::vector data, TMaybe checkpoint) { diff --git a/ydb/tests/fq/pq_async_io/ut_helpers.h b/ydb/tests/fq/pq_async_io/ut_helpers.h index 114b678686a5..b685412b5d01 100644 --- a/ydb/tests/fq/pq_async_io/ut_helpers.h +++ b/ydb/tests/fq/pq_async_io/ut_helpers.h @@ -124,7 +124,7 @@ void AddReadRule( NYdb::TDriver& driver, const TString& streamName); +std::vector> UVPairParser(const NUdf::TUnboxedValue& item); std::vector UVParser(const NUdf::TUnboxedValue& item); -std::vector UVParserWithMetadatafields(const NUdf::TUnboxedValue& item); } diff --git a/ydb/tests/fq/pq_async_io/ya.make b/ydb/tests/fq/pq_async_io/ya.make index 487500ce4bab..75efb71a2b4f 100644 --- a/ydb/tests/fq/pq_async_io/ya.make +++ b/ydb/tests/fq/pq_async_io/ya.make @@ -1,12 +1,14 @@ LIBRARY() SRCS( + mock_pq_gateway.cpp ut_helpers.cpp ) PEERDIR( yql/essentials/minikql/computation/llvm14 ydb/library/yql/providers/common/ut_helpers + ydb/library/yql/providers/pq/gateway/dummy ydb/public/sdk/cpp/client/ydb_topic ) diff --git a/ydb/tests/fq/yds/test_row_dispatcher.py b/ydb/tests/fq/yds/test_row_dispatcher.py index 97bd610c189e..ebd6ef7a914f 100644 --- a/ydb/tests/fq/yds/test_row_dispatcher.py +++ b/ydb/tests/fq/yds/test_row_dispatcher.py @@ -5,6 +5,7 @@ import pytest import logging import time +import json from ydb.tests.tools.fq_runner.kikimr_utils import yq_v1 from ydb.tests.tools.datastreams_helpers.test_yds_base import TestYdsBase @@ -21,12 +22,13 @@ import ydb.public.api.protos.draft.fq_pb2 as fq YDS_CONNECTION = "yds" +COMPUTE_NODE_COUNT = 3 @pytest.fixture def kikimr(request): kikimr_conf = StreamingOverKikimrConfig( - cloud_mode=True, node_count={"/cp": TenantConfig(1), "/compute": TenantConfig(3)} + cloud_mode=True, node_count={"/cp": TenantConfig(1), "/compute": TenantConfig(COMPUTE_NODE_COUNT)} ) kikimr = StreamingOverKikimr(kikimr_conf) kikimr.compute_plane.fq_config['row_dispatcher']['enabled'] = True @@ -80,6 +82,23 @@ def wait_row_dispatcher_sensor_value(kikimr, sensor, expected_count, exact_match pass +def wait_public_sensor_value(kikimr, query_id, sensor, expected_value): + deadline = time.time() + 60 + cloud_id = "mock_cloud" + folder_id = "my_folder" + while True: + count = 0 + for node_index in kikimr.compute_plane.kikimr_cluster.nodes: + value = kikimr.compute_plane.get_sensors(node_index, "yq_public").find_sensor( + {"cloud_id": cloud_id, "folder_id": folder_id, "query_id": query_id, "name": sensor}) + count += value if value is not None else 0 + if count >= expected_value: + break + assert time.time() < deadline, f"Waiting sensor {sensor} value failed, current count {count}" + time.sleep(1) + pass + + class TestPqRowDispatcher(TestYdsBase): def run_and_check(self, kikimr, client, sql, input, output, expected_predicate): @@ -146,7 +165,6 @@ def test_simple_not_null(self, kikimr, client): query_id = start_yds_query(kikimr, client, sql) wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 1) - time.sleep(10) data = [ '{"time": 101, "data": "hello1", "event": "event1"}', @@ -167,6 +185,35 @@ def test_simple_not_null(self, kikimr, client): assert len(read_rules) == 0, read_rules wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 0) + @yq_v1 + def test_metadatafields(self, kikimr, client): + client.create_yds_connection( + YDS_CONNECTION, os.getenv("YDB_DATABASE"), os.getenv("YDB_ENDPOINT"), shared_reading=True + ) + self.init_topics("test_metadatafields") + + # Its not completely clear why metadatafields appear in this request( + sql = Rf''' + PRAGMA FeatureR010="prototype"; + PRAGMA config.flags("TimeOrderRecoverDelay", "-10"); + PRAGMA config.flags("TimeOrderRecoverAhead", "10"); + INSERT INTO {YDS_CONNECTION}.`{self.output_topic}` + SELECT ToBytes(Unwrap(Json::SerializeJson(Yson::From(TableRow())))) FROM {YDS_CONNECTION}.`{self.input_topic}` + WITH (format=json_each_row, SCHEMA (time Int32 NOT NULL)) + MATCH_RECOGNIZE( + ORDER BY CAST(time as Timestamp) + MEASURES LAST(A.time) as b_key + PATTERN (A ) + DEFINE A as A.time > 4 + );''' + + query_id = start_yds_query(kikimr, client, sql) + wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 1) + + self.write_stream(['{"time": 100}', '{"time": 120}']) + assert len(self.read_stream(1, topic_path=self.output_topic)) == 1 + stop_yds_query(client, query_id) + @yq_v1 def test_simple_optional(self, kikimr, client): client.create_yds_connection( @@ -216,7 +263,7 @@ def test_scheme_error(self, kikimr, client): client.wait_query_status(query_id, fq.QueryMeta.FAILED) issues = str(client.describe_query(query_id).result.query.issue) - assert "Cannot parse JSON string" in issues, "Incorrect Issues: " + issues + assert "Failed to parse json message for offset" in issues, "Incorrect Issues: " + issues wait_actor_count(kikimr, "DQ_PQ_READ_ACTOR", 0) wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 0) @@ -878,6 +925,38 @@ def test_many_partitions(self, kikimr, client): stop_yds_query(client, query_id) wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 0) + @yq_v1 + def test_2_connection(self, kikimr, client): + client.create_yds_connection("name1", os.getenv("YDB_DATABASE"), os.getenv("YDB_ENDPOINT"), shared_reading=True) + client.create_yds_connection("name2", os.getenv("YDB_DATABASE"), os.getenv("YDB_ENDPOINT"), shared_reading=True) + self.init_topics("test_2_connection", partitions_count=2) + create_read_rule(self.input_topic, "best") + + sql1 = Rf''' + INSERT INTO `name1`.`{self.output_topic}` + SELECT Cast(time as String) FROM `name1`.`{self.input_topic}` WITH (format=json_each_row, SCHEMA (time Int32 NOT NULL));''' + sql2 = Rf''' + INSERT INTO `name2`.`{self.output_topic}` + SELECT Cast(time as String) FROM `name2`.`{self.input_topic}` WITH (format=json_each_row, SCHEMA (time Int32 NOT NULL));''' + query_id1 = start_yds_query(kikimr, client, sql1) + query_id2 = start_yds_query(kikimr, client, sql2) + wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 4) + + input_messages1 = [Rf'''{{"time": {c}}}''' for c in range(100, 110)] + write_stream(self.input_topic, input_messages1, "partition_key1") + + input_messages2 = [Rf'''{{"time": {c}}}''' for c in range(110, 120)] + write_stream(self.input_topic, input_messages2, "partition_key2") + + expected = [Rf'''{c}''' for c in range(100, 120)] + expected = [item for item in expected for i in range(2)] + assert sorted(self.read_stream(len(expected), topic_path=self.output_topic)) == expected + + stop_yds_query(client, query_id1) + stop_yds_query(client, query_id2) + + wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 0) + @yq_v1 def test_sensors(self, kikimr, client): client.create_yds_connection( @@ -897,12 +976,23 @@ def test_sensors(self, kikimr, client): wait_actor_count(kikimr, "DQ_PQ_READ_ACTOR", 1) wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 1) + wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_COMPILE_SERVICE", COMPUTE_NODE_COUNT) + wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_FORMAT_HANDLER", 1) wait_row_dispatcher_sensor_value(kikimr, "ClientsCount", 1) wait_row_dispatcher_sensor_value(kikimr, "RowsSent", 1, exact_match=False) wait_row_dispatcher_sensor_value(kikimr, "IncomingRequests", 1, exact_match=False) + wait_public_sensor_value(kikimr, query_id, "query.input_filtered_bytes", 1) + wait_public_sensor_value(kikimr, query_id, "query.source_input_filtered_records", 1) + wait_public_sensor_value(kikimr, query_id, "query.input_queued_bytes", 0) + wait_public_sensor_value(kikimr, query_id, "query.source_input_queued_records", 0) stop_yds_query(client, query_id) wait_actor_count(kikimr, "DQ_PQ_READ_ACTOR", 0) wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 0) wait_row_dispatcher_sensor_value(kikimr, "ClientsCount", 0) + + stat = json.loads(client.describe_query(query_id).result.query.statistics.json) + filtered_bytes = stat['Graph=0']['IngressFilteredBytes']['sum'] + filtered_rows = stat['Graph=0']['IngressFilteredRows']['sum'] + assert filtered_bytes > 1 and filtered_rows > 0 diff --git a/ydb/tests/tools/fq_runner/kikimr_runner.py b/ydb/tests/tools/fq_runner/kikimr_runner.py index 6af8dac66a35..9ed1d2fcfd1d 100644 --- a/ydb/tests/tools/fq_runner/kikimr_runner.py +++ b/ydb/tests/tools/fq_runner/kikimr_runner.py @@ -478,6 +478,13 @@ def fill_config(self, control_plane): fq_config['test_connection'] = {'enabled': True} fq_config['common']['keep_internal_errors'] = True + fq_config['common']['ydb_driver_config'] = {} + fq_config['common']['ydb_driver_config']['network_threads_num'] = 1 + fq_config['common']['ydb_driver_config']['client_threads_num'] = 1 + + fq_config['common']['topic_client_handlers_executor_threads_num'] = 1 + fq_config['common']['topic_client_compression_executor_threads_num'] = 1 + if self.mvp_mock_port is not None: fq_config['common']['ydb_mvp_cloud_endpoint'] = "localhost:" + str(self.mvp_mock_port) diff --git a/ydb/tests/tools/fqrun/.gitignore b/ydb/tests/tools/fqrun/.gitignore new file mode 100644 index 000000000000..ee3ed5154ff0 --- /dev/null +++ b/ydb/tests/tools/fqrun/.gitignore @@ -0,0 +1,9 @@ +sync_dir + +*.log +*.sql +*.conf +*.parquet +*.json +*.svg +*.txt diff --git a/ydb/tests/tools/fqrun/README.md b/ydb/tests/tools/fqrun/README.md new file mode 100644 index 000000000000..dcaee95ccc7c --- /dev/null +++ b/ydb/tests/tools/fqrun/README.md @@ -0,0 +1,44 @@ +# FQ run tool + +Tool can be used to execute streaming queries by using FQ proxy infrastructure. + +For profiling memory allocations build fqrun with ya make flags `-D PROFILE_MEMORY_ALLOCATIONS -D CXXFLAGS=-DPROFILE_MEMORY_ALLOCATIONS`. + +## Scripts + +* `flame_graph.sh` - script for collecting flame graphs in svg format, usage: + ```(bash) + ./flame_graph.sh [graph collection time in seconds] + ``` + +## Examples + +### Queries + +* Run select 42: + ```(bash) + ./fqrun -s "SELECT 42" + ``` + +### Logs + +* Setup log settings: + ```(bash) + ./fqrun -s "SELECT 42" --log-default=warn --log FQ_RUN_ACTOR=trace --log-file query.log + ``` + +### Cluster + +* Embedded UI: + ```(bash) + ./fqrun -M 32000 + ``` + + Monitoring endpoint: http://localhost:32000 + +* gRPC endpoint: + ```(bash) + ./fqrun -G 32000 + ``` + + Connect with ydb CLI: `ydb -e grpc://localhost:32000 -d /Root` diff --git a/ydb/tests/tools/fqrun/configuration/as_config.conf b/ydb/tests/tools/fqrun/configuration/as_config.conf new file mode 100644 index 000000000000..c1f7d09e7f6a --- /dev/null +++ b/ydb/tests/tools/fqrun/configuration/as_config.conf @@ -0,0 +1,46 @@ +Executor { + Type: BASIC + Threads: 1 + SpinThreshold: 10 + Name: "System" +} +Executor { + Type: BASIC + Threads: 6 + SpinThreshold: 1 + Name: "User" +} +Executor { + Type: BASIC + Threads: 1 + SpinThreshold: 1 + Name: "Batch" +} +Executor { + Type: IO + Threads: 1 + Name: "IO" +} +Executor { + Type: BASIC + Threads: 2 + SpinThreshold: 10 + Name: "IC" + TimePerMailboxMicroSecs: 100 +} + +Scheduler { + Resolution: 64 + SpinThreshold: 0 + ProgressThreshold: 10000 +} + +SysExecutor: 0 +UserExecutor: 1 +IoExecutor: 3 +BatchExecutor: 2 + +ServiceExecutor { + ServiceName: "Interconnect" + ExecutorId: 4 +} diff --git a/ydb/tests/tools/fqrun/configuration/fq_config.conf b/ydb/tests/tools/fqrun/configuration/fq_config.conf new file mode 100644 index 000000000000..576b0d394bdb --- /dev/null +++ b/ydb/tests/tools/fqrun/configuration/fq_config.conf @@ -0,0 +1,307 @@ +Enabled: true +EnableDynamicNameservice: true +EnableTaskCounters: true + +CheckpointCoordinator { + Enabled: true + CheckpointingPeriodMillis: 30000 + MaxInflight: 1 + + CheckpointGarbageConfig { + Enabled: true + } + + Storage { + TablePrefix: "yq/checkpoints" + ClientTimeoutSec: 70 + OperationTimeoutSec: 60 + CancelAfterSec: 60 + } +} + +Common { + YdbMvpCloudEndpoint: "https://ydbc.ydb.cloud.yandex.net:8789/ydbc/cloud-prod" + MdbGateway: "https://mdb.api.cloud.yandex.net:443" + MdbTransformHost: false + ObjectStorageEndpoint: "https://storage.yandexcloud.net" + IdsPrefix: "kr" + QueryArtifactsCompressionMethod: "zstd_6" + MonitoringEndpoint: "monitoring.api.cloud.yandex.net" + KeepInternalErrors: true + UseNativeProtocolForClickHouse: true + ShowQueryTimeline: true + MaxTasksPerOperation: 400 + MaxTasksPerStage: 50 + PqReconnectPeriod: "30m" + + YdbDriverConfig { + ClientThreadsNum: 6 + NetworkThreadsNum: 6 + } +} + +ControlPlaneProxy { + Enabled: true +} + +ControlPlaneStorage { + Enabled: true + StatsMode: STATS_MODE_PROFILE + DumpRawStatistics: true + TasksBatchSize: 100 + NumTasksProportion: 4 + AnalyticsRetryCounterLimit: 20 + StreamingRetryCounterLimit: 20 + AnalyticsRetryCounterUpdateTime: "1d" + StreamingRetryCounterUpdateTime: "1d" + TaskLeaseTtl: "30s" + DisableCurrentIam: false + + AvailableConnection: "OBJECT_STORAGE" + AvailableConnection: "DATA_STREAMS" + AvailableConnection: "MONITORING" + AvailableConnection: "POSTGRESQL_CLUSTER" + AvailableConnection: "CLICKHOUSE_CLUSTER" + AvailableConnection: "YDB_DATABASE" + AvailableConnection: "GREENPLUM_CLUSTER" + AvailableConnection: "MYSQL_CLUSTER" + + AvailableStreamingConnection: "OBJECT_STORAGE" + AvailableStreamingConnection: "DATA_STREAMS" + AvailableStreamingConnection: "MONITORING" + AvailableStreamingConnection: "POSTGRESQL_CLUSTER" + AvailableStreamingConnection: "CLICKHOUSE_CLUSTER" + AvailableStreamingConnection: "YDB_DATABASE" + AvailableStreamingConnection: "GREENPLUM_CLUSTER" + AvailableStreamingConnection: "MYSQL_CLUSTER" + + AvailableBinding: "OBJECT_STORAGE" + AvailableBinding: "DATA_STREAMS" + + Storage { + TablePrefix: "yq/control_plane" + ClientTimeoutSec: 70 + OperationTimeoutSec: 60 + CancelAfterSec: 60 + } +} + +DbPool { + Enabled: true + + Storage { + TablePrefix: "yq/db_pool" + ClientTimeoutSec: 70 + OperationTimeoutSec: 60 + CancelAfterSec: 60 + } +} + +Gateways { + Enabled: true + + Dq { + DefaultSettings { + Name: "HashShuffleTasksRatio" + Value: "1" + } + DefaultSettings { + Name: "UseFinalizeByKey" + Value: "true" + } + } + + Generic { + MdbGateway: "https://mdb.api.cloud.yandex.net:443" + + Connector { + UseSsl: false + + Endpoint { + host: "localhost" + port: 2130 + } + } + + DefaultSettings { + Name: "DateTimeFormat" + Value: "string" + } + } + + HttpGateway { + BuffersSizePerStream: 5000000 + ConnectionTimeoutSeconds: 15 + LowSpeedBytesLimit: 1024 + LowSpeedTimeSeconds: 20 + MaxInFlightCount: 2000 + MaxSimulatenousDownloadsSize: 2000000000 + RequestTimeoutSeconds: 0 + } + + Pq { + ClusterMapping { + Name: "pq" + Endpoint: "localhost:2135" + Database: "local" + ClusterType: CT_DATA_STREAMS + UseSsl: true + SharedReading: true + ReadGroup: "fqrun" + } + } + + Solomon { + DefaultSettings { + Name: "_EnableReading" + Value: "true" + } + } + + S3 { + AllowConcurrentListings: true + AllowLocalFiles: true + FileSizeLimit: 100000000000 + GeneratorPathsLimit: 50000 + ListingCallbackPerThreadQueueSize: 100 + ListingCallbackThreadCount: 1 + MaxDirectoriesAndFilesPerQuery: 500000 + MaxDiscoveryFilesPerQuery: 1000 + MaxFilesPerQuery: 500000 + MaxInflightListsPerQuery: 100 + MinDesiredDirectoriesOfFilesPerQuery: 1000 + RegexpCacheSize: 100 + + FormatSizeLimit { + Name: "parquet" + FileSizeLimit: 52428800 + } + FormatSizeLimit { + Name: "raw" + FileSizeLimit: 52428800 + } + + DefaultSettings { + Name: "AtomicUploadCommit" + Value: "true" + } + DefaultSettings { + Name: "UseBlocksSource" + Value: "true" + } + } + + YqlCore { + Flags { + Name: "_EnableMatchRecognize" + } + Flags { + Name: "_EnableStreamLookupJoin" + } + } +} + +Health { + Enabled: true +} + +NodesManager { + Enabled: true +} + +PendingFetcher { + Enabled: true +} + +PrivateApi { + Enabled: true + Loopback: true +} + +PrivateProxy { + Enabled: true +} + +QuotasManager { + Enabled: true + QuotaDescriptions { + SubjectType: "cloud" + MetricName: "yq.cpuPercent.count" + HardLimit: 7500 + DefaultLimit: 3500 + } + QuotaDescriptions { + SubjectType: "cloud" + MetricName: "yq.streamingQueryDurationMinutes.count" + DefaultLimit: 10080 + } + QuotaDescriptions { + SubjectType: "cloud" + MetricName: "yq.analyticsQueryDurationMinutes.count" + HardLimit: 1440 + DefaultLimit: 30 + } +} + +RateLimiter { + Enabled: true + ControlPlaneEnabled: true + DataPlaneEnabled: true + + Database { + TablePrefix: "yq/rate_limiter" + ClientTimeoutSec: 70 + OperationTimeoutSec: 60 + CancelAfterSec: 60 + } + + Limiters { + CoordinationNodePath: "limiter_alpha" + } +} + +ReadActorsFactoryConfig { + PqReadActorFactoryConfig { + CookieCommitMode: false + } +} + +ResourceManager { + Enabled: true + MkqlInitialMemoryLimit: 16777216 + MkqlTotalMemoryLimit: 193273528320 + MkqlAllocSize: 16777216 + MkqlTaskHardMemoryLimit: 24696061952 +} + +RowDispatcher { + Enabled: true + SendStatusPeriodSec: 10 + TimeoutBeforeStartSessionSec: 0 + MaxSessionUsedMemory: 16000000 + WithoutConsumer: false + + CompileService { + ParallelCompilationLimit: 20 + } + + Coordinator { + CoordinationNodePath: "yq/row_dispatcher" + + Database { + TablePrefix: "yq/row_dispatcher" + ClientTimeoutSec: 70 + OperationTimeoutSec: 60 + CancelAfterSec: 60 + } + } + + JsonParser { + BatchSizeBytes: 1048576 + BatchCreationTimeoutMs: 1000 + } +} + +TestConnection { + Enabled: true +} diff --git a/ydb/tests/tools/fqrun/flame_graph.sh b/ydb/tests/tools/fqrun/flame_graph.sh new file mode 100755 index 000000000000..b6f8fe6da4f0 --- /dev/null +++ b/ydb/tests/tools/fqrun/flame_graph.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash + +set -eux + +function cleanup { + sudo rm ./profdata + rm ./profdata.txt +} +trap cleanup EXIT + +if [ $# -gt 1 ]; then + echo "Too many arguments" + exit -1 +fi + +fqrun_pid=$(pgrep -u $USER fqrun) + +sudo perf record -F 50 --call-graph dwarf -g --proc-map-timeout=10000 --pid $fqrun_pid -v -o profdata -- sleep ${1:-'30'} +sudo perf script -i profdata > profdata.txt + +SCRIPT_DIR=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) + +flame_graph_tool="$SCRIPT_DIR/../../../../contrib/tools/flame-graph/" + +${flame_graph_tool}/stackcollapse-perf.pl profdata.txt | ${flame_graph_tool}/flamegraph.pl > profdata.svg diff --git a/ydb/tests/tools/fqrun/fqrun.cpp b/ydb/tests/tools/fqrun/fqrun.cpp new file mode 100644 index 000000000000..99f597d95c58 --- /dev/null +++ b/ydb/tests/tools/fqrun/fqrun.cpp @@ -0,0 +1,330 @@ +#include +#include + +#include + +#include +#include +#include +#include +#include + +#ifdef PROFILE_MEMORY_ALLOCATIONS +#include +#endif + +using namespace NKikimrRun; + +namespace NFqRun { + +namespace { + +struct TExecutionOptions { + TString Query; + std::vector Connections; + std::vector Bindings; + + bool HasResults() const { + return !Query.empty(); + } + + TRequestOptions GetQueryOptions() const { + return { + .Query = Query + }; + } + + void Validate(const TRunnerOptions& runnerOptions) const { + if (!Query && !runnerOptions.FqSettings.MonitoringEnabled && !runnerOptions.FqSettings.GrpcEnabled) { + ythrow yexception() << "Nothing to execute and is not running as daemon"; + } + } +}; + +void RunArgumentQueries(const TExecutionOptions& executionOptions, TFqRunner& runner) { + NColorizer::TColors colors = NColorizer::AutoColors(Cout); + + if (!executionOptions.Connections.empty()) { + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Creating connections..." << colors.Default() << Endl; + if (!runner.CreateConnections(executionOptions.Connections)) { + ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Failed to create connections"; + } + } + + if (!executionOptions.Bindings.empty()) { + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Creating bindings..." << colors.Default() << Endl; + if (!runner.CreateBindings(executionOptions.Bindings)) { + ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Failed to create bindings"; + } + } + + if (executionOptions.Query) { + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Executing query..." << colors.Default() << Endl; + if (!runner.ExecuteStreamQuery(executionOptions.GetQueryOptions())) { + ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Query execution failed"; + } + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Fetching query results..." << colors.Default() << Endl; + if (!runner.FetchQueryResults()) { + ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Fetch query results failed"; + } + } + + if (executionOptions.HasResults()) { + try { + runner.PrintQueryResults(); + } catch (...) { + ythrow yexception() << "Failed to print script results, reason:\n" << CurrentExceptionMessage(); + } + } +} + +void RunAsDaemon() { + NColorizer::TColors colors = NColorizer::AutoColors(Cout); + + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Initialization finished" << colors.Default() << Endl; + while (true) { + Sleep(TDuration::Seconds(1)); + } +} + +void RunScript(const TExecutionOptions& executionOptions, const TRunnerOptions& runnerOptions) { + NColorizer::TColors colors = NColorizer::AutoColors(Cout); + + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Initialization of fq runner..." << colors.Default() << Endl; + TFqRunner runner(runnerOptions); + + try { + RunArgumentQueries(executionOptions, runner); + } catch (const yexception& exception) { + if (runnerOptions.FqSettings.MonitoringEnabled) { + Cerr << colors.Red() << CurrentExceptionMessage() << colors.Default() << Endl; + } else { + throw exception; + } + } + + if (runnerOptions.FqSettings.MonitoringEnabled || runnerOptions.FqSettings.GrpcEnabled) { + RunAsDaemon(); + } + + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Finalization of fq runner..." << colors.Default() << Endl; +} + +class TMain : public TMainBase { + using EVerbose = TFqSetupSettings::EVerbose; + +protected: + void RegisterOptions(NLastGetopt::TOpts& options) override { + options.SetTitle("FqRun -- tool to execute stream queries through FQ proxy"); + options.AddHelpOption('h'); + options.SetFreeArgsNum(0); + + // Inputs + + options.AddLongOption('p', "query", "Query to execute") + .RequiredArgument("file") + .StoreMappedResult(&ExecutionOptions.Query, &LoadFile); + + options.AddLongOption('s', "sql", "Query SQL text to execute") + .RequiredArgument("str") + .StoreResult(&ExecutionOptions.Query); + options.MutuallyExclusive("query", "sql"); + + options.AddLongOption('c', "connection", "External datasource connection protobuf FederatedQuery::ConnectionContent") + .RequiredArgument("file") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + auto& connection = ExecutionOptions.Connections.emplace_back(); + const TString file(TString(option->CurValOrDef())); + if (!google::protobuf::TextFormat::ParseFromString(LoadFile(file), &connection)) { + ythrow yexception() << "Bad format of FQ connection in file '" << file << "'"; + } + SetupAcl(connection.mutable_acl()); + }); + + options.AddLongOption('b', "binding", "External datasource binding protobuf FederatedQuery::BindingContent") + .RequiredArgument("file") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + auto& binding = ExecutionOptions.Bindings.emplace_back(); + const TString file(TString(option->CurValOrDef())); + if (!google::protobuf::TextFormat::ParseFromString(LoadFile(file), &binding)) { + ythrow yexception() << "Bad format of FQ binding in file '" << file << "'"; + } + SetupAcl(binding.mutable_acl()); + }); + + options.AddLongOption("fq-cfg", "File with FQ config (NFq::NConfig::TConfig for FQ proxy)") + .RequiredArgument("file") + .DefaultValue("./configuration/fq_config.conf") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + if (!google::protobuf::TextFormat::ParseFromString(LoadFile(TString(option->CurValOrDef())), &RunnerOptions.FqSettings.FqConfig)) { + ythrow yexception() << "Bad format of FQ configuration"; + } + }); + + options.AddLongOption("as-cfg", "File with actor system config (TActorSystemConfig), use '-' for default") + .RequiredArgument("file") + .DefaultValue("./configuration/as_config.conf") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + const TString file(option->CurValOrDef()); + if (file == "-") { + return; + } + + RunnerOptions.FqSettings.ActorSystemConfig = NKikimrConfig::TActorSystemConfig(); + if (!google::protobuf::TextFormat::ParseFromString(LoadFile(file), &(*RunnerOptions.FqSettings.ActorSystemConfig))) { + ythrow yexception() << "Bad format of actor system configuration"; + } + }); + + options.AddLongOption("emulate-s3", "Enable readings by s3 provider from files, `bucket` value in connection - path to folder with files") + .NoArgument() + .SetFlag(&RunnerOptions.FqSettings.EmulateS3); + + options.AddLongOption("emulate-pq", "Emulate YDS with local file, accepts list of tables to emulate with following format: topic@file (can be used in query from cluster `pq`)") + .RequiredArgument("topic@file") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + TStringBuf topicName, others; + TStringBuf(option->CurVal()).Split('@', topicName, others); + + TStringBuf path, partitionCountStr; + TStringBuf(others).Split(':', path, partitionCountStr); + size_t partitionCount = !partitionCountStr.empty() ? FromString(partitionCountStr) : 1; + if (!partitionCount) { + ythrow yexception() << "Topic partition count should be at least one"; + } + if (topicName.empty() || path.empty()) { + ythrow yexception() << "Incorrect PQ file mapping, expected form topic@path[:partitions_count]"; + } + if (!PqFilesMapping.emplace(topicName, NYql::TDummyTopic("pq", TString(topicName), TString(path), partitionCount)).second) { + ythrow yexception() << "Got duplicated topic name: " << topicName; + } + }); + + options.AddLongOption("cancel-on-file-finish", "Cancel emulate YDS topics when topic file finished") + .RequiredArgument("topic") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + TopicsSettings[option->CurVal()].CancelOnFileFinish = true; + }); + + // Outputs + + options.AddLongOption("result-file", "File with query results (use '-' to write in stdout)") + .RequiredArgument("file") + .DefaultValue("-") + .StoreMappedResultT(&RunnerOptions.ResultOutput, &GetDefaultOutput); + + TChoices resultFormat({ + {"rows", EResultOutputFormat::RowsJson}, + {"full-json", EResultOutputFormat::FullJson}, + {"full-proto", EResultOutputFormat::FullProto} + }); + options.AddLongOption('R', "result-format", "Query result format") + .RequiredArgument("result-format") + .DefaultValue("rows") + .Choices(resultFormat.GetChoices()) + .StoreMappedResultT(&RunnerOptions.ResultOutputFormat, resultFormat); + + // Pipeline settings + + options.AddLongOption("verbose", TStringBuilder() << "Common verbose level (max level " << static_cast(EVerbose::Max) - 1 << ")") + .RequiredArgument("uint") + .DefaultValue(static_cast(EVerbose::Info)) + .StoreMappedResultT(&RunnerOptions.FqSettings.VerboseLevel, [](ui8 value) { + return static_cast(std::min(value, static_cast(EVerbose::Max))); + }); + + RegisterKikimrOptions(options, RunnerOptions.FqSettings); + } + + int DoRun(NLastGetopt::TOptsParseResult&&) override { + ExecutionOptions.Validate(RunnerOptions); + + RunnerOptions.FqSettings.YqlToken = GetEnv(YQL_TOKEN_VARIABLE); + RunnerOptions.FqSettings.FunctionRegistry = CreateFunctionRegistry().Get(); + + auto& gatewayConfig = *RunnerOptions.FqSettings.FqConfig.mutable_gateways(); + FillTokens(gatewayConfig.mutable_pq()); + FillTokens(gatewayConfig.mutable_s3()); + FillTokens(gatewayConfig.mutable_generic()); + FillTokens(gatewayConfig.mutable_ydb()); + FillTokens(gatewayConfig.mutable_solomon()); + + auto& logConfig = RunnerOptions.FqSettings.LogConfig; + logConfig.SetDefaultLevel(NActors::NLog::EPriority::PRI_CRIT); + FillLogConfig(logConfig); + + if (!PqFilesMapping.empty()) { + auto fileGateway = MakeIntrusive(); + for (auto [_, topic] : PqFilesMapping) { + if (const auto it = TopicsSettings.find(topic.TopicName); it != TopicsSettings.end()) { + topic.CancelOnFileFinish = it->second.CancelOnFileFinish; + TopicsSettings.erase(it); + } + fileGateway->AddDummyTopic(topic); + } + RunnerOptions.FqSettings.PqGatewayFactory = CreatePqFileGatewayFactory(fileGateway); + } + if (!TopicsSettings.empty()) { + ythrow yexception() << "Found topic settings for not existing topic: '" << TopicsSettings.begin()->first << "'"; + } + +#ifdef PROFILE_MEMORY_ALLOCATIONS + if (RunnerOptions.FqSettings.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors.Cyan() << "Starting profile memory allocations" << CoutColors.Default() << Endl; + } + NAllocProfiler::StartAllocationSampling(true); +#else + if (ProfileAllocationsOutput) { + ythrow yexception() << "Profile memory allocations disabled, please rebuild fqrun with flag `-D PROFILE_MEMORY_ALLOCATIONS`"; + } +#endif + + RunScript(ExecutionOptions, RunnerOptions); + +#ifdef PROFILE_MEMORY_ALLOCATIONS + if (RunnerOptions.FqSettings.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors.Cyan() << "Finishing profile memory allocations" << CoutColors.Default() << Endl; + } + FinishProfileMemoryAllocations(); +#endif + + return 0; + } + +private: + template + void FillTokens(TGatewayConfig* gateway) const { + for (auto& cluster : *gateway->mutable_clustermapping()) { + if (!cluster.GetToken()) { + cluster.SetToken(RunnerOptions.FqSettings.YqlToken); + } + } + } + +private: + TExecutionOptions ExecutionOptions; + TRunnerOptions RunnerOptions; + std::unordered_map PqFilesMapping; + + struct TTopicSettings { + bool CancelOnFileFinish = false; + }; + std::unordered_map TopicsSettings; +}; + +} // anonymous namespace + +} // namespace NFqRun + +int main(int argc, const char* argv[]) { + SetupSignalActions(); + + try { + NFqRun::TMain().Run(argc, argv); + } catch (...) { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cerr << colors.Red() << CurrentExceptionMessage() << colors.Default() << Endl; + return 1; + } +} diff --git a/ydb/tests/tools/fqrun/src/common.cpp b/ydb/tests/tools/fqrun/src/common.cpp new file mode 100644 index 000000000000..7d0fc1804532 --- /dev/null +++ b/ydb/tests/tools/fqrun/src/common.cpp @@ -0,0 +1,11 @@ +#include "common.h" + +namespace NFqRun { + +void SetupAcl(FederatedQuery::Acl* acl) { + if (acl->visibility() == FederatedQuery::Acl::VISIBILITY_UNSPECIFIED) { + acl->set_visibility(FederatedQuery::Acl::SCOPE); + } +} + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/common.h b/ydb/tests/tools/fqrun/src/common.h new file mode 100644 index 000000000000..5c139d83f5cc --- /dev/null +++ b/ydb/tests/tools/fqrun/src/common.h @@ -0,0 +1,51 @@ +#pragma once + +#include + +#include +#include +#include +#include + +#include + +namespace NFqRun { + +constexpr char YQL_TOKEN_VARIABLE[] = "YQL_TOKEN"; +constexpr i64 MAX_RESULT_SET_ROWS = 1000; + +struct TFqSetupSettings : public NKikimrRun::TServerSettings { + enum class EVerbose { + None, + Info, + QueriesText, + InitLogs, + Max + }; + + bool EmulateS3 = false; + + EVerbose VerboseLevel = EVerbose::Info; + + TString YqlToken; + NYql::IPqGatewayFactory::TPtr PqGatewayFactory; + TIntrusivePtr FunctionRegistry; + NFq::NConfig::TConfig FqConfig; + NKikimrConfig::TLogConfig LogConfig; + std::optional ActorSystemConfig; +}; + +struct TRunnerOptions { + IOutputStream* ResultOutput = nullptr; + NKikimrRun::EResultOutputFormat ResultOutputFormat = NKikimrRun::EResultOutputFormat::RowsJson; + + TFqSetupSettings FqSettings; +}; + +struct TRequestOptions { + TString Query; +}; + +void SetupAcl(FederatedQuery::Acl* acl); + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/fq_runner.cpp b/ydb/tests/tools/fqrun/src/fq_runner.cpp new file mode 100644 index 000000000000..33995c17fac8 --- /dev/null +++ b/ydb/tests/tools/fqrun/src/fq_runner.cpp @@ -0,0 +1,201 @@ +#include "fq_runner.h" +#include "fq_setup.h" + +#include + +using namespace NKikimrRun; + +namespace NFqRun { + +class TFqRunner::TImpl { + using EVerbose = TFqSetupSettings::EVerbose; + + static constexpr TDuration REFRESH_PERIOD = TDuration::Seconds(1); + +public: + explicit TImpl(const TRunnerOptions& options) + : Options(options) + , VerboseLevel(options.FqSettings.VerboseLevel) + , FqSetup(options.FqSettings) + , CerrColors(NColorizer::AutoColors(Cerr)) + , CoutColors(NColorizer::AutoColors(Cout)) + {} + + bool ExecuteStreamQuery(const TRequestOptions& query) { + if (VerboseLevel >= EVerbose::QueriesText) { + Cout << CoutColors.Cyan() << "Starting stream request:\n" << CoutColors.Default() << query.Query << Endl; + } + + const TRequestResult status = FqSetup.StreamRequest(query, StreamQueryId); + + if (!status.IsSuccess()) { + Cerr << CerrColors.Red() << "Failed to start stream request execution, reason:" << CerrColors.Default() << Endl << status.ToString() << Endl; + return false; + } + + return WaitStreamQuery(); + } + + bool FetchQueryResults() { + ResultSets.clear(); + ResultSets.resize(ExecutionMeta.ResultSetSizes.size()); + for (i32 resultSetId = 0; resultSetId < static_cast(ExecutionMeta.ResultSetSizes.size()); ++resultSetId) { + const auto rowsCount = ExecutionMeta.ResultSetSizes[resultSetId]; + if (rowsCount > MAX_RESULT_SET_ROWS) { + Cerr << CerrColors.Red() << "Result set with id " << resultSetId << " have " << rowsCount << " rows, it is larger than allowed limit " << MAX_RESULT_SET_ROWS << ", results will be truncated" << CerrColors.Default() << Endl; + } + + const TRequestResult status = FqSetup.FetchQueryResults(StreamQueryId, resultSetId, ResultSets[resultSetId]); + if (!status.IsSuccess()) { + Cerr << CerrColors.Red() << "Failed to fetch result set with id " << resultSetId << ", reason:" << CerrColors.Default() << Endl << status.ToString() << Endl; + return false; + } + } + + return true; + } + + void PrintQueryResults() { + if (Options.ResultOutput) { + Cout << CoutColors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Writing query results..." << CoutColors.Default() << Endl; + for (size_t i = 0; i < ResultSets.size(); ++i) { + if (ResultSets.size() > 1 && VerboseLevel >= EVerbose::Info) { + *Options.ResultOutput << CoutColors.Cyan() << "Result set " << i + 1 << ":" << CoutColors.Default() << Endl; + } + PrintResultSet(Options.ResultOutputFormat, *Options.ResultOutput, ResultSets[i]); + } + } + } + + bool CreateConnections(const std::vector& connections) { + for (const auto& connection : connections) { + if (VerboseLevel >= EVerbose::QueriesText) { + Cout << CoutColors.Cyan() << "Creating connection:\n" << CoutColors.Default() << Endl << connection.DebugString() << Endl; + } + + TString connectionId; + const TRequestResult status = FqSetup.CreateConnection(connection, connectionId); + + if (!status.IsSuccess()) { + Cerr << CerrColors.Red() << "Failed to create connection '" << connection.name() << "', reason:" << CerrColors.Default() << Endl << status.ToString() << Endl; + return false; + } + + if (!ConnectionNameToId.emplace(connection.name(), connectionId).second) { + Cerr << CerrColors.Red() << "Got duplicated connection name '" << connection.name() << "'" << CerrColors.Default() << Endl; + return false; + } + } + + return true; + } + + bool CreateBindings(const std::vector& bindings) const { + for (auto binding : bindings) { + if (VerboseLevel >= EVerbose::QueriesText) { + Cout << CoutColors.Cyan() << "Creating binding:\n" << CoutColors.Default() << Endl << binding.DebugString() << Endl; + } + + const auto it = ConnectionNameToId.find(binding.connection_id()); + if (it == ConnectionNameToId.end()) { + Cerr << CerrColors.Red() << "Failed to create binding '" << binding.name() << "', connection with name '" << binding.connection_id() << "' not found" << CerrColors.Default() << Endl; + return false; + } + + binding.set_connection_id(it->second); + const TRequestResult status = FqSetup.CreateBinding(binding); + + if (!status.IsSuccess()) { + Cerr << CerrColors.Red() << "Failed to create binding '" << binding.name() << "', reason:" << CerrColors.Default() << Endl << status.ToString() << Endl; + return false; + } + } + + return true; + } + +private: + static bool IsFinalStatus(FederatedQuery::QueryMeta::ComputeStatus status) { + using EStatus = FederatedQuery::QueryMeta; + return IsIn({EStatus::FAILED, EStatus::COMPLETED, EStatus::ABORTED_BY_USER, EStatus::ABORTED_BY_SYSTEM}, status); + } + + bool WaitStreamQuery() { + StartTime = TInstant::Now(); + + while (true) { + TExecutionMeta meta; + const TRequestResult status = FqSetup.DescribeQuery(StreamQueryId, meta); + + if (meta.TransientIssues.Size() != ExecutionMeta.TransientIssues.Size() && VerboseLevel >= EVerbose::Info) { + Cerr << CerrColors.Red() << "Query transient issues updated:" << CerrColors.Default() << Endl << meta.TransientIssues.ToString() << Endl; + } + ExecutionMeta = meta; + + if (IsFinalStatus(ExecutionMeta.Status)) { + break; + } + + if (!status.IsSuccess()) { + Cerr << CerrColors.Red() << "Failed to describe query, reason:" << CerrColors.Default() << Endl << status.ToString() << Endl; + return false; + } + + Sleep(REFRESH_PERIOD); + } + + if (VerboseLevel >= EVerbose::Info) { + Cout << CoutColors.Cyan() << "Query finished. Duration: " << TInstant::Now() - StartTime << CoutColors.Default() << Endl; + } + + if (ExecutionMeta.Status != FederatedQuery::QueryMeta::COMPLETED) { + Cerr << CerrColors.Red() << "Failed to execute query, invalid final status " << FederatedQuery::QueryMeta::ComputeStatus_Name(ExecutionMeta.Status) << ", issues:" << CerrColors.Default() << Endl << ExecutionMeta.Issues.ToString() << Endl; + return false; + } + + if (ExecutionMeta.Issues) { + Cerr << CerrColors.Red() << "Query finished with issues:" << CerrColors.Default() << Endl << ExecutionMeta.Issues.ToString() << Endl; + } + + return true; + } + +private: + const TRunnerOptions Options; + const EVerbose VerboseLevel; + const TFqSetup FqSetup; + const NColorizer::TColors CerrColors; + const NColorizer::TColors CoutColors; + + TString StreamQueryId; + TInstant StartTime; + TExecutionMeta ExecutionMeta; + std::vector ResultSets; + std::unordered_map ConnectionNameToId; +}; + +TFqRunner::TFqRunner(const TRunnerOptions& options) + : Impl(new TImpl(options)) +{} + +bool TFqRunner::ExecuteStreamQuery(const TRequestOptions& query) const { + return Impl->ExecuteStreamQuery(query); +} + +bool TFqRunner::FetchQueryResults() const { + return Impl->FetchQueryResults(); +} + +void TFqRunner::PrintQueryResults() const { + Impl->PrintQueryResults(); +} + +bool TFqRunner::CreateConnections(const std::vector& connections) const { + return Impl->CreateConnections(connections); +} + +bool TFqRunner::CreateBindings(const std::vector& bindings) const { + return Impl->CreateBindings(bindings); +} + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/fq_runner.h b/ydb/tests/tools/fqrun/src/fq_runner.h new file mode 100644 index 000000000000..b6e85db142bc --- /dev/null +++ b/ydb/tests/tools/fqrun/src/fq_runner.h @@ -0,0 +1,28 @@ +#pragma once + +#include "common.h" + +#include + +namespace NFqRun { + +class TFqRunner { +public: + explicit TFqRunner(const TRunnerOptions& options); + + bool ExecuteStreamQuery(const TRequestOptions& query) const; + + bool FetchQueryResults() const; + + void PrintQueryResults() const; + + bool CreateConnections(const std::vector& connections) const; + + bool CreateBindings(const std::vector& bindings) const; + +private: + class TImpl; + std::shared_ptr Impl; +}; + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/fq_setup.cpp b/ydb/tests/tools/fqrun/src/fq_setup.cpp new file mode 100644 index 000000000000..e870d3cff211 --- /dev/null +++ b/ydb/tests/tools/fqrun/src/fq_setup.cpp @@ -0,0 +1,305 @@ +#include "fq_setup.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +using namespace NKikimrRun; + +namespace NFqRun { + +namespace { + +TRequestResult GetStatus(const NYql::TIssues& issues) { + return TRequestResult(issues ? Ydb::StatusIds::BAD_REQUEST : Ydb::StatusIds::SUCCESS, issues); +} + +} // anonymous namespace + +class TFqSetup::TImpl { + using EVerbose = TFqSetupSettings::EVerbose; + +private: + TAutoPtr CreateLogBackend() const { + if (Settings.LogOutputFile) { + return NActors::CreateFileBackend(Settings.LogOutputFile); + } else { + return NActors::CreateStderrBackend(); + } + } + + void SetLoggerSettings(NKikimr::Tests::TServerSettings& serverSettings) const { + auto loggerInitializer = [this](NActors::TTestActorRuntime& runtime) { + InitLogSettings(Settings.LogConfig, runtime); + runtime.SetLogBackendFactory([this]() { return CreateLogBackend(); }); + }; + + serverSettings.SetLoggerInitializer(loggerInitializer); + } + + void SetFunctionRegistry(NKikimr::Tests::TServerSettings& serverSettings) const { + if (Settings.FunctionRegistry) { + serverSettings.SetFrFactory([this](const NKikimr::NScheme::TTypeRegistry&) { + return Settings.FunctionRegistry.Get(); + }); + } + } + + NKikimr::Tests::TServerSettings GetServerSettings(ui32 grpcPort) { + NKikimr::Tests::TServerSettings serverSettings(PortManager.GetPort()); + + serverSettings.SetDomainName(Settings.DomainName); + serverSettings.SetVerbose(Settings.VerboseLevel >= EVerbose::InitLogs); + + NKikimrConfig::TAppConfig config; + *config.MutableLogConfig() = Settings.LogConfig; + if (Settings.ActorSystemConfig) { + *config.MutableActorSystemConfig() = *Settings.ActorSystemConfig; + } + serverSettings.SetAppConfig(config); + + SetLoggerSettings(serverSettings); + SetFunctionRegistry(serverSettings); + + if (Settings.MonitoringEnabled) { + serverSettings.InitKikimrRunConfig(); + serverSettings.SetMonitoringPortOffset(Settings.MonitoringPortOffset, true); + serverSettings.SetNeedStatsCollectors(true); + } + + serverSettings.SetGrpcPort(grpcPort); + serverSettings.SetEnableYqGrpc(true); + + return serverSettings; + } + + void InitializeServer(ui32 grpcPort) { + const auto& serverSettings = GetServerSettings(grpcPort); + + Server = MakeIntrusive(serverSettings); + Server->GetRuntime()->SetDispatchTimeout(TDuration::Max()); + + Server->EnableGRpc(NYdbGrpc::TServerOptions() + .SetHost("localhost") + .SetPort(grpcPort) + .SetLogger(NYdbGrpc::CreateActorSystemLogger(*GetRuntime()->GetActorSystem(0), NKikimrServices::GRPC_SERVER)) + .SetGRpcShutdownDeadline(TDuration::Zero()) + ); + + Client = std::make_unique(serverSettings); + Client->InitRootScheme(); + } + + NFq::NConfig::TConfig GetFqProxyConfig(ui32 grpcPort) const { + auto fqConfig = Settings.FqConfig; + + fqConfig.MutableControlPlaneStorage()->AddSuperUsers(BUILTIN_ACL_ROOT); + fqConfig.MutablePrivateProxy()->AddGrantedUsers(BUILTIN_ACL_ROOT); + + const TString endpoint = TStringBuilder() << "localhost:" << grpcPort; + const TString database = NKikimr::CanonizePath(Settings.DomainName); + const auto fillStorageConfig = [endpoint, database](NFq::NConfig::TYdbStorageConfig* config) { + config->SetEndpoint(endpoint); + config->SetDatabase(database); + }; + fillStorageConfig(fqConfig.MutableControlPlaneStorage()->MutableStorage()); + fillStorageConfig(fqConfig.MutableDbPool()->MutableStorage()); + fillStorageConfig(fqConfig.MutableCheckpointCoordinator()->MutableStorage()); + fillStorageConfig(fqConfig.MutableRateLimiter()->MutableDatabase()); + fillStorageConfig(fqConfig.MutableRowDispatcher()->MutableCoordinator()->MutableDatabase()); + + auto* privateApiConfig = fqConfig.MutablePrivateApi(); + privateApiConfig->SetTaskServiceEndpoint(endpoint); + privateApiConfig->SetTaskServiceDatabase(database); + + auto* nodesMenagerConfig = fqConfig.MutableNodesManager(); + nodesMenagerConfig->SetPort(grpcPort); + nodesMenagerConfig->SetHost("localhost"); + + auto* healthConfig = fqConfig.MutableHealth(); + healthConfig->SetPort(grpcPort); + healthConfig->SetDatabase(database); + + if (Settings.EmulateS3) { + fqConfig.MutableCommon()->SetObjectStorageEndpoint("file://"); + } + + return fqConfig; + } + + void InitializeFqProxy(ui32 grpcPort) { + const auto& fqConfig = GetFqProxyConfig(grpcPort); + const auto counters = GetRuntime()->GetAppData().Counters->GetSubgroup("counters", "yq"); + YqSharedResources = NFq::CreateYqSharedResources(fqConfig, NKikimr::CreateYdbCredentialsProviderFactory, counters); + + const auto actorRegistrator = [runtime = GetRuntime()](NActors::TActorId serviceActorId, NActors::IActor* actor) { + auto actorId = runtime->Register(actor, 0, runtime->GetAppData().UserPoolId); + runtime->RegisterService(serviceActorId, actorId); + }; + + const auto folderServiceFactory = [](auto& config) { + return NKikimr::NFolderService::CreateMockFolderServiceAdapterActor(config, ""); + }; + + NFq::Init( + fqConfig, GetRuntime()->GetNodeId(), actorRegistrator, &GetRuntime()->GetAppData(), + Settings.DomainName, nullptr, YqSharedResources, folderServiceFactory, 0, {}, Settings.PqGatewayFactory + ); + YqSharedResources->Init(GetRuntime()->GetActorSystem(0)); + } + +public: + explicit TImpl(const TFqSetupSettings& settings) + : Settings(settings) + { + const ui32 grpcPort = Settings.GrpcPort ? Settings.GrpcPort : PortManager.GetPort(); + InitializeServer(grpcPort); + InitializeFqProxy(grpcPort); + + if (Settings.MonitoringEnabled && Settings.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors.Cyan() << "Monitoring port: " << CoutColors.Default() << GetRuntime()->GetMonPort() << Endl; + } + + if (Settings.GrpcEnabled && Settings.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors.Cyan() << "Domain gRPC port: " << CoutColors.Default() << grpcPort << Endl; + } + } + + ~TImpl() { + if (YqSharedResources) { + YqSharedResources->Stop(); + } + } + + NFq::TEvControlPlaneProxy::TEvCreateQueryResponse::TPtr StreamRequest(const TRequestOptions& query) const { + FederatedQuery::CreateQueryRequest request; + request.set_execute_mode(FederatedQuery::ExecuteMode::RUN); + + auto& content = *request.mutable_content(); + content.set_type(FederatedQuery::QueryContent::STREAMING); + content.set_text(query.Query); + SetupAcl(content.mutable_acl()); + + return RunControlPlaneProxyRequest(request); + } + + NFq::TEvControlPlaneProxy::TEvDescribeQueryResponse::TPtr DescribeQuery(const TString& queryId) const { + FederatedQuery::DescribeQueryRequest request; + request.set_query_id(queryId); + + return RunControlPlaneProxyRequest(request); + } + + NFq::TEvControlPlaneProxy::TEvGetResultDataResponse::TPtr FetchQueryResults(const TString& queryId, i32 resultSetId) const { + FederatedQuery::GetResultDataRequest request; + request.set_query_id(queryId); + request.set_result_set_index(resultSetId); + request.set_limit(MAX_RESULT_SET_ROWS); + + return RunControlPlaneProxyRequest(request); + } + + NFq::TEvControlPlaneProxy::TEvCreateConnectionResponse::TPtr CreateConnection(const FederatedQuery::ConnectionContent& connection) const { + FederatedQuery::CreateConnectionRequest request; + *request.mutable_content() = connection; + + return RunControlPlaneProxyRequest(request); + } + + NFq::TEvControlPlaneProxy::TEvCreateBindingResponse::TPtr CreateBinding(const FederatedQuery::BindingContent& binding) const { + FederatedQuery::CreateBindingRequest request; + *request.mutable_content() = binding; + + return RunControlPlaneProxyRequest(request); + } + +private: + NActors::TTestActorRuntime* GetRuntime() const { + return Server->GetRuntime(); + } + + template + typename TResponse::TPtr RunControlPlaneProxyRequest(const TProto& request) const { + auto event = std::make_unique("yandexcloud://fqrun", request, BUILTIN_ACL_ROOT, Settings.YqlToken ? Settings.YqlToken : "fqrun", TVector{}); + return RunControlPlaneProxyRequest(std::move(event)); + } + + template + typename TResponse::TPtr RunControlPlaneProxyRequest(std::unique_ptr event) const { + NActors::TActorId edgeActor = GetRuntime()->AllocateEdgeActor(); + NActors::TActorId controlPlaneProxy = NFq::ControlPlaneProxyActorId(); + + GetRuntime()->Send(controlPlaneProxy, edgeActor, event.release()); + + return GetRuntime()->GrabEdgeEvent(edgeActor); + } + +private: + const TFqSetupSettings Settings; + const NColorizer::TColors CoutColors; + + NKikimr::Tests::TServer::TPtr Server; + std::unique_ptr Client; + NFq::IYqSharedResources::TPtr YqSharedResources; + TPortManager PortManager; +}; + +TFqSetup::TFqSetup(const TFqSetupSettings& settings) + : Impl(new TImpl(settings)) +{} + +TRequestResult TFqSetup::StreamRequest(const TRequestOptions& query, TString& queryId) const { + const auto response = Impl->StreamRequest(query); + + queryId = response->Get()->Result.query_id(); + + return GetStatus(response->Get()->Issues); +} + +TRequestResult TFqSetup::DescribeQuery(const TString& queryId, TExecutionMeta& meta) const { + const auto response = Impl->DescribeQuery(queryId); + + const auto& result = response->Get()->Result.query(); + meta.Status = result.meta().status(); + NYql::IssuesFromMessage(result.issue(), meta.Issues); + NYql::IssuesFromMessage(result.transient_issue(), meta.TransientIssues); + + meta.ResultSetSizes.clear(); + for (const auto& resultMeta : result.result_set_meta()) { + meta.ResultSetSizes.emplace_back(resultMeta.rows_count()); + } + + return GetStatus(response->Get()->Issues); +} + +TRequestResult TFqSetup::FetchQueryResults(const TString& queryId, i32 resultSetId, Ydb::ResultSet& resultSet) const { + const auto response = Impl->FetchQueryResults(queryId, resultSetId); + + resultSet = response->Get()->Result.result_set(); + + return GetStatus(response->Get()->Issues); +} + +TRequestResult TFqSetup::CreateConnection(const FederatedQuery::ConnectionContent& connection, TString& connectionId) const { + const auto response = Impl->CreateConnection(connection); + + connectionId = response->Get()->Result.connection_id(); + + return GetStatus(response->Get()->Issues); +} + +TRequestResult TFqSetup::CreateBinding(const FederatedQuery::BindingContent& binding) const { + const auto response = Impl->CreateBinding(binding); + return GetStatus(response->Get()->Issues); +} + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/fq_setup.h b/ydb/tests/tools/fqrun/src/fq_setup.h new file mode 100644 index 000000000000..2ef7733bf028 --- /dev/null +++ b/ydb/tests/tools/fqrun/src/fq_setup.h @@ -0,0 +1,37 @@ +#pragma once + +#include "common.h" + +#include + +namespace NFqRun { + +struct TExecutionMeta { + FederatedQuery::QueryMeta::ComputeStatus Status; + NYql::TIssues Issues; + NYql::TIssues TransientIssues; + std::vector ResultSetSizes; +}; + +class TFqSetup { + using TRequestResult = NKikimrRun::TRequestResult; + +public: + explicit TFqSetup(const TFqSetupSettings& settings); + + TRequestResult StreamRequest(const TRequestOptions& query, TString& queryId) const; + + TRequestResult DescribeQuery(const TString& queryId, TExecutionMeta& meta) const; + + TRequestResult FetchQueryResults(const TString& queryId, i32 resultSetId, Ydb::ResultSet& resultSet) const; + + TRequestResult CreateConnection(const FederatedQuery::ConnectionContent& connection, TString& connectionId) const; + + TRequestResult CreateBinding(const FederatedQuery::BindingContent& binding) const; + +private: + class TImpl; + std::shared_ptr Impl; +}; + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/ya.make b/ydb/tests/tools/fqrun/src/ya.make new file mode 100644 index 000000000000..a5a54edf3a87 --- /dev/null +++ b/ydb/tests/tools/fqrun/src/ya.make @@ -0,0 +1,27 @@ +LIBRARY() + +SRCS( + common.cpp + fq_runner.cpp + fq_setup.cpp +) + +PEERDIR( + library/cpp/colorizer + library/cpp/testing/unittest + ydb/core/fq/libs/config/protos + ydb/core/fq/libs/control_plane_proxy/events + ydb/core/fq/libs/init + ydb/core/fq/libs/mock + ydb/core/testlib + ydb/library/folder_service/mock + ydb/library/grpc/server/actors + ydb/library/security + ydb/library/yql/providers/pq/provider + ydb/tests/tools/kqprun/runlib + yql/essentials/minikql +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/ydb/tests/tools/fqrun/ya.make b/ydb/tests/tools/fqrun/ya.make new file mode 100644 index 000000000000..a80046acebed --- /dev/null +++ b/ydb/tests/tools/fqrun/ya.make @@ -0,0 +1,31 @@ +PROGRAM(fqrun) + +IF (PROFILE_MEMORY_ALLOCATIONS) + MESSAGE("Enabled profile memory allocations") + ALLOCATOR(LF_DBG) +ENDIF() + +SRCS( + fqrun.cpp +) + +PEERDIR( + library/cpp/colorizer + library/cpp/getopt + library/cpp/lfalloc/alloc_profiler + ydb/core/blob_depot + ydb/library/yql/providers/pq/gateway/dummy + ydb/tests/tools/fqrun/src + ydb/tests/tools/kqprun/runlib + yql/essentials/parser/pg_wrapper + yql/essentials/sql/pg +) + +PEERDIR( + yql/essentials/udfs/common/compress_base + yql/essentials/udfs/common/re2 +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/ydb/tests/tools/kqprun/.gitignore b/ydb/tests/tools/kqprun/.gitignore index a4578942a676..86fc65c21b31 100644 --- a/ydb/tests/tools/kqprun/.gitignore +++ b/ydb/tests/tools/kqprun/.gitignore @@ -1,6 +1,8 @@ +storage sync_dir example udfs + *.log *.json *.sql diff --git a/ydb/tests/tools/kqprun/README.md b/ydb/tests/tools/kqprun/README.md new file mode 100644 index 000000000000..10e283294135 --- /dev/null +++ b/ydb/tests/tools/kqprun/README.md @@ -0,0 +1,62 @@ +# KqpRun tool + +Tool can be used to execute queries by using kikimr provider. + +For profiling memory allocations build kqprun with ya make flag `-D PROFILE_MEMORY_ALLOCATIONS -D CXXFLAGS=-DPROFILE_MEMORY_ALLOCATIONS`. + +## Examples + +### Queries + +* Run select 42: + ```(bash) + ./kqprun --sql "SELECT 42" + ``` + +* Queries shooting: + ```(bash) + ./kqprun --sql "SELECT 42" -C async --loop-count 0 --loop-delay 100 --inflight-limit 10 + ``` + +### Logs + +* Setup log settings (`-C query` for clear logs): + ```(bash) + ./kqprun --sql "SELECT 42" -C query --log-default=warn --log KQP_YQL=trace --log-file query.log + ``` + +* Trace opt: + ```(bash) + ./kqprun --sql "SELECT 42" -C query -T script + ``` + +* Runtime statistics: + ```(bash) + ./kqprun --sql "SELECT 42" --script-statistics stats.log --script-timeline-file timeline.svg + ``` + +### Cluster + +* Embedded UI: + ```(bash) + ./kqprun -M 32000 + ``` + + Monitoring endpoint: http://localhost:32000 + +* gRPC endpoint: + ```(bash) + ./kqprun -G 32000 + ``` + + Connect with ydb CLI: `ydb -e grpc://localhost:32000 -d /Root` + +* Static storage: + ```(bash) + ./kqprun -M 32000 --storage-path ./storage --storage-size 32 + ``` + +* Create serverless domain and execute query in this domain: + ```(bash) + ./kqprun -M 32000 --shared my-shared --serverless my-serverless --sql "SELECT 42" -D /Root/my-serverless + ``` diff --git a/ydb/tests/tools/kqprun/configuration/app_config.conf b/ydb/tests/tools/kqprun/configuration/app_config.conf index 749a24bf9520..ba4ff05f187b 100644 --- a/ydb/tests/tools/kqprun/configuration/app_config.conf +++ b/ydb/tests/tools/kqprun/configuration/app_config.conf @@ -46,6 +46,7 @@ ActorSystemConfig { ColumnShardConfig { DisabledOnSchemeShard: false + WritingInFlightRequestBytesLimit: 104857600 } FeatureFlags { @@ -53,6 +54,7 @@ FeatureFlags { EnableScriptExecutionOperations: true EnableExternalSourceSchemaInference: true EnableTempTables: true + EnableReplaceIfExistsForExternalEntities: true } KQPConfig { @@ -97,7 +99,7 @@ QueryServiceConfig { Endpoint { host: "localhost" - port: 50051 + port: 2130 } } } diff --git a/ydb/tests/tools/kqprun/flame_graph.sh b/ydb/tests/tools/kqprun/flame_graph.sh index ead8de30683a..61749c81b21e 100755 --- a/ydb/tests/tools/kqprun/flame_graph.sh +++ b/ydb/tests/tools/kqprun/flame_graph.sh @@ -1,11 +1,12 @@ #!/usr/bin/env bash -# For svg graph download https://github.com/brendangregg/FlameGraph -# and run `FlameGraph/stackcollapse-perf.pl profdata.txt | FlameGraph/flamegraph.pl > profdata.svg` +set -eux -pid=$(pgrep -u $USER kqprun) +kqprun_pid=$(pgrep -u $USER kqprun) -echo "Target process id: ${pid}" - -sudo perf record -F 50 --call-graph dwarf -g --proc-map-timeout=10000 --pid $pid -v -o profdata -- sleep 30 +sudo perf record -F 50 --call-graph dwarf -g --proc-map-timeout=10000 --pid $kqprun_pid -v -o profdata -- sleep ${1:-'30'} sudo perf script -i profdata > profdata.txt + +flame_graph_tool="../../../../contrib/tools/flame-graph/" + +${flame_graph_tool}/stackcollapse-perf.pl profdata.txt | ${flame_graph_tool}/flamegraph.pl > profdata.svg diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index 69ddc226f7f9..a9ba125897e1 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -1,7 +1,3 @@ -#include "src/kqp_runner.h" - -#include - #include #include @@ -11,16 +7,26 @@ #include #include -#include - +#include #include #include -#include +#include +#include +#include + #include #include #include -#include +#ifdef PROFILE_MEMORY_ALLOCATIONS +#include +#endif + +using namespace NKikimrRun; + +namespace NKqpRun { + +namespace { struct TExecutionOptions { enum class EExecutionCase { @@ -77,11 +83,8 @@ struct TExecutionOptions { return GetValue(index, ScriptQueryActions, NKikimrKqp::EQueryAction::QUERY_ACTION_EXECUTE); } - NKqpRun::TRequestOptions GetSchemeQueryOptions() const { + TRequestOptions GetSchemeQueryOptions() const { TString sql = SchemeQuery; - if (UseTemplates) { - ReplaceYqlTokenTemplate(sql); - } return { .Query = sql, @@ -89,17 +92,16 @@ struct TExecutionOptions { .TraceId = DefaultTraceId, .PoolId = "", .UserSID = BUILTIN_ACL_ROOT, - .Database = "", + .Database = GetValue(0, Databases, TString()), .Timeout = TDuration::Zero() }; } - NKqpRun::TRequestOptions GetScriptQueryOptions(size_t index, size_t queryId, TInstant startTime) const { + TRequestOptions GetScriptQueryOptions(size_t index, size_t queryId, TInstant startTime) const { Y_ABORT_UNLESS(index < ScriptQueries.size()); TString sql = ScriptQueries[index]; if (UseTemplates) { - ReplaceYqlTokenTemplate(sql); SubstGlobal(sql, "${QUERY_ID}", ToString(queryId)); } @@ -110,40 +112,46 @@ struct TExecutionOptions { .PoolId = GetValue(index, PoolIds, TString()), .UserSID = GetValue(index, UserSIDs, TString(BUILTIN_ACL_ROOT)), .Database = GetValue(index, Databases, TString()), - .Timeout = GetValue(index, Timeouts, TDuration::Zero()) + .Timeout = GetValue(index, Timeouts, TDuration::Zero()), + .QueryId = queryId }; } - void Validate(const NKqpRun::TRunnerOptions& runnerOptions) const { + void Validate(const TRunnerOptions& runnerOptions) const { if (!SchemeQuery && ScriptQueries.empty() && !runnerOptions.YdbSettings.MonitoringEnabled && !runnerOptions.YdbSettings.GrpcEnabled) { ythrow yexception() << "Nothing to execute and is not running as daemon"; } - ValidateOptionsSizes(); + ValidateOptionsSizes(runnerOptions); ValidateSchemeQueryOptions(runnerOptions); ValidateScriptExecutionOptions(runnerOptions); ValidateAsyncOptions(runnerOptions.YdbSettings.AsyncQueriesSettings); - ValidateTraceOpt(runnerOptions.TraceOptType); + ValidateTraceOpt(runnerOptions); + ValidateStorageSettings(runnerOptions.YdbSettings); } private: - void ValidateOptionsSizes() const { - const auto checker = [numberQueries = ScriptQueries.size()](size_t checkSize, const TString& optionName) { - if (checkSize > numberQueries) { - ythrow yexception() << "Too many " << optionName << ". Specified " << checkSize << ", when number of queries is " << numberQueries; + void ValidateOptionsSizes(const TRunnerOptions& runnerOptions) const { + const auto checker = [numberQueries = ScriptQueries.size()](size_t checkSize, const TString& optionName, bool useInSchemeQuery = false) { + if (checkSize > std::max(numberQueries, static_cast(useInSchemeQuery ? 1 : 0))) { + ythrow yexception() << "Too many " << optionName << ". Specified " << checkSize << ", when number of script queries is " << numberQueries; } }; checker(ExecutionCases.size(), "execution cases"); checker(ScriptQueryActions.size(), "script query actions"); - checker(Databases.size(), "databases"); + checker(Databases.size(), "databases", true); checker(TraceIds.size(), "trace ids"); checker(PoolIds.size(), "pool ids"); checker(UserSIDs.size(), "user SIDs"); checker(Timeouts.size(), "timeouts"); + checker(runnerOptions.ScriptQueryAstOutputs.size(), "ast output files"); + checker(runnerOptions.ScriptQueryPlanOutputs.size(), "plan output files"); + checker(runnerOptions.ScriptQueryTimelineFiles.size(), "timeline files"); + checker(runnerOptions.InProgressStatisticsOutputFiles.size(), "statistics files"); } - void ValidateSchemeQueryOptions(const NKqpRun::TRunnerOptions& runnerOptions) const { + void ValidateSchemeQueryOptions(const TRunnerOptions& runnerOptions) const { if (SchemeQuery) { return; } @@ -152,7 +160,7 @@ struct TExecutionOptions { } } - void ValidateScriptExecutionOptions(const NKqpRun::TRunnerOptions& runnerOptions) const { + void ValidateScriptExecutionOptions(const TRunnerOptions& runnerOptions) const { if (runnerOptions.YdbSettings.SameSession && HasExecutionCase(EExecutionCase::AsyncQuery)) { ythrow yexception() << "Same session can not be used with async quries"; } @@ -175,7 +183,7 @@ struct TExecutionOptions { if (ResultsRowsLimit) { ythrow yexception() << "Result rows limit can not be used without script queries"; } - if (runnerOptions.InProgressStatisticsOutputFile) { + if (!runnerOptions.InProgressStatisticsOutputFiles.empty()) { ythrow yexception() << "Script statistics can not be used without script queries"; } @@ -183,10 +191,10 @@ struct TExecutionOptions { if (HasExecutionCase(EExecutionCase::YqlScript)) { return; } - if (runnerOptions.ScriptQueryAstOutput) { + if (!runnerOptions.ScriptQueryAstOutputs.empty()) { ythrow yexception() << "Script query AST output can not be used without script/yql queries"; } - if (runnerOptions.ScriptQueryPlanOutput) { + if (!runnerOptions.ScriptQueryPlanOutputs.empty()) { ythrow yexception() << "Script query plan output can not be used without script/yql queries"; } if (runnerOptions.YdbSettings.SameSession) { @@ -194,7 +202,7 @@ struct TExecutionOptions { } } - void ValidateAsyncOptions(const NKqpRun::TAsyncQueriesSettings& asyncQueriesSettings) const { + void ValidateAsyncOptions(const TAsyncQueriesSettings& asyncQueriesSettings) const { if (asyncQueriesSettings.InFlightLimit && !HasExecutionCase(EExecutionCase::AsyncQuery)) { ythrow yexception() << "In flight limit can not be used without async queries"; } @@ -205,51 +213,61 @@ struct TExecutionOptions { } } - void ValidateTraceOpt(NKqpRun::TRunnerOptions::ETraceOptType traceOptType) const { - switch (traceOptType) { - case NKqpRun::TRunnerOptions::ETraceOptType::Scheme: { + void ValidateTraceOpt(const TRunnerOptions& runnerOptions) const { + NColorizer::TColors colors = NColorizer::AutoColors(Cout); + switch (runnerOptions.TraceOptType) { + case TRunnerOptions::ETraceOptType::Scheme: { if (!SchemeQuery) { ythrow yexception() << "Trace opt type scheme cannot be used without scheme query"; } break; } - case NKqpRun::TRunnerOptions::ETraceOptType::Script: { + case TRunnerOptions::ETraceOptType::Script: { if (ScriptQueries.empty()) { ythrow yexception() << "Trace opt type script cannot be used without script queries"; } } - case NKqpRun::TRunnerOptions::ETraceOptType::All: { + case TRunnerOptions::ETraceOptType::All: { if (!SchemeQuery && ScriptQueries.empty()) { ythrow yexception() << "Trace opt type all cannot be used without any queries"; } } - case NKqpRun::TRunnerOptions::ETraceOptType::Disabled: { + case TRunnerOptions::ETraceOptType::Disabled: { break; } } - } -private: - template - static TValue GetValue(size_t index, const std::vector& values, TValue defaultValue) { - if (values.empty()) { - return defaultValue; + if (const auto traceOptId = runnerOptions.TraceOptScriptId) { + if (runnerOptions.TraceOptType != TRunnerOptions::ETraceOptType::Script) { + ythrow yexception() << "Trace opt id allowed only for trace opt type script (used " << runnerOptions.TraceOptType << ")"; + } + + const ui64 scriptNumber = ScriptQueries.size() * LoopCount; + if (*traceOptId >= scriptNumber) { + ythrow yexception() << "Invalid trace opt id " << *traceOptId << ", it should be less than number of script queries " << scriptNumber; + } + if (scriptNumber == 1) { + Cout << colors.Red() << "Warning: trace opt id is not necessary for single script mode" << Endl; + } } - return values[std::min(index, values.size() - 1)]; } - static void ReplaceYqlTokenTemplate(TString& sql) { - const TString variableName = TStringBuilder() << "${" << NKqpRun::YQL_TOKEN_VARIABLE << "}"; - if (const TString& yqlToken = GetEnv(NKqpRun::YQL_TOKEN_VARIABLE)) { - SubstGlobal(sql, variableName, yqlToken); - } else if (sql.Contains(variableName)) { - ythrow yexception() << "Failed to replace ${YQL_TOKEN} template, please specify YQL_TOKEN environment variable\n"; + static void ValidateStorageSettings(const TYdbSetupSettings& ydbSettings) { + if (ydbSettings.DisableDiskMock) { + if (ydbSettings.NodeCount + ydbSettings.Tenants.size() > 1) { + ythrow yexception() << "Disable disk mock cannot be used for multi node clusters (already disabled)"; + } else if (ydbSettings.PDisksPath) { + ythrow yexception() << "Disable disk mock cannot be used with real PDisks (already disabled)"; + } + } + if (ydbSettings.FormatStorage && !ydbSettings.PDisksPath) { + ythrow yexception() << "Cannot format storage without real PDisks, please use --storage-path"; } } }; -void RunArgumentQuery(size_t index, size_t queryId, TInstant startTime, const TExecutionOptions& executionOptions, NKqpRun::TKqpRunner& runner) { +void RunArgumentQuery(size_t index, size_t queryId, TInstant startTime, const TExecutionOptions& executionOptions, TKqpRunner& runner) { NColorizer::TColors colors = NColorizer::AutoColors(Cout); switch (executionOptions.GetExecutionCase(index)) { @@ -292,7 +310,7 @@ void RunArgumentQuery(size_t index, size_t queryId, TInstant startTime, const TE } -void RunArgumentQueries(const TExecutionOptions& executionOptions, NKqpRun::TKqpRunner& runner) { +void RunArgumentQueries(const TExecutionOptions& executionOptions, TKqpRunner& runner) { NColorizer::TColors colors = NColorizer::AutoColors(Cout); if (executionOptions.SchemeQuery) { @@ -354,11 +372,11 @@ void RunAsDaemon() { } -void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunnerOptions& runnerOptions) { +void RunScript(const TExecutionOptions& executionOptions, const TRunnerOptions& runnerOptions) { NColorizer::TColors colors = NColorizer::AutoColors(Cout); Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Initialization of kqp runner..." << colors.Default() << Endl; - NKqpRun::TKqpRunner runner(runnerOptions); + TKqpRunner runner(runnerOptions); try { RunArgumentQueries(executionOptions, runner); @@ -378,82 +396,18 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner } -TIntrusivePtr CreateFunctionRegistry(const TString& udfsDirectory, TVector udfsPaths, bool excludeLinkedUdfs) { - if (!udfsDirectory.empty() || !udfsPaths.empty()) { - NColorizer::TColors colors = NColorizer::AutoColors(Cout); - Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Fetching udfs..." << colors.Default() << Endl; - } - - NKikimr::NMiniKQL::FindUdfsInDir(udfsDirectory, &udfsPaths); - auto functionRegistry = NKikimr::NMiniKQL::CreateFunctionRegistry(&PrintBackTrace, NKikimr::NMiniKQL::CreateBuiltinRegistry(), false, udfsPaths)->Clone(); - - if (excludeLinkedUdfs) { - for (const auto& wrapper : NYql::NUdf::GetStaticUdfModuleWrapperList()) { - auto [name, ptr] = wrapper(); - if (!functionRegistry->IsLoadedUdfModule(name)) { - functionRegistry->AddModule(TString(NKikimr::NMiniKQL::StaticModulePrefix) + name, name, std::move(ptr)); - } - } - } else { - NKikimr::NMiniKQL::FillStaticModules(*functionRegistry); - } - - return functionRegistry; -} - +class TMain : public TMainBase { + using EVerbose = TYdbSetupSettings::EVerbose; -class TMain : public TMainClassArgs { - inline static const TString YqlToken = GetEnv(NKqpRun::YQL_TOKEN_VARIABLE); - inline static std::vector> FileHolders; + inline static const TString YqlToken = GetEnv(YQL_TOKEN_VARIABLE); TExecutionOptions ExecutionOptions; - NKqpRun::TRunnerOptions RunnerOptions; + TRunnerOptions RunnerOptions; + std::unordered_map Templates; THashMap TablesMapping; - TVector UdfsPaths; - TString UdfsDirectory; - bool ExcludeLinkedUdfs = false; bool EmulateYt = false; - static TString LoadFile(const TString& file) { - return TFileInput(file).ReadAll(); - } - - static IOutputStream* GetDefaultOutput(const TString& file) { - if (file == "-") { - return &Cout; - } - if (file) { - FileHolders.emplace_back(new TFileOutput(file)); - return FileHolders.back().get(); - } - return nullptr; - } - - template - class TChoices { - public: - explicit TChoices(std::map choicesMap) - : ChoicesMap(std::move(choicesMap)) - {} - - TResult operator()(const TString& choice) const { - return ChoicesMap.at(choice); - } - - TVector GetChoices() const { - TVector choices; - choices.reserve(ChoicesMap.size()); - for (const auto& [choice, _] : ChoicesMap) { - choices.emplace_back(choice); - } - return choices; - } - - private: - const std::map ChoicesMap; - }; - protected: void RegisterOptions(NLastGetopt::TOpts& options) override { options.SetTitle("KqpRun -- tool to execute queries by using kikimr provider (instead of dq provider in DQrun tool)"); @@ -467,15 +421,46 @@ class TMain : public TMainClassArgs { .Handler1([this](const NLastGetopt::TOptsParser* option) { ExecutionOptions.SchemeQuery = LoadFile(option->CurVal()); }); + options.AddLongOption('p', "script-query", "Script query to execute (typically DML query)") .RequiredArgument("file") .Handler1([this](const NLastGetopt::TOptsParser* option) { ExecutionOptions.ScriptQueries.emplace_back(LoadFile(option->CurVal())); }); + + options.AddLongOption("sql", "Script query SQL text to execute (typically DML query)") + .RequiredArgument("str") + .AppendTo(&ExecutionOptions.ScriptQueries); + options.AddLongOption("templates", "Enable templates for -s and -p queries, such as ${YQL_TOKEN} and ${QUERY_ID}") .NoArgument() .SetFlag(&ExecutionOptions.UseTemplates); + options.AddLongOption("var-template", "Add template from environment variables or file for -s and -p queries (use variable@file for files)") + .RequiredArgument("variable") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + TStringBuf variable; + TStringBuf filePath; + TStringBuf(option->CurVal()).Split('@', variable, filePath); + if (variable.empty()) { + ythrow yexception() << "Variable name should not be empty"; + } + + TString value; + if (!filePath.empty()) { + value = LoadFile(TString(filePath)); + } else { + value = GetEnv(TString(variable)); + if (!value) { + ythrow yexception() << "Invalid env template, can not find value for variable '" << variable << "'"; + } + } + + if (!Templates.emplace(variable, value).second) { + ythrow yexception() << "Got duplicated template variable name '" << variable << "'"; + } + }); + options.AddLongOption('t', "table", "File with input table (can be used by YT with -E flag), table@file") .RequiredArgument("table@file") .Handler1([this](const NLastGetopt::TOptsParser* option) { @@ -483,12 +468,11 @@ class TMain : public TMainClassArgs { TStringBuf filePath; TStringBuf(option->CurVal()).Split('@', tableName, filePath); if (tableName.empty() || filePath.empty()) { - ythrow yexception() << "Incorrect table mapping, expected form table@file, e.g. yt.Root/plato.Input@input.txt"; + ythrow yexception() << "Incorrect table mapping, expected form table@file, e. g. yt.Root/plato.Input@input.txt"; } - if (TablesMapping.contains(tableName)) { - ythrow yexception() << "Got duplicate table name: " << tableName; + if (!TablesMapping.emplace(tableName, filePath).second) { + ythrow yexception() << "Got duplicated table name: " << tableName; } - TablesMapping[tableName] = filePath; }); options.AddLongOption('c', "app-config", "File with app config (TAppConfig for ydb tenant)") @@ -504,33 +488,15 @@ class TMain : public TMainClassArgs { } }); - options.AddLongOption('u', "udf", "Load shared library with UDF by given path") - .RequiredArgument("file") - .EmplaceTo(&UdfsPaths); - options.AddLongOption("udfs-dir", "Load all shared libraries with UDFs found in given directory") - .RequiredArgument("directory") - .StoreResult(&UdfsDirectory); - options.AddLongOption("exclude-linked-udfs", "Exclude linked udfs when same udf passed from -u or --udfs-dir") - .NoArgument() - .SetFlag(&ExcludeLinkedUdfs); - // Outputs - options.AddLongOption("log-file", "File with execution logs (writes in stderr if empty)") - .RequiredArgument("file") - .StoreResult(&RunnerOptions.YdbSettings.LogOutputFile) - .Handler1([](const NLastGetopt::TOptsParser* option) { - if (const TString& file = option->CurVal()) { - std::remove(file.c_str()); - } - }); - TChoices traceOpt({ - {"all", NKqpRun::TRunnerOptions::ETraceOptType::All}, - {"scheme", NKqpRun::TRunnerOptions::ETraceOptType::Scheme}, - {"script", NKqpRun::TRunnerOptions::ETraceOptType::Script}, - {"disabled", NKqpRun::TRunnerOptions::ETraceOptType::Disabled} + TChoices traceOpt({ + {"all", TRunnerOptions::ETraceOptType::All}, + {"scheme", TRunnerOptions::ETraceOptType::Scheme}, + {"script", TRunnerOptions::ETraceOptType::Script}, + {"disabled", TRunnerOptions::ETraceOptType::Disabled} }); - options.AddLongOption('T', "trace-opt", "print AST in the begin of each transformation") + options.AddLongOption('T', "trace-opt", "Print AST in the begin of each transformation") .RequiredArgument("trace-opt-query") .DefaultValue("disabled") .Choices(traceOpt.GetChoices()) @@ -539,6 +505,11 @@ class TMain : public TMainClassArgs { RunnerOptions.YdbSettings.TraceOptEnabled = traceOptType != NKqpRun::TRunnerOptions::ETraceOptType::Disabled; return traceOptType; }); + + options.AddLongOption('I', "trace-opt-index", "Index of -p query to use --trace-opt, starts from zero") + .RequiredArgument("uint") + .StoreResult(&RunnerOptions.TraceOptScriptId); + options.AddLongOption("trace-id", "Trace id for -p queries") .RequiredArgument("id") .EmplaceTo(&ExecutionOptions.TraceIds); @@ -547,14 +518,16 @@ class TMain : public TMainClassArgs { .RequiredArgument("file") .DefaultValue("-") .StoreMappedResultT(&RunnerOptions.ResultOutput, &GetDefaultOutput); + options.AddLongOption('L', "result-rows-limit", "Rows limit for script execution results") .RequiredArgument("uint") .DefaultValue(0) .StoreResult(&ExecutionOptions.ResultsRowsLimit); - TChoices resultFormat({ - {"rows", NKqpRun::TRunnerOptions::EResultOutputFormat::RowsJson}, - {"full-json", NKqpRun::TRunnerOptions::EResultOutputFormat::FullJson}, - {"full-proto", NKqpRun::TRunnerOptions::EResultOutputFormat::FullProto} + + TChoices resultFormat({ + {"rows", EResultOutputFormat::RowsJson}, + {"full-json", EResultOutputFormat::FullJson}, + {"full-proto", EResultOutputFormat::FullProto} }); options.AddLongOption('R', "result-format", "Script query result format") .RequiredArgument("result-format") @@ -568,19 +541,26 @@ class TMain : public TMainClassArgs { options.AddLongOption("script-ast-file", "File with script query ast (use '-' to write in stdout)") .RequiredArgument("file") - .StoreMappedResultT(&RunnerOptions.ScriptQueryAstOutput, &GetDefaultOutput); + .Handler1([this](const NLastGetopt::TOptsParser* option) { + RunnerOptions.ScriptQueryAstOutputs.emplace_back(GetDefaultOutput(TString(option->CurValOrDef()))); + }); options.AddLongOption("script-plan-file", "File with script query plan (use '-' to write in stdout)") .RequiredArgument("file") - .StoreMappedResultT(&RunnerOptions.ScriptQueryPlanOutput, &GetDefaultOutput); + .Handler1([this](const NLastGetopt::TOptsParser* option) { + RunnerOptions.ScriptQueryPlanOutputs.emplace_back(GetDefaultOutput(TString(option->CurValOrDef()))); + }); + options.AddLongOption("script-statistics", "File with script inprogress statistics") .RequiredArgument("file") - .StoreMappedResultT(&RunnerOptions.InProgressStatisticsOutputFile, [](const TString& file) { + .Handler1([this](const NLastGetopt::TOptsParser* option) { + const TString file(option->CurValOrDef()); if (file == "-") { ythrow yexception() << "Script in progress statistics cannot be printed to stdout, please specify file name"; } - return file; + RunnerOptions.InProgressStatisticsOutputFiles.emplace_back(file); }); + TChoices planFormat({ {"pretty", NYdb::NConsoleClient::EDataFormat::Pretty}, {"table", NYdb::NConsoleClient::EDataFormat::PrettyTable}, @@ -594,11 +574,12 @@ class TMain : public TMainClassArgs { options.AddLongOption("script-timeline-file", "File with script query timline in svg format") .RequiredArgument("file") - .StoreMappedResultT(&RunnerOptions.ScriptQueryTimelineFile, [](const TString& file) { + .Handler1([this](const NLastGetopt::TOptsParser* option) { + const TString file(option->CurValOrDef()); if (file == "-") { ythrow yexception() << "Script timline cannot be printed to stdout, please specify file name"; } - return file; + RunnerOptions.ScriptQueryTimelineFiles.emplace_back(file); }); // Pipeline settings @@ -616,13 +597,22 @@ class TMain : public TMainClassArgs { TString choice(option->CurValOrDef()); ExecutionOptions.ExecutionCases.emplace_back(executionCase(choice)); }); + options.AddLongOption("inflight-limit", "In flight limit for async queries (use 0 for unlimited)") .RequiredArgument("uint") .DefaultValue(0) .StoreResult(&RunnerOptions.YdbSettings.AsyncQueriesSettings.InFlightLimit); - TChoices verbose({ - {"each-query", NKqpRun::TAsyncQueriesSettings::EVerbose::EachQuery}, - {"final", NKqpRun::TAsyncQueriesSettings::EVerbose::Final} + + options.AddLongOption("verbose", TStringBuilder() << "Common verbose level (max level " << static_cast(EVerbose::Max) - 1 << ")") + .RequiredArgument("uint") + .DefaultValue(static_cast(EVerbose::Info)) + .StoreMappedResultT(&RunnerOptions.YdbSettings.VerboseLevel, [](ui8 value) { + return static_cast(std::min(value, static_cast(EVerbose::Max))); + }); + + TChoices verbose({ + {"each-query", TAsyncQueriesSettings::EVerbose::EachQuery}, + {"final", TAsyncQueriesSettings::EVerbose::Final} }); options.AddLongOption("async-verbose", "Verbose type for async queries") .RequiredArgument("type") @@ -660,10 +650,12 @@ class TMain : public TMainClassArgs { .RequiredArgument("uint") .DefaultValue(ExecutionOptions.LoopCount) .StoreResult(&ExecutionOptions.LoopCount); + options.AddLongOption("loop-delay", "Delay in milliseconds between loop steps") .RequiredArgument("uint") .DefaultValue(0) .StoreMappedResultT(&ExecutionOptions.LoopDelay, &TDuration::MilliSeconds); + options.AddLongOption("continue-after-fail", "Don't not stop requests execution after fails") .NoArgument() .SetFlag(&ExecutionOptions.ContinueAfterFail); @@ -696,89 +688,108 @@ class TMain : public TMainClassArgs { return nodeCount; }); - options.AddLongOption('M', "monitoring", "Embedded UI port (use 0 to start on random free port), if used kqprun will be run as daemon") - .RequiredArgument("uint") - .Handler1([this](const NLastGetopt::TOptsParser* option) { - if (const TString& port = option->CurVal()) { - RunnerOptions.YdbSettings.MonitoringEnabled = true; - RunnerOptions.YdbSettings.MonitoringPortOffset = FromString(port); - } - }); + options.AddLongOption('E', "emulate-yt", "Emulate YT tables (use file gateway instead of native gateway)") + .NoArgument() + .SetFlag(&EmulateYt); - options.AddLongOption('G', "grpc", "gRPC port (use 0 to start on random free port), if used kqprun will be run as daemon") + options.AddLongOption('H', "health-check", TStringBuilder() << "Level of health check before start (max level " << static_cast(TYdbSetupSettings::EHealthCheck::Max) - 1 << ")") .RequiredArgument("uint") - .Handler1([this](const NLastGetopt::TOptsParser* option) { - if (const TString& port = option->CurVal()) { - RunnerOptions.YdbSettings.GrpcEnabled = true; - RunnerOptions.YdbSettings.GrpcPort = FromString(port); - } + .DefaultValue(static_cast(TYdbSetupSettings::EHealthCheck::NodesCount)) + .StoreMappedResultT(&RunnerOptions.YdbSettings.HealthCheckLevel, [](ui8 value) { + return static_cast(std::min(value, static_cast(TYdbSetupSettings::EHealthCheck::Max))); }); - options.AddLongOption('E', "emulate-yt", "Emulate YT tables (use file gateway instead of native gateway)") - .NoArgument() - .SetFlag(&EmulateYt); + options.AddLongOption("health-check-timeout", "Health check timeout in seconds") + .RequiredArgument("uint") + .DefaultValue(10) + .StoreMappedResultT(&RunnerOptions.YdbSettings.HealthCheckTimeout, &TDuration::Seconds); + + const auto addTenant = [this](const TString& type, TStorageMeta::TTenant::EType protoType, const NLastGetopt::TOptsParser* option) { + TStringBuf tenant; + TStringBuf nodesCountStr; + TStringBuf(option->CurVal()).Split(':', tenant, nodesCountStr); + if (tenant.empty()) { + ythrow yexception() << type << " tenant name should not be empty"; + } - options.AddLongOption("domain", "Test cluster domain name") - .RequiredArgument("name") - .DefaultValue(RunnerOptions.YdbSettings.DomainName) - .StoreResult(&RunnerOptions.YdbSettings.DomainName); + TStorageMeta::TTenant tenantInfo; + tenantInfo.SetType(protoType); + tenantInfo.SetNodesCount(nodesCountStr ? FromString(nodesCountStr) : 1); + if (tenantInfo.GetNodesCount() == 0) { + ythrow yexception() << type << " tenant should have at least one node"; + } - options.AddLongOption("dedicated", "Dedicated tenant path, relative inside domain") + if (!RunnerOptions.YdbSettings.Tenants.emplace(tenant, tenantInfo).second) { + ythrow yexception() << "Got duplicated tenant name: " << tenant; + } + }; + options.AddLongOption("dedicated", "Dedicated tenant path, relative inside domain (for node count use dedicated-name:node-count)") .RequiredArgument("path") - .InsertTo(&RunnerOptions.YdbSettings.DedicatedTenants); + .Handler1(std::bind(addTenant, "Dedicated", TStorageMeta::TTenant::DEDICATED, std::placeholders::_1)); - options.AddLongOption("shared", "Shared tenant path, relative inside domain") + options.AddLongOption("shared", "Shared tenant path, relative inside domain (for node count use dedicated-name:node-count)") .RequiredArgument("path") - .InsertTo(&RunnerOptions.YdbSettings.SharedTenants); + .Handler1(std::bind(addTenant, "Shared", TStorageMeta::TTenant::SHARED, std::placeholders::_1)); options.AddLongOption("serverless", "Serverless tenant path, relative inside domain (use string serverless-name@shared-name to specify shared database)") .RequiredArgument("path") - .InsertTo(&RunnerOptions.YdbSettings.ServerlessTenants); + .Handler1([this](const NLastGetopt::TOptsParser* option) { + TStringBuf serverless; + TStringBuf shared; + TStringBuf(option->CurVal()).Split('@', serverless, shared); + if (serverless.empty()) { + ythrow yexception() << "Serverless tenant name should not be empty"; + } + + TStorageMeta::TTenant tenantInfo; + tenantInfo.SetType(TStorageMeta::TTenant::SERVERLESS); + tenantInfo.SetSharedTenant(TString(shared)); + if (!RunnerOptions.YdbSettings.Tenants.emplace(serverless, tenantInfo).second) { + ythrow yexception() << "Got duplicated tenant name: " << serverless; + } + }); - options.AddLongOption("storage-size", "Domain storage size in gigabytes") + options.AddLongOption("storage-size", TStringBuilder() << "Domain storage size in gigabytes (" << NKikimr::NBlobDepot::FormatByteSize(DEFAULT_STORAGE_SIZE) << " by default)") .RequiredArgument("uint") - .DefaultValue(32) .StoreMappedResultT(&RunnerOptions.YdbSettings.DiskSize, [](ui32 diskSize) { return static_cast(diskSize) << 30; }); - options.AddLongOption("real-pdisks", "Use real PDisks instead of in memory PDisks (also disable disk mock)") + options.AddLongOption("storage-path", "Use real PDisks by specified path instead of in memory PDisks (also disable disk mock), use '-' to use temp directory") + .RequiredArgument("directory") + .StoreResult(&RunnerOptions.YdbSettings.PDisksPath); + + options.AddLongOption("format-storage", "Clear storage if it exists on --storage-path") .NoArgument() - .SetFlag(&RunnerOptions.YdbSettings.UseRealPDisks); + .SetFlag(&RunnerOptions.YdbSettings.FormatStorage); options.AddLongOption("disable-disk-mock", "Disable disk mock on single node cluster") .NoArgument() .SetFlag(&RunnerOptions.YdbSettings.DisableDiskMock); - TChoices> backtrace({ - {"heavy", &NKikimr::EnableYDBBacktraceFormat}, - {"light", []() { SetFormatBackTraceFn(FormatBackTrace); }} - }); - options.AddLongOption("backtrace", "Default backtrace format function") - .RequiredArgument("backtrace-type") - .DefaultValue("heavy") - .Choices(backtrace.GetChoices()) - .Handler1([backtrace](const NLastGetopt::TOptsParser* option) { - TString choice(option->CurValOrDef()); - backtrace(choice)(); - }); + RegisterKikimrOptions(options, RunnerOptions.YdbSettings); } int DoRun(NLastGetopt::TOptsParseResult&&) override { ExecutionOptions.Validate(RunnerOptions); - if (RunnerOptions.YdbSettings.DisableDiskMock && RunnerOptions.YdbSettings.NodeCount + RunnerOptions.YdbSettings.SharedTenants.size() + RunnerOptions.YdbSettings.DedicatedTenants.size() > 1) { - ythrow yexception() << "Disable disk mock cannot be used for multi node clusters"; + ReplaceTemplates(ExecutionOptions.SchemeQuery); + for (auto& sql : ExecutionOptions.ScriptQueries) { + ReplaceTemplates(sql); } RunnerOptions.YdbSettings.YqlToken = YqlToken; - RunnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry(UdfsDirectory, UdfsPaths, ExcludeLinkedUdfs).Get(); + RunnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry().Get(); + + auto& appConfig = RunnerOptions.YdbSettings.AppConfig; if (ExecutionOptions.ResultsRowsLimit) { - RunnerOptions.YdbSettings.AppConfig.MutableQueryServiceConfig()->SetScriptResultRowsLimit(ExecutionOptions.ResultsRowsLimit); + appConfig.MutableQueryServiceConfig()->SetScriptResultRowsLimit(ExecutionOptions.ResultsRowsLimit); } + FillLogConfig(*appConfig.MutableLogConfig()); + if (EmulateYt) { - const auto& fileStorageConfig = RunnerOptions.YdbSettings.AppConfig.GetQueryServiceConfig().GetFileStorage(); + const auto& fileStorageConfig = appConfig.GetQueryServiceConfig().GetFileStorage(); auto fileStorage = WithAsync(CreateFileStorage(fileStorageConfig, {MakeYtDownloader(fileStorageConfig)})); auto ytFileServices = NYql::NFile::TYtFileServices::Make(RunnerOptions.YdbSettings.FunctionRegistry.Get(), TablesMapping, fileStorage); RunnerOptions.YdbSettings.YtGateway = NYql::CreateYtFileGateway(ytFileServices); @@ -787,40 +798,54 @@ class TMain : public TMainClassArgs { ythrow yexception() << "Tables mapping is not supported without emulate YT mode"; } - RunScript(ExecutionOptions, RunnerOptions); - return 0; - } -}; - - -void KqprunTerminateHandler() { - NColorizer::TColors colors = NColorizer::AutoColors(Cerr); - - Cerr << colors.Red() << "======= terminate() call stack ========" << colors.Default() << Endl; - FormatBackTrace(&Cerr); - Cerr << colors.Red() << "=======================================" << colors.Default() << Endl; +#ifdef PROFILE_MEMORY_ALLOCATIONS + if (RunnerOptions.YdbSettings.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors.Cyan() << "Starting profile memory allocations" << CoutColors.Default() << Endl; + } + NAllocProfiler::StartAllocationSampling(true); +#else + if (ProfileAllocationsOutput) { + ythrow yexception() << "Profile memory allocations disabled, please rebuild kqprun with flag `-D PROFILE_MEMORY_ALLOCATIONS`"; + } +#endif - abort(); -} + RunScript(ExecutionOptions, RunnerOptions); +#ifdef PROFILE_MEMORY_ALLOCATIONS + if (RunnerOptions.YdbSettings.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors.Cyan() << "Finishing profile memory allocations" << CoutColors.Default() << Endl; + } + FinishProfileMemoryAllocations(); +#endif -void SegmentationFaultHandler(int) { - NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + return 0; + } - Cerr << colors.Red() << "======= segmentation fault call stack ========" << colors.Default() << Endl; - FormatBackTrace(&Cerr); - Cerr << colors.Red() << "==============================================" << colors.Default() << Endl; +private: + void ReplaceTemplates(TString& sql) const { + for (const auto& [variable, value] : Templates) { + SubstGlobal(sql, TStringBuilder() << "${" << variable <<"}", value); + } + if (ExecutionOptions.UseTemplates) { + const TString tokenVariableName = TStringBuilder() << "${" << YQL_TOKEN_VARIABLE << "}"; + if (const TString& yqlToken = GetEnv(YQL_TOKEN_VARIABLE)) { + SubstGlobal(sql, tokenVariableName, yqlToken); + } else if (sql.Contains(tokenVariableName)) { + ythrow yexception() << "Failed to replace ${YQL_TOKEN} template, please specify YQL_TOKEN environment variable"; + } + } + } +}; - abort(); -} +} // anonymous namespace +} // namespace NKqpRun int main(int argc, const char* argv[]) { - std::set_terminate(KqprunTerminateHandler); - signal(SIGSEGV, &SegmentationFaultHandler); + SetupSignalActions(); try { - TMain().Run(argc, argv); + NKqpRun::TMain().Run(argc, argv); } catch (...) { NColorizer::TColors colors = NColorizer::AutoColors(Cerr); diff --git a/ydb/tests/tools/kqprun/runlib/application.cpp b/ydb/tests/tools/kqprun/runlib/application.cpp new file mode 100644 index 000000000000..6b27760b5b51 --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/application.cpp @@ -0,0 +1,164 @@ +#include "application.h" +#include "utils.h" + +#include +#include + +#include + +#include +#include + +#include +#include + +#ifdef PROFILE_MEMORY_ALLOCATIONS +#include +#endif + +namespace NKikimrRun { + +#ifdef PROFILE_MEMORY_ALLOCATIONS +void TMainBase::FinishProfileMemoryAllocations() { + if (ProfileAllocationsOutput) { + NAllocProfiler::StopAllocationSampling(*ProfileAllocationsOutput); + } else { + TString output; + TStringOutput stream(output); + NAllocProfiler::StopAllocationSampling(stream); + + Cout << CoutColors.Red() << "Warning: profile memory allocations output is not specified, please use flag `--profile-output` for writing profile info (dump size " << NKikimr::NBlobDepot::FormatByteSize(output.size()) << ")" << CoutColors.Default() << Endl; + } +} +#endif + +void TMainBase::RegisterKikimrOptions(NLastGetopt::TOpts& options, TServerSettings& settings) { + options.AddLongOption('u', "udf", "Load shared library with UDF by given path") + .RequiredArgument("file") + .EmplaceTo(&UdfsPaths); + + options.AddLongOption("udfs-dir", "Load all shared libraries with UDFs found in given directory") + .RequiredArgument("directory") + .StoreResult(&UdfsDirectory); + + options.AddLongOption("exclude-linked-udfs", "Exclude linked udfs when same udf passed from -u or --udfs-dir") + .NoArgument() + .SetFlag(&ExcludeLinkedUdfs); + + options.AddLongOption("log-file", "File with execution logs (writes in stderr if empty)") + .RequiredArgument("file") + .StoreResult(&settings.LogOutputFile) + .Handler1([](const NLastGetopt::TOptsParser* option) { + if (const TString& file = option->CurVal()) { + std::remove(file.c_str()); + } + }); + + RegisterLogOptions(options); + + options.AddLongOption("profile-output", "File with profile memory allocations output (use '-' to write in stdout)") + .RequiredArgument("file") + .StoreMappedResultT(&ProfileAllocationsOutput, &GetDefaultOutput); + + options.AddLongOption('M', "monitoring", "Embedded UI port (use 0 to start on random free port), if used will be run as daemon") + .RequiredArgument("uint") + .Handler1([&settings](const NLastGetopt::TOptsParser* option) { + if (const TString& port = option->CurVal()) { + settings.MonitoringEnabled = true; + settings.MonitoringPortOffset = FromString(port); + } + }); + + options.AddLongOption('G', "grpc", "gRPC port (use 0 to start on random free port), if used will be run as daemon") + .RequiredArgument("uint") + .Handler1([&settings](const NLastGetopt::TOptsParser* option) { + if (const TString& port = option->CurVal()) { + settings.GrpcEnabled = true; + settings.GrpcPort = FromString(port); + } + }); + + options.AddLongOption("domain", "Test cluster domain name") + .RequiredArgument("name") + .DefaultValue(settings.DomainName) + .StoreResult(&settings.DomainName); + + TChoices> backtrace({ + {"heavy", &NKikimr::EnableYDBBacktraceFormat}, + {"light", []() { SetFormatBackTraceFn(FormatBackTrace); }} + }); + options.AddLongOption("backtrace", "Default backtrace format function") + .RequiredArgument("backtrace-type") + .DefaultValue("heavy") + .Choices(backtrace.GetChoices()) + .Handler1([backtrace](const NLastGetopt::TOptsParser* option) { + TString choice(option->CurValOrDef()); + backtrace(choice)(); + }); +} + +void TMainBase::RegisterLogOptions(NLastGetopt::TOpts& options) { + options.AddLongOption("log-default", "Default log priority") + .RequiredArgument("priority") + .StoreMappedResultT(&DefaultLogPriority, GetLogPrioritiesMap("log-default")); + + options.AddLongOption("log", "Component log priority in format = (e. g. KQP_YQL=trace)") + .RequiredArgument("component priority") + .Handler1([this, logPriority = GetLogPrioritiesMap("log")](const NLastGetopt::TOptsParser* option) { + TStringBuf component; + TStringBuf priority; + TStringBuf(option->CurVal()).Split('=', component, priority); + if (component.empty() || priority.empty()) { + ythrow yexception() << "Incorrect log setting, expected form component=priority, e. g. KQP_YQL=trace"; + } + + const auto service = GetLogService(TString(component)); + if (!LogPriorities.emplace(service, logPriority(TString(priority))).second) { + ythrow yexception() << "Got duplicated log service name: " << component; + } + }); +} + +void TMainBase::FillLogConfig(NKikimrConfig::TLogConfig& config) const { + if (DefaultLogPriority) { + config.SetDefaultLevel(*DefaultLogPriority); + } + ModifyLogPriorities(LogPriorities, config); +} + +IOutputStream* TMainBase::GetDefaultOutput(const TString& file) { + if (file == "-") { + return &Cout; + } + if (file) { + FileHolders.emplace_back(new TFileOutput(file)); + return FileHolders.back().get(); + } + return nullptr; +} + +TIntrusivePtr TMainBase::CreateFunctionRegistry() const { + if (!UdfsDirectory.empty() || !UdfsPaths.empty()) { + NColorizer::TColors colors = NColorizer::AutoColors(Cout); + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Fetching udfs..." << colors.Default() << Endl; + } + + auto paths = UdfsPaths; + NKikimr::NMiniKQL::FindUdfsInDir(UdfsDirectory, &paths); + auto functionRegistry = NKikimr::NMiniKQL::CreateFunctionRegistry(&PrintBackTrace, NKikimr::NMiniKQL::CreateBuiltinRegistry(), false, paths)->Clone(); + + if (ExcludeLinkedUdfs) { + for (const auto& wrapper : NYql::NUdf::GetStaticUdfModuleWrapperList()) { + auto [name, ptr] = wrapper(); + if (!functionRegistry->IsLoadedUdfModule(name)) { + functionRegistry->AddModule(TString(NKikimr::NMiniKQL::StaticModulePrefix) + name, name, std::move(ptr)); + } + } + } else { + NKikimr::NMiniKQL::FillStaticModules(*functionRegistry); + } + + return functionRegistry; +} + +} // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/application.h b/ydb/tests/tools/kqprun/runlib/application.h new file mode 100644 index 000000000000..fd67eb2ac0d9 --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/application.h @@ -0,0 +1,50 @@ +#pragma once + +#include "settings.h" + +#include +#include + +#include + +#include +#include +#include + +#include + +namespace NKikimrRun { + +class TMainBase : public TMainClassArgs { +#ifdef PROFILE_MEMORY_ALLOCATIONS +public: + static void FinishProfileMemoryAllocations(); +#endif + +protected: + void RegisterKikimrOptions(NLastGetopt::TOpts& options, TServerSettings& settings); + + virtual void RegisterLogOptions(NLastGetopt::TOpts& options); + + void FillLogConfig(NKikimrConfig::TLogConfig& config) const; + + static IOutputStream* GetDefaultOutput(const TString& file); + + TIntrusivePtr CreateFunctionRegistry() const; + +protected: + inline static NColorizer::TColors CoutColors = NColorizer::AutoColors(Cout); + inline static IOutputStream* ProfileAllocationsOutput = nullptr; + +private: + inline static std::vector> FileHolders; + + std::optional DefaultLogPriority; + std::unordered_map LogPriorities; + + TString UdfsDirectory; + TVector UdfsPaths; + bool ExcludeLinkedUdfs; +}; + +} // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/settings.h b/ydb/tests/tools/kqprun/runlib/settings.h new file mode 100644 index 000000000000..c77caf1876f2 --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/settings.h @@ -0,0 +1,25 @@ +#pragma once + +#include + +namespace NKikimrRun { + +struct TServerSettings { + TString DomainName = "Root"; + + bool MonitoringEnabled = false; + ui16 MonitoringPortOffset = 0; + + bool GrpcEnabled = false; + ui16 GrpcPort = 0; + + TString LogOutputFile; +}; + +enum class EResultOutputFormat { + RowsJson, // Rows in json format + FullJson, // Columns, rows and types in json format + FullProto, // Columns, rows and types in proto string format +}; + +} // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/utils.cpp b/ydb/tests/tools/kqprun/runlib/utils.cpp new file mode 100644 index 000000000000..abd1865dde83 --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/utils.cpp @@ -0,0 +1,304 @@ +#include "utils.h" +#include "application.h" + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include + +#include + +namespace NKikimrRun { + +namespace { + +void TerminateHandler() { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cerr << colors.Red() << "======= terminate() call stack ========" << colors.Default() << Endl; + FormatBackTrace(&Cerr); + Cerr << colors.Red() << "=======================================" << colors.Default() << Endl; + + abort(); +} + + +void SegmentationFaultHandler(int) { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cerr << colors.Red() << "======= segmentation fault call stack ========" << colors.Default() << Endl; + FormatBackTrace(&Cerr); + Cerr << colors.Red() << "==============================================" << colors.Default() << Endl; + + abort(); +} + +void FloatingPointExceptionHandler(int) { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cerr << colors.Red() << "======= floating point exception call stack ========" << colors.Default() << Endl; + FormatBackTrace(&Cerr); + Cerr << colors.Red() << "====================================================" << colors.Default() << Endl; + + abort(); +} + +#ifdef PROFILE_MEMORY_ALLOCATIONS +void InterruptHandler(int) { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cout << colors.Red() << "Execution interrupted, finishing profile memory allocations..." << colors.Default() << Endl; + TMainBase::FinishProfileMemoryAllocations(); + + abort(); +} +#endif + +} // nonymous namespace + + +TRequestResult::TRequestResult() + : Status(Ydb::StatusIds::STATUS_CODE_UNSPECIFIED) +{} + +TRequestResult::TRequestResult(Ydb::StatusIds::StatusCode status, const NYql::TIssues& issues) + : Status(status) + , Issues(issues) +{} + +TRequestResult::TRequestResult(Ydb::StatusIds::StatusCode status, const google::protobuf::RepeatedPtrField& issues) + : Status(status) +{ + NYql::IssuesFromMessage(issues, Issues); +} + +bool TRequestResult::IsSuccess() const { + return Status == Ydb::StatusIds::SUCCESS; +} + +TString TRequestResult::ToString() const { + return TStringBuilder() << "Request finished with status: " << Status << "\nIssues:\n" << Issues.ToString() << "\n"; +} + +TStatsPrinter::TStatsPrinter(NYdb::NConsoleClient::EDataFormat planFormat) + : PlanFormat(planFormat) + , StatProcessor(NFq::CreateStatProcessor("stat_full")) +{} + +void TStatsPrinter::PrintPlan(const TString& plan, IOutputStream& output) const { + if (!plan) { + return; + } + + NJson::TJsonValue planJson; + NJson::ReadJsonTree(plan, &planJson, true); + if (!planJson.GetMapSafe().contains("meta")) { + return; + } + + NYdb::NConsoleClient::TQueryPlanPrinter printer(PlanFormat, true, output); + printer.Print(plan); +} + +void TStatsPrinter::PrintInProgressStatistics(const TString& plan, IOutputStream& output) const { + output << TInstant::Now().ToIsoStringLocal() << " Script in progress statistics" << Endl; + + auto convertedPlan = plan; + try { + convertedPlan = StatProcessor->ConvertPlan(plan); + } catch (const NJson::TJsonException& ex) { + output << "Error plan conversion: " << ex.what() << Endl; + return; + } + + try { + double cpuUsage = 0.0; + auto fullStat = StatProcessor->GetQueryStat(convertedPlan, cpuUsage, nullptr); + auto flatStat = StatProcessor->GetFlatStat(convertedPlan); + auto publicStat = StatProcessor->GetPublicStat(fullStat); + + output << "\nCPU usage: " << cpuUsage << Endl; + PrintStatistics(fullStat, flatStat, publicStat, output); + } catch (const NJson::TJsonException& ex) { + output << "Error stat conversion: " << ex.what() << Endl; + return; + } + + output << "\nPlan visualization:" << Endl; + PrintPlan(convertedPlan, output); +} + +void TStatsPrinter::PrintTimeline(const TString& plan, IOutputStream& output) { + TPlanVisualizer planVisualizer; + planVisualizer.LoadPlans(plan); + output.Write(planVisualizer.PrintSvg()); +} + +void TStatsPrinter::PrintStatistics(const TString& fullStat, const THashMap& flatStat, const NFq::TPublicStat& publicStat, IOutputStream& output) { + output << "\nFlat statistics:" << Endl; + for (const auto& [propery, value] : flatStat) { + TString valueString = ToString(value); + if (propery.Contains("Bytes")) { + valueString = NKikimr::NBlobDepot::FormatByteSize(value); + } else if (propery.Contains("TimeUs")) { + valueString = NFq::FormatDurationUs(value); + } else if (propery.Contains("TimeMs")) { + valueString = NFq::FormatDurationMs(value); + } else { + valueString = FormatNumber(value); + } + output << propery << " = " << valueString << Endl; + } + + output << "\nPublic statistics:" << Endl; + if (auto memoryUsageBytes = publicStat.MemoryUsageBytes) { + output << "MemoryUsage = " << NKikimr::NBlobDepot::FormatByteSize(*memoryUsageBytes) << Endl; + } + if (auto cpuUsageUs = publicStat.CpuUsageUs) { + output << "CpuUsage = " << NFq::FormatDurationUs(*cpuUsageUs) << Endl; + } + if (auto inputBytes = publicStat.InputBytes) { + output << "InputSize = " << NKikimr::NBlobDepot::FormatByteSize(*inputBytes) << Endl; + } + if (auto outputBytes = publicStat.OutputBytes) { + output << "OutputSize = " << NKikimr::NBlobDepot::FormatByteSize(*outputBytes) << Endl; + } + if (auto sourceInputRecords = publicStat.SourceInputRecords) { + output << "SourceInputRecords = " << FormatNumber(*sourceInputRecords) << Endl; + } + if (auto sinkOutputRecords = publicStat.SinkOutputRecords) { + output << "SinkOutputRecords = " << FormatNumber(*sinkOutputRecords) << Endl; + } + if (auto runningTasks = publicStat.RunningTasks) { + output << "RunningTasks = " << FormatNumber(*runningTasks) << Endl; + } + + output << "\nFull statistics:" << Endl; + NJson::TJsonValue statsJson; + NJson::ReadJsonTree(fullStat, &statsJson); + NJson::WriteJson(&output, &statsJson, true, true, true); + output << Endl; +} + +TString TStatsPrinter::FormatNumber(i64 number) { + struct TSeparator : public std::numpunct { + char do_thousands_sep() const final { + return '.'; + } + + std::string do_grouping() const final { + return "\03"; + } + }; + + std::ostringstream stream; + stream.imbue(std::locale(stream.getloc(), new TSeparator())); + stream << number; + return stream.str(); +} + +TString LoadFile(const TString& file) { + return TFileInput(file).ReadAll(); +} + +NKikimrServices::EServiceKikimr GetLogService(const TString& serviceName) { + NKikimrServices::EServiceKikimr service; + if (!NKikimrServices::EServiceKikimr_Parse(serviceName, &service)) { + ythrow yexception() << "Invalid kikimr service name " << serviceName; + } + return service; +} + +void ModifyLogPriorities(std::unordered_map logPriorities, NKikimrConfig::TLogConfig& logConfig) { + for (auto& entry : *logConfig.MutableEntry()) { + const auto it = logPriorities.find(GetLogService(entry.GetComponent())); + if (it != logPriorities.end()) { + entry.SetLevel(it->second); + logPriorities.erase(it); + } + } + for (const auto& [service, priority] : logPriorities) { + auto* entry = logConfig.AddEntry(); + entry->SetComponent(NKikimrServices::EServiceKikimr_Name(service)); + entry->SetLevel(priority); + } +} + +void InitLogSettings(const NKikimrConfig::TLogConfig& logConfig, NActors::TTestActorRuntimeBase& runtime) { + if (logConfig.HasDefaultLevel()) { + auto priority = NActors::NLog::EPriority(logConfig.GetDefaultLevel()); + auto descriptor = NKikimrServices::EServiceKikimr_descriptor(); + for (int i = 0; i < descriptor->value_count(); ++i) { + runtime.SetLogPriority(static_cast(descriptor->value(i)->number()), priority); + } + } + + for (const auto& setting : logConfig.get_arr_entry()) { + runtime.SetLogPriority(GetLogService(setting.GetComponent()), NActors::NLog::EPriority(setting.GetLevel())); + } +} + +TChoices GetLogPrioritiesMap(const TString& optionName) { + return TChoices({ + {"emerg", NActors::NLog::EPriority::PRI_EMERG}, + {"alert", NActors::NLog::EPriority::PRI_ALERT}, + {"crit", NActors::NLog::EPriority::PRI_CRIT}, + {"error", NActors::NLog::EPriority::PRI_ERROR}, + {"warn", NActors::NLog::EPriority::PRI_WARN}, + {"notice", NActors::NLog::EPriority::PRI_NOTICE}, + {"info", NActors::NLog::EPriority::PRI_INFO}, + {"debug", NActors::NLog::EPriority::PRI_DEBUG}, + {"trace", NActors::NLog::EPriority::PRI_TRACE}, + }, optionName, false); +} + +void SetupSignalActions() { + std::set_terminate(&TerminateHandler); + signal(SIGSEGV, &SegmentationFaultHandler); + signal(SIGFPE, &FloatingPointExceptionHandler); + +#ifdef PROFILE_MEMORY_ALLOCATIONS + signal(SIGINT, &InterruptHandler); +#endif +} + +void PrintResultSet(EResultOutputFormat format, IOutputStream& output, const Ydb::ResultSet& resultSet) { + switch (format) { + case EResultOutputFormat::RowsJson: { + NYdb::TResultSet result(resultSet); + NYdb::TResultSetParser parser(result); + while (parser.TryNextRow()) { + NJsonWriter::TBuf writer(NJsonWriter::HEM_UNSAFE, &output); + writer.SetWriteNanAsString(true); + NYdb::FormatResultRowJson(parser, result.GetColumnsMeta(), writer, NYdb::EBinaryStringEncoding::Unicode); + output << Endl; + } + break; + } + + case EResultOutputFormat::FullJson: { + resultSet.PrintJSON(output); + output << Endl; + break; + } + + case EResultOutputFormat::FullProto: { + TString resultSetString; + google::protobuf::TextFormat::Printer printer; + printer.SetSingleLineMode(false); + printer.SetUseUtf8StringEscaping(true); + printer.PrintToString(resultSet, &resultSetString); + output << resultSetString; + break; + } + } +} + +} // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/utils.h b/ydb/tests/tools/kqprun/runlib/utils.h new file mode 100644 index 000000000000..4bf873e48cae --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/utils.h @@ -0,0 +1,113 @@ +#pragma once + +#include "settings.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace NKikimrRun { + +struct TRequestResult { + Ydb::StatusIds::StatusCode Status; + NYql::TIssues Issues; + + TRequestResult(); + + TRequestResult(Ydb::StatusIds::StatusCode status, const NYql::TIssues& issues); + + TRequestResult(Ydb::StatusIds::StatusCode status, const google::protobuf::RepeatedPtrField& issues); + + bool IsSuccess() const; + + TString ToString() const; +}; + +template +class TChoices { +public: + explicit TChoices(std::map choicesMap, const TString& optionName = "", bool checkRegister = true) + : ChoicesMap(std::move(choicesMap)) + , OptionName(optionName) + , CheckRegister(checkRegister) + {} + + TResult operator()(TString choice) const { + if (!CheckRegister) { + std::for_each(choice.begin(), choice.vend(), [](char& c) { c = std::tolower(c); }); + } + + const auto it = ChoicesMap.find(choice); + if (it == ChoicesMap.end()) { + auto error = yexception() << "Value '" << choice << "' is not allowed " << (OptionName ? TStringBuilder() << "for option " << OptionName : TStringBuilder()) << ", available variants:"; + for (const auto& [value, _] : ChoicesMap) { + error << " " << value; + } + throw error; + } + return it->second; + } + + TVector GetChoices() const { + TVector choices; + choices.reserve(ChoicesMap.size()); + for (const auto& [choice, _] : ChoicesMap) { + choices.emplace_back(choice); + } + return choices; + } + + bool Contains(const TString& choice) const { + return ChoicesMap.contains(choice); + } + +private: + const std::map ChoicesMap; + const TString OptionName; + const bool CheckRegister; +}; + +class TStatsPrinter { +public: + explicit TStatsPrinter(NYdb::NConsoleClient::EDataFormat planFormat); + + void PrintPlan(const TString& plan, IOutputStream& output) const; + + void PrintInProgressStatistics(const TString& plan, IOutputStream& output) const; + + static void PrintTimeline(const TString& plan, IOutputStream& output); + + static void PrintStatistics(const TString& fullStat, const THashMap& flatStat, const NFq::TPublicStat& publicStat, IOutputStream& output); + + // Function adds thousands separators + // 123456789 -> 123.456.789 + static TString FormatNumber(i64 number); + +private: + const NYdb::NConsoleClient::EDataFormat PlanFormat; + const std::unique_ptr StatProcessor; +}; + +TString LoadFile(const TString& file); + +NKikimrServices::EServiceKikimr GetLogService(const TString& serviceName); + +void ModifyLogPriorities(std::unordered_map logPriorities, NKikimrConfig::TLogConfig& logConfig); + +void InitLogSettings(const NKikimrConfig::TLogConfig& logConfig, NActors::TTestActorRuntimeBase& runtime); + +TChoices GetLogPrioritiesMap(const TString& optionName); + +void SetupSignalActions(); + +void PrintResultSet(EResultOutputFormat format, IOutputStream& output, const Ydb::ResultSet& resultSet); + +} // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/ya.make b/ydb/tests/tools/kqprun/runlib/ya.make new file mode 100644 index 000000000000..39bd1328bfe7 --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/ya.make @@ -0,0 +1,30 @@ +LIBRARY() + +SRCS( + application.cpp + utils.cpp +) + +PEERDIR( + library/cpp/colorizer + library/cpp/getopt + library/cpp/json + ydb/core/base + ydb/core/blob_depot + ydb/core/fq/libs/compute/common + ydb/core/protos + ydb/library/actors/core + ydb/library/actors/testlib + ydb/library/services + ydb/public/api/protos + ydb/public/lib/json_value + ydb/public/lib/ydb_cli/common + yql/essentials/minikql + yql/essentials/minikql/invoke_builtins + yql/essentials/public/issue + yql/essentials/public/udf +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/ydb/tests/tools/kqprun/src/actors.cpp b/ydb/tests/tools/kqprun/src/actors.cpp index 66b28b109388..c548e0025660 100644 --- a/ydb/tests/tools/kqprun/src/actors.cpp +++ b/ydb/tests/tools/kqprun/src/actors.cpp @@ -14,6 +14,7 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped promise, TProgressCallback progressCallback) : TargetNode_(request.TargetNode) + , QueryId_(request.QueryId) , Request_(std::move(request.Event)) , Promise_(promise) , ResultRowsLimit_(std::numeric_limits::max()) @@ -82,13 +83,20 @@ class TRunScriptActorMock : public NActors::TActorBootstrappedGet()->Record); + try { + if (ProgressCallback_) { + ProgressCallback_(QueryId_, ev->Get()->Record); + } + } catch (...) { + Cerr << CerrColors_.Red() << "Got unexpected exception during progress callback: " << CurrentExceptionMessage() << CerrColors_.Default() << Endl; } } private: - ui32 TargetNode_ = 0; + const ui32 TargetNode_ = 0; + const size_t QueryId_ = 0; + const NColorizer::TColors CerrColors_ = NColorizer::AutoColors(Cerr); + std::unique_ptr Request_; NThreading::TPromise Promise_; ui64 ResultRowsLimit_; @@ -226,76 +234,157 @@ class TAsyncQueryRunnerActor : public NActors::TActor { }; class TResourcesWaiterActor : public NActors::TActorBootstrapped { + using IRetryPolicy = IRetryPolicy; + using EVerbose = TYdbSetupSettings::EVerbose; + using EHealthCheck = TYdbSetupSettings::EHealthCheck; + static constexpr TDuration REFRESH_PERIOD = TDuration::MilliSeconds(10); public: - TResourcesWaiterActor(NThreading::TPromise promise, i32 expectedNodeCount) - : ExpectedNodeCount_(expectedNodeCount) + TResourcesWaiterActor(NThreading::TPromise promise, const TWaitResourcesSettings& settings) + : Settings_(settings) + , RetryPolicy_(IRetryPolicy::GetExponentialBackoffPolicy( + &TResourcesWaiterActor::Retryable, REFRESH_PERIOD, + TDuration::MilliSeconds(100), TDuration::Seconds(1), + std::numeric_limits::max(), std::max(2 * REFRESH_PERIOD, Settings_.HealthCheckTimeout) + )) , Promise_(promise) {} void Bootstrap() { Become(&TResourcesWaiterActor::StateFunc); - CheckResourcesPublish(); + + HealthCheckStage_ = EHealthCheck::NodesCount; + DoHealthCheck(); } - void Handle(NActors::TEvents::TEvWakeup::TPtr&) { - CheckResourcesPublish(); + void DoHealthCheck() { + if (Settings_.HealthCheckLevel < HealthCheckStage_) { + Finish(); + return; + } + + switch (HealthCheckStage_) { + case TYdbSetupSettings::EHealthCheck::NodesCount: + CheckResourcesPublish(); + break; + + case TYdbSetupSettings::EHealthCheck::ScriptRequest: + StartScriptQuery(); + break; + + case TYdbSetupSettings::EHealthCheck::None: + case TYdbSetupSettings::EHealthCheck::Max: + Finish(); + break; + } } void Handle(TEvPrivate::TEvResourcesInfo::TPtr& ev) { - if (ev->Get()->NodeCount == ExpectedNodeCount_) { - Promise_.SetValue(); - PassAway(); + const auto nodeCount = ev->Get()->NodeCount; + if (nodeCount == Settings_.ExpectedNodeCount) { + HealthCheckStage_ = EHealthCheck::ScriptRequest; + DoHealthCheck(); + return; + } + + Retry(TStringBuilder() << "invalid node count, got " << nodeCount << ", expected " << Settings_.ExpectedNodeCount, true); + } + + void Handle(NKikimr::NKqp::TEvKqp::TEvScriptResponse::TPtr& ev) { + const auto status = ev->Get()->Status; + if (status == Ydb::StatusIds::SUCCESS) { + Finish(); return; } - Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup()); + Retry(TStringBuilder() << "script creation fail with status " << status << ", reason:\n" << CoutColors_.Default() << ev->Get()->Issues.ToString(), true); } STRICT_STFUNC(StateFunc, - hFunc(NActors::TEvents::TEvWakeup, Handle); + sFunc(NActors::TEvents::TEvWakeup, DoHealthCheck); hFunc(TEvPrivate::TEvResourcesInfo, Handle); + hFunc(NKikimr::NKqp::TEvKqp::TEvScriptResponse, Handle); ) private: void CheckResourcesPublish() { - GetResourceManager(); - if (!ResourceManager_) { - Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup()); - return; + ResourceManager_ = NKikimr::NKqp::TryGetKqpResourceManager(SelfId().NodeId()); } - UpdateResourcesInfo(); - } - - void GetResourceManager() { - if (ResourceManager_) { + if (!ResourceManager_) { + Retry("uninitialized resource manager", true); return; } - ResourceManager_ = NKikimr::NKqp::TryGetKqpResourceManager(SelfId().NodeId()); - } - void UpdateResourcesInfo() const { ResourceManager_->RequestClusterResourcesInfo( [selfId = SelfId(), actorContext = ActorContext()](TVector&& resources) { actorContext.Send(selfId, new TEvPrivate::TEvResourcesInfo(resources.size())); }); } + void StartScriptQuery() { + auto event = MakeHolder(); + event->Record.SetUserToken(NACLib::TUserToken("", BUILTIN_ACL_ROOT, {}).SerializeAsString()); + + auto request = event->Record.MutableRequest(); + request->SetQuery("SELECT 42"); + request->SetType(NKikimrKqp::QUERY_TYPE_SQL_GENERIC_SCRIPT); + request->SetAction(NKikimrKqp::EQueryAction::QUERY_ACTION_EXECUTE); + request->SetDatabase(Settings_.Database); + + Send(NKikimr::NKqp::MakeKqpProxyID(SelfId().NodeId()), event.Release()); + } + + void Retry(const TString& message, bool shortRetry) { + if (RetryState_ == nullptr) { + RetryState_ = RetryPolicy_->CreateRetryState(); + } + + if (auto delay = RetryState_->GetNextRetryDelay(shortRetry)) { + if (Settings_.VerboseLevel >= EVerbose::InitLogs) { + Cout << CoutColors_.Cyan() << "Retry in " << *delay << " " << message << CoutColors_.Default() << Endl; + } + Schedule(*delay, new NActors::TEvents::TEvWakeup()); + } else { + Fail(TStringBuilder() << "Health check timeout " << Settings_.HealthCheckTimeout << " exceeded, use --health-check-timeout for increasing it or check out health check logs by using --verbose " << static_cast(EVerbose::InitLogs)); + } + } + + void Finish() { + Promise_.SetValue(); + PassAway(); + } + + void Fail(const TString& error) { + Promise_.SetException(error); + PassAway(); + } + + static ERetryErrorClass Retryable(bool shortRetry) { + return shortRetry ? ERetryErrorClass::ShortRetry : ERetryErrorClass::LongRetry; + } + private: - const i32 ExpectedNodeCount_; + const TWaitResourcesSettings Settings_; + const NColorizer::TColors CoutColors_ = NColorizer::AutoColors(Cout); + const IRetryPolicy::TPtr RetryPolicy_; + IRetryPolicy::IRetryState::TPtr RetryState_ = nullptr; NThreading::TPromise Promise_; + EHealthCheck HealthCheckStage_ = EHealthCheck::None; std::shared_ptr ResourceManager_; }; class TSessionHolderActor : public NActors::TActorBootstrapped { + using EVerbose = TYdbSetupSettings::EVerbose; + public: TSessionHolderActor(TCreateSessionRequest request, NThreading::TPromise openPromise, NThreading::TPromise closePromise) : TargetNode_(request.TargetNode) , TraceId_(request.Event->Record.GetTraceId()) + , VerboseLevel_(request.VerboseLevel) , Request_(std::move(request.Event)) , OpenPromise_(openPromise) , ClosePromise_(closePromise) @@ -314,7 +403,9 @@ class TSessionHolderActor : public NActors::TActorBootstrapped= EVerbose::Info) { + Cout << CoutColors_.Cyan() << "Created new session on node " << TargetNode_ << " with id " << SessionId_ << "\n"; + } PingSession(); } @@ -390,6 +481,7 @@ class TSessionHolderActor : public NActors::TActorBootstrapped Request_; @@ -408,8 +500,8 @@ NActors::IActor* CreateAsyncQueryRunnerActor(const TAsyncQueriesSettings& settin return new TAsyncQueryRunnerActor(settings); } -NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise promise, i32 expectedNodeCount) { - return new TResourcesWaiterActor(promise, expectedNodeCount); +NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise promise, const TWaitResourcesSettings& settings) { + return new TResourcesWaiterActor(promise, settings); } NActors::IActor* CreateSessionHolderActor(TCreateSessionRequest request, NThreading::TPromise openPromise, NThreading::TPromise closePromise) { diff --git a/ydb/tests/tools/kqprun/src/actors.h b/ydb/tests/tools/kqprun/src/actors.h index cd477239ddb4..410794491b90 100644 --- a/ydb/tests/tools/kqprun/src/actors.h +++ b/ydb/tests/tools/kqprun/src/actors.h @@ -18,11 +18,21 @@ struct TQueryRequest { ui32 TargetNode; ui64 ResultRowsLimit; ui64 ResultSizeLimit; + size_t QueryId; }; struct TCreateSessionRequest { std::unique_ptr Event; ui32 TargetNode; + TYdbSetupSettings::EVerbose VerboseLevel; +}; + +struct TWaitResourcesSettings { + i32 ExpectedNodeCount; + TYdbSetupSettings::EHealthCheck HealthCheckLevel; + TDuration HealthCheckTimeout; + TYdbSetupSettings::EVerbose VerboseLevel; + TString Database; }; struct TEvPrivate { @@ -75,13 +85,13 @@ struct TEvPrivate { }; }; -using TProgressCallback = std::function; +using TProgressCallback = std::function; NActors::IActor* CreateRunScriptActorMock(TQueryRequest request, NThreading::TPromise promise, TProgressCallback progressCallback); NActors::IActor* CreateAsyncQueryRunnerActor(const TAsyncQueriesSettings& settings); -NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise promise, i32 expectedNodeCount); +NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise promise, const TWaitResourcesSettings& settings); NActors::IActor* CreateSessionHolderActor(TCreateSessionRequest request, NThreading::TPromise openPromise, NThreading::TPromise closePromise); diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h index 0c55b7bc04de..427dc6c03947 100644 --- a/ydb/tests/tools/kqprun/src/common.h +++ b/ydb/tests/tools/kqprun/src/common.h @@ -1,19 +1,25 @@ #pragma once #include -#include #include +#include +#include +#include +#include +#include + +#include #include #include -#include -#include +#include namespace NKqpRun { constexpr char YQL_TOKEN_VARIABLE[] = "YQL_TOKEN"; +constexpr ui64 DEFAULT_STORAGE_SIZE = 32_GB; struct TAsyncQueriesSettings { enum class EVerbose { @@ -25,27 +31,35 @@ struct TAsyncQueriesSettings { EVerbose Verbose = EVerbose::EachQuery; }; -struct TYdbSetupSettings { +struct TYdbSetupSettings : public NKikimrRun::TServerSettings { + enum class EVerbose { + None, + Info, + QueriesText, + InitLogs, + Max + }; + + enum class EHealthCheck { + None, + NodesCount, + ScriptRequest, + Max + }; + ui32 NodeCount = 1; - TString DomainName = "Root"; - std::unordered_set DedicatedTenants; - std::unordered_set SharedTenants; - std::unordered_set ServerlessTenants; - TDuration InitializationTimeout = TDuration::Seconds(10); + std::map Tenants; + TDuration HealthCheckTimeout = TDuration::Seconds(10); + EHealthCheck HealthCheckLevel = EHealthCheck::NodesCount; bool SameSession = false; bool DisableDiskMock = false; - bool UseRealPDisks = false; - ui64 DiskSize = 32_GB; - - bool MonitoringEnabled = false; - ui16 MonitoringPortOffset = 0; - - bool GrpcEnabled = false; - ui16 GrpcPort = 0; + bool FormatStorage = false; + std::optional PDisksPath; + std::optional DiskSize; bool TraceOptEnabled = false; - TString LogOutputFile; + EVerbose VerboseLevel = EVerbose::Info; TString YqlToken; TIntrusivePtr FunctionRegistry; @@ -64,22 +78,17 @@ struct TRunnerOptions { All, }; - enum class EResultOutputFormat { - RowsJson, // Rows in json format - FullJson, // Columns, rows and types in json format - FullProto, // Columns, rows and types in proto string format - }; - IOutputStream* ResultOutput = nullptr; IOutputStream* SchemeQueryAstOutput = nullptr; - IOutputStream* ScriptQueryAstOutput = nullptr; - IOutputStream* ScriptQueryPlanOutput = nullptr; - TString ScriptQueryTimelineFile; - TString InProgressStatisticsOutputFile; + std::vector ScriptQueryAstOutputs; + std::vector ScriptQueryPlanOutputs; + std::vector ScriptQueryTimelineFiles; + std::vector InProgressStatisticsOutputFiles; - EResultOutputFormat ResultOutputFormat = EResultOutputFormat::RowsJson; + NKikimrRun::EResultOutputFormat ResultOutputFormat = NKikimrRun::EResultOutputFormat::RowsJson; NYdb::NConsoleClient::EDataFormat PlanOutputFormat = NYdb::NConsoleClient::EDataFormat::Default; ETraceOptType TraceOptType = ETraceOptType::Disabled; + std::optional TraceOptScriptId; TDuration ScriptCancelAfter; @@ -95,6 +104,15 @@ struct TRequestOptions { TString UserSID; TString Database; TDuration Timeout; + size_t QueryId = 0; }; +template +TValue GetValue(size_t index, const std::vector& values, TValue defaultValue) { + if (values.empty()) { + return defaultValue; + } + return values[std::min(index, values.size() - 1)]; +} + } // namespace NKqpRun diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.cpp b/ydb/tests/tools/kqprun/src/kqp_runner.cpp index e880554d2c34..bbc432be16e4 100644 --- a/ydb/tests/tools/kqprun/src/kqp_runner.cpp +++ b/ydb/tests/tools/kqprun/src/kqp_runner.cpp @@ -2,91 +2,16 @@ #include "ydb_setup.h" #include -#include - -#include -#include - -#include -#include -#include +using namespace NKikimrRun; namespace NKqpRun { -namespace { - -// Function adds thousands separators -// 123456789 -> 123.456.789 -TString FormatNumber(i64 number) { - struct TSeparator : public std::numpunct { - char do_thousands_sep() const final { - return '.'; - } - - std::string do_grouping() const final { - return "\03"; - } - }; - - std::ostringstream stream; - stream.imbue(std::locale(stream.getloc(), new TSeparator())); - stream << number; - return stream.str(); -} - -void PrintStatistics(const TString& fullStat, const THashMap& flatStat, const NFq::TPublicStat& publicStat, IOutputStream& output) { - output << "\nFlat statistics:" << Endl; - for (const auto& [propery, value] : flatStat) { - TString valueString = ToString(value); - if (propery.Contains("Bytes")) { - valueString = NKikimr::NBlobDepot::FormatByteSize(value); - } else if (propery.Contains("TimeUs")) { - valueString = NFq::FormatDurationUs(value); - } else if (propery.Contains("TimeMs")) { - valueString = NFq::FormatDurationMs(value); - } else { - valueString = FormatNumber(value); - } - output << propery << " = " << valueString << Endl; - } - - output << "\nPublic statistics:" << Endl; - if (auto memoryUsageBytes = publicStat.MemoryUsageBytes) { - output << "MemoryUsage = " << NKikimr::NBlobDepot::FormatByteSize(*memoryUsageBytes) << Endl; - } - if (auto cpuUsageUs = publicStat.CpuUsageUs) { - output << "CpuUsage = " << NFq::FormatDurationUs(*cpuUsageUs) << Endl; - } - if (auto inputBytes = publicStat.InputBytes) { - output << "InputSize = " << NKikimr::NBlobDepot::FormatByteSize(*inputBytes) << Endl; - } - if (auto outputBytes = publicStat.OutputBytes) { - output << "OutputSize = " << NKikimr::NBlobDepot::FormatByteSize(*outputBytes) << Endl; - } - if (auto sourceInputRecords = publicStat.SourceInputRecords) { - output << "SourceInputRecords = " << FormatNumber(*sourceInputRecords) << Endl; - } - if (auto sinkOutputRecords = publicStat.SinkOutputRecords) { - output << "SinkOutputRecords = " << FormatNumber(*sinkOutputRecords) << Endl; - } - if (auto runningTasks = publicStat.RunningTasks) { - output << "RunningTasks = " << FormatNumber(*runningTasks) << Endl; - } - - output << "\nFull statistics:" << Endl; - NJson::TJsonValue statsJson; - NJson::ReadJsonTree(fullStat, &statsJson); - NJson::WriteJson(&output, &statsJson, true, true, true); - output << Endl; -} - -} // anonymous namespace - - //// TKqpRunner::TImpl class TKqpRunner::TImpl { + using EVerbose = TYdbSetupSettings::EVerbose; + public: enum class EQueryType { ScriptQuery, @@ -96,8 +21,9 @@ class TKqpRunner::TImpl { explicit TImpl(const TRunnerOptions& options) : Options_(options) + , VerboseLevel_(Options_.YdbSettings.VerboseLevel) , YdbSetup_(options.YdbSettings) - , StatProcessor_(NFq::CreateStatProcessor("stat_full")) + , StatsPrinter_(Options_.PlanOutputFormat) , CerrColors_(NColorizer::AutoColors(Cerr)) , CoutColors_(NColorizer::AutoColors(Cout)) {} @@ -105,6 +31,10 @@ class TKqpRunner::TImpl { bool ExecuteSchemeQuery(const TRequestOptions& query) const { StartSchemeTraceOpt(); + if (VerboseLevel_ >= EVerbose::QueriesText) { + Cout << CoutColors_.Cyan() << "Starting scheme request:\n" << CoutColors_.Default() << query.Query << Endl; + } + TSchemeMeta meta; TRequestResult status = YdbSetup_.SchemeQueryRequest(query, meta); TYdbSetup::StopTraceOpt(); @@ -120,7 +50,11 @@ class TKqpRunner::TImpl { } bool ExecuteScript(const TRequestOptions& script) { - StartScriptTraceOpt(); + StartScriptTraceOpt(script.QueryId); + + if (VerboseLevel_ >= EVerbose::QueriesText) { + Cout << CoutColors_.Cyan() << "Starting script request:\n" << CoutColors_.Default() << script.Query << Endl; + } TRequestResult status = YdbSetup_.ScriptRequest(script, ExecutionOperation_); @@ -132,13 +66,17 @@ class TKqpRunner::TImpl { ExecutionMeta_ = TExecutionMeta(); ExecutionMeta_.Database = script.Database; - return WaitScriptExecutionOperation(); + return WaitScriptExecutionOperation(script.QueryId); } bool ExecuteQuery(const TRequestOptions& query, EQueryType queryType) { - StartScriptTraceOpt(); + StartScriptTraceOpt(query.QueryId); StartTime_ = TInstant::Now(); + if (VerboseLevel_ >= EVerbose::QueriesText) { + Cout << CoutColors_.Cyan() << "Starting query request:\n" << CoutColors_.Default() << query.Query << Endl; + } + TString queryTypeStr; TQueryMeta meta; TRequestResult status; @@ -164,9 +102,9 @@ class TKqpRunner::TImpl { meta.Plan = ExecutionMeta_.Plan; } - PrintScriptAst(meta.Ast); - PrintScriptProgress(meta.Plan); - PrintScriptPlan(meta.Plan); + PrintScriptAst(query.QueryId, meta.Ast); + PrintScriptProgress(query.QueryId, meta.Plan); + PrintScriptPlan(query.QueryId, meta.Plan); PrintScriptFinish(meta, queryTypeStr); if (!status.IsSuccess()) { @@ -224,16 +162,16 @@ class TKqpRunner::TImpl { if (Options_.ResultOutput) { Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Writing script query results..." << CoutColors_.Default() << Endl; for (size_t i = 0; i < ResultSets_.size(); ++i) { - if (ResultSets_.size() > 1) { + if (ResultSets_.size() > 1 && VerboseLevel_ >= EVerbose::Info) { *Options_.ResultOutput << CoutColors_.Cyan() << "Result set " << i + 1 << ":" << CoutColors_.Default() << Endl; } - PrintScriptResult(ResultSets_[i]); + PrintResultSet(Options_.ResultOutputFormat, *Options_.ResultOutput, ResultSets_[i]); } } } private: - bool WaitScriptExecutionOperation() { + bool WaitScriptExecutionOperation(ui64 queryId) { StartTime_ = TInstant::Now(); TDuration getOperationPeriod = TDuration::Seconds(1); @@ -244,7 +182,7 @@ class TKqpRunner::TImpl { TRequestResult status; while (true) { status = YdbSetup_.GetScriptExecutionOperationRequest(ExecutionMeta_.Database, ExecutionOperation_, ExecutionMeta_); - PrintScriptProgress(ExecutionMeta_.Plan); + PrintScriptProgress(queryId, ExecutionMeta_.Plan); if (ExecutionMeta_.Ready) { break; @@ -267,9 +205,11 @@ class TKqpRunner::TImpl { Sleep(getOperationPeriod); } - PrintScriptAst(ExecutionMeta_.Ast); - PrintScriptProgress(ExecutionMeta_.Plan); - PrintScriptPlan(ExecutionMeta_.Plan); + TYdbSetup::StopTraceOpt(); + + PrintScriptAst(queryId, ExecutionMeta_.Ast); + PrintScriptProgress(queryId, ExecutionMeta_.Plan); + PrintScriptPlan(queryId, ExecutionMeta_.Plan); PrintScriptFinish(ExecutionMeta_, "Script"); if (!status.IsSuccess() || ExecutionMeta_.ExecutionStatus != NYdb::NQuery::EExecStatus::Completed) { @@ -290,127 +230,70 @@ class TKqpRunner::TImpl { } } - void StartScriptTraceOpt() const { - if (Options_.TraceOptType == TRunnerOptions::ETraceOptType::All || Options_.TraceOptType == TRunnerOptions::ETraceOptType::Script) { + void StartScriptTraceOpt(size_t queryId) const { + bool startTraceOpt = Options_.TraceOptType == TRunnerOptions::ETraceOptType::All; + + if (Options_.TraceOptType == TRunnerOptions::ETraceOptType::Script) { + startTraceOpt |= !Options_.TraceOptScriptId || *Options_.TraceOptScriptId == queryId; + } + + if (startTraceOpt) { YdbSetup_.StartTraceOpt(); } } void PrintSchemeQueryAst(const TString& ast) const { if (Options_.SchemeQueryAstOutput) { - Cout << CoutColors_.Cyan() << "Writing scheme query ast" << CoutColors_.Default() << Endl; + if (VerboseLevel_ >= EVerbose::Info) { + Cout << CoutColors_.Cyan() << "Writing scheme query ast" << CoutColors_.Default() << Endl; + } Options_.SchemeQueryAstOutput->Write(ast); } } - void PrintScriptAst(const TString& ast) const { - if (Options_.ScriptQueryAstOutput) { - Cout << CoutColors_.Cyan() << "Writing script query ast" << CoutColors_.Default() << Endl; - Options_.ScriptQueryAstOutput->Write(ast); - } - } - - void PrintPlan(const TString& plan, IOutputStream* output) const { - if (!plan) { - return; - } - - NJson::TJsonValue planJson; - NJson::ReadJsonTree(plan, &planJson, true); - if (!planJson.GetMapSafe().contains("meta")) { - return; + void PrintScriptAst(size_t queryId, const TString& ast) const { + if (const auto output = GetValue(queryId, Options_.ScriptQueryAstOutputs, nullptr)) { + if (VerboseLevel_ >= EVerbose::Info) { + Cout << CoutColors_.Cyan() << "Writing script query ast" << CoutColors_.Default() << Endl; + } + output->Write(ast); } - - NYdb::NConsoleClient::TQueryPlanPrinter printer(Options_.PlanOutputFormat, true, *output); - printer.Print(plan); } - void PrintScriptPlan(const TString& plan) const { - if (Options_.ScriptQueryPlanOutput) { - Cout << CoutColors_.Cyan() << "Writing script query plan" << CoutColors_.Default() << Endl; - PrintPlan(plan, Options_.ScriptQueryPlanOutput); + void PrintScriptPlan(size_t queryId, const TString& plan) const { + if (const auto output = GetValue(queryId, Options_.ScriptQueryPlanOutputs, nullptr)) { + if (VerboseLevel_ >= EVerbose::Info) { + Cout << CoutColors_.Cyan() << "Writing script query plan" << CoutColors_.Default() << Endl; + } + StatsPrinter_.PrintPlan(plan, *output); } } - void PrintScriptProgress(const TString& plan) const { - if (Options_.InProgressStatisticsOutputFile) { - TFileOutput outputStream(Options_.InProgressStatisticsOutputFile); - outputStream << TInstant::Now().ToIsoStringLocal() << " Script in progress statistics" << Endl; - - auto convertedPlan = plan; - try { - convertedPlan = StatProcessor_->ConvertPlan(plan); - } catch (const NJson::TJsonException& ex) { - outputStream << "Error plan conversion: " << ex.what() << Endl; - } - - try { - double cpuUsage = 0.0; - auto fullStat = StatProcessor_->GetQueryStat(convertedPlan, cpuUsage, nullptr); - auto flatStat = StatProcessor_->GetFlatStat(convertedPlan); - auto publicStat = StatProcessor_->GetPublicStat(fullStat); - - outputStream << "\nCPU usage: " << cpuUsage << Endl; - PrintStatistics(fullStat, flatStat, publicStat, outputStream); - } catch (const NJson::TJsonException& ex) { - outputStream << "Error stat conversion: " << ex.what() << Endl; - } - - outputStream << "\nPlan visualization:" << Endl; - PrintPlan(convertedPlan, &outputStream); - + void PrintScriptProgress(size_t queryId, const TString& plan) const { + if (const auto& output = GetValue(queryId, Options_.InProgressStatisticsOutputFiles, {})) { + TFileOutput outputStream(output); + StatsPrinter_.PrintInProgressStatistics(plan, outputStream); outputStream.Finish(); } - if (Options_.ScriptQueryTimelineFile) { - TFileOutput outputStream(Options_.ScriptQueryTimelineFile); - - TPlanVisualizer planVisualizer; - planVisualizer.LoadPlans(plan); - outputStream.Write(planVisualizer.PrintSvg()); - + if (const auto& output = GetValue(queryId, Options_.ScriptQueryTimelineFiles, {})) { + TFileOutput outputStream(output); + StatsPrinter_.PrintTimeline(plan, outputStream); outputStream.Finish(); } } TProgressCallback GetProgressCallback() { - return [this](const NKikimrKqp::TEvExecuterProgress& executerProgress) mutable { + return [this](ui64 queryId, const NKikimrKqp::TEvExecuterProgress& executerProgress) mutable { const TString& plan = executerProgress.GetQueryPlan(); ExecutionMeta_.Plan = plan; - PrintScriptProgress(plan); + PrintScriptProgress(queryId, plan); }; } - void PrintScriptResult(const Ydb::ResultSet& resultSet) const { - switch (Options_.ResultOutputFormat) { - case TRunnerOptions::EResultOutputFormat::RowsJson: { - NYdb::TResultSet result(resultSet); - NYdb::TResultSetParser parser(result); - while (parser.TryNextRow()) { - NJsonWriter::TBuf writer(NJsonWriter::HEM_UNSAFE, Options_.ResultOutput); - writer.SetWriteNanAsString(true); - NYdb::FormatResultRowJson(parser, result.GetColumnsMeta(), writer, NYdb::EBinaryStringEncoding::Unicode); - *Options_.ResultOutput << Endl; - } - break; - } - - case TRunnerOptions::EResultOutputFormat::FullJson: - resultSet.PrintJSON(*Options_.ResultOutput); - *Options_.ResultOutput << Endl; - break; - - case TRunnerOptions::EResultOutputFormat::FullProto: - TString resultSetString; - google::protobuf::TextFormat::Printer printer; - printer.SetSingleLineMode(false); - printer.SetUseUtf8StringEscaping(true); - printer.PrintToString(resultSet, &resultSetString); - *Options_.ResultOutput << resultSetString; - break; - } - } - void PrintScriptFinish(const TQueryMeta& meta, const TString& queryType) const { + if (Options_.YdbSettings.VerboseLevel < EVerbose::Info) { + return; + } Cout << CoutColors_.Cyan() << queryType << " request finished."; if (meta.TotalDuration) { Cout << " Total duration: " << meta.TotalDuration; @@ -422,9 +305,10 @@ class TKqpRunner::TImpl { private: TRunnerOptions Options_; + EVerbose VerboseLevel_; TYdbSetup YdbSetup_; - std::unique_ptr StatProcessor_; + TStatsPrinter StatsPrinter_; NColorizer::TColors CerrColors_; NColorizer::TColors CoutColors_; diff --git a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto new file mode 100644 index 000000000000..ee4a7c4162a2 --- /dev/null +++ b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto @@ -0,0 +1,22 @@ +syntax = "proto3"; + +package NKqpRun; + +message TStorageMeta { + message TTenant { + enum EType { + DEDICATED = 0; + SHARED = 1; + SERVERLESS = 2; + } + + EType Type = 1; + uint32 NodesCount = 2; + string SharedTenant = 3; // Only for serverless tenants + } + + uint64 StorageGeneration = 1; + uint64 StorageSize = 2; + string DomainName = 3; + map Tenants = 4; +} diff --git a/ydb/tests/tools/kqprun/src/proto/ya.make b/ydb/tests/tools/kqprun/src/proto/ya.make new file mode 100644 index 000000000000..8cbcbe198594 --- /dev/null +++ b/ydb/tests/tools/kqprun/src/proto/ya.make @@ -0,0 +1,9 @@ +PROTO_LIBRARY() + +ONLY_TAGS(CPP_PROTO) + +SRCS( + storage_meta.proto +) + +END() diff --git a/ydb/tests/tools/kqprun/src/ya.make b/ydb/tests/tools/kqprun/src/ya.make index 8b74f61204d0..ed5f5a619bd5 100644 --- a/ydb/tests/tools/kqprun/src/ya.make +++ b/ydb/tests/tools/kqprun/src/ya.make @@ -8,8 +8,13 @@ SRCS( PEERDIR( ydb/core/testlib + + ydb/tests/tools/kqprun/runlib + ydb/tests/tools/kqprun/src/proto ) +GENERATE_ENUM_SERIALIZATION(common.h) + YQL_LAST_ABI_VERSION() END() diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp index 97ab3a0b9b5f..8ff68914a656 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp +++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp @@ -2,15 +2,19 @@ #include +#include #include #include - #include #include #include + +#include + #include +using namespace NKikimrRun; namespace NKqpRun { @@ -65,7 +69,7 @@ class TStaticSecuredCredentialsFactory : public NYql::ISecuredServiceAccountCred class TSessionState { public: - explicit TSessionState(NActors::TTestActorRuntime* runtime, ui32 targetNodeIndex, const TString& database, const TString& traceId) + explicit TSessionState(NActors::TTestActorRuntime* runtime, ui32 targetNodeIndex, const TString& database, const TString& traceId, TYdbSetupSettings::EVerbose verboseLevel) : Runtime_(runtime) , TargetNodeIndex_(targetNodeIndex) { @@ -78,7 +82,8 @@ class TSessionState { auto closePromise = NThreading::NewPromise(); SessionHolderActor_ = Runtime_->Register(CreateSessionHolderActor(TCreateSessionRequest{ .Event = std::move(event), - .TargetNode = Runtime_->GetNodeId(targetNodeIndex) + .TargetNode = Runtime_->GetNodeId(targetNodeIndex), + .VerboseLevel = verboseLevel }, openPromise, closePromise)); SessionId_ = openPromise.GetFuture().GetValueSync(); @@ -130,6 +135,8 @@ void FillQueryMeta(TQueryMeta& meta, const NKikimrKqp::TQueryResponse& response) //// TYdbSetup::TImpl class TYdbSetup::TImpl { + using EVerbose = TYdbSetupSettings::EVerbose; + private: TAutoPtr CreateLogBackend() const { if (Settings_.LogOutputFile) { @@ -141,23 +148,7 @@ class TYdbSetup::TImpl { void SetLoggerSettings(NKikimr::Tests::TServerSettings& serverSettings) const { auto loggerInitializer = [this](NActors::TTestActorRuntime& runtime) { - if (Settings_.AppConfig.GetLogConfig().HasDefaultLevel()) { - auto priority = NActors::NLog::EPriority(Settings_.AppConfig.GetLogConfig().GetDefaultLevel()); - auto descriptor = NKikimrServices::EServiceKikimr_descriptor(); - for (int i = 0; i < descriptor->value_count(); ++i) { - runtime.SetLogPriority(static_cast(descriptor->value(i)->number()), priority); - } - } - - for (auto setting : Settings_.AppConfig.GetLogConfig().get_arr_entry()) { - NKikimrServices::EServiceKikimr service; - if (!NKikimrServices::EServiceKikimr_Parse(setting.GetComponent(), &service)) { - ythrow yexception() << "Invalid kikimr service name " << setting.GetComponent(); - } - - runtime.SetLogPriority(service, NActors::NLog::EPriority(setting.GetLevel())); - } - + InitLogSettings(Settings_.AppConfig.GetLogConfig(), runtime); runtime.SetLogBackendFactory([this]() { return CreateLogBackend(); }); }; @@ -176,16 +167,68 @@ class TYdbSetup::TImpl { serverSettings.SetFrFactory(functionRegistryFactory); } - void SetStorageSettings(NKikimr::Tests::TServerSettings& serverSettings) const { + void SetStorageSettings(NKikimr::Tests::TServerSettings& serverSettings) { + TFsPath diskPath; + if (Settings_.PDisksPath && *Settings_.PDisksPath != "-") { + diskPath = *Settings_.PDisksPath; + if (!diskPath) { + ythrow yexception() << "Storage directory path should not be empty"; + } + if (diskPath.IsRelative()) { + diskPath = TFsPath::Cwd() / diskPath; + } + diskPath.Fix(); + diskPath.MkDir(); + if (Settings_.VerboseLevel >= EVerbose::InitLogs) { + Cout << CoutColors_.Cyan() << "Setup storage by path: " << diskPath.GetPath() << CoutColors_.Default() << Endl; + } + } + + bool formatDisk = true; + if (diskPath) { + StorageMetaPath_ = TFsPath(diskPath).Child("kqprun_storage_meta.conf"); + if (StorageMetaPath_.Exists() && !Settings_.FormatStorage) { + if (!google::protobuf::TextFormat::ParseFromString(TFileInput(StorageMetaPath_.GetPath()).ReadAll(), &StorageMeta_)) { + ythrow yexception() << "Storage meta is corrupted, please use --format-storage"; + } + StorageMeta_.SetStorageGeneration(StorageMeta_.GetStorageGeneration() + 1); + formatDisk = false; + } + + if (Settings_.DiskSize && StorageMeta_.GetStorageSize() != *Settings_.DiskSize) { + if (!formatDisk) { + ythrow yexception() << "Cannot change disk size without formatting storage, current disk size " << NKikimr::NBlobDepot::FormatByteSize(StorageMeta_.GetStorageSize()) << ", please use --format-storage"; + } + StorageMeta_.SetStorageSize(*Settings_.DiskSize); + } else if (!StorageMeta_.GetStorageSize()) { + StorageMeta_.SetStorageSize(DEFAULT_STORAGE_SIZE); + } + + const TString& domainName = NKikimr::CanonizePath(Settings_.DomainName); + if (!StorageMeta_.GetDomainName()) { + StorageMeta_.SetDomainName(domainName); + } else if (StorageMeta_.GetDomainName() != domainName) { + ythrow yexception() << "Cannot change domain name without formatting storage, current name " << StorageMeta_.GetDomainName() << ", please use --format-storage"; + } + + UpdateStorageMeta(); + } + + TString storagePath = diskPath.GetPath(); + SlashFolderLocal(storagePath); + const NKikimr::NFake::TStorage storage = { - .UseDisk = Settings_.UseRealPDisks, + .UseDisk = !!Settings_.PDisksPath, .SectorSize = NKikimr::TTestStorageFactory::SECTOR_SIZE, - .ChunkSize = Settings_.UseRealPDisks ? NKikimr::TTestStorageFactory::CHUNK_SIZE : NKikimr::TTestStorageFactory::MEM_CHUNK_SIZE, - .DiskSize = Settings_.DiskSize + .ChunkSize = Settings_.PDisksPath ? NKikimr::TTestStorageFactory::CHUNK_SIZE : NKikimr::TTestStorageFactory::MEM_CHUNK_SIZE, + .DiskSize = Settings_.DiskSize ? *Settings_.DiskSize : 32_GB, + .FormatDisk = formatDisk, + .DiskPath = storagePath }; - serverSettings.SetEnableMockOnSingleNode(!Settings_.DisableDiskMock && !Settings_.UseRealPDisks); + serverSettings.SetEnableMockOnSingleNode(!Settings_.DisableDiskMock && !Settings_.PDisksPath); serverSettings.SetCustomDiskParams(storage); + serverSettings.SetStorageGeneration(StorageMeta_.GetStorageGeneration()); } NKikimr::Tests::TServerSettings GetServerSettings(ui32 grpcPort) { @@ -210,7 +253,7 @@ class TYdbSetup::TImpl { serverSettings.SetYtGateway(Settings_.YtGateway); serverSettings.S3ActorsFactory = NYql::NDq::CreateS3ActorsFactory(); serverSettings.SetInitializeFederatedQuerySetupFactory(true); - serverSettings.SetVerbose(false); + serverSettings.SetVerbose(Settings_.VerboseLevel >= EVerbose::InitLogs); SetLoggerSettings(serverSettings); SetFunctionRegistry(serverSettings); @@ -226,28 +269,54 @@ class TYdbSetup::TImpl { serverSettings.SetGrpcPort(grpcPort); } - if (!Settings_.SharedTenants.empty() || !Settings_.DedicatedTenants.empty()) { - serverSettings.SetDynamicNodeCount(Settings_.SharedTenants.size() + Settings_.DedicatedTenants.size()); - for (const TString& dedicatedTenant : Settings_.DedicatedTenants) { - serverSettings.AddStoragePoolType(dedicatedTenant); - } - for (const auto& sharedTenant : Settings_.SharedTenants) { - serverSettings.AddStoragePoolType(sharedTenant); + for (const auto& [tenantPath, tenantInfo] : StorageMeta_.GetTenants()) { + Settings_.Tenants.emplace(tenantPath, tenantInfo); + } + + ui32 dynNodesCount = 0; + for (const auto& [tenantPath, tenantInfo] : Settings_.Tenants) { + if (tenantInfo.GetType() != TStorageMeta::TTenant::SERVERLESS) { + serverSettings.AddStoragePoolType(tenantPath); + dynNodesCount += tenantInfo.GetNodesCount(); } } + serverSettings.SetDynamicNodeCount(dynNodesCount); return serverSettings; } - void CreateTenant(Ydb::Cms::CreateDatabaseRequest&& request, const TString& type) const { - const auto path = request.path(); - Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Creating " << type << " tenant " << path << "..." << CoutColors_.Default() << Endl; - Tenants_->CreateTenant(std::move(request)); + void CreateTenant(Ydb::Cms::CreateDatabaseRequest&& request, const TString& relativePath, const TString& type, TStorageMeta::TTenant tenantInfo) { + const auto absolutePath = request.path(); + const auto [it, inserted] = StorageMeta_.MutableTenants()->emplace(relativePath, tenantInfo); + if (inserted) { + if (Settings_.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Creating " << type << " tenant " << absolutePath << "..." << CoutColors_.Default() << Endl; + } + Tenants_->CreateTenant(std::move(request), tenantInfo.GetNodesCount()); + UpdateStorageMeta(); + } else { + if (it->second.GetType() != tenantInfo.GetType()) { + ythrow yexception() << "Can not change tenant " << absolutePath << " type without formatting storage, current type " << TStorageMeta::TTenant::EType_Name(it->second.GetType()) << ", please use --format-storage"; + } + if (it->second.GetSharedTenant() != tenantInfo.GetSharedTenant()) { + ythrow yexception() << "Can not change tenant " << absolutePath << " shared resources without formatting storage from '" << it->second.GetSharedTenant() << "', please use --format-storage"; + } + if (it->second.GetNodesCount() != tenantInfo.GetNodesCount()) { + it->second.SetNodesCount(tenantInfo.GetNodesCount()); + UpdateStorageMeta(); + } + if (Settings_.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Starting " << type << " tenant " << absolutePath << "..." << CoutColors_.Default() << Endl; + } + if (!request.has_serverless_resources()) { + Tenants_->Run(absolutePath, tenantInfo.GetNodesCount()); + } + } if (Settings_.MonitoringEnabled) { - ui32 nodeIndex = GetNodeIndexForDatabase(path); + ui32 nodeIndex = GetNodeIndexForDatabase(absolutePath); NActors::TActorId edgeActor = GetRuntime()->AllocateEdgeActor(nodeIndex); - GetRuntime()->Register(NKikimr::CreateBoardPublishActor(NKikimr::MakeEndpointsBoardPath(path), "", edgeActor, 0, true), nodeIndex, GetRuntime()->GetAppData(nodeIndex).UserPoolId); + GetRuntime()->Register(NKikimr::CreateBoardPublishActor(NKikimr::MakeEndpointsBoardPath(absolutePath), "", edgeActor, 0, true), nodeIndex, GetRuntime()->GetAppData(nodeIndex).UserPoolId); } } @@ -257,39 +326,50 @@ class TYdbSetup::TImpl { } void CreateTenants() { - for (const TString& dedicatedTenant : Settings_.DedicatedTenants) { + std::set sharedTenants; + std::map serverlessTenants; + for (const auto& [tenantPath, tenantInfo] : Settings_.Tenants) { Ydb::Cms::CreateDatabaseRequest request; - request.set_path(GetTenantPath(dedicatedTenant)); - AddTenantStoragePool(request.mutable_resources()->add_storage_units(), dedicatedTenant); - CreateTenant(std::move(request), "dedicated"); + request.set_path(GetTenantPath(tenantPath)); + + switch (tenantInfo.GetType()) { + case TStorageMeta::TTenant::DEDICATED: + AddTenantStoragePool(request.mutable_resources()->add_storage_units(), tenantPath); + CreateTenant(std::move(request), tenantPath, "dedicated", tenantInfo); + break; + + case TStorageMeta::TTenant::SHARED: + sharedTenants.emplace(tenantPath); + AddTenantStoragePool(request.mutable_shared_resources()->add_storage_units(), tenantPath); + CreateTenant(std::move(request), tenantPath, "shared", tenantInfo); + break; + + case TStorageMeta::TTenant::SERVERLESS: + serverlessTenants.emplace(tenantPath, tenantInfo); + break; + + default: + ythrow yexception() << "Unexpected tenant type: " << TStorageMeta::TTenant::EType_Name(tenantInfo.GetType()); + break; + } } - for (const TString& sharedTenant : Settings_.SharedTenants) { - Ydb::Cms::CreateDatabaseRequest request; - request.set_path(GetTenantPath(sharedTenant)); - AddTenantStoragePool(request.mutable_shared_resources()->add_storage_units(), sharedTenant); - CreateTenant(std::move(request), "shared"); - } + for (auto [tenantPath, tenantInfo] : serverlessTenants) { + if (!tenantInfo.GetSharedTenant()) { + if (sharedTenants.empty()) { + ythrow yexception() << "Can not create serverless tenant, there is no shared tenants, please use `--shared `"; + } + if (sharedTenants.size() > 1) { + ythrow yexception() << "Can not create serverless tenant, there is more than one shared tenant, please use `--serverless " << tenantPath << "@`"; + } + tenantInfo.SetSharedTenant(*sharedTenants.begin()); + } - ServerlessToShared_.reserve(Settings_.ServerlessTenants.size()); - for (const TString& serverlessTenant : Settings_.ServerlessTenants) { Ydb::Cms::CreateDatabaseRequest request; - if (serverlessTenant.Contains('@')) { - TStringBuf serverless; - TStringBuf shared; - TStringBuf(serverlessTenant).Split('@', serverless, shared); - - request.set_path(GetTenantPath(TString(serverless))); - request.mutable_serverless_resources()->set_shared_database_path(GetTenantPath(TString(shared))); - } else if (!Settings_.SharedTenants.empty()) { - request.set_path(GetTenantPath(serverlessTenant)); - request.mutable_serverless_resources()->set_shared_database_path(GetTenantPath(*Settings_.SharedTenants.begin())); - } else { - ythrow yexception() << "Can not create serverless tenant " << serverlessTenant << ", there is no shared tenants"; - } + request.set_path(GetTenantPath(tenantPath)); + request.mutable_serverless_resources()->set_shared_database_path(GetTenantPath(tenantInfo.GetSharedTenant())); ServerlessToShared_[request.path()] = request.serverless_resources().shared_database_path(); - - CreateTenant(std::move(request), "serverless"); + CreateTenant(std::move(request), tenantPath, "serverless", tenantInfo); } } @@ -315,30 +395,23 @@ class TYdbSetup::TImpl { return; } - bool found = false; - for (auto& entry : *Settings_.AppConfig.MutableLogConfig()->MutableEntry()) { - if (entry.GetComponent() == "KQP_YQL") { - entry.SetLevel(NActors::NLog::PRI_TRACE); - found = true; - break; - } - } - - if (!found) { - auto entry = Settings_.AppConfig.MutableLogConfig()->AddEntry(); - entry->SetComponent("KQP_YQL"); - entry->SetLevel(NActors::NLog::PRI_TRACE); - } - + ModifyLogPriorities({{NKikimrServices::EServiceKikimr::KQP_YQL, NActors::NLog::PRI_TRACE}}, *Settings_.AppConfig.MutableLogConfig()); NYql::NLog::InitLogger(NActors::CreateNullBackend()); } void WaitResourcesPublishing() const { auto promise = NThreading::NewPromise(); - GetRuntime()->Register(CreateResourcesWaiterActor(promise, Settings_.NodeCount), 0, GetRuntime()->GetAppData().SystemPoolId); + const TWaitResourcesSettings settings = { + .ExpectedNodeCount = static_cast(Settings_.NodeCount), + .HealthCheckLevel = Settings_.HealthCheckLevel, + .HealthCheckTimeout = Settings_.HealthCheckTimeout, + .VerboseLevel = Settings_.VerboseLevel, + .Database = NKikimr::CanonizePath(Settings_.DomainName) + }; + GetRuntime()->Register(CreateResourcesWaiterActor(promise, settings), 0, GetRuntime()->GetAppData().SystemPoolId); try { - promise.GetFuture().GetValue(Settings_.InitializationTimeout); + promise.GetFuture().GetValue(2 * Settings_.HealthCheckTimeout); } catch (...) { ythrow yexception() << "Failed to initialize all resources: " << CurrentExceptionMessage(); } @@ -355,13 +428,23 @@ class TYdbSetup::TImpl { InitializeServer(grpcPort); WaitResourcesPublishing(); - if (Settings_.MonitoringEnabled) { + if (Settings_.MonitoringEnabled && Settings_.VerboseLevel >= EVerbose::Info) { for (ui32 nodeIndex = 0; nodeIndex < Settings_.NodeCount; ++nodeIndex) { - Cout << CoutColors_.Cyan() << "Monitoring port" << (Settings_.NodeCount > 1 ? TStringBuilder() << " for node " << nodeIndex + 1 : TString()) << ": " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(nodeIndex) << Endl; + Cout << CoutColors_.Cyan() << "Monitoring port" << (Server_->StaticNodes() + Server_->DynamicNodes() > 1 ? TStringBuilder() << " for static node " << nodeIndex + 1 : TString()) << ": " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(nodeIndex) << Endl; } + const auto printTenantNodes = [this](const std::pair& tenantInfo) { + if (tenantInfo.second.GetType() == TStorageMeta::TTenant::SERVERLESS) { + return; + } + const auto& nodes = Tenants_->List(GetTenantPath(tenantInfo.first)); + for (auto it = nodes.rbegin(); it != nodes.rend(); ++it) { + Cout << CoutColors_.Cyan() << "Monitoring port for dynamic node " << *it + 1 << " [" << tenantInfo.first << "]: " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(*it) << Endl; + } + }; + std::for_each(Settings_.Tenants.rbegin(), Settings_.Tenants.rend(), std::bind(printTenantNodes, std::placeholders::_1)); } - if (Settings_.GrpcEnabled) { + if (Settings_.GrpcEnabled && Settings_.VerboseLevel >= EVerbose::Info) { Cout << CoutColors_.Cyan() << "Domain gRPC port: " << CoutColors_.Default() << grpcPort << Endl; } } @@ -516,7 +599,7 @@ class TYdbSetup::TImpl { if (Settings_.SameSession) { if (!SessionState_) { - SessionState_ = TSessionState(GetRuntime(), targetNodeIndex, database, query.TraceId); + SessionState_ = TSessionState(GetRuntime(), targetNodeIndex, database, query.TraceId, Settings_.VerboseLevel); } request->SetSessionId(SessionState_->GetSessionId()); } @@ -535,7 +618,8 @@ class TYdbSetup::TImpl { .Event = std::move(event), .TargetNode = GetRuntime()->GetNodeId(targetNodeIndex), .ResultRowsLimit = Settings_.AppConfig.GetQueryServiceConfig().GetScriptResultRowsLimit(), - .ResultSizeLimit = Settings_.AppConfig.GetQueryServiceConfig().GetScriptResultSizeLimit() + .ResultSizeLimit = Settings_.AppConfig.GetQueryServiceConfig().GetScriptResultSizeLimit(), + .QueryId = query.QueryId }; } @@ -544,12 +628,12 @@ class TYdbSetup::TImpl { } TString GetDatabasePath(const TString& database) const { - return NKikimr::CanonizePath(database ? database : Settings_.DomainName); + return NKikimr::CanonizePath(database ? database : GetDefaultDatabase()); } ui32 GetNodeIndexForDatabase(const TString& path) const { - auto canonizedPath = NKikimr::CanonizePath(path); - if (canonizedPath.empty() || canonizedPath == NKikimr::CanonizePath(Settings_.DomainName)) { + auto canonizedPath = NKikimr::CanonizePath(path ? path : GetDefaultDatabase()); + if (canonizedPath == NKikimr::CanonizePath(Settings_.DomainName)) { return RandomNumber(Settings_.NodeCount); } @@ -564,6 +648,27 @@ class TYdbSetup::TImpl { ythrow yexception() << "Unknown tenant '" << canonizedPath << "'"; } + TString GetDefaultDatabase() const { + if (StorageMeta_.TenantsSize() > 1) { + ythrow yexception() << "Can not choose default database, there is more than one tenants, please use `-D `"; + } + if (StorageMeta_.TenantsSize() == 1) { + return StorageMeta_.GetTenants().begin()->first; + } + return Settings_.DomainName; + } + + void UpdateStorageMeta() const { + if (StorageMetaPath_) { + TString storageMetaStr; + google::protobuf::TextFormat::PrintToString(StorageMeta_, &storageMetaStr); + + TFileOutput storageMetaOutput(StorageMetaPath_.GetPath()); + storageMetaOutput.Write(storageMetaStr); + storageMetaOutput.Finish(); + } + } + private: TYdbSetupSettings Settings_; NColorizer::TColors CoutColors_; @@ -576,35 +681,11 @@ class TYdbSetup::TImpl { std::unordered_map ServerlessToShared_; std::optional AsyncQueryRunnerActorId_; std::optional SessionState_; + TFsPath StorageMetaPath_; + NKqpRun::TStorageMeta StorageMeta_; }; -//// TRequestResult - -TRequestResult::TRequestResult() - : Status(Ydb::StatusIds::STATUS_CODE_UNSPECIFIED) -{} - -TRequestResult::TRequestResult(Ydb::StatusIds::StatusCode status, const NYql::TIssues& issues) - : Status(status) - , Issues(issues) -{} - -TRequestResult::TRequestResult(Ydb::StatusIds::StatusCode status, const google::protobuf::RepeatedPtrField& issues) - : Status(status) -{ - NYql::IssuesFromMessage(issues, Issues); -} - -bool TRequestResult::IsSuccess() const { - return Status == Ydb::StatusIds::SUCCESS; -} - -TString TRequestResult::ToString() const { - return TStringBuilder() << "Request finished with status: " << Status << "\nIssues:\n" << Issues.ToString() << "\n"; -} - - //// TYdbSetup TYdbSetup::TYdbSetup(const TYdbSetupSettings& settings) diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.h b/ydb/tests/tools/kqprun/src/ydb_setup.h index 393873b511bf..4edd849b9d26 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.h +++ b/ydb/tests/tools/kqprun/src/ydb_setup.h @@ -4,6 +4,7 @@ #include "actors.h" #include +#include namespace NKqpRun { @@ -30,23 +31,9 @@ struct TExecutionMeta : public TQueryMeta { }; -struct TRequestResult { - Ydb::StatusIds::StatusCode Status; - NYql::TIssues Issues; - - TRequestResult(); - - TRequestResult(Ydb::StatusIds::StatusCode status, const NYql::TIssues& issues); - - TRequestResult(Ydb::StatusIds::StatusCode status, const google::protobuf::RepeatedPtrField& issues); - - bool IsSuccess() const; - - TString ToString() const; -}; - - class TYdbSetup { + using TRequestResult = NKikimrRun::TRequestResult; + public: explicit TYdbSetup(const TYdbSetupSettings& settings); diff --git a/ydb/tests/tools/kqprun/ya.make b/ydb/tests/tools/kqprun/ya.make index 1ed5ff21cdd1..b39159f0ea40 100644 --- a/ydb/tests/tools/kqprun/ya.make +++ b/ydb/tests/tools/kqprun/ya.make @@ -1,5 +1,10 @@ PROGRAM(kqprun) +IF (PROFILE_MEMORY_ALLOCATIONS) + MESSAGE("Enabled profile memory allocations") + ALLOCATOR(LF_DBG) +ENDIF() + SRCS( kqprun.cpp ) @@ -11,6 +16,7 @@ PEERDIR( yt/yql/providers/yt/gateway/file yql/essentials/sql/pg + ydb/tests/tools/kqprun/runlib ydb/tests/tools/kqprun/src ) diff --git a/ydb/tests/tools/ya.make b/ydb/tests/tools/ya.make index b7586673bba3..bcae74c00dc2 100644 --- a/ydb/tests/tools/ya.make +++ b/ydb/tests/tools/ya.make @@ -2,6 +2,7 @@ RECURSE( canondata_sync datastreams_helpers fq_runner + fqrun idx_test kqprun mdb_mock diff --git a/yql/essentials/core/type_ann/type_ann_join.cpp b/yql/essentials/core/type_ann/type_ann_join.cpp index 1de6af8fec47..140b15eca299 100644 --- a/yql/essentials/core/type_ann/type_ann_join.cpp +++ b/yql/essentials/core/type_ann/type_ann_join.cpp @@ -475,7 +475,7 @@ namespace NTypeAnnImpl { IGraphTransformer::TStatus MapJoinCoreWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) { Y_UNUSED(output); - + if (!EnsureArgsCount(*input, 9, ctx.Expr)) { return IGraphTransformer::TStatus::Error; } diff --git a/yql/essentials/core/yql_join.cpp b/yql/essentials/core/yql_join.cpp index 044e49873e84..1206419d6965 100644 --- a/yql/essentials/core/yql_join.cpp +++ b/yql/essentials/core/yql_join.cpp @@ -319,8 +319,16 @@ namespace { } } else if (option.IsAtom("forceSortedMerge") || option.IsAtom("forceStreamLookup")) { - if (!EnsureTupleSize(*child, 1, ctx)) { - return IGraphTransformer::TStatus::Error; + if (option.IsAtom("forceStreamLookup")) { + if (child->ChildrenSize() % 2 == 0) { + ctx.AddError(TIssue(ctx.GetPosition(option.Pos()), TStringBuilder() << + "streamlookup() expects KEY VALUE... pairs")); + return IGraphTransformer::TStatus::Error; + } + } else { + if (!EnsureTupleSize(*child, 1, ctx)) { + return IGraphTransformer::TStatus::Error; + } } if (hasJoinStrategyHint) { ctx.AddError(TIssue(ctx.GetPosition(option.Pos()), TStringBuilder() << @@ -1351,9 +1359,14 @@ TEquiJoinLinkSettings GetEquiJoinLinkSettings(const TExprNode& linkSettings) { } result.ForceSortedMerge = HasSetting(linkSettings, "forceSortedMerge"); - - if (HasSetting(linkSettings, "forceStreamLookup")) { + + if (auto streamlookup = GetSetting(linkSettings, "forceStreamLookup")) { + YQL_ENSURE(result.JoinAlgoOptions.empty()); result.JoinAlgo = EJoinAlgoType::StreamLookupJoin; + auto size = streamlookup->ChildrenSize(); + for (decltype(size) i = 1; i < size; ++i) { + result.JoinAlgoOptions.push_back(TString(streamlookup->Child(i)->Content())); + } } if (HasSetting(linkSettings, "compact")) { diff --git a/yql/essentials/core/yql_join.h b/yql/essentials/core/yql_join.h index e4fe1e985b02..d7df5f988b85 100644 --- a/yql/essentials/core/yql_join.h +++ b/yql/essentials/core/yql_join.h @@ -148,6 +148,7 @@ struct TEquiJoinLinkSettings { // JOIN implementation may ignore this flags if SortedMerge strategy is not supported bool ForceSortedMerge = false; bool Compact = false; + TVector JoinAlgoOptions; }; TEquiJoinLinkSettings GetEquiJoinLinkSettings(const TExprNode& linkSettings); diff --git a/yql/essentials/minikql/computation/mkql_computation_node.h b/yql/essentials/minikql/computation/mkql_computation_node.h index da109d4234b3..85462824fd1e 100644 --- a/yql/essentials/minikql/computation/mkql_computation_node.h +++ b/yql/essentials/minikql/computation/mkql_computation_node.h @@ -77,6 +77,7 @@ struct TComputationMutables { std::vector SerializableValues; // Indices of values that need to be saved in IComputationGraph::SaveGraphState() and restored in IComputationGraph::LoadGraphState(). ui32 CurWideFieldsIndex = 0U; std::vector WideFieldInitialize; + std::vector CachedValues; // Indices of values that holds temporary cached data and unreachable by dependencies void DeferWideFieldsInit(ui32 count, std::set used) { Y_DEBUG_ABORT_UNLESS(AllOf(used, [count](ui32 i) { return i < count; })); @@ -251,7 +252,8 @@ class IComputationGraph { virtual IComputationExternalNode* GetEntryPoint(size_t index, bool require) = 0; virtual const TArrowKernelsTopology* GetKernelsTopology() = 0; virtual const TComputationNodePtrDeque& GetNodes() const = 0; - virtual void Invalidate() = 0; + virtual void Invalidate() = 0; // Invalidate all mutable values in graph (may lead to udf recreation) + virtual void InvalidateCaches() = 0; // Invalidate only cached values virtual TMemoryUsageInfo& GetMemInfo() const = 0; virtual const THolderFactory& GetHolderFactory() const = 0; virtual ITerminator* GetTerminator() const = 0; diff --git a/yql/essentials/minikql/computation/mkql_computation_node_graph.cpp b/yql/essentials/minikql/computation/mkql_computation_node_graph.cpp index 67e3546a1c28..f9209ad23971 100644 --- a/yql/essentials/minikql/computation/mkql_computation_node_graph.cpp +++ b/yql/essentials/minikql/computation/mkql_computation_node_graph.cpp @@ -696,6 +696,12 @@ class TComputationGraph final : public IComputationGraph { std::fill_n(Ctx->MutableValues.get(), PatternNodes->GetMutables().CurValueIndex, NUdf::TUnboxedValue(NUdf::TUnboxedValuePod::Invalid())); } + void InvalidateCaches() override { + for (const auto cachedIndex : Ctx->Mutables.CachedValues) { + Ctx->MutableValues[cachedIndex] = NUdf::TUnboxedValuePod::Invalid(); + } + } + const TComputationNodePtrDeque& GetNodes() const override { return PatternNodes->GetNodes(); } diff --git a/yql/essentials/minikql/computation/mkql_computation_node_holders_codegen.cpp b/yql/essentials/minikql/computation/mkql_computation_node_holders_codegen.cpp index 2d8ea64c10a1..3fe7772ba9d1 100644 --- a/yql/essentials/minikql/computation/mkql_computation_node_holders_codegen.cpp +++ b/yql/essentials/minikql/computation/mkql_computation_node_holders_codegen.cpp @@ -10,6 +10,7 @@ TContainerCacheOnContext::TContainerCacheOnContext(TComputationMutables& mutable : Index(mutables.CurValueIndex++) { ++++mutables.CurValueIndex; + mutables.CachedValues.insert(mutables.CachedValues.end(), {Index, Index + 1, Index + 2}); } NUdf::TUnboxedValuePod TContainerCacheOnContext::NewArray(TComputationContext& ctx, ui64 size, NUdf::TUnboxedValue*& items) const { diff --git a/yql/essentials/minikql/computation/mkql_computation_node_impl.cpp b/yql/essentials/minikql/computation/mkql_computation_node_impl.cpp index 8a0cbf65ba98..49f686dd9e0c 100644 --- a/yql/essentials/minikql/computation/mkql_computation_node_impl.cpp +++ b/yql/essentials/minikql/computation/mkql_computation_node_impl.cpp @@ -7,7 +7,7 @@ namespace NMiniKQL { void ThrowNotSupportedImplForClass(const TString& className, const char *func) { THROW yexception() << "Unsupported access to '" << func << "' method of: " << className; -} +} template void TRefCountedComputationNode::Ref() { @@ -100,7 +100,7 @@ Y_NO_INLINE void TStatefulComputationNodeBase::AddDependenceImpl(const IComputat Dependencies.emplace_back(node); } -Y_NO_INLINE void TStatefulComputationNodeBase::CollectDependentIndexesImpl(const IComputationNode* self, const IComputationNode* owner, +Y_NO_INLINE void TStatefulComputationNodeBase::CollectDependentIndexesImpl(const IComputationNode* self, const IComputationNode* owner, IComputationNode::TIndexesMap& dependencies, bool stateless) const { if (self == owner) return; @@ -188,7 +188,7 @@ Y_NO_INLINE TStatefulFlowComputationNodeBase::TStatefulFlowComputationNodeBase(u , StateKind(stateKind) {} -Y_NO_INLINE void TStatefulFlowComputationNodeBase::CollectDependentIndexesImpl(const IComputationNode* self, const IComputationNode* owner, +Y_NO_INLINE void TStatefulFlowComputationNodeBase::CollectDependentIndexesImpl(const IComputationNode* self, const IComputationNode* owner, IComputationNode::TIndexesMap& dependencies, const IComputationNode* dependence) const { if (self == owner) return; @@ -205,7 +205,7 @@ Y_NO_INLINE TPairStateFlowComputationNodeBase::TPairStateFlowComputationNodeBase , SecondKind(secondKind) {} -Y_NO_INLINE void TPairStateFlowComputationNodeBase::CollectDependentIndexesImpl(const IComputationNode* self, const IComputationNode* owner, +Y_NO_INLINE void TPairStateFlowComputationNodeBase::CollectDependentIndexesImpl(const IComputationNode* self, const IComputationNode* owner, IComputationNode::TIndexesMap& dependencies, const IComputationNode* dependence) const { if (self == owner) return; @@ -221,7 +221,7 @@ Y_NO_INLINE ui32 TStatelessWideFlowComputationNodeBase::GetIndexImpl() const { THROW yexception() << "Failed to get stateless node index."; } -Y_NO_INLINE void TStatelessWideFlowComputationNodeBase::CollectDependentIndexesImpl(const IComputationNode* self, const IComputationNode* owner, +Y_NO_INLINE void TStatelessWideFlowComputationNodeBase::CollectDependentIndexesImpl(const IComputationNode* self, const IComputationNode* owner, IComputationNode::TIndexesMap& dependencies, const IComputationNode* dependence) const { if (self == owner) return; @@ -263,7 +263,7 @@ Y_NO_INLINE TPairStateWideFlowComputationNodeBase::TPairStateWideFlowComputation {} Y_NO_INLINE void TPairStateWideFlowComputationNodeBase::CollectDependentIndexesImpl( - const IComputationNode* self, const IComputationNode* owner, + const IComputationNode* self, const IComputationNode* owner, IComputationNode::TIndexesMap& dependencies, const IComputationNode* dependence) const { if (self == owner) return; @@ -308,7 +308,9 @@ void TExternalComputationNode::CollectDependentIndexes(const IComputationNode*, TExternalComputationNode::TExternalComputationNode(TComputationMutables& mutables, EValueRepresentation kind) : TStatefulComputationNode(mutables, kind) -{} +{ + mutables.CachedValues.push_back(ValueIndex); +} NUdf::TUnboxedValue TExternalComputationNode::GetValue(TComputationContext& ctx) const { return Getter ? Getter(ctx) : ValueRef(ctx); diff --git a/yql/essentials/providers/common/proto/gateways_config.proto b/yql/essentials/providers/common/proto/gateways_config.proto index 875ccaa22f66..977ab2e7b9ef 100644 --- a/yql/essentials/providers/common/proto/gateways_config.proto +++ b/yql/essentials/providers/common/proto/gateways_config.proto @@ -2,8 +2,6 @@ package NYql; option java_package = "ru.yandex.yql.proto"; import "yql/essentials/protos/clickhouse.proto"; -import "ydb/library/yql/providers/generic/connector/api/common/data_source.proto"; -import "ydb/library/yql/providers/generic/connector/api/common/endpoint.proto"; /////////////////////////////// common /////////////////////////////// @@ -552,29 +550,140 @@ message TDbToolConfig { /////////// Generic gateway for the external data sources //////////// +// TGenericEndpoint represents the network address of a generic data source instance +message TGenericEndpoint { + optional string host = 1; + optional uint32 port = 2; +} + + +// TGenericCredentials represents various ways of user authentication in the data source instance +message TGenericCredentials { + message TBasic { + optional string username = 1; + optional string password = 2; + } + + message TToken { + optional string type = 1; + optional string value = 2; + } + + oneof payload { + TBasic basic = 1; + TToken token = 2; + } +} + +// EGenericDataSourceKind enumerates the external data sources +// supported by the federated query system +enum EGenericDataSourceKind { + DATA_SOURCE_KIND_UNSPECIFIED = 0; + CLICKHOUSE = 1; + POSTGRESQL = 2; + S3 = 3; + YDB = 4; + MYSQL = 5; + MS_SQL_SERVER = 6; + GREENPLUM = 7; + ORACLE = 8; + LOGGING = 9; +} + +// EGenericProtocol generalizes various kinds of network protocols supported by different databases. +enum EGenericProtocol { + PROTOCOL_UNSPECIFIED = 0; + NATIVE = 1; // CLICKHOUSE, POSTGRESQL + HTTP = 2; // CLICKHOUSE, S3 +} + +// TPostgreSQLDataSourceOptions represents settings specific to PostgreSQL +message TPostgreSQLDataSourceOptions { + // PostgreSQL schema + optional string schema = 1; +} + +// TClickhouseDataSourceOptions represents settings specific to Clickhouse +message TClickhouseDataSourceOptions { +} + +// TS3DataSourceOptions represents settings specific to S3 (Simple Storage Service) +message TS3DataSourceOptions { + // the region where data is stored + optional string region = 1; + // the bucket the object belongs to + optional string bucket = 2; +} + +// TGreenplumDataSourceOptions represents settings specific to Greenplum +message TGreenplumDataSourceOptions { + // Greenplum schema + optional string schema = 1; +} + +// TOracleDataSourceOptions represents settings specific to Oracle +message TOracleDataSourceOptions { + // Oracle service_name - alias to SID of oracle INSTANCE, or SID, or PDB. + // More about connection options in Oracle docs: + // https://docs.oracle.com/en/database/other-databases/essbase/21/essoa/connection-string-formats.html + optional string service_name = 1; +} + +// TLoggingDataSourceOptions represents settings specific to Logging +message TLoggingDataSourceOptions { + optional string folder_id = 1; +} + +// TGenericDataSourceInstance helps to identify the instance of a data source to redirect request to. +message TGenericDataSourceInstance { + // Data source kind + optional EGenericDataSourceKind kind = 1; + // Network address + optional TGenericEndpoint endpoint = 2; + // Database name + optional string database = 3; + // Credentials to access database + optional TGenericCredentials credentials = 4; + // If true, Connector server will use secure connections to access remote data sources. + // Certificates will be obtained from the standard system paths. + optional bool use_tls = 5; + // Allows to specify network protocol that should be used between + // during the connection between Connector and the remote data source + optional EGenericProtocol protocol = 6; + // Options specific to various data sources + oneof options { + TPostgreSQLDataSourceOptions pg_options = 7; + TClickhouseDataSourceOptions ch_options = 8; + TS3DataSourceOptions s3_options = 9; + TGreenplumDataSourceOptions gp_options = 10; + TOracleDataSourceOptions oracle_options = 11; + TLoggingDataSourceOptions logging_options = 12; + } +} + message TGenericClusterConfig { // Cluster name optional string Name = 1; // Data source kind - optional NYql.NConnector.NApi.EDataSourceKind Kind = 8; + optional EGenericDataSourceKind Kind = 8; // Location represents the network address of a data source instance we want to connect oneof Location { // Endpoint must be used for on-premise deployments. - NYql.NConnector.NApi.TEndpoint Endpoint = 9; + TGenericEndpoint Endpoint = 9; // DatabaseId must be used when the data source is deployed in cloud. // Data source FQDN and port will be resolved by MDB service. string DatabaseId = 4; } // Credentials used to access data source instance - optional NYql.NConnector.NApi.TCredentials Credentials = 10; + optional TGenericCredentials Credentials = 10; // Credentials used to access managed databases APIs. // When working with external data source instances deployed in clouds, - // one should either set (ServiceAccountId, ServiceAccountIdSignature) pair - // that will be resolved into IAM Token via Token Accessor, + // one should either set (ServiceAccountId, ServiceAccountIdSignature) pair + // that will be resolved into IAM Token via Token Accessor, // or provide IAM Token directly. optional string ServiceAccountId = 6; optional string ServiceAccountIdSignature = 7; @@ -588,7 +697,7 @@ message TGenericClusterConfig { optional string DatabaseName = 13; // Transport protocol used to establish a network connection with database - optional NYql.NConnector.NApi.EProtocol Protocol = 14; + optional EGenericProtocol Protocol = 14; // Data source options specific to various data sources map DataSourceOptions = 15; @@ -601,7 +710,7 @@ message TGenericConnectorConfig { // 1. Set address statically via `Endpoint.Host` and `Endpoint.Port`; // 2. Ask YDB to set `Endpoint.Port` to the value of expression `./ydbd --ic-port + OffsetFromIcPort`, // while Connector's hostname will still be taken from `Endpoint.Host` (with 'localhost' as a default value). - optional NYql.NConnector.NApi.TEndpoint Endpoint = 3; + optional TGenericEndpoint Endpoint = 3; optional uint32 OffsetFromIcPort = 6; // If true, Connector GRPC Client will use TLS encryption. @@ -638,7 +747,7 @@ message TGenericGatewayConfig { message TDbResolverConfig { // Ydb / Yds MVP endpoint. - // Expected format: + // Expected format: // [http|https]://host:port/ydbc/cloud-prod/ optional string YdbMvpEndpoint = 2; } diff --git a/yql/essentials/providers/common/proto/ya.make b/yql/essentials/providers/common/proto/ya.make index 6ce6a6c99cac..eefd4c6b72b7 100644 --- a/yql/essentials/providers/common/proto/ya.make +++ b/yql/essentials/providers/common/proto/ya.make @@ -7,7 +7,6 @@ SRCS( PEERDIR( yql/essentials/protos - ydb/library/yql/providers/generic/connector/api/common ) IF (NOT PY_PROTOS_FOR) diff --git a/yql/essentials/public/purecalc/common/interface.h b/yql/essentials/public/purecalc/common/interface.h index 6e56c9aa3f9b..4096195bd35f 100644 --- a/yql/essentials/public/purecalc/common/interface.h +++ b/yql/essentials/public/purecalc/common/interface.h @@ -684,6 +684,11 @@ namespace NYql { * Get time provider */ virtual ITimeProvider* GetTimeProvider() const = 0; + + /** + * Release all input data from worker state + */ + virtual void Invalidate() = 0; }; /** diff --git a/yql/essentials/public/purecalc/common/worker.cpp b/yql/essentials/public/purecalc/common/worker.cpp index f670458c728c..a61e206d0fc0 100644 --- a/yql/essentials/public/purecalc/common/worker.cpp +++ b/yql/essentials/public/purecalc/common/worker.cpp @@ -340,6 +340,17 @@ void TWorker::Release() { } } +template +void TWorker::Invalidate() { + auto& ctx = Graph_.ComputationGraph_->GetContext(); + for (const auto* selfNode : Graph_.SelfNodes_) { + if (selfNode) { + selfNode->InvalidateValue(ctx); + } + } + Graph_.ComputationGraph_->InvalidateCaches(); +} + TPullStreamWorker::~TPullStreamWorker() { auto guard = Guard(GetScopedAlloc()); Output_.Clear(); diff --git a/yql/essentials/public/purecalc/common/worker.h b/yql/essentials/public/purecalc/common/worker.h index 07b8dfa2e79e..39f381c9eaf9 100644 --- a/yql/essentials/public/purecalc/common/worker.h +++ b/yql/essentials/public/purecalc/common/worker.h @@ -104,6 +104,7 @@ namespace NYql { const TString& GetLLVMSettings() const override; ui64 GetNativeYtTypeFlags() const override; ITimeProvider* GetTimeProvider() const override; + void Invalidate() override; protected: void Release() override; }; diff --git a/yql/essentials/public/purecalc/ut/test_mixed_allocators.cpp b/yql/essentials/public/purecalc/ut/test_mixed_allocators.cpp index 797f3c5b5125..2932538f78f4 100644 --- a/yql/essentials/public/purecalc/ut/test_mixed_allocators.cpp +++ b/yql/essentials/public/purecalc/ut/test_mixed_allocators.cpp @@ -51,7 +51,7 @@ namespace { // Clear graph after each object because // values allocated on another allocator and should be released - Worker_->GetGraph().Invalidate(); + Worker_->Invalidate(); } } diff --git a/yql/essentials/sql/v1/join.cpp b/yql/essentials/sql/v1/join.cpp index de789569c766..11044d559de0 100644 --- a/yql/essentials/sql/v1/join.cpp +++ b/yql/essentials/sql/v1/join.cpp @@ -502,7 +502,11 @@ class TEquiJoin: public TJoinBase { if (TJoinLinkSettings::EStrategy::SortedMerge == descr.LinkSettings.Strategy) { linkOptions = L(linkOptions, Q(Y(Q("forceSortedMerge")))); } else if (TJoinLinkSettings::EStrategy::StreamLookup == descr.LinkSettings.Strategy) { - linkOptions = L(linkOptions, Q(Y(Q("forceStreamLookup")))); + auto streamlookup = Y(Q("forceStreamLookup")); + for (auto&& option: descr.LinkSettings.Values) { + streamlookup = L(streamlookup, Q(option)); + } + linkOptions = L(linkOptions, Q(streamlookup)); } else if (TJoinLinkSettings::EStrategy::ForceMap == descr.LinkSettings.Strategy) { linkOptions = L(linkOptions, Q(Y(Q("join_algo"), Q("MapJoin")))); } else if (TJoinLinkSettings::EStrategy::ForceGrace == descr.LinkSettings.Strategy) { diff --git a/yql/essentials/sql/v1/source.h b/yql/essentials/sql/v1/source.h index ba904d6c21e8..20a7da80f908 100644 --- a/yql/essentials/sql/v1/source.h +++ b/yql/essentials/sql/v1/source.h @@ -172,6 +172,7 @@ namespace NSQLTranslationV1 { ForceGrace }; EStrategy Strategy = EStrategy::Default; + TVector Values; bool Compact = false; }; diff --git a/yql/essentials/sql/v1/sql_select.cpp b/yql/essentials/sql/v1/sql_select.cpp index 4a06f8e51b14..c27aa8f126ab 100644 --- a/yql/essentials/sql/v1/sql_select.cpp +++ b/yql/essentials/sql/v1/sql_select.cpp @@ -43,6 +43,7 @@ bool CollectJoinLinkSettings(TPosition pos, TJoinLinkSettings& linkSettings, TCo if (TJoinLinkSettings::EStrategy::Default == linkSettings.Strategy) { linkSettings.Strategy = newStrategy; + linkSettings.Values = hint.Values; } else if (newStrategy == linkSettings.Strategy) { ctx.Error() << "Duplicate join strategy hint"; return false; diff --git a/yt/yql/providers/yt/lib/lambda_builder/lambda_builder.cpp b/yt/yql/providers/yt/lib/lambda_builder/lambda_builder.cpp index dd6a8bbb205a..a2f7bf601f12 100644 --- a/yt/yql/providers/yt/lib/lambda_builder/lambda_builder.cpp +++ b/yt/yql/providers/yt/lib/lambda_builder/lambda_builder.cpp @@ -130,6 +130,9 @@ class TComputationGraphProxy: public IComputationGraph { void Invalidate() final { return Graph->Invalidate(); } + void InvalidateCaches() final { + return Graph->InvalidateCaches(); + } TMemoryUsageInfo& GetMemInfo() const final { return Graph->GetMemInfo(); }