Add timeout parameter support

Convert timeout values to milliseconds if less than 10 seconds.
anti-social · Dec 11, 2024 · 741d5fd · 741d5fd
1 parent dd7f060
commit 741d5fd
Showing 6 changed files with 107 additions and 6 deletions.
diff --git a/elasticmagic/api/elasticmagic.api b/elasticmagic/api/elasticmagic.api
@@ -4793,6 +4793,10 @@ public final class dev/evo/elasticmagic/types/ValueSerializationException : java
 	public synthetic fun <init> (Ljava/lang/Object;Ljava/lang/Throwable;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 }
 
+public final class dev/evo/elasticmagic/util/HelperKt {
+	public static final fun toTimeoutString-LRDsOJo (J)Ljava/lang/String;
+}
+
 public final class dev/evo/elasticmagic/util/OrderedMap {
 	public fun <init> ()V
 	public fun <init> ([Lkotlin/Pair;)V

diff --git a/elasticmagic/src/commonMain/kotlin/dev/evo/elasticmagic/ElasticsearchCluster.kt b/elasticmagic/src/commonMain/kotlin/dev/evo/elasticmagic/ElasticsearchCluster.kt
@@ -2,27 +2,35 @@ package dev.evo.elasticmagic
 
 import dev.evo.elasticmagic.bulk.Action
 import dev.evo.elasticmagic.compile.ActionCompiler
-import dev.evo.elasticmagic.compile.PreparedBulk
 import dev.evo.elasticmagic.compile.CompilerSet
+import dev.evo.elasticmagic.compile.PreparedBulk
 import dev.evo.elasticmagic.compile.PreparedCreateIndex
 import dev.evo.elasticmagic.compile.PreparedUpdateMapping
 import dev.evo.elasticmagic.doc.BaseDocSource
 import dev.evo.elasticmagic.doc.Document
 import dev.evo.elasticmagic.serde.Serde
+import dev.evo.elasticmagic.transport.ApiRequest
 import dev.evo.elasticmagic.transport.BulkRequest
 import dev.evo.elasticmagic.transport.ElasticsearchException
 import dev.evo.elasticmagic.transport.ElasticsearchTransport
-import dev.evo.elasticmagic.transport.ApiRequest
 import dev.evo.elasticmagic.transport.Method
 import dev.evo.elasticmagic.transport.Parameters
+import dev.evo.elasticmagic.util.toTimeoutString
+import kotlinx.coroutines.CompletableDeferred
+import kotlin.time.Duration
 import kotlin.time.ExperimentalTime
 import kotlin.time.measureTimedValue
-import kotlinx.coroutines.CompletableDeferred
 
 internal fun Params.toRequestParameters(): Parameters {
     val entries = this.map { (key, value) ->
         key to when (value) {
             is ToValue<*> -> value.toValue()
+            is Duration -> {
+                if (key == "timeout")
+                    value.toTimeoutString()
+                else value
+            }
+
             else -> value
         }
     }
@@ -42,7 +50,7 @@ class ElasticsearchCluster(
         transport: ElasticsearchTransport,
         serde: Serde.OneLineJson,
         compilers: CompilerSet? = null,
-    ): this(
+    ) : this(
         transport,
         apiSerde = serde,
         bulkSerde = serde,

diff --git a/elasticmagic/src/commonMain/kotlin/dev/evo/elasticmagic/compile/SearchQueryCompiler.kt b/elasticmagic/src/commonMain/kotlin/dev/evo/elasticmagic/compile/SearchQueryCompiler.kt
@@ -38,6 +38,7 @@ import dev.evo.elasticmagic.transport.ApiRequest
 import dev.evo.elasticmagic.transport.BulkRequest
 import dev.evo.elasticmagic.transport.Method
 import dev.evo.elasticmagic.transport.Parameters
+import dev.evo.elasticmagic.util.toTimeoutString
 
 abstract class BaseSearchQueryCompiler(
     features: ElasticsearchFeatures,
@@ -91,12 +92,15 @@ abstract class BaseSearchQueryCompiler(
             is ObjExpression -> ctx.obj {
                 visit(this, value)
             }
+
             is ArrayExpression -> {
                 visit(ctx, value)
             }
+
             is ToValue<*> -> {
                 ctx.value(value.toValue())
             }
+
             else -> super.dispatch(ctx, value)
         }
     }
@@ -106,12 +110,15 @@ abstract class BaseSearchQueryCompiler(
             is ObjExpression -> ctx.obj(name) {
                 visit(this, value)
             }
+
             is ArrayExpression -> ctx.array(name) {
                 visit(this, value)
             }
+
             is ToValue<*> -> {
                 ctx.field(name, value.toValue())
             }
+
             else -> super.dispatch(ctx, name, value)
         }
     }
@@ -131,7 +138,7 @@ open class SearchQueryCompiler(
         }
 
         if (searchQuery.timeout != null) {
-            ctx.field("timeout", searchQuery.timeout.inWholeSeconds.toString() + "s")
+            ctx.field("timeout", searchQuery.timeout.toTimeoutString())
         }
 
         if (searchQuery.rescores.isNotEmpty()) {

diff --git a/elasticmagic/src/commonMain/kotlin/dev/evo/elasticmagic/util/Helper.kt b/elasticmagic/src/commonMain/kotlin/dev/evo/elasticmagic/util/Helper.kt
@@ -0,0 +1,11 @@
+package dev.evo.elasticmagic.util
+
+import kotlin.time.Duration
+
+private const val USE_MILLISECONDS_WHILE_SECONDS_LESS_THAN = 10
+
+fun Duration.toTimeoutString() = if (inWholeSeconds > USE_MILLISECONDS_WHILE_SECONDS_LESS_THAN) {
+    "${inWholeSeconds}s"
+} else {
+    "${inWholeMilliseconds}ms"
+}
diff --git a/elasticmagic/src/commonTest/kotlin/dev/evo/elasticmagic/SearchQueryTimeoutTests.kt b/elasticmagic/src/commonTest/kotlin/dev/evo/elasticmagic/SearchQueryTimeoutTests.kt
@@ -0,0 +1,71 @@
+package dev.evo.elasticmagic
+
+import dev.evo.elasticmagic.compile.BaseCompilerTest
+import dev.evo.elasticmagic.compile.SearchQueryCompiler
+import dev.evo.elasticmagic.doc.Document
+import io.kotest.matchers.maps.shouldContainExactly
+import io.kotest.matchers.shouldBe
+import kotlin.test.Test
+import kotlin.time.Duration.Companion.seconds
+import kotlin.time.ExperimentalTime
+
+@OptIn(ExperimentalTime::class)
+class SearchQueryTimeoutTests : BaseCompilerTest<SearchQueryCompiler>(::SearchQueryCompiler) {
+    @Test
+    fun timeoutInSearchQuery() = testWithCompiler {
+        val userDoc = object : Document() {
+            val login by keyword()
+            val isActive by boolean()
+        }
+
+        val sq1 = SearchQuery()
+            .filter(userDoc.login.eq("root"))
+            .setTimeout(4.seconds)
+
+        sq1.prepareSearch().let {
+            it.size shouldBe null
+            it.filters.size shouldBe 1
+            it.timeout shouldBe 4.seconds
+        }
+
+        compile(sq1).body shouldBe mapOf(
+            "query" to mapOf(
+                "bool" to mapOf(
+                    "filter" to listOf(
+                        mapOf("term" to mapOf("login" to "root"))
+                    )
+                )
+            ),
+            "timeout" to "4000ms"
+        )
+    }
+
+    @Test
+    fun timeoutInParams() = testWithCompiler {
+        val userDoc = object : Document() {
+            val login by keyword()
+            val isActive by boolean()
+        }
+
+        val sq1 = SearchQuery(params = Params("timeout" to 4.seconds))
+            .filter(userDoc.login.eq("root"))
+
+
+        sq1.prepareSearch().let {
+            it.size shouldBe null
+            it.filters.size shouldBe 1
+            it.timeout shouldBe null
+        }
+        val compiled = compile(sq1)
+        compiled.body shouldBe mapOf(
+            "query" to mapOf(
+                "bool" to mapOf(
+                    "filter" to listOf(
+                        mapOf("term" to mapOf("login" to "root"))
+                    )
+                )
+            )
+        )
+        compiled.params shouldContainExactly mapOf("timeout" to listOf("4000ms"))
+    }
+}
diff --git a/elasticmagic/src/commonTest/kotlin/dev/evo/elasticmagic/compile/SearchQueryCompilerTests.kt b/elasticmagic/src/commonTest/kotlin/dev/evo/elasticmagic/compile/SearchQueryCompilerTests.kt
@@ -87,7 +87,7 @@ class SearchQueryCompilerTests : BaseCompilerTest<SearchQueryCompiler>(::SearchQ
                     )
                 )
             ),
-            "timeout" to "10s"
+            "timeout" to "10000ms"
         )
     }