Change Keys method in kv interface to Scan (guacsec#1558)

Keys was a placeholder, as an external keyvalue may have more keys than we want to read at once. Additionally we don't want to overload the external store with an operation that will block for a long time. Now the Scan interface will only get some number of keys at a time. Implementations for Redis and TiKV have been added. Signed-off-by: Jeff Mendoza <jlm@jlm.name>
mrizzi · Dec 5, 2023 · 18ad0d0 · 18ad0d0
1 parent 30218ea
commit 18ad0d0
Show file tree

Hide file tree

Showing 28 changed files with 500 additions and 265 deletions.
diff --git a/internal/testing/stablememmap/stablememmap.go b/internal/testing/stablememmap/stablememmap.go
@@ -17,35 +17,47 @@ package stablememmap
 
 import (
 	"context"
+	"errors"
 	"slices"
 
 	"github.com/guacsec/guac/pkg/assembler/kv"
 	"github.com/guacsec/guac/pkg/assembler/kv/memmap"
 )
 
-type Store struct {
+type store struct {
 	mm kv.Store
 }
 
 func GetStore() kv.Store {
-	return &Store{
+	return &store{
 		mm: memmap.GetStore(),
 	}
 }
 
-func (s *Store) Get(ctx context.Context, c, k string, v any) error {
+func (s *store) Get(ctx context.Context, c, k string, v any) error {
 	return s.mm.Get(ctx, c, k, v)
 }
 
-func (s *Store) Set(ctx context.Context, c, k string, v any) error {
+func (s *store) Set(ctx context.Context, c, k string, v any) error {
 	return s.mm.Set(ctx, c, k, v)
 }
 
-func (s *Store) Keys(ctx context.Context, c string) ([]string, error) {
-	keys, err := s.mm.Keys(ctx, c)
+func (s *store) Keys(c string) kv.Scanner {
+	return &scanner{mms: s.mm.Keys(c)}
+}
+
+type scanner struct {
+	mms kv.Scanner
+}
+
+func (s *scanner) Scan(ctx context.Context) ([]string, bool, error) {
+	keys, done, err := s.mms.Scan(ctx)
 	if err != nil {
-		return nil, err
+		return nil, false, err
+	}
+	if !done {
+		return nil, false, errors.New("Expect memmap to always return all keys at once")
 	}
 	slices.Sort(keys)
-	return keys, nil
+	return keys, true, nil
 }
diff --git a/pkg/assembler/backends/keyvalue/artifact.go b/pkg/assembler/backends/keyvalue/artifact.go
@@ -239,28 +239,33 @@ func (c *demoClient) Artifacts(ctx context.Context, artifactSpec *model.Artifact
 	algorithm := strings.ToLower(nilToEmpty(artifactSpec.Algorithm))
 	digest := strings.ToLower(nilToEmpty(artifactSpec.Digest))
 	var rv []*model.Artifact
-	artKeys, err := c.kv.Keys(ctx, artCol)
-	if err != nil {
-		return nil, err
-	}
-	for _, ak := range artKeys {
-		a, err := byKeykv[*artStruct](ctx, artCol, ak, c)
+	var done bool
+	scn := c.kv.Keys(artCol)
+	for !done {
+		var artKeys []string
+		artKeys, done, err = scn.Scan(ctx)
 		if err != nil {
 			return nil, err
 		}
-
-		matchAlgorithm := false
-		if algorithm == "" || algorithm == a.Algorithm {
-			matchAlgorithm = true
-		}
-
-		matchDigest := false
-		if digest == "" || digest == a.Digest {
-			matchDigest = true
-		}
-
-		if matchDigest && matchAlgorithm {
-			rv = append(rv, c.convArtifact(a))
+		for _, ak := range artKeys {
+			a, err := byKeykv[*artStruct](ctx, artCol, ak, c)
+			if err != nil {
+				return nil, err
+			}
+
+			matchAlgorithm := false
+			if algorithm == "" || algorithm == a.Algorithm {
+				matchAlgorithm = true
+			}
+
+			matchDigest := false
+			if digest == "" || digest == a.Digest {
+				matchDigest = true
+			}
+
+			if matchDigest && matchAlgorithm {
+				rv = append(rv, c.convArtifact(a))
+			}
 		}
 	}
 	return rv, nil

diff --git a/pkg/assembler/backends/keyvalue/builder.go b/pkg/assembler/backends/keyvalue/builder.go
@@ -124,16 +124,21 @@ func (c *demoClient) Builders(ctx context.Context, builderSpec *model.BuilderSpe
 		return []*model.Builder{c.convBuilder(b)}, nil
 	}
 	var builders []*model.Builder
-	bKeys, err := c.kv.Keys(ctx, builderCol)
-	if err != nil {
-		return nil, err
-	}
-	for _, bk := range bKeys {
-		b, err := byKeykv[*builderStruct](ctx, builderCol, bk, c)
+	var done bool
+	scn := c.kv.Keys(builderCol)
+	for !done {
+		var bKeys []string
+		bKeys, done, err = scn.Scan(ctx)
 		if err != nil {
 			return nil, err
 		}
-		builders = append(builders, c.convBuilder(b))
+		for _, bk := range bKeys {
+			b, err := byKeykv[*builderStruct](ctx, builderCol, bk, c)
+			if err != nil {
+				return nil, err
+			}
+			builders = append(builders, c.convBuilder(b))
+		}
 	}
 	return builders, nil
 }

diff --git a/pkg/assembler/backends/keyvalue/certifyBad.go b/pkg/assembler/backends/keyvalue/certifyBad.go
@@ -244,18 +244,24 @@ func (c *demoClient) CertifyBad(ctx context.Context, filter *model.CertifyBadSpe
 			}
 		}
 	} else {
-		cgKeys, err := c.kv.Keys(ctx, cbCol)
-		if err != nil {
-			return nil, err
-		}
-		for _, cgk := range cgKeys {
-			link, err := byKeykv[*badLink](ctx, cbCol, cgk, c)
+		var done bool
+		scn := c.kv.Keys(cbCol)
+		for !done {
+			var cgKeys []string
+			var err error
+			cgKeys, done, err = scn.Scan(ctx)
 			if err != nil {
 				return nil, err
 			}
-			out, err = c.addCBIfMatch(ctx, out, filter, link)
-			if err != nil {
-				return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+			for _, cgk := range cgKeys {
+				link, err := byKeykv[*badLink](ctx, cbCol, cgk, c)
+				if err != nil {
+					return nil, err
+				}
+				out, err = c.addCBIfMatch(ctx, out, filter, link)
+				if err != nil {
+					return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+				}
 			}
 		}
 	}

diff --git a/pkg/assembler/backends/keyvalue/certifyGood.go b/pkg/assembler/backends/keyvalue/certifyGood.go
@@ -243,18 +243,24 @@ func (c *demoClient) CertifyGood(ctx context.Context, filter *model.CertifyGoodS
 			}
 		}
 	} else {
-		cgKeys, err := c.kv.Keys(ctx, cgCol)
-		if err != nil {
-			return nil, err
-		}
-		for _, cgk := range cgKeys {
-			link, err := byKeykv[*goodLink](ctx, cgCol, cgk, c)
+		var done bool
+		scn := c.kv.Keys(cgCol)
+		for !done {
+			var cgKeys []string
+			var err error
+			cgKeys, done, err = scn.Scan(ctx)
 			if err != nil {
 				return nil, err
 			}
-			out, err = c.addCGIfMatch(ctx, out, filter, link)
-			if err != nil {
-				return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+			for _, cgk := range cgKeys {
+				link, err := byKeykv[*goodLink](ctx, cgCol, cgk, c)
+				if err != nil {
+					return nil, err
+				}
+				out, err = c.addCGIfMatch(ctx, out, filter, link)
+				if err != nil {
+					return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+				}
 			}
 		}
 	}

diff --git a/pkg/assembler/backends/keyvalue/certifyLegal.go b/pkg/assembler/backends/keyvalue/certifyLegal.go
@@ -343,18 +343,24 @@ func (c *demoClient) CertifyLegal(ctx context.Context, filter *model.CertifyLega
 			}
 		}
 	} else {
-		clKeys, err := c.kv.Keys(ctx, clCol)
-		if err != nil {
-			return nil, err
-		}
-		for _, clk := range clKeys {
-			link, err := byKeykv[*certifyLegalStruct](ctx, clCol, clk, c)
+		var done bool
+		scn := c.kv.Keys(clCol)
+		for !done {
+			var clKeys []string
+			var err error
+			clKeys, done, err = scn.Scan(ctx)
 			if err != nil {
 				return nil, err
 			}
-			out, err = c.addLegalIfMatch(ctx, out, filter, link)
-			if err != nil {
-				return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+			for _, clk := range clKeys {
+				link, err := byKeykv[*certifyLegalStruct](ctx, clCol, clk, c)
+				if err != nil {
+					return nil, err
+				}
+				out, err = c.addLegalIfMatch(ctx, out, filter, link)
+				if err != nil {
+					return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+				}
 			}
 		}
 	}

diff --git a/pkg/assembler/backends/keyvalue/certifyScorecard.go b/pkg/assembler/backends/keyvalue/certifyScorecard.go
@@ -183,18 +183,24 @@ func (c *demoClient) Scorecards(ctx context.Context, filter *model.CertifyScorec
 			}
 		}
 	} else {
-		cscKeys, err := c.kv.Keys(ctx, cscCol)
-		if err != nil {
-			return nil, err
-		}
-		for _, csck := range cscKeys {
-			link, err := byKeykv[*scorecardLink](ctx, cscCol, csck, c)
+		var done bool
+		scn := c.kv.Keys(cscCol)
+		for !done {
+			var cscKeys []string
+			var err error
+			cscKeys, done, err = scn.Scan(ctx)
 			if err != nil {
 				return nil, err
 			}
-			out, err = c.addSCIfMatch(ctx, out, filter, link)
-			if err != nil {
-				return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+			for _, csck := range cscKeys {
+				link, err := byKeykv[*scorecardLink](ctx, cscCol, csck, c)
+				if err != nil {
+					return nil, err
+				}
+				out, err = c.addSCIfMatch(ctx, out, filter, link)
+				if err != nil {
+					return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+				}
 			}
 		}
 	}

diff --git a/pkg/assembler/backends/keyvalue/certifyVEXStatement.go b/pkg/assembler/backends/keyvalue/certifyVEXStatement.go
@@ -254,18 +254,24 @@ func (c *demoClient) CertifyVEXStatement(ctx context.Context, filter *model.Cert
 			}
 		}
 	} else {
-		keys, err := c.kv.Keys(ctx, cVEXCol)
-		if err != nil {
-			return nil, err
-		}
-		for _, key := range keys {
-			link, err := byKeykv[*vexLink](ctx, cVEXCol, key, c)
+		var done bool
+		scn := c.kv.Keys(cVEXCol)
+		for !done {
+			var keys []string
+			var err error
+			keys, done, err = scn.Scan(ctx)
 			if err != nil {
 				return nil, err
 			}
-			out, err = c.addVexIfMatch(ctx, out, filter, link)
-			if err != nil {
-				return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+			for _, key := range keys {
+				link, err := byKeykv[*vexLink](ctx, cVEXCol, key, c)
+				if err != nil {
+					return nil, err
+				}
+				out, err = c.addVexIfMatch(ctx, out, filter, link)
+				if err != nil {
+					return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+				}
 			}
 		}
 	}

diff --git a/pkg/assembler/backends/keyvalue/certifyVuln.go b/pkg/assembler/backends/keyvalue/certifyVuln.go
@@ -227,18 +227,24 @@ func (c *demoClient) CertifyVuln(ctx context.Context, filter *model.CertifyVulnS
 			}
 		}
 	} else {
-		keys, err := c.kv.Keys(ctx, cVulnCol)
-		if err != nil {
-			return nil, err
-		}
-		for _, key := range keys {
-			link, err := byKeykv[*certifyVulnerabilityLink](ctx, cVulnCol, key, c)
+		var done bool
+		scn := c.kv.Keys(cVulnCol)
+		for !done {
+			var keys []string
+			var err error
+			keys, done, err = scn.Scan(ctx)
 			if err != nil {
 				return nil, err
 			}
-			out, err = c.addCVIfMatch(ctx, out, filter, link)
-			if err != nil {
-				return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+			for _, key := range keys {
+				link, err := byKeykv[*certifyVulnerabilityLink](ctx, cVulnCol, key, c)
+				if err != nil {
+					return nil, err
+				}
+				out, err = c.addCVIfMatch(ctx, out, filter, link)
+				if err != nil {
+					return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+				}
 			}
 		}
 	}

diff --git a/pkg/assembler/backends/keyvalue/hasMetadata.go b/pkg/assembler/backends/keyvalue/hasMetadata.go
@@ -253,18 +253,24 @@ func (c *demoClient) HasMetadata(ctx context.Context, filter *model.HasMetadataS
 			}
 		}
 	} else {
-		hmk, err := c.kv.Keys(ctx, hasMDCol)
-		if err != nil {
-			return nil, err
-		}
-		for _, hk := range hmk {
-			link, err := byKeykv[*hasMetadataLink](ctx, hasMDCol, hk, c)
+		var done bool
+		scn := c.kv.Keys(hasMDCol)
+		for !done {
+			var hmk []string
+			var err error
+			hmk, done, err = scn.Scan(ctx)
 			if err != nil {
 				return nil, err
 			}
-			out, err = c.addHMIfMatch(ctx, out, filter, link)
-			if err != nil {
-				return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+			for _, hk := range hmk {
+				link, err := byKeykv[*hasMetadataLink](ctx, hasMDCol, hk, c)
+				if err != nil {
+					return nil, err
+				}
+				out, err = c.addHMIfMatch(ctx, out, filter, link)
+				if err != nil {
+					return nil, gqlerror.Errorf("%v :: %v", funcName, err)
+				}
 			}
 		}
 	}