1
package sqlite3
2
 
3
import (
4
	"bytes"
5
	"context"
6
	"database/sql"
7
	"fmt"
8
 
9
	"go.miragespace.co/specter/spec/protocol"
10
)
11
 
12
const removeKeysBatchSize = 200
13
 
14
func importedSimpleValue(val *protocol.KVTransfer) ([]byte, bool) {
15
	if val == nil {
16
		return nil, false
17
	}
18
	if val.SimpleValue != nil {
19
		return val.SimpleValue, true
20
	}
21
	// Legacy senders omitted empty simple values. Retain their standalone-empty
22
	// interpretation; live keys with no simple value have prefix or lease state.
23
	// An empty value sharing that state needs an explicit field from the sender.
24
	if len(val.PrefixChildren) == 0 && val.LeaseToken == 0 {
25
		return []byte{}, true
26
	}
27
	return nil, false
28
}
29
 
30
func (s *SqliteKV) ListKeys(ctx context.Context, prefix []byte) ([]*protocol.KeyComposite, error) {
31
	keys := make([]*protocol.KeyComposite, 0)
32
 
33
	rows, err := s.stmts.listKeys.QueryContext(ctx)
34
	if err != nil {
35
		return nil, err
36
	}
37
	defer rows.Close()
38
 
39
	for rows.Next() {
40
		var (
41
			key   []byte
42
			flags uint8
43
		)
44
		if err := rows.Scan(&key, &flags); err != nil {
45
			return nil, err
46
		}
47
		if !bytes.HasPrefix(key, prefix) {
48
			continue
49
		}
50
		if flags&SimpleFlag != 0 {
51
			keys = append(keys, &protocol.KeyComposite{
52
				Type: protocol.KeyComposite_SIMPLE,
53
				Key:  key,
54
			})
55
		}
56
		if flags&PrefixFlag != 0 {
57
			keys = append(keys, &protocol.KeyComposite{
58
				Type: protocol.KeyComposite_PREFIX,
59
				Key:  key,
60
			})
61
		}
62
		if flags&LeaseFlag != 0 {
63
			keys = append(keys, &protocol.KeyComposite{
64
				Type: protocol.KeyComposite_LEASE,
65
				Key:  key,
66
			})
67
		}
68
	}
69
	if err := rows.Err(); err != nil {
70
		return nil, err
71
	}
72
 
73
	return keys, nil
74
}
75
 
76
func (s *SqliteKV) Import(ctx context.Context, keys [][]byte, values []*protocol.KVTransfer) error {
77
	if len(keys) != len(values) {
78
		return fmt.Errorf("keys and values length mismatch: %d != %d", len(keys), len(values))
79
	}
80
	return withWriteTx(ctx, s.writer, func(tx *sql.Tx) error {
81
		for i, key := range keys {
82
			var (
83
				val        = values[i]
84
				flag uint8 = 0
85
			)
86
 
87
			if val == nil {
88
				return fmt.Errorf("values[%d] is nil", i)
89
			}
90
 
91
			if simpleValue, ok := importedSimpleValue(val); ok {
92
				_, err := tx.StmtContext(ctx, s.stmts.simplePut).Exec(key, simpleValue)
93
				if err != nil {
94
					return err
95
				}
96
				flag |= SimpleFlag
97
			}
98
 
99
			// there's no point to override prefix values, but we need to handle conflict
100
			children := val.GetPrefixChildren()
101
			for _, child := range children {
102
				_, err := tx.StmtContext(ctx, s.stmts.prefixAppend).Exec(key, child)
103
				if err != nil {
104
					return err
105
				}
106
			}
107
			if len(children) > 0 {
108
				flag |= PrefixFlag
109
			}
110
 
111
			// override lease token when importing
112
			if val.GetLeaseToken() != 0 {
113
				_, err := tx.StmtContext(ctx, s.stmts.leaseImport).Exec(key, bindUint64AsInt64(val.GetLeaseToken()))
114
				if err != nil {
115
					return err
116
				}
117
				flag |= LeaseFlag
118
			}
119
			if err := s.updateKeyTracker(ctx, tx, key, flag, 0); err != nil {
120
				return err
121
			}
122
		}
123
		return nil
124
	})
125
}
126
 
127
func (s *SqliteKV) Export(ctx context.Context, keys [][]byte) ([]*protocol.KVTransfer, error) {
128
	vals := make([]*protocol.KVTransfer, len(keys))
129
	err := withReadTx(ctx, s.reader, func(tx *sql.Tx) error {
130
		for i, key := range keys {
131
			var (
132
				simpleValue []byte
133
				prefix      [][]byte
134
				leaseToken  int64
135
			)
136
 
137
			// simple value
138
			simpleErr := tx.StmtContext(ctx, s.stmts.exportSimpleGet).QueryRow(key).Scan(&simpleValue)
139
			if simpleErr != nil && simpleErr != sql.ErrNoRows {
140
				return simpleErr
141
			}
142
			// Preserve behavior: distinguish between missing row (nil) and empty blob
143
			if simpleErr == nil && simpleValue == nil {
144
				simpleValue = []byte{}
145
			}
146
 
147
			// prefix children
148
			prefixRows, err := tx.StmtContext(ctx, s.stmts.exportPrefixList).Query(key)
149
			if err != nil {
150
				return err
151
			}
152
			prefix = make([][]byte, 0)
153
			for prefixRows.Next() {
154
				var child []byte
155
				if err := prefixRows.Scan(&child); err != nil {
156
					prefixRows.Close()
157
					return err
158
				}
159
				prefix = append(prefix, child)
160
			}
161
			if err := prefixRows.Err(); err != nil {
162
				prefixRows.Close()
163
				return err
164
			}
165
			prefixRows.Close()
166
 
167
			// lease token
168
			leaseErr := tx.StmtContext(ctx, s.stmts.exportLeaseGet).QueryRow(key).Scan(&leaseToken)
169
			if leaseErr != nil && leaseErr != sql.ErrNoRows {
170
				return leaseErr
171
			}
172
 
173
			val := &protocol.KVTransfer{
174
				SimpleValue:    simpleValue,
175
				PrefixChildren: prefix,
176
				LeaseToken:     scanInt64AsUint64(leaseToken),
177
			}
178
			vals[i] = val
179
		}
180
 
181
		return nil
182
	})
183
	if err != nil {
184
		return nil, err
185
	}
186
	return vals, nil
187
}
188
 
189
func (s *SqliteKV) RemoveKeys(ctx context.Context, keys [][]byte) error {
190
	if len(keys) == 0 {
191
		return nil
192
	}
193
	return withWriteTx(ctx, s.writer, func(tx *sql.Tx) error {
194
		for start := 0; start < len(keys); start += removeKeysBatchSize {
195
			end := min(start+removeKeysBatchSize, len(keys))
196
			batch := keys[start:end]
197
			ph := placeholders(len(batch))
198
			args := make([]any, len(batch))
199
			for i, k := range batch {
200
				args[i] = k
201
			}
202
 
203
			if _, err := tx.Exec("DELETE FROM `simple_entries` WHERE `key` IN ("+ph+")", args...); err != nil {
204
				return err
205
			}
206
			if _, err := tx.Exec("DELETE FROM `prefix_entries` WHERE `prefix` IN ("+ph+")", args...); err != nil {
207
				return err
208
			}
209
			if _, err := tx.Exec("DELETE FROM `lease_entries` WHERE `owner` IN ("+ph+")", args...); err != nil {
210
				return err
211
			}
212
			if _, err := tx.Exec("DELETE FROM `key_trackers` WHERE `key` IN ("+ph+")", args...); err != nil {
213
				return err
214
			}
215
		}
216
		return nil
217
	})
218
}
219
 
220
func (s *SqliteKV) RangeKeys(ctx context.Context, low uint64, high uint64) ([][]byte, error) {
221
	keys := make([][]byte, 0)
222
 
223
	var args []any
224
 
225
	lowI := bindUint64AsInt64(low)
226
	highI := bindUint64AsInt64(high)
227
 
228
	var stmt *sql.Stmt
229
	if high > low {
230
		stmt = s.stmts.rangeKeysNorm
231
	} else {
232
		stmt = s.stmts.rangeKeysWrap
233
	}
234
	args = []any{lowI, highI, highI}
235
 
236
	rows, err := stmt.QueryContext(ctx, args...)
237
	if err != nil {
238
		return nil, err
239
	}
240
	defer rows.Close()
241
 
242
	for rows.Next() {
243
		var key []byte
244
		if err := rows.Scan(&key); err != nil {
245
			return nil, err
246
		}
247
		keys = append(keys, key)
248
	}
249
	return keys, rows.Err()
250
}