diff options
31 files changed, 1747 insertions, 0 deletions
diff --git a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp index aeb6c03a13e..dc98ca5978e 100644 --- a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp +++ b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp @@ -229,6 +229,8 @@ #include <ydb/services/ext_index/common/config.h> #include <ydb/services/ext_index/service/executor.h> +#include <ydb/services/udf_store/service.h> + #include <ydb/library/actors/protos/services_common.pb.h> #include <ydb/library/actors/core/actor_bootstrapped.h> @@ -3387,4 +3389,23 @@ void TNbsServiceInitializer::InitializeServices(NActors::TActorSystemSetup *setu #endif +TUdfStoreInitializer::TUdfStoreInitializer(const TKikimrRunConfig& runConfig, TIntrusivePtr<NMiniKQL::IMutableFunctionRegistry> functionRegistry) + : IKikimrServicesInitializer(runConfig) + , FunctionRegistry(functionRegistry) { +} + +void TUdfStoreInitializer::InitializeServices(NActors::TActorSystemSetup* setup, const NKikimr::TAppData* appData) { + if (!Config.HasUdfStoreConfig()) { + return; + } + auto service = NUdfStore::CreateService(Config.GetUdfStoreConfig(), FunctionRegistry); + if (!service) { + return; + } + setup->LocalServices.push_back(std::make_pair( + NUdfStore::MakeServiceId(NodeId), + TActorSetupCmd(service, TMailboxType::HTSwap, appData->UserPoolId))); +} + + } // namespace NKikimr::NKikimrServicesInitializers diff --git a/ydb/core/driver_lib/run/kikimr_services_initializers.h b/ydb/core/driver_lib/run/kikimr_services_initializers.h index 77fe78000f0..b3f2d0abdd3 100644 --- a/ydb/core/driver_lib/run/kikimr_services_initializers.h +++ b/ydb/core/driver_lib/run/kikimr_services_initializers.h @@ -690,5 +690,12 @@ public: }; #endif +class TUdfStoreInitializer: public IKikimrServicesInitializer { + TIntrusivePtr<NMiniKQL::IMutableFunctionRegistry> FunctionRegistry; +public: + TUdfStoreInitializer(const TKikimrRunConfig& runConfig, TIntrusivePtr<NMiniKQL::IMutableFunctionRegistry> functionRegistry); + void InitializeServices(NActors::TActorSystemSetup* setup, const NKikimr::TAppData* appData) override; +}; + } // namespace NKikimrServicesInitializers } // namespace NKikimr diff --git a/ydb/core/driver_lib/run/run.cpp b/ydb/core/driver_lib/run/run.cpp index 43d7fc4e9f3..d9edf5ecec4 100644 --- a/ydb/core/driver_lib/run/run.cpp +++ b/ydb/core/driver_lib/run/run.cpp @@ -2259,6 +2259,10 @@ TIntrusivePtr<TServiceInitializersList> TKikimrRunner::CreateServiceInitializers } #endif + if (serviceMask.EnableUdfStore) { + sil->AddServiceInitializer(new TUdfStoreInitializer(runConfig, FunctionRegistry)); + } + return sil; } diff --git a/ydb/core/driver_lib/run/service_mask.h b/ydb/core/driver_lib/run/service_mask.h index dfac4f22cfa..0db11ca100a 100644 --- a/ydb/core/driver_lib/run/service_mask.h +++ b/ydb/core/driver_lib/run/service_mask.h @@ -85,6 +85,7 @@ union TBasicKikimrServicesMask { bool EnableOverloadManager : 1; bool EnableCountersInfoProvider : 1; bool EnableNBSService : 1; + bool EnableUdfStore : 1; }; struct { diff --git a/ydb/core/driver_lib/run/ya.make b/ydb/core/driver_lib/run/ya.make index 0dc10bfd3a3..990e5c38515 100644 --- a/ydb/core/driver_lib/run/ya.make +++ b/ydb/core/driver_lib/run/ya.make @@ -174,6 +174,7 @@ PEERDIR( ydb/services/maintenance ydb/services/metadata ydb/services/metadata/ds_table + ydb/services/udf_store ydb/services/monitoring ydb/services/persqueue_cluster_discovery ydb/services/persqueue_v1 diff --git a/ydb/core/grpc_services/rpc_keyvalue.cpp b/ydb/core/grpc_services/rpc_keyvalue.cpp index 9c6cd8e4410..089004b5431 100644 --- a/ydb/core/grpc_services/rpc_keyvalue.cpp +++ b/ydb/core/grpc_services/rpc_keyvalue.cpp @@ -1527,4 +1527,9 @@ void DoGetStorageChannelStatusKeyValueV2(std::unique_ptr<IRequestNoOpCtx> p, con TActivationContext::AsActorContext().Register(new TGetStorageChannelStatusRequest<false>(p.release())); } +template<> +IActor* TEvCreateVolumeKeyValueRequest::CreateRpcActor(NKikimr::NGRpcService::IRequestOpCtx* msg) { + return new TCreateVolumeRequest(msg); +} + } // namespace NKikimr::NGRpcService diff --git a/ydb/core/protos/config.proto b/ydb/core/protos/config.proto index 7b68d7a174f..77df843e45a 100644 --- a/ydb/core/protos/config.proto +++ b/ydb/core/protos/config.proto @@ -3255,6 +3255,13 @@ message TOpaqueConfig { // Empty } +message TUdfStoreConfig { + optional bool Enabled = 1 [default = false]; + optional string KvStorageMedia = 2 [default = "ssd"]; + optional bool EnableUnsafeNativeUdf = 3 [default = false]; + optional string UnsafeNativeUdfDir = 4; +} + message TAppConfig { option (NMarkers.Root) = true; optional TActorSystemConfig ActorSystemConfig = 1; @@ -3345,6 +3352,7 @@ message TAppConfig { optional TGroupedMemoryLimiterConfig CompGroupedMemoryLimiterConfig = 92; optional TBlockstoreConfig BlockstoreConfig = 93; optional TOpaqueConfig PrivateDatabaseConfig = 94 [(NMarkers.AllowInDatabaseConfig) = true, (NMarkers.OpaqueConfig) = true]; + optional TUdfStoreConfig UdfStoreConfig = 95; repeated TNamedConfig NamedConfigs = 100; optional string ClusterYamlConfig = 101; diff --git a/ydb/services/metadata/request/common.h b/ydb/services/metadata/request/common.h index c3bd022542f..8a7af055ed8 100644 --- a/ydb/services/metadata/request/common.h +++ b/ydb/services/metadata/request/common.h @@ -7,6 +7,7 @@ #include <ydb/library/accessor/accessor.h> #include <ydb/public/api/protos/ydb_table.pb.h> #include <ydb/public/api/protos/ydb_scheme.pb.h> +#include <ydb/public/api/protos/ydb_keyvalue.pb.h> #include <ydb/core/base/events.h> namespace NKikimr::NMetadata::NRequest { @@ -52,6 +53,10 @@ enum EEvents { EvDeleteSessionInternalResponse, EvDeleteSessionResponse, + EvCreateKvVolumeRequest, + EvCreateKvVolumeInternalResponse, + EvCreateKvVolumeResponse, + EvEnd }; @@ -94,6 +99,8 @@ using TDialogCreateSession = TDialogPolicyImpl<Ydb::Table::CreateSessionRequest, EEvents::EvCreateSessionRequest, EEvents::EvCreateSessionInternalResponse, EEvents::EvCreateSessionResponse>; using TDialogDeleteSession = TDialogPolicyImpl<Ydb::Table::DeleteSessionRequest, Ydb::Table::DeleteSessionResponse, EEvents::EvDeleteSessionRequest, EEvents::EvDeleteSessionInternalResponse, EEvents::EvDeleteSessionResponse>; +using TDialogCreateKvVolume = TDialogPolicyImpl<Ydb::KeyValue::CreateVolumeRequest, Ydb::KeyValue::CreateVolumeResponse, + EEvents::EvCreateKvVolumeRequest, EEvents::EvCreateKvVolumeInternalResponse, EEvents::EvCreateKvVolumeResponse>; template <ui32 evResult = EEvents::EvGeneralYQLResponse> using TCustomDialogYQLRequest = TDialogPolicyImpl<Ydb::Table::ExecuteDataQueryRequest, Ydb::Table::ExecuteDataQueryResponse, @@ -140,6 +147,15 @@ public: } }; +template <> +class TOperatorChecker<Ydb::KeyValue::CreateVolumeResponse> { +public: + static bool IsSuccess(const Ydb::KeyValue::CreateVolumeResponse& r) { + return r.operation().status() == Ydb::StatusIds::SUCCESS || + r.operation().status() == Ydb::StatusIds::ALREADY_EXISTS; + } +}; + template <class TDialogPolicy> class TEvRequestResult: public NActors::TEventLocal<TEvRequestResult<TDialogPolicy>, TDialogPolicy::EvResult> { private: diff --git a/ydb/services/metadata/request/ya.make b/ydb/services/metadata/request/ya.make index 81c695bc050..6b63e70d4f2 100644 --- a/ydb/services/metadata/request/ya.make +++ b/ydb/services/metadata/request/ya.make @@ -13,6 +13,7 @@ PEERDIR( ydb/core/grpc_services/local_rpc ydb/core/grpc_services/base ydb/core/grpc_services + ydb/public/api/protos ) END() diff --git a/ydb/services/udf_store/events.h b/ydb/services/udf_store/events.h new file mode 100644 index 00000000000..7f95d59de6d --- /dev/null +++ b/ydb/services/udf_store/events.h @@ -0,0 +1,44 @@ +#pragma once + +#include <ydb/library/actors/core/event_local.h> +#include <ydb/library/actors/core/events.h> + +namespace NKikimr::NUdfStore { + +enum EEv { + EvStoreInitialized = EventSpaceBegin(NActors::TEvents::ES_PRIVATE), + EvStoreInitFailed, + EvReadBodyResponse, + EvEnd +}; + +// Sent to the parent actor when the meta table and KV volume are ready. +struct TEvStoreInitialized : public NActors::TEventLocal<TEvStoreInitialized, EvStoreInitialized> { + TEvStoreInitialized(const TString& kvVolumePath) + : KvVolumePath(kvVolumePath) + {} + TString KvVolumePath; +}; + +// Sent to the parent actor when infrastructure initialization fails. +struct TEvStoreInitFailed : public NActors::TEventLocal<TEvStoreInitFailed, EvStoreInitFailed> { + explicit TEvStoreInitFailed(TString errorMessage) + : ErrorMessage(std::move(errorMessage)) + {} + TString ErrorMessage; +}; + + +struct TEvReadBodyResponse : public NActors::TEventLocal<TEvReadBodyResponse, EvReadBodyResponse> { + bool Success; + TString Name; + TString ErrorMessage; + + TEvReadBodyResponse(bool success, const TString& name, const TString& errorMessage = {}) + : Success(success) + , Name(name) + , ErrorMessage(errorMessage) + {} +}; + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/kv_body_store.cpp b/ydb/services/udf_store/kv_body_store.cpp new file mode 100644 index 00000000000..cf1c634e9b3 --- /dev/null +++ b/ydb/services/udf_store/kv_body_store.cpp @@ -0,0 +1,295 @@ +#include "kv_body_store.h" + +#include <ydb/library/services/services.pb.h> +#include <ydb/library/actors/core/log.h> + +#include <util/folder/dirut.h> +#include <util/folder/path.h> +#include <util/system/fs.h> + +namespace NKikimr::NUdfStore { + +// +// TKvBodyReadActor +// + +void TKvBodyReadActor::Bootstrap() { + if (ExpectedSize == 0) { + ReplyError(TStringBuilder() << "UDF '" << Md5Key << "' has zero size in metadata"); + return; + } + Become(&TKvBodyReadActor::StateResolve); + SendNavigateRequest(); +} + +void TKvBodyReadActor::SendNavigateRequest() { + auto req = MakeHolder<NSchemeCache::TSchemeCacheNavigate>(); + auto& entry = req->ResultSet.emplace_back(); + entry.Path = SplitPath(VolumePath); + entry.RequestType = NSchemeCache::TSchemeCacheNavigate::TEntry::ERequestType::ByPath; + entry.ShowPrivatePath = true; + entry.SyncVersion = false; + Send(MakeSchemeCacheID(), new TEvTxProxySchemeCache::TEvNavigateKeySet(req.Release())); +} + +void TKvBodyReadActor::HandleNavigateResult(TEvTxProxySchemeCache::TEvNavigateKeySetResult::TPtr& ev) { + NSchemeCache::TSchemeCacheNavigate* request = ev->Get()->Request.Get(); + + if (request->ResultSet.size() != 1) { + ReplyError("SchemeCache returned unexpected result set size"); + return; + } + + auto& entry = request->ResultSet[0]; + if (entry.Status != NSchemeCache::TSchemeCacheNavigate::EStatus::Ok) { + ReplyError(TStringBuilder() << "SchemeCache resolve failed for path '" << VolumePath + << "', status: " << static_cast<int>(entry.Status)); + return; + } + + if (!entry.SolomonVolumeInfo) { + ReplyError(TStringBuilder() << "Path '" << VolumePath << "' is not a KeyValue volume"); + return; + } + + const auto& desc = entry.SolomonVolumeInfo->Description; + if (desc.PartitionsSize() == 0) { + ReplyError("KeyValue volume has no partitions"); + return; + } + + // Use partition 0 + KVTabletId = desc.GetPartitions(0).GetTabletId(); + if (!KVTabletId) { + ReplyError("Failed to get tablet ID for partition 0"); + return; + } + + Become(&TKvBodyReadActor::StateRead); + CreatePipeAndSendRead(); +} + +void TKvBodyReadActor::CreatePipeAndSendRead() { + // Prepare the temp file for incremental writing. + if (!NFs::Exists(OutputDir)) { + NFs::MakeDirectoryRecursive(OutputDir); + } + + TFsPath finalPath = TFsPath(OutputDir) / Md5Key; + TmpFilePath = finalPath.GetPath() + ".tmp"; + + try { + TmpFile = MakeHolder<TFile>(TmpFilePath, CreateAlways | WrOnly | Seq); + } catch (const yexception& e) { + ReplyError(TStringBuilder() << "Failed to open tmp file '" << TmpFilePath << "': " << e.what()); + return; + } + + CurrentOffset = 0; + Md5Ctx = MD5(); + + NTabletPipe::TClientConfig cfg; + cfg.RetryPolicy = { + .RetryLimitCount = 3u + }; + PipeClient = Register(NTabletPipe::CreateClient(SelfId(), KVTabletId, cfg)); + + SendNextChunkRead(); +} + +void TKvBodyReadActor::SendNextChunkRead() { + auto req = std::make_unique<TEvKeyValue::TEvRead>(); + req->Record.set_tablet_id(KVTabletId); + req->Record.set_key(Md5Key); + req->Record.set_offset(CurrentOffset); + req->Record.set_size(ReadChunkSize); + NTabletPipe::SendData(SelfId(), PipeClient, req.release()); + ALS_DEBUG(NKikimrServices::METADATA_PROVIDER) + << "TKvBodyReadActor: sending read for key='" << Md5Key + << "' offset=" << CurrentOffset << " size=" << ReadChunkSize; +} + +void TKvBodyReadActor::HandleReadResponse(TEvKeyValue::TEvReadResponse::TPtr& ev) { + const auto& record = ev->Get()->Record; + const auto status = record.status(); + + // RSTATUS_NOT_FOUND at a non-zero offset means we've read past the end of + // a value whose size is an exact multiple of ReadChunkSize. Treat it as EOF. + if (status == NKikimrKeyValue::Statuses::RSTATUS_NOT_FOUND && CurrentOffset > 0) { + FinalizeAndSave(); + return; + } + + if (status != NKikimrKeyValue::Statuses::RSTATUS_OK) { + TString msg = TStringBuilder() << "KV read failed with status: " + << static_cast<int>(status); + if (!record.msg().empty()) { + msg += " - " + record.msg(); + } + CleanupTmpFile(); + ReplyError(msg); + return; + } + + // Extract the chunk data. + TString chunk; + if (ev->Get()->IsPayload()) { + TRope rope = ev->Get()->GetBuffer(); + const TContiguousSpan span = rope.GetContiguousSpan(); + chunk = TString(span.data(), span.size()); + } else { + chunk = record.value(); + } + + if (!chunk.empty()) { + try { + TmpFile->Write(chunk.data(), chunk.size()); + } catch (const yexception& e) { + CleanupTmpFile(); + ReplyError(TStringBuilder() << "Failed to write to tmp file: " << e.what()); + return; + } + Md5Ctx.Update(chunk.data(), chunk.size()); + CurrentOffset += chunk.size(); + } + + // If we got fewer bytes than requested (or zero), this is the last chunk. + if (chunk.size() < ReadChunkSize) { + FinalizeAndSave(); + } else { + SendNextChunkRead(); + } +} + +void TKvBodyReadActor::FinalizeAndSave() { + // Close the tmp file before renaming it to the final path. + try { + TmpFile->Close(); + TmpFile.Reset(); + } catch (const yexception& e) { + CleanupTmpFile(); + ReplyError(TStringBuilder() << "Failed to close tmp file: " << e.what()); + return; + } + + if (CurrentOffset != ExpectedSize) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TKvBodyReadActor: size mismatch for UDF '" << Md5Key + << "': expected=" << ExpectedSize << ", actual=" << CurrentOffset; + NFs::Remove(TmpFilePath); + Send(ReplyTo, new TEvReadBodyResponse(false, Md5Key, + TStringBuilder() << "Size mismatch: expected=" << ExpectedSize + << ", actual=" << CurrentOffset)); + PassAway(); + return; + } + + // Verify MD5. + char md5Buf[33]; + TString computedMd5 = Md5Ctx.End(md5Buf); + + if (!Md5Key.empty() && computedMd5 != Md5Key) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TKvBodyReadActor: MD5 mismatch for UDF '" << Md5Key + << "': stored=" << Md5Key << ", computed=" << computedMd5; + NFs::Remove(TmpFilePath); + Send(ReplyTo, new TEvReadBodyResponse(false, Md5Key, + TStringBuilder() << "MD5 mismatch: stored=" << Md5Key << ", computed=" << computedMd5)); + PassAway(); + return; + } + + TString finalPath = TFsPath(OutputDir) / Md5Key; + if (!NFs::Rename(TmpFilePath, finalPath)) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TKvBodyReadActor: failed to rename tmp file for UDF '" << Md5Key << "'"; + NFs::Remove(TmpFilePath); + Send(ReplyTo, new TEvReadBodyResponse(false, Md5Key, + TStringBuilder() << "Failed to rename tmp file to '" << finalPath << "'")); + PassAway(); + return; + } + + if (!LoadUdfIntoRegistry(finalPath)) { + NFs::Remove(finalPath); + Send(ReplyTo, new TEvReadBodyResponse(false, Md5Key, + TStringBuilder() << "Failed to load UDF '" << Md5Key << "' into function registry")); + PassAway(); + return; + } + + ALS_INFO(NKikimrServices::METADATA_PROVIDER) + << "TKvBodyReadActor: saved UDF '" << Md5Key + << "' (" << CurrentOffset << " bytes) to " << finalPath; + + Send(ReplyTo, new TEvReadBodyResponse(true, Md5Key)); + PassAway(); +} + +bool TKvBodyReadActor::LoadUdfIntoRegistry(const TString& finalPath) const { + if (!FunctionRegistry) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TKvBodyReadActor: function registry is not available for UDF '" << Md5Key << "'"; + return false; + } + + NMiniKQL::TUdfModuleRemappings remappings; + THashSet<TString> modules; + try { + FunctionRegistry->LoadUdfs(finalPath, remappings, 0, {}, &modules); + } catch (const std::exception& e) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TKvBodyReadActor: failed to load UDF '" << Md5Key + << "' into function registry: " << e.what(); + return false; + } + + if (modules.empty()) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TKvBodyReadActor: no UDF modules were registered from '" << finalPath << "'"; + return false; + } + + ALS_INFO(NKikimrServices::METADATA_PROVIDER) + << "TKvBodyReadActor: loaded UDF '" << Md5Key + << "' into function registry from " << finalPath; + return true; +} + +void TKvBodyReadActor::CleanupTmpFile() { + if (TmpFile) { + try { TmpFile->Close(); } catch (...) {} + TmpFile.Reset(); + } + if (!TmpFilePath.empty()) { + NFs::Remove(TmpFilePath); + } +} + +void TKvBodyReadActor::HandlePipeConnected(TEvTabletPipe::TEvClientConnected::TPtr& ev) { + if (ev->Get()->Status != NKikimrProto::OK) { + CleanupTmpFile(); + ReplyError(TStringBuilder() << "Failed to connect to KV tablet " << KVTabletId); + } +} + +void TKvBodyReadActor::HandlePipeDestroyed(TEvTabletPipe::TEvClientDestroyed::TPtr& /*ev*/) { + CleanupTmpFile(); + ReplyError(TStringBuilder() << "Connection to KV tablet " << KVTabletId << " was lost"); +} + +void TKvBodyReadActor::ReplyError(const TString& message) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) << "TKvBodyReadActor: " << message; + Send(ReplyTo, new TEvReadBodyResponse(false, Md5Key, message)); + PassAway(); +} + +void TKvBodyReadActor::PassAway() { + if (PipeClient) { + NTabletPipe::CloseClient(SelfId(), PipeClient); + PipeClient = {}; + } + NActors::TActorBootstrapped<TKvBodyReadActor>::PassAway(); +} + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/kv_body_store.h b/ydb/services/udf_store/kv_body_store.h new file mode 100644 index 00000000000..8324af1f69b --- /dev/null +++ b/ydb/services/udf_store/kv_body_store.h @@ -0,0 +1,101 @@ +#pragma once +#include "events.h" +#include <ydb/library/actors/core/actor_bootstrapped.h> +#include <ydb/library/actors/core/hfunc.h> +#include <ydb/core/base/tablet_pipe.h> +#include <ydb/core/keyvalue/keyvalue_events.h> +#include <ydb/core/tx/scheme_cache/scheme_cache.h> + +#include <yql/essentials/minikql/mkql_function_registry.h> + +#include <library/cpp/digest/md5/md5.h> +#include <util/system/file.h> + +namespace NKikimr::NUdfStore { + + +// Actor that reads a UDF body from a KV tablet via the actor system, +// then saves it to disk and loads it into the function registry. +// Flow: +// 1. Resolves VolumePath via SchemeCache → gets tablet ID from SolomonVolumeInfo +// 2. Opens tablet pipe via NTabletPipe::CreateClient +// 3. Sends TEvKeyValue::TEvRead requests in ReadChunkSize-byte pages until EOF +// 4. Each chunk is appended to a temp file; MD5 is computed incrementally +// 5. On EOF: verifies body size against metadata, MD5, renames tmp→final, +// loads UDF into FunctionRegistry and checks that modules were registered +// 6. Sends TEvReadBodyResponse(success/failure) to ReplyTo +class TKvBodyReadActor : public NActors::TActorBootstrapped<TKvBodyReadActor> { +private: + NActors::TActorId ReplyTo; + + TString VolumePath; + TString Md5Key; + TString OutputDir; + ui64 ExpectedSize = 0; + + TIntrusivePtr<NMiniKQL::IMutableFunctionRegistry> FunctionRegistry; + + ui64 KVTabletId = 0; + NActors::TActorId PipeClient; + + // Chunked-read state + // KV tablet hard-caps a single read response at ~25 MiB (TotalSizeLimit). + // We use 16 MiB chunks to stay well within that limit. + static constexpr ui64 ReadChunkSize = 16ull << 20; // 16 MiB + ui64 CurrentOffset = 0; + TString TmpFilePath; + THolder<TFile> TmpFile; + MD5 Md5Ctx; + + void SendNavigateRequest(); + void CreatePipeAndSendRead(); // opens pipe, kicks off first chunk read + void SendNextChunkRead(); // sends TEvRead for [CurrentOffset, +ReadChunkSize) + void FinalizeAndSave(); // called when last chunk received + void CleanupTmpFile(); // closes + removes the tmp file on error + bool LoadUdfIntoRegistry(const TString& finalPath) const; + void ReplyError(const TString& message); + +public: + TKvBodyReadActor(const NActors::TActorId& replyTo, + const TString& md5Key, + const TString& volumePath, + const TString& outputDir, + TIntrusivePtr<NMiniKQL::IMutableFunctionRegistry> functionRegistry, + ui64 expectedSize) + : ReplyTo(replyTo) + , VolumePath(volumePath) + , Md5Key(md5Key) + , OutputDir(outputDir) + , ExpectedSize(expectedSize) + , FunctionRegistry(std::move(functionRegistry)) + {} + + void Bootstrap(); + + void HandleNavigateResult(TEvTxProxySchemeCache::TEvNavigateKeySetResult::TPtr& ev); + void HandleReadResponse(TEvKeyValue::TEvReadResponse::TPtr& ev); + void HandlePipeConnected(TEvTabletPipe::TEvClientConnected::TPtr& ev); + void HandlePipeDestroyed(TEvTabletPipe::TEvClientDestroyed::TPtr& ev); + + void PassAway() override; + + STFUNC(StateResolve) { + switch (ev->GetTypeRewrite()) { + hFunc(TEvTxProxySchemeCache::TEvNavigateKeySetResult, HandleNavigateResult); + default: + break; + } + } + + STFUNC(StateRead) { + switch (ev->GetTypeRewrite()) { + hFunc(TEvKeyValue::TEvReadResponse, HandleReadResponse); + hFunc(TEvTabletPipe::TEvClientConnected, HandlePipeConnected); + hFunc(TEvTabletPipe::TEvClientDestroyed, HandlePipeDestroyed); + default: + break; + } + } +}; + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/metadata_subscription/fetcher.h b/ydb/services/udf_store/metadata_subscription/fetcher.h new file mode 100644 index 00000000000..8852f816451 --- /dev/null +++ b/ydb/services/udf_store/metadata_subscription/fetcher.h @@ -0,0 +1,17 @@ +#pragma once + +#include "snapshot.h" +#include <ydb/services/metadata/abstract/common.h> +#include <ydb/library/accessor/accessor.h> + +namespace NKikimr::NUdfStore { + +class TSnapshotsFetcher: public NMetadata::NFetcher::TSnapshotsFetcher<TSnapshot> { + virtual std::vector<NMetadata::IClassBehaviour::TPtr> DoGetManagers() const override { + return { + TUdfMeta::GetBehaviour() + }; + } +}; + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/metadata_subscription/snapshot.cpp b/ydb/services/udf_store/metadata_subscription/snapshot.cpp new file mode 100644 index 00000000000..8a37e843cae --- /dev/null +++ b/ydb/services/udf_store/metadata_subscription/snapshot.cpp @@ -0,0 +1,39 @@ +#include "snapshot.h" + +namespace NKikimr::NUdfStore { + +bool TSnapshot::DoDeserializeFromResultSet(const Ydb::Table::ExecuteQueryResult& rawDataResult) { + Y_ABORT_UNLESS(rawDataResult.result_sets().size() == 1); + ParseSnapshotObjects<TUdfMeta>(rawDataResult.result_sets()[0], [this](TUdfMeta&& u) { + Udfs.emplace(u.GetMd5(), std::move(u)); + }); + return true; +} + +TString TSnapshot::DoSerializeToString() const { + TStringBuilder sb; + sb << "UDFS:"; + for (auto&& [md5, udf] : Udfs) { + sb << udf.SerializeToString(); + } + return sb; +} + +const TUdfMeta* TSnapshot::GetUdfByMd5(const TString& name) const { + auto it = Udfs.find(name); + if (it == Udfs.end()) { + return nullptr; + } + return &it->second; +} + +std::vector<TString> TSnapshot::GetUdfMd5s() const { + std::vector<TString> result; + result.reserve(Udfs.size()); + for (auto&& [md5, _] : Udfs) { + result.emplace_back(md5); + } + return result; +} + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/metadata_subscription/snapshot.h b/ydb/services/udf_store/metadata_subscription/snapshot.h new file mode 100644 index 00000000000..68259e2193e --- /dev/null +++ b/ydb/services/udf_store/metadata_subscription/snapshot.h @@ -0,0 +1,24 @@ +#pragma once +#include "udf_meta.h" + +#include <ydb/services/metadata/abstract/common.h> +#include <ydb/library/accessor/accessor.h> + +namespace NKikimr::NUdfStore { + +class TSnapshot: public NMetadata::NFetcher::ISnapshot { +private: + using TBase = NMetadata::NFetcher::ISnapshot; + using TUdfs = std::map<TString, TUdfMeta>; + YDB_READONLY_DEF(TUdfs, Udfs); +protected: + virtual bool DoDeserializeFromResultSet(const Ydb::Table::ExecuteQueryResult& rawData) override; + virtual TString DoSerializeToString() const override; +public: + using TBase::TBase; + + const TUdfMeta* GetUdfByMd5(const TString& name) const; + std::vector<TString> GetUdfMd5s() const; +}; + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/metadata_subscription/udf_behaviour.h b/ydb/services/udf_store/metadata_subscription/udf_behaviour.h new file mode 100644 index 00000000000..3658db10b65 --- /dev/null +++ b/ydb/services/udf_store/metadata_subscription/udf_behaviour.h @@ -0,0 +1,31 @@ +#pragma once +#include "udf_meta.h" +#include <ydb/services/metadata/abstract/initialization.h> +#include <ydb/services/metadata/manager/abstract.h> + +namespace NKikimr::NUdfStore { + +class TUdfBehaviour: public NMetadata::TClassBehaviour<TUdfMeta> { + virtual NMetadata::NInitializer::IInitializationBehaviour::TPtr ConstructInitializer() const override { + return {}; + } + virtual NMetadata::NModifications::IOperationsManager::TPtr ConstructOperationsManager() const override { + return {}; + } + virtual TString GetInternalStorageTablePath() const override { + return "udf_store/meta"; + } + virtual TString GetTypeId() const override { + return "UdfMeta"; + } + +public: + TUdfBehaviour() = default; + static IClassBehaviour::TPtr GetInstance() { + static std::shared_ptr<TUdfBehaviour> result = std::make_shared<TUdfBehaviour>(); + return result; + } + +}; + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/metadata_subscription/udf_meta.cpp b/ydb/services/udf_store/metadata_subscription/udf_meta.cpp new file mode 100644 index 00000000000..eefc2bb31a6 --- /dev/null +++ b/ydb/services/udf_store/metadata_subscription/udf_meta.cpp @@ -0,0 +1,89 @@ +#include "udf_meta.h" +#include "udf_behaviour.h" +#include <ydb/services/metadata/manager/ydb_value_operator.h> + +namespace NKikimr::NUdfStore { + +TUdfMeta::TDecoder::TDecoder(const Ydb::ResultSet& rawData) { + Md5Idx = GetFieldIndex(rawData, Md5ColName); + SizeIdx = GetFieldIndex(rawData, SizeColName); + NameIdx = GetFieldIndex(rawData, NameColName); + TypeIdx = GetFieldIndex(rawData, TypeColName); + ManifestIdx = GetFieldIndex(rawData, ManifestColName); +} + +bool TUdfMeta::TDecoder::Read(const i32 columnIdx, EUdfType& result, const Ydb::Value& r) const { + if (columnIdx >= (i32)r.items().size() || columnIdx < 0) { + return false; + } + auto& pValue = r.items()[columnIdx]; + if (!pValue.has_text_value()) { + return false; + } + // String values are fixed for backward compatibility + if (pValue.text_value() == "NATIVE_UNSAFE") { + result = EUdfType::NATIVE_UNSAFE; + } else if (pValue.text_value() == "WASM") { + result = EUdfType::WASM; + } else { + return false; + } + return true; +}; + +TVector<NKikimrSchemeOp::TColumnDescription> TUdfMeta::GetColumnDescription(){ + auto makeCol = [](const TString& name, const char* type) { + NKikimrSchemeOp::TColumnDescription col; + col.SetName(name); + col.SetType(type); + return col; + }; + return { + makeCol(Md5ColName, "Utf8"), + makeCol(SizeColName, "Uint64"), + makeCol(NameColName, "Utf8"), + makeCol(TypeColName, "Utf8"), + makeCol(ManifestColName, "Json"), + }; +} + +TVector<TString> TUdfMeta::GetPk() { + return {Md5ColName}; +} + + +bool TUdfMeta::DeserializeFromRecord(const TDecoder& decoder, const Ydb::Value& rawValue) { + if (!decoder.Read(decoder.GetMd5Idx(), Md5, rawValue)) { + return false; + } + if (!decoder.Read(decoder.GetSizeIdx(), Size, rawValue)) { + return false; + } + if (Size == 0) { + return false; + } + if (!decoder.Read(decoder.GetNameIdx(), Name, rawValue)) { + return false; + } + if (!decoder.Read(decoder.GetTypeIdx(), Type, rawValue)) { + return false; + } + if (decoder.GetManifestIdx() >= 0) { + decoder.Read(decoder.GetManifestIdx(), Manifest, rawValue); + } + return true; +} + +NMetadata::NInternal::TTableRecord TUdfMeta::SerializeToRecord() const { + return {}; +} + +TString TUdfMeta::SerializeToString() const { + return TStringBuilder() << "{" << "Md5: " << Md5 << ", Size: " << Size << ", Name: " << Name << ", Type: " << Type << ", Manifest: " << Manifest << "}"; +} + +NMetadata::IClassBehaviour::TPtr TUdfMeta::GetBehaviour() { + return TUdfBehaviour::GetInstance(); +} + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/metadata_subscription/udf_meta.h b/ydb/services/udf_store/metadata_subscription/udf_meta.h new file mode 100644 index 00000000000..b3a7f6946a2 --- /dev/null +++ b/ydb/services/udf_store/metadata_subscription/udf_meta.h @@ -0,0 +1,64 @@ +#pragma once +#include <ydb/library/accessor/accessor.h> +#include <ydb/core/base/appdata.h> + +#include <ydb/services/metadata/abstract/decoder.h> +#include <ydb/services/metadata/manager/object.h> +#include <ydb/services/metadata/manager/preparation_controller.h> + +namespace NKikimr::NUdfStore { + +enum class EUdfType { + NATIVE_UNSAFE, + WASM +}; + +class TUdfMeta: public NMetadata::NModifications::TObject<TUdfMeta> { +public: + static inline const TString Md5ColName = "md5"; //Utf8 (PK) + static inline const TString SizeColName = "size"; // Uint64 + static inline const TString NameColName = "name"; //Utf8 - informative + static inline const TString TypeColName = "type"; //Utf8 (allowed values from EUdfType) + static inline const TString ManifestColName = "manifest"; //Json + +private: + YDB_ACCESSOR_DEF(TString, Md5); + YDB_ACCESSOR_DEF(ui64, Size); + YDB_ACCESSOR_DEF(TString, Name); + YDB_ACCESSOR_DEF(EUdfType, Type); + YDB_ACCESSOR_DEF(TString, Manifest); +public: + static NMetadata::IClassBehaviour::TPtr GetBehaviour(); + static TVector<NKikimrSchemeOp::TColumnDescription> GetColumnDescription(); + static TVector<TString> GetPk(); + + // TDecoder maps table columns to field indices. + // Note: Body is NOT a table column; it is stored in a KV tablet. + class TDecoder: public NMetadata::NInternal::TDecoderBase { + using TBase = NMetadata::NInternal::TDecoderBase; + private: + YDB_ACCESSOR(i32, Md5Idx, -1); + YDB_ACCESSOR(i32, SizeIdx, -1); + YDB_ACCESSOR(i32, NameIdx, -1); + YDB_ACCESSOR(i32, TypeIdx, -1); + YDB_ACCESSOR(i32, ManifestIdx, -1); + + public: + TDecoder(const Ydb::ResultSet& rawData); + using TBase::Read; + bool Read(const i32 columnIdx, EUdfType& result, const Ydb::Value& r) const; + }; + + bool DeserializeFromRecord(const TDecoder& decoder, const Ydb::Value& rawValue); + NMetadata::NInternal::TTableRecord SerializeToRecord() const; + TString SerializeToString() const; + + bool operator<(const TUdfMeta& other) const { + return Md5 < other.Md5; + } + bool operator==(const TUdfMeta& other) const { + return Md5 == other.Md5; + } +}; + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/metadata_subscription/ya.make b/ydb/services/udf_store/metadata_subscription/ya.make new file mode 100644 index 00000000000..dd52d827af1 --- /dev/null +++ b/ydb/services/udf_store/metadata_subscription/ya.make @@ -0,0 +1,26 @@ +LIBRARY() +YQL_LAST_ABI_VERSION() + +SRCS( + udf_meta.cpp + snapshot.cpp +) + +GENERATE_ENUM_SERIALIZATION(udf_meta.h) + +PEERDIR( + ydb/library/actors/core + ydb/core/base + ydb/core/keyvalue + ydb/core/tx/scheme_cache + ydb/library/aclib + ydb/library/table_creator + ydb/services/metadata/request + ydb/services/metadata/abstract + ydb/services/metadata/manager + ydb/services/metadata + yql/essentials/minikql + library/cpp/digest/md5 +) + +END() diff --git a/ydb/services/udf_store/service.cpp b/ydb/services/udf_store/service.cpp new file mode 100644 index 00000000000..74b01cbbf1a --- /dev/null +++ b/ydb/services/udf_store/service.cpp @@ -0,0 +1,196 @@ +#include "service.h" +#include "metadata_subscription/fetcher.h" + + +#include <ydb/services/metadata/service.h> +#include <ydb/core/base/appdata.h> + +#include <util/folder/path.h> +#include <util/system/fs.h> + +#include <algorithm> + +namespace NKikimr::NUdfStore { + +bool TUdfStoreService::IsMd5Pending(const TString& md5) const { + return std::any_of(PendingUdfs.begin(), PendingUdfs.end(), + [&](const TPendingUdf& pending) { return pending.Md5 == md5; }); +} + +void TUdfStoreService::EnqueueUdfIfNeeded(const TString& md5, ui64 expectedSize) { + if (expectedSize == 0) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: UDF '" << md5 << "' has zero size in metadata, skipping fetch"; + return; + } + if (LoadedUdfs.contains(md5) || IsMd5Pending(md5)) { + return; + } + PendingUdfs.push_back({md5, expectedSize}); +} + +void TUdfStoreService::Bootstrap() { + Become(&TUdfStoreService::StateMain); + Register(new TUdfStoreInitializer(SelfId(), KvStorageMedia)); +} + +void TUdfStoreService::Handle(TEvStoreInitialized::TPtr& ev) { + KvVolumePath = ev->Get()->KvVolumePath; + ALS_INFO(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: infrastructure initialized, KV Volume path: " << KvVolumePath; + Send(NMetadata::NProvider::MakeServiceId(SelfId().NodeId()), + new NMetadata::NProvider::TEvSubscribeExternal(std::make_shared<TSnapshotsFetcher>())); +} + +void TUdfStoreService::Handle(TEvStoreInitFailed::TPtr& ev) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: infrastructure initialization failed: " << ev->Get()->ErrorMessage; + PassAway(); +} + +void TUdfStoreService::Handle(NMetadata::NProvider::TEvRefreshSubscriberData::TPtr& ev) { + auto snapshot = ev->Get()->GetSnapshotPtrAs<TSnapshot>(); + if (!snapshot) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: received non-UDF snapshot"; + return; + } + ALS_INFO(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: received UDF snapshot"; + + for (const auto& [md5, udf] : snapshot->GetUdfs()) { + const TUdfMeta* existing = CurrentSnapshot ? CurrentSnapshot->GetUdfByMd5(md5) : nullptr; + const bool isNew = !existing; + + if (isNew) { + ALS_INFO(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: UDF added" + << ", md5=" << md5 + << ", name=" << udf.GetName() + << ", type=" << udf.GetType() + << ", size=" << udf.GetSize(); + } else if (existing->GetSize() != udf.GetSize()) { + ALS_INFO(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: UDF size changed" + << ", md5=" << md5 + << ", old_size=" << existing->GetSize() + << ", new_size=" << udf.GetSize(); + LoadedUdfs.erase(md5); + FetchRetryCounts.erase(md5); + } + switch(udf.GetType()) { + case EUdfType::NATIVE_UNSAFE: + if (!EnableUnsafeNativeUdfFlag) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: EnableUnsafeNativeUdf is not set," + << " skipping UDF '" << md5 << "' with name '" << udf.GetName() << "'"; + break; + } + if (UnsafeNativeUdfDir.empty()) { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: EnableUnsafeNativeUdf is set but UnsafeNativeUdfDir is empty," + << " skipping UDF '" << md5 << "' with name '" << udf.GetName() << "'"; + break; + } + if (!LoadedUdfs.contains(md5)) { + FetchRetryCounts.erase(md5); + } + EnqueueUdfIfNeeded(md5, udf.GetSize()); + break; + case EUdfType::WASM: + //TODO implement me + break; + } + } + + // Log removed UDFs + if (CurrentSnapshot) { + for (const auto& [md5, udf] : CurrentSnapshot->GetUdfs()) { + if (!snapshot->GetUdfByMd5(md5)) { + ALS_INFO(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: UDF removed" + << ": md5=" << md5 + << ", type=" << udf.GetType() + << ", name=" << udf.GetName(); + LoadedUdfs.erase(md5); + FetchRetryCounts.erase(md5); + } + //TODO unregister removed UDFs + } + } + + CurrentSnapshot = snapshot; + + if (!FetchInProgress) { + FetchNextBody(); + } +} + +void TUdfStoreService::FetchNextBody() { + if (PendingUdfs.empty()) { + FetchInProgress = false; + return; + } + + FetchInProgress = true; + const auto& pending = PendingUdfs.front(); + + // Pass Md5 as the KV key; TKvBodyReadActor writes the binary to OutputDir/<md5>. + Register(new TKvBodyReadActor( + SelfId(), + pending.Md5, + KvVolumePath, + UnsafeNativeUdfDir, + FunctionRegistry, + pending.ExpectedSize)); +} + +void TUdfStoreService::Handle(TEvReadBodyResponse::TPtr& ev) { + if (PendingUdfs.empty()) { + ALS_WARN(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: received unexpected TEvReadBodyResponse for UDF '" + << ev->Get()->Name << "' with no pending fetches"; + return; + } + + auto pending = std::move(PendingUdfs.front()); + PendingUdfs.pop_front(); + + if (ev->Get()->Success) { + LoadedUdfs.insert(pending.Md5); + FetchRetryCounts.erase(pending.Md5); + ALS_INFO(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: native UDF '" << pending.Md5 + << "' saved to " << UnsafeNativeUdfDir; + } else { + ui32& retryCount = FetchRetryCounts[pending.Md5]; + if (retryCount < MaxFetchRetries) { + ++retryCount; + EnqueueUdfIfNeeded(pending.Md5, pending.ExpectedSize); + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: failed to save native UDF '" << pending.Md5 + << "' (retry " << retryCount << "/" << MaxFetchRetries + << "): " << ev->Get()->ErrorMessage; + } else { + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreService: giving up on native UDF '" << pending.Md5 + << "' after " << MaxFetchRetries << " retries: " << ev->Get()->ErrorMessage; + } + } + + FetchNextBody(); +} + +NActors::TActorId MakeServiceId(ui32 nodeId) { + return NActors::TActorId(nodeId, "SrvcUdfStore"); +} + +NActors::IActor* CreateService(const NKikimrConfig::TUdfStoreConfig& serviceConfig, TIntrusivePtr<NMiniKQL::IMutableFunctionRegistry> functionRegistry) { + if (!serviceConfig.GetEnabled()) { + return nullptr; + } + return new TUdfStoreService(serviceConfig, std::move(functionRegistry)); +} + + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/service.h b/ydb/services/udf_store/service.h new file mode 100644 index 00000000000..31938371d6a --- /dev/null +++ b/ydb/services/udf_store/service.h @@ -0,0 +1,82 @@ +#pragma once +#include "metadata_subscription/snapshot.h" +#include "kv_body_store.h" +#include "store_initializer.h" + +#include <ydb/services/metadata/abstract/common.h> +#include <ydb/library/actors/core/actor_bootstrapped.h> +#include <ydb/library/actors/core/hfunc.h> +#include <ydb/library/accessor/accessor.h> + +#include <yql/essentials/minikql/mkql_function_registry.h> + +namespace NKikimr::NUdfStore { + +struct TPendingUdf { + TString Md5; + ui64 ExpectedSize = 0; // body size from metadata (must be > 0) +}; + +// Actor that subscribes to UDF metadata snapshot updates and schedules body fetches +// for NATIVE_UNSAFE UDFs that are not yet loaded on disk. +// Bodies are read from a KV tablet (not from the metadata table) by TKvBodyReadActor, +// which saves them under UnsafeNativeUdfDir/<md5>, verifies size/MD5 and loads them +// into FunctionRegistry. This actor owns the fetch queue, retry policy and loaded set. +// Infrastructure (meta table + KV volume) is initialized by the store initializer actor +// launched from Bootstrap. +class TUdfStoreService: public NActors::TActorBootstrapped<TUdfStoreService> { +private: + using TBase = NActors::TActorBootstrapped<TUdfStoreService>; + TIntrusivePtr<NMiniKQL::IMutableFunctionRegistry> FunctionRegistry; + TString KvStorageMedia; + YDB_READONLY_FLAG(EnableUnsafeNativeUdf, false); + TString KvVolumePath; + TString UnsafeNativeUdfDir; + std::shared_ptr<TSnapshot> CurrentSnapshot; + + + bool FetchInProgress = false; + THashSet<TString> LoadedUdfs; + THashMap<TString, ui32> FetchRetryCounts; + + static constexpr ui32 MaxFetchRetries = 5; + + // Queue of UDFs whose bodies are being fetched from the KV tablet + std::deque<TPendingUdf> PendingUdfs; + + bool IsMd5Pending(const TString& md5) const; + void EnqueueUdfIfNeeded(const TString& md5, ui64 expectedSize); + void FetchNextBody(); + +protected: + void Handle(TEvStoreInitialized::TPtr& ev); + void Handle(TEvStoreInitFailed::TPtr& ev); + void Handle(NMetadata::NProvider::TEvRefreshSubscriberData::TPtr& ev); + void Handle(TEvReadBodyResponse::TPtr& ev); + +public: + TUdfStoreService(const NKikimrConfig::TUdfStoreConfig& config, TIntrusivePtr<NMiniKQL::IMutableFunctionRegistry> functionRegistry) + : FunctionRegistry(std::move(functionRegistry)) + , KvStorageMedia(config.GetKvStorageMedia()) + , EnableUnsafeNativeUdfFlag(config.GetEnableUnsafeNativeUdf()) + , UnsafeNativeUdfDir(config.GetUnsafeNativeUdfDir()) + {} + + void Bootstrap(); + + STATEFN(StateMain) { + switch (ev->GetTypeRewrite()) { + hFunc(TEvStoreInitialized, Handle); + hFunc(TEvStoreInitFailed, Handle); + hFunc(NMetadata::NProvider::TEvRefreshSubscriberData, Handle); + hFunc(TEvReadBodyResponse, Handle); + default: + break; + } + } +}; + +NActors::IActor* CreateService(const NKikimrConfig::TUdfStoreConfig& serviceConfig, TIntrusivePtr<NMiniKQL::IMutableFunctionRegistry> functionRegistry); +NActors::TActorId MakeServiceId(ui32 nodeId); + +} diff --git a/ydb/services/udf_store/store_initializer.cpp b/ydb/services/udf_store/store_initializer.cpp new file mode 100644 index 00000000000..3feaa35ce2c --- /dev/null +++ b/ydb/services/udf_store/store_initializer.cpp @@ -0,0 +1,84 @@ +#include "store_initializer.h" +#include "metadata_subscription/udf_meta.h" + +#include <ydb/library/actors/core/log.h> +#include <ydb/library/aclib/aclib.h> +#include <ydb/core/base/appdata.h> + +namespace NKikimr::NUdfStore { + +void TUdfStoreInitializer::Bootstrap() { + Become(&TUdfStoreInitializer::StateFunc); + const auto& path = NKikimr::SplitPath(TUdfMeta::GetBehaviour()->GetStorageTablePath()); + auto it = cbegin(path); + while (it != path.end() && *it != NMetadata::NProvider::TServiceOperator::GetPath()) { + ++it; + } + AFL_VERIFY(it != cend(path)); + Register(CreateTableCreator( + {it, cend(path)}, + TUdfMeta::GetColumnDescription(), + TUdfMeta::GetPk(), + NKikimrServices::METADATA_PROVIDER, + Nothing(), + {}, + /* isSystemUser */ true + )); +} + +void TUdfStoreInitializer::HandleTableCreated(TEvTableCreator::TEvCreateTableResponse::TPtr& ev) { + if (!ev->Get()->Success) { + const TString errorMessage = TStringBuilder() + << "failed to create meta table: " << ev->Get()->Issues.ToString(); + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreInitializer: " << errorMessage; + Send(ParentId, new TEvStoreInitFailed(errorMessage)); + PassAway(); + return; + } + + + auto tablePath = SplitPath(TUdfMeta::GetBehaviour()->GetStorageTablePath()); + AFL_VERIFY(!tablePath.empty()); + tablePath.pop_back(); + tablePath.push_back("binaries"); + KvVolumePath = NKikimr::CombinePath(cbegin(tablePath), cend(tablePath)); + + ALS_INFO(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreInitializer: meta table created, creating KV volume at " << KvVolumePath; + + NACLib::TUserToken userToken("metadata@system", {}); + + Ydb::KeyValue::CreateVolumeRequest kvRequest; + kvRequest.set_path(KvVolumePath); + kvRequest.set_partition_count(1); + + auto* storageConfig = kvRequest.mutable_storage_config(); + for (int i = 0; i < 3; ++i) { + storageConfig->add_channel()->set_media(KvStorageMedia); + } + + auto controller = std::make_shared<NMetadata::NRequest::TNaiveExternalController<NMetadata::NRequest::TDialogCreateKvVolume>>(SelfId()); + NMetadata::NRequest::TYDBOneRequestSender<NMetadata::NRequest::TDialogCreateKvVolume> sender(kvRequest, userToken, controller); + sender.Start(); +} + +void TUdfStoreInitializer::HandleKvVolumeCreated( + NMetadata::NRequest::TEvRequestResult<NMetadata::NRequest::TDialogCreateKvVolume>::TPtr& /*ev*/) +{ + ALS_INFO(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreInitializer: KV volume '" << KvVolumePath << "' created successfully"; + Send(ParentId, new TEvStoreInitialized{KvVolumePath}); + PassAway(); +} + +void TUdfStoreInitializer::HandleRequestFailed(NMetadata::NRequest::TEvRequestFailed::TPtr& ev) { + const TString errorMessage = TStringBuilder() + << "failed to create KV volume: " << ev->Get()->GetErrorMessage(); + ALS_ERROR(NKikimrServices::METADATA_PROVIDER) + << "TUdfStoreInitializer: " << errorMessage; + Send(ParentId, new TEvStoreInitFailed(errorMessage)); + PassAway(); +} + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/store_initializer.h b/ydb/services/udf_store/store_initializer.h new file mode 100644 index 00000000000..867ccb253a2 --- /dev/null +++ b/ydb/services/udf_store/store_initializer.h @@ -0,0 +1,48 @@ +#pragma once + +#include "events.h" +#include <ydb/services/metadata/request/common.h> +#include <ydb/services/metadata/request/request_actor_cb.h> +#include <ydb/library/table_creator/table_creator.h> +#include <ydb/library/actors/core/actor_bootstrapped.h> +#include <ydb/library/actors/core/hfunc.h> + +namespace NKikimr::NUdfStore { + +// Actor that sequentially initializes UDF store infrastructure: +// 1. Creates the metadata table .metadata/udf_store/meta via CreateTableCreator. +// 2. On success, creates the KV volume .metadata/udf_store/binaries with 3 channels. +// On success, sends TEvStoreInitialized to ParentId; on failure, TEvStoreInitFailed. +// Dies after both steps complete (success or failure). +class TUdfStoreInitializer : public NActors::TActorBootstrapped<TUdfStoreInitializer> { +private: + using TBase = NActors::TActorBootstrapped<TUdfStoreInitializer>; + + NActors::TActorId ParentId; + TString KvStorageMedia; + TString KvVolumePath; + + void HandleTableCreated(TEvTableCreator::TEvCreateTableResponse::TPtr& ev); + void HandleKvVolumeCreated(NMetadata::NRequest::TEvRequestResult<NMetadata::NRequest::TDialogCreateKvVolume>::TPtr& ev); + void HandleRequestFailed(NMetadata::NRequest::TEvRequestFailed::TPtr& ev); + +public: + explicit TUdfStoreInitializer(const NActors::TActorId& parentId, const TString& kvStorageMedia = "ssd") + : ParentId(parentId) + , KvStorageMedia(kvStorageMedia) + {} + + void Bootstrap(); + + STATEFN(StateFunc) { + switch (ev->GetTypeRewrite()) { + hFunc(TEvTableCreator::TEvCreateTableResponse, HandleTableCreated); + hFunc(NMetadata::NRequest::TEvRequestResult<NMetadata::NRequest::TDialogCreateKvVolume>, HandleKvVolumeCreated); + hFunc(NMetadata::NRequest::TEvRequestFailed, HandleRequestFailed); + default: + break; + } + } +}; + +} // namespace NKikimr::NUdfStore diff --git a/ydb/services/udf_store/ya.make b/ydb/services/udf_store/ya.make new file mode 100644 index 00000000000..88da8981715 --- /dev/null +++ b/ydb/services/udf_store/ya.make @@ -0,0 +1,26 @@ +LIBRARY() +YQL_LAST_ABI_VERSION() + +SRCS( + service.cpp + store_initializer.cpp + kv_body_store.cpp +) + +PEERDIR( + ydb/library/actors/core + ydb/core/base + ydb/core/keyvalue + ydb/core/tx/scheme_cache + ydb/library/aclib + ydb/library/table_creator + ydb/services/udf_store/metadata_subscription + ydb/services/metadata/request + ydb/services/metadata/abstract + ydb/services/metadata/manager + ydb/services/metadata + yql/essentials/minikql + library/cpp/digest/md5 +) + +END() diff --git a/ydb/services/ya.make b/ydb/services/ya.make index 7165164de00..495018a67c2 100644 --- a/ydb/services/ya.make +++ b/ydb/services/ya.make @@ -17,6 +17,7 @@ RECURSE( local_discovery maintenance metadata + udf_store monitoring persqueue_cluster_discovery persqueue_v1 diff --git a/ydb/tests/functional/udf_store/lib/constants.py b/ydb/tests/functional/udf_store/lib/constants.py new file mode 100644 index 00000000000..7e81975e9eb --- /dev/null +++ b/ydb/tests/functional/udf_store/lib/constants.py @@ -0,0 +1,7 @@ +# -*- coding: utf-8 -*- +# Shared constants for the UDF store. +# Must match TUdfBehaviour::GetInternalStorageTablePath() under .metadata/. + +UDF_STORE_PATH = ".metadata/udf_store" +UDF_TABLE_META_PATH = UDF_STORE_PATH + "/meta" +UDF_KV_BINARIES_PATH = UDF_STORE_PATH + "/binaries" diff --git a/ydb/tests/functional/udf_store/lib/ya.make b/ydb/tests/functional/udf_store/lib/ya.make new file mode 100644 index 00000000000..8d5a79856d6 --- /dev/null +++ b/ydb/tests/functional/udf_store/lib/ya.make @@ -0,0 +1,7 @@ +PY3_LIBRARY() + +PY_SRCS( + constants.py +) + +END() diff --git a/ydb/tests/functional/udf_store/test_udf_store.py b/ydb/tests/functional/udf_store/test_udf_store.py new file mode 100644 index 00000000000..a4c832b15d8 --- /dev/null +++ b/ydb/tests/functional/udf_store/test_udf_store.py @@ -0,0 +1,324 @@ +# -*- coding: utf-8 -*- +import hashlib +import logging +import os +import shutil +import subprocess +import time + +import pytest +import yatest.common + +from ydb.tests.library.harness.kikimr_runner import KiKiMR +from ydb.tests.library.harness.kikimr_config import KikimrConfigGenerator +from ydb.tests.oss.ydb_sdk_import import ydb +from ydb.tests.functional.udf_store.lib.constants import UDF_TABLE_META_PATH, UDF_KV_BINARIES_PATH + +logger = logging.getLogger(__name__) + +CLUSTER_CONFIG = dict( + additional_log_configs={ + "METADATA_PROVIDER": 7, # DEBUG + } +) + +UDF_OUTPUT_DIR = yatest.common.output_path("ydb_udfs") + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + +def _run_query(config, query): + with ydb.Driver(config) as driver: + with ydb.QuerySessionPool(driver, size=1) as pool: + return pool.execute_with_retries(query) + + +def _wait_for_condition(condition_fn, timeout_seconds=60, poll_interval=1, description="condition"): + """Poll until condition_fn() returns True or timeout expires.""" + deadline = time.time() + timeout_seconds + while time.time() < deadline: + if condition_fn(): + return True + logger.info("Waiting for %s... (%.0fs remaining)", description, deadline - time.time()) + time.sleep(poll_interval) + return False + + +def _kv_volume_tool(): + return yatest.common.binary_path(os.environ["YDB_KV_VOLUME_TOOL_PATH"]) + + +def _run_kv_tool(endpoint, database, path, command, *extra_args): + """Run kv_volume_tool with the given command; raise RuntimeError on failure.""" + cmd = [_kv_volume_tool(), command, "-e", endpoint, "-d", database, "-p", path, "-v", *extra_args] + logger.info("Running kv_volume_tool: %s", " ".join(cmd)) + result = subprocess.run(cmd, capture_output=True, text=True, timeout=60) + logger.info("kv_volume_tool stdout: %s", result.stdout) + if result.stderr: + logger.info("kv_volume_tool stderr: %s", result.stderr) + if result.returncode != 0: + raise RuntimeError( + f"kv_volume_tool {command} failed (rc={result.returncode}): " + f"stdout={result.stdout}, stderr={result.stderr}" + ) + return result + + +def _table_exists(config, database, table_path=UDF_TABLE_META_PATH): + """Return True if the UDF metadata table can be queried.""" + try: + result = _run_query( + config, + "SELECT COUNT(*) AS cnt FROM `{database}/{path}`".format(database=database, path=table_path), + ) + return bool(result and result[0].rows) + except Exception as e: + logger.debug("UDF metadata table not ready yet: %s", e) + return False + + +def _kv_volume_exists(endpoint, database, path=UDF_KV_BINARIES_PATH): + """Return True if the KV volume responds to 'describe'.""" + try: + _run_kv_tool(endpoint, database, path, "describe") + return True + except RuntimeError as e: + logger.debug("KV volume not ready yet: %s", e) + return False + + [email protected]("enable_udf_store", [True, False], ids=["flag_on", "flag_off"]) +def test_udf_store_feature_flag(enable_udf_store): + """ + When udf_store_config.enabled=true → both the UDF metadata table and the KV volume must be created. + When udf_store_config is absent or disabled → neither must appear. + """ + database = "/Root/test" + cluster = _make_cluster(enable_udf_store=enable_udf_store) + db_nodes = _create_database(cluster, database) + try: + node = cluster.nodes[1] + driver_config = ydb.DriverConfig( + endpoint="%s:%s" % (node.host, node.port), + database=database, + ) + grpc_endpoint = "grpc://%s:%s" % (node.host, node.port) + timeout = _SETTLE_TIMEOUT if enable_udf_store else _ABSENT_TIMEOUT + + table_appeared = _wait_for_condition( + lambda: _table_exists(driver_config, database), + timeout_seconds=timeout, + description="UDF metadata table (enable_udf_store=%s)" % enable_udf_store, + ) + kv_appeared = _wait_for_condition( + lambda: _kv_volume_exists(grpc_endpoint, database), + timeout_seconds=timeout, + description="KV volume (enable_udf_store=%s)" % enable_udf_store, + ) + + if enable_udf_store: + assert table_appeared, ( + "UDF metadata table `%s` was NOT created within %ds when udf_store_config.enabled=true" + % (UDF_TABLE_META_PATH, _SETTLE_TIMEOUT) + ) + assert kv_appeared, ( + "KV volume `%s` was NOT created within %ds when udf_store_config.enabled=true" + % (UDF_KV_BINARIES_PATH, _SETTLE_TIMEOUT) + ) + else: + assert not table_appeared, ( + "UDF metadata table `%s` appeared even though udf_store_config is disabled" % UDF_TABLE_META_PATH + ) + assert not kv_appeared, ( + "KV volume `%s` appeared even though udf_store_config is disabled" % UDF_KV_BINARIES_PATH + ) + finally: + cluster.remove_database(database) + cluster.unregister_and_stop_slots(db_nodes) + cluster.stop() + + +def _upload_udf_binary(): + return yatest.common.binary_path(os.environ["YDB_UPLOAD_UDF_PATH"]) + + +def _run_upload_udf(endpoint, database, udf_so_path): + """ + Invoke the upload_udf binary as a subprocess. + + Returns the md5 hex-digest printed by the binary on stdout. + Raises RuntimeError if the binary exits with a non-zero code. + """ + cmd = [ + _upload_udf_binary(), + "--endpoint", endpoint, + "--database", database, + "--udf-file", udf_so_path, + ] + # Resolve YDB_KV_VOLUME_TOOL_PATH to an absolute path so the subprocess + # can find the binary regardless of its working directory. + env = os.environ.copy() + env["YDB_KV_VOLUME_TOOL_PATH"] = _kv_volume_tool() + logger.info("Running upload_udf: %s", " ".join(cmd)) + result = subprocess.run(cmd, capture_output=True, text=True, timeout=120, env=env) + if result.stderr: + logger.info("upload_udf stderr:\n%s", result.stderr.strip()) + if result.returncode != 0: + raise RuntimeError( + f"upload_udf failed (rc={result.returncode}): {result.stderr}" + ) + return result.stdout.strip() + + +def test_using_native_unsafe_udf(): + """ + 1. Use the pre-built dicts UDF shared library as the binary to upload. + 2. Delegate upload + metadata registration to the upload_udf helper binary. + 3. TUdfStoreService detects the new metadata row, fetches the binary + from the KV store, and writes it to UnsafeNativeUdfDir/<md5>. + 4. Assert that the file exists, its size and md5 match metadata. + """ + udf_output_dir = UDF_OUTPUT_DIR + database = "/Root/test" + cluster = _make_cluster(enable_udf_store=True, enable_native_udf=True, native_udf_dir=udf_output_dir) + db_nodes = _create_database(cluster, database) + try: + node = cluster.nodes[1] + driver_config = ydb.DriverConfig( + endpoint="%s:%s" % (node.host, node.port), + database=database, + ) + endpoint = "grpc://%s:%s" % (node.host, node.port) + + # --- Step 0: Wait for UDF metadata table --- + assert _wait_for_condition( + lambda: _table_exists(driver_config, database), + timeout_seconds=60, + description="UDF metadata table creation at startup", + ), "UDF metadata table was not created at startup within timeout" + + # --- Step 1: Clean output directory so we start from a known state --- + if os.path.exists(udf_output_dir): + shutil.rmtree(udf_output_dir) + + # --- Step 2: Resolve the pre-built dicts UDF path --- + udf_so_path = yatest.common.binary_path(os.environ["YDB_DICTS_UDF_PATH"]) + logger.info("Dicts UDF binary path: %s", udf_so_path) + + # --- Step 3: Wait for KV volume --- + assert _wait_for_condition( + lambda: _kv_volume_exists(endpoint, database), + timeout_seconds=60, + description="KV volume creation at startup", + ), f"KV volume at {UDF_KV_BINARIES_PATH} was not created at startup within timeout" + + # --- Step 4+5: Upload binary and register metadata (with size) via upload_udf --- + udf_md5 = _run_upload_udf(endpoint, database, udf_so_path) + logger.info("upload_udf reported md5=%s", udf_md5) + + # --- Step 6: Wait for binary to appear in UnsafeNativeUdfDir --- + # TKvBodyReadActor names the output file after the md5 checksum. + expected_file_path = os.path.join(udf_output_dir, udf_md5) + assert _wait_for_condition( + lambda: os.path.isfile(expected_file_path), + timeout_seconds=120, + description=f"native UDF file {expected_file_path}", + ), ( + f"Native UDF file was not created at {expected_file_path} within timeout. " + f"Expected TUdfStoreService to fetch the binary from KV and write it to " + f"UnsafeNativeUdfDir='{udf_output_dir}' under filename=md5='{udf_md5}'." + ) + + # --- Step 7: Verify file size and md5 --- + CHUNK_SIZE = 4 * 1024 * 1024 # 4 MiB + binary_size = os.path.getsize(udf_so_path) + saved_size = os.path.getsize(expected_file_path) + assert saved_size == binary_size, ( + f"File size mismatch: expected {binary_size}, got {saved_size}" + ) + file_md5_ctx = hashlib.md5() + with open(expected_file_path, "rb") as f: + while True: + chunk = f.read(CHUNK_SIZE) + if not chunk: + break + file_md5_ctx.update(chunk) + saved_md5 = file_md5_ctx.hexdigest() + assert saved_md5 == udf_md5, ( + f"MD5 mismatch: expected {udf_md5}, got {saved_md5}" + ) + + # --- Step 8: Execute a query using the loaded UDF and verify the result --- + # TKvBodyReadActor calls LoadUdfs() after writing the file to disk, but the + # function registry update may not be visible on the query layer immediately. + # Poll until the query succeeds (the UDF module may take a moment to register). + UDF_QUERY = 'SELECT Dicts::StrToInt("Sorted");' + udf_query_result = [None] + + def try_udf_query(): + try: + udf_query_result[0] = _run_query(driver_config, UDF_QUERY) + return True + except Exception as e: + logger.debug("UDF query not ready yet: %s", e) + return False + + assert _wait_for_condition( + try_udf_query, + timeout_seconds=60, + description="Dicts UDF query execution", + ), "UDF query did not succeed within timeout after the binary was written to disk" + + # Dicts::StrToInt("Sorted") returns a dict mapping number-word strings to ints, + # e.g. {b'zero': 0, b'one': 1, ..., b'nine': 9}. + rows = udf_query_result[0][0].rows + assert len(rows) == 1, "UDF query returned wrong number of rows" + result_value = list(rows[0].values())[0] + assert isinstance(result_value, dict), ( + f"Dicts::StrToInt('Sorted') expected a dict, got {type(result_value)}: {result_value!r}" + ) + assert result_value.get(b'zero') == 0, ( + f"Expected result_value[b'zero'] == 0, got {result_value!r}" + ) + assert result_value.get(b'nine') == 9, ( + f"Expected result_value[b'nine'] == 9, got {result_value!r}" + ) + logger.info("Test passed: dicts UDF (md5=%s) appeared at %s and query returned %s", + udf_md5, expected_file_path, result_value) + + finally: + cluster.remove_database(database) + cluster.unregister_and_stop_slots(db_nodes) + cluster.stop() + + +# --------------------------------------------------------------------------- +# Feature-flag test: parametrised over udf_store_config enabled / disabled +# --------------------------------------------------------------------------- + +_SETTLE_TIMEOUT = 60 # seconds to wait when flag is ON +_ABSENT_TIMEOUT = 15 # seconds to confirm absence when flag is OFF + + +def _make_cluster(enable_udf_store: bool, enable_native_udf: bool = False, native_udf_dir: str = ""): + configurator = KikimrConfigGenerator( + additional_log_configs={"METADATA_PROVIDER": 7}, + ) + if enable_udf_store: + udf_store_config = {"enabled": True, "kv_storage_media": "hdd"} + if enable_native_udf: + udf_store_config["enable_unsafe_native_udf"] = True + udf_store_config["unsafe_native_udf_dir"] = native_udf_dir + configurator.yaml_config["udf_store_config"] = udf_store_config + cluster = KiKiMR(configurator=configurator) + cluster.start() + return cluster + + +def _create_database(cluster, database): + cluster.create_database(database, storage_pool_units_count={"hdd": 1}) + nodes = cluster.register_and_start_slots(database, count=1) + cluster.wait_tenant_up(database) + return nodes diff --git a/ydb/tests/functional/udf_store/upload_udf/__main__.py b/ydb/tests/functional/udf_store/upload_udf/__main__.py new file mode 100644 index 00000000000..b3151900f0e --- /dev/null +++ b/ydb/tests/functional/udf_store/upload_udf/__main__.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Standalone helper: upload a UDF binary to the YDB KV-body store and register +it in the UDF metadata table (md5, size, name, type). + +Usage: + upload_udf \\ + --endpoint grpc://localhost:2136 \\ + --database /Root/mydb \\ + --udf-file /path/to/libdicts_udf.so + +The kv_volume_tool binary is located via the YDB_KV_VOLUME_TOOL_PATH +environment variable (set automatically when running under `ya test`). + +On success the program prints the md5 hex-digest of the uploaded binary to +stdout and exits with code 0. On any error it prints a message to stderr and +exits with a non-zero code. +""" + +import argparse +import hashlib +import os +import subprocess +import sys + +import ydb + +from ydb.tests.functional.udf_store.lib.constants import UDF_KV_BINARIES_PATH, UDF_TABLE_META_PATH +_CHUNK_SIZE = 4 * 1024 * 1024 # 4 MiB + + +def _kv_tool() -> str: + path = os.environ.get("YDB_KV_VOLUME_TOOL_PATH") + if not path: + raise RuntimeError("YDB_KV_VOLUME_TOOL_PATH environment variable is not set") + return path + + +def _compute_md5(path: str) -> tuple: + """Return (hex_md5, file_size_bytes) computed by streaming the file.""" + ctx = hashlib.md5() + size = 0 + with open(path, "rb") as f: + while True: + chunk = f.read(_CHUNK_SIZE) + if not chunk: + break + ctx.update(chunk) + size += len(chunk) + return ctx.hexdigest(), size + + +def _upload_to_kv(endpoint: str, database: str, udf_file: str, md5: str) -> None: + """Upload the UDF binary to the KV body store with key=md5.""" + full_volume_path = "{}/{}".format(database, UDF_KV_BINARIES_PATH) + cmd = [ + _kv_tool(), "upload", + "-e", endpoint, + "-d", database, + "-p", full_volume_path, + "-v", + "--partition-id", "0", + "--key", md5, + "--file", udf_file, + ] + result = subprocess.run(cmd, capture_output=True, text=True, timeout=120) + if result.returncode != 0: + raise RuntimeError( + "kv_volume_tool upload failed (rc={}): stdout={!r}, stderr={!r}".format( + result.returncode, result.stdout, result.stderr) + ) + + +def _upsert_udf_row(endpoint: str, database: str, md5: str, name: str, size: int) -> None: + """Insert/update a row in the UDF metadata table via the YDB Python SDK.""" + full_table = "{}/{}".format(database, UDF_TABLE_META_PATH) + query = ( + 'UPSERT INTO `{}` (md5, size, name, type) ' + 'VALUES ("{}", {}, "{}", "NATIVE_UNSAFE")' + ).format(full_table, md5, size, name) + + with ydb.Driver(ydb.DriverConfig(endpoint=endpoint, database=database)) as driver: + driver.wait(timeout=30, fail_fast=True) + with ydb.QuerySessionPool(driver, size=1) as pool: + pool.execute_with_retries(query) + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Upload a UDF binary to the YDB KV store and register it in the metadata table." + ) + parser.add_argument("--endpoint", required=True, + help="YDB gRPC endpoint, e.g. grpc://localhost:2136") + parser.add_argument("--database", required=True, + help="YDB database path, e.g. /Root/mydb") + parser.add_argument("--udf-file", required=True, + help="Path to the UDF shared library (.so) to upload") + args = parser.parse_args() + + # Derive the UDF name from the filename, stripping known suffixes. + udf_basename = os.path.basename(args.udf_file) + udf_name = udf_basename.removesuffix(".so") + + try: + md5, size = _compute_md5(args.udf_file) + print("[upload_udf] file={} size={} md5={}".format(args.udf_file, size, md5), file=sys.stderr) + + _upload_to_kv(args.endpoint, args.database, args.udf_file, md5) + print("[upload_udf] binary uploaded to KV volume with key={}".format(md5), file=sys.stderr) + + _upsert_udf_row(args.endpoint, args.database, md5, udf_name, size) + print("[upload_udf] metadata row inserted: name={} md5={} size={}".format(udf_name, md5, size), file=sys.stderr) + + except Exception as exc: + print("[upload_udf] ERROR: {}".format(exc), file=sys.stderr) + return 1 + + # Print md5 to stdout so callers can capture it. + print(md5) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/ydb/tests/functional/udf_store/upload_udf/ya.make b/ydb/tests/functional/udf_store/upload_udf/ya.make new file mode 100644 index 00000000000..7bf85be0153 --- /dev/null +++ b/ydb/tests/functional/udf_store/upload_udf/ya.make @@ -0,0 +1,16 @@ +PY3_PROGRAM(upload_udf) + +PY_SRCS( + __main__.py +) + +PEERDIR( + ydb/tests/functional/udf_store/lib + ydb/tests/oss/ydb_sdk_import +) + +DEPENDS( + ydb/tests/stress/kv_volume_tool +) + +END() diff --git a/ydb/tests/functional/udf_store/ya.make b/ydb/tests/functional/udf_store/ya.make new file mode 100644 index 00000000000..7261be6af8c --- /dev/null +++ b/ydb/tests/functional/udf_store/ya.make @@ -0,0 +1,37 @@ +PY3TEST() + +INCLUDE(${ARCADIA_ROOT}/ydb/tests/harness_dep.inc) + +ENV(YDB_KV_VOLUME_TOOL_PATH="ydb/tests/stress/kv_volume_tool/kv_volume_tool") +ENV(YDB_DICTS_UDF_PATH="yql/essentials/udfs/examples/dicts/libdicts_udf.so") +ENV(YDB_UPLOAD_UDF_PATH="ydb/tests/functional/udf_store/upload_udf/upload_udf") + +TEST_SRCS( + test_udf_store.py +) + +SPLIT_FACTOR(10) + +DEPENDS( + ydb/tests/stress/kv_volume_tool + yql/essentials/udfs/examples/dicts + ydb/tests/functional/udf_store/upload_udf +) + +PEERDIR( + ydb/tests/functional/udf_store/lib + ydb/tests/library + ydb/tests/oss/ydb_sdk_import +) + +FORK_SUBTESTS() + +IF (SANITIZER_TYPE) + SIZE(LARGE) + INCLUDE(${ARCADIA_ROOT}/ydb/tests/large.inc) + REQUIREMENTS(ram:10 cpu:16) +ELSE() + SIZE(MEDIUM) +ENDIF() + +END() |