TImportClient::ListObjectsInS3Export in C++ SDK (#18366)

UgnineSirdis · Gazizonoki · commit 91c43f35f1b6 · 2025-05-23T18:56:18.000+03:00
diff --git a/include/ydb-cpp-sdk/client/import/import.h b/include/ydb-cpp-sdk/client/import/import.h
@@ -5,6 +5,10 @@
 
 #include <ydb-cpp-sdk/client/types/s3_settings.h>
 
+namespace Ydb::Import {
+class ListObjectsInS3ExportResult;
+}
+
 namespace NYdb::inline V3 {
 namespace NImport {
 
@@ -77,6 +81,49 @@ class TImportFromS3Response : public TOperation {
     TMetadata Metadata_;
 };
 
+using TAsyncImportFromS3Response = NThreading::TFuture<TImportFromS3Response>;
+
+struct TListObjectsInS3ExportSettings : public TOperationRequestSettings<TListObjectsInS3ExportSettings>,
+                                        public TS3Settings<TListObjectsInS3ExportSettings> {
+    using TSelf = TListObjectsInS3ExportSettings;
+
+    struct TItem {
+        // Database object path.
+        std::string Path = {};
+    };
+
+    FLUENT_SETTING_VECTOR(TItem, Item);
+    FLUENT_SETTING_OPTIONAL(uint32_t, NumberOfRetries);
+    FLUENT_SETTING_OPTIONAL(std::string, Prefix);
+    FLUENT_SETTING_OPTIONAL(std::string, SymmetricKey);
+};
+
+class TListObjectsInS3ExportResult : public TStatus {
+public:
+    struct TItem {
+        // S3 object prefix
+        std::string Prefix;
+
+        // Database object path
+        std::string Path;
+
+        void Out(IOutputStream& out) const;
+    };
+
+    TListObjectsInS3ExportResult(TStatus&& status, const ::Ydb::Import::ListObjectsInS3ExportResult& proto);
+
+    const std::vector<TItem>& GetItems() const;
+    const std::string& NextPageToken() const { return NextPageToken_; }
+
+    void Out(IOutputStream& out) const;
+
+private:
+    std::vector<TItem> Items_;
+    std::string NextPageToken_;
+};
+
+using TAsyncListObjectsInS3ExportResult = NThreading::TFuture<TListObjectsInS3ExportResult>;
+
 /// Data
 struct TImportYdbDumpDataSettings : public TOperationRequestSettings<TImportYdbDumpDataSettings> {
     using TSelf = TImportYdbDumpDataSettings;
@@ -99,7 +146,9 @@ class TImportClient {
 public:
     TImportClient(const TDriver& driver, const TCommonClientSettings& settings = TCommonClientSettings());
 
-    NThreading::TFuture<TImportFromS3Response> ImportFromS3(const TImportFromS3Settings& settings);
+    TAsyncImportFromS3Response ImportFromS3(const TImportFromS3Settings& settings);
+
+    TAsyncListObjectsInS3ExportResult ListObjectsInS3Export(const TListObjectsInS3ExportSettings& settings, std::int64_t pageSize = 0, const std::string& pageToken = {});
 
     // ydb dump format
     TAsyncImportDataResult ImportData(const std::string& table, std::string&& data, const TImportYdbDumpDataSettings& settings);
@@ -111,8 +160,3 @@ class TImportClient {
 
 } // namespace NImport
 } // namespace NYdb
-
-template<>
-inline void Out<NYdb::NImport::TImportFromS3Response>(IOutputStream& o, const NYdb::NImport::TImportFromS3Response& x) {
-    return x.Out(o);
-}
diff --git a/src/api/grpc/ydb_import_v1.proto b/src/api/grpc/ydb_import_v1.proto
@@ -11,6 +11,9 @@ service ImportService {
     // Method starts an asynchronous operation that can be cancelled while it is in progress.
     rpc ImportFromS3(Import.ImportFromS3Request) returns (Import.ImportFromS3Response);
 
+    // List objects from existing export stored in S3 bucket
+    rpc ListObjectsInS3Export(Import.ListObjectsInS3ExportRequest) returns (Import.ListObjectsInS3ExportResponse);
+
     // Writes data to a table.
     // Method accepts serialized data in the selected format and writes it non-transactionally.
     rpc ImportData(Import.ImportDataRequest) returns (Import.ImportDataResponse);
diff --git a/src/api/protos/ydb_import.proto b/src/api/protos/ydb_import.proto
@@ -121,6 +121,87 @@ message ImportFromS3Response {
     Ydb.Operations.Operation operation = 1;
 }
 
+message ListObjectsInS3ExportSettings {
+    message Item {
+        // Database object path
+        // Recursive for directories
+        string path = 1;
+    }
+
+    string endpoint = 1 [(required) = true];
+    ImportFromS3Settings.Scheme scheme = 2; // HTTPS if not specified
+    string bucket = 3 [(required) = true];
+    string access_key = 4 [(required) = true];
+    string secret_key = 5 [(required) = true];
+    repeated Item items = 6;
+    uint32 number_of_retries = 7;
+
+    // Region to use in requests
+    string region = 8;
+
+    // disables virtual hosting style buckets aws s3 feature
+    // it changes the way bucket appended to url. e.g. https//bucket_name.example.com/ vs https://example.com/bucket_name
+    // details: https://docs.aws.amazon.com/AmazonS3/latest/userguide/VirtualHosting.html
+    // it is especially useful for custom s3 implementations
+    bool disable_virtual_addressing = 9;
+
+    // A default path prefix for all items,
+    // determines that the import works with the list of objects in the SchemaMapping file.
+    // Must be provided for encrypted exports.
+    string prefix = 10;
+
+    // Settings how data is encrypted.
+    // If encryption_settings field is not specified,
+    // the resulting data is considered not encrypted.
+    Ydb.Export.EncryptionSettings encryption_settings = 11;
+}
+
+message ListObjectsInS3ExportResult {
+    message Item {
+        /* YDB database objects in S3 are stored in one or more S3 objects (see ydb_export.proto).
+          The S3 object name begins with a prefix, followed by:
+            * '/data_PartNumber', where 'PartNumber' represents the index of the part, starting at zero;
+            * '/scheme.pb' - object with information about scheme, indexes, etc;
+            * '/permissions.pb' - object with information about ACL and owner.
+        */
+
+        // S3 object prefix
+        string prefix = 1;
+
+        // Database object path
+        string path = 2;
+    }
+
+    repeated Item items = 1;
+
+    // This token allows you to get the next page of results for ListObjectsInS3Export requests,
+    // if the number of results is larger than `page_size` specified in the request.
+    // To get the next page, specify the value of `next_page_token` as a value for
+    // the `page_token` parameter in the next ListObjectsInS3Export request. Subsequent ListObjectsInS3Export
+    // requests will have their own `next_page_token` to continue paging through the results.
+    string next_page_token = 2;
+}
+
+message ListObjectsInS3ExportRequest {
+    Ydb.Operations.OperationParams operation_params = 1;
+    ListObjectsInS3ExportSettings settings = 2 [(required) = true];
+
+    // The maximum number of results per page that should be returned. If the number of available
+    // results is larger than `page_size`, the service returns a `next_page_token` that can be used
+    // to get the next page of results in subsequent ListObjectsInS3Export requests.
+    // 0 means that server returns all objects.
+    int64 page_size = 3 [(value) = "<= 10000"];
+
+    // Page token. Set `page_token` to the `next_page_token` returned by a previous ListObjectsInS3Export
+    // request to get the next page of results.
+    string page_token = 4;
+}
+
+message ListObjectsInS3ExportResponse {
+    // operation.result = ListObjectsInS3ExportResult
+    Ydb.Operations.Operation operation = 1;
+}
+
 /// Data
 message YdbDumpFormat {
     repeated string columns = 1;
diff --git a/src/client/import/import.cpp b/src/client/import/import.cpp
@@ -10,6 +10,8 @@
 #include <src/client/common_client/impl/client.h>
 #include <ydb-cpp-sdk/client/proto/accessor.h>
 
+#include <util/string/join.h>
+
 namespace NYdb::inline V3 {
 namespace NImport {
 
@@ -34,6 +36,25 @@ std::vector<TImportItemProgress> ItemsProgressFromProto(const google::protobuf::
     return result;
 }
 
+template <class TS3SettingsProto, class TSettings>
+void FillS3Settings(TS3SettingsProto& proto, const TSettings& settings) {
+    proto.set_endpoint(TStringType{settings.Endpoint_});
+    proto.set_scheme(TProtoAccessor::GetProto<ImportFromS3Settings>(settings.Scheme_));
+    proto.set_bucket(TStringType{settings.Bucket_});
+    proto.set_access_key(TStringType{settings.AccessKey_});
+    proto.set_secret_key(TStringType{settings.SecretKey_});
+
+    if (settings.NumberOfRetries_) {
+        proto.set_number_of_retries(settings.NumberOfRetries_.value());
+    }
+
+    if (settings.SymmetricKey_) {
+        proto.mutable_encryption_settings()->mutable_symmetric_key()->set_key(*settings.SymmetricKey_);
+    }
+
+    proto.set_disable_virtual_addressing(!settings.UseVirtualAddressing_);
+}
+
 } // anonymous
 
 /// S3
@@ -66,6 +87,36 @@ const TImportFromS3Response::TMetadata& TImportFromS3Response::Metadata() const
     return Metadata_;
 }
 
+TListObjectsInS3ExportResult::TListObjectsInS3ExportResult(TStatus&& status, const ::Ydb::Import::ListObjectsInS3ExportResult& proto)
+    : TStatus(std::move(status))
+{
+    Items_.reserve(proto.items_size());
+    for (const auto& item : proto.items()) {
+        Items_.emplace_back(TItem{
+            .Prefix = item.prefix(),
+            .Path = item.path()
+        });
+    }
+    NextPageToken_ = proto.next_page_token();
+}
+
+const std::vector<TListObjectsInS3ExportResult::TItem>& TListObjectsInS3ExportResult::GetItems() const {
+    return Items_;
+}
+
+void TListObjectsInS3ExportResult::Out(IOutputStream& out) const {
+    if (IsSuccess()) {
+        out << "{ items: [" << JoinSeq(", ", Items_) << "], next_page_token: \"" << NextPageToken_ << "\" }";
+    } else {
+        return TStatus::Out(out);
+    }
+}
+
+void TListObjectsInS3ExportResult::TItem::Out(IOutputStream& out) const {
+    out << "{ prefix: \"" << Prefix << "\""
+        << ", path: \"" << Path << "\" }";
+}
+
 /// Data
 TImportDataResult::TImportDataResult(TStatus&& status)
     : TStatus(std::move(status))
@@ -80,13 +131,37 @@ class TImportClient::TImpl : public TClientImplCommon<TImportClient::TImpl> {
     {
     }
 
-    TFuture<TImportFromS3Response> ImportFromS3(ImportFromS3Request&& request, const TImportFromS3Settings& settings) {
+    TAsyncImportFromS3Response ImportFromS3(ImportFromS3Request&& request, const TImportFromS3Settings& settings) {
         return RunOperation<V1::ImportService, ImportFromS3Request, ImportFromS3Response, TImportFromS3Response>(
             std::move(request),
             &V1::ImportService::Stub::AsyncImportFromS3,
             TRpcRequestSettings::Make(settings));
     }
 
+    TAsyncListObjectsInS3ExportResult ListObjectsInS3Export(ListObjectsInS3ExportRequest&& request, const TListObjectsInS3ExportSettings& settings) {
+        auto promise = NThreading::NewPromise<TListObjectsInS3ExportResult>();
+
+        auto extractor = [promise]
+            (google::protobuf::Any* any, TPlainStatus status) mutable {
+                ListObjectsInS3ExportResult result;
+                if (any) {
+                    any->UnpackTo(&result);
+                }
+
+                promise.SetValue(TListObjectsInS3ExportResult(TStatus(std::move(status)), result));
+            };
+
+        Connections_->RunDeferred<V1::ImportService, ListObjectsInS3ExportRequest, ListObjectsInS3ExportResponse>(
+            std::move(request),
+            extractor,
+            &V1::ImportService::Stub::AsyncListObjectsInS3Export,
+            DbDriverState_,
+            INITIAL_DEFERRED_CALL_DELAY,
+            TRpcRequestSettings::Make(settings));
+
+        return promise.GetFuture();
+    }
+
     template <typename TSettings>
     TAsyncImportDataResult ImportData(ImportDataRequest&& request, const TSettings& settings) {
         auto promise = NThreading::NewPromise<TImportDataResult>();
@@ -131,22 +206,18 @@ TImportClient::TImportClient(const TDriver& driver, const TCommonClientSettings&
 {
 }
 
-TFuture<TImportFromS3Response> TImportClient::ImportFromS3(const TImportFromS3Settings& settings) {
+TAsyncImportFromS3Response TImportClient::ImportFromS3(const TImportFromS3Settings& settings) {
     auto request = MakeOperationRequest<ImportFromS3Request>(settings);
-
-    request.mutable_settings()->set_endpoint(TStringType{settings.Endpoint_});
-    request.mutable_settings()->set_scheme(TProtoAccessor::GetProto<ImportFromS3Settings>(settings.Scheme_));
-    request.mutable_settings()->set_bucket(TStringType{settings.Bucket_});
-    request.mutable_settings()->set_access_key(TStringType{settings.AccessKey_});
-    request.mutable_settings()->set_secret_key(TStringType{settings.SecretKey_});
+    Ydb::Import::ImportFromS3Settings& settingsProto = *request.mutable_settings();
+    FillS3Settings(settingsProto, settings);
 
     for (const auto& item : settings.Item_) {
         if (!item.Src.empty() && !item.SrcPath.empty()) {
             throw TContractViolation(
                 TStringBuilder() << "Invalid item: both source prefix and source path are set: \"" << item.Src << "\" and \"" << item.SrcPath << "\"");
         }
 
-        auto& protoItem = *request.mutable_settings()->mutable_items()->Add();
+        auto& protoItem = *settingsProto.mutable_items()->Add();
         if (!item.Src.empty()) {
             protoItem.set_source_prefix(item.Src);
         }
@@ -157,32 +228,47 @@ TFuture<TImportFromS3Response> TImportClient::ImportFromS3(const TImportFromS3Se
     }
 
     if (settings.Description_) {
-        request.mutable_settings()->set_description(TStringType{settings.Description_.value()});
-    }
-
-    if (settings.NumberOfRetries_) {
-        request.mutable_settings()->set_number_of_retries(settings.NumberOfRetries_.value());
+        settingsProto.set_description(TStringType{settings.Description_.value()});
     }
 
     if (settings.NoACL_) {
-        request.mutable_settings()->set_no_acl(settings.NoACL_.value());
+        settingsProto.set_no_acl(settings.NoACL_.value());
     }
 
     if (settings.SourcePrefix_) {
-        request.mutable_settings()->set_source_prefix(settings.SourcePrefix_.value());
+        settingsProto.set_source_prefix(settings.SourcePrefix_.value());
     }
 
     if (settings.DestinationPath_) {
-        request.mutable_settings()->set_destination_path(settings.DestinationPath_.value());
+        settingsProto.set_destination_path(settings.DestinationPath_.value());
     }
 
-    if (settings.SymmetricKey_) {
-        request.mutable_settings()->mutable_encryption_settings()->mutable_symmetric_key()->set_key(*settings.SymmetricKey_);
+    return Impl_->ImportFromS3(std::move(request), settings);
+}
+
+TAsyncListObjectsInS3ExportResult TImportClient::ListObjectsInS3Export(const TListObjectsInS3ExportSettings& settings, std::int64_t pageSize, const std::string& pageToken) {
+    auto request = MakeOperationRequest<ListObjectsInS3ExportRequest>(settings);
+    Ydb::Import::ListObjectsInS3ExportSettings& settingsProto = *request.mutable_settings();
+    FillS3Settings(settingsProto, settings);
+
+    if (settings.Prefix_) {
+        settingsProto.set_prefix(settings.Prefix_.value());
     }
 
-    request.mutable_settings()->set_disable_virtual_addressing(!settings.UseVirtualAddressing_);
+    for (const auto& item : settings.Item_) {
+        if (item.Path.empty()) {
+            throw TContractViolation(
+                TStringBuilder() << "Invalid item: path is not set");
+        }
 
-    return Impl_->ImportFromS3(std::move(request), settings);
+        settingsProto.add_items()->set_path(item.Path);
+    }
+
+    // Paging
+    request.set_page_size(pageSize);
+    request.set_page_token(pageToken);
+
+    return Impl_->ListObjectsInS3Export(std::move(request), settings);
 }
 
 TAsyncImportDataResult TImportClient::ImportData(const std::string& table, std::string&& data, const TImportYdbDumpDataSettings& settings) {
diff --git a/src/client/import/out.cpp b/src/client/import/out.cpp
@@ -3,3 +3,15 @@
 Y_DECLARE_OUT_SPEC(, NYdb::NImport::TImportDataResult, o, x) {
     return x.Out(o);
 }
+
+Y_DECLARE_OUT_SPEC(, NYdb::NImport::TImportFromS3Response, o, x) {
+    return x.Out(o);
+}
+
+Y_DECLARE_OUT_SPEC(, NYdb::NImport::TListObjectsInS3ExportResult, o, x) {
+    return x.Out(o);
+}
+
+Y_DECLARE_OUT_SPEC(, NYdb::NImport::TListObjectsInS3ExportResult::TItem, o, x) {
+    return x.Out(o);
+}
