diff --git a/src/commands.zig b/src/commands.zig index 62320c2..4df57fc 100644 --- a/src/commands.zig +++ b/src/commands.zig @@ -2296,34 +2296,6 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { } else if (std.mem.eql(u8, stage_name, "$limit")) { const n = try stage_count(reply, stage[0].value, "$limit") orelse return; end = @min(end, start + n); - } else if (std.mem.eql(u8, stage_name, "$project")) { - const pp = doc_arg(stage[0].value) orelse return bad_value(reply, "$project requires a document"); - if (try refuse_unprojectable(reply, pp)) return; - // Applied here rather than remembered for the emit. It used to set - // a variable that only the last `$project` in a pipeline could win - // and that no later stage could see -- so a `$match` after a - // `$project` still matched on a field the projection had removed. - const arena = reply.arena_alloc(); - var projected: std.ArrayListUnmanaged(*const bson.Document) = .empty; - errdefer projected.deinit(ctx.gpa); - try projected.ensureTotalCapacity(ctx.gpa, end - start); - if (in_trees) { - for (trees.items[start..end]) |d| { - projected.appendAssumeCapacity(try projected_tree(arena, d, pp)); - } - } else { - for (offs.items[start..end]) |off| { - projected.appendAssumeCapacity(try projected_tree(arena, try doc_tree(arena, coll, off), pp)); - } - } - offs.deinit(ctx.gpa); - offs = .empty; - trees.deinit(ctx.gpa); - trees = projected; - projected = .empty; - in_trees = true; - start = 0; - end = trees.items.len; } else if (std.mem.eql(u8, stage_name, "$group")) { const gp = doc_arg(stage[0].value) orelse return bad_value(reply, "$group requires a document"); const src: Stream = if (in_trees) @@ -2343,7 +2315,7 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { end = trees.items.len; } else if (std.mem.eql(u8, stage_name, "$addFields") or std.mem.eql(u8, stage_name, "$set") or std.mem.eql(u8, stage_name, "$unset") or std.mem.eql(u8, stage_name, "$replaceRoot") or - std.mem.eql(u8, stage_name, "$unwind")) + std.mem.eql(u8, stage_name, "$unwind") or std.mem.eql(u8, stage_name, "$project")) { // The stages that rewrite a document, all one shape: read the // window, build a new list, replace the stream. The design review @@ -2512,55 +2484,18 @@ fn count_only_pipeline(reply: *wire.Reply, stages: []const bson.Value) !?CountSh return .{ .filter = filter, .id_value = id_value, .accs = accs.items }; } -/// Refuse a `$project` this engine cannot carry out, answering the client. -/// Returns true when it did. +/// Refuse a `$project` that mixes inclusion with exclusion, which mongod +/// refuses too -- so this is parity rather than a limitation of this server. +/// Judged on the *flattened* flags, so a nested spec is treated the same way a +/// dotted one is. /// -/// The stage takes inclusion and exclusion flags and nothing else. A computed -/// field (`{y: {$literal: 5}}`) needs the expression evaluator that does not -/// exist yet, and a nested spec (`{a: {b: 1}}`) needs a narrowing -/// `query.project` does not do. Both used to be read as *falsy*, which put the -/// whole projection into its exclusion branch: `{$project: {y: {$literal: 5}}}` -/// returned every document with `y` removed, where mongod adds a computed `y`. -/// Measured on a live server, not inferred. -/// -/// Codes read off mongod 8.3.7. It refuses the empty and the mixed forms too, -/// so those two are parity rather than a limitation of this server. -fn refuse_unprojectable(reply: *wire.Reply, pp: []const bson.Pair) !bool { - if (pp.len == 0) { - try reply.put_error( - @intFromEnum(ErrorCode.location_project_empty), - "Location51272", - "Invalid $project :: caused by :: projection specification must have at least one field", - ); - return true; - } +/// The rest of what this used to refuse -- computed fields and nested specs -- +/// is implemented now. They were refused because both read as falsy, which put +/// the whole projection into its exclusion branch and returned the entire +/// document minus that field. +fn refuse_mixed_projection(reply: *wire.Reply, flags: []const bson.Pair) !bool { var include: ?bool = null; - for (pp) |p| { - switch (p.value) { - .bool, .int32, .int64, .double => {}, - else => { - // A `$`-led document is an expression mongod evaluates and - // names in its own message; anything else is a nested spec, - // where the honest message is the one this server can stand - // behind. - const detail = if (p.value == .doc and p.value.doc.len > 0 and - p.value.doc[0].key.len > 0 and p.value.doc[0].key[0] == '$') - try std.fmt.allocPrint( - reply.arena_alloc(), - "Invalid $project :: caused by :: Unknown expression {s}", - .{p.value.doc[0].key}, - ) - else - try std.fmt.allocPrint( - reply.arena_alloc(), - "Invalid $project :: caused by :: field '{s}' must be an inclusion or " ++ - "exclusion flag: this server projects no computed or nested fields", - .{p.key}, - ); - try reply.put_error(@intFromEnum(ErrorCode.location_project_unknown_expression), "Location31325", detail); - return true; - }, - } + for (flags) |p| { // `_id` is the one field that may be excluded from an inclusion // projection, so it never decides which kind this is. if (std.mem.eql(u8, p.key, "_id")) continue; @@ -3002,6 +2937,83 @@ fn compile_switch(reply: *wire.Reply, arena: std.mem.Allocator, spec: bson.Value const EvalError = error{ DivideByZero, SwitchNoDefault, NonNumericArithmetic, ReplaceRootNotDocument } || std.mem.Allocator.Error || error{ EndOfStream, Overflow, InvalidBson }; +/// A `$project` taken apart: the inclusion/exclusion flags flattened to dotted +/// paths, and the computed fields as expressions. +/// +/// Flattening is what lets a nested spec work without touching +/// `query.project`, which `find` shares: `{n: {x: 1}}` *is* `{"n.x": 1}`, and a +/// dotted path is something the projection already narrows correctly. The +/// nested form used to read as *falsy*, which flipped the whole projection into +/// its exclusion branch and returned every document minus that field. +const ProjectSpec = struct { + flags: []const bson.Pair, + computed: []const Expr.Field, + /// A projection that only computes -- `{$project: {b: }}` -- keeps + /// `_id` and nothing else. `query.project` cannot say that: with no + /// non-`_id` flag it reads the spec as an *exclusion* and returns the whole + /// document, so this case is built directly instead. + id_only: bool = false, +}; + +fn compile_project( + reply: *wire.Reply, + arena: std.mem.Allocator, + pp: []const bson.Pair, +) !?ProjectSpec { + var flags: std.ArrayListUnmanaged(bson.Pair) = .empty; + var computed: std.ArrayListUnmanaged(Expr.Field) = .empty; + if (try flatten_project(reply, arena, pp, "", &flags, &computed)) return null; + if (flags.items.len == 0 and computed.items.len == 0) { + try reply.put_error( + @intFromEnum(ErrorCode.location_project_empty), + "Location51272", + "Invalid $project :: caused by :: projection specification must have at least one field", + ); + return null; + } + return ProjectSpec{ + .flags = flags.items, + .computed = computed.items, + .id_only = flags.items.len == 0, + }; +} + +/// Walk a projection spec into flags and computed fields. Returns true when it +/// has already answered the client. +fn flatten_project( + reply: *wire.Reply, + arena: std.mem.Allocator, + pp: []const bson.Pair, + prefix: []const u8, + flags: *std.ArrayListUnmanaged(bson.Pair), + computed: *std.ArrayListUnmanaged(Expr.Field), +) CompileError!bool { + for (pp) |p| { + const key = if (prefix.len == 0) + p.key + else + try std.fmt.allocPrint(arena, "{s}.{s}", .{ prefix, p.key }); + switch (p.value) { + .bool, .int32, .int64, .double => try flags.append(arena, .{ .key = key, .value = p.value }), + .doc => |d| { + // `{a: {$op: ...}}` computes; `{a: {b: 1}}` narrows. + if (d.len > 0 and d[0].key.len > 0 and d[0].key[0] == '$') { + const e = (try compile_expr(reply, arena, p.value, "a $project value")) orelse return true; + try computed.append(arena, .{ .key = key, .value = e }); + } else if (try flatten_project(reply, arena, d, key, flags, computed)) { + return true; + } + }, + // A bare path is a rename, which is a computed field like any other. + else => { + const e = (try compile_expr(reply, arena, p.value, "a $project value")) orelse return true; + try computed.append(arena, .{ .key = key, .value = e }); + }, + } + } + return false; +} + /// A compiled document-rewriting stage. const Rewrite = union(enum) { /// `$addFields` and `$set`, which are the same stage under two names. @@ -3009,6 +3021,7 @@ const Rewrite = union(enum) { unset: []const []const u8, replace_root: Expr, unwind: UnwindSpec, + project: ProjectSpec, }; fn compile_rewrite( @@ -3019,6 +3032,17 @@ fn compile_rewrite( spec: bson.Value, ) !?Rewrite { _ = ctx; + if (std.mem.eql(u8, name, "$project")) { + const pp = doc_arg(spec) orelse { + try bad_value(reply, "$project requires a document"); + return null; + }; + const parsed = (try compile_project(reply, arena, pp)) orelse return null; + // Mixed inclusion and exclusion is judged on the *flattened* flags, so + // a nested spec is treated the same way a dotted one is. + if (try refuse_mixed_projection(reply, parsed.flags)) return null; + return Rewrite{ .project = parsed }; + } if (std.mem.eql(u8, name, "$unwind")) { return Rewrite{ .unwind = (try unwind_spec(reply, spec)) orelse return null }; } @@ -3081,6 +3105,23 @@ fn apply_rewrite( }; try out.append(gpa, try tree_of(ec.arena, pairs)); }, + .project => |spec| { + var kept: std.ArrayListUnmanaged(bson.Pair) = .empty; + if (spec.id_only) { + if (bson.get_pair(doc.pairs, "_id")) |id| try kept.append(ec.arena, .{ .key = "_id", .value = id }); + } else { + query.project(ec.arena, doc, &.{ .arena = undefined, .pairs = spec.flags }, &kept) catch + return error.OutOfMemory; + } + var pairs: []const bson.Pair = kept.items; + for (spec.computed) |f| { + // Same rule as `$addFields`: an expression that resolves to + // nothing leaves the field out rather than setting it null. + const v = (try eval_expr(ec, f.value)) orelse continue; + pairs = try set_path(ec.arena, pairs, f.key, v); + } + try out.append(gpa, try tree_of(ec.arena, pairs)); + }, .unwind => |spec| { const found = try eval_expr(ec, .{ .path = spec.path }); const items: []const bson.Value = switch (found orelse bson.Value.null) { @@ -5301,15 +5342,13 @@ test "$project is a stage, not a note about how to print the answer" { }, .code = 31254, }, + // A computed field and a nested spec stood here until the expression + // evaluator and the flattening landed; both work now. What is left is + // the one shape mongod refuses too. .{ - .name = "a computed field", - .spec = &.{.{ .key = "y", .value = .{ .doc = &.{.{ .key = "$literal", .value = .{ .int32 = 5 } }} } }}, - .code = 31325, - }, - .{ - .name = "a nested spec", - .spec = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "b", .value = .{ .int32 = 1 } }} } }}, - .code = 31325, + .name = "an unknown expression", + .spec = &.{.{ .key = "y", .value = .{ .doc = &.{.{ .key = "$bogusExpr", .value = .{ .int32 = 5 } }} } }}, + .code = 168, }, }; for (cases) |c| { diff --git a/tests/spec/aggregate/README.md b/tests/spec/aggregate/README.md index 829f3a4..33d2709 100644 --- a/tests/spec/aggregate/README.md +++ b/tests/spec/aggregate/README.md @@ -64,7 +64,7 @@ accumulators in: ``` group-accumulators.json 19 pass 0 fail 0 skip expressions.json 27 pass 0 fail 0 skip -document-stages.json 21 pass 3 fail 0 skip +document-stages.json 24 pass 0 fail 0 skip ``` `group-accumulators` found its first real disagreement on the way to 18: @@ -73,11 +73,8 @@ that counted documents rather than numbers would have passed every test anybody would think to write by hand. `expressions.json` was recorded before the evaluator was written and read -1 pass / 26 fail against it; it is green now. `document-stages.json` was recorded at -0 pass / 24 fail; `$addFields`/`$set`, `$unset`, `$replaceRoot` and `$unwind` -took it to 21. The three left are `$project`'s computed fields, its renames and -its nested inclusions, which want the shared `query.project` and are their own -change. +1 pass / 26 fail against it; it is green now. `document-stages.json` was recorded at 0 pass / 24 fail and is green. The whole +corpus is: 70 cases, every answer byte-identical to mongod 8.3.7. What recording *that* settled: