Create common db package. Move original message to fetch to db.

cmd/dumporiginal/dumporiginal.go

@@ -1,47 +1,19 @@
 package main
 
 import (
-	"database/sql"
 	"flag"
 	"fmt"
 	"os"
+	"xinu.tv/email/db"
 
 	"github.com/golang/glog"
-	_ "github.com/lib/pq"
 )
 
-func fetchMessage(db *sql.DB, hash string) error {
-	rows, err := db.Query(`
-SELECT
-    blob
-FROM
-    original
-WHERE
-    hash = $1
-`, hash)
-	if err != nil {
-		return err
-	}
-	var b []byte
-
-	for rows.Next() {
-		if err := rows.Scan(&b); err != nil {
-			return err
-		}
-		fmt.Println(string(b))
-	}
-	if err := rows.Err(); err != nil {
-		return err
-	}
-	return nil
-}
-
 func main() {
 	defer glog.Flush()
 	flag.Parse()
 
-	// TODO(wathiede): make a set of flags.
+	c, err := db.NewConn("")
-	db, err := sql.Open("postgres", "user=gomail dbname=gomail sslmode=disable")
 	if err != nil {
 		glog.Fatal(err)
 	}
@@ -51,9 +23,11 @@ func main() {
 		os.Exit(1)
 	}
 
+	var blob []byte
 	for _, hash := range flag.Args() {
-		if err := fetchMessage(db, hash); err != nil {
+		if err := c.OriginalBlobByHash(hash, &blob); err != nil {
 			glog.Fatal(err)
 		}
+		fmt.Println(string(blob))
 	}
 }

db/util.go

@@ -0,0 +1,43 @@
+package db
+
+import (
+	"database/sql"
+	"flag"
+	"fmt"
+
+	_ "github.com/lib/pq"
+)
+
+var (
+	user    = flag.String("dbuser", "gomail", "PostgreSQL user name")
+	name    = flag.String("dbname", "gomail", "PostgreSQL DB name")
+	sslmode = flag.String("dbsslmode", "disable", "PostgreSQL sslmode setting")
+)
+
+type conn struct {
+	*sql.DB
+}
+
+// NewDB creates connection to PostgreSQL DB.  If dsn is empty, the flags
+// -dbuser, -dbname, -dbsslmode are used.
+func NewConn(dsn string) (*conn, error) {
+	if dsn == "" {
+		dsn = fmt.Sprintf("user=%s dbname=%s sslmode=%s", *user, *name,
+			*sslmode)
+	}
+	db, err := sql.Open("postgres", dsn)
+	return &conn{DB: db}, err
+}
+
+func (c *conn) OriginalBlobByHash(hash string, blob *[]byte) error {
+	row := c.QueryRow(`
+SELECT
+    blob
+FROM
+    original
+WHERE
+    hash = $1
+`, hash)
+
+	return row.Scan(blob)
+}