// CollectionExists returns true if the given intent's collection exists.
func (restore *MongoRestore) CollectionExists(intent *intents.Intent) (bool, error) {
restore.knownCollectionsMutex.Lock()
defer restore.knownCollectionsMutex.Unlock()
// make sure the map exists
if restore.knownCollections == nil {
restore.knownCollections = map[string][]string{}
}
// first check if we haven't done listCollections for this database already
if restore.knownCollections[intent.DB] == nil {
// if the database name isn't in the cache, grab collection
// names from the server
session, err := restore.SessionProvider.GetSession()
if err != nil {
return false, fmt.Errorf("error establishing connection: %v", err)
}
defer session.Close()
collections, err := session.DB(intent.DB).CollectionNames()
if err != nil {
return false, err
}
// update the cache
restore.knownCollections[intent.DB] = collections
}
// now check the cache for the given collection name
exists := util.StringSliceContains(restore.knownCollections[intent.DB], intent.C)
return exists, nil
}
// tokensToBSON reads in slice of records - along with ordered column names -
// and returns a BSON document for the record.
func tokensToBSON(colSpecs []ColumnSpec, tokens []string, numProcessed uint64, ignoreBlanks bool) (bson.D, error) {
log.Logvf(log.DebugHigh, "got line: %v", tokens)
var parsedValue interface{}
document := bson.D{}
for index, token := range tokens {
if token == "" && ignoreBlanks {
continue
}
if index < len(colSpecs) {
parsedValue, err := colSpecs[index].Parser.Parse(token)
if err != nil {
log.Logvf(log.DebugHigh, "parse failure in document #%d for column '%s',"+
"could not parse token '%s' to type %s",
numProcessed, colSpecs[index].Name, token, colSpecs[index].TypeName)
switch colSpecs[index].ParseGrace {
case pgAutoCast:
parsedValue = autoParse(token)
case pgSkipField:
continue
case pgSkipRow:
log.Logvf(log.Always, "skipping row #%d: %v", numProcessed, tokens)
return nil, coercionError{}
case pgStop:
return nil, fmt.Errorf("type coercion failure in document #%d for column '%s', "+
"could not parse token '%s' to type %s",
numProcessed, colSpecs[index].Name, token, colSpecs[index].TypeName)
}
}
if strings.Index(colSpecs[index].Name, ".") != -1 {
setNestedValue(colSpecs[index].Name, parsedValue, &document)
} else {
document = append(document, bson.DocElem{Name: colSpecs[index].Name, Value: parsedValue})
}
} else {
parsedValue = autoParse(token)
key := "field" + strconv.Itoa(index)
if util.StringSliceContains(ColumnNames(colSpecs), key) {
return nil, fmt.Errorf("duplicate field name - on %v - for token #%v ('%v') in document #%v",
key, index+1, parsedValue, numProcessed)
}
document = append(document, bson.DocElem{Name: key, Value: parsedValue})
}
}
return document, nil
}
// tokensToBSON reads in slice of records - along with ordered fields names -
// and returns a BSON document for the record.
func tokensToBSON(fields, tokens []string, numProcessed uint64) (bson.D, error) {
log.Logf(log.DebugHigh, "got line: %v", tokens)
var parsedValue interface{}
document := bson.D{}
for index, token := range tokens {
parsedValue = getParsedValue(token)
if index < len(fields) {
if strings.Index(fields[index], ".") != -1 {
setNestedValue(fields[index], parsedValue, &document)
} else {
document = append(document, bson.DocElem{fields[index], parsedValue})
}
} else {
key := "field" + strconv.Itoa(index)
if util.StringSliceContains(fields, key) {
return nil, fmt.Errorf("duplicate field name - on %v - for token #%v ('%v') in document #%v",
key, index+1, parsedValue, numProcessed)
}
document = append(document, bson.DocElem{key, parsedValue})
}
}
return document, nil
}
// getCursor returns a cursor that can be iterated over to get all the documents
// to export, based on the options given to mongoexport. Also returns the
// associated session, so that it can be closed once the cursor is used up.
func (exp *MongoExport) getCursor() (*mgo.Iter, *mgo.Session, error) {
sortFields := []string{}
if exp.InputOpts != nil && exp.InputOpts.Sort != "" {
sortD, err := getSortFromArg(exp.InputOpts.Sort)
if err != nil {
return nil, nil, err
}
sortFields, err = bsonutil.MakeSortString(sortD)
if err != nil {
return nil, nil, err
}
}
query := map[string]interface{}{}
if exp.InputOpts != nil && exp.InputOpts.HasQuery() {
var err error
content, err := exp.InputOpts.GetQuery()
if err != nil {
return nil, nil, err
}
query, err = getObjectFromByteArg(content)
if err != nil {
return nil, nil, err
}
}
flags := 0
if len(query) == 0 && exp.InputOpts != nil &&
exp.InputOpts.ForceTableScan != true && exp.InputOpts.Sort == "" {
flags = flags | db.Snapshot
}
session, err := exp.SessionProvider.GetSession()
if err != nil {
return nil, nil, err
}
skip := 0
if exp.InputOpts != nil {
skip = exp.InputOpts.Skip
}
limit := 0
if exp.InputOpts != nil {
limit = exp.InputOpts.Limit
}
if exp.InputOpts.AssertExists {
collNames, err := session.DB(exp.ToolOptions.Namespace.DB).CollectionNames()
if err != nil {
return nil, session, err
}
if !util.StringSliceContains(collNames, exp.ToolOptions.Namespace.Collection) {
return nil, session, fmt.Errorf("collection '%s' does not exist",
exp.ToolOptions.Namespace.Collection)
}
}
// build the query
q := session.DB(exp.ToolOptions.Namespace.DB).
C(exp.ToolOptions.Namespace.Collection).Find(query).Sort(sortFields...).
Skip(skip).Limit(limit)
if len(exp.OutputOpts.Fields) > 0 {
q.Select(makeFieldSelector(exp.OutputOpts.Fields))
}
q = db.ApplyFlags(q, session, flags)
return q.Iter(), session, nil
}