Change cache backend from sqlite to bbolt

5 files changed, 79 insertions, 81 deletions

diff --git a/modules/cache/boltstore.go b/modules/cache/boltstore.go
new file mode 100644
index 0000000..7176445
--- /dev/null
+++ b/modules/cache/boltstore.go
@@ -0,0 +1,74 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package cache
+
+import (
+	"errors"
+	"log"
+	"os"
+
+	"go.etcd.io/bbolt"
+)
+
+var cache = []byte("cache")
+
+func NewBoltStore(file string) *BoltStore {
+	db, err := bbolt.Open(file, 0644, nil)
+	if err != nil {
+		log.Printf("cache: failed to create database file %q: %v\n", file, err)
+		os.Exit(1)
+	}
+
+	c := &BoltStore{db: db}
+
+	return c
+}
+
+type BoltStore struct {
+	db *bbolt.DB
+}
+
+func (s *BoltStore) Get(key string) ([]byte, bool) {
+	var value []byte
+
+	err := s.db.View(func(tx *bbolt.Tx) error {
+		bucket := tx.Bucket(cache)
+		if bucket == nil {
+			return errors.New("bucket not found")
+		}
+
+		v := bucket.Get([]byte(key))
+		if v == nil {
+			return errors.New("key not found")
+		}
+
+		value = make([]byte, len(v))
+		copy(value, v)
+
+		return nil
+	})
+	if err != nil {
+		return nil, false
+	}
+	return value, true
+}
+
+func (s *BoltStore) Set(key string, value []byte) {
+	err := s.db.Update(func(tx *bbolt.Tx) error {
+		bucket, err := tx.CreateBucketIfNotExists(cache)
+		if err != nil {
+			return err
+		}
+
+		return bucket.Put([]byte(key), value)
+	})
+	if err != nil {
+		log.Printf("cache: failed to insert cache key %q: %v\n", key, err)
+	}
+}
+
+func (s *BoltStore) Close() {
+	s.db.Close()
+}
diff --git a/modules/cache/sqlitestore_test.go b/modules/cache/boltstore_test.go
index 769a69f..0af954f 100644
--- a/modules/cache/sqlitestore_test.go
+++ b/modules/cache/boltstore_test.go
@@ -12,12 +12,12 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestSQLiteStore(t *testing.T) {
-	dir, err := os.MkdirTemp("", "sqlitestore")
+func TestBoltStore(t *testing.T) {
+	dir, err := os.MkdirTemp("", "boltstore")
 	require.NoError(t, err)
 	defer os.RemoveAll(dir)
 
-	store := cache.NewSQLiteStore(dir + "/test.db")
+	store := cache.NewBoltStore(dir + "/test.db")
 
 	v, ok := store.Get("foo")
 	require.Nil(t, v)
diff --git a/modules/cache/cache.go b/modules/cache/cache.go
index 401aa49..7164506 100644
--- a/modules/cache/cache.go
+++ b/modules/cache/cache.go
@@ -34,15 +34,12 @@ func (Module) ModuleInfo() flyscrape.ModuleInfo {
 
 func (m *Module) Provision(ctx flyscrape.Context) {
 	switch {
-	case m.Cache == "memory":
-		m.store = NewMemStore()
-
 	case m.Cache == "file":
 		file := replaceExt(ctx.ScriptName(), ".cache")
-		m.store = NewSQLiteStore(file)
+		m.store = NewBoltStore(file)
 
 	case strings.HasPrefix(m.Cache, "file:"):
-		m.store = NewSQLiteStore(strings.TrimPrefix(m.Cache, "file:"))
+		m.store = NewBoltStore(strings.TrimPrefix(m.Cache, "file:"))
 	}
 }
 
diff --git a/modules/cache/memstore.go b/modules/cache/memstore.go
deleted file mode 100644
index 0f4d9e2..0000000
--- a/modules/cache/memstore.go
+++ /dev/null
@@ -1,13 +0,0 @@
-// This Source Code Form is subject to the terms of the Mozilla Public
-// License, v. 2.0. If a copy of the MPL was not distributed with this
-// file, You can obtain one at http://mozilla.org/MPL/2.0/.
-
-package cache
-
-import (
-	"github.com/cornelk/hashmap"
-)
-
-func NewMemStore() *hashmap.Map[string, []byte] {
-	return hashmap.New[string, []byte]()
-}
diff --git a/modules/cache/sqlitestore.go b/modules/cache/sqlitestore.go
deleted file mode 100644
index 778699b..0000000
--- a/modules/cache/sqlitestore.go
+++ /dev/null
@@ -1,60 +0,0 @@
-// This Source Code Form is subject to the terms of the Mozilla Public
-// License, v. 2.0. If a copy of the MPL was not distributed with this
-// file, You can obtain one at http://mozilla.org/MPL/2.0/.
-
-package cache
-
-import (
-	"database/sql"
-	"fmt"
-	"log"
-	"os"
-
-	_ "github.com/mattn/go-sqlite3"
-)
-
-func NewSQLiteStore(file string) *SQLiteStore {
-	db, err := sql.Open("sqlite3", fmt.Sprintf("file:%s?_timeout=5000&_journal=WAL", file))
-	if err != nil {
-		log.Printf("cache: failed to create database file %q: %v\n", file, err)
-		os.Exit(1)
-	}
-
-	c := &SQLiteStore{db: db}
-	c.migrate()
-
-	return c
-}
-
-type SQLiteStore struct {
-	db *sql.DB
-}
-
-func (s *SQLiteStore) Get(key string) ([]byte, bool) {
-	var value []byte
-	if err := s.db.QueryRow(`SELECT value FROM cache WHERE key = ? LIMIT 1`, key).Scan(&value); err != nil {
-		return nil, false
-	}
-	return value, true
-}
-
-func (s *SQLiteStore) Set(key string, value []byte) {
-	if _, err := s.db.Exec(`INSERT INTO cache (key, value) VALUES (?, ?)`, key, value); err != nil {
-		log.Printf("cache: failed to insert cache key %q: %v\n", key, err)
-	}
-}
-
-func (s *SQLiteStore) Close() {
-	s.db.Close()
-}
-
-func (s *SQLiteStore) migrate() {
-	if _, err := s.db.Exec(`CREATE TABLE IF NOT EXISTS cache (key TEXT, value BLOB)`); err != nil {
-		log.Printf("cache: failed to create cache table: %v\n", err)
-		os.Exit(1)
-	}
-	if _, err := s.db.Exec(`CREATE UNIQUE INDEX IF NOT EXISTS cache_key_idx ON cache(key)`); err != nil {
-		log.Printf("cache: failed to create cache index: %v\n", err)
-		os.Exit(1)
-	}
-}