implement additional driver.Rows interfaces

Author: hugo.rut

driver.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"reflect"
 	"strconv"
 	"sync"
 
@@ -249,7 +250,7 @@ func (c ocConn) Query(query string, args []driver.Value) (rows driver.Rows, err
 			return nil, err
 		}
 
-		return ocRows{parent: rows, ctx: ctx, options: c.options}, nil
+		return wrapRows(rows, ctx, c.options), nil
 	}
 
 	return nil, driver.ErrSkip
@@ -287,7 +288,7 @@ func (c ocConn) QueryContext(ctx context.Context, query string, args []driver.Na
 			return nil, err
 		}
 
-		return ocRows{parent: rows, ctx: ctx, options: c.options}, nil
+		return wrapRows(rows, ctx, c.options), nil
 	}
 
 	return nil, driver.ErrSkip
@@ -529,7 +530,7 @@ func (s ocStmt) Query(args []driver.Value) (rows driver.Rows, err error) {
 	if err != nil {
 		return nil, err
 	}
-	rows, err = ocRows{parent: rows, ctx: ctx, options: s.options}, nil
+	rows, err = wrapRows(rows, ctx, s.options), nil
 	return
 }
 
@@ -603,17 +604,92 @@ func (s ocStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (row
 	if err != nil {
 		return nil, err
 	}
-	rows, err = ocRows{parent: rows, ctx: ctx, options: s.options}, nil
+	rows, err = wrapRows(rows, ctx, s.options), nil
 	return
 }
 
+
+// RowsColumnTypeScanType is a duplicate interface for driver.RowsColumnTypeScanType but
+// with the driver.Rows composition removed.
+//
+// This is used to embed a anonymous struct without running into ambiguous method errors.
+type RowsColumnTypeScanType interface {
+	ColumnTypeScanType(index int) reflect.Type
+}
+
 // ocRows implements driver.Rows.
 type ocRows struct {
 	parent  driver.Rows
 	ctx     context.Context
 	options TraceOptions
 }
 
+// HasNextResultSet calls the implements the driver.RowsNextResultSet for ocRows.
+// It returns the the underlying result of HasNextResultSet from the ocRows.parent
+// if the parent implements driver.RowsNextResultSet.
+func (r ocRows) HasNextResultSet() bool {
+	if v, ok := r.parent.(driver.RowsNextResultSet); ok {
+		return v.HasNextResultSet()
+	}
+
+	return false
+}
+
+// NextResultsSet calls the implements the driver.RowsNextResultSet for ocRows.
+// It returns the the underlying result of NextResultSet from the ocRows.parent
+// if the parent implements driver.RowsNextResultSet.
+func (r ocRows) NextResultSet() error {
+	if v, ok := r.parent.(driver.RowsNextResultSet); ok {
+		return v.NextResultSet()
+	}
+
+	return io.EOF
+}
+
+// ColumnTypeDatabaseTypeName calls the implements the driver.RowsColumnTypeDatabaseTypeName for ocRows.
+// It returns the the underlying result of ColumnTypeDatabaseTypeName from the ocRows.parent
+// if the parent implements driver.RowsColumnTypeDatabaseTypeName.
+func (r ocRows) ColumnTypeDatabaseTypeName(index int) string {
+	if v, ok := r.parent.(driver.RowsColumnTypeDatabaseTypeName); ok {
+		return v.ColumnTypeDatabaseTypeName(index)
+	}
+
+	return ""
+}
+
+// ColumnTypeLength calls the implements the driver.RowsColumnTypeLength for ocRows.
+// It returns the the underlying result of ColumnTypeLength from the ocRows.parent
+// if the parent implements driver.RowsColumnTypeLength.
+func (r ocRows) ColumnTypeLength(index int) (length int64, ok bool) {
+	if v, ok := r.parent.(driver.RowsColumnTypeLength); ok {
+		return v.ColumnTypeLength(index)
+	}
+
+	return 0, false
+}
+
+// ColumnTypeNullable calls the implements the driver.RowsColumnTypeNullable for ocRows.
+// It returns the the underlying result of ColumnTypeNullable from the ocRows.parent
+// if the parent implements driver.RowsColumnTypeNullable.
+func (r ocRows) ColumnTypeNullable(index int) (nullable, ok bool) {
+	if v, ok := r.parent.(driver.RowsColumnTypeNullable); ok {
+		return v.ColumnTypeNullable(index)
+	}
+
+	return false, false
+}
+
+// ColumnTypePrecisionScale calls the implements the driver.RowsColumnTypePrecisionScale for ocRows.
+// It returns the the underlying result of ColumnTypePrecisionScale from the ocRows.parent
+// if the parent implements driver.RowsColumnTypePrecisionScale.
+func (r ocRows) ColumnTypePrecisionScale(index int) (precision, scale int64, ok bool) {
+	if v, ok := r.parent.(driver.RowsColumnTypePrecisionScale); ok {
+		return v.ColumnTypePrecisionScale(index)
+	}
+
+	return 0, 0, false
+}
+
 func (r ocRows) Columns() []string {
 	return r.parent.Columns()
 }
@@ -655,6 +731,30 @@ func (r ocRows) Next(dest []driver.Value) (err error) {
 	return
 }
 
+// wrapRows returns a struct which conforms to the driver.Rows interface.
+// It checks if the parent adheres to any additional driver interfaces and returns a matching
+// implementation accordingly.
+func wrapRows(parent driver.Rows, ctx context.Context, options TraceOptions) driver.Rows {
+	var (
+		ts, hasColumnTypeScan = parent.(driver.RowsColumnTypeScanType)
+	)
+
+	r := ocRows{
+		parent:  parent,
+		ctx:     ctx,
+		options: options,
+	}
+
+	if hasColumnTypeScan {
+		return struct {
+			driver.Rows
+			RowsColumnTypeScanType
+		}{r, ts}
+	}
+
+	return r
+}
+
 // ocTx implemens driver.Tx
 type ocTx struct {
 	parent  driver.Tx

driver_test.go

@@ -0,0 +1,119 @@
+package ocsql_test
+
+import (
+	"database/sql"
+	"database/sql/driver"
+	"reflect"
+	"testing"
+
+	"github.com/opencensus-integrations/ocsql"
+)
+
+type stubRows struct{}
+
+func (stubRows) Columns() []string              { return []string{} }
+func (stubRows) Close() error                   { return nil }
+func (stubRows) Next(dest []driver.Value) error { return nil }
+
+type stubScanType struct {
+	toReturn reflect.Type
+}
+
+func (s stubScanType) ColumnTypeScanType(index int) reflect.Type { return s.toReturn }
+
+type stubDriver struct {
+	rows driver.Rows
+}
+
+func (d stubDriver) Open(name string) (driver.Conn, error) {
+	return stubConnection{rows: d.rows}, nil
+}
+
+type stubConnection struct {
+	rows driver.Rows
+}
+
+func (c stubConnection) Prepare(query string) (driver.Stmt, error) {
+	return stubStmt{rows: c.rows}, nil
+}
+
+func (stubConnection) Close() error              { return nil }
+func (stubConnection) Begin() (driver.Tx, error) { return &sql.Tx{}, nil }
+
+type stubStmt struct {
+	rows driver.Rows
+}
+
+func (stubStmt) Close() error                                    { return nil }
+func (stubStmt) NumInput() int                                   { return 0 }
+func (stubStmt) Exec(args []driver.Value) (driver.Result, error) { return stubResult{}, nil }
+
+func (s stubStmt) Query(args []driver.Value) (driver.Rows, error) {
+	return s.rows, nil
+}
+
+type stubResult struct{}
+
+func (stubResult) LastInsertId() (int64, error) { return 0, nil }
+func (stubResult) RowsAffected() (int64, error) { return 0, nil }
+
+type testFunc func(t *testing.T, rows driver.Rows)
+
+var testNotAssignableToScanTypeInterface testFunc = func(t *testing.T, rows driver.Rows) {
+	if _, ok := rows.(driver.RowsColumnTypeScanType); ok {
+		t.Error("expected output to not be assignable to type: RowsColumnTypeLength")
+	}
+}
+
+var testAssignableToScanTypeInterface testFunc = func(t *testing.T, rows driver.Rows) {
+	if _, ok := rows.(driver.RowsColumnTypeScanType); !ok {
+		t.Error("expected output to be assignable to type: RowsColumnTypeLength")
+	}
+}
+
+func TestRowsAreWrappedWithCorrectInterfaceType(t *testing.T) {
+	type test struct {
+		name      string
+		input     driver.Rows
+		testFunc testFunc
+	}
+
+	tests := []test{
+		{
+			input: stubRows{},
+			name:  "test non scan type parent is not wrapped with scan type interface",
+			testFunc: testNotAssignableToScanTypeInterface,
+		},
+		{
+			input: struct {
+				driver.Rows
+				ocsql.RowsColumnTypeScanType
+			}{ stubRows{}, stubScanType{}},
+			name:  "test wraps rows with scan type interface",
+			testFunc: testAssignableToScanTypeInterface,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var d = ocsql.Wrap(stubDriver{
+				rows: tt.input,
+			}, ocsql.WithAllTraceOptions())
+			var c, _ = d.Open("fake-connection")
+
+			s, err := c.Prepare("SELECT * FROM test;")
+			if err != nil {
+				t.Errorf("connection.Prepare returned unexpected err: %v", err)
+				return
+			}
+
+			rows, err := s.Query([]driver.Value{})
+			if err != nil {
+				t.Errorf("stmt.Query returned unexpected err: %v", err)
+				return
+			}
+
+			tt.testFunc(t, rows)
+		})
+	}
+}