Clean up mongodb

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)
 }

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()
 }

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) {
+func StartSession(conn *net.TCPConn, conf *mailhog.Config, mongo *storage.MongoDB) {
-	conv := &Session{conn, "", conf, ESTABLISH, &data.SMTPMessage{}}
+	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

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) {
+type MongoDB struct {
-	msg := data.ParseSMTPMessage(c, m)
+	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)
+		log.Fatalf("Error connecting to MongoDB: %s", err)
-		return "", err
+		return nil
 	}
-	defer session.Close()
+	return &MongoDB{
-	err = session.DB(c.MongoDb).C(c.MongoColl).Insert(msg)
+		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) {
+func (mongo *MongoDB) List(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()
 	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 {
+func (mongo *MongoDB) DeleteOne(id string) error {
-	session, err := mgo.Dial(c.MongoUri)
+	_, err := mongo.Collection.RemoveAll(bson.M{"id": id})
-	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})
 	return err
 }
 
-func DeleteAll(c *mailhog.Config) error {
+func (mongo *MongoDB) DeleteAll() error {
-	session, err := mgo.Dial(c.MongoUri)
+	_, err := mongo.Collection.RemoveAll(bson.M{})
-	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{})
 	return err
 }
 
-func Load(c *mailhog.Config, id string) (*data.Message, error) {
+func (mongo *MongoDB) Load(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()
 	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

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)
 	}
 }

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)