1 files changed, 53 insertions, 24 deletions

diff --git a/setup.go b/setup.go
index 1a69d45..e43cbdd 100644
--- a/setup.go
+++ b/setup.go
@@ -37,7 +37,7 @@ func isDatabaseEmpty(db *sql.DB) tea.Cmd {
 		}
 
 		// Only populate the database if it is empty.
-		return isDBEmptyMsg(count == 0)
+		return isDictionaryEmptyMsg(count == 0)
 	}
 }
 
@@ -78,10 +78,29 @@ type SenseForDictionaryEntry struct {
 	Example string
 }
 
-func populateDictionary(rawDictionary string, db *sql.DB) tea.Cmd {
+// dictionaryPopulator contains all the information required to populate the
+// SQLite dictionary from the raw JSONL data. This is in a struct so that we can
+// report progress back to the UI, then resume where we left off.
+type dictionaryPopulator struct {
+	db                *sql.DB
+	rawDictionaryPath string
+	langCode          string
+
+	tx      *sql.Tx
+	stmt    *sql.Stmt
+	tmpl    *template.Template
+	scanner *bufio.Scanner
+
+	totalLines  int
+	currentLine int
+}
+
+func setupPopulator(dp *dictionaryPopulator) tea.Cmd {
 	return func() tea.Msg {
+		var err error
+
 		// Set up the template
-		tmpl, err := template.New("entry").Parse(
+		dp.tmpl, err = template.New("entry").Parse(
 			`<p>{{ .Word }} {{ .Sound }} <i>{{ .POS }} {{ .Gender }}</i></p>
 <ol>{{ range .Senses}}
 	<li class=sense>{{ .Sense }}<br>
@@ -89,38 +108,44 @@ func populateDictionary(rawDictionary string, db *sql.DB) tea.Cmd {
 {{ end }}</ol>
 {{ if .Etymology }}<p><i>Étymologie: {{ .Etymology }}</i>{{ end }}`)
 		if err != nil {
-			panic(err)
+			return errMsg(fmt.Errorf("preparing template: %w", err))
 		}
 
-		tx, err := db.Begin()
+		dp.tx, err = dp.db.Begin()
 		if err != nil {
 			return errMsg(fmt.Errorf("starting transaction: %w", err))
 		}
 
 		// Set up a prepared statement
-		stmt, err := tx.Prepare("insert into words(word, definition) values(?, ?)")
+		dp.stmt, err = dp.tx.Prepare("insert into words(word, definition) values(?, ?)")
 		if err != nil {
 			return errMsg(fmt.Errorf("preparing statement: %w", err))
 		}
-		defer stmt.Close()
 
-		file, err := os.Open(rawDictionary)
+		file, err := os.Open(dp.rawDictionaryPath)
 		if err != nil {
 			return errMsg(fmt.Errorf("opening: %w", err))
 		}
-		defer file.Close()
 
-		var wordsAdded int
-		scanner := bufio.NewScanner(file)
+		dp.scanner = bufio.NewScanner(file)
 
 		maxCapacity := 2_000_000
+
 		buf := make([]byte, maxCapacity)
-		scanner.Buffer(buf, maxCapacity)
+		dp.scanner.Buffer(buf, maxCapacity)
+
+		return populatingDictionaryMsg(dp)
+	}
+}
+
+func populateDictionary(dp *dictionaryPopulator) tea.Cmd {
+	return func() tea.Msg {
+		for dp.scanner.Scan() {
+			dp.currentLine++
 
-		for scanner.Scan() {
 			var result rawDictionaryEntry
-			json.Unmarshal([]byte(scanner.Text()), &result)
-			if result.LangCode != "fr" {
+			json.Unmarshal([]byte(dp.scanner.Text()), &result)
+			if result.LangCode != dp.langCode {
 				continue
 			}
 
@@ -171,33 +196,37 @@ func populateDictionary(rawDictionary string, db *sql.DB) tea.Cmd {
 			}
 
 			formattedDefinition := strings.Builder{}
-			err := tmpl.Execute(&formattedDefinition, entry)
+			err := dp.tmpl.Execute(&formattedDefinition, entry)
 			if err != nil {
 				return errMsg(fmt.Errorf("failed to render: %w", err))
 			}
 
 			// Insert the entry
-			_, err = stmt.Exec(entry.Word, formattedDefinition.String())
+			_, err = dp.stmt.Exec(entry.Word, formattedDefinition.String())
 			if err != nil {
 				return errMsg(fmt.Errorf("inserting '%s': %w", entry.Word, err))
 			}
 
-			wordsAdded++
-
+			// Report status every once in a while by breaking out to the caller
+			if dp.currentLine%10000 == 0 {
+				return populatingDictionaryMsg(dp)
+			}
 		}
-		if err := scanner.Err(); err != nil {
+
+		// If we're outside of the loop, we either encountered an error, or it's
+		// time to commit the changes.
+		if err := dp.scanner.Err(); err != nil {
 			return errMsg(fmt.Errorf("scanning: %w", err))
 		}
 
-		if err := tx.Commit(); err != nil {
+		if err := dp.tx.Commit(); err != nil {
 			return errMsg(fmt.Errorf("committing: %w", err))
 		}
-
-		_, err = db.Exec("create index wordindex on words(word);")
+		_, err := dp.db.Exec("create index wordindex on words(word);")
 		if err != nil {
 			return errMsg(fmt.Errorf("creating index: %s", err))
 		}
 
-		return isDBEmptyMsg(false)
+		return isDictionaryEmptyMsg(false) // We're done!
 	}
 }