external storage type usage

7 files changed, 59 insertions, 35 deletions

diff --git a/ydb/core/protos/flat_scheme_op.proto b/ydb/core/protos/flat_scheme_op.proto
index 826e62b31ff..ae6158d65a4 100644
--- a/ydb/core/protos/flat_scheme_op.proto
+++ b/ydb/core/protos/flat_scheme_op.proto
@@ -879,7 +879,7 @@ message TS3Settings {
     optional string ObjectKeyPattern = 4; // dst for backup, src for restore
     optional string AccessKey = 5;
     optional string SecretKey = 6;
-    optional Ydb.Export.ExportToS3Settings.StorageClass StorageClass = 7;
+    optional Ydb.Export.ExportToS3Settings.StorageClass StorageClass = 7 [default = STANDARD];
     optional bool VerifySSL = 8;
     optional string ProxyHost = 9;
     optional uint32 ProxyPort = 10;
diff --git a/ydb/core/tx/datashard/extstorage_usage_config.cpp b/ydb/core/tx/datashard/extstorage_usage_config.cpp
index 20a196e290a..5dce5bad6cd 100644
--- a/ydb/core/tx/datashard/extstorage_usage_config.cpp
+++ b/ydb/core/tx/datashard/extstorage_usage_config.cpp
@@ -1,34 +1,11 @@
 #include "extstorage_usage_config.h"
+#include <ydb/core/wrappers/s3_storage_config.h>
 
 #ifndef KIKIMR_DISABLE_S3_OPS
 namespace NKikimr::NDataShard {
 
 Aws::S3::Model::StorageClass TS3Settings::GetStorageClass() const {
-    using ExportToS3Settings = Ydb::Export::ExportToS3Settings;
-    using AwsStorageClass = Aws::S3::Model::StorageClass;
-
-    switch (StorageClass) {
-        case ExportToS3Settings::STORAGE_CLASS_UNSPECIFIED:
-            return AwsStorageClass::NOT_SET;
-        case ExportToS3Settings::STANDARD:
-            return AwsStorageClass::STANDARD;
-        case ExportToS3Settings::REDUCED_REDUNDANCY:
-            return AwsStorageClass::REDUCED_REDUNDANCY;
-        case ExportToS3Settings::STANDARD_IA:
-            return AwsStorageClass::STANDARD_IA;
-        case ExportToS3Settings::ONEZONE_IA:
-            return AwsStorageClass::ONEZONE_IA;
-        case ExportToS3Settings::INTELLIGENT_TIERING:
-            return AwsStorageClass::INTELLIGENT_TIERING;
-        case ExportToS3Settings::GLACIER:
-            return AwsStorageClass::GLACIER;
-        case ExportToS3Settings::DEEP_ARCHIVE:
-            return AwsStorageClass::DEEP_ARCHIVE;
-        case ExportToS3Settings::OUTPOSTS:
-            return AwsStorageClass::OUTPOSTS;
-        default:
-            return AwsStorageClass::NOT_SET;
-    }
+    return NKikimr::NWrappers::NExternalStorage::TS3ExternalStorageConfig::ConvertStorageClass(StorageClass);
 }
 
 }
diff --git a/ydb/core/tx/tiering/s3_actor.cpp b/ydb/core/tx/tiering/s3_actor.cpp
index 566931e2a44..abe1eb28a5f 100644
--- a/ydb/core/tx/tiering/s3_actor.cpp
+++ b/ydb/core/tx/tiering/s3_actor.cpp
@@ -419,7 +419,6 @@ private:
     void SendPutObject(const TString& key, TString&& data) const {
         auto request = Aws::S3::Model::PutObjectRequest()
             .WithKey(key);
-            //.WithStorageClass(Aws::S3::Model::StorageClass::STANDARD_IA); // TODO: move to config
 #if 0
         Aws::Map<Aws::String, Aws::String> metadata;
         metadata.emplace("Content-Type", "application/x-compressed");
diff --git a/ydb/core/wrappers/s3_storage.cpp b/ydb/core/wrappers/s3_storage.cpp
index 55912060365..b4d08cd5694 100644
--- a/ydb/core/wrappers/s3_storage.cpp
+++ b/ydb/core/wrappers/s3_storage.cpp
@@ -27,9 +27,10 @@ protected:
     using TRequest = typename TEvRequest::TRequest;
     using TOutcome = typename TEvResponse::TOutcome;
 public:
-    explicit TCommonContextBase(const TActorSystem* sys, const TActorId& sender, IRequestContext::TPtr requestContext)
+    explicit TCommonContextBase(const TActorSystem* sys, const TActorId& sender, IRequestContext::TPtr requestContext, const Aws::S3::Model::StorageClass storageClass)
         : AsyncCallerContext()
         , RequestContext(requestContext)
+        , StorageClass(storageClass)
         , ActorSystem(sys)
         , Sender(sender)
     {
@@ -54,6 +55,7 @@ protected:
 
     mutable bool Replied = false;
     IRequestContext::TPtr RequestContext;
+    const Aws::S3::Model::StorageClass StorageClass;
 private:
     const TActorSystem* ActorSystem;
     const TActorId Sender;
@@ -200,6 +202,8 @@ private:
 
 template <typename TEvRequest, typename TEvResponse>
 class TInputStreamContext: public TContextBase<TEvRequest, TEvResponse> {
+private:
+    using TBase = TContextBase<TEvRequest, TEvResponse>;
 protected:
     using TRequest = typename TEvRequest::TRequest;
     using TOutcome = typename TEvResponse::TOutcome;
@@ -221,11 +225,10 @@ private:
     };
 
 public:
-    using TContextBase<TEvRequest, TEvResponse>::TContextBase;
+    using TBase::TBase;
 
     const TRequest& PrepareRequest(typename TEvRequest::TPtr& ev) override {
         auto& request = ev->Get()->MutableRequest();
-
         Buffer = std::move(ev->Get()->Body);
         request.SetBody(MakeShared<DefaultUnderlyingStream>("StreamContext",
             MakeUnique<TInputStreamBuf>("StreamContext", Buffer)));
@@ -238,6 +241,20 @@ private:
 
 }; // TInputStreamContext
 
+template <typename TEvRequest, typename TEvResponse>
+class TPutInputStreamContext: public TInputStreamContext<TEvRequest, TEvResponse> {
+private:
+    using TBase = TInputStreamContext<TEvRequest, TEvResponse>;
+public:
+    using TBase::TBase;
+
+    const typename TBase::TRequest& PrepareRequest(typename TEvRequest::TPtr& ev) override {
+        auto& request = ev->Get()->MutableRequest();
+        request.WithStorageClass(TBase::StorageClass);
+        return TBase::PrepareRequest(ev);
+    }
+}; // TPutInputStreamContext
+
 } // anonymous
 
 TS3ExternalStorage::~TS3ExternalStorage() {
@@ -267,7 +284,7 @@ void TS3ExternalStorage::Execute(TEvHeadObjectRequest::TPtr& ev) const {
 }
 
 void TS3ExternalStorage::Execute(TEvPutObjectRequest::TPtr& ev) const {
-    Call<TEvPutObjectRequest, TEvPutObjectResponse, TInputStreamContext>(
+    Call<TEvPutObjectRequest, TEvPutObjectResponse, TPutInputStreamContext>(
         ev, &S3Client::PutObjectAsync);
 }
 
diff --git a/ydb/core/wrappers/s3_storage.h b/ydb/core/wrappers/s3_storage.h
index 4317570bb97..e452704e6f4 100644
--- a/ydb/core/wrappers/s3_storage.h
+++ b/ydb/core/wrappers/s3_storage.h
@@ -33,6 +33,7 @@ private:
     const Aws::Client::ClientConfiguration Config;
     const Aws::Auth::AWSCredentials Credentials;
     const TString Bucket;
+    const Aws::S3::Model::StorageClass StorageClass = Aws::S3::Model::StorageClass::STANDARD;
 
     template <typename TRequest, typename TOutcome>
     using THandler = std::function<void(const Aws::S3::S3Client*, const TRequest&, const TOutcome&, const std::shared_ptr<const Aws::Client::AsyncCallerContext>&)>;
@@ -45,7 +46,7 @@ private:
         using TCtx = TContext<TEvRequest, TEvResponse>;
         ev->Get()->MutableRequest().WithBucket(Bucket);
 
-        auto ctx = std::make_shared<TCtx>(TlsActivationContext->ActorSystem(), ev->Sender, ev->Get()->GetRequestContext());
+        auto ctx = std::make_shared<TCtx>(TlsActivationContext->ActorSystem(), ev->Sender, ev->Get()->GetRequestContext(), StorageClass);
         auto callback = [](
             const Aws::S3::S3Client*,
             const typename TEvRequest::TRequest& request,
@@ -66,11 +67,14 @@ private:
     }
 
 public:
-    TS3ExternalStorage(const Aws::Client::ClientConfiguration& config, const Aws::Auth::AWSCredentials& credentials, const TString& bucket)
+    TS3ExternalStorage(const Aws::Client::ClientConfiguration& config,
+        const Aws::Auth::AWSCredentials& credentials,
+        const TString& bucket, const Aws::S3::Model::StorageClass storageClass)
         : Client(new Aws::S3::S3Client(credentials, config))
         , Config(config)
         , Credentials(credentials)
         , Bucket(bucket)
+        , StorageClass(storageClass)
     {
     }
 
diff --git a/ydb/core/wrappers/s3_storage_config.cpp b/ydb/core/wrappers/s3_storage_config.cpp
index 0032e45a519..3a1926026b1 100644
--- a/ydb/core/wrappers/s3_storage_config.cpp
+++ b/ydb/core/wrappers/s3_storage_config.cpp
@@ -183,7 +183,7 @@ TString TS3ExternalStorageConfig::DoGetStorageId() const {
 }
 
 IExternalStorageOperator::TPtr TS3ExternalStorageConfig::DoConstructStorageOperator() const {
-    return std::make_shared<TS3ExternalStorage>(Config, Credentials, Bucket);
+    return std::make_shared<TS3ExternalStorage>(Config, Credentials, Bucket, StorageClass);
 }
 
 TS3ExternalStorageConfig::TS3ExternalStorageConfig(const Ydb::Import::ImportFromS3Settings& settings): Config(ConfigFromSettings(settings))
@@ -203,10 +203,35 @@ TS3ExternalStorageConfig::TS3ExternalStorageConfig(const Aws::Auth::AWSCredentia
 TS3ExternalStorageConfig::TS3ExternalStorageConfig(const NKikimrSchemeOp::TS3Settings& settings)
     : Config(ConfigFromSettings(settings))
     , Credentials(CredentialsFromSettings(settings))
+    , StorageClass(ConvertStorageClass(settings.GetStorageClass()))
 {
     Bucket = settings.GetBucket();
 }
 
+Aws::S3::Model::StorageClass TS3ExternalStorageConfig::ConvertStorageClass(const Ydb::Export::ExportToS3Settings::StorageClass storage) {
+    switch (storage) {
+        case Ydb::Export::ExportToS3Settings::STANDARD:
+            return Aws::S3::Model::StorageClass::STANDARD;
+        case Ydb::Export::ExportToS3Settings::STANDARD_IA:
+            return Aws::S3::Model::StorageClass::STANDARD_IA;
+        case Ydb::Export::ExportToS3Settings::REDUCED_REDUNDANCY:
+            return Aws::S3::Model::StorageClass::REDUCED_REDUNDANCY;
+        case Ydb::Export::ExportToS3Settings::ONEZONE_IA:
+            return Aws::S3::Model::StorageClass::ONEZONE_IA;
+        case Ydb::Export::ExportToS3Settings::INTELLIGENT_TIERING:
+            return Aws::S3::Model::StorageClass::INTELLIGENT_TIERING;
+        case Ydb::Export::ExportToS3Settings::GLACIER:
+            return Aws::S3::Model::StorageClass::GLACIER;
+        case Ydb::Export::ExportToS3Settings::DEEP_ARCHIVE:
+            return Aws::S3::Model::StorageClass::DEEP_ARCHIVE;
+        case Ydb::Export::ExportToS3Settings::OUTPOSTS:
+            return Aws::S3::Model::StorageClass::OUTPOSTS;
+        case Ydb::Export::ExportToS3Settings::STORAGE_CLASS_UNSPECIFIED:
+        default:
+            return Aws::S3::Model::StorageClass::NOT_SET;
+    }
+}
+
 }
 
 #endif // KIKIMR_DISABLE_S3_OPS
diff --git a/ydb/core/wrappers/s3_storage_config.h b/ydb/core/wrappers/s3_storage_config.h
index 07fb87b9fb3..c78ea13aa54 100644
--- a/ydb/core/wrappers/s3_storage_config.h
+++ b/ydb/core/wrappers/s3_storage_config.h
@@ -28,16 +28,18 @@ private:
     YDB_READONLY_DEF(TString, Bucket);
     Aws::Client::ClientConfiguration Config;
     const Aws::Auth::AWSCredentials Credentials;
+    YDB_READONLY(Aws::S3::Model::StorageClass, StorageClass, Aws::S3::Model::StorageClass::STANDARD);
 
     static Aws::Client::ClientConfiguration ConfigFromSettings(const NKikimrSchemeOp::TS3Settings& settings);
     static Aws::Auth::AWSCredentials CredentialsFromSettings(const NKikimrSchemeOp::TS3Settings& settings);
     static Aws::Client::ClientConfiguration ConfigFromSettings(const Ydb::Import::ImportFromS3Settings& settings);
     static Aws::Auth::AWSCredentials CredentialsFromSettings(const Ydb::Import::ImportFromS3Settings& settings);
-
 protected:
     virtual TString DoGetStorageId() const override;
     virtual IExternalStorageOperator::TPtr DoConstructStorageOperator() const override;
 public:
+    static Aws::S3::Model::StorageClass ConvertStorageClass(const Ydb::Export::ExportToS3Settings::StorageClass storage);
+
     const Aws::Client::ClientConfiguration& GetConfig() const {
         return Config;
     }