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 <workerd/jsg/exception.h>
8
 
9
#include <sqlite3.h>
10
 
11
namespace workerd {
12
 
13
void SqliteKvRegulator::onError(kj::Maybe<int> sqliteErrorCode, kj::StringPtr message) const {
14
  KJ_IF_SOME(ec, sqliteErrorCode) {
15
    switch (ec) {
16
      case SQLITE_TOOBIG:
17
        // We want to return SQLITE_TOOBIG to the user since it's usually because of user error.
18
        JSG_ASSERT(false, Error, message);
19
        return;
20
      // We don't want to return other errors to the user since they're usually our fault.
21
      // In that case we do nothing because the contract of onError is not to handle the error in
22
      // its entirety, but instead to optionally handle it, and do nothing otherwise.
23
      // When onError does nothing, the code calling into onError is still responsible for
24
      // handling the error by other means, usually by throwing a KJ exception itself.
25
      default:
26
        return;
27
    }
28
    KJ_UNREACHABLE;
29
  }
30
}
31
 
32
SqliteKv::SqliteKv(SqliteDatabase& db): ResetListener(db) {
33
  if (db.run("SELECT name FROM sqlite_master WHERE type='table' AND name='_cf_KV'").isDone()) {
34
    // The _cf_KV table doesn't exist. Defer initialization.
35
    state.init<Uninitialized>(Uninitialized{});
36
  } else {
37
    // The KV table was initialized in the past. We can go ahead and prepare our statements.
38
    // (We don't call ensureInitialized() here because the `CREATE TABLE IF NOT EXISTS` query it
39
    // executes would be redundant.)
40
    tableCreated = true;
41
    state.init<Initialized>(db);
42
  }
43
}
44
 
45
SqliteKv::~SqliteKv() noexcept(false) {
46
  cancelCurrentCursor();
47
}
48
 
49
void SqliteKv::cancelCurrentCursor() {
50
  KJ_IF_SOME(c, currentCursor) {
51
    c.state = kj::none;
52
    c.canceled = true;
53
  }
54
}
55
 
56
SqliteKv::Initialized& SqliteKv::ensureInitialized(bool allowUnconfirmed) {
57
  if (!tableCreated) {
58
    db.run(SqliteDatabase::QueryOptions{.regulator = SqliteDatabase::TRUSTED,
59
             .allowUnconfirmed = allowUnconfirmed},
60
        R"(
61
      CREATE TABLE IF NOT EXISTS _cf_KV (
62
        key TEXT PRIMARY KEY,
63
        value BLOB
64
      ) WITHOUT ROWID;
65
    )");
66
 
67
    tableCreated = true;
68
 
69
    // If we're in a transaction and it gets rolled back, we better mark that the table is actually
70
    // not created anymore.
71
    db.onRollback([this]() { tableCreated = false; });
72
  }
73
 
74
  KJ_SWITCH_ONEOF(state) {
75
    KJ_CASE_ONEOF(uninitialized, Uninitialized) {
76
      return state.init<Initialized>(db);
77
    }
78
    KJ_CASE_ONEOF(initialized, Initialized) {
79
      return initialized;
80
    }
81
  }
82
  KJ_UNREACHABLE;
83
}
84
 
85
kj::Own<SqliteKv::ListCursor> SqliteKv::list(
86
    KeyPtr begin, kj::Maybe<KeyPtr> end, kj::Maybe<uint> limit, Order order) {
87
  if (!tableCreated) return kj::heap<ListCursor>(nullptr);
88
  auto& stmts = KJ_UNWRAP_OR(state.tryGet<Initialized>(), return kj::heap<ListCursor>(nullptr));
89
 
90
  if (order == Order::FORWARD) {
91
    KJ_IF_SOME(e, end) {
92
      KJ_IF_SOME(l, limit) {
93
        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListEndLimit, begin, e,
94
            static_cast<int64_t>(l));
95
      } else {
96
        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListEnd, begin, e);
97
      }
98
    } else {
99
      KJ_IF_SOME(l, limit) {
100
        return kj::heap<ListCursor>(
101
            kj::Badge<SqliteKv>(), *this, stmts.stmtListLimit, begin, static_cast<int64_t>(l));
102
      } else {
103
        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtList, begin);
104
      }
105
    }
106
  } else {
107
    KJ_IF_SOME(e, end) {
108
      KJ_IF_SOME(l, limit) {
109
        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListEndLimitReverse,
110
            begin, e, static_cast<int64_t>(l));
111
      } else {
112
        return kj::heap<ListCursor>(
113
            kj::Badge<SqliteKv>(), *this, stmts.stmtListEndReverse, begin, e);
114
      }
115
    } else {
116
      KJ_IF_SOME(l, limit) {
117
        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListLimitReverse, begin,
118
            static_cast<int64_t>(l));
119
      } else {
120
        return kj::heap<ListCursor>(kj::Badge<SqliteKv>(), *this, stmts.stmtListReverse, begin);
121
      }
122
    }
123
  }
124
}
125
 
126
kj::Maybe<SqliteKv::ListCursor::KeyValuePair> SqliteKv::ListCursor::next() {
127
  auto& state = KJ_UNWRAP_OR(this->state, return kj::none);
128
  if (first) {
129
    first = false;
130
  } else {
131
    state.query.nextRow();
132
  }
133
  if (state.query.isDone()) {
134
    this->state = kj::none;
135
    return kj::none;
136
  }
137
 
138
  return KeyValuePair{state.query.getText(0), state.query.getBlob(1)};
139
}
140
 
141
void SqliteKv::put(KeyPtr key, ValuePtr value) {
142
  put(key, value, {});
143
}
144
 
145
void SqliteKv::put(KeyPtr key, ValuePtr value, WriteOptions options) {
146
  ensureInitialized(options.allowUnconfirmed)
147
      .stmtPut.run({.allowUnconfirmed = options.allowUnconfirmed}, key, value);
148
}
149
 
150
bool SqliteKv::delete_(KeyPtr key) {
151
  return delete_(key, {});
152
}
153
 
154
bool SqliteKv::delete_(KeyPtr key, WriteOptions options) {
155
  auto query = ensureInitialized(options.allowUnconfirmed)
156
                   .stmtDelete.run({.allowUnconfirmed = options.allowUnconfirmed}, key);
157
  return query.changeCount() > 0;
158
}
159
 
160
uint SqliteKv::deleteAll() {
161
  // TODO(perf): Consider introducing a compatibility flag that causes deleteAll() to always return
162
  //   1. Apps almost certainly don't care about the return value but historically we returned the
163
  //   count of keys deleted, so now we're stuck counting the table size for no good reason.
164
  uint count = tableCreated ? ensureInitialized(false).stmtCountKeys.run().getInt(0) : 0;
165
  db.reset();
166
  return count;
167
}
168
 
169
void SqliteKv::beforeSqliteReset() {
170
  // We'll need to recreate the table on the next operation.
171
  tableCreated = false;
172
}
173
 
174
void SqliteKv::rollbackMultiPut(Initialized& stmts, WriteOptions options) {
175
  KJ_IF_SOME(e, kj::runCatchingExceptions([&]() {
176
    // This should be rare, so we don't prepare a statement for it.
177
    stmts.db.run({.regulator = stmts.regulator, .allowUnconfirmed = options.allowUnconfirmed},
178
        kj::str("ROLLBACK TO _cf_put_multiple_savepoint"));
179
    stmts.stmtMultiPutRelease.run({.allowUnconfirmed = options.allowUnconfirmed});
180
  })) {
181
    KJ_LOG(WARNING, "silencing exception encountered while rolling back multi-put", e);
182
  }
183
}
184
 
185
}  // namespace workerd