statement_cache.go

package dbx

import (
	"context"
	"database/sql"
	"sync"
)

type statementCache struct {
	mu    sync.Mutex
	max   int
	order []string
	items map[string]*sql.Stmt
}

func newStatementCache(max int) *statementCache {
	if max < 1 {
		max = 1
	}
	return &statementCache{
		max:   max,
		order: make([]string, 0, max),
		items: make(map[string]*sql.Stmt, max),
	}
}

func (c *statementCache) GetOrPrepare(ctx context.Context, db *sql.DB, query string) (*sql.Stmt, error) {
	c.mu.Lock()
	if stmt, ok := c.items[query]; ok {
		c.mu.Unlock()
		return stmt, nil
	}
	c.mu.Unlock()

	if ctx == nil {
		ctx = context.Background()
	}
	stmt, err := db.PrepareContext(ctx, query)
	if err != nil {
		return nil, err
	}

	c.mu.Lock()
	defer c.mu.Unlock()
	if existing, ok := c.items[query]; ok {
		_ = stmt.Close()
		return existing, nil
	}
	c.items[query] = stmt
	c.order = append(c.order, query)
	if len(c.order) > c.max {
		old := c.order[0]
		c.order = c.order[1:]
		if s, ok := c.items[old]; ok {
			delete(c.items, old)
			_ = s.Close()
		}
	}
	return stmt, nil
}

func (c *statementCache) Close() {
	c.mu.Lock()
	defer c.mu.Unlock()
	for _, stmt := range c.items {
		_ = stmt.Close()
	}
	c.items = map[string]*sql.Stmt{}
	c.order = c.order[:0]
}