Remove the callback library causing race issue

go.mod

@@ -7,6 +7,7 @@ require (
 	cloud.google.com/go/firestore v1.3.0 // indirect
 	contrib.go.opencensus.io/exporter/ocagent v0.7.1-0.20200615190824-f8c219d2d895 // indirect
 	contrib.go.opencensus.io/exporter/prometheus v0.2.1-0.20200609204449-6bcf6f8577f0 // indirect
+	contrib.go.opencensus.io/integrations/ocsql v0.1.6
 	firebase.google.com/go v3.13.0+incompatible
 	github.com/Microsoft/go-winio v0.4.15-0.20190919025122-fc70bd9a86b5 // indirect
 	github.com/asaskevich/govalidator v0.0.0-20200108200545-475eaeb16496 // indirect
@@ -23,6 +24,7 @@ require (
 	github.com/jinzhu/gorm v1.9.16
 	github.com/jinzhu/now v1.1.1 // indirect
 	github.com/kelseyhightower/envconfig v1.4.0 // indirect
+	github.com/lib/pq v1.8.0
 	github.com/mattn/go-sqlite3 v2.0.1+incompatible // indirect
 	github.com/mikehelmick/go-chaff v0.3.0
 	github.com/opencensus-integrations/redigo v2.0.1+incompatible

go.sum

@@ -59,6 +59,8 @@ contrib.go.opencensus.io/exporter/prometheus v0.2.1-0.20200609204449-6bcf6f8577f
 contrib.go.opencensus.io/exporter/prometheus v0.2.1-0.20200609204449-6bcf6f8577f0/go.mod h1:MjHoxkI7Ny27toPeFkRbXbzVjzIGkwOAptrAy8Mxtm8=
 contrib.go.opencensus.io/exporter/stackdriver v0.13.4 h1:ksUxwH3OD5sxkjzEqGxNTl+Xjsmu3BnC/300MhSVTSc=
 contrib.go.opencensus.io/exporter/stackdriver v0.13.4/go.mod h1:aXENhDJ1Y4lIg4EUaVTwzvYETVNZk10Pu26tevFKLUc=
+contrib.go.opencensus.io/integrations/ocsql v0.1.6 h1:9qmZJBlnMtffShflmfhW4EZK7M+CujIDG4bEwUrg+ms=
+contrib.go.opencensus.io/integrations/ocsql v0.1.6/go.mod h1:8DsSdjz3F+APR+0z0WkU1aRorQCFfRxvqjUUPMbF3fE=
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
 firebase.google.com/go v3.13.0+incompatible h1:3TdYC3DDi6aHn20qoRkxwGqNgdjtblwVAyRLQwGn/+4=
 firebase.google.com/go v3.13.0+incompatible/go.mod h1:xlah6XbEyW6tbfSklcfe5FHJIwjt8toICdV5Wh9ptHs=

pkg/database/database.go

@@ -17,6 +17,7 @@ package database
 
 import (
 	"context"
+	"database/sql"
 	"encoding/base64"
 	"errors"
 	"fmt"
@@ -32,7 +33,8 @@ import (
 	"go.uber.org/zap"
 
 	// ensure the postgres dialiect is compiled in.
-	_ "github.com/jinzhu/gorm/dialects/postgres"
+	"contrib.go.opencensus.io/integrations/ocsql"
+	postgres "github.com/lib/pq"
 )
 
 // Database is a handle to the database layer for the Exposure Notifications
@@ -53,6 +55,20 @@ type Database struct {
 
 	// secretManager is used to resolve secrets.
 	secretManager secrets.SecretManager
+
+	statsCloser func()
+}
+
+// Overrides the postgresql driver with
+func init() {
+	const driverName = "ocsql"
+	for _, v := range sql.Drivers() {
+		if v == driverName {
+			return
+		}
+	}
+	ocsql.RegisterAllViews()
+	sql.Register(driverName, ocsql.Wrap(&postgres.Driver{}))
 }
 
 // SupportsPerRealmSigning returns true if the configuration supports
@@ -129,19 +145,37 @@ func (db *Database) OpenWithCacher(ctx context.Context, cacher cache.Cacher) err
 
 	// Establish a connection to the database. We use this later to register
 	// opencenusus stats.
+	var rawSQL *sql.DB
+	if err := withRetries(ctx, func(ctx context.Context) error {
+		var err error
+		rawSQL, err = sql.Open("ocsql", c.ConnectionString())
+		if err != nil {
+			return retry.RetryableError(err)
+		}
+		db.statsCloser = ocsql.RecordStats(rawSQL, 5*time.Second)
+		return nil
+	}); err != nil {
+		return fmt.Errorf("failed to create sql connection: %w", err)
+	}
+	if rawSQL == nil {
+		return fmt.Errorf("failed to create database connection")
+	}
+
 	var rawDB *gorm.DB
 	if err := withRetries(ctx, func(ctx context.Context) error {
 		var err error
-		rawDB, err = gorm.Open("postgres", c.ConnectionString())
+		// Need to give postgres dialect as otherwise gorm starts running
+		// in compatibility mode
+		rawDB, err = gorm.Open("postgres", rawSQL)
 		if err != nil {
 			return retry.RetryableError(err)
 		}
 		return nil
 	}); err != nil {
-		return fmt.Errorf("failed to connect to database: %w", err)
+		return fmt.Errorf("failed to initialize gorm: %w", err)
 	}
 	if rawDB == nil {
-		return fmt.Errorf("failed to create database connection")
+		return fmt.Errorf("failed to initialize gorm")
 	}
 
 	// Set connection configuration.
@@ -194,6 +228,7 @@ func (db *Database) OpenWithCacher(ctx context.Context, cacher cache.Cacher) err
 
 // Close will close the database connection. Should be deferred right after Open.
 func (db *Database) Close() error {
+	db.statsCloser()
 	return db.db.Close()
 }