SERVER-110342 Implement query shape and query stats store key for replacement update

Co-authored-by: Denis Grebennicov <denis.grebennicov@mongodb.com>
GitOrigin-RevId: caa63bb62eef65ddf375d75d90b463da4b10f32f

src/mongo/db/local_catalog/shard_role_api/shard_role.h

@@ -215,6 +215,16 @@ public:
 
     const NamespaceString& nss() const;
 
+    query_shape::CollectionType getCollectionType() const {
+        if (!exists()) {
+            return query_shape::CollectionType::kNonExistent;
+        }
+        if (getCollectionPtr()->isNewTimeseriesWithoutView()) {
+            return query_shape::CollectionType::kTimeseries;
+        }
+        return query_shape::CollectionType::kCollection;
+    }
+
     /**
      * Returns whether the acquisition found a collection or the collection didn't exist.
      */
@@ -332,14 +342,7 @@ public:
                 return query_shape::CollectionType::kTimeseries;
             return query_shape::CollectionType::kView;
         }
-        const auto& collection = getCollection();
+        return getCollection().getCollectionType();
-        if (!collection.exists()) {
-            return query_shape::CollectionType::kNonExistent;
-        }
-        if (collection.getCollectionPtr()->isNewTimeseriesWithoutView()) {
-            return query_shape::CollectionType::kTimeseries;
-        }
-        return query_shape::CollectionType::kCollection;
     }
 
     bool collectionExists() const {

src/mongo/db/query/query_shape/BUILD.bazel

@@ -61,6 +61,24 @@ mongo_cc_library(
     ],
 )
 
+mongo_cc_library(
+    name = "update_cmd_shape",
+    srcs = [
+        "update_cmd_shape.cpp",
+        "//src/mongo/db/query:count_request.h",
+        "//src/mongo/db/query/compiler/logical_model/projection:projection_ast_util.h",
+    ],
+    hdrs = [
+        "let_shape_component.h",
+        "update_cmd_shape.h",
+    ],
+    deps = [
+        ":query_shape_common",
+        "//src/mongo:base",
+        "//src/mongo/db/query/write_ops:parsed_update_and_delete",
+    ],
+)
+
 idl_generator(
     name = "query_shape_hash_gen",
     src = "query_shape_hash.idl",
@@ -121,11 +139,13 @@ mongo_cc_unit_test(
         "find_cmd_shape_test.cpp",
         "let_shape_component_test.cpp",
         "serialization_options_test.cpp",
+        "update_cmd_shape_test.cpp",
     ],
     tags = [
         "mongo_unittest_sixth_group",
     ],
     deps = [
+        ":update_cmd_shape",
         "//src/mongo/db/pipeline:expression_context_for_test",
     ],
 )

src/mongo/db/query/query_shape/README.md

@@ -43,6 +43,7 @@ The structure is as follows:
   - [`DistinctCmdShapeComponents`](distinct_cmd_shape.h)
   - [`FindCmdShapeComponents`](find_cmd_shape.h)
   - [`LetShapeComponent`](let_shape_component.h)
+  - [`UpdateShapeComponent`](update_cmd_shape.h)
 
 See more information for the different shapes in their respective classes, structured as follows:
 
@@ -51,6 +52,7 @@ See more information for the different shapes in their respective classes, struc
   - [`CountCmdShape`](count_cmd_shape.h)
   - [`DistinctCmdShape`](distinct_cmd_shape.h)
   - [`FindCmdShape`](find_cmd_shape.h)
+  - [`UpdateCmdShape`](update_cmd_shape.h)
 
 ## Serialization Options
 

src/mongo/db/query/query_shape/update_cmd_shape.cpp

@@ -0,0 +1,194 @@
+/**
+ *    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.
+ */
+
+#include "mongo/db/query/query_shape/update_cmd_shape.h"
+
+#include "mongo/bson/bsonelement.h"
+#include "mongo/bson/bsonobj.h"
+#include "mongo/bson/bsonobjbuilder.h"
+#include "mongo/db/pipeline/expression_context_builder.h"
+#include "mongo/db/query/query_shape/let_shape_component.h"
+#include "mongo/db/query/query_shape/serialization_options.h"
+#include "mongo/db/query/query_shape/shape_helpers.h"
+#include "mongo/db/query/write_ops/update_request.h"
+#include "mongo/db/query/write_ops/write_ops_parsers.h"
+
+#include <absl/hash/hash.h>
+#include <boost/smart_ptr/intrusive_ptr.hpp>
+
+namespace mongo::query_shape {
+namespace {
+
+BSONObj shapifyQuery(const ParsedUpdate& parsedUpdate, const SerializationOptions& opts) {
+    tassert(11034203, "query is required to be parsed", parsedUpdate.hasParsedQuery());
+    auto cq = parsedUpdate.getParsedQuery();
+    auto matchExpr = cq->getPrimaryMatchExpression();
+    return matchExpr ? matchExpr->serialize(opts) : BSONObj{};
+}
+
+Value shapifyUpdateOp(const write_ops::UpdateModification& modification,
+                      const SerializationOptions& opts =
+                          SerializationOptions::kRepresentativeQueryShapeSerializeOptions) {
+    tassert(11034201,
+            "Unsupported type of update modification",
+            modification.type() == write_ops::UpdateModification::Type::kReplacement);
+
+    if (modification.type() == write_ops::UpdateModification::Type::kReplacement) {
+        return opts.serializeLiteral(modification.getUpdateReplacement());
+    }
+    return {};
+}
+
+}  // namespace
+
+UpdateCmdShapeComponents::UpdateCmdShapeComponents(const ParsedUpdate& parsedUpdate,
+                                                   LetShapeComponent let,
+                                                   const SerializationOptions& opts)
+    : representativeQ(shapifyQuery(parsedUpdate, opts)),
+      _representativeUObj(
+          shapifyUpdateOp(parsedUpdate.getRequest()->getUpdateModification(), opts).wrap(""_sd)),
+      multi(parsedUpdate.getRequest()->getMulti()),
+      upsert(parsedUpdate.getRequest()->isUpsert()),
+      let(let) {
+    // TODO(SERVER-110343): Suppport storing 'representativeC' when shapifying pipeline udpates.
+    // TODO(SERVER-110344): Support representativeArrayFilters when shapifying update modifiers.
+}
+
+void UpdateCmdShapeComponents::HashValue(absl::HashState state) const {
+    state = absl::HashState::combine(
+        std::move(state), simpleHash(representativeQ), simpleHash(_representativeUObj));
+    state = absl::HashState::combine(
+        std::move(state), representativeC.has_value(), representativeArrayFilters.has_value());
+    if (representativeC) {
+        // TODO(SERVER-110343): Revisit here when supporting storing 'representativeC' when
+        // shapifying pipeline udpates.
+        state = absl::HashState::combine(std::move(state), simpleHash(*representativeC));
+    }
+    if (representativeArrayFilters) {
+        // TODO(SERVER-110344): Revisit here when supporting representativeArrayFilters when
+        // shapifying update modifiers.
+        for (const auto& filter : *representativeArrayFilters) {
+            state = absl::HashState::combine(std::move(state), simpleHash(filter));
+        }
+    }
+    state = absl::HashState::combine(std::move(state), multi, upsert, let);
+}
+
+void UpdateCmdShapeComponents::appendTo(
+    BSONObjBuilder& bob,
+    const SerializationOptions& opts,
+    const boost::intrusive_ptr<ExpressionContext>& expCtx) const {
+    bob.append("command", "update");
+
+    bob.append(write_ops::UpdateOpEntry::kQFieldName, representativeQ);
+    bob.appendAs(getRepresentativeU(), write_ops::UpdateOpEntry::kUFieldName);
+
+    if (representativeC.has_value()) {
+        bob.append(write_ops::UpdateOpEntry::kCFieldName, *representativeC);
+    }
+    if (representativeArrayFilters.has_value()) {
+        bob.append(write_ops::UpdateOpEntry::kArrayFiltersFieldName, *representativeArrayFilters);
+    }
+
+    bob.append(write_ops::UpdateOpEntry::kMultiFieldName, bool(multi));
+    bob.append(write_ops::UpdateOpEntry::kUpsertFieldName, bool(upsert));
+
+    let.appendTo(bob, opts, expCtx);
+}
+
+// As part of the size, we must track the allocation of elements in the representative shapes.
+size_t UpdateCmdShapeComponents::size() const {
+    return sizeof(UpdateCmdShapeComponents) + representativeQ.objsize() +
+        _representativeUObj.objsize() + (representativeC ? representativeC->objsize() : 0) +
+        (representativeArrayFilters ? shape_helpers::containerSize(*representativeArrayFilters)
+                                    : 0) +
+        let.size() - sizeof(LetShapeComponent);
+}
+
+UpdateCmdShape::UpdateCmdShape(const write_ops::UpdateCommandRequest& updateCommand,
+                               const ParsedUpdate& parsedUpdate,
+                               const boost::intrusive_ptr<ExpressionContext>& expCtx)
+    : Shape(updateCommand.getNamespace(), parsedUpdate.getRequest()->getCollation()),
+      _components(parsedUpdate, LetShapeComponent(updateCommand.getLet(), expCtx)) {}
+
+const CmdSpecificShapeComponents& UpdateCmdShape::specificComponents() const {
+    return _components;
+}
+
+size_t UpdateCmdShape::extraSize() const {
+    // To account for possible padding, we calculate the extra space with the difference instead of
+    // using sizeof(bool);
+    return sizeof(UpdateCmdShape) - sizeof(Shape) - sizeof(UpdateCmdShapeComponents);
+}
+
+void UpdateCmdShape::appendCmdSpecificShapeComponents(BSONObjBuilder& bob,
+                                                      OperationContext* opCtx,
+                                                      const SerializationOptions& opts) const {
+    tassert(11034200,
+            "We don't support serializing to the unmodified shape here, since we have already "
+            "shapified and stored the representative query - we've lost the original literals",
+            !opts.isKeepingLiteralsUnchanged());
+
+    auto expCtx = makeBlankExpressionContext(opCtx, nssOrUUID, _components.let.shapifiedLet);
+    if (opts == SerializationOptions::kRepresentativeQueryShapeSerializeOptions) {
+        // We have this copy stored already!
+        _components.appendTo(bob, opts, expCtx);
+        return;
+    }
+
+    // Slow path: we need to re-parse from our representative shapes and re-shapify with 'opts'.
+
+    // Prepare UpdateOpEntry and UpdateRequest to reconstruct ParsedUpdate.
+    write_ops::UpdateOpEntry op;
+    op.setQ(_components.representativeQ);
+    op.setU(_components.getRepresentativeU());
+    op.setC(_components.representativeC);
+    op.setArrayFilters(_components.representativeArrayFilters);
+    op.setMulti(_components.multi);
+    op.setUpsert(_components.upsert);
+
+    UpdateRequest updateRequest(op);
+    tassert(11034202,
+            "nssOrUUID for an update must be a namespace string",
+            nssOrUUID.isNamespaceString());
+    updateRequest.setNamespaceString(nssOrUUID.nss());
+    if (_components.let.hasLet) {
+        updateRequest.setLetParameters(_components.let.shapifiedLet);
+    }
+
+    ParsedUpdate parsedUpdate(opCtx,
+                              &updateRequest,
+                              CollectionPtr::null /*CollectionPtr*/,
+                              false /*forgoOpCounterIncrements*/);
+    uassertStatusOK(parsedUpdate.parseRequest());
+
+    UpdateCmdShapeComponents{parsedUpdate, _components.let, opts}.appendTo(bob, opts, expCtx);
+}
+
+}  // namespace mongo::query_shape

src/mongo/db/query/query_shape/update_cmd_shape.h

@@ -0,0 +1,120 @@
+/**
+ *    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.
+ */
+
+#pragma once
+
+#include "mongo/bson/bsonelement.h"
+#include "mongo/db/pipeline/expression_context.h"
+#include "mongo/db/query/query_shape/let_shape_component.h"
+#include "mongo/db/query/query_shape/query_shape.h"
+#include "mongo/db/query/write_ops/parsed_update.h"
+#include "mongo/db/query/write_ops/write_ops_gen.h"
+#include "mongo/util/modules.h"
+
+#include <boost/intrusive_ptr.hpp>
+
+namespace mongo::query_shape {
+
+/**
+ * A struct representing the update command's specific components that are to be considered part
+ * of the query shape.
+ *
+ * This struct stores the shapified compoments (e.g., 'representativeQ') in BSON as a memory
+ * optimization. We choose to store them in BSON rather than in their parsed objects (e.g.,
+ * ParsedUpdate) because those parsed versions often still need these BSONs to survive as backing
+ * memory. So we store the representative components so that we are able to parse those components
+ * again if we need to compute a different shape.
+ */
+struct UpdateCmdShapeComponents : public query_shape::CmdSpecificShapeComponents {
+    UpdateCmdShapeComponents(const ParsedUpdate& parsedUpdate,
+                             LetShapeComponent let,
+                             const SerializationOptions& opts =
+                                 SerializationOptions::kRepresentativeQueryShapeSerializeOptions);
+
+    size_t size() const final;
+
+    void appendTo(BSONObjBuilder&,
+                  const SerializationOptions&,
+                  const boost::intrusive_ptr<ExpressionContext>&) const;
+
+    void HashValue(absl::HashState state) const final;
+
+    /**
+     * Returns the representative value corresponding to the 'u' (update) parameter to the update
+     * statement.
+     *
+     * Since u could be of various types. For example, it may be a BSONObj for a modification update
+     * or a BSONArray for a pipeline update, we use BSONElement to represent 'u' because it is
+     * capable of representing various types. But BSONElement does not own the data itself, we need
+     * to keep the backing BSONObj '_representativeUObj'.
+     */
+    inline BSONElement getRepresentativeU() const {
+        return _representativeUObj.firstElement();
+    }
+
+    // Representative shapes serialized with kRepresentativeQueryShapeSerializeOptions.
+    BSONObj representativeQ;
+    BSONObj _representativeUObj;  // The backing BSONObj for u. It always only has one element.
+    boost::optional<BSONObj> representativeC;
+    boost::optional<std::vector<BSONObj>> representativeArrayFilters;
+
+    bool multi;
+    bool upsert;
+
+    LetShapeComponent let;
+};
+
+/**
+ * A class representing the query shape of an aggregate command. The components are listed above.
+ * This class knows how to utilize those components to serialize to BSON with any
+ * SerializationOptions. Mostly this involves correctly setting up an ExpressionContext to re-parse
+ * the request if needed.
+ */
+class UpdateCmdShape : public Shape {
+public:
+    UpdateCmdShape(const write_ops::UpdateCommandRequest&,
+                   const ParsedUpdate&,
+                   const boost::intrusive_ptr<ExpressionContext>&);
+
+    const CmdSpecificShapeComponents& specificComponents() const final;
+
+    size_t extraSize() const final;
+
+protected:
+    void appendCmdSpecificShapeComponents(BSONObjBuilder&,
+                                          OperationContext*,
+                                          const SerializationOptions&) const final;
+
+private:
+    UpdateCmdShapeComponents _components;
+};
+static_assert(sizeof(UpdateCmdShape) == sizeof(Shape) + sizeof(UpdateCmdShapeComponents),
+              "If the class' members have changed, this assert and the extraSize() calculation may "
+              "need to be updated with a new value.");
+}  // namespace mongo::query_shape

src/mongo/db/query/query_shape/update_cmd_shape_test.cpp

@@ -0,0 +1,281 @@
+/**
+ *    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.
+ */
+
+#include "mongo/db/query/query_shape/update_cmd_shape.h"
+
+#include "mongo/bson/json.h"
+#include "mongo/db/pipeline/expression_context_for_test.h"
+#include "mongo/db/query/query_shape/let_shape_component.h"
+#include "mongo/db/query/query_test_service_context.h"
+#include "mongo/db/query/write_ops/update_request.h"
+#include "mongo/unittest/unittest.h"
+
+namespace mongo::query_shape {
+namespace {
+
+using write_ops::UpdateCommandRequest;
+
+static const NamespaceString kDefaultTestNss =
+    NamespaceString::createNamespaceString_forTest("testDB.testColl");
+
+class UpdateCmdShapeTest : public unittest::Test {
+public:
+    void setUp() final {
+        _queryTestServiceContext = std::make_unique<QueryTestServiceContext>();
+        _operationContext = _queryTestServiceContext->makeOperationContext();
+        _expCtx = make_intrusive<ExpressionContextForTest>();
+    }
+
+    std::vector<UpdateCmdShape> makeShapesFromUpdate(StringData updateCmd) {
+        auto updateRequest = UpdateCommandRequest::parseOwned(fromjson(updateCmd));
+
+        std::vector<UpdateCmdShape> shapes;
+        for (const auto& op : updateRequest.getUpdates()) {
+            UpdateRequest request(op);
+            request.setNamespaceString(kDefaultTestNss);
+            if (updateRequest.getLet()) {
+                request.setLetParameters(*updateRequest.getLet());
+            }
+
+            ParsedUpdate parsedUpdate(
+                _operationContext.get(), &request, CollectionPtr::null, false);
+            ASSERT_OK(parsedUpdate.parseRequest());
+            shapes.emplace_back(updateRequest, parsedUpdate, _expCtx);
+        }
+        return shapes;
+    }
+
+    UpdateCmdShape makeOneShapeFromUpdate(StringData updateCmd) {
+        auto shapes = makeShapesFromUpdate(updateCmd);
+        ASSERT_EQ(shapes.size(), 1);
+        return shapes.front();
+    }
+
+    std::unique_ptr<QueryTestServiceContext> _queryTestServiceContext;
+
+    ServiceContext::UniqueOperationContext _operationContext;
+    boost::intrusive_ptr<ExpressionContext> _expCtx;
+};
+
+TEST_F(UpdateCmdShapeTest, BasicReplacementUpdateShape) {
+    auto shape = makeOneShapeFromUpdate(R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" }, multi: false, upsert: false } ],
+        "$db": "testDB"
+    })"_sd);
+    ASSERT_BSONOBJ_EQ_AUTO(  // NOLINT
+        R"({
+            "cmdNs": {
+                "db": "testDB",
+                "coll": "testColl"
+            },
+            "command": "update",
+            "q": {
+                "x": {
+                    "$eq": "?number"
+                }
+            },
+            "u": "?object",
+            "multi": false,
+            "upsert": false 
+        })",
+        shape.toBson(_operationContext.get(),
+                     SerializationOptions::kDebugQueryShapeSerializeOptions,
+                     SerializationContext::stateDefault()));
+
+    ASSERT_BSONOBJ_EQ_AUTO(  // NOLINT
+        R"({
+            "cmdNs": {
+                "db": "testDB",
+                "coll": "testColl"
+            },
+            "command": "update",
+            "q": {
+                "x": {
+                    "$eq": 1
+                }
+            },
+            "u": {
+                "?": "?"
+            },
+            "multi": false,
+            "upsert": false
+        })",
+        shape.toBson(_operationContext.get(),
+                     SerializationOptions::kRepresentativeQueryShapeSerializeOptions,
+                     SerializationContext::stateDefault()));
+}
+
+TEST_F(UpdateCmdShapeTest, BatchReplacementUpdateShape) {
+    auto shapes = makeShapesFromUpdate(R"({
+        update: "testColl",
+        updates: [
+          { q: { x: {$eq: 3} }, u: { foo: "bar" }, multi: false, upsert: false },
+          { q: { x: {$gt: 3}, y: "foo" }, u: { x: {y: 100}, z: false }, multi: false, upsert: true }
+        ],
+        "$db": "testDB"
+    })"_sd);
+    ASSERT_EQ(shapes.size(), 2);
+    ASSERT_BSONOBJ_EQ_AUTO(  // NOLINT
+        R"({
+            "cmdNs": {
+                "db": "testDB",
+                "coll": "testColl"
+            },
+            "command": "update",
+            "q": {
+                "x": {
+                    "$eq": "?number"
+                }
+            },
+            "u": "?object",
+            "multi": false,
+            "upsert": false
+        })",
+        shapes[0].toBson(_operationContext.get(),
+                         SerializationOptions::kDebugQueryShapeSerializeOptions,
+                         SerializationContext::stateDefault()));
+    ASSERT_BSONOBJ_EQ_AUTO(  // NOLINT
+        R"({
+            "cmdNs": {
+                "db": "testDB",
+                "coll": "testColl"
+            },
+            "command": "update",
+            "q": {
+                "$and": [
+                    {
+                        "y": {
+                            "$eq": "?string"
+                        }
+                    },
+                    {
+                        "x": {
+                            "$gt": "?number"
+                        }
+                    }
+                ]
+            },
+            "u": "?object",
+            "multi": false,
+            "upsert": true
+        })",
+        shapes[1].toBson(_operationContext.get(),
+                         SerializationOptions::kDebugQueryShapeSerializeOptions,
+                         SerializationContext::stateDefault()));
+}
+
+TEST_F(UpdateCmdShapeTest, IncludesOptionalValues) {
+    // Test setting optional values such as 'multi' and 'upsert' to verify that they are included in
+    // the query shape.
+    auto shape = makeOneShapeFromUpdate(R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" }, multi: false, upsert: true } ],
+        ordered: false,
+        bypassDocumentValidation: true,
+        let: {x: 4, y: "abc"},
+        "$db": "testDB"
+    })"_sd);
+    ASSERT_BSONOBJ_EQ_AUTO(  // NOLINT
+        R"({
+            "cmdNs": {
+                "db": "testDB",
+                "coll": "testColl"
+            },
+            "command": "update",
+            "q": {
+                "x": {
+                    "$eq": "?number"
+                }
+            },
+            "u": "?object",
+            "multi": false,
+            "upsert": true,
+            "let": {
+                "x": "?number",
+                "y": "?string"
+            }
+        })",
+        shape.toBson(_operationContext.get(),
+                     SerializationOptions::kDebugQueryShapeSerializeOptions,
+                     SerializationContext::stateDefault()));
+}
+
+// Verifies that "update" command shape hash value is stable (does not change between the
+// versions of the server).
+TEST_F(UpdateCmdShapeTest, StableQueryShapeHashValue) {
+    auto shape = makeOneShapeFromUpdate(R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" }, multi: false, upsert: false } ],
+        "$db": "testDB"
+    })"_sd);
+
+    auto serializationContext = SerializationContext::stateCommandRequest();
+    const auto hash = shape.sha256Hash(_operationContext.get(), serializationContext);
+    ASSERT_EQ("56593B6B2CE6C3968E03CC55DFED93AE728CA730A21E0659390360636BD96B15",
+              hash.toHexString());
+}
+
+TEST_F(UpdateCmdShapeTest, SizeOfUpdateCmdShapeComponents) {
+    auto shape = makeOneShapeFromUpdate(R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" }, multi: false, upsert: false } ],
+        "$db": "testDB"
+    })"_sd);
+    auto updateComponents =
+        static_cast<const UpdateCmdShapeComponents&>(shape.specificComponents());
+
+    // The sizes of any members of UpdateCmdShapeComponents are typically accounted for by
+    // sizeof(UpdateCmdShapeComponents). The important part of the test here is to ensure that
+    // any additional memory allocations are also included in the size() operation.
+    const auto letSize = updateComponents.let.size();
+
+    ASSERT_EQ(updateComponents.size(),
+              sizeof(UpdateCmdShapeComponents) + updateComponents.representativeQ.objsize() +
+                  updateComponents._representativeUObj.objsize() + letSize -
+                  sizeof(LetShapeComponent));
+}
+
+TEST_F(UpdateCmdShapeTest, EquivalentUpdateCmdShapeSizes) {
+    auto shape = makeOneShapeFromUpdate(R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" }, multi: false, upsert: false } ],
+        "$db": "testDB"
+    })"_sd);
+    auto shapeOptionalValues = makeOneShapeFromUpdate(R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" }, multi: false, upsert: true } ],
+        "$db": "testDB",
+        ordered: false,
+        bypassDocumentValidation: true
+    })"_sd);
+    ASSERT_EQ(shape.size(), shapeOptionalValues.size());
+}
+}  // namespace
+}  // namespace mongo::query_shape

src/mongo/db/query/query_stats/BUILD.bazel

@@ -28,7 +28,6 @@ mongo_cc_library(
     name = "query_stats",
     srcs = [
         "aggregated_metric.cpp",
-        "key.cpp",
         "optimizer_metrics_stats_entry.cpp",
         "query_stats.cpp",
         "query_stats_entry.cpp",
@@ -38,7 +37,6 @@ mongo_cc_library(
     ],
     hdrs = [
         "aggregated_metric.h",
-        "key.h",
         "optimizer_metrics_stats_entry.h",
         "query_stats.h",
         "query_stats_entry.h",
@@ -50,6 +48,7 @@ mongo_cc_library(
         "//src/mongo/client:read_preference",  # TODO(SERVER-93876): Remove.
     ],
     deps = [
+        ":key",
         ":rate_limiting",
         "//src/mongo/db:api_parameters",
         "//src/mongo/db:commands",
@@ -57,6 +56,39 @@ mongo_cc_library(
     ],
 )
 
+mongo_cc_library(
+    name = "key",
+    srcs = [
+        "key.cpp",
+    ],
+    hdrs = [
+        "key.h",
+    ],
+    header_deps = [
+        "//src/mongo/client:read_preference",  # TODO(SERVER-93876): Remove.
+    ],
+    deps = [
+        ":rate_limiting",
+        "//src/mongo/db:api_parameters",
+        "//src/mongo/db:commands",
+        "//src/mongo/util:buildinfo",
+    ],
+)
+
+mongo_cc_library(
+    name = "update_key",
+    srcs = [
+        "update_key.cpp",
+    ],
+    hdrs = [
+        "update_key.h",
+    ],
+    deps = [
+        ":key",
+        "//src/mongo/db/query/query_shape:update_cmd_shape",
+    ],
+)
+
 idl_generator(
     name = "transform_algorithm_gen",
     src = "transform_algorithm.idl",
@@ -84,6 +116,7 @@ mongo_cc_unit_test(
         "query_stats_test.cpp",
         "rate_limiting_test.cpp",
         "supplemental_metrics_test.cpp",
+        "update_key_test.cpp",
     ],
     header_deps = [
         "//src/mongo/db/pipeline:expression_context_for_test",
@@ -92,6 +125,7 @@ mongo_cc_unit_test(
     tags = ["mongo_unittest_fourth_group"],
     deps = [
         ":query_stats",
+        ":update_key",
         "//src/mongo/db:service_context_d_test_fixture",
         "//src/mongo/db/query:query_test_service_context",
         "//src/mongo/util:version_impl",

src/mongo/db/query/query_stats/update_key.cpp

@@ -0,0 +1,84 @@
+/**
+ *    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.
+ */
+
+#include "mongo/db/query/query_stats/update_key.h"
+
+#include "mongo/db/pipeline/pipeline.h"
+#include "mongo/db/query/query_shape/query_shape.h"
+#include "mongo/db/query/query_shape/serialization_options.h"
+
+#include <memory>
+
+#include <absl/container/node_hash_set.h>
+#include <boost/cstdint.hpp>
+#include <boost/move/utility_core.hpp>
+#include <boost/optional/optional.hpp>
+#include <boost/smart_ptr/intrusive_ptr.hpp>
+
+namespace mongo::query_stats {
+
+UpdateCmdComponents::UpdateCmdComponents(const write_ops::UpdateCommandRequest& request)
+    : _ordered(request.getOrdered()),
+      _bypassDocumentValidation(request.getBypassDocumentValidation()) {}
+
+
+void UpdateCmdComponents::HashValue(absl::HashState state) const {
+    state = absl::HashState::combine(std::move(state), _ordered, _bypassDocumentValidation);
+}
+
+void UpdateCmdComponents::appendTo(BSONObjBuilder& bob, const SerializationOptions& opts) const {
+    bob.append(write_ops::UpdateCommandRequest::kOrderedFieldName, _ordered);
+
+    bob.append(write_ops::UpdateCommandRequest::kBypassDocumentValidationFieldName,
+               _bypassDocumentValidation);
+}
+
+size_t UpdateCmdComponents::size() const {
+    return sizeof(UpdateCmdComponents);
+}
+
+void UpdateKey::appendCommandSpecificComponents(BSONObjBuilder& bob,
+                                                const SerializationOptions& opts) const {
+    return _components.appendTo(bob, opts);
+}
+
+UpdateKey::UpdateKey(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                     const write_ops::UpdateCommandRequest& request,
+                     const boost::optional<BSONObj>& hint,
+                     std::unique_ptr<query_shape::Shape> updateShape,
+                     query_shape::CollectionType collectionType)
+    : Key(expCtx->getOperationContext(),
+          std::move(updateShape),
+          hint,
+          request.getReadConcern(),
+          request.getMaxTimeMS().has_value(),
+          collectionType),
+      _components(request) {}
+
+}  // namespace mongo::query_stats

src/mongo/db/query/query_stats/update_key.h

@@ -0,0 +1,109 @@
+/**
+ *    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.
+ */
+
+#pragma once
+
+#include "mongo/bson/bsonobj.h"
+#include "mongo/bson/bsonobjbuilder.h"
+#include "mongo/db/local_catalog/collection_type.h"
+#include "mongo/db/pipeline/expression_context.h"
+#include "mongo/db/pipeline/pipeline.h"
+#include "mongo/db/pipeline/variables.h"
+#include "mongo/db/query/query_stats/key.h"
+#include "mongo/db/query/write_ops/write_ops_gen.h"
+#include "mongo/util/modules.h"
+
+#include <utility>
+
+#include <absl/container/node_hash_map.h>
+#include <boost/move/utility_core.hpp>
+#include <boost/none.hpp>
+#include <boost/optional/optional.hpp>
+#include <boost/smart_ptr/intrusive_ptr.hpp>
+
+namespace mongo::query_stats {
+
+/**
+ * Struct representing the update command's unique arguments which should be included in the
+ * query stats key.
+ */
+struct UpdateCmdComponents : public SpecificKeyComponents {
+    UpdateCmdComponents(const write_ops::UpdateCommandRequest&);
+
+    void HashValue(absl::HashState state) const final;
+
+    void appendTo(BSONObjBuilder& bob, const SerializationOptions& opts) const;
+
+    size_t size() const override;
+
+    bool _ordered;
+    bool _bypassDocumentValidation;
+};
+
+/**
+ * An implementation of the query stats store key for the update command. This class is a wrapper
+ * around the base 'Key' class and 'QueryCmdShape', and it includes UpdateCmdComponents for those
+ * update-specific components (e.g., ordered, bypassDocumentValidation).
+ */
+class UpdateKey final : public Key {
+public:
+    UpdateKey(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+              const write_ops::UpdateCommandRequest& request,
+              const boost::optional<BSONObj>& hint,
+              std::unique_ptr<query_shape::Shape> updateShape,
+              query_shape::CollectionType collectionType = query_shape::CollectionType::kUnknown);
+
+    const SpecificKeyComponents& specificComponents() const final {
+        return _components;
+    }
+
+    // The default implementation of hashing for smart pointers is not a good one for our purposes.
+    // Here we overload them to actually take the hash of the object, rather than hashing the
+    // pointer itself.
+    template <typename H>
+    friend H AbslHashValue(H h, const std::unique_ptr<const UpdateKey>& key) {
+        return H::combine(std::move(h), *key);
+    }
+    template <typename H>
+    friend H AbslHashValue(H h, const std::shared_ptr<const UpdateKey>& key) {
+        return H::combine(std::move(h), *key);
+    }
+
+
+protected:
+    void appendCommandSpecificComponents(BSONObjBuilder& bob,
+                                         const SerializationOptions& opts) const final;
+
+private:
+    const UpdateCmdComponents _components;
+};
+static_assert(
+    sizeof(UpdateKey) == sizeof(Key) + sizeof(UpdateCmdComponents),
+    "If the class' members have changed, this assert may need to be updated with a new value.");
+}  // namespace mongo::query_stats

src/mongo/db/query/query_stats/update_key_test.cpp

@@ -0,0 +1,206 @@
+/**
+ *    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.
+ */
+
+#include "mongo/db/query/query_stats/update_key.h"
+
+#include "mongo/bson/json.h"
+#include "mongo/db/pipeline/expression_context.h"
+#include "mongo/db/pipeline/expression_context_for_test.h"
+#include "mongo/db/query/query_shape/update_cmd_shape.h"
+#include "mongo/db/query/write_ops/update_request.h"
+#include "mongo/db/service_context_test_fixture.h"
+#include "mongo/unittest/unittest.h"
+#include "mongo/util/intrusive_counter.h"
+
+#include <boost/smart_ptr/intrusive_ptr.hpp>
+
+namespace mongo::query_stats {
+
+namespace {
+
+using write_ops::UpdateCommandRequest;
+
+static const NamespaceStringOrUUID kDefaultTestNss =
+    NamespaceStringOrUUID{NamespaceString::createNamespaceString_forTest("testDB.testColl")};
+
+static constexpr auto collectionType = query_shape::CollectionType::kCollection;
+
+class UpdateKeyTest : public ServiceContextTest {
+public:
+    std::vector<std::unique_ptr<const Key>> makeUpdateKeys(
+        const boost::intrusive_ptr<ExpressionContext>& expCtx, StringData cmd) {
+        auto ucr = UpdateCommandRequest::parseOwned(fromjson(cmd));
+
+        std::vector<std::unique_ptr<const Key>> keys;
+        for (const auto& op : ucr.getUpdates()) {
+            UpdateRequest request(op);
+            request.setNamespaceString(kDefaultTestNss.nss());
+            if (ucr.getLet()) {
+                request.setLetParameters(*ucr.getLet());
+            }
+
+            ParsedUpdate parsedUpdate(
+                expCtx->getOperationContext(), &request, CollectionPtr::null, false);
+            ASSERT_OK(parsedUpdate.parseRequest());
+            auto shape = std::make_unique<query_shape::UpdateCmdShape>(ucr, parsedUpdate, expCtx);
+
+            auto key = std::make_unique<UpdateKey>(expCtx,
+                                                   ucr,
+                                                   parsedUpdate.getRequest()->getHint(),
+                                                   std::move(shape),
+                                                   collectionType);
+            keys.push_back(std::move(key));
+        }
+        return keys;
+    }
+
+    std::unique_ptr<const Key> makeOneUpdateKey(
+        const boost::intrusive_ptr<ExpressionContext>& expCtx, StringData cmd) {
+        auto keys = makeUpdateKeys(expCtx, cmd);
+        ASSERT_EQ(keys.size(), 1);
+        return std::move(keys.front());
+    }
+};
+
+TEST_F(UpdateKeyTest, ReplacementUpdateCmdComponents) {
+    // Create a request that none of the optional values are set.
+    auto ucr = UpdateCommandRequest::parseOwned(fromjson(R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" } } ],
+        "$db": "testDB"
+    })"_sd));
+
+    auto updateComponents = std::make_unique<UpdateCmdComponents>(ucr);
+
+    // Confirm all the optional values are still present. They are provided with their default
+    // values specified in the IDL.
+    BSONObjBuilder bob;
+    updateComponents->appendTo(bob,
+                               SerializationOptions::kRepresentativeQueryShapeSerializeOptions);
+    ASSERT_BSONOBJ_EQ(fromjson(R"({ ordered: true, bypassDocumentValidation: false })"), bob.obj());
+}
+
+TEST_F(UpdateKeyTest, IncludesOptionalValues) {
+    // Create a request that all the optional values are included.
+    auto ucr = UpdateCommandRequest::parseOwned(fromjson(R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" }, multi: true, upsert: true } ],
+        bypassDocumentValidation: true,
+        ordered: false,
+        "$db": "testDB"
+    })"_sd));
+
+    auto updateComponents = std::make_unique<UpdateCmdComponents>(ucr);
+
+    // Verify that the optional values are reflected in the stats key components.
+    BSONObjBuilder bob;
+    updateComponents->appendTo(bob,
+                               SerializationOptions::kRepresentativeQueryShapeSerializeOptions);
+    ASSERT_BSONOBJ_EQ(fromjson(R"({ ordered: false, bypassDocumentValidation: true })"), bob.obj());
+}
+
+TEST_F(UpdateKeyTest, SizeOfUpdateCmdComponents) {
+    auto update = fromjson(R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" }, multi: false, upsert: false } ],
+        "$db": "testDB"
+    })"_sd);
+    auto ucr = UpdateCommandRequest::parse(std::move(update));
+
+    auto updateComponents = std::make_unique<UpdateCmdComponents>(ucr);
+
+    const auto minimumSize = sizeof(SpecificKeyComponents) + 2 /*size for bools*/;
+    ASSERT_GTE(updateComponents->size(), minimumSize);
+    ASSERT_LTE(updateComponents->size(), minimumSize + 8 /*padding*/);
+}
+
+TEST_F(UpdateKeyTest, EquivalentUpdateCmdComponentSizes) {
+    // Create a request that has no values set.
+    auto ucrNoSetValues = UpdateCommandRequest::parseOwned(fromjson(R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" } } ],
+        "$db": "testDB"
+    })"_sd));
+
+    auto updateComponentsNoValues = std::make_unique<UpdateCmdComponents>(ucrNoSetValues);
+
+    // Create a request that has all values set. None of these should affect the size.
+    auto ucrAllValues = UpdateCommandRequest::parseOwned(fromjson(R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" }, multi: true, upsert: true } ],
+        bypassDocumentValidation: true,
+        ordered: false,
+        "$db": "testDB"
+    })"_sd));
+
+    auto updateComponentsAllValues = std::make_unique<UpdateCmdComponents>(ucrAllValues);
+
+    // Verify their sizes are equal. This is because the optional parameters such as
+    // 'bypassDocumentValidation' and 'ordered' are always provided with their default values when
+    // they are not set from command requests.
+    ASSERT_EQ(updateComponentsAllValues->size(), updateComponentsNoValues->size());
+}
+
+// Testing item in opCtx that should impact key size.
+TEST_F(UpdateKeyTest, SizeOfUpdateKeyWithAndWithoutComment) {
+    auto cmd = R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" } } ],
+        "$db": "testDB"
+    })"_sd;
+    auto expCtx = make_intrusive<ExpressionContextForTest>(kDefaultTestNss.nss());
+    auto keyWithoutComment = makeOneUpdateKey(expCtx, cmd);
+
+    expCtx->getOperationContext()->setComment(BSON("comment" << " foo"));
+    auto keyWithComment = makeOneUpdateKey(expCtx, cmd);
+    ASSERT_LT(keyWithoutComment->size(), keyWithComment->size());
+}
+
+// Testing item in command request that should impact key size.
+TEST_F(UpdateKeyTest, SizeOfUpdateKeyWithAndWithoutReadConcern) {
+    auto expCtx = make_intrusive<ExpressionContextForTest>(kDefaultTestNss.nss());
+    auto keyWithoutReadConcern = makeOneUpdateKey(expCtx, R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" } } ],
+        "$db": "testDB"
+    })"_sd);
+
+    auto keyWithReadConcern = makeOneUpdateKey(expCtx, R"({
+        update: "testColl",
+        updates: [ { q: { x: {$eq: 3} }, u: { foo: "bar" } } ],
+        "$db": "testDB",
+        readConcern: {
+            afterClusterTime: Timestamp(1654272333, 13),
+            level: "majority"
+        }
+    })"_sd);
+    ASSERT_LT(keyWithoutReadConcern->size(), keyWithReadConcern->size());
+}
+}  // namespace
+}  // namespace mongo::query_stats

src/mongo/db/query/write_ops/BUILD.bazel

@@ -230,7 +230,9 @@ mongo_cc_library(
         "//src/mongo/db/query:explain_diagnostic_printer",
         "//src/mongo/db/query:shard_key_diagnostic_printer",
         "//src/mongo/db/query/query_settings:query_settings_service",
+        "//src/mongo/db/query/query_shape:update_cmd_shape",
         "//src/mongo/db/query/query_stats",
+        "//src/mongo/db/query/query_stats:update_key",
         "//src/mongo/db/repl:oplog",
         "//src/mongo/db/repl:repl_coordinator_interface",
         "//src/mongo/db/s:query_analysis_writer",

src/mongo/db/query/write_ops/write_ops_exec.cpp

@@ -34,6 +34,7 @@
 #include "mongo/base/string_data.h"
 #include "mongo/bson/bsonelement.h"
 #include "mongo/bson/bsonelement_comparator.h"
+#include "mongo/bson/bsonobjbuilder.h"
 #include "mongo/bson/bsontypes.h"
 #include "mongo/db/auth/action_type.h"
 #include "mongo/db/auth/authorization_session.h"
@@ -83,6 +84,11 @@
 #include "mongo/db/query/plan_explainer.h"
 #include "mongo/db/query/plan_summary_stats.h"
 #include "mongo/db/query/plan_yield_policy.h"
+#include "mongo/db/query/query_shape/query_shape.h"
+#include "mongo/db/query/query_shape/shape_helpers.h"
+#include "mongo/db/query/query_shape/update_cmd_shape.h"
+#include "mongo/db/query/query_stats/query_stats.h"
+#include "mongo/db/query/query_stats/update_key.h"
 #include "mongo/db/query/shard_key_diagnostic_printer.h"
 #include "mongo/db/query/write_ops/delete_request_gen.h"
 #include "mongo/db/query/write_ops/insert.h"
@@ -1388,6 +1394,67 @@ static SingleWriteResult performSingleUpdateOpNoRetry(OperationContext* opCtx,
     return result;
 }
 
+void registerRequestForQueryStats(OperationContext* opCtx,
+                                  const NamespaceString& ns,
+                                  const CollectionAcquisition& collection,
+                                  const write_ops::UpdateCommandRequest& wholeOp,
+                                  const ParsedUpdate& parsedUpdate) {
+    // Skip registering for query stats when the feature flag is disabled.
+    if (!feature_flags::gFeatureFlagQueryStatsWriteCommand.isEnabledUseLastLTSFCVWhenUninitialized(
+            VersionContext::getDecoration(opCtx),
+            serverGlobalParams.featureCompatibility.acquireFCVSnapshot())) {
+        return;
+    }
+
+    // TODO(SERVER-111930): Support recording query stats for updates with simple ID query
+    // Skip if the parse query is unavailable. This could happen if the query is a simple Id query:
+    // an exact-match query on _id.
+    if (!parsedUpdate.hasParsedQuery()) {
+        return;
+    }
+
+    // Skip registering the request with encrypted fields as indicated by the inclusion of
+    // encryptionInformation. It is important to do this before canonicalizing and optimizing the
+    // query, each of which would alter the query shape.
+    if (wholeOp.getEncryptionInformation()) {
+        return;
+    }
+
+    // Skip unsupported update types.
+    // TODO(SERVER-110343) and TODO(SERVER-110344) Support pipeline and modifier updates.
+    if (parsedUpdate.getRequest()->getUpdateModification().type() !=
+        write_ops::UpdateModification::Type::kReplacement) {
+        return;
+    }
+
+    // Compute QueryShapeHash and record it in CurOp.
+    query_shape::DeferredQueryShape deferredShape{[&]() {
+        return shape_helpers::tryMakeShape<query_shape::UpdateCmdShape>(
+            wholeOp, parsedUpdate, parsedUpdate.expCtx());
+    }};
+
+    // QueryShapeHash(QSH) will be recorded in CurOp, but it is not being used for anything else
+    // downstream yet until we support updates in PQS. Using std::ignore to indicate that discarding
+    // the returned QSH is intended.
+    std::ignore = CurOp::get(opCtx)->debug().ensureQueryShapeHash(opCtx, [&]() {
+        return shape_helpers::computeQueryShapeHash(
+            parsedUpdate.expCtx(), deferredShape, wholeOp.getNamespace());
+    });
+
+
+    // Register query stats collection.
+    query_stats::registerRequest(opCtx, ns, [&]() {
+        uassertStatusOKWithContext(deferredShape->getStatus(), "Failed to compute query shape");
+        return std::make_unique<query_stats::UpdateKey>(parsedUpdate.expCtx(),
+                                                        wholeOp,
+                                                        parsedUpdate.getRequest()->getHint(),
+                                                        std::move(deferredShape->getValue()),
+                                                        collection.getCollectionType());
+    });
+
+    // TODO(SERVER-110348) Support collecting data-bearing node metrics here.
+}
+
 /**
  * Performs a single update, sometimes retrying failure due to WriteConflictException.
  */
@@ -1395,6 +1462,7 @@ static SingleWriteResult performSingleUpdateOp(
     OperationContext* opCtx,
     const NamespaceString& ns,
     const timeseries::CollectionPreConditions& preConditions,
+    const write_ops::UpdateCommandRequest& wholeOp,
     UpdateRequest* updateRequest,
     OperationSource source,
     bool forgoOpCounterIncrements,
@@ -1478,6 +1546,11 @@ static SingleWriteResult performSingleUpdateOp(
                               updateRequest->source() == OperationSource::kTimeseriesUpdate);
     uassertStatusOK(parsedUpdate.parseRequest());
 
+    // Register query shape here once we obtain the ParsedUpdate, before executing the update
+    // command. After parsedUpdate.parseRequest(), the parsed query and the update driver become
+    // available for computing query shape.
+    registerRequestForQueryStats(opCtx, ns, collection, wholeOp, parsedUpdate);
+
     // Create an RAII object that prints useful information about the ExpressionContext in the case
     // of a tassert or crash.
     ScopedDebugInfo expCtxDiagnostics(
@@ -1527,6 +1600,7 @@ static SingleWriteResult performSingleUpdateOpWithDupKeyRetry(
     OperationContext* opCtx,
     const NamespaceString& ns,
     const std::vector<StmtId>& stmtIds,
+    const write_ops::UpdateCommandRequest& wholeOp,
     const write_ops::UpdateOpEntry& op,
     const timeseries::CollectionPreConditions& preConditions,
     LegacyRuntimeConstants runtimeConstants,
@@ -1581,6 +1655,7 @@ static SingleWriteResult performSingleUpdateOpWithDupKeyRetry(
             auto ret = performSingleUpdateOp(opCtx,
                                              ns,
                                              preConditions,
+                                             wholeOp,
                                              &request,
                                              source,
                                              forgoOpCounterIncrements,
@@ -1820,6 +1895,7 @@ WriteResult performUpdates(
                 performSingleUpdateOpWithDupKeyRetry(opCtx,
                                                      ns,
                                                      stmtIds,
+                                                     wholeOp,
                                                      singleOp,
                                                      preConditions,
                                                      runtimeConstants,