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) KVList(ctx context.Context) (map[string]any, error) {
29	rows, err := db.read.QueryContext(ctx, "SELECT key, value FROM kv_store")
30	if err != nil {
31		return nil, err
32	}
33	defer rows.Close()
34
35	result := make(map[string]any)
36	for rows.Next() {
37		var key string
38		var value any
39		if err := rows.Scan(&key, &value); err != nil {
40			return nil, err
41		}
42		result[key] = value
43	}
44	return result, rows.Err()
45}
46
47func (db *DB) KVInc(ctx context.Context, key string, delta int64) (int64, error) {
48	tx, err := db.GetWriteTx(ctx)
49	if err != nil {
50		return 0, err
51	}
52
53	var val any
54	err = tx.QueryRowContext(ctx, "SELECT value FROM kv_store WHERE key = ?", key).Scan(&val)
55	if err != nil && !errors.Is(err, sql.ErrNoRows) {
56		_ = tx.Rollback()
57		return 0, err
58	}
59
60	if errors.Is(err, sql.ErrNoRows) {
61		_, err = tx.ExecContext(ctx, "INSERT INTO kv_store (key, value, updated_at) VALUES (?, ?, CURRENT_TIMESTAMP)", key, delta)
62		if err != nil {
63			_ = tx.Rollback()
64			return 0, err
65		}
66		return delta, tx.Commit()
67	}
68
69	if _, ok := val.(int64); !ok {
70		_ = tx.Rollback()
71		return 0, errors.New("KVInc: value is not an integer")
72	}
73
74	query := `
75INSERT INTO kv_store (key, value, updated_at)
76VALUES (?, ?, CURRENT_TIMESTAMP)
77ON CONFLICT(key) DO UPDATE SET
78	value = CAST(kv_store.value AS INTEGER) + excluded.value,
79	updated_at = CURRENT_TIMESTAMP
80RETURNING value;`
81
82	var newCount int64
83	err = tx.QueryRowContext(ctx, query, key, delta).Scan(&newCount)
84	if err != nil {
85		_ = tx.Rollback()
86		return 0, err
87	}
88
89	return newCount, tx.Commit()
90}
91
92func (db *DB) KVPop(ctx context.Context, key string) (any, error) {
93	var rawVal any
94	query := `DELETE FROM kv_store WHERE key = ? RETURNING kv_store.value;`
95	err := db.write.QueryRowContext(ctx, query, key).Scan(&rawVal)
96	return rawVal, err
97}