TOOLS-716: Mongodump: add option to read query from file

mongodump/mongodump.go

@@ -63,6 +63,8 @@ func (dump *MongoDump) ValidateOptions() error {
 		return fmt.Errorf("cannot dump a collection without a specified database")
 	case dump.InputOptions.Query != "" && dump.ToolOptions.Namespace.Collection == "":
 		return fmt.Errorf("cannot dump using a query without a specified collection")
+	case dump.InputOptions.QueryFile != "" && dump.ToolOptions.Namespace.Collection == "":
+		return fmt.Errorf("cannot dump using a queryFile without a specified collection")
 	case dump.OutputOptions.DumpDBUsersAndRoles && dump.ToolOptions.Namespace.DB == "":
 		return fmt.Errorf("must specify a database when running with dumpDbUsersAndRoles")
 	case dump.OutputOptions.DumpDBUsersAndRoles && dump.ToolOptions.Namespace.Collection != "":
@@ -79,6 +81,8 @@ func (dump *MongoDump) ValidateOptions() error {
 		return fmt.Errorf("--db is required when --excludeCollectionsWithPrefix is specified")
 	case dump.OutputOptions.Repair && dump.InputOptions.Query != "":
 		return fmt.Errorf("cannot run a query with --repair enabled")
+	case dump.OutputOptions.Repair && dump.InputOptions.QueryFile != "":
+		return fmt.Errorf("cannot run a queryFile with --repair enabled")
 	case dump.OutputOptions.Out != "" && dump.OutputOptions.Archive != "":
 		return fmt.Errorf("--out not allowed when --archive is specified")
 	case dump.OutputOptions.Out == "-" && dump.OutputOptions.Gzip:
@@ -129,10 +133,18 @@ func (dump *MongoDump) Init() error {
 
 // Dump handles some final options checking and executes MongoDump.
 func (dump *MongoDump) Dump() (err error) {
-	if dump.InputOptions.Query != "" {
+	if dump.InputOptions.Query != "" && dump.InputOptions.QueryFile != ""{
+		return fmt.Errorf("either query or queryFile can be specified as a query option")
+	}
+
+	if dump.InputOptions.HasQuery() {
 		// parse JSON then convert extended JSON values
 		var asJSON interface{}
-		err = json.Unmarshal([]byte(dump.InputOptions.Query), &asJSON)
+		content, err := dump.InputOptions.GetQuery()
+		if err != nil {
+			return err
+		}
+		err = json.Unmarshal(content, &asJSON)
 		if err != nil {
 			return fmt.Errorf("error parsing query as json: %v", err)
 		}

mongodump/mongodump_test.go

@@ -466,9 +466,9 @@ func TestMongoDumpBSON(t *testing.T) {
 			So(err, ShouldBeNil)
 			jsonQueryBytes, err := json.Marshal(jsonQuery)
 			So(err, ShouldBeNil)
-			md.InputOptions.Query = string(jsonQueryBytes)
 
-			Convey("for all the collections in the database", func() {
+			Convey("using --query for all the collections in the database", func() {
+				md.InputOptions.Query = string(jsonQueryBytes)
 				md.ToolOptions.Namespace.DB = testDB
 				md.OutputOptions.Out = "dump"
 
@@ -516,6 +516,56 @@ func TestMongoDumpBSON(t *testing.T) {
 
 			})
 
+			Convey("using --queryFile for all the collections in the database", func() {
+				ioutil.WriteFile("example.json", jsonQueryBytes, 0777)
+				md.InputOptions.QueryFile = "example.json"
+				md.ToolOptions.Namespace.DB = testDB
+				md.OutputOptions.Out = "dump"
+
+				origDB := session.DB(testDB)
+				restoredDB := session.DB(testRestoreDB)
+
+				// we can only dump using query per collection
+				for _, testCollName := range testCollectionNames {
+					md.ToolOptions.Namespace.Collection = testCollName
+
+					err = md.Init()
+					So(err, ShouldBeNil)
+
+					err = md.Dump()
+					So(err, ShouldBeNil)
+				}
+
+				path, err := os.Getwd()
+				So(err, ShouldBeNil)
+
+				dumpDir := util.ToUniversalPath(filepath.Join(path, "dump"))
+				dumpDBDir := util.ToUniversalPath(filepath.Join(dumpDir, testDB))
+				So(fileDirExists(dumpDir), ShouldBeTrue)
+				So(fileDirExists(dumpDBDir), ShouldBeTrue)
+
+				err = readBSONIntoDatabase(dumpDBDir, testRestoreDB)
+				So(err, ShouldBeNil)
+
+				for _, testCollName := range testCollectionNames {
+					// count filtered docs
+					numDocs1, err := origDB.C(testCollName).Find(bsonQuery).Count()
+					So(err, ShouldBeNil)
+
+					// count number of all restored documents
+					numDocs2, err := restoredDB.C(testCollName).Find(nil).Count()
+					So(err, ShouldBeNil)
+
+					So(numDocs1, ShouldEqual, numDocs2)
+				}
+
+				Reset(func() {
+					So(session.DB(testRestoreDB).DropDatabase(), ShouldBeNil)
+					So(os.RemoveAll(dumpDir), ShouldBeNil)
+					So(os.Remove("example.json"), ShouldBeNil)
+				})
+
+			})
 		})
 
 		Reset(func() {

mongodump/options.go

@@ -1,4 +1,8 @@
 package mongodump
+import (
+	"fmt"
+	"io/ioutil"
+)
 
 var Usage = `<options>
 
@@ -10,15 +14,34 @@ See http://docs.mongodb.org/manual/reference/program/mongodump/ for more informa
 
 // InputOptions defines the set of options to use in retrieving data from the server.
 type InputOptions struct {
-	Query     string `long:"query" short:"q" description:"query filter, as a JSON string, e.g., '{x:{$gt:1}}'"`
-	TableScan bool   `long:"forceTableScan" description:"force a table scan"`
+	Query     	string `long:"query" short:"q" description:"query filter, as a JSON string, e.g., '{x:{$gt:1}}'"`
+	QueryFile	string `long:"queryFile" description:"query filter, as a JSON file"`
+	TableScan 	bool   `long:"forceTableScan" description:"force a table scan"`
 }
 
 // Name returns a human-readable group name for input options.
 func (_ *InputOptions) Name() string {
 	return "query"
 }
 
+func (inputOptions *InputOptions) HasQuery() (bool){
+	return inputOptions.Query != "" || inputOptions.QueryFile != ""
+}
+
+func (inputOptions *InputOptions) GetQuery() ([]byte, error) {
+	if inputOptions.Query != "" {
+		return []byte(inputOptions.Query), nil
+	} else if inputOptions.QueryFile != "" {
+		content, err := ioutil.ReadFile(inputOptions.QueryFile)
+		if err != nil{
+			fmt.Errorf("error reading queryFile: %v", err)
+		}
+		return content, err
+	} else {
+		return nil, fmt.Errorf("GetQuery can return valid values only for query or queryFile input")
+	}
+}
+
 // OutputOptions defines the set of options for writing dump data.
 type OutputOptions struct {
 	Out                        string   `long:"out" short:"o" description:"output directory, or '-' for stdout (defaults to 'dump')" default-mask:"-"`

mongoexport/mongoexport.go

@@ -94,11 +94,19 @@ func (exp *MongoExport) ValidateSettings() error {
 		return fmt.Errorf("invalid output type '%v', choose 'json' or 'csv'", exp.OutputOpts.Type)
 	}
 
-	if exp.InputOpts != nil && exp.InputOpts.Query != "" {
-		_, err := getObjectFromArg(exp.InputOpts.Query)
+	if exp.InputOpts.Query != "" && exp.InputOpts.QueryFile != ""{
+		return fmt.Errorf("either query or queryFile can be specified as a query option")
+	}
+
+	if exp.InputOpts != nil && exp.InputOpts.HasQuery() {
+		content, err := exp.InputOpts.GetQuery()
 		if err != nil {
 			return err
 		}
+		_, err2 := getObjectFromByteArg(content)
+		if err2 != nil {
+			return err2
+		}
 	}
 
 	if exp.InputOpts != nil && exp.InputOpts.Sort != "" {
@@ -192,7 +200,6 @@ func (exp *MongoExport) getCount() (c int, err error) {
 // to export, based on the options given to mongoexport. Also returns the
 // associated session, so that it can be closed once the cursor is used up.
 func (exp *MongoExport) getCursor() (*mgo.Iter, *mgo.Session, error) {
-
 	sortFields := []string{}
 	if exp.InputOpts != nil && exp.InputOpts.Sort != "" {
 		sortD, err := getSortFromArg(exp.InputOpts.Sort)
@@ -206,9 +213,13 @@ func (exp *MongoExport) getCursor() (*mgo.Iter, *mgo.Session, error) {
 	}
 
 	query := map[string]interface{}{}
-	if exp.InputOpts != nil && exp.InputOpts.Query != "" {
+	if exp.InputOpts != nil && exp.InputOpts.HasQuery() {
 		var err error
-		query, err = getObjectFromArg(exp.InputOpts.Query)
+		content, err := exp.InputOpts.GetQuery()
+		if err != nil {
+			return nil, nil, err
+		}
+		query, err = getObjectFromByteArg(content)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -368,12 +379,12 @@ func (exp *MongoExport) getExportOutput(out io.Writer) (ExportOutput, error) {
 	return NewJSONExportOutput(exp.OutputOpts.JSONArray, exp.OutputOpts.Pretty, out), nil
 }
 
-// getObjectFromArg takes an object in extended JSON, and converts it to an object that
+// getObjectFromByteArg takes an object in extended JSON, and converts it to an object that
 // can be passed straight to db.collection.find(...) as a query or sort critera.
 // Returns an error if the string is not valid JSON, or extended JSON.
-func getObjectFromArg(queryRaw string) (map[string]interface{}, error) {
+func getObjectFromByteArg(queryRaw []byte) (map[string]interface{}, error) {
 	parsedJSON := map[string]interface{}{}
-	err := json.Unmarshal([]byte(queryRaw), &parsedJSON)
+	err := json.Unmarshal(queryRaw, &parsedJSON)
 	if err != nil {
 		return nil, fmt.Errorf("query '%v' is not valid JSON: %v", queryRaw, err)
 	}

mongoexport/options.go

@@ -1,4 +1,8 @@
 package mongoexport
+import (
+	"io/ioutil"
+	"fmt"
+)
 
 var Usage = `<options>
 
@@ -35,6 +39,7 @@ func (*OutputFormatOptions) Name() string {
 // InputOptions defines the set of options to use in retrieving data from the server.
 type InputOptions struct {
 	Query          string `long:"query" short:"q" description:"query filter, as a JSON string, e.g., '{x:{$gt:1}}'"`
+	QueryFile      string `long:"queryFile" description:"query filter, as a JSON file"`
 	SlaveOk        bool   `long:"slaveOk" short:"k" description:"allow secondary reads if available (default true)" default:"true" default-mask:"-"`
 	ForceTableScan bool   `long:"forceTableScan" description:"force a table scan (do not use $snapshot)"`
 	Skip           int    `long:"skip" description:"number of documents to skip"`
@@ -46,3 +51,21 @@ type InputOptions struct {
 func (*InputOptions) Name() string {
 	return "querying"
 }
+
+func (inputOptions *InputOptions) HasQuery() (bool) {
+	return inputOptions.Query != "" || inputOptions.QueryFile != ""
+}
+
+func (inputOptions *InputOptions) GetQuery() ([]byte, error) {
+	if inputOptions.Query != "" {
+		return []byte(inputOptions.Query), nil
+	} else if inputOptions.QueryFile != "" {
+		content, err := ioutil.ReadFile(inputOptions.QueryFile)
+		if err != nil{
+			fmt.Errorf("error reading queryFile: %v", err)
+		}
+		return content, err
+	} else {
+		return nil, fmt.Errorf("GetQuery can return valid values only for query or queryFile input")
+	}
+}