Add beforeExecute hook for SearchQuery

anti-social · Jan 4, 2024 · 5581d20 · 5581d20
1 parent ba7b350
commit 5581d20
Show file tree

Hide file tree

Showing 13 changed files with 78 additions and 57 deletions.
diff --git a/elasticmagic/api/elasticmagic.api b/elasticmagic/api/elasticmagic.api
@@ -27,6 +27,7 @@ public abstract class dev/evo/elasticmagic/BaseSearchQuery {
 	public final fun aggs (Ldev/evo/elasticmagic/SearchQuery$CLEAR;)Ldev/evo/elasticmagic/BaseSearchQuery;
 	public final fun aggs (Ljava/util/Map;)Ldev/evo/elasticmagic/BaseSearchQuery;
 	public final fun aggs ([Lkotlin/Pair;)Ldev/evo/elasticmagic/BaseSearchQuery;
+	public fun beforeExecute ()V
 	public final fun clone ()Ldev/evo/elasticmagic/BaseSearchQuery;
 	public final fun docvalueFields (Ldev/evo/elasticmagic/SearchQuery$CLEAR;)Ldev/evo/elasticmagic/BaseSearchQuery;
 	public final fun docvalueFields (Ljava/util/List;)Ldev/evo/elasticmagic/BaseSearchQuery;
@@ -516,6 +517,8 @@ public class dev/evo/elasticmagic/SearchQuery : dev/evo/elasticmagic/BaseSearchQ
 	public static synthetic fun execute$default (Ldev/evo/elasticmagic/SearchQuery;Ldev/evo/elasticmagic/ElasticsearchIndex;Ljava/util/Map;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public synthetic fun new (Lkotlin/jvm/functions/Function1;)Ldev/evo/elasticmagic/BaseSearchQuery;
 	protected fun new (Lkotlin/jvm/functions/Function1;)Ldev/evo/elasticmagic/SearchQuery;
+	public final fun search (Ldev/evo/elasticmagic/ElasticsearchIndex;Ljava/util/Map;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static synthetic fun search$default (Ldev/evo/elasticmagic/SearchQuery;Ldev/evo/elasticmagic/ElasticsearchIndex;Ljava/util/Map;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public final fun update (Ldev/evo/elasticmagic/ElasticsearchIndex;Ldev/evo/elasticmagic/query/Script;Ldev/evo/elasticmagic/Refresh;Ldev/evo/elasticmagic/Conflicts;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/util/Map;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static synthetic fun update$default (Ldev/evo/elasticmagic/SearchQuery;Ldev/evo/elasticmagic/ElasticsearchIndex;Ldev/evo/elasticmagic/query/Script;Ldev/evo/elasticmagic/Refresh;Ldev/evo/elasticmagic/Conflicts;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/util/Map;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public final fun updateAsync (Ldev/evo/elasticmagic/ElasticsearchIndex;Ldev/evo/elasticmagic/query/Script;Ldev/evo/elasticmagic/Refresh;Ldev/evo/elasticmagic/Conflicts;Ljava/lang/Integer;Ljava/lang/Integer;Ljava/util/Map;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;

diff --git a/elasticmagic/src/commonMain/kotlin/dev/evo/elasticmagic/SearchQuery.kt b/elasticmagic/src/commonMain/kotlin/dev/evo/elasticmagic/SearchQuery.kt
@@ -689,6 +689,11 @@ abstract class BaseSearchQuery<S : BaseDocSource, T : BaseSearchQuery<S, T>>(
         params.putNotNullOrRemove("seq_no_primary_term", seqNoPrimaryTerm)
     }
 
+    /**
+     * Method is called before query execution. Can be overriden in subclasses of [SearchQuery].
+     */
+    open fun beforeExecute() {}
+
     /**
      * Makes an immutable view of the search query. Be careful when using this method.
      *
@@ -812,10 +817,19 @@ open class SearchQuery<S : BaseDocSource>(
     /**
      * Executes the search query using an [index].
      */
-    suspend fun execute(index: ElasticsearchIndex, params: Params? = null): SearchQueryResult<S> {
+    suspend fun search(index: ElasticsearchIndex, params: Params? = null): SearchQueryResult<S> {
+        beforeExecute()
         return index.search(prepareSearch(params))
     }
 
+    /**
+     * Alias of [SearchQuery.search]. Executes the search query using an [index].
+     */
+    @Deprecated("Replaced with `search` method", replaceWith = ReplaceWith("search"))
+    suspend fun execute(index: ElasticsearchIndex, params: Params? = null): SearchQueryResult<S> {
+        return search(index, params)
+    }
+
     /**
      * Retrieves a number of hits for the search query using an [index].
      *
@@ -838,6 +852,7 @@ open class SearchQuery<S : BaseDocSource>(
         scrollSize: Int? = null,
         params: Params? = null,
     ): DeleteByQueryResult {
+        beforeExecute()
         return index.deleteByQuery(
             prepareDelete(
                 Params(
@@ -859,6 +874,7 @@ open class SearchQuery<S : BaseDocSource>(
         scrollSize: Int? = null,
         params: Params? = null,
     ): AsyncResult<DeleteByQueryPartialResult, DeleteByQueryResult?> {
+        beforeExecute()
         return index.deleteByQueryAsync(
             prepareDelete(
                 Params(
@@ -886,6 +902,7 @@ open class SearchQuery<S : BaseDocSource>(
         scrollSize: Int? = null,
         params: Params? = null,
     ): UpdateByQueryResult {
+        beforeExecute()
         return index.updateByQuery(
             prepareUpdate(
                 script,
@@ -909,6 +926,7 @@ open class SearchQuery<S : BaseDocSource>(
         scrollSize: Int? = null,
         params: Params? = null,
     ): AsyncResult<UpdateByQueryPartialResult, UpdateByQueryResult?> {
+        beforeExecute()
         return index.updateByQueryAsync(
             prepareUpdate(
                 script,

diff --git a/integ-tests/src/commonTest/kotlin/dev/evo/elasticmagic/DynamicTemplatesTests.kt b/integ-tests/src/commonTest/kotlin/dev/evo/elasticmagic/DynamicTemplatesTests.kt
@@ -49,17 +49,17 @@ class DynamicTemplatesTests : ElasticsearchTestBase() {
                 // Field does not support fielddata
                 SearchQuery()
                     .sort(companyIdField.id)
-                    .execute(index)
+                    .search(index)
             }
             val totalHits = SearchQuery()
                 .filter(companyIdField.id eq "10")
-                .execute(index)
+                .search(index)
                 .totalHits
             totalHits shouldBe 1L
 
             val hits = SearchQuery()
                 .sort(companyIdField.asc())
-                .execute(index)
+                .search(index)
                 .hits
             hits.size shouldBe 2
             hits[0].id shouldBe "2"

diff --git a/integ-tests/src/commonTest/kotlin/dev/evo/elasticmagic/ExplanationTest.kt b/integ-tests/src/commonTest/kotlin/dev/evo/elasticmagic/ExplanationTest.kt
@@ -24,7 +24,7 @@ class ExplanationTest : ElasticsearchTestBase() {
     fun withoutExplanation() = runTestWithSerdes {
         withFixtures(ItemDoc, FIXTURES) {
             val searchQuery = SearchQuery()
-            val result = searchQuery.execute(index)
+            val result = searchQuery.search(index)
             result.totalHits shouldBe 8
             result.hits.mapNotNull { it.explanation } shouldBe emptyList()
         }
@@ -34,7 +34,7 @@ class ExplanationTest : ElasticsearchTestBase() {
     fun withExplanationButEmptySearchQuery() = runTestWithSerdes {
         withFixtures(ItemDoc, FIXTURES) {
             val searchQuery = SearchQuery()
-            val result = searchQuery.execute(index, Params("explain" to true))
+            val result = searchQuery.search(index, Params("explain" to true))
             result.totalHits shouldBe 8
             val explanations = result.hits.mapNotNull { it.explanation }
 
@@ -50,7 +50,7 @@ class ExplanationTest : ElasticsearchTestBase() {
     fun withExplanation() = runTestWithSerdes {
         withFixtures(ItemDoc, FIXTURES) {
             val searchQuery = SearchQuery(StringField("model").match("Galaxy Note 10"))
-            val result = searchQuery.execute(index, Params("explain" to true))
+            val result = searchQuery.search(index, Params("explain" to true))
             result.totalHits shouldBe 5
             val explanations = result.hits.mapNotNull { it.explanation }
 

diff --git a/integ-tests/src/commonTest/kotlin/dev/evo/elasticmagic/ParentChildTests.kt b/integ-tests/src/commonTest/kotlin/dev/evo/elasticmagic/ParentChildTests.kt
@@ -110,7 +110,7 @@ class ParentChildTests : ElasticsearchTestBase() {
                 "question" to ::QuestionDocSource,
                 "answer" to ::AnswerDocSource
             )
-            SearchQuery(sourceFactory).execute(index).let { searchResult ->
+            SearchQuery(sourceFactory).search(index).let { searchResult ->
                 searchResult.hits.toSet() shouldContainExactly setOf(
                     SearchHit(
                         index = index.name,
@@ -173,7 +173,7 @@ class ParentChildTests : ElasticsearchTestBase() {
 
             SearchQuery(sourceFactory)
                 .filter(AnswerDoc.join.eq("answer"))
-                .execute(index)
+                .search(index)
                 .let { searchResult ->
                     searchResult.hits.toSet() shouldContainExactly setOf(
                         SearchHit(
@@ -225,7 +225,7 @@ class ParentChildTests : ElasticsearchTestBase() {
                         KnowledgeDoc.join.parent("question")
                     )
                 )
-                .execute(index)
+                .search(index)
                 .let { searchResult ->
                     val parentsAgg = searchResult.agg<TermsAggResult<String>>("parents")
                     parentsAgg.buckets shouldBe listOf(
@@ -236,7 +236,7 @@ class ParentChildTests : ElasticsearchTestBase() {
 
             SearchQuery(sourceFactory)
                 .filter(HasParent(QuestionDoc.text.match("What"), "question"))
-                .execute(index)
+                .search(index)
                 .let { searchResult ->
                     searchResult.hits.toSet() shouldContainExactly setOf(
                         SearchHit(
@@ -260,7 +260,7 @@ class ParentChildTests : ElasticsearchTestBase() {
 
             SearchQuery(sourceFactory)
                 .filter(HasChild(MatchAll, "answer", minChildren = 1))
-                .execute(index)
+                .search(index)
                 .let { searchResult ->
                     searchResult.hits.toSet() shouldContainExactly setOf(
                         SearchHit(
@@ -284,7 +284,7 @@ class ParentChildTests : ElasticsearchTestBase() {
 
             SearchQuery(sourceFactory)
                 .filter(HasChild(MatchAll, "answer", maxChildren = 2))
-                .execute(index)
+                .search(index)
                 .let { searchResult ->
                     searchResult.hits.toSet() shouldContainExactly setOf(
                         SearchHit(
@@ -300,7 +300,7 @@ class ParentChildTests : ElasticsearchTestBase() {
 
             SearchQuery(sourceFactory)
                 .filter(HasChild(MatchAll, "answer", minChildren = 3))
-                .execute(index)
+                .search(index)
                 .let { searchResult ->
                     searchResult.hits.toSet() shouldContainExactly setOf(
                         SearchHit(

diff --git a/integ-tests/src/commonTest/kotlin/dev/evo/elasticmagic/SearchQueryTests.kt b/integ-tests/src/commonTest/kotlin/dev/evo/elasticmagic/SearchQueryTests.kt
@@ -236,7 +236,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
                             scriptType = "unknown",
                         )
                     )
-                    .execute(index)
+                    .search(index)
             }
         }
 
@@ -251,17 +251,17 @@ class SearchQueryTests : ElasticsearchTestBase() {
             when (val version = cluster.getVersion()) {
                 is Version.Opensearch -> {
                     shouldThrow<ElasticsearchException.BadRequest> {
-                        query.execute(index)
+                        query.search(index)
                     }
                 }
                 is Version.Elasticsearch -> {
                     if (version.major < 7) {
                         shouldThrow<ElasticsearchException.Internal> {
-                            query.execute(index)
+                            query.search(index)
                         }
                     } else {
                         shouldThrow<ElasticsearchException.BadRequest> {
-                            query.execute(index)
+                            query.search(index)
                         }
                     }
                 }
@@ -275,7 +275,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
             karlssonsJam, karlssonsBestDonuts, karlssonsJustDonuts, littleBrotherDogStuff
         )) {
             val searchResult = SearchQuery(::OrderDocSource)
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 4
             searchResult.maxScore shouldBe 1.0F
@@ -289,7 +289,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
         withFixtures(OrderDoc, listOf(karlssonsJam)) {
             val searchResult = SearchQuery(::OrderDocSource)
                 .source(excludes=listOf(OrderDoc.comment))
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 1
             searchResult.maxScore shouldBe 1.0F
@@ -308,7 +308,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
         withFixtures(OrderDoc, listOf(karlssonsJam)) {
             val searchResult = SearchQuery(::OrderDocSource)
                 .source(false)
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 1
             searchResult.maxScore shouldBe 1.0F
@@ -325,7 +325,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
             val exc = shouldThrow<IllegalStateException> {
                 SearchQuery(::OrderDocSource)
                     .source(OrderDoc.comment)
-                    .execute(index)
+                    .search(index)
             }
             exc.message shouldBe "Field [user] is required"
         }
@@ -336,7 +336,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
         withFixtures(OrderDoc, listOf(karlssonsJam)) {
             val searchResult = SearchQuery(::OrderDocSource)
                 .docvalueFields(OrderDoc.status, OrderDoc.dateCreated.format("YYYY"))
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 1
             searchResult.maxScore shouldBe 1.0F
@@ -358,7 +358,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
                 ::OrderDocSource,
                 OrderDoc.user.id.eq(1)
             )
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 1
             searchResult.maxScore shouldBe 1.0F
@@ -394,7 +394,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
         )) {
             val searchResult = SearchQuery(::OrderDocSource)
                 .filter(OrderDoc.status.eq(OrderStatus.NEW))
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 3
             searchResult.maxScore.shouldNotBeNull() shouldBe 0.0F
@@ -410,7 +410,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
         )) {
             val searchResult = SearchQuery(::OrderDocSource)
                 .filter(OrderDoc.dateCreated.lt(LocalDate(2020, 1, 1).atStartOfDayIn(TimeZone.UTC)))
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 1
             searchResult.maxScore.shouldNotBeNull() shouldBe 0.0F
@@ -432,7 +432,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
         )) {
             val searchResult = SearchQuery(::OrderDocSource)
                 .filter(Ids(listOf("105", "102")))
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 2
             searchResult.maxScore shouldBe 0.0F
@@ -459,7 +459,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
                         order = Sort.Order.DESC,
                     )
                 )
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 4
             searchResult.maxScore shouldBe null
@@ -487,7 +487,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
                         nested = Sort.Nested(OrderDoc.items)
                     )
                 )
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 4
             searchResult.maxScore shouldBe null
@@ -513,7 +513,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
                     ),
                 )
                 .size(0)
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 4
             // Elasticsearch 6.x has max score 0.0, but 7.x has null
@@ -551,7 +551,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
                     ),
                 )
                 .size(0)
-                .execute(index)
+                .search(index)
 
             val ordersByYear = searchResult.agg<DateHistogramAggResult<Instant>>("orders_by_year")
             ordersByYear.buckets.shouldHaveSize(3)
@@ -589,7 +589,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
                     )
                 )
                 .size(0)
-                .execute(index)
+                .search(index)
 
             searchResult.totalHits shouldBe 4
             searchResult.maxScore ?: 0.0 shouldBe 0.0
@@ -649,7 +649,7 @@ class SearchQueryTests : ElasticsearchTestBase() {
                 .aggs("days_of_week" to TermsAgg(dayOfWeekField))
                 .fields(dayOfWeekField, statusStrField)
                 .sort(dayOfWeekField)
-                .execute(index)
+                .search(index)
 
             searchResult.hits.size shouldBe 4
             searchResult.hits.flatMap { hit -> hit.fields[dayOfWeekField] } shouldBe listOf(

diff --git a/integ-tests/src/commonTest/kotlin/dev/evo/elasticmagic/UpdateMappingTests.kt b/integ-tests/src/commonTest/kotlin/dev/evo/elasticmagic/UpdateMappingTests.kt
@@ -44,7 +44,7 @@ class UpdateMappingTests : ElasticsearchTestBase() {
             SearchQuery(
                 ::UserDocSource, UserV2Doc.name.match("boy")
             )
-                .execute(index)
+                .search(index)
                 .totalHits shouldBe 0
 
             cluster.updateMapping(
@@ -63,7 +63,7 @@ class UpdateMappingTests : ElasticsearchTestBase() {
             SearchQuery(
                 ::UserDocSource, UserV2Doc.name.match("boy")
             )
-                .execute(index)
+                .search(index)
                 .totalHits shouldBe 1
         }
     }