stardb/sql_runtime.go

package stardb

import (
	internalsqlruntime "b612.me/stardb/internal/sqlruntime"
	"context"
	"time"
)

// SQLBeforeHook runs before a SQL statement is executed.
type SQLBeforeHook func(ctx context.Context, query string, args []interface{})

// SQLAfterHook runs after a SQL statement is executed.
type SQLAfterHook func(ctx context.Context, query string, args []interface{}, duration time.Duration, err error)

// PlaceholderStyle controls SQL placeholder format conversion.
type PlaceholderStyle int

const (
	// PlaceholderQuestion keeps '?' placeholders unchanged.
	PlaceholderQuestion PlaceholderStyle = iota
	// PlaceholderDollar converts '?' placeholders to '$1,$2,...'.
	PlaceholderDollar
)

// SQLFingerprintMode controls SQL fingerprint generation strategy.
type SQLFingerprintMode int

const (
	// SQLFingerprintBasic lowercases SQL and collapses whitespace.
	SQLFingerprintBasic SQLFingerprintMode = iota
	// SQLFingerprintMaskLiterals also masks numeric/string literals and $n placeholders.
	SQLFingerprintMaskLiterals
)

type sqlHookMetaKey struct{}

// SQLHookMeta contains extra hook metadata attached to context.
type SQLHookMeta struct {
	Fingerprint string
}

// SQLHookMetaFromContext extracts SQL hook metadata from context.
func SQLHookMetaFromContext(ctx context.Context) (SQLHookMeta, bool) {
	if ctx == nil {
		return SQLHookMeta{}, false
	}
	meta, ok := ctx.Value(sqlHookMetaKey{}).(SQLHookMeta)
	return meta, ok
}

func withSQLHookMeta(ctx context.Context, meta SQLHookMeta) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithValue(ctx, sqlHookMetaKey{}, meta)
}

type sqlRuntime struct {
	state internalsqlruntime.State
}

func normalizePlaceholderStyle(style PlaceholderStyle) PlaceholderStyle {
	return PlaceholderStyle(internalsqlruntime.NormalizePlaceholderStyle(int(style)))
}

func normalizeSQLFingerprintMode(mode SQLFingerprintMode) SQLFingerprintMode {
	return SQLFingerprintMode(internalsqlruntime.NormalizeFingerprintMode(int(mode)))
}

func cloneHookArgs(args []interface{}) []interface{} {
	return internalsqlruntime.CloneHookArgs(args)
}

func (s *StarDB) runtimeOptions() (SQLBeforeHook, SQLAfterHook, PlaceholderStyle, time.Duration) {
	if s == nil {
		return nil, nil, PlaceholderQuestion, 0
	}

	beforeAny, afterAny, rawStyle, slowThreshold := s.runtime.state.Options()
	var before SQLBeforeHook
	if b, ok := beforeAny.(SQLBeforeHook); ok {
		before = b
	}
	var after SQLAfterHook
	if a, ok := afterAny.(SQLAfterHook); ok {
		after = a
	}

	return before, after, normalizePlaceholderStyle(PlaceholderStyle(rawStyle)), slowThreshold
}

func (s *StarDB) sqlHooks() (SQLBeforeHook, SQLAfterHook, time.Duration) {
	before, after, _, slowThreshold := s.runtimeOptions()
	return before, after, slowThreshold
}

func (s *StarDB) prepareSQLCall(query string, args []interface{}) (string, SQLBeforeHook, SQLAfterHook, []interface{}, time.Duration) {
	before, after, style, slowThreshold := s.runtimeOptions()
	query = ConvertPlaceholders(query, style)
	if before == nil && after == nil {
		return query, nil, nil, nil, slowThreshold
	}
	return query, before, after, cloneHookArgs(args), slowThreshold
}

func (s *StarDB) hookContext(ctx context.Context, query string, before SQLBeforeHook, after SQLAfterHook) context.Context {
	if s == nil {
		return ctx
	}

	needCounter := s.SQLFingerprintCounterEnabled()
	needMeta := (before != nil || after != nil) && s.SQLFingerprintEnabled()
	if !needCounter && !needMeta {
		return ctx
	}

	mode := s.SQLFingerprintMode()
	keepComments := s.SQLFingerprintKeepComments()
	fingerprint := internalsqlruntime.FingerprintSQL(query, int(mode), keepComments)
	if fingerprint == "" {
		return ctx
	}

	if needCounter {
		s.runtime.state.IncFingerprintCount(fingerprint)
	}
	if !needMeta {
		return ctx
	}
	return withSQLHookMeta(ctx, SQLHookMeta{Fingerprint: fingerprint})
}

func shouldRunAfterHook(after SQLAfterHook, slowThreshold, duration time.Duration, err error) bool {
	return internalsqlruntime.ShouldRunAfterHook(after != nil, slowThreshold, duration, err)
}

// SetSQLHooks sets SQL before/after hooks.
func (s *StarDB) SetSQLHooks(before SQLBeforeHook, after SQLAfterHook) {
	if s == nil {
		return
	}
	s.runtime.state.SetHooks(before, after)
}

// SetSQLBeforeHook sets SQL before hook.
func (s *StarDB) SetSQLBeforeHook(before SQLBeforeHook) {
	if s == nil {
		return
	}
	s.runtime.state.SetBeforeHook(before)
}

// SetSQLAfterHook sets SQL after hook.
func (s *StarDB) SetSQLAfterHook(after SQLAfterHook) {
	if s == nil {
		return
	}
	s.runtime.state.SetAfterHook(after)
}

// SetPlaceholderStyle sets placeholder conversion style.
func (s *StarDB) SetPlaceholderStyle(style PlaceholderStyle) {
	if s == nil {
		return
	}
	s.runtime.state.SetPlaceholderStyle(int(style))
}

// SetSQLSlowThreshold sets minimum duration for triggering after hook.
// When threshold > 0, after hook runs only for statements slower than threshold or those with error.
func (s *StarDB) SetSQLSlowThreshold(threshold time.Duration) {
	if s == nil {
		return
	}
	s.runtime.state.SetSlowThreshold(threshold)
}

// SQLSlowThreshold returns current slow SQL threshold.
func (s *StarDB) SQLSlowThreshold() time.Duration {
	if s == nil {
		return 0
	}
	return s.runtime.state.SlowThreshold()
}

// PlaceholderStyle returns current placeholder style.
func (s *StarDB) PlaceholderStyle() PlaceholderStyle {
	if s == nil {
		return PlaceholderQuestion
	}
	style := PlaceholderStyle(s.runtime.state.PlaceholderStyle())
	return normalizePlaceholderStyle(style)
}

// SetSQLFingerprintEnabled toggles SQL fingerprint metadata generation for hooks.
func (s *StarDB) SetSQLFingerprintEnabled(enabled bool) {
	if s == nil {
		return
	}
	s.runtime.state.SetFingerprintEnabled(enabled)
}

// SQLFingerprintEnabled reports whether SQL fingerprint metadata generation is enabled.
func (s *StarDB) SQLFingerprintEnabled() bool {
	if s == nil {
		return false
	}
	return s.runtime.state.FingerprintEnabled()
}

// SetSQLFingerprintMode sets SQL fingerprint generation mode.
func (s *StarDB) SetSQLFingerprintMode(mode SQLFingerprintMode) {
	if s == nil {
		return
	}
	s.runtime.state.SetFingerprintMode(int(mode))
}

// SQLFingerprintMode returns SQL fingerprint generation mode.
func (s *StarDB) SQLFingerprintMode() SQLFingerprintMode {
	if s == nil {
		return SQLFingerprintBasic
	}
	mode := SQLFingerprintMode(s.runtime.state.FingerprintMode())
	return normalizeSQLFingerprintMode(mode)
}

// SetSQLFingerprintKeepComments controls whether comments are preserved in SQL fingerprints.
// Default is false.
func (s *StarDB) SetSQLFingerprintKeepComments(keep bool) {
	if s == nil {
		return
	}
	s.runtime.state.SetFingerprintKeepComments(keep)
}

// SQLFingerprintKeepComments reports whether SQL fingerprints preserve comments.
func (s *StarDB) SQLFingerprintKeepComments() bool {
	if s == nil {
		return false
	}
	return s.runtime.state.FingerprintKeepComments()
}

// SetSQLFingerprintCounterEnabled toggles in-memory SQL fingerprint hit counter.
// Default is false.
func (s *StarDB) SetSQLFingerprintCounterEnabled(enabled bool) {
	if s == nil {
		return
	}
	s.runtime.state.SetFingerprintCounterEnabled(enabled)
}

// SQLFingerprintCounterEnabled reports whether in-memory SQL fingerprint hit counter is enabled.
func (s *StarDB) SQLFingerprintCounterEnabled() bool {
	if s == nil {
		return false
	}
	return s.runtime.state.FingerprintCounterEnabled()
}

// SQLFingerprintCounters returns a snapshot of fingerprint hit counters.
func (s *StarDB) SQLFingerprintCounters() map[string]uint64 {
	if s == nil {
		return map[string]uint64{}
	}
	return s.runtime.state.FingerprintCountsSnapshot()
}

// ResetSQLFingerprintCounters clears all in-memory fingerprint hit counters.
func (s *StarDB) ResetSQLFingerprintCounters() {
	if s == nil {
		return
	}
	s.runtime.state.ResetFingerprintCounts()
}