From 1061780515ba987e469764a02e9c483fc1dd73e6 Mon Sep 17 00:00:00 2001
From: "Ryan Br..." <ryan.broughan@gmail.com>
Date: Wed, 8 Jan 2025 16:26:35 -0800
Subject: [PATCH] chore: config aws s3 sdk to avoid connection closed error 2
 (#50995)

---
 airbyte-cdk/bulk/toolkits/load-s3/build.gradle               | 3 ++-
 .../src/main/kotlin/io/airbyte/cdk/load/file/s3/S3Client.kt  | 5 +++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/airbyte-cdk/bulk/toolkits/load-s3/build.gradle b/airbyte-cdk/bulk/toolkits/load-s3/build.gradle
index e19bc337b342..ec1667e9b2ca 100644
--- a/airbyte-cdk/bulk/toolkits/load-s3/build.gradle
+++ b/airbyte-cdk/bulk/toolkits/load-s3/build.gradle
@@ -5,5 +5,6 @@ dependencies {
     api project(':airbyte-cdk:bulk:toolkits:bulk-cdk-toolkit-load-object-storage')
 
     testFixturesApi(testFixtures(project(":airbyte-cdk:bulk:toolkits:bulk-cdk-toolkit-load-object-storage")))
-    implementation("aws.sdk.kotlin:s3:1.3.94")
+    implementation("aws.sdk.kotlin:s3:1.3.98")
+    implementation("aws.smithy.kotlin:http-client-engine-okhttp:1.3.31")
 }
diff --git a/airbyte-cdk/bulk/toolkits/load-s3/src/main/kotlin/io/airbyte/cdk/load/file/s3/S3Client.kt b/airbyte-cdk/bulk/toolkits/load-s3/src/main/kotlin/io/airbyte/cdk/load/file/s3/S3Client.kt
index e3d36350b9b3..48558739a029 100644
--- a/airbyte-cdk/bulk/toolkits/load-s3/src/main/kotlin/io/airbyte/cdk/load/file/s3/S3Client.kt
+++ b/airbyte-cdk/bulk/toolkits/load-s3/src/main/kotlin/io/airbyte/cdk/load/file/s3/S3Client.kt
@@ -18,6 +18,7 @@ import aws.sdk.kotlin.services.s3.model.PutObjectRequest
 import aws.smithy.kotlin.runtime.auth.awscredentials.CredentialsProvider
 import aws.smithy.kotlin.runtime.content.ByteStream
 import aws.smithy.kotlin.runtime.content.toInputStream
+import aws.smithy.kotlin.runtime.http.engine.okhttp.OkHttpEngine
 import aws.smithy.kotlin.runtime.net.url.Url
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings
 import io.airbyte.cdk.load.command.aws.AWSAccessKeyConfigurationProvider
@@ -38,6 +39,7 @@ import jakarta.inject.Singleton
 import java.io.ByteArrayOutputStream
 import java.io.InputStream
 import java.io.OutputStream
+import kotlin.time.Duration.Companion.milliseconds
 import kotlinx.coroutines.flow.flow
 
 data class S3Object(override val key: String, override val storageConfig: S3BucketConfiguration) :
@@ -244,6 +246,9 @@ class S3ClientFactory(
                             Url.parse(it)
                         } else null
                     }
+                // Fix for connection reset issue:
+                // https://github.com/awslabs/aws-sdk-kotlin/issues/1214#issuecomment-2464831817
+                httpClient(OkHttpEngine) { connectionIdlePollingInterval = 200.milliseconds }
             }
 
         return S3Client(