From 9b12dfffc279221d3da1e83a458124c5ca3d70b0 Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Mon, 17 Mar 2025 14:13:53 +0100
Subject: [PATCH] Add option to use prepared statements for SQLite PutMany
---
base/database/storage/sqlite/prepared.go | 125 ++++++++++++++++++
base/database/storage/sqlite/prepared_test.go | 86 ++++++++++++
base/database/storage/sqlite/sqlite.go | 5 +
3 files changed, 216 insertions(+)
create mode 100644 base/database/storage/sqlite/prepared.go
create mode 100644 base/database/storage/sqlite/prepared_test.go
diff --git a/base/database/storage/sqlite/prepared.go b/base/database/storage/sqlite/prepared.go
new file mode 100644
index 00000000..11136dd8
--- /dev/null
+++ b/base/database/storage/sqlite/prepared.go
@@ -0,0 +1,125 @@
+package sqlite
+
+import (
+ "context"
+ "strconv"
+
+ "github.com/stephenafamo/bob"
+ "github.com/stephenafamo/bob/dialect/sqlite/im"
+ "github.com/stephenafamo/bob/expr"
+
+ "github.com/safing/portmaster/base/database/record"
+ "github.com/safing/portmaster/base/database/storage/sqlite/models"
+ "github.com/safing/structures/dsd"
+)
+
+var UsePreparedStatements bool = true
+
+// PutMany stores many records in the database.
+func (db *SQLite) putManyWithPreparedStmts(shadowDelete bool) (chan<- record.Record, <-chan error) {
+ batch := make(chan record.Record, 100)
+ errs := make(chan error, 1)
+
+ // Simulate upsert with custom selection on conflict.
+ rawQuery, _, err := models.Records.Insert(
+ im.Into("records", "key", "format", "value", "created", "modified", "expires", "deleted", "secret", "crownjewel"),
+ im.Values(expr.Arg("key"), expr.Arg("format"), expr.Arg("value"), expr.Arg("created"), expr.Arg("modified"), expr.Arg("expires"), expr.Arg("deleted"), expr.Arg("secret"), expr.Arg("crownjewel")),
+ im.OnConflict("key").DoUpdate(
+ im.SetExcluded("format", "value", "created", "modified", "expires", "deleted", "secret", "crownjewel"),
+ ),
+ ).Build(db.ctx)
+ if err != nil {
+ errs <- err
+ return batch, errs
+ }
+
+ // Start transaction.
+ tx, err := db.bob.BeginTx(db.ctx, nil)
+ if err != nil {
+ errs <- err
+ return batch, errs
+ }
+
+ // Create prepared statement WITHIN TRANSACTION.
+ preparedStmt, err := tx.PrepareContext(db.ctx, rawQuery)
+ if err != nil {
+ errs <- err
+ return batch, errs
+ }
+
+ // start handler
+ go func() {
+ // Read all put records.
+ writeBatch:
+ for {
+ select {
+ case r := <-batch:
+ if r != nil {
+ // Write record.
+ err := writeWithPreparedStatement(db.ctx, &preparedStmt, r)
+ if err != nil {
+ errs <- err
+ break writeBatch
+ }
+ } else {
+ // Finalize transcation.
+ errs <- tx.Commit()
+ return
+ }
+
+ case <-db.ctx.Done():
+ break writeBatch
+ }
+ }
+
+ // Rollback transaction.
+ errs <- tx.Rollback()
+ }()
+
+ return batch, errs
+}
+
+func writeWithPreparedStatement(ctx context.Context, pStmt *bob.StdPrepared, r record.Record) error {
+ r.Lock()
+ defer r.Unlock()
+
+ // Serialize to JSON.
+ data, err := r.MarshalDataOnly(r, dsd.JSON)
+ if err != nil {
+ return err
+ }
+
+ // Get Meta.
+ m := r.Meta()
+
+ // Insert.
+ if len(data) > 0 {
+ format := strconv.Itoa(dsd.JSON)
+ _, err = pStmt.ExecContext(
+ ctx,
+ r.DatabaseKey(),
+ format,
+ data,
+ m.Created,
+ m.Modified,
+ m.Expires,
+ m.Deleted,
+ m.IsSecret(),
+ m.IsCrownJewel(),
+ )
+ } else {
+ _, err = pStmt.ExecContext(
+ ctx,
+ r.DatabaseKey(),
+ nil,
+ nil,
+ m.Created,
+ m.Modified,
+ m.Expires,
+ m.Deleted,
+ m.IsSecret(),
+ m.IsCrownJewel(),
+ )
+ }
+ return err
+}
diff --git a/base/database/storage/sqlite/prepared_test.go b/base/database/storage/sqlite/prepared_test.go
new file mode 100644
index 00000000..8a90bcb9
--- /dev/null
+++ b/base/database/storage/sqlite/prepared_test.go
@@ -0,0 +1,86 @@
+package sqlite
+
+import (
+ "strconv"
+ "testing"
+)
+
+func BenchmarkPutMany(b *testing.B) {
+ // Configure prepared statement usage.
+ origSetting := UsePreparedStatements
+ UsePreparedStatements = false
+ defer func() {
+ UsePreparedStatements = origSetting
+ }()
+
+ // Run benchmark.
+ benchPutMany(b)
+}
+
+func BenchmarkPutManyPreparedStmt(b *testing.B) {
+ // Configure prepared statement usage.
+ origSetting := UsePreparedStatements
+ UsePreparedStatements = true
+ defer func() {
+ UsePreparedStatements = origSetting
+ }()
+
+ // Run benchmark.
+ benchPutMany(b)
+}
+
+func benchPutMany(b *testing.B) { //nolint:thelper
+ // Start database.
+ testDir := b.TempDir()
+ db, err := openSQLite("test", testDir, false)
+ if err != nil {
+ b.Fatal(err)
+ }
+ defer func() {
+ // shutdown
+ err = db.Shutdown()
+ if err != nil {
+ b.Fatal(err)
+ }
+ }()
+
+ // Start benchmarking.
+ b.ResetTimer()
+
+ // Benchmark PutMany.
+ records, errs := db.PutMany(false)
+ for i := range b.N {
+ // Create test record.
+ newTestRecord := &TestRecord{
+ S: "banana",
+ I: 42,
+ I8: 42,
+ I16: 42,
+ I32: 42,
+ I64: 42,
+ UI: 42,
+ UI8: 42,
+ UI16: 42,
+ UI32: 42,
+ UI64: 42,
+ F32: 42.42,
+ F64: 42.42,
+ B: true,
+ }
+ newTestRecord.UpdateMeta()
+ newTestRecord.SetKey("test:" + strconv.Itoa(i))
+
+ select {
+ case records <- newTestRecord:
+ case err := <-errs:
+ b.Fatal(err)
+ }
+ }
+
+ // Finalize.
+ close(records)
+ err = <-errs
+ if err != nil {
+ b.Fatal(err)
+ }
+}
diff --git a/base/database/storage/sqlite/sqlite.go b/base/database/storage/sqlite/sqlite.go
index 76e1554a..c6288322 100644
--- a/base/database/storage/sqlite/sqlite.go
+++ b/base/database/storage/sqlite/sqlite.go
@@ -205,6 +205,11 @@ func (db *SQLite) PutMany(shadowDelete bool) (chan<- record.Record, <-chan error
db.wg.Add(1)
defer db.wg.Done()
+ // Check if we should use prepared statement optimized inserting.
+ if UsePreparedStatements {
+ return db.putManyWithPreparedStmts(shadowDelete)
+ }
+
batch := make(chan record.Record, 100)
errs := make(chan error, 1)