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() }