sift/internal/dcache/pdf.go

package dcache

import (
	"crypto/sha256"
	"encoding/hex"
	"time"
)

// PDFEntry is one cached PDF: how to revalidate it, its metadata (the pdftext.Meta JSON, opaque
// here) and the text of each extracted page (Pages[i] is page i+1).
type PDFEntry struct {
	ContentHash  string
	ETag         string
	LastModified string
	FetchedAt    time.Time
	Meta         []byte
	Pages        []string
}

// PDFKey is the cache key for a document: a hash of the URL and the caller class, like the
// other keys here (the URL never appears in the clear).
func PDFKey(url, callerClass string) string {
	if callerClass == "" {
		callerClass = ClassFleet
	}
	sum := sha256.Sum256([]byte(url + "\x00" + callerClass))
	return hex.EncodeToString(sum[:])
}

// PutPDF writes the document and its pages in one transaction, deleting the old pages first so a
// failed replace never leaves a partial document behind, then calls evict.
func (d *DB) PutPDF(key string, e PDFEntry) {
	if d == nil {
		return
	}
	nbytes := len(e.Meta)
	for _, p := range e.Pages {
		nbytes += len(p)
	}
	tx, err := d.sql.Begin()
	if err != nil {
		if d.log != nil {
			d.log.Warn("pdf cache begin", "err", err)
		}
		return
	}
	if _, err := tx.Exec(`DELETE FROM pdf_page WHERE k = ?`, key); err != nil {
		_ = tx.Rollback()
		if d.log != nil {
			d.log.Warn("pdf cache pages del", "err", err)
		}
		return
	}
	if _, err := tx.Exec(`INSERT OR REPLACE INTO pdf_doc(k, content_hash, etag, last_modified, fetched_at, meta, last_hit, nbytes) VALUES(?,?,?,?,?,?,?,?)`,
		key, e.ContentHash, e.ETag, e.LastModified, e.FetchedAt.Unix(), e.Meta, time.Now().Unix(), nbytes); err != nil {
		_ = tx.Rollback()
		if d.log != nil {
			d.log.Warn("pdf cache doc", "err", err)
		}
		return
	}
	for i, p := range e.Pages {
		if _, err := tx.Exec(`INSERT INTO pdf_page(k, page, text) VALUES(?,?,?)`, key, i, p); err != nil {
			_ = tx.Rollback()
			if d.log != nil {
				d.log.Warn("pdf cache page put", "err", err)
			}
			return
		}
	}
	if err := tx.Commit(); err != nil {
		if d.log != nil {
			d.log.Warn("pdf cache commit", "err", err)
		}
		return
	}
	d.evict()
}

// GetPDF returns the pages in page order (an empty page keeps its place), updates last_hit, and
// counts the hit; a missing document is a miss.
func (d *DB) GetPDF(key string) (PDFEntry, bool) {
	if d == nil {
		return PDFEntry{}, false
	}
	var e PDFEntry
	var fetchedAt int64
	var meta []byte
	err := d.sql.QueryRow(`
		SELECT content_hash, etag, last_modified, fetched_at, meta
		FROM pdf_doc
		WHERE k = ?`, key).Scan(&e.ContentHash, &e.ETag, &e.LastModified, &fetchedAt, &meta)
	if err != nil {
		d.mu.Lock()
		d.PDFM++
		d.mu.Unlock()
		return PDFEntry{}, false
	}
	e.FetchedAt = time.Unix(fetchedAt, 0).UTC()
	e.Meta = meta

	rows, err := d.sql.Query(`SELECT page, text FROM pdf_page WHERE k = ? ORDER BY page ASC`, key)
	if err != nil {
		if d.log != nil {
			d.log.Warn("pdf cache pages", "err", err)
		}
		d.mu.Lock()
		d.PDFM++
		d.mu.Unlock()
		return PDFEntry{}, false
	}
	defer rows.Close()
	pages := make([]string, 0)
	for rows.Next() {
		var pg int
		var text string
		if err := rows.Scan(&pg, &text); err != nil {
			if d.log != nil {
				d.log.Warn("pdf cache page", "err", err)
			}
			d.mu.Lock()
			d.PDFM++
			d.mu.Unlock()
			return PDFEntry{}, false
		}
		pages = append(pages, text)
	}
	if err := rows.Err(); err != nil {
		if d.log != nil {
			d.log.Warn("pdf cache pages", "err", err)
		}
		d.mu.Lock()
		d.PDFM++
		d.mu.Unlock()
		return PDFEntry{}, false
	}
	e.Pages = pages

	_, _ = d.sql.Exec(`UPDATE pdf_doc SET last_hit = ? WHERE k = ?`, time.Now().Unix(), key)
	d.mu.Lock()
	d.PDFH++
	d.mu.Unlock()
	return e, true
}

// TouchPDF moves only the freshness timestamp after a 304; a miss changes nothing.
func (d *DB) TouchPDF(key string, fetchedAt time.Time) {
	if d == nil {
		return
	}
	if _, err := d.sql.Exec(`UPDATE pdf_doc SET fetched_at = ? WHERE k = ?`, fetchedAt.Unix(), key); err != nil {
		if d.log != nil {
			d.log.Warn("pdf cache touch", "err", err)
		}
	}
}