shell-history/internal/importer/import.go

124 lines
2.6 KiB
Go

package importer
import (
"bufio"
"database/sql"
"fmt"
"os"
"path/filepath"
"git.meatbag.se/vlv/shell-history/internal/model"
)
func Run(d *sql.DB, args []string) error {
if len(args) < 1 {
fmt.Fprintln(os.Stderr, "usage: shist import DIR")
os.Exit(1)
}
return importDir(d, args[0])
}
func importDir(d *sql.DB, dir string) error {
files, err := filepath.Glob(filepath.Join(dir, "zsh-history-*.log"))
if err != nil {
return err
}
var totalRead, totalSkipped, totalInserted int
for _, f := range files {
read, skipped, inserted, err := importFile(d, f)
if err != nil {
return fmt.Errorf("import %s: %w", f, err)
}
totalRead += read
totalSkipped += skipped
totalInserted += inserted
}
// Set last_export_id to max id so we don't re-export on first sync
var maxID int64
err = d.QueryRow("SELECT COALESCE(MAX(id), 0) FROM history").Scan(&maxID)
if err != nil {
return err
}
_, err = d.Exec(
`INSERT OR REPLACE INTO sync_meta (key, value) VALUES ('last_export_id', ?)`,
fmt.Sprintf("%d", maxID),
)
if err != nil {
return err
}
fmt.Printf("Import complete: %d lines read, %d skipped, %d inserted, %d duplicates\n",
totalRead, totalSkipped, totalInserted, totalRead-totalSkipped-totalInserted)
return nil
}
func importFile(d *sql.DB, path string) (read, skipped, inserted int, err error) {
f, err := os.Open(path)
if err != nil {
return 0, 0, 0, err
}
defer f.Close()
tx, err := d.Begin()
if err != nil {
return 0, 0, 0, err
}
defer tx.Rollback()
stmt, err := tx.Prepare(
`INSERT OR IGNORE INTO history (timestamp, hostname, working_dir, command) VALUES (?, ?, ?, ?)`,
)
if err != nil {
return 0, 0, 0, err
}
defer stmt.Close()
scanner := bufio.NewScanner(f)
scanner.Buffer(make([]byte, 0, 1024*1024), 1024*1024)
batch := 0
for scanner.Scan() {
read++
entry, err := model.ParseLine(scanner.Text())
if err != nil {
skipped++
continue
}
res, err := stmt.Exec(entry.Timestamp, entry.Hostname, entry.WorkingDir, entry.Command)
if err != nil {
return read, skipped, inserted, err
}
n, _ := res.RowsAffected()
inserted += int(n)
batch++
if batch >= 1000 {
if err := tx.Commit(); err != nil {
return read, skipped, inserted, err
}
tx, err = d.Begin()
if err != nil {
return read, skipped, inserted, err
}
stmt, err = tx.Prepare(
`INSERT OR IGNORE INTO history (timestamp, hostname, working_dir, command) VALUES (?, ?, ?, ?)`,
)
if err != nil {
return read, skipped, inserted, err
}
batch = 0
}
}
if err := scanner.Err(); err != nil {
return read, skipped, inserted, err
}
return read, skipped, inserted, tx.Commit()
}