sqlcomment: expose the underlying driver QueryContext/ExecContext methods (#5)

Author: a8m

driver.go

@@ -9,26 +9,29 @@ import (
 )
 
 type (
-	commenter struct {
-		options
-	}
-	// Driver is a driver that adds an SQL comment (see https://google.github.io/sqlcommenter/).
+	// Driver is a driver that adds an SQL comment.
+	// See: https://google.github.io/sqlcommenter.
 	Driver struct {
 		dialect.Driver // underlying driver.
 		commenter
 	}
+
 	// Tx is a transaction implementation that adds an SQL comment.
 	Tx struct {
 		dialect.Tx                 // underlying transaction.
 		ctx        context.Context // underlying transaction context.
 		commenter
 	}
+
+	commenter struct {
+		options
+	}
 )
 
 // NewDriver decorates the given driver and adds an SQL comment to every query.
 func NewDriver(drv dialect.Driver, options ...Option) dialect.Driver {
-	defaultCommenters := []Tagger{contextTagger{}}
-	opts := buildOptions(append(options, WithTagger(defaultCommenters...)))
+	taggers := []Tagger{contextTagger{}}
+	opts := buildOptions(append(options, WithTagger(taggers...)))
 	return &Driver{drv, commenter{opts}}
 }
 
@@ -40,14 +43,36 @@ func (c commenter) withComment(ctx context.Context, query string) string {
 	return fmt.Sprintf("%s /*%s*/", query, tags.Marshal())
 }
 
+// Query adds an SQL comment to the original query and calls the underlying driver Query method.
+func (d *Driver) Query(ctx context.Context, query string, args, v interface{}) error {
+	return d.Driver.Query(ctx, d.withComment(ctx, query), args, v)
+}
+
+// QueryContext calls QueryContext of the underlying driver, or fails if it is not supported.
+func (d *Driver) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
+	drv, ok := d.Driver.(interface {
+		QueryContext(context.Context, string, ...interface{}) (*sql.Rows, error)
+	})
+	if !ok {
+		return nil, fmt.Errorf("Driver.QueryContext is not supported")
+	}
+	return drv.QueryContext(ctx, d.withComment(ctx, query), args...)
+}
+
 // Exec adds an SQL comment to the original query and calls the underlying driver Exec method.
 func (d *Driver) Exec(ctx context.Context, query string, args, v interface{}) error {
 	return d.Driver.Exec(ctx, d.withComment(ctx, query), args, v)
 }
 
-// Query adds an SQL comment to the original query and calls the underlying driver Query method.
-func (d *Driver) Query(ctx context.Context, query string, args, v interface{}) error {
-	return d.Driver.Query(ctx, d.withComment(ctx, query), args, v)
+// ExecContext calls ExecContext of the underlying driver, or fails if it is not supported.
+func (d *Driver) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
+	drv, ok := d.Driver.(interface {
+		ExecContext(context.Context, string, ...interface{}) (sql.Result, error)
+	})
+	if !ok {
+		return nil, fmt.Errorf("Driver.ExecContext is not supported")
+	}
+	return drv.ExecContext(ctx, d.withComment(ctx, query), args...)
 }
 
 // Tx wraps the underlying Tx command with a commenter.
@@ -79,11 +104,33 @@ func (d *Tx) Exec(ctx context.Context, query string, args, v interface{}) error
 	return d.Tx.Exec(ctx, d.withComment(ctx, query), args, v)
 }
 
+// ExecContext logs its params and calls the underlying transaction ExecContext method if it is supported.
+func (d *Tx) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
+	tx, ok := d.Tx.(interface {
+		ExecContext(context.Context, string, ...interface{}) (sql.Result, error)
+	})
+	if !ok {
+		return nil, fmt.Errorf("Tx.ExecContext is not supported")
+	}
+	return tx.ExecContext(ctx, d.withComment(ctx, query), args...)
+}
+
 // Query adds an SQL comment and calls the underlying transaction Query method.
 func (d *Tx) Query(ctx context.Context, query string, args, v interface{}) error {
 	return d.Tx.Query(ctx, d.withComment(ctx, query), args, v)
 }
 
+// QueryContext logs its params and calls the underlying transaction QueryContext method if it is supported.
+func (d *Tx) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
+	tx, ok := d.Tx.(interface {
+		QueryContext(context.Context, string, ...interface{}) (*sql.Rows, error)
+	})
+	if !ok {
+		return nil, fmt.Errorf("Tx.QueryContext is not supported")
+	}
+	return tx.QueryContext(ctx, d.withComment(ctx, query), args...)
+}
+
 // Commit commits the underlying Tx.
 func (d *Tx) Commit() error {
 	return d.Tx.Commit()

go.mod

@@ -1,11 +1,11 @@
 module ariga.io/sqlcomment
 
-go 1.17
+go 1.19
 
 require (
-	entgo.io/ent v0.9.1
-	github.com/mattn/go-sqlite3 v1.14.8
-	github.com/stretchr/testify v1.7.0
+	entgo.io/ent v0.11.3-0.20220816070906-2b54aadcce3a
+	github.com/mattn/go-sqlite3 v1.14.14
+	github.com/stretchr/testify v1.8.0
 	go.opencensus.io v0.23.0
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.25.0
 	go.opentelemetry.io/otel v1.0.1
@@ -14,14 +14,24 @@ require (
 )
 
 require (
+	ariga.io/atlas v0.6.0 // indirect
+	github.com/agext/levenshtein v1.2.1 // indirect
+	github.com/apparentlymart/go-textseg/v13 v13.0.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/felixge/httpsnoop v1.0.2 // indirect
-	github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e // indirect
+	github.com/go-openapi/inflect v0.19.0 // indirect
+	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
+	github.com/google/go-cmp v0.5.6 // indirect
 	github.com/google/uuid v1.3.0 // indirect
+	github.com/hashicorp/hcl/v2 v2.13.0 // indirect
+	github.com/mitchellh/go-wordwrap v0.0.0-20150314170334-ad45545899c7 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/zclconf/go-cty v1.8.0 // indirect
 	go.opentelemetry.io/otel/internal/metric v0.24.0 // indirect
 	go.opentelemetry.io/otel/metric v0.24.0 // indirect
 	go.opentelemetry.io/otel/trace v1.0.1 // indirect
-	golang.org/x/sys v0.0.0-20210510120138-977fb7262007 // indirect
-	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
+	golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4 // indirect
+	golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f // indirect
+	golang.org/x/text v0.3.7 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )