src/commands.zig

@@ -73,6 +73,9 @@ pub const ErrorCode = enum(i32) {
     // error-code parity every milestone is held to. The message names the
     // construct, so the answer is a bug report rather than a wrong number.
     invalid_pipeline_operator = 168,
+    location_project_empty = 51272,
+    location_project_mixed = 31254,
+    location_project_unknown_expression = 31325,
     location_unknown_group_operator = 15952,
     location_group_needs_id = 15955,
     location_accumulator_not_object = 40234,
@@ -2155,7 +2158,6 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
     var start: usize = 0;
     var end: usize = offs.items.len;
     var count_stage: ?[]const u8 = null;
-    var proj_pairs: ?[]const bson.Pair = null;
 
     for (stages) |stage_v| {
         const stage = switch (stage_v) {
@@ -2229,7 +2231,33 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
             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")) {
-            proj_pairs = doc_arg(stage[0].value);
+            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)
@@ -2274,12 +2302,12 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
         const arena = reply.arena_alloc();
         if (in_trees) {
             const window = trees.items[start..end];
-            try emit_first_batch(ctx, reply, db_name, coll_name, proj_pairs, window, batch_size);
+            try emit_first_batch(ctx, reply, db_name, coll_name, null, window, batch_size);
         } else {
             var page: std.ArrayListUnmanaged(*const bson.Document) = .empty;
             for (offs.items[start..end]) |off| try page.append(arena, try doc_tree(arena, coll, off));
             const window = page.items;
-            try emit_first_batch(ctx, reply, db_name, coll_name, proj_pairs, window, batch_size);
+            try emit_first_batch(ctx, reply, db_name, coll_name, null, window, batch_size);
         }
     }
     try reply.put_ok();
@@ -2350,6 +2378,93 @@ 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.
+///
+/// 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;
+    }
+    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;
+            },
+        }
+        // `_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;
+        const flag = query.truthy(p.value);
+        if (include) |want| {
+            if (want != flag) {
+                const detail = try std.fmt.allocPrint(
+                    reply.arena_alloc(),
+                    "Invalid $project :: caused by :: Cannot do {s} on field {s} in {s} projection",
+                    .{
+                        if (flag) "inclusion" else "exclusion",
+                        p.key,
+                        if (want) "inclusion" else "exclusion",
+                    },
+                );
+                try reply.put_error(@intFromEnum(ErrorCode.location_project_mixed), "Location31254", detail);
+                return true;
+            }
+        } else include = flag;
+    }
+    return false;
+}
+
+/// One document put through a projection, as a tree the rest of the pipeline
+/// can read. This is what makes `$project` a stage rather than a note about how
+/// to print the answer.
+fn projected_tree(
+    arena: std.mem.Allocator,
+    doc: *const bson.Document,
+    pp: []const bson.Pair,
+) !*const bson.Document {
+    var out: std.ArrayListUnmanaged(bson.Pair) = .empty;
+    try query.project(arena, doc, &.{ .arena = undefined, .pairs = pp }, &out);
+    const projected = try arena.create(bson.Document);
+    projected.* = .{ .arena = undefined, .pairs = out.items };
+    return projected;
+}
+
 /// Where a pipeline stage reads its input.
 ///
 /// A pipeline starts as slab offsets -- matched and reordered in place, never
@@ -4034,6 +4149,121 @@ test "aggregate $sort without a preceding $group sorts and frees correctly" {
     }
 }
 
+test "$project is a stage, not a note about how to print the answer" {
+    // It used to set a variable applied once, at the emit. Three consequences,
+    // all measured on a live server before this changed: only the *last*
+    // `$project` in a pipeline had any effect; a `$match` after one still saw
+    // the field it had removed; and `{y: {$literal: 5}}` read as falsy, which
+    // flipped the whole projection into its exclusion branch and returned every
+    // document minus `y`.
+    //
+    // Every answer below is byte-identical to mongod 8.3.7 on the same input.
+    //
+    // Mutation check: move the projection back to the emit -- pass `pp` to
+    // `emit_first_batch` instead of transforming the stream -- and the first
+    // two cases go red.
+    var threaded = std.Io.Threaded.init(testing.allocator, .{});
+    defer threaded.deinit();
+    const io = threaded.io();
+
+    var tdb = try TestDb.init(io);
+    defer tdb.deinit();
+    try dispatch_insert(&tdb, io, "pj", &.{
+        .{ .doc = &.{ .{ .key = "_id", .value = .{ .int32 = 1 } }, .{ .key = "g", .value = .{ .string = "a" } }, .{ .key = "x", .value = .{ .int32 = 10 } } } },
+        .{ .doc = &.{ .{ .key = "_id", .value = .{ .int32 = 2 } }, .{ .key = "g", .value = .{ .string = "b" } }, .{ .key = "x", .value = .{ .int32 = 20 } } } },
+    });
+    var ctx = tdb.ctx(io);
+
+    const keep_g = bson.Value{ .doc = &.{.{ .key = "$project", .value = .{ .doc = &.{
+        .{ .key = "g", .value = .{ .int32 = 1 } },
+    } } }} };
+
+    // A $match after a $project cannot see what the projection removed.
+    {
+        const match_x = bson.Value{ .doc = &.{.{ .key = "$match", .value = .{ .doc = &.{
+            .{ .key = "x", .value = .{ .int32 = 10 } },
+        } } }} };
+        const stages = [_]bson.Value{ keep_g, match_x };
+        var msg = try parse_fake_msg("aggregate", .{ .string = "pj" }, &.{
+            .{ .key = "pipeline", .value = .{ .array = &stages } },
+            .{ .key = "cursor", .value = .{ .doc = &.{} } },
+        });
+        defer msg.deinit();
+        var reply = wire.Reply.init(testing.allocator);
+        defer reply.deinit();
+        try dispatch(&ctx, &msg, &reply);
+        const cur = bson.get_pair(reply.pairs.items, "cursor").?;
+        try testing.expectEqual(@as(usize, 0), bson.get_pair(cur.doc, "firstBatch").?.array.len);
+    }
+
+    // And a $group after one reads the projected document, not the stored one.
+    {
+        const group_g = bson.Value{ .doc = &.{.{ .key = "$group", .value = .{ .doc = &.{
+            .{ .key = "_id", .value = .null },
+            .{ .key = "n", .value = .{ .doc = &.{.{ .key = "$sum", .value = .{ .int32 = 1 } }} } },
+            .{ .key = "x", .value = .{ .doc = &.{.{ .key = "$sum", .value = .{ .string = "$x" } }} } },
+        } } }} };
+        const stages = [_]bson.Value{ keep_g, group_g };
+        var msg = try parse_fake_msg("aggregate", .{ .string = "pj" }, &.{
+            .{ .key = "pipeline", .value = .{ .array = &stages } },
+            .{ .key = "cursor", .value = .{ .doc = &.{} } },
+        });
+        defer msg.deinit();
+        var reply = wire.Reply.init(testing.allocator);
+        defer reply.deinit();
+        try dispatch(&ctx, &msg, &reply);
+        const cur = bson.get_pair(reply.pairs.items, "cursor").?;
+        const batch = bson.get_pair(cur.doc, "firstBatch").?.array;
+        try testing.expectEqual(@as(usize, 1), batch.len);
+        try testing.expectEqual(@as(i32, 2), bson.get_pair(batch[0].doc, "n").?.int32);
+        // `x` was projected away, so summing it is summing nothing.
+        try testing.expectEqual(@as(i32, 0), bson.get_pair(batch[0].doc, "x").?.int32);
+    }
+
+    // The three shapes mongod refuses, refused with its codes.
+    const Case = struct { name: []const u8, spec: []const bson.Pair, code: i32 };
+    const cases = [_]Case{
+        .{ .name = "empty", .spec = &.{}, .code = 51272 },
+        .{
+            .name = "mixed inclusion and exclusion",
+            .spec = &.{
+                .{ .key = "g", .value = .{ .int32 = 1 } },
+                .{ .key = "x", .value = .{ .int32 = 0 } },
+            },
+            .code = 31254,
+        },
+        .{
+            .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,
+        },
+    };
+    for (cases) |c| {
+        const stages = [_]bson.Value{.{ .doc = &.{.{ .key = "$project", .value = .{ .doc = c.spec } }} }};
+        var msg = try parse_fake_msg("aggregate", .{ .string = "pj" }, &.{
+            .{ .key = "pipeline", .value = .{ .array = &stages } },
+            .{ .key = "cursor", .value = .{ .doc = &.{} } },
+        });
+        defer msg.deinit();
+        var reply = wire.Reply.init(testing.allocator);
+        defer reply.deinit();
+        try dispatch(&ctx, &msg, &reply);
+        testing.expectEqual(@as(f64, 0.0), bson.get_pair(reply.pairs.items, "ok").?.double) catch |err| {
+            std.debug.print("  {s}: answered ok:1\n", .{c.name});
+            return err;
+        };
+        testing.expectEqual(c.code, bson.get_pair(reply.pairs.items, "code").?.int32) catch |err| {
+            std.debug.print("  {s}: wrong code\n", .{c.name});
+            return err;
+        };
+    }
+}
+
 test "a pipeline may group what an earlier stage generated" {
     // A remote crash, reachable by any client with a two-stage pipeline and no
     // authentication in front of it: `$group` took `[]const u64` and was handed