From 868f4154c27c4820655dd17b88f81f07f545fd01 Mon Sep 17 00:00:00 2001
From: zombee0 <ewang2027@gmail.com>
Date: Tue, 5 Nov 2024 14:01:39 +0800
Subject: [PATCH] [BugFix]don't set NETWORK_CONNECTION error when received
 response from oss (#52569)

Signed-off-by: zombee0 <ewang2027@gmail.com>
---
 be/src/fs/s3/poco_http_client.cpp    | 13 ++-----------
 be/test/fs/poco_http_client_test.cpp | 26 ++++++++++++++++++++++++++
 be/test/io/s3_input_stream_test.cpp  |  9 +++++++++
 3 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/be/src/fs/s3/poco_http_client.cpp b/be/src/fs/s3/poco_http_client.cpp
index 50cfd63a38d7b..8c1dbf0ece4f8 100644
--- a/be/src/fs/s3/poco_http_client.cpp
+++ b/be/src/fs/s3/poco_http_client.cpp
@@ -182,17 +182,8 @@ void PocoHttpClient::MakeRequestInternal(Aws::Http::HttpRequest& request,
                     // pass
                     // TODO, throttling
                 }
-                if (status_code >= ResponseCode::SUCCESS_RESPONSE_MIN &&
-                    status_code <= ResponseCode::SUCCESS_RESPONSE_MAX) {
-                    if (!poco_response.isBodyFilled()) {
-                        Poco::StreamCopier::copyStream(response_body_stream, response->GetResponseBody());
-                    }
-                } else {
-                    std::string error_message;
-                    Poco::StreamCopier::copyToString(response_body_stream, error_message);
-
-                    response->SetClientErrorType(Aws::Client::CoreErrors::NETWORK_CONNECTION);
-                    response->SetClientErrorMessage(error_message);
+                if (!poco_response.isBodyFilled()) {
+                    Poco::StreamCopier::copyStream(response_body_stream, response->GetResponseBody());
                 }
             }
 
diff --git a/be/test/fs/poco_http_client_test.cpp b/be/test/fs/poco_http_client_test.cpp
index 0e1f8d0f1ce72..522c62584990c 100644
--- a/be/test/fs/poco_http_client_test.cpp
+++ b/be/test/fs/poco_http_client_test.cpp
@@ -170,4 +170,30 @@ TEST_F(PocoHttpClientTest, TestErrorAkSk) {
     EXPECT_TRUE(r.status().message().find("SdkResponseCode=403") != std::string::npos);
 }
 
+TEST_F(PocoHttpClientTest, TestNotFoundKey) {
+    Aws::Client::ClientConfiguration config;
+    config.endpointOverride = config::object_storage_endpoint.empty() ? getenv("STARROCKS_UT_S3_ENDPOINT")
+                                                                      : config::object_storage_endpoint;
+    // Create a custom retry strategy
+    int maxRetries = 2;
+    long scaleFactor = 25;
+    std::shared_ptr<Aws::Client::RetryStrategy> retryStrategy =
+            std::make_shared<Aws::Client::DefaultRetryStrategy>(maxRetries, scaleFactor);
+
+    // Create a client configuration object and set the custom retry strategy
+    config.retryStrategy = retryStrategy;
+
+    auto credentials = std::make_shared<Aws::Auth::SimpleAWSCredentialsProvider>(ak, sk);
+
+    auto client = std::make_shared<Aws::S3::S3Client>(std::move(credentials), std::move(config),
+                                                      Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::Never, true);
+
+    auto stream = std::make_unique<starrocks::io::S3InputStream>(client, s_bucket_name, "not_found_key");
+    char buf[6];
+    auto r = stream->read(buf, sizeof(buf));
+    EXPECT_TRUE(r.status().message().find("SdkResponseCode=404") != std::string::npos);
+    // ErrorCode 16 means RESOURCE_NOT_FOUND
+    EXPECT_TRUE(r.status().message().find("SdkErrorType=16") != std::string::npos);
+}
+
 } // namespace starrocks::poco
diff --git a/be/test/io/s3_input_stream_test.cpp b/be/test/io/s3_input_stream_test.cpp
index 2f2f38d936956..9a15ef3f2a6b5 100644
--- a/be/test/io/s3_input_stream_test.cpp
+++ b/be/test/io/s3_input_stream_test.cpp
@@ -138,6 +138,15 @@ TEST_F(S3InputStreamTest, test_read) {
     ASSERT_EQ(10, *f->position());
 }
 
+TEST_F(S3InputStreamTest, test_not_found) {
+    auto f = std::make_unique<S3InputStream>(g_s3client, s_bucket_name, "key_not_found");
+    char buf[6];
+    auto r = f->read(buf, sizeof(buf));
+    EXPECT_TRUE(r.status().message().find("SdkResponseCode=404") != std::string::npos);
+    // ErrorCode 16 means RESOURCE_NOT_FOUND
+    EXPECT_TRUE(r.status().message().find("SdkErrorType=16") != std::string::npos);
+}
+
 TEST_F(S3InputStreamTest, test_skip) {
     auto f = new_random_access_file();
     char buf[6];