sjs994 · sjs994 · Nov 21, 2022 · Nov 21, 2022 · Nov 22, 2022 · Nov 22, 2022
diff --git a/go/core/core.go b/go/core/core.go
@@ -31,17 +31,29 @@ const (
 	Controller  string = "controller"
 	Action      string = "action"
 	Framework   string = "framework"
-	Driver      string = "driver"
+	Driver      string = "db_driver"
 	Traceparent string = "traceparent"
+	Application string = "application"
 )
 
-type CommenterOptions struct {
+type CommenterConfig struct {
 	EnableDBDriver    bool
 	EnableRoute       bool
 	EnableFramework   bool
 	EnableController  bool
 	EnableAction      bool
 	EnableTraceparent bool
+	EnableApplication bool
+}
+
+type StaticTags struct {
+	Application string
+	DriverName  string
+}
+
+type CommenterOptions struct {
+	Config CommenterConfig
+	Tags   StaticTags
 }
 
 func encodeURL(k string) string {

diff --git a/go/database/sql/connection.go b/go/database/sql/connection.go
@@ -0,0 +1,186 @@
+package sql
+
+import (
+	"context"
+	"database/sql/driver"
+	"fmt"
+	"runtime/debug"
+	"strings"
+
+	"github.com/google/sqlcommenter/go/core"
+)
+
+var attemptedToAutosetApplication = false
+
+type sqlCommenterConn struct {
+	driver.Conn
+	options core.CommenterOptions
+}
+
+func newSQLCommenterConn(conn driver.Conn, options core.CommenterOptions) *sqlCommenterConn {
+	return &sqlCommenterConn{
+		Conn:    conn,
+		options: options,
+	}
+}
+
+/*
+TODO: Check whether to implement or not ?
+type Conn interface {
+	// Prepare returns a prepared statement, bound to this connection.
+	Prepare(query string) (Stmt, error)
+
+	// Close invalidates and potentially stops any current
+	// prepared statements and transactions, marking this
+	// connection as no longer in use.
+	//
+	// Because the sql package maintains a free pool of
+	// connections and only calls Close when there's a surplus of
+	// idle connections, it shouldn't be necessary for drivers to
+	// do their own connection caching.
+	//
+	// Drivers must ensure all network calls made by Close
+	// do not block indefinitely (e.g. apply a timeout).
+	Close() error
+
+	// Begin starts and returns a new transaction.
+	//
+	// Deprecated: Drivers should implement ConnBeginTx instead (or additionally).
+	Begin() (Tx, error)
+}
+
+// ConnPrepareContext enhances the Conn interface with context.
+type ConnPrepareContext interface {
+	// PrepareContext returns a prepared statement, bound to this connection.
+	// context is for the preparation of the statement,
+	// it must not store the context within the statement itself.
+	PrepareContext(ctx context.Context, query string) (Stmt, error)
+}
+
+// ConnBeginTx enhances the Conn interface with context and TxOptions.
+type ConnBeginTx interface {
+	// BeginTx starts and returns a new transaction.
+	// If the context is canceled by the user the sql package will
+	// call Tx.Rollback before discarding and closing the connection.
+	//
+	// This must check opts.Isolation to determine if there is a set
+	// isolation level. If the driver does not support a non-default
+	// level and one is set or if there is a non-default isolation level
+	// that is not supported, an error must be returned.
+	//
+	// This must also check opts.ReadOnly to determine if the read-only
+	// value is true to either set the read-only transaction property if supported
+	// or return an error if it is not supported.
+	BeginTx(ctx context.Context, opts TxOptions) (Tx, error)
+}
+
+// SessionResetter may be implemented by Conn to allow drivers to reset the
+// session state associated with the connection and to signal a bad connection.
+type SessionResetter interface {
+	// ResetSession is called prior to executing a query on the connection
+	// if the connection has been used before. If the driver returns ErrBadConn
+	// the connection is discarded.
+	ResetSession(ctx context.Context) error
+}
+*/
+
+func (s *sqlCommenterConn) Query(query string, args []driver.Value) (driver.Rows, error) {
+	queryer, ok := s.Conn.(driver.Queryer)
+	if !ok {
+		return nil, driver.ErrSkip
+	}
+	ctx := context.Background()
+	extendedQuery := s.withComment(ctx, query)
+	return queryer.Query(extendedQuery, args)
+}
+
+func (s *sqlCommenterConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+	queryer, ok := s.Conn.(driver.QueryerContext)
+	if !ok {
+		return nil, driver.ErrSkip
+	}
+	extendedQuery := s.withComment(ctx, query)
+	return queryer.QueryContext(ctx, extendedQuery, args)
+}
+
+func (s *sqlCommenterConn) Exec(query string, args []driver.Value) (driver.Result, error) {
+	execor, ok := s.Conn.(driver.Execer)
+	if !ok {
+		return nil, driver.ErrSkip
+	}
+	ctx := context.Background()
+	extendedQuery := s.withComment(ctx, query)
+	return execor.Exec(extendedQuery, args)
+}
+
+func (s *sqlCommenterConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
+	execor, ok := s.Conn.(driver.ExecerContext)
+	if !ok {
+		return nil, driver.ErrSkip
+	}
+	extendedQuery := s.withComment(ctx, query)
+	return execor.ExecContext(ctx, extendedQuery, args)
+}
+
+func (s *sqlCommenterConn) Raw() driver.Conn {
+	return s.Conn
+}
+
+// ***** Commenter Functions *****
+
+func (conn *sqlCommenterConn) withComment(ctx context.Context, query string) string {
+	var commentsMap = map[string]string{}
+	query = strings.TrimSpace(query)
+	config := conn.options.Config
+
+	// Sorted alphabetically
+	if config.EnableAction && (ctx.Value(core.Action) != nil) {
+		commentsMap[core.Action] = ctx.Value(core.Action).(string)
+	}
+
+	// `driver` information should not be coming from framework.
+	// So, explicitly adding that here.
+	if config.EnableDBDriver {
+		commentsMap[core.Driver] = fmt.Sprintf("database/sql:%s", conn.options.Tags.DriverName)
+	}
+
+	if config.EnableFramework && (ctx.Value(core.Framework) != nil) {
+		commentsMap[core.Framework] = ctx.Value(core.Framework).(string)
+	}
+
+	if config.EnableRoute && (ctx.Value(core.Route) != nil) {
+		commentsMap[core.Route] = ctx.Value(core.Route).(string)
+	}
+
+	if config.EnableTraceparent {
+		carrier := core.ExtractTraceparent(ctx)
+		if val, ok := carrier["traceparent"]; ok {
+			commentsMap[core.Traceparent] = val
+		}
+	}
+
+	if config.EnableApplication {
+		if !attemptedToAutosetApplication && conn.options.Tags.Application == "" {
+			attemptedToAutosetApplication = true
+			bi, ok := debug.ReadBuildInfo()
+			if ok {
+				conn.options.Tags.Application = bi.Path
+			}
+		}
+		commentsMap[core.Application] = conn.options.Tags.Application
+	}
+
+	var commentsString string = ""
+	if len(commentsMap) > 0 { // Converts comments map to string and appends it to query
+		commentsString = fmt.Sprintf("/*%s*/", core.ConvertMapToComment(commentsMap))
+	}
+
+	// A semicolon at the end of the SQL statement means the query ends there.
+	// We need to insert the comment before that to be considered as part of the SQL statemtent.
+	if query[len(query)-1:] == ";" {
+		return fmt.Sprintf("%s%s;", strings.TrimSuffix(query, ";"), commentsString)
+	}
+	return fmt.Sprintf("%s%s", query, commentsString)
+}
+
+// ***** Commenter Functions *****
diff --git a/go/database/sql/go-sql.go b/go/database/sql/go-sql.go
@@ -17,97 +17,104 @@ package sql
 import (
 	"context"
 	"database/sql"
-	"fmt"
-	"strings"
+	"database/sql/driver"
 
 	"github.com/google/sqlcommenter/go/core"
 )
 
-type DB struct {
-	*sql.DB
+var (
+	_ driver.Driver        = (*sqlCommenterDriver)(nil)
+	_ driver.DriverContext = (*sqlCommenterDriver)(nil)
+	_ driver.Connector     = (*sqlCommenterConnector)(nil)
+)
+
+// SQLCommenterDriver returns a driver object that contains SQLCommenter drivers.
+type sqlCommenterDriver struct {
+	driver  driver.Driver
 	options core.CommenterOptions
 }
 
-func Open(driverName string, dataSourceName string, options core.CommenterOptions) (*DB, error) {
-	db, err := sql.Open(driverName, dataSourceName)
-	return &DB{DB: db, options: options}, err
+func newSQLCommenterDriver(dri driver.Driver, options core.CommenterOptions) *sqlCommenterDriver {
+	return &sqlCommenterDriver{driver: dri, options: options}
 }
 
-// ***** Query Functions *****
-
-func (db *DB) Query(query string, args ...any) (*sql.Rows, error) {
-	return db.DB.Query(db.withComment(context.Background(), query), args...)
+func (d *sqlCommenterDriver) Open(name string) (driver.Conn, error) {
+	rawConn, err := d.driver.Open(name)
+	if err != nil {
+		return nil, err
+	}
+	return newSQLCommenterConn(rawConn, d.options), nil
 }
 
-func (db *DB) QueryRow(query string, args ...interface{}) *sql.Row {
-	return db.DB.QueryRow(db.withComment(context.Background(), query), args...)
+func (d *sqlCommenterDriver) OpenConnector(name string) (driver.Connector, error) {
+	rawConnector, err := d.driver.(driver.DriverContext).OpenConnector(name)
+	if err != nil {
+		return nil, err
+	}
+	return newConnector(rawConnector, d, d.options), err
 }
 
-func (db *DB) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
-	return db.DB.QueryContext(ctx, db.withComment(ctx, query), args...)
+type sqlCommenterConnector struct {
+	driver.Connector
+	driver  *sqlCommenterDriver
+	options core.CommenterOptions
 }
 
-func (db *DB) Exec(query string, args ...any) (sql.Result, error) {
-	return db.DB.Exec(db.withComment(context.Background(), query), args...)
+func newConnector(connector driver.Connector, driver *sqlCommenterDriver, options core.CommenterOptions) *sqlCommenterConnector {
+	return &sqlCommenterConnector{
+		Connector: connector,
+		driver:    driver,
+		options:   options,
+	}
 }
 
-func (db *DB) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
-	return db.DB.ExecContext(ctx, db.withComment(ctx, query), args...)
+func (c *sqlCommenterConnector) Connect(ctx context.Context) (connection driver.Conn, err error) {
+	connection, err = c.Connector.Connect(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return newSQLCommenterConn(connection, c.options), nil
 }
 
-func (db *DB) Prepare(query string) (*sql.Stmt, error) {
-	return db.DB.Prepare(db.withComment(context.Background(), query))
+func (c *sqlCommenterConnector) Driver() driver.Driver {
+	return c.driver
 }
 
-func (db *DB) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) {
-	return db.DB.PrepareContext(ctx, db.withComment(ctx, query))
+type dsnConnector struct {
+	dsn    string
+	driver driver.Driver
 }
 
-// ***** Query Functions *****
-
-// ***** Commenter Functions *****
-
-func (db *DB) withComment(ctx context.Context, query string) string {
-	var commentsMap = map[string]string{}
-	query = strings.TrimSpace(query)
+func (t dsnConnector) Connect(_ context.Context) (driver.Conn, error) {
+	return t.driver.Open(t.dsn)
+}
 
-	// Sorted alphabetically
-	if db.options.EnableAction && (ctx.Value(core.Action) != nil) {
-		commentsMap[core.Action] = ctx.Value(core.Action).(string)
-	}
+func (t dsnConnector) Driver() driver.Driver {
+	return t.driver
+}
 
-	// `driver` information should not be coming from framework.
-	// So, explicitly adding that here.
-	if db.options.EnableDBDriver {
-		commentsMap[core.Driver] = "database/sql"
+// Open is a wrapper over sql.Open with OTel instrumentation.
+func Open(driverName, dataSourceName string, options core.CommenterOptions) (*sql.DB, error) {
+	// Retrieve the driver implementation we need to wrap with instrumentation
+	db, err := sql.Open(driverName, "")
+	if err != nil {
+		return nil, err
 	}
-
-	if db.options.EnableFramework && (ctx.Value(core.Framework) != nil) {
-		commentsMap[core.Framework] = ctx.Value(core.Framework).(string)
+	d := db.Driver()
+	if err = db.Close(); err != nil {
+		return nil, err
 	}
 
-	if db.options.EnableRoute && (ctx.Value(core.Route) != nil) {
-		commentsMap[core.Route] = ctx.Value(core.Route).(string)
-	}
+	options.Tags.DriverName = driverName
+	sqlCommenterDriver := newSQLCommenterDriver(d, options)
 
-	if db.options.EnableTraceparent {
-		carrier := core.ExtractTraceparent(ctx)
-		if val, ok := carrier["traceparent"]; ok {
-			commentsMap[core.Traceparent] = val
+	if _, ok := d.(driver.DriverContext); ok {
+		connector, err := sqlCommenterDriver.OpenConnector(dataSourceName)
+		if err != nil {
+			return nil, err
 		}
+		return sql.OpenDB(connector), nil
 	}
 
-	var commentsString string = ""
-	if len(commentsMap) > 0 { // Converts comments map to string and appends it to query
-		commentsString = fmt.Sprintf("/*%s*/", core.ConvertMapToComment(commentsMap))
-	}
-
-	// A semicolon at the end of the SQL statement means the query ends there.
-	// We need to insert the comment before that to be considered as part of the SQL statemtent.
-	if query[len(query)-1:] == ";" {
-		return fmt.Sprintf("%s%s;", strings.TrimSuffix(query, ";"), commentsString)
-	}
-	return fmt.Sprintf("%s%s", query, commentsString)
+	return sql.OpenDB(dsnConnector{dsn: dataSourceName, driver: sqlCommenterDriver}), nil
 }
-
-// ***** Commenter Functions *****