diff --git a/plugin/hydrate_call.go b/plugin/hydrate_call.go index 989d943a..c92ada60 100644 --- a/plugin/hydrate_call.go +++ b/plugin/hydrate_call.go @@ -16,10 +16,10 @@ type HydrateCall struct { Name string } -func newHydrateCall(hydrateFunc HydrateFunc, config *HydrateConfig) *HydrateCall { +func newHydrateCall(config *HydrateConfig) *HydrateCall { res := &HydrateCall{ - Name: helpers.GetFunctionName(hydrateFunc), - Func: hydrateFunc, + Name: helpers.GetFunctionName(config.Func), + Func: config.Func, Config: config, } for _, f := range config.Depends { diff --git a/plugin/plugin.go b/plugin/plugin.go index f612141c..9a65489a 100644 --- a/plugin/plugin.go +++ b/plugin/plugin.go @@ -21,7 +21,6 @@ import ( "github.com/turbot/steampipe-plugin-sdk/v4/version" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" - "golang.org/x/exp/maps" "golang.org/x/sync/semaphore" "log" "runtime/debug" @@ -435,7 +434,7 @@ func (p *Plugin) executeForConnection(ctx context.Context, req *proto.ExecuteReq log.Printf("[INFO] queryCacheGet returned CACHE MISS (%s)", connectionCallId) // NOTE: update the cache request to include ALL the columns which will be fetched, not just those requested // this means subsequent queries requesting other columns from same hydrate func(s) can be served from the cache - cacheRequest.Columns = maps.Keys(queryData.columns) + cacheRequest.Columns = queryData.getColumnNames() p.queryCache.StartSet(ctx, cacheRequest) } else { log.Printf("[INFO] Cache DISABLED connectionCallId: %s", connectionCallId) diff --git a/plugin/query_data.go b/plugin/query_data.go index 0414a2c0..5569a44f 100644 --- a/plugin/query_data.go +++ b/plugin/query_data.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "github.com/turbot/steampipe-plugin-sdk/v4/error_helpers" + "golang.org/x/exp/maps" "log" "runtime/debug" "sync" @@ -741,3 +742,8 @@ func (d *QueryData) getCacheQualMap() map[string]*proto.Quals { } return res } + +// return the names of all columns that will be returned, adding in the _ctx column +func (d *QueryData) getColumnNames() []string { + return append(maps.Keys(d.columns), ContextColumnName) +} diff --git a/plugin/required_hydrate_calls.go b/plugin/required_hydrate_calls.go index e463307d..91b5ea2b 100644 --- a/plugin/required_hydrate_calls.go +++ b/plugin/required_hydrate_calls.go @@ -31,7 +31,7 @@ func (c requiredHydrateCallBuilder) Add(hydrateFunc HydrateFunc) { // get the config for this hydrate function config := c.table.hydrateConfigMap[hydrateName] - c.requiredHydrateCalls[hydrateName] = newHydrateCall(hydrateFunc, config) + c.requiredHydrateCalls[hydrateName] = newHydrateCall(config) // now add dependencies (we have already checked for circular dependencies so recursion is fine for _, dep := range config.Depends {