Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/repository/server_timing_sql.go

@@ -0,0 +1,311 @@
+package repository
+
+import (
+	"context"
+	"database/sql/driver"
+	"errors"
+	"io"
+	"reflect"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
+)
+
+type serverTimingConnector struct {
+	base driver.Connector
+}
+
+func newServerTimingConnector(base driver.Connector) driver.Connector {
+	return &serverTimingConnector{base: base}
+}
+
+func (c *serverTimingConnector) Connect(ctx context.Context) (driver.Conn, error) {
+	startedAt := time.Now()
+	conn, err := c.base.Connect(ctx)
+	servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
+	if err != nil {
+		return nil, err
+	}
+	return &serverTimingConn{Conn: conn}, nil
+}
+
+func (c *serverTimingConnector) Driver() driver.Driver {
+	return c.base.Driver()
+}
+
+type serverTimingConn struct {
+	driver.Conn
+}
+
+func (c *serverTimingConn) Prepare(query string) (driver.Stmt, error) {
+	stmt, err := c.Conn.Prepare(query)
+	if err != nil {
+		return nil, err
+	}
+	return &serverTimingStmt{Stmt: stmt}, nil
+}
+
+func (c *serverTimingConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
+	startedAt := time.Now()
+	var (
+		stmt driver.Stmt
+		err  error
+	)
+	if preparer, ok := c.Conn.(driver.ConnPrepareContext); ok {
+		stmt, err = preparer.PrepareContext(ctx, query)
+	} else {
+		stmt, err = c.Conn.Prepare(query)
+	}
+	servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
+	if err != nil {
+		return nil, err
+	}
+	return &serverTimingStmt{Stmt: stmt}, nil
+}
+
+func (c *serverTimingConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
+	execer, ok := c.Conn.(driver.ExecerContext)
+	if !ok {
+		return nil, driver.ErrSkip
+	}
+	startedAt := time.Now()
+	result, err := execer.ExecContext(ctx, query, args)
+	servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
+	return result, err
+}
+
+func (c *serverTimingConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+	queryer, ok := c.Conn.(driver.QueryerContext)
+	if !ok {
+		return nil, driver.ErrSkip
+	}
+	startedAt := time.Now()
+	rows, err := queryer.QueryContext(ctx, query, args)
+	servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
+	if err != nil || rows == nil {
+		return rows, err
+	}
+	return newServerTimingRows(ctx, rows), nil
+}
+
+func (c *serverTimingConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	startedAt := time.Now()
+	var (
+		tx  driver.Tx
+		err error
+	)
+	if beginner, ok := c.Conn.(driver.ConnBeginTx); ok {
+		tx, err = beginner.BeginTx(ctx, opts)
+	} else {
+		if opts.Isolation != driver.IsolationLevel(0) {
+			return nil, errors.New("driver does not support non-default isolation")
+		}
+		if opts.ReadOnly {
+			return nil, errors.New("driver does not support read-only transactions")
+		}
+		// The wrapper exposes ConnBeginTx, so it must retain database/sql's
+		// legacy fallback for drivers that only implement Conn.Begin.
+		tx, err = c.Conn.Begin() //nolint:staticcheck // Required driver compatibility fallback.
+	}
+	servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
+	if err != nil || tx == nil {
+		return tx, err
+	}
+	return &serverTimingTx{Tx: tx, ctx: ctx}, nil
+}
+
+func (c *serverTimingConn) Ping(ctx context.Context) error {
+	if pinger, ok := c.Conn.(driver.Pinger); ok {
+		startedAt := time.Now()
+		err := pinger.Ping(ctx)
+		servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
+		return err
+	}
+	return nil
+}
+
+func (c *serverTimingConn) ResetSession(ctx context.Context) error {
+	if resetter, ok := c.Conn.(driver.SessionResetter); ok {
+		startedAt := time.Now()
+		err := resetter.ResetSession(ctx)
+		servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
+		return err
+	}
+	return nil
+}
+
+func (c *serverTimingConn) IsValid() bool {
+	if validator, ok := c.Conn.(driver.Validator); ok {
+		return validator.IsValid()
+	}
+	return true
+}
+
+func (c *serverTimingConn) CheckNamedValue(value *driver.NamedValue) error {
+	if checker, ok := c.Conn.(driver.NamedValueChecker); ok {
+		return checker.CheckNamedValue(value)
+	}
+	return driver.ErrSkip
+}
+
+type serverTimingStmt struct {
+	driver.Stmt
+}
+
+func (s *serverTimingStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
+	startedAt := time.Now()
+	var (
+		result driver.Result
+		err    error
+	)
+	if execer, ok := s.Stmt.(driver.StmtExecContext); ok {
+		result, err = execer.ExecContext(ctx, args)
+	} else {
+		var values []driver.Value
+		values, err = namedValues(args)
+		if err == nil {
+			// The wrapper exposes StmtExecContext and must preserve the fallback
+			// database/sql would use for a legacy driver statement.
+			result, err = s.Stmt.Exec(values) //nolint:staticcheck // Required driver compatibility fallback.
+		}
+	}
+	servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
+	return result, err
+}
+
+func (s *serverTimingStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
+	startedAt := time.Now()
+	var (
+		rows driver.Rows
+		err  error
+	)
+	if queryer, ok := s.Stmt.(driver.StmtQueryContext); ok {
+		rows, err = queryer.QueryContext(ctx, args)
+	} else {
+		var values []driver.Value
+		values, err = namedValues(args)
+		if err == nil {
+			// The wrapper exposes StmtQueryContext and must preserve the fallback
+			// database/sql would use for a legacy driver statement.
+			rows, err = s.Stmt.Query(values) //nolint:staticcheck // Required driver compatibility fallback.
+		}
+	}
+	servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
+	if err != nil || rows == nil {
+		return rows, err
+	}
+	return newServerTimingRows(ctx, rows), nil
+}
+
+func (s *serverTimingStmt) CheckNamedValue(value *driver.NamedValue) error {
+	if checker, ok := s.Stmt.(driver.NamedValueChecker); ok {
+		return checker.CheckNamedValue(value)
+	}
+	return driver.ErrSkip
+}
+
+func namedValues(args []driver.NamedValue) ([]driver.Value, error) {
+	values := make([]driver.Value, len(args))
+	for i, arg := range args {
+		if arg.Name != "" {
+			return nil, errors.New("named parameters are not supported")
+		}
+		values[i] = arg.Value
+	}
+	return values, nil
+}
+
+type serverTimingRows struct {
+	driver.Rows
+	ctx context.Context
+}
+
+func newServerTimingRows(ctx context.Context, rows driver.Rows) *serverTimingRows {
+	return &serverTimingRows{Rows: rows, ctx: ctx}
+}
+
+func (r *serverTimingRows) Close() error {
+	startedAt := time.Now()
+	err := r.Rows.Close()
+	servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now())
+	return err
+}
+
+func (r *serverTimingRows) Next(dest []driver.Value) error {
+	startedAt := time.Now()
+	err := r.Rows.Next(dest)
+	servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now())
+	return err
+}
+
+func (r *serverTimingRows) HasNextResultSet() bool {
+	if rows, ok := r.Rows.(driver.RowsNextResultSet); ok {
+		return rows.HasNextResultSet()
+	}
+	return false
+}
+
+func (r *serverTimingRows) NextResultSet() error {
+	rows, ok := r.Rows.(driver.RowsNextResultSet)
+	if !ok {
+		return io.EOF
+	}
+	startedAt := time.Now()
+	err := rows.NextResultSet()
+	servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now())
+	return err
+}
+
+func (r *serverTimingRows) ColumnTypeScanType(index int) reflect.Type {
+	if rows, ok := r.Rows.(driver.RowsColumnTypeScanType); ok {
+		return rows.ColumnTypeScanType(index)
+	}
+	return reflect.TypeOf(new(any)).Elem()
+}
+
+func (r *serverTimingRows) ColumnTypeDatabaseTypeName(index int) string {
+	if rows, ok := r.Rows.(driver.RowsColumnTypeDatabaseTypeName); ok {
+		return rows.ColumnTypeDatabaseTypeName(index)
+	}
+	return ""
+}
+
+func (r *serverTimingRows) ColumnTypeLength(index int) (int64, bool) {
+	if rows, ok := r.Rows.(driver.RowsColumnTypeLength); ok {
+		return rows.ColumnTypeLength(index)
+	}
+	return 0, false
+}
+
+func (r *serverTimingRows) ColumnTypeNullable(index int) (bool, bool) {
+	if rows, ok := r.Rows.(driver.RowsColumnTypeNullable); ok {
+		return rows.ColumnTypeNullable(index)
+	}
+	return false, false
+}
+
+func (r *serverTimingRows) ColumnTypePrecisionScale(index int) (int64, int64, bool) {
+	if rows, ok := r.Rows.(driver.RowsColumnTypePrecisionScale); ok {
+		return rows.ColumnTypePrecisionScale(index)
+	}
+	return 0, 0, false
+}
+
+type serverTimingTx struct {
+	driver.Tx
+	ctx context.Context
+}
+
+func (t *serverTimingTx) Commit() error {
+	startedAt := time.Now()
+	err := t.Tx.Commit()
+	servertiming.RecordInterval(t.ctx, servertiming.MetricDatabase, startedAt, time.Now())
+	return err
+}
+
+func (t *serverTimingTx) Rollback() error {
+	startedAt := time.Now()
+	err := t.Tx.Rollback()
+	servertiming.RecordInterval(t.ctx, servertiming.MetricDatabase, startedAt, time.Now())
+	return err
+}