cacheの設計の見直し (#15)

速度を少し犠牲にし、安全性を手に入れた

Author: pirosiki197

template/cache_test.go

@@ -3,8 +3,6 @@ package template
 import (
 	"context"
 	"database/sql/driver"
-	"errors"
-	"fmt"
 	"os"
 	"testing"
 
@@ -33,7 +31,7 @@ func TestCacheRows(t *testing.T) {
 	assert.NoError(t, err)
 	defer conn.Close()
 
-	var cacheRows driver.Rows
+	var cacheRows *CacheRows
 	err = conn.Raw(func(driverConn any) error {
 		conn := driverConn.(driver.Conn)
 		stmt, err := conn.Prepare("SELECT id, name FROM users")
@@ -46,31 +44,10 @@ func TestCacheRows(t *testing.T) {
 		if err != nil {
 			return err
 		}
-		cacheRows = NewCachedRows(rows)
+		cacheRows = NewCacheRows(rows)
 		defer cacheRows.Close()
 
-		dest := make([]driver.Value, 2)
-		if err := cacheRows.Next(dest); err != nil {
-			return err
-		}
-
-		id, ok := dest[0].(int64)
-		if !ok {
-			return fmt.Errorf("dest[0] is not int, got=%T", dest[0])
-		}
-		if id != 1 {
-			return errors.New("id != 1")
-		}
-
-		name, ok := dest[1].([]byte)
-		if !ok {
-			return fmt.Errorf("dest[1] is not string, got=%T", dest[1])
-		}
-		if string(name) != "Alice" {
-			return errors.New("name != Alice")
-		}
-
-		return nil
+		return cacheRows.createCache()
 	})
 	if err != nil {
 		t.Error(err)

template/driver.go

@@ -153,7 +153,7 @@ func (r *CacheRows) Clone() *CacheRows {
 	}
 }
 
-func NewCachedRows(inner driver.Rows) *CacheRows {
+func NewCacheRows(inner driver.Rows) *CacheRows {
 	return &CacheRows{inner: inner}
 }
 
@@ -208,7 +208,6 @@ func (r *CacheRows) Close() error {
 		r.rows.reset()
 		return nil
 	}
-	r.cached = true
 	return r.inner.Close()
 }
 
@@ -244,3 +243,19 @@ func (r *CacheRows) Next(dest []driver.Value) error {
 
 	return nil
 }
+
+func (r *CacheRows) createCache() error {
+	columns := r.Columns()
+	dest := make([]driver.Value, len(columns))
+	for {
+		err := r.Next(dest)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return err
+		}
+	}
+	r.Close()
+	return nil
+}

template/driver.tmpl

@@ -152,7 +152,7 @@ func (r *CacheRows) Clone() *CacheRows {
 	}
 }
 
-func NewCachedRows(inner driver.Rows) *CacheRows {
+func NewCacheRows(inner driver.Rows) *CacheRows {
 	return &CacheRows{inner: inner}
 }
 
@@ -207,7 +207,6 @@ func (r *CacheRows) Close() error {
 		r.rows.reset()
 		return nil
 	}
-	r.cached = true
 	return r.inner.Close()
 }
 
@@ -243,3 +242,19 @@ func (r *CacheRows) Next(dest []driver.Value) error {
 
 	return nil
 }
+
+func (r *CacheRows) createCache() error {
+	columns := r.Columns()
+	dest := make([]driver.Value, len(columns))
+	for {
+		err := r.Next(dest)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return err
+		}
+	}
+	r.Close()
+	return nil
+}

template/stmt.go

@@ -175,6 +175,8 @@ func cacheKey(args []driver.Value) string {
 }
 
 func replaceFn(ctx context.Context, key string) (*CacheRows, error) {
+	var res *CacheRows
+
 	queryerCtx, ok := ctx.Value(queryerCtxKey{}).(driver.QueryerContext)
 	if ok {
 		query := ctx.Value(queryKey{}).(string)
@@ -183,16 +185,20 @@ func replaceFn(ctx context.Context, key string) (*CacheRows, error) {
 		if err != nil {
 			return nil, err
 		}
-		return NewCachedRows(rows), nil
+		res = NewCacheRows(rows)
+	} else {
+		stmt := ctx.Value(stmtKey{}).(*CustomCacheStatement)
+		args := ctx.Value(argsKey{}).([]driver.Value)
+		rows, err := stmt.inner.Query(args)
+		if err != nil {
+			return nil, err
+		}
+		res = NewCacheRows(rows)
 	}
 
-	stmt := ctx.Value(stmtKey{}).(*CustomCacheStatement)
-
-	args := ctx.Value(argsKey{}).([]driver.Value)
-	rows, err := stmt.inner.Query(args)
-	if err != nil {
+	if err := res.createCache(); err != nil {
 		return nil, err
 	}
 
-	return NewCachedRows(rows), nil
+	return res, nil
 }

template/stmt.tmpl

@@ -175,6 +175,8 @@ func cacheKey(args []driver.Value) string {
 }
 
 func replaceFn(ctx context.Context, key string) (*CacheRows, error) {
+	var res *CacheRows
+
 	queryerCtx, ok := ctx.Value(queryerCtxKey{}).(driver.QueryerContext)
 	if ok {
 		query := ctx.Value(queryKey{}).(string)
@@ -183,16 +185,20 @@ func replaceFn(ctx context.Context, key string) (*CacheRows, error) {
 		if err != nil {
 			return nil, err
 		}
-		return NewCachedRows(rows), nil
+		res = NewCacheRows(rows)
+	} else {
+		stmt := ctx.Value(stmtKey{}).(*CustomCacheStatement)
+		args := ctx.Value(argsKey{}).([]driver.Value)
+		rows, err := stmt.inner.Query(args)
+		if err != nil {
+			return nil, err
+		}
+		res = NewCacheRows(rows)
 	}
 
-	stmt := ctx.Value(stmtKey{}).(*CustomCacheStatement)
-
-	args := ctx.Value(argsKey{}).([]driver.Value)
-	rows, err := stmt.inner.Query(args)
-	if err != nil {
+	if err := res.createCache(); err != nil {
 		return nil, err
 	}
 
-	return NewCachedRows(rows), nil
+	return res, nil
 }