1
package sqlite3
2
 
3
import (
4
	"database/sql"
5
	"fmt"
6
	"os"
7
	"path/filepath"
8
	"runtime"
9
 
10
	"go.miragespace.co/specter/spec/chord"
11
 
12
	"go.uber.org/zap"
13
)
14
 
15
type SimpleEntry struct {
16
	Key   []byte
17
	Value []byte
18
}
19
 
20
type PrefixEntry struct {
21
	Prefix []byte
22
	Child  []byte
23
}
24
 
25
type LeaseEntry struct {
26
	Owner []byte
27
	Token uint64
28
}
29
 
30
type Config struct {
31
	Logger  *zap.Logger
32
	HashFn  chord.HashFn
33
	DataDir string
34
}
35
 
36
type SqliteKV struct {
37
	logger *zap.Logger
38
	hashFn chord.HashFn
39
	reader *sql.DB
40
	writer *sql.DB
41
	stmts  *statements
42
}
43
 
44
func (c Config) validate() error {
45
	if c.Logger == nil {
46
		return fmt.Errorf("nil Logger is invalid")
47
	}
48
	if c.HashFn == nil {
49
		return fmt.Errorf("nil HashFn is invalid")
50
	}
51
	if c.DataDir == "" {
52
		return fmt.Errorf("empty DataDir is invalid")
53
	}
54
	return nil
55
}
56
 
57
func New(cfg Config) (*SqliteKV, error) {
58
	if err := cfg.validate(); err != nil {
59
		return nil, err
60
	}
61
 
62
	dbDir := filepath.Join(cfg.DataDir, "sqlite3")
63
	if err := os.MkdirAll(dbDir, 0750); err != nil {
64
		return nil, err
65
	}
66
	dbPath := filepath.Join(dbDir, "db")
67
 
68
	var (
69
		reader *sql.DB
70
		writer *sql.DB
71
		stmts  *statements
72
		err    error
73
	)
74
 
75
	defer func() {
76
		if err == nil {
77
			return
78
		}
79
		if stmts != nil {
80
			stmts.close()
81
		}
82
		if reader != nil {
83
			_ = reader.Close()
84
		}
85
		if writer != nil {
86
			_ = writer.Close()
87
		}
88
	}()
89
 
90
	reader, err = openSQLite(cfg.Logger, dbPath)
91
	if err != nil {
92
		return nil, err
93
	}
94
	writer, err = openSQLite(cfg.Logger, dbPath)
95
	if err != nil {
96
		return nil, err
97
	}
98
 
99
	reader.SetMaxOpenConns(max(4, runtime.NumCPU()))
100
	// prevent SQLITE_BUSY
101
	writer.SetMaxOpenConns(1)
102
 
103
	if err := migrate(writer); err != nil {
104
		return nil, fmt.Errorf("migrating database: %w", err)
105
	}
106
 
107
	stmts, err = prepareStatements(reader, writer)
108
	if err != nil {
109
		return nil, fmt.Errorf("preparing statements: %w", err)
110
	}
111
 
112
	return &SqliteKV{
113
		logger: cfg.Logger,
114
		hashFn: cfg.HashFn,
115
		reader: reader,
116
		writer: writer,
117
		stmts:  stmts,
118
	}, nil
119
}
120
 
121
func (s *SqliteKV) Close() {
122
	if s == nil {
123
		return
124
	}
125
	if s.stmts != nil {
126
		s.stmts.close()
127
	}
128
	if s.reader != nil {
129
		_ = s.reader.Close()
130
	}
131
	if s.writer != nil {
132
		_ = s.writer.Close()
133
	}
134
}
135
 
136
var _ chord.KVProvider = (*SqliteKV)(nil)