mirror
/
mongo
mirror of https://github.com/mongodb/mongo
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

SERVER-100957 Provide different QuerySettingsService implementations for router and shard (#35027) 

Co-authored-by: Catalin Sumanaru <catalin.sumanaru@mongodb.com>
GitOrigin-RevId: 8bd8a3e7a3756108a2893b3457353a53778aa77a
This commit is contained in:
Denis Grebennicov 2025-04-25 15:41:43 +02:00 committed by MongoDB Bot
parent 3de1221ad8
commit 9f8a5dc5bd
13 changed files with 1144 additions and 838 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
src/mongo/db/commands/query_cmd/query_settings_cmds.cpp
View File

@ -213,17 +213,19 @@ void validateAndSimplifyQuerySettings(
const boost::optional<const RepresentativeQueryInfo&>& representativeQueryInfo,
const boost::optional<QueryInstance>& previousRepresentativeQuery,
QuerySettings& querySettings) {
auto& service = QuerySettingsService::get(opCtx);
// In case the representative query was not provided but the previous representative query is
// available, assert that query settings will be set on a valid query.
if (!representativeQueryInfo && previousRepresentativeQuery) {
validateQueryCompatibleWithAnyQuerySettings(
service.validateQueryCompatibleWithAnyQuerySettings(
createRepresentativeInfo(opCtx, *previousRepresentativeQuery, tenantId));
}
if (representativeQueryInfo) {
validateQueryCompatibleWithQuerySettings(*representativeQueryInfo, querySettings);
service.validateQueryCompatibleWithQuerySettings(*representativeQueryInfo, querySettings);
}
simplifyQuerySettings(querySettings);
validateQuerySettings(querySettings);
service.simplifyQuerySettings(querySettings);
service.validateQuerySettings(querySettings);
}
class SetQuerySettingsCommand final : public TypedCommand<SetQuerySettingsCommand> {
@ -312,7 +314,8 @@ public:
// Assert that query settings will be set on a valid query.
if (representativeQuery) {
validateQueryCompatibleWithAnyQuerySettings(*representativeQueryInfo);
QuerySettingsService::get(opCtx).validateQueryCompatibleWithAnyQuerySettings(
*representativeQueryInfo);
}
SetQuerySettingsCommandReply reply;

2
src/mongo/db/query/query_settings/BUILD.bazel
View File

@ -56,6 +56,7 @@ idl_generator(
mongo_cc_library(
name = "query_settings_service",
srcs = [
"query_settings_cluster_parameter_impl.cpp",
"query_settings_hash.cpp",
"query_settings_manager.cpp",
"query_settings_manager.h",
@ -90,6 +91,7 @@ mongo_cc_unit_test(
srcs = [
"index_hints_serialization_test.cpp",
"query_framework_serialization_test.cpp",
"query_settings_cluster_parameter_impl_test.cpp",
"query_settings_hash_test.cpp",
"query_settings_manager_test.cpp",
"query_settings_service_test.cpp",

147
src/mongo/db/query/query_settings/query_settings_cluster_parameter_impl.cpp Normal file
View File

@ -0,0 +1,147 @@
/**
* 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 <algorithm>
#include <boost/move/utility_core.hpp>
#include <boost/optional.hpp>
#include "mongo/base/error_codes.h"
#include "mongo/base/status.h"
#include "mongo/base/string_data.h"
#include "mongo/bson/bsonelement.h"
#include "mongo/bson/bsonobjbuilder.h"
#include "mongo/bson/bsontypes.h"
#include "mongo/db/client.h"
#include "mongo/db/commands/server_status.h"
#include "mongo/db/logical_time.h"
#include "mongo/db/operation_context.h"
#include "mongo/db/query/query_settings/query_settings_cluster_parameter_gen.h"
#include "mongo/db/query/query_settings/query_settings_service.h"
#include "mongo/db/service_context.h"
#include "mongo/idl/idl_parser.h"
namespace mongo::query_settings {
class QuerySettingsServerStatusSection final : public ServerStatusSection {
public:
using ServerStatusSection::ServerStatusSection;
bool includeByDefault() const override {
// Only include if Query Settings are enabled.
// We need to use isEnabledUseLatestFCVWhenUninitialized() instead of isEnabled() because
// this could run during startup while the FCV is still uninitialized.
return feature_flags::gFeatureFlagQuerySettings.isEnabledUseLatestFCVWhenUninitialized(
serverGlobalParams.featureCompatibility.acquireFCVSnapshot());
}
BSONObj generateSection(OperationContext* opCtx,
const BSONElement& configElement) const override {
stdx::lock_guard<stdx::mutex> lk(_mutex);
return BSON("count" << _count << "size" << _size << "rejectCount"
<< _numSettingsWithReject);
}
void record(int count, int size, int numSettingsWithReject) {
stdx::lock_guard<stdx::mutex> lk(_mutex);
_count = count;
_size = size;
_numSettingsWithReject = numSettingsWithReject;
}
private:
int _count = 0;
int _size = 0;
int _numSettingsWithReject = 0;
mutable stdx::mutex _mutex;
};
auto& querySettingsServerStatusSection =
*ServerStatusSectionBuilder<QuerySettingsServerStatusSection>("querySettings");
void QuerySettingsClusterParameter::append(OperationContext* opCtx,
BSONObjBuilder* bob,
StringData name,
const boost::optional<TenantId>& tenantId) {
auto& querySettingsService = QuerySettingsService::get(getGlobalServiceContext());
auto config = querySettingsService.getAllQueryShapeConfigurations(tenantId);
bob->append(QuerySettingsClusterParameterValue::k_idFieldName,
querySettingsService.getQuerySettingsClusterParameterName());
BSONArrayBuilder arrayBuilder(
bob->subarrayStart(QuerySettingsClusterParameterValue::kSettingsArrayFieldName));
for (auto&& item : config.queryShapeConfigurations) {
arrayBuilder.append(item.toBSON());
}
arrayBuilder.done();
bob->append(QuerySettingsClusterParameterValue::kClusterParameterTimeFieldName,
config.clusterParameterTime.asTimestamp());
}
Status QuerySettingsClusterParameter::set(const BSONElement& newValueElement,
const boost::optional<TenantId>& tenantId) {
auto& querySettingsService = QuerySettingsService::get(getGlobalServiceContext());
auto newSettings = QuerySettingsClusterParameterValue::parse(
IDLParserContext("querySettingsParameterValue",
boost::none /* vts */,
tenantId,
SerializationContext::stateDefault()),
newValueElement.Obj());
auto& settingsArray = newSettings.getSettingsArray();
// TODO SERVER-97546 Remove PQS index hint sanitization.
querySettingsService.sanitizeQuerySettingsHints(settingsArray);
size_t rejectCount = 0;
for (const auto& config : settingsArray) {
if (config.getSettings().getReject()) {
++rejectCount;
}
}
querySettingsServerStatusSection.record(
/* count */ static_cast<int>(settingsArray.size()),
/* size */ static_cast<int>(newValueElement.valuesize()),
/* numSettingsWithReject */ static_cast<int>(rejectCount));
querySettingsService.setAllQueryShapeConfigurations(
{std::move(settingsArray), newSettings.getClusterParameterTime()}, tenantId);
return Status::OK();
}
Status QuerySettingsClusterParameter::reset(const boost::optional<TenantId>& tenantId) {
auto& querySettingsService = QuerySettingsService::get(getGlobalServiceContext());
querySettingsService.removeAllQueryShapeConfigurations(tenantId);
return Status::OK();
}
LogicalTime QuerySettingsClusterParameter::getClusterParameterTime(
const boost::optional<TenantId>& tenantId) const {
auto& querySettingsService = QuerySettingsService::get(getGlobalServiceContext());
return querySettingsService.getClusterParameterTime(tenantId);
}
}; // namespace mongo::query_settings

343
src/mongo/db/query/query_settings/query_settings_cluster_parameter_impl_test.cpp Normal file
View File

@ -0,0 +1,343 @@
/**
* 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 <boost/move/utility_core.hpp>
#include <boost/optional.hpp>
#include <boost/optional/optional.hpp>
#include "mongo/base/error_codes.h"
#include "mongo/base/string_data.h"
#include "mongo/bson/bsonelement.h"
#include "mongo/bson/bsonobjbuilder.h"
#include "mongo/bson/bsontypes.h"
#include "mongo/bson/timestamp.h"
#include "mongo/db/client.h"
#include "mongo/db/logical_time.h"
#include "mongo/db/matcher/expression_parser.h"
#include "mongo/db/operation_context.h"
#include "mongo/db/query/query_knobs_gen.h"
#include "mongo/db/query/query_settings/query_settings_cluster_parameter_gen.h"
#include "mongo/db/query/query_settings/query_settings_gen.h"
#include "mongo/db/query/query_settings/query_settings_service.h"
#include "mongo/db/server_parameter.h"
#include "mongo/db/service_context.h"
#include "mongo/db/service_context_test_fixture.h"
#include "mongo/db/tenant_id.h"
#include "mongo/unittest/unittest.h"
#include "mongo/util/serialization_context.h"
namespace mongo::query_settings {
static bool operator==(const QueryShapeConfiguration& lhs, const QueryShapeConfiguration& rhs) {
return SimpleBSONObjComparator::kInstance.compare(lhs.toBSON(), rhs.toBSON()) == 0;
}
static bool operator==(const QueryShapeConfigurationsWithTimestamp& lhs,
const QueryShapeConfigurationsWithTimestamp& rhs) {
return lhs.clusterParameterTime == rhs.clusterParameterTime &&
lhs.queryShapeConfigurations == rhs.queryShapeConfigurations;
}
namespace {
static auto const kSerializationContext =
SerializationContext{SerializationContext::Source::Command,
SerializationContext::CallerType::Request,
SerializationContext::Prefix::ExcludePrefix};
auto makeDbName(StringData dbName) {
return DatabaseNameUtil::deserialize(
boost::none /*tenantId=*/, dbName, SerializationContext::stateDefault());
}
NamespaceSpec makeNsSpec(StringData collName) {
NamespaceSpec ns;
ns.setDb(makeDbName("testDbA"));
ns.setColl(collName);
return ns;
}
QuerySettings makeQuerySettings(const IndexHintSpecs& indexHints, bool setFramework = false) {
QuerySettings settings;
if (!indexHints.empty()) {
settings.setIndexHints(indexHints);
}
if (setFramework) {
settings.setQueryFramework(mongo::QueryFrameworkControlEnum::kTrySbeEngine);
}
return settings;
}
class QuerySettingsClusterParameterTest : public ServiceContextTest {
public:
void setUp() final {
_opCtx = cc().makeOperationContext();
}
OperationContext* opCtx() {
return _opCtx.get();
}
QuerySettingsService& service() {
return QuerySettingsService::get(opCtx());
}
QueryShapeConfiguration makeQueryShapeConfiguration(
const BSONObj& cmdBSON,
const QuerySettings& querySettings,
boost::optional<TenantId> tenantId = boost::none) {
auto queryShapeHash = createRepresentativeInfo(opCtx(), cmdBSON, tenantId).queryShapeHash;
QueryShapeConfiguration config(queryShapeHash, querySettings);
config.setRepresentativeQuery(cmdBSON);
return config;
}
BSONObj makeQuerySettingsClusterParameter(const QueryShapeConfigurationsWithTimestamp& config) {
BSONArrayBuilder bob;
for (const auto& c : config.queryShapeConfigurations) {
bob.append(c.toBSON());
}
return BSON("_id" << QuerySettingsService::getQuerySettingsClusterParameterName()
<< QuerySettingsClusterParameterValue::kSettingsArrayFieldName
<< bob.arr()
<< QuerySettingsClusterParameterValue::kClusterParameterTimeFieldName
<< config.clusterParameterTime.asTimestamp());
}
void assertTransformInvalidQuerySettings(
const std::list<std::pair<const IndexHintSpecs, const IndexHintSpecs>>&
listOfIndexHintSpecs) {
auto sp = std::make_unique<QuerySettingsClusterParameter>(
QuerySettingsService::getQuerySettingsClusterParameterName(),
ServerParameterType::kClusterWide);
QueryShapeConfigurationsWithTimestamp configsWithTs;
std::vector<QueryShapeConfiguration> expectedQueryShapeConfigurations;
size_t collIndex = 0;
// Create the 'querySettingsClusterParamValue' as BSON.
for (const auto& [initialIndexHintSpecs, expectedIndexHintSpecs] : listOfIndexHintSpecs) {
QueryInstance representativeQuery =
BSON("find" << "coll_" + std::to_string(collIndex) << "$db"
<< "foo");
configsWithTs.queryShapeConfigurations.emplace_back(makeQueryShapeConfiguration(
representativeQuery, makeQuerySettings(initialIndexHintSpecs)));
if (!expectedIndexHintSpecs.empty()) {
expectedQueryShapeConfigurations.emplace_back(makeQueryShapeConfiguration(
representativeQuery, makeQuerySettings(expectedIndexHintSpecs)));
}
collIndex++;
}
// Set the new 'clusterParamValue' for "querySettings" cluster parameter.
const auto clusterParamValue = makeQuerySettingsClusterParameter(configsWithTs);
ASSERT_OK(
sp->set(BSON("" << clusterParamValue).firstElement(), boost::none /* tenantId */));
// Assert that parsing after transforming invalid settings (if any) works.
ASSERT_EQ(expectedQueryShapeConfigurations,
service()
.getAllQueryShapeConfigurations(/* tenantId */ boost::none)
.queryShapeConfigurations);
}
private:
ServiceContext::UniqueOperationContext _opCtx;
};
/**
* Tests that set/reset methods over the QuerySettings cluster parameter correctly populate the data
* in QuerySettingsService, as well as that QuerySettings data is correctly populated when
* serializing cluster parameter.
*/
TEST_F(QuerySettingsClusterParameterTest, QuerySettingsClusterParameterSetReset) {
boost::optional<TenantId> tenantId;
auto sp = std::make_unique<QuerySettingsClusterParameter>(
QuerySettingsService::getQuerySettingsClusterParameterName(),
ServerParameterType::kClusterWide);
// Ensure no clusterParameterTime is specified for "querySettings" cluster parameter, if no
// settings are specified.
ASSERT_EQ(sp->getClusterParameterTime(tenantId), LogicalTime());
ASSERT_EQ(service().getClusterParameterTime(tenantId), LogicalTime());
{
BSONObjBuilder bob;
sp->append(
opCtx(), &bob, QuerySettingsService::getQuerySettingsClusterParameterName(), tenantId);
ASSERT_BSONOBJ_EQ(
bob.done(), makeQuerySettingsClusterParameter(QueryShapeConfigurationsWithTimestamp()));
}
// Set query shape configuration.
QuerySettings settings;
settings.setQueryFramework(QueryFrameworkControlEnum::kTrySbeEngine);
QueryInstance findCmdBSON = BSON("find" << "exampleColl"
<< "$db"
<< "foo");
auto config = makeQueryShapeConfiguration(findCmdBSON, settings);
LogicalTime clusterParameterTime(Timestamp(1, 2));
QueryShapeConfigurationsWithTimestamp configsWithTs{{config}, clusterParameterTime};
// Ensure that after parameter is set, the query shape configurations are present in the
// QuerySettingsService.
const auto clusterParamValue = makeQuerySettingsClusterParameter(configsWithTs);
ASSERT_OK(sp->set(BSON("" << clusterParamValue).firstElement(), tenantId));
ASSERT_EQ(service().getAllQueryShapeConfigurations(tenantId), configsWithTs);
ASSERT_EQ(sp->getClusterParameterTime(tenantId), clusterParameterTime);
ASSERT_EQ(service().getClusterParameterTime(tenantId), clusterParameterTime);
// Ensure the serialized parameter value contains 'settingsArray' with 'config' as value as well
// parameter id and clusterParameterTime.
{
BSONObjBuilder bob;
sp->append(
opCtx(), &bob, QuerySettingsService::getQuerySettingsClusterParameterName(), tenantId);
ASSERT_BSONOBJ_EQ(bob.done(), clusterParamValue);
}
// Ensure that after parameter is reset, no query shape configurations are present in the
// QuerySettingsService and clusterParameterTime is reset.
ASSERT_OK(sp->reset(tenantId));
ASSERT(service().getAllQueryShapeConfigurations(tenantId).queryShapeConfigurations.empty());
ASSERT_EQ(sp->getClusterParameterTime(tenantId), LogicalTime());
ASSERT_EQ(service().getClusterParameterTime(tenantId), LogicalTime());
{
BSONObjBuilder bob;
sp->append(
opCtx(), &bob, QuerySettingsService::getQuerySettingsClusterParameterName(), tenantId);
ASSERT_BSONOBJ_EQ(
bob.done(), makeQuerySettingsClusterParameter(QueryShapeConfigurationsWithTimestamp()));
}
}
/**
* Tests that query settings with invalid key-pattern index hints leads to no query settings after
* sanitization.
*/
TEST_F(QuerySettingsClusterParameterTest,
QuerySettingsClusterParameterSetInvalidQSWithSanitization) {
boost::optional<TenantId> tenantId;
auto sp = std::make_unique<QuerySettingsClusterParameter>(
QuerySettingsService::getQuerySettingsClusterParameterName(),
ServerParameterType::kClusterWide);
IndexHintSpecs initialIndexHintSpec{
IndexHintSpec(makeNsSpec("testCollC"_sd),
{IndexHint(BSON("a" << 1 << "$natural" << 1)), IndexHint(BSONObj{})}),
IndexHintSpec(makeNsSpec("testCollD"_sd),
{IndexHint(BSON("b" << "some-string"
<< "a" << 1))})};
QuerySettings querySettings;
querySettings.setIndexHints(initialIndexHintSpec);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 9646001);
// Create the 'querySettingsClusterParamValue' as BSON.
const auto findCmdBSON = BSON("find" << "bar"
<< "$db"
<< "foo");
const auto config = makeQueryShapeConfiguration(findCmdBSON, querySettings);
// Assert that parsing after transforming invalid settings (if any) works.
const auto clusterParamValue =
makeQuerySettingsClusterParameter({{config}, LogicalTime(Timestamp(3, 4))});
ASSERT_OK(sp->set(BSON("" << clusterParamValue).firstElement(), tenantId));
ASSERT(service().getAllQueryShapeConfigurations(tenantId).queryShapeConfigurations.empty());
}
/**
* Sets the value of QuerySettingsClusterParameter and asserts the resulting index hint spec is
* sanitized.
*/
TEST_F(QuerySettingsClusterParameterTest, SetValidClusterParameterAndAssertResultIsSanitized) {
IndexHintSpecs indexHintSpec{IndexHintSpec(
makeNsSpec("testCollA"), {IndexHint(BSON("a" << 1)), IndexHint(BSON("b" << -1.0))})};
assertTransformInvalidQuerySettings({{indexHintSpec, indexHintSpec}});
}
/**
* Same as above test, but assert that the validation would fail with invalid key-pattern indexes.
*/
TEST_F(QuerySettingsClusterParameterTest, SetClusterParameterAndAssertResultIsSanitized) {
IndexHintSpecs initialIndexHintSpec{
IndexHintSpec(makeNsSpec("testCollA"_sd),
{IndexHint(BSON("a" << 1.0)), IndexHint(BSON("b" << "-1.0"))}),
IndexHintSpec(makeNsSpec("testCollB"_sd),
{IndexHint(BSON("a" << 2)), IndexHint(BSONObj{})})};
IndexHintSpecs expectedIndexHintSpec{
IndexHintSpec(makeNsSpec("testCollA"_sd), {IndexHint(BSON("a" << 1))}),
IndexHintSpec(makeNsSpec("testCollB"_sd), {IndexHint(BSON("a" << 2))})};
ASSERT_THROWS_CODE(service().validateQuerySettings(makeQuerySettings(initialIndexHintSpec)),
DBException,
9646001);
assertTransformInvalidQuerySettings({{initialIndexHintSpec, expectedIndexHintSpec}});
}
/**
* Same as above test, but with multiple QuerySettings set.
*/
TEST_F(QuerySettingsClusterParameterTest, SetClusterParameterAndAssertResultIsSanitizedMultipleQS) {
// First pair of index hints.
IndexHintSpecs initialIndexHintSpec1{
IndexHintSpec(makeNsSpec("testCollA"_sd),
{IndexHint(BSON("a" << 2)), IndexHint(BSONObj{})}),
IndexHintSpec(makeNsSpec("testCollB"_sd),
{IndexHint(BSON("a" << 1.0)),
IndexHint(BSON("b" << 3)),
IndexHint("a_1"),
IndexHint(NaturalOrderHint(NaturalOrderHint::Direction::kForward))})};
IndexHintSpecs expectedIndexHintSpec1{
IndexHintSpec(makeNsSpec("testCollA"_sd), {IndexHint(BSON("a" << 2.0))}),
IndexHintSpec(makeNsSpec("testCollB"_sd),
{IndexHint(BSON("a" << 1.0)),
IndexHint(BSON("b" << 3)),
IndexHint("a_1"),
IndexHint(NaturalOrderHint(NaturalOrderHint::Direction::kForward))})};
// Second pair of index hints.
IndexHintSpecs initialIndexHintSpec2{
IndexHintSpec(makeNsSpec("testCollC"_sd),
{IndexHint(BSON("a" << 1 << "$natural" << 1)), IndexHint(BSONObj{})}),
IndexHintSpec(makeNsSpec("testCollD"_sd),
{IndexHint(BSON("b" << "some-string"
<< "a" << 1)),
IndexHint(BSONObj{})})};
ASSERT_THROWS_CODE(service().validateQuerySettings(makeQuerySettings(initialIndexHintSpec1)),
DBException,
9646000);
ASSERT_THROWS_CODE(service().validateQuerySettings(makeQuerySettings(initialIndexHintSpec2)),
DBException,
9646001);
assertTransformInvalidQuerySettings(
{{initialIndexHintSpec1, expectedIndexHintSpec1}, {initialIndexHintSpec2, {}}});
}
} // namespace
} // namespace mongo::query_settings

21
src/mongo/db/query/query_settings/query_settings_lookup_bm.cpp
View File

@ -264,11 +264,11 @@ public:
benchmark::Counter::OneK::kIs1024);
// Update the qurey shape configurations present in the system.
query_settings::setAllQueryShapeConfigurations(
opCtx.get(),
QueryShapeConfigurationsWithTimestamp{std::move(queryShapeConfigs),
LogicalTime(Timestamp(1))},
tenantId);
QuerySettingsService::get(opCtx.get())
.setAllQueryShapeConfigurations(
QueryShapeConfigurationsWithTimestamp{std::move(queryShapeConfigs),
LogicalTime(Timestamp(1))},
tenantId);
benchmark::ClobberMemory();
}
@ -311,11 +311,12 @@ public:
auto&& queryShapeConfigurations =
queryShapeConfigurationsWithTimestamp.queryShapeConfigurations;
queryShapeConfigurations.push_back(hitQueryShapeConfiguration);
query_settings::setAllQueryShapeConfigurations(
opCtx.get(),
QueryShapeConfigurationsWithTimestamp{std::move(queryShapeConfigurations),
LogicalTime(Timestamp(2))},
tid);
QuerySettingsService::get(opCtx.get())
.setAllQueryShapeConfigurations(
QueryShapeConfigurationsWithTimestamp{std::move(queryShapeConfigurations),
LogicalTime(Timestamp(2))},
tid);
}
return parsedFindRequest;

198
src/mongo/db/query/query_settings/query_settings_manager.cpp
View File

@ -32,67 +32,15 @@
#include <algorithm>
#include <boost/move/utility_core.hpp>
#include <boost/optional.hpp>
#include <boost/optional/optional.hpp>
#include "mongo/base/error_codes.h"
#include "mongo/base/status.h"
#include "mongo/base/string_data.h"
#include "mongo/bson/bsonelement.h"
#include "mongo/bson/bsonobjbuilder.h"
#include "mongo/bson/bsontypes.h"
#include "mongo/db/client.h"
#include "mongo/db/commands/server_status.h"
#include "mongo/db/logical_time.h"
#include "mongo/db/operation_context.h"
#include "mongo/db/query/query_settings/query_settings_cluster_parameter_gen.h"
#include "mongo/db/query/query_settings/query_settings_gen.h"
#include "mongo/db/service_context.h"
#include "mongo/idl/idl_parser.h"
#include "mongo/util/decorable.h"
#include "mongo/db/query/query_settings/query_settings_service.h"
namespace mongo::query_settings {
namespace {
const auto getQuerySettingsManager =
ServiceContext::declareDecoration<std::unique_ptr<QuerySettingsManager>>();
class QuerySettingsServerStatusSection final : public ServerStatusSection {
public:
using ServerStatusSection::ServerStatusSection;
bool includeByDefault() const override {
// Only include if Query Settings are enabled.
// We need to use isEnabledUseLatestFCVWhenUninitialized instead of isEnabled because
// this could run during startup while the FCV is still uninitialized.
return feature_flags::gFeatureFlagQuerySettings.isEnabledUseLatestFCVWhenUninitialized(
serverGlobalParams.featureCompatibility.acquireFCVSnapshot());
}
BSONObj generateSection(OperationContext* opCtx,
const BSONElement& configElement) const override {
stdx::lock_guard<stdx::mutex> lk(_mutex);
return BSON("count" << _count << "size" << _size << "rejectCount"
<< _numSettingsWithReject);
}
void record(int count, int size, int numSettingsWithReject) {
stdx::lock_guard<stdx::mutex> lk(_mutex);
_count = count;
_size = size;
_numSettingsWithReject = numSettingsWithReject;
}
private:
int _count = 0;
int _size = 0;
int _numSettingsWithReject = 0;
mutable stdx::mutex _mutex;
};
auto& querySettingsServerStatusSection =
*ServerStatusSectionBuilder<QuerySettingsServerStatusSection>("querySettings");
auto computeTenantConfiguration(std::vector<QueryShapeConfiguration>&& settingsArray) {
QueryShapeConfigurationsMap queryShapeConfigurationMap;
queryShapeConfigurationMap.reserve(settingsArray.size());
@ -105,22 +53,6 @@ auto computeTenantConfiguration(std::vector<QueryShapeConfiguration>&& settingsA
}
} // namespace
QuerySettingsManager& QuerySettingsManager::get(ServiceContext* service) {
return *getQuerySettingsManager(service);
}
QuerySettingsManager& QuerySettingsManager::get(OperationContext* opCtx) {
return *getQuerySettingsManager(opCtx->getServiceContext());
}
void QuerySettingsManager::create(
ServiceContext* service,
std::function<void(OperationContext*)> clusterParameterRefreshFn,
std::function<void(std::vector<QueryShapeConfiguration>&)> sanitizeQuerySettingsHintsFn) {
getQuerySettingsManager(service) = std::make_unique<QuerySettingsManager>(
service, clusterParameterRefreshFn, sanitizeQuerySettingsHintsFn);
}
boost::optional<QuerySettings> QuerySettingsManager::getQuerySettingsForQueryShapeHash(
const query_shape::QueryShapeHash& queryShapeHash,
const boost::optional<TenantId>& tenantId) const {
@ -145,13 +77,12 @@ boost::optional<QuerySettings> QuerySettingsManager::getQuerySettingsForQuerySha
return queryShapeConfigurationsIt->second.first;
}
void QuerySettingsManager::setQueryShapeConfigurations(
std::vector<QueryShapeConfiguration>&& settingsArray,
LogicalTime parameterClusterTime,
const boost::optional<TenantId>& tenantId) {
void QuerySettingsManager::setAllQueryShapeConfigurations(
QueryShapeConfigurationsWithTimestamp&& config, const boost::optional<TenantId>& tenantId) {
// Build new query shape configurations.
VersionedQueryShapeConfigurations newQueryShapeConfigurations{
computeTenantConfiguration(std::move(settingsArray)), parameterClusterTime};
computeTenantConfiguration(std::move(config.queryShapeConfigurations)),
config.clusterParameterTime};
// Install the query shape configurations.
{
@ -171,6 +102,27 @@ void QuerySettingsManager::setQueryShapeConfigurations(
}
}
void QuerySettingsManager::removeAllQueryShapeConfigurations(
const boost::optional<TenantId>& tenantId) {
// Previous query shape configurations for destruction outside the critical section.
VersionedQueryShapeConfigurations previousQueryShapeConfigurations;
{
auto writeLock = _mutex.writeLock();
const auto versionedQueryShapeConfigurationsIt =
_tenantIdToVersionedQueryShapeConfigurationsMap.find(tenantId);
if (_tenantIdToVersionedQueryShapeConfigurationsMap.end() !=
versionedQueryShapeConfigurationsIt) {
// Swap the configurations to minimize the time the lock is held in exclusive mode by
// deferring the destruction of the previous version of the query shape configurations
// to the time when the lock is not held.
std::swap(versionedQueryShapeConfigurationsIt->second,
previousQueryShapeConfigurations);
_tenantIdToVersionedQueryShapeConfigurationsMap.erase(
versionedQueryShapeConfigurationsIt);
}
}
}
QueryShapeConfigurationsWithTimestamp QuerySettingsManager::getAllQueryShapeConfigurations(
const boost::optional<TenantId>& tenantId) const {
auto readLock = _mutex.readLock();
@ -200,32 +152,6 @@ std::vector<QueryShapeConfiguration> QuerySettingsManager::getAllQueryShapeConfi
return configurations;
}
void QuerySettingsManager::refreshQueryShapeConfigurations(OperationContext* opCtx) {
if (_clusterParameterRefreshFn)
_clusterParameterRefreshFn(opCtx);
}
void QuerySettingsManager::removeAllQueryShapeConfigurations(
const boost::optional<TenantId>& tenantId) {
// Previous query shape configurations for destruction outside the critical section.
VersionedQueryShapeConfigurations previousQueryShapeConfigurations;
{
auto writeLock = _mutex.writeLock();
const auto versionedQueryShapeConfigurationsIt =
_tenantIdToVersionedQueryShapeConfigurationsMap.find(tenantId);
if (_tenantIdToVersionedQueryShapeConfigurationsMap.end() !=
versionedQueryShapeConfigurationsIt) {
// Swap the configurations to minimize the time the lock is held in exclusive mode by
// deferring the destruction of the previous version of the query shape configurations
// to the time when the lock is not held.
std::swap(versionedQueryShapeConfigurationsIt->second,
previousQueryShapeConfigurations);
_tenantIdToVersionedQueryShapeConfigurationsMap.erase(
versionedQueryShapeConfigurationsIt);
}
}
}
LogicalTime QuerySettingsManager::getClusterParameterTime(
const boost::optional<TenantId>& tenantId) const {
auto readLock = _mutex.readLock();
@ -242,76 +168,4 @@ LogicalTime QuerySettingsManager::getClusterParameterTime_inlock(
}
return versionedQueryShapeConfigurationsIt->second.clusterParameterTime;
}
void QuerySettingsManager::appendQuerySettingsClusterParameterValue(
BSONObjBuilder* bob, const boost::optional<TenantId>& tenantId) {
auto readLock = _mutex.readLock();
bob->append("_id"_sd, QuerySettingsManager::kQuerySettingsClusterParameterName);
BSONArrayBuilder arrayBuilder(
bob->subarrayStart(QuerySettingsClusterParameterValue::kSettingsArrayFieldName));
for (auto&& item : getAllQueryShapeConfigurations_inlock(tenantId)) {
arrayBuilder.append(item.toBSON());
}
arrayBuilder.done();
bob->append(QuerySettingsClusterParameterValue::kClusterParameterTimeFieldName,
getClusterParameterTime_inlock(tenantId).asTimestamp());
}
void QuerySettingsClusterParameter::append(OperationContext* opCtx,
BSONObjBuilder* bob,
StringData name,
const boost::optional<TenantId>& tenantId) {
auto& querySettingsManager = QuerySettingsManager::get(getGlobalServiceContext());
querySettingsManager.appendQuerySettingsClusterParameterValue(bob, tenantId);
}
Status QuerySettingsClusterParameter::set(const BSONElement& newValueElement,
const boost::optional<TenantId>& tenantId) {
auto& querySettingsManager = QuerySettingsManager::get(getGlobalServiceContext());
auto newSettings = QuerySettingsClusterParameterValue::parse(
IDLParserContext("querySettingsParameterValue",
boost::none /* vts */,
tenantId,
SerializationContext::stateDefault()),
newValueElement.Obj());
auto& settingsArray = newSettings.getSettingsArray();
// TODO SERVER-97546 Remove PQS index hint sanitization.
querySettingsManager.sanitizeQuerySettingsHints(settingsArray);
size_t rejectCount = 0;
for (const auto& config : settingsArray) {
if (config.getSettings().getReject()) {
++rejectCount;
}
}
querySettingsServerStatusSection.record(
/* count */ static_cast<int>(settingsArray.size()),
/* size */ static_cast<int>(newValueElement.valuesize()),
/* numSettingsWithReject */ static_cast<int>(rejectCount));
querySettingsManager.setQueryShapeConfigurations(
std::move(settingsArray), newSettings.getClusterParameterTime(), tenantId);
return Status::OK();
}
Status QuerySettingsClusterParameter::reset(const boost::optional<TenantId>& tenantId) {
auto& querySettingsManager = QuerySettingsManager::get(getGlobalServiceContext());
querySettingsManager.removeAllQueryShapeConfigurations(tenantId);
return Status::OK();
}
LogicalTime QuerySettingsClusterParameter::getClusterParameterTime(
const boost::optional<TenantId>& tenantId) const {
auto& querySettingsManager = QuerySettingsManager::get(getGlobalServiceContext());
return querySettingsManager.getClusterParameterTime(tenantId);
}
// TODO SERVER-97546 Remove PQS index hint sanitization.
void QuerySettingsManager::sanitizeQuerySettingsHints(
std::vector<QueryShapeConfiguration>& queryShapeConfigs) {
if (_sanitizeQuerySettingsHintsFn) {
_sanitizeQuerySettingsHintsFn(queryShapeConfigs);
}
}
}; // namespace mongo::query_settings

54
src/mongo/db/query/query_settings/query_settings_manager.h
View File

@ -94,38 +94,20 @@ struct VersionedQueryShapeConfigurations {
};
/**
* Class responsible for managing in-memory storage and fetching of query settings. The in-memory
* storage is eventually consistent with the query settings on other cluster nodes and is updated
* based on OpObserver call performed when executing setClusterParameter command.
*
* Query settings in-memory storage is maintained separately for each tenant. In dedicated
* environments the 'tenantId' argument passed to the methods must be boost::none.
*
* Query settings should only be retrieved through this class.
* Class responsible for managing in-memory storage and fetching of query settings. Query settings
* in-memory storage is maintained separately for each tenant. In dedicated environments the
* 'tenantId' argument passed to the methods must be boost::none.
*/
class QuerySettingsManager {
public:
static constexpr auto kQuerySettingsClusterParameterName = "querySettings"_sd;
QuerySettingsManager(
ServiceContext* service,
std::function<void(OperationContext*)> clusterParameterRefreshFn,
std::function<void(std::vector<QueryShapeConfiguration>&)> sanitizeQuerySettingsHintsFn)
: _clusterParameterRefreshFn(clusterParameterRefreshFn),
_sanitizeQuerySettingsHintsFn(sanitizeQuerySettingsHintsFn) {}
QuerySettingsManager() = default;
~QuerySettingsManager() = default;
QuerySettingsManager(QuerySettingsManager&&) = delete;
QuerySettingsManager(const QuerySettingsManager&) = delete;
QuerySettingsManager& operator=(QuerySettingsManager&&) = delete;
QuerySettingsManager& operator=(const QuerySettingsManager&) = delete;
static void create(
ServiceContext* service,
std::function<void(OperationContext*)> clusterParameterRefreshFn,
std::function<void(std::vector<QueryShapeConfiguration>&)> sanitizeQuerySettingsHintsFn);
static QuerySettingsManager& get(ServiceContext* service);
static QuerySettingsManager& get(OperationContext* opCtx);
/**
* Returns QuerySettings associated with a query which query shape hash is 'queryShapeHash' for
* the given tenant.
@ -145,15 +127,8 @@ public:
* Sets the QueryShapeConfigurations by replacing an existing VersionedQueryShapeConfigurations
* with the newly built one.
*/
void setQueryShapeConfigurations(std::vector<QueryShapeConfiguration>&& settings,
LogicalTime parameterClusterTime,
const boost::optional<TenantId>& tenantId);
/**
* Updates the QueryShapeConfiguration cache by calling the 'clusterParameterRefreshFn'. This
* can be a no-op if no 'clusterParameterRefreshFn' was provided in the constructor.
*/
void refreshQueryShapeConfigurations(OperationContext* opCtx);
void setAllQueryShapeConfigurations(QueryShapeConfigurationsWithTimestamp&& config,
const boost::optional<TenantId>& tenantId);
/**
* Removes all query settings documents for the given tenant.
@ -166,16 +141,6 @@ public:
*/
LogicalTime getClusterParameterTime(const boost::optional<TenantId>& tenantId) const;
/**
* Appends the QuerySettingsClusterParameterValue maintained as
* VersionedQueryShapeConfigurations for the given tenant.
*/
void appendQuerySettingsClusterParameterValue(BSONObjBuilder* bob,
const boost::optional<TenantId>& tenantId);
// TODO SERVER-97546 Remove PQS index hint sanitization.
void sanitizeQuerySettingsHints(std::vector<QueryShapeConfiguration>& queryShapeConfigs);
private:
std::vector<QueryShapeConfiguration> getAllQueryShapeConfigurations_inlock(
const boost::optional<TenantId>& tenantId) const;
@ -185,9 +150,6 @@ private:
mutable WriteRarelyRWMutex _mutex;
absl::flat_hash_map<boost::optional<TenantId>, VersionedQueryShapeConfigurations>
_tenantIdToVersionedQueryShapeConfigurationsMap;
std::function<void(OperationContext*)> _clusterParameterRefreshFn;
// TODO SERVER-97546 Remove PQS index hint sanitization.
std::function<void(std::vector<QueryShapeConfiguration>&)> _sanitizeQuerySettingsHintsFn;
};
}; // namespace query_settings
} // namespace mongo

424
src/mongo/db/query/query_settings/query_settings_manager_test.cpp
View File

@ -27,31 +27,24 @@
* it in the license file.
*/
#include <algorithm>
#include <iterator>
#include <vector>
#include <boost/move/utility_core.hpp>
#include <boost/none.hpp>
#include <boost/optional/optional.hpp>
#include <boost/smart_ptr/intrusive_ptr.hpp>
#include <vector>
#include "mongo/base/string_data.h"
#include "mongo/bson/bsonmisc.h"
#include "mongo/bson/bsonobj.h"
#include "mongo/bson/oid.h"
#include "mongo/bson/simple_bsonobj_comparator.h"
#include "mongo/bson/timestamp.h"
#include "mongo/db/client.h"
#include "mongo/db/logical_time.h"
#include "mongo/db/pipeline/expression_context.h"
#include "mongo/db/query/index_hint.h"
#include "mongo/db/query/query_knobs_gen.h"
#include "mongo/db/query/query_settings/query_settings_cluster_parameter_gen.h"
#include "mongo/db/query/query_settings/query_settings_gen.h"
#include "mongo/db/query/query_settings/query_settings_manager.h"
#include "mongo/db/query/query_settings/query_settings_service.h"
#include "mongo/db/query/query_shape/query_shape.h"
#include "mongo/db/server_parameter.h"
#include "mongo/db/service_context_test_fixture.h"
#include "mongo/db/tenant_id.h"
@ -62,67 +55,19 @@
namespace mongo::query_settings {
namespace {
QueryShapeConfiguration makeQueryShapeConfiguration(const QuerySettings& settings,
QueryInstance query,
OperationContext* opCtx,
boost::optional<TenantId> tenantId) {
auto queryShapeHash = createRepresentativeInfo(opCtx, query, tenantId).queryShapeHash;
QueryShapeConfiguration result(queryShapeHash, settings);
result.setRepresentativeQuery(query);
return result;
bool operator==(const QuerySettings& lhs, const QuerySettings& rhs) {
return SimpleBSONObjComparator::kInstance.compare(lhs.toBSON(), rhs.toBSON()) == 0;
}
// QueryShapeConfiguration is not comparable, therefore comparing the corresponding
// BSONObj encoding.
void assertQueryShapeConfigurationsEquals(
const std::vector<QueryShapeConfiguration>& expectedQueryShapeConfigurations,
const std::vector<QueryShapeConfiguration>& actualQueryShapeConfigurations) {
std::vector<BSONObj> lhs, rhs;
std::transform(expectedQueryShapeConfigurations.begin(),
expectedQueryShapeConfigurations.end(),
std::back_inserter(lhs),
[](auto x) { return x.toBSON(); });
std::transform(actualQueryShapeConfigurations.begin(),
actualQueryShapeConfigurations.end(),
std::back_inserter(rhs),
[](auto x) { return x.toBSON(); });
std::sort(lhs.begin(), lhs.end(), SimpleBSONObjComparator::kInstance.makeLessThan());
std::sort(rhs.begin(), rhs.end(), SimpleBSONObjComparator::kInstance.makeLessThan());
ASSERT(std::equal(
lhs.begin(), lhs.end(), rhs.begin(), SimpleBSONObjComparator::kInstance.makeEqualTo()));
}
BSONObj makeSettingsClusterParameter(const BSONArray& settings) {
LogicalTime clusterParameterTime(Timestamp(113, 59));
return BSON("_id" << QuerySettingsManager::kQuerySettingsClusterParameterName
<< QuerySettingsClusterParameterValue::kSettingsArrayFieldName << settings
<< QuerySettingsClusterParameterValue::kClusterParameterTimeFieldName
<< clusterParameterTime.asTimestamp());
}
QuerySettings makeQuerySettings(const IndexHintSpecs& indexHints, bool setFramework = true) {
QuerySettings makeQuerySettings(const IndexHintSpecs& indexHints) {
QuerySettings settings;
if (!indexHints.empty()) {
settings.setIndexHints(indexHints);
}
if (setFramework) {
settings.setQueryFramework(mongo::QueryFrameworkControlEnum::kTrySbeEngine);
}
settings.setQueryFramework(mongo::QueryFrameworkControlEnum::kTrySbeEngine);
return settings;
}
auto makeDbName(StringData dbName) {
return DatabaseNameUtil::deserialize(
boost::none /*tenantId=*/, dbName, SerializationContext::stateDefault());
}
NamespaceSpec makeNsSpec(StringData collName) {
NamespaceSpec ns;
ns.setDb(makeDbName("testDbA"));
ns.setColl(collName);
return ns;
}
} // namespace
static auto const kSerializationContext =
@ -135,8 +80,7 @@ public:
static constexpr StringData kCollName = "exampleCol"_sd;
static constexpr StringData kDbName = "foo"_sd;
static std::vector<QueryShapeConfiguration> getExampleQueryShapeConfigurations(
OperationContext* opCtx, TenantId tenantId) {
std::vector<QueryShapeConfiguration> getExampleQueryShapeConfigurations(TenantId tenantId) {
NamespaceSpec ns;
ns.setDb(DatabaseNameUtil::deserialize(tenantId, kDbName, kSerializationContext));
ns.setColl(kCollName);
@ -146,28 +90,29 @@ public:
BSON("find" << kCollName << "$db" << kDbName << "filter" << BSON("a" << 2));
QueryInstance queryB =
BSON("find" << kCollName << "$db" << kDbName << "filter" << BSON("a" << BSONNULL));
return {makeQueryShapeConfiguration(settings, queryA, opCtx, tenantId),
makeQueryShapeConfiguration(settings, queryB, opCtx, tenantId)};
return {makeQueryShapeConfiguration(settings, queryA, tenantId),
makeQueryShapeConfiguration(settings, queryB, tenantId)};
}
QueryShapeConfiguration makeQueryShapeConfiguration(const QuerySettings& settings,
QueryInstance query,
boost::optional<TenantId> tenantId) {
auto queryShapeHash = createRepresentativeInfo(opCtx(), query, tenantId).queryShapeHash;
QueryShapeConfiguration result(queryShapeHash, settings);
result.setRepresentativeQuery(query);
return result;
}
void setUp() final {
QuerySettingsManager::create(
getServiceContext(), {}, query_settings::sanitizeQuerySettingsHints);
_opCtx = cc().makeOperationContext();
_expCtx = ExpressionContext::makeBlankExpressionContext(opCtx(), {NamespaceString()});
}
OperationContext* opCtx() {
return _opCtx.get();
}
boost::intrusive_ptr<ExpressionContext> expCtx() {
return _expCtx;
}
QuerySettingsManager& manager() {
return QuerySettingsManager::get(opCtx());
return _manager;
}
static NamespaceString nss(boost::optional<TenantId> tenantId) {
@ -180,331 +125,32 @@ public:
tenantId, kDbName, kCollName, kSerializationContext);
}
void assertTransformInvalidQuerySettings(
const std::list<std::pair<const IndexHintSpecs, const IndexHintSpecs>>&
listOfIndexHintSpecs) {
auto sp = std::make_unique<QuerySettingsClusterParameter>(
"querySettings", ServerParameterType::kClusterWide);
BSONArrayBuilder bob;
std::vector<QueryShapeConfiguration> expectedQueryShapeConfigurations;
size_t collIndex = 0;
// Create the 'querySettingsClusterParamValue' as BSON.
for (const auto& [initialIndexHintSpecs, expectedIndexHintSpecs] : listOfIndexHintSpecs) {
QueryInstance representativeQuery =
BSON("find" << "coll_" + std::to_string(collIndex) << "$db"
<< "foo");
bob.append(makeQueryShapeConfiguration(makeQuerySettings(initialIndexHintSpecs),
representativeQuery,
opCtx(),
/* tenantId */ boost::none)
.toBSON());
expectedQueryShapeConfigurations.emplace_back(
makeQueryShapeConfiguration(makeQuerySettings(expectedIndexHintSpecs),
representativeQuery,
opCtx(),
/* tenantId */ boost::none));
collIndex++;
}
// Set the cluster param value.
const auto clusterParamValues = BSON_ARRAY(makeSettingsClusterParameter(bob.arr()));
// Assert that parsing after transforming invalid settings (if any) works.
ASSERT_OK(sp->set(clusterParamValues.firstElement(), boost::none));
assertQueryShapeConfigurationsEquals(
expectedQueryShapeConfigurations,
manager()
.getAllQueryShapeConfigurations(/* tenantId */ boost::none)
.queryShapeConfigurations);
}
void assertSanitizeInvalidIndexHints(const IndexHintSpecs& initialSpec,
const IndexHintSpecs& expectedSpec) {
QueryInstance query = BSON("find" << "exampleColl"
<< "$db"
<< "foo");
auto initSettings = makeQueryShapeConfiguration(makeQuerySettings(initialSpec),
query,
opCtx(),
/* tenantId */ boost::none);
std::vector<QueryShapeConfiguration> queryShapeConfigs{initSettings};
const auto expectedSettings = makeQuerySettings(expectedSpec);
ASSERT_DOES_NOT_THROW(sanitizeQuerySettingsHints(queryShapeConfigs));
ASSERT_BSONOBJ_EQ(queryShapeConfigs[0].getSettings().toBSON(), expectedSettings.toBSON());
}
private:
ServiceContext::UniqueOperationContext _opCtx;
boost::intrusive_ptr<ExpressionContext> _expCtx;
QuerySettingsManager _manager;
};
TEST_F(QuerySettingsManagerTest, QuerySettingsClusterParameterSerialization) {
// Set query shape configuration.
QuerySettings settings;
settings.setQueryFramework(QueryFrameworkControlEnum::kTrySbeEngine);
QueryInstance query = BSON("find" << "exampleColl"
<< "$db"
<< "foo");
auto config = makeQueryShapeConfiguration(settings, query, opCtx(), /* tenantId */ boost::none);
LogicalTime clusterParameterTime(Timestamp(113, 59));
manager().setQueryShapeConfigurations(
{config}, clusterParameterTime, /* tenantId */ boost::none);
// Ensure the serialized parameter value contains 'settingsArray' with 'config' as value as well
// parameter id and clusterParameterTime.
BSONObjBuilder bob;
manager().appendQuerySettingsClusterParameterValue(&bob, /* tenantId */ boost::none);
ASSERT_BSONOBJ_EQ(
bob.done(),
BSON("_id" << QuerySettingsManager::kQuerySettingsClusterParameterName
<< QuerySettingsClusterParameterValue::kSettingsArrayFieldName
<< BSON_ARRAY(config.toBSON())
<< QuerySettingsClusterParameterValue::kClusterParameterTimeFieldName
<< clusterParameterTime.asTimestamp()));
}
TEST_F(QuerySettingsManagerTest, QuerySettingsSetAndReset) {
RAIIServerParameterControllerForTest multitenanyController("multitenancySupport", true);
LogicalTime firstWriteTime(Timestamp(1, 0)), secondWriteTime(Timestamp(2, 0));
TenantId tenantId(OID::fromTerm(1)), otherTenantId(OID::fromTerm(2));
auto firstConfig = getExampleQueryShapeConfigurations(opCtx(), tenantId)[0];
auto secondConfig = getExampleQueryShapeConfigurations(opCtx(), otherTenantId)[1];
// Ensure that the maintained in-memory query shape configurations equal to the
// configurations specified in the parameter for both tenants.
manager().setQueryShapeConfigurations({firstConfig}, firstWriteTime, tenantId);
manager().setQueryShapeConfigurations({firstConfig}, firstWriteTime, otherTenantId);
assertQueryShapeConfigurationsEquals(
{firstConfig}, manager().getAllQueryShapeConfigurations(tenantId).queryShapeConfigurations);
assertQueryShapeConfigurationsEquals(
{firstConfig},
manager().getAllQueryShapeConfigurations(otherTenantId).queryShapeConfigurations);
ASSERT_EQ(manager().getClusterParameterTime(tenantId), firstWriteTime);
ASSERT_EQ(manager().getClusterParameterTime(otherTenantId), firstWriteTime);
// Update query settings for tenant with 'tenantId'. Ensure its query shape configurations and
// parameter cluster time are updated accordingly.
manager().setQueryShapeConfigurations({secondConfig}, secondWriteTime, tenantId);
assertQueryShapeConfigurationsEquals(
{secondConfig},
manager().getAllQueryShapeConfigurations(tenantId).queryShapeConfigurations);
assertQueryShapeConfigurationsEquals(
{firstConfig},
manager().getAllQueryShapeConfigurations(otherTenantId).queryShapeConfigurations);
ASSERT_EQ(manager().getClusterParameterTime(tenantId), secondWriteTime);
ASSERT_EQ(manager().getClusterParameterTime(otherTenantId), firstWriteTime);
// Reset the parameter value and ensure that the in-memory storage is cleared for tenant with
// 'tenantId'.
manager().removeAllQueryShapeConfigurations(tenantId);
assertQueryShapeConfigurationsEquals(
{}, manager().getAllQueryShapeConfigurations(tenantId).queryShapeConfigurations);
// Attempt to remove QueryShapeConfigurations for a tenant that does not have any.
manager().removeAllQueryShapeConfigurations(tenantId);
assertQueryShapeConfigurationsEquals(
{}, manager().getAllQueryShapeConfigurations(tenantId).queryShapeConfigurations);
// Verify that QueryShapeConfigurations for tenant with 'otherTenantId' were not be affected.
assertQueryShapeConfigurationsEquals(
{firstConfig},
manager().getAllQueryShapeConfigurations(otherTenantId).queryShapeConfigurations);
ASSERT_EQ(manager().getClusterParameterTime(tenantId), LogicalTime::kUninitialized);
ASSERT_EQ(manager().getClusterParameterTime(otherTenantId), firstWriteTime);
}
TEST_F(QuerySettingsManagerTest, QuerySettingsLookup) {
using Result = boost::optional<QuerySettings>;
RAIIServerParameterControllerForTest multitenanyController("multitenancySupport", true);
TenantId tenantId1(OID::fromTerm(1));
TenantId tenantId2(OID::fromTerm(2));
// Helper function for ensuring that two
// 'QuerySettingsManager::getQuerySettingsForQueryShapeHash()' results are identical.
auto assertResultsEq = [](Result r0, Result r1) {
// Ensure that either both results are there, or both are missing.
ASSERT_EQ(r0.has_value(), r1.has_value());
// Early exit if any of them are missing.
if (!r0.has_value() || !r1.has_value()) {
return;
}
ASSERT_BSONOBJ_EQ(r0->toBSON(), r1->toBSON());
};
TenantId tenantId(OID::fromTerm(1));
auto configs = getExampleQueryShapeConfigurations(opCtx(), tenantId);
manager().setQueryShapeConfigurations(
std::vector<QueryShapeConfiguration>(configs), LogicalTime(), tenantId);
auto configs = getExampleQueryShapeConfigurations(tenantId1);
manager().setAllQueryShapeConfigurations({{configs}, LogicalTime()}, tenantId1);
// Ensure QuerySettingsManager returns boost::none when QuerySettings are not found.
assertResultsEq(
manager().getQuerySettingsForQueryShapeHash(query_shape::QueryShapeHash(), tenantId),
boost::none);
ASSERT_EQ(manager().getQuerySettingsForQueryShapeHash(query_shape::QueryShapeHash(), tenantId1),
boost::none);
// Ensure QuerySettingsManager returns a valid QuerySettings on lookup.
assertResultsEq(
manager().getQuerySettingsForQueryShapeHash(configs[1].getQueryShapeHash(), tenantId),
ASSERT_EQ(
manager().getQuerySettingsForQueryShapeHash(configs[1].getQueryShapeHash(), tenantId1),
configs[1].getSettings());
// Ensure QuerySettingsManager does not return a valid QuerySettings of 'tenantId1', when
// performing lookup as 'tenantId2'.
ASSERT_EQ(
manager().getQuerySettingsForQueryShapeHash(configs[1].getQueryShapeHash(), tenantId2),
boost::none);
}
/**
* Tests that valid index hint specs are the same before and after index hint sanitization.
*/
TEST_F(QuerySettingsManagerTest, ValidIndexHintsAreTheSameBeforeAndAfterSanitization) {
IndexHintSpecs indexHintSpec{IndexHintSpec(
makeNsSpec("testCollA"_sd), {IndexHint(BSON("a" << 1)), IndexHint(BSON("b" << -1.0))})};
assertSanitizeInvalidIndexHints(indexHintSpec, indexHintSpec);
}
/**
* Tests that invalid key-pattern are removed after sanitization.
*/
TEST_F(QuerySettingsManagerTest, InvalidKeyPatternIndexesAreRemovedAfterSanitization) {
IndexHintSpecs indexHintSpec{IndexHintSpec(makeNsSpec("testCollA"_sd),
{IndexHint(BSON("a" << 1 << "c"
<< "invalid")),
IndexHint(BSON("b" << -1.0))})};
IndexHintSpecs expectedHintSpec{
IndexHintSpec(makeNsSpec("testCollA"_sd), {IndexHint(BSON("b" << -1.0))})};
assertSanitizeInvalidIndexHints(indexHintSpec, expectedHintSpec);
}
/**
* Same as the above test but with more complex examples.
*/
TEST_F(QuerySettingsManagerTest, InvalidKeyPatternIndexesAreRemovedAfterSanitizationComplex) {
IndexHintSpecs indexHintSpec{IndexHintSpec(makeNsSpec("testCollA"_sd),
{
IndexHint(BSON("a" << 1 << "c"
<< "invalid")),
IndexHint(BSON("b" << -1.0)),
IndexHint(BSON("c" << -2.0 << "b" << 4)),
IndexHint("index_name"),
IndexHint(BSON("$natural" << 1)),
IndexHint(BSON("$natural" << -1 << "a" << 2)),
IndexHint(BSON("a" << -1 << "$natural" << 1)),
})};
IndexHintSpecs expectedHintSpec{IndexHintSpec(makeNsSpec("testCollA"_sd),
{IndexHint(BSON("b" << -1.0)),
IndexHint(BSON("c" << -2.0 << "b" << 4)),
IndexHint("index_name")})};
assertSanitizeInvalidIndexHints(indexHintSpec, expectedHintSpec);
}
/**
* Tests that invalid key-pattern are removed after sanitization and the resulted index hints are
* empty.
*/
TEST_F(QuerySettingsManagerTest, InvalidKeyPatternIndexesAreRemovedAfterSanitizationEmptyHints) {
IndexHintSpecs indexHintSpec{IndexHintSpec(makeNsSpec("testCollA"_sd),
{
IndexHint(BSON("a" << 1 << "c"
<< "invalid")),
})};
IndexHintSpecs expectedHintSpec;
assertSanitizeInvalidIndexHints(indexHintSpec, expectedHintSpec);
}
/**
* Sets the value of QuerySettingsClusterParameter and asserts the resulting index hint spec is
* sanitized.
*/
TEST_F(QuerySettingsManagerTest, SetValidClusterParameterAndAssertResultIsSanitized) {
IndexHintSpecs indexHintSpec{IndexHintSpec(
makeNsSpec("testCollA"), {IndexHint(BSON("a" << 1)), IndexHint(BSON("b" << -1.0))})};
assertTransformInvalidQuerySettings({{indexHintSpec, indexHintSpec}});
}
/**
* Same as above test, but assert that the validation would fail with invalid key-pattern indexes.
*/
TEST_F(QuerySettingsManagerTest, SetClusterParameterAndAssertResultIsSanitized) {
IndexHintSpecs initialIndexHintSpec{
IndexHintSpec(makeNsSpec("testCollA"_sd),
{IndexHint(BSON("a" << 1.0)), IndexHint(BSON("b" << "-1.0"))}),
IndexHintSpec(makeNsSpec("testCollB"_sd),
{IndexHint(BSON("a" << 2)), IndexHint(BSONObj{})})};
IndexHintSpecs expectedIndexHintSpec{
IndexHintSpec(makeNsSpec("testCollA"_sd), {IndexHint(BSON("a" << 1))}),
IndexHintSpec(makeNsSpec("testCollB"_sd), {IndexHint(BSON("a" << 2))})};
ASSERT_THROWS_CODE(
validateQuerySettings(makeQuerySettings(initialIndexHintSpec)), DBException, 9646001);
assertTransformInvalidQuerySettings({{initialIndexHintSpec, expectedIndexHintSpec}});
}
/**
* Same as above test, but with multiple QuerySettings set.
*/
TEST_F(QuerySettingsManagerTest, SetClusterParameterAndAssertResultIsSanitizedMultipleQS) {
// First pair of index hints.
IndexHintSpecs initialIndexHintSpec1{
IndexHintSpec(makeNsSpec("testCollA"_sd),
{IndexHint(BSON("a" << 2)), IndexHint(BSONObj{})}),
IndexHintSpec(makeNsSpec("testCollB"_sd),
{IndexHint(BSON("a" << 1.0)),
IndexHint(BSON("b" << 3)),
IndexHint("a_1"),
IndexHint(NaturalOrderHint(NaturalOrderHint::Direction::kForward))})};
IndexHintSpecs expectedIndexHintSpec1{
IndexHintSpec(makeNsSpec("testCollA"_sd), {IndexHint(BSON("a" << 2.0))}),
IndexHintSpec(makeNsSpec("testCollB"_sd),
{IndexHint(BSON("a" << 1.0)),
IndexHint(BSON("b" << 3)),
IndexHint("a_1"),
IndexHint(NaturalOrderHint(NaturalOrderHint::Direction::kForward))})};
// Second pair of index hints.
IndexHintSpecs initialIndexHintSpec2{
IndexHintSpec(makeNsSpec("testCollC"_sd),
{IndexHint(BSON("a" << 1 << "$natural" << 1)), IndexHint(BSONObj{})}),
IndexHintSpec(makeNsSpec("testCollD"_sd),
{IndexHint(BSON("b" << "some-string"
<< "a" << 1)),
IndexHint(BSONObj{})})};
ASSERT_THROWS_CODE(
validateQuerySettings(makeQuerySettings(initialIndexHintSpec1)), DBException, 9646000);
ASSERT_THROWS_CODE(
validateQuerySettings(makeQuerySettings(initialIndexHintSpec2)), DBException, 9646001);
assertTransformInvalidQuerySettings(
{{initialIndexHintSpec1, expectedIndexHintSpec1}, {initialIndexHintSpec2, {}}});
}
/**
* Tests that query settings with invalid key-pattern index hints leads to no query settings after
* sanitization.
*/
TEST_F(QuerySettingsManagerTest, SetClusterParameterInvalidQSSanitization) {
auto sp = std::make_unique<QuerySettingsClusterParameter>("querySettings",
ServerParameterType::kClusterWide);
IndexHintSpecs initialIndexHintSpec{
IndexHintSpec(makeNsSpec("testCollC"_sd),
{IndexHint(BSON("a" << 1 << "$natural" << 1)), IndexHint(BSONObj{})}),
IndexHintSpec(makeNsSpec("testCollD"_sd),
{IndexHint(BSON("b" << "some-string"
<< "a" << 1))})};
const auto querySettings = makeQuerySettings(initialIndexHintSpec, false);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 9646001);
// Create the 'querySettingsClusterParamValue' as BSON.
const auto config = makeQueryShapeConfiguration(querySettings,
BSON("find" << "bar"
<< "$db"
<< "foo"),
opCtx(),
/* tenantId */ boost::none);
// Assert that parsing after transforming invalid settings (if any) works.
const auto clusterParamValues =
BSON_ARRAY(makeSettingsClusterParameter(BSON_ARRAY(config.toBSON())));
ASSERT_OK(sp->set(clusterParamValues.firstElement(), boost::none));
auto res = manager().getQuerySettingsForQueryShapeHash(config.getQueryShapeHash(),
/* tenantId */ boost::none);
ASSERT(!res.has_value());
ASSERT(manager()
.getAllQueryShapeConfigurations(/* tenantId */ boost::none)
.queryShapeConfigurations.empty());
}
} // namespace mongo::query_settings

275
src/mongo/db/query/query_settings/query_settings_service.cpp
View File

@ -75,6 +75,11 @@ const stdx::unordered_set<StringData, StringMapHasher> rejectionIncompatibleStag
"$operationMetrics"_sd,
};
const auto getQuerySettingsService =
ServiceContext::declareDecoration<std::unique_ptr<QuerySettingsService>>();
static constexpr auto kQuerySettingsClusterParameterName = "querySettings"_sd;
MONGO_FAIL_POINT_DEFINE(allowAllSetQuerySettings);
/**
@ -380,6 +385,149 @@ void sanitizeKeyPatternIndexHints(QueryShapeConfiguration& queryShapeItem) {
}
} // namespace
class QuerySettingsRouterService : public QuerySettingsService {
public:
QuerySettingsRouterService() : QuerySettingsService() {}
QueryShapeConfigurationsWithTimestamp getAllQueryShapeConfigurations(
const boost::optional<TenantId>& tenantId) const final {
return _manager.getAllQueryShapeConfigurations(tenantId);
}
void setAllQueryShapeConfigurations(QueryShapeConfigurationsWithTimestamp&& config,
const boost::optional<TenantId>& tenantId) final {
_manager.setAllQueryShapeConfigurations(std::move(config), tenantId);
}
void removeAllQueryShapeConfigurations(const boost::optional<TenantId>& tenantId) final {
_manager.removeAllQueryShapeConfigurations(tenantId);
}
LogicalTime getClusterParameterTime(const boost::optional<TenantId>& tenantId) const final {
return _manager.getClusterParameterTime(tenantId);
}
QuerySettings lookupQuerySettingsWithRejectionCheck(
const boost::intrusive_ptr<ExpressionContext>& expCtx,
const query_shape::DeferredQueryShape& deferredShape,
const NamespaceString& nss,
const boost::optional<QuerySettings>& querySettingsFromOriginalCommand =
boost::none) const override {
QuerySettings settings = [&]() {
try {
// No query settings lookup for IDHACK queries.
if (expCtx->isIdHackQuery()) {
return QuerySettings();
}
// No query settings for queries with encryption information.
if (expCtx->isFleQuery()) {
return QuerySettings();
}
// No query settings lookup on internal dbs or system collections in user dbs.
if (nss.isOnInternalDb() || nss.isSystem()) {
return QuerySettings();
}
// Force shape computation and early exit with empty settings if shape computation
// has failed.
const auto& shapePtr = deferredShape();
if (!shapePtr.isOK()) {
return QuerySettings();
}
// Compute the QueryShapeHash and store it in CurOp.
auto* opCtx = expCtx->getOperationContext();
QueryShapeHash hash = shapePtr.getValue()->sha256Hash(opCtx, kSerializationContext);
setQueryShapeHash(opCtx, hash);
// Return the found query settings or an empty one.
return _manager.getQuerySettingsForQueryShapeHash(hash, nss.tenantId())
.get_value_or({});
} catch (const DBException& ex) {
LOGV2_WARNING_OPTIONS(10153400,
{logv2::LogComponent::kQuery},
"Failed to perform query settings lookup",
"error"_attr = ex.toString());
return QuerySettings();
}
}();
// Fail the current command, if 'reject: true' flag is present.
failIfRejectedBySettings(expCtx, settings);
return settings;
}
void refreshQueryShapeConfigurations(OperationContext* opCtx) override {
// QuerySettingsManager modifies a cluster-wide parameter and thus a refresh of
// the parameter after that modification should observe results of preceeding
// writes.
const bool kEnsureReadYourWritesConsistency = true;
auto refreshStatus = ClusterServerParameterRefresher::get(opCtx)->refreshParameters(
opCtx, kEnsureReadYourWritesConsistency);
if (!refreshStatus.isOK()) {
LOGV2_WARNING(8472500,
"Error occurred when fetching the latest version of query settings",
"error_code"_attr = refreshStatus.code(),
"reason"_attr = refreshStatus.reason());
}
}
private:
QuerySettingsManager _manager;
};
class QuerySettingsShardService : public QuerySettingsRouterService {
public:
QuerySettingsShardService() : QuerySettingsRouterService() {}
QuerySettings lookupQuerySettingsWithRejectionCheck(
const boost::intrusive_ptr<ExpressionContext>& expCtx,
const query_shape::DeferredQueryShape& queryShape,
const NamespaceString& nss,
const boost::optional<QuerySettings>& querySettingsFromOriginalCommand =
boost::none) const override {
// No query settings lookup for IDHACK queries.
if (expCtx->isIdHackQuery()) {
return QuerySettings();
}
auto* opCtx = expCtx->getOperationContext();
if (requestComesFromRouterOrSentDirectlyToShard(opCtx->getClient()) ||
querySettingsFromOriginalCommand.has_value()) {
return querySettingsFromOriginalCommand.get_value_or(QuerySettings());
}
// The underlying shard does not belong to a sharded cluster, therefore proceed by
// performing the lookup as a router.
return QuerySettingsRouterService::lookupQuerySettingsWithRejectionCheck(
expCtx, queryShape, nss);
}
void refreshQueryShapeConfigurations(OperationContext* opCtx) override {
/* no-op */
}
};
QuerySettingsService& QuerySettingsService::get(ServiceContext* service) {
return *getQuerySettingsService(service);
}
QuerySettingsService& QuerySettingsService::get(OperationContext* opCtx) {
return *getQuerySettingsService(opCtx->getServiceContext());
}
std::string QuerySettingsService::getQuerySettingsClusterParameterName() {
return kQuerySettingsClusterParameterName.toString();
}
const stdx::unordered_set<StringData, StringMapHasher>&
QuerySettingsService::getRejectionIncompatibleStages() {
return rejectionIncompatibleStages;
};
RepresentativeQueryInfo createRepresentativeInfo(OperationContext* opCtx,
const BSONObj& cmd,
const boost::optional<TenantId>& tenantId) {
@ -397,81 +545,25 @@ RepresentativeQueryInfo createRepresentativeInfo(OperationContext* opCtx,
}
void initializeForRouter(ServiceContext* serviceContext) {
QuerySettingsManager::create(
serviceContext,
[](OperationContext* opCtx) {
// QuerySettingsManager modifies a cluster-wide parameter and thus a refresh of the
// parameter after that modification should observe results of preceeding writes.
const bool kEnsureReadYourWritesConsistency = true;
auto refreshStatus = ClusterServerParameterRefresher::get(opCtx)->refreshParameters(
opCtx, kEnsureReadYourWritesConsistency);
if (!refreshStatus.isOK()) {
LOGV2_WARNING(8472500,
"Error occurred when fetching the latest version of query settings",
"error_code"_attr = refreshStatus.code(),
"reason"_attr = refreshStatus.reason());
}
},
sanitizeQuerySettingsHints);
getQuerySettingsService(serviceContext) = std::make_unique<QuerySettingsRouterService>();
}
void initializeForShard(ServiceContext* serviceContext) {
query_settings::QuerySettingsManager::create(serviceContext, {}, sanitizeQuerySettingsHints);
getQuerySettingsService(serviceContext) = std::make_unique<QuerySettingsShardService>();
}
void initializeForTest(ServiceContext* serviceContext) {
query_settings::QuerySettingsManager::create(serviceContext, {}, {});
initializeForShard(serviceContext);
}
QuerySettings lookupQuerySettingsWithRejectionCheckOnRouter(
const boost::intrusive_ptr<ExpressionContext>& expCtx,
const query_shape::DeferredQueryShape& deferredShape,
const NamespaceString& nss) {
QuerySettings settings = [&]() {
try {
// No query settings lookup for IDHACK queries.
if (expCtx->isIdHackQuery()) {
return QuerySettings();
}
// No query settings for queries with encryption information.
if (expCtx->isFleQuery()) {
return QuerySettings();
}
// No query settings lookup on internal dbs or system collections in user dbs.
if (nss.isOnInternalDb() || nss.isSystem()) {
return QuerySettings();
}
// Force shape computation and early exit with empty settings if shape computation has
// failed.
const auto& shapePtr = deferredShape();
if (!shapePtr.isOK()) {
return QuerySettings();
}
// Compute the QueryShapeHash and store it in CurOp.
auto* opCtx = expCtx->getOperationContext();
QueryShapeHash hash = shapePtr.getValue()->sha256Hash(opCtx, kSerializationContext);
setQueryShapeHash(opCtx, hash);
// Return the found query settings or an empty one.
auto& manager = QuerySettingsManager::get(opCtx);
return manager.getQuerySettingsForQueryShapeHash(hash, nss.tenantId()).get_value_or({});
} catch (const DBException& ex) {
LOGV2_WARNING_OPTIONS(10153400,
{logv2::LogComponent::kQuery},
"Failed to perform query settings lookup",
"error"_attr = ex.toString());
return QuerySettings();
}
}();
// Fail the current command, if 'reject: true' flag is present.
failIfRejectedBySettings(expCtx, settings);
return settings;
auto* service =
getQuerySettingsService(expCtx->getOperationContext()->getServiceContext()).get();
dassert(dynamic_cast<QuerySettingsRouterService*>(service));
return service->lookupQuerySettingsWithRejectionCheck(expCtx, deferredShape, nss);
}
QuerySettings lookupQuerySettingsWithRejectionCheckOnShard(
@ -479,49 +571,30 @@ QuerySettings lookupQuerySettingsWithRejectionCheckOnShard(
const query_shape::DeferredQueryShape& deferredShape,
const NamespaceString& nss,
const boost::optional<QuerySettings>& querySettingsFromOriginalCommand) {
// No query settings lookup for IDHACK queries.
if (expCtx->isIdHackQuery()) {
return QuerySettings();
}
auto* opCtx = expCtx->getOperationContext();
if (requestComesFromRouterOrSentDirectlyToShard(opCtx->getClient()) ||
querySettingsFromOriginalCommand.has_value()) {
return querySettingsFromOriginalCommand.get_value_or(QuerySettings());
}
// The underlying shard does not belong to a sharded cluster, therefore proceed by performing
// the lookup as a router.
return lookupQuerySettingsWithRejectionCheckOnRouter(expCtx, deferredShape, nss);
auto* service =
getQuerySettingsService(expCtx->getOperationContext()->getServiceContext()).get();
dassert(dynamic_cast<QuerySettingsShardService*>(service));
return service->lookupQuerySettingsWithRejectionCheck(
expCtx, deferredShape, nss, querySettingsFromOriginalCommand);
}
QueryShapeConfigurationsWithTimestamp getAllQueryShapeConfigurations(
OperationContext* opCtx, const boost::optional<TenantId>& tenantId) {
return QuerySettingsManager::get(opCtx).getAllQueryShapeConfigurations(tenantId);
}
void setAllQueryShapeConfigurations(OperationContext* opCtx,
QueryShapeConfigurationsWithTimestamp&& config,
const boost::optional<TenantId>& tenantId) {
QuerySettingsManager::get(opCtx).setQueryShapeConfigurations(
std::move(config.queryShapeConfigurations), config.clusterParameterTime, tenantId);
return QuerySettingsService::get(opCtx).getAllQueryShapeConfigurations(tenantId);
}
void refreshQueryShapeConfigurations(OperationContext* opCtx) {
QuerySettingsManager::get(opCtx).refreshQueryShapeConfigurations(opCtx);
QuerySettingsService::get(opCtx).refreshQueryShapeConfigurations(opCtx);
}
std::string getQuerySettingsClusterParameterName() {
return QuerySettingsManager::kQuerySettingsClusterParameterName.toString();
return kQuerySettingsClusterParameterName.toString();
}
const stdx::unordered_set<StringData, StringMapHasher>& getRejectionIncompatibleStages() {
return rejectionIncompatibleStages;
};
bool canPipelineBeRejected(const std::vector<BSONObj>& pipeline) {
return pipeline.empty() ||
!getRejectionIncompatibleStages().contains(pipeline.at(0).firstElementFieldName());
!QuerySettingsService::getRejectionIncompatibleStages().contains(
pipeline.at(0).firstElementFieldName());
}
bool allowQuerySettingsFromClient(Client* client) {
@ -543,7 +616,7 @@ bool isDefault(const QuerySettings& settings) {
return !(settings.getQueryFramework() || settings.getIndexHints() || settings.getReject());
}
void validateQuerySettings(const QuerySettings& querySettings) {
void QuerySettingsService::validateQuerySettings(const QuerySettings& querySettings) const {
// Validates that the settings field for query settings is not empty.
uassert(7746604,
"the resulting settings cannot be empty or contain only default values",
@ -552,8 +625,8 @@ void validateQuerySettings(const QuerySettings& querySettings) {
validateQuerySettingsIndexHints(querySettings.getIndexHints());
}
void validateQueryCompatibleWithAnyQuerySettings(
const RepresentativeQueryInfo& representativeQueryInfo) {
void QuerySettingsService::validateQueryCompatibleWithAnyQuerySettings(
const RepresentativeQueryInfo& representativeQueryInfo) const {
if (MONGO_unlikely(allowAllSetQuerySettings.shouldFail())) {
return;
}
@ -573,8 +646,8 @@ void validateQueryCompatibleWithAnyQuerySettings(
!representativeQueryInfo.isIdHackQuery);
}
void validateQueryCompatibleWithQuerySettings(
const RepresentativeQueryInfo& representativeQueryInfo, const QuerySettings& settings) {
void QuerySettingsService::validateQueryCompatibleWithQuerySettings(
const RepresentativeQueryInfo& representativeQueryInfo, const QuerySettings& settings) const {
if (MONGO_unlikely(allowAllSetQuerySettings.shouldFail())) {
return;
}
@ -584,7 +657,7 @@ void validateQueryCompatibleWithQuerySettings(
!(settings.getReject() && representativeQueryInfo.systemStage.has_value()));
}
void simplifyQuerySettings(QuerySettings& settings) {
void QuerySettingsService::simplifyQuerySettings(QuerySettings& settings) const {
// If reject is present, but is false, set to an empty optional.
if (settings.getReject().has_value() && !settings.getReject()) {
settings.setReject({});
@ -610,8 +683,9 @@ void simplifyQuerySettings(QuerySettings& settings) {
}
// TODO SERVER-97546 Remove PQS index hint sanitization.
void sanitizeQuerySettingsHints(std::vector<QueryShapeConfiguration>& queryShapeConfigs) {
std::erase_if(queryShapeConfigs, [](QueryShapeConfiguration& queryShapeItem) {
void QuerySettingsService::sanitizeQuerySettingsHints(
std::vector<QueryShapeConfiguration>& queryShapeConfigs) const {
std::erase_if(queryShapeConfigs, [&](QueryShapeConfiguration& queryShapeItem) {
auto& settings = queryShapeItem.getSettings();
sanitizeKeyPatternIndexHints(queryShapeItem);
simplifyQuerySettings(settings);
@ -626,5 +700,4 @@ void sanitizeQuerySettingsHints(std::vector<QueryShapeConfiguration>& queryShape
return false;
});
}
} // namespace mongo::query_settings

175
src/mongo/db/query/query_settings/query_settings_service.h
View File

@ -70,6 +70,126 @@ RepresentativeQueryInfo createRepresentativeInfo(OperationContext* opCtx,
const QueryInstance& queryInstance,
const boost::optional<TenantId>& tenantId);
class QuerySettingsService {
public:
/**
* Gets the instance of the class using the service context.
*/
static QuerySettingsService& get(ServiceContext* serviceContext);
/**
* Gets the instance of the class using the operation context.
*/
static QuerySettingsService& get(OperationContext* opCtx);
/**
* Returns the name of the cluster parameter that stores QuerySettings for all QueryShapes.
*/
static std::string getQuerySettingsClusterParameterName();
/**
* Returns a set of system and administrative aggregation pipeline stages that, if used as the
* initial stage, prevent the query from being rejected via query settings.
*
* Query settings module is responsible for maintaining the information about what aggregation
* stages can be rejected.
*/
static const stdx::unordered_set<StringData, StringMapHasher>& getRejectionIncompatibleStages();
virtual ~QuerySettingsService() = default;
/**
* Returns the appropriate QuerySettings:
*
* - On router and shard in replica set deployment performs QuerySettings lookup for a specified
* 'queryShape'. If no settings are found or if the 'queryShape' is ineligible (e.g., IDHACK
* queries), returns empty QuerySettings. If the QuerySettings include 'reject: true' and is not
* run in explain, a uassert is thrown with the QueryRejectedBySettings error code, rejecting
* the query. Additionally, records the QueryShapeHash within the CurOp
*
* - On shard in sharded cluster returns 'querySettingsFromOriginalCommand'. This corresponds to
* the QuerySettings looked up on the router and passed to shards as part of the command.
* Rejection check is not performed here, as queries with 'reject: true' in their QuerySettings
* would already have been rejected by the router.
*/
virtual QuerySettings lookupQuerySettingsWithRejectionCheck(
const boost::intrusive_ptr<ExpressionContext>& expCtx,
const query_shape::DeferredQueryShape& queryShape,
const NamespaceString& nss,
const boost::optional<QuerySettings>& querySettingsFromOriginalCommand =
boost::none) const = 0;
/**
* Returns all the query shape configurations and the timestamp of the last modification.
*/
virtual QueryShapeConfigurationsWithTimestamp getAllQueryShapeConfigurations(
const boost::optional<TenantId>& tenantId) const = 0;
/**
* Sets all the query shape configurations with the given timestamp.
*/
virtual void setAllQueryShapeConfigurations(
QueryShapeConfigurationsWithTimestamp&& queryShapeConfigurations,
const boost::optional<TenantId>& tenantId) = 0;
/**
* Removes all query shape configurations.
*/
virtual void removeAllQueryShapeConfigurations(const boost::optional<TenantId>& tenantId) = 0;
/**
* Returns the LogicalTime of the 'querySettings' cluster parameter.
*/
virtual LogicalTime getClusterParameterTime(
const boost::optional<TenantId>& tenantId) const = 0;
/**
* Refreshes the local copy of all QueryShapeConfiguration by fetching the latest version from
* the configsvr. Is a no-op when called on shard.
*/
virtual void refreshQueryShapeConfigurations(OperationContext* opCtx) = 0;
/**
* Validates that 'querySettings' do not have:
* - empty settings or settings with default values
* - index hints specified without namespace information
* - index hints specified for the same namespace more than once
*
* Throws a uassert if compatibility checks fail, indicating that 'querySettings' cannot be set.
*/
void validateQuerySettings(const QuerySettings& querySettings) const;
/**
* Validates that QuerySettings can be applied to the query represented by 'queryInfo'.
* Throws a uassert if compatibility checks fail, indicating that 'querySettings' cannot be set.
*/
void validateQueryCompatibleWithAnyQuerySettings(
const RepresentativeQueryInfo& queryInfo) const;
/**
* Validates that 'querySettings' can be applied to the query represented by 'queryInfo'.
* Throws a uassert if compatibility checks fail, indicating that 'querySettings' cannot be set.
*/
void validateQueryCompatibleWithQuerySettings(const RepresentativeQueryInfo& queryInfo,
const QuerySettings& querySettings) const;
/**
* Simplifies 'querySettings' in-place by:
* - resetting the 'reject' field to boost::none if it contains a false value
* - removing index hints that specify empty 'allowedIndexes', potentially resetting
* 'indexHints' to boost::none if all 'allowedIndexes' are empty.
*/
void simplifyQuerySettings(QuerySettings& querySettings) const;
/**
* Sanitizes the 'queryShapeConfigurations' removing those hints that contain invalid index key
* pattern. In case the underlying query settings object contains only the default settings, the
* corresponding QueryShapeConfiguration is removed.
*/
void sanitizeQuerySettingsHints(
std::vector<QueryShapeConfiguration>& queryShapeConfigurations) const;
};
void initializeForRouter(ServiceContext* serviceContext);
void initializeForShard(ServiceContext* serviceContext);
void initializeForTest(ServiceContext* serviceContext);
@ -112,14 +232,6 @@ QuerySettings lookupQuerySettingsWithRejectionCheckOnShard(
QueryShapeConfigurationsWithTimestamp getAllQueryShapeConfigurations(
OperationContext* opCtx, const boost::optional<TenantId>& tenantId);
/**
* Sets all the query shape configurations with the given timestamp.
*/
void setAllQueryShapeConfigurations(
OperationContext* opCtx,
QueryShapeConfigurationsWithTimestamp&& queryShapeConfigurations,
const boost::optional<TenantId>& tenantId);
/**
* Refreshes the local copy of all QueryShapeConfiguration by fetching the latest version from the
* configsvr. Is a no-op when called on shard.
@ -131,15 +243,6 @@ void refreshQueryShapeConfigurations(OperationContext* opCtx);
*/
std::string getQuerySettingsClusterParameterName();
/**
* Returns a set of system and administrative aggregation pipeline stages that, if used as the
* initial stage, prevent the query from being rejected via query settings.
*
* Query settings module is responsible for maintaining the information about what aggregation
* stages can be rejected.
*/
const stdx::unordered_set<StringData, StringMapHasher>& getRejectionIncompatibleStages();
/**
* Returns true if the aggregation pipeline 'pipeline' does not start with rejection incompatible
* stage, and therefore can be rejected.
@ -156,42 +259,4 @@ bool allowQuerySettingsFromClient(Client* client);
* Returns true if given QuerySettings instance contains only default values.
*/
bool isDefault(const QuerySettings& querySettings);
/**
* Validates that 'querySettings' do not have:
* - empty settings or settings with default values
* - index hints specified without namespace information
* - index hints specified for the same namespace more than once
*
* Throws a uassert if compatibility checks fail, indicating that 'querySettings' cannot be set.
*/
void validateQuerySettings(const QuerySettings& querySettings);
/**
* Validates that QuerySettings can be applied to the query represented by 'queryInfo'.
* Throws a uassert if compatibility checks fail, indicating that 'querySettings' cannot be set.
*/
void validateQueryCompatibleWithAnyQuerySettings(const RepresentativeQueryInfo& queryInfo);
/**
* Validates that 'querySettings' can be applied to the query represented by 'queryInfo'.
* Throws a uassert if compatibility checks fail, indicating that 'querySettings' cannot be set.
*/
void validateQueryCompatibleWithQuerySettings(const RepresentativeQueryInfo& queryInfo,
const QuerySettings& querySettings);
/**
* Simplifies 'querySettings' in-place by:
* - resetting the 'reject' field to boost::none if it contains a false value
* - removing index hints that specify empty 'allowedIndexes', potentially resetting 'indexHints'
* to boost::none if all 'allowedIndexes' are empty.
*/
void simplifyQuerySettings(QuerySettings& querySettings);
/**
* Sanitizes the 'queryShapeConfigurations' removing those hints that contain invalid index key
* pattern. In case the underlying query settings object contains only the default settings, the
* corresponding QueryShapeConfiguration is removed.
*/
void sanitizeQuerySettingsHints(std::vector<QueryShapeConfiguration>& queryShapeConfigurations);
} // namespace mongo::query_settings

260
src/mongo/db/query/query_settings/query_settings_service_test.cpp
View File

@ -42,6 +42,7 @@
#include "mongo/db/pipeline/expression_context.h"
#include "mongo/db/query/query_knobs_gen.h"
#include "mongo/db/query/query_request_helper.h"
#include "mongo/db/query/query_settings/query_settings_cluster_parameter_gen.h"
#include "mongo/db/query/query_settings/query_settings_gen.h"
#include "mongo/db/query/query_settings/query_settings_service.h"
#include "mongo/db/query/query_shape/agg_cmd_shape.h"
@ -55,12 +56,65 @@
#include "mongo/util/serialization_context.h"
namespace mongo::query_settings {
static bool operator==(const QueryShapeConfiguration& lhs, const QueryShapeConfiguration& rhs) {
return SimpleBSONObjComparator::kInstance.compare(lhs.toBSON(), rhs.toBSON()) == 0;
}
static bool operator==(const QueryShapeConfigurationsWithTimestamp& lhs,
const QueryShapeConfigurationsWithTimestamp& rhs) {
return lhs.clusterParameterTime == rhs.clusterParameterTime &&
lhs.queryShapeConfigurations == rhs.queryShapeConfigurations;
}
namespace {
static auto const kSerializationContext =
SerializationContext{SerializationContext::Source::Command,
SerializationContext::CallerType::Request,
SerializationContext::Prefix::ExcludePrefix};
BSONObj makeSettingsClusterParameter(const QueryShapeConfigurationsWithTimestamp& config) {
BSONArrayBuilder settings;
for (const auto& queryShapeConfig : config.queryShapeConfigurations) {
settings.append(queryShapeConfig.toBSON());
}
return BSON("_id" << QuerySettingsService::getQuerySettingsClusterParameterName()
<< QuerySettingsClusterParameterValue::kSettingsArrayFieldName
<< settings.arr()
<< QuerySettingsClusterParameterValue::kClusterParameterTimeFieldName
<< config.clusterParameterTime.asTimestamp());
}
BSONObj makeSettingsClusterParameter(
const BSONArray& settings, LogicalTime clusterParameterTime = LogicalTime(Timestamp(113, 59))) {
return BSON("_id" << QuerySettingsService::getQuerySettingsClusterParameterName()
<< QuerySettingsClusterParameterValue::kSettingsArrayFieldName << settings
<< QuerySettingsClusterParameterValue::kClusterParameterTimeFieldName
<< clusterParameterTime.asTimestamp());
}
QuerySettings makeQuerySettings(const IndexHintSpecs& indexHints, bool setFramework = true) {
QuerySettings settings;
if (!indexHints.empty()) {
settings.setIndexHints(indexHints);
}
if (setFramework) {
settings.setQueryFramework(mongo::QueryFrameworkControlEnum::kTrySbeEngine);
}
return settings;
}
auto makeDbName(StringData dbName) {
return DatabaseNameUtil::deserialize(
boost::none /*tenantId=*/, dbName, SerializationContext::stateDefault());
}
NamespaceSpec makeNsSpec(StringData collName) {
NamespaceSpec ns;
ns.setDb(makeDbName("testDbA"));
ns.setColl(collName);
return ns;
}
class QuerySettingsScope {
public:
QuerySettingsScope(OperationContext* opCtx,
@ -71,16 +125,14 @@ public:
LogicalTime newTime = _previousQueryShapeConfigurationsWithTimestamp.clusterParameterTime;
newTime.addTicks(1);
setAllQueryShapeConfigurations(
_opCtx,
QuerySettingsService::get(_opCtx).setAllQueryShapeConfigurations(
QueryShapeConfigurationsWithTimestamp{queryShapeConfigurations, newTime},
boost::none /* tenantId */);
}
~QuerySettingsScope() {
setAllQueryShapeConfigurations(_opCtx,
std::move(_previousQueryShapeConfigurationsWithTimestamp),
boost::none /* tenantId */);
QuerySettingsService::get(_opCtx).setAllQueryShapeConfigurations(
std::move(_previousQueryShapeConfigurationsWithTimestamp), boost::none /* tenantId */);
}
private:
@ -127,17 +179,14 @@ private:
boost::optional<ExplainOptions::Verbosity> _previousExplain;
};
class QuerySettingsLoookupTest : public ServiceContextTest {
class QuerySettingsServiceTest : public ServiceContextTest {
public:
static constexpr StringData kCollName = "exampleColl"_sd;
static constexpr StringData kDbName = "foo"_sd;
void setUp() final {
// Initialize the query settings.
initializeForTest(getServiceContext());
_opCtx = cc().makeOperationContext();
_expCtx = ExpressionContext::makeBlankExpressionContext(opCtx(), {NamespaceString()});
}
OperationContext* opCtx() {
@ -145,34 +194,66 @@ public:
}
boost::intrusive_ptr<ExpressionContext> expCtx() {
if (!_expCtx) {
_expCtx = ExpressionContext::makeBlankExpressionContext(opCtx(), {NamespaceString()});
}
return _expCtx;
}
QuerySettingsService& service() {
return QuerySettingsService::get(opCtx());
}
static NamespaceString nss() {
return NamespaceStringUtil::deserialize(
boost::none, kDbName, kCollName, kSerializationContext);
}
QueryShapeConfiguration makeQueryShapeConfiguration(const BSONObj& cmdBSON,
const QuerySettings& querySettings) {
auto queryShapeHash =
createRepresentativeInfo(opCtx(), cmdBSON, boost::none /* tenantId */).queryShapeHash;
QueryShapeConfiguration makeQueryShapeConfiguration(
const BSONObj& cmdBSON,
const QuerySettings& querySettings,
boost::optional<TenantId> tenantId = boost::none) {
auto queryShapeHash = createRepresentativeInfo(opCtx(), cmdBSON, tenantId).queryShapeHash;
QueryShapeConfiguration config(queryShapeHash, querySettings);
config.setRepresentativeQuery(cmdBSON);
return config;
}
std::vector<QueryShapeConfiguration> getExampleQueryShapeConfigurations(TenantId tenantId) {
NamespaceSpec ns;
ns.setDb(DatabaseNameUtil::deserialize(tenantId, kDbName, kSerializationContext));
ns.setColl(kCollName);
const QuerySettings settings = makeQuerySettings({IndexHintSpec(ns, {IndexHint("a_1")})});
QueryInstance queryA =
BSON("find" << kCollName << "$db" << kDbName << "filter" << BSON("a" << 2));
QueryInstance queryB =
BSON("find" << kCollName << "$db" << kDbName << "filter" << BSON("a" << BSONNULL));
return {makeQueryShapeConfiguration(queryA, settings, tenantId),
makeQueryShapeConfiguration(queryB, settings, tenantId)};
}
void assertQuerySettingsLookup(const BSONObj& cmdBSON,
const query_shape::DeferredQueryShape& deferredShape,
const NamespaceString& nss) {
assertQuerySettingsLookupWithoutRejectionCheck(cmdBSON, deferredShape, nss);
assertQuerySettingsLookupWithRejectionCheck(cmdBSON, deferredShape, nss);
{
initializeForRouter(getServiceContext());
assertQuerySettingsLookupWithoutRejectionCheckForRouter(cmdBSON, deferredShape, nss);
assertQuerySettingsLookupWithRejectionCheckForRouter(cmdBSON, deferredShape, nss);
}
{
initializeForShard(getServiceContext());
assertQuerySettingsLookupWithoutRejectionCheckForShard(cmdBSON, deferredShape, nss);
assertQuerySettingsLookupWithRejectionCheckForShard(cmdBSON, deferredShape, nss);
}
}
/**
* Ensures that QuerySettings lookup returns the correct settings for the corresponding query.
*/
void assertQuerySettingsLookupWithoutRejectionCheck(
void assertQuerySettingsLookupWithoutRejectionCheckForRouter(
const BSONObj& cmdBSON,
const query_shape::DeferredQueryShape& deferredShape,
const NamespaceString& nss) {
@ -186,10 +267,7 @@ public:
useSbeEngineSettings.setQueryFramework(QueryFrameworkControlEnum::kTrySbeEngine);
// Ensure empty settings are returned if no settings are present in the system.
ASSERT_EQ(lookupQuerySettingsWithRejectionCheckOnRouter(expCtx(), deferredShape, nss),
QuerySettings());
ASSERT_EQ(lookupQuerySettingsWithRejectionCheckOnShard(
expCtx(), deferredShape, nss, QuerySettings()),
ASSERT_EQ(service().lookupQuerySettingsWithRejectionCheck(expCtx(), deferredShape, nss),
QuerySettings());
// Set { queryFramework: 'classic' } settings to 'cmdForSettingsBSON'.
@ -198,12 +276,39 @@ public:
// Ensure that 'forceClassicEngineSettings' are returned during the lookup, after query
// settings have been populated.
ASSERT_EQ(lookupQuerySettingsWithRejectionCheckOnRouter(expCtx(), deferredShape, nss),
ASSERT_EQ(service().lookupQuerySettingsWithRejectionCheck(expCtx(), deferredShape, nss),
forceClassicEngineSettings);
}
/**
* Ensures that QuerySettings lookup returns the correct settings for the corresponding query.
*/
void assertQuerySettingsLookupWithoutRejectionCheckForShard(
const BSONObj& cmdBSON,
const query_shape::DeferredQueryShape& deferredShape,
const NamespaceString& nss) {
const bool isExplain = cmdBSON.firstElementFieldNameStringData() == "explain";
BSONObj cmdForSettingsBSON = isExplain ? cmdBSON.firstElement().Obj() : cmdBSON;
QuerySettings forceClassicEngineSettings;
forceClassicEngineSettings.setQueryFramework(
QueryFrameworkControlEnum::kForceClassicEngine);
QuerySettings useSbeEngineSettings;
useSbeEngineSettings.setQueryFramework(QueryFrameworkControlEnum::kTrySbeEngine);
// Ensure empty settings are returned if no settings are present in the system.
ASSERT_EQ(service().lookupQuerySettingsWithRejectionCheck(
expCtx(), deferredShape, nss, QuerySettings()),
QuerySettings());
// Set { queryFramework: 'classic' } settings to 'cmdForSettingsBSON'.
QuerySettingsScope forceClassicEngineQuerySettingsScope(
opCtx(), {makeQueryShapeConfiguration(cmdForSettingsBSON, forceClassicEngineSettings)});
// Ensure that in case of a replica set case, a regular query settings lookup is performed.
ASSERT_EQ(
lookupQuerySettingsWithRejectionCheckOnShard(
service().lookupQuerySettingsWithRejectionCheck(
expCtx(), deferredShape, nss, boost::none /* querySettingsFromOriginalCommand */),
forceClassicEngineSettings);
@ -213,13 +318,13 @@ public:
// Ensure that settings passed to the method are being returned as opposed to performing
// the QuerySettings lookup.
ASSERT_EQ(lookupQuerySettingsWithRejectionCheckOnShard(
ASSERT_EQ(service().lookupQuerySettingsWithRejectionCheck(
expCtx(), deferredShape, nss, useSbeEngineSettings),
useSbeEngineSettings);
// Ensure that empty settings are returned if original command did not have any settings
// specified as opposed to performing QuerySettings lookup.
ASSERT_EQ(lookupQuerySettingsWithRejectionCheckOnShard(
ASSERT_EQ(service().lookupQuerySettingsWithRejectionCheck(
expCtx(),
deferredShape,
nss,
@ -232,7 +337,7 @@ public:
* Ensures that QuerySettings lookup rejects the queries if the associated settings contain
* 'reject: true'.
*/
void assertQuerySettingsLookupWithRejectionCheck(
void assertQuerySettingsLookupWithRejectionCheckForRouter(
const BSONObj& cmdBSON,
const query_shape::DeferredQueryShape& deferredShape,
const NamespaceString& nss) {
@ -250,13 +355,37 @@ public:
if (isExplain) {
ASSERT_DOES_NOT_THROW(
lookupQuerySettingsWithRejectionCheckOnRouter(expCtx(), deferredShape, nss));
ASSERT_DOES_NOT_THROW(lookupQuerySettingsWithRejectionCheckOnShard(
expCtx(), deferredShape, nss, boost::none /* querySettingsFromOriginalCommand */));
} else {
ASSERT_THROWS_CODE(
lookupQuerySettingsWithRejectionCheckOnRouter(expCtx(), deferredShape, nss),
DBException,
ErrorCodes::QueryRejectedBySettings);
}
}
/**
* Ensures that QuerySettings lookup rejects the queries if the associated settings contain
* 'reject: true'.
*/
void assertQuerySettingsLookupWithRejectionCheckForShard(
const BSONObj& cmdBSON,
const query_shape::DeferredQueryShape& deferredShape,
const NamespaceString& nss) {
const bool isExplain = cmdBSON.firstElementFieldNameStringData() == "explain";
BSONObj cmdForSettingsBSON = isExplain ? cmdBSON.firstElement().Obj() : cmdBSON;
// Set { reject: true } settings to 'cmdForSettingsBSON'.
QuerySettings querySettingsWithReject;
querySettingsWithReject.setReject(true);
QuerySettingsScope rejectQuerySettingsScope(
opCtx(), {makeQueryShapeConfiguration(cmdForSettingsBSON, querySettingsWithReject)});
// Ensure query is not rejected if an explain query is run, otherwise is rejected by
// throwing an exception with QueryRejectedBySettings error code.
if (isExplain) {
ASSERT_DOES_NOT_THROW(lookupQuerySettingsWithRejectionCheckOnShard(
expCtx(), deferredShape, nss, boost::none /* querySettingsFromOriginalCommand */));
} else {
ASSERT_THROWS_CODE(lookupQuerySettingsWithRejectionCheckOnShard(
expCtx(),
deferredShape,
@ -276,12 +405,25 @@ public:
}
}
void assertSanitizeInvalidIndexHints(const IndexHintSpecs& initialSpec,
const IndexHintSpecs& expectedSpec) {
QueryInstance query = BSON("find" << "exampleColl"
<< "$db"
<< "foo");
auto initSettings = makeQueryShapeConfiguration(query, makeQuerySettings(initialSpec));
std::vector<QueryShapeConfiguration> queryShapeConfigs{initSettings};
const auto expectedSettings = makeQuerySettings(expectedSpec);
ASSERT_DOES_NOT_THROW(service().sanitizeQuerySettingsHints(queryShapeConfigs));
ASSERT_BSONOBJ_EQ(queryShapeConfigs[0].getSettings().toBSON(), expectedSettings.toBSON());
}
private:
ServiceContext::UniqueOperationContext _opCtx;
boost::intrusive_ptr<ExpressionContext> _expCtx;
};
TEST_F(QuerySettingsLoookupTest, QuerySettingsLookupForFind) {
TEST_F(QuerySettingsServiceTest, QuerySettingsLookupForFind) {
auto findCmdStr = "{find: 'exampleColl', '$db': 'foo'}"_sd;
auto findCmdBSON = fromjson(findCmdStr);
auto findCmd = query_request_helper::makeFromFindCommandForTests(findCmdBSON, nss());
@ -298,7 +440,7 @@ TEST_F(QuerySettingsLoookupTest, QuerySettingsLookupForFind) {
}
}
TEST_F(QuerySettingsLoookupTest, QuerySettingsLookupForAgg) {
TEST_F(QuerySettingsServiceTest, QuerySettingsLookupForAgg) {
auto aggCmdStr =
"{aggregate: 'exampleColl', pipeline: [{$match: {_id: 0}}], cursor: {}, '$db': 'foo'}"_sd;
auto aggCmdBSON = fromjson(aggCmdStr);
@ -317,7 +459,7 @@ TEST_F(QuerySettingsLoookupTest, QuerySettingsLookupForAgg) {
}
}
TEST_F(QuerySettingsLoookupTest, QuerySettingsLookupForDistinct) {
TEST_F(QuerySettingsServiceTest, QuerySettingsLookupForDistinct) {
auto distinctCmdBSON = fromjson("{distinct: 'exampleColl', key: 'x', $db: 'foo'}");
auto distinctCmd = std::make_unique<DistinctCommandRequest>(
DistinctCommandRequest::parse(IDLParserContext("distinctCommandRequest",
@ -342,4 +484,62 @@ TEST_F(QuerySettingsLoookupTest, QuerySettingsLookupForDistinct) {
}
}
/**
* Tests that valid index hint specs are the same before and after index hint sanitization.
*/
TEST_F(QuerySettingsServiceTest, ValidIndexHintsAreTheSameBeforeAndAfterSanitization) {
IndexHintSpecs indexHintSpec{IndexHintSpec(
makeNsSpec("testCollA"_sd), {IndexHint(BSON("a" << 1)), IndexHint(BSON("b" << -1.0))})};
assertSanitizeInvalidIndexHints(indexHintSpec, indexHintSpec);
}
/**
* Tests that invalid key-pattern are removed after sanitization.
*/
TEST_F(QuerySettingsServiceTest, InvalidKeyPatternIndexesAreRemovedAfterSanitization) {
IndexHintSpecs indexHintSpec{IndexHintSpec(makeNsSpec("testCollA"_sd),
{IndexHint(BSON("a" << 1 << "c"
<< "invalid")),
IndexHint(BSON("b" << -1.0))})};
IndexHintSpecs expectedHintSpec{
IndexHintSpec(makeNsSpec("testCollA"_sd), {IndexHint(BSON("b" << -1.0))})};
assertSanitizeInvalidIndexHints(indexHintSpec, expectedHintSpec);
}
/**
* Same as the above test but with more complex examples.
*/
TEST_F(QuerySettingsServiceTest, InvalidKeyPatternIndexesAreRemovedAfterSanitizationComplex) {
IndexHintSpecs indexHintSpec{IndexHintSpec(makeNsSpec("testCollA"_sd),
{
IndexHint(BSON("a" << 1 << "c"
<< "invalid")),
IndexHint(BSON("b" << -1.0)),
IndexHint(BSON("c" << -2.0 << "b" << 4)),
IndexHint("index_name"),
IndexHint(BSON("$natural" << 1)),
IndexHint(BSON("$natural" << -1 << "a" << 2)),
IndexHint(BSON("a" << -1 << "$natural" << 1)),
})};
IndexHintSpecs expectedHintSpec{IndexHintSpec(makeNsSpec("testCollA"_sd),
{IndexHint(BSON("b" << -1.0)),
IndexHint(BSON("c" << -2.0 << "b" << 4)),
IndexHint("index_name")})};
assertSanitizeInvalidIndexHints(indexHintSpec, expectedHintSpec);
}
/**
* Tests that invalid key-pattern are removed after sanitization and the resulted index hints are
* empty.
*/
TEST_F(QuerySettingsServiceTest, InvalidKeyPatternIndexesAreRemovedAfterSanitizationEmptyHints) {
IndexHintSpecs indexHintSpec{IndexHintSpec(makeNsSpec("testCollA"_sd),
{
IndexHint(BSON("a" << 1 << "c"
<< "invalid")),
})};
IndexHintSpecs expectedHintSpec;
assertSanitizeInvalidIndexHints(indexHintSpec, expectedHintSpec);
}
} // namespace
} // namespace mongo::query_settings

66
src/mongo/db/query/query_settings/query_settings_validation_test.cpp
View File

@ -46,6 +46,11 @@ protected:
ShardingState::create(getServiceContext());
expCtx = make_intrusive<ExpressionContextForTest>();
gAggregateOperationResourceConsumptionMetrics = true;
query_settings::initializeForTest(getServiceContext());
}
QuerySettingsService& service() {
return QuerySettingsService::get(getServiceContext());
}
boost::intrusive_ptr<ExpressionContext> expCtx;
@ -59,7 +64,8 @@ void assertInvalidQueryWithAnyQuerySettings(OperationContext* opCtx,
size_t errorCode) {
auto representativeQueryInfo =
createRepresentativeInfo(opCtx, representativeQuery, boost::none);
ASSERT_THROWS_CODE(validateQueryCompatibleWithAnyQuerySettings(representativeQueryInfo),
ASSERT_THROWS_CODE(QuerySettingsService::get(opCtx).validateQueryCompatibleWithAnyQuerySettings(
representativeQueryInfo),
DBException,
errorCode);
}
@ -70,10 +76,10 @@ void assertInvalidQueryAndQuerySettingsCombination(OperationContext* opCtx,
size_t errorCode) {
auto representativeQueryInfo =
createRepresentativeInfo(opCtx, representativeQuery, boost::none);
ASSERT_THROWS_CODE(
validateQueryCompatibleWithQuerySettings(representativeQueryInfo, querySettings),
DBException,
errorCode);
ASSERT_THROWS_CODE(QuerySettingsService::get(opCtx).validateQueryCompatibleWithQuerySettings(
representativeQueryInfo, querySettings),
DBException,
errorCode);
}
NamespaceSpec makeNamespace(StringData dbName, StringData collName) {
@ -137,7 +143,7 @@ TEST_F(QuerySettingsValidationTestFixture,
"$operationMetrics"_sd,
};
for (auto&& stage : getRejectionIncompatibleStages()) {
for (auto&& stage : QuerySettingsService::getRejectionIncompatibleStages()) {
// Avoid testing these stages, as they require more complex setup.
if (stage == "$listLocalSessions" || stage == "$listSessions" ||
stage == "$listSampledQueries") {
@ -186,21 +192,21 @@ TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndicesCannotReferToSame
auto indexSpecA = IndexHintSpec(ns, {IndexHint("sku")});
auto indexSpecB = IndexHintSpec(ns, {IndexHint("uks")});
querySettings.setIndexHints({{indexSpecA, indexSpecB}});
simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 7746608);
service().simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 7746608);
}
TEST_F(QuerySettingsValidationTestFixture, QuerySettingsCannotBeEmpty) {
QuerySettings querySettings;
simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 7746604);
service().simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 7746604);
}
TEST_F(QuerySettingsValidationTestFixture, QuerySettingsCannotHaveDefaultValues) {
QuerySettings querySettings;
querySettings.setReject(false);
simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 7746604);
service().simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 7746604);
}
TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithNoDbSpecified) {
@ -208,8 +214,8 @@ TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithNoDbSpecif
NamespaceSpec ns;
ns.setColl("collName"_sd);
querySettings.setIndexHints({{IndexHintSpec(ns, {IndexHint("a")})}});
simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 8727500);
service().simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 8727500);
}
TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithNoCollSpecified) {
@ -218,17 +224,17 @@ TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithNoCollSpec
ns.setDb(DatabaseNameUtil::deserialize(
boost::none /* tenantId */, "dbName"_sd, SerializationContext::stateDefault()));
querySettings.setIndexHints({{IndexHintSpec(ns, {IndexHint("a")})}});
simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 8727501);
service().simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 8727501);
}
TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithEmptyAllowedIndexes) {
QuerySettings querySettings;
auto ns = makeNamespace("testDB", "testColl");
querySettings.setIndexHints({{IndexHintSpec(ns, {})}});
simplifyQuerySettings(querySettings);
service().simplifyQuerySettings(querySettings);
ASSERT_EQUALS(querySettings.getIndexHints(), boost::none);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 7746604);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 7746604);
}
TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithAllEmptyAllowedIndexes) {
@ -238,9 +244,9 @@ TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithAllEmptyAl
IndexHintSpec(makeNamespace("testDB", "testColl2"), {}),
IndexHintSpec(makeNamespace("testDB", "testColl3"), {}),
}});
simplifyQuerySettings(querySettings);
service().simplifyQuerySettings(querySettings);
ASSERT_EQUALS(querySettings.getIndexHints(), boost::none);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 7746604);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 7746604);
}
TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithSomeEmptyAllowedIndexes) {
@ -251,7 +257,7 @@ TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithSomeEmptyA
}});
const auto expectedIndexHintSpec =
IndexHintSpec(makeNamespace("testDB", "testColl2"), {IndexHint("a")});
simplifyQuerySettings(querySettings);
service().simplifyQuerySettings(querySettings);
const auto simplifiedIndexHints = querySettings.getIndexHints();
ASSERT_NE(simplifiedIndexHints, boost::none);
@ -259,7 +265,7 @@ TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithSomeEmptyA
ASSERT_EQ(indexHintsList.size(), 1);
const auto& actualIndexHintSpec = indexHintsList[0];
ASSERT_BSONOBJ_EQ(expectedIndexHintSpec.toBSON(), actualIndexHintSpec.toBSON());
ASSERT_DOES_NOT_THROW(validateQuerySettings(querySettings));
ASSERT_DOES_NOT_THROW(service().validateQuerySettings(querySettings));
}
TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithEmptyKeyPattern) {
@ -267,8 +273,8 @@ TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithEmptyKeyPa
querySettings.setIndexHints({{
IndexHintSpec(makeNamespace("testDB", "testColl"), {IndexHint(BSONObj{})}),
}});
simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 9646000);
service().simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 9646000);
}
TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithInvalidKeyPattern) {
@ -278,8 +284,8 @@ TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithInvalidKey
{IndexHint(BSON("a" << 1 << "b"
<< "some-string"))}),
}});
simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 9646001);
service().simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 9646001);
}
TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithInvalidNaturalHint) {
@ -288,8 +294,8 @@ TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithInvalidNat
IndexHintSpec(makeNamespace("testDB", "testColl"),
{IndexHint(BSON("$natural" << 1 << "b" << 2))}),
}});
simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 9646001);
service().simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 9646001);
}
TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithInvalidNaturalHintInverse) {
@ -298,8 +304,8 @@ TEST_F(QuerySettingsValidationTestFixture, QuerySettingsIndexHintsWithInvalidNat
IndexHintSpec(makeNamespace("testDB", "testColl"),
{IndexHint(BSON("b" << 2 << "$natural" << 1))}),
}});
simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(validateQuerySettings(querySettings), DBException, 9646001);
service().simplifyQuerySettings(querySettings);
ASSERT_THROWS_CODE(service().validateQuerySettings(querySettings), DBException, 9646001);
}
} // namespace

4
src/mongo/dbtests/dbtests.cpp
View File

@ -54,6 +54,7 @@
#include "mongo/db/index/index_descriptor.h"
#include "mongo/db/index_builds/multi_index_block.h"
#include "mongo/db/query/client_cursor/cursor_manager.h"
#include "mongo/db/query/query_settings/query_settings_service.h"
#include "mongo/db/repl/member_state.h"
#include "mongo/db/repl/repl_settings.h"
#include "mongo/db/repl/replication_coordinator.h"
@ -277,6 +278,9 @@ int dbtestsMain(int argc, char** argv) {
AuthorizationManager::get(service->getService())->setAuthEnabled(false);
ScriptEngine::setup(ExecutionEnvironment::Server);
query_settings::initializeForTest(service);
return mongo::dbtests::runDbTests(argc, argv);
}