21123: Fix inadvertent use of run-tests db in dispatcher tests.

[arvados.git] / lib / ctrlctx / db.go

diff --git a/lib/ctrlctx/db.go b/lib/ctrlctx/db.go
index a76420860604b9a6fb9823bdc6b3775c70f85ff4..d33fd8ab53cfe1525a65329d1a5f0cfa2b9c7214 100644 (file)
--- a/lib/ctrlctx/db.go
+++ b/lib/ctrlctx/db.go
@@ -10,6 +10,7 @@ import (
        "sync"
 
        "git.arvados.org/arvados.git/lib/controller/api"
+       "git.arvados.org/arvados.git/sdk/go/arvados"
        "git.arvados.org/arvados.git/sdk/go/ctxlog"
        "github.com/jmoiron/sqlx"
 
@@ -68,7 +69,7 @@ type finishFunc func(*error)
 // commit or rollback the transaction, if any.
 //
 //     func example(ctx context.Context) (err error) {
-//             ctx, finishtx := New(ctx, dber)
+//             ctx, finishtx := New(ctx, getdb)
 //             defer finishtx(&err)
 //             // ...
 //             tx, err := CurrentTx(ctx)
@@ -142,3 +143,45 @@ func CurrentTx(ctx context.Context) (*sqlx.Tx, error) {
        })
        return txn.tx, txn.err
 }
+
+var errDBConnection = errors.New("database connection error")
+
+type DBConnector struct {
+       PostgreSQL arvados.PostgreSQL
+       pgdb       *sqlx.DB
+       mtx        sync.Mutex
+}
+
+func (dbc *DBConnector) GetDB(ctx context.Context) (*sqlx.DB, error) {
+       dbc.mtx.Lock()
+       defer dbc.mtx.Unlock()
+       if dbc.pgdb != nil {
+               return dbc.pgdb, nil
+       }
+       db, err := sqlx.Open("postgres", dbc.PostgreSQL.Connection.String())
+       if err != nil {
+               ctxlog.FromContext(ctx).WithError(err).Error("postgresql connect failed")
+               return nil, errDBConnection
+       }
+       if p := dbc.PostgreSQL.ConnectionPool; p > 0 {
+               db.SetMaxOpenConns(p)
+       }
+       if err := db.Ping(); err != nil {
+               ctxlog.FromContext(ctx).WithError(err).Error("postgresql connect succeeded but ping failed")
+               db.Close()
+               return nil, errDBConnection
+       }
+       dbc.pgdb = db
+       return db, nil
+}
+
+func (dbc *DBConnector) Close() error {
+       dbc.mtx.Lock()
+       defer dbc.mtx.Unlock()
+       var err error
+       if dbc.pgdb != nil {
+               err = dbc.pgdb.Close()
+               dbc.pgdb = nil
+       }
+       return err
+}