Add mode to load message from stdin.

If no maildir specified stdin is treated as a message. This program is useful
for loading email into the DB as procmail rule.
This commit is contained in:
Bill Thiede 2014-03-29 21:56:57 -07:00
parent 12199604c1
commit c266e04053

View File

@ -6,7 +6,6 @@ import (
"database/sql" "database/sql"
"expvar" "expvar"
"flag" "flag"
"fmt"
"io" "io"
"io/ioutil" "io/ioutil"
"os" "os"
@ -38,9 +37,38 @@ var (
CRCR = []byte("\n\n") CRCR = []byte("\n\n")
// Maximum bytes we parse when decoding // Maximum bytes we parse when decoding
// email message, 50MB. // email message, 50MB.
maxMessageSize int64 = 50 << 20 maxMessageSize int64 = 50 << 20
insertOriginalQuery = `INSERT INTO
original (uid, hash, header_size, total_size, blob)
VALUES (
$1, $2, $3, $4, $5
);`
) )
func LoadReader(db *sql.DB, uid int, r io.Reader) error {
stmt, err := db.Prepare(insertOriginalQuery)
if err != nil {
return err
}
b, err := ioutil.ReadAll(&io.LimitedReader{R: r, N: maxMessageSize})
if err != nil {
return err
}
hdr_size := bytes.Index(b, CRCR)
chksum, err := email.HashReader(bytes.NewReader(b))
if err != nil {
return err
}
n := len(b)
if _, err := stmt.Exec(uid, chksum, hdr_size, n, b); err != nil {
return err
}
return nil
}
func Load(db *sql.DB, uid int, root string, skip *set.StringSet) error { func Load(db *sql.DB, uid int, root string, skip *set.StringSet) error {
dup := set.NewStrings() dup := set.NewStrings()
start := time.Now() start := time.Now()
@ -54,8 +82,7 @@ func Load(db *sql.DB, uid int, root string, skip *set.StringSet) error {
return err return err
} }
stmt, err := txn.Prepare("INSERT INTO original (uid, hash, header_size, total_size, blob) VALUES ($1, $2, $3, $4, $5);") stmt, err := txn.Prepare(insertOriginalQuery)
//stmt, err := txn.Prepare(pq.CopyIn("original", "uid", "hash", "total_size", "header", "body"))
if err != nil { if err != nil {
if err := txn.Rollback(); err != nil { if err := txn.Rollback(); err != nil {
glog.Errorln("txn.Prepare stmt error rolling back", err) glog.Errorln("txn.Prepare stmt error rolling back", err)
@ -63,7 +90,6 @@ func Load(db *sql.DB, uid int, root string, skip *set.StringSet) error {
return err return err
} }
//hstmt, err := txn.Prepare(pq.CopyIn("files", "hash", "path"))
hstmt, err := txn.Prepare("INSERT INTO files (hash, path) VALUES ($1, $2);") hstmt, err := txn.Prepare("INSERT INTO files (hash, path) VALUES ($1, $2);")
if err != nil { if err != nil {
if err := txn.Rollback(); err != nil { if err := txn.Rollback(); err != nil {
@ -193,10 +219,13 @@ func main() {
glog.Infoln("Using uid", uid, "for", *username) glog.Infoln("Using uid", uid, "for", *username)
if *maildir == "" { if *maildir == "" {
fmt.Println("Must specify Maildir with -maildir") if err := LoadReader(db, uid, os.Stdin); err != nil {
os.Exit(1) glog.Fatal(err)
}
return
} }
// Load all files in maildir
skip := set.NewStrings(strings.Split(*skipFiles, ",")...) skip := set.NewStrings(strings.Split(*skipFiles, ",")...)
glog.Infoln("Skip files", skip) glog.Infoln("Skip files", skip)