Пример #1
0
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)
	}
}
Пример #2
0
// 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
}
Пример #3
0
// 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
}