SERVER-68751 Add analyzeShardKey and configureQueryAnalyzer commands on mongos

3 files changed, 248 insertions, 0 deletions

diff --git a/src/mongo/s/commands/SConscript b/src/mongo/s/commands/SConscript
index 0d60bad7378..05cd828b268 100644
--- a/src/mongo/s/commands/SConscript
+++ b/src/mongo/s/commands/SConscript
@@ -199,9 +199,11 @@ env.Library(
     source=[
         'cluster_add_shard_cmd.cpp',
         'cluster_add_shard_to_zone_cmd.cpp',
+        'cluster_analyze_shard_key_cmd.cpp',
         'cluster_balancer_collection_status_cmd.cpp',
         'cluster_clear_jumbo_flag_cmd.cpp',
         'cluster_configure_collection_balancing.cpp',
+        'cluster_configure_query_analyzer_cmd.cpp',
         'cluster_control_balancer_cmd.cpp',
         'cluster_list_shards_cmd.cpp',
         'cluster_merge_chunks_cmd.cpp',
@@ -224,6 +226,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/auth/saslauth',
         '$BUILD_DIR/mongo/db/commands',
         '$BUILD_DIR/mongo/db/commands/servers',
+        '$BUILD_DIR/mongo/db/s/analyze_shard_key',
         '$BUILD_DIR/mongo/db/server_base',
         '$BUILD_DIR/mongo/rpc/client_metadata',
         '$BUILD_DIR/mongo/rpc/rewrite_state_change_errors',
diff --git a/src/mongo/s/commands/cluster_analyze_shard_key_cmd.cpp b/src/mongo/s/commands/cluster_analyze_shard_key_cmd.cpp
new file mode 100644
index 00000000000..073dcdbea7e
--- /dev/null
+++ b/src/mongo/s/commands/cluster_analyze_shard_key_cmd.cpp
@@ -0,0 +1,137 @@
+/**
+ *    Copyright (C) 2022-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/platform/basic.h"
+
+#include "mongo/db/auth/authorization_session.h"
+#include "mongo/db/commands.h"
+#include "mongo/db/s/analyze_shard_key_cmd_gen.h"
+#include "mongo/db/s/analyze_shard_key_feature_flag_gen.h"
+#include "mongo/logv2/log.h"
+#include "mongo/s/cluster_commands_helpers.h"
+#include "mongo/s/grid.h"
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kCommand
+
+
+namespace mongo {
+
+namespace {
+
+ShardId getRandomOwningShardId(const ChunkManager& cm, const CachedDatabaseInfo& dbInfo) {
+    if (cm.isSharded()) {
+        std::set<ShardId> shardIds;
+        cm.getAllShardIds(&shardIds);
+        auto it = shardIds.begin();
+        std::advance(it, std::rand() % shardIds.size());
+        return *it;
+    }
+    return dbInfo->getPrimary();
+}
+
+class AnalyzeShardKeyCmd : public TypedCommand<AnalyzeShardKeyCmd> {
+public:
+    using Request = AnalyzeShardKey;
+    using Response = AnalyzeShardKeyResponse;
+
+    class Invocation final : public InvocationBase {
+    public:
+        using InvocationBase::InvocationBase;
+
+        Response typedRun(OperationContext* opCtx) {
+            const auto& nss = ns();
+            const auto& catalogCache = Grid::get(opCtx)->catalogCache();
+            const auto cm = uassertStatusOK(catalogCache->getCollectionRoutingInfo(opCtx, nss));
+            const auto dbInfo = uassertStatusOK(catalogCache->getDatabase(opCtx, nss.db()));
+
+            auto shardId = getRandomOwningShardId(cm, dbInfo);
+            uassert(ErrorCodes::IllegalOperation,
+                    "Cannot analyze a shard key for a collection on the config server",
+                    shardId != ShardId::kConfigServerId);
+            const auto shard =
+                uassertStatusOK(Grid::get(opCtx)->shardRegistry()->getShard(opCtx, shardId));
+
+            auto cmdObj = CommandHelpers::filterCommandRequestForPassthrough(request().toBSON({}));
+            if (cm.isSharded()) {
+                cmdObj = appendShardVersion(cmdObj, cm.getVersion(shardId));
+            } else {
+                cmdObj = appendShardVersion(cmdObj, ChunkVersion::UNSHARDED());
+                cmdObj = appendDbVersionIfPresent(cmdObj, dbInfo->getVersion());
+            }
+
+            auto swResponse = shard->runCommandWithFixedRetryAttempts(
+                opCtx,
+                ReadPreferenceSetting{ReadPreference::Nearest},
+                NamespaceString::kAdminDb.toString(),
+                cmdObj,
+                Shard::RetryPolicy::kIdempotent);
+            uassertStatusOK(Shard::CommandResponse::getEffectiveStatus(swResponse));
+
+            auto response = AnalyzeShardKeyResponse::parse(
+                IDLParserContext("clusterAnalyzeShardKey"), swResponse.getValue().response);
+            return response;
+        }
+
+    private:
+        NamespaceString ns() const override {
+            return request().getCommandParameter();
+        }
+
+        bool supportsWriteConcern() const override {
+            return false;
+        }
+
+        void doCheckAuthorization(OperationContext* opCtx) const override {
+            uassert(ErrorCodes::Unauthorized,
+                    "Unauthorized",
+                    AuthorizationSession::get(opCtx->getClient())
+                        ->isAuthorizedForActionsOnResource(ResourcePattern::forExactNamespace(ns()),
+                                                           ActionType::shardCollection));
+        }
+    };
+
+    AllowedOnSecondary secondaryAllowed(ServiceContext*) const override {
+        return AllowedOnSecondary::kNever;
+    }
+
+    bool adminOnly() const override {
+        return true;
+    }
+
+    std::string help() const override {
+        return "Returns metrics for evaluating a shard key for a collection.";
+    }
+};
+
+MONGO_REGISTER_FEATURE_FLAGGED_COMMAND(AnalyzeShardKeyCmd,
+                                       analyze_shard_key::gFeatureFlagAnalyzeShardKey);
+
+}  // namespace
+
+}  // namespace mongo
diff --git a/src/mongo/s/commands/cluster_configure_query_analyzer_cmd.cpp b/src/mongo/s/commands/cluster_configure_query_analyzer_cmd.cpp
new file mode 100644
index 00000000000..eaacc2ff522
--- /dev/null
+++ b/src/mongo/s/commands/cluster_configure_query_analyzer_cmd.cpp
@@ -0,0 +1,108 @@
+/**
+ *    Copyright (C) 2022-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/platform/basic.h"
+
+#include "mongo/db/auth/authorization_session.h"
+#include "mongo/db/commands.h"
+#include "mongo/db/s/analyze_shard_key_feature_flag_gen.h"
+#include "mongo/db/s/configure_query_analyzer_cmd_gen.h"
+#include "mongo/logv2/log.h"
+#include "mongo/s/grid.h"
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kCommand
+
+
+namespace mongo {
+
+namespace {
+
+class ConfigureQueryAnalyzerCmd : public TypedCommand<ConfigureQueryAnalyzerCmd> {
+public:
+    using Request = ConfigureQueryAnalyzer;
+    using Response = ConfigureQueryAnalyzerResponse;
+
+    class Invocation final : public InvocationBase {
+    public:
+        using InvocationBase::InvocationBase;
+
+        Response typedRun(OperationContext* opCtx) {
+            const auto configShard = Grid::get(opCtx)->shardRegistry()->getConfigShard();
+
+            auto swResponse = configShard->runCommandWithFixedRetryAttempts(
+                opCtx,
+                ReadPreferenceSetting{ReadPreference::PrimaryOnly},
+                NamespaceString::kAdminDb.toString(),
+                request().toBSON({}),
+                Shard::RetryPolicy::kIdempotent);
+            uassertStatusOK(Shard::CommandResponse::getEffectiveStatus(swResponse));
+
+            auto response = ConfigureQueryAnalyzerResponse::parse(
+                IDLParserContext("clusterConfigureQueryAnalyzer"), swResponse.getValue().response);
+            return response;
+        }
+
+    private:
+        NamespaceString ns() const override {
+            return request().getCommandParameter();
+        }
+
+        bool supportsWriteConcern() const override {
+            return false;
+        }
+
+        void doCheckAuthorization(OperationContext* opCtx) const override {
+            uassert(ErrorCodes::Unauthorized,
+                    "Unauthorized",
+                    AuthorizationSession::get(opCtx->getClient())
+                        ->isAuthorizedForActionsOnResource(ResourcePattern::forExactNamespace(ns()),
+                                                           ActionType::createCollection));
+        }
+    };
+
+    AllowedOnSecondary secondaryAllowed(ServiceContext*) const override {
+        return AllowedOnSecondary::kNever;
+    }
+
+    bool adminOnly() const override {
+        return true;
+    }
+
+    std::string help() const override {
+        return "Starts or stops collecting metrics about read and write queries against "
+               "collection.";
+    }
+};
+
+MONGO_REGISTER_FEATURE_FLAGGED_COMMAND(ConfigureQueryAnalyzerCmd,
+                                       analyze_shard_key::gFeatureFlagAnalyzeShardKey);
+
+}  // namespace
+
+}  // namespace mongo