From b4d5b024029586c8e77d8b46297ed8f1efa6c097 Mon Sep 17 00:00:00 2001 From: dobarx Date: Tue, 5 Mar 2024 11:23:46 +0200 Subject: [PATCH] internal/elastic: rename plugin to elastic & extend arguments --- .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 +- tools/docgen/main.go | 4 +- 11 files changed, 105 insertions(+), 43 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/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),