diff --git a/go.mod b/go.mod index 51dacabfb..c2921d853 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 16d16881f..3db6a16da 100644 --- a/go.sum +++ b/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= diff --git a/pkg/database/database.go b/pkg/database/database.go index bda0c7a13..a2e3c695c 100644 --- a/pkg/database/database.go +++ b/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() }