diff --git a/cmd/md2pq/md2pq.go b/cmd/md2pq/md2pq.go index 6883dbc..adc98b1 100644 --- a/cmd/md2pq/md2pq.go +++ b/cmd/md2pq/md2pq.go @@ -6,7 +6,6 @@ import ( "database/sql" "expvar" "flag" - "fmt" "io" "io/ioutil" "os" @@ -38,9 +37,38 @@ var ( CRCR = []byte("\n\n") // Maximum bytes we parse when decoding // 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 { dup := set.NewStrings() start := time.Now() @@ -54,8 +82,7 @@ func Load(db *sql.DB, uid int, root string, skip *set.StringSet) error { 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(pq.CopyIn("original", "uid", "hash", "total_size", "header", "body")) + stmt, err := txn.Prepare(insertOriginalQuery) if err != nil { if err := txn.Rollback(); err != nil { 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 } - //hstmt, err := txn.Prepare(pq.CopyIn("files", "hash", "path")) hstmt, err := txn.Prepare("INSERT INTO files (hash, path) VALUES ($1, $2);") if err != nil { if err := txn.Rollback(); err != nil { @@ -193,10 +219,13 @@ func main() { glog.Infoln("Using uid", uid, "for", *username) if *maildir == "" { - fmt.Println("Must specify Maildir with -maildir") - os.Exit(1) + if err := LoadReader(db, uid, os.Stdin); err != nil { + glog.Fatal(err) + } + return } + // Load all files in maildir skip := set.NewStrings(strings.Split(*skipFiles, ",")...) glog.Infoln("Skip files", skip)