SERVER-63479 Remove ExpressionFieldPathNonSharded

6 files changed, 143 insertions, 40 deletions

diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript
index da74b908548..7b19ed2b917 100644
--- a/src/mongo/db/pipeline/SConscript
+++ b/src/mongo/db/pipeline/SConscript
@@ -256,6 +256,7 @@ pipelineEnv.Library(
         'document_source_union_with.cpp',
         'document_source_unwind.cpp',
         'pipeline.cpp',
+        'search_helper.cpp',
         'semantic_analysis.cpp',
         'sequential_document_cache.cpp',
         'tee_buffer.cpp',
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index 23e203d20bf..6574830d26a 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -2097,14 +2097,6 @@ intrusive_ptr<ExpressionFieldPath> ExpressionFieldPath::parse(
         const StringData varName = fieldPath.substr(0, fieldPath.find('.'));
         Variables::uassertValidNameForUserRead(varName);
         auto varId = vps.getVariable(varName);
-        if (varName.compare(Variables::getBuiltinVariableName(Variables::kSearchMetaId)) == 0) {
-            return new ExpressionFieldPathNonSharded(
-                expCtx.get(),
-                fieldPath.toString(),
-                varId,
-                std::string("Search queries accessing $$SEARCH_META are not supported in sharded "
-                            "pipelines"));
-        }
         return new ExpressionFieldPath(expCtx, fieldPath.toString(), varId);
     } else {
         return new ExpressionFieldPath(expCtx,
@@ -2149,7 +2141,7 @@ void ExpressionFieldPath::_doAddDependencies(DepsTracker* deps) const {
         } else {
             deps->fields.insert(_fieldPath.tail().fullPath());
         }
-    } else if (Variables::isUserDefinedVariable(_variable)) {
+    } else {
         deps->vars.insert(_variable);
     }
 }
@@ -2222,12 +2214,6 @@ Value ExpressionFieldPath::serialize(bool explain) const {
     }
 }
 
-Value ExpressionFieldPathNonSharded::evaluate(const Document& root, Variables* variables) const {
-    uassert(
-        5858100, _errMsg, !getExpressionContext()->needsMerge && !getExpressionContext()->inMongos);
-    return ExpressionFieldPath::evaluate(root, variables);
-}
-
 Expression::ComputedPaths ExpressionFieldPath::getComputedPaths(const std::string& exprFieldPath,
                                                                 Variables::Id renamingVar) const {
     // An expression field path is either considered a rename or a computed path. We need to find
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index bbb3ef1c968..95ba8820523 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -1421,23 +1421,6 @@ private:
     const Variables::Id _variable;
 };
 
-/**
- * A version of ExpressionFieldPath that will throw if evaluated in a sharded pipeline.
- */
-class ExpressionFieldPathNonSharded : public ExpressionFieldPath {
-public:
-    Value evaluate(const Document& root, Variables* variables) const final;
-
-    ExpressionFieldPathNonSharded(ExpressionContext* const expCtx,
-                                  const std::string& fieldPath,
-                                  Variables::Id variable,
-                                  std::string errMsg)
-        : ExpressionFieldPath(expCtx, fieldPath, variable), _errMsg(std::move(errMsg)) {}
-
-private:
-    std::string _errMsg;
-};
-
 class ExpressionFilter final : public Expression {
 public:
     boost::intrusive_ptr<Expression> optimize() final;
diff --git a/src/mongo/db/pipeline/search_helper.cpp b/src/mongo/db/pipeline/search_helper.cpp
new file mode 100644
index 00000000000..1a1fb731fcf
--- /dev/null
+++ b/src/mongo/db/pipeline/search_helper.cpp
@@ -0,0 +1,41 @@
+/**
+ *    Copyright (C) 2022-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/db/pipeline/search_helper.h"
+
+namespace mongo {
+ServiceContext::Decoration<std::unique_ptr<SearchDefaultHelperFunctions>> getSearchHelpers =
+    ServiceContext::declareDecoration<std::unique_ptr<SearchDefaultHelperFunctions>>();
+
+ServiceContext::ConstructorActionRegisterer searchQueryHelperRegisterer{
+    "searchQueryHelperRegisterer", [](ServiceContext* context) {
+        invariant(context);
+        getSearchHelpers(context) = std::make_unique<SearchDefaultHelperFunctions>();
+    }};
+}  // namespace mongo
diff --git a/src/mongo/db/pipeline/search_helper.h b/src/mongo/db/pipeline/search_helper.h
new file mode 100644
index 00000000000..a045eaff8cb
--- /dev/null
+++ b/src/mongo/db/pipeline/search_helper.h
@@ -0,0 +1,58 @@
+/**
+ *    Copyright (C) 2022-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/db/pipeline/pipeline.h"
+
+namespace mongo {
+
+/**
+ * A class that contains any functions needed to run $seach queries when the enterprise module
+ * is compiled in. The enterprise module will override these functions, these are just stubs.
+ */
+class SearchDefaultHelperFunctions {
+public:
+    virtual ~SearchDefaultHelperFunctions() {}
+    /**
+     * Check if the passed in pipeline is valid in a sharded collection. If it is not, return the
+     * error message, otherwise return boost::none.
+     */
+    virtual boost::optional<std::string> validatePipelineForShardedCollection(
+        const Pipeline& pipeline) {
+        return boost::none;
+    }
+};
+
+/**
+ * A 'ServiceContext' decorator that allows enterprise to set its own version of the above class.
+ */
+extern ServiceContext::Decoration<std::unique_ptr<SearchDefaultHelperFunctions>> getSearchHelpers;
+
+}  // namespace mongo
diff --git a/src/mongo/db/pipeline/sharded_agg_helpers.cpp b/src/mongo/db/pipeline/sharded_agg_helpers.cpp
index 0fb583813f9..7f67767ee4d 100644
--- a/src/mongo/db/pipeline/sharded_agg_helpers.cpp
+++ b/src/mongo/db/pipeline/sharded_agg_helpers.cpp
@@ -47,6 +47,7 @@
 #include "mongo/db/pipeline/document_source_sort.h"
 #include "mongo/db/pipeline/document_source_unwind.h"
 #include "mongo/db/pipeline/lite_parsed_pipeline.h"
+#include "mongo/db/pipeline/search_helper.h"
 #include "mongo/db/pipeline/semantic_analysis.h"
 #include "mongo/logv2/log.h"
 #include "mongo/rpc/get_status_from_command_result.h"
@@ -56,6 +57,7 @@
 #include "mongo/s/query/document_source_merge_cursors.h"
 #include "mongo/s/query/document_source_update_on_add_shard.h"
 #include "mongo/s/query/establish_cursors.h"
+#include "mongo/s/stale_exception.h"
 #include "mongo/s/transaction_router.h"
 #include "mongo/util/fail_point.h"
 
@@ -849,6 +851,17 @@ DispatchShardPipelineResults dispatchShardPipeline(
     boost::optional<ShardedExchangePolicy> exchangeSpec;
     boost::optional<SplitPipeline> splitPipelines;
 
+    // If set, the pipeline is not valid to be run if the collection is sharded. The given string
+    // is the error message to print if the collection is sharded.
+    auto pipelinePtr = pipeline.get();
+    const auto failOnShardedCollection = [opCtx, pipelinePtr]() {
+        if (opCtx->getServiceContext() && pipelinePtr) {
+            return getSearchHelpers(opCtx->getServiceContext())
+                ->validatePipelineForShardedCollection(*pipelinePtr);
+        }
+        return boost::optional<std::string>{};
+    }();
+
     if (needsSplit) {
         LOGV2_DEBUG(20906,
                     5,
@@ -931,14 +944,35 @@ DispatchShardPipelineResults dispatchShardPipeline(
                                                            collationObj);
         }
     } else {
-        cursors = establishShardCursors(opCtx,
-                                        expCtx->mongoProcessInterface->taskExecutor,
-                                        expCtx->ns,
-                                        hasChangeStream,
-                                        executionNsRoutingInfo,
-                                        shardIds,
-                                        targetedCommand,
-                                        ReadPreferenceSetting::get(opCtx));
+        try {
+            cursors = establishShardCursors(opCtx,
+                                            expCtx->mongoProcessInterface->taskExecutor,
+                                            expCtx->ns,
+                                            mustRunOnAll,
+                                            executionNsRoutingInfo,
+                                            shardIds,
+                                            targetedCommand,
+                                            ReadPreferenceSetting::get(opCtx));
+
+        } catch (const StaleConfigException& e) {
+            // Check to see if the command failed because of a stale shard version or something
+            // else.
+            auto staleInfo = e.extraInfo<StaleConfigInfo>();
+            uassert(6441003, "StaleConfigInfo was null during sharded aggregation", staleInfo);
+            if (failOnShardedCollection && staleInfo->getVersionWanted() &&
+                staleInfo->getVersionWanted() != ChunkVersion::UNSHARDED()) {
+                // If we thought the collection was not sharded, we were wrong. Collection must be
+                // sharded.
+                uassert(5858100, *failOnShardedCollection, executionNsRoutingInfo->cm());
+            }
+            throw;
+        }
+        // If we thought the collection was sharded and the shard confirmed this, fail if the query
+        // isn't valid on a sharded collection.
+        uassert(6347900,
+                *failOnShardedCollection,
+                !failOnShardedCollection || !executionNsRoutingInfo->cm());
+
         invariant(cursors.size() % shardIds.size() == 0,
                   str::stream() << "Number of cursors (" << cursors.size()
                                 << ") is not a multiple of producers (" << shardIds.size() << ")");