Add tracing

ntpdb/tx.go

@@ -0,0 +1,64 @@
+package ntpdb
+
+import (
+	"context"
+	"database/sql"
+)
+
+type QuerierTx interface {
+	Querier
+
+	Begin(ctx context.Context) (QuerierTx, error)
+	Commit(ctx context.Context) error
+	Rollback(ctx context.Context) error
+}
+
+type Beginner interface {
+	Begin(context.Context) (sql.Tx, error)
+}
+
+type Tx interface {
+	Begin(context.Context) (sql.Tx, error)
+	Commit(ctx context.Context) error
+	Rollback(ctx context.Context) error
+}
+
+func (q *Queries) Begin(ctx context.Context) (QuerierTx, error) {
+	tx, err := q.db.(Beginner).Begin(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return &Queries{db: &tx}, nil
+}
+
+func (q *Queries) Commit(ctx context.Context) error {
+	tx, ok := q.db.(Tx)
+	if !ok {
+		return sql.ErrTxDone
+	}
+	return tx.Commit(ctx)
+}
+
+func (q *Queries) Rollback(ctx context.Context) error {
+	tx, ok := q.db.(Tx)
+	if !ok {
+		return sql.ErrTxDone
+	}
+	return tx.Rollback(ctx)
+}
+
+type WrappedQuerier struct {
+	QuerierTxWithTracing
+}
+
+func NewWrappedQuerier(q QuerierTx) QuerierTx {
+	return &WrappedQuerier{NewQuerierTxWithTracing(q, "")}
+}
+
+func (wq *WrappedQuerier) Begin(ctx context.Context) (QuerierTx, error) {
+	q, err := wq.QuerierTxWithTracing.Begin(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return NewWrappedQuerier(q), nil
+}