summaryrefslogtreecommitdiff
path: root/modules
diff options
context:
space:
mode:
Diffstat (limited to 'modules')
-rw-r--r--modules/cache/cache.go52
-rw-r--r--modules/cache/cache_test.go38
-rw-r--r--modules/cache/memstore.go13
-rw-r--r--modules/cache/sqlitestore.go60
-rw-r--r--modules/cache/sqlitestore_test.go32
-rw-r--r--modules/ratelimit/ratelimit.go16
-rw-r--r--modules/ratelimit/ratelimit_test.go8
7 files changed, 154 insertions, 65 deletions
diff --git a/modules/cache/cache.go b/modules/cache/cache.go
index 1a321be..10762f9 100644
--- a/modules/cache/cache.go
+++ b/modules/cache/cache.go
@@ -9,8 +9,9 @@ import (
"bytes"
"net/http"
"net/http/httputil"
+ "path/filepath"
+ "strings"
- "github.com/cornelk/hashmap"
"github.com/philippta/flyscrape"
)
@@ -21,7 +22,7 @@ func init() {
type Module struct {
Cache string `json:"cache"`
- cache *hashmap.Map[string, []byte]
+ store Store
}
func (Module) ModuleInfo() flyscrape.ModuleInfo {
@@ -31,24 +32,28 @@ func (Module) ModuleInfo() flyscrape.ModuleInfo {
}
}
-func (m *Module) Provision(flyscrape.Context) {
- if m.disabled() {
- return
- }
- if m.cache == nil {
- m.cache = hashmap.New[string, []byte]()
+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)
+
+ case strings.HasPrefix(m.Cache, "file:"):
+ m.store = NewSQLiteStore(strings.TrimPrefix(m.Cache, "file:"))
}
}
func (m *Module) AdaptTransport(t http.RoundTripper) http.RoundTripper {
- if m.disabled() {
+ if m.store == nil {
return t
}
-
return flyscrape.RoundTripFunc(func(r *http.Request) (*http.Response, error) {
- key := cacheKey(r)
+ key := r.Method + " " + r.URL.String()
- if b, ok := m.cache.Get(key); ok {
+ if b, ok := m.store.Get(key); ok {
if resp, err := http.ReadResponse(bufio.NewReader(bytes.NewReader(b)), r); err == nil {
return resp, nil
}
@@ -64,15 +69,28 @@ func (m *Module) AdaptTransport(t http.RoundTripper) http.RoundTripper {
return resp, err
}
- m.cache.Set(key, encoded)
+ m.store.Set(key, encoded)
return resp, nil
})
}
-func (m *Module) disabled() bool {
- return m.Cache == ""
+func (m *Module) Finalize() {
+ if v, ok := m.store.(interface{ Close() }); ok {
+ v.Close()
+ }
+}
+
+func replaceExt(filePath string, newExt string) string {
+ ext := filepath.Ext(filePath)
+ if ext != "" {
+ fileNameWithoutExt := filePath[:len(filePath)-len(ext)]
+ newFilePath := fileNameWithoutExt + newExt
+ return newFilePath
+ }
+ return filePath + newExt
}
-func cacheKey(r *http.Request) string {
- return r.Method + " " + r.URL.String()
+type Store interface {
+ Get(key string) ([]byte, bool)
+ Set(key string, value []byte)
}
diff --git a/modules/cache/cache_test.go b/modules/cache/cache_test.go
deleted file mode 100644
index 4565e00..0000000
--- a/modules/cache/cache_test.go
+++ /dev/null
@@ -1,38 +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 (
- "net/http"
- "testing"
-
- "github.com/philippta/flyscrape"
- "github.com/philippta/flyscrape/modules/cache"
- "github.com/philippta/flyscrape/modules/hook"
- "github.com/philippta/flyscrape/modules/starturl"
- "github.com/stretchr/testify/require"
-)
-
-func TestCache(t *testing.T) {
- cachemod := &cache.Module{Cache: "memory"}
- calls := 0
-
- for i := 0; i < 2; i++ {
- scraper := flyscrape.NewScraper()
- scraper.LoadModule(&starturl.Module{URL: "http://www.example.com"})
- scraper.LoadModule(hook.Module{
- AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
- return flyscrape.RoundTripFunc(func(r *http.Request) (*http.Response, error) {
- calls++
- return flyscrape.MockResponse(200, "foo")
- })
- },
- })
- scraper.LoadModule(cachemod)
- scraper.Run()
- }
-
- require.Equal(t, 1, calls)
-}
diff --git a/modules/cache/memstore.go b/modules/cache/memstore.go
new file mode 100644
index 0000000..0f4d9e2
--- /dev/null
+++ b/modules/cache/memstore.go
@@ -0,0 +1,13 @@
+// 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
new file mode 100644
index 0000000..50c8007
--- /dev/null
+++ b/modules/cache/sqlitestore.go
@@ -0,0 +1,60 @@
+// 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, value)
+ }
+}
+
+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
new file mode 100644
index 0000000..769a69f
--- /dev/null
+++ b/modules/cache/sqlitestore_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 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/ratelimit/ratelimit.go b/modules/ratelimit/ratelimit.go
index 9588db3..b23cd7a 100644
--- a/modules/ratelimit/ratelimit.go
+++ b/modules/ratelimit/ratelimit.go
@@ -5,6 +5,7 @@
package ratelimit
import (
+ "net/http"
"time"
"github.com/philippta/flyscrape"
@@ -45,11 +46,14 @@ func (m *Module) Provision(v flyscrape.Context) {
}()
}
-func (m *Module) BuildRequest(_ *flyscrape.Request) {
+func (m *Module) AdaptTransport(t http.RoundTripper) http.RoundTripper {
if m.disabled() {
- return
+ return t
}
- <-m.semaphore
+ return flyscrape.RoundTripFunc(func(r *http.Request) (*http.Response, error) {
+ <-m.semaphore
+ return t.RoundTrip(r)
+ })
}
func (m *Module) Finalize() {
@@ -64,7 +68,7 @@ func (m *Module) disabled() bool {
}
var (
- _ flyscrape.RequestBuilder = (*Module)(nil)
- _ flyscrape.Provisioner = (*Module)(nil)
- _ flyscrape.Finalizer = (*Module)(nil)
+ _ flyscrape.TransportAdapter = (*Module)(nil)
+ _ flyscrape.Provisioner = (*Module)(nil)
+ _ flyscrape.Finalizer = (*Module)(nil)
)
diff --git a/modules/ratelimit/ratelimit_test.go b/modules/ratelimit/ratelimit_test.go
index ffd061c..1fe22b1 100644
--- a/modules/ratelimit/ratelimit_test.go
+++ b/modules/ratelimit/ratelimit_test.go
@@ -23,17 +23,17 @@ func TestRatelimit(t *testing.T) {
scraper := flyscrape.NewScraper()
scraper.LoadModule(&starturl.Module{URL: "http://www.example.com"})
scraper.LoadModule(&followlinks.Module{})
- scraper.LoadModule(&ratelimit.Module{
- Rate: 100,
- })
scraper.LoadModule(hook.Module{
AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
return flyscrape.MockTransport(200, `<a href="foo">foo</a>`)
},
- BuildRequestFn: func(r *flyscrape.Request) {
+ ReceiveResponseFn: func(r *flyscrape.Response) {
times = append(times, time.Now())
},
})
+ scraper.LoadModule(&ratelimit.Module{
+ Rate: 100,
+ })
start := time.Now()
scraper.Run()
generated by cgit v1.2.3 (git 2.25.1) at 2025-12-06 15:30:51 +0000