Skip to content
rachel/workerd
File

Blob: src/workerd/api/sql.h

cpp362 lines
1// Copyright (c) 2017-2022 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 <workerd/api/actor-state.h>
8#include <workerd/io/compatibility-date.capnp.h>
9#include <workerd/io/io-context.h>
10#include <workerd/jsg/jsg.h>
11#include <workerd/util/sqlite.h>
12 
13namespace workerd::api {
14 
15class SqlStorage final: public jsg::Object, private SqliteDatabase::Regulator {
16 public:
17  SqlStorage(jsg::Ref<DurableObjectStorage> storage);
18  ~SqlStorage();
19 
20  using BindingValue = kj::Maybe<kj::OneOf<kj::Array<const byte>, kj::String, double>>;
21 
22  class Cursor;
23  class Statement;
24  struct IngestResult;
25 
26  // One value returned from SQL. Note that we intentionally return StringPtr instead of String
27  // because we know that the underlying buffer returned by SQLite will be valid long enough to be
28  // converted by JSG into a V8 string. For byte arrays, on the other hand, we pass ownership to
29  // JSG, which does not need to make a copy.
30  using SqlValue = kj::Maybe<kj::OneOf<kj::Array<byte>, kj::StringPtr, double>>;
31 
32  jsg::Ref<Cursor> exec(jsg::Lock& js, jsg::JsString query, jsg::Arguments<BindingValue> bindings);
33  IngestResult ingest(jsg::Lock& js, kj::String query);
34  void setMaxPageCountForTest(jsg::Lock& js, int count);
35 
36  jsg::Ref<Statement> prepare(jsg::Lock& js, jsg::JsString query);
37 
38  double getDatabaseSize(jsg::Lock& js);
39 
40  JSG_RESOURCE_TYPE(SqlStorage, CompatibilityFlags::Reader flags) {
41    JSG_METHOD(exec);
42 
43    if (flags.getWorkerdExperimental()) {
44      // Prepared statement API is experimental-only and deprecated. exec() will automatically
45      // handle caching prepared statements, so apps don't need to worry about it.
46      JSG_METHOD(prepare);
47 
48      // 'ingest' functionality is still experimental-only
49      JSG_METHOD(ingest);
50 
51      JSG_METHOD(setMaxPageCountForTest);
52    }
53 
54    JSG_READONLY_PROTOTYPE_PROPERTY(databaseSize, getDatabaseSize);
55 
56    JSG_NESTED_TYPE(Cursor);
57    JSG_NESTED_TYPE(Statement);
58 
59    JSG_TS_OVERRIDE({
60      exec<T extends Record<string, SqlStorageValue>>(query: string, ...bindings: any[]): SqlStorageCursor<T>
61    });
62  }
63 
64  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const;
65 
66 private:
67  void visitForGc(jsg::GcVisitor& visitor) {
68    visitor.visit(storage);
69  }
70 
71  bool isAllowedName(kj::StringPtr name) const override;
72  bool isAllowedTrigger(kj::StringPtr name) const override;
73  void onError(kj::Maybe<int> sqliteErrorCode, kj::StringPtr message) const override;
74  bool allowTransactions() const override;
75  bool shouldAddQueryStats() const override;
76 
77  SqliteDatabase& getDb(jsg::Lock& js) {
78    return storage->getSqliteDb(js);
79  }
80 
81  jsg::Ref<DurableObjectStorage> storage;
82 
83  kj::Maybe<uint> pageSize;
84  kj::Maybe<IoOwn<SqliteDatabase::Statement>> pragmaPageCount;
85  kj::Maybe<IoOwn<SqliteDatabase::Statement>> pragmaGetMaxPageCount;
86 
87  // A statement in the statement cache.
88  struct CachedStatement: public kj::Refcounted {
89    jsg::HashableV8Ref<v8::String> query;
90    size_t statementSize;
91    SqliteDatabase::Statement statement;
92    kj::ListLink<CachedStatement> lruLink;
93    uint useCount = 0;
94 
95    CachedStatement(jsg::Lock& js,
96        SqlStorage& sqlStorage,
97        SqliteDatabase& db,
98        jsg::JsString jsQuery,
99        kj::String kjQuery)
100        : query(js.v8Isolate, jsQuery),
101          statementSize(kjQuery.size()),
102          statement(db.prepareMulti(sqlStorage, kj::mv(kjQuery))) {}
103  };
104 
105  class StatementCacheCallbacks {
106   public:
107    inline const jsg::HashableV8Ref<v8::String>& keyForRow(
108        const kj::Rc<CachedStatement>& entry) const {
109      return entry->query;
110    }
111 
112    inline bool matches(const kj::Rc<CachedStatement>& entry, jsg::JsString key) const {
113      return entry->query == key;
114    }
115    inline bool matches(
116        const kj::Rc<CachedStatement>& entry, const jsg::HashableV8Ref<v8::String>& key) const {
117      return entry->query == key;
118    }
119 
120    inline auto hashCode(jsg::JsString key) const {
121      return key.hashCode();
122    }
123    inline auto hashCode(const jsg::HashableV8Ref<v8::String>& key) const {
124      return key.hashCode();
125    }
126  };
127 
128  // We can't quite just use kj::HashMap here because we want the table key to be
129  // `CachedStatement::query`, which is a member of the refcounted object.
130  using StatementMap = kj::Table<kj::Rc<CachedStatement>, kj::HashIndex<StatementCacheCallbacks>>;
131 
132  struct StatementCache {
133    StatementMap map;
134    kj::List<CachedStatement, &CachedStatement::lruLink> lru;
135    size_t totalSize = 0;
136 
137    ~StatementCache() noexcept(false);
138  };
139  IoOwn<StatementCache> statementCache;
140 
141  template <size_t size, typename... Params>
142  SqliteDatabase::Query execMemoized(SqliteDatabase& db,
143      kj::Maybe<IoOwn<SqliteDatabase::Statement>>& slot,
144      const char (&sqlCode)[size],
145      Params&&... params) {
146    // Run a (trusted) statement, preparing it on the first call and reusing the prepared version
147    // for future calls.
148 
149    SqliteDatabase::Statement* stmt;
150    KJ_IF_SOME(s, slot) {
151      stmt = &*s;
152    } else {
153      stmt = &*slot.emplace(IoContext::current().addObject(kj::heap(db.prepare(sqlCode))));
154    }
155    return stmt->run(kj::fwd<Params>(params)...);
156  }
157 
158  uint64_t getPageSize(SqliteDatabase& db) {
159    KJ_IF_SOME(p, pageSize) {
160      return p;
161    } else {
162      return pageSize.emplace(db.run("PRAGMA page_size;").getInt64(0));
163    }
164  }
165 
166  // Utility functions to convert SqlValue to a JS value. We can't just return the C++ values and
167  // let JSG do the work because we're trying to avoid having to make a copy of string contents out
168  // of SQLite's buffer when the conversion to JS is just going to make another copy. We can't use
169  // jsg::TypeHandler because SqlValue contains StringPtr, which doesn't support unwrapping. We
170  // don't actually ever use unwrapping, but requesting a TypeHandler forces JSG to try to generate
171  // the code for unwrapping, leading to compiler errors.
172  //
173  // TODO(cleanup): Think hard about how to make JSG support this better. Part of the problem is
174  //   that we're being too clever with optimizations to avoid copying strings when we don't need
175  //   to.
176  static jsg::JsValue wrapSqlValue(jsg::Lock& js, SqlValue value);
177};
178 
179class SqlStorage::Cursor final: public jsg::Object {
180 public:
181  template <typename... Params>
182  Cursor(jsg::Lock& js, kj::Maybe<kj::Function<void(Cursor&)>> doneCb, Params&&... params)
183      : doneCallback(kj::mv(doneCb)) {
184    auto stateObj = kj::heap<State>(kj::fwd<Params>(params)...);
185    initColumnNames(js, *stateObj);
186    if (stateObj->query.isDone()) {
187      endQuery(*stateObj);
188    } else {
189      state = IoContext::current().addObject(kj::mv(stateObj));
190    }
191  }
192  ~Cursor() noexcept(false);
193 
194  double getRowsRead();
195  double getRowsWritten();
196 
197  jsg::JsArray getColumnNames(jsg::Lock& js);
198  JSG_RESOURCE_TYPE(Cursor, CompatibilityFlags::Reader flags) {
199    JSG_METHOD(next);
200    JSG_METHOD(toArray);
201    JSG_METHOD(one);
202 
203    JSG_ITERABLE(rows);
204    JSG_METHOD(raw);
205    JSG_READONLY_PROTOTYPE_PROPERTY(columnNames, getColumnNames);
206    JSG_READONLY_PROTOTYPE_PROPERTY(rowsRead, getRowsRead);
207    JSG_READONLY_PROTOTYPE_PROPERTY(rowsWritten, getRowsWritten);
208 
209    JSG_TS_DEFINE(type SqlStorageValue = ArrayBuffer | string | number | null);
210    JSG_TS_OVERRIDE(<T extends Record<string, SqlStorageValue>> {
211      [Symbol.iterator](): IterableIterator<T>;
212      raw<U extends SqlStorageValue[]>(): IterableIterator<U>;
213      next(): { done?: false, value: T } | { done: true, value?: never };
214      toArray(): T[];
215      one(): T;
216      columnNames: string[];
217    });
218 
219    if (flags.getWorkerdExperimental()) {
220      JSG_READONLY_PROTOTYPE_PROPERTY(reusedCachedQueryForTest, getReusedCachedQueryForTest);
221    }
222  }
223 
224  JSG_ITERATOR(RowIterator, rows, jsg::JsObject, jsg::Ref<Cursor>, rowIteratorNext);
225  JSG_ITERATOR(RawIterator, raw, jsg::JsArray, jsg::Ref<Cursor>, rawIteratorNext);
226 
227  RowIterator::Next next(jsg::Lock& js);
228  jsg::JsArray toArray(jsg::Lock& js);
229  jsg::JsValue one(jsg::Lock& js);
230 
231  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
232    if (state != kj::none) {
233      tracker.trackFieldWithSize("IoOwn<State>", sizeof(IoOwn<State>));
234    }
235    tracker.trackField("columnNames", columnNames);
236  }
237 
238  bool getReusedCachedQueryForTest() {
239    return reusedCachedQuery;
240  }
241 
242 private:
243  struct State {
244    kj::Maybe<kj::Rc<CachedStatement>> cachedStatement;
245 
246    // The bindings that were used to construct `query`. We have to keep these alive until the query
247    // is done since it might contain pointers into strings and blobs.
248    kj::Array<BindingValue> bindings;
249 
250    SqliteDatabase::Query query;
251 
252    State(SqliteDatabase& db,
253        SqliteDatabase::Regulator& regulator,
254        kj::StringPtr sqlCode,
255        kj::Array<BindingValue> bindings);
256 
257    State(kj::Rc<CachedStatement> cachedStatement, kj::Array<BindingValue> bindings);
258  };
259 
260  // Nulled out when query is done or canceled.
261  kj::Maybe<IoOwn<State>> state;
262 
263  // Called when the query is done or canceled.
264  kj::Maybe<kj::Function<void(Cursor&)>> doneCallback;
265 
266  // True if the cursor was canceled by a new call to the same statement. This is used only to
267  // flag an error if the application tries to reuse the cursor.
268  bool canceled = false;
269 
270  // Did we reuse a query from the query cache? Tracked for testing purposes.
271  bool reusedCachedQuery = false;
272 
273  // Reference to a weak reference that might point back to this object. If so, null it out at
274  // destruction. Used by Statement to invalidate past cursors when the statement is
275  // executed again.
276  kj::Maybe<kj::Maybe<Cursor&>&> selfRef;
277 
278  // Row IO counts. These are updated as the query runs. We keep these outside the State so they
279  // remain available even after the query is done or canceled.
280  uint64_t rowsRead = 0;
281  // Row IO counts. These are updated as the query runs. We keep these outside the State so they
282  // remain available even after the query is done or canceled.
283  uint64_t rowsWritten = 0;
284 
285  jsg::JsRef<jsg::JsArray> columnNames;
286 
287  // Invoke when `query.isDone()`, or when we want to prematurely cancel the query. This records
288  // row counters and then sets `state` to `none` to drop the query and return the prepared
289  // statement to the statement cache.
290  void endQuery(State& stateRef);
291 
292  // Initialize `columnNames` from the state object.
293  void initColumnNames(jsg::Lock& js, State& stateRef);
294 
295  static kj::Array<const SqliteDatabase::Query::ValuePtr> mapBindings(
296      kj::ArrayPtr<BindingValue> values);
297 
298  static kj::Maybe<jsg::JsObject> rowIteratorNext(jsg::Lock& js, jsg::Ref<Cursor>& obj);
299  static kj::Maybe<jsg::JsArray> rawIteratorNext(jsg::Lock& js, jsg::Ref<Cursor>& obj);
300  static kj::Maybe<v8::LocalVector<v8::Value>> iteratorImpl(jsg::Lock& js, jsg::Ref<Cursor>& obj);
301 
302  friend class Statement;
303 
304  void visitForGc(jsg::GcVisitor& visitor) {
305    visitor.visit(columnNames);
306  }
307};
308 
309// The prepared statement API is supported only for backwards compatibility for certain early
310// internal users of SQLite-backed DOs. This API was not released because we chose instead to
311// implement automatic prepared statement caching via the simple `exec()` API. Since this is
312// a compatibility shim only, to simplify things, it is actually just a wrapper around `exec()`.
313class SqlStorage::Statement final: public jsg::Object {
314 public:
315  Statement(jsg::Lock& js, jsg::Ref<SqlStorage> sqlStorage, jsg::JsString query)
316      : sqlStorage(kj::mv(sqlStorage)),
317        // Internalize the string before constructing the statement so that it doesn't have to
318        // re-lookup the internalized string for every invocation.
319        query(js.v8Isolate, query.internalize(js)) {}
320 
321  jsg::Ref<Cursor> run(jsg::Lock& js, jsg::Arguments<BindingValue> bindings);
322 
323  JSG_RESOURCE_TYPE(Statement) {
324    JSG_CALLABLE(run);
325  }
326 
327  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
328    tracker.trackField("sqlStorage", sqlStorage);
329    tracker.trackField("query", query);
330  }
331 
332 private:
333  jsg::Ref<SqlStorage> sqlStorage;
334  jsg::V8Ref<v8::String> query;
335 
336  friend class Cursor;
337};
338 
339struct SqlStorage::IngestResult {
340  IngestResult(kj::String remainder, double rowsRead, double rowsWritten, double statementCount)
341      : remainder(kj::mv(remainder)),
342        rowsRead(rowsRead),
343        rowsWritten(rowsWritten),
344        statementCount(statementCount) {}
345 
346  kj::String remainder;
347  double rowsRead;
348  double rowsWritten;
349  double statementCount;
350 
351  JSG_STRUCT(remainder, rowsRead, rowsWritten, statementCount);
352};
353 
354#define EW_SQL_ISOLATE_TYPES                                                                       \
355  api::SqlStorage, api::SqlStorage::Statement, api::SqlStorage::Cursor,                            \
356      api::SqlStorage::IngestResult, api::SqlStorage::Cursor::RowIterator,                         \
357      api::SqlStorage::Cursor::RowIterator::Next, api::SqlStorage::Cursor::RawIterator,            \
358      api::SqlStorage::Cursor::RawIterator::Next
359// The list of sql.h types that are added to worker.c++'s JSG_DECLARE_ISOLATE_TYPE
360 
361}  // namespace workerd::api