func main() { // initialize command-line opts usageStr := " --host myhost --db my_cms --collection docs < mydocfile." + "json \n\nImport CSV, TSV or JSON data into MongoDB.\n\nWhen importing " + "JSON documents, each document must be a separate line of the input file." opts := commonopts.New("mongoimport", "0.0.1", usageStr) inputOpts := &options.InputOptions{} opts.AddOptions(inputOpts) ingestOpts := &options.IngestOptions{} opts.AddOptions(ingestOpts) _, err := opts.Parse() if err != nil { util.Panicf("error parsing command line options: %v", err) } // print help, if specified if opts.PrintHelp() { return } // print version, if specified if opts.PrintVersion() { return } // create a session provider to connect to the db sessionProvider, err := db.InitSessionProvider(*opts) if err != nil { util.Panicf("error initializing database session: %v", err) } importer := mongoimport.MongoImport{ ToolOptions: opts, InputOptions: inputOpts, IngestOptions: ingestOpts, SessionProvider: sessionProvider, } if err = importer.ValidateSettings(); err != nil { util.PrintfTimeStamped("Error validating settings: %v\n", err) os.Exit(1) } numDocs, err := importer.ImportDocuments() if !opts.Quiet { message := fmt.Sprintf("imported 1 object\n") if numDocs != 1 { message = fmt.Sprintf("imported %v objects\n", numDocs) } util.PrintfTimeStamped(message) } if err != nil { fmt.Fprintf(os.Stderr, "Error importing documents: %v\n", err) os.Exit(1) } }
// ValidateSettings ensures that the tool specific options supplied for // MongoImport are valid func (mongoImport *MongoImport) ValidateSettings() error { // Namespace must have a valid database if mongoImport.ToolOptions.Namespace.DB == "" { return fmt.Errorf("must specify a database") } // use JSON as default input type if mongoImport.InputOptions.Type == "" { mongoImport.InputOptions.Type = JSON } else { if !(mongoImport.InputOptions.Type == TSV || mongoImport.InputOptions.Type == JSON || mongoImport.InputOptions.Type == CSV) { return fmt.Errorf("don't know what type [\"%v\"] is", mongoImport.InputOptions.Type) } } // ensure headers are supplied for CSV/TSV if mongoImport.InputOptions.Type == CSV || mongoImport.InputOptions.Type == TSV { if !mongoImport.InputOptions.HeaderLine { if mongoImport.InputOptions.Fields == "" && mongoImport.InputOptions.FieldFile == "" { return fmt.Errorf("You need to specify fields or have a " + "header line to import this file type") } } } // ensure we have a valid string to use for the collection if mongoImport.ToolOptions.Namespace.Collection == "" { if mongoImport.InputOptions.File == "" { return fmt.Errorf("must specify a collection or filename") } fileBaseName := filepath.Base(mongoImport.InputOptions.File) lastDotIndex := strings.LastIndex(fileBaseName, ".") if lastDotIndex != -1 { fileBaseName = fileBaseName[0:lastDotIndex] } mongoImport.ToolOptions.Namespace.Collection = fileBaseName util.PrintlnTimeStamped("no collection specified!") util.PrintfTimeStamped("using filename '%v' as collection\n", fileBaseName) } return nil }
// importDocuments is a helper to ImportDocuments and does all the ingestion // work by taking data from the 'importInput' source and writing it to the // appropriate namespace func (mongoImport *MongoImport) importDocuments(importInput ImportInput) ( int64, error) { session, err := mongoImport.SessionProvider.GetSession() if err != nil { return 0, err } defer session.Close() connUrl := mongoImport.ToolOptions.Host if mongoImport.ToolOptions.Port != "" { connUrl = connUrl + ":" + mongoImport.ToolOptions.Port } fmt.Fprintf(os.Stdout, "connected to: %v\n", connUrl) collection := session.DB(mongoImport.ToolOptions.DB). C(mongoImport.ToolOptions.Collection) // drop the database if necessary if mongoImport.IngestOptions.Drop { util.PrintfTimeStamped("dropping: %v.%v\n", mongoImport.ToolOptions.DB, mongoImport.ToolOptions.Collection) if err := collection.DropCollection(); err != nil { // this is hacky but necessary :( if err.Error() != errNsNotFound.Error() { return 0, err } } } // trim upsert fields if supplied var upsertFields []string if mongoImport.IngestOptions.Upsert { if len(mongoImport.IngestOptions.UpsertFields) != 0 { upsertFields = strings.Split(mongoImport.IngestOptions.UpsertFields, ",") } } docsCount := int64(0) for { document, err := importInput.ImportDocument() if err != nil { if err == io.EOF { return docsCount, nil } if mongoImport.IngestOptions.StopOnError { return docsCount, err } if document == nil { return docsCount, err } continue } // ignore blank fields if specified if mongoImport.IngestOptions.IgnoreBlanks && mongoImport.InputOptions.Type != JSON { document = removeBlankFields(document) } // if upsert is specified without any fields, default to inserts if mongoImport.IngestOptions.Upsert { selector := constructUpsertDocument(upsertFields, document) if selector == nil { err = collection.Insert(document) } else { _, err = collection.Upsert(selector, document) } } else { err = collection.Insert(document) } if err != nil { if mongoImport.IngestOptions.StopOnError { return docsCount, err } fmt.Fprintf(os.Stderr, "error inserting document: %v\n", err) continue } docsCount++ } return docsCount, nil }