SERVER-119723 Allow a jstest to access the TPC-H dataset (#48158)

GitOrigin-RevId: 16528783332a63273179fbb602936be8b57dca18

etc/evergreen_yml_components/definitions.yml

@@ -3689,3 +3689,22 @@ functions:
   "save tracing data":
     - *tar_tracing_data
     - *archive_tracing_data
+
+  "fetch tpch dataset":
+    - *f_expansions_write
+    - command: s3.get
+      params:
+        role_arn: arn:aws:iam::579766882180:role/evergreen.mongo-db-master
+        bucket: query-benchmark-data
+        region: us-east-1
+        remote_file: tpc-h/tpch-${scale}-normalized.archive.gz
+        local_file: tpc-h/tpch-${scale}-normalized.archive.gz
+        require_checksum_sha256: ${checksum}
+
+  "fetch mongodb database tools":
+    - *f_expansions_write
+    - command: subprocess.exec
+      params:
+        binary: bash
+        args:
+          - "./src/evergreen/fetch_mongodb_database_tools.sh"

evergreen/BUILD.bazel

@@ -159,6 +159,11 @@ sh_binary(
     srcs = ["feature_flag_tags_check.sh"],
 )
 
+sh_binary(
+    name = "fetch_mongodb_database_tools",
+    srcs = ["fetch_mongodb_database_tools.sh"],
+)
+
 sh_binary(
     name = "garasign_gpg_crypt_sign",
     srcs = ["garasign_gpg_crypt_sign.sh"],

evergreen/OWNERS.yml

@@ -60,3 +60,6 @@ filters:
   - "verify_all_extensions_visibility.sh":
     approvers:
       - 10gen/query-integration-extensions-api
+  - "fetch_mongodb_database_tools.sh":
+    approvers:
+      - 10gen/query-optimization-correctness

evergreen/fetch_mongodb_database_tools.sh

@@ -0,0 +1,17 @@
+#
+# Download the MongoDB Database tools so that they are available for use in jstests
+#
+
+set -ex
+
+DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
+source "$DIR/functions/get_mongodb_tools_url.sh"
+
+mkdir -p mongodb_database_tools
+pushd mongodb_database_tools
+
+database_tools_url="$(get_mongodb_tools_url 100.14.1)" || exit 1
+# Place the tools under mongodb_database_tools/bin in the root evergreen directory
+curl ${database_tools_url} | tar xvz --strip-components=1
+
+popd

evergreen/functions/BUILD.bazel

@@ -47,6 +47,11 @@ sh_binary(
     srcs = ["get_mongot_version.sh"],
 )
 
+sh_binary(
+    name = "get_mongodb_tools_url",
+    srcs = ["get_mongodb_tools_url.sh"],
+)
+
 sh_binary(
     name = "modified_patch_files_get_all",
     srcs = ["modified_patch_files_get_all.sh"],

evergreen/functions/get_mongodb_tools_url.sh

@@ -0,0 +1,36 @@
+get_mongodb_tools_url() {
+    local arch=$(uname -m)
+    local mongodb_tools_version="$1"
+    local database_tools_url
+
+    if [ -f /etc/os-release ]; then
+        . /etc/os-release
+        if [ "$ID" == "amzn" ]; then
+            case $arch in
+            "x86_64" | "aarch64")
+                case $VERSION_ID in
+                "2" | "2023")
+                    database_tools_url="https://fastdl.mongodb.org/tools/db/mongodb-database-tools-amazon${VERSION_ID}-${arch}-${mongodb_tools_version}.tgz"
+                    ;;
+                *)
+                    echo "Unsupported Amazon Linux version: $VERSION_ID"
+                    return 1
+                    ;;
+                esac
+                ;;
+            *)
+                echo "Unsupported architecture: $arch"
+                return 1
+                ;;
+            esac
+        else
+            echo "Unsupported Linux distribution: $ID"
+            return 1
+        fi
+    else
+        echo "Unable to determine Linux distribution"
+        return 1
+    fi
+
+    echo "$database_tools_url"
+}

evergreen/gen_supplementary_data.sh

@@ -1,5 +1,6 @@
 DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
 . "$DIR/prelude.sh"
+source "$DIR/functions/get_mongodb_tools_url.sh"
 
 cd src
 
@@ -21,37 +22,7 @@ if [ -z "${build_patch_id}" ] || [ -z "${reuse_compile_from}" ] || [ "${is_patch
     mkdir -p mongodb/server_params
     cp ./all_server_params.txt mongodb/server_params
 
-    # Download mongo tools
-    arch=$(uname -m)
-    if [ -f /etc/os-release ]; then
-        . /etc/os-release
-        if [ "$ID" == "amzn" ]; then
-            case $arch in
-            "x86_64" | "aarch64")
-                case $VERSION_ID in
-                "2" | "2023")
-                    binary_url="https://fastdl.mongodb.org/tools/db/mongodb-database-tools-amazon${VERSION_ID}-${arch}-100.9.4.tgz"
-                    ;;
-                *)
-                    echo "Unsupported Amazon Linux version: $VERSION_ID"
-                    exit 1
-                    ;;
-                esac
-                ;;
-            *)
-                echo "Unsupported architecture: $arch"
-                exit 1
-                ;;
-            esac
-        else
-            echo "Unsupported Linux distribution: $ID"
-            exit 1
-        fi
-    else
-        echo "Unable to determine Linux distribution"
-        exit 1
-    fi
-
+    binary_url="$(get_mongodb_tools_url 100.9.4)" || exit 1
     wget "$binary_url" -O mongo-tools.tar.gz
     tar -xzvf mongo-tools.tar.gz -C mongodb/ --strip-components=1 "mong*/bin"
 

jstests/libs/OWNERS.yml

@@ -111,3 +111,6 @@ filters:
   - "cluster_server_parameter_utils.js":
     approvers:
       - 10gen/server-catalog-and-routing-routing-and-topology
+  - "mongodb_database_tools.js":
+    approvers:
+      - 10gen/query-optimization-correctness

jstests/libs/mongodb_database_tools.js

@@ -0,0 +1,56 @@
+/**
+ * Allows the execution of the MongoDB Database Tools from within a jstest.
+ *
+ * The "fetch database tools" evergreen command makes those tools available
+ * in the environment.
+ *
+ * @class
+ */
+export class Mongorestore {
+    constructor() {
+        this.uri = "mongodb://" + db.getMongo().host;
+    }
+
+    execute({
+        archive,
+        nsFrom = undefined,
+        nsTo = undefined,
+        drop = true,
+        maintainInsertionOrder = true,
+        gzip = true,
+    } = {}) {
+        if (archive === undefined) {
+            throw new Error("Archive must be provided to Mongorestore.execute()");
+        }
+
+        let args = [
+            TestData.inEvergreen ? "../mongodb_database_tools/bin/mongorestore" : "mongorestore",
+            "--uri",
+            this.uri,
+        ];
+
+        if (nsFrom) {
+            args.push(`--nsFrom=${nsFrom}`);
+        }
+
+        if (nsTo) {
+            args.push(`--nsTo=${nsTo}`);
+        }
+
+        if (maintainInsertionOrder) {
+            args.push("--maintainInsertionOrder");
+        }
+
+        if (gzip) {
+            args.push("--gzip");
+        }
+
+        if (drop) {
+            args.push("--drop");
+        }
+
+        args.push(`--archive=${archive}`);
+
+        assert.eq(runNonMongoProgram(...args), 0);
+    }
+}