Skip to content

Commit

Permalink
fixed tracing over database/sql
Browse files Browse the repository at this point in the history
  • Loading branch information
@asmyasnikov
asmyasnikov committed Dec 11, 2024
1 parent 33ae952 commit 09192ee
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 9 deletions.
74 changes: 65 additions & 9 deletions internal/xsql/conn.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import (
"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
internalTable "github.com/ydb-platform/ydb-go-sdk/v3/internal/table"
"github.com/ydb-platform/ydb-go-sdk/v3/internal/table/config"
"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
"github.com/ydb-platform/ydb-go-sdk/v3/internal/xslices"
"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/conn"
Expand All @@ -37,6 +38,7 @@ type (

cc conn.Conn
currentTx *txWrapper
ctx context.Context //nolint:containedctx

connector *Connector
lastUsage xsync.LastUsage
Expand All @@ -60,10 +62,13 @@ func (c *connWrapper) CheckNamedValue(value *driver.NamedValue) error {
return nil
}

func (c *connWrapper) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
if c.currentTx != nil {
return nil, xerrors.WithStackTrace(xerrors.AlreadyHasTx(c.currentTx.ID()))
}
func BeginTx(ctx context.Context, c *connWrapper, opts driver.TxOptions) (_ driver.Tx, finalErr error) {
onDone := trace.DatabaseSQLOnConnBegin(c.connector.trace, &ctx,
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql.BeginTx"),
)
defer func() {
onDone(c.currentTx, finalErr)
}()

tx, err := c.cc.BeginTx(ctx, opts)
if err != nil {
Expand All @@ -79,7 +84,22 @@ func (c *connWrapper) BeginTx(ctx context.Context, opts driver.TxOptions) (drive
return c.currentTx, nil
}

func (c *connWrapper) Close() error {
func (c *connWrapper) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
if c.currentTx != nil {
return nil, xerrors.WithStackTrace(xerrors.AlreadyHasTx(c.currentTx.ID()))
}

return BeginTx(ctx, c, opts)
}

func Close(c *connWrapper) (finalErr error) {
onDone := trace.DatabaseSQLOnConnClose(c.connector.Trace(), &c.ctx,
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql.Close"),
)
defer func() {
onDone(finalErr)
}()

err := c.cc.Close()
if err != nil {
return xerrors.WithStackTrace(err)
Expand All @@ -88,6 +108,10 @@ func (c *connWrapper) Close() error {
return nil
}

func (c *connWrapper) Close() (finalErr error) {
return Close(c)
}

func (c *connWrapper) Connector() *Connector {
return c.connector
}
Expand Down Expand Up @@ -139,9 +163,9 @@ func (c *connWrapper) Prepare(string) (driver.Stmt, error) {
return nil, errDeprecated
}

func (c *connWrapper) PrepareContext(ctx context.Context, sql string) (_ driver.Stmt, finalErr error) {
func PrepareContext(ctx context.Context, c *connWrapper, sql string) (_ driver.Stmt, finalErr error) {
onDone := trace.DatabaseSQLOnConnPrepare(c.connector.Trace(), &ctx,
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql.(*connWrapper).PrepareContext"),
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql.PrepareContext"),
sql,
)
defer func() {
Expand All @@ -160,6 +184,10 @@ func (c *connWrapper) PrepareContext(ctx context.Context, sql string) (_ driver.
}, nil
}

func (c *connWrapper) PrepareContext(ctx context.Context, sql string) (_ driver.Stmt, finalErr error) {
return PrepareContext(ctx, c, sql)
}

func (c *connWrapper) LastUsage() time.Time {
return c.lastUsage.Get()
}
Expand All @@ -178,7 +206,17 @@ func (c *connWrapper) toYdb(sql string, args ...driver.NamedValue) (yql string,
return yql, &params, nil
}

func (c *connWrapper) QueryContext(ctx context.Context, sql string, args []driver.NamedValue) (driver.Rows, error) {
func QueryContext(ctx context.Context, c *connWrapper, sql string, args []driver.NamedValue) (
_ driver.Rows, finalErr error,
) {
onDone := trace.DatabaseSQLOnConnQuery(c.connector.Trace(), &ctx,
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql.QueryContext"),
sql, c.connector.processor.String(), xcontext.IsIdempotent(ctx), c.connector.clock.Since(c.lastUsage.Get()),
)
defer func() {
onDone(finalErr)
}()

done := c.lastUsage.Start()
defer done()

Expand All @@ -203,7 +241,21 @@ func (c *connWrapper) QueryContext(ctx context.Context, sql string, args []drive
return c.cc.Query(ctx, sql, params)
}

func (c *connWrapper) ExecContext(ctx context.Context, sql string, args []driver.NamedValue) (driver.Result, error) {
func (c *connWrapper) QueryContext(ctx context.Context, sql string, args []driver.NamedValue) (driver.Rows, error) {
return QueryContext(ctx, c, sql, args)
}

func ExecContext(ctx context.Context, c *connWrapper, sql string, args []driver.NamedValue) (
_ driver.Result, finalErr error,
) {
onDone := trace.DatabaseSQLOnConnExec(c.connector.Trace(), &ctx,
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql.ExecContext"),
sql, c.connector.processor.String(), xcontext.IsIdempotent(ctx), c.connector.clock.Since(c.lastUsage.Get()),
)
defer func() {
onDone(finalErr)
}()

done := c.lastUsage.Start()
defer done()

Expand All @@ -219,6 +271,10 @@ func (c *connWrapper) ExecContext(ctx context.Context, sql string, args []driver
return c.cc.Exec(ctx, sql, params)
}

func (c *connWrapper) ExecContext(ctx context.Context, sql string, args []driver.NamedValue) (driver.Result, error) {
return ExecContext(ctx, c, sql, args)
}

func (c *connWrapper) GetDatabaseName() string {
return c.connector.Name()
}
Expand Down
13 changes: 13 additions & 0 deletions internal/xsql/connector.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,17 @@ type (
}
)

func (e Engine) String() string {
switch e {
case TABLE_SERVICE:
return "TABLE_SERVICE"
case QUERY_SERVICE:
return "QUERY_SERVICE"
default:
return "UNKNOWN"
}
}

func (c *Connector) RetryBudget() budget.Budget {
return c.retryBudget
}
Expand Down Expand Up @@ -130,6 +141,7 @@ func (c *Connector) Connect(ctx context.Context) (driver.Conn, error) {
c.conns.Delete(id)
}))...,
),
ctx: ctx,
connector: c,
lastUsage: xsync.NewLastUsage(xsync.WithClock(c.Clock())),
}
Expand All @@ -153,6 +165,7 @@ func (c *Connector) Connect(ctx context.Context) (driver.Conn, error) {
c.conns.Delete(id)
}))...,
),
ctx: ctx,
connector: c,
lastUsage: xsync.NewLastUsage(xsync.WithClock(c.Clock())),
}
Expand Down

0 comments on commit 09192ee

Please sign in to comment.