From 26b5835ad1f8de16733df560bcc8a7c65e5f3c46 Mon Sep 17 00:00:00 2001
From: Jye Cusch <jye.cusch@nitric.io>
Date: Wed, 28 Feb 2024 22:00:37 +1100
Subject: [PATCH] implement kv keys method for gcp

---
 cloud/gcp/runtime/keyvalue/firestore.go | 55 ++++++++++++++++++++++++-
 1 file changed, 54 insertions(+), 1 deletion(-)

diff --git a/cloud/gcp/runtime/keyvalue/firestore.go b/cloud/gcp/runtime/keyvalue/firestore.go
index cb360d142..f18ca2fea 100644
--- a/cloud/gcp/runtime/keyvalue/firestore.go
+++ b/cloud/gcp/runtime/keyvalue/firestore.go
@@ -17,11 +17,13 @@ package keyvalue
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/nitrictech/nitric/core/pkg/decorators/keyvalue"
 	grpc_errors "github.com/nitrictech/nitric/core/pkg/grpc/errors"
 	v1 "github.com/nitrictech/nitric/core/pkg/proto/keyvalue/v1"
 
+	"google.golang.org/api/iterator"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/protobuf/types/known/structpb"
 
@@ -164,6 +166,53 @@ func (s *FirestoreDocService) Delete(ctx context.Context, req *v1.KeyValueDelete
 	return &v1.KeyValueDeleteResponse{}, nil
 }
 
+func (s *FirestoreDocService) Keys(req *v1.KeyValueKeysRequest, stream v1.KeyValue_KeysServer) error {
+	newErr := grpc_errors.ErrorsWithScope("FirestoreDocService.Keys")
+	storeName := req.GetStore().GetName()
+
+	if storeName == "" {
+		return newErr(
+			codes.InvalidArgument,
+			"store name is required",
+			nil,
+		)
+	}
+
+	iter := s.getCollectionRef(storeName).DocumentRefs(stream.Context())
+
+	for {
+		doc, err := iter.Next()
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			return newErr(
+				codes.Internal,
+				"error iterating over firestore collection",
+				err,
+			)
+		}
+
+		// Where <= doesn't appear to work where querying based on document ID.
+		// so instead we filter the results here
+		if !strings.HasPrefix(doc.ID, req.Prefix) {
+			continue
+		}
+
+		if err := stream.Send(&v1.KeyValueKeysResponse{
+			Key: doc.ID,
+		}); err != nil {
+			return newErr(
+				codes.Internal,
+				"failed to send response",
+				err,
+			)
+		}
+	}
+
+	return nil
+}
+
 func New() (v1.KeyValueServer, error) {
 	ctx := context.Background()
 
@@ -189,5 +238,9 @@ func NewWithClient(client *firestore.Client) (v1.KeyValueServer, error) {
 }
 
 func (s *FirestoreDocService) getDocRef(ref *v1.ValueRef) *firestore.DocumentRef {
-	return s.client.Collection(ref.Store).Doc(ref.Key)
+	return s.getCollectionRef(ref.Store).Doc(ref.Key)
+}
+
+func (s *FirestoreDocService) getCollectionRef(store string) *firestore.CollectionRef {
+	return s.client.Collection(store)
 }