Golang util.ToUniversalPath函數代碼示例

// GetBSONReader opens and returns an io.ReadCloser for the BSONFileName in BSONDumpOptions
// or nil if none is set. The caller is responsible for closing it.
func (bdo *BSONDumpOptions) GetBSONReader() (io.ReadCloser, error) {
	if bdo.BSONFileName != "" {
		file, err := os.Open(util.ToUniversalPath(bdo.BSONFileName))
		if err != nil {
			return nil, fmt.Errorf("couldn't open BSON file: %v", err)
		}
		return file, nil
	}
	return ReadNopCloser{os.Stdin}, nil
}

func testQuery(md *MongoDump, session *mgo.Session) string {
	origDB := session.DB(testDB)
	restoredDB := session.DB(testRestoreDB)

	// query to test --query* flags
	bsonQuery := bson.M{"age": bson.M{"$lt": 10}}

	// 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)
	}
	return dumpDir
}

// GetWriter opens and returns an io.WriteCloser for the OutFileName in BSONDumpOptions
// or nil if none is set. The caller is responsible for closing it.
func (bdo *BSONDumpOptions) GetWriter() (io.WriteCloser, error) {
	if bdo.OutFileName != "" {
		file, err := os.Create(util.ToUniversalPath(bdo.OutFileName))
		if err != nil {
			return nil, err
		}
		return file, nil
	}

	return WriteNopCloser{os.Stdout}, nil
}

func TestMongoDumpKerberos(t *testing.T) {
	testutil.VerifyTestType(t, testutil.KerberosTestType)

	Convey("Should be able to run mongodump with Kerberos auth", t, func() {
		opts, err := testutil.GetKerberosOptions()

		So(err, ShouldBeNil)

		mongoDump := MongoDump{
			ToolOptions:  opts,
			InputOptions: &InputOptions{},
			OutputOptions: &OutputOptions{
				NumParallelCollections: 1,
			},
		}

		mongoDump.OutputOptions.Out = KerberosDumpDirectory

		err = mongoDump.Init()
		So(err, ShouldBeNil)
		err = mongoDump.Dump()
		So(err, ShouldBeNil)
		path, err := os.Getwd()
		So(err, ShouldBeNil)

		dumpDir := util.ToUniversalPath(filepath.Join(path, KerberosDumpDirectory))
		dumpDBDir := util.ToUniversalPath(filepath.Join(dumpDir, opts.Namespace.DB))
		So(fileDirExists(dumpDir), ShouldBeTrue)
		So(fileDirExists(dumpDBDir), ShouldBeTrue)

		dumpCollectionFile := util.ToUniversalPath(filepath.Join(dumpDBDir, opts.Namespace.Collection+".bson"))
		So(fileDirExists(dumpCollectionFile), ShouldBeTrue)

		countColls, err := countNonIndexBSONFiles(dumpDBDir)
		So(err, ShouldBeNil)
		So(countColls, ShouldEqual, 1)
	})
}

// GetOutputWriter opens and returns an io.WriteCloser for the output
// options or nil if none is set. The caller is responsible for closing it.
func (exp *MongoExport) GetOutputWriter() (io.WriteCloser, error) {
	if exp.OutputOpts.OutputFile != "" {
		// If the directory in which the output file is to be
		// written does not exist, create it
		fileDir := filepath.Dir(exp.OutputOpts.OutputFile)
		err := os.MkdirAll(fileDir, 0750)
		if err != nil {
			return nil, err
		}

		file, err := os.Create(util.ToUniversalPath(exp.OutputOpts.OutputFile))
		if err != nil {
			return nil, err
		}
		return file, err
	}
	// No writer, so caller should assume Stdout (or some other reasonable default)
	return nil, nil
}

// getSourceReader returns an io.Reader to read from the input source. Also
// returns a progress.Progressor which can be used to track progress if the
// reader supports it.
func (imp *MongoImport) getSourceReader() (io.ReadCloser, int64, error) {
	if imp.InputOptions.File != "" {
		file, err := os.Open(util.ToUniversalPath(imp.InputOptions.File))
		if err != nil {
			return nil, -1, err
		}
		fileStat, err := file.Stat()
		if err != nil {
			return nil, -1, err
		}
		log.Logvf(log.Info, "filesize: %v bytes", fileStat.Size())
		return file, int64(fileStat.Size()), err
	}

	log.Logvf(log.Info, "reading from stdin")

	// Stdin has undefined max size, so return 0
	return os.Stdin, 0, nil
}

func TestCreateIntentsForCollection(t *testing.T) {
	var mr *MongoRestore
	var buff bytes.Buffer

	testutil.VerifyTestType(t, testutil.UnitTestType)

	Convey("With a test MongoRestore", t, func() {
		buff = bytes.Buffer{}
		mr = &MongoRestore{
			manager:      intents.NewIntentManager(),
			ToolOptions:  &commonOpts.ToolOptions{Namespace: &commonOpts.Namespace{}},
			InputOptions: &InputOptions{},
		}
		log.SetWriter(&buff)

		Convey("running CreateIntentForCollection on a file without metadata", func() {
			ddl, err := newActualPath(util.ToUniversalPath("testdata/testdirs/db1/c2.bson"))
			So(err, ShouldBeNil)
			err = mr.CreateIntentForCollection("myDB", "myC", ddl)
			So(err, ShouldBeNil)
			mr.manager.Finalize(intents.Legacy)

			Convey("should create one intent with 'myDb' and 'myC' fields", func() {
				i0 := mr.manager.Pop()
				So(i0, ShouldNotBeNil)
				So(i0.DB, ShouldEqual, "myDB")
				So(i0.C, ShouldEqual, "myC")
				ddl, err := newActualPath(util.ToUniversalPath("testdata/testdirs/db1/c2.bson"))
				So(err, ShouldBeNil)
				So(i0.Location, ShouldEqual, ddl.Path())
				i1 := mr.manager.Pop()
				So(i1, ShouldBeNil)

				Convey("and no Metadata path", func() {
					So(i0.MetadataLocation, ShouldEqual, "")
					logs := buff.String()
					So(strings.Contains(logs, "without metadata"), ShouldEqual, true)
				})
			})
		})

		Convey("running CreateIntentForCollection on a file *with* metadata", func() {
			ddl, err := newActualPath(util.ToUniversalPath("testdata/testdirs/db1/c1.bson"))
			So(err, ShouldBeNil)
			err = mr.CreateIntentForCollection("myDB", "myC", ddl)
			So(err, ShouldBeNil)
			mr.manager.Finalize(intents.Legacy)

			Convey("should create one intent with 'myDb' and 'myC' fields", func() {
				i0 := mr.manager.Pop()
				So(i0, ShouldNotBeNil)
				So(i0.DB, ShouldEqual, "myDB")
				So(i0.C, ShouldEqual, "myC")
				So(i0.Location, ShouldEqual, util.ToUniversalPath("testdata/testdirs/db1/c1.bson"))
				i1 := mr.manager.Pop()
				So(i1, ShouldBeNil)

				Convey("and a set Metadata path", func() {
					So(i0.MetadataLocation, ShouldEqual, util.ToUniversalPath("testdata/testdirs/db1/c1.metadata.json"))
					logs := buff.String()
					So(strings.Contains(logs, "found metadata"), ShouldEqual, true)
				})
			})
		})

		Convey("running CreateIntentForCollection on a non-existent file", func() {
			_, err := newActualPath("aaaaaaaaaaaaaa.bson")
			Convey("should fail", func() {
				So(err, ShouldNotBeNil)
			})
		})

		Convey("running CreateIntentForCollection on a directory", func() {
			ddl, err := newActualPath("testdata")
			So(err, ShouldBeNil)
			err = mr.CreateIntentForCollection(
				"myDB", "myC", ddl)

			Convey("should fail", func() {
				So(err, ShouldNotBeNil)
			})
		})

		Convey("running CreateIntentForCollection on non-bson file", func() {
			ddl, err := newActualPath("testdata/testdirs/db1/c1.metadata.json")
			So(err, ShouldBeNil)
			err = mr.CreateIntentForCollection(
				"myDB", "myC", ddl)

			Convey("should fail", func() {
				So(err, ShouldNotBeNil)
			})
		})

	})
}

func TestMongoDumpMetaData(t *testing.T) {
	testutil.VerifyTestType(t, testutil.IntegrationTestType)
	log.SetWriter(ioutil.Discard)

	Convey("With a MongoDump instance", t, func() {
		err := setUpMongoDumpTestData()
		So(err, ShouldBeNil)

		Convey("testing that the dumped directory contains information about indexes", func() {
			md := simpleMongoDumpInstance()
			md.OutputOptions.Out = "dump"
			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)

			Convey("having one metadata file per collection", func() {
				c1, err := countNonIndexBSONFiles(dumpDBDir)
				So(err, ShouldBeNil)

				c2, err := countMetaDataFiles(dumpDBDir)
				So(err, ShouldBeNil)

				So(c1, ShouldEqual, c2)

				Convey("and that the JSON in a metadata file is valid", func() {
					metaFiles, err := getMatchingFiles(dumpDBDir, ".*\\.metadata\\.json")
					So(err, ShouldBeNil)
					So(len(metaFiles), ShouldBeGreaterThan, 0)

					oneMetaFile, err := os.Open(util.ToUniversalPath(filepath.Join(dumpDBDir, metaFiles[0])))
					So(err, ShouldBeNil)
					contents, err := ioutil.ReadAll(oneMetaFile)
					var jsonResult map[string]interface{}
					err = json.Unmarshal(contents, &jsonResult)
					So(err, ShouldBeNil)

					Convey("and contains an 'indexes' key", func() {
						_, ok := jsonResult["indexes"]
						So(ok, ShouldBeTrue)
						So(oneMetaFile.Close(), ShouldBeNil)
					})

				})

			})

			Reset(func() {
				So(os.RemoveAll(dumpDir), ShouldBeNil)
			})
		})

		Reset(func() {
			So(tearDownMongoDumpTestData(), ShouldBeNil)
		})

	})
}

func TestMongoDumpBSON(t *testing.T) {
	testutil.VerifyTestType(t, testutil.IntegrationTestType)
	log.SetWriter(ioutil.Discard)

	Convey("With a MongoDump instance", t, func() {
		err := setUpMongoDumpTestData()
		So(err, ShouldBeNil)

		Convey("testing that using MongoDump WITHOUT giving a query dumps everything in the database and/or collection", func() {
			md := simpleMongoDumpInstance()
			md.InputOptions.Query = ""

			Convey("and that for a particular collection", func() {
				md.ToolOptions.Namespace.Collection = testCollectionNames[0]
				err = md.Init()
				So(err, ShouldBeNil)

				Convey("it dumps to the default output directory", func() {
					// we don't have to set this manually if parsing options via command line
					md.OutputOptions.Out = "dump"
					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)

					session, err := getBareSession()
					So(err, ShouldBeNil)

					countColls, err := countNonIndexBSONFiles(dumpDBDir)
					So(err, ShouldBeNil)
					So(countColls, ShouldEqual, 1)

					collOriginal := session.DB(testDB).C(testCollectionNames[0])
					collRestore := session.DB(testRestoreDB).C(testCollectionNames[0])

					Convey("with the correct number of documents", func() {
						numDocsOrig, err := collOriginal.Count()
						So(err, ShouldBeNil)

						numDocsRestore, err := collRestore.Count()
						So(err, ShouldBeNil)

						So(numDocsOrig, ShouldEqual, numDocsRestore)
					})

					Convey("that are the same as the documents in the test database", func() {
						iter := collOriginal.Find(nil).Iter()

						var result bson.M
						for iter.Next(&result) {
							restoredCount, err := collRestore.Find(result).Count()
							So(err, ShouldBeNil)
							So(restoredCount, ShouldNotEqual, 0)
						}
						So(iter.Close(), ShouldBeNil)
					})

					Reset(func() {
						So(session.DB(testRestoreDB).DropDatabase(), ShouldBeNil)
						So(os.RemoveAll(dumpDir), ShouldBeNil)
					})
				})

				Convey("it dumps to a user-specified output directory", func() {
					md.OutputOptions.Out = "dump_user"
					err = md.Dump()
					So(err, ShouldBeNil)
					path, err := os.Getwd()
					So(err, ShouldBeNil)

					dumpDir := util.ToUniversalPath(filepath.Join(path, "dump_user"))
					dumpDBDir := util.ToUniversalPath(filepath.Join(dumpDir, testDB))
					So(fileDirExists(dumpDir), ShouldBeTrue)
					So(fileDirExists(dumpDBDir), ShouldBeTrue)

					countColls, err := countNonIndexBSONFiles(dumpDBDir)
					So(err, ShouldBeNil)
					So(countColls, ShouldEqual, 1)

					Reset(func() {
						So(os.RemoveAll(dumpDir), ShouldBeNil)
					})

				})

				Convey("it dumps to standard output", func() {
					md.OutputOptions.Out = "-"
					stdoutBuf := &bytes.Buffer{}
					md.stdout = stdoutBuf
					err = md.Dump()
					So(err, ShouldBeNil)
					var count int
//.........這裏部分代碼省略.........

//.........這裏部分代碼省略.........
				str, err := mf.Run(false)
				So(err, ShouldBeNil)
				So(len(str), ShouldNotEqual, 0)

				testFile, err := os.Open("testfile1")
				So(err, ShouldBeNil)
				defer testFile.Close()

				// pretty small file; so read all
				testFile1Bytes, err := ioutil.ReadAll(testFile)
				So(err, ShouldBeNil)
				So(len(testFile1Bytes), ShouldEqual, bytesExpected[0])
			})

			Reset(func() {
				// remove 'testfile1' or 'testfile1copy'
				if fileExists("testfile1") {
					err = os.Remove("testfile1")
				}
				So(err, ShouldBeNil)
				if fileExists("testfile1copy") {
					err = os.Remove("testfile1copy")
				}
				So(err, ShouldBeNil)
			})
		})

		Convey("Testing the 'put' command by putting some lorem ipsum file with 287613 bytes should", func() {
			args := []string{"put", "lorem_ipsum_287613_bytes.txt"}

			mf, err := simpleMongoFilesInstance(args)
			So(err, ShouldBeNil)
			So(mf, ShouldNotBeNil)
			mf.StorageOptions.LocalFileName = util.ToUniversalPath("testdata/lorem_ipsum_287613_bytes.txt")

			Convey("insert the file by creating two chunks (ceil(287,613 / 255 * 1024)) in GridFS", func() {
				str, err := mf.Run(false)
				So(err, ShouldBeNil)
				So(len(str), ShouldNotEqual, 0)

				Convey("and should have exactly 287613 bytes", func() {
					args = []string{"list", ""}

					mfAfter, err := simpleMongoFilesInstance(args)
					So(err, ShouldBeNil)
					So(mf, ShouldNotBeNil)

					str, err = mfAfter.Run(false)
					So(err, ShouldBeNil)

					lines := cleanAndTokenizeTestOutput(str)
					filesGotten, _ := getFilesAndBytesFromLines(lines)
					So(len(lines), ShouldEqual, len(filesExpected)+1)
					So(filesGotten, ShouldContain, "lorem_ipsum_287613_bytes.txt")
				})

				Convey("and should have exactly the same content as the original file", func() {
					args = []string{"get", "lorem_ipsum_287613_bytes.txt"}
					So(err, ShouldBeNil)
					mfAfter, err := simpleMongoFilesInstance(args)
					So(err, ShouldBeNil)
					So(mf, ShouldNotBeNil)

					mfAfter.StorageOptions.LocalFileName = "lorem_ipsum_copy.txt"
					str, err = mfAfter.Run(false)
					So(err, ShouldBeNil)

func main() {
	// initialize command-line opts
	opts := options.New("mongorestore", mongorestore.Usage,
		options.EnabledOptions{Auth: true, Connection: true})
	nsOpts := &mongorestore.NSOptions{}
	opts.AddOptions(nsOpts)
	inputOpts := &mongorestore.InputOptions{}
	opts.AddOptions(inputOpts)
	outputOpts := &mongorestore.OutputOptions{}
	opts.AddOptions(outputOpts)

	extraArgs, err := opts.Parse()
	if err != nil {
		log.Logvf(log.Always, "error parsing command line options: %v", err)
		log.Logvf(log.Always, "try 'mongorestore --help' for more information")
		os.Exit(util.ExitBadOptions)
	}

	// Allow the db connector to fall back onto the current database when no
	// auth database is given; the standard -d/-c options go into nsOpts now
	opts.Namespace = &options.Namespace{DB: nsOpts.DB}

	// print help or version info, if specified
	if opts.PrintHelp(false) {
		return
	}

	if opts.PrintVersion() {
		return
	}

	log.SetVerbosity(opts.Verbosity)

	targetDir, err := getTargetDirFromArgs(extraArgs, inputOpts.Directory)
	if err != nil {
		log.Logvf(log.Always, "%v", err)
		log.Logvf(log.Always, "try 'mongorestore --help' for more information")
		os.Exit(util.ExitBadOptions)
	}
	targetDir = util.ToUniversalPath(targetDir)

	// connect directly, unless a replica set name is explicitly specified
	_, setName := util.ParseConnectionString(opts.Host)
	opts.Direct = (setName == "")
	opts.ReplicaSetName = setName

	provider, err := db.NewSessionProvider(*opts)
	if err != nil {
		log.Logvf(log.Always, "error connecting to host: %v", err)
		os.Exit(util.ExitError)
	}
	provider.SetBypassDocumentValidation(outputOpts.BypassDocumentValidation)

	// disable TCP timeouts for restore jobs
	provider.SetFlags(db.DisableSocketTimeout)
	restore := mongorestore.MongoRestore{
		ToolOptions:     opts,
		OutputOptions:   outputOpts,
		InputOptions:    inputOpts,
		NSOptions:       nsOpts,
		TargetDirectory: targetDir,
		SessionProvider: provider,
	}

	finishedChan := signals.HandleWithInterrupt(restore.HandleInterrupt)
	defer close(finishedChan)

	if err = restore.Restore(); err != nil {
		log.Logvf(log.Always, "Failed: %v", err)
		if err == util.ErrTerminated {
			os.Exit(util.ExitKill)
		}
		os.Exit(util.ExitError)
	}
}

func main() {
	// initialize command-line opts
	opts := options.New("mongorestore", mongorestore.Usage,
		options.EnabledOptions{Auth: true, Connection: true, Namespace: true})
	inputOpts := &mongorestore.InputOptions{}
	opts.AddOptions(inputOpts)
	outputOpts := &mongorestore.OutputOptions{}
	opts.AddOptions(outputOpts)

	extraArgs, err := opts.Parse()
	if err != nil {
		log.Logf(log.Always, "error parsing command line options: %v", err)
		log.Logf(log.Always, "try 'mongorestore --help' for more information")
		os.Exit(util.ExitBadOptions)
	}

	// print help or version info, if specified
	if opts.PrintHelp(false) {
		return
	}

	if opts.PrintVersion() {
		return
	}

	log.SetVerbosity(opts.Verbosity)

	targetDir, err := getTargetDirFromArgs(extraArgs, inputOpts.Directory)
	if err != nil {
		log.Logf(log.Always, "%v", err)
		log.Logf(log.Always, "try 'mongorestore --help' for more information")
		os.Exit(util.ExitBadOptions)
	}
	targetDir = util.ToUniversalPath(targetDir)

	// connect directly, unless a replica set name is explicitly specified
	_, setName := util.ParseConnectionString(opts.Host)
	opts.Direct = (setName == "")
	opts.ReplicaSetName = setName

	provider, err := db.NewSessionProvider(*opts)
	if err != nil {
		log.Logf(log.Always, "error connecting to host: %v", err)
		os.Exit(util.ExitError)
	}
	provider.SetBypassDocumentValidation(outputOpts.BypassDocumentValidation)

	// disable TCP timeouts for restore jobs
	provider.SetFlags(db.DisableSocketTimeout)
	restore := mongorestore.MongoRestore{
		ToolOptions:     opts,
		OutputOptions:   outputOpts,
		InputOptions:    inputOpts,
		TargetDirectory: targetDir,
		SessionProvider: provider,
	}

	if err = restore.Restore(); err != nil {
		log.Logf(log.Always, "Failed: %v", err)
		if err == util.ErrTerminated {
			os.Exit(util.ExitKill)
		}
		os.Exit(util.ExitError)
	}
}