<commit_before>/**
 * (C) Copyright 2014, Deft Labs
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at:
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package dlshared

import (
	"time"
	"strings"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type MongoDataSource struct {
	DbName string
	CollectionName string
	Mongo *Mongo
	Logger
}

func ObjectIdHex(objectIdHex string) *bson.ObjectId {
	id := bson.ObjectIdHex(objectIdHex)
	return &id
}

func (self *MongoDataSource) NewObjectId() *bson.ObjectId {
	id := bson.NewObjectId()
	return &id
}

// Insert a document into a collection with the base configured write concern.
func (self *MongoDataSource) Insert(doc interface{}) error {
	if err := self.Mongo.Collection(self.DbName, self.CollectionName).Insert(doc); err != nil {
		return NewStackError("Unable to Insert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Upsert a document in a collection with the base configured write concern.
func (self *MongoDataSource) Upsert(selector interface{}, change interface{}) error {
	if _, err := self.Mongo.Collection(self.DbName, self.CollectionName).Upsert(selector, change); err != nil {
		return NewStackError("Unable to Upsert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Upsert a document into a collection with the passed write concern.
func (self *MongoDataSource) UpsertSafe(selector interface{}, change interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	if _, err := self.Mongo.Collection(self.DbName, self.CollectionName).Upsert(selector, change); err != nil {
		return NewStackError("Unable to Upsert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Updates a document into a collection with the passed write concern.
func (self *MongoDataSource) UpdateSafe(selector interface{}, change interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	if err := self.Mongo.Collection(self.DbName, self.CollectionName).Update(selector, change); err != nil {
		return NewStackError("Unable to Update - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Insert a document into a collection with the passed write concern.
func (self *MongoDataSource) InsertSafe(doc interface{}) error {
	session := self.SessionClone()

	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)
	if err := session.DB(self.DbName).C(self.CollectionName).Insert(doc); err != nil {
		return NewStackError("Unable to InsertSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Unset a field with the default safe enabled - uses $unset
func (self *MongoDataSource) UnsetFieldSafe(query interface{}, field string) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$unset": &bson.M{ field: nadaStr } }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to UnsetFieldSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Set fields using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) SetFieldsSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$set": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to SetFieldsSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Pull from an array call using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) PullSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$pull": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to PullSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Find the distinct string fields. Do not use this on datasets with a large amount of distinct values or
// you will blow out memory. The selector can be nil.
func (self *MongoDataSource) FindDistinctStrs(selector interface{}, fieldName string) ([]string, error) {
	var result []string
	if err := self.Mongo.Collection(self.DbName, self.CollectionName).Find(selector).Distinct(fieldName, &result); err != nil { return nil, err }

	return result, nil
}

// The caller must close the cursor when done. Use: defer cursor.Clse()
func (self *MongoDataSource) FindManyWithBatchSize(selector interface{}, batchSize int) *mgo.Iter {
	return self.Mongo.Collection(self.DbName, self.CollectionName).Find(selector).Batch(batchSize).Iter()
}

// Push to an array call using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) PushSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$push": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to PushSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Set a property using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) SetFieldSafe(query interface{}, field string, value interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$set": &bson.M{ field: value } }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to SetFieldSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Find by the _id. Returns false if not found.
func (self *MongoDataSource) FindById(id interface{}, result interface{}) error {
	return self.FindOne(&bson.M{ "_id": id }, result)
}

// Returns nil if this is a NOT a document not found error.
func (self *MongoDataSource) RemoveNotFoundErr(err error) error {
	if self.NotFoundErr(err) {
		return nil
	}
	return err
}

// Returns true if this is a document not found error.
func (self *MongoDataSource) NotFoundErr(err error) (bool) {
	return err == mgo.ErrNotFound
}

// Finds one document or returns false.
func (self *MongoDataSource) FindOne(query *bson.M, result interface{}) error {
	return self.Collection().Find(query).One(result)
}

// Delete one document from the collection. If the document is not found, no error is returned.
func (self *MongoDataSource) DeleteOne(selector interface{}) error {
	if err := self.RemoveNotFoundErr(self.Collection().Remove(selector)); err != nil {
		return NewStackError("Unable to DeleteOne - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Delete one or more documents from the collection. If the document(s) is/are not found, no error
// is returned.
func (self *MongoDataSource) Delete(selector interface{}) error {
	if _, err := self.Collection().RemoveAll(selector); err != nil {
		return NewStackError("Unable to Delete - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a unique, non-sparse index is created. This does not create in the background. This does
// NOT drop duplicates if they exist. Duplicates will cause an error.
func (self *MongoDataSource) EnsureUniqueIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: true,
		DropDups: false,
		Background: false,
		Sparse: false,
	}); err != nil {
		return NewStackError("Unable to EnsureUniqueIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a non-unique, non-sparse index is created. This does not create in the background.
func (self *MongoDataSource) EnsureIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: false,
	}); err != nil {
		return NewStackError("Unable to EnsureIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a non-unique, sparse index is created. This does not create in the background.
func (self *MongoDataSource) EnsureSparseIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: true,
	}); err != nil {
		return NewStackError("Unable to EnsureSparseIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Create a capped collection.
func (self *MongoDataSource) CreateCappedCollection(sizeInBytes int) error {
	if err := self.Collection().Create(&mgo.CollectionInfo{ DisableIdIndex: false, ForceIdIndex: true, Capped: true, MaxBytes: sizeInBytes }); err != nil {

		// This is a bit of a hack, but the error returned does not provide any codes.
		msg := strings.ToLower(err.Error())
		if strings.Contains(msg, "already") || strings.Contains(msg, "exists") { return nil }

		return NewStackError("Unable to CreateCappedCollection - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)

	} else {
		return nil
	}
}

// Create a ttl index.
func (self *MongoDataSource) EnsureTtlIndex(field string, expireAfterSeconds int) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: []string{ field },
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: false,
		ExpireAfter: time.Duration(expireAfterSeconds) * time.Second,
	}); err != nil {
		return NewStackError("Unable to EnsureTtlIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

func (self *MongoDataSource) Now() *time.Time {
	now := time.Now()
	return &now
}

// Ensure a unique, sparse index is created. This does not create in the background. This does
// NOT drop duplicates if they exist. Duplicates will cause an error.
func (self *MongoDataSource) EnsureUniqueSparseIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: true,
		DropDups: false,
		Background: false,
		Sparse: true,
	}); err != nil {
		return NewStackError("Unable to EnsureUniqueSparseIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Returns the collection from the session.
func (self *MongoDataSource) Collection() *mgo.Collection { return self.Mongo.Collection(self.DbName, self.CollectionName) }

// Returns the database from the session.
func (self *MongoDataSource) Db() *mgo.Database { return self.Mongo.Db(self.DbName) }

// Returns the session struct.
func (self *MongoDataSource) Session() *mgo.Session { return self.Mongo.session }

// Returns a clone of the session struct.
func (self *MongoDataSource) SessionClone() *mgo.Session { return self.Mongo.session.Clone() }

// Returns a copy of the session struct.
func (self *MongoDataSource) SessionCopy() *mgo.Session { return self.Mongo.session.Clone() }

<commit_msg>added delete by id conv method.<commit_after>/**
 * (C) Copyright 2014, Deft Labs
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at:
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package dlshared

import (
	"time"
	"strings"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type MongoDataSource struct {
	DbName string
	CollectionName string
	Mongo *Mongo
	Logger
}

func ObjectIdHex(objectIdHex string) *bson.ObjectId {
	id := bson.ObjectIdHex(objectIdHex)
	return &id
}

func (self *MongoDataSource) NewObjectId() *bson.ObjectId {
	id := bson.NewObjectId()
	return &id
}

// Insert a document into a collection with the base configured write concern.
func (self *MongoDataSource) Insert(doc interface{}) error {
	if err := self.Mongo.Collection(self.DbName, self.CollectionName).Insert(doc); err != nil {
		return NewStackError("Unable to Insert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Upsert a document in a collection with the base configured write concern.
func (self *MongoDataSource) Upsert(selector interface{}, change interface{}) error {
	if _, err := self.Mongo.Collection(self.DbName, self.CollectionName).Upsert(selector, change); err != nil {
		return NewStackError("Unable to Upsert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Upsert a document into a collection with the passed write concern.
func (self *MongoDataSource) UpsertSafe(selector interface{}, change interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	if _, err := self.Mongo.Collection(self.DbName, self.CollectionName).Upsert(selector, change); err != nil {
		return NewStackError("Unable to Upsert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Updates a document into a collection with the passed write concern.
func (self *MongoDataSource) UpdateSafe(selector interface{}, change interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	if err := self.Mongo.Collection(self.DbName, self.CollectionName).Update(selector, change); err != nil {
		return NewStackError("Unable to Update - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Insert a document into a collection with the passed write concern.
func (self *MongoDataSource) InsertSafe(doc interface{}) error {
	session := self.SessionClone()

	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)
	if err := session.DB(self.DbName).C(self.CollectionName).Insert(doc); err != nil {
		return NewStackError("Unable to InsertSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Unset a field with the default safe enabled - uses $unset
func (self *MongoDataSource) UnsetFieldSafe(query interface{}, field string) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$unset": &bson.M{ field: nadaStr } }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to UnsetFieldSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Set fields using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) SetFieldsSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$set": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to SetFieldsSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Pull from an array call using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) PullSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$pull": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to PullSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Find the distinct string fields. Do not use this on datasets with a large amount of distinct values or
// you will blow out memory. The selector can be nil.
func (self *MongoDataSource) FindDistinctStrs(selector interface{}, fieldName string) ([]string, error) {
	var result []string
	if err := self.Mongo.Collection(self.DbName, self.CollectionName).Find(selector).Distinct(fieldName, &result); err != nil { return nil, err }

	return result, nil
}

// The caller must close the cursor when done. Use: defer cursor.Clse()
func (self *MongoDataSource) FindManyWithBatchSize(selector interface{}, batchSize int) *mgo.Iter {
	return self.Mongo.Collection(self.DbName, self.CollectionName).Find(selector).Batch(batchSize).Iter()
}

// Push to an array call using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) PushSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$push": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to PushSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Set a property using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) SetFieldSafe(query interface{}, field string, value interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$set": &bson.M{ field: value } }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to SetFieldSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Find by the _id. Returns false if not found.
func (self *MongoDataSource) FindById(id interface{}, result interface{}) error {
	return self.FindOne(&bson.M{ "_id": id }, result)
}

// Returns nil if this is a NOT a document not found error.
func (self *MongoDataSource) RemoveNotFoundErr(err error) error {
	if self.NotFoundErr(err) {
		return nil
	}
	return err
}

// Returns true if this is a document not found error.
func (self *MongoDataSource) NotFoundErr(err error) (bool) {
	return err == mgo.ErrNotFound
}

// Finds one document or returns false.
func (self *MongoDataSource) FindOne(query *bson.M, result interface{}) error {
	return self.Collection().Find(query).One(result)
}

// Delete one document by the _id.
func (self *MongoDataSource) DeleteById(id interface{}) error { return self.DeleteOne(&bson.M{"_id": id }) }

// Delete one document from the collection. If the document is not found, no error is returned.
func (self *MongoDataSource) DeleteOne(selector interface{}) error {
	if err := self.RemoveNotFoundErr(self.Collection().Remove(selector)); err != nil {
		return NewStackError("Unable to DeleteOne - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Delete one or more documents from the collection. If the document(s) is/are not found, no error
// is returned.
func (self *MongoDataSource) Delete(selector interface{}) error {
	if _, err := self.Collection().RemoveAll(selector); err != nil {
		return NewStackError("Unable to Delete - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a unique, non-sparse index is created. This does not create in the background. This does
// NOT drop duplicates if they exist. Duplicates will cause an error.
func (self *MongoDataSource) EnsureUniqueIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: true,
		DropDups: false,
		Background: false,
		Sparse: false,
	}); err != nil {
		return NewStackError("Unable to EnsureUniqueIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a non-unique, non-sparse index is created. This does not create in the background.
func (self *MongoDataSource) EnsureIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: false,
	}); err != nil {
		return NewStackError("Unable to EnsureIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a non-unique, sparse index is created. This does not create in the background.
func (self *MongoDataSource) EnsureSparseIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: true,
	}); err != nil {
		return NewStackError("Unable to EnsureSparseIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Create a capped collection.
func (self *MongoDataSource) CreateCappedCollection(sizeInBytes int) error {
	if err := self.Collection().Create(&mgo.CollectionInfo{ DisableIdIndex: false, ForceIdIndex: true, Capped: true, MaxBytes: sizeInBytes }); err != nil {

		// This is a bit of a hack, but the error returned does not provide any codes.
		msg := strings.ToLower(err.Error())
		if strings.Contains(msg, "already") || strings.Contains(msg, "exists") { return nil }

		return NewStackError("Unable to CreateCappedCollection - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)

	} else {
		return nil
	}
}

// Create a ttl index.
func (self *MongoDataSource) EnsureTtlIndex(field string, expireAfterSeconds int) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: []string{ field },
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: false,
		ExpireAfter: time.Duration(expireAfterSeconds) * time.Second,
	}); err != nil {
		return NewStackError("Unable to EnsureTtlIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

func (self *MongoDataSource) Now() *time.Time {
	now := time.Now()
	return &now
}

// Ensure a unique, sparse index is created. This does not create in the background. This does
// NOT drop duplicates if they exist. Duplicates will cause an error.
func (self *MongoDataSource) EnsureUniqueSparseIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: true,
		DropDups: false,
		Background: false,
		Sparse: true,
	}); err != nil {
		return NewStackError("Unable to EnsureUniqueSparseIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Returns the collection from the session.
func (self *MongoDataSource) Collection() *mgo.Collection { return self.Mongo.Collection(self.DbName, self.CollectionName) }

// Returns the database from the session.
func (self *MongoDataSource) Db() *mgo.Database { return self.Mongo.Db(self.DbName) }

// Returns the session struct.
func (self *MongoDataSource) Session() *mgo.Session { return self.Mongo.session }

// Returns a clone of the session struct.
func (self *MongoDataSource) SessionClone() *mgo.Session { return self.Mongo.session.Clone() }

// Returns a copy of the session struct.
func (self *MongoDataSource) SessionCopy() *mgo.Session { return self.Mongo.session.Clone() }

<|endoftext|>


<commit_before>/**
 * (C) Copyright 2014, Deft Labs
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at:
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package dlshared

import (
	"time"
	"strings"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type MongoDataSource struct {
	DbName string
	CollectionName string
	Mongo *Mongo
	Logger
}

func ObjectIdHex(objectIdHex string) *bson.ObjectId {
	id := bson.ObjectIdHex(objectIdHex)
	return &id
}

func (self *MongoDataSource) NewObjectId() *bson.ObjectId {
	id := bson.NewObjectId()
	return &id
}

// Insert a document into a collection with the base configured write concern.
func (self *MongoDataSource) Insert(doc interface{}) error {
	if err := self.Mongo.Collection(self.DbName, self.CollectionName).Insert(doc); err != nil {
		return NewStackError("Unable to Insert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Upsert a document in a collection with the base configured write concern.
func (self *MongoDataSource) Upsert(selector interface{}, change interface{}) error {
	if _, err := self.Mongo.Collection(self.DbName, self.CollectionName).Upsert(selector, change); err != nil {
		return NewStackError("Unable to Upsert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Upsert a document into a collection with the passed write concern.
func (self *MongoDataSource) UpsertSafe(selector interface{}, change interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	if _, err := session.DB(self.DbName).C(self.CollectionName).Upsert(selector, change); err != nil {
		return NewStackError("Unable to Upsert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Updates a document into a collection with the passed write concern.
func (self *MongoDataSource) UpdateSafe(selector interface{}, change interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	if err := session.DB(self.DbName).C(self.CollectionName).Update(selector, change); err != nil {
		return NewStackError("Unable to Update - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Insert a document into a collection with the passed write concern.
func (self *MongoDataSource) InsertSafe(doc interface{}) error {
	session := self.SessionClone()

	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)
	if err := session.DB(self.DbName).C(self.CollectionName).Insert(doc); err != nil {
		return NewStackError("Unable to InsertSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Unset a field with the default safe enabled - uses $unset
func (self *MongoDataSource) UnsetFieldSafe(query interface{}, field string) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$unset": &bson.M{ field: nadaStr } }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to UnsetFieldSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Set fields using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) SetFieldsSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$set": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to SetFieldsSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Pull from an array call using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) PullSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$pull": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to PullSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Find the distinct string fields. Do not use this on datasets with a large amount of distinct values or
// you will blow out memory. The selector can be nil.
func (self *MongoDataSource) FindDistinctStrs(selector interface{}, fieldName string) ([]string, error) {
	var result []string
	if err := self.Mongo.Collection(self.DbName, self.CollectionName).Find(selector).Distinct(fieldName, &result); err != nil { return nil, err }

	return result, nil
}

// The caller must close the cursor when done. Use: defer cursor.Clse()
func (self *MongoDataSource) FindManyWithBatchSize(selector interface{}, batchSize int) *mgo.Iter {
	return self.Mongo.Collection(self.DbName, self.CollectionName).Find(selector).Batch(batchSize).Iter()
}

// Push to an array call using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) PushSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$push": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to PushSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Set a property using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) SetFieldSafe(query interface{}, field string, value interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$set": &bson.M{ field: value } }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to SetFieldSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Find by the _id. Returns false if not found.
func (self *MongoDataSource) FindById(id interface{}, result interface{}) error {
	return self.FindOne(&bson.M{ "_id": id }, result)
}

// Returns nil if this is a NOT a document not found error.
func (self *MongoDataSource) RemoveNotFoundErr(err error) error {
	if self.NotFoundErr(err) { return nil }
	return err
}

// Returns true if this is a document not found error.
func (self *MongoDataSource) NotFoundErr(err error) bool { return err == mgo.ErrNotFound }

// Finds one document or returns false.
func (self *MongoDataSource) FindOne(query *bson.M, result interface{}) error { return self.Collection().Find(query).One(result) }

// Delete one document by the _id.
func (self *MongoDataSource) DeleteById(id interface{}) error { return self.DeleteOne(&bson.M{"_id": id }) }

// Delete one document from the collection. If the document is not found, no error is returned.
func (self *MongoDataSource) DeleteOne(selector interface{}) error {
	if err := self.RemoveNotFoundErr(self.Collection().Remove(selector)); err != nil {
		return NewStackError("Unable to DeleteOne - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Delete one or more documents from the collection. If the document(s) is/are not found, no error
// is returned.
func (self *MongoDataSource) Delete(selector interface{}) error {
	if _, err := self.Collection().RemoveAll(selector); err != nil {
		return NewStackError("Unable to Delete - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a unique, non-sparse index is created. This does not create in the background. This does
// NOT drop duplicates if they exist. Duplicates will cause an error.
func (self *MongoDataSource) EnsureUniqueIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: true,
		DropDups: false,
		Background: false,
		Sparse: false,
	}); err != nil {
		return NewStackError("Unable to EnsureUniqueIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a non-unique, non-sparse index is created. This does not create in the background.
func (self *MongoDataSource) EnsureIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: false,
	}); err != nil {
		return NewStackError("Unable to EnsureIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a non-unique, sparse index is created. This does not create in the background.
func (self *MongoDataSource) EnsureSparseIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: true,
	}); err != nil {
		return NewStackError("Unable to EnsureSparseIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Create a capped collection.
func (self *MongoDataSource) CreateCappedCollection(sizeInBytes int) error {
	if err := self.Collection().Create(&mgo.CollectionInfo{ DisableIdIndex: false, ForceIdIndex: true, Capped: true, MaxBytes: sizeInBytes }); err != nil {

		// This is a bit of a hack, but the error returned does not provide any codes.
		msg := strings.ToLower(err.Error())
		if strings.Contains(msg, "already") || strings.Contains(msg, "exists") { return nil }

		return NewStackError("Unable to CreateCappedCollection - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)

	} else {
		return nil
	}
}

// Create a ttl index.
func (self *MongoDataSource) EnsureTtlIndex(field string, expireAfterSeconds int) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: []string{ field },
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: false,
		ExpireAfter: time.Duration(expireAfterSeconds) * time.Second,
	}); err != nil {
		return NewStackError("Unable to EnsureTtlIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

func (self *MongoDataSource) Now() *time.Time {
	now := time.Now()
	return &now
}

// Ensure a unique, sparse index is created. This does not create in the background. This does
// NOT drop duplicates if they exist. Duplicates will cause an error.
func (self *MongoDataSource) EnsureUniqueSparseIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: true,
		DropDups: false,
		Background: false,
		Sparse: true,
	}); err != nil {
		return NewStackError("Unable to EnsureUniqueSparseIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Returns the collection from the session.
func (self *MongoDataSource) Collection() *mgo.Collection { return self.Mongo.Collection(self.DbName, self.CollectionName) }

// Returns the database from the session.
func (self *MongoDataSource) Db() *mgo.Database { return self.Mongo.Db(self.DbName) }

// Returns the session struct.
func (self *MongoDataSource) Session() *mgo.Session { return self.Mongo.session }

// Returns a clone of the session struct.
func (self *MongoDataSource) SessionClone() *mgo.Session { return self.Mongo.session.Clone() }

// Returns a copy of the session struct.
func (self *MongoDataSource) SessionCopy() *mgo.Session { return self.Mongo.session.Clone() }

<commit_msg>added IsDupErr check and nil check for not found err check<commit_after>/**
 * (C) Copyright 2014, Deft Labs
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at:
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package dlshared

import (
	"time"
	"strings"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type MongoDataSource struct {
	DbName string
	CollectionName string
	Mongo *Mongo
	Logger
}

func ObjectIdHex(objectIdHex string) *bson.ObjectId {
	id := bson.ObjectIdHex(objectIdHex)
	return &id
}

func (self *MongoDataSource) NewObjectId() *bson.ObjectId {
	id := bson.NewObjectId()
	return &id
}

// Insert a document into a collection with the base configured write concern.
func (self *MongoDataSource) Insert(doc interface{}) error {
	if err := self.Mongo.Collection(self.DbName, self.CollectionName).Insert(doc); err != nil {
		return NewStackError("Unable to Insert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Upsert a document in a collection with the base configured write concern.
func (self *MongoDataSource) Upsert(selector interface{}, change interface{}) error {
	if _, err := self.Mongo.Collection(self.DbName, self.CollectionName).Upsert(selector, change); err != nil {
		return NewStackError("Unable to Upsert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Upsert a document into a collection with the passed write concern.
func (self *MongoDataSource) UpsertSafe(selector interface{}, change interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	if _, err := session.DB(self.DbName).C(self.CollectionName).Upsert(selector, change); err != nil {
		return NewStackError("Unable to Upsert - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Updates a document into a collection with the passed write concern.
func (self *MongoDataSource) UpdateSafe(selector interface{}, change interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	if err := session.DB(self.DbName).C(self.CollectionName).Update(selector, change); err != nil {
		return NewStackError("Unable to Update - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Insert a document into a collection with the passed write concern.
func (self *MongoDataSource) InsertSafe(doc interface{}) error {
	session := self.SessionClone()

	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)
	if err := session.DB(self.DbName).C(self.CollectionName).Insert(doc); err != nil {
		return NewStackError("Unable to InsertSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Unset a field with the default safe enabled - uses $unset
func (self *MongoDataSource) UnsetFieldSafe(query interface{}, field string) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$unset": &bson.M{ field: nadaStr } }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to UnsetFieldSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Set fields using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) SetFieldsSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$set": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to SetFieldsSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Pull from an array call using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) PullSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$pull": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to PullSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Find the distinct string fields. Do not use this on datasets with a large amount of distinct values or
// you will blow out memory. The selector can be nil.
func (self *MongoDataSource) FindDistinctStrs(selector interface{}, fieldName string) ([]string, error) {
	var result []string
	if err := self.Mongo.Collection(self.DbName, self.CollectionName).Find(selector).Distinct(fieldName, &result); err != nil { return nil, err }

	return result, nil
}

// The caller must close the cursor when done. Use: defer cursor.Clse()
func (self *MongoDataSource) FindManyWithBatchSize(selector interface{}, batchSize int) *mgo.Iter {
	return self.Mongo.Collection(self.DbName, self.CollectionName).Find(selector).Batch(batchSize).Iter()
}

// Push to an array call using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) PushSafe(query interface{}, fieldsDoc interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$push": fieldsDoc }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to PushSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Set a property using a "safe" operation. If this is a standalone mongo or a mongos, it will use: WMode: "majority".
// If this is a standalone mongo, it will use: w: 1
func (self *MongoDataSource) SetFieldSafe(query interface{}, field string, value interface{}) error {
	session := self.SessionClone()
	defer session.Close()

	session.SetSafe(self.Mongo.DefaultSafe)

	update := &bson.M{ "$set": &bson.M{ field: value } }

	if err := self.RemoveNotFoundErr(session.DB(self.DbName).C(self.CollectionName).Update(query, update)); err != nil {
		return NewStackError("Unable to SetFieldSafe - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Find by the _id. Returns false if not found.
func (self *MongoDataSource) FindById(id interface{}, result interface{}) error {
	return self.FindOne(&bson.M{ "_id": id }, result)
}

// Returns nil if this is a NOT a document not found error.
func (self *MongoDataSource) RemoveNotFoundErr(err error) error {
	if self.NotFoundErr(err) { return nil }
	return err
}

// Returns true if this is a dup error. If nil is passed, false is returned.
func (self *MongoDataSource) IsDupErr(err error) bool {
	if err == nil { return false }
	return mgo.IsDup(err)
}

// Returns true if this is a document not found error. If nil is passed, false is returned.
func (self *MongoDataSource) NotFoundErr(err error) bool {
	if err == nil { return false }
	return err == mgo.ErrNotFound
}

// Finds one document or returns false.
func (self *MongoDataSource) FindOne(query *bson.M, result interface{}) error { return self.Collection().Find(query).One(result) }

// Delete one document by the _id.
func (self *MongoDataSource) DeleteById(id interface{}) error { return self.DeleteOne(&bson.M{"_id": id }) }

// Delete one document from the collection. If the document is not found, no error is returned.
func (self *MongoDataSource) DeleteOne(selector interface{}) error {
	if err := self.RemoveNotFoundErr(self.Collection().Remove(selector)); err != nil {
		return NewStackError("Unable to DeleteOne - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Delete one or more documents from the collection. If the document(s) is/are not found, no error
// is returned.
func (self *MongoDataSource) Delete(selector interface{}) error {
	if _, err := self.Collection().RemoveAll(selector); err != nil {
		return NewStackError("Unable to Delete - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a unique, non-sparse index is created. This does not create in the background. This does
// NOT drop duplicates if they exist. Duplicates will cause an error.
func (self *MongoDataSource) EnsureUniqueIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: true,
		DropDups: false,
		Background: false,
		Sparse: false,
	}); err != nil {
		return NewStackError("Unable to EnsureUniqueIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a non-unique, non-sparse index is created. This does not create in the background.
func (self *MongoDataSource) EnsureIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: false,
	}); err != nil {
		return NewStackError("Unable to EnsureIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Ensure a non-unique, sparse index is created. This does not create in the background.
func (self *MongoDataSource) EnsureSparseIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: true,
	}); err != nil {
		return NewStackError("Unable to EnsureSparseIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Create a capped collection.
func (self *MongoDataSource) CreateCappedCollection(sizeInBytes int) error {
	if err := self.Collection().Create(&mgo.CollectionInfo{ DisableIdIndex: false, ForceIdIndex: true, Capped: true, MaxBytes: sizeInBytes }); err != nil {

		// This is a bit of a hack, but the error returned does not provide any codes.
		msg := strings.ToLower(err.Error())
		if strings.Contains(msg, "already") || strings.Contains(msg, "exists") { return nil }

		return NewStackError("Unable to CreateCappedCollection - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)

	} else {
		return nil
	}
}

// Create a ttl index.
func (self *MongoDataSource) EnsureTtlIndex(field string, expireAfterSeconds int) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: []string{ field },
		Unique: false,
		DropDups: false,
		Background: false,
		Sparse: false,
		ExpireAfter: time.Duration(expireAfterSeconds) * time.Second,
	}); err != nil {
		return NewStackError("Unable to EnsureTtlIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

func (self *MongoDataSource) Now() *time.Time {
	now := time.Now()
	return &now
}

// Ensure a unique, sparse index is created. This does not create in the background. This does
// NOT drop duplicates if they exist. Duplicates will cause an error.
func (self *MongoDataSource) EnsureUniqueSparseIndex(fields []string) error {
	if err := self.Collection().EnsureIndex(mgo.Index{
		Key: fields,
		Unique: true,
		DropDups: false,
		Background: false,
		Sparse: true,
	}); err != nil {
		return NewStackError("Unable to EnsureUniqueSparseIndex - db: %s - collection: %s - error: %v", self.DbName, self.CollectionName, err)
	}

	return nil
}

// Returns the collection from the session.
func (self *MongoDataSource) Collection() *mgo.Collection { return self.Mongo.Collection(self.DbName, self.CollectionName) }

// Returns the database from the session.
func (self *MongoDataSource) Db() *mgo.Database { return self.Mongo.Db(self.DbName) }

// Returns the session struct.
func (self *MongoDataSource) Session() *mgo.Session { return self.Mongo.session }

// Returns a clone of the session struct.
func (self *MongoDataSource) SessionClone() *mgo.Session { return self.Mongo.session.Clone() }

// Returns a copy of the session struct.
func (self *MongoDataSource) SessionCopy() *mgo.Session { return self.Mongo.session.Clone() }

<|endoftext|>


<commit_before>// Package build outputs a function's zip to stdout.
package build

import (
	"errors"
	"io"
	"os"

	"github.com/spf13/cobra"

	"github.com/apex/apex/cmd/apex/root"
	"github.com/apex/apex/stats"
	"github.com/apex/apex/utils"
)

// name of function.
var name string

// env supplied.
var env []string

// example output.
const example = `  Build zip output for a function
  $ apex build foo > /tmp/out.zip`

// Command config.
var Command = &cobra.Command{
	Use:     "build <name>",
	Short:   "Build a function",
	Example: example,
	PreRunE: preRun,
	RunE:    run,
}

// Initialize.
func init() {
	root.Register(Command)

	f := Command.Flags()
	f.StringSliceVarP(&env, "set", "s", nil, "Set environment variable")
}

// PreRun errors if argument is missing.
func preRun(c *cobra.Command, args []string) error {
	if len(args) < 1 {
		return errors.New("Missing name argument")
	}

	name = args[0]
	return nil
}

// Run command.
func run(c *cobra.Command, args []string) error {
	if err := root.Project.LoadFunctions(name); err != nil {
		return err
	}

	fn := root.Project.Functions[0]

	stats.Track("Build", map[string]interface{}{
		"runtime": fn.Runtime,
		"env":     len(env),
	})

	vars, err := utils.ParseEnv(env)
	if err != nil {
		return err
	}

	for k, v := range vars {
		root.Project.Setenv(k, v)
	}

	zip, err := fn.Build()
	if err != nil {
		return err
	}

	if err := fn.Clean(); err != nil {
		return err
	}

	_, err = io.Copy(os.Stdout, zip)
	return err
}
<commit_msg>add missing --env-file for build command. Closes #452<commit_after>// Package build outputs a function's zip to stdout.
package build

import (
	"errors"
	"fmt"
	"io"
	"os"

	"github.com/spf13/cobra"

	"github.com/apex/apex/cmd/apex/root"
	"github.com/apex/apex/stats"
	"github.com/apex/apex/utils"
)

// name of function.
var name string

// env file.
var envFile string

// env supplied.
var env []string

// example output.
const example = `  Build zip output for a function
  $ apex build foo > /tmp/out.zip`

// Command config.
var Command = &cobra.Command{
	Use:     "build <name>",
	Short:   "Build a function",
	Example: example,
	PreRunE: preRun,
	RunE:    run,
}

// Initialize.
func init() {
	root.Register(Command)

	f := Command.Flags()
	f.StringVarP(&envFile, "env-file", "E", "", "Set environment variables from JSON file")
	f.StringSliceVarP(&env, "set", "s", nil, "Set environment variable")
}

// PreRun errors if argument is missing.
func preRun(c *cobra.Command, args []string) error {
	if len(args) < 1 {
		return errors.New("Missing name argument")
	}

	name = args[0]
	return nil
}

// Run command.
func run(c *cobra.Command, args []string) error {
	if err := root.Project.LoadFunctions(name); err != nil {
		return err
	}

	fn := root.Project.Functions[0]

	stats.Track("Build", map[string]interface{}{
		"runtime":      fn.Runtime,
		"env":          len(env),
		"has_env_file": envFile != "",
	})

	if envFile != "" {
		if err := root.Project.LoadEnvFromFile(envFile); err != nil {
			return fmt.Errorf("reading env file %q: %s", envFile, err)
		}
	}

	vars, err := utils.ParseEnv(env)
	if err != nil {
		return err
	}

	for k, v := range vars {
		root.Project.Setenv(k, v)
	}

	zip, err := fn.Build()
	if err != nil {
		return err
	}

	if err := fn.Clean(); err != nil {
		return err
	}

	_, err = io.Copy(os.Stdout, zip)
	return err
}
<|endoftext|>


<commit_before>package entrypoint

import (
	"context"
	"fmt"
	"os"
	"time"

	"github.com/datawire/dlib/dcontext"
	"github.com/datawire/dlib/dgroup"
)

func runEnvoy(ctx context.Context, envoyHUP chan os.Signal) error {
	// Wait until we get a SIGHUP to start envoy.
	//var bootstrap string
	select {
	case <-envoyHUP:
		/*bytes, err := ioutil.ReadFile(GetEnvoyBootstrapFile())
		if err != nil {
			panic(err)
		}
		bootstrap = string(bytes)*/
	case <-ctx.Done():
		return nil
	}

	// Try to run envoy directly, but fallback to running it inside docker if there is
	// no envoy executable available.
	if IsEnvoyAvailable() {
		cmd := subcommand(ctx, "envoy", GetEnvoyFlags()...)
		if envbool("DEV_SHUTUP_ENVOY") {
			cmd.Stdout = nil
			cmd.Stderr = nil
		}
		return cmd.Run()
	} else {
		// For some reason docker only sometimes passes the signal onto the process inside
		// the container, so we setup this cleanup function so that in the docker case we
		// can do a docker kill, just to be sure it is really dead and we don't leave an
		// envoy lying around.

		// Create a label unique to this invocation so we can use it to do a docker
		// kill for cleanup.
		label := fmt.Sprintf("amb-envoy-label-%d", os.Getpid())

		snapdir := GetSnapshotDir()

		group := dgroup.NewGroup(ctx, dgroup.GroupConfig{
			ShutdownOnNonError: true,
		})
		group.Go("envoy", func(ctx context.Context) error {
			// XXX: will host networking work on a mac? (probably not)
			cmd := subcommand(ctx, "docker", append([]string{"run", "-l", label, "--rm", "--network", "host",
				"-v", fmt.Sprintf("%s:%s", snapdir, snapdir),
				"-v", fmt.Sprintf("%s:%s", GetEnvoyBootstrapFile(), GetEnvoyBootstrapFile()),
				"--entrypoint", "envoy", "docker.io/datawire/aes:1.6.2"},
				GetEnvoyFlags()...)...)
			if envbool("DEV_SHUTUP_ENVOY") {
				cmd.Stdout = nil
				cmd.Stderr = nil
			}
			return cmd.Run()
		})
		group.Go("cleanup", func(ctx context.Context) error {
			<-ctx.Done()
			ctx = dcontext.HardContext(ctx)

			for {
				cids := cidsForLabel(ctx, label)
				if len(cids) == 0 {
					return nil
				}

				// Give the container one second to exit
				tctx, _ := context.WithTimeout(ctx, 1*time.Second)
				wait := subcommand(tctx, "docker", append([]string{"wait"}, cids...)...)
				wait.Stdout = nil
				logExecError("docker wait", wait.Run())

				cids = cidsForLabel(ctx, label)
				if len(cids) == 0 {
					return nil
				}

				kill := subcommand(ctx, "docker", append([]string{"kill"}, cids...)...)
				kill.Stdout = nil
				logExecError("docker kill", kill.Run())
			}
		})
		return group.Wait()
	}
}
<commit_msg>(from AES) Add env var to control whether envoy is run as privileged.<commit_after>package entrypoint

import (
	"context"
	"fmt"
	"os"
	"time"

	"github.com/datawire/dlib/dcontext"
	"github.com/datawire/dlib/dgroup"
)

func runEnvoy(ctx context.Context, envoyHUP chan os.Signal) error {
	// Wait until we get a SIGHUP to start envoy.
	//var bootstrap string
	select {
	case <-envoyHUP:
		/*bytes, err := ioutil.ReadFile(GetEnvoyBootstrapFile())
		if err != nil {
			panic(err)
		}
		bootstrap = string(bytes)*/
	case <-ctx.Done():
		return nil
	}

	// Try to run envoy directly, but fallback to running it inside docker if there is
	// no envoy executable available.
	if IsEnvoyAvailable() {
		cmd := subcommand(ctx, "envoy", GetEnvoyFlags()...)
		if envbool("DEV_SHUTUP_ENVOY") {
			cmd.Stdout = nil
			cmd.Stderr = nil
		}
		return cmd.Run()
	} else {
		// For some reason docker only sometimes passes the signal onto the process inside
		// the container, so we setup this cleanup function so that in the docker case we
		// can do a docker kill, just to be sure it is really dead and we don't leave an
		// envoy lying around.

		// Create a label unique to this invocation so we can use it to do a docker
		// kill for cleanup.
		label := fmt.Sprintf("amb-envoy-label-%d", os.Getpid())

		snapdir := GetSnapshotDir()

		group := dgroup.NewGroup(ctx, dgroup.GroupConfig{
			ShutdownOnNonError: true,
		})
		group.Go("envoy", func(ctx context.Context) error {
			// XXX: will host networking work on a mac? (probably not)
			dockerArgs := []string{
				"run", "-l", label, "--rm", "--network", "host",
				"-v", fmt.Sprintf("%s:%s", snapdir, snapdir),
				"-v", fmt.Sprintf("%s:%s", GetEnvoyBootstrapFile(), GetEnvoyBootstrapFile()),
			}

			if envbool("DEV_ENVOY_DOCKER_PRIVILEGED") {
				dockerArgs = append(dockerArgs, "--privileged")
			}

			dockerArgs = append(dockerArgs, "--entrypoint", "envoy", "docker.io/datawire/aes:1.6.2")
			cmd := subcommand(ctx, "docker", append(dockerArgs, GetEnvoyFlags()...)...)
			if envbool("DEV_SHUTUP_ENVOY") {
				cmd.Stdout = nil
				cmd.Stderr = nil
			}
			return cmd.Run()
		})
		group.Go("cleanup", func(ctx context.Context) error {
			<-ctx.Done()
			ctx = dcontext.HardContext(ctx)

			for {
				cids := cidsForLabel(ctx, label)
				if len(cids) == 0 {
					return nil
				}

				// Give the container one second to exit
				tctx, _ := context.WithTimeout(ctx, 1*time.Second)
				wait := subcommand(tctx, "docker", append([]string{"wait"}, cids...)...)
				wait.Stdout = nil
				logExecError("docker wait", wait.Run())

				cids = cidsForLabel(ctx, label)
				if len(cids) == 0 {
					return nil
				}

				kill := subcommand(ctx, "docker", append([]string{"kill"}, cids...)...)
				kill.Stdout = nil
				logExecError("docker kill", kill.Run())
			}
		})
		return group.Wait()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"path/filepath"
	"flag"

	"github.com/constabulary/gb"
	"github.com/constabulary/gb/cmd"
	"github.com/constabulary/gb/cmd/gb-vendor/vendor"
)

var (
	// gb vendor update flags

	// update all dependencies
	updateAll bool
)

func init() {
	registerCommand("update", UpdateCmd)
}

func addUpdateFlags(fs *flag.FlagSet){
	fs.BoolVar(&updateAll, "all", false, "update all dependencies")
}

var UpdateCmd = &cmd.Command{
	ShortDesc: "updates a local dependency",
	Run: func(ctx *gb.Context, args []string) error {
		if len(args) != 1 && !updateAll{
			return fmt.Errorf("update: import path or --all flag is missing")
		} else if len(args) == 1 && updateAll {
			return fmt.Errorf("update: you cannot specify path and --all flag at once")
		}

		m, err := vendor.ReadManifest(manifestFile(ctx))
		if err != nil {
			return fmt.Errorf("could not load manifest: %T %v", err, err)
		}

		var dependencies []vendor.Dependency
		if updateAll {
			dependencies = make([]vendor.Dependency, len(m.Dependencies))
			copy(dependencies, m.Dependencies)
		} else {
			p := args[0]
			dependency, err := m.GetDependencyForImportpath(p)
			if err != nil {
				return fmt.Errorf("could not get dependency: %T %v", err, err)
			}
			dependencies = append(dependencies, dependency)
		}

		for _, d := range dependencies {
			url := d.Repository
			path := d.Importpath

			err = m.RemoveDependency(d)
			if err != nil {
				return fmt.Errorf("dependency could not be deleted from manifest: %T %v", err, err)
			}

			localClone := vendor.GitClone{
				Path: filepath.Join(ctx.Projectdir(), "vendor", "src", path),
			}
			err = localClone.Destroy()
			if err != nil {
				return fmt.Errorf("dependency could not be deleted: %T %v", err, err)
			}

			repo := vendor.GitRepo{
				URL: url,
			}

			wc, err := repo.Clone()
			if err != nil {
				return err
			}

			rev, err := wc.Revision()
			if err != nil {
				return err
			}

			branch, err := wc.Branch()
			if err != nil {
				return err
			}

			dep := vendor.Dependency{
				Importpath: path,
				Repository: url,
				Revision:   rev,
				Branch:     branch,
				Path:       "",
			}

			if err := m.AddDependency(dep); err != nil {
				return err
			}

			dst := filepath.Join(ctx.Projectdir(), "vendor", "src", dep.Importpath)
			src := filepath.Join(wc.Dir(), dep.Path)

			if err := copypath(dst, src); err != nil {
				return err
			}

			if err := wc.Destroy(); err != nil {
				return err
			}
			fmt.Println(dependencies)
		}

		return vendor.WriteManifest(manifestFile(ctx), m)
	},
	AddFlags: addUpdateFlags,
}
<commit_msg>Fixed formatting<commit_after>package main

import (
	"flag"
	"fmt"
	"path/filepath"

	"github.com/constabulary/gb"
	"github.com/constabulary/gb/cmd"
	"github.com/constabulary/gb/cmd/gb-vendor/vendor"
)

var (
	// gb vendor update flags

	// update all dependencies
	updateAll bool
)

func init() {
	registerCommand("update", UpdateCmd)
}

func addUpdateFlags(fs *flag.FlagSet) {
	fs.BoolVar(&updateAll, "all", false, "update all dependencies")
}

var UpdateCmd = &cmd.Command{
	ShortDesc: "updates a local dependency",
	Run: func(ctx *gb.Context, args []string) error {
		if len(args) != 1 && !updateAll {
			return fmt.Errorf("update: import path or --all flag is missing")
		} else if len(args) == 1 && updateAll {
			return fmt.Errorf("update: you cannot specify path and --all flag at once")
		}

		m, err := vendor.ReadManifest(manifestFile(ctx))
		if err != nil {
			return fmt.Errorf("could not load manifest: %T %v", err, err)
		}

		var dependencies []vendor.Dependency
		if updateAll {
			dependencies = make([]vendor.Dependency, len(m.Dependencies))
			copy(dependencies, m.Dependencies)
		} else {
			p := args[0]
			dependency, err := m.GetDependencyForImportpath(p)
			if err != nil {
				return fmt.Errorf("could not get dependency: %T %v", err, err)
			}
			dependencies = append(dependencies, dependency)
		}

		for _, d := range dependencies {
			url := d.Repository
			path := d.Importpath

			err = m.RemoveDependency(d)
			if err != nil {
				return fmt.Errorf("dependency could not be deleted from manifest: %T %v", err, err)
			}

			localClone := vendor.GitClone{
				Path: filepath.Join(ctx.Projectdir(), "vendor", "src", path),
			}
			err = localClone.Destroy()
			if err != nil {
				return fmt.Errorf("dependency could not be deleted: %T %v", err, err)
			}

			repo := vendor.GitRepo{
				URL: url,
			}

			wc, err := repo.Clone()
			if err != nil {
				return err
			}

			rev, err := wc.Revision()
			if err != nil {
				return err
			}

			branch, err := wc.Branch()
			if err != nil {
				return err
			}

			dep := vendor.Dependency{
				Importpath: path,
				Repository: url,
				Revision:   rev,
				Branch:     branch,
				Path:       "",
			}

			if err := m.AddDependency(dep); err != nil {
				return err
			}

			dst := filepath.Join(ctx.Projectdir(), "vendor", "src", dep.Importpath)
			src := filepath.Join(wc.Dir(), dep.Path)

			if err := copypath(dst, src); err != nil {
				return err
			}

			if err := wc.Destroy(); err != nil {
				return err
			}
			fmt.Println(dependencies)
		}

		return vendor.WriteManifest(manifestFile(ctx), m)
	},
	AddFlags: addUpdateFlags,
}
<|endoftext|>


<commit_before>/*
 Copyright 2021 The GoPlus Authors (goplus.org)

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
*/

// Package run implements the ``gop run'' command.
package run

import (
	"os"
	"os/exec"
	"path/filepath"

	"github.com/goplus/gop/ast"
	"github.com/goplus/gop/cl"
	"github.com/goplus/gop/cmd/internal/base"
	"github.com/goplus/gop/parser"
	"github.com/goplus/gop/token"
	"github.com/goplus/gox"
	"github.com/qiniu/x/log"
)

// -----------------------------------------------------------------------------

// Cmd - gop run
var Cmd = &base.Command{
	UsageLine: "gop run [-asm -quiet -debug -prof] <gopSrcDir|gopSrcFile>",
	Short:     "Run a Go+ program",
}

var (
	flag      = &Cmd.Flag
	flagAsm   = flag.Bool("asm", false, "generates `asm` code of Go+ bytecode backend")
	flagQuiet = flag.Bool("quiet", false, "don't generate any compiling stage log")
	flagDebug = flag.Bool("debug", false, "print debug information")
	flagProf  = flag.Bool("prof", false, "do profile and generate profile report")
)

func init() {
	Cmd.Run = runCmd
}

func saveGoFile(dir string, pkg *gox.Package) error {
	err := os.MkdirAll(dir, 0777)
	if err != nil {
		return err
	}
	return gox.WriteFile(dir+"/gop_autogen.go", pkg)
}

func runCmd(cmd *base.Command, args []string) {
	flag.Parse(args)
	if flag.NArg() < 1 {
		cmd.Usage(os.Stderr)
	}

	log.SetFlags(log.Ldefault &^ log.LstdFlags)
	if *flagQuiet {
		log.SetOutputLevel(0x7000)
	} else if *flagDebug {
		log.SetOutputLevel(log.Ldebug)
	}
	if *flagProf {
		panic("TODO: profile not impl")
	}
	fset := token.NewFileSet()

	target, _ := filepath.Abs(flag.Arg(0))
	isDir, err := IsDir(target)
	if err != nil {
		log.Fatalln("input arg check failed:", err)
	}
	var pkgs map[string]*ast.Package
	if isDir {
		pkgs, err = parser.ParseDir(fset, target, nil, 0)
	} else {
		pkgs, err = parser.Parse(fset, target, nil, 0)
	}
	if err != nil {
		log.Fatalln("parser.Parse failed:", err)
	}

	conf := &cl.Config{}
	out, err := cl.NewPackage("", pkgs["main"], fset, conf)
	if err != nil {
		log.Fatalln("cl.NewPackage failed:", err)
	}
	if *flagAsm {
		panic("TODO: gop run -asm not impl")
	}
	var gofile string
	if isDir {
		err = saveGoFile(target, out)
		gofile = filepath.Join(target, "*.go")
	} else {
		dir, file := filepath.Split(target)
		gofile = filepath.Join(dir, file[:len(file)-len(filepath.Ext(file))]+".go")
		err = gox.WriteFile(gofile, out)
	}
	if err != nil {
		log.Fatalln("saveGoFile failed:", err)
	}
	err = exec.Command("go", "run", gofile).Run()
	if err != nil {
		log.Fatalln("go run failed:", err)
	}
	if *flagProf {
		panic("TODO: profile not impl")
	}
}

// IsDir checks a target path is dir or not.
func IsDir(target string) (bool, error) {
	fi, err := os.Stat(target)
	if err != nil {
		return false, err
	}
	return fi.IsDir(), nil
}

// -----------------------------------------------------------------------------
<commit_msg>gop run<commit_after>/*
 Copyright 2021 The GoPlus Authors (goplus.org)

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
*/

// Package run implements the ``gop run'' command.
package run

import (
	"os"
	"os/exec"
	"path/filepath"

	"github.com/goplus/gop/ast"
	"github.com/goplus/gop/cl"
	"github.com/goplus/gop/cmd/internal/base"
	"github.com/goplus/gop/parser"
	"github.com/goplus/gop/token"
	"github.com/goplus/gox"
	"github.com/qiniu/x/log"
)

// -----------------------------------------------------------------------------

// Cmd - gop run
var Cmd = &base.Command{
	UsageLine: "gop run [-asm -quiet -debug -prof] <gopSrcDir|gopSrcFile>",
	Short:     "Run a Go+ program",
}

var (
	flag      = &Cmd.Flag
	flagAsm   = flag.Bool("asm", false, "generates `asm` code of Go+ bytecode backend")
	flagQuiet = flag.Bool("quiet", false, "don't generate any compiling stage log")
	flagDebug = flag.Bool("debug", false, "print debug information")
	flagProf  = flag.Bool("prof", false, "do profile and generate profile report")
)

func init() {
	Cmd.Run = runCmd
}

func saveGoFile(dir string, pkg *gox.Package) error {
	err := os.MkdirAll(dir, 0777)
	if err != nil {
		return err
	}
	return gox.WriteFile(dir+"/gop_autogen.go", pkg)
}

func runCmd(cmd *base.Command, args []string) {
	flag.Parse(args)
	if flag.NArg() < 1 {
		cmd.Usage(os.Stderr)
	}

	log.SetFlags(log.Ldefault &^ log.LstdFlags)
	if *flagQuiet {
		log.SetOutputLevel(0x7000)
	} else if *flagDebug {
		log.SetOutputLevel(log.Ldebug)
	}
	if *flagProf {
		panic("TODO: profile not impl")
	}
	fset := token.NewFileSet()

	target, _ := filepath.Abs(flag.Arg(0))
	isDir, err := IsDir(target)
	if err != nil {
		log.Fatalln("input arg check failed:", err)
	}
	var pkgs map[string]*ast.Package
	if isDir {
		pkgs, err = parser.ParseDir(fset, target, nil, 0)
	} else {
		pkgs, err = parser.Parse(fset, target, nil, 0)
	}
	if err != nil {
		log.Fatalln("parser.Parse failed:", err)
	}

	conf := &cl.Config{}
	out, err := cl.NewPackage("", pkgs["main"], fset, conf)
	if err != nil {
		log.Fatalln("cl.NewPackage failed:", err)
	}
	if *flagAsm {
		panic("TODO: gop run -asm not impl")
	}
	var gofile string
	if isDir {
		err = saveGoFile(target, out)
		gofile = target
	} else {
		dir, file := filepath.Split(target)
		gofile = filepath.Join(dir, file[:len(file)-len(filepath.Ext(file))]+".go")
		err = gox.WriteFile(gofile, out)
	}
	if err != nil {
		log.Fatalln("saveGoFile failed:", err)
	}
	err = goRun(gofile)
	if err != nil {
		switch e := err.(type) {
		case *exec.ExitError:
			os.Stderr.Write(e.Stderr)
		default:
			log.Fatalln("go run failed:", err)
		}
	}
	if *flagProf {
		panic("TODO: profile not impl")
	}
}

// IsDir checks a target path is dir or not.
func IsDir(target string) (bool, error) {
	fi, err := os.Stat(target)
	if err != nil {
		return false, err
	}
	return fi.IsDir(), nil
}

func goRun(target string) error {
	cmd := exec.Command("go", "run", target)
	gorun, err := cmd.Output()
	os.Stdout.Write(gorun)
	return err
}

// -----------------------------------------------------------------------------
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	stderrors "errors"
	"fmt"

	"launchpad.net/gnuflag"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/environs/storage"
	"launchpad.net/juju-core/environs/sync"
	envtools "launchpad.net/juju-core/environs/tools"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/rpc"
	coretools "launchpad.net/juju-core/tools"
	"launchpad.net/juju-core/version"
)

// UpgradeJujuCommand upgrades the agents in a juju installation.
type UpgradeJujuCommand struct {
	cmd.EnvCommandBase
	vers        string
	Version     version.Number
	UploadTools bool
	Series      []string
}

var upgradeJujuDoc = `
The upgrade-juju command upgrades a running environment by setting a version
number for all juju agents to run. By default, it chooses the most recent
supported version compatible with the command-line tools version.

A development version is defined to be any version with an odd minor
version or a nonzero build component (for example version 2.1.1, 3.3.0
and 2.0.0.1 are development versions; 2.0.3 and 3.4.1 are not). A
development version may be chosen in two cases:

 - when the current agent version is a development one and there is
   a more recent version available with the same major.minor numbers;
 - when an explicit --version major.minor is given (e.g. --version 1.17,
   or 1.17.2, but not just 1)

For development use, the --upload-tools flag specifies that the juju tools will
packaged (or compiled locally, if no jujud binaries exists, for which you will
need the golang packages installed) and uploaded before the version is set.
Currently the tools will be uploaded as if they had the version of the current
juju tool, unless specified otherwise by the --version flag.

When run without arguments. upgrade-juju will try to upgrade to the
following versions, in order of preference, depending on the current
value of the environment's agent-version setting:

 - The highest patch.build version of the *next* stable major.minor version.
 - The highest patch.build version of the *current* major.minor version.

Both of these depend on tools availability, which some situations (no
outgoing internet access) and provider types (such as maas) require that
you manage yourself; see the documentation for "sync-tools".
`

func (c *UpgradeJujuCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "upgrade-juju",
		Purpose: "upgrade the tools in a juju environment",
		Doc:     upgradeJujuDoc,
	}
}

func (c *UpgradeJujuCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.StringVar(&c.vers, "version", "", "upgrade to specific version")
	f.BoolVar(&c.UploadTools, "upload-tools", false, "upload local version of tools")
	f.Var(seriesVar{&c.Series}, "series", "upload tools for supplied comma-separated series list")
}

func (c *UpgradeJujuCommand) Init(args []string) error {
	if c.vers != "" {
		vers, err := version.Parse(c.vers)
		if err != nil {
			return err
		}
		if vers.Major != version.Current.Major {
			return fmt.Errorf("cannot upgrade to version incompatible with CLI")
		}
		if c.UploadTools && vers.Build != 0 {
			// TODO(fwereade): when we start taking versions from actual built
			// code, we should disable --version when used with --upload-tools.
			// For now, it's the only way to experiment with version upgrade
			// behaviour live, so the only restriction is that Build cannot
			// be used (because its value needs to be chosen internally so as
			// not to collide with existing tools).
			return fmt.Errorf("cannot specify build number when uploading tools")
		}
		c.Version = vers
	}
	if len(c.Series) > 0 && !c.UploadTools {
		return fmt.Errorf("--series requires --upload-tools")
	}
	return cmd.CheckEmpty(args)
}

var errUpToDate = stderrors.New("no upgrades available")

// Run changes the version proposed for the juju envtools.
func (c *UpgradeJujuCommand) Run(_ *cmd.Context) (err error) {
	client, err := juju.NewAPIClientFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer client.Close()
	defer func() {
		if err == errUpToDate {
			log.Noticef(err.Error())
			err = nil
		}
	}()

	// Determine the version to upgrade to, uploading tools if necessary.
	attrs, err := client.EnvironmentGet()
	if rpc.IsNoSuchRequest(err) {
		return c.run1dot16()
	}
	if err != nil {
		return err
	}
	cfg, err := config.New(config.NoDefaults, attrs)
	if err != nil {
		return err
	}
	env, err := environs.New(cfg)
	if err != nil {
		return err
	}
	v, err := c.initVersions(cfg, env)
	if err != nil {
		return err
	}
	if c.UploadTools {
		series := getUploadSeries(cfg, c.Series)
		if err := v.uploadTools(env.Storage(), series); err != nil {
			return err
		}
	}
	if err := v.validate(); err != nil {
		return err
	}
	log.Infof("upgrade version chosen: %s", v.chosen)
	// TODO(fwereade): this list may be incomplete, pending envtools.Upload change.
	log.Infof("available tools: %s", v.tools)

	if err := client.SetEnvironAgentVersion(v.chosen); err != nil {
		return err
	}
	log.Noticef("started upgrade to %s", v.chosen)
	return nil
}

// run1dot16 implements the command without access to the API. This is
// needed for compatibility, so 1.16 can be upgraded to newer
// releases. It should be removed in 1.18.
func (c *UpgradeJujuCommand) run1dot16() error {
	log.Infof("running in 1.16 compatibility mode")
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()
	defer func() {
		if err == errUpToDate {
			log.Noticef(err.Error())
			err = nil
		}
	}()

	// Determine the version to upgrade to, uploading tools if necessary.
	env := conn.Environ
	cfg, err := conn.State.EnvironConfig()
	if err != nil {
		return err
	}
	v, err := c.initVersions(cfg, env)
	if err != nil {
		return err
	}
	if c.UploadTools {
		series := getUploadSeries(cfg, c.Series)
		if err := v.uploadTools(env.Storage(), series); err != nil {
			return err
		}
	}
	if err := v.validate(); err != nil {
		return err
	}
	log.Infof("upgrade version chosen: %s", v.chosen)
	// TODO(fwereade): this list may be incomplete, pending envtools.Upload change.
	log.Infof("available tools: %s", v.tools)

	if err := conn.State.SetEnvironAgentVersion(v.chosen); err != nil {
		return err
	}
	log.Noticef("started upgrade to %s", v.chosen)
	return nil
}

// initVersions collects state relevant to an upgrade decision. The returned
// agent and client versions, and the list of currently available tools, will
// always be accurate; the chosen version, and the flag indicating development
// mode, may remain blank until uploadTools or validate is called.
func (c *UpgradeJujuCommand) initVersions(cfg *config.Config, env environs.Environ) (*upgradeVersions, error) {
	agent, ok := cfg.AgentVersion()
	if !ok {
		// Can't happen. In theory.
		return nil, fmt.Errorf("incomplete environment configuration")
	}
	if c.Version == agent {
		return nil, errUpToDate
	}
	client := version.Current.Number
	// TODO use an API call rather than requiring the environment,
	// so that we can restrict access to the provider secrets
	// while still allowing users to upgrade.
	available, err := envtools.FindTools(env, client.Major, -1, coretools.Filter{}, envtools.DoNotAllowRetry)
	if err != nil {
		if !errors.IsNotFoundError(err) {
			return nil, err
		}
		if !c.UploadTools {
			// No tools found and we shouldn't upload any, so pretend
			// there is no more recent version available.
			if c.Version == version.Zero {
				return nil, errUpToDate
			}
			return nil, err
		}
	}
	return &upgradeVersions{
		agent:  agent,
		client: client,
		chosen: c.Version,
		tools:  available,
	}, nil
}

// upgradeVersions holds the version information for making upgrade decisions.
type upgradeVersions struct {
	agent  version.Number
	client version.Number
	chosen version.Number
	tools  coretools.List
}

// uploadTools compiles jujud from $GOPATH and uploads it into the supplied
// storage. If no version has been explicitly chosen, the version number
// reported by the built tools will be based on the client version number.
// In any case, the version number reported will have a build component higher
// than that of any otherwise-matching available envtools.
// uploadTools resets the chosen version and replaces the available tools
// with the ones just uploaded.
func (v *upgradeVersions) uploadTools(storage storage.Storage, series []string) error {
	// TODO(fwereade): this is kinda crack: we should not assume that
	// version.Current matches whatever source happens to be built. The
	// ideal would be:
	//  1) compile jujud from $GOPATH into some build dir
	//  2) get actual version with `jujud version`
	//  3) check actual version for compatibility with CLI tools
	//  4) generate unique build version with reference to available tools
	//  5) force-version that unique version into the dir directly
	//  6) archive and upload the build dir
	// ...but there's no way we have time for that now. In the meantime,
	// considering the use cases, this should work well enough; but it
	// won't detect an incompatible major-version change, which is a shame.
	if v.chosen == version.Zero {
		v.chosen = v.client
	}
	v.chosen = uploadVersion(v.chosen, v.tools)

	// TODO(fwereade): envtools.Upload should return envtools.List, and should
	// include all the extra series we build, so we can set *that* onto
	// v.available and maybe one day be able to check that a given upgrade
	// won't leave out-of-date machines lying around, starved of envtools.
	uploaded, err := sync.Upload(storage, &v.chosen, series...)
	if err != nil {
		return err
	}
	v.tools = coretools.List{uploaded}
	return nil
}

// validate chooses an upgrade version, if one has not already been chosen,
// and ensures the tools list contains no entries that do not have that version.
// If validate returns no error, the environment agent-version can be set to
// the value of the chosen field.
func (v *upgradeVersions) validate() (err error) {
	if v.chosen == version.Zero {
		// No explicitly specified version, so find the next available
		// stable release to upgrade to, starting from the current agent
		// version and doing major.minor+1 or +2 as needed.
		nextStable := v.agent
		if v.agent.IsDev() {
			nextStable.Minor += 1
		} else {
			nextStable.Minor += 2
		}

		newestNextStable, found := v.tools.NewestCompatible(nextStable)
		if found {
			log.Debugf("found a more recent stable version %s", newestNextStable)
			v.chosen = newestNextStable
		} else {
			newestCurrent, found := v.tools.NewestCompatible(v.agent)
			if found {
				log.Debugf("found more recent current version %s", newestCurrent)
				v.chosen = newestCurrent
			} else {
				return fmt.Errorf("no more recent supported versions available")
			}
		}
	} else {
		// If not completely specified already, pick a single tools version.
		filter := coretools.Filter{Number: v.chosen, Released: !v.chosen.IsDev()}
		if v.tools, err = v.tools.Match(filter); err != nil {
			return err
		}
		v.chosen, v.tools = v.tools.Newest()
	}
	if v.chosen == v.agent {
		return errUpToDate
	}

	// Major version upgrade
	if v.chosen.Major < v.agent.Major {
		// TODO(fwereade): I'm a bit concerned about old agent/CLI tools even
		// *connecting* to environments with higher agent-versions; but ofc they
		// have to connect in order to discover they shouldn't. However, once
		// any of our tools detect an incompatible version, they should act to
		// minimize damage: the CLI should abort politely, and the agents should
		// run an Upgrader but no other tasks.
		return fmt.Errorf("cannot change major version from %d to %d", v.agent.Major, v.chosen.Major)
	} else if v.chosen.Major > v.agent.Major {
		return fmt.Errorf("major version upgrades are not supported yet")
	}

	return nil
}

// uploadVersion returns a copy of the supplied version with a build number
// higher than any of the supplied tools that share its major, minor and patch.
func uploadVersion(vers version.Number, existing coretools.List) version.Number {
	vers.Build++
	for _, t := range existing {
		if t.Version.Major != vers.Major || t.Version.Minor != vers.Minor || t.Version.Patch != vers.Patch {
			continue
		}
		if t.Version.Build >= vers.Build {
			vers.Build = t.Version.Build + 1
		}
	}
	return vers
}
<commit_msg>Infof to Warningf<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	stderrors "errors"
	"fmt"

	"launchpad.net/gnuflag"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/environs/storage"
	"launchpad.net/juju-core/environs/sync"
	envtools "launchpad.net/juju-core/environs/tools"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/rpc"
	coretools "launchpad.net/juju-core/tools"
	"launchpad.net/juju-core/version"
)

// UpgradeJujuCommand upgrades the agents in a juju installation.
type UpgradeJujuCommand struct {
	cmd.EnvCommandBase
	vers        string
	Version     version.Number
	UploadTools bool
	Series      []string
}

var upgradeJujuDoc = `
The upgrade-juju command upgrades a running environment by setting a version
number for all juju agents to run. By default, it chooses the most recent
supported version compatible with the command-line tools version.

A development version is defined to be any version with an odd minor
version or a nonzero build component (for example version 2.1.1, 3.3.0
and 2.0.0.1 are development versions; 2.0.3 and 3.4.1 are not). A
development version may be chosen in two cases:

 - when the current agent version is a development one and there is
   a more recent version available with the same major.minor numbers;
 - when an explicit --version major.minor is given (e.g. --version 1.17,
   or 1.17.2, but not just 1)

For development use, the --upload-tools flag specifies that the juju tools will
packaged (or compiled locally, if no jujud binaries exists, for which you will
need the golang packages installed) and uploaded before the version is set.
Currently the tools will be uploaded as if they had the version of the current
juju tool, unless specified otherwise by the --version flag.

When run without arguments. upgrade-juju will try to upgrade to the
following versions, in order of preference, depending on the current
value of the environment's agent-version setting:

 - The highest patch.build version of the *next* stable major.minor version.
 - The highest patch.build version of the *current* major.minor version.

Both of these depend on tools availability, which some situations (no
outgoing internet access) and provider types (such as maas) require that
you manage yourself; see the documentation for "sync-tools".
`

func (c *UpgradeJujuCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "upgrade-juju",
		Purpose: "upgrade the tools in a juju environment",
		Doc:     upgradeJujuDoc,
	}
}

func (c *UpgradeJujuCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.StringVar(&c.vers, "version", "", "upgrade to specific version")
	f.BoolVar(&c.UploadTools, "upload-tools", false, "upload local version of tools")
	f.Var(seriesVar{&c.Series}, "series", "upload tools for supplied comma-separated series list")
}

func (c *UpgradeJujuCommand) Init(args []string) error {
	if c.vers != "" {
		vers, err := version.Parse(c.vers)
		if err != nil {
			return err
		}
		if vers.Major != version.Current.Major {
			return fmt.Errorf("cannot upgrade to version incompatible with CLI")
		}
		if c.UploadTools && vers.Build != 0 {
			// TODO(fwereade): when we start taking versions from actual built
			// code, we should disable --version when used with --upload-tools.
			// For now, it's the only way to experiment with version upgrade
			// behaviour live, so the only restriction is that Build cannot
			// be used (because its value needs to be chosen internally so as
			// not to collide with existing tools).
			return fmt.Errorf("cannot specify build number when uploading tools")
		}
		c.Version = vers
	}
	if len(c.Series) > 0 && !c.UploadTools {
		return fmt.Errorf("--series requires --upload-tools")
	}
	return cmd.CheckEmpty(args)
}

var errUpToDate = stderrors.New("no upgrades available")

// Run changes the version proposed for the juju envtools.
func (c *UpgradeJujuCommand) Run(_ *cmd.Context) (err error) {
	client, err := juju.NewAPIClientFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer client.Close()
	defer func() {
		if err == errUpToDate {
			log.Noticef(err.Error())
			err = nil
		}
	}()

	// Determine the version to upgrade to, uploading tools if necessary.
	attrs, err := client.EnvironmentGet()
	if rpc.IsNoSuchRequest(err) {
		return c.run1dot16()
	}
	if err != nil {
		return err
	}
	cfg, err := config.New(config.NoDefaults, attrs)
	if err != nil {
		return err
	}
	env, err := environs.New(cfg)
	if err != nil {
		return err
	}
	v, err := c.initVersions(cfg, env)
	if err != nil {
		return err
	}
	if c.UploadTools {
		series := getUploadSeries(cfg, c.Series)
		if err := v.uploadTools(env.Storage(), series); err != nil {
			return err
		}
	}
	if err := v.validate(); err != nil {
		return err
	}
	log.Infof("upgrade version chosen: %s", v.chosen)
	// TODO(fwereade): this list may be incomplete, pending envtools.Upload change.
	log.Infof("available tools: %s", v.tools)

	if err := client.SetEnvironAgentVersion(v.chosen); err != nil {
		return err
	}
	log.Noticef("started upgrade to %s", v.chosen)
	return nil
}

// run1dot16 implements the command without access to the API. This is
// needed for compatibility, so 1.16 can be upgraded to newer
// releases. It should be removed in 1.18.
func (c *UpgradeJujuCommand) run1dot16() error {
	log.Warningf("running in 1.16 compatibility mode")
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()
	defer func() {
		if err == errUpToDate {
			log.Noticef(err.Error())
			err = nil
		}
	}()

	// Determine the version to upgrade to, uploading tools if necessary.
	env := conn.Environ
	cfg, err := conn.State.EnvironConfig()
	if err != nil {
		return err
	}
	v, err := c.initVersions(cfg, env)
	if err != nil {
		return err
	}
	if c.UploadTools {
		series := getUploadSeries(cfg, c.Series)
		if err := v.uploadTools(env.Storage(), series); err != nil {
			return err
		}
	}
	if err := v.validate(); err != nil {
		return err
	}
	log.Infof("upgrade version chosen: %s", v.chosen)
	// TODO(fwereade): this list may be incomplete, pending envtools.Upload change.
	log.Infof("available tools: %s", v.tools)

	if err := conn.State.SetEnvironAgentVersion(v.chosen); err != nil {
		return err
	}
	log.Noticef("started upgrade to %s", v.chosen)
	return nil
}

// initVersions collects state relevant to an upgrade decision. The returned
// agent and client versions, and the list of currently available tools, will
// always be accurate; the chosen version, and the flag indicating development
// mode, may remain blank until uploadTools or validate is called.
func (c *UpgradeJujuCommand) initVersions(cfg *config.Config, env environs.Environ) (*upgradeVersions, error) {
	agent, ok := cfg.AgentVersion()
	if !ok {
		// Can't happen. In theory.
		return nil, fmt.Errorf("incomplete environment configuration")
	}
	if c.Version == agent {
		return nil, errUpToDate
	}
	client := version.Current.Number
	// TODO use an API call rather than requiring the environment,
	// so that we can restrict access to the provider secrets
	// while still allowing users to upgrade.
	available, err := envtools.FindTools(env, client.Major, -1, coretools.Filter{}, envtools.DoNotAllowRetry)
	if err != nil {
		if !errors.IsNotFoundError(err) {
			return nil, err
		}
		if !c.UploadTools {
			// No tools found and we shouldn't upload any, so pretend
			// there is no more recent version available.
			if c.Version == version.Zero {
				return nil, errUpToDate
			}
			return nil, err
		}
	}
	return &upgradeVersions{
		agent:  agent,
		client: client,
		chosen: c.Version,
		tools:  available,
	}, nil
}

// upgradeVersions holds the version information for making upgrade decisions.
type upgradeVersions struct {
	agent  version.Number
	client version.Number
	chosen version.Number
	tools  coretools.List
}

// uploadTools compiles jujud from $GOPATH and uploads it into the supplied
// storage. If no version has been explicitly chosen, the version number
// reported by the built tools will be based on the client version number.
// In any case, the version number reported will have a build component higher
// than that of any otherwise-matching available envtools.
// uploadTools resets the chosen version and replaces the available tools
// with the ones just uploaded.
func (v *upgradeVersions) uploadTools(storage storage.Storage, series []string) error {
	// TODO(fwereade): this is kinda crack: we should not assume that
	// version.Current matches whatever source happens to be built. The
	// ideal would be:
	//  1) compile jujud from $GOPATH into some build dir
	//  2) get actual version with `jujud version`
	//  3) check actual version for compatibility with CLI tools
	//  4) generate unique build version with reference to available tools
	//  5) force-version that unique version into the dir directly
	//  6) archive and upload the build dir
	// ...but there's no way we have time for that now. In the meantime,
	// considering the use cases, this should work well enough; but it
	// won't detect an incompatible major-version change, which is a shame.
	if v.chosen == version.Zero {
		v.chosen = v.client
	}
	v.chosen = uploadVersion(v.chosen, v.tools)

	// TODO(fwereade): envtools.Upload should return envtools.List, and should
	// include all the extra series we build, so we can set *that* onto
	// v.available and maybe one day be able to check that a given upgrade
	// won't leave out-of-date machines lying around, starved of envtools.
	uploaded, err := sync.Upload(storage, &v.chosen, series...)
	if err != nil {
		return err
	}
	v.tools = coretools.List{uploaded}
	return nil
}

// validate chooses an upgrade version, if one has not already been chosen,
// and ensures the tools list contains no entries that do not have that version.
// If validate returns no error, the environment agent-version can be set to
// the value of the chosen field.
func (v *upgradeVersions) validate() (err error) {
	if v.chosen == version.Zero {
		// No explicitly specified version, so find the next available
		// stable release to upgrade to, starting from the current agent
		// version and doing major.minor+1 or +2 as needed.
		nextStable := v.agent
		if v.agent.IsDev() {
			nextStable.Minor += 1
		} else {
			nextStable.Minor += 2
		}

		newestNextStable, found := v.tools.NewestCompatible(nextStable)
		if found {
			log.Debugf("found a more recent stable version %s", newestNextStable)
			v.chosen = newestNextStable
		} else {
			newestCurrent, found := v.tools.NewestCompatible(v.agent)
			if found {
				log.Debugf("found more recent current version %s", newestCurrent)
				v.chosen = newestCurrent
			} else {
				return fmt.Errorf("no more recent supported versions available")
			}
		}
	} else {
		// If not completely specified already, pick a single tools version.
		filter := coretools.Filter{Number: v.chosen, Released: !v.chosen.IsDev()}
		if v.tools, err = v.tools.Match(filter); err != nil {
			return err
		}
		v.chosen, v.tools = v.tools.Newest()
	}
	if v.chosen == v.agent {
		return errUpToDate
	}

	// Major version upgrade
	if v.chosen.Major < v.agent.Major {
		// TODO(fwereade): I'm a bit concerned about old agent/CLI tools even
		// *connecting* to environments with higher agent-versions; but ofc they
		// have to connect in order to discover they shouldn't. However, once
		// any of our tools detect an incompatible version, they should act to
		// minimize damage: the CLI should abort politely, and the agents should
		// run an Upgrader but no other tasks.
		return fmt.Errorf("cannot change major version from %d to %d", v.agent.Major, v.chosen.Major)
	} else if v.chosen.Major > v.agent.Major {
		return fmt.Errorf("major version upgrades are not supported yet")
	}

	return nil
}

// uploadVersion returns a copy of the supplied version with a build number
// higher than any of the supplied tools that share its major, minor and patch.
func uploadVersion(vers version.Number, existing coretools.List) version.Number {
	vers.Build++
	for _, t := range existing {
		if t.Version.Major != vers.Major || t.Version.Minor != vers.Minor || t.Version.Patch != vers.Patch {
			continue
		}
		if t.Version.Build >= vers.Build {
			vers.Build = t.Version.Build + 1
		}
	}
	return vers
}
<|endoftext|>