Added try catch

GrigoriyPA · Jan 23, 2024 · 5819339 · 5819339
1 parent ebea693
commit 5819339
Showing 1 changed file with 26 additions and 2 deletions.
diff --git a/ydb/library/yql/providers/s3/actors/yql_s3_write_actor.cpp b/ydb/library/yql/providers/s3/actors/yql_s3_write_actor.cpp
@@ -139,7 +139,7 @@ class TS3FileWriteActor : public TActorBootstrapped<TS3FileWriteActor> {
         const TString& token)
         : TxId(txId)
         , Gateway(std::move(gateway))
-        , AuthInfo(crdentials.GetAuthInfo())
+        , Credentials(std::move(crdentials))
         , RetryPolicy(retryPolicy)
         , ActorSystem(TActivationContext::ActorSystem())
         , Key(key)
@@ -155,6 +155,9 @@ class TS3FileWriteActor : public TActorBootstrapped<TS3FileWriteActor> {
     void Bootstrap(const TActorId& parentId) {
         ParentId = parentId;
         LOG_D("TS3FileWriteActor", "Bootstrap by " << ParentId << " for Key: [" << Key << "], Url: [" << Url << "], request id: [" << RequestId << "]");
+        if (!UpdateAuthInfo()) {
+            return;
+        }
         if (DirtyWrite && Parts->IsSealed() && Parts->Size() <= 1) {
             Become(&TS3FileWriteActor::SinglepartWorkingStateFunc);
             const size_t size = Max<size_t>(Parts->Volume(), 1);
@@ -359,13 +362,27 @@ class TS3FileWriteActor : public TActorBootstrapped<TS3FileWriteActor> {
         }
     }
 
+    bool UpdateAuthInfo() {
+        try {
+            AuthInfo = Credentials.GetAuthInfo();
+            return true;
+        }
+        catch (const yexception& ex) {
+            Send(ParentId, new TEvPrivate::TEvUploadError(NYql::NDqProto::StatusIds::BAD_REQUEST, TStringBuilder() << "Failed to get auth info: " << ex.what()));
+            return false;
+        }
+    }
+
     void StartUploadParts() {
         while (auto part = Parts->Pop()) {
             const auto size = part.size();
             const auto index = Tags.size();
             Tags.emplace_back();
             InFlight += size;
             SentSize += size;
+            if (!UpdateAuthInfo()) {
+                return;
+            }
             Gateway->Upload(Url + "?partNumber=" + std::to_string(index + 1) + "&uploadId=" + UploadId,
                 IHTTPGateway::MakeYcHeaders(RequestId, AuthInfo.GetToken(), {}, AuthInfo.GetAwsUserPwd(), AuthInfo.GetAwsSigV4()),
                 std::move(part),
@@ -393,6 +410,9 @@ class TS3FileWriteActor : public TActorBootstrapped<TS3FileWriteActor> {
         for (const auto& tag : Tags)
             xml << "<Part><PartNumber>" << ++i << "</PartNumber><ETag>" << tag << "</ETag></Part>" << Endl;
         xml << "</CompleteMultipartUpload>" << Endl;
+        if (!UpdateAuthInfo()) {
+            return;
+        }
         Gateway->Upload(Url + "?uploadId=" + UploadId,
             IHTTPGateway::MakeYcHeaders(RequestId, AuthInfo.GetToken(), "application/xml", AuthInfo.GetAwsUserPwd(), AuthInfo.GetAwsSigV4()),
             xml,
@@ -409,6 +429,9 @@ class TS3FileWriteActor : public TActorBootstrapped<TS3FileWriteActor> {
             return;
         }
 
+        if (!UpdateAuthInfo()) {
+            return;
+        }
         Gateway->Delete(Url + "?uploadId=" + UploadId,
             IHTTPGateway::MakeYcHeaders(RequestId, AuthInfo.GetToken(), "application/xml", AuthInfo.GetAwsUserPwd(), AuthInfo.GetAwsSigV4()),
             std::bind(&TS3FileWriteActor::OnMultipartUploadAbort, ActorSystem, SelfId(), TxId, RequestId, std::placeholders::_1),
@@ -421,7 +444,7 @@ class TS3FileWriteActor : public TActorBootstrapped<TS3FileWriteActor> {
 
     const TTxId TxId;
     const IHTTPGateway::TPtr Gateway;
-    const TS3Credentials::TAuthInfo AuthInfo;
+    const TS3Credentials Credentials;
     const IHTTPGateway::TRetryPolicy::TPtr RetryPolicy;
 
     TActorSystem* const ActorSystem;
@@ -437,6 +460,7 @@ class TS3FileWriteActor : public TActorBootstrapped<TS3FileWriteActor> {
     TString UploadId;
     bool DirtyWrite;
     TString Token;
+    TS3Credentials::TAuthInfo AuthInfo;
 };
 
 class TS3WriteActor : public TActorBootstrapped<TS3WriteActor>, public IDqComputeActorAsyncOutput {