SQLite: Handle automated expiration

This commit is contained in:
Tim Wojtulewicz 2024-08-12 13:23:24 -07:00
parent e95784db16
commit ec49f5d550
5 changed files with 88 additions and 9 deletions

View file

@ -42,7 +42,7 @@ ErrorResult SQLite::DoOpen(RecordValPtr options) {
}
std::string create = "create table if not exists " + table_name + " (";
create.append("key_str text primary key, value_str text not null);");
create.append("key_str text primary key, value_str text not null, expire_time real);");
char* errorMsg = nullptr;
if ( int res = sqlite3_exec(db, create.c_str(), NULL, NULL, &errorMsg); res != SQLITE_OK ) {
@ -54,12 +54,15 @@ ErrorResult SQLite::DoOpen(RecordValPtr options) {
}
static std::map<std::string, std::string> statements =
{{"put", util::fmt("insert into %s (key_str, value_str) values(?, ?)", table_name.c_str())},
{"put_update", util::fmt("insert into %s (key_str, value_str) values(?, ?) ON CONFLICT(key_str) "
{{"put", util::fmt("insert into %s (key_str, value_str, expire_time) values(?, ?, ?)", table_name.c_str())},
{"put_update",
util::fmt("insert into %s (key_str, value_str, expire_time) values(?, ?, ?) ON CONFLICT(key_str) "
"DO UPDATE SET value_str=?",
table_name.c_str())},
{"get", util::fmt("select value_str from %s where key_str=?", table_name.c_str())},
{"erase", util::fmt("delete from %s where key_str=?", table_name.c_str())}};
{"erase", util::fmt("delete from %s where key_str=?", table_name.c_str())},
{"expire", util::fmt("delete from %s where expire_time > 0 and expire_time != 0 and expire_time <= ?",
table_name.c_str())}};
for ( const auto& [key, stmt] : statements ) {
sqlite3_stmt* ps;
@ -123,8 +126,13 @@ ErrorResult SQLite::DoPut(ValPtr key, ValPtr value, bool overwrite, double expir
return res;
}
if ( auto res = checkError(sqlite3_bind_double(stmt, 3, expiration_time)); res.has_value() ) {
sqlite3_reset(stmt);
return res;
}
if ( overwrite ) {
if ( auto res = checkError(sqlite3_bind_text(stmt, 3, value_str.data(), value_str.size(), SQLITE_STATIC));
if ( auto res = checkError(sqlite3_bind_text(stmt, 4, value_str.data(), value_str.size(), SQLITE_STATIC));
res.has_value() ) {
sqlite3_reset(stmt);
return res;
@ -155,7 +163,7 @@ ValResult SQLite::DoGet(ValPtr key, ValResultCallback* cb) {
if ( auto res = checkError(sqlite3_bind_text(stmt, 1, key_str.data(), key_str.size(), SQLITE_STATIC));
res.has_value() ) {
sqlite3_reset(stmt);
return nonstd::unexpected<std::string>(res.value());
return zeek::unexpected<std::string>(res.value());
}
int errorcode = sqlite3_step(stmt);
@ -200,6 +208,23 @@ ErrorResult SQLite::DoErase(ValPtr key, ErrorResultCallback* cb) {
return std::nullopt;
}
/**
* Removes any entries in the backend that have expired. Can be overridden by
* derived classes.
*/
void SQLite::Expire() {
auto stmt = prepared_stmts["expire"];
if ( auto res = checkError(sqlite3_bind_double(stmt, 1, util::current_time())); res.has_value() ) {
sqlite3_reset(stmt);
// TODO: do something with the error here?
}
if ( auto res = checkError(sqlite3_step(stmt)); res.has_value() ) {
// TODO: do something with the error here?
}
}
// returns true in case of error
ErrorResult SQLite::checkError(int code) {
if ( code != SQLITE_OK && code != SQLITE_DONE ) {

View file

@ -48,7 +48,11 @@ public:
*/
ErrorResult DoErase(ValPtr key, ErrorResultCallback* cb = nullptr) override;
// TODO: add support for checking for expired data
/**
* Removes any entries in the backend that have expired. Can be overridden by
* derived classes.
*/
void Expire() override;
private:
ErrorResult checkError(int code);

View file

@ -0,0 +1,3 @@
### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63.
error in /Users/tim/Desktop/projects/storage-framework/testing/btest/.tmp/scripts.base.frameworks.storage.expiration/expiration.zeek, line 20: Failed to retrieve data: Failed to find row for key: no more rows available (Storage::get(backend, to_any_coerce key, F))
received termination signal

View file

@ -0,0 +1,5 @@
### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63.
put result, T
get result, value7890
get result same as inserted, T
get result, F

View file

@ -0,0 +1,42 @@
# @TEST-DOC: Automatic expiration of stored data
# @TEST-EXEC: zcat <$TRACES/echo-connections.pcap.gz | zeek -b %INPUT > out
# @TEST-EXEC: btest-diff out
# @TEST-EXEC: btest-diff .stderr
@load base/frameworks/storage
@load policy/frameworks/storage/backend/sqlite
redef Storage::expire_interval = 5 secs;
redef exit_only_after_terminate = T;
# Create a typename here that can be passed down into get().
type str: string;
global backend: opaque of Storage::BackendHandle;
global key: string = "key1234";
global value: string = "value7890";
event check_removed() {
local res2 = Storage::get(backend, key, F);
print "get result", res2;
Storage::close_backend(backend);
terminate();
}
event zeek_init() {
local opts : Storage::Backend::SQLite::Options;
opts$database_path = "storage-test.sqlite";
opts$table_name = "testing";
backend = Storage::open_backend(Storage::SQLITE, opts, str, str);
local res = Storage::put(backend, [$key=key, $value=value, $overwrite=T, $expire_time=2 secs, $async_mode=F]);
print "put result", res;
local res2 = Storage::get(backend, key, F);
print "get result", res2;
print "get result same as inserted", value == (res2 as string);
schedule 5 secs { check_removed() };
}