Implement MsgCount for Tigris (#928)

Closes #771.
FerretDB · Jul 26, 2022 · 4446f44 · 4446f44
1 parent 3d53f60
commit 4446f44
Show file tree

Hide file tree

Showing 8 changed files with 189 additions and 14 deletions.
diff --git a/Taskfile.yml b/Taskfile.yml
@@ -130,19 +130,25 @@ tasks:
     desc: "Run integration tests for PostgreSQL handler"
     dir: integration
     cmds:
-      - go test -count=1 {{if ne OS "windows"}}-race{{end}} -shuffle=on -coverprofile=integration-pg.txt -coverpkg=../... -handler=pg
+      - >
+        go test -count=1 {{if ne OS "windows"}}-race{{end}} -shuffle=on -coverpkg=../...
+        -coverprofile=integration-pg.txt . -handler=pg
 
   test-integration-tigris:
     desc: "Run integration tests for Tigris handler"
-    dir: integration/tigris
+    dir: integration
     cmds:
-      - go test -count=1 {{if ne OS "windows"}}-race{{end}} -tags=ferretdb_tigris -shuffle=on -coverprofile=../integration-tigris.txt -coverpkg=../... -handler=tigris
+      - >
+        go test -count=1 {{if ne OS "windows"}}-race{{end}} -shuffle=on -coverpkg=../...
+        -coverprofile=integration-tigris.txt -tags=ferretdb_tigris ./tigris/... -handler=tigris
 
   test-integration-mongodb:
     desc: "Run integration tests for MongoDB"
     dir: integration
     cmds:
-      - go test -count=1 {{if ne OS "windows"}}-race{{end}} -shuffle=on -coverprofile=integration-mongodb.txt -coverpkg=../... -target-port=37017 -compat-port=0
+      - >
+        go test -count=1 {{if ne OS "windows"}}-race{{end}} -shuffle=on -coverpkg=../...
+        -coverprofile=integration-mongodb.txt . -target-port=37017 -compat-port=0
 
   bench-short:
     desc: "Benchmark for about 20 seconds (with default BENCHTIME)"

diff --git a/go.mod b/go.mod
@@ -18,7 +18,6 @@ require (
 	golang.org/x/exp v0.0.0-20220713135740-79cabaa25d75
 	golang.org/x/net v0.0.0-20220708220712-1185a9018129
 	golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8
-	google.golang.org/grpc v1.48.0
 )
 
 require (
@@ -59,6 +58,7 @@ require (
 	golang.org/x/text v0.3.7 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/genproto v0.0.0-20220526192754-51939a95c655 // indirect
+	google.golang.org/grpc v1.48.0 // indirect
 	google.golang.org/protobuf v1.28.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

diff --git a/integration/tigris/smoke_test.go b/integration/tigris/smoke_test.go
@@ -15,8 +15,10 @@
 package tigris
 
 import (
+	"math"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.mongodb.org/mongo-driver/bson"
 
@@ -34,3 +36,54 @@ func TestSmoke(t *testing.T) {
 	require.NoError(t, err)
 	integration.AssertEqualDocuments(t, bson.D{{"_id", "fixed_double"}, {"double_value", 42.13}}, doc)
 }
+
+// TestSmokeMsgCount implements simple smoke tests for MsgCount.
+// TODO Implement proper testing: https://github.com/FerretDB/FerretDB/issues/931.
+func TestSmokeMsgCount(t *testing.T) {
+	t.Parallel()
+
+	// As Tigris require different fields for different types,
+	// for this smoke test we only use the fixed scalars.
+	ctx, collection := setup.Setup(t, shareddata.FixedScalars)
+
+	for name, tc := range map[string]struct {
+		command  any
+		response int32
+	}{
+		"CountAllFixedScalars": {
+			command:  bson.D{{"count", collection.Name()}},
+			response: 6,
+		},
+		"CountExactlyOneDocument": {
+			command: bson.D{
+				{"count", collection.Name()},
+				{"query", bson.D{{"double_value", math.MaxFloat64}}},
+			},
+			response: 1,
+		},
+		"CountNonExistingCollection": {
+			command: bson.D{
+				{"count", "doesnotexist"},
+				{"query", bson.D{{"v", true}}},
+			},
+			response: 0,
+		},
+	} {
+		name, tc := name, tc
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			var actual bson.D
+			err := collection.Database().RunCommand(ctx, tc.command).Decode(&actual)
+			require.NoError(t, err)
+
+			m := actual.Map()
+
+			assert.Equal(t, float64(1), m["ok"])
+
+			keys := integration.CollectKeys(t, actual)
+			assert.Contains(t, keys, "n")
+			assert.Equal(t, tc.response, m["n"])
+		})
+	}
+}
diff --git a/internal/handlers/tigris/fetch.go b/internal/handlers/tigris/fetch.go
@@ -18,6 +18,7 @@ import (
 	"context"
 
 	"github.com/tigrisdata/tigris-client-go/driver"
+	"go.uber.org/zap"
 
 	"github.com/FerretDB/FerretDB/internal/tjson"
 	"github.com/FerretDB/FerretDB/internal/types"
@@ -37,7 +38,19 @@ func (h *Handler) fetch(ctx context.Context, param fetchParam) ([]*types.Documen
 	db := h.driver.UseDatabase(param.db)
 
 	collection, err := db.DescribeCollection(ctx, param.collection)
-	if err != nil {
+	switch err := err.(type) {
+	case nil:
+		// do nothing
+	case *driver.Error:
+		if isNotFound(err) {
+			h.L.Debug(
+				"Collection doesn't exist, handling a case to deal with a non-existing collection (return empty list)",
+				zap.String("db", param.db), zap.String("collection", param.collection),
+			)
+			return []*types.Document{}, nil
+		}
+		return nil, lazyerrors.Error(err)
+	default:
 		return nil, lazyerrors.Error(err)
 	}
 

diff --git a/internal/handlers/tigris/msg_count.go b/internal/handlers/tigris/msg_count.go
@@ -16,13 +16,97 @@ package tigris
 
 import (
 	"context"
+	"fmt"
 
+	"github.com/FerretDB/FerretDB/internal/handlers/common"
+	"github.com/FerretDB/FerretDB/internal/types"
+	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
 	"github.com/FerretDB/FerretDB/internal/util/must"
 	"github.com/FerretDB/FerretDB/internal/wire"
 )
 
 // MsgCount implements HandlerInterface.
 func (h *Handler) MsgCount(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg, error) {
-	// TODO https://github.com/FerretDB/FerretDB/issues/771
-	return nil, notImplemented(must.NotFail(msg.Document()).Command())
+	document, err := msg.Document()
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	unimplementedFields := []string{
+		"skip",
+		"collation",
+	}
+	if err := common.Unimplemented(document, unimplementedFields...); err != nil {
+		return nil, err
+	}
+	ignoredFields := []string{
+		"hint",
+		"readConcern",
+		"comment",
+	}
+	common.Ignored(document, h.L, ignoredFields...)
+
+	var filter *types.Document
+	if filter, err = common.GetOptionalParam(document, "query", filter); err != nil {
+		return nil, err
+	}
+
+	var limit int64
+	if l, _ := document.Get("limit"); l != nil {
+		if limit, err = common.GetWholeNumberParam(l); err != nil {
+			return nil, err
+		}
+	}
+
+	var fp fetchParam
+	if fp.db, err = common.GetRequiredParam[string](document, "$db"); err != nil {
+		return nil, err
+	}
+	collectionParam, err := document.Get(document.Command())
+	if err != nil {
+		return nil, err
+	}
+	var ok bool
+	if fp.collection, ok = collectionParam.(string); !ok {
+		return nil, common.NewErrorMsg(
+			common.ErrBadValue,
+			fmt.Sprintf("collection name has invalid type %s", common.AliasFromType(collectionParam)),
+		)
+	}
+
+	fetchedDocs, err := h.fetch(ctx, fp)
+	if err != nil {
+		return nil, err
+	}
+
+	resDocs := make([]*types.Document, 0, 16)
+	for _, doc := range fetchedDocs {
+		matches, err := common.FilterDocument(doc, filter)
+		if err != nil {
+			return nil, err
+		}
+
+		if !matches {
+			continue
+		}
+
+		resDocs = append(resDocs, doc)
+	}
+
+	if resDocs, err = common.LimitDocuments(resDocs, limit); err != nil {
+		return nil, err
+	}
+
+	var reply wire.OpMsg
+	err = reply.SetSections(wire.OpMsgSection{
+		Documents: []*types.Document{must.NotFail(types.NewDocument(
+			"n", int32(len(resDocs)),
+			"ok", float64(1),
+		))},
+	})
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	return &reply, nil
 }
diff --git a/internal/handlers/tigris/msg_drop.go b/internal/handlers/tigris/msg_drop.go
@@ -17,8 +17,7 @@ package tigris
 import (
 	"context"
 
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
+	"github.com/tigrisdata/tigris-client-go/driver"
 
 	"github.com/FerretDB/FerretDB/internal/handlers/common"
 	"github.com/FerretDB/FerretDB/internal/types"
@@ -46,11 +45,17 @@ func (h *Handler) MsgDrop(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg, er
 		return nil, err
 	}
 
-	if err = h.driver.UseDatabase(db).DropCollection(ctx, collection); err != nil {
-		if status.Code(err) == codes.NotFound {
+	err = h.driver.UseDatabase(db).DropCollection(ctx, collection)
+	switch err := err.(type) {
+	case nil:
+		// do nothing
+	case *driver.Error:
+		if isNotFound(err) {
 			return nil, common.NewErrorMsg(common.ErrNamespaceNotFound, "ns not found")
 		}
 		return nil, lazyerrors.Error(err)
+	default:
+		return nil, lazyerrors.Error(err)
 	}
 
 	var reply wire.OpMsg

diff --git a/internal/handlers/tigris/msg_dropdatabase.go b/internal/handlers/tigris/msg_dropdatabase.go
@@ -17,7 +17,6 @@ package tigris
 import (
 	"context"
 
-	api "github.com/tigrisdata/tigris-client-go/api/server/v1"
 	"github.com/tigrisdata/tigris-client-go/driver"
 
 	"github.com/FerretDB/FerretDB/internal/handlers/common"
@@ -47,7 +46,7 @@ func (h *Handler) MsgDropDatabase(ctx context.Context, msg *wire.OpMsg) (*wire.O
 	case nil:
 		res.Set("dropped", db)
 	case *driver.Error:
-		if err.Code != api.Code_NOT_FOUND {
+		if !isNotFound(err) {
 			return nil, lazyerrors.Error(err)
 		}
 		// nothing otherwise

diff --git a/internal/handlers/tigris/tigris.go b/internal/handlers/tigris/tigris.go
@@ -19,6 +19,7 @@ import (
 	"context"
 	"time"
 
+	api "github.com/tigrisdata/tigris-client-go/api/server/v1"
 	"github.com/tigrisdata/tigris-client-go/config"
 	"github.com/tigrisdata/tigris-client-go/driver"
 	"go.uber.org/zap"
@@ -69,6 +70,20 @@ func (h *Handler) Close() {
 	h.driver.Close()
 }
 
+// isNotFound returns true if the error is a "not found" error.
+// This function is implemented to keep nolint in a single place.
+func isNotFound(err *driver.Error) bool {
+	if err == nil {
+		return false
+	}
+
+	//nolint:nosnakecase // Tigris named their const that way
+	if err.Code == api.Code_NOT_FOUND {
+		return true
+	}
+	return false
+}
+
 // check interfaces
 var (
 	_ handlers.Interface = (*Handler)(nil)