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