sdk/go/arvadostest/db.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package arvadostest
   6
   7 import (
   8         "context"
   9
  10         "git.arvados.org/arvados.git/lib/ctrlctx"
  11         "git.arvados.org/arvados.git/sdk/go/arvados"
  12         "github.com/jmoiron/sqlx"
  13         // sqlx needs lib/pq to talk to PostgreSQL
  14         _ "github.com/lib/pq"
  15         "gopkg.in/check.v1"
  16 )
  17
  18 // DB returns a DB connection for the given cluster config.
  19 func DB(c *check.C, cluster *arvados.Cluster) *sqlx.DB {
  20         db, err := sqlx.Open("postgres", cluster.PostgreSQL.Connection.String())
  21         c.Assert(err, check.IsNil)
  22         return db
  23 }
  24
  25 // TransactionContext returns a context suitable for running a test
  26 // case in a new transaction, and a rollback func which the caller
  27 // should call after the test.
  28 func TransactionContext(c *check.C, db *sqlx.DB) (ctx context.Context, rollback func()) {
  29         tx, err := db.Beginx()
  30         c.Assert(err, check.IsNil)
  31         return ctrlctx.NewWithTransaction(context.Background(), tx), func() {
  32                 c.Check(tx.Rollback(), check.IsNil)
  33         }
  34 }