Skip to content

Commit

Permalink
Add query timeout
Browse files Browse the repository at this point in the history 
https://www.elastic.co/guide/en/elasticsearch/reference/current/search-your-data.html#search-timeout
  • Loading branch information
@UsenkoArtem @anti-social
UsenkoArtem authored and anti-social committed Dec 11, 2024
1 parent e8aa865 commit 836cfc2
Show file tree
Hide file tree
Showing 5 changed files with 88 additions and 46 deletions.
13 changes: 9 additions & 4 deletions elasticmagic/api/elasticmagic.api
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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
Expand Down
19 changes: 19 additions & 0 deletions elasticmagic/src/commonMain/kotlin/dev/evo/elasticmagic/SearchQuery.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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.
*
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down
5 changes: 5 additions & 0 deletions elasticmagic/src/commonMain/kotlin/dev/evo/elasticmagic/compile/SearchQueryCompiler.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,11 @@ open class SearchQueryCompiler(
visit(this, searchQuery.aggregations)
}
}

if (searchQuery.timeout != null) {
ctx.field("timeout", searchQuery.timeout.inWholeSeconds.toString() + "s")
}

if (searchQuery.rescores.isNotEmpty()) {
ctx.array("rescore") {
visit(this, searchQuery.rescores)
Expand Down
6 changes: 6 additions & 0 deletions elasticmagic/src/commonTest/kotlin/dev/evo/elasticmagic/SearchQueryTests.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,9 @@ 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
import kotlin.time.ExperimentalTime
import kotlin.time.Duration.Companion.seconds

private data class SimpleExtension(override val name: String) : SearchExt {
override fun clone() = copy()
Expand All @@ -18,6 +21,7 @@ private data class SimpleExtension(override val name: String) : SearchExt {

}

@OptIn(ExperimentalTime::class)
class SearchQueryTests {
@Test
fun cloning() {
Expand All @@ -29,6 +33,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))
Expand All @@ -38,6 +43,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
Expand Down
91 changes: 49 additions & 42 deletions elasticmagic/src/commonTest/kotlin/dev/evo/elasticmagic/compile/SearchQueryCompilerTests.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ import dev.evo.elasticmagic.doc.BoundField
import dev.evo.elasticmagic.doc.BoundRuntimeField
import dev.evo.elasticmagic.doc.Document
import dev.evo.elasticmagic.doc.RootFieldSet
import dev.evo.elasticmagic.types.SimpleFieldType
import dev.evo.elasticmagic.doc.SubDocument
import dev.evo.elasticmagic.doc.datetime
import dev.evo.elasticmagic.query.Bool
Expand All @@ -26,23 +25,22 @@ import dev.evo.elasticmagic.query.MatchAll
import dev.evo.elasticmagic.query.MinimumShouldMatch
import dev.evo.elasticmagic.query.MultiMatch
import dev.evo.elasticmagic.query.NodeHandle
import dev.evo.elasticmagic.query.Script
import dev.evo.elasticmagic.query.Sort
import dev.evo.elasticmagic.query.QueryExpressionNode
import dev.evo.elasticmagic.query.QueryRescore
import dev.evo.elasticmagic.query.Script
import dev.evo.elasticmagic.query.SearchExt
import dev.evo.elasticmagic.query.Sort
import dev.evo.elasticmagic.query.match
import dev.evo.elasticmagic.serde.Serializer
import dev.evo.elasticmagic.types.FloatType
import dev.evo.elasticmagic.types.KeywordType
import dev.evo.elasticmagic.types.LongType

import dev.evo.elasticmagic.types.SimpleFieldType
import io.kotest.assertions.throwables.shouldThrow
import io.kotest.matchers.maps.shouldContainExactly

import kotlinx.datetime.LocalDateTime

import kotlin.test.Test
import kotlin.time.Duration.Companion.seconds

class AnyField(name: String) : BoundField<Any, Any>(
name,
Expand Down Expand Up @@ -79,7 +77,7 @@ class SearchQueryCompilerTests : BaseCompilerTest<SearchQueryCompiler>(::SearchQ
}

val compiled = compile(
SearchQuery(userDoc.lastLoggedAt.gt(LocalDateTime(2020, 12, 31, 23, 0)))
SearchQuery(userDoc.lastLoggedAt.gt(LocalDateTime(2020, 12, 31, 23, 0))).setTimeout(10.seconds)
)
compiled.body shouldContainExactly mapOf(
"query" to mapOf(
Expand All @@ -88,7 +86,8 @@ class SearchQueryCompilerTests : BaseCompilerTest<SearchQueryCompiler>(::SearchQ
"gt" to "2020-12-31T23:00:00Z"
)
)
)
),
"timeout" to "10s"
)
}

Expand Down Expand Up @@ -186,17 +185,19 @@ class SearchQueryCompilerTests : BaseCompilerTest<SearchQueryCompiler>(::SearchQ
)
}

query.aggs(mapOf(
"date_created_hist" to DateHistogramAgg(
AnyField("date_created"),
interval = DateHistogramAgg.Interval.Fixed(FixedInterval.Hours(2)),
aggs = mapOf(
"min_price" to MinAgg(
AnyField("price")
query.aggs(
mapOf(
"date_created_hist" to DateHistogramAgg(
AnyField("date_created"),
interval = DateHistogramAgg.Interval.Fixed(FixedInterval.Hours(2)),
aggs = mapOf(
"min_price" to MinAgg(
AnyField("price")
)
)
)
)
))
)
compile(query).let { compiled ->
compiled.body shouldContainExactly mapOf(
"aggs" to mapOf(
Expand Down Expand Up @@ -457,23 +458,25 @@ class SearchQueryCompilerTests : BaseCompilerTest<SearchQueryCompiler>(::SearchQ
)
)

query.rescore(listOf(
QueryRescore(
FunctionScore(
functions = listOf(
FunctionScore.ScriptScore(
script = Script.Source(
"Math.log10(doc[params.field].value + 2)",
params = mapOf(
"field" to AnyField("likes")
query.rescore(
listOf(
QueryRescore(
FunctionScore(
functions = listOf(
FunctionScore.ScriptScore(
script = Script.Source(
"Math.log10(doc[params.field].value + 2)",
params = mapOf(
"field" to AnyField("likes")
)
)
)
)
)
),
scoreMode = QueryRescore.ScoreMode.MULTIPLY,
),
scoreMode = QueryRescore.ScoreMode.MULTIPLY,
)
)
))
)
compile(query).body shouldContainExactly mapOf(
"rescore" to listOf(
mapOf(
Expand Down Expand Up @@ -806,15 +809,17 @@ class SearchQueryCompilerTests : BaseCompilerTest<SearchQueryCompiler>(::SearchQ
)
)

query.runtimeMappings(listOf(
BoundRuntimeField(
"attr_ids", LongType,
Script.Source(
"emit(doc[attr_id].value >>> 32)",
),
RootFieldSet
query.runtimeMappings(
listOf(
BoundRuntimeField(
"attr_ids", LongType,
Script.Source(
"emit(doc[attr_id].value >>> 32)",
),
RootFieldSet
)
)
))
)
compile(query).body shouldContainExactly mapOf(
"runtime_mappings" to mapOf(
"day_of_week" to mapOf(
Expand Down Expand Up @@ -945,11 +950,13 @@ class SearchQueryCompilerTests : BaseCompilerTest<SearchQueryCompiler>(::SearchQ
@Test
fun idsQuery() = testWithCompiler {
val query = SearchQuery(
Ids(listOf(
"order~3",
"order~2",
"order~1",
))
Ids(
listOf(
"order~3",
"order~2",
"order~1",
)
)
)
compile(query).body shouldContainExactly mapOf(
"query" to mapOf(
Expand Down

0 comments on commit 836cfc2

Please sign in to comment.