TOOLS-4148 Convert mode_upsert_id_subdoc.js to Go

mongoimport/mongoimport_test.go

@@ -3411,6 +3411,116 @@ func runImportOpts(
 	return err
 }
 
+// TestImportModeUpsertIDSubdoc verifies that --mode=upsert uses the full
+// subdocument _id as the upsert key, preserving field order through a
+// round-trip of export → upsert-replace → re-import.
+func TestImportModeUpsertIDSubdoc(t *testing.T) {
+	testtype.SkipUnlessTestType(t, testtype.IntegrationTestType)
+
+	const (
+		dbName   = "mongoimport_upsert_subdoc_test"
+		collName = "c"
+	)
+
+	sessionProvider, _, err := testutil.GetBareSessionProvider()
+	require.NoError(t, err)
+	client, err := sessionProvider.GetSession()
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = client.Database(dbName).Drop(context.Background())
+	})
+
+	coll := client.Database(dbName).Collection(collName)
+	ns := &options.Namespace{DB: dbName, Collection: collName}
+
+	origDocs := subdocIDDocs("string")
+	insertDocs := make([]any, len(origDocs))
+	for i, d := range origDocs {
+		insertDocs[i] = d
+	}
+	_, err = coll.InsertMany(t.Context(), insertDocs)
+	require.NoError(t, err)
+
+	exportedFile := exportCollectionToFile(t, ns)
+	str2File := writeSubdocIDFile(t, "str2")
+
+	t.Run("upsert with replacement data updates all docs in place", func(t *testing.T) {
+		require.NoError(t, runImportOpts(t, ns, str2File, IngestOptions{Mode: modeUpsert}))
+		n, err := coll.CountDocuments(t.Context(), bson.D{})
+		require.NoError(t, err)
+		assert.EqualValues(t, 20, n, "count should be unchanged after upsert")
+		n, err = coll.CountDocuments(t.Context(), bson.D{{"x", "str2"}})
+		require.NoError(t, err)
+		assert.EqualValues(t, 20, n, "all docs should have x=str2 after upsert")
+	})
+
+	t.Run("re-import original export reverts all docs", func(t *testing.T) {
+		require.NoError(t, runImportOpts(t, ns, exportedFile, IngestOptions{Mode: modeUpsert}))
+		n, err := coll.CountDocuments(t.Context(), bson.D{})
+		require.NoError(t, err)
+		assert.EqualValues(t, 20, n, "count should be unchanged after re-import")
+		n, err = coll.CountDocuments(t.Context(), bson.D{{"x", "string"}})
+		require.NoError(t, err)
+		assert.EqualValues(t, 20, n, "all docs should have x=string after re-import")
+	})
+}
+
+func exportCollectionToFile(t *testing.T, ns *options.Namespace) string {
+	t.Helper()
+	exportFile, err := os.CreateTemp(t.TempDir(), "export-*.json")
+	require.NoError(t, err)
+	exportToolOptions, err := testutil.GetToolOptions()
+	require.NoError(t, err)
+	exportToolOptions.Namespace = ns
+	me, err := mongoexport.New(mongoexport.Options{
+		ToolOptions: exportToolOptions,
+		OutputFormatOptions: &mongoexport.OutputFormatOptions{
+			Type:       "json",
+			JSONFormat: "canonical",
+		},
+		InputOptions: &mongoexport.InputOptions{},
+	})
+	require.NoError(t, err)
+	defer me.Close()
+	_, err = me.Export(exportFile)
+	require.NoError(t, err)
+	require.NoError(t, exportFile.Close())
+	return exportFile.Name()
+}
+
+func writeSubdocIDFile(t *testing.T, xFieldValue string) string {
+	t.Helper()
+	f, err := os.CreateTemp(t.TempDir(), "subdoc-*.json")
+	require.NoError(t, err)
+	for _, doc := range subdocIDDocs(xFieldValue) {
+		b, err := bson.MarshalExtJSON(doc, true, false)
+		require.NoError(t, err)
+		_, err = f.Write(b)
+		require.NoError(t, err)
+		_, err = f.Write([]byte("\n"))
+		require.NoError(t, err)
+	}
+	require.NoError(t, f.Close())
+	return f.Name()
+}
+
+func subdocIDDocs(xFieldValue string) []bson.D {
+	docs := make([]bson.D, 0, 20)
+	for i := range 4 {
+		for j := range 5 {
+			docs = append(docs, bson.D{
+				{"_id", bson.D{
+					{"a", i},
+					{"b", bson.A{0, 1, 2, bson.D{{"c", j}, {"d", "foo"}}}},
+					{"e", "bar"},
+				}},
+				{"x", xFieldValue},
+			})
+		}
+	}
+	return docs
+}
+
 // writeJSONLinesFile marshals each doc in docs as a JSON object and writes them
 // as newline-separated lines to a file in dir named name. Returns the file path.
 func writeJSONLinesFile(t *testing.T, dir, name string, docs []map[string]any) string {