orchard

  1package main
  2
  3import (
  4	"crypto/hmac"
  5	"crypto/rand"
  6	"crypto/sha256"
  7	"database/sql"
  8	"encoding/hex"
  9	"encoding/json"
 10	"os"
 11	"path/filepath"
 12	"strings"
 13	"time"
 14
 15	_ "modernc.org/sqlite"
 16)
 17
 18// Store holds the page archive and the search result cache.
 19//
 20// It deliberately never stores a question. Pages and passages are an archive of
 21// public articles, so they are kept in full and indexed for search, while the
 22// only thing that has to persist per query is a lookup key, and that is an HMAC
 23// so the database cannot be read back as a history of what was asked.
 24type Store struct {
 25	db     *sql.DB
 26	secret []byte
 27}
 28
 29const schema = `
 30CREATE TABLE IF NOT EXISTS pages (
 31  id         INTEGER PRIMARY KEY,
 32  url        TEXT NOT NULL UNIQUE,
 33  title      TEXT NOT NULL DEFAULT '',
 34  site       TEXT NOT NULL DEFAULT '',
 35  published  TEXT NOT NULL DEFAULT '',
 36  markdown   TEXT NOT NULL,
 37  fetched_at INTEGER NOT NULL
 38);
 39
 40CREATE TABLE IF NOT EXISTS passages (
 41  id      INTEGER PRIMARY KEY,
 42  page_id INTEGER NOT NULL REFERENCES pages(id) ON DELETE CASCADE,
 43  ord     INTEGER NOT NULL,
 44  text    TEXT NOT NULL
 45);
 46CREATE INDEX IF NOT EXISTS passages_page ON passages(page_id);
 47
 48CREATE VIRTUAL TABLE IF NOT EXISTS passages_fts USING fts5(
 49  text, content='passages', content_rowid='id', tokenize='porter unicode61'
 50);
 51
 52CREATE TRIGGER IF NOT EXISTS passages_ai AFTER INSERT ON passages BEGIN
 53  INSERT INTO passages_fts(rowid, text) VALUES (new.id, new.text);
 54END;
 55CREATE TRIGGER IF NOT EXISTS passages_ad AFTER DELETE ON passages BEGIN
 56  INSERT INTO passages_fts(passages_fts, rowid, text) VALUES('delete', old.id, old.text);
 57END;
 58
 59-- Outbound links found in a page's article body, so a cache hit still has the
 60-- candidates an entity link is resolved from.
 61CREATE TABLE IF NOT EXISTS links (
 62  id      INTEGER PRIMARY KEY,
 63  page_id INTEGER NOT NULL REFERENCES pages(id) ON DELETE CASCADE,
 64  url     TEXT NOT NULL,
 65  text    TEXT NOT NULL
 66);
 67CREATE INDEX IF NOT EXISTS links_page ON links(page_id);
 68
 69-- key is HMAC(secret, normalized query). The query text itself is never here.
 70CREATE TABLE IF NOT EXISTS serp (
 71  key        TEXT PRIMARY KEY,
 72  results    TEXT NOT NULL,
 73  fetched_at INTEGER NOT NULL
 74);
 75`
 76
 77func OpenStore(dataDir string) (*Store, error) {
 78	if err := os.MkdirAll(dataDir, 0o755); err != nil {
 79		return nil, err
 80	}
 81	db, err := sql.Open("sqlite", filepath.Join(dataDir, "search.db")+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)&_pragma=foreign_keys(1)")
 82	if err != nil {
 83		return nil, err
 84	}
 85	if _, err := db.Exec(schema); err != nil {
 86		return nil, err
 87	}
 88	secret, err := loadSecret(dataDir)
 89	if err != nil {
 90		return nil, err
 91	}
 92	return &Store{db: db, secret: secret}, nil
 93}
 94
 95// loadSecret keeps the HMAC key beside the database. It has to persist or every
 96// restart would miss every cached search, and it has to be secret or the keys
 97// could be checked against a dictionary of guessed queries.
 98func loadSecret(dataDir string) ([]byte, error) {
 99	path := filepath.Join(dataDir, "cache.key")
100	if b, err := os.ReadFile(path); err == nil && len(b) >= 32 {
101		return b, nil
102	}
103	b := make([]byte, 32)
104	if _, err := rand.Read(b); err != nil {
105		return nil, err
106	}
107	if err := os.WriteFile(path, b, 0o600); err != nil {
108		return nil, err
109	}
110	return b, nil
111}
112
113func (s *Store) Close() error { return s.db.Close() }
114
115func (s *Store) key(query string) string {
116	m := hmac.New(sha256.New, s.secret)
117	m.Write([]byte(strings.ToLower(strings.Join(strings.Fields(query), " "))))
118	return hex.EncodeToString(m.Sum(nil))
119}
120
121// CachedSERP returns a stored result list for a query, or nil past its TTL.
122func (s *Store) CachedSERP(query string, ttl time.Duration) []Result {
123	var blob string
124	var at int64
125	err := s.db.QueryRow(`SELECT results, fetched_at FROM serp WHERE key = ?`, s.key(query)).Scan(&blob, &at)
126	if err != nil || time.Since(time.Unix(at, 0)) > ttl {
127		return nil
128	}
129	var out []Result
130	if json.Unmarshal([]byte(blob), &out) != nil {
131		return nil
132	}
133	return out
134}
135
136func (s *Store) PutSERP(query string, results []Result) error {
137	blob, err := json.Marshal(results)
138	if err != nil {
139		return err
140	}
141	_, err = s.db.Exec(
142		`INSERT INTO serp(key, results, fetched_at) VALUES(?,?,?)
143		 ON CONFLICT(key) DO UPDATE SET results=excluded.results, fetched_at=excluded.fetched_at`,
144		s.key(query), string(blob), time.Now().Unix())
145	return err
146}
147
148// CachedPage returns a stored page if it was fetched inside the TTL. Articles
149// do not change, so this can be generous.
150func (s *Store) CachedPage(url string, ttl time.Duration) *Page {
151	var p Page
152	var at int64
153	err := s.db.QueryRow(
154		`SELECT url, title, site, published, markdown, fetched_at FROM pages WHERE url = ?`, url,
155	).Scan(&p.URL, &p.Title, &p.Site, &p.Published, &p.Markdown, &at)
156	if err != nil || time.Since(time.Unix(at, 0)) > ttl {
157		return nil
158	}
159	return &p
160}
161
162// PutPage stores a page and replaces its passages.
163func (s *Store) PutPage(p *Page) (int64, error) {
164	tx, err := s.db.Begin()
165	if err != nil {
166		return 0, err
167	}
168	defer tx.Rollback()
169
170	_, err = tx.Exec(
171		`INSERT INTO pages(url, title, site, published, markdown, fetched_at) VALUES(?,?,?,?,?,?)
172		 ON CONFLICT(url) DO UPDATE SET title=excluded.title, site=excluded.site,
173		   published=excluded.published, markdown=excluded.markdown, fetched_at=excluded.fetched_at`,
174		p.URL, p.Title, p.Site, p.Published, p.Markdown, time.Now().Unix())
175	if err != nil {
176		return 0, err
177	}
178	var id int64
179	if err := tx.QueryRow(`SELECT id FROM pages WHERE url = ?`, p.URL).Scan(&id); err != nil {
180		return 0, err
181	}
182	if _, err := tx.Exec(`DELETE FROM passages WHERE page_id = ?`, id); err != nil {
183		return 0, err
184	}
185	if _, err := tx.Exec(`DELETE FROM links WHERE page_id = ?`, id); err != nil {
186		return 0, err
187	}
188	for _, l := range p.Links {
189		if _, err := tx.Exec(`INSERT INTO links(page_id, url, text) VALUES(?,?,?)`, id, l.URL, l.Text); err != nil {
190			return 0, err
191		}
192	}
193	for i, text := range Chunk(p.Markdown) {
194		if _, err := tx.Exec(`INSERT INTO passages(page_id, ord, text) VALUES(?,?,?)`, id, i, text); err != nil {
195			return 0, err
196		}
197	}
198	return id, tx.Commit()
199}
200
201// PageLinks returns the outbound links stored for a page.
202func (s *Store) PageLinks(pageID int64) []Link {
203	rows, err := s.db.Query(`SELECT url, text FROM links WHERE page_id = ?`, pageID)
204	if err != nil {
205		return nil
206	}
207	defer rows.Close()
208	var out []Link
209	for rows.Next() {
210		var l Link
211		if rows.Scan(&l.URL, &l.Text) == nil {
212			out = append(out, l)
213		}
214	}
215	return out
216}
217
218// LocalHit is a passage found in the archive rather than on the web.
219type LocalHit struct {
220	URL   string
221	Title string
222	Text  string
223}
224
225// SearchLocal queries the passage index. This is what lets a question reuse
226// pages fetched for an unrelated question that shared none of its words.
227func (s *Store) SearchLocal(query string, limit int) []LocalHit {
228	rows, err := s.db.Query(`
229		SELECT p.url, p.title, x.text
230		FROM passages_fts f
231		JOIN passages x ON x.id = f.rowid
232		JOIN pages p ON p.id = x.page_id
233		WHERE passages_fts MATCH ?
234		ORDER BY bm25(passages_fts) LIMIT ?`, ftsQuery(query), limit)
235	if err != nil {
236		return nil
237	}
238	defer rows.Close()
239	var out []LocalHit
240	for rows.Next() {
241		var h LocalHit
242		if rows.Scan(&h.URL, &h.Title, &h.Text) == nil {
243			out = append(out, h)
244		}
245	}
246	return out
247}
248
249// ftsQuery turns a plain question into an OR query of its content words. FTS5
250// treats most punctuation as syntax, so anything not alphanumeric is dropped
251// rather than escaped.
252func ftsQuery(q string) string {
253	var words []string
254	for _, f := range strings.Fields(strings.ToLower(q)) {
255		clean := strings.Map(func(r rune) rune {
256			if r >= 'a' && r <= 'z' || r >= '0' && r <= '9' {
257				return r
258			}
259			return -1
260		}, f)
261		if len(clean) > 2 && !stopword[clean] {
262			words = append(words, clean)
263		}
264	}
265	if len(words) == 0 {
266		return "\"\""
267	}
268	return strings.Join(words, " OR ")
269}
270
271var stopword = map[string]bool{
272	"the": true, "and": true, "for": true, "are": true, "was": true, "what": true,
273	"how": true, "why": true, "who": true, "does": true, "did": true, "can": true,
274	"you": true, "with": true, "from": true, "that": true, "this": true, "has": true,
275}
276
277func (s *Store) Stats() (pages, passages, sites int) {
278	s.db.QueryRow(`SELECT count(*) FROM pages`).Scan(&pages)
279	s.db.QueryRow(`SELECT count(*) FROM passages`).Scan(&passages)
280	s.db.QueryRow(`SELECT count(DISTINCT site) FROM pages WHERE site != ''`).Scan(&sites)
281	return
282}
283
284// Chunk splits markdown into passages a model can be handed one at a time.
285// Paragraph boundaries are the split points, and short ones get merged so a
286// passage is a claim rather than a heading.
287func Chunk(md string) []string {
288	const target = 1200
289	var out []string
290	var cur strings.Builder
291	for _, para := range strings.Split(md, "\n\n") {
292		p := strings.TrimSpace(para)
293		if p == "" {
294			continue
295		}
296		if cur.Len() > 0 && cur.Len()+len(p) > target {
297			out = append(out, cur.String())
298			cur.Reset()
299		}
300		if cur.Len() > 0 {
301			cur.WriteString("\n\n")
302		}
303		cur.WriteString(p)
304	}
305	if strings.TrimSpace(cur.String()) != "" {
306		out = append(out, cur.String())
307	}
308	return out
309}
310
311// RankPassages orders one page's chunks by how well they answer the question.
312//
313// Taking the first chunks of a page is what made recipe answers useless: the
314// ingredients and the method sit well down the document, behind the story about
315// the author's trip to Oaxaca. bm25 over the chunks the page just produced puts
316// them back on top.
317func (s *Store) RankPassages(pageID int64, question string, limit int) []string {
318	rows, err := s.db.Query(`
319		SELECT x.text
320		FROM passages_fts f
321		JOIN passages x ON x.id = f.rowid
322		WHERE passages_fts MATCH ? AND x.page_id = ?
323		ORDER BY bm25(passages_fts) LIMIT ?`, ftsQuery(question), pageID, limit)
324	if err != nil {
325		return nil
326	}
327	defer rows.Close()
328	var out []string
329	for rows.Next() {
330		var t string
331		if rows.Scan(&t) == nil {
332			out = append(out, t)
333		}
334	}
335	return out
336}
337
338// PageChunks returns a page's chunks in document order, which is the fallback
339// when nothing matches the query terms.
340func (s *Store) PageChunks(pageID int64, limit int) []string {
341	rows, err := s.db.Query(`SELECT text FROM passages WHERE page_id = ? ORDER BY ord LIMIT ?`, pageID, limit)
342	if err != nil {
343		return nil
344	}
345	defer rows.Close()
346	var out []string
347	for rows.Next() {
348		var t string
349		if rows.Scan(&t) == nil {
350			out = append(out, t)
351		}
352	}
353	return out
354}
355
356// PageID resolves a stored page's row id, needed to rank its chunks after a
357// cache hit skipped the insert.
358func (s *Store) PageID(url string) int64 {
359	var id int64
360	s.db.QueryRow(`SELECT id FROM pages WHERE url = ?`, url).Scan(&id)
361	return id
362}