TOOLS-4148 Convert test/qa-tests/jstests/export/fields_csv.js to Go

mongoexport/mongoexport_test.go

@@ -9,6 +9,7 @@ package mongoexport
 import (
 	"bytes"
 	"context"
+	"encoding/csv"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -233,6 +234,142 @@ func TestMongoExportTOOLS1952(t *testing.T) {
 	})
 }
 
+// TestExportNestedFieldsCSV verifies that mongoexport correctly handles nested
+// field paths and $ projection in --fields with --csv output.
+func TestExportNestedFieldsCSV(t *testing.T) {
+	testtype.SkipUnlessTestType(t, testtype.IntegrationTestType)
+	log.SetWriter(io.Discard)
+
+	const dbName = "mongoexport_nestedfieldscsv_test"
+	const collName = "source"
+
+	sessionProvider, _, err := testutil.GetBareSessionProvider()
+	require.NoError(t, err)
+	client, err := sessionProvider.GetSession()
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		if err := client.Database(dbName).Drop(context.Background()); err != nil {
+			t.Errorf("dropping test database: %v", err)
+		}
+	})
+
+	coll := client.Database(dbName).Collection(collName)
+	_, err = coll.InsertMany(t.Context(), []any{
+		bson.D{{"a", bson.A{1, 2, 3, 4, 5}}, {"b", bson.D{{"c", bson.A{-1, -2, -3, -4}}}}},
+		bson.D{
+			{"a", int32(1)},
+			{"b", int32(2)},
+			{"c", int32(3)},
+			{"d", bson.D{{"e", bson.A{4, 5, 6}}}},
+		},
+		bson.D{
+			{"a", int32(1)},
+			{"b", int32(2)},
+			{"c", int32(3)},
+			{"d", int32(5)},
+			{"e", bson.D{{"0", bson.A{"foo", "bar", "baz"}}}},
+		},
+		bson.D{
+			{"a", int32(1)},
+			{"b", int32(2)},
+			{"c", int32(3)},
+			{"d", bson.A{4, 5, 6}},
+			{"e", bson.A{bson.D{{"0", 0}, {"1", 1}}, bson.D{{"2", 2}, {"3", 3}}}},
+		},
+	})
+	require.NoError(t, err)
+
+	toolOptions, err := testutil.GetToolOptions()
+	require.NoError(t, err)
+	toolOptions.Namespace = &options.Namespace{DB: dbName, Collection: collName}
+
+	rows := parseCSVRows(t, exportNestedCSV(t, toolOptions, "d.e.2", ""))
+	assert.True(
+		t, rowsContainValue(rows, "d.e.2", "6"),
+		"d.e.2 should select the third element of d.e array",
+	)
+
+	rows = parseCSVRows(t, exportNestedCSV(t, toolOptions, "e.0.0", ""))
+	assert.True(
+		t, rowsContainValue(rows, "e.0.0", "foo"),
+		"e.0.0 should select nested numeric array value",
+	)
+
+	rows = parseCSVRows(t, exportNestedCSV(t, toolOptions, "b,d.1,e.1.3", ""))
+	assert.True(t, rowsContainValue(rows, "b", "2"), "b column should contain 2")
+	assert.True(t, rowsContainValue(rows, "d.1", "5"), "d.1 column should contain 5")
+	assert.True(t, rowsContainValue(rows, "e.1.3", "3"), "e.1.3 column should contain 3")
+
+	// $ projection strips the trailing .$ from the header name and wraps the result in [].
+	rows = parseCSVRows(t, exportNestedCSV(t, toolOptions, "d.$", `{"d": 4}`))
+	assert.True(
+		t, rowsContainValue(rows, "d", "[4]"),
+		"d.$ with query {d:4} should select matching element",
+	)
+
+	rows = parseCSVRows(t, exportNestedCSV(t, toolOptions, "a.$", `{"a": {"$gt": 1}}`))
+	assert.True(
+		t, rowsContainValue(rows, "a", "[2]"),
+		"a.$ with query {a:{$gt:1}} should select matching element",
+	)
+
+	rows = parseCSVRows(t, exportNestedCSV(t, toolOptions, "b.c.$", `{"b.c": -1}`))
+	assert.True(
+		t, rowsContainValue(rows, "b.c", "[-1]"),
+		"b.c.$ with query {b.c:-1} should select matching element",
+	)
+}
+
+func exportNestedCSV(t *testing.T, toolOptions *options.ToolOptions, fields, query string) string {
+	t.Helper()
+	me, err := New(Options{
+		ToolOptions: toolOptions,
+		OutputFormatOptions: &OutputFormatOptions{
+			Type:       "csv",
+			JSONFormat: "canonical",
+			Fields:     fields,
+		},
+		InputOptions: &InputOptions{Query: query},
+	})
+	require.NoError(t, err)
+	defer me.Close()
+	var buf bytes.Buffer
+	_, err = me.Export(&buf)
+	require.NoError(t, err)
+	return buf.String()
+}
+
+func parseCSVRows(t *testing.T, output string) []map[string]string {
+	t.Helper()
+	r := csv.NewReader(strings.NewReader(output))
+	records, err := r.ReadAll()
+	require.NoError(t, err)
+	if len(records) == 0 {
+		return nil
+	}
+	headers := records[0]
+	rows := make([]map[string]string, 0, len(records)-1)
+	for _, record := range records[1:] {
+		row := make(map[string]string, len(headers))
+		for i, h := range headers {
+			if i < len(record) {
+				row[h] = record[i]
+			}
+		}
+		rows = append(rows, row)
+	}
+	return rows
+}
+
+func rowsContainValue(rows []map[string]string, col, val string) bool {
+	for _, row := range rows {
+		if row[col] == val {
+			return true
+		}
+	}
+	return false
+}
+
 func TestBadOptions(t *testing.T) {
 	testtype.SkipUnlessTestType(t, testtype.IntegrationTestType)
 

mongoimport/mongoimport_test.go

@@ -31,6 +31,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/mongo"
 	mopt "go.mongodb.org/mongo-driver/v2/mongo/options"
 )
 
@@ -1886,3 +1887,113 @@ func TestRoundTripFieldFile(t *testing.T) {
 	require.NoError(t, err)
 	assert.EqualValues(t, 0, n, "c=3 should not have been exported (not in fieldFile)")
 }
+
+// TestRoundTripFieldsCSV verifies that mongoexport --csv --fields limits which
+// fields are exported, and that mongoimport correctly restores the filtered data.
+func TestRoundTripFieldsCSV(t *testing.T) {
+	testtype.SkipUnlessTestType(t, testtype.IntegrationTestType)
+
+	const dbName = "mongoimport_roundtrip_fieldscsv_test"
+
+	sessionProvider, _, err := testutil.GetBareSessionProvider()
+	require.NoError(t, err)
+	client, err := sessionProvider.GetSession()
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		if err := client.Database(dbName).Drop(context.Background()); err != nil {
+			t.Errorf("dropping test database: %v", err)
+		}
+	})
+
+	db := client.Database(dbName)
+	_, err = db.Collection("source").InsertMany(t.Context(), []any{
+		bson.D{{"a", 1}},
+		bson.D{{"a", 1}, {"b", 1}},
+		bson.D{{"a", 1}, {"b", 2}, {"c", 3}},
+	})
+	require.NoError(t, err)
+
+	exportCSVAndImport(t, dbName, "a", db)
+	dest := db.Collection("dest")
+	n, err := dest.CountDocuments(t.Context(), bson.D{{"a", 1}})
+	require.NoError(t, err)
+	assert.EqualValues(t, 3, n, "3 documents should have a=1")
+	n, err = dest.CountDocuments(t.Context(), bson.D{{"b", 1}})
+	require.NoError(t, err)
+	assert.EqualValues(t, 0, n, "b=1 should not have been exported")
+	n, err = dest.CountDocuments(t.Context(), bson.D{{"b", 2}})
+	require.NoError(t, err)
+	assert.EqualValues(t, 0, n, "b=2 should not have been exported")
+	n, err = dest.CountDocuments(t.Context(), bson.D{{"c", 3}})
+	require.NoError(t, err)
+	assert.EqualValues(t, 0, n, "c=3 should not have been exported")
+
+	exportCSVAndImport(t, dbName, "a,b,c", db)
+	n, err = dest.CountDocuments(t.Context(), bson.D{{"a", 1}})
+	require.NoError(t, err)
+	assert.EqualValues(t, 3, n, "3 documents should have a=1")
+	n, err = dest.CountDocuments(t.Context(), bson.D{{"b", 1}})
+	require.NoError(t, err)
+	assert.EqualValues(t, 1, n, "1 document should have b=1")
+	n, err = dest.CountDocuments(t.Context(), bson.D{{"b", 2}})
+	require.NoError(t, err)
+	assert.EqualValues(t, 1, n, "1 document should have b=2")
+	n, err = dest.CountDocuments(t.Context(), bson.D{{"c", 3}})
+	require.NoError(t, err)
+	assert.EqualValues(t, 1, n, "1 document should have c=3")
+
+	var fromSource, fromDest bson.M
+	q := bson.D{{"a", 1}, {"b", 1}}
+	err = db.Collection("source").FindOne(t.Context(), q).Decode(&fromSource)
+	require.NoError(t, err)
+	err = dest.FindOne(t.Context(), q).Decode(&fromDest)
+	require.NoError(t, err)
+	assert.NotEqual(t, fromSource["_id"], fromDest["_id"], "_id should not have been exported")
+}
+
+func exportCSVAndImport(t *testing.T, dbName, exportFields string, db *mongo.Database) {
+	t.Helper()
+	require.NoError(t, db.Collection("dest").Drop(t.Context()))
+
+	exportTarget, err := os.CreateTemp(t.TempDir(), "export-*.csv")
+	require.NoError(t, err)
+	require.NoError(t, exportTarget.Close())
+
+	exportToolOptions, err := testutil.GetToolOptions()
+	require.NoError(t, err)
+	exportToolOptions.Namespace = &options.Namespace{DB: dbName, Collection: "source"}
+	me, err := mongoexport.New(mongoexport.Options{
+		ToolOptions: exportToolOptions,
+		OutputFormatOptions: &mongoexport.OutputFormatOptions{
+			Type:       "csv",
+			JSONFormat: "canonical",
+			Fields:     exportFields,
+		},
+		InputOptions: &mongoexport.InputOptions{},
+	})
+	require.NoError(t, err)
+	defer me.Close()
+	f, err := os.OpenFile(exportTarget.Name(), os.O_WRONLY, 0o644)
+	require.NoError(t, err)
+	_, err = me.Export(f)
+	require.NoError(t, err)
+	require.NoError(t, f.Close())
+
+	importFields := "a,b,c"
+	importToolOptions, err := testutil.GetToolOptions()
+	require.NoError(t, err)
+	importToolOptions.Namespace = &options.Namespace{DB: dbName, Collection: "dest"}
+	mi, err := New(Options{
+		ToolOptions: importToolOptions,
+		InputOptions: &InputOptions{
+			File:       exportTarget.Name(),
+			Type:       "csv",
+			Fields:     &importFields,
+			ParseGrace: "stop",
+		},
+		IngestOptions: &IngestOptions{},
+	})
+	require.NoError(t, err)
+	_, _, err = mi.ImportDocuments()
+	require.NoError(t, err)
+}