ctxGPT/database/db.go

package database

import (
	"context"
	"fmt"
	"log"
	"zombiezen.com/go/sqlite"
	"zombiezen.com/go/sqlite/sqlitemigration"
	"zombiezen.com/go/sqlite/sqlitex"
)

const dbLocation = "./context.db"

var schema = sqlitemigration.Schema{
	// Each element of the Migrations slice is applied in sequence. When you
	// want to change the schema, add a new SQL script to this list.
	//
	// Existing databases will pick up at the same position in the Migrations
	// slice as they last left off.
	Migrations: []string{

		// sqlite create a table called context_store with an id, key, and content column
		"CREATE TABLE IF NOT EXISTS context_store ( id INTEGER PRIMARY KEY, key_name TEXT NOT NULL, value TEXT );",

		// sqlite create a virtual table called context using fts5 with a key and value column
		"CREATE VIRTUAL TABLE context_search USING fts5 ( key_name, value );",

		`CREATE TRIGGER context_ai AFTER INSERT ON context_store BEGIN
			INSERT INTO context_search(rowid, key_name, value) VALUES (new.id, new.key_name, new.value);
		END;`,
		`CREATE TRIGGER context_ad AFTER DELETE ON context_store BEGIN
			DELETE FROM context_search WHERE key_name = old.key_name AND value = old.value;
		END;`,
		`CREATE TRIGGER context_au AFTER UPDATE ON context_store BEGIN
			UPDATE context_search SET key_name = new.key_name, value = new.value WHERE id = old.id;
		END;`,
	},
}

type DB struct {
	db *sqlitex.Pool
}

func NewDB() (*DB, error) {

	err := runMigrations()
	if err != nil {
		return nil, fmt.Errorf("error running migrations | %w", err)
	}

	dbpool, err := sqlitex.Open(dbLocation, 0, 10)
	if err != nil {
		return nil, fmt.Errorf("error opening db | %w", err)
	}
	return &DB{
		db: dbpool,
	}, nil
}

func (d *DB) Close() error {
	return d.db.Close()
}

func runMigrations() error {
	// Open a pool. This does not block, and will start running any migrations
	// asynchronously.
	pool := sqlitemigration.NewPool(dbLocation, schema, sqlitemigration.Options{
		Flags: sqlite.OpenReadWrite | sqlite.OpenCreate,
		PrepareConn: func(conn *sqlite.Conn) error {
			// Enable foreign keys. See https://sqlite.org/foreignkeys.html
			return sqlitex.ExecuteTransient(conn, "PRAGMA foreign_keys = ON;", nil)
		},
		OnError: func(e error) {
			log.Println(e)
		},
	})
	defer pool.Close()

	// Get a connection. This blocks until the migration completes.
	conn, err := pool.Get(context.Background())
	if err != nil {
		return fmt.Errorf("error getting db connection | %w", err)
	}
	defer pool.Put(conn)

	// Print the list of schema objects created.
	const listSchemaQuery = `SELECT "type", "name" FROM sqlite_master ORDER BY 1, 2;`
	err = sqlitex.ExecuteTransient(conn, listSchemaQuery, &sqlitex.ExecOptions{
		ResultFunc: func(stmt *sqlite.Stmt) error {
			//fmt.Printf("%-5s %s\n", stmt.ColumnText(0), stmt.ColumnText(1))
			return nil
		},
	})
	if err != nil {
		return fmt.Errorf("error executing query | %w", err)
	}
	return nil
}

func (d *DB) Get(ctx context.Context, search string) ([]string, error) {
	conn := d.db.Get(ctx)
	if conn == nil {
		return []string{}, fmt.Errorf("error getting db connection")
	}
	defer d.db.Put(conn)

	query := "SELECT * FROM context_search WHERE context_search MATCH $search;"

	stmt := conn.Prep(query)
	defer stmt.Finalize()
	stmt.SetText("$search", search)
	for {
		if hasRow, err := stmt.Step(); err != nil {
			return []string{}, fmt.Errorf("error getting value from db | %w", err)
		} else if !hasRow {
			break
		}
		keyName := stmt.GetText("key_name")
		value := stmt.GetText("value")
		return []string{keyName, value}, nil
	}
	return []string{}, nil
}

func (d *DB) Save(ctx context.Context, key, value string) error {
	conn := d.db.Get(ctx)
	if conn == nil {
		return fmt.Errorf("error getting db connection to save info")
	}
	defer d.db.Put(conn)

	query := "INSERT INTO context_store (key_name, value) VALUES ($key, $value);"
	stmt := conn.Prep(query)
	defer stmt.Finalize()
	stmt.SetText("$key", key)
	stmt.SetText("$value", value)

	_, err := stmt.Step()
	if err != nil {
		return fmt.Errorf("error inserting new context | %w", err)
	}
	return nil
}