TOOLS-657: properly project dotted fields

4 files changed, 91 insertions, 14 deletions

diff --git a/mongoexport/mongoexport.go b/mongoexport/mongoexport.go
index 99bcb56f4c9..8e74ef6bd76 100644
--- a/mongoexport/mongoexport.go
+++ b/mongoexport/mongoexport.go
@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"github.com/mongodb/mongo-tools/common/bsonutil"
 	"github.com/mongodb/mongo-tools/common/db"
-	sloppyjson "github.com/mongodb/mongo-tools/common/json"
+	"github.com/mongodb/mongo-tools/common/json"
 	"github.com/mongodb/mongo-tools/common/log"
 	"github.com/mongodb/mongo-tools/common/options"
 	"github.com/mongodb/mongo-tools/common/util"
@@ -132,17 +132,25 @@ func (exp *MongoExport) GetOutputWriter() (io.WriteCloser, error) {
 }
 
 // Take a comma-delimited set of field names and build a selector doc for query projection.
-// e.g. "a,b,c.d" -> {a:1, b:1, "c.d":1}
+// For fields containing a dot '.', we project the entire top-level portion.
+// e.g. "a,b,c.d.e,f.$" -> {a:1, b:1, "c":1, "f.$": 1}.
 func makeFieldSelector(fields string) bson.M {
-	r := bson.M{"_id": 1}
+	selector := bson.M{"_id": 1}
 	if fields == "" {
-		return r
+		return selector
 	}
-	f := strings.Split(fields, ",")
-	for _, v := range f {
-		r[v] = 1
+
+	for _, field := range strings.Split(fields, ",") {
+		// Projections like "a.0" work fine for nested documents not for arrays
+		// - if passed directly to mongod. To handle this, we have to retrieve
+		// the entire top-level document and then filter afterwards. An exception
+		// is made for '$' projections - which are passed directly to mongod.
+		if i := strings.LastIndex(field, "."); i != -1 && field[i+1:] != "$" {
+			field = field[:strings.Index(field, ".")]
+		}
+		selector[field] = 1
 	}
-	return r
+	return selector
 }
 
 // getCursor returns a cursor that can be iterated over to get all the documents
@@ -244,6 +252,7 @@ func (exp *MongoExport) exportInternal(out io.Writer) (int64, error) {
 	var result bson.M
 
 	docsCount := int64(0)
+
 	// Write document content
 	for cursor.Next(&result) {
 		err := exportOutput.ExportDocument(result)
@@ -291,7 +300,17 @@ func (exp *MongoExport) getExportOutput(out io.Writer) (ExportOutput, error) {
 		} else {
 			return nil, fmt.Errorf("CSV mode requires a field list")
 		}
-		return NewCSVExportOutput(fields, out), nil
+
+		exportFields := make([]string, 0, len(fields))
+		for _, field := range fields {
+			// for '$' field projections, exclude '.$' from the field name
+			if i := strings.LastIndex(field, "."); i != -1 && field[i+1:] == "$" {
+				exportFields = append(exportFields, field[:i])
+			} else {
+				exportFields = append(exportFields, field)
+			}
+		}
+		return NewCSVExportOutput(exportFields, out), nil
 	}
 	return NewJSONExportOutput(exp.OutputOpts.JSONArray, exp.OutputOpts.Pretty, out), nil
 }
@@ -301,7 +320,7 @@ func (exp *MongoExport) getExportOutput(out io.Writer) (ExportOutput, error) {
 // Returns an error if the string is not valid JSON, or extended JSON.
 func getObjectFromArg(queryRaw string) (map[string]interface{}, error) {
 	parsedJSON := map[string]interface{}{}
-	err := sloppyjson.Unmarshal([]byte(queryRaw), &parsedJSON)
+	err := json.Unmarshal([]byte(queryRaw), &parsedJSON)
 	if err != nil {
 		return nil, fmt.Errorf("query '%v' is not valid JSON: %v", queryRaw, err)
 	}
@@ -317,7 +336,7 @@ func getObjectFromArg(queryRaw string) (map[string]interface{}, error) {
 // object which preserves the ordering of the keys as they appear in the input.
 func getSortFromArg(queryRaw string) (bson.D, error) {
 	parsedJSON := bson.D{}
-	err := sloppyjson.Unmarshal([]byte(queryRaw), &parsedJSON)
+	err := json.Unmarshal([]byte(queryRaw), &parsedJSON)
 	if err != nil {
 		return nil, fmt.Errorf("query '%v' is not valid JSON: %v", queryRaw, err)
 	}
diff --git a/mongoexport/mongoexport_test.go b/mongoexport/mongoexport_test.go
index c11488bd858..f5d0e7e2afa 100644
--- a/mongoexport/mongoexport_test.go
+++ b/mongoexport/mongoexport_test.go
@@ -39,6 +39,6 @@ func TestFieldSelect(t *testing.T) {
 	Convey("Using makeFieldSelector should return correct projection doc", t, func() {
 		So(makeFieldSelector("a,b"), ShouldResemble, bson.M{"_id": 1, "a": 1, "b": 1})
 		So(makeFieldSelector(""), ShouldResemble, bson.M{"_id": 1})
-		So(makeFieldSelector("x,foo.baz"), ShouldResemble, bson.M{"_id": 1, "foo.baz": 1, "x": 1})
+		So(makeFieldSelector("x,foo.baz"), ShouldResemble, bson.M{"_id": 1, "foo": 1, "x": 1})
 	})
 }
diff --git a/mongoexport/options.go b/mongoexport/options.go
index 93b4da148f4..40c1013332a 100644
--- a/mongoexport/options.go
+++ b/mongoexport/options.go
@@ -28,7 +28,7 @@ type OutputFormatOptions struct {
 }
 
 // Name returns a human-readable group name for output format options.
-func (_ *OutputFormatOptions) Name() string {
+func (*OutputFormatOptions) Name() string {
 	return "output"
 }
 
@@ -43,6 +43,6 @@ type InputOptions struct {
 }
 
 // Name returns a human-readable group name for input options.
-func (_ *InputOptions) Name() string {
+func (*InputOptions) Name() string {
 	return "querying"
 }
diff --git a/test/qa-tests/jstests/export/fields_csv.js b/test/qa-tests/jstests/export/fields_csv.js
index 9e0f3ade926..3f1991aaae6 100644
--- a/test/qa-tests/jstests/export/fields_csv.js
+++ b/test/qa-tests/jstests/export/fields_csv.js
@@ -90,6 +90,64 @@
     var fromDest = destColl.findOne({ a: 1, b: 1 });
     assert.neq(fromSource._id, fromDest._id);
 
+
+    /* Test passing positional arguments to --fields */
+
+    // outputMatchesExpected takes an output string and returns
+    // a boolean indicating if any line of the output matched
+    // the expected string.
+    var outputMatchesExpected = function(output, expected) {
+        var found = false;
+        output.split('\n').forEach(function(line) {
+          if (line.match(expected)) found = true;
+        });
+        return found;
+    }
+
+    // remove the export, clear the destination collection
+    removeFile(exportTarget);
+    sourceColl.remove({});
+
+    // ensure source collection is empty
+    assert.eq(0, sourceColl.count());
+
+    // insert some data
+    sourceColl.insert({ a: [1, 2, 3, 4, 5], b: { c: [-1, -2, -3, -4] } });
+    sourceColl.insert({ a: 1, b: 2, c: 3, d: { e: [4, 5, 6] } });
+    sourceColl.insert({ a: 1, b: 2, c: 3, d: 5, e: {"0": ["foo", "bar", "baz"] } });
+    sourceColl.insert({ a: 1, b: 2, c: 3, d: [4, 5, 6], e: [ {"0": 0, "1": 1}, {"2": 2, "3": 3} ] });
+
+    // ensure the insertion worked
+    assert.eq(4, sourceColl.count());
+
+    // use the following fields as filters:
+    var cases = [
+        { field: 'd.e.2', expected: /6/ }, // specify nested field with array value
+        { field: 'e.0.0', expected: /foo/ }, // specify nested field with numeric array value
+        { field: 'b,d.1,e.1.3', expected: /2,5,3/ }, // specify varying levels of field nesting
+    ];
+
+    for (var i = 0; i < cases.length; i++) {
+        assert.eq(0, toolTest.runTool.apply(toolTest,['export', '--fields', cases[i].field, '--out', exportTarget, '--db', 'test', '--collection', 'source', '--csv'].concat(commonToolArgs)));
+        var output = cat(exportTarget);
+        jsTest.log("Fields Test " + (i + 1) + ": \n" + output);
+        assert.eq(outputMatchesExpected(output, cases[i].expected), true);
+    }
+
+    // test with $ projection and query
+    cases = [
+        { query: '{ d: 4 }', field: 'd.$', expected: /[4]/ },
+        { query: '{ a: { $gt: 1 } }', field: 'a.$', expected: /[2]/ },
+        { query: '{ "b.c": -1 }', field: 'b.c.$', expected: /[-1]/ },
+    ];
+
+    for (var i = 0; i < cases.length; i++) {
+        assert.eq(0, toolTest.runTool.apply(toolTest,['export', '--query', cases[i].query, '--fields', cases[i].field, '--out', exportTarget, '--db', 'test', '--collection', 'source', '--csv'].concat(commonToolArgs)));
+        var output = cat(exportTarget);
+        jsTest.log("Fields + Query Test " + (i + 1) + ": \n" + output);
+        assert.eq(outputMatchesExpected(output, cases[i].expected), true);
+    }
+
     // success
     toolTest.stop();