internal/elastic: rename plugin to elastic & extend arguments

.goreleaser-dev.yaml

@@ -21,9 +21,9 @@ builds:
 
   # Plugins
 
-  - id: elasticsearch
-    main: ./internal/elasticsearch/cmd
-    binary: "plugins/blackstork/elasticsearch@{{ .Version }}"
+  - id: elastic
+    main: ./internal/elastic/cmd
+    binary: "plugins/blackstork/elastic@{{ .Version }}"
     ldflags: "-X main.version={{.Version}}"
     no_unique_dist_dir: true
 

.goreleaser.yaml

@@ -29,9 +29,9 @@ builds:
   # Plugins
   # TODO: generate this list with custom script or use Premium goreleaser to template it
 
-  - id: plugin_elasticsearch
-    main: ./internal/elasticsearch/cmd
-    binary: "elasticsearch@{{ .Version }}"
+  - id: plugin_elastic
+    main: ./internal/elastic/cmd
+    binary: "plugins/blackstork/elastic@{{ .Version }}"
     flags: "-trimpath"
     hooks:
       post:
@@ -192,12 +192,12 @@ archives:
   # Plugins
   # TODO: generate this list with custom script or use Premium goreleaser to template it
 
-  - id: plugin_elasticsearch
+  - id: plugin_elastic
     format: tar.gz
     builds:
-      - plugin_elasticsearch
+      - plugin_elastic
     name_template: >-
-      plugin_elasticsearch_
+      plugin_elastic_
       {{- .Os }}_
       {{- if eq .Arch "amd64" }}x86_64
       {{- else if eq .Arch "386" }}i386

docs/plugins/elasticsearch.md → docs/plugins/elastic.md

@@ -1,10 +1,10 @@
 ---
-title: blackstork/elasticsearch
+title: blackstork/elastic
 weight: 20
 type: docs
 ---
 
-# `blackstork/elasticsearch` plugin
+# `blackstork/elastic` plugin
 
 ## Installation
 
@@ -13,7 +13,7 @@ To install the plugin, add it to `plugin_versions` map in the Fabric global conf
 ```hcl
 fabric {
   plugin_versions = {
-    "blackstork/elasticsearch" = "=> v0.0.0-dev"
+    "blackstork/elastic" = "=> v0.0.0-dev"
   }
 }
 ```

internal/elasticsearch/cmd/main.go → internal/elastic/cmd/main.go

@@ -1,14 +1,14 @@
 package main
 
 import (
-	"github.com/blackstork-io/fabric/internal/elasticsearch"
+	"github.com/blackstork-io/fabric/internal/elastic"
 	pluginapiv1 "github.com/blackstork-io/fabric/plugin/pluginapi/v1"
 )
 
 var version string
 
 func main() {
 	pluginapiv1.Serve(
-		elasticsearch.Plugin(version),
+		elastic.Plugin(version),
 	)
 }

internal/elasticsearch/data_elasticsearch.go → internal/elastic/data_elasticsearch.go

@@ -1,4 +1,4 @@
-package elasticsearch
+package elastic
 
 import (
 	"context"
@@ -76,6 +76,16 @@ func makeElasticSearchDataSource() *plugin.DataSource {
 				Type:     cty.Map(cty.DynamicPseudoType),
 				Required: false,
 			},
+			"aggs": &hcldec.AttrSpec{
+				Name:     "aggs",
+				Type:     cty.Map(cty.DynamicPseudoType),
+				Required: false,
+			},
+			"only_hits": &hcldec.AttrSpec{
+				Name:     "only_hits",
+				Type:     cty.Bool,
+				Required: false,
+			},
 			"fields": &hcldec.AttrSpec{
 				Name:     "fields",
 				Type:     cty.List(cty.String),
@@ -99,6 +109,23 @@ func fetchElasticSearchData(ctx context.Context, params *plugin.RetrieveDataPara
 			Detail:   err.Error(),
 		}}
 	}
+	var diags hcl.Diagnostics
+	if (params.Args.GetAttr("only_hits").IsNull() || params.Args.GetAttr("only_hits").True()) &&
+		!params.Args.GetAttr("aggs").IsNull() {
+		if params.Args.GetAttr("query").IsNull() && params.Args.GetAttr("query_string").IsNull() {
+			return nil, hcl.Diagnostics{{
+				Severity: hcl.DiagError,
+				Summary:  "Invalid arguments",
+				Detail:   "Aggregations are not supported without a query or query_string",
+			}}
+		}
+		diags = diags.Append(&hcl.Diagnostic{
+			Severity: hcl.DiagWarning,
+			Summary:  "Aggregations are not supported",
+			Detail:   "Aggregations are not supported when only_hits is true",
+		})
+	}
+
 	id := params.Args.GetAttr("id")
 	var data plugin.Data
 	if !id.IsNull() {
@@ -107,11 +134,11 @@ func fetchElasticSearchData(ctx context.Context, params *plugin.RetrieveDataPara
 		data, err = search(client.Search, params.Args)
 	}
 	if err != nil {
-		return nil, hcl.Diagnostics{{
+		return nil, diags.Extend(hcl.Diagnostics{{
 			Severity: hcl.DiagError,
 			Summary:  "Failed to get data",
 			Detail:   err.Error(),
-		}}
+		}})
 	}
-	return data, nil
+	return data, diags
 }

internal/elasticsearch/data_elasticsearch_test.go → internal/elastic/data_elasticsearch_test.go

@@ -1,4 +1,4 @@
-package elasticsearch
+package elastic
 
 import (
 	"bytes"
@@ -116,6 +116,8 @@ func (s *IntegrationTestSuite) TestSearchDefaults() {
 		"index":        cty.StringVal("test_index"),
 		"query":        cty.NullVal(cty.DynamicPseudoType),
 		"query_string": cty.NullVal(cty.String),
+		"only_hits":    cty.NullVal(cty.Bool),
+		"aggs":         cty.NullVal(cty.DynamicPseudoType),
 		"fields":       cty.NullVal(cty.String),
 		"size":         cty.NullVal(cty.Number),
 	})
@@ -124,11 +126,10 @@ func (s *IntegrationTestSuite) TestSearchDefaults() {
 		Args:   args,
 	})
 	s.Require().Nil(diags)
-	m := data.(plugin.MapData)
-	raw, err := json.MarshalIndent(m["hits"], "", "  ")
+	m := data.(plugin.ListData)
+	raw, err := json.MarshalIndent(m, "", "  ")
 	s.Require().NoError(err, "failed to marshal data: %v", err)
-	s.JSONEq(`{
-		"hits": [
+	s.JSONEq(`[
 		  {
 			"_id": "54f7a815-eac5-4f7c-a339-5fefd0f54967",
 			"_index": "test_index",
@@ -165,13 +166,7 @@ func (s *IntegrationTestSuite) TestSearchDefaults() {
 			  "type": "foo"
 			}
 		  }
-		],
-		"max_score": 1,
-		"total": {
-		  "relation": "eq",
-		  "value": 3
-		}
-	  }`, string(raw))
+		]`, string(raw))
 }
 
 func (s *IntegrationTestSuite) TestSearchFields() {
@@ -180,6 +175,8 @@ func (s *IntegrationTestSuite) TestSearchFields() {
 		"index":        cty.StringVal("test_index"),
 		"query":        cty.NullVal(cty.DynamicPseudoType),
 		"query_string": cty.NullVal(cty.String),
+		"only_hits":    cty.BoolVal(false),
+		"aggs":         cty.NullVal(cty.DynamicPseudoType),
 		"fields":       cty.ListVal([]cty.Value{cty.StringVal("name"), cty.StringVal("age")}),
 		"size":         cty.NullVal(cty.Number),
 	})
@@ -235,6 +232,8 @@ func (s *IntegrationTestSuite) TestSearchQueryString() {
 		"index":        cty.StringVal("test_index"),
 		"query":        cty.NullVal(cty.DynamicPseudoType),
 		"query_string": cty.StringVal("type:foo"),
+		"only_hits":    cty.BoolVal(false),
+		"aggs":         cty.NullVal(cty.DynamicPseudoType),
 		"fields":       cty.NullVal(cty.String),
 		"size":         cty.NullVal(cty.Number),
 	})
@@ -289,6 +288,8 @@ func (s *IntegrationTestSuite) TestSearchQuery() {
 			"match_all": cty.MapValEmpty(cty.DynamicPseudoType),
 		}),
 		"query_string": cty.NullVal(cty.String),
+		"only_hits":    cty.BoolVal(false),
+		"aggs":         cty.NullVal(cty.DynamicPseudoType),
 		"fields":       cty.NullVal(cty.String),
 		"size":         cty.NullVal(cty.Number),
 	})
@@ -353,6 +354,8 @@ func (s *IntegrationTestSuite) TestSearchSize() {
 		"index":        cty.StringVal("test_index"),
 		"query":        cty.NullVal(cty.DynamicPseudoType),
 		"query_string": cty.NullVal(cty.String),
+		"only_hits":    cty.BoolVal(false),
+		"aggs":         cty.NullVal(cty.DynamicPseudoType),
 		"fields":       cty.NullVal(cty.String),
 		"size":         cty.NumberIntVal(1),
 	})
@@ -393,6 +396,8 @@ func (s *IntegrationTestSuite) TestGetByID() {
 		"index":        cty.StringVal("test_index"),
 		"query":        cty.NullVal(cty.DynamicPseudoType),
 		"query_string": cty.NullVal(cty.String),
+		"only_hits":    cty.BoolVal(false),
+		"aggs":         cty.NullVal(cty.DynamicPseudoType),
 		"fields":       cty.NullVal(cty.String),
 	})
 	data, diags := s.plugin.RetrieveData(s.ctx, "elasticsearch", &plugin.RetrieveDataParams{
@@ -426,6 +431,8 @@ func (s *IntegrationTestSuite) TestGetByIDFields() {
 		"index":        cty.StringVal("test_index"),
 		"query":        cty.NullVal(cty.DynamicPseudoType),
 		"query_string": cty.NullVal(cty.String),
+		"only_hits":    cty.BoolVal(false),
+		"aggs":         cty.NullVal(cty.DynamicPseudoType),
 		"fields":       cty.ListVal([]cty.Value{cty.StringVal("name"), cty.StringVal("age")}),
 	})
 	data, diags := s.plugin.RetrieveData(s.ctx, "elasticsearch", &plugin.RetrieveDataParams{
@@ -456,6 +463,8 @@ func (s *IntegrationTestSuite) TestGetByIDNotFound() {
 		"index":        cty.StringVal("test_index"),
 		"query":        cty.NullVal(cty.DynamicPseudoType),
 		"query_string": cty.NullVal(cty.String),
+		"only_hits":    cty.BoolVal(false),
+		"aggs":         cty.NullVal(cty.DynamicPseudoType),
 		"fields":       cty.NullVal(cty.String),
 	})
 	data, diags := s.plugin.RetrieveData(s.ctx, "elasticsearch", &plugin.RetrieveDataParams{

internal/elasticsearch/plugin.go → internal/elastic/plugin.go

@@ -1,4 +1,4 @@
-package elasticsearch
+package elastic
 
 import (
 	"bytes"
@@ -21,7 +21,7 @@ const (
 
 func Plugin(version string) *plugin.Schema {
 	return &plugin.Schema{
-		Name:    "blackstork/elasticsearch",
+		Name:    "blackstork/elastic",
 		Version: version,
 		DataSources: plugin.DataSources{
 			"elasticsearch": makeElasticSearchDataSource(),
@@ -106,6 +106,9 @@ func getByID(fn esapi.Get, args cty.Value) (plugin.Data, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to unmarshal search result: %s", err)
 	}
+	if onlyHits := args.GetAttr("only_hits"); onlyHits.IsNull() || onlyHits.True() {
+		return extractHits(data)
+	}
 	return data, nil
 }
 
@@ -120,14 +123,19 @@ func search(fn esapi.Search, args cty.Value) (plugin.Data, error) {
 	if queryString := args.GetAttr("query_string"); !queryString.IsNull() {
 		opts = append(opts, fn.WithQuery(queryString.AsString()))
 	}
+	body := map[string]any{}
 	if query := args.GetAttr("query"); !query.IsNull() {
-		queryRaw, err := json.Marshal(map[string]any{
-			"query": query.AsValueMap(),
-		})
+		body["query"] = query.AsValueMap()
+	}
+	if aggs := args.GetAttr("aggs"); !aggs.IsNull() {
+		body["aggs"] = aggs.AsValueMap()
+	}
+	if len(body) > 0 {
+		rawBody, err := json.Marshal(body)
 		if err != nil {
 			return nil, fmt.Errorf("failed to marshal query: %s", err)
 		}
-		opts = append(opts, fn.WithBody(bytes.NewReader(queryRaw)))
+		opts = append(opts, fn.WithBody(bytes.NewReader(rawBody)))
 	}
 	if size := args.GetAttr("size"); !size.IsNull() {
 		n, _ := size.AsBigFloat().Int64()
@@ -144,7 +152,6 @@ func search(fn esapi.Search, args cty.Value) (plugin.Data, error) {
 		}
 		opts = append(opts, fn.WithSource(fieldStrings...))
 	}
-
 	res, err := fn(opts...)
 	if err != nil {
 		return nil, err
@@ -159,5 +166,24 @@ func search(fn esapi.Search, args cty.Value) (plugin.Data, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to unmarshal search result: %s", err)
 	}
+	if onlyHits := args.GetAttr("only_hits"); onlyHits.IsNull() || onlyHits.True() {
+		return extractHits(data)
+	}
 	return data, nil
 }
+
+func extractHits(data plugin.Data) (plugin.Data, error) {
+	m, ok := data.(plugin.MapData)
+	if !ok {
+		return nil, fmt.Errorf("unexpected search result type: %T", data)
+	}
+	data, ok = m["hits"]
+	if !ok {
+		return nil, fmt.Errorf("unexpected search result type: %T", data)
+	}
+	m, ok = data.(plugin.MapData)
+	if !ok {
+		return nil, fmt.Errorf("unexpected search result type: %T", data)
+	}
+	return m["hits"], nil
+}

internal/elasticsearch/plugin_test.go → internal/elastic/plugin_test.go

@@ -1,4 +1,4 @@
-package elasticsearch
+package elastic
 
 import (
 	"testing"
@@ -8,7 +8,7 @@ import (
 
 func TestPlugin_Schema(t *testing.T) {
 	schema := Plugin("1.2.3")
-	assert.Equal(t, "blackstork/elasticsearch", schema.Name)
+	assert.Equal(t, "blackstork/elastic", schema.Name)
 	assert.Equal(t, "1.2.3", schema.Version)
 	assert.NotNil(t, schema.DataSources["elasticsearch"])
 }

internal/plugin_validity_test.go

@@ -9,7 +9,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/blackstork-io/fabric/internal/builtin"
-	"github.com/blackstork-io/fabric/internal/elasticsearch"
+	"github.com/blackstork-io/fabric/internal/elastic"
 	"github.com/blackstork-io/fabric/internal/github"
 	"github.com/blackstork-io/fabric/internal/graphql"
 	"github.com/blackstork-io/fabric/internal/hackerone"
@@ -29,7 +29,7 @@ func TestAllPluginSchemaValidity(t *testing.T) {
 	ver := "1.2.3"
 	plugins := []*plugin.Schema{
 		builtin.Plugin(ver),
-		elasticsearch.Plugin(ver),
+		elastic.Plugin(ver),
 		github.Plugin(ver, nil),
 		graphql.Plugin(ver),
 		openai.Plugin(ver, nil),

tools/docgen/main.go

@@ -12,7 +12,7 @@ import (
 	"github.com/spf13/pflag"
 
 	"github.com/blackstork-io/fabric/internal/builtin"
-	"github.com/blackstork-io/fabric/internal/elasticsearch"
+	"github.com/blackstork-io/fabric/internal/elastic"
 	"github.com/blackstork-io/fabric/internal/github"
 	"github.com/blackstork-io/fabric/internal/graphql"
 	"github.com/blackstork-io/fabric/internal/hackerone"
@@ -49,7 +49,7 @@ func main() {
 	// load all plugins
 	plugins := []*plugin.Schema{
 		builtin.Plugin(version),
-		elasticsearch.Plugin(version),
+		elastic.Plugin(version),
 		github.Plugin(version, nil),
 		graphql.Plugin(version),
 		openai.Plugin(version, nil),