From 3c550b938efb0713664f6ee70d010aad5752d236 Mon Sep 17 00:00:00 2001 From: Ralph Caraveo Date: Mon, 8 Jan 2024 14:57:53 -0800 Subject: [PATCH] Support prefix filtering with multiple OR prefixes such as: ./tips "foo | bar" --- cmd/utility.go | 6 +++--- cmd/utility_test.go | 2 +- pkg/config_ctx.go | 36 +++++++++++++++++++++++++++++++++++- pkg/config_ctx_test.go | 26 ++++++++++++++++++++++++++ pkg/db_generic.go | 32 ++++++++++++++++++-------------- pkg/repository_cached.go | 4 ++-- 6 files changed, 85 insertions(+), 21 deletions(-) diff --git a/cmd/utility.go b/cmd/utility.go index a841e33..a25c8ee 100644 --- a/cmd/utility.go +++ b/cmd/utility.go @@ -22,12 +22,12 @@ func packageCfg(args []string) (*pkg.ConfigCtx, error) { // The 0th arg is the Primary filter, if nothing was specified we consider it to represent: @ for all if len(args) > 0 { if strings.TrimSpace(args[0]) == allFilterCLI { - cfgCtx.PrefixFilter = allFilter + cfgCtx.PrefixFilter = pkg.ParsePrefixFilter("*") } else { - cfgCtx.PrefixFilter = args[0] + cfgCtx.PrefixFilter = pkg.ParsePrefixFilter(args[0]) } } else { - cfgCtx.PrefixFilter = allFilter + cfgCtx.PrefixFilter = pkg.ParsePrefixFilter("*") } // The 1st arg along with the rest - [1:] when provided is a remote command to execute. diff --git a/cmd/utility_test.go b/cmd/utility_test.go index c079c74..d00e6af 100644 --- a/cmd/utility_test.go +++ b/cmd/utility_test.go @@ -28,7 +28,7 @@ func TestPackageCfg(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, cfg) - assert.Equal(t, cfg.PrefixFilter, "*") + assert.True(t, cfg.PrefixFilter.IsAll()) assert.Equal(t, cfg.RemoteCmd, "echo 'hello world' && sleep 0.5 && ps aux | grep foo") } diff --git a/pkg/config_ctx.go b/pkg/config_ctx.go index 9e3f5a9..34c9997 100644 --- a/pkg/config_ctx.go +++ b/pkg/config_ctx.go @@ -26,6 +26,7 @@ SOFTWARE. package pkg import ( + "slices" "strconv" "github.com/charmbracelet/log" @@ -113,7 +114,7 @@ type ConfigCtx struct { JsonOutput bool NoCache bool NoColor bool - PrefixFilter string // We'll do prefix filtering in Boltdb + PrefixFilter *PrefixFilter RemoteCmd string Slice *SliceCfg SortOrder []SortSpec @@ -148,3 +149,36 @@ func ParseColumns(s string) mapset.Set[string] { return m } + +type PrefixFilter struct { + originalQuery string + orPrefixes []string +} + +func (p *PrefixFilter) IsAll() bool { + return len(p.orPrefixes) == 1 && (p.orPrefixes[0] == "*" || p.orPrefixes[0] == "@") +} + +func (p *PrefixFilter) Count() int { + return len(p.orPrefixes) +} + +func (p *PrefixFilter) PrefixAt(idx int) string { + return p.orPrefixes[idx] +} + +func ParsePrefixFilter(s string) *PrefixFilter { + var prefixes []string + parts := strings.Split(s, "|") + for _, p := range parts { + prefixes = append(prefixes, strings.TrimSpace(p)) + } + + slices.Sort(prefixes) + + pf := &PrefixFilter{ + originalQuery: s, + orPrefixes: prefixes, + } + return pf +} diff --git a/pkg/config_ctx_test.go b/pkg/config_ctx_test.go index 9d39a8d..55d46b5 100644 --- a/pkg/config_ctx_test.go +++ b/pkg/config_ctx_test.go @@ -2,6 +2,8 @@ package pkg import ( "testing" + + "github.com/stretchr/testify/assert" ) func TestParseSlice(t *testing.T) { @@ -64,3 +66,27 @@ func TestParseSlice(t *testing.T) { t.Error("expected slice to be defined") } } + +func TestParsePrefixFilter(t *testing.T) { + // The prefix for everything prefix + pf := ParsePrefixFilter("*") + assert.NotNil(t, pf) + assert.True(t, pf.IsAll()) + + // The prefix for everything prefix (the same thing) + pf = ParsePrefixFilter("@") + assert.NotNil(t, pf) + assert.True(t, pf.IsAll()) + + // The prefix with multiple OR conditions + pf = ParsePrefixFilter("foo|bar | baz") + assert.NotNil(t, pf) + assert.False(t, pf.IsAll()) + + assert.Equal(t, pf.Count(), 3) + + expectedPrefixOrder := []string{"bar", "baz", "foo"} + for i := 0; i < pf.Count(); i++ { + assert.Equal(t, pf.PrefixAt(i), expectedPrefixOrder[i]) + } +} diff --git a/pkg/db_generic.go b/pkg/db_generic.go index 4fa6234..ec068aa 100644 --- a/pkg/db_generic.go +++ b/pkg/db_generic.go @@ -145,9 +145,9 @@ func (d *DB2[T]) IndexOpaqueItems(ctx context.Context, bucketName string, items return nil } -type DBQuery2 struct { - PrefixFilter string - PrimaryKeys []string +type DBQuery struct { + PrefixFilters *PrefixFilter + PrimaryKeys []string } func (d *DB2[T]) LookupOpaqueItem(ctx context.Context, bucketName, primaryKey string) (*T, error) { @@ -181,10 +181,10 @@ func (d *DB2[T]) lookupOpaqueItem(bucket *bolt.Bucket, primaryKey string) (*T, e // 1. Using one or more primary keys, in which case this is a direct lookup (not technically a search) // 2. Using the * (all/everything) construct, this is just a full table scan really. // 3. Using a prefix scan, this is a seek to a segment of the index and should be fast assuming good selectivity. -func (d *DB2[T]) SearchOpaqueItems(ctx context.Context, bucketName string, query DBQuery2) ([]T, error) { +func (d *DB2[T]) SearchOpaqueItems(ctx context.Context, bucketName string, query DBQuery) ([]T, error) { //cfg := CtxAsConfig(ctx, CtxKeyConfig) - if strings.TrimSpace(query.PrefixFilter) == "" { + if query.PrefixFilters.Count() == 0 { panic("query.PrefixFilter must never be empty, in the case of all it must be: *") } @@ -196,7 +196,6 @@ func (d *DB2[T]) SearchOpaqueItems(ctx context.Context, bucketName string, query return errors.New("bucket is unknown: " + bucketName) } - c := b.Cursor() // Search by primary keys, this a direct lookup, the fastest. if len(query.PrimaryKeys) > 0 { for _, pk := range query.PrimaryKeys { @@ -206,7 +205,8 @@ func (d *DB2[T]) SearchOpaqueItems(ctx context.Context, bucketName string, query } items = append(items, *item) } - } else if query.PrefixFilter == "*" { + } else if query.PrefixFilters.IsAll() { + c := b.Cursor() // Search by everything, linear (full-table scan) for k, v := c.First(); k != nil; k, v = c.Next() { var item T @@ -216,14 +216,18 @@ func (d *DB2[T]) SearchOpaqueItems(ctx context.Context, bucketName string, query items = append(items, item) } } else { - // Prefix scan (much faster) when a prefix is present. - prefix := []byte(query.PrefixFilter) - for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { - var item T - if err := json.Unmarshal(v, &item); err != nil { - return err + // Since Prefix Filters support OR filters: "foo|bar" we do this in a loop with a new cursor for each prefix. + // Most of the time, just one iteration occurs. + for i := 0; i < query.PrefixFilters.Count(); i++ { + c := b.Cursor() + prefix := []byte(query.PrefixFilters.PrefixAt(i)) + for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { + var item T + if err := json.Unmarshal(v, &item); err != nil { + return err + } + items = append(items, item) } - items = append(items, item) } } diff --git a/pkg/repository_cached.go b/pkg/repository_cached.go index a9b22c2..08a5118 100644 --- a/pkg/repository_cached.go +++ b/pkg/repository_cached.go @@ -67,7 +67,7 @@ func (c *CachedRepository) DevicesResource(ctx context.Context) ([]*WrappedDevic // Care is taken to measure just cache retrieval time. cachedStartTime := time.Now() - devList, err := deviceIndexedRepo.SearchOpaqueItems(ctx, DevicesBucket, DBQuery2{PrefixFilter: cfg.PrefixFilter}) + devList, err := deviceIndexedRepo.SearchOpaqueItems(ctx, DevicesBucket, DBQuery{PrefixFilters: cfg.PrefixFilter}) if err != nil { return nil, err } @@ -105,7 +105,7 @@ func (c *CachedRepository) DevicesResource(ctx context.Context) ([]*WrappedDevic // 4. Return the data from the db because the db can utilize the index on prefix filters. // In the future it may also do other heavyweight filters that we don't have to do in "user space" - devList, err = deviceIndexedRepo.SearchOpaqueItems(ctx, DevicesBucket, DBQuery2{PrefixFilter: cfg.PrefixFilter}) + devList, err = deviceIndexedRepo.SearchOpaqueItems(ctx, DevicesBucket, DBQuery{PrefixFilters: cfg.PrefixFilter}) if err != nil { return nil, err }