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         _ "github.com/lib/pq"
  14         "gopkg.in/check.v1"
  15 )
  16
  17 // DB returns a DB connection for the given cluster config.
  18 func DB(c *check.C, cluster *arvados.Cluster) *sqlx.DB {
  19         db, err := sqlx.Open("postgres", cluster.PostgreSQL.Connection.String())
  20         c.Assert(err, check.IsNil)
  21         return db
  22 }
  23
  24 // TransactionContext returns a context suitable for running a test
  25 // case in a new transaction, and a rollback func which the caller
  26 // should call after the test.
  27 func TransactionContext(c *check.C, db *sqlx.DB) (ctx context.Context, rollback func()) {
  28         tx, err := db.Beginx()
  29         c.Assert(err, check.IsNil)
  30         return ctrlctx.NewWithTransaction(context.Background(), tx), func() {
  31                 c.Check(tx.Rollback(), check.IsNil)
  32         }
  33 }