From 175cb0700b86134d3b04eab23d100dce2c3b6d56 Mon Sep 17 00:00:00 2001 From: Philipp Tanlak Date: Sun, 24 Nov 2024 23:16:31 +0100 Subject: Change cache backend from sqlite to bbolt --- modules/cache/boltstore.go | 74 +++++++++++++++++++++++++++++++++++++++ modules/cache/boltstore_test.go | 32 +++++++++++++++++ modules/cache/cache.go | 7 ++-- modules/cache/memstore.go | 13 ------- modules/cache/sqlitestore.go | 60 ------------------------------- modules/cache/sqlitestore_test.go | 32 ----------------- modules/cookies/cookies.go | 4 --- 7 files changed, 108 insertions(+), 114 deletions(-) create mode 100644 modules/cache/boltstore.go create mode 100644 modules/cache/boltstore_test.go delete mode 100644 modules/cache/memstore.go delete mode 100644 modules/cache/sqlitestore.go delete mode 100644 modules/cache/sqlitestore_test.go (limited to 'modules') 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/boltstore_test.go b/modules/cache/boltstore_test.go new file mode 100644 index 0000000..0af954f --- /dev/null +++ b/modules/cache/boltstore_test.go @@ -0,0 +1,32 @@ +// 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_test + +import ( + "os" + "testing" + + "github.com/philippta/flyscrape/modules/cache" + "github.com/stretchr/testify/require" +) + +func TestBoltStore(t *testing.T) { + dir, err := os.MkdirTemp("", "boltstore") + require.NoError(t, err) + defer os.RemoveAll(dir) + + store := cache.NewBoltStore(dir + "/test.db") + + v, ok := store.Get("foo") + require.Nil(t, v) + require.False(t, ok) + + store.Set("foo", []byte("bar")) + + v, ok = store.Get("foo") + require.NotNil(t, v) + require.True(t, ok) + require.Equal(t, []byte("bar"), 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) - } -} diff --git a/modules/cache/sqlitestore_test.go b/modules/cache/sqlitestore_test.go deleted file mode 100644 index 769a69f..0000000 --- a/modules/cache/sqlitestore_test.go +++ /dev/null @@ -1,32 +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_test - -import ( - "os" - "testing" - - "github.com/philippta/flyscrape/modules/cache" - "github.com/stretchr/testify/require" -) - -func TestSQLiteStore(t *testing.T) { - dir, err := os.MkdirTemp("", "sqlitestore") - require.NoError(t, err) - defer os.RemoveAll(dir) - - store := cache.NewSQLiteStore(dir + "/test.db") - - v, ok := store.Get("foo") - require.Nil(t, v) - require.False(t, ok) - - store.Set("foo", []byte("bar")) - - v, ok = store.Get("foo") - require.NotNil(t, v) - require.True(t, ok) - require.Equal(t, []byte("bar"), v) -} diff --git a/modules/cookies/cookies.go b/modules/cookies/cookies.go index 2af2a27..2f57a3f 100644 --- a/modules/cookies/cookies.go +++ b/modules/cookies/cookies.go @@ -51,10 +51,6 @@ func (m Module) AdaptTransport(t http.RoundTripper) http.RoundTripper { return flyscrape.RoundTripFunc(func(r *http.Request) (*http.Response, error) { for _, store := range stores { for _, cookie := range store.Cookies(r.URL) { - // Unquote cookie value - if len(cookie.Value) >= 2 && cookie.Value[0] == '"' && cookie.Value[len(cookie.Value)-1] == '"' { - cookie.Value = cookie.Value[1 : len(cookie.Value)-2] - } r.AddCookie(cookie) } } -- cgit v1.2.3