From 5a25adf84f126baca6f54af22630f9a2d52d3963 Mon Sep 17 00:00:00 2001 From: miraclesu Date: Thu, 11 May 2017 18:08:25 +0800 Subject: [PATCH] =?UTF-8?q?pkg:=20=E6=9B=BF=E6=8D=A2=20mgo=20=E5=8C=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- conf/conf.go | 9 ++++- conf/files/db.json.sample | 8 +++- models/common.go | 6 ++- models/db/mgo.go | 81 +++++++++++++++++++++++++++++++++++++++ models/db/mid/auto_inc.go | 60 +++++++++++++++++++++++++++++ models/job_log.go | 6 +-- models/mdb.go | 6 +-- 7 files changed, 164 insertions(+), 12 deletions(-) create mode 100644 models/db/mgo.go create mode 100644 models/db/mid/auto_inc.go diff --git a/conf/conf.go b/conf/conf.go index 19c2068..54619e9 100644 --- a/conf/conf.go +++ b/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) diff --git a/conf/files/db.json.sample b/conf/files/db.json.sample index f8326dc..acb302e 100644 --- a/conf/files/db.json.sample +++ b/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 } \ No newline at end of file diff --git a/models/common.go b/models/common.go index 84759a5..7703d60 100644 --- a/models/common.go +++ b/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 diff --git a/models/db/mgo.go b/models/db/mgo.go new file mode 100644 index 0000000..63c7d0a --- /dev/null +++ b/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) + }) +} diff --git a/models/db/mid/auto_inc.go b/models/db/mid/auto_inc.go new file mode 100644 index 0000000..70f8c62 --- /dev/null +++ b/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 +} diff --git a/models/job_log.go b/models/job_log.go index c8ae7c3..179c118 100644 --- a/models/job_log.go +++ b/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 } diff --git a/models/mdb.go b/models/mdb.go index 3c9191b..a318aee 100644 --- a/models/mdb.go +++ b/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 }