Skip to content

Commit cc2472e

Browse files
autarchautarch
committed
Convert test/qa-tests/jstests/export/fields_csv.js to Go
1 parent 9c9dbd5 commit cc2472e

3 files changed

Lines changed: 252 additions & 173 deletions

File tree

‎mongoexport/mongoexport_test.go‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,15 @@ package mongoexport
88

99
import (
1010
"bytes"
11+
"context"
12+
"encoding/csv"
1113
"encoding/json"
1214
"errors"
1315
"fmt"
1416
"io"
1517
"os"
1618
"runtime"
19+
"strings"
1720
"testing"
1821

1922
"github.com/mongodb/mongo-tools/common/bsonutil"
@@ -229,6 +232,143 @@ func TestMongoExportTOOLS1952(t *testing.T) {
229232
})
230233
}
231234

235+
// TestExportNestedFieldsCSV verifies that mongoexport correctly handles nested
236+
// field paths and $ projection in --fields with --csv output.
237+
// Covers the nested field and projection scenarios of fields_csv.js.
238+
func TestExportNestedFieldsCSV(t *testing.T) {
239+
testtype.SkipUnlessTestType(t, testtype.IntegrationTestType)
240+
log.SetWriter(io.Discard)
241+
242+
const dbName = "mongoexport_nestedfieldscsv_test"
243+
const collName = "source"
244+
245+
sessionProvider, _, err := testutil.GetBareSessionProvider()
246+
require.NoError(t, err)
247+
client, err := sessionProvider.GetSession()
248+
require.NoError(t, err)
249+
t.Cleanup(func() {
250+
if err := client.Database(dbName).Drop(context.Background()); err != nil {
251+
t.Errorf("dropping test database: %v", err)
252+
}
253+
})
254+
255+
coll := client.Database(dbName).Collection(collName)
256+
_, err = coll.InsertMany(t.Context(), []any{
257+
bson.D{{"a", bson.A{1, 2, 3, 4, 5}}, {"b", bson.D{{"c", bson.A{-1, -2, -3, -4}}}}},
258+
bson.D{
259+
{"a", int32(1)},
260+
{"b", int32(2)},
261+
{"c", int32(3)},
262+
{"d", bson.D{{"e", bson.A{4, 5, 6}}}},
263+
},
264+
bson.D{
265+
{"a", int32(1)},
266+
{"b", int32(2)},
267+
{"c", int32(3)},
268+
{"d", int32(5)},
269+
{"e", bson.D{{"0", bson.A{"foo", "bar", "baz"}}}},
270+
},
271+
bson.D{
272+
{"a", int32(1)},
273+
{"b", int32(2)},
274+
{"c", int32(3)},
275+
{"d", bson.A{4, 5, 6}},
276+
{"e", bson.A{bson.D{{"0", 0}, {"1", 1}}, bson.D{{"2", 2}, {"3", 3}}}},
277+
},
278+
})
279+
require.NoError(t, err)
280+
281+
toolOptions, err := testutil.GetToolOptions()
282+
require.NoError(t, err)
283+
toolOptions.Namespace = &options.Namespace{DB: dbName, Collection: collName}
284+
285+
rows := parseCSVRows(t, exportNestedCSV(t, toolOptions, "d.e.2", ""))
286+
assert.True(
287+
t, rowsContainValue(rows, "d.e.2", "6"),
288+
"d.e.2 should select the third element of d.e array",
289+
)
290+
291+
rows = parseCSVRows(t, exportNestedCSV(t, toolOptions, "e.0.0", ""))
292+
assert.True(
293+
t, rowsContainValue(rows, "e.0.0", "foo"),
294+
"e.0.0 should select nested numeric array value",
295+
)
296+
297+
rows = parseCSVRows(t, exportNestedCSV(t, toolOptions, "b,d.1,e.1.3", ""))
298+
assert.True(t, rowsContainValue(rows, "b", "2"), "b column should contain 2")
299+
assert.True(t, rowsContainValue(rows, "d.1", "5"), "d.1 column should contain 5")
300+
assert.True(t, rowsContainValue(rows, "e.1.3", "3"), "e.1.3 column should contain 3")
301+
302+
// $ projection strips the trailing .$ from the header name and wraps the result in [].
303+
rows = parseCSVRows(t, exportNestedCSV(t, toolOptions, "d.$", `{"d": 4}`))
304+
assert.True(
305+
t, rowsContainValue(rows, "d", "[4]"),
306+
"d.$ with query {d:4} should select matching element",
307+
)
308+
309+
rows = parseCSVRows(t, exportNestedCSV(t, toolOptions, "a.$", `{"a": {"$gt": 1}}`))
310+
assert.True(
311+
t, rowsContainValue(rows, "a", "[2]"),
312+
"a.$ with query {a:{$gt:1}} should select matching element",
313+
)
314+
315+
rows = parseCSVRows(t, exportNestedCSV(t, toolOptions, "b.c.$", `{"b.c": -1}`))
316+
assert.True(
317+
t, rowsContainValue(rows, "b.c", "[-1]"),
318+
"b.c.$ with query {b.c:-1} should select matching element",
319+
)
320+
}
321+
322+
func exportNestedCSV(t *testing.T, toolOptions *options.ToolOptions, fields, query string) string {
323+
t.Helper()
324+
me, err := New(Options{
325+
ToolOptions: toolOptions,
326+
OutputFormatOptions: &OutputFormatOptions{
327+
Type: "csv",
328+
JSONFormat: "canonical",
329+
Fields: fields,
330+
},
331+
InputOptions: &InputOptions{Query: query},
332+
})
333+
require.NoError(t, err)
334+
defer me.Close()
335+
var buf bytes.Buffer
336+
_, err = me.Export(&buf)
337+
require.NoError(t, err)
338+
return buf.String()
339+
}
340+
341+
func parseCSVRows(t *testing.T, output string) []map[string]string {
342+
t.Helper()
343+
r := csv.NewReader(strings.NewReader(output))
344+
records, err := r.ReadAll()
345+
require.NoError(t, err)
346+
if len(records) == 0 {
347+
return nil
348+
}
349+
headers := records[0]
350+
rows := make([]map[string]string, 0, len(records)-1)
351+
for _, record := range records[1:] {
352+
row := make(map[string]string, len(headers))
353+
for i, h := range headers {
354+
if i < len(record) {
355+
row[h] = record[i]
356+
}
357+
}
358+
rows = append(rows, row)
359+
}
360+
return rows
361+
}
362+
363+
func rowsContainValue(rows []map[string]string, col, val string) bool {
364+
for _, row := range rows {
365+
if row[col] == val {
366+
return true
367+
}
368+
}
369+
return false
370+
}
371+
232372
func TestBadOptions(t *testing.T) {
233373
testtype.SkipUnlessTestType(t, testtype.IntegrationTestType)
234374

‎mongoimport/mongoimport_test.go‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import (
3131
"github.com/stretchr/testify/assert"
3232
"github.com/stretchr/testify/require"
3333
"go.mongodb.org/mongo-driver/v2/bson"
34+
"go.mongodb.org/mongo-driver/v2/mongo"
3435
mopt "go.mongodb.org/mongo-driver/v2/mongo/options"
3536
)
3637

@@ -1886,3 +1887,114 @@ func TestRoundTripFieldFile(t *testing.T) {
18861887
require.NoError(t, err)
18871888
assert.EqualValues(t, 0, n, "c=3 should not have been exported (not in fieldFile)")
18881889
}
1890+
1891+
// TestRoundTripFieldsCSV verifies that mongoexport --csv --fields limits which
1892+
// fields are exported, and that mongoimport correctly restores the filtered data.
1893+
// Covers the first two scenarios of fields_csv.js.
1894+
func TestRoundTripFieldsCSV(t *testing.T) {
1895+
testtype.SkipUnlessTestType(t, testtype.IntegrationTestType)
1896+
1897+
const dbName = "mongoimport_roundtrip_fieldscsv_test"
1898+
1899+
sessionProvider, _, err := testutil.GetBareSessionProvider()
1900+
require.NoError(t, err)
1901+
client, err := sessionProvider.GetSession()
1902+
require.NoError(t, err)
1903+
t.Cleanup(func() {
1904+
if err := client.Database(dbName).Drop(context.Background()); err != nil {
1905+
t.Errorf("dropping test database: %v", err)
1906+
}
1907+
})
1908+
1909+
db := client.Database(dbName)
1910+
_, err = db.Collection("source").InsertMany(t.Context(), []any{
1911+
bson.D{{"a", int32(1)}},
1912+
bson.D{{"a", int32(1)}, {"b", int32(1)}},
1913+
bson.D{{"a", int32(1)}, {"b", int32(2)}, {"c", int32(3)}},
1914+
})
1915+
require.NoError(t, err)
1916+
1917+
exportCSVAndImport(t, dbName, "a", db)
1918+
dest := db.Collection("dest")
1919+
n, err := dest.CountDocuments(t.Context(), bson.D{{"a", int32(1)}})
1920+
require.NoError(t, err)
1921+
assert.EqualValues(t, 3, n, "3 documents should have a=1")
1922+
n, err = dest.CountDocuments(t.Context(), bson.D{{"b", int32(1)}})
1923+
require.NoError(t, err)
1924+
assert.EqualValues(t, 0, n, "b=1 should not have been exported")
1925+
n, err = dest.CountDocuments(t.Context(), bson.D{{"b", int32(2)}})
1926+
require.NoError(t, err)
1927+
assert.EqualValues(t, 0, n, "b=2 should not have been exported")
1928+
n, err = dest.CountDocuments(t.Context(), bson.D{{"c", int32(3)}})
1929+
require.NoError(t, err)
1930+
assert.EqualValues(t, 0, n, "c=3 should not have been exported")
1931+
1932+
exportCSVAndImport(t, dbName, "a,b,c", db)
1933+
n, err = dest.CountDocuments(t.Context(), bson.D{{"a", int32(1)}})
1934+
require.NoError(t, err)
1935+
assert.EqualValues(t, 3, n, "3 documents should have a=1")
1936+
n, err = dest.CountDocuments(t.Context(), bson.D{{"b", int32(1)}})
1937+
require.NoError(t, err)
1938+
assert.EqualValues(t, 1, n, "1 document should have b=1")
1939+
n, err = dest.CountDocuments(t.Context(), bson.D{{"b", int32(2)}})
1940+
require.NoError(t, err)
1941+
assert.EqualValues(t, 1, n, "1 document should have b=2")
1942+
n, err = dest.CountDocuments(t.Context(), bson.D{{"c", int32(3)}})
1943+
require.NoError(t, err)
1944+
assert.EqualValues(t, 1, n, "1 document should have c=3")
1945+
1946+
var fromSource, fromDest bson.M
1947+
q := bson.D{{"a", int32(1)}, {"b", int32(1)}}
1948+
err = db.Collection("source").FindOne(t.Context(), q).Decode(&fromSource)
1949+
require.NoError(t, err)
1950+
err = dest.FindOne(t.Context(), q).Decode(&fromDest)
1951+
require.NoError(t, err)
1952+
assert.NotEqual(t, fromSource["_id"], fromDest["_id"], "_id should not have been exported")
1953+
}
1954+
1955+
func exportCSVAndImport(t *testing.T, dbName, exportFields string, db *mongo.Database) {
1956+
t.Helper()
1957+
require.NoError(t, db.Collection("dest").Drop(t.Context()))
1958+
1959+
exportTarget, err := os.CreateTemp(t.TempDir(), "export-*.csv")
1960+
require.NoError(t, err)
1961+
require.NoError(t, exportTarget.Close())
1962+
1963+
exportToolOptions, err := testutil.GetToolOptions()
1964+
require.NoError(t, err)
1965+
exportToolOptions.Namespace = &options.Namespace{DB: dbName, Collection: "source"}
1966+
me, err := mongoexport.New(mongoexport.Options{
1967+
ToolOptions: exportToolOptions,
1968+
OutputFormatOptions: &mongoexport.OutputFormatOptions{
1969+
Type: "csv",
1970+
JSONFormat: "canonical",
1971+
Fields: exportFields,
1972+
},
1973+
InputOptions: &mongoexport.InputOptions{},
1974+
})
1975+
require.NoError(t, err)
1976+
defer me.Close()
1977+
f, err := os.OpenFile(exportTarget.Name(), os.O_WRONLY, 0o644)
1978+
require.NoError(t, err)
1979+
_, err = me.Export(f)
1980+
require.NoError(t, err)
1981+
require.NoError(t, f.Close())
1982+
1983+
importFields := "a,b,c"
1984+
importToolOptions, err := testutil.GetToolOptions()
1985+
require.NoError(t, err)
1986+
importToolOptions.Namespace = &options.Namespace{DB: dbName, Collection: "dest"}
1987+
mi, err := New(Options{
1988+
ToolOptions: importToolOptions,
1989+
InputOptions: &InputOptions{
1990+
File: exportTarget.Name(),
1991+
Type: "csv",
1992+
Fields: &importFields,
1993+
ParseGrace: "stop",
1994+
},
1995+
IngestOptions: &IngestOptions{},
1996+
})
1997+
require.NoError(t, err)
1998+
_, _, err = mi.ImportDocuments()
1999+
require.NoError(t, err)
2000+
}

0 commit comments

Comments
 (0)