Golang log.Writer函數代碼示例

// Init performs preliminary setup operations for MongoDump.
func (dump *MongoDump) Init() error {
	err := dump.ValidateOptions()
	if err != nil {
		return fmt.Errorf("bad option: %v", err)
	}
	if dump.stdout == nil {
		dump.stdout = os.Stdout
	}
	dump.sessionProvider, err = db.NewSessionProvider(*dump.ToolOptions)
	if err != nil {
		return fmt.Errorf("can't create session: %v", err)
	}

	// temporarily allow secondary reads for the isMongos check
	dump.sessionProvider.SetReadPreference(mgo.Nearest)
	dump.isMongos, err = dump.sessionProvider.IsMongos()
	if err != nil {
		return err
	}

	if dump.isMongos && dump.OutputOptions.Oplog {
		return fmt.Errorf("can't use --oplog option when dumping from a mongos")
	}

	var mode mgo.Mode
	if dump.ToolOptions.ReplicaSetName != "" || dump.isMongos {
		mode = mgo.Primary
	} else {
		mode = mgo.Nearest
	}
	var tags bson.D

	if dump.InputOptions.ReadPreference != "" {
		mode, tags, err = db.ParseReadPreference(dump.InputOptions.ReadPreference)
		if err != nil {
			return fmt.Errorf("error parsing --readPreference : %v", err)
		}
		if len(tags) > 0 {
			dump.sessionProvider.SetTags(tags)
		}
	}

	// warn if we are trying to dump from a secondary in a sharded cluster
	if dump.isMongos && mode != mgo.Primary {
		log.Logf(log.Always, db.WarningNonPrimaryMongosConnection)
	}

	dump.sessionProvider.SetReadPreference(mode)
	dump.sessionProvider.SetTags(tags)
	dump.sessionProvider.SetFlags(db.DisableSocketTimeout)

	// return a helpful error message for mongos --repair
	if dump.OutputOptions.Repair && dump.isMongos {
		return fmt.Errorf("--repair flag cannot be used on a mongos")
	}

	dump.manager = intents.NewIntentManager()
	dump.progressManager = progress.NewProgressBarManager(log.Writer(0), progressBarWaitTime)
	return nil
}

// ImportDocuments is used to write input data to the database. It returns the
// number of documents successfully imported to the appropriate namespace and
// any error encountered in doing this
func (imp *MongoImport) ImportDocuments() (uint64, error) {
	source, fileSize, err := imp.getSourceReader()
	if err != nil {
		return 0, err
	}
	defer source.Close()

	inputReader, err := imp.getInputReader(source)
	if err != nil {
		return 0, err
	}

	if imp.InputOptions.HeaderLine {
		if imp.InputOptions.ColumnsHaveTypes {
			err = inputReader.ReadAndValidateTypedHeader(ParsePG(imp.InputOptions.ParseGrace))
		} else {
			err = inputReader.ReadAndValidateHeader()
		}
		if err != nil {
			return 0, err
		}
	}

	bar := &progress.Bar{
		Name:      fmt.Sprintf("%v.%v", imp.ToolOptions.DB, imp.ToolOptions.Collection),
		Watching:  &fileSizeProgressor{fileSize, inputReader},
		Writer:    log.Writer(0),
		BarLength: progressBarLength,
		IsBytes:   true,
	}
	bar.Start()
	defer bar.Stop()
	return imp.importDocuments(inputReader)
}

func simpleMongoDumpInstance() *MongoDump {
	ssl := testutil.GetSSLOptions()
	auth := testutil.GetAuthOptions()
	namespace := &options.Namespace{
		DB: testDB,
	}
	connection := &options.Connection{
		Host: testServer,
		Port: testPort,
	}
	toolOptions := &options.ToolOptions{
		SSL:           &ssl,
		Namespace:     namespace,
		Connection:    connection,
		Auth:          &auth,
		HiddenOptions: &options.HiddenOptions{},
		Verbosity:     &options.Verbosity{},
	}
	outputOptions := &OutputOptions{}
	inputOptions := &InputOptions{}

	log.SetVerbosity(toolOptions.Verbosity)

	return &MongoDump{
		ToolOptions:     toolOptions,
		InputOptions:    inputOptions,
		OutputOptions:   outputOptions,
		ProgressManager: progress.NewProgressBarManager(log.Writer(0), progressBarWaitTime),
		HandleSignals:   HandleSignals,
	}
}

// Init performs preliminary setup operations for MongoDump.
func (dump *MongoDump) Init() error {
	err := dump.ValidateOptions()
	if err != nil {
		return fmt.Errorf("bad option: %v", err)
	}
	dump.sessionProvider, err = db.NewSessionProvider(*dump.ToolOptions)
	if err != nil {
		return fmt.Errorf("can't create session: %v", err)
	}

	// allow secondary reads for the isMongos check
	dump.sessionProvider.SetFlags(db.Monotonic)
	dump.isMongos, err = dump.sessionProvider.IsMongos()
	if err != nil {
		return err
	}

	// ensure we allow secondary reads on mongods and disable TCP timeouts
	flags := db.DisableSocketTimeout
	if dump.isMongos {
		log.Logf(log.Info, "connecting to mongos; secondary reads disabled")
	} else {
		flags |= db.Monotonic
	}
	dump.sessionProvider.SetFlags(flags)

	// return a helpful error message for mongos --repair
	if dump.OutputOptions.Repair && dump.isMongos {
		return fmt.Errorf("--repair flag cannot be used on a mongos")
	}

	dump.manager = intents.NewIntentManager()
	dump.progressManager = progress.NewProgressBarManager(log.Writer(0), progressBarWaitTime)
	return nil
}

// RestoreIntents iterates through all of the intents stored in the IntentManager, and restores them.
func (restore *MongoRestore) RestoreIntents() error {
	// start up the progress bar manager
	restore.progressManager = progress.NewProgressBarManager(log.Writer(0), progressBarWaitTime)
	restore.progressManager.Start()
	defer restore.progressManager.Stop()

	log.Logf(log.DebugLow, "restoring up to %v collections in parallel", restore.OutputOptions.NumParallelCollections)

	if restore.OutputOptions.NumParallelCollections > 0 {
		resultChan := make(chan error)

		// start a goroutine for each job thread
		for i := 0; i < restore.OutputOptions.NumParallelCollections; i++ {
			go func(id int) {
				log.Logf(log.DebugHigh, "starting restore routine with id=%v", id)
				for {
					intent := restore.manager.Pop()
					if intent == nil {
						log.Logf(log.DebugHigh, "ending restore routine with id=%v, no more work to do", id)
						resultChan <- nil // done
						return
					}
					err := restore.RestoreIntent(intent)
					if err != nil {
						resultChan <- fmt.Errorf("%v: %v", intent.Namespace(), err)
						return
					}
					restore.manager.Finish(intent)
				}
			}(i)
		}

		// wait until all goroutines are done or one of them errors out
		for i := 0; i < restore.OutputOptions.NumParallelCollections; i++ {
			if err := <-resultChan; err != nil {
				return err
			}
		}
		return nil
	}

	// single-threaded
	for {
		intent := restore.manager.Pop()
		if intent == nil {
			return nil
		}
		err := restore.RestoreIntent(intent)
		if err != nil {
			return fmt.Errorf("%v: %v", intent.Namespace(), err)
		}
		restore.manager.Finish(intent)
	}
	return nil
}

// Internal function that handles exporting to the given writer. Used primarily
// for testing, because it bypasses writing to the file system.
func (exp *MongoExport) exportInternal(out io.Writer) (int64, error) {

	max, err := exp.getCount()
	if err != nil {
		return 0, err
	}

	progressManager := progress.NewProgressBarManager(log.Writer(0), progressBarWaitTime)
	progressManager.Start()
	defer progressManager.Stop()

	watchProgressor := progress.NewCounter(int64(max))
	bar := &progress.Bar{
		Name:      fmt.Sprintf("%v.%v", exp.ToolOptions.Namespace.DB, exp.ToolOptions.Namespace.Collection),
		Watching:  watchProgressor,
		BarLength: progressBarLength,
	}
	progressManager.Attach(bar)
	defer progressManager.Detach(bar)

	exportOutput, err := exp.getExportOutput(out)
	if err != nil {
		return 0, err
	}

	cursor, session, err := exp.getCursor()
	if err != nil {
		return 0, err
	}
	defer session.Close()
	defer cursor.Close()

	connURL := exp.ToolOptions.Host
	if connURL == "" {
		connURL = util.DefaultHost
	}
	if exp.ToolOptions.Port != "" {
		connURL = connURL + ":" + exp.ToolOptions.Port
	}
	log.Logf(log.Always, "connected to: %v", connURL)

	// Write headers
	err = exportOutput.WriteHeader()
	if err != nil {
		return 0, err
	}

	var result bson.M

	docsCount := int64(0)

	// Write document content
	for cursor.Next(&result) {
		err := exportOutput.ExportDocument(result)
		if err != nil {
			return docsCount, err
		}
		docsCount++
		if docsCount%watchProgressorUpdateFrequency == 0 {
			watchProgressor.Set(docsCount)
		}
	}
	watchProgressor.Set(docsCount)
	if err := cursor.Err(); err != nil {
		return docsCount, err
	}

	// Write footers
	err = exportOutput.WriteFooter()
	if err != nil {
		return docsCount, err
	}
	exportOutput.Flush()
	return docsCount, nil
}

// RestoreOplog attempts to restore a MongoDB oplog.
func (restore *MongoRestore) RestoreOplog() error {
	log.Log(log.Always, "replaying oplog")
	intent := restore.manager.Oplog()
	if intent == nil {
		// this should not be reached
		log.Log(log.Always, "no oplog.bson file in root of the dump directory, skipping oplog application")
		return nil
	}
	if err := intent.BSONFile.Open(); err != nil {
		return err
	}
	defer intent.BSONFile.Close()
	bsonSource := db.NewDecodedBSONSource(db.NewBSONSource(intent.BSONFile))
	defer bsonSource.Close()

	entryArray := make([]interface{}, 0, 1024)
	rawOplogEntry := &bson.Raw{}

	var totalOps int64
	var entrySize, bufferedBytes int

	oplogProgressor := progress.NewCounter(intent.BSONSize)
	bar := progress.Bar{
		Name:      "oplog",
		Watching:  oplogProgressor,
		WaitTime:  3 * time.Second,
		Writer:    log.Writer(0),
		BarLength: progressBarLength,
		IsBytes:   true,
	}
	bar.Start()
	defer bar.Stop()

	session, err := restore.SessionProvider.GetSession()
	if err != nil {
		return fmt.Errorf("error establishing connection: %v", err)
	}
	defer session.Close()

	// To restore the oplog, we iterate over the oplog entries,
	// filling up a buffer. Once the buffer reaches max document size,
	// apply the current buffered ops and reset the buffer.
	for bsonSource.Next(rawOplogEntry) {
		entrySize = len(rawOplogEntry.Data)
		if bufferedBytes+entrySize > oplogMaxCommandSize {
			err = restore.ApplyOps(session, entryArray)
			if err != nil {
				return fmt.Errorf("error applying oplog: %v", err)
			}
			entryArray = make([]interface{}, 0, 1024)
			bufferedBytes = 0
		}

		entryAsOplog := db.Oplog{}
		err = bson.Unmarshal(rawOplogEntry.Data, &entryAsOplog)
		if err != nil {
			return fmt.Errorf("error reading oplog: %v", err)
		}
		if entryAsOplog.Operation == "n" {
			//skip no-ops
			continue
		}
		if !restore.TimestampBeforeLimit(entryAsOplog.Timestamp) {
			log.Logf(
				log.DebugLow,
				"timestamp %v is not below limit of %v; ending oplog restoration",
				entryAsOplog.Timestamp,
				restore.oplogLimit,
			)
			break
		}

		totalOps++
		bufferedBytes += entrySize
		oplogProgressor.Inc(int64(entrySize))
		entryArray = append(entryArray, entryAsOplog)
	}
	// finally, flush the remaining entries
	if len(entryArray) > 0 {
		err = restore.ApplyOps(session, entryArray)
		if err != nil {
			return fmt.Errorf("error applying oplog: %v", err)
		}
	}

	log.Logf(log.Info, "applied %v ops", totalOps)
	return nil

}

func main() {
	// initialize command-line opts
	opts := options.New("mongodump", mongodump.Usage, options.EnabledOptions{true, true, true})

	inputOpts := &mongodump.InputOptions{}
	opts.AddOptions(inputOpts)
	outputOpts := &mongodump.OutputOptions{}
	opts.AddOptions(outputOpts)

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

	if len(args) > 0 {
		log.Logf(log.Always, "positional arguments not allowed: %v", args)
		log.Logf(log.Always, "try 'mongodump --help' for more information")
		os.Exit(util.ExitBadOptions)
	}

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

	// print version, if specified
	if opts.PrintVersion() {
		return
	}

	// init logger
	log.SetVerbosity(opts.Verbosity)

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

	dump := mongodump.MongoDump{
		ToolOptions:     opts,
		OutputOptions:   outputOpts,
		InputOptions:    inputOpts,
		ProgressManager: progress.NewProgressBarManager(log.Writer(0), progressBarWaitTime),
		HandleSignals:   mongodump.HandleSignals,
	}

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

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

// dumpQueryToWriter takes an mgo Query and a writer, performs the query,
// and writes the raw bson results to the writer.
func (dump *MongoDump) dumpQueryToWriter(query *mgo.Query, writer io.Writer) error {
	var dumpCounter int
	total, err := query.Count()
	if err != nil {
		return fmt.Errorf("error reading from db: %v", err)
	}
	log.Logf(1, "\t%v documents", total)

	bar := progress.ProgressBar{
		Max:        total,
		CounterPtr: &dumpCounter,
		WaitTime:   3 * time.Second,
		Writer:     log.Writer(0),
		BarLength:  ProgressBarLength,
	}
	bar.Start()
	defer bar.Stop()

	cursor := query.Iter()
	defer cursor.Close()

	// We run the result iteration in its own goroutine,
	// this allows disk i/o to not block reads from the db,
	// which gives a slight speedup on benchmarks
	buffChan := make(chan []byte)
	go func() {
		for {
			raw := &bson.Raw{}
			if err := cursor.Err(); err != nil {
				log.Logf(0, "error reading from db: %v", err)
			}
			next := cursor.Next(raw)
			if !next {
				close(buffChan)
				return
			}
			buffChan <- raw.Data
		}
	}()

	// wrap writer in buffer to reduce load on disk
	// TODO extensive optimization on buffer size
	w := bufio.NewWriterSize(writer, 1024*32)

	// while there are still results in the database,
	// grab results from the goroutine and write them to filesystem
	for {
		buff, alive := <-buffChan
		if !alive {
			if cursor.Err() != nil {
				return fmt.Errorf("error reading collection: %v", cursor.Err())
			}
			break
		}
		_, err := w.Write(buff)
		if err != nil {
			return fmt.Errorf("error writing to file: %v", err)
		}
		dumpCounter++
	}
	err = w.Flush()
	if err != nil {
		return fmt.Errorf("error flushing file writer: %v", err)
	}
	return nil
}

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)
	defer provider.Close()
	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)

	// start up the progress bar manager
	progressManager := progress.NewBarWriter(log.Writer(0), progressBarWaitTime, progressBarLength, true)
	progressManager.Start()
	defer progressManager.Stop()

	restore := mongorestore.MongoRestore{
		ToolOptions:     opts,
		OutputOptions:   outputOpts,
		InputOptions:    inputOpts,
		NSOptions:       nsOpts,
		TargetDirectory: targetDir,
		SessionProvider: provider,
		ProgressManager: progressManager,
	}

	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)
	}
}