1package main
2
3import (
4 "context"
5 "database/sql"
6 "fmt"
7 "strings"
8
9 "github.com/mattn/go-sqlite3"
10)
11
12type (
13 DB struct {
14 db *sql.DB
15 source string // for backup
16 }
17
18 Word struct {
19 Word string
20 Line string
21 }
22)
23
24func Open(filename string) (*DB, error) {
25 sql.Register("sqlite3_with_extensions", &sqlite3.SQLiteDriver{
26 ConnectHook: func(conn *sqlite3.SQLiteConn) error {
27 return conn.LoadExtension("ext/libsqlite3ext", "sqlite3_spellfix_init")
28 },
29 })
30
31 db, err := sql.Open("sqlite3_with_extensions", filename)
32 if err != nil {
33 return nil, err
34 }
35
36 return &DB{
37 db: db,
38 source: filename,
39 }, nil
40}
41
42func (d *DB) Migrate(ctx context.Context) error {
43 _, err := d.db.ExecContext(
44 ctx,
45 `CREATE VIRTUAL TABLE IF NOT EXISTS words USING fts5 (word, line);
46 CREATE VIRTUAL TABLE IF NOT EXISTS words_terms USING fts4aux(words);
47 CREATE VIRTUAL TABLE IF NOT EXISTS spell USING spellfix1;
48 `,
49 )
50 return err
51}
52
53func (d *DB) SelectDict(ctx context.Context, query string, limit int) ([]*Word, error) {
54 rows, err := d.db.QueryContext(
55 ctx,
56 `SELECT
57 word, line
58 FROM words
59 WHERE word MATCH ?
60 ORDER BY rank;`,
61 query, limit,
62 )
63 if err != nil {
64 return nil, err
65 }
66
67 words := make([]*Word, 0)
68 for rows.Next() {
69 w := Word{}
70 err := rows.Scan(&w.Word, &w.Line)
71 if err != nil {
72 return nil, err
73 }
74 words = append(words, &w)
75 }
76
77 return words, err
78
79}
80
81func (d *DB) SelectSpell(ctx context.Context, query string) ([]string, error) {
82 rows, err := d.db.QueryContext(
83 ctx,
84 `SELECT
85 word
86 FROM spell
87 WHERE word MATCH ?;`,
88 query,
89 )
90 if err != nil {
91 return nil, err
92 }
93
94 words := make([]string, 0)
95 for rows.Next() {
96 w := ""
97 err := rows.Scan(&w)
98 if err != nil {
99 return nil, err
100 }
101 words = append(words, w)
102 }
103
104 return words, err
105
106}
107
108func (d *DB) InsertLine(ctx context.Context, line string) error {
109 p := strings.SplitN(line, "\t", 2)
110
111 _, err := d.db.ExecContext(
112 ctx,
113 `INSERT INTO words (WORD, LINE) VALUES(?, ?);`,
114 p[0], strings.ReplaceAll(p[1], "\t", " "),
115 )
116 if err != nil {
117 return err
118 }
119 return err
120}
121
122func (d *DB) Consolidade(ctx context.Context) error {
123 _, err := d.db.ExecContext(
124 ctx,
125 `INSERT INTO spell(word,rank)
126 SELECT term, documents FROM words_terms WHERE col='*'`,
127 )
128 if err != nil {
129 return err
130 }
131 return err
132}
133
134func (d *DB) Backup(ctx context.Context, name string) error {
135 destDb, err := sql.Open("sqlite3_with_extensions", name)
136 if err != nil {
137 return err
138 }
139 defer destDb.Close()
140
141 return Copy(ctx, d.db, destDb)
142}
143
144func (d *DB) Restore(ctx context.Context, name string) error {
145 srcDb, err := sql.Open("sqlite3_with_extensions", name)
146 if err != nil {
147 return err
148 }
149 defer srcDb.Close()
150
151 return Copy(ctx, srcDb, d.db)
152}
153
154func Copy(ctx context.Context, srcDb *sql.DB, destDb *sql.DB) error {
155 destConn, err := destDb.Conn(ctx)
156 if err != nil {
157 return err
158 }
159 defer destConn.Close()
160
161 srcConn, err := srcDb.Conn(ctx)
162 if err != nil {
163 return err
164 }
165 defer srcConn.Close()
166
167 return destConn.Raw(func(destConn interface{}) error {
168 return srcConn.Raw(func(srcConn interface{}) error {
169 destSQLiteConn, ok := destConn.(*sqlite3.SQLiteConn)
170 if !ok {
171 return fmt.Errorf("can't convert destination connection to SQLiteConn")
172 }
173
174 srcSQLiteConn, ok := srcConn.(*sqlite3.SQLiteConn)
175 if !ok {
176 return fmt.Errorf("can't convert source connection to SQLiteConn")
177 }
178
179 b, err := destSQLiteConn.Backup("main", srcSQLiteConn, "main")
180 if err != nil {
181 return fmt.Errorf("error initializing SQLite backup: %w", err)
182 }
183
184 done, err := b.Step(-1)
185 if !done {
186 return fmt.Errorf("step of -1, but not done")
187 }
188 if err != nil {
189 return fmt.Errorf("error in stepping backup: %w", err)
190 }
191
192 err = b.Finish()
193 if err != nil {
194 return fmt.Errorf("error finishing backup: %w", err)
195 }
196
197 return err
198 })
199 })
200
201}