proxy multi database support

Author: walterwanderley

cmd/sqlite-http-proxy/main.go

@@ -20,6 +20,7 @@ import (
 
 var (
 	port       uint
+	dbParams   string
 	allowHTTP2 bool
 	verbose    bool
 
@@ -35,6 +36,7 @@ var (
 
 func main() {
 	flag.UintVar(&port, "p", 8080, "Server port")
+	flag.StringVar(&dbParams, "db-params", "_journal=WAL&_sync=NORMAL&_timeout=5000&_txlock=immediate", "Database connection params")
 	flag.BoolVar(&verbose, "v", false, "Enable verbose mode")
 	flag.BoolVar(&allowHTTP2, "h2", false, "Allow HTTP2")
 	flag.UintVar(&ttl, "ttl", 0, "Time to Live in seconds (0 is infinite time)")
@@ -45,36 +47,57 @@ func main() {
 	flag.BoolVar(&readOnly, "ro", false, "Read Only mode. Do not store new HTTP responses")
 	flag.Parse()
 
-	if len(flag.Args()) != 1 {
-		log.Fatalf("Usage: %s <flags> [DSN]\n\nExample:\n\t%s file:example.db\n", os.Args[0], os.Args[0])
+	if len(flag.Args()) == 0 {
+		log.Fatalf("Usage: %s <flags> [DatabasePath1] [DatabasePathN\n\nExample:\n\t%s example.db example2.db example3.db\n", os.Args[0], os.Args[0])
 	}
-	dsn := flag.Args()[0]
 
-	sqlDB, err := sql.Open("sqlite3", dsn)
-	if err != nil {
-		log.Fatalf("open db error: %v", err)
-	}
-	defer sqlDB.Close()
+	dbs := make([]*sql.DB, 0)
+	var (
+		repository db.Repository
+		tableList  []string
+		err        error
+	)
+	for _, file := range flag.Args() {
+		var dsn string
+		if file == ":memory:" {
+			dsn = file + "?cache=shared"
+		} else {
+			dsn = fmt.Sprintf("file:%s?%s", file, dbParams)
+		}
 
-	var tableList []string
-	if responseTables == "" {
-		tableList, err = db.ResponseTables(sqlDB)
+		sqlDB, err := sql.Open("sqlite3", dsn)
 		if err != nil {
-			log.Fatalf("discovery response tables: %v", err)
+			log.Fatalf("open db error: %v", err)
 		}
-	} else {
-		tableList = strings.Split(responseTables, ",")
-		if forceCreateTables {
-			err := db.CreateResponseTables(sqlDB, tableList...)
+		defer sqlDB.Close()
+
+		dbs = append(dbs, sqlDB)
+
+		if responseTables == "" {
+			tableList, err = db.ResponseTables(sqlDB)
 			if err != nil {
-				log.Fatalf("force create tables: %v", err)
+				log.Fatalf("discovery response tables: %v", err)
+			}
+		} else {
+			tableList = strings.Split(responseTables, ",")
+			if forceCreateTables {
+				err := db.CreateResponseTables(sqlDB, tableList...)
+				if err != nil {
+					log.Fatalf("force create tables: %v", err)
+				}
 			}
 		}
 	}
-
-	repository, err := db.NewRepository(sqlDB, tableList...)
-	if err != nil {
-		log.Fatalf("new repository: %v", err)
+	if len(dbs) == 1 {
+		repository, err = db.NewRepository(dbs[0], tableList...)
+		if err != nil {
+			log.Fatalf("new repository: %v", err)
+		}
+	} else {
+		repository, err = db.NewMultiDatabaseRepository(dbs)
+		if err != nil {
+			log.Fatalf("new multi database repository: %v", err)
+		}
 	}
 	defer repository.Close()
 

cmd/sqlite-http-proxy/request.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"errors"
+	"fmt"
 	"log/slog"
 	"net/http"
 	"time"
@@ -33,13 +34,13 @@ func (h *requestHandler) Handle(r *http.Request, ctx *goproxy.ProxyCtx) (*http.R
 			slog.Error("database query", "error", err.Error())
 		}
 		// tell the responseHandler to save the new response data
-		ctx.UserData = ""
+		ctx.UserData = ":-1"
 		return r, nil
 	}
 
 	if !readOnly && uint(time.Since(resp.Timestamp).Seconds()) > ttl {
 		// data is too old, tell the responseHandler to save the new data
-		ctx.UserData = resp.TableName
+		ctx.UserData = fmt.Sprintf("%s:%d", resp.TableName, resp.DatabaseID)
 		return r, nil
 	}
 	if verbose {

cmd/sqlite-http-proxy/response.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"log/slog"
 	"net/http"
+	"strconv"
+	"strings"
 
 	"github.com/elazarl/goproxy"
 	"github.com/walterwanderley/sqlite-http-cache/db"
@@ -26,7 +28,12 @@ func (h *responseHandler) Handle(resp *http.Response, ctx *goproxy.ProxyCtx) *ht
 		if err != nil {
 			slog.Error("adapter response body", "error", err)
 		} else {
-			responseDB.TableName = ctx.UserData.(string)
+			userData := ctx.UserData.(string)
+			tableName, databaseID, ok := strings.Cut(userData, ":")
+			if ok {
+				responseDB.DatabaseID, _ = strconv.Atoi(databaseID)
+			}
+			responseDB.TableName = tableName
 			err := h.writer.Write(context.Background(), ctx.Req.URL.String(), responseDB)
 			if err != nil {
 				slog.Error("recording response", "error", err, "url", ctx.Req.URL.String(), "status", resp.StatusCode)

db/concurrent_repository.go

@@ -16,9 +16,13 @@ type concurrentRepository struct {
 	// roundRobin strategy to choose one writer
 	currentWriter int
 	muWriter      sync.Mutex
+
+	// workers for MultiDatabaseRepository
+	source chan work
+	quit   chan struct{}
 }
 
-func newConcurrentRepository(db *sql.DB, tableNames ...string) (*concurrentRepository, error) {
+func newConcurrentRepository(db *sql.DB, databaseID int, tableNames ...string) (*concurrentRepository, error) {
 	globalQuit := make(chan struct{})
 
 	size := len(tableNames)
@@ -29,7 +33,7 @@ func newConcurrentRepository(db *sql.DB, tableNames ...string) (*concurrentRepos
 		if err != nil {
 			return nil, fmt.Errorf("prepare read query for %q: %w", tableName, err)
 		}
-		q := newQuerier(readStmt, tableName)
+		q := newQuerier(readStmt, databaseID, tableName)
 		queriers[i] = q
 		q.start()
 
@@ -111,27 +115,51 @@ func (r *concurrentRepository) Write(ctx context.Context, url string, resp *Resp
 }
 
 func (r *concurrentRepository) Close() error {
+	if r.quit != nil {
+		close(r.quit)
+	}
 	close(r.globalQuit)
 	return nil
 }
 
+func (r *concurrentRepository) start() {
+	r.source = make(chan work, 10)
+	go func() {
+		for {
+			select {
+			case unitWork := <-r.source:
+				resp, err := r.FindByURL(unitWork.ctx, unitWork.url)
+				if err != nil {
+					unitWork.resp <- nil
+					continue
+				}
+				unitWork.resp <- resp
+			case <-r.quit:
+				return
+			}
+		}
+	}()
+}
+
 type work struct {
 	ctx  context.Context
 	url  string
 	resp chan *Response
 }
 
 type querier struct {
-	source    chan work
-	quit      chan struct{}
-	stmt      *sql.Stmt
-	tableName string
+	source     chan work
+	quit       chan struct{}
+	stmt       *sql.Stmt
+	tableName  string
+	databaseID int
 }
 
-func newQuerier(stmt *sql.Stmt, tableName string) *querier {
+func newQuerier(stmt *sql.Stmt, databaseID int, tableName string) *querier {
 	return &querier{
-		stmt:      stmt,
-		tableName: tableName,
+		stmt:       stmt,
+		tableName:  tableName,
+		databaseID: databaseID,
 	}
 }
 
@@ -147,6 +175,7 @@ func (q *querier) start() {
 					continue
 				}
 				response.TableName = q.tableName
+				response.DatabaseID = q.databaseID
 				unitWork.resp <- response
 
 			case <-q.quit:

db/multi_database.go

@@ -0,0 +1,95 @@
+package db
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"sync"
+)
+
+type MultiDatabaseRepository struct {
+	dbs                    []*sql.DB
+	concurrentRepositories []*concurrentRepository
+	// roundRobin strategy to choose one writer
+	currentWriter int
+	muWriter      sync.Mutex
+}
+
+func NewMultiDatabaseRepository(dbs []*sql.DB) (*MultiDatabaseRepository, error) {
+	concurrentRepositories := make([]*concurrentRepository, len(dbs))
+	for i, db := range dbs {
+		tables, err := ResponseTables(db)
+		if err != nil {
+			return nil, err
+		}
+		cr, err := newConcurrentRepository(db, i, tables...)
+		if err != nil {
+			return nil, err
+		}
+		cr.start()
+		concurrentRepositories[i] = cr
+	}
+	return &MultiDatabaseRepository{
+		dbs:                    dbs,
+		concurrentRepositories: concurrentRepositories,
+	}, nil
+}
+
+func (r *MultiDatabaseRepository) FindByURL(ctx context.Context, url string) (*Response, error) {
+	size := len(r.concurrentRepositories)
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	respCh := make(chan *Response, size)
+	unitWork := work{
+		ctx:  ctx,
+		url:  url,
+		resp: respCh,
+	}
+
+	go func() {
+		for _, cr := range r.concurrentRepositories {
+			cr.source <- unitWork
+		}
+	}()
+	var count int
+	for {
+		select {
+		case resp := <-respCh:
+			count++
+			if resp != nil {
+				return resp, nil
+			}
+			if count == size {
+				return nil, sql.ErrNoRows
+			}
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	}
+}
+
+func (r *MultiDatabaseRepository) Write(ctx context.Context, url string, resp *Response) error {
+	var cr Repository
+	if resp.DatabaseID == -1 {
+		r.muWriter.Lock()
+		defer r.muWriter.Unlock()
+		r.currentWriter++
+		if r.currentWriter >= len(r.concurrentRepositories) {
+			r.currentWriter = 0
+		}
+		cr = r.concurrentRepositories[r.currentWriter]
+	} else {
+		cr = r.concurrentRepositories[resp.DatabaseID]
+	}
+	return cr.Write(ctx, url, resp)
+}
+
+func (r *MultiDatabaseRepository) Close() error {
+	var err error
+	for _, cr := range r.concurrentRepositories {
+		err = errors.Join(err, cr.Close())
+	}
+	return err
+}

db/reository.go

@@ -20,11 +20,12 @@ type Repository interface {
 }
 
 type Response struct {
-	Status    int
-	Body      io.ReadCloser
-	Headers   map[string][]string
-	Timestamp time.Time
-	TableName string
+	Status     int
+	Body       io.ReadCloser
+	Headers    map[string][]string
+	Timestamp  time.Time
+	DatabaseID int
+	TableName  string
 }
 
 func NewRepository(db *sql.DB, tableNames ...string) (Repository, error) {
@@ -43,7 +44,7 @@ func NewRepository(db *sql.DB, tableNames ...string) (Repository, error) {
 		return newSingleRepository(db, tableNames[0])
 	}
 
-	return newConcurrentRepository(db, tableNames...)
+	return newConcurrentRepository(db, 0, tableNames...)
 }
 
 func CreateResponseTables(db *sql.DB, tableNames ...string) error {

http/ro_transport.go

@@ -44,7 +44,7 @@ func (t *readOnlyTransport) RoundTrip(req *http.Request) (*http.Response, error)
 
 	url := req.URL.String()
 	resp, err := t.querier.FindByURL(req.Context(), url)
-	if err != nil || time.Since(resp.Timestamp) > t.ttl {
+	if err != nil || (t.ttl > 0 && time.Since(resp.Timestamp) > t.ttl) {
 		return t.base.RoundTrip(req)
 	}
 

http/rw_transport.go

@@ -43,7 +43,7 @@ func (t *readWriteTransport) RoundTrip(req *http.Request) (*http.Response, error
 
 	url := req.URL.String()
 	respDB, err := t.querier.FindByURL(req.Context(), url)
-	if err != nil || time.Since(respDB.Timestamp) > t.ttl {
+	if err != nil || (t.ttl > 0 && time.Since(respDB.Timestamp) > t.ttl) {
 		resp, err := t.base.RoundTrip(req)
 		if err == nil {
 			newRespDB, err := db.HttpToResponse(resp)