2 files changed, 89 insertions, 2 deletions
diff --git a/cmd/internal/cmdconfig/cmdconfig.go b/cmd/internal/cmdconfig/cmdconfig.go
index e9d915b6..9600f925 100644
--- a/cmd/internal/cmdconfig/cmdconfig.go
+++ b/cmd/internal/cmdconfig/cmdconfig.go
@@ -98,8 +98,8 @@ func ExperimentGetter(ctx context.Context, cfg *config.Config) middleware.Experi
 func OpenDB(ctx context.Context, cfg *config.Config, bypassLicenseCheck bool) (_ *postgres.DB, err error) {
 	defer derrors.Wrap(&err, "cmdconfig.OpenDB(ctx, cfg)")
 
-	// Wrap the postgres driver with OpenCensus instrumentation.
-	ocDriver, err := ocsql.Register(cfg.DBDriver, ocsql.WithAllTraceOptions())
+	// Wrap the postgres driver with our own wrapper, which adds OpenCensus instrumentation.
+	ocDriver, err := database.RegisterOCWrapper(cfg.DBDriver, ocsql.WithAllTraceOptions())
 	if err != nil {
 		return nil, fmt.Errorf("unable to register the ocsql driver: %v", err)
 	}
diff --git a/internal/database/driver.go b/internal/database/driver.go
new file mode 100644
index 00000000..44158423
--- /dev/null
+++ b/internal/database/driver.go
@@ -0,0 +1,87 @@
+// Copyright 2021 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package database
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"errors"
+
+	"contrib.go.opencensus.io/integrations/ocsql"
+)
+
+// RegisterOCWrapper registers a driver that wraps the OpenCensus driver, which in
+// turn wraps the driver named as the first argument.
+func RegisterOCWrapper(driverName string, opts ...ocsql.TraceOption) (string, error) {
+	// Get the driver to wrap.
+	db, err := sql.Open(driverName, "")
+	if err != nil {
+		return "", err
+	}
+	dri := db.Driver()
+	if err := db.Close(); err != nil {
+		return "", err
+	}
+	name := "ocWrapper-" + driverName
+	sql.Register(name, &wrapOCDriver{dri, opts})
+	return name, nil
+}
+
+type wrapOCDriver struct {
+	underlying driver.Driver
+	opts       []ocsql.TraceOption
+}
+
+// Open implements database/sql/driver.Driver.
+func (d *wrapOCDriver) Open(name string) (driver.Conn, error) {
+	conn, err := d.underlying.Open(name)
+	if err != nil {
+		return nil, err
+	}
+	oc := ocsql.WrapConn(conn, d.opts...)
+	return &wrapConn{conn, oc.(iconn)}, nil
+}
+
+type iconn interface {
+	driver.Pinger
+	driver.ExecerContext
+	driver.QueryerContext
+	driver.Conn
+	driver.ConnPrepareContext
+	driver.ConnBeginTx
+}
+
+// A wrapConn knows about both the underlying Conn, and the OpenCensus Conn that wraps it.
+// It delegates all calls to the OpenCensus Conn, but the underlying conn is available
+// to this package.
+type wrapConn struct {
+	underlying driver.Conn
+	oc         iconn
+}
+
+// Ping and all the following methods implement driver.Conn and related interfaces,
+// listed in the iconn interface above.
+func (c *wrapConn) Ping(ctx context.Context) error { return c.oc.Ping(ctx) }
+
+func (c *wrapConn) Prepare(query string) (driver.Stmt, error) { return c.oc.Prepare(query) }
+func (c *wrapConn) Close() error                              { return c.oc.Close() }
+func (c *wrapConn) Begin() (driver.Tx, error)                 { return nil, errors.New("unimplmented") }
+
+func (c *wrapConn) ExecContext(ctx context.Context, q string, args []driver.NamedValue) (driver.Result, error) {
+	return c.oc.ExecContext(ctx, q, args)
+}
+
+func (c *wrapConn) QueryContext(ctx context.Context, q string, args []driver.NamedValue) (driver.Rows, error) {
+	return c.oc.QueryContext(ctx, q, args)
+}
+
+func (c *wrapConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
+	return c.oc.PrepareContext(ctx, query)
+}
+
+func (c *wrapConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	return c.oc.BeginTx(ctx, opts)
+}