From 6fdecf81e3be395879224f474a917541d661bb42 Mon Sep 17 00:00:00 2001 From: Ian Kent Date: Sun, 20 Apr 2014 20:33:42 +0100 Subject: [PATCH] Clean up mongodb --- mailhog/http/api/v1.go | 12 ++++--- mailhog/http/server.go | 5 +-- mailhog/smtp/session.go | 7 ++-- mailhog/storage/mongodb.go | 65 ++++++++++++++++---------------------- main.go | 8 +++-- main_test.go | 3 +- 6 files changed, 49 insertions(+), 51 deletions(-) diff --git a/mailhog/http/api/v1.go b/mailhog/http/api/v1.go index e5d6e93..93de6f3 100644 --- a/mailhog/http/api/v1.go +++ b/mailhog/http/api/v1.go @@ -14,14 +14,16 @@ type APIv1 struct { config *mailhog.Config exitChannel chan int server *http.Server + mongo *storage.MongoDB } -func CreateAPIv1(exitCh chan int, conf *mailhog.Config, server *http.Server) *APIv1 { +func CreateAPIv1(exitCh chan int, conf *mailhog.Config, server *http.Server, mongo *storage.MongoDB) *APIv1 { log.Println("Creating API v1") apiv1 := &APIv1{ config: conf, exitChannel: exitCh, server: server, + mongo: mongo, } server.Handler.(*handler.RegexpHandler).HandleFunc(regexp.MustCompile("^/api/v1/messages/?$"), apiv1.messages) @@ -36,7 +38,7 @@ func (apiv1 *APIv1) messages(w http.ResponseWriter, r *http.Request, route *hand log.Println("[APIv1] GET /api/v1/messages") // TODO start, limit - messages, _ := storage.List(apiv1.config, 0, 1000) + messages, _ := apiv1.mongo.List(0, 1000) bytes, _ := json.Marshal(messages) w.Header().Set("Content-Type", "text/json") w.Write(bytes) @@ -47,7 +49,7 @@ func (apiv1 *APIv1) message(w http.ResponseWriter, r *http.Request, route *handl id := match[1] log.Printf("[APIv1] GET /api/v1/messages/%s\n", id) - message, _ := storage.Load(apiv1.config, id) + message, _ := apiv1.mongo.Load(id) bytes, _ := json.Marshal(message) w.Header().Set("Content-Type", "text/json") w.Write(bytes) @@ -57,7 +59,7 @@ func (apiv1 *APIv1) delete_all(w http.ResponseWriter, r *http.Request, route *ha log.Println("[APIv1] POST /api/v1/messages/delete") w.Header().Set("Content-Type", "text/json") - storage.DeleteAll(apiv1.config) + apiv1.mongo.DeleteAll() } func (apiv1 *APIv1) delete_one(w http.ResponseWriter, r *http.Request, route *handler.Route) { @@ -66,5 +68,5 @@ func (apiv1 *APIv1) delete_one(w http.ResponseWriter, r *http.Request, route *ha log.Printf("[APIv1] POST /api/v1/messages/%s/delete\n", id) w.Header().Set("Content-Type", "text/json") - storage.DeleteOne(apiv1.config, id) + apiv1.mongo.DeleteOne(id) } diff --git a/mailhog/http/server.go b/mailhog/http/server.go index 78c3dd9..d08c232 100644 --- a/mailhog/http/server.go +++ b/mailhog/http/server.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/ian-kent/MailHog/mailhog" "github.com/ian-kent/MailHog/mailhog/templates" + "github.com/ian-kent/MailHog/mailhog/storage" "github.com/ian-kent/MailHog/mailhog/templates/images" "github.com/ian-kent/MailHog/mailhog/templates/js" "github.com/ian-kent/MailHog/mailhog/http/api" @@ -43,7 +44,7 @@ func web_headers(w http.ResponseWriter) { w.Header().Set("Content-Type", "text/html") } -func Start(exitCh chan int, conf *mailhog.Config) { +func Start(exitCh chan int, conf *mailhog.Config, mongo *storage.MongoDB) { exitChannel = exitCh config = conf @@ -57,7 +58,7 @@ func Start(exitCh chan int, conf *mailhog.Config) { server.Handler.(*handler.RegexpHandler).HandleFunc(regexp.MustCompile("^/images/hog.png$"), web_imgcontroller) server.Handler.(*handler.RegexpHandler).HandleFunc(regexp.MustCompile("^/$"), web_index) - api.CreateAPIv1(exitCh, conf, server) + api.CreateAPIv1(exitCh, conf, server, mongo) server.ListenAndServe() } diff --git a/mailhog/smtp/session.go b/mailhog/smtp/session.go index fed3a56..d018490 100644 --- a/mailhog/smtp/session.go +++ b/mailhog/smtp/session.go @@ -18,6 +18,7 @@ type Session struct { conf *mailhog.Config state int message *data.SMTPMessage + mongo *storage.MongoDB } const ( @@ -30,8 +31,8 @@ const ( // TODO replace ".." lines with . in data -func StartSession(conn *net.TCPConn, conf *mailhog.Config) { - conv := &Session{conn, "", conf, ESTABLISH, &data.SMTPMessage{}} +func StartSession(conn *net.TCPConn, conf *mailhog.Config, mongo *storage.MongoDB) { + conv := &Session{conn, "", conf, ESTABLISH, &data.SMTPMessage{}, mongo} conv.log("Starting session") conv.Write("220", conv.conf.Hostname + " ESMTP Go-MailHog") conv.Read() @@ -76,7 +77,7 @@ func (c *Session) Parse() { c.message.Data += parts[0] + "\n" if(strings.HasSuffix(c.message.Data, "\r\n.\r\n")) { c.message.Data = strings.TrimSuffix(c.message.Data, "\r\n.\r\n") - id, err := storage.Store(c.conf, c.message) + id, err := c.mongo.Store(c.message) c.state = DONE if err != nil { // FIXME diff --git a/mailhog/storage/mongodb.go b/mailhog/storage/mongodb.go index 296b7ee..8261717 100644 --- a/mailhog/storage/mongodb.go +++ b/mailhog/storage/mongodb.go @@ -8,15 +8,28 @@ import ( "github.com/ian-kent/MailHog/mailhog" ) -func Store(c *mailhog.Config, m *data.SMTPMessage) (string, error) { - msg := data.ParseSMTPMessage(c, m) +type MongoDB struct { + Session *mgo.Session + Config *mailhog.Config + Collection *mgo.Collection +} + +func CreateMongoDB(c *mailhog.Config) *MongoDB { session, err := mgo.Dial(c.MongoUri) if(err != nil) { - log.Printf("Error connecting to MongoDB: %s", err) - return "", err + log.Fatalf("Error connecting to MongoDB: %s", err) + return nil } - defer session.Close() - err = session.DB(c.MongoDb).C(c.MongoColl).Insert(msg) + return &MongoDB{ + Session: session, + Config: c, + Collection: session.DB(c.MongoDb).C(c.MongoColl), + } +} + +func (mongo *MongoDB) Store(m *data.SMTPMessage) (string, error) { + msg := data.ParseSMTPMessage(mongo.Config, m) + err := mongo.Collection.Insert(msg) if err != nil { log.Printf("Error inserting message: %s", err) return "", err @@ -24,15 +37,9 @@ func Store(c *mailhog.Config, m *data.SMTPMessage) (string, error) { return msg.Id, nil } -func List(c *mailhog.Config, start int, limit int) (*data.Messages, error) { - session, err := mgo.Dial(c.MongoUri) - if(err != nil) { - log.Printf("Error connecting to MongoDB: %s", err) - return nil, err - } - defer session.Close() +func (mongo *MongoDB) List(start int, limit int) (*data.Messages, error) { messages := &data.Messages{} - err = session.DB(c.MongoDb).C(c.MongoColl).Find(bson.M{}).Skip(start).Limit(limit).All(messages) + err := mongo.Collection.Find(bson.M{}).Skip(start).Limit(limit).All(messages) if err != nil { log.Printf("Error loading messages: %s", err) return nil, err @@ -40,37 +47,19 @@ func List(c *mailhog.Config, start int, limit int) (*data.Messages, error) { return messages, nil; } -func DeleteOne(c *mailhog.Config, id string) error { - session, err := mgo.Dial(c.MongoUri) - if(err != nil) { - log.Printf("Error connecting to MongoDB: %s", err) - return err - } - defer session.Close() - _, err = session.DB(c.MongoDb).C(c.MongoColl).RemoveAll(bson.M{"id": id}) +func (mongo *MongoDB) DeleteOne(id string) error { + _, err := mongo.Collection.RemoveAll(bson.M{"id": id}) return err } -func DeleteAll(c *mailhog.Config) error { - session, err := mgo.Dial(c.MongoUri) - if(err != nil) { - log.Printf("Error connecting to MongoDB: %s", err) - return err - } - defer session.Close() - _, err = session.DB(c.MongoDb).C(c.MongoColl).RemoveAll(bson.M{}) +func (mongo *MongoDB) DeleteAll() error { + _, err := mongo.Collection.RemoveAll(bson.M{}) return err } -func Load(c *mailhog.Config, id string) (*data.Message, error) { - session, err := mgo.Dial(c.MongoUri) - if(err != nil) { - log.Printf("Error connecting to MongoDB: %s", err) - return nil, err - } - defer session.Close() +func (mongo *MongoDB) Load(id string) (*data.Message, error) { result := &data.Message{} - err = session.DB(c.MongoDb).C(c.MongoColl).Find(bson.M{"id": id}).One(&result) + err := mongo.Collection.Find(bson.M{"id": id}).One(&result) if err != nil { log.Printf("Error loading message: %s", err) return nil, err diff --git a/main.go b/main.go index a7ff668..774f155 100644 --- a/main.go +++ b/main.go @@ -5,12 +5,14 @@ import ( "github.com/ian-kent/MailHog/mailhog" "github.com/ian-kent/MailHog/mailhog/http" "github.com/ian-kent/MailHog/mailhog/smtp" + "github.com/ian-kent/MailHog/mailhog/storage" "log" "net" "os" ) var conf *mailhog.Config +var mongo *storage.MongoDB var exitCh chan int func config() { @@ -33,6 +35,8 @@ func config() { MongoDb: mongodb, MongoColl: mongocoll, } + + mongo = storage.CreateMongoDB(conf) } func main() { @@ -53,7 +57,7 @@ func main() { func web_listen() { log.Printf("[HTTP] Binding to address: %s\n", conf.HTTPBindAddr) - http.Start(exitCh, conf) + http.Start(exitCh, conf, mongo) } func smtp_listen() *net.TCPListener { @@ -72,6 +76,6 @@ func smtp_listen() *net.TCPListener { } defer conn.Close() - go smtp.StartSession(conn.(*net.TCPConn), conf) + go smtp.StartSession(conn.(*net.TCPConn), conf, mongo) } } diff --git a/main_test.go b/main_test.go index 6b12526..78ab117 100644 --- a/main_test.go +++ b/main_test.go @@ -88,7 +88,8 @@ func TestBasicHappyPath(t *testing.T) { assert.Nil(t, err) assert.Equal(t, string(buf[0:n]), "221 Bye\n") - message, err := storage.Load(mailhog.DefaultConfig(), match[1]) + s := storage.CreateMongoDB(mailhog.DefaultConfig()) + message, err := s.Load(match[1]) assert.Nil(t, err) assert.NotNil(t, message)