feat: support keys_only in QueryRecords

api/open_saves.proto

@@ -346,7 +346,7 @@ message GetRecordRequest {
 // Multiple conditions are AND'ed together.
 message QueryRecordsRequest {
   // store_key is the primary key of the store.
-  // Optional and the method queries all stores when omitted.
+  // Optional. The method queries all stores when omitted.
   string store_key = 1;
 
   // List of filters
@@ -364,6 +364,10 @@ message QueryRecordsRequest {
 
   // the limit of the number of records to return.
   int32 limit = 6;
+
+  // If keys_only is set to true, the server will only return records.key 
+  // and store_keys in the QueryRecordsResponse message.
+  bool keys_only = 7;
 }
 
 // FilterOperator has a list of comperators.

docs/reference.md

@@ -593,12 +593,13 @@ Multiple conditions are AND'ed together.
 
 | Field | Type | Label | Description |
 | ----- | ---- | ----- | ----------- |
-| store_key | [string](#string) |  | store_key is the primary key of the store. Optional and the method queries all stores when omitted. |
+| store_key | [string](#string) |  | store_key is the primary key of the store. Optional. The method queries all stores when omitted. |
 | filters | [QueryFilter](#opensaves-QueryFilter) | repeated | List of filters |
 | tags | [string](#string) | repeated | List of tags |
 | owner_id | [string](#string) |  | owner_id is the owner of records, represented as an external user ID. |
 | sort_orders | [SortOrder](#opensaves-SortOrder) | repeated | List of sort orders to return records. These SortOrders are applied in sequence. |
 | limit | [int32](#int32) |  | the limit of the number of records to return. |
+| keys_only | [bool](#bool) |  | If keys_only is set to true, the server will only return records.key and store_keys in the QueryRecordsResponse message. |
 
 
 

internal/app/server/open_saves.go

@@ -214,15 +214,17 @@ func (s *openSavesServer) UpdateRecord(ctx context.Context, req *pb.UpdateRecord
 }
 
 func (s *openSavesServer) QueryRecords(ctx context.Context, req *pb.QueryRecordsRequest) (*pb.QueryRecordsResponse, error) {
-	records, storeKeys, err := s.metaDB.QueryRecords(ctx, req)
+	records, err := s.metaDB.QueryRecords(ctx, req)
 	if err != nil {
 		log.Warnf("QueryRecords failed for store(%s), filters(%+v): %v",
 			req.StoreKey, req.Filters, err)
 		return nil, err
 	}
 	var rr []*pb.Record
+	var storeKeys []string
 	for _, r := range records {
 		rr = append(rr, r.ToProto())
+		storeKeys = append(storeKeys, r.StoreKey)
 	}
 	return &pb.QueryRecordsResponse{
 		Records:   rr,

internal/app/server/open_saves_test.go

@@ -836,11 +836,19 @@ func TestOpenSaves_QueryRecords_EqualityFilter(t *testing.T) {
 	resp, err := client.QueryRecords(ctx, queryReq)
 	require.NoError(t, err)
 	// Only one record matches the query.
-	require.Equal(t, 1, len(resp.Records))
-	require.Equal(t, 1, len(resp.StoreKeys))
+	require.Len(t, resp.Records, 1)
+	require.Len(t, resp.StoreKeys, 1)
 
 	assert.Equal(t, storeKey, resp.StoreKeys[0])
 	assert.Equal(t, resp.Records[0].Properties["prop1"].Value, stringVal1)
+
+	// Test KeysOnly
+	queryReq.KeysOnly = true
+	resp, err = client.QueryRecords(ctx, queryReq)
+	require.NoError(t, err)
+	require.Len(t, resp.Records, 1)
+	assertEqualRecord(t, &pb.Record{Key: recordKey1}, resp.Records[0])
+	assert.Equal(t, []string{storeKey}, resp.GetStoreKeys())
 }
 
 // NOTE: this test requires composite indexes to be created. You can find the corresponding index
@@ -1117,7 +1125,6 @@ func TestOpenSaves_QueryRecords_Limit(t *testing.T) {
 	// Make sure the query only returns 1 record.
 	require.Equal(t, 1, len(resp.Records))
 	require.Equal(t, 1, len(resp.StoreKeys))
-
 }
 
 func TestOpenSaves_CreateChunkedBlobNonExistent(t *testing.T) {

internal/pkg/metadb/metadb.go

@@ -262,6 +262,7 @@ func (m *MetaDB) DeleteStore(ctx context.Context, key string) error {
 // Returns error if there is already a record with the same key.
 func (m *MetaDB) InsertRecord(ctx context.Context, storeKey string, record *record.Record) (*record.Record, error) {
 	record.Timestamps = timestamps.New()
+	record.StoreKey = storeKey
 	rkey := m.createRecordKey(storeKey, record.Key)
 	_, err := m.client.RunInTransaction(ctx, func(tx *ds.Transaction) error {
 		dskey := m.createStoreKey(storeKey)
@@ -711,9 +712,9 @@ func addPropertyFilter(q *ds.Query, f *pb.QueryFilter) (*ds.Query, error) {
 	return q.Filter(filter, record.ExtractValue(f.Value)), nil
 }
 
-// QueryRecords returns a list of records that match the given filters and their stores.
+// QueryRecords returns a list of records that match the given filters.
 // TODO(https://github.com/googleforgames/open-saves/issues/339): consider refactoring this to fewer arguments.
-func (m *MetaDB) QueryRecords(ctx context.Context, req *pb.QueryRecordsRequest) ([]*record.Record, []string, error) {
+func (m *MetaDB) QueryRecords(ctx context.Context, req *pb.QueryRecordsRequest) ([]*record.Record, error) {
 	query := m.newQuery(recordKind)
 	if req.GetStoreKey() != "" {
 		dsKey := m.createStoreKey(req.GetStoreKey())
@@ -725,7 +726,7 @@ func (m *MetaDB) QueryRecords(ctx context.Context, req *pb.QueryRecordsRequest)
 	for _, f := range req.GetFilters() {
 		q, err := addPropertyFilter(query, f)
 		if err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 		query = q
 	}
@@ -741,11 +742,11 @@ func (m *MetaDB) QueryRecords(ctx context.Context, req *pb.QueryRecordsRequest)
 			property = "Timestamps.UpdatedAt"
 		case pb.SortOrder_USER_PROPERTY:
 			if s.UserPropertyName == "" {
-				return nil, nil, status.Error(codes.InvalidArgument, "got empty user sort property")
+				return nil, status.Error(codes.InvalidArgument, "got empty user sort property")
 			}
 			property = fmt.Sprintf("%s.%s", propertiesField, s.UserPropertyName)
 		default:
-			return nil, nil, status.Errorf(codes.InvalidArgument, "got invalid SortOrder property value: %v", s.Property)
+			return nil, status.Errorf(codes.InvalidArgument, "got invalid SortOrder property value: %v", s.Property)
 		}
 
 		switch s.Direction {
@@ -754,30 +755,36 @@ func (m *MetaDB) QueryRecords(ctx context.Context, req *pb.QueryRecordsRequest)
 		case pb.SortOrder_DESC:
 			query = query.Order("-" + strconv.Quote(property))
 		default:
-			return nil, nil, status.Errorf(codes.InvalidArgument, "got invalid SortOrder direction value: %v", s.Direction)
+			return nil, status.Errorf(codes.InvalidArgument, "got invalid SortOrder direction value: %v", s.Direction)
 		}
 	}
 	if limit := req.GetLimit(); limit > 0 {
 		query = query.Limit(int(limit))
 	}
+	if req.GetKeysOnly() {
+		query = query.KeysOnly()
+	}
 	iter := m.client.Run(ctx, query)
 
 	var match []*record.Record
-	var keys []string
 	for {
 		var r record.Record
 		key, err := iter.Next(&r)
 		if err == iterator.Done {
 			break
 		}
 		if err != nil {
-			return nil, nil, status.Errorf(codes.Internal,
+			return nil, status.Errorf(codes.Internal,
 				"metadb QueryRecords: %v", err)
 		}
+		if req.GetKeysOnly() {
+			if err := r.LoadKey(key); err != nil {
+				return nil, status.Errorf(codes.Internal, "metadb QueryRecords LoadKey: %v", err)
+			}
+		}
 		match = append(match, &r)
-		keys = append(keys, key.Parent.Name)
 	}
-	return match, keys, nil
+	return match, nil
 }
 
 func (m *MetaDB) findChunkRefsByNumber(ctx context.Context, tx *ds.Transaction, storeKey, recordKey string, blobKey uuid.UUID, number int32) ([]*chunkref.ChunkRef, error) {

internal/pkg/metadb/metadb_test.go

@@ -20,6 +20,8 @@ import (
 	"time"
 
 	"cloud.google.com/go/datastore"
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/google/uuid"
 	pb "github.com/googleforgames/open-saves/api"
 	m "github.com/googleforgames/open-saves/internal/pkg/metadb"
@@ -31,6 +33,7 @@ import (
 	"github.com/googleforgames/open-saves/internal/pkg/metadb/store"
 	"github.com/googleforgames/open-saves/internal/pkg/metadb/timestamps"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/api/iterator"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -126,18 +129,23 @@ func setupTestStoreRecord(ctx context.Context, t *testing.T, metaDB *m.MetaDB, s
 	metadbtest.AssertEqualStore(t, store, newStore, "CreateStore should return the created store.")
 	var newRecord *record.Record
 	if r != nil {
-		newRecord, err = metaDB.InsertRecord(ctx, newStore.Key, r)
-		if err != nil {
-			t.Fatalf("Could not create a new record: %v", err)
-		}
-		t.Cleanup(func() {
-			metaDB.DeleteRecord(ctx, newStore.Key, newRecord.Key)
-		})
-		metadbtest.AssertEqualRecord(t, r, newRecord, "GetRecord should return the exact same record.")
+		newRecord = setupTestRecord(ctx, t, metaDB, newStore.Key, r)
 	}
 	return newStore, newRecord
 }
 
+func setupTestRecord(ctx context.Context, t *testing.T, metaDB *m.MetaDB, storeKey string, r *record.Record) *record.Record {
+	t.Helper()
+
+	newRecord, err := metaDB.InsertRecord(ctx, storeKey, r)
+	require.NoError(t, err, "InsertRecord")
+	t.Cleanup(func() {
+		metaDB.DeleteRecord(ctx, storeKey, newRecord.Key)
+	})
+	metadbtest.AssertEqualRecord(t, r, newRecord, "GetRecord should return the exact same record.")
+	return newRecord
+}
+
 func setupTestBlobRef(ctx context.Context, t *testing.T, metaDB *m.MetaDB, blob *blobref.BlobRef) *blobref.BlobRef {
 	t.Helper()
 	newBlob, err := metaDB.InsertBlobRef(ctx, blob)
@@ -1210,3 +1218,104 @@ func TestMetaDB_ListChunkRefsByStatus(t *testing.T) {
 		assert.Nil(t, b)
 	}
 }
+
+func TestMetaDB_QueryRecords(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
+	metaDB := newMetaDB(ctx, t)
+
+	stores := make([]*store.Store, 2)
+	for i := range stores {
+		stores[i], _ = setupTestStoreRecord(ctx, t, metaDB, &store.Store{Key: newStoreKey()}, nil)
+	}
+	records := []*record.Record{
+		{
+			Key:        newRecordKey(),
+			OwnerID:    "abc",
+			Tags:       []string{"tag1", "tag2"},
+			Properties: make(record.PropertyMap),
+		},
+		{
+			Key:        newRecordKey(),
+			OwnerID:    "cba",
+			Tags:       []string{"gat1", "gat2"},
+			Properties: make(record.PropertyMap),
+		},
+		{
+			Key:        newRecordKey(),
+			OwnerID:    "xyz",
+			Tags:       []string{"tag1", "tag2"},
+			Properties: make(record.PropertyMap),
+		},
+	}
+	records[0] = setupTestRecord(ctx, t, metaDB, stores[0].Key, records[0])
+	records[1] = setupTestRecord(ctx, t, metaDB, stores[0].Key, records[1])
+	records[2] = setupTestRecord(ctx, t, metaDB, stores[1].Key, records[2])
+
+	testCases := []struct {
+		name        string
+		req         *pb.QueryRecordsRequest
+		wantRecords []*record.Record
+		wantCode    codes.Code
+	}{
+		{
+			"OwnerId",
+			&pb.QueryRecordsRequest{
+				StoreKey: stores[0].Key,
+				OwnerId:  "abc",
+			},
+			[]*record.Record{records[0]}, codes.OK,
+		},
+		{
+			"Tags No Result",
+			&pb.QueryRecordsRequest{
+				StoreKey: stores[1].Key,
+				Tags:     []string{"tag1", "gat2"},
+			},
+			nil, codes.OK,
+		},
+		{
+			"Tags Multiple Records",
+			&pb.QueryRecordsRequest{
+				Tags: []string{"tag1"},
+			},
+			[]*record.Record{records[0], records[2]}, codes.OK,
+		},
+		{
+			"Limit",
+			&pb.QueryRecordsRequest{
+				StoreKey:   stores[0].Key,
+				SortOrders: []*pb.SortOrder{{Property: pb.SortOrder_CREATED_AT}},
+				Limit:      1,
+			},
+			[]*record.Record{records[0]}, codes.OK,
+		},
+		{
+			"Keys Only",
+			&pb.QueryRecordsRequest{
+				StoreKey: stores[0].Key,
+				OwnerId:  "abc",
+				Tags:     []string{"tag1"},
+				KeysOnly: true,
+			},
+			[]*record.Record{{Key: records[0].Key, StoreKey: records[0].StoreKey}}, codes.OK,
+		},
+	}
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			rr, err := metaDB.QueryRecords(ctx, tc.req)
+			assert.Empty(t, cmp.Diff(rr, tc.wantRecords,
+				cmpopts.SortSlices(func(x, y *record.Record) bool {
+					return x.Key < y.Key
+				}),
+				cmpopts.EquateEmpty()))
+			assert.Equal(t, tc.wantCode, status.Code(err))
+			if tc.wantCode == codes.OK {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}

internal/pkg/metadb/record/record.go

@@ -22,7 +22,6 @@ import (
 	"github.com/googleforgames/open-saves/internal/pkg/metadb/checksums"
 	"github.com/googleforgames/open-saves/internal/pkg/metadb/timestamps"
 	"github.com/vmihailenco/msgpack/v5"
-	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 // Assert PropertyMap implements PropertyLoadSave.
@@ -116,8 +115,8 @@ func (r *Record) ToProto() *pb.Record {
 		Tags:         r.Tags,
 		Properties:   r.Properties.ToProto(),
 		OpaqueString: r.OpaqueString,
-		CreatedAt:    timestamppb.New(r.Timestamps.CreatedAt),
-		UpdatedAt:    timestamppb.New(r.Timestamps.UpdatedAt),
+		CreatedAt:    timestamps.TimeToProto(r.Timestamps.CreatedAt),
+		UpdatedAt:    timestamps.TimeToProto(r.Timestamps.UpdatedAt),
 		Signature:    r.Timestamps.Signature[:],
 	}
 	return ret

internal/pkg/metadb/store/store.go

@@ -18,7 +18,6 @@ import (
 	"cloud.google.com/go/datastore"
 	pb "github.com/googleforgames/open-saves/api"
 	"github.com/googleforgames/open-saves/internal/pkg/metadb/timestamps"
-	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 // Store represents a Open Saves store in the metadata database.
@@ -45,8 +44,8 @@ func (s *Store) ToProto() *pb.Store {
 		Name:      s.Name,
 		Tags:      s.Tags,
 		OwnerId:   s.OwnerID,
-		CreatedAt: timestamppb.New(s.Timestamps.CreatedAt),
-		UpdatedAt: timestamppb.New(s.Timestamps.UpdatedAt),
+		CreatedAt: timestamps.TimeToProto(s.Timestamps.CreatedAt),
+		UpdatedAt: timestamps.TimeToProto(s.Timestamps.UpdatedAt),
 	}
 }
 

internal/pkg/metadb/timestamps/timestamps.go

@@ -19,6 +19,7 @@ import (
 
 	"cloud.google.com/go/datastore"
 	"github.com/google/uuid"
+	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 const (
@@ -59,6 +60,15 @@ func New() Timestamps {
 	}
 }
 
+// TimeToProto converts time.Time to timestamppb.Timestamp.
+// If t is zero, it returns nil.
+func TimeToProto(t time.Time) *timestamppb.Timestamp {
+	if t.IsZero() {
+		return nil
+	}
+	return timestamppb.New(t)
+}
+
 // Update updates the UpdatedAt and Signature fields with time.Now() and uuid.New().
 func (t *Timestamps) Update() {
 	t.UpdatedAt = time.Now().UTC().Truncate(Precision)