Skip to content

Commit 8479210

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

3 files changed

Lines changed: 248 additions & 173 deletions

File tree

‎mongoexport/mongoexport_test.go‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ package mongoexport
99
import (
1010
"bytes"
1111
"context"
12+
"encoding/csv"
1213
"encoding/json"
1314
"errors"
1415
"fmt"
@@ -233,6 +234,142 @@ func TestMongoExportTOOLS1952(t *testing.T) {
233234
})
234235
}
235236

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

‎mongoimport/mongoimport_test.go‎

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

0 commit comments

Comments
 (0)