From b99453ba321b480b0f580a5df49ff129bb699690 Mon Sep 17 00:00:00 2001
From: Mike Smith <mike.smith@deliveroo.co.uk>
Date: Wed, 24 Jul 2024 18:29:21 +0100
Subject: [PATCH] CBB-1079: add source disable

---
 aggs_metric.go      | 16 ++++++++++++++--
 aggs_metric_test.go | 36 ++++++++++++++++++++++++++++++++++++
 common.go           |  1 +
 search.go           | 16 +++++++++++++---
 search_test.go      |  8 ++++++++
 5 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/aggs_metric.go b/aggs_metric.go
index 8505a40..a4d944d 100644
--- a/aggs_metric.go
+++ b/aggs_metric.go
@@ -452,6 +452,12 @@ func (agg *TopHitsAgg) Sorts(sorts ...map[string]Sort) *TopHitsAgg {
 	return agg
 }
 
+// SourceDisabled when set will prevent source fields returning
+func (agg *TopHitsAgg) SourceDisabled(b bool) *TopHitsAgg {
+	agg.source.disabled = b
+	return agg
+}
+
 // SourceIncludes sets the keys to return from the top matching documents.
 func (agg *TopHitsAgg) SourceIncludes(keys ...string) *TopHitsAgg {
 	agg.source.includes = keys
@@ -472,8 +478,14 @@ func (agg *TopHitsAgg) Map() map[string]interface{} {
 	if len(agg.sorts) > 0 {
 		innerMap["sort"] = agg.sorts
 	}
-	if len(agg.source.includes) > 0 {
-		innerMap["_source"] = agg.source.Map()
+
+	if agg.source.disabled {
+		innerMap["_source"] = false
+	} else {
+		source := agg.source.Map()
+		if len(source) > 0 {
+			innerMap["_source"] = source
+		}
 	}
 
 	return map[string]interface{}{
diff --git a/aggs_metric_test.go b/aggs_metric_test.go
index e61050d..3db7cac 100644
--- a/aggs_metric_test.go
+++ b/aggs_metric_test.go
@@ -187,5 +187,41 @@ func TestMetricAggs(t *testing.T) {
 				},
 			},
 		},
+		{
+			"top_hits agg with no source",
+			TopHits("top_hits").
+				Sorts(
+					Sorts{
+						{
+							"field_1": {
+								Order: OrderDesc,
+							},
+						},
+						{
+							"field_2": {
+								Order: OrderAsc,
+							},
+						},
+					}...,
+				).
+				SourceDisabled(true),
+			map[string]interface{}{
+				"top_hits": map[string]interface{}{
+					"sort": []map[string]interface{}{
+						{
+							"field_1": map[string]interface{}{
+								"order": OrderDesc,
+							},
+						},
+						{
+							"field_2": map[string]interface{}{
+								"order": OrderAsc,
+							},
+						},
+					},
+					"_source": false,
+				},
+			},
+		},
 	})
 }
diff --git a/common.go b/common.go
index b18c73f..6398cae 100644
--- a/common.go
+++ b/common.go
@@ -5,6 +5,7 @@ package query
 type Source struct {
 	includes []string
 	excludes []string
+	disabled bool
 }
 
 // Map returns a map representation of the Source object.
diff --git a/search.go b/search.go
index af685ca..b7cb207 100644
--- a/search.go
+++ b/search.go
@@ -89,6 +89,12 @@ func (req *SearchRequest) Timeout(dur time.Duration) *SearchRequest {
 	return req
 }
 
+// SourceDisabled when set will prevent source fields returning
+func (req *SearchRequest) SourceDisabled(b bool) *SearchRequest {
+	req.source.disabled = b
+	return req
+}
+
 // SourceIncludes sets the keys to return from the matching documents.
 func (req *SearchRequest) SourceIncludes(keys ...string) *SearchRequest {
 	req.source.includes = keys
@@ -147,9 +153,13 @@ func (req *SearchRequest) Map() map[string]interface{} {
 		m["search_after"] = req.searchAfter
 	}
 
-	source := req.source.Map()
-	if len(source) > 0 {
-		m["_source"] = source
+	if req.source.disabled {
+		m["_source"] = false
+	} else {
+		source := req.source.Map()
+		if len(source) > 0 {
+			m["_source"] = source
+		}
 	}
 
 	return m
diff --git a/search_test.go b/search_test.go
index c748648..75471c6 100644
--- a/search_test.go
+++ b/search_test.go
@@ -24,6 +24,14 @@ func TestSearchMaps(t *testing.T) {
 				"size": 20,
 			},
 		},
+		{
+			"a simple query with no source",
+			Search().SearchAfter("_id", "name").SourceDisabled(true),
+			map[string]interface{}{
+				"search_after": []string{"_id", "name"},
+				"_source":      false,
+			},
+		},
 		{
 			"a complex query with an aggregation and various other options",
 			Search().