void ChunkManager::getShardIdsForQuery(OperationContext* txn,
const BSONObj& query,
set<ShardId>* shardIds) const {
auto statusWithCQ =
CanonicalQuery::canonicalize(NamespaceString(_ns), query, ExtensionsCallbackNoop());
uassertStatusOK(statusWithCQ.getStatus());
unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
// Query validation
if (QueryPlannerCommon::hasNode(cq->root(), MatchExpression::GEO_NEAR)) {
uassert(13501, "use geoNear command rather than $near query", false);
}
// Fast path for targeting equalities on the shard key.
auto shardKeyToFind = _keyPattern.extractShardKeyFromQuery(*cq);
if (shardKeyToFind.isOK() && !shardKeyToFind.getValue().isEmpty()) {
auto chunk = findIntersectingChunk(txn, shardKeyToFind.getValue());
shardIds->insert(chunk->getShardId());
return;
}
// Transforms query into bounds for each field in the shard key
// for example :
// Key { a: 1, b: 1 },
// Query { a : { $gte : 1, $lt : 2 },
// b : { $gte : 3, $lt : 4 } }
// => Bounds { a : [1, 2), b : [3, 4) }
IndexBounds bounds = getIndexBoundsForQuery(_keyPattern.toBSON(), *cq);
// Transforms bounds for each shard key field into full shard key ranges
// for example :
// Key { a : 1, b : 1 }
// Bounds { a : [1, 2), b : [3, 4) }
// => Ranges { a : 1, b : 3 } => { a : 2, b : 4 }
BoundList ranges = _keyPattern.flattenBounds(bounds);
for (BoundList::const_iterator it = ranges.begin(); it != ranges.end(); ++it) {
getShardIdsForRange(*shardIds, it->first /*min*/, it->second /*max*/);
// once we know we need to visit all shards no need to keep looping
if (shardIds->size() == _shardIds.size())
break;
}
// SERVER-4914 Some clients of getShardIdsForQuery() assume at least one shard will be
// returned. For now, we satisfy that assumption by adding a shard with no matches rather
// than return an empty set of shards.
if (shardIds->empty()) {
massert(16068, "no chunk ranges available", !_chunkRanges.ranges().empty());
shardIds->insert(_chunkRanges.ranges().begin()->second->getShardId());
}
}
Status ClusterAggregate::runAggregate(OperationContext* txn,
const Namespaces& namespaces,
BSONObj cmdObj,
int options,
BSONObjBuilder* result) {
auto dbname = namespaces.executionNss.db().toString();
auto status = grid.catalogCache()->getDatabase(txn, dbname);
if (!status.isOK()) {
appendEmptyResultSet(*result, status.getStatus(), namespaces.requestedNss.ns());
return Status::OK();
}
std::shared_ptr<DBConfig> conf = status.getValue();
if (!conf->isShardingEnabled()) {
return aggPassthrough(txn, namespaces, conf, cmdObj, result, options);
}
auto request = AggregationRequest::parseFromBSON(namespaces.executionNss, cmdObj);
if (!request.isOK()) {
return request.getStatus();
}
// Determine the appropriate collation and 'resolve' involved namespaces to make the
// ExpressionContext.
// We won't try to execute anything on a mongos, but we still have to populate this map so that
// any $lookups, etc. will be able to have a resolved view definition. It's okay that this is
// incorrect, we will repopulate the real resolved namespace map on the mongod. Note that we
// need to check if any involved collections are sharded before forwarding an aggregation
// command on an unsharded collection.
StringMap<ExpressionContext::ResolvedNamespace> resolvedNamespaces;
LiteParsedPipeline liteParsedPipeline(request.getValue());
for (auto&& ns : liteParsedPipeline.getInvolvedNamespaces()) {
uassert(28769, str::stream() << ns.ns() << " cannot be sharded", !conf->isSharded(ns.ns()));
resolvedNamespaces[ns.coll()] = {ns, std::vector<BSONObj>{}};
}
if (!conf->isSharded(namespaces.executionNss.ns())) {
return aggPassthrough(txn, namespaces, conf, cmdObj, result, options);
}
auto chunkMgr = conf->getChunkManager(txn, namespaces.executionNss.ns());
std::unique_ptr<CollatorInterface> collation;
if (!request.getValue().getCollation().isEmpty()) {
collation = uassertStatusOK(CollatorFactoryInterface::get(txn->getServiceContext())
->makeFromBSON(request.getValue().getCollation()));
} else if (chunkMgr->getDefaultCollator()) {
collation = chunkMgr->getDefaultCollator()->clone();
}
boost::intrusive_ptr<ExpressionContext> mergeCtx = new ExpressionContext(
txn, request.getValue(), std::move(collation), std::move(resolvedNamespaces));
mergeCtx->inRouter = true;
// explicitly *not* setting mergeCtx->tempDir
// Parse and optimize the pipeline specification.
auto pipeline = Pipeline::parse(request.getValue().getPipeline(), mergeCtx);
if (!pipeline.isOK()) {
return pipeline.getStatus();
}
pipeline.getValue()->optimizePipeline();
// If the first $match stage is an exact match on the shard key (with a simple collation or
// no string matching), we only have to send it to one shard, so send the command to that
// shard.
BSONObj firstMatchQuery = pipeline.getValue()->getInitialQuery();
BSONObj shardKeyMatches;
shardKeyMatches = uassertStatusOK(
chunkMgr->getShardKeyPattern().extractShardKeyFromQuery(txn, firstMatchQuery));
bool singleShard = false;
if (!shardKeyMatches.isEmpty()) {
auto chunk = chunkMgr->findIntersectingChunk(
txn, shardKeyMatches, request.getValue().getCollation());
if (chunk.isOK()) {
singleShard = true;
}
}
// Don't need to split pipeline if the first $match is an exact match on shard key, unless
// there is a stage that needs to be run on the primary shard.
const bool needPrimaryShardMerger = pipeline.getValue()->needsPrimaryShardMerger();
const bool needSplit = !singleShard || needPrimaryShardMerger;
// Split the pipeline into pieces for mongod(s) and this mongos. If needSplit is true,
// 'pipeline' will become the merger side.
boost::intrusive_ptr<Pipeline> shardPipeline(needSplit ? pipeline.getValue()->splitForSharded()
: pipeline.getValue());
// Create the command for the shards. The 'fromRouter' field means produce output to be
// merged.
MutableDocument commandBuilder(request.getValue().serializeToCommandObj());
commandBuilder[AggregationRequest::kPipelineName] = Value(shardPipeline->serialize());
if (needSplit) {
commandBuilder[AggregationRequest::kFromRouterName] = Value(true);
commandBuilder[AggregationRequest::kCursorName] =
Value(DOC(AggregationRequest::kBatchSizeName << 0));
}
// These fields are not part of the AggregationRequest since they are not handled by the
// aggregation subsystem, so we serialize them separately.
const std::initializer_list<StringData> fieldsToPropagateToShards = {
"$queryOptions", "readConcern", QueryRequest::cmdOptionMaxTimeMS,
};
for (auto&& field : fieldsToPropagateToShards) {
commandBuilder[field] = Value(cmdObj[field]);
}
BSONObj shardedCommand = commandBuilder.freeze().toBson();
BSONObj shardQuery = shardPipeline->getInitialQuery();
// Run the command on the shards
// TODO need to make sure cursors are killed if a retry is needed
std::vector<Strategy::CommandResult> shardResults;
Strategy::commandOp(txn,
dbname,
shardedCommand,
options,
namespaces.executionNss.ns(),
shardQuery,
request.getValue().getCollation(),
&shardResults);
if (mergeCtx->isExplain) {
// This must be checked before we start modifying result.
uassertAllShardsSupportExplain(shardResults);
if (needSplit) {
*result << "needsPrimaryShardMerger" << needPrimaryShardMerger << "splitPipeline"
<< DOC("shardsPart" << shardPipeline->writeExplainOps() << "mergerPart"
<< pipeline.getValue()->writeExplainOps());
} else {
*result << "splitPipeline" << BSONNULL;
}
BSONObjBuilder shardExplains(result->subobjStart("shards"));
for (size_t i = 0; i < shardResults.size(); i++) {
shardExplains.append(shardResults[i].shardTargetId,
BSON("host" << shardResults[i].target.toString() << "stages"
<< shardResults[i].result["stages"]));
}
return Status::OK();
}
if (!needSplit) {
invariant(shardResults.size() == 1);
invariant(shardResults[0].target.getServers().size() == 1);
auto executorPool = grid.getExecutorPool();
const BSONObj reply =
uassertStatusOK(storePossibleCursor(shardResults[0].target.getServers()[0],
shardResults[0].result,
namespaces.requestedNss,
executorPool->getArbitraryExecutor(),
grid.getCursorManager()));
result->appendElements(reply);
return getStatusFromCommandResult(reply);
}
pipeline.getValue()->addInitialSource(
DocumentSourceMergeCursors::create(parseCursors(shardResults), mergeCtx));
MutableDocument mergeCmd(request.getValue().serializeToCommandObj());
mergeCmd["pipeline"] = Value(pipeline.getValue()->serialize());
mergeCmd["cursor"] = Value(cmdObj["cursor"]);
if (cmdObj.hasField("$queryOptions")) {
mergeCmd["$queryOptions"] = Value(cmdObj["$queryOptions"]);
}
if (cmdObj.hasField(QueryRequest::cmdOptionMaxTimeMS)) {
mergeCmd[QueryRequest::cmdOptionMaxTimeMS] =
Value(cmdObj[QueryRequest::cmdOptionMaxTimeMS]);
}
mergeCmd.setField("writeConcern", Value(cmdObj["writeConcern"]));
mergeCmd.setField("readConcern", Value(cmdObj["readConcern"]));
// If the user didn't specify a collation already, make sure there's a collation attached to
// the merge command, since the merging shard may not have the collection metadata.
if (mergeCmd.peek()["collation"].missing()) {
mergeCmd.setField("collation",
mergeCtx->getCollator()
? Value(mergeCtx->getCollator()->getSpec().toBSON())
: Value(Document{CollationSpec::kSimpleSpec}));
}
std::string outputNsOrEmpty;
if (DocumentSourceOut* out =
dynamic_cast<DocumentSourceOut*>(pipeline.getValue()->getSources().back().get())) {
outputNsOrEmpty = out->getOutputNs().ns();
}
// Run merging command on random shard, unless a stage needs the primary shard. Need to use
// ShardConnection so that the merging mongod is sent the config servers on connection init.
auto& prng = txn->getClient()->getPrng();
const auto& mergingShardId = (needPrimaryShardMerger || internalQueryAlwaysMergeOnPrimaryShard)
? conf->getPrimaryId()
: shardResults[prng.nextInt32(shardResults.size())].shardTargetId;
const auto mergingShard = uassertStatusOK(grid.shardRegistry()->getShard(txn, mergingShardId));
ShardConnection conn(mergingShard->getConnString(), outputNsOrEmpty);
BSONObj mergedResults =
aggRunCommand(conn.get(), namespaces, mergeCmd.freeze().toBson(), options);
conn.done();
if (auto wcErrorElem = mergedResults["writeConcernError"]) {
appendWriteConcernErrorToCmdResponse(mergingShardId, wcErrorElem, *result);
}
// Copy output from merging (primary) shard to the output object from our command.
// Also, propagates errmsg and code if ok == false.
result->appendElementsUnique(mergedResults);
return getStatusFromCommandResult(result->asTempObj());
}
