Merge branch 'master' into sequence-should-support-update

cognitedata · Nov 15, 2024 · 52cde6e · 52cde6e
2 parents 56806d1 + d595a40
commit 52cde6e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/src/main/scala/cognite/spark/v1/FileContentRelation.scala b/src/main/scala/cognite/spark/v1/FileContentRelation.scala
@@ -42,7 +42,9 @@ class FileContentRelation(config: RelationConfig, fileExternalId: String, inferS
       backend <- Resource.make(
         IO(
           AsyncHttpClientFs2Backend
-            .usingClient[IO](SttpClientBackendFactory.create("file content download"), dispatcher))
+            .usingClient[IO](
+              SttpClientBackendFactory.create("file content download", Some(3600000)),
+              dispatcher))
       )(backend => backend.close())
     } yield backend
 

diff --git a/src/main/scala/cognite/spark/v1/SttpClientBackendFactory.scala b/src/main/scala/cognite/spark/v1/SttpClientBackendFactory.scala
@@ -6,7 +6,7 @@ import org.asynchttpclient.AsyncHttpClient
 import sttp.client3.SttpBackendOptions
 
 object SttpClientBackendFactory {
-  def create(prefix: String = "Cdf-Spark"): AsyncHttpClient = {
+  def create(prefix: String = "Cdf-Spark", requestTimeoutMs: Option[Int] = None): AsyncHttpClient = {
     // It's important that the threads made by the async http client is daemon threads,
     // so that we don't hang applications using our library during exit.
     // See for more info https://github.com/cognitedata/cdp-spark-datasource/pull/415/files#r396774391
@@ -22,10 +22,14 @@ object SttpClientBackendFactory {
         .build()
     AsyncHttpClientBackend.clientWithModifiedOptions(
       SttpBackendOptions.Default,
-      options =>
+      options => {
         options
           .setThreadFactory(clientThreadFactory)
           .setNettyTimer(new HashedWheelTimer(timerThreadFactory))
+        //Timeout override for potentially long stream operation
+        requestTimeoutMs.foreach(options.setRequestTimeout)
+        options
+      }
     )
   }
 }