Nicer progress bar for first run

1 files changed, 39 insertions, 3 deletions

diff --git a/setup.go b/setup.go
index e43cbdd..4b4ce91 100644
--- a/setup.go
+++ b/setup.go
@@ -2,9 +2,11 @@ package main
 
 import (
 	"bufio"
+	"bytes"
 	"database/sql"
 	"fmt"
 	"html/template"
+	"io"
 	"os"
 	"strings"
 
@@ -89,6 +91,7 @@ type dictionaryPopulator struct {
 	tx      *sql.Tx
 	stmt    *sql.Stmt
 	tmpl    *template.Template
+	fh      *os.File
 	scanner *bufio.Scanner
 
 	totalLines  int
@@ -122,15 +125,25 @@ func setupPopulator(dp *dictionaryPopulator) tea.Cmd {
 			return errMsg(fmt.Errorf("preparing statement: %w", err))
 		}
 
-		file, err := os.Open(dp.rawDictionaryPath)
+		dp.fh, err = os.Open(dp.rawDictionaryPath)
 		if err != nil {
 			return errMsg(fmt.Errorf("opening: %w", err))
 		}
 
-		dp.scanner = bufio.NewScanner(file)
+		// Figure out how many lines the file has, for reporting import
+		// progress.
+		lines, err := lineCounter(dp.fh)
+		if err != nil {
+			return errMsg(fmt.Errorf("reading lines from file: %w", err))
+		}
+		dp.totalLines = lines
 
-		maxCapacity := 2_000_000
+		// We've just read through the whole file, reset the read position to
+		// the beginning because we're about to set up a scanner on it.
+		dp.fh.Seek(0, 0)
 
+		dp.scanner = bufio.NewScanner(dp.fh)
+		maxCapacity := 2_000_000
 		buf := make([]byte, maxCapacity)
 		dp.scanner.Buffer(buf, maxCapacity)
 
@@ -138,6 +151,25 @@ func setupPopulator(dp *dictionaryPopulator) tea.Cmd {
 	}
 }
 
+func lineCounter(r io.Reader) (int, error) {
+	buf := make([]byte, 64*1024)
+	count := 0
+	lineSep := []byte{'\n'}
+
+	for {
+		c, err := r.Read(buf)
+		count += bytes.Count(buf[:c], lineSep)
+
+		switch {
+		case err == io.EOF:
+			return count, nil
+
+		case err != nil:
+			return count, err
+		}
+	}
+}
+
 func populateDictionary(dp *dictionaryPopulator) tea.Cmd {
 	return func() tea.Msg {
 		for dp.scanner.Scan() {
@@ -227,6 +259,10 @@ func populateDictionary(dp *dictionaryPopulator) tea.Cmd {
 			return errMsg(fmt.Errorf("creating index: %s", err))
 		}
 
+		// Clean up resources
+		dp.stmt.Close()
+		dp.fh.Close()
+
 		return isDictionaryEmptyMsg(false) // We're done!
 	}
 }