From 09acfd6a3dbe8afd91169334ec7da8301a0913df Mon Sep 17 00:00:00 2001
From: Hardy <caoxiaozhen@secnium.com>
Date: Tue, 2 Mar 2021 11:51:41 +0800
Subject: [PATCH] feat: add support for search after

---
 go.sum         |  4 ----
 search.go      | 31 ++++++++++++++++++++-----------
 search_test.go |  7 +++++++
 3 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/go.sum b/go.sum
index 4108027..fae67ff 100644
--- a/go.sum
+++ b/go.sum
@@ -1,9 +1,5 @@
-github.com/elastic/go-elasticsearch v0.0.0 h1:Pd5fqOuBxKxv83b0+xOAJDAkziWYwFinWnBO0y+TZaA=
-github.com/elastic/go-elasticsearch v0.0.0/go.mod h1:TkBSJBuTyFdBnrNqoPc54FN0vKf5c04IdM4zuStJ7xg=
 github.com/elastic/go-elasticsearch/v7 v7.6.0 h1:sYpGLpEFHgLUKLsZUBfuaVI9QgHjS3JdH9fX4/z8QI8=
 github.com/elastic/go-elasticsearch/v7 v7.6.0/go.mod h1:OJ4wdbtDNk5g503kvlHLyErCgQwwzmDtaFC4XyOxXA4=
-github.com/elastic/go-elasticsearch/v8 v8.0.0-20200210103600-aff00e5adfde h1:Y9SZx8RQqFycLxi5W5eFmxMqnmijULVc3LMjBTtZQdM=
-github.com/elastic/go-elasticsearch/v8 v8.0.0-20200210103600-aff00e5adfde/go.mod h1:xe9a/L2aeOgFKKgrO3ibQTnMdpAeL0GC+5/HpGScSa4=
 github.com/fatih/structs v1.1.0 h1:Q7juDM0QtcnhCpeyLGQKyg4TOIghuNXrkL32pHAUMxo=
 github.com/fatih/structs v1.1.0/go.mod h1:9NiDSp5zOcgEDl+j00MP/WkGVPOlPRLejGD8Ga6PJ7M=
 github.com/jgroeneveld/schema v1.0.0 h1:J0E10CrOkiSEsw6dfb1IfrDJD14pf6QLVJ3tRPl/syI=
diff --git a/search.go b/search.go
index 39dabf6..03e934f 100644
--- a/search.go
+++ b/search.go
@@ -15,16 +15,17 @@ import (
 // Not all features of the search API are currently supported, but a request can
 // currently include a query, aggregations, and more.
 type SearchRequest struct {
-	aggs       []Aggregation
-	explain    *bool
-	from       *uint64
-	highlight  Mappable
-	postFilter Mappable
-	query      Mappable
-	size       *uint64
-	sort       Sort
-	source     Source
-	timeout    *time.Duration
+	aggs        []Aggregation
+	explain     *bool
+	from        *uint64
+	highlight   Mappable
+	searchAfter []string
+	postFilter  Mappable
+	query       Mappable
+	size        *uint64
+	sort        Sort
+	source      Source
+	timeout     *time.Duration
 }
 
 // Search creates a new SearchRequest object, to be filled via method chaining.
@@ -74,6 +75,12 @@ func (req *SearchRequest) Sort(name string, order Order) *SearchRequest {
 	return req
 }
 
+// SearchAfter retrieve the sorted result
+func (req *SearchRequest) SearchAfter(s ...string) *SearchRequest {
+	req.searchAfter = append(req.searchAfter, s...)
+	return req
+}
+
 // Explain sets whether the ElasticSearch API should return an explanation for
 // how each hit's score was calculated.
 func (req *SearchRequest) Explain(b bool) *SearchRequest {
@@ -105,7 +112,6 @@ func (req *SearchRequest) Highlight(highlight Mappable) *SearchRequest {
 	return req
 }
 
-
 // Map implements the Mappable interface. It converts the request to into a
 // nested map[string]interface{}, as expected by the go-elasticsearch library.
 func (req *SearchRequest) Map() map[string]interface{} {
@@ -142,6 +148,9 @@ func (req *SearchRequest) Map() map[string]interface{} {
 	if req.highlight != nil {
 		m["highlight"] = req.highlight.Map()
 	}
+	if req.searchAfter != nil {
+		m["search_after"] = req.searchAfter
+	}
 
 	source := req.source.Map()
 	if len(source) > 0 {
diff --git a/search_test.go b/search_test.go
index ad8a1ce..1bcf948 100644
--- a/search_test.go
+++ b/search_test.go
@@ -7,6 +7,13 @@ import (
 
 func TestSearchMaps(t *testing.T) {
 	runMapTests(t, []mapTest{
+		{
+			"a simple query with search after",
+			Search().SearchAfter("_id", "name"),
+			map[string]interface{}{
+				"search_after": []string{"_id", "name"},
+			},
+		},
 		{
 			"a simple match_all query with a size and no aggs",
 			Search().Query(MatchAll()).Size(20),