Implement simple query pushdown for Tigris

.github/workflows/integration.yml

@@ -31,10 +31,33 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - {name: "PostgreSQL", task: "pg"}
-          - {name: "Tigris", task: "tigris"}
-          - {name: "Tigris main", task: "tigris", tigris_dockerfile: "tigris_main"}
-          - {name: "MongoDB", task: "mongodb"}
+          - name: "MongoDB"
+            task: "mongodb"
+            buildtags: "ferretdb_tigris"
+
+          # with pushdown
+          - name: "PostgreSQL"
+            task: "pg"
+            buildtags: "ferretdb_tigris"
+          - name: "Tigris"
+            task: "tigris"
+            buildtags: "ferretdb_tigris"
+          - name: "Tigris main"
+            task: "tigris"
+            buildtags: "ferretdb_tigris"
+            tigris_dockerfile: "tigris_main"
+
+          # without pushdown
+          - name: "PostgreSQL (no pushdown)"
+            task: "pg"
+            buildtags: "ferretdb_tigris,ferretdb_nopushdown"
+          - name: "Tigris (no pushdown)"
+            task: "tigris"
+            buildtags: "ferretdb_tigris,ferretdb_nopushdown"
+          - name: "Tigris main (no pushdown)"
+            task: "tigris"
+            buildtags: "ferretdb_tigris,ferretdb_nopushdown"
+            tigris_dockerfile: "tigris_main"
 
     steps:
       - name: Checkout code
@@ -63,7 +86,7 @@ jobs:
         run: bin/task env-setup
 
       - name: Run ${{ matrix.task }} tests
-        run: bin/task test-integration-${{ matrix.task }}
+        run: bin/task test-integration-${{ matrix.task }} BUILDTAGS=${{ matrix.buildtags }}
 
       # The token is not required but should make uploads more stable.
       # If secrets are unavailable (for example, for a pull request from a fork), it fallbacks to the tokenless uploads.

Taskfile.yml

@@ -8,6 +8,8 @@ vars:
   BENCHTIME: 5s
   FUZZTIME: 15s
   FUZZCORPUS: ../fuzz-corpus
+  RACEFLAG: -race={{ne OS "windows"}}
+  BUILDTAGS: ferretdb_tigris
 
 tasks:
   # invoked when `task` is run without arguments
@@ -61,7 +63,7 @@ tasks:
   env-setup:
     deps: [gen-version]
     cmds:
-      - go run {{if ne OS "windows"}}-race{{end}} ./cmd/envtool/main.go
+      - go run {{.RACEFLAG}} ./cmd/envtool/main.go
 
   env-up-detach:
     cmds:
@@ -110,14 +112,14 @@ tasks:
   test-unit-short:
     desc: "Run short unit tests (with caching)"
     cmds:
-      - go test -short {{if ne OS "windows"}}-race{{end}} -shuffle=on -coverprofile=cover.txt -coverpkg=./... ./...
+      - go test -short {{.RACEFLAG}} -tags={{.BUILDTAGS}} -shuffle=on -coverprofile=cover.txt -coverpkg=./... ./...
       - bin/task{{exeExt}} -d tools test-unit-short
 
   test-unit:
     desc: "Run all unit tests"
     cmds:
-      - go test -count=1 {{if ne OS "windows"}}-race{{end}} -shuffle=on -coverprofile=cover.txt -coverpkg=./... ./...
-      - go test -count=1 {{if ne OS "windows"}}-race{{end}} -shuffle=on -bench=. -benchtime=1x ./...
+      - go test -count=1 {{.RACEFLAG}} -tags={{.BUILDTAGS}} -shuffle=on -coverprofile=cover.txt -coverpkg=./... ./...
+      - go test -count=1 {{.RACEFLAG}} -tags={{.BUILDTAGS}} -shuffle=on -bench=. -benchtime=1x ./...
       - bin/task{{exeExt}} -d tools test-unit-short
 
   test-integration:
@@ -132,27 +134,24 @@ tasks:
     dir: integration
     cmds:
       - >
-        go test -count=1 {{if ne OS "windows"}}-race{{end}} -shuffle=on -coverpkg=../...
-        -coverprofile=integration-pg.txt
-        . -handler=pg
+        go test -count=1 {{.RACEFLAG}} -tags={{.BUILDTAGS}} -shuffle=on -coverpkg=../...
+        -coverprofile=integration-pg.txt . -handler=pg
 
   test-integration-tigris:
     desc: "Run integration tests for Tigris handler"
     dir: integration
     cmds:
       - >
-        go test -count=1 {{if ne OS "windows"}}-race{{end}} -shuffle=on -coverpkg=../...
-        -coverprofile=integration-tigris.txt -tags=ferretdb_tigris
-        . -handler=tigris
+        go test -count=1 {{.RACEFLAG}} -tags={{.BUILDTAGS}} -shuffle=on -coverpkg=../...
+        -coverprofile=integration-tigris.txt . -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 -coverpkg=../...
-        -coverprofile=integration-mongodb.txt
-        . -target-port=37017 -compat-port=0
+        go test -count=1 {{.RACEFLAG}} -tags={{.BUILDTAGS}} -shuffle=on -coverpkg=../...
+        -coverprofile=integration-mongodb.txt . -target-port=37017 -compat-port=0
 
   bench-short:
     desc: "Benchmark for about 20 seconds (with default BENCHTIME)"
@@ -191,16 +190,16 @@ tasks:
   fuzz-corpus:
     desc: "Sync seed and generated fuzz corpora with FUZZCORPUS"
     cmds:
-      - go run {{if ne OS "windows"}}-race{{end}} ./cmd/fuzztool/fuzztool.go -dst={{.FUZZCORPUS}} -src=generated
-      - go run {{if ne OS "windows"}}-race{{end}} ./cmd/fuzztool/fuzztool.go -dst={{.FUZZCORPUS}} -src=seed
-      - go run {{if ne OS "windows"}}-race{{end}} ./cmd/fuzztool/fuzztool.go -src={{.FUZZCORPUS}} -dst=generated
+      - go run {{.RACEFLAG}} ./cmd/fuzztool/fuzztool.go -dst={{.FUZZCORPUS}} -src=generated
+      - go run {{.RACEFLAG}} ./cmd/fuzztool/fuzztool.go -dst={{.FUZZCORPUS}} -src=seed
+      - go run {{.RACEFLAG}} ./cmd/fuzztool/fuzztool.go -src={{.FUZZCORPUS}} -dst=generated
 
   build-testcover:
     desc: "Build bin/ferretdb-testcover"
     run: once
     deps: [gen-version]
     cmds:
-      - go test -c -o=bin/ferretdb-testcover{{exeExt}} -trimpath -tags=ferretdb_testcover,ferretdb_tigris {{if ne OS "windows"}}-race{{end}} -coverpkg=./... ./cmd/ferretdb
+      - go test -c -o=bin/ferretdb-testcover{{exeExt}} -trimpath {{.RACEFLAG}} -tags=ferretdb_testcover,{{.BUILDTAGS}} -coverpkg=./... ./cmd/ferretdb
 
   run:
     desc: "Run FerretDB"

docker-compose.yml

@@ -26,6 +26,8 @@ services:
     container_name: ferretdb_tigris
     ports:
       - 8081:8081
+    environment:
+      - TIGRIS_SERVER_LOG_LEVEL=info
 
   # for proxy mode and mongosh
   mongodb:

integration/Taskfile.yml

@@ -2,17 +2,20 @@
 ---
 version: 3
 
+vars:
+  RACEFLAG: -race={{ne OS "windows"}}
+
 tasks:
   env-data:
     cmds:
       - >
-        go test -count=1 {{if ne OS "windows"}}-race{{end}} -run=TestEnvData -v
+        go test -count=1 {{.RACEFLAG}} -run=TestEnvData -v
         . -handler=pg
       - >
-        go test -count=1 {{if ne OS "windows"}}-race{{end}} -run=TestEnvData -v
+        go test -count=1 {{.RACEFLAG}} -run=TestEnvData -v
         -tags=ferretdb_tigris . -handler=tigris
       - >
-        go test -count=1 {{if ne OS "windows"}}-race{{end}} -run=TestEnvData -v
+        go test -count=1 {{.RACEFLAG}} -run=TestEnvData -v
         . -target-port=37017 -compat-port=0
 
   gen:

integration/delete_compat_test.go

@@ -36,6 +36,8 @@ type deleteCompatTestCase struct {
 }
 
 func TestDeleteCompat(t *testing.T) {
+	t.Parallel()
+
 	testCases := map[string]deleteCompatTestCase{
 		"Empty": {
 			filters:    []bson.D{},

integration/query_compat_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo/options"
 
 	"github.com/FerretDB/FerretDB/integration/setup"
@@ -33,6 +34,24 @@ type queryCompatTestCase struct {
 	skip       string                   // skips test if non-empty
 }
 
+func TestQueryCompat(t *testing.T) {
+	t.Parallel()
+
+	testCases := map[string]queryCompatTestCase{
+		"Empty": {
+			filter: bson.D{},
+		},
+		"IDString": {
+			filter: bson.D{{"_id", "string"}},
+		},
+		"IDObjectID": {
+			filter: bson.D{{"_id", primitive.NilObjectID}},
+		},
+	}
+
+	testQueryCompat(t, testCases)
+}
+
 // testQueryCompat tests query compatibility test cases.
 func testQueryCompat(t *testing.T, testCases map[string]queryCompatTestCase) {
 	t.Helper()
@@ -53,7 +72,7 @@ func testQueryCompat(t *testing.T, testCases map[string]queryCompatTestCase) {
 			t.Parallel()
 
 			filter := tc.filter
-			require.NotNil(t, filter)
+			require.NotNil(t, filter, "filter should be set")
 
 			sort := tc.sort
 			if sort == nil {

integration/shareddata/scalars.go

@@ -228,3 +228,13 @@ var Unsets = &Values[string]{
 		"unset": unset,
 	},
 }
+
+// ObjectIDKeys contains documents with ObjectID keys for tests.
+var ObjectIDKeys = &Values[primitive.ObjectID]{
+	name:     "ObjectIDKeys",
+	handlers: []string{"pg", "tigris"},
+	data: map[primitive.ObjectID]any{
+		{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x10, 0x11}: "objectid",
+		primitive.NilObjectID: "objectid-empty",
+	},
+}

integration/shareddata/shareddata.go

@@ -54,6 +54,7 @@ func AllProviders() []Provider {
 		// Timestamps, TODO https://github.com/FerretDB/FerretDB/issues/1007
 		Int64s,
 		// Unsets, TODO https://github.com/FerretDB/FerretDB/issues/1023
+		ObjectIDKeys,
 
 		Composites,
 

internal/handlers/tigris/msg_delete.go

@@ -189,7 +189,7 @@ func (h *Handler) delete(ctx context.Context, fp tigrisdb.FetchParam, docs []*ty
 	ids := make([]map[string]any, len(docs))
 	for i, doc := range docs {
 		id := must.NotFail(tjson.Marshal(must.NotFail(doc.Get("_id"))))
-		ids[i] = map[string]any{"_id": map[string]json.RawMessage{"$eq": id}}
+		ids[i] = map[string]any{"_id": json.RawMessage(id)}
 	}
 
 	var f driver.Filter

internal/handlers/tigris/msg_find.go

@@ -90,7 +90,9 @@ func (h *Handler) MsgFind(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg, er
 		}
 	}
 
-	var fp tigrisdb.FetchParam
+	fp := tigrisdb.FetchParam{
+		Filter: filter,
+	}
 
 	if fp.DB, err = common.GetRequiredParam[string](document, "$db"); err != nil {
 		return nil, err

internal/handlers/tigris/tigrisdb/pushdown_disabled.go

@@ -0,0 +1,20 @@
+// Copyright 2021 FerretDB Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build ferretdb_nopushdown
+
+package tigrisdb
+
+// pushdownEnabled is false when "ferretdb_nopushdown" build tag is provided.
+const pushdownEnabled = false

internal/handlers/tigris/tigrisdb/pushdown_enabled.go

@@ -0,0 +1,20 @@
+// Copyright 2021 FerretDB Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build !ferretdb_nopushdown
+
+package tigrisdb
+
+// pushdownEnabled is true when no "ferretdb_nopushdown" build tag was provided.
+const pushdownEnabled = true

internal/handlers/tigris/tigrisdb/query.go

@@ -16,19 +16,24 @@ package tigrisdb
 
 import (
 	"context"
+	"encoding/json"
 
 	"github.com/tigrisdata/tigris-client-go/driver"
 	"go.uber.org/zap"
 
 	"github.com/FerretDB/FerretDB/internal/tjson"
 	"github.com/FerretDB/FerretDB/internal/types"
 	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
+	"github.com/FerretDB/FerretDB/internal/util/must"
 )
 
 // FetchParam represents options/parameters used by the fetch/query.
 type FetchParam struct {
 	DB         string
 	Collection string
+
+	// Query filter for possible pushdown; may be ignored in part or entirely.
+	Filter *types.Document
 }
 
 // QueryDocuments fetches documents from the given collection.
@@ -59,7 +64,10 @@ func (tdb *TigrisDB) QueryDocuments(ctx context.Context, param FetchParam) ([]*t
 		return nil, lazyerrors.Error(err)
 	}
 
-	iter, err := db.Read(ctx, param.Collection, driver.Filter(`{}`), nil)
+	filter := tdb.queryDocumentsFilter(param.Filter)
+	tdb.L.Sugar().Debugf("Read filter: %s", filter)
+
+	iter, err := db.Read(ctx, param.Collection, filter, nil)
 	if err != nil {
 		return nil, lazyerrors.Error(err)
 	}
@@ -79,3 +87,38 @@ func (tdb *TigrisDB) QueryDocuments(ctx context.Context, param FetchParam) ([]*t
 
 	return res, iter.Err()
 }
+
+// queryDocumentsFilter returns Tigris filter expression that may cover a part of the given filter.
+//
+// FerretDB always filters data itself, so that should be a purely performance optimization.
+func (tdb *TigrisDB) queryDocumentsFilter(filter *types.Document) driver.Filter {
+	if !pushdownEnabled || filter.Len() == 0 {
+		return driver.Filter(`{}`)
+	}
+
+	for k, v := range filter.Map() {
+		// filter only by _id for now
+		if k != "_id" {
+			continue
+		}
+
+		switch v.(type) {
+		case string:
+			// filtering by string values is complicated if the storage supports encodings, collations, etc,
+			// but Tigris doesn't support any of these
+		case types.ObjectID:
+			// filtering by ObjectID is always safe
+		default:
+			// skip other types for now
+			continue
+		}
+
+		// filter by the exact _id value
+		id := must.NotFail(tjson.Marshal(v))
+		f := map[string]any{"_id": json.RawMessage(id)}
+
+		return must.NotFail(json.Marshal(f))
+	}
+
+	return driver.Filter(`{}`)
+}

tools/Taskfile.yml

@@ -2,6 +2,9 @@
 ---
 version: 3
 
+vars:
+  RACEFLAG: -race={{ne OS "windows"}}
+
 tasks:
   lint-go-consistent:
     cmds:
@@ -11,4 +14,4 @@ tasks:
 
   test-unit-short:
     cmds:
-      - go test -short {{if ne OS "windows"}}-race{{end}} -shuffle=on -coverprofile=cover.txt -coverpkg=./... ./checkswitch/...
+      - go test -short {{.RACEFLAG}} -shuffle=on -coverprofile=cover.txt -coverpkg=./... ./checkswitch/...