diff --git a/integration/explain_command_test.go b/integration/explain_command_test.go index c9d45820e7b5..24d9ceaf825c 100644 --- a/integration/explain_command_test.go +++ b/integration/explain_command_test.go @@ -126,3 +126,20 @@ func TestExplainNonExistentCollection(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, res) } + +func TestExplainLimitInt(t *testing.T) { + t.Parallel() + + ctx, collection := setup.Setup(t) + + var res bson.D + err := collection.Database().RunCommand(ctx, bson.D{ + {"explain", bson.D{ + {"find", collection.Name()}, + {"limit", int32(1)}, + }}, + }).Decode(&res) + + assert.NoError(t, err) + assert.NotNil(t, res) +} diff --git a/integration/query_test.go b/integration/query_test.go index d2df76ad73d9..d274ed3c3978 100644 --- a/integration/query_test.go +++ b/integration/query_test.go @@ -821,8 +821,6 @@ func TestQueryCommandSingleBatch(t *testing.T) { } func TestQueryCommandLimitPushDown(t *testing.T) { - setup.SkipForPostgreSQL(t, "https://github.com/FerretDB/FerretDB/issues/3416") - t.Parallel() // must use a collection of documents which does not support query pushdown to test limit pushdown @@ -844,34 +842,29 @@ func TestQueryCommandLimitPushDown(t *testing.T) { failsForSQLite string // optional, if set, the case is expected to fail for SQLite due to given issue }{ "Simple": { - limit: 1, - len: 1, - limitPushdown: true, - failsForSQLite: "https://github.com/FerretDB/FerretDB/issues/3181", + limit: 1, + len: 1, + limitPushdown: true, }, "AlmostAll": { - limit: int64(len(shareddata.Composites.Docs()) - 1), - len: len(shareddata.Composites.Docs()) - 1, - limitPushdown: true, - failsForSQLite: "https://github.com/FerretDB/FerretDB/issues/3181", + limit: int64(len(shareddata.Composites.Docs()) - 1), + len: len(shareddata.Composites.Docs()) - 1, + limitPushdown: true, }, "All": { - limit: int64(len(shareddata.Composites.Docs())), - len: len(shareddata.Composites.Docs()), - limitPushdown: true, - failsForSQLite: "https://github.com/FerretDB/FerretDB/issues/3181", + limit: int64(len(shareddata.Composites.Docs())), + len: len(shareddata.Composites.Docs()), + limitPushdown: true, }, "More": { - limit: int64(len(shareddata.Composites.Docs()) + 1), - len: len(shareddata.Composites.Docs()), - limitPushdown: true, - failsForSQLite: "https://github.com/FerretDB/FerretDB/issues/3181", + limit: int64(len(shareddata.Composites.Docs()) + 1), + len: len(shareddata.Composites.Docs()), + limitPushdown: true, }, "Big": { - limit: 1000, - len: len(shareddata.Composites.Docs()), - limitPushdown: true, - failsForSQLite: "https://github.com/FerretDB/FerretDB/issues/3181", + limit: 1000, + len: len(shareddata.Composites.Docs()), + limitPushdown: true, }, "Zero": { limit: 0, diff --git a/internal/backends/collection.go b/internal/backends/collection.go index 1795d3221ed4..91686700fd87 100644 --- a/internal/backends/collection.go +++ b/internal/backends/collection.go @@ -80,6 +80,7 @@ type QueryParams struct { // TODO https://github.com/FerretDB/FerretDB/issues/3235 Filter *types.Document Sort *SortField + Limit int64 // if 0 no limit pushdown is applied OnlyRecordIDs bool // TODO https://github.com/FerretDB/FerretDB/issues/3490 Comment string // TODO https://github.com/FerretDB/FerretDB/issues/3573 } @@ -207,6 +208,7 @@ type ExplainParams struct { // TODO https://github.com/FerretDB/FerretDB/issues/3235 Filter *types.Document Sort *SortField + Limit int64 // if 0 no limit pushdown is applied } // ExplainResult represents the results of Collection.Explain method. @@ -215,6 +217,7 @@ type ExplainResult struct { // TODO https://github.com/FerretDB/FerretDB/issues/3235 QueryPushdown bool SortPushdown bool + LimitPushdown bool } // Explain return a backend-specific execution plan for the given query. diff --git a/internal/backends/postgresql/collection.go b/internal/backends/postgresql/collection.go index 6447905a27ff..ef04c31bd4d7 100644 --- a/internal/backends/postgresql/collection.go +++ b/internal/backends/postgresql/collection.go @@ -101,6 +101,11 @@ func (c *collection) Query(ctx context.Context, params *backends.QueryParams) (* args = append(args, sortArgs...) } + if params.Limit != 0 { + q += fmt.Sprintf(` LIMIT %s`, placeholder.Next()) + args = append(args, params.Limit) + } + rows, err := p.Query(ctx, q, args...) if err != nil { return nil, lazyerrors.Error(err) @@ -326,6 +331,12 @@ func (c *collection) Explain(ctx context.Context, params *backends.ExplainParams q += where + if params.Limit != 0 { + q += fmt.Sprintf(` LIMIT %s`, placeholder.Next()) + args = append(args, params.Limit) + res.LimitPushdown = true + } + var b []byte err = p.QueryRow(ctx, q, args...).Scan(&b) diff --git a/internal/backends/sqlite/collection.go b/internal/backends/sqlite/collection.go index b2e00d69bc93..b779203bb29f 100644 --- a/internal/backends/sqlite/collection.go +++ b/internal/backends/sqlite/collection.go @@ -64,12 +64,16 @@ func (c *collection) Query(ctx context.Context, params *backends.QueryParams) (* }, nil } + if params == nil { + params = new(backends.QueryParams) + } + var whereClause string var args []any // that logic should exist in one place // TODO https://github.com/FerretDB/FerretDB/issues/3235 - if params != nil && params.Filter.Len() == 1 { + if params.Filter.Len() == 1 { v, _ := params.Filter.Get("_id") switch v.(type) { case string, types.ObjectID: @@ -82,6 +86,11 @@ func (c *collection) Query(ctx context.Context, params *backends.QueryParams) (* q := fmt.Sprintf(`SELECT %s FROM %q`+whereClause, metadata.DefaultColumn, meta.TableName) + if params.Limit != 0 { + q += ` LIMIT ?` + args = append(args, params.Limit) + } + rows, err := db.QueryContext(ctx, q, args...) if err != nil { return nil, lazyerrors.Error(err) @@ -238,13 +247,17 @@ func (c *collection) Explain(ctx context.Context, params *backends.ExplainParams }, nil } + if params == nil { + params = new(backends.ExplainParams) + } + var queryPushdown bool var whereClause string var args []any // that logic should exist in one place // TODO https://github.com/FerretDB/FerretDB/issues/3235 - if params != nil && params.Filter.Len() == 1 { + if params.Filter.Len() == 1 { v, _ := params.Filter.Get("_id") switch v.(type) { case string, types.ObjectID: @@ -258,6 +271,14 @@ func (c *collection) Explain(ctx context.Context, params *backends.ExplainParams q := fmt.Sprintf(`EXPLAIN QUERY PLAN SELECT %s FROM %q`+whereClause, metadata.DefaultColumn, meta.TableName) + var limitPushdown bool + + if params.Limit != 0 { + q += ` LIMIT ?` + args = append(args, params.Limit) + limitPushdown = true + } + rows, err := db.QueryContext(ctx, q, args...) if err != nil { return nil, lazyerrors.Error(err) @@ -290,6 +311,7 @@ func (c *collection) Explain(ctx context.Context, params *backends.ExplainParams return &backends.ExplainResult{ QueryPlanner: must.NotFail(types.NewDocument("Plan", queryPlan)), QueryPushdown: queryPushdown, + LimitPushdown: limitPushdown, }, nil } diff --git a/internal/handlers/common/explain.go b/internal/handlers/common/explain.go index 6f5a0163b623..991de7a8bed4 100644 --- a/internal/handlers/common/explain.go +++ b/internal/handlers/common/explain.go @@ -91,7 +91,7 @@ func GetExplainParams(document *types.Document, l *zap.Logger) (*ExplainParams, var limit, skip int64 - if limit, err = GetOptionalParam(explain, "limit", limit); err != nil { + if limit, err = GetLimitParam(explain); err != nil { return nil, err } diff --git a/internal/handlers/common/params.go b/internal/handlers/common/params.go index 76739d2965bc..3f8d0882cd06 100644 --- a/internal/handlers/common/params.go +++ b/internal/handlers/common/params.go @@ -100,6 +100,26 @@ func AssertType[T types.Type](value any) (T, error) { return res, nil } +// GetLimitParam returns limit value from provided query document. +func GetLimitParam(doc *types.Document) (int64, error) { + v, _ := doc.Get("limit") + if v == nil { + return 0, nil + } + + res, err := commonparams.GetWholeNumberParam(v) + if err != nil { + msg := fmt.Sprintf( + `BSON field '%s' is the wrong type '%s', expected type '%s'`, + "limit", commonparams.AliasFromType(v), commonparams.AliasFromType(res), + ) + + return res, commonerrors.NewCommandErrorMsgWithArgument(commonerrors.ErrTypeMismatch, msg, "limit") + } + + return res, nil +} + // GetLimitStageParam returns $limit stage argument from the provided value. // It returns the proper error if value doesn't meet requirements. func GetLimitStageParam(value any) (int64, error) { diff --git a/internal/handlers/sqlite/msg_explain.go b/internal/handlers/sqlite/msg_explain.go index 2ee08213a3d0..c29244ffe0e7 100644 --- a/internal/handlers/sqlite/msg_explain.go +++ b/internal/handlers/sqlite/msg_explain.go @@ -101,6 +101,15 @@ func (h *Handler) MsgExplain(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg, } } + // Limit pushdown is not applied if: + // - `filter` is set, it must fetch all documents to filter them in memory; + // - `sort` is set but `EnableSortPushdown` is not set, it must fetch all documents + // and sort them in memory; + // - `skip` is non-zero value, skip pushdown is not supported yet. + if params.Filter.Len() == 0 && (params.Sort.Len() == 0 || h.EnableSortPushdown) && params.Skip == 0 { + qp.Limit = params.Limit + } + res, err := coll.Explain(ctx, &qp) if err != nil { return nil, lazyerrors.Error(err) @@ -118,7 +127,7 @@ func (h *Handler) MsgExplain(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg, // TODO https://github.com/FerretDB/FerretDB/issues/3235 "pushdown", res.QueryPushdown, "sortingPushdown", res.SortPushdown, - "limitPushdown", false, + "limitPushdown", res.LimitPushdown, "ok", float64(1), ))}, diff --git a/internal/handlers/sqlite/msg_find.go b/internal/handlers/sqlite/msg_find.go index 4d4a7774ca43..0ba863104751 100644 --- a/internal/handlers/sqlite/msg_find.go +++ b/internal/handlers/sqlite/msg_find.go @@ -98,6 +98,15 @@ func (h *Handler) MsgFind(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg, er } } + // Limit pushdown is not applied if: + // - `filter` is set, it must fetch all documents to filter them in memory; + // - `sort` is set but `EnableSortPushdown` is not set, it must fetch all documents + // and sort them in memory; + // - `skip` is non-zero value, skip pushdown is not supported yet. + if params.Filter.Len() == 0 && (params.Sort.Len() == 0 || h.EnableSortPushdown) && params.Skip == 0 { + qp.Limit = params.Limit + } + cancel := func() {} if params.MaxTimeMS != 0 { // It is not clear if maxTimeMS affects only find, or both find and getMore (as the current code does).