Implement file based caching

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()