diff --git a/scripts/policy/frameworks/storage/backend/sqlite/main.zeek b/scripts/policy/frameworks/storage/backend/sqlite/main.zeek index 7557d73121..90a512e454 100644 --- a/scripts/policy/frameworks/storage/backend/sqlite/main.zeek +++ b/scripts/policy/frameworks/storage/backend/sqlite/main.zeek @@ -20,10 +20,12 @@ export { ## different between the two. table_name: string; - ## Key/value table for passing tuning parameters when opening the - ## database. These must be pairs that can be passed to the ``pragma`` - ## command in sqlite. - tuning_params: table[string] of string &default=table( + ## Key/value table for passing pragma commands when opening the database. + ## These must be pairs that can be passed to the ``pragma`` command in + ## sqlite. The ``integrity_check`` pragma is run automatically and does + ## not need to be included here. For pragmas without a second argument, + ## set the value to an empty string. + pragma_commands: table[string] of string &default=table( ["busy_timeout"] = "5000", ["journal_mode"] = "WAL", ["synchronous"] = "normal", diff --git a/src/storage/backend/sqlite/SQLite.cc b/src/storage/backend/sqlite/SQLite.cc index 11bb2fb138..a852138824 100644 --- a/src/storage/backend/sqlite/SQLite.cc +++ b/src/storage/backend/sqlite/SQLite.cc @@ -11,6 +11,41 @@ namespace zeek::storage::backend::sqlite { +OperationResult SQLite::RunPragma(std::string_view name, std::optional value) { + char* errorMsg = nullptr; + + std::string cmd = util::fmt("pragma %.*s", static_cast(name.size()), name.data()); + if ( value && ! value->empty() ) + cmd += util::fmt(" = %.*s", static_cast(value->size()), value->data()); + + while ( attempts < 5 ) { + int res = sqlite3_exec(db, cmd.c_str(), NULL, NULL, &errorMsg); + if ( res == SQLITE_OK ) { + break; + } + else if ( res == SQLITE_BUSY ) { + // If we got back that the database is busy, it likely means that another process is trying to + // do their pragmas at startup too. Sleep for a little bit and try again. + sqlite3_free(errorMsg); + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + ++attempts; + } + else { + std::string err = util::fmt("Error while executing pragma '%s': %s", cmd.c_str(), errorMsg); + sqlite3_free(errorMsg); + return {ReturnCode::INITIALIZATION_FAILED, std::move(err)}; + } + } + + if ( attempts == 5 ) { + std::string err = + util::fmt("Database was busy while executing %.*s pragma", static_cast(name.size()), name.data()); + return {ReturnCode::INITIALIZATION_FAILED, std::move(err)}; + } + + return {ReturnCode::SUCCESS}; +} + storage::BackendPtr SQLite::Instantiate() { return make_intrusive(); } /** @@ -47,74 +82,23 @@ OperationResult SQLite::DoOpen(OpenResultCallback* cb, RecordValPtr options) { char* errorMsg = nullptr; - int attempts = 0; - while ( attempts < 5 ) { - int res = sqlite3_exec(db, "pragma integrity_check", NULL, NULL, &errorMsg); - if ( res == SQLITE_OK ) { - break; - } - else if ( res == SQLITE_BUSY ) { - // If we got back that the database is busy, it likely means that another process is trying to - // do their pragmas at startup too. Sleep for a little bit and try again. - sqlite3_free(errorMsg); - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - ++attempts; - } - else { - std::string err = util::fmt("Error executing integrity check (%d): %s", res, errorMsg); - Error(err.c_str()); - sqlite3_free(errorMsg); - Close(nullptr); - return {ReturnCode::INITIALIZATION_FAILED, std::move(err)}; - } + OperationResult pragma_res = RunPragma("integrity_check"); + if ( pragma_res.code != ReturnCode::SUCCESS ) { + Error(pragma_res.err_str.c_str()); + return pragma_res; } - if ( attempts == 5 ) { - std::string err = util::fmt("Database was busy while attempting integrity checks"); - Error(err.c_str()); - Close(nullptr); - return {ReturnCode::INITIALIZATION_FAILED, std::move(err)}; - } - - auto tuning_params = backend_options->GetField("tuning_params")->ToMap(); - for ( const auto& [k, v] : tuning_params ) { - attempts = 0; + auto pragmas = backend_options->GetField("pragma_commands")->ToMap(); + for ( const auto& [k, v] : pragmas ) { auto ks = k->AsListVal()->Idx(0)->AsStringVal(); auto ks_sv = ks->ToStdStringView(); auto vs = v->AsStringVal(); auto vs_sv = vs->ToStdStringView(); - while ( attempts < 5 ) { - std::string cmd = util::fmt("pragma %.*s = %.*s", static_cast(ks_sv.size()), ks_sv.data(), - static_cast(vs_sv.size()), vs_sv.data()); - - int res = sqlite3_exec(db, cmd.c_str(), NULL, NULL, &errorMsg); - if ( res == SQLITE_OK ) { - break; - } - else if ( res == SQLITE_BUSY ) { - // If we got back that the database is busy, it likely means that another process is trying to - // do their pragmas at startup too. Sleep for a little bit and try again. - sqlite3_free(errorMsg); - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - ++attempts; - } - else { - std::string err = util::fmt("Error executing %.*s pragma statement for (%d): %s", - static_cast(ks_sv.size()), ks_sv.data(), res, errorMsg); - Error(err.c_str()); - sqlite3_free(errorMsg); - Close(nullptr); - return {ReturnCode::INITIALIZATION_FAILED, std::move(err)}; - } - } - - if ( attempts == 5 ) { - std::string err = util::fmt("Database was busy while executing %.*s pragma", static_cast(ks_sv.size()), - ks_sv.data()); - Error(err.c_str()); - Close(nullptr); - return {ReturnCode::INITIALIZATION_FAILED, std::move(err)}; + pragma_res = RunPragma(ks_sv, vs_sv); + if ( pragma_res.code != ReturnCode::SUCCESS ) { + Error(pragma_res.err_str.c_str()); + return pragma_res; } } @@ -175,7 +159,8 @@ OperationResult SQLite::DoClose(ResultCallback* cb) { expire_stmt.reset(); char* errmsg; - if ( int res = sqlite3_exec(db, "pragma optimize", NULL, NULL, &errmsg); res != SQLITE_OK ) { + if ( int res = sqlite3_exec(db, "pragma optimize", NULL, NULL, &errmsg); + res != SQLITE_OK && res != SQLITE_BUSY ) { // We're shutting down so capture the error message here for informational // reasons, but don't do anything else with it. op_res = {ReturnCode::DISCONNECTION_FAILED, util::fmt("Sqlite failed to optimize at shutdown: %s", errmsg)}; diff --git a/src/storage/backend/sqlite/SQLite.h b/src/storage/backend/sqlite/SQLite.h index 6f011611b4..598db99a3f 100644 --- a/src/storage/backend/sqlite/SQLite.h +++ b/src/storage/backend/sqlite/SQLite.h @@ -45,6 +45,11 @@ private: */ OperationResult Step(sqlite3_stmt* stmt, bool parse_value = false); + /** + * Helper utility for running pragmas on the database. + */ + OperationResult RunPragma(std::string_view name, std::optional value = std::nullopt); + sqlite3* db = nullptr; using stmt_deleter = std::function; diff --git a/testing/btest/Baseline/scripts.base.frameworks.storage.sqlite-basic/out b/testing/btest/Baseline/scripts.base.frameworks.storage.sqlite-basic/out index 0c0634a0c0..d0934b75b8 100644 --- a/testing/btest/Baseline/scripts.base.frameworks.storage.sqlite-basic/out +++ b/testing/btest/Baseline/scripts.base.frameworks.storage.sqlite-basic/out @@ -1,5 +1,5 @@ ### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63. -Storage::backend_opened, Storage::STORAGE_BACKEND_SQLITE, [serializer=Storage::STORAGE_SERIALIZER_JSON, sqlite=[database_path=test.sqlite, table_name=testing, tuning_params={ +Storage::backend_opened, Storage::STORAGE_BACKEND_SQLITE, [serializer=Storage::STORAGE_SERIALIZER_JSON, sqlite=[database_path=test.sqlite, table_name=testing, pragma_commands={ [synchronous] = normal, [journal_mode] = WAL, [temp_store] = memory,