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
#pragma once
6
 
7
#include "sqlite.h"
8
 
9
#include <kj/debug.h>
10
#include <kj/exception.h>
11
 
12
namespace workerd {
13
 
14
// Small class which is used to customize certain aspects of the underlying sql operations
15
// In this case we just customize the error reporting to emit JSG user visible errors instead
16
// of KJ exceptions which become internal errors.
17
class SqliteKvRegulator: public SqliteDatabase::Regulator {
18
  void onError(kj::Maybe<int> sqliteErrorCode, kj::StringPtr message) const override;
19
 
20
  // We bill for KV operations as rows read/written.
21
  virtual bool shouldAddQueryStats() const override {
22
    return true;
23
  }
24
};
25
 
26
// Class which implements KV storage on top of SQLite. This is intended to be used for Durable
27
// Object storage.
28
//
29
// The table is named `_cf_KV`. The naming is designed so that if the application is allowed to
30
// perform direct SQL queries, we can block it from accessing any table prefixed with `_cf_`.
31
// (Ideally this class would allow configuring the table name, but this would require a somewhat
32
// obnoxious amount of string allocation.)
33
class SqliteKv: private SqliteDatabase::ResetListener {
34
 public:
35
  explicit SqliteKv(SqliteDatabase& db);
36
  ~SqliteKv() noexcept(false);
37
 
38
  using KeyPtr = kj::StringPtr;
39
  using ValuePtr = kj::ArrayPtr<const kj::byte>;
40
 
41
  // Search for a match for the given key. Calls the callback function with the result (a ValuePtr)
42
  // if found. This is intended to avoid the need to copy the bytes, if the caller would just parse
43
  // them and drop them immediately anyway. Returns true if there was a match, false if not.
44
  template <typename Func>
45
  bool get(KeyPtr key, Func&& callback);
46
 
47
  enum Order { FORWARD, REVERSE };
48
 
49
  // Search for all known keys and values in a range, calling the callback (with KeyPtr and
50
  // ValuePtr parameters) for each one seen. `end` and `limit` can be null to request no constraint
51
  // be enforced.
52
  template <typename Func>
53
  uint list(
54
      KeyPtr begin, kj::Maybe<KeyPtr> end, kj::Maybe<uint> limit, Order order, Func&& callback);
55
 
56
  // List returning a cursor which can be iterated one at a time.
57
  class ListCursor;
58
  kj::Own<ListCursor> list(KeyPtr begin, kj::Maybe<KeyPtr> end, kj::Maybe<uint> limit, Order order);
59
 
60
  struct WriteOptions {
61
    bool allowUnconfirmed = false;
62
  };
63
 
64
  // Store a value into the table.
65
  void put(KeyPtr key, ValuePtr value);
66
  void put(KeyPtr key, ValuePtr value, WriteOptions options);
67
 
68
  // Atomically store multiple values into the table.
69
  //
70
  // ArrayOfKeyValuePair should be a type that allows iteration of a struct that has two members,
71
  // key and value, that can be coerced into KeyPtr and ValuePtr, respectively.  I'm using a
72
  // template so that we don't have to transform (by copy) the values passed in from higher levels
73
  // while also preventing this module from taking a dependency on types from higher levels.
74
  template <typename ArrayOfKeyValuePair>
75
  void put(ArrayOfKeyValuePair& pairs, WriteOptions options);
76
 
77
  // Delete the key and return whether it was matched.
78
  bool delete_(KeyPtr key);
79
  bool delete_(KeyPtr key, WriteOptions options);
80
 
81
  uint deleteAll();
82
 
83
  // TODO(perf): Should we provide multi-get, multi-put, and multi-delete? It's a bit tricky to
84
  //   implement them as single SQL queries, while still using prepared statements. The c-array
85
  //   extension might help here, though it can only support arrays of NUL-terminated strings, not
86
  //   byte blobs or strings containing NUL bytes.
87
 
88
 private:
89
  struct Uninitialized {};
90
 
91
  struct Initialized {
92
    // This reference is redundant but storing it here makes the prepared statement code below
93
    // easier to manage.
94
    SqliteDatabase& db;
95
 
96
    SqliteKvRegulator regulator;
97
 
98
    SqliteDatabase::Statement stmtGet = db.prepare(regulator, R"(
99
      SELECT value FROM _cf_KV WHERE key = ?
100
    )");
101
    SqliteDatabase::Statement stmtPut = db.prepare(regulator, R"(
102
      INSERT INTO _cf_KV VALUES(?, ?)
103
        ON CONFLICT DO UPDATE SET value = excluded.value;
104
    )");
105
    SqliteDatabase::Statement stmtDelete = db.prepare(regulator, R"(
106
      DELETE FROM _cf_KV WHERE key = ?
107
    )");
108
    SqliteDatabase::Statement stmtList = db.prepare(regulator, R"(
109
      SELECT * FROM _cf_KV
110
      WHERE key >= ?
111
      ORDER BY key
112
    )");
113
    SqliteDatabase::Statement stmtListEnd = db.prepare(regulator, R"(
114
      SELECT * FROM _cf_KV
115
      WHERE key >= ? AND key < ?
116
      ORDER BY key
117
    )");
118
    SqliteDatabase::Statement stmtListLimit = db.prepare(regulator, R"(
119
      SELECT * FROM _cf_KV
120
      WHERE key >= ?
121
      ORDER BY key
122
      LIMIT ?
123
    )");
124
    SqliteDatabase::Statement stmtListEndLimit = db.prepare(regulator, R"(
125
      SELECT * FROM _cf_KV
126
      WHERE key >= ? AND key < ?
127
      ORDER BY key
128
      LIMIT ?
129
    )");
130
    SqliteDatabase::Statement stmtListReverse = db.prepare(regulator, R"(
131
      SELECT * FROM _cf_KV
132
      WHERE key >= ?
133
      ORDER BY key DESC
134
    )");
135
    SqliteDatabase::Statement stmtListEndReverse = db.prepare(regulator, R"(
136
      SELECT * FROM _cf_KV
137
      WHERE key >= ? AND key < ?
138
      ORDER BY key DESC
139
    )");
140
    SqliteDatabase::Statement stmtListLimitReverse = db.prepare(regulator, R"(
141
      SELECT * FROM _cf_KV
142
      WHERE key >= ?
143
      ORDER BY key DESC
144
      LIMIT ?
145
    )");
146
    SqliteDatabase::Statement stmtListEndLimitReverse = db.prepare(regulator, R"(
147
      SELECT * FROM _cf_KV
148
      WHERE key >= ? AND key < ?
149
      ORDER BY key DESC
150
      LIMIT ?
151
    )");
152
    SqliteDatabase::Statement stmtCountKeys = db.prepare(regulator, R"(
153
      SELECT count(*) FROM _cf_KV
154
    )");
155
    SqliteDatabase::Statement stmtMultiPutSavepoint = db.prepare(regulator, R"(
156
      SAVEPOINT _cf_put_multiple_savepoint
157
    )");
158
    SqliteDatabase::Statement stmtMultiPutRelease = db.prepare(regulator, R"(
159
      RELEASE _cf_put_multiple_savepoint
160
    )");
161
 
162
    Initialized(SqliteDatabase& db): db(db) {}
163
  };
164
 
165
  kj::OneOf<Uninitialized, Initialized> state;
166
 
167
  // Has the _cf_KV table been created? This is separate from Uninitialized/Initialized since it
168
  // has to be repeated after a reset, whereas the statements do not need to be recreated.
169
  bool tableCreated = false;
170
 
171
  kj::Maybe<ListCursor&> currentCursor;
172
 
173
  void cancelCurrentCursor();
174
 
175
  Initialized& ensureInitialized(bool allowUnconfirmed);
176
  // Make sure the KV table is created and prepared statements are ready. Not called until the
177
  // first write.
178
 
179
  void beforeSqliteReset() override;
180
 
181
  // Helper function that rolls back a multi-put statement and swallows any exceptions that may
182
  // occur during the rollback.
183
  void rollbackMultiPut(Initialized& stmts, WriteOptions options);
184
};
185
 
186
// Iterator over list results.
187
class SqliteKv::ListCursor {
188
 public:
189
  template <typename... Params>
190
  ListCursor(kj::Badge<SqliteKv>, SqliteKv& parent, Params&&... params) {
191
    parent.cancelCurrentCursor();
192
    state.emplace(parent, kj::fwd<Params>(params)...);
193
    parent.currentCursor = *this;
194
  }
195
  ListCursor(decltype(nullptr)) {}
196
 
197
  template <typename Func>
198
  uint forEach(Func&& callback) {
199
    auto& query = KJ_UNWRAP_OR(state, return 0).query;
200
    size_t count = 0;
201
    while (!query.isDone()) {
202
      callback(query.getText(0), query.getBlob(1));
203
      query.nextRow();
204
      ++count;
205
    }
206
    return count;
207
  };
208
 
209
  struct KeyValuePair {
210
    kj::StringPtr key;
211
    kj::ArrayPtr<const byte> value;
212
  };
213
  kj::Maybe<KeyValuePair> next();
214
 
215
  // If true, the cursor was canceled due to a new list() operation starting. Only one list() is
216
  // allowed at a time.
217
  bool wasCanceled() {
218
    return canceled;
219
  }
220
 
221
 private:
222
  struct State {
223
    SqliteKv& parent;
224
    SqliteDatabase::Query query;
225
 
226
    template <typename... Params>
227
    State(SqliteKv& parent, SqliteDatabase::Statement& stmt, Params&&... params)
228
        : parent(parent),
229
          query(stmt.run(kj::fwd<Params>(params)...)) {}
230
    ~State() noexcept(false) {
231
      parent.currentCursor = kj::none;
232
    }
233
  };
234
 
235
  kj::Maybe<State> state;
236
 
237
  // Are we at the beginning of the list?
238
  bool first = true;
239
 
240
  bool canceled = false;
241
 
242
  friend class SqliteKv;
243
};
244
 
245
// =======================================================================================
246
// inline implementation details
247
//
248
// We define these two methods as templates rather than use kj::Function since they're not too
249
// complicated and avoiding the virtual call is nice. Plus in list()'s case, the actual call sites
250
// pass constants for `order` so the `order ==` branch can be eliminated.
251
 
252
template <typename Func>
253
bool SqliteKv::get(KeyPtr key, Func&& callback) {
254
  if (!tableCreated) return false;
255
  auto& stmts = KJ_UNWRAP_OR(state.tryGet<Initialized>(), return false);
256
 
257
  auto query = stmts.stmtGet.run(key);
258
 
259
  if (query.isDone()) {
260
    return false;
261
  } else {
262
    callback(query.getBlob(0));
263
    return true;
264
  }
265
}
266
 
267
template <typename Func>
268
uint SqliteKv::list(
269
    KeyPtr begin, kj::Maybe<KeyPtr> end, kj::Maybe<uint> limit, Order order, Func&& callback) {
270
  return list(begin, end, limit, order)->forEach(kj::fwd<Func>(callback));
271
}
272
 
273
template <typename ArrayOfKeyValuePair>
274
void SqliteKv::put(ArrayOfKeyValuePair& pairs, WriteOptions options) {
275
  // TODO(cleanup): This code is very similar to DurableObjectStorage::transactionSync.  Perhaps the
276
  // general structure can be shared somehow?
277
  auto& stmts = ensureInitialized(options.allowUnconfirmed);
278
  stmts.stmtMultiPutSavepoint.run({.allowUnconfirmed = options.allowUnconfirmed});
279
 
280
  {
281
    // If any of the puts throw an exception, rollback the transaction and re-throw the exception
282
    // from the put that failed.
283
    KJ_ON_SCOPE_FAILURE(rollbackMultiPut(stmts, options));
284
    for (const auto& pair: pairs) {
285
      put(pair.key, pair.value, {.allowUnconfirmed = options.allowUnconfirmed});
286
    }
287
  }
288
  stmts.stmtMultiPutRelease.run({.allowUnconfirmed = options.allowUnconfirmed});
289
}
290
 
291
}  // namespace workerd