From c5e16dbe45aa10d5fec4141eb8f168a5997fb763 Mon Sep 17 00:00:00 2001
From: dobarx <111326505+dobarx@users.noreply.github.com>
Date: Tue, 5 Mar 2024 22:04:26 +0200
Subject: [PATCH] Rename elastic plugin to elastic & extend arguments (#121)

---
 .goreleaser-dev.yaml                          |  6 +--
 .goreleaser.yaml                              | 12 +++---
 docs/plugins/{elasticsearch.md => elastic.md} |  6 +--
 .../{elasticsearch => elastic}/cmd/main.go    |  4 +-
 .../data_elasticsearch.go                     | 35 ++++++++++++++--
 .../data_elasticsearch_test.go                | 33 +++++++++------
 internal/{elasticsearch => elastic}/plugin.go | 40 +++++++++++++++----
 .../{elasticsearch => elastic}/plugin_test.go |  4 +-
 .../testdata/data.json                        |  0
 internal/plugin_validity_test.go              |  4 +-
 plugin/resolver/source_local_test.go          |  6 +--
 tools/docgen/main.go                          |  4 +-
 12 files changed, 108 insertions(+), 46 deletions(-)
 rename docs/plugins/{elasticsearch.md => elastic.md} (91%)
 rename internal/{elasticsearch => elastic}/cmd/main.go (62%)
 rename internal/{elasticsearch => elastic}/data_elasticsearch.go (73%)
 rename internal/{elasticsearch => elastic}/data_elasticsearch_test.go (93%)
 rename internal/{elasticsearch => elastic}/plugin.go (82%)
 rename internal/{elasticsearch => elastic}/plugin_test.go (74%)
 rename internal/{elasticsearch => elastic}/testdata/data.json (100%)

diff --git a/.goreleaser-dev.yaml b/.goreleaser-dev.yaml
index af24e57a..eb7f37b6 100644
--- a/.goreleaser-dev.yaml
+++ b/.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
 
diff --git a/.goreleaser.yaml b/.goreleaser.yaml
index a3dcc651..4b587039 100644
--- a/.goreleaser.yaml
+++ b/.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
diff --git a/docs/plugins/elasticsearch.md b/docs/plugins/elastic.md
similarity index 91%
rename from docs/plugins/elasticsearch.md
rename to docs/plugins/elastic.md
index d40f5c76..b7e33bdd 100644
--- a/docs/plugins/elasticsearch.md
+++ b/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"
   }
 }
 ```
diff --git a/internal/elasticsearch/cmd/main.go b/internal/elastic/cmd/main.go
similarity index 62%
rename from internal/elasticsearch/cmd/main.go
rename to internal/elastic/cmd/main.go
index 987128b2..32548429 100644
--- a/internal/elasticsearch/cmd/main.go
+++ b/internal/elastic/cmd/main.go
@@ -1,7 +1,7 @@
 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"
 )
 
@@ -9,6 +9,6 @@ var version string
 
 func main() {
 	pluginapiv1.Serve(
-		elasticsearch.Plugin(version),
+		elastic.Plugin(version),
 	)
 }
diff --git a/internal/elasticsearch/data_elasticsearch.go b/internal/elastic/data_elasticsearch.go
similarity index 73%
rename from internal/elasticsearch/data_elasticsearch.go
rename to internal/elastic/data_elasticsearch.go
index a5a578ae..781dc762 100644
--- a/internal/elasticsearch/data_elasticsearch.go
+++ b/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
 }
diff --git a/internal/elasticsearch/data_elasticsearch_test.go b/internal/elastic/data_elasticsearch_test.go
similarity index 93%
rename from internal/elasticsearch/data_elasticsearch_test.go
rename to internal/elastic/data_elasticsearch_test.go
index 7a6f3375..3fd3c59f 100644
--- a/internal/elasticsearch/data_elasticsearch_test.go
+++ b/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{
diff --git a/internal/elasticsearch/plugin.go b/internal/elastic/plugin.go
similarity index 82%
rename from internal/elasticsearch/plugin.go
rename to internal/elastic/plugin.go
index ead7e531..b63858dd 100644
--- a/internal/elasticsearch/plugin.go
+++ b/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
+}
diff --git a/internal/elasticsearch/plugin_test.go b/internal/elastic/plugin_test.go
similarity index 74%
rename from internal/elasticsearch/plugin_test.go
rename to internal/elastic/plugin_test.go
index 4b18c3aa..ca79a5f2 100644
--- a/internal/elasticsearch/plugin_test.go
+++ b/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"])
 }
diff --git a/internal/elasticsearch/testdata/data.json b/internal/elastic/testdata/data.json
similarity index 100%
rename from internal/elasticsearch/testdata/data.json
rename to internal/elastic/testdata/data.json
diff --git a/internal/plugin_validity_test.go b/internal/plugin_validity_test.go
index b93b0863..0452f462 100644
--- a/internal/plugin_validity_test.go
+++ b/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),
diff --git a/plugin/resolver/source_local_test.go b/plugin/resolver/source_local_test.go
index 5208c93d..b066617d 100644
--- a/plugin/resolver/source_local_test.go
+++ b/plugin/resolver/source_local_test.go
@@ -24,13 +24,13 @@ func mockFileDir(t *testing.T, files []mockFile) string {
 	tmpDir := t.TempDir()
 	for _, file := range files {
 		if file.isDir {
-			err := os.MkdirAll(filepath.Join(tmpDir, file.path), 0755)
+			err := os.MkdirAll(filepath.Join(tmpDir, file.path), 0o755)
 			require.NoError(t, err)
 			continue
 		}
-		err := os.MkdirAll(filepath.Dir(filepath.Join(tmpDir, file.path)), 0755)
+		err := os.MkdirAll(filepath.Dir(filepath.Join(tmpDir, file.path)), 0o755)
 		require.NoError(t, err)
-		err = os.WriteFile(filepath.Join(tmpDir, file.path), []byte(file.content), 0644)
+		err = os.WriteFile(filepath.Join(tmpDir, file.path), []byte(file.content), 0o644)
 		require.NoError(t, err)
 	}
 	return tmpDir
diff --git a/tools/docgen/main.go b/tools/docgen/main.go
index 97286374..68fff488 100644
--- a/tools/docgen/main.go
+++ b/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),