From 70d0db7b6ad7c49a091980e0c648eb9ec861515d Mon Sep 17 00:00:00 2001
From: "anjali.agarwal" <anjali.aggarwal@gojek.com>
Date: Wed, 4 Oct 2023 16:06:47 +0530
Subject: [PATCH] feat: enable included fields feature for search assets api

---
 core/asset/discovery.go                           |  3 +++
 internal/server/v1beta1/search.go                 | 15 ++++++++-------
 .../elasticsearch/discovery_search_repository.go  | 14 +++++++++++---
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/core/asset/discovery.go b/core/asset/discovery.go
index cc6f5ad5..ba2fd2c0 100644
--- a/core/asset/discovery.go
+++ b/core/asset/discovery.go
@@ -70,6 +70,9 @@ type SearchConfig struct {
 	// Offset parameter defines the offset from the first result you want to fetch
 	// Note that MaxResults + Offset can not be more than the `index.max_result_window` index setting in ES cluster, which defaults to 10,000
 	Offset int
+
+	// IncludeFields specifies the fields to return in response
+	IncludeFields []string
 }
 
 // SearchResult represents an item/result in a list of search results
diff --git a/internal/server/v1beta1/search.go b/internal/server/v1beta1/search.go
index 1c33c8ca..b35a851f 100644
--- a/internal/server/v1beta1/search.go
+++ b/internal/server/v1beta1/search.go
@@ -23,13 +23,14 @@ func (server *APIServer) SearchAssets(ctx context.Context, req *compassv1beta1.S
 	}
 
 	cfg := asset.SearchConfig{
-		Text:       text,
-		MaxResults: int(req.GetSize()),
-		Filters:    filterConfigFromValues(req.GetFilter()),
-		RankBy:     req.GetRankby(),
-		Queries:    req.GetQuery(),
-		Flags:      getSearchFlagsFromFlags(req.GetFlags()),
-		Offset:     int(req.GetOffset()),
+		Text:          text,
+		MaxResults:    int(req.GetSize()),
+		Filters:       filterConfigFromValues(req.GetFilter()),
+		RankBy:        req.GetRankby(),
+		Queries:       req.GetQuery(),
+		Flags:         getSearchFlagsFromFlags(req.GetFlags()),
+		Offset:        int(req.GetOffset()),
+		IncludeFields: req.GetIncludeFields(),
 	}
 
 	results, err := server.assetService.SearchAssets(ctx, cfg)
diff --git a/internal/store/elasticsearch/discovery_search_repository.go b/internal/store/elasticsearch/discovery_search_repository.go
index cff8ff45..b2d18e39 100644
--- a/internal/store/elasticsearch/discovery_search_repository.go
+++ b/internal/store/elasticsearch/discovery_search_repository.go
@@ -22,13 +22,13 @@ const (
 	suggesterName                      = "name-phrase-suggest"
 )
 
-var returnedAssetFieldsResult = []string{"id", "urn", "type", "service", "name", "description", "data", "labels", "created_at", "updated_at"}
-
 // Search the asset store
 func (repo *DiscoveryRepository) Search(ctx context.Context, cfg asset.SearchConfig) (results []asset.SearchResult, err error) {
 	if strings.TrimSpace(cfg.Text) == "" {
 		return nil, asset.DiscoveryError{Op: "Search", Err: errors.New("search text cannot be empty")}
 	}
+	var returnedAssetFieldsResult []string
+
 	maxResults := cfg.MaxResults
 	if maxResults <= 0 {
 		maxResults = defaultMaxResults
@@ -38,6 +38,13 @@ func (repo *DiscoveryRepository) Search(ctx context.Context, cfg asset.SearchCon
 		offset = 0
 	}
 
+	if len(cfg.IncludeFields) == 0 {
+		returnedAssetFieldsResult = []string{"id", "urn", "type", "service", "name", "description", "data", "labels",
+			"created_at", "updated_at"}
+	} else {
+		returnedAssetFieldsResult = cfg.IncludeFields
+	}
+
 	defer func(start time.Time) {
 		const op = "search"
 		repo.cli.instrumentOp(ctx, instrumentParams{
@@ -80,7 +87,7 @@ func (repo *DiscoveryRepository) Search(ctx context.Context, cfg asset.SearchCon
 		return nil, asset.DiscoveryError{Op: "Search", Err: fmt.Errorf("decode search response: %w", err)}
 	}
 
-	return toSearchResults(response.Hits.Hits), nil
+	return toSearchResults(response.Hits.Hits, cfg.IncludeFields), nil
 }
 
 func (repo *DiscoveryRepository) GroupAssets(ctx context.Context, cfg asset.GroupConfig) (results []asset.GroupResult, err error) {
@@ -350,6 +357,7 @@ func toSearchResults(hits []searchHit) []asset.SearchResult {
 				"_highlight": hit.HighLight,
 			}
 		}
+
 		results[i] = asset.SearchResult{
 			Type:        r.Type.String(),
 			ID:          id,