SERVER-108487 Integrate the data generator with jstests (#39510)

GitOrigin-RevId: 4feb3238cfe3d6672b5c51f57bf642af60519b28
2025-08-05 16:06:08 +02:00 · 2025-08-05 16:06:08 +02:00 · 8e69410fa0
commit 8e69410fa0
parent 189bdb2dfe
8 changed files with 224 additions and 2 deletions
--- a/etc/evergreen_yml_components/tasks/compile_tasks_shared.yml
+++ b/etc/evergreen_yml_components/tasks/compile_tasks_shared.yml
@ -214,6 +214,7 @@ tasks:
            - "src/src/mongo/db/modules/enterprise/docs/**"
            - "src/src/mongo/db/modules/enterprise/jstests/**"
            - "src/src/mongo/db/modules/subscription/jstests/**"
+            - "src/src/mongo/db/query/benchmark/data_generator/**"
            - "src/src/mongo/db/query/query_tester/tests/**"
            - "src/src/mongo/util/options_parser/test_config_files/**"
            - "src/src/mongo/util/version/**"
--- a/jstests/libs/query/data_generator.js
+++ b/jstests/libs/query/data_generator.js
@ -0,0 +1,78 @@
+
+import {getPython3Binary} from "jstests/libs/python.js";
+import {extractUUIDFromObject} from "jstests/libs/uuid_util.js";
+
+/**
+ * Allows the execution of the python data generator from
+ * src/mongo/db/query/benchmark/data_generator from within a MongoDB jstest.
+ *
+ * For more information on the generator itself, see
+ * src/mongo/db/query/benchmark/data_generator/README.md
+ *
+ * For an example of how to use this class, see
+ * jstests/noPassthroughWithMongod/query/data_generator_integration.js
+ * @class
+ */
+export class DataGenerator {
+    static PROGRAM_PATH = 'src/mongo/db/query/benchmark/data_generator/driver.py';
+    constructor({
+        db = null,
+        module = null,
+        seed = null,
+    } = {}) {
+        if (db == null) {
+            throw new Error('A db object must be provided to the DataGenerator constructor.');
+        } else {
+            this.dbName = db.getName();
+            this.uri = 'mongodb://' + db.getMongo().host;
+        }
+
+        const tmpDir = (_getEnv("TMPDIR") || _getEnv("TMP_DIR") || "/tmp");
+        const randomUUID = extractUUIDFromObject(UUID());
+        this.out = tmpDir + '/data_generator_' + randomUUID;
+
+        if (module == null) {
+            throw new Error(
+                'A data generator module name must be provided to DataGenerator constructor.');
+        } else {
+            this.module = module;
+        }
+
+        this.seed = seed;
+    }
+
+    execute({spec = null, size = null, indices = null} = {}) {
+        let args = [
+            getPython3Binary(),
+            DataGenerator.PROGRAM_PATH,
+            '--uri',
+            this.uri,
+            '--db',
+            this.dbName,
+            '--out',
+            this.out
+        ];
+
+        if (spec == null || size == null) {
+            throw new Error('Both specs and size must be provided to DataGenerator.execute()');
+        }
+
+        args.push("--size", size);
+
+        if (indices !== null) {
+            args.push('--indices', indices);
+        }
+
+        if (this.seed !== null) {
+            args.push('--seed', this.seed);
+        }
+
+        args.push(this.module, spec);
+
+        assert.eq(runNonMongoProgram(...args), 0);
+    }
+
+    cleanup() {
+        removeFile(this.out);
+    }
+}
--- a/jstests/noPassthroughWithMongod/query/data_generator_integration.js
+++ b/jstests/noPassthroughWithMongod/query/data_generator_integration.js
@ -0,0 +1,21 @@
+/**
+ * A simple test for the DataGenerator class shows how to use it to call
+ * the data generator Python script located at src/mongo/db/query/benchmark/data_generator
+ *
+ * See src/mongo/db/query/benchmark/data_generator/README.md for more information.
+ */
+
+import {DataGenerator} from "jstests/libs/query/data_generator.js";
+
+const dg = new DataGenerator({db: db, module: 'specs.test', seed: 1});
+try {
+    const collName = 'Test';
+    const size = 10;
+
+    dg.execute({spec: collName, size: size, indices: 'test_index'});
+
+    assert.eq(db[collName].find({i: {$exists: true}}).count(), size);
+    assert.eq(db[collName].getIndexes().length, 2);  // _id and i_idx test index
+} finally {
+    dg.cleanup();
+}
--- a/poetry.lock
+++ b/poetry.lock
@ -2314,6 +2314,72 @@ files = [
    {file = "nodeenv-1.9.1.tar.gz", hash = "sha256:6ec12890a2dab7946721edbfbcd91f3319c6ccc9aec47be7c7e6b7011ee6645f"},
 ]

+[[package]]
+name = "numpy"
+version = "2.2.6"
+description = "Fundamental package for array computing in Python"
+optional = false
+python-versions = ">=3.10"
+groups = ["main"]
+markers = "platform_machine != \"s390x\" and platform_machine != \"ppc64le\" or platform_machine == \"s390x\" or platform_machine == \"ppc64le\""
+files = [
+    {file = "numpy-2.2.6-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b412caa66f72040e6d268491a59f2c43bf03eb6c96dd8f0307829feb7fa2b6fb"},
+    {file = "numpy-2.2.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8e41fd67c52b86603a91c1a505ebaef50b3314de0213461c7a6e99c9a3beff90"},
+    {file = "numpy-2.2.6-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:37e990a01ae6ec7fe7fa1c26c55ecb672dd98b19c3d0e1d1f326fa13cb38d163"},
+    {file = "numpy-2.2.6-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:5a6429d4be8ca66d889b7cf70f536a397dc45ba6faeb5f8c5427935d9592e9cf"},
+    {file = "numpy-2.2.6-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:efd28d4e9cd7d7a8d39074a4d44c63eda73401580c5c76acda2ce969e0a38e83"},
+    {file = "numpy-2.2.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc7b73d02efb0e18c000e9ad8b83480dfcd5dfd11065997ed4c6747470ae8915"},
+    {file = "numpy-2.2.6-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:74d4531beb257d2c3f4b261bfb0fc09e0f9ebb8842d82a7b4209415896adc680"},
+    {file = "numpy-2.2.6-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8fc377d995680230e83241d8a96def29f204b5782f371c532579b4f20607a289"},
+    {file = "numpy-2.2.6-cp310-cp310-win32.whl", hash = "sha256:b093dd74e50a8cba3e873868d9e93a85b78e0daf2e98c6797566ad8044e8363d"},
+    {file = "numpy-2.2.6-cp310-cp310-win_amd64.whl", hash = "sha256:f0fd6321b839904e15c46e0d257fdd101dd7f530fe03fd6359c1ea63738703f3"},
+    {file = "numpy-2.2.6-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f9f1adb22318e121c5c69a09142811a201ef17ab257a1e66ca3025065b7f53ae"},
+    {file = "numpy-2.2.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c820a93b0255bc360f53eca31a0e676fd1101f673dda8da93454a12e23fc5f7a"},
+    {file = "numpy-2.2.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:3d70692235e759f260c3d837193090014aebdf026dfd167834bcba43e30c2a42"},
+    {file = "numpy-2.2.6-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:481b49095335f8eed42e39e8041327c05b0f6f4780488f61286ed3c01368d491"},
+    {file = "numpy-2.2.6-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b64d8d4d17135e00c8e346e0a738deb17e754230d7e0810ac5012750bbd85a5a"},
+    {file = "numpy-2.2.6-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba10f8411898fc418a521833e014a77d3ca01c15b0c6cdcce6a0d2897e6dbbdf"},
+    {file = "numpy-2.2.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:bd48227a919f1bafbdda0583705e547892342c26fb127219d60a5c36882609d1"},
+    {file = "numpy-2.2.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:9551a499bf125c1d4f9e250377c1ee2eddd02e01eac6644c080162c0c51778ab"},
+    {file = "numpy-2.2.6-cp311-cp311-win32.whl", hash = "sha256:0678000bb9ac1475cd454c6b8c799206af8107e310843532b04d49649c717a47"},
+    {file = "numpy-2.2.6-cp311-cp311-win_amd64.whl", hash = "sha256:e8213002e427c69c45a52bbd94163084025f533a55a59d6f9c5b820774ef3303"},
+    {file = "numpy-2.2.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:41c5a21f4a04fa86436124d388f6ed60a9343a6f767fced1a8a71c3fbca038ff"},
+    {file = "numpy-2.2.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:de749064336d37e340f640b05f24e9e3dd678c57318c7289d222a8a2f543e90c"},
+    {file = "numpy-2.2.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:894b3a42502226a1cac872f840030665f33326fc3dac8e57c607905773cdcde3"},
+    {file = "numpy-2.2.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:71594f7c51a18e728451bb50cc60a3ce4e6538822731b2933209a1f3614e9282"},
+    {file = "numpy-2.2.6-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2618db89be1b4e05f7a1a847a9c1c0abd63e63a1607d892dd54668dd92faf87"},
+    {file = "numpy-2.2.6-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd83c01228a688733f1ded5201c678f0c53ecc1006ffbc404db9f7a899ac6249"},
+    {file = "numpy-2.2.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:37c0ca431f82cd5fa716eca9506aefcabc247fb27ba69c5062a6d3ade8cf8f49"},
+    {file = "numpy-2.2.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fe27749d33bb772c80dcd84ae7e8df2adc920ae8297400dabec45f0dedb3f6de"},
+    {file = "numpy-2.2.6-cp312-cp312-win32.whl", hash = "sha256:4eeaae00d789f66c7a25ac5f34b71a7035bb474e679f410e5e1a94deb24cf2d4"},
+    {file = "numpy-2.2.6-cp312-cp312-win_amd64.whl", hash = "sha256:c1f9540be57940698ed329904db803cf7a402f3fc200bfe599334c9bd84a40b2"},
+    {file = "numpy-2.2.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0811bb762109d9708cca4d0b13c4f67146e3c3b7cf8d34018c722adb2d957c84"},
+    {file = "numpy-2.2.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:287cc3162b6f01463ccd86be154f284d0893d2b3ed7292439ea97eafa8170e0b"},
+    {file = "numpy-2.2.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:f1372f041402e37e5e633e586f62aa53de2eac8d98cbfb822806ce4bbefcb74d"},
+    {file = "numpy-2.2.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:55a4d33fa519660d69614a9fad433be87e5252f4b03850642f88993f7b2ca566"},
+    {file = "numpy-2.2.6-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f92729c95468a2f4f15e9bb94c432a9229d0d50de67304399627a943201baa2f"},
+    {file = "numpy-2.2.6-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1bc23a79bfabc5d056d106f9befb8d50c31ced2fbc70eedb8155aec74a45798f"},
+    {file = "numpy-2.2.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e3143e4451880bed956e706a3220b4e5cf6172ef05fcc397f6f36a550b1dd868"},
+    {file = "numpy-2.2.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b4f13750ce79751586ae2eb824ba7e1e8dba64784086c98cdbbcc6a42112ce0d"},
+    {file = "numpy-2.2.6-cp313-cp313-win32.whl", hash = "sha256:5beb72339d9d4fa36522fc63802f469b13cdbe4fdab4a288f0c441b74272ebfd"},
+    {file = "numpy-2.2.6-cp313-cp313-win_amd64.whl", hash = "sha256:b0544343a702fa80c95ad5d3d608ea3599dd54d4632df855e4c8d24eb6ecfa1c"},
+    {file = "numpy-2.2.6-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0bca768cd85ae743b2affdc762d617eddf3bcf8724435498a1e80132d04879e6"},
+    {file = "numpy-2.2.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:fc0c5673685c508a142ca65209b4e79ed6740a4ed6b2267dbba90f34b0b3cfda"},
+    {file = "numpy-2.2.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:5bd4fc3ac8926b3819797a7c0e2631eb889b4118a9898c84f585a54d475b7e40"},
+    {file = "numpy-2.2.6-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:fee4236c876c4e8369388054d02d0e9bb84821feb1a64dd59e137e6511a551f8"},
+    {file = "numpy-2.2.6-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e1dda9c7e08dc141e0247a5b8f49cf05984955246a327d4c48bda16821947b2f"},
+    {file = "numpy-2.2.6-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f447e6acb680fd307f40d3da4852208af94afdfab89cf850986c3ca00562f4fa"},
+    {file = "numpy-2.2.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:389d771b1623ec92636b0786bc4ae56abafad4a4c513d36a55dce14bd9ce8571"},
+    {file = "numpy-2.2.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8e9ace4a37db23421249ed236fdcdd457d671e25146786dfc96835cd951aa7c1"},
+    {file = "numpy-2.2.6-cp313-cp313t-win32.whl", hash = "sha256:038613e9fb8c72b0a41f025a7e4c3f0b7a1b5d768ece4796b674c8f3fe13efff"},
+    {file = "numpy-2.2.6-cp313-cp313t-win_amd64.whl", hash = "sha256:6031dd6dfecc0cf9f668681a37648373bddd6421fff6c66ec1624eed0180ee06"},
+    {file = "numpy-2.2.6-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:0b605b275d7bd0c640cad4e5d30fa701a8d59302e127e5f79138ad62762c3e3d"},
+    {file = "numpy-2.2.6-pp310-pypy310_pp73-macosx_14_0_x86_64.whl", hash = "sha256:7befc596a7dc9da8a337f79802ee8adb30a552a94f792b9c9d18c840055907db"},
+    {file = "numpy-2.2.6-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce47521a4754c8f4593837384bd3424880629f718d87c5d44f8ed763edd63543"},
+    {file = "numpy-2.2.6-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:d042d24c90c41b54fd506da306759e06e568864df8ec17ccc17e9e884634fd00"},
+    {file = "numpy-2.2.6.tar.gz", hash = "sha256:e29554e2bef54a90aa5cc07da6ce955accb83f21ab5de01a62c8478897b264fd"},
+]
+
 [[package]]
 name = "oauthlib"
 version = "3.2.2"
@ -5497,4 +5563,4 @@ libdeps = ["cxxfilt", "eventlet", "flask", "flask-cors", "gevent", "lxml", "prog
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.10,<4.0"
-content-hash = "6d093a5822e5564b90067b272d6fd8b054d44b4071565bbd811d7e12d6e342ca"
+content-hash = "35359c99056176eb8c0e366f8e66a05bc2d47b6aa8be7c328babe0a01c943925"
--- a/pyproject.toml
+++ b/pyproject.toml
@ -29,6 +29,7 @@ dependencies = [
    "pympler",
    "faker (==37.1.0)",
    "pymongo (==4.12.0)",
+    "numpy (==2.2.6)",
 ]

 [tool.poetry.dependencies]
--- a/src/mongo/db/query/benchmark/data_generator/driver.py
+++ b/src/mongo/db/query/benchmark/data_generator/driver.py
@ -252,7 +252,7 @@ async def main():
                if hasattr(module, index_set_name):
                    index_set = getattr(module, index_set_name)
                    indices = index_set() if callable(index_set) else index_set
-                    database_instance.database.get_collection(collection_name).create_indexes(
+                    await database_instance.database.get_collection(collection_name).create_indexes(
                        indices
                    )
                else:
--- a/src/mongo/db/query/benchmark/data_generator/specs/BUILD.bazel
+++ b/src/mongo/db/query/benchmark/data_generator/specs/BUILD.bazel
@ -9,3 +9,8 @@ py_library(
    name = "employee",
    srcs = ["employee.py"],
 )
+
+py_library(
+    name = "test",
+    srcs = ["test.py"],
+)
--- a/src/mongo/db/query/benchmark/data_generator/specs/test.py
+++ b/src/mongo/db/query/benchmark/data_generator/specs/test.py
@ -0,0 +1,50 @@
+# Copyright (C) 2025-present MongoDB, Inc.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the Server Side Public License, version 1,
+# as published by MongoDB, Inc.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# Server Side Public License for more details.
+#
+# You should have received a copy of the Server Side Public License
+# along with this program. If not, see
+# <http://www.mongodb.com/licensing/server-side-public-license>.
+#
+# As a special exception, the copyright holders give permission to link the
+# code of portions of this program with the OpenSSL library under certain
+# conditions as described in each individual source file and distribute
+# linked combinations including the program with the OpenSSL library. You
+# must comply with the Server Side Public License in all respects for
+# all of the code used other than as permitted herein. If you modify file(s)
+# with this exception, you may extend this exception to your version of the
+# file(s), but you are not obligated to do so. If you do not wish to do so,
+# delete this exception statement from your version. If you delete this
+# exception statement from all source files in the program, then also delete
+# it in the license file.
+
+"""
+Example collection spec used for smoke tests of the generator itself.
+"""
+
+import dataclasses
+
+import faker
+import pymongo
+from datagen.util import Specification
+
+
+@dataclasses.dataclass
+class Test:
+    i: Specification(int)
+
+    @staticmethod
+    def compute_il(fkr: faker.proxy.Faker) -> float:
+        return 1
+
+def test_index() -> list[pymongo.IndexModel]:
+    return [
+        pymongo.IndexModel(keys="i", name="i_idx"),
+    ]