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
}
