Enhance GetDocuments API by adding bulk retrieval

api/yorkie/v1/admin.proto

@@ -117,6 +117,7 @@ message GetDocumentResponse {
 message GetDocumentsRequest {
   string project_name = 1;
   repeated string document_keys = 2;
+  bool include_snapshot = 3;
 }
 
 message GetDocumentsResponse {

server/backend/database/database.go

@@ -164,6 +164,13 @@ type Database interface {
 		docKey key.Key,
 	) (*DocInfo, error)
 
+	// FindDocInfosByKeys finds the documents of the given keys.
+	FindDocInfosByKeys(
+		ctx context.Context,
+		projectID types.ID,
+		docKeys []key.Key,
+	) ([]*DocInfo, error)
+
 	// FindDocInfoByKeyAndOwner finds the document of the given key. If the
 	// createDocIfNotExist condition is true, create the document if it does not
 	// exist.

server/backend/database/memory/database.go

@@ -756,6 +756,31 @@ func (d *DB) FindDocInfoByKey(
 	return info.DeepCopy(), nil
 }
 
+// FindDocInfosByKeys finds the documents of the given keys.
+func (d *DB) FindDocInfosByKeys(
+	_ context.Context,
+	projectID types.ID,
+	keys []key.Key,
+) ([]*database.DocInfo, error) {
+	txn := d.db.Txn(false)
+	defer txn.Abort()
+
+	var infos []*database.DocInfo
+	for _, k := range keys {
+		info, err := d.findDocInfoByKey(txn, projectID, k)
+		if err != nil {
+			return nil, fmt.Errorf("find doc info by key: %w", err)
+		}
+		if info == nil {
+			continue
+		}
+
+		infos = append(infos, info.DeepCopy())
+	}
+
+	return infos, nil
+}
+
 // FindDocInfoByRefKey finds a docInfo of the given refKey.
 func (d *DB) FindDocInfoByRefKey(
 	_ context.Context,

server/backend/database/memory/database_test.go

@@ -48,6 +48,10 @@ func TestDB(t *testing.T) {
 		testcases.RunFindDocInfoTest(t, db, projectID)
 	})
 
+	t.Run("RunFindDocInfosByKeys test", func(t *testing.T) {
+		testcases.RunFindDocInfosByKeysTest(t, db, projectID)
+	})
+
 	t.Run("RunFindDocInfosByQuery test", func(t *testing.T) {
 		testcases.RunFindDocInfosByQueryTest(t, db, projectOneID)
 	})

server/backend/database/mongo/client.go

@@ -763,6 +763,35 @@ func (c *Client) FindDocInfoByKey(
 	return &docInfo, nil
 }
 
+// FindDocInfosByKeys finds the documents of the given keys.
+func (c *Client) FindDocInfosByKeys(
+	ctx context.Context,
+	projectID types.ID,
+	docKeys []key.Key,
+) ([]*database.DocInfo, error) {
+	filter := bson.M{
+		"project_id": projectID,
+		"key": bson.M{
+			"$in": docKeys,
+		},
+		"removed_at": bson.M{
+			"$exists": false,
+		},
+	}
+
+	cursor, err := c.collection(ColDocuments).Find(ctx, filter)
+	if err != nil {
+		return nil, fmt.Errorf("find documents: %w", err)
+	}
+
+	var docInfos []*database.DocInfo
+	if err := cursor.All(ctx, &docInfos); err != nil {
+		return nil, fmt.Errorf("fetch documents: %w", err)
+	}
+
+	return docInfos, nil
+}
+
 // FindDocInfoByRefKey finds a docInfo of the given refKey.
 func (c *Client) FindDocInfoByRefKey(
 	ctx context.Context,

server/backend/database/mongo/client_test.go

@@ -63,6 +63,10 @@ func TestClient(t *testing.T) {
 		testcases.RunFindDocInfoTest(t, cli, dummyProjectID)
 	})
 
+	t.Run("RunFindDocInfosByKeys test", func(t *testing.T) {
+		testcases.RunFindDocInfosByKeysTest(t, cli, dummyProjectID)
+	})
+
 	t.Run("RunFindDocInfosByQuery test", func(t *testing.T) {
 		t.Skip("TODO(hackerwins): the order of docInfos is different with memDB")
 		testcases.RunFindDocInfosByQueryTest(t, cli, projectOneID)

server/backend/database/testcases/testcases.go

@@ -93,6 +93,71 @@ func RunFindDocInfoTest(
 	})
 }
 
+// RunFindDocInfosByKeysTest runs the FindDocInfosByKeys test for the given db.
+func RunFindDocInfosByKeysTest(
+	t *testing.T,
+	db database.Database,
+	projectID types.ID,
+) {
+	t.Run("find docInfos by keys test", func(t *testing.T) {
+		ctx := context.Background()
+		clientInfo, err := db.ActivateClient(ctx, projectID, t.Name())
+		assert.NoError(t, err)
+
+		// 01. Create documents
+		docKeys := []key.Key{
+			"test", "test$3", "test123", "test$0",
+			"search$test", "abcde", "test abc",
+		}
+		for _, docKey := range docKeys {
+			_, err := db.FindDocInfoByKeyAndOwner(ctx, clientInfo.RefKey(), docKey, true)
+			assert.NoError(t, err)
+		}
+
+		// 02. Find documents
+		infos, err := db.FindDocInfosByKeys(ctx, projectID, docKeys)
+		assert.NoError(t, err)
+
+		actualKeys := make([]key.Key, len(infos))
+		for i, info := range infos {
+			actualKeys[i] = info.Key
+		}
+
+		assert.ElementsMatch(t, docKeys, actualKeys)
+		assert.Len(t, infos, len(docKeys))
+	})
+
+	t.Run("find docInfos by keys where some keys are not found test", func(t *testing.T) {
+		ctx := context.Background()
+		clientInfo, err := db.ActivateClient(ctx, projectID, t.Name())
+		assert.NoError(t, err)
+
+		// 01. Create documents
+		docKeys := []key.Key{
+			"exist-key1", "exist-key2", "exist-key3",
+		}
+		for _, docKey := range docKeys {
+			_, err := db.FindDocInfoByKeyAndOwner(ctx, clientInfo.RefKey(), docKey, true)
+			assert.NoError(t, err)
+		}
+
+		// 02. append a key that does not exist
+		docKeysWithNonExistKey := append(docKeys, "non-exist-key")
+
+		// 03. Find documents
+		infos, err := db.FindDocInfosByKeys(ctx, projectID, docKeysWithNonExistKey)
+		assert.NoError(t, err)
+
+		actualKeys := make([]key.Key, len(infos))
+		for i, info := range infos {
+			actualKeys[i] = info.Key
+		}
+
+		assert.ElementsMatch(t, docKeys, actualKeys)
+		assert.Len(t, infos, len(docKeys))
+	})
+}
+
 // RunFindProjectInfoBySecretKeyTest runs the FindProjectInfoBySecretKey test for the given db.
 func RunFindProjectInfoBySecretKeyTest(
 	t *testing.T,

server/documents/documents.go

@@ -97,16 +97,7 @@ func GetDocumentSummary(
 	project *types.Project,
 	k key.Key,
 ) (*types.DocumentSummary, error) {
-	// TODO(hackerwins): Split FindDocInfoByKeyAndOwner into upsert and find.
-	docInfo, err := be.DB.FindDocInfoByKeyAndOwner(
-		ctx,
-		types.ClientRefKey{
-			ProjectID: project.ID,
-			ClientID:  types.IDFromActorID(time.InitialActorID),
-		},
-		k,
-		false,
-	)
+	docInfo, err := be.DB.FindDocInfoByKey(ctx, project.ID, k)
 	if err != nil {
 		return nil, err
 	}
@@ -129,16 +120,36 @@ func GetDocumentSummary(
 // GetDocumentSummaries returns a list of document summaries.
 func GetDocumentSummaries(
 	ctx context.Context,
-	b *backend.Backend,
+	be *backend.Backend,
 	project *types.Project,
 	keys []key.Key,
+	includeSnapshot bool,
 ) ([]*types.DocumentSummary, error) {
-	// TODO(hackerwins): Resolve the N+1 problem.
+	docInfos, err := be.DB.FindDocInfosByKeys(ctx, project.ID, keys)
+	if err != nil {
+		return nil, err
+	}
+
 	var summaries []*types.DocumentSummary
-	for _, k := range keys {
-		summary, err := GetDocumentSummary(ctx, b, project, k)
-		if err != nil {
-			return nil, err
+	for _, docInfo := range docInfos {
+		snapshot := ""
+		if includeSnapshot {
+			// TODO(hackerwins, kokodak): Resolve the N+1 problem.
+			doc, err := packs.BuildDocumentForServerSeq(ctx, be, docInfo, docInfo.ServerSeq)
+			if err != nil {
+				return nil, err
+			}
+
+			snapshot = doc.Marshal()
+		}
+
+		summary := &types.DocumentSummary{
+			ID:         docInfo.ID,
+			Key:        docInfo.Key,
+			CreatedAt:  docInfo.CreatedAt,
+			AccessedAt: docInfo.AccessedAt,
+			UpdatedAt:  docInfo.UpdatedAt,
+			Snapshot:   snapshot,
 		}
 
 		summaries = append(summaries, summary)

server/rpc/admin_server.go

@@ -225,6 +225,7 @@ func (s *adminServer) GetDocuments(
 		s.backend,
 		project,
 		keys,
+		req.Msg.IncludeSnapshot,
 	)
 	if err != nil {
 		return nil, err

test/sharding/mongo_client_test.go

@@ -72,6 +72,10 @@ func TestClientWithShardedDB(t *testing.T) {
 		testcases.RunFindDocInfoTest(t, cli, dummyProjectID)
 	})
 
+	t.Run("RunFindDocInfosByKeys test", func(t *testing.T) {
+		testcases.RunFindDocInfosByKeysTest(t, cli, dummyProjectID)
+	})
+
 	t.Run("RunFindDocInfosByQuery test", func(t *testing.T) {
 		t.Skip("TODO(hackerwins): the order of docInfos is different with memDB")
 		testcases.RunFindDocInfosByQueryTest(t, cli, projectOneID)