SERVER-70966 [CQF] Unit test auto-updating support

5 files changed, 216 insertions, 16 deletions

diff --git a/src/mongo/db/query/optimizer/explain.cpp b/src/mongo/db/query/optimizer/explain.cpp
index b1bb27568d5..5e9df5df3bb 100644
--- a/src/mongo/db/query/optimizer/explain.cpp
+++ b/src/mongo/db/query/optimizer/explain.cpp
@@ -2542,8 +2542,12 @@ static void printBSONstr(PrinterType& printer,
     }
 }
 
-std::string ExplainGenerator::printBSON(const sbe::value::TypeTags tag,
-                                        const sbe::value::Value val) {
+std::string ExplainGenerator::explainBSONStr(const ABT& node,
+                                             bool displayProperties,
+                                             const cascades::MemoExplainInterface* memoInterface,
+                                             const NodeToGroupPropsMap& nodeMap) {
+    const auto [tag, val] = explainBSON(node, displayProperties, memoInterface, nodeMap);
+    sbe::value::ValueGuard vg(tag, val);
     ExplainPrinterImpl<ExplainVersion::V2> printer;
     printBSONstr(printer, tag, val);
     return printer.str();
diff --git a/src/mongo/db/query/optimizer/explain.h b/src/mongo/db/query/optimizer/explain.h
index ad62dd54126..19e9221f8dd 100644
--- a/src/mongo/db/query/optimizer/explain.h
+++ b/src/mongo/db/query/optimizer/explain.h
@@ -94,7 +94,10 @@ public:
                                   const cascades::MemoExplainInterface* memoInterface = nullptr,
                                   const NodeToGroupPropsMap& nodeMap = {});
 
-    static std::string printBSON(sbe::value::TypeTags tag, sbe::value::Value val);
+    static std::string explainBSONStr(const ABT& node,
+                                      bool displayProperties = false,
+                                      const cascades::MemoExplainInterface* memoInterface = nullptr,
+                                      const NodeToGroupPropsMap& nodeMap = {});
 
     static std::string explainLogicalProps(const std::string& description,
                                            const properties::LogicalProps& props);
diff --git a/src/mongo/db/query/optimizer/optimizer_test.cpp b/src/mongo/db/query/optimizer/optimizer_test.cpp
index d52b84470a1..423a6a399af 100644
--- a/src/mongo/db/query/optimizer/optimizer_test.cpp
+++ b/src/mongo/db/query/optimizer/optimizer_test.cpp
@@ -29,18 +29,45 @@
 
 #include "mongo/db/query/optimizer/explain.h"
 #include "mongo/db/query/optimizer/node.h"
-#include "mongo/db/query/optimizer/opt_phase_manager.h"
 #include "mongo/db/query/optimizer/reference_tracker.h"
 #include "mongo/db/query/optimizer/rewrites/const_eval.h"
 #include "mongo/db/query/optimizer/syntax/syntax.h"
-#include "mongo/db/query/optimizer/syntax/syntax_fwd_declare.h"
 #include "mongo/db/query/optimizer/utils/unit_test_utils.h"
 #include "mongo/db/query/optimizer/utils/utils.h"
 #include "mongo/unittest/unittest.h"
 
+
 namespace mongo::optimizer {
 namespace {
 
+TEST(Optimizer, AutoUpdateExplain) {
+    ABT tree = make<BinaryOp>(Operations::Add,
+                              Constant::int64(1),
+                              make<Variable>("very very very very very very very very very very "
+                                             "very very long variable name with \"quotes\""));
+
+    /**
+     * To exercise the auto-updating behavior:
+     *   1. Change the flag "kAutoUpdateOnFailure" to "true".
+     *   2. Induce a failure: change something in the expected output.
+     *   3. Recompile and run the test binary as normal.
+     *   4. Observe after the run the test file is updated with the correct output.
+     */
+    ASSERT_EXPLAIN_V2_AUTO(  // NOLINT (test auto-update)
+        "BinaryOp [Add]\n"
+        "|   Variable [very very very very very very very very very very very very long variable "
+        "name with \"quotes\"]\n"
+        "Const [1]\n",
+        tree);
+
+    // Test for short constant. It should not be inlined. The nolint comment on the string constant
+    // itself is auto-generated.
+    ABT tree1 = make<Variable>("short name");
+    ASSERT_EXPLAIN_V2_AUTO(         // NOLINT (test auto-update)
+        "Variable [short name]\n",  // NOLINT (test auto-update)
+        tree1);
+}
+
 Constant* constEval(ABT& tree) {
     auto env = VariableEnvironment::build(tree);
     ConstEval evaluator{env};
@@ -746,13 +773,14 @@ TEST(Explain, ExplainV2Compact) {
 
 TEST(Explain, ExplainBsonForConstant) {
     ABT cNode = Constant::int64(3);
-    auto [tag, val] = ExplainGenerator::explainBSON(cNode);
-    sbe::value::ValueGuard vg(tag, val);
-    ASSERT_EQ(
-        "{\n    nodeType: \"Const\", \n"
+
+    ASSERT_EXPLAIN_BSON(
+        "{\n"
+        "    nodeType: \"Const\", \n"
         "    tag: \"NumberInt64\", \n"
-        "    value: 3\n}\n",
-        ExplainGenerator::printBSON(tag, val));
+        "    value: 3\n"
+        "}\n",
+        cNode);
 }
 
 }  // namespace
diff --git a/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp b/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp
index e7ee776d217..538a4b1f036 100644
--- a/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp
+++ b/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp
@@ -29,6 +29,8 @@
 
 #include "mongo/db/query/optimizer/utils/unit_test_utils.h"
 
+#include <fstream>
+
 #include "mongo/db/pipeline/abt/utils.h"
 #include "mongo/db/query/cost_model/cost_estimator.h"
 #include "mongo/db/query/cost_model/cost_model_manager.h"
@@ -37,27 +39,162 @@
 #include "mongo/db/query/optimizer/metadata.h"
 #include "mongo/db/query/optimizer/node.h"
 #include "mongo/db/query/optimizer/rewrites/const_eval.h"
-#include "mongo/unittest/unittest.h"
+#include "mongo/util/str_escape.h"
 
 
 namespace mongo::optimizer {
 
 static constexpr bool kDebugAsserts = false;
 
+// DO NOT COMMIT WITH "TRUE".
+static constexpr bool kAutoUpdateOnFailure = false;
+static constexpr const char* kTempFileSuffix = ".tmp.txt";
+
+// Map from file name to a list of updates. We keep track of how many lines are added or deleted at
+// a particular line of a source file.
+using LineDeltaVector = std::vector<std::pair<uint64_t, int64_t>>;
+std::map<std::string, LineDeltaVector> gLineDeltaMap;
+
+
 void maybePrintABT(const ABT& abt) {
     // Always print using the supported versions to make sure we don't crash.
     const std::string strV1 = ExplainGenerator::explain(abt);
     const std::string strV2 = ExplainGenerator::explainV2(abt);
     const std::string strV2Compact = ExplainGenerator::explainV2Compact(abt);
-    auto [tag, val] = ExplainGenerator::explainBSON(abt);
-    sbe::value::ValueGuard vg(tag, val);
+    const std::string strBSON = ExplainGenerator::explainBSONStr(abt);
 
     if constexpr (kDebugAsserts) {
         std::cout << "V1: " << strV1 << "\n";
         std::cout << "V2: " << strV2 << "\n";
         std::cout << "V2Compact: " << strV2Compact << "\n";
-        std::cout << "BSON: " << ExplainGenerator::printBSON(tag, val) << "\n";
+        std::cout << "BSON: " << strBSON << "\n";
+    }
+}
+
+static std::vector<std::string> formatStr(const std::string& str) {
+    std::vector<std::string> replacementLines;
+    std::istringstream lineInput(str);
+
+    // Account for maximum line length after linting. We need to indent, add quotes, etc.
+    static constexpr size_t kEscapedLength = 88;
+
+    std::string line;
+    while (std::getline(lineInput, line)) {
+        // Read the string line by line and format it to match the test file's expected format. We
+        // have an initial indentation, followed by quotes and the escaped string itself.
+
+        std::string escaped = mongo::str::escapeForJSON(line);
+        for (;;) {
+            // If the line is estimated to exceed the maximum length allowed by the linter, break it
+            // up and make sure to insert '\n' only at the end of the last segment.
+            const bool breakupLine = escaped.size() > kEscapedLength;
+
+            std::ostringstream os;
+            os << "        \"" << escaped.substr(0, kEscapedLength);
+            if (!breakupLine) {
+                os << "\\n";
+            }
+            os << "\"\n";
+            replacementLines.push_back(os.str());
+
+            if (breakupLine) {
+                escaped = escaped.substr(kEscapedLength);
+            } else {
+                break;
+            }
+        }
+    }
+
+    if (!replacementLines.empty() && !replacementLines.back().empty()) {
+        // Account for the fact that we need an extra comma after the string constant in the macro.
+        auto& lastLine = replacementLines.back();
+        lastLine.insert(lastLine.size() - 1, ",");
+
+        if (replacementLines.size() == 1) {
+            // For single lines, add a 'nolint' comment to prevent the linter from inlining the
+            // single line with the macro itself.
+            lastLine.insert(lastLine.size() - 1, "  // NOLINT (test auto-update)");
+        }
     }
+
+    return replacementLines;
+}
+
+bool handleAutoUpdate(const std::string& expected,
+                      const std::string& actual,
+                      const std::string& fileName,
+                      const size_t lineNumber) {
+    if (expected == actual) {
+        return true;
+    }
+    if constexpr (!kAutoUpdateOnFailure) {
+        std::cout << "Auto-updating is disabled.\n";
+        return false;
+    }
+
+    const auto expectedFormatted = formatStr(expected);
+    const auto actualFormatted = formatStr(actual);
+
+    std::cout << "Updating expected result in file '" << fileName << "', line: " << lineNumber
+              << ".\n";
+    std::cout << "Replacement:\n";
+    for (const auto& line : actualFormatted) {
+        std::cout << line;
+    }
+
+    // Compute the total number of lines added or removed before the current macro line.
+    auto& lineDeltas = gLineDeltaMap.emplace(fileName, LineDeltaVector{}).first->second;
+    int64_t totalDelta = 0;
+    for (const auto& [line, delta] : lineDeltas) {
+        if (line < lineNumber) {
+            totalDelta += delta;
+        }
+    }
+
+    const size_t replacementEndLine = lineNumber + totalDelta;
+    // Treat an empty string as needing one line. Adjust for line delta.
+    const size_t replacementStartLine =
+        replacementEndLine - (expectedFormatted.empty() ? 1 : expectedFormatted.size());
+
+    const std::string tempFileName = fileName + kTempFileSuffix;
+    std::string line;
+    size_t lineIndex = 0;
+
+    try {
+        std::ifstream in;
+        in.open(fileName);
+        std::ofstream out;
+        out.open(tempFileName);
+
+        // Generate a new test file, updated with the replacement string.
+        while (std::getline(in, line)) {
+            lineIndex++;
+
+            if (lineIndex < replacementStartLine || lineIndex >= replacementEndLine) {
+                out << line << "\n";
+            } else if (lineIndex == replacementStartLine) {
+                for (const auto& line1 : actualFormatted) {
+                    out << line1;
+                }
+            }
+        }
+
+        out.close();
+        in.close();
+
+        std::rename(tempFileName.c_str(), fileName.c_str());
+    } catch (const std::exception& ex) {
+        // Print and re-throw exception.
+        std::cout << "Caught an exception while manipulating files: " << ex.what();
+        throw ex;
+    }
+
+    // Add the current delta.
+    lineDeltas.emplace_back(
+        lineNumber, static_cast<int64_t>(actualFormatted.size()) - expectedFormatted.size());
+
+    // Do not assert in order to allow multiple tests to be updated.
+    return true;
 }
 
 ABT makeIndexPath(FieldPathType fieldPath, bool isMultiKey) {
diff --git a/src/mongo/db/query/optimizer/utils/unit_test_utils.h b/src/mongo/db/query/optimizer/utils/unit_test_utils.h
index 69c1f3ce8db..e358dc13495 100644
--- a/src/mongo/db/query/optimizer/utils/unit_test_utils.h
+++ b/src/mongo/db/query/optimizer/utils/unit_test_utils.h
@@ -39,6 +39,11 @@ namespace mongo::optimizer {
 
 void maybePrintABT(const ABT& abt);
 
+bool handleAutoUpdate(const std::string& expected,
+                      const std::string& actual,
+                      const std::string& fileName,
+                      size_t lineNumber);
+
 #define ASSERT_EXPLAIN(expected, abt) \
     maybePrintABT(abt);               \
     ASSERT_EQ(expected, ExplainGenerator::explain(abt))
@@ -47,13 +52,36 @@ void maybePrintABT(const ABT& abt);
     maybePrintABT(abt);                  \
     ASSERT_EQ(expected, ExplainGenerator::explainV2(abt))
 
+/**
+ * Auto update result back in the source file if the assert fails.
+ * The expected result must be a multi-line string in the following form:
+ *
+ * ASSERT_EXPLAIN_V2_AUTO(     // NOLINT
+ *       "BinaryOp [Add]\n"
+ *       "|   Const [2]\n"
+ *       "Const [1]\n",
+ *       tree);
+ *
+ * Limitations:
+ *      1. There should not be any comments or other formatting inside the multi-line string
+ *      constant other than 'NOLINT'. If we have a single-line constant, the auto-updating will
+ *      generate a 'NOLINT' at the end of the line.
+ *      2. The expression which we are explaining ('tree' in the example above) must fit on a single
+ *      line. The macro should be indented by 4 spaces.
+ *
+ * TODO: SERVER-71004: Extend the usability of the auto-update macro.
+ */
+#define ASSERT_EXPLAIN_V2_AUTO(expected, abt) \
+    maybePrintABT(abt);                       \
+    ASSERT(handleAutoUpdate(expected, ExplainGenerator::explainV2(abt), __FILE__, __LINE__))
+
 #define ASSERT_EXPLAIN_V2Compact(expected, abt) \
     maybePrintABT(abt);                         \
     ASSERT_EQ(expected, ExplainGenerator::explainV2Compact(abt))
 
 #define ASSERT_EXPLAIN_BSON(expected, abt) \
     maybePrintABT(abt);                    \
-    ASSERT_EQ(expected, ExplainGenerator::explainBSON(abt))
+    ASSERT_EQ(expected, ExplainGenerator::explainBSONStr(abt))
 
 #define ASSERT_EXPLAIN_PROPS_V2(expected, phaseManager)                              \
     ASSERT_EQ(expected,                                                              \