1package db
 2
 3import (
 4	"context"
 5	"database/sql"
 6	"errors"
 7)
 8
 9func (db *DB) KVSet(ctx context.Context, key string, value any) error {
10	// Upsert into SQLite
11	query := `
12        INSERT INTO kv_store (key, value, updated_at) 
13        VALUES (?, ?, CURRENT_TIMESTAMP)
14        ON CONFLICT(key) DO UPDATE SET 
15            value = excluded.value, 
16            updated_at = CURRENT_TIMESTAMP;`
17
18	_, err := db.write.ExecContext(ctx, query, key, value)
19	return err
20}
21
22func (db *DB) KVGet(ctx context.Context, key string) (any, error) {
23	var rawVal any
24	err := db.read.QueryRowContext(ctx, "SELECT value FROM kv_store WHERE key = ?", key).Scan(&rawVal)
25	return rawVal, err
26}
27
28func (db *DB) KVInc(ctx context.Context, key string, delta int64) (int64, error) {
29	tx, err := db.GetWriteTx(ctx)
30	if err != nil {
31		return 0, err
32	}
33
34	var val any
35	err = tx.QueryRowContext(ctx, "SELECT value FROM kv_store WHERE key = ?", key).Scan(&val)
36	if err != nil && !errors.Is(err, sql.ErrNoRows) {
37		_ = tx.Rollback()
38		return 0, err
39	}
40
41	if errors.Is(err, sql.ErrNoRows) {
42		_, err = tx.ExecContext(ctx, "INSERT INTO kv_store (key, value, updated_at) VALUES (?, ?, CURRENT_TIMESTAMP)", key, delta)
43		if err != nil {
44			_ = tx.Rollback()
45			return 0, err
46		}
47		return delta, tx.Commit()
48	}
49
50	if _, ok := val.(int64); !ok {
51		_ = tx.Rollback()
52		return 0, errors.New("KVInc: value is not an integer")
53	}
54
55	query := `
56INSERT INTO kv_store (key, value, updated_at)
57VALUES (?, ?, CURRENT_TIMESTAMP)
58ON CONFLICT(key) DO UPDATE SET
59	value = CAST(kv_store.value AS INTEGER) + excluded.value,
60	updated_at = CURRENT_TIMESTAMP
61RETURNING value;`
62
63	var newCount int64
64	err = tx.QueryRowContext(ctx, query, key, delta).Scan(&newCount)
65	if err != nil {
66		_ = tx.Rollback()
67		return 0, err
68	}
69
70	return newCount, tx.Commit()
71}
72
73func (db *DB) KVPop(ctx context.Context, key string) (any, error) {
74	var rawVal any
75	query := `DELETE FROM kv_store WHERE key = ? RETURNING kv_store.value;`
76	err := db.write.QueryRowContext(ctx, query, key).Scan(&rawVal)
77	return rawVal, err
78}