// Save the DataModel to DataStore
func SaveModel(mgo_db string, m DataModel, conn *mgo.Session) (err error) {
if conn == nil {
conn, err = MgoConnGet(mgo_db)
if err != nil {
return
}
defer conn.Close()
}
if conn != nil {
bsonMid := m.MidGet()
c := conn.DB(mgo_db).C(m.Type())
//Debug("SAVING ", mgo_db, " type=", m.Type(), " Mid=", bsonMid)
if len(bsonMid) < 5 {
m.MidSet(bson.NewObjectId())
if err = c.Insert(m); err != nil {
Log(ERROR, "MGOU ERROR on insert ", err, " TYPE=", m.Type(), " ", m.MidGet())
} else {
//Log(DEBUG, "successfully inserted!!!!!! ", m.MidGet(), " oid=", m.OidGet())
}
} else {
// YOU MUST NOT SEND Mid "_id" to Mongo
mid := m.MidGet()
m.MidSet("") // omitempty means it doesn't get sent
if err = c.Update(bson.M{"_id": bson.ObjectId(bsonMid)}, m); err != nil {
Log(ERROR, "MGOU ERROR on update ", err, " ", bsonMid, " MID=?", m.MidGet())
}
m.MidSet(mid)
}
} else {
Log(ERROR, "MGOU Nil connection")
return errors.New("no db connection")
}
return
}
// Save the DataModel to DataStore
func Insert(mgo_db string, m DataModel, conn *mgo.Session) (err error) {
if conn == nil {
conn, err = MgoConnGet(mgo_db)
if err != nil {
return
}
defer conn.Close()
}
if conn != nil {
c := conn.DB(mgo_db).C(m.Type())
if len(m.MidGet()) == 0 {
m.MidSet(bson.NewObjectId())
}
if err = c.Insert(m); err != nil {
Log(ERROR, "MGOU ERROR on insert ", err, " TYPE=", m.Type(), " ", m.MidGet())
} else {
//Log(DEBUG, "successfully inserted!!!!!! ", m.MidGet(), " oid=", m.OidGet())
}
} else {
Log(ERROR, "MGOU Nil connection")
return errors.New("no db connection")
}
return
}
func newServer() (*coretesting.MgoInstance, error) {
inst := &coretesting.MgoInstance{Params: []string{"--replSet", name}}
err := inst.Start(true)
if err != nil {
return nil, fmt.Errorf("Error starting mongo server: %s", err.Error())
}
// by dialing right now, we'll wait until it's running
strategy := utils.AttemptStrategy{Total: time.Second * 5, Delay: time.Millisecond * 100}
attempt := strategy.Start()
for attempt.Next() {
var session *mgo.Session
session, err = inst.DialDirect()
if err != nil {
err = fmt.Errorf("Error dialing mongo server %q: %s", inst.Addr(), err.Error())
} else {
session.SetMode(mgo.Monotonic, true)
err = session.Ping()
if err != nil {
err = fmt.Errorf("Error pinging mongo server %q: %s", inst.Addr(), err.Error())
}
session.Close()
}
if err == nil || !attempt.HasNext() {
break
}
}
return inst, err
}
// CloseSession puts the connection back into the pool
func CloseSession(sessionId string, mongoSession *mgo.Session) {
defer helper.CatchPanic(nil, sessionId, "CloseSession")
tracelog.Started(sessionId, "CloseSession")
mongoSession.Close()
tracelog.Completed(sessionId, "CloseSession")
}
func CloseSession(s *mgo.Session) {
defer func() {
if err := recover(); err != nil {
log.Print("[MGO2_CLOSE_SESSION_RECOVER] close session panic", err)
}
}()
s.Close()
}
// TODO remove dead session from freecon
func (this *Client) checkServerStatus(wg *sync.WaitGroup, sess *mgo.Session) {
defer wg.Done()
err := sess.Ping()
if err != nil {
log.Error("mongodb err: %v %s", sess.LiveServers(), err)
sess.Close()
}
}
// CloseSession puts the connection back into the pool
func CloseSession(sessionId string, mongoSession *mgo.Session) {
defer helper.CatchPanic(nil, sessionId, "CloseSession")
tracelog.STARTED(sessionId, "CloseSession")
mongoSession.Close()
tracelog.COMPLETED(sessionId, "CloseSession")
}
// insertworker calls func insertdoc for specified insert count
func insertworker(s *mgo.Session, ch chan int) {
c := s.DB("test").C("mongotest")
defer s.Close()
for i := 0; i < Insert_Count; i++ {
err := insertdoc(c, i)
if err != nil {
panic(err)
}
}
ch <- 1
}
func (this *Client) checkServerStatus(wg *sync.WaitGroup, sess *mgo.Session) {
defer wg.Done()
err := sess.Ping()
if err != nil {
// TODO show mongodb uri in log
log.Error("mongodb killed for: %s", err)
sess.Close()
this.killConn(sess)
}
}
func DeleteTodos(resp http.ResponseWriter, req *http.Request,
b *Broker, session *mgo.Session) {
defer session.Close()
var todos []string
if ReadJson(req.Body, &todos) == nil {
collection := TodoCollection(session)
collection.RemoveAll(bson.M{"guid": bson.M{"$in": todos}})
for _, todo := range todos {
message := &Message{"delete", todo}
b.messages <- message
}
}
}
func SaveTodo(resp http.ResponseWriter, req *http.Request,
b *Broker, session *mgo.Session) {
defer session.Close()
var todo = new(Todo)
if err := todo.ReadJson(req.Body); err != nil {
http.Error(resp, "Bad Request", http.StatusBadRequest)
} else {
collection := TodoCollection(session)
collection.Upsert(bson.M{"guid": todo.Guid}, todo)
message := &Message{"update", todo.ToJsonString()}
b.messages <- message
}
}
func (this *Client) putFreeConn(uri string, sess *mgo.Session) {
this.lk.Lock()
defer this.lk.Unlock()
if this.freeconns == nil {
this.freeconns = make(map[string][]*mgo.Session)
}
freelist := this.freeconns[uri]
if len(freelist) >= this.conf.MaxIdleConnsPerServer {
sess.Close()
return
}
this.freeconns[uri] = append(this.freeconns[uri], sess)
}
func QueueTodos(session *mgo.Session, messageChan chan *Message) {
var todos []Todo
defer session.Close()
collection := TodoCollection(session)
iter := collection.Find(nil).Iter()
err := iter.All(&todos)
if err != nil {
return
}
for _, todo := range todos {
message := &Message{"update", todo.ToJsonString()}
messageChan <- message
}
}
// Serve sshkeys at http://:33845/{boxname}
func main() {
defer func() {
if err := recover(); err != nil {
panic(err)
}
}()
var session *mgo.Session
var err error
go func() {
for session == nil {
log.Println("Connecting to mongo..")
session, err = mgo.Dial(os.Getenv("CU_DB"))
if err != nil {
if session != nil {
session.Close()
session = nil
}
log.Printf("Database connection failed (%q), retrying..", err)
time.Sleep(10 * time.Second)
}
}
log.Println("Connected to mongo.")
}()
http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
u := r.URL.Path[1:]
if session == nil {
log.Printf("%v StatusServiceUnavailable session == nil", u)
w.WriteHeader(http.StatusServiceUnavailable)
return
}
boxUsers, allUsers, keys := getKeys(session, u)
log.Printf("%s:%v:%q:%q", u, strings.Count(keys, "\n"), boxUsers, allUsers)
fmt.Fprint(w, keys)
})
log.Println("Serving..")
err = http.ListenAndServe("localhost:33845", nil)
check(err)
}
func releaseMongo(session *mgo.Session, db *mgo.Database) {
mongo, _ := Application.Config["mongo"].(map[string]interface{})
keep_db, ok := mongo["keep_test_db"].(bool)
if !ok {
keep_db = false
}
if !keep_db {
fmt.Println("Testing finished... Dropping database")
db.DropDatabase()
} else {
fmt.Println("Keep test database is true, not dropping database after unit test")
}
session.Close()
}
func CloseSession(session *mgo.Session) bool {
session.Close()
return true
}
// Run is the block's main loop. Here we listen on the different channels we set up.
func (b *ToMongoDB) Run() {
var collectionname string
var dbname string
var collection *mgo.Collection
var session *mgo.Session
var host = ""
var err error
var batch = 0
var count = 0
var maxindex = 0
var list []interface{}
for {
select {
case msgI := <-b.inrule:
// set host string for MongoDB server
host, err = util.ParseRequiredString(msgI, "Host")
if err != nil {
b.Error(err.Error())
continue
}
// set database name
dbname, err = util.ParseRequiredString(msgI, "Database")
if err != nil {
b.Error(err.Error())
continue
}
// set collection name
collectionname, err = util.ParseRequiredString(msgI, "Collection")
if err != nil {
b.Error(err.Error())
continue
}
// set number of records to insert at a time
batch, err = util.ParseInt(msgI, "BatchSize")
if err != nil || batch < 0 {
b.Error(errors.New("Error parsing batch size....setting to 0"))
batch = 0
} else {
if batch > 1 {
list = make([]interface{}, batch, batch)
// set maxindex to 1 minus batch size
// use maxindex for looping everywhere for consistency
maxindex = batch - 1
}
}
// create MongoDB connection
session, err = mgo.Dial(host)
if err != nil {
// swallowing a panic from mgo here - streamtools must not die
b.Error(errors.New("Could not initiate connection with MongoDB service"))
continue
}
// use the specified DB and collection
collection = session.DB(dbname).C(collectionname)
case <-b.quit:
// close connection to MongoDB and quit
if session != nil {
session.Close()
}
return
case msg := <-b.in:
// deal with inbound data
if session != nil {
// mgo is so cool - it will check if the message can be serialized to valid bson.
// So, no need to do a json.Marshal on the inbound. just append to the list and
// batch insert
if maxindex >= 1 {
if count <= maxindex {
list[count] = msg
count = count + 1
}
if count == maxindex {
// insert batch if count reaches batch size
err = collection.Insert(list...)
if err != nil {
b.Error(err.Error())
}
// reset list and count
list = make([]interface{}, batch, batch)
count = 0
}
} else {
// mgo coolness again. No need to do a json.Marshal on the inbound.
err = collection.Insert(msg)
if err != nil {
b.Error(err.Error())
}
}
} else {
b.Error(errors.New("MongoDB connection not initated or lost. Please check your MongoDB server or block settings."))
}
case respChan := <-b.queryrule:
// deal with a query request
respChan <- map[string]interface{}{
"Collection": collectionname,
"Database": dbname,
"Host": host,
"BatchSize": batch,
}
}
}
}
func main() {
flag.Parse()
defer func() {
if err := recover(); err != nil {
panic(err)
}
}()
var session *mgo.Session
var err error
go func() {
for session == nil {
log.Println("Connecting to mongo..")
session, err = mgo.Dial(os.Getenv("CU_DB"))
if err != nil {
if session != nil {
session.Close()
session = nil
}
log.Printf("Database connection failed (%q), retrying..", err)
time.Sleep(10 * time.Second)
}
}
log.Println("Connected to mongo.")
}()
http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
// Query boxes
splitted := strings.SplitN(r.URL.Path[1:], "/", 2)
if len(splitted) < 2 {
w.WriteHeader(http.StatusBadRequest)
return
}
if session == nil {
log.Printf("%v StatusServiceUnavailable session == nil", splitted[0])
w.WriteHeader(http.StatusServiceUnavailable)
return
}
validTokenQuery := bson.M{
"$and": []bson.M{
bson.M{"name": splitted[0]},
bson.M{"boxJSON.publish_token": splitted[1]}}}
err := ErrDbTimeout
var n int
done := make(chan struct{})
go func() {
this_session := session.Copy()
defer this_session.Close()
db := this_session.DB("")
n, err = db.C("boxes").Find(validTokenQuery).Count()
close(done)
}()
select {
case <-time.After(5 * time.Second):
case <-done:
}
if err != nil {
log.Printf("%v StatusServiceUnavailable %q", splitted[0], err)
w.WriteHeader(http.StatusServiceUnavailable)
return
}
if n == 0 {
log.Printf("%v StatusForbidden", splitted[0])
w.WriteHeader(http.StatusForbidden)
return
}
w.WriteHeader(http.StatusOK)
})
println("Listening...")
err = http.ListenAndServe(":23423", nil)
check(err)
}
// Tail sends mongodb operations for the namespace on the specified channel.
// Interrupts tailing if exit chan closes.
func Tail(session *mgo.Session, ns string, initial bool, lastTs *Timestamp, opc chan<- *Operation, exit chan bool) error {
defer close(opc)
defer session.Close()
// Always do initial import in case a previous optime doesn't exist.
if lastTs == nil || int64(*lastTs) == 0 {
initial = true
}
if initial {
// If we are doing an intitial import, replace the oplog timestamp with the most current
// as it doesn't make sense to apply the same objects multiple times.
if ts, err := Optime(session); err != nil {
return err
} else {
lastTs = ts
}
nsParts := strings.Split(ns, ".")
if len(nsParts) != 2 {
return errors.New("Exected namespace provided as database.collection")
}
col := session.DB(nsParts[0]).C(nsParts[1])
iter := col.Find(nil).Iter()
initialDone := make(chan bool)
go func() {
log.Println("Doing initial import, this may take a while...")
var count uint64
for {
var result bson.M
if iter.Next(&result) {
select {
case opc <- &Operation{
Namespace: ns,
Op: Insert,
Object: result,
}:
count++
case <-exit:
break
}
} else {
break
}
}
log.Println("Initial import object count:", count)
close(initialDone)
}()
waitInitialSync:
for {
select {
case <-initialDone:
if err := iter.Close(); err != nil {
return err
}
break waitInitialSync
case <-exit:
log.Println("Initial import was interrupted")
err := iter.Close()
<-initialDone
return err
}
}
log.Println("Initial import has completed")
}
// Start tailing oplog
col := session.DB("local").C("oplog.rs")
log.Println("Resuming oplog from timestamp:", *lastTs)
log.Println("It could take a moment for MongoDB to scan through the oplog collection...")
query := bson.M{"ns": ns, "ts": bson.M{"$gt": *lastTs}}
// Start tailing, sorted by forward natural order by default in capped collections.
iter := col.Find(query).Tail(-1)
iterClosed := make(chan bool)
go func() {
for {
var result Operation
if iter.Next(&result) {
select {
case opc <- &result:
case <-exit:
break
}
} else {
break
}
}
close(iterClosed)
}()
// Block until we are supposed to exit, close the iterator when that happens
select {
case <-exit:
case <-iterClosed:
}
err := iter.Close()
// Make sure iterator has stoped pumping into opc since it will be closed on defered func
<-iterClosed
return err
}
