summaryrefslogtreecommitdiff
path: root/modules
diff options
context:
space:
mode:
Diffstat (limited to 'modules')
-rw-r--r--modules/cache/boltstore.go74
-rw-r--r--modules/cache/boltstore_test.go (renamed from modules/cache/sqlitestore_test.go)6
-rw-r--r--modules/cache/cache.go7
-rw-r--r--modules/cache/memstore.go13
-rw-r--r--modules/cache/sqlitestore.go60
-rw-r--r--modules/cookies/cookies.go4
6 files changed, 79 insertions, 85 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)
- }
-}
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)
}
}