feat: implement GroupAssets API with ES composite aggregation

- Add GroupConfig, GroupResult, GroupField types to core/asset/discovery.go
- Add GroupAssets to DiscoveryRepository interface
- Implement GroupAssets in ES using composite aggregation with top_hits
- Add GroupAssets gRPC handler with filter and include_fields support
- Add GroupAssets to AssetService interface and service delegate
- Update mocks for new interface methods

Author: ravisuhag

core/asset/discovery.go

@@ -14,6 +14,28 @@ type DiscoveryRepository interface {
 	SoftDeleteByURN(ctx context.Context, ns *namespace.Namespace, assetURN string) error
 	Search(ctx context.Context, cfg SearchConfig) (results []SearchResult, err error)
 	Suggest(ctx context.Context, cfg SearchConfig) (suggestions []string, err error)
+	GroupAssets(ctx context.Context, cfg GroupConfig) ([]GroupResult, error)
+}
+
+// GroupConfig represents configuration for grouping assets
+type GroupConfig struct {
+	GroupBy       []string
+	Filters       SearchFilter
+	IncludeFields []string
+	Size          int
+	Namespace     *namespace.Namespace
+}
+
+// GroupResult represents a single group of assets
+type GroupResult struct {
+	Fields []GroupField
+	Assets []Asset
+}
+
+// GroupField represents a key-value pair identifying a group
+type GroupField struct {
+	Key   string
+	Value string
 }
 
 // SearchFilter is a filter intended to be used as a search

core/asset/mocks/discovery_repository.go

@@ -195,6 +195,10 @@ func (s *Service) GetTypes(ctx context.Context, flt Filter) (map[Type]int, error
 func (s *Service) SearchAssets(ctx context.Context, cfg SearchConfig) (results []SearchResult, err error) {
 	return s.discoveryRepository.Search(ctx, cfg)
 }
+
+func (s *Service) GroupAssets(ctx context.Context, cfg GroupConfig) ([]GroupResult, error) {
+	return s.discoveryRepository.GroupAssets(ctx, cfg)
+}
 func (s *Service) SuggestAssets(ctx context.Context, cfg SearchConfig) (suggestions []string, err error) {
 	return s.discoveryRepository.Suggest(ctx, cfg)
 }

core/asset/service.go

@@ -36,6 +36,8 @@ type AssetService interface {
 	SuggestAssets(ctx context.Context, cfg asset.SearchConfig) (suggestions []string, err error)
 
 	AddProbe(ctx context.Context, ns *namespace.Namespace, assetURN string, probe *asset.Probe) error
+
+	GroupAssets(ctx context.Context, cfg asset.GroupConfig) ([]asset.GroupResult, error)
 }
 
 func (server *APIServer) GetAllAssets(ctx context.Context, req *compassv1beta1.GetAllAssetsRequest) (*compassv1beta1.GetAllAssetsResponse, error) {

internal/server/v1beta1/asset.go

@@ -80,6 +80,60 @@ func (server *APIServer) SuggestAssets(ctx context.Context, req *compassv1beta1.
 	}, nil
 }
 
+func (server *APIServer) GroupAssets(ctx context.Context, req *compassv1beta1.GroupAssetsRequest) (*compassv1beta1.GroupAssetsResponse, error) {
+	if err := req.ValidateAll(); err != nil {
+		return nil, status.Error(codes.InvalidArgument, bodyParserErrorMsg(err))
+	}
+	ns := grpc_interceptor.FetchNamespaceFromContext(ctx)
+	if _, err := server.validateUserInCtx(ctx, ns); err != nil {
+		return nil, err
+	}
+
+	if len(req.GetGroupby()) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "groupby must be specified")
+	}
+
+	cfg := asset.GroupConfig{
+		GroupBy:       req.GetGroupby(),
+		Filters:       filterConfigFromValues(req.GetFilter()),
+		IncludeFields: req.GetIncludeFields(),
+		Size:          int(req.GetSize()),
+		Namespace:     ns,
+	}
+
+	results, err := server.assetService.GroupAssets(ctx, cfg)
+	if err != nil {
+		return nil, internalServerError(server.logger, fmt.Sprintf("error grouping assets: %s", err.Error()))
+	}
+
+	var groups []*compassv1beta1.AssetGroup
+	for _, gr := range results {
+		var fields []*compassv1beta1.GroupField
+		for _, f := range gr.Fields {
+			fields = append(fields, &compassv1beta1.GroupField{
+				GroupKey:   f.Key,
+				GroupValue: f.Value,
+			})
+		}
+		var assets []*compassv1beta1.Asset
+		for _, a := range gr.Assets {
+			ap, err := assetToProto(a, false)
+			if err != nil {
+				return nil, internalServerError(server.logger, fmt.Sprintf("error converting asset to proto: %s", err.Error()))
+			}
+			assets = append(assets, ap)
+		}
+		groups = append(groups, &compassv1beta1.AssetGroup{
+			GroupFields: fields,
+			Assets:      assets,
+		})
+	}
+
+	return &compassv1beta1.GroupAssetsResponse{
+		AssetGroups: groups,
+	}, nil
+}
+
 func getSearchFlagsFromProto(flags *compassv1beta1.SearchFlags) asset.SearchFlags {
 	if flags == nil {
 		return asset.SearchFlags{}

internal/server/v1beta1/mocks/asset_service.go

@@ -330,3 +330,152 @@ func (repo *DiscoveryRepository) toSuggestions(response searchResponse) (results
 
 	return
 }
+
+const defaultGroupsSize = 10
+
+// GroupAssets groups assets by specified fields using ES composite aggregation
+func (repo *DiscoveryRepository) GroupAssets(ctx context.Context, cfg asset.GroupConfig) ([]asset.GroupResult, error) {
+	if cfg.Namespace == nil {
+		return nil, asset.DiscoveryError{Err: fmt.Errorf("namespace cannot be empty")}
+	}
+
+	size := cfg.Size
+	if size <= 0 {
+		size = defaultGroupsSize
+	}
+
+	// Build composite aggregation sources from group-by fields
+	sources := make([]map[string]interface{}, 0, len(cfg.GroupBy))
+	for _, field := range cfg.GroupBy {
+		sources = append(sources, map[string]interface{}{
+			field: map[string]interface{}{
+				"terms": map[string]interface{}{
+					"field": fmt.Sprintf("%s.keyword", field),
+				},
+			},
+		})
+	}
+
+	// Build filter query
+	boolQuery := elastic.NewBoolQuery()
+	// Ensure group-by fields exist
+	for _, field := range cfg.GroupBy {
+		boolQuery.Filter(elastic.NewExistsQuery(fmt.Sprintf("%s.keyword", field)))
+	}
+	repo.buildFilterTermQueries(boolQuery, cfg.Filters)
+
+	includedFields := defaultIncludedFields
+	if len(cfg.IncludeFields) > 0 {
+		includedFields = cfg.IncludeFields
+	}
+
+	// Build the aggregation query manually as JSON
+	payload := map[string]interface{}{
+		"size": 0,
+		"query": map[string]interface{}{
+			"bool": map[string]interface{}{},
+		},
+		"aggs": map[string]interface{}{
+			"group_result": map[string]interface{}{
+				"composite": map[string]interface{}{
+					"size":    size,
+					"sources": sources,
+				},
+				"aggs": map[string]interface{}{
+					"top_assets": map[string]interface{}{
+						"top_hits": map[string]interface{}{
+							"size":     size,
+							"_source":  includedFields,
+						},
+					},
+				},
+			},
+		},
+	}
+
+	// Use the bool query
+	boolSrc, err := boolQuery.Source()
+	if err != nil {
+		return nil, asset.DiscoveryError{Err: fmt.Errorf("error building query: %w", err)}
+	}
+	payload["query"] = boolSrc
+
+	body, err := json.Marshal(payload)
+	if err != nil {
+		return nil, asset.DiscoveryError{Err: fmt.Errorf("error encoding query: %w", err)}
+	}
+
+	res, err := repo.cli.client.Search(
+		repo.cli.client.Search.WithBody(strings.NewReader(string(body))),
+		repo.cli.client.Search.WithIndex(BuildAliasNameFromNamespace(cfg.Namespace)),
+		repo.cli.client.Search.WithIgnoreUnavailable(true),
+		repo.cli.client.Search.WithContext(ctx),
+	)
+	if err != nil {
+		return nil, asset.DiscoveryError{Err: fmt.Errorf("error executing group query: %w", err)}
+	}
+	defer res.Body.Close()
+
+	if res.IsError() {
+		return nil, asset.DiscoveryError{Err: fmt.Errorf("error response from elasticsearch: %s", errorReasonFromResponse(res))}
+	}
+
+	var response groupResponse
+	if err := json.NewDecoder(res.Body).Decode(&response); err != nil {
+		return nil, asset.DiscoveryError{Err: fmt.Errorf("error decoding group response: %w", err)}
+	}
+
+	return repo.toGroupResults(response, cfg.GroupBy), nil
+}
+
+type groupResponse struct {
+	Aggregations struct {
+		GroupResult struct {
+			Buckets []groupBucket `json:"buckets"`
+		} `json:"group_result"`
+	} `json:"aggregations"`
+}
+
+type groupBucket struct {
+	Key       map[string]string `json:"key"`
+	DocCount  int               `json:"doc_count"`
+	TopAssets struct {
+		Hits struct {
+			Hits []searchHit `json:"hits"`
+		} `json:"hits"`
+	} `json:"top_assets"`
+}
+
+func (repo *DiscoveryRepository) toGroupResults(response groupResponse, groupBy []string) []asset.GroupResult {
+	var results []asset.GroupResult
+	for _, bucket := range response.Aggregations.GroupResult.Buckets {
+		var fields []asset.GroupField
+		for _, key := range groupBy {
+			fields = append(fields, asset.GroupField{
+				Key:   key,
+				Value: bucket.Key[key],
+			})
+		}
+
+		var assets []asset.Asset
+		for _, hit := range bucket.TopAssets.Hits.Hits {
+			r := hit.Source
+			assets = append(assets, asset.Asset{
+				ID:          r.ID,
+				URN:         r.URN,
+				Type:        r.Type,
+				Name:        r.Name,
+				Service:     r.Service,
+				Description: r.Description,
+				Data:        r.Data,
+				Labels:      r.Labels,
+			})
+		}
+
+		results = append(results, asset.GroupResult{
+			Fields: fields,
+			Assets: assets,
+		})
+	}
+	return results
+}