Add query timeout

elasticmagic/api/elasticmagic.api

@@ -62,6 +62,7 @@ public abstract class dev/evo/elasticmagic/BaseSearchQuery {
 	protected final fun getSource ()Ldev/evo/elasticmagic/query/Source;
 	protected final fun getStoredFields ()Ljava/util/List;
 	protected final fun getTerminateAfter ()Ljava/lang/Integer;
+	protected final fun getTimeout-FghU774 ()Lkotlin/time/Duration;
 	protected final fun getTrackScores ()Ljava/lang/Boolean;
 	protected final fun getTrackTotalHits ()Ljava/lang/Boolean;
 	protected abstract fun new (Lkotlin/jvm/functions/Function1;)Ldev/evo/elasticmagic/BaseSearchQuery;
@@ -102,6 +103,8 @@ public abstract class dev/evo/elasticmagic/BaseSearchQuery {
 	protected final fun setSize (Ljava/lang/Integer;)V
 	protected final fun setSource (Ldev/evo/elasticmagic/query/Source;)V
 	protected final fun setTerminateAfter (Ljava/lang/Integer;)V
+	public final fun setTimeout-BwNAW2A (Lkotlin/time/Duration;)Ldev/evo/elasticmagic/BaseSearchQuery;
+	protected final fun setTimeout-BwNAW2A (Lkotlin/time/Duration;)V
 	protected final fun setTrackScores (Ljava/lang/Boolean;)V
 	protected final fun setTrackTotalHits (Ljava/lang/Boolean;)V
 	public final fun size (Ljava/lang/Integer;)Ldev/evo/elasticmagic/BaseSearchQuery;
@@ -586,7 +589,7 @@ public final class dev/evo/elasticmagic/SearchQuery$Delete$Companion {
 
 public final class dev/evo/elasticmagic/SearchQuery$Search : dev/evo/elasticmagic/PreparedSearchQuery {
 	public static final field Companion Ldev/evo/elasticmagic/SearchQuery$Search$Companion;
-	public fun <init> (Lkotlin/jvm/functions/Function1;Ldev/evo/elasticmagic/query/QueryExpression;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/List;Ljava/util/List;Ljava/lang/Boolean;Ljava/lang/Boolean;Ldev/evo/elasticmagic/query/Source;Ljava/util/List;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/Map;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/util/List;Ljava/util/Map;)V
+	public synthetic fun <init> (Lkotlin/jvm/functions/Function1;Ldev/evo/elasticmagic/query/QueryExpression;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/List;Ljava/util/List;Ljava/lang/Boolean;Ljava/lang/Boolean;Ldev/evo/elasticmagic/query/Source;Ljava/util/List;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/Map;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/util/List;Lkotlin/time/Duration;Ljava/util/Map;Lkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun component1 ()Lkotlin/jvm/functions/Function1;
 	public final fun component10 ()Ldev/evo/elasticmagic/query/Source;
 	public final fun component11 ()Ljava/util/List;
@@ -599,16 +602,17 @@ public final class dev/evo/elasticmagic/SearchQuery$Search : dev/evo/elasticmagi
 	public final fun component18 ()Ljava/lang/Integer;
 	public final fun component19 ()Ljava/util/List;
 	public final fun component2 ()Ldev/evo/elasticmagic/query/QueryExpression;
-	public final fun component20 ()Ljava/util/Map;
+	public final fun component20-FghU774 ()Lkotlin/time/Duration;
+	public final fun component21 ()Ljava/util/Map;
 	public final fun component3 ()Ljava/util/List;
 	public final fun component4 ()Ljava/util/List;
 	public final fun component5 ()Ljava/util/Map;
 	public final fun component6 ()Ljava/util/List;
 	public final fun component7 ()Ljava/util/List;
 	public final fun component8 ()Ljava/lang/Boolean;
 	public final fun component9 ()Ljava/lang/Boolean;
-	public final fun copy (Lkotlin/jvm/functions/Function1;Ldev/evo/elasticmagic/query/QueryExpression;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/List;Ljava/util/List;Ljava/lang/Boolean;Ljava/lang/Boolean;Ldev/evo/elasticmagic/query/Source;Ljava/util/List;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/Map;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/util/List;Ljava/util/Map;)Ldev/evo/elasticmagic/SearchQuery$Search;
-	public static synthetic fun copy$default (Ldev/evo/elasticmagic/SearchQuery$Search;Lkotlin/jvm/functions/Function1;Ldev/evo/elasticmagic/query/QueryExpression;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/List;Ljava/util/List;Ljava/lang/Boolean;Ljava/lang/Boolean;Ldev/evo/elasticmagic/query/Source;Ljava/util/List;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/Map;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/util/List;Ljava/util/Map;ILjava/lang/Object;)Ldev/evo/elasticmagic/SearchQuery$Search;
+	public final fun copy-Q9Qo52k (Lkotlin/jvm/functions/Function1;Ldev/evo/elasticmagic/query/QueryExpression;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/List;Ljava/util/List;Ljava/lang/Boolean;Ljava/lang/Boolean;Ldev/evo/elasticmagic/query/Source;Ljava/util/List;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/Map;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/util/List;Lkotlin/time/Duration;Ljava/util/Map;)Ldev/evo/elasticmagic/SearchQuery$Search;
+	public static synthetic fun copy-Q9Qo52k$default (Ldev/evo/elasticmagic/SearchQuery$Search;Lkotlin/jvm/functions/Function1;Ldev/evo/elasticmagic/query/QueryExpression;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/List;Ljava/util/List;Ljava/lang/Boolean;Ljava/lang/Boolean;Ldev/evo/elasticmagic/query/Source;Ljava/util/List;Ljava/util/List;Ljava/util/List;Ljava/util/Map;Ljava/util/Map;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/util/List;Lkotlin/time/Duration;Ljava/util/Map;ILjava/lang/Object;)Ldev/evo/elasticmagic/SearchQuery$Search;
 	public fun equals (Ljava/lang/Object;)Z
 	public final fun getAggregations ()Ljava/util/Map;
 	public final fun getDocSourceFactory ()Lkotlin/jvm/functions/Function1;
@@ -628,6 +632,7 @@ public final class dev/evo/elasticmagic/SearchQuery$Search : dev/evo/elasticmagi
 	public final fun getSource ()Ldev/evo/elasticmagic/query/Source;
 	public final fun getStoredFields ()Ljava/util/List;
 	public fun getTerminateAfter ()Ljava/lang/Integer;
+	public final fun getTimeout-FghU774 ()Lkotlin/time/Duration;
 	public final fun getTrackScores ()Ljava/lang/Boolean;
 	public final fun getTrackTotalHits ()Ljava/lang/Boolean;
 	public fun hashCode ()I
@@ -4788,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

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,

elasticmagic/src/commonMain/kotlin/dev/evo/elasticmagic/SearchQuery.kt

@@ -16,6 +16,7 @@ import dev.evo.elasticmagic.query.Sort
 import dev.evo.elasticmagic.query.Source
 import dev.evo.elasticmagic.query.collect
 import dev.evo.elasticmagic.serde.Deserializer
+import kotlin.time.Duration
 
 enum class SearchType : ToValue<String> {
     QUERY_THEN_FETCH, DFS_QUERY_THEN_FETCH;
@@ -57,6 +58,7 @@ abstract class BaseSearchQuery<S : BaseDocSource, T : BaseSearchQuery<S, T>>(
     protected var size: Int? = null
     protected var from: Int? = null
     protected var terminateAfter: Int? = null
+    protected var timeout: Duration? = null
 
     protected val extensions: MutableList<SearchExt> = mutableListOf()
 
@@ -105,6 +107,7 @@ abstract class BaseSearchQuery<S : BaseDocSource, T : BaseSearchQuery<S, T>>(
         cloned.size = size
         cloned.from = from
         cloned.terminateAfter = terminateAfter
+        cloned.timeout = timeout
         cloned.params.putAll(params)
         return cloned
     }
@@ -694,6 +697,20 @@ abstract class BaseSearchQuery<S : BaseDocSource, T : BaseSearchQuery<S, T>>(
      */
     open fun beforeExecute() {}
 
+    /**
+     * Sets a timeout for the search query.
+     *
+     * @param timeout the maximum time to wait for the search query to complete.
+     *                If the search takes longer than this time, it will be terminated.
+     *
+     * @return the current instance of the search query with the updated timeout.
+     *
+     * @see <https://www.elastic.co/guide/en/elasticsearch/reference/current/search-search.html#search-timeout>
+     */
+    fun setTimeout(timeout: Duration?): T = self {
+        this.timeout = timeout
+    }
+
     /**
      * Makes an immutable view of the search query. Be careful when using this method.
      *
@@ -730,6 +747,7 @@ abstract class BaseSearchQuery<S : BaseDocSource, T : BaseSearchQuery<S, T>>(
             from = from,
             terminateAfter = terminateAfter,
             extensions = extensions,
+            timeout = timeout,
             params = Params(
                 PreparedSearchQuery.filteredParams(this.params, SearchQuery.Search.ALLOWED_PARAMS),
                 params
@@ -758,7 +776,7 @@ abstract class BaseSearchQuery<S : BaseDocSource, T : BaseSearchQuery<S, T>>(
             terminateAfter = terminateAfter,
             script = script,
             params = Params(
-                PreparedSearchQuery.filteredParams(this.params, SearchQuery.Delete.ALLOWED_PARAMS),
+                PreparedSearchQuery.filteredParams(this.params, SearchQuery.Update.ALLOWED_PARAMS),
                 params
             )
         )
@@ -965,6 +983,7 @@ open class SearchQuery<S : BaseDocSource>(
         val from: Int?,
         override val terminateAfter: Int?,
         val extensions: List<SearchExt>,
+        val timeout: Duration?,
         val params: Params,
     ) : PreparedSearchQuery {
         companion object {
@@ -996,6 +1015,7 @@ open class SearchQuery<S : BaseDocSource>(
                 "track_total_hits",
                 "typed_keys",
                 "version",
+                "timeout",
             )
         }
     }

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)
         }
     }
@@ -129,6 +136,11 @@ open class SearchQueryCompiler(
                 visit(this, searchQuery.aggregations)
             }
         }
+
+        if (searchQuery.timeout != null) {
+            ctx.field("timeout", searchQuery.timeout.toTimeoutString())
+        }
+
         if (searchQuery.rescores.isNotEmpty()) {
             ctx.array("rescore") {
                 visit(this, searchQuery.rescores)

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"
+}

elasticmagic/src/commonTest/kotlin/dev/evo/elasticmagic/SearchQueryTests.kt

@@ -9,6 +9,8 @@ import dev.evo.elasticmagic.query.SearchExt
 import dev.evo.elasticmagic.serde.Serializer
 import io.kotest.matchers.shouldBe
 import kotlin.test.Test
+import kotlin.time.Duration.Companion.seconds
+import kotlin.time.ExperimentalTime
 
 private data class SimpleExtension(override val name: String) : SearchExt {
     override fun clone() = copy()
@@ -18,6 +20,7 @@ private data class SimpleExtension(override val name: String) : SearchExt {
 
 }
 
+@OptIn(ExperimentalTime::class)
 class SearchQueryTests {
     @Test
     fun cloning() {
@@ -29,6 +32,7 @@ class SearchQueryTests {
         val sq1 = SearchQuery()
             .filter(userDoc.login.eq("root"))
             .ext(SimpleExtension("test"))
+            .setTimeout(10.seconds)
 
         val sq2 = sq1.clone()
             .filter(userDoc.isActive.eq(true))
@@ -38,6 +42,7 @@ class SearchQueryTests {
             it.size shouldBe null
             it.filters.size shouldBe 1
             it.extensions.size shouldBe 1
+            it.timeout shouldBe 10.seconds
         }
         sq2.prepareSearch().let {
             it.size shouldBe 1

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"))
+    }
+}