pkg: 替换 mgo 包

conf/conf.go

@@ -9,10 +9,10 @@ import (
 	"github.com/fsnotify/fsnotify"
 	"github.com/go-gomail/gomail"
 
-	"sunteng/commons/db/imgo"
 	"sunteng/commons/log"
 
 	"github.com/shunfei/cronsun/event"
+	"github.com/shunfei/cronsun/models/db"
 	"github.com/shunfei/cronsun/utils"
 )
 
@@ -66,7 +66,7 @@ type Conf struct {
 
 	Log  *log.Config
 	Etcd client.Config
-	Mgo  *imgo.Config
+	Mgo  *db.Config
 	Web  webConfig
 	Mail *MailConf
 
@@ -131,6 +131,11 @@ func (c *Conf) parse() error {
 	if c.Mail.Keepalive <= 0 {
 		c.Mail.Keepalive = 30
 	}
+	if c.Mgo.Timeout <= 0 {
+		c.Mgo.Timeout = 10 * time.Second
+	} else {
+		c.Mgo.Timeout *= time.Second
+	}
 	log.InitConf(c.Log)
 
 	c.Node = cleanKeyPrefix(c.Node)

conf/files/db.json.sample

@@ -1,4 +1,8 @@
 {
-  "Host": "192.168.11.16:3000",
-  "Database": "cronsun"
+  "Hosts": [
+    "192.168.11.16:3000"
+  ],
+  "Database": "cronsun",
+  "#Timeout": "connect timeout duration/second",
+  "Timeout": 15
 }

models/common.go

@@ -1,8 +1,8 @@
 package models
 
 import (
-	"sunteng/commons/db/imgo"
 	"github.com/shunfei/cronsun/conf"
+	"github.com/shunfei/cronsun/models/db"
 )
 
 var (
@@ -30,7 +30,9 @@ func Init() (err error) {
 	}
 
 	// init mongoDB
-	mgoDB = imgo.NewMdbWithConf(conf.Config.Mgo)
+	if mgoDB, err = db.NewMdb(conf.Config.Mgo); err != nil {
+		return
+	}
 
 	initialized = true
 	return

models/db/mgo.go

@@ -0,0 +1,81 @@
+package db
+
+import (
+	"strings"
+	"time"
+
+	"gopkg.in/mgo.v2"
+	"gopkg.in/mgo.v2/bson"
+)
+
+type Config struct {
+	Hosts    []string
+	UserName string
+	Password string
+	Database string
+	Timeout  time.Duration // second
+}
+
+type Mdb struct {
+	*Config
+	*mgo.Session
+}
+
+func NewMdb(c *Config) (*Mdb, error) {
+	m := &Mdb{
+		Config: c,
+	}
+	return m, m.connect()
+}
+
+func (m *Mdb) connect() error {
+	// url: [mongodb://][user:pass@]host1[:port1][,host2[:port2],...][/database][?options]
+	url := strings.Join(m.Config.Hosts, ",")
+	if len(m.Config.UserName) > 0 && len(m.Config.Password) > 0 {
+		url = m.Config.UserName + ":" + m.Config.Password + "@" + url
+	}
+
+	if len(m.Config.Database) > 0 {
+		url += "/" + m.Config.Database
+	}
+
+	session, err := mgo.DialWithTimeout(url, m.Config.Timeout)
+	if err != nil {
+		return err
+	}
+
+	m.Session = session
+	return nil
+}
+
+func (m *Mdb) WithC(collection string, job func(*mgo.Collection) error) error {
+	s := m.Session.New()
+	err := job(s.DB(m.Config.Database).C(collection))
+	s.Close()
+	return err
+}
+
+func (self *Mdb) Upsert(collection string, selector interface{}, change interface{}) error {
+	return self.WithC(collection, func(c *mgo.Collection) error {
+		_, err := c.Upsert(selector, change)
+		return err
+	})
+}
+
+func (self *Mdb) Insert(collection string, data ...interface{}) error {
+	return self.WithC(collection, func(c *mgo.Collection) error {
+		return c.Insert(data...)
+	})
+}
+
+func (self *Mdb) FindId(collection string, id interface{}, result interface{}) error {
+	return self.WithC(collection, func(c *mgo.Collection) error {
+		return c.Find(bson.M{"_id": id}).One(result)
+	})
+}
+
+func (self *Mdb) FindOne(collection string, query interface{}, result interface{}) error {
+	return self.WithC(collection, func(c *mgo.Collection) error {
+		return c.Find(query).One(result)
+	})
+}

models/db/mid/auto_inc.go

@@ -0,0 +1,60 @@
+package mid
+
+import (
+	"fmt"
+
+	"gopkg.in/mgo.v2"
+	"gopkg.in/mgo.v2/bson"
+)
+
+var (
+	field = &Field{
+		Id:         "seq",
+		Collection: "_id",
+	}
+)
+
+type Field struct {
+	Id         string
+	Collection string
+}
+
+// 如果不设置，则用默认设置
+func SetFieldName(id, collection string) {
+	field.Id = id
+	field.Collection = collection
+}
+
+//使collection 为 name 的 id 自增 1 并返回当前 id 的值
+func AutoInc(c *mgo.Collection, name string) (id int, err error) {
+	return incr(c, name, 1)
+}
+
+//批量申请一段id
+func ApplyBatchIds(c *mgo.Collection, name string, amount int) (id int, err error) {
+	return incr(c, name, amount)
+}
+
+func incr(c *mgo.Collection, name string, step int) (id int, err error) {
+	result := make(map[string]interface{})
+	change := mgo.Change{
+		Update:    bson.M{"$inc": bson.M{field.Id: step}},
+		Upsert:    true,
+		ReturnNew: true,
+	}
+	_, err = c.Find(bson.M{field.Collection: name}).Apply(change, result)
+	if err != nil {
+		return
+	}
+	id, ok := result[field.Id].(int)
+	if ok {
+		return
+	}
+	id64, ok := result[field.Id].(int64)
+	if !ok {
+		err = fmt.Errorf("%s is ont int or int64", field.Id)
+		return
+	}
+	id = int(id64)
+	return
+}

models/job_log.go

@@ -3,7 +3,7 @@ package models
 import (
 	"time"
 
-	mgo "gopkg.in/mgo.v2"
+	"gopkg.in/mgo.v2"
 	"gopkg.in/mgo.v2/bson"
 
 	"sunteng/commons/log"
@@ -139,11 +139,11 @@ type StatExecuted struct {
 }
 
 func JobLogStat() (s *StatExecuted, err error) {
-	err = mgoDB.One(Coll_Stat, bson.M{"name": "job"}, &s)
+	err = mgoDB.FindOne(Coll_Stat, bson.M{"name": "job"}, &s)
 	return
 }
 
 func JobLogDayStat(day time.Time) (s *StatExecuted, err error) {
-	err = mgoDB.One(Coll_Stat, bson.M{"name": "job-day", "date": day.Format("2006-01-02")}, &s)
+	err = mgoDB.FindOne(Coll_Stat, bson.M{"name": "job-day", "date": day.Format("2006-01-02")}, &s)
 	return
 }

models/mdb.go

@@ -1,13 +1,13 @@
 package models
 
 import (
-	"sunteng/commons/db/imgo"
+	"github.com/shunfei/cronsun/models/db"
 )
 
 var (
-	mgoDB *imgo.Mdb
+	mgoDB *db.Mdb
 )
 
-func GetDb() *imgo.Mdb {
+func GetDb() *db.Mdb {
 	return mgoDB
 }