Pushdown Tigris queries with dot notation (#1908)

Closes #1841.
FerretDB · Feb 9, 2023 · 8a41070 · 8a41070
1 parent 31858ee
commit 8a41070
Show file tree

Hide file tree

Showing 5 changed files with 201 additions and 15 deletions.
diff --git a/integration/query_comparison_compat_test.go b/integration/query_comparison_compat_test.go
@@ -52,10 +52,9 @@ func testQueryComparisonCompatImplicit() map[string]queryCompatTestCase {
 			resultPushdown: true,
 		},
 		"DocumentDotNotationNoSuchField": {
-			filter:             bson.D{{"no-such-field.some", 42}},
-			resultType:         emptyResult,
-			resultPushdown:     true,
-			skipTigrisPushdown: true,
+			filter:         bson.D{{"no-such-field.some", 42}},
+			resultType:     emptyResult,
+			resultPushdown: true,
 		},
 		"ArrayNoSuchField": {
 			filter:     bson.D{{"no-such-field", bson.A{42}}},

diff --git a/internal/handlers/pg/pgdb/query.go b/internal/handlers/pg/pgdb/query.go
@@ -236,18 +236,18 @@ func prepareWhereClause(sqlFilters *types.Document) (string, []any) {
 	var p Placeholder
 
 	for k, v := range sqlFilters.Map() {
-		if len(k) != 0 && k[0] == '$' {
-			// skip $comment
-			continue
-		}
-
 		keyOperator := "->" // keyOperator is the operator that is used to access the field. (->/#>)
 
-		var key any = k   // key can be either a string '"v"' or path '{v,foo}'
+		var key any = k   // key can be either a string '"v"' or PostgreSQL path '{v,foo}'
 		var prefix string // prefix is the first key in path, if the filter key is not a path - the prefix is empty
 
-		// If the key is in dot notation use path operator (#>)
-		if len(k) != 0 {
+		if k != "" {
+			// skip $comment
+			if k[0] == '$' {
+				continue
+			}
+
+			// If the key is in dot notation use path operator (#>)
 			if path := types.NewPathFromString(k); path.Len() > 1 {
 				keyOperator = "#>"
 				key = path.Slice()     // '{v,foo}'

diff --git a/internal/handlers/pg/pgdb/query_test.go b/internal/handlers/pg/pgdb/query_test.go
@@ -16,6 +16,7 @@ package pgdb
 
 import (
 	"context"
+	"math"
 	"testing"
 
 	"github.com/jackc/pgx/v4"
@@ -281,3 +282,80 @@ func TestGetDocuments(t *testing.T) {
 		require.NoError(t, err)
 	})
 }
+
+func TestPrepareWhereClause(t *testing.T) {
+	t.Parallel()
+	objectID := types.ObjectID{0x62, 0x56, 0xc5, 0xba, 0x0b, 0xad, 0xc0, 0xff, 0xee, 0xff, 0xff, 0xff}
+
+	// WHERE clauses occurring frequently in tests
+	whereEq := " WHERE ((_jsonb->$1)::jsonb = $2)"
+	whereEqOrContain := whereEq + " OR (_jsonb->$1)::jsonb @> $2"
+
+	for name, tc := range map[string]struct {
+		filter   *types.Document
+		expected string
+	}{
+		"String": {
+			filter:   must.NotFail(types.NewDocument("v", "foo")),
+			expected: whereEqOrContain,
+		},
+		"EmptyString": {
+			filter:   must.NotFail(types.NewDocument("v", "")),
+			expected: whereEqOrContain,
+		},
+		"Int32": {
+			filter:   must.NotFail(types.NewDocument("v", int32(42))),
+			expected: whereEqOrContain,
+		},
+		"Int64": {
+			filter:   must.NotFail(types.NewDocument("v", int64(42))),
+			expected: whereEqOrContain,
+		},
+		"Float64": {
+			filter:   must.NotFail(types.NewDocument("v", float64(42.13))),
+			expected: whereEqOrContain,
+		},
+		"MaxFloat64": {
+			filter:   must.NotFail(types.NewDocument("v", math.MaxFloat64)),
+			expected: whereEqOrContain,
+		},
+		"Bool": {
+			filter: must.NotFail(types.NewDocument("v", true)),
+		},
+		"Comment": {
+			filter: must.NotFail(types.NewDocument("$comment", "I'm comment")),
+		},
+		"ObjectID": {
+			filter:   must.NotFail(types.NewDocument("v", objectID)),
+			expected: whereEqOrContain,
+		},
+		"IDObjectID": {
+			filter:   must.NotFail(types.NewDocument("_id", objectID)),
+			expected: whereEq,
+		},
+		"IDString": {
+			filter:   must.NotFail(types.NewDocument("_id", "foo")),
+			expected: whereEq,
+		},
+		"IDDotNotation": {
+			filter:   must.NotFail(types.NewDocument("_id.doc", "foo")),
+			expected: " WHERE ((_jsonb#>$1)::jsonb = $2)",
+		},
+		"DotNotation": {
+			filter:   must.NotFail(types.NewDocument("v.doc", "foo")),
+			expected: " WHERE ((_jsonb#>$1)::jsonb = $2) OR (_jsonb#>$1)::jsonb @> $2",
+		},
+		"DotNotationArrayIndex": {
+			filter:   must.NotFail(types.NewDocument("v.arr.0", "foo")),
+			expected: " WHERE ((_jsonb#>$1)::jsonb = $2) OR (_jsonb#>$1)::jsonb @> $2",
+		},
+	} {
+		name, tc := name, tc
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			actual, _ := prepareWhereClause(tc.filter)
+			assert.Equal(t, tc.expected, actual)
+		})
+	}
+}
diff --git a/internal/handlers/tigris/tigrisdb/query.go b/internal/handlers/tigris/tigrisdb/query.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"strconv"
 	"time"
 
 	"github.com/tigrisdata/tigris-client-go/driver"
@@ -87,17 +88,37 @@ func (tdb *TigrisDB) BuildFilter(filter *types.Document) driver.Filter {
 	res := map[string]any{}
 
 	for k, v := range filter.Map() {
+		key := k // key can be either a single key string '"v"' or Tigris dot notation '"v.foo"'
+
 		// TODO https://github.com/FerretDB/FerretDB/issues/1940
 		if v == "" {
 			continue
 		}
 
 		if k != "" {
 			// don't pushdown $comment, it's attached to query in handlers
-			// TODO https://github.com/FerretDB/FerretDB/issues/1841
-			if k[0] == '$' || types.NewPathFromString(k).Len() > 1 {
+			if k[0] == '$' {
 				continue
 			}
+
+			// If the key is in dot notation translate it to a tigris dot notation
+			if path := types.NewPathFromString(k); path.Len() > 1 {
+				indexSearch := false
+
+				// TODO https://github.com/FerretDB/FerretDB/issues/1914
+				for _, k := range path.Slice() {
+					if _, err := strconv.Atoi(k); err == nil {
+						indexSearch = true
+						break
+					}
+				}
+
+				if indexSearch {
+					continue
+				}
+
+				key = path.String() // '"v.foo"'
+			}
 		}
 
 		switch v.(type) {
@@ -106,7 +127,7 @@ func (tdb *TigrisDB) BuildFilter(filter *types.Document) driver.Filter {
 			continue
 		case float64, string, types.ObjectID, int32, int64:
 			rawValue := must.NotFail(tjson.Marshal(v))
-			res[k] = json.RawMessage(rawValue)
+			res[key] = json.RawMessage(rawValue)
 		default:
 			panic(fmt.Sprintf("Unexpected type of field %s: %T", k, v))
 		}

diff --git a/internal/handlers/tigris/tigrisdb/query_test.go b/internal/handlers/tigris/tigrisdb/query_test.go
@@ -18,9 +18,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 	"strings"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/tigrisdata/tigris-client-go/config"
 	"github.com/tigrisdata/tigris-client-go/driver"
@@ -220,6 +222,92 @@ func TestQueryDocuments(t *testing.T) {
 	})
 }
 
+func TestBuildFilter(t *testing.T) {
+	t.Parallel()
+	tdb := TigrisDB{}
+	objectID := types.ObjectID{0x62, 0x56, 0xc5, 0xba, 0x0b, 0xad, 0xc0, 0xff, 0xee, 0xff, 0xff, 0xff}
+
+	for name, tc := range map[string]struct {
+		filter   *types.Document
+		expected string
+		skip     string // defaults to `{}`
+	}{
+		"String": {
+			filter:   must.NotFail(types.NewDocument("v", "foo")),
+			expected: `{"v":"foo"}`,
+		},
+		"EmptyString": {
+			filter:   must.NotFail(types.NewDocument("v", "")),
+			expected: `{"v":""}`,
+			skip:     "https://github.com/FerretDB/FerretDB/issues/1940",
+		},
+		"Int32": {
+			filter:   must.NotFail(types.NewDocument("v", int32(42))),
+			expected: `{"v":42}`,
+		},
+		"Int64": {
+			filter:   must.NotFail(types.NewDocument("v", int64(42))),
+			expected: `{"v":42}`,
+		},
+		"Float64": {
+			filter:   must.NotFail(types.NewDocument("v", float64(42.13))),
+			expected: `{"v":42.13}`,
+		},
+		"MaxFloat64": {
+			filter:   must.NotFail(types.NewDocument("v", math.MaxFloat64)),
+			expected: `{"v":1.7976931348623157e+308}`,
+		},
+		"Bool": {
+			filter: must.NotFail(types.NewDocument("v", true)),
+		},
+		"Comment": {
+			filter: must.NotFail(types.NewDocument("$comment", "I'm comment")),
+		},
+		"ObjectID": {
+			filter:   must.NotFail(types.NewDocument("v", objectID)),
+			expected: `{"v":"YlbFugutwP/u////"}`,
+		},
+		"IDObjectID": {
+			filter:   must.NotFail(types.NewDocument("_id", objectID)),
+			expected: `{"_id":"YlbFugutwP/u////"}`,
+		},
+		"IDString": {
+			filter:   must.NotFail(types.NewDocument("_id", "foo")),
+			expected: `{"_id":"foo"}`,
+		},
+		"IDDotNotation": {
+			filter:   must.NotFail(types.NewDocument("_id.doc", "foo")),
+			expected: `{"_id.doc":"foo"}`,
+		},
+		"DotNotation": {
+			filter:   must.NotFail(types.NewDocument("v.doc", "foo")),
+			expected: `{"v.doc":"foo"}`,
+		},
+		"DotNotationArrayIndex": {
+			filter: must.NotFail(types.NewDocument("v.arr.0", "foo")),
+		},
+	} {
+		name, tc := name, tc
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			if tc.skip != "" {
+				t.Skip(tc.skip)
+			}
+
+			// replace default value with default json
+			if tc.expected == "" {
+				tc.expected = "{}"
+			}
+
+			expected := driver.Filter(tc.expected)
+			actual := tdb.BuildFilter(tc.filter)
+
+			assert.Equal(t, expected, actual)
+		})
+	}
+}
+
 func setup(t *testing.T) (string, string, context.Context, *TigrisDB) {
 	t.Helper()