feat: add context to cache interface

Author: fr3fou

cache/cache.go

@@ -1,6 +1,7 @@
 package cache
 
 import (
+	"context"
 	"database/sql/driver"
 	"time"
 )
@@ -17,7 +18,7 @@ type Cacher interface {
 	// Get must return a pointer to the item, a boolean representing whether
 	// item is present or not, and an error (must be nil when key is not
 	// present).
-	Get(key string) (*Item, bool, error)
+	Get(ctx context.Context, key string) (*Item, bool, error)
 	// Set sets the item into cache with the given TTL.
-	Set(key string, item *Item, ttl time.Duration) error
+	Set(ctx context.Context, key string, item *Item, ttl time.Duration) error
 }

interceptor.go

@@ -84,15 +84,17 @@ func (i *Interceptor) Disable() {
 }
 
 // StmtQueryContext intecepts database/sql's stmt.QueryContext calls from a prepared statement.
-func (i *Interceptor) StmtQueryContext(ctx context.Context, conn driver.StmtQueryContext, query string, args []driver.NamedValue) (driver.Rows, error) {
+func (i *Interceptor) StmtQueryContext(ctx context.Context, conn driver.StmtQueryContext, query string, args []driver.NamedValue) (context.Context, driver.Rows, error) {
 
 	if i.disabled {
-		return conn.QueryContext(ctx, args)
+		r , err := conn.QueryContext(ctx, args)
+		return ctx, r, err
 	}
 
 	attrs := getAttrs(query)
 	if attrs == nil {
-		return conn.QueryContext(ctx, args)
+		r , err := conn.QueryContext(ctx, args)
+		return ctx, r, err
 	}
 
 	hash, err := i.hashFunc(query, args)
@@ -101,20 +103,21 @@ func (i *Interceptor) StmtQueryContext(ctx context.Context, conn driver.StmtQuer
 		if i.onErr != nil {
 			i.onErr(fmt.Errorf("HashFunc failed: %w", err))
 		}
-		return conn.QueryContext(ctx, args)
+		r , err := conn.QueryContext(ctx, args)
+		return ctx, r, err
 	}
 
-	if cached := i.checkCache(hash); cached != nil {
-		return cached, nil
+	if cached := i.checkCache(ctx, hash); cached != nil {
+		return ctx, cached, nil
 	}
 
 	rows, err := conn.QueryContext(ctx, args)
 	if err != nil {
-		return rows, err
+		return ctx, rows, err
 	}
 
 	cacheSetter := func(item *cache.Item) {
-		err := i.c.Set(hash, item, time.Duration(attrs.ttl)*time.Second)
+		err := i.c.Set(ctx, hash, item, time.Duration(attrs.ttl)*time.Second)
 		if err != nil {
 			atomic.AddUint64(&i.stats.Errors, 1)
 			if i.onErr != nil {
@@ -123,19 +126,21 @@ func (i *Interceptor) StmtQueryContext(ctx context.Context, conn driver.StmtQuer
 		}
 	}
 
-	return newRowsRecorder(cacheSetter, rows, attrs.maxRows), err
+	return ctx, newRowsRecorder(cacheSetter, rows, attrs.maxRows), err
 }
 
 // ConnQueryContext intecepts database/sql's DB.QueryContext Conn.QueryContext calls.
-func (i *Interceptor) ConnQueryContext(ctx context.Context, conn driver.QueryerContext, query string, args []driver.NamedValue) (driver.Rows, error) {
+func (i *Interceptor) ConnQueryContext(ctx context.Context, conn driver.QueryerContext, query string, args []driver.NamedValue) (context.Context, driver.Rows, error) {
 
 	if i.disabled {
-		return conn.QueryContext(ctx, query, args)
+		r, err := conn.QueryContext(ctx, query, args)
+		return ctx, r, err
 	}
 
 	attrs := getAttrs(query)
 	if attrs == nil {
-		return conn.QueryContext(ctx, query, args)
+		r, err := conn.QueryContext(ctx, query, args)
+		return ctx, r, err
 	}
 
 	hash, err := i.hashFunc(query, args)
@@ -144,20 +149,21 @@ func (i *Interceptor) ConnQueryContext(ctx context.Context, conn driver.QueryerC
 		if i.onErr != nil {
 			i.onErr(fmt.Errorf("HashFunc failed: %w", err))
 		}
-		return conn.QueryContext(ctx, query, args)
+		r, err := conn.QueryContext(ctx, query, args)
+		return ctx, r, err
 	}
 
-	if cached := i.checkCache(hash); cached != nil {
-		return cached, nil
+	if cached := i.checkCache(ctx, hash); cached != nil {
+		return ctx, cached, nil
 	}
 
 	rows, err := conn.QueryContext(ctx, query, args)
 	if err != nil {
-		return rows, err
+		return ctx, rows, err
 	}
 
 	cacheSetter := func(item *cache.Item) {
-		err := i.c.Set(hash, item, time.Duration(attrs.ttl)*time.Second)
+		err := i.c.Set(ctx, hash, item, time.Duration(attrs.ttl)*time.Second)
 		if err != nil {
 			atomic.AddUint64(&i.stats.Errors, 1)
 			if i.onErr != nil {
@@ -166,11 +172,11 @@ func (i *Interceptor) ConnQueryContext(ctx context.Context, conn driver.QueryerC
 		}
 	}
 
-	return newRowsRecorder(cacheSetter, rows, attrs.maxRows), err
+	return ctx, newRowsRecorder(cacheSetter, rows, attrs.maxRows), err
 }
 
-func (i *Interceptor) checkCache(hash string) driver.Rows {
-	item, ok, err := i.c.Get(hash)
+func (i *Interceptor) checkCache(ctx context.Context, hash string) driver.Rows {
+	item, ok, err := i.c.Get(ctx, hash)
 	if err != nil {
 		atomic.AddUint64(&i.stats.Errors, 1)
 		if i.onErr != nil {