1
// Copyright (c) 2023 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
#include "sqlite-kv.h"
6
 
7
#include <kj/test.h>
8
 
9
namespace workerd {
10
namespace {
11
 
12
struct GlobalInit {
13
  GlobalInit() {
14
    installSqliteCustomAllocator();
15
  }
16
};
17
 
18
static GlobalInit init;
19
 
20
KJ_TEST("SQLite-KV") {
21
  class TestSqliteObserver: public SqliteObserver {
22
   public:
23
    void addQueryStats(uint64_t read, uint64_t written) override {
24
      rowsRead += read;
25
      rowsWritten += written;
26
    }
27
 
28
    uint64_t rowsRead = 0;
29
    uint64_t rowsWritten = 0;
30
  };
31
 
32
  auto dir = kj::newInMemoryDirectory(kj::nullClock());
33
  SqliteDatabase::Vfs vfs(*dir);
34
  TestSqliteObserver sqliteObserver;
35
  SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY,
36
      /*sqliteMaxMemoryBytes=*/kj::maxValue, sqliteObserver);
37
  SqliteKv kv(db);
38
 
39
  kv.put("foo", "abc"_kj.asBytes());
40
  kv.put("bar", "def"_kj.asBytes());
41
  kv.put("baz", "123"_kj.asBytes());
42
  kv.put("qux", "321"_kj.asBytes());
43
 
44
  KJ_EXPECT(sqliteObserver.rowsWritten == 4);
45
  KJ_EXPECT(sqliteObserver.rowsRead == 0);
46
 
47
  {
48
    bool called = false;
49
    KJ_EXPECT(kv.get("foo", [&](kj::ArrayPtr<const byte> value) {
50
      KJ_EXPECT(kj::str(value.asChars()) == "abc");
51
      called = true;
52
    }));
53
    KJ_EXPECT(called);
54
  }
55
 
56
  KJ_EXPECT(sqliteObserver.rowsWritten == 4);
57
  KJ_EXPECT(sqliteObserver.rowsRead == 1);
58
 
59
  {
60
    bool called = false;
61
    KJ_EXPECT(kv.get("bar", [&](kj::ArrayPtr<const byte> value) {
62
      KJ_EXPECT(kj::str(value.asChars()) == "def");
63
      called = true;
64
    }));
65
    KJ_EXPECT(called);
66
  }
67
 
68
  KJ_EXPECT(sqliteObserver.rowsWritten == 4);
69
  KJ_EXPECT(sqliteObserver.rowsRead == 2);
70
 
71
  KJ_EXPECT(!kv.get("corge", [&](kj::ArrayPtr<const byte> value) {
72
    KJ_FAIL_EXPECT("should not call callback when no match", value.asChars());
73
  }));
74
 
75
  KJ_EXPECT(sqliteObserver.rowsWritten == 4);
76
  KJ_EXPECT(sqliteObserver.rowsRead == 2);
77
 
78
  auto list = [&](auto&&... params) {
79
    kj::Vector<kj::String> results;
80
    auto callback = [&](kj::StringPtr key, kj::ArrayPtr<const byte> value) {
81
      results.add(kj::str(key, "=", value.asChars()));
82
    };
83
 
84
    auto n = kv.list(params..., callback);
85
    KJ_EXPECT(results.size() == n);
86
    return kj::strArray(results, ", ");
87
  };
88
 
89
  constexpr auto F = SqliteKv::FORWARD;
90
  constexpr auto R = SqliteKv::REVERSE;
91
 
92
  KJ_EXPECT(list(nullptr, kj::none, kj::none, F) == "bar=def, baz=123, foo=abc, qux=321");
93
  KJ_EXPECT(list("cat"_kj, kj::none, kj::none, F) == "foo=abc, qux=321");
94
  KJ_EXPECT(list("foo"_kj, kj::none, kj::none, F) == "foo=abc, qux=321");
95
  KJ_EXPECT(list("fop"_kj, kj::none, kj::none, F) == "qux=321");
96
  KJ_EXPECT(list("foo "_kj, kj::none, kj::none, F) == "qux=321");
97
 
98
  KJ_EXPECT(list(nullptr, "cat"_kj, kj::none, F) == "bar=def, baz=123");
99
  KJ_EXPECT(list(nullptr, "foo"_kj, kj::none, F) == "bar=def, baz=123");
100
  KJ_EXPECT(list(nullptr, "fop"_kj, kj::none, F) == "bar=def, baz=123, foo=abc");
101
 
102
  KJ_EXPECT(list(nullptr, kj::none, 2, F) == "bar=def, baz=123");
103
  KJ_EXPECT(list(nullptr, kj::none, 3, F) == "bar=def, baz=123, foo=abc");
104
  KJ_EXPECT(list("baz"_kj, kj::none, 2, F) == "baz=123, foo=abc");
105
  KJ_EXPECT(list(nullptr, "foo"_kj, 1, F) == "bar=def");
106
  KJ_EXPECT(list(nullptr, "foo"_kj, 2, F) == "bar=def, baz=123");
107
  KJ_EXPECT(list(nullptr, "foo"_kj, 3, F) == "bar=def, baz=123");
108
 
109
  KJ_EXPECT(list(nullptr, kj::none, kj::none, R) == "qux=321, foo=abc, baz=123, bar=def");
110
  KJ_EXPECT(list("foo"_kj, kj::none, kj::none, R) == "qux=321, foo=abc");
111
  KJ_EXPECT(list(nullptr, "foo"_kj, kj::none, R) == "baz=123, bar=def");
112
  KJ_EXPECT(list(nullptr, kj::none, 2, R) == "qux=321, foo=abc");
113
  KJ_EXPECT(list(nullptr, "foo"_kj, 1, R) == "baz=123");
114
 
115
  KJ_EXPECT(kv.delete_("baz"));
116
  KJ_EXPECT(!kv.delete_("corge"));
117
 
118
  KJ_EXPECT(list(nullptr, kj::none, kj::none, F) == "bar=def, foo=abc, qux=321");
119
 
120
  // Put can overwrite.
121
  kv.put("foo", "hello"_kj.asBytes());
122
  KJ_EXPECT(list(nullptr, kj::none, kj::none, F) == "bar=def, foo=hello, qux=321");
123
 
124
  // deleteAll()
125
  KJ_EXPECT(kv.deleteAll() == 3);
126
  KJ_EXPECT(list(nullptr, kj::none, kj::none, F) == "");
127
 
128
  KJ_EXPECT(!kv.get("bar", [&](kj::ArrayPtr<const byte> value) {
129
    KJ_FAIL_EXPECT("should not call callback when no match", value.asChars());
130
  }));
131
 
132
  kv.put("bar", "ghi"_kj.asBytes());
133
  kv.put("corge", "garply"_kj.asBytes());
134
 
135
  KJ_EXPECT(list(nullptr, kj::none, kj::none, F) == "bar=ghi, corge=garply");
136
 
137
  {
138
    bool called = false;
139
    KJ_EXPECT(kv.get("bar", [&](kj::ArrayPtr<const byte> value) {
140
      KJ_EXPECT(kj::str(value.asChars()) == "ghi");
141
      called = true;
142
    }));
143
    KJ_EXPECT(called);
144
  }
145
}
146
 
147
KJ_TEST("large key") {
148
  auto dir = kj::newInMemoryDirectory(kj::nullClock());
149
  SqliteDatabase::Vfs vfs(*dir);
150
  SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY);
151
  SqliteKv kv(db);
152
 
153
  // 2MB because we document a 2MB limit for SQLite Durable Objects
154
  kj::String closeToLimitString = kj::heapString(2000000);
155
  kv.put(closeToLimitString, "hello"_kj.asBytes());
156
 
157
  // Actual limit is 2.2MB, so we test more than that to see if it throws
158
  kj::String tooBigString = kj::heapString(2400000);
159
 
160
  KJ_EXPECT_THROW_MESSAGE(
161
      "string or blob too big: SQLITE_TOOBIG", kv.put(tooBigString, "hello"_kj.asBytes()));
162
}
163
 
164
KJ_TEST("SQLite-KV multi-put") {
165
  class TestSqliteObserver: public SqliteObserver {
166
   public:
167
    void addQueryStats(uint64_t read, uint64_t written) override {
168
      rowsRead += read;
169
      rowsWritten += written;
170
    }
171
 
172
    uint64_t rowsRead = 0;
173
    uint64_t rowsWritten = 0;
174
  };
175
 
176
  auto dir = kj::newInMemoryDirectory(kj::nullClock());
177
  SqliteDatabase::Vfs vfs(*dir);
178
  TestSqliteObserver sqliteObserver;
179
  SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY,
180
      /*sqliteMaxMemoryBytes=*/kj::maxValue, sqliteObserver);
181
  SqliteKv kv(db);
182
 
183
  // Test basic multi-put with a simple struct
184
  struct KeyValue {
185
    kj::StringPtr key;
186
    kj::ArrayPtr<const byte> value;
187
  };
188
 
189
  kj::Vector<KeyValue> pairs;
190
  pairs.add(KeyValue{"foo"_kj, "abc"_kj.asBytes()});
191
  pairs.add(KeyValue{"bar"_kj, "def"_kj.asBytes()});
192
  pairs.add(KeyValue{"baz"_kj, "123"_kj.asBytes()});
193
 
194
  kv.put(pairs, {.allowUnconfirmed = false});
195
 
196
  KJ_EXPECT(sqliteObserver.rowsWritten == 3);
197
 
198
  // Verify all values were written correctly
199
  bool called = false;
200
  KJ_EXPECT(kv.get("foo", [&](kj::ArrayPtr<const byte> value) {
201
    KJ_EXPECT(kj::str(value.asChars()) == "abc");
202
    called = true;
203
  }));
204
  KJ_EXPECT(called);
205
 
206
  called = false;
207
  KJ_EXPECT(kv.get("bar", [&](kj::ArrayPtr<const byte> value) {
208
    KJ_EXPECT(kj::str(value.asChars()) == "def");
209
    called = true;
210
  }));
211
  KJ_EXPECT(called);
212
 
213
  called = false;
214
  KJ_EXPECT(kv.get("baz", [&](kj::ArrayPtr<const byte> value) {
215
    KJ_EXPECT(kj::str(value.asChars()) == "123");
216
    called = true;
217
  }));
218
  KJ_EXPECT(called);
219
 
220
  // Test multi-put overwrites existing values
221
  kj::Vector<KeyValue> pairs2;
222
  pairs2.add(KeyValue{"foo"_kj, "xyz"_kj.asBytes()});
223
  pairs2.add(KeyValue{"bar"_kj, "uvw"_kj.asBytes()});
224
 
225
  kv.put(pairs2, {.allowUnconfirmed = false});
226
 
227
  called = false;
228
  KJ_EXPECT(kv.get("foo", [&](kj::ArrayPtr<const byte> value) {
229
    KJ_EXPECT(kj::str(value.asChars()) == "xyz");
230
    called = true;
231
  }));
232
  KJ_EXPECT(called);
233
 
234
  called = false;
235
  KJ_EXPECT(kv.get("bar", [&](kj::ArrayPtr<const byte> value) {
236
    KJ_EXPECT(kj::str(value.asChars()) == "uvw");
237
    called = true;
238
  }));
239
  KJ_EXPECT(called);
240
 
241
  // Verify other key unchanged
242
  called = false;
243
  KJ_EXPECT(kv.get("baz", [&](kj::ArrayPtr<const byte> value) {
244
    KJ_EXPECT(kj::str(value.asChars()) == "123");
245
    called = true;
246
  }));
247
  KJ_EXPECT(called);
248
 
249
  // Test empty multi-put (should succeed)
250
  kj::Vector<KeyValue> emptyPairs;
251
  kv.put(emptyPairs, {.allowUnconfirmed = false});
252
 
253
  // Verify database unchanged
254
  called = false;
255
  KJ_EXPECT(kv.get("foo", [&](kj::ArrayPtr<const byte> value) {
256
    KJ_EXPECT(kj::str(value.asChars()) == "xyz");
257
    called = true;
258
  }));
259
  KJ_EXPECT(called);
260
}
261
 
262
KJ_TEST("SQLite-KV multi-put rollback on error") {
263
  auto dir = kj::newInMemoryDirectory(kj::nullClock());
264
  SqliteDatabase::Vfs vfs(*dir);
265
  SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY);
266
  SqliteKv kv(db);
267
 
268
  // Pre-populate with some data
269
  kv.put("existing", "value"_kj.asBytes());
270
 
271
  struct KeyValue {
272
    kj::StringPtr key;
273
    kj::ArrayPtr<const byte> value;
274
  };
275
 
276
  // Create a multi-put that will fail due to a key being too large
277
  kj::Vector<KeyValue> pairs;
278
  pairs.add(KeyValue{"key1"_kj, "value1"_kj.asBytes()});
279
  pairs.add(KeyValue{"key2"_kj, "value2"_kj.asBytes()});
280
 
281
  // Add a key that exceeds the limit (2.2MB actual limit)
282
  kj::String tooBigString = kj::heapString(2400000);
283
  pairs.add(KeyValue{tooBigString, "value3"_kj.asBytes()});
284
 
285
  // The multi-put should throw
286
  KJ_EXPECT_THROW_MESSAGE(
287
      "string or blob too big: SQLITE_TOOBIG", kv.put(pairs, {.allowUnconfirmed = false}));
288
 
289
  // Verify that the first two keys were NOT written (transaction rolled back)
290
  KJ_EXPECT(!kv.get("key1", [&](kj::ArrayPtr<const byte> value) {
291
    KJ_FAIL_EXPECT("key1 should not exist after rollback");
292
  }));
293
 
294
  KJ_EXPECT(!kv.get("key2", [&](kj::ArrayPtr<const byte> value) {
295
    KJ_FAIL_EXPECT("key2 should not exist after rollback");
296
  }));
297
 
298
  // Verify existing data is unchanged
299
  bool called = false;
300
  KJ_EXPECT(kv.get("existing", [&](kj::ArrayPtr<const byte> value) {
301
    KJ_EXPECT(kj::str(value.asChars()) == "value");
302
    called = true;
303
  }));
304
  KJ_EXPECT(called);
305
}
306
 
307
KJ_TEST("SQLite-KV multi-put with allowUnconfirmed") {
308
  auto dir = kj::newInMemoryDirectory(kj::nullClock());
309
  SqliteDatabase::Vfs vfs(*dir);
310
  SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY);
311
  SqliteKv kv(db);
312
 
313
  struct KeyValue {
314
    kj::StringPtr key;
315
    kj::ArrayPtr<const byte> value;
316
  };
317
 
318
  kj::Vector<KeyValue> pairs;
319
  pairs.add(KeyValue{"foo"_kj, "abc"_kj.asBytes()});
320
  pairs.add(KeyValue{"bar"_kj, "def"_kj.asBytes()});
321
 
322
  // Test with allowUnconfirmed = true
323
  kv.put(pairs, {.allowUnconfirmed = true});
324
 
325
  // Verify values were written
326
  bool called = false;
327
  KJ_EXPECT(kv.get("foo", [&](kj::ArrayPtr<const byte> value) {
328
    KJ_EXPECT(kj::str(value.asChars()) == "abc");
329
    called = true;
330
  }));
331
  KJ_EXPECT(called);
332
 
333
  called = false;
334
  KJ_EXPECT(kv.get("bar", [&](kj::ArrayPtr<const byte> value) {
335
    KJ_EXPECT(kj::str(value.asChars()) == "def");
336
    called = true;
337
  }));
338
  KJ_EXPECT(called);
339
}
340
 
341
}  // namespace
342
}  // namespace workerd