SERVER-57334 Create desugaring namespace for DocumentSourceDensify

4 files changed, 407 insertions, 0 deletions

diff --git a/jstests/aggregation/sources/densify/desugar.js b/jstests/aggregation/sources/densify/desugar.js
new file mode 100644
index 00000000000..b1f7c11154d
--- /dev/null
+++ b/jstests/aggregation/sources/densify/desugar.js
@@ -0,0 +1,140 @@
+/**
+ * Test how $densify desugars.
+ *
+ * @tags: [
+ *   # Needed as $densify is a 51 feature.
+ *   requires_fcv_51,
+ *   # We're testing the explain plan, not the query results, so the facet passthrough would fail.
+ *   do_not_wrap_aggregations_in_facets,
+ * ]
+ */
+(function() {
+"use strict";
+
+const featureEnabled =
+    assert.commandWorked(db.adminCommand({getParameter: 1, featureFlagDensify: 1}))
+        .featureFlagDensify.value;
+if (!featureEnabled) {
+    jsTestLog("Skipping test because the densify feature flag is disabled");
+    return;
+}
+
+load("jstests/libs/fixture_helpers.js");
+
+const coll = db[jsTestName()];
+coll.insert({});
+
+// Use .explain() to see what the stage desugars to.
+// The result is formatted as explain-output, which differs from MQL syntax in some cases:
+// for example {$sort: {a: 1}} explains as {$sort: {sortKey: {a: 1}}}.
+function desugar(stage) {
+    const result = coll.explain().aggregate([
+        // prevent stages from being absorbed into the .find() layer
+        {$_internalInhibitOptimization: {}},
+        stage,
+    ]);
+
+    assert.commandWorked(result);
+    // We proceed by cases based on topology.
+    if (!FixtureHelpers.isMongos(db)) {
+        assert(Array.isArray(result.stages), result);
+        // The first two stages should be the .find() cursor and the inhibit-optimization stage;
+        // the rest of the stages are what the user's 'stage' expanded to.
+        assert(result.stages[0].$cursor, result);
+        assert(result.stages[1].$_internalInhibitOptimization, result);
+        return result.stages.slice(2);
+    } else {
+        if (result.splitPipeline) {
+            assert(result.splitPipeline.shardsPart[0].$_internalInhibitOptimization, result);
+            assert.eq(result.splitPipeline.shardsPart.length, 2);
+            assert.eq(result.splitPipeline.mergerPart.length, 1);
+            return [result.splitPipeline.shardsPart[1], result.splitPipeline.mergerPart[0]];
+        } else if (result.stages) {
+            // Required for aggregation_mongos_passthrough.
+            assert(Array.isArray(result.stages), result);
+            // The first two stages should be the .find() cursor and the inhibit-optimization stage;
+            // the rest of the stages are what the user's 'stage' expanded to.
+            assert(result.stages[0].$cursor, result);
+            assert(result.stages[1].$_internalInhibitOptimization, result);
+            return result.stages.slice(2);
+        } else {
+            // Required for aggregation_one_shard_sharded_collections.
+            assert(Array.isArray(result.shards["shard-rs0"].stages), result);
+            assert(result.shards["shard-rs0"].stages[0].$cursor, result);
+            assert(result.shards["shard-rs0"].stages[1].$_internalInhibitOptimization, result);
+            return result.shards["shard-rs0"].stages.slice(2);
+        }
+    }
+}
+
+// Implicit partition fields and sort are generated.
+assert.eq(desugar({$densify: {field: "a", range: {step: 1.0, bounds: "full"}}}), [
+    {$sort: {sortKey: {a: 1}}},
+    {$_internalDensify: {field: "a", partitionByFields: [], range: {step: 1.0, bounds: "full"}}},
+]);
+
+// PartitionByFields are prepended to the sortKey if "partition" is specified.
+assert.eq(
+    desugar({
+        $densify:
+            {field: "a", partitionByFields: ["b", "c"], range: {step: 1.0, bounds: "partition"}}
+    }),
+    [
+        {$sort: {sortKey: {b: 1, c: 1, a: 1}}},
+        {
+            $_internalDensify:
+                {field: "a", partitionByFields: ["b", "c"], range: {step: 1.0, bounds: "partition"}}
+        },
+    ]);
+
+// PartitionByFields are not prepended to the sortKey if "full" is specified.
+assert.eq(
+    desugar({
+        $densify: {field: "a", partitionByFields: ["b", "c"], range: {step: 1.0, bounds: "full"}}
+    }),
+    [
+        {$sort: {sortKey: {a: 1}}},
+        {
+            $_internalDensify:
+                {field: "a", partitionByFields: ["b", "c"], range: {step: 1.0, bounds: "full"}}
+        },
+    ]);
+
+// PartitionByFields are prepended to the sortKey if numeric bounds are specified.
+assert.eq(
+    desugar({
+        $densify: {field: "a", partitionByFields: ["b", "c"], range: {step: 1.0, bounds: [-10, 0]}}
+    }),
+    [
+        {$sort: {sortKey: {b: 1, c: 1, a: 1}}},
+        {
+            $_internalDensify:
+                {field: "a", partitionByFields: ["b", "c"], range: {step: 1.0, bounds: [-10, 0]}}
+        },
+    ]);
+
+// PartitionByFields are prepended to the sortKey if date bounds are specified.
+assert.eq(
+    desugar({
+        $densify: {
+            field: "a",
+            partitionByFields: ["b", "c"],
+            range:
+                {step: 1.0, bounds: [new Date("2020-01-03"), new Date("2020-01-04")], unit: "day"}
+        }
+    }),
+    [
+        {$sort: {sortKey: {b: 1, c: 1, a: 1}}},
+        {
+            $_internalDensify: {
+                field: "a",
+                partitionByFields: ["b", "c"],
+                range: {
+                    step: 1.0,
+                    bounds: [new Date("2020-01-03"), new Date("2020-01-04")],
+                    unit: "day"
+                }
+            }
+        },
+    ]);
+})();
diff --git a/jstests/aggregation/sources/densify/internal_parse.js b/jstests/aggregation/sources/densify/internal_parse.js
new file mode 100644
index 00000000000..d837c765af4
--- /dev/null
+++ b/jstests/aggregation/sources/densify/internal_parse.js
@@ -0,0 +1,39 @@
+/**
+ * Test the syntax of $_internalDensify.
+ * @tags: [
+ *     # Needed as $densify is a 51 feature.
+ *     requires_fcv_51,
+ *     # Our method of connecting via an internal client requries an unsharded topology.
+ *     assumes_unsharded_collection,
+ *     assumes_against_mongod_not_mongos,
+ *     assumes_read_preference_unchanged,
+ * ]
+ */
+
+(function() {
+"use strict";
+
+load("jstests/aggregation/sources/densify/libs/parse_util.js");
+
+const dbName = jsTestName();
+const testDB = db.getSiblingDB(dbName);
+testDB.dropDatabase();
+const collName = jsTestName();
+
+const testInternalClient = (function createInternalClient() {
+    const connInternal = new Mongo(testDB.getMongo().host);
+    const curDB = connInternal.getDB(dbName);
+    assert.commandWorked(curDB.runCommand({
+        ["hello"]: 1,
+        internalClient: {minWireVersion: NumberInt(0), maxWireVersion: NumberInt(7)}
+    }));
+    return connInternal;
+})();
+
+const internalDB = testInternalClient.getDB(dbName);
+const internalColl = internalDB[collName];
+
+parseUtil(internalDB, internalColl, "$_internalDensify", {
+    writeConcern: {w: "majority"},
+});
+})();
diff --git a/jstests/aggregation/sources/densify/libs/parse_util.js b/jstests/aggregation/sources/densify/libs/parse_util.js
new file mode 100644
index 00000000000..9e51df53445
--- /dev/null
+++ b/jstests/aggregation/sources/densify/libs/parse_util.js
@@ -0,0 +1,210 @@
+/**
+ * Utility for testing the parsing of densify-like commands. I.e. $_internalDensify and $densify.
+ */
+
+let parseUtil = (function(db, coll, stageName, options = {}) {
+    function run(stage, extraCommandArgs = options) {
+        return coll.runCommand(Object.merge(
+            {aggregate: coll.getName(), pipeline: [stage], cursor: {}}, extraCommandArgs));
+    }
+
+    function runTest(stageName) {
+        const featureEnabled =
+            assert.commandWorked(db.adminCommand({getParameter: 1, featureFlagDensify: 1}))
+                .featureFlagDensify.value;
+        if (!featureEnabled) {
+            jsTestLog("Skipping test because the densify feature flag is disabled");
+            return;
+        }
+
+        // Required fields.
+        const kIDLRequiredFieldErrorCode = 40414;
+        assert.commandFailedWithCode(
+            run({[stageName]: {field: "a", range: {step: 1.0}}}),
+            kIDLRequiredFieldErrorCode,
+            "BSON field '$densify.range.bounds' is missing but a required field");
+        assert.commandFailedWithCode(run({[stageName]: {range: {step: 1.0, bounds: "full"}}}),
+                                     kIDLRequiredFieldErrorCode,
+                                     "BSON field '$densify.field' is missing but a required field");
+        assert.commandFailedWithCode(
+            run({[stageName]: {field: "a", range: {bounds: "full"}}}),
+            kIDLRequiredFieldErrorCode,
+            "BSON field '$densify.range.step' is missing but a required field");
+
+        // Wrong types
+        assert.commandFailedWithCode(
+            run({[stageName]: {field: 1.0, range: {step: 1.0, bounds: "full"}}}),
+            ErrorCodes.TypeMismatch,
+            "BSON field '$densify.field' is the wrong type 'double', expected type 'string'");
+        assert.commandFailedWithCode(
+            run({[stageName]: {field: "a", range: {step: "invalid", bounds: "full"}}}),
+            ErrorCodes.TypeMismatch,
+            "BSON field '$densify.range.step' is the wrong type 'string', expected types '[int, decimal, double, long']");
+        assert.commandFailedWithCode(
+            run({
+                [stageName]:
+                    {field: "a", partitionByFields: "incorrect", range: {step: 1.0, bounds: "full"}}
+            }),
+            ErrorCodes.TypeMismatch,
+            "BSON field '$densify.partitionByFields' is the wrong type 'string', expected type 'array'");
+        assert.commandFailedWithCode(
+            run({
+                [stageName]: {
+                    field: "a",
+                    range: {
+                        step: 1.0,
+                        bounds: [new Date("2020-01-01"), new Date("2020-01-02")],
+                        unit: 1000
+                    }
+                }
+            }),
+            ErrorCodes.TypeMismatch,
+            "BSON field '$densify.range.unit' is the wrong type 'double', expected type 'string'");
+
+        // Logical errors
+        // Too few bounds
+        assert.commandFailedWithCode(
+            run({
+                [stageName]: {
+                    field: "a",
+                    range: {step: 1.0, bounds: [new Date("2020-01-01")], unit: "second"}
+                }
+            }),
+            5733403,
+            "a bounding array in a range statement must have exactly two elements");
+        // Too many elements
+        assert.commandFailedWithCode(
+            run({[stageName]: {field: "a", range: {step: 1.0, bounds: [0, 1, 2]}}}),
+            5733403,
+            "a bounding array in a range statement must have exactly two elements");
+        // Negative step
+        assert.commandFailedWithCode(
+            run({[stageName]: {field: "a", range: {step: -1.0, bounds: [0, 1]}}}),
+            5733401,
+            "the step parameter in a range statement must be a strictly positive numeric value");
+        // Field path expression instead of field path
+        assert.commandFailedWithCode(
+            run({[stageName]: {field: "$a", range: {step: 1.0, bounds: [0, 1]}}}),
+            16410,
+            "FieldPath field names may not start with '$'. Consider using $getField or $setField.");
+        // Field path expression instead of field path
+        assert.commandFailedWithCode(
+            run({
+                [stageName]:
+                    {field: "a", partitionByFields: ["$b"], range: {step: 1.0, bounds: [0, 1]}}
+            }),
+            16410,
+            "FieldPath field names may not start with '$'. Consider using $getField or $setField.");
+        // Partition bounds but not partitionByFields
+        assert.commandFailedWithCode(
+            run({
+                [stageName]:
+                    {field: "a", partitionByFields: [], range: {step: 1.0, bounds: "partition"}}
+            }),
+            5733408,
+            "one may not specify the bounds as 'partition' without specifying a non-empty array of partitionByFields. You may have meant to specify 'full' bounds.");
+        assert.commandFailedWithCode(
+            run({[stageName]: {field: "a", range: {step: 1.0, bounds: "partition"}}}),
+            5733408,
+            "one may not specify the bounds as 'partition' without specifying a non-empty array of partitionByFields. You may have meant to specify 'full' bounds.");
+        // Left bound greater than right bound
+        assert.commandFailedWithCode(
+            run({[stageName]: {field: "a", range: {step: 1.0, bounds: [2, 1]}}}),
+            5733402,
+            "the bounds in a range statement must be the string 'full', 'partition', or an ascending array of two numbers or two dates");
+        assert.commandFailedWithCode(
+            run({
+                [stageName]: {
+                    field: "a",
+                    range: {
+                        step: 1.0,
+                        bounds: [new Date("2020-01-01"), new Date("2019-01-01")],
+                        unit: "second"
+                    }
+                }
+            }),
+            5733402,
+            "the bounds in a range statement must be the string 'full', 'partition', or an ascending array of two numbers or two dates");
+        // Unit with numeric bounds
+        assert.commandFailedWithCode(
+            run({[stageName]: {field: "a", range: {step: 1.0, bounds: [1, 2], unit: "second"}}}),
+            5733409,
+            "numeric bounds may not have unit parameter");
+        // Mixed numeric and date bounds
+        assert.commandFailedWithCode(
+            run({
+                [stageName]: {field: "a", range: {step: 1.0, bounds: [1, new Date("2020-01-01")]}}
+            }),
+            5733406,
+            "a bounding array must contain either both dates or both numeric types");
+        assert.commandFailedWithCode(
+            run({
+                [stageName]: {
+                    field: "a",
+                    range: {step: 1.0, bounds: [new Date("2020-01-01"), 1], unit: "second"}
+                }
+            }),
+            5733402,
+            "a bounding array must be an ascending array of either two dates or two numbers");
+
+        // Positive test cases
+        assert.commandWorked(run({[stageName]: {field: "a", range: {step: 1.0, bounds: [1, 2]}}}));
+        // TODO SERVER-57337: Enable this parsing test.
+        /* assert.commandWorked(run({
+         *     [stageName]: {
+         *         field: "a",
+         *         range: {
+         *             step: 1.0,
+         *             bounds: [new Date("2020-01-01"), new Date("2021-01-01")],
+         *             unit: "second"
+         *         }
+         *     }
+         * })); */
+        // TODO SERVER-57337 SERVER-57344: Enable this parsing test.
+        /* assert.commandWorked(run({
+         *     [stageName]: {
+         *         field: "a",
+         *         partitionByFields: ["b", "c"],
+         *         range: {
+         *             step: 1.0,
+         *             bounds: [new Date("2020-01-01"), new Date("2021-01-01")],
+         *             unit: "second"
+         *         }
+         *     }
+         * })); */
+        // TODO SERVER-57337 SERVER-57344: Enable this parsing test.
+        /* assert.commandWorked(run({
+         *     [stageName]: {
+         *         field: "a",
+         *         partitionByFields: ["b", "c"],
+         *         range: {step: 1.0, bounds: "partition", unit: "second"}
+         *     }
+         * })); */
+        // TODO SERVER-57337 SERVER-57344: Enable this parsing test.
+        /* assert.commandWorked(run({
+         *     [stageName]: {
+         *         field: "a",
+         *         partitionByFields: ["b", "c"],
+         *         range: {step: 1.0, bounds: "full", unit: "second"}
+         *     }
+         * })); */
+        // TODO SERVER-57344: Enable this parsing test.
+        /* assert.commandWorked(run({
+         *     [stageName]:
+         *         {field: "a", partitionByFields: ["b", "c"], range: {step: 1.0, bounds: "full"}}
+         * })); */
+        // TODO SERVER-57344: Enable this parsing test.
+        /* assert.commandWorked(run({
+            [stageName]: {
+                field: "a",
+                partitionByFields: [
+                    "b",
+                ],
+                range: { step: 1.0, bounds: "partition" }
+            }
+           })); */
+        assert.commandWorked(run({[stageName]: {field: "a", range: {step: 1.0, bounds: "full"}}}));
+    }
+
+    runTest(stageName);
+});
diff --git a/jstests/aggregation/sources/densify/parse.js b/jstests/aggregation/sources/densify/parse.js
new file mode 100644
index 00000000000..58154ed47c7
--- /dev/null
+++ b/jstests/aggregation/sources/densify/parse.js
@@ -0,0 +1,18 @@
+/**
+ * Test the syntax of $densify.
+ * @tags: [
+ *   # Needed as $densify is a 51 feature.
+ *   requires_fcv_51,
+ * ]
+ */
+
+(function() {
+"use strict";
+
+load("jstests/aggregation/sources/densify/libs/parse_util.js");
+
+const coll = db.densify_parse;
+coll.drop();
+
+parseUtil(db, coll, "$densify");
+})();