deprecate access to underlying *sql.DB via internal/db.DB() (#506)

internal/auth/db_test.go

@@ -1,6 +1,7 @@
 package auth
 
 import (
+	"context"
 	"testing"
 
 	"github.com/hashicorp/boundary/internal/db"
@@ -40,31 +41,45 @@ select count(*) from test_auth_method where public_id = $1;
 
 	assert, require := assert.New(t), require.New(t)
 
+	ctx := context.Background()
 	conn, _ := db.TestSetup(t, "postgres")
 	wrapper := db.TestWrapper(t)
-	db := conn.DB()
-	_, err := db.Exec(createTable)
+	rw := db.New(conn)
+	_, err := rw.Exec(ctx, createTable, nil)
 	require.NoError(err)
 
-	_, err = db.Exec(addTriggers)
+	_, err = rw.Exec(ctx, addTriggers, nil)
 	require.NoError(err)
 
 	org, _ := iam.TestScopes(t, iam.TestRepo(t, conn, wrapper))
 
 	id := "l1Ocw0TpHn800CekIxIXlmQqRDgFDfYl"
-	_, err = db.Query(insert, id, org.GetPublicId())
+	inserted, err := rw.Exec(ctx, insert, []interface{}{id, org.GetPublicId()})
 	require.NoError(err)
+	require.Equal(1, inserted)
 
-	var count int
-	err = db.QueryRow(baseTableQuery, id).Scan(&count)
+	type c struct {
+		Count int
+	}
+	var count c
+	rows, err := rw.Query(ctx, baseTableQuery, []interface{}{id})
 	require.NoError(err)
-	assert.Equal(1, count)
+	defer rows.Close()
+	for rows.Next() {
+		err := rw.ScanRows(rows, &count)
+		require.NoError(err)
+	}
+	assert.Equal(1, count.Count)
 
-	count = 0
-
-	err = db.QueryRow(testTableQuery, id).Scan(&count)
+	count.Count = 0
+	rows, err = rw.Query(ctx, testTableQuery, []interface{}{id})
 	require.NoError(err)
-	assert.Equal(1, count)
+	defer rows.Close()
+	for rows.Next() {
+		err := rw.ScanRows(rows, &count)
+		require.NoError(err)
+	}
+	assert.Equal(1, count.Count)
 }
 
 func TestDB_AuthAccountIDTrigger(t *testing.T) {
@@ -105,33 +120,46 @@ values
 
 	assert, require := assert.New(t), require.New(t)
 
+	ctx := context.Background()
 	conn, _ := db.TestSetup(t, "postgres")
 	wrapper := db.TestWrapper(t)
 	repo := iam.TestRepo(t, conn, wrapper)
-	db := conn.DB()
-	_, err := db.Exec(createTable)
+	rw := db.New(conn)
+	_, err := rw.Exec(ctx, createTable, nil)
 	require.NoError(err)
 
-	_, err = db.Exec(addTriggers)
+	_, err = rw.Exec(ctx, addTriggers, nil)
 	require.NoError(err)
 
 	org, _ := iam.TestScopes(t, repo)
 	meth_id := "31Ocw0TpHn800CekIxIXlmQqRDgFDfYl"
-	_, err = db.Query(insertAuthMethod, meth_id, org.GetPublicId())
+	_, err = rw.Exec(ctx, insertAuthMethod, []interface{}{meth_id, org.GetPublicId()})
 	require.NoError(err)
 
 	id := "l1Ocw0TpHn800CekIxIXlmQqRDgFDfYl"
-	_, err = db.Query(insert, id, meth_id)
+	_, err = rw.Exec(ctx, insert, []interface{}{id, meth_id})
 	require.NoError(err)
 
-	var count int
-	err = db.QueryRow(baseTableQuery, id).Scan(&count)
+	type c struct {
+		Count int
+	}
+	var count c
+	rows, err := rw.Query(ctx, baseTableQuery, []interface{}{id})
 	require.NoError(err)
-	assert.Equal(1, count)
+	defer rows.Close()
+	for rows.Next() {
+		err := rw.ScanRows(rows, &count)
+		require.NoError(err)
+	}
+	assert.Equal(1, count.Count)
 
-	count = 0
-
-	err = db.QueryRow(testTableQuery, id).Scan(&count)
+	count.Count = 0
+	rows, err = rw.Query(ctx, testTableQuery, []interface{}{id})
 	require.NoError(err)
-	assert.Equal(1, count)
+	defer rows.Close()
+	for rows.Next() {
+		err := rw.ScanRows(rows, &count)
+		require.NoError(err)
+	}
+	assert.Equal(1, count.Count)
 }

internal/auth/password/repository_account_test.go

@@ -908,7 +908,7 @@ func TestRepository_UpdateAccount(t *testing.T) {
 			assert.Equal(tt.wantCount, gotCount, "row count")
 			assert.NotSame(tt.orig, got)
 			assert.Equal(tt.orig.AuthMethodId, got.AuthMethodId)
-			dbassert := dbassert.New(t, rw)
+			dbassert := dbassert.New(t, conn.DB())
 			if tt.want.Name == "" {
 				dbassert.IsNull(got, "name")
 				return

internal/auth/password/repository_authmethod_test.go

@@ -695,7 +695,7 @@ func TestRepository_UpdateAuthMethod(t *testing.T) {
 			require.NoError(err)
 			assert.Empty(cmp.Diff(updatedAM, foundAuthMethod, protocmp.Transform()))
 
-			dbassert := dbassert.New(t, rw)
+			dbassert := dbassert.New(t, conn.DB())
 			if amToUpdate.Name == "" && contains(tt.args.fieldMaskPaths, "name") {
 				dbassert.IsNull(foundAuthMethod, "name")
 			}

internal/auth/password/repository_configuration.go

@@ -134,11 +134,8 @@ func (r *Repository) currentConfig(ctx context.Context, authMethodId string) (*c
 
 func (r *Repository) currentConfigForAccount(ctx context.Context, accountId string) (*currentConfig, error) {
 	var confs []currentConfig
-	tx, err := r.reader.DB()
-	if err != nil {
-		return nil, err
-	}
-	rows, err := tx.Query(currentConfigForAccountQuery, accountId)
+
+	rows, err := r.reader.Query(ctx, currentConfigForAccountQuery, []interface{}{accountId})
 	if err != nil {
 		return nil, err
 	}

internal/auth/password/repository_password.go

@@ -203,11 +203,7 @@ func (r *Repository) ChangePassword(ctx context.Context, scopeId, accountId, old
 func (r *Repository) authenticate(ctx context.Context, scopeId, authMethodId, loginName, password string) (*authAccount, error) {
 	var accts []authAccount
 
-	tx, err := r.reader.DB()
-	if err != nil {
-		return nil, err
-	}
-	rows, err := tx.Query(authenticateQuery, authMethodId, loginName)
+	rows, err := r.reader.Query(ctx, authenticateQuery, []interface{}{authMethodId, loginName})
 	if err != nil {
 		return nil, err
 	}

internal/db/README.md

@@ -24,10 +24,9 @@ Just some high-level usage highlights to get you started.  Read the godocs for a
     err = rw.LookupByPublicId(context.Background(), foundUser)
 
     // There's reader ScanRows that facilitates scanning rows from 
-    // a "raw" SQL query into your Gorm struct
-    tx, err := rw.DB()
+    // a query into your Gorm struct
     where := "select * from test_users where name in ($1, $2)"
-    rows, err := tx.Query(where, "alice", "bob")
+    rows, err := rw.Query(context.Background(), where, []interface{}{"alice", "bob"})
 	defer rows.Close()
 	for rows.Next() {
         user := db_test.NewTestUser()

internal/db/assert/asserts.go

@@ -1,9 +1,10 @@
 package dbassert
 
 import (
+	"database/sql"
+
 	dbassert "github.com/hashicorp/dbassert"
 
-	"github.com/hashicorp/boundary/internal/db"
 	gormAssert "github.com/hashicorp/dbassert/gorm"
 	"github.com/stretchr/testify/assert"
 )
@@ -14,13 +15,10 @@ type DbAsserts struct {
 }
 
 // New creates a new DbAsserts.
-func New(t dbassert.TestingT, r db.Reader) *DbAsserts {
+func New(t dbassert.TestingT, r *sql.DB) *DbAsserts {
 	assert.NotNil(t, r, "db.Reader is nill")
-
-	db, err := r.DB()
-	assert.NoError(t, err)
 	return &DbAsserts{
-		asserts: gormAssert.New(t, db, "postgres"),
+		asserts: gormAssert.New(t, r, "postgres"),
 	}
 }
 

internal/db/assert/asserts_test.go

@@ -6,7 +6,6 @@ import (
 	dbassert "github.com/hashicorp/dbassert"
 	gormAssert "github.com/hashicorp/dbassert/gorm"
 
-	"github.com/hashicorp/boundary/internal/db"
 	"github.com/jinzhu/gorm"
 	"github.com/stretchr/testify/assert"
 )
@@ -22,12 +21,8 @@ func Test_FieldDomain(t *testing.T) {
 			t.Error(err)
 		}
 	}()
-	assert := assert.New(t)
-	gormDb, err := gorm.Open("postgres", conn)
-	assert.NoError(err)
-	r := db.New(gormDb)
 	mockery := new(dbassert.MockTesting)
-	dbassert := New(mockery, r)
+	dbassert := New(mockery, conn)
 
 	dbassert.Domain(&gormAssert.TestModel{}, "PublicId", "dbasserts_public_id")
 	mockery.AssertNoError(t)
@@ -48,12 +43,8 @@ func Test_FieldNullable(t *testing.T) {
 			t.Error(err)
 		}
 	}()
-	assert := assert.New(t)
-	gormDb, err := gorm.Open("postgres", conn)
-	assert.NoError(err)
-	r := db.New(gormDb)
 	mockery := new(dbassert.MockTesting)
-	dbassert := New(mockery, r)
+	dbassert := New(mockery, conn)
 
 	dbassert.Nullable(&gormAssert.TestModel{}, "Nullable")
 	mockery.AssertNoError(t)
@@ -77,9 +68,8 @@ func Test_FieldIsNull(t *testing.T) {
 	assert := assert.New(t)
 	gormDb, err := gorm.Open("postgres", conn)
 	assert.NoError(err)
-	r := db.New(gormDb)
 	mockery := new(dbassert.MockTesting)
-	dbassert := New(mockery, r)
+	dbassert := New(mockery, conn)
 
 	v := 1
 	m := gormAssert.CreateTestModel(t, gormDb, nil, &v)
@@ -106,9 +96,8 @@ func Test_FieldNotNull(t *testing.T) {
 	assert := assert.New(t)
 	gormDb, err := gorm.Open("postgres", conn)
 	assert.NoError(err)
-	r := db.New(gormDb)
 	mockery := new(dbassert.MockTesting)
-	dbassert := New(mockery, r)
+	dbassert := New(mockery, conn)
 
 	v := 1
 	m := gormAssert.CreateTestModel(t, gormDb, nil, &v)

internal/db/read_writer.go

@@ -50,13 +50,10 @@ type Reader interface {
 	// operate within the context of any ongoing transaction for the db.Reader.  The
 	// caller must close the returned *sql.Rows. Query can/should be used in
 	// combination with ScanRows.
-	Query(sql string, values []interface{}, opt ...Option) (*sql.Rows, error)
+	Query(ctx context.Context, sql string, values []interface{}, opt ...Option) (*sql.Rows, error)
 
 	// ScanRows will scan sql rows into the interface provided
 	ScanRows(rows *sql.Rows, result interface{}) error
-
-	// DB returns the sql.DB
-	DB() (*sql.DB, error)
 }
 
 // Writer interface defines create, update and retryable transaction handlers
@@ -106,13 +103,10 @@ type Writer interface {
 	// should be to rollback. Delete returns the number of rows deleted or an error.
 	DeleteItems(ctx context.Context, deleteItems []interface{}, opt ...Option) (int, error)
 
-	// DB returns the sql.DB
-	DB() (*sql.DB, error)
-
 	// Exec will execute the sql with the values as parameters. The int returned
 	// is the number of rows affected by the sql. No options are currently
 	// supported.
-	Exec(sql string, values []interface{}, opt ...Option) (int, error)
+	Exec(ctx context.Context, sql string, values []interface{}, opt ...Option) (int, error)
 
 	// GetTicket returns an oplog ticket for the aggregate root of "i" which can
 	// be used to WriteOplogEntryWith for that aggregate root.
@@ -190,18 +184,10 @@ func New(underlying *gorm.DB) *Db {
 	return &Db{underlying: underlying}
 }
 
-// DB returns the sql.DB
-func (rw *Db) DB() (*sql.DB, error) {
-	if rw.underlying == nil {
-		return nil, fmt.Errorf("missing underlying db: %w", ErrInvalidParameter)
-	}
-	return rw.underlying.DB(), nil
-}
-
 // Exec will execute the sql with the values as parameters. The int returned
 // is the number of rows affected by the sql. No options are currently
 // supported.
-func (rw *Db) Exec(sql string, values []interface{}, opt ...Option) (int, error) {
+func (rw *Db) Exec(ctx context.Context, sql string, values []interface{}, opt ...Option) (int, error) {
 	if sql == "" {
 		return NoRowsAffected, fmt.Errorf("missing sql: %w", ErrInvalidParameter)
 	}
@@ -216,7 +202,7 @@ func (rw *Db) Exec(sql string, values []interface{}, opt ...Option) (int, error)
 // operate within the context of any ongoing transaction for the db.Reader.  The
 // caller must close the returned *sql.Rows. Query can/should be used in
 // combination with ScanRows.
-func (rw *Db) Query(sql string, values []interface{}, opt ...Option) (*sql.Rows, error) {
+func (rw *Db) Query(ctx context.Context, sql string, values []interface{}, opt ...Option) (*sql.Rows, error) {
 	if sql == "" {
 		return nil, fmt.Errorf("raw missing sql: %w", ErrInvalidParameter)
 	}

internal/db/read_writer_test.go

@@ -1069,28 +1069,6 @@ func TestDb_SearchWhere(t *testing.T) {
 	}
 }
 
-func TestDb_DB(t *testing.T) {
-	t.Parallel()
-	db, _ := TestSetup(t, "postgres")
-	t.Run("valid", func(t *testing.T) {
-		require := require.New(t)
-		w := Db{underlying: db}
-		d, err := w.DB()
-		require.NoError(err)
-		require.NotNil(d)
-		err = d.Ping()
-		require.NoError(err)
-	})
-	t.Run("nil-tx", func(t *testing.T) {
-		assert, require := assert.New(t), require.New(t)
-		w := Db{underlying: nil}
-		d, err := w.DB()
-		require.Error(err)
-		assert.Nil(d)
-		assert.Equal("missing underlying db: invalid parameter", err.Error())
-	})
-}
-
 func TestDb_Exec(t *testing.T) {
 	t.Parallel()
 	t.Run("update", func(t *testing.T) {
@@ -1104,7 +1082,7 @@ func TestDb_Exec(t *testing.T) {
 		err = w.Create(context.Background(), user)
 		require.NoError(err)
 		require.NotEmpty(user.Id)
-		rowsAffected, err := w.Exec("update db_test_user set name = ? where public_id = ?", []interface{}{"updated-" + id, user.PublicId})
+		rowsAffected, err := w.Exec(context.Background(), "update db_test_user set name = ? where public_id = ?", []interface{}{"updated-" + id, user.PublicId})
 		require.NoError(err)
 		require.Equal(1, rowsAffected)
 	})
@@ -1506,11 +1484,8 @@ func TestDb_ScanRows(t *testing.T) {
 		err = w.Create(context.Background(), user)
 		require.NoError(err)
 		assert.NotEmpty(user.Id)
-
-		tx, err := w.DB()
-		require.NoError(err)
 		where := "select * from db_test_user where name in ($1, $2)"
-		rows, err := tx.Query(where, "alice", "bob")
+		rows, err := w.Query(context.Background(), where, []interface{}{"alice", "bob"})
 		require.NoError(err)
 		defer func() { err := rows.Close(); assert.NoError(err) }()
 		for rows.Next() {
@@ -1540,7 +1515,7 @@ func TestDb_Query(t *testing.T) {
 		assert.Equal("alice", user.Name)
 
 		where := "select * from db_test_user where name in ($1, $2)"
-		rows, err := rw.Query(where, []interface{}{"alice", "bob"})
+		rows, err := rw.Query(context.Background(), where, []interface{}{"alice", "bob"})
 		require.NoError(err)
 		defer func() { err := rows.Close(); assert.NoError(err) }()
 		for rows.Next() {

internal/host/static/repository_host_catalog_test.go

@@ -461,7 +461,7 @@ func TestRepository_UpdateCatalog(t *testing.T) {
 			assert.Equal(tt.wantCount, gotCount, "row count")
 			assert.NotSame(tt.orig, got)
 			assert.Equal(tt.orig.ScopeId, got.ScopeId)
-			dbassert := dbassert.New(t, rw)
+			dbassert := dbassert.New(t, conn.DB())
 			if tt.want.Name == "" {
 				dbassert.IsNull(got, "name")
 				return

internal/host/static/repository_host_set_member.go

@@ -321,17 +321,12 @@ func (r *Repository) changes(ctx context.Context, setId string, hostIds []string
 	}
 	query := fmt.Sprintf(setChangesQuery, inClause)
 
-	tx, err := r.reader.DB()
-	if err != nil {
-		return nil, fmt.Errorf("changes: unable to get DB: %w", err)
-	}
-
 	var params []interface{}
 	params = append(params, setId)
 	for _, v := range hostIds {
 		params = append(params, v)
 	}
-	rows, err := tx.QueryContext(ctx, query, params...)
+	rows, err := r.reader.Query(ctx, query, params)
 	if err != nil {
 		return nil, fmt.Errorf("changes: query failed: %w", err)
 	}

internal/host/static/repository_host_set_test.go

@@ -493,7 +493,7 @@ func TestRepository_UpdateSet(t *testing.T) {
 			assert.Equal(tt.wantCount, gotCount, "row count")
 			assert.NotSame(tt.orig, got)
 			assert.Equal(tt.orig.CatalogId, got.CatalogId)
-			dbassert := dbassert.New(t, rw)
+			dbassert := dbassert.New(t, conn.DB())
 			if tt.want.Name == "" {
 				dbassert.IsNull(got, "name")
 				return

internal/host/static/repository_host_test.go

@@ -609,7 +609,7 @@ func TestRepository_UpdateHost(t *testing.T) {
 			assert.Equal(tt.wantCount, gotCount, "row count")
 			assert.NotSame(tt.orig, got)
 			assert.Equal(tt.orig.CatalogId, got.CatalogId)
-			dbassert := dbassert.New(t, rw)
+			dbassert := dbassert.New(t, conn.DB())
 			if tt.want.Name == "" {
 				dbassert.IsNull(got, "name")
 				return

internal/iam/account_test.go

@@ -79,12 +79,11 @@ func TestAccount_Clone(t *testing.T) {
 	})
 	t.Run("not-equal", func(t *testing.T) {
 		assert := assert.New(t)
-		rw := db.New(conn)
 		u := TestUser(t, repo, org.PublicId)
 		authMethodPublicId := testAuthMethod(t, conn, org.PublicId)
 		acct := testAccount(t, conn, org.PublicId, authMethodPublicId, u.PublicId)
 		acct2 := testAccount(t, conn, org.PublicId, authMethodPublicId, "")
-		dbassert := dbassert.New(t, rw)
+		dbassert := dbassert.New(t, conn.DB())
 		dbassert.IsNull(acct2, "IamUserId")
 		cp := acct.Clone()
 		assert.True(!proto.Equal(cp.(*authAccount).Account, acct2.Account))

internal/iam/group_test.go

@@ -353,7 +353,7 @@ func Test_GroupUpdate(t *testing.T) {
 			require.NoError(err)
 			assert.True(proto.Equal(updateGrp, foundGrp))
 			if len(tt.args.nullPaths) != 0 {
-				dbassert := dbassert.New(t, rw)
+				dbassert := dbassert.New(t, conn.DB())
 				for _, f := range tt.args.nullPaths {
 					dbassert.IsNull(&foundGrp, f)
 				}

internal/iam/repository_group.go

@@ -503,7 +503,7 @@ func groupMemberChanges(ctx context.Context, reader db.Reader, groupId string, u
 		params = append(params, v)
 	}
 	// fmt.Println(query, params)
-	rows, err := reader.Query(query, params)
+	rows, err := reader.Query(ctx, query, params)
 	if err != nil {
 		return nil, nil, fmt.Errorf("changes: query failed: %w", err)
 	}

internal/iam/repository_group_test.go

@@ -451,7 +451,7 @@ func TestRepository_UpdateGroup(t *testing.T) {
 			foundGrp, _, err := repo.LookupGroup(context.Background(), u.PublicId)
 			require.NoError(err)
 			assert.True(proto.Equal(groupAfterUpdate, foundGrp))
-			dbassert := dbassert.New(t, rw)
+			dbassert := dbassert.New(t, conn.DB())
 			if tt.args.name == "" {
 				dbassert.IsNull(foundGrp, "name")
 			}

internal/iam/repository_role_grant.go

@@ -448,11 +448,7 @@ select role_scope as scope_id, role_grant as grant from final;
 	}
 
 	var grants []perms.GrantPair
-	tx, err := r.reader.DB()
-	if err != nil {
-		return nil, err
-	}
-	rows, err := tx.Query(query, userId)
+	rows, err := r.reader.Query(ctx, query, []interface{}{userId})
 	if err != nil {
 		return nil, err
 	}

internal/iam/repository_role_test.go

@@ -448,7 +448,7 @@ func TestRepository_UpdateRole(t *testing.T) {
 			foundRole, _, _, err := repo.LookupRole(context.Background(), r.PublicId)
 			assert.NoError(err)
 			assert.True(proto.Equal(roleAfterUpdate, foundRole))
-			dbassert := dbassert.New(t, rw)
+			dbassert := dbassert.New(t, conn.DB())
 			if tt.args.name == "" {
 				assert.Equal(foundRole.Name, "")
 				dbassert.IsNull(foundRole, "name")

internal/iam/repository_user.go

@@ -304,11 +304,7 @@ func (r *Repository) getUserWithAccount(ctx context.Context, withAccountId strin
 	if withAccountId == "" {
 		return nil, fmt.Errorf("missing account id %w", db.ErrInvalidParameter)
 	}
-	underlying, err := r.reader.DB()
-	if err != nil {
-		return nil, fmt.Errorf("unable to get underlying db for account search: %w", err)
-	}
-	rows, err := underlying.Query(whereUserAccount, withAccountId)
+	rows, err := r.reader.Query(ctx, whereUserAccount, []interface{}{withAccountId})
 	if err != nil {
 		return nil, fmt.Errorf("unable to query account %s", withAccountId)
 	}
@@ -762,7 +758,7 @@ func associationChanges(ctx context.Context, reader db.Reader, userId string, ac
 	for _, v := range accountIds {
 		params = append(params, v)
 	}
-	rows, err := reader.Query(query, params)
+	rows, err := reader.Query(ctx, query, params)
 	if err != nil {
 		return nil, nil, fmt.Errorf("changes: query failed: %w", err)
 	}

internal/iam/repository_user_test.go

@@ -368,7 +368,7 @@ func TestRepository_UpdateUser(t *testing.T) {
 			sort.Strings(foundAccountIds)
 			assert.Equal(accountIds, foundAccountIds)
 
-			dbassert := dbassert.New(t, rw)
+			dbassert := dbassert.New(t, conn.DB())
 			if tt.args.name == "" {
 				dbassert.IsNull(foundUser, "name")
 			}
@@ -645,7 +645,7 @@ func TestRepository_LookupUserWithLogin(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			assert, require := assert.New(t), require.New(t)
-			dbassert := dbassert.New(t, rw)
+			dbassert := dbassert.New(t, conn.DB())
 			got, err := repo.LookupUserWithLogin(context.Background(), tt.args.withAccountId, tt.args.opt...)
 			if tt.wantErr {
 				require.Error(err)
@@ -963,7 +963,7 @@ func TestRepository_dissociateUserWithAccount(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			assert, require := assert.New(t), require.New(t)
-			dbassert := dbassert.New(t, rw)
+			dbassert := dbassert.New(t, conn.DB())
 			err := dissociateUserFromAccounts(context.Background(), kms, rw, rw, tt.args.Ids.user, tt.args.Ids.accts, tt.args.opt...)
 			if tt.wantErr {
 				require.Error(err)

internal/iam/role_test.go

@@ -442,7 +442,7 @@ func Test_RoleUpdate(t *testing.T) {
 			}
 			assert.True(proto.Equal(updateRole, foundRole))
 			if len(tt.args.nullPaths) != 0 {
-				dbassert := dbassert.New(t, rw)
+				dbassert := dbassert.New(t, conn.DB())
 				for _, f := range tt.args.nullPaths {
 					dbassert.IsNull(&foundRole, f)
 				}

internal/iam/testing.go

@@ -264,7 +264,7 @@ func testAccount(t *testing.T, conn *gorm.DB, scopeId, authMethodId, userId stri
 	require.NotEmpty(acct.PublicId)
 
 	if userId == "" {
-		dbassert := dbassert.New(t, rw)
+		dbassert := dbassert.New(t, conn.DB())
 		dbassert.IsNull(acct, "IamUserId")
 	}
 	return acct

internal/servers/repository.go

@@ -95,24 +95,17 @@ func (r *Repository) UpsertServer(ctx context.Context, server *Server, opt ...Op
 		address = $5,
 		update_time = $6;
 	`
-	underlying, err := r.writer.DB()
-	if err != nil {
-		return nil, db.NoRowsAffected, fmt.Errorf("error fetching underlying DB for upsert operation: %w", err)
-	}
-	result, err := underlying.ExecContext(ctx, q,
-		server.PrivateId,
-		server.Type,
-		server.Name,
-		server.Description,
-		server.Address,
-		time.Now().Format(time.RFC3339))
+
+	rowsAffected, err := r.writer.Exec(ctx, q,
+		[]interface{}{server.PrivateId,
+			server.Type,
+			server.Name,
+			server.Description,
+			server.Address,
+			time.Now().Format(time.RFC3339)})
 	if err != nil {
 		return nil, db.NoRowsAffected, fmt.Errorf("error performing status upsert: %w", err)
 	}
-	rowsAffected, err := result.RowsAffected()
-	if err != nil {
-		return nil, db.NoRowsAffected, fmt.Errorf("unable to fetch number of rows affected from query: %w", err)
-	}
 	// If updating a controller, done
 	if server.Type == resource.Controller.String() {
 		return nil, int(rowsAffected), nil

internal/session/repository_session.go

@@ -207,12 +207,7 @@ func (r *Repository) ListSessions(ctx context.Context, opt ...Option) ([]*Sessio
 	q := sessionList
 	query := fmt.Sprintf(q, limit, whereClause, opts.withOrder)
 
-	tx, err := r.reader.DB()
-	if err != nil {
-		return nil, fmt.Errorf("list sessions: unable to get DB: %w", err)
-	}
-
-	rows, err := tx.QueryContext(ctx, query, args...)
+	rows, err := r.reader.Query(ctx, query, args)
 	if err != nil {
 		return nil, fmt.Errorf("changes: query failed: %w", err)
 	}
@@ -310,7 +305,7 @@ func (r *Repository) TerminateSession(ctx context.Context, sessionId string, ses
 		db.StdRetryCnt,
 		db.ExpBackoff{},
 		func(reader db.Reader, w db.Writer) error {
-			rowsAffected, err := w.Exec(terminateSessionCte, []interface{}{sessionId, sessionVersion})
+			rowsAffected, err := w.Exec(ctx, terminateSessionCte, []interface{}{sessionId, sessionVersion})
 			if err != nil {
 				return fmt.Errorf("unable to terminate session %s: %w", sessionId, err)
 			}
@@ -362,7 +357,7 @@ func (r *Repository) AuthorizeConnection(ctx context.Context, sessionId string)
 		db.StdRetryCnt,
 		db.ExpBackoff{},
 		func(reader db.Reader, w db.Writer) error {
-			rowsAffected, err := w.Exec(authorizeConnectionCte, []interface{}{sessionId, connectionId})
+			rowsAffected, err := w.Exec(ctx, authorizeConnectionCte, []interface{}{sessionId, connectionId})
 			if err != nil {
 				return status.Errorf(codes.Internal, "unable to authorize connection %s: %v", sessionId, err)
 			}
@@ -396,11 +391,7 @@ type ConnectionAuthzSummary struct {
 }
 
 func (r *Repository) sessionAuthzSummary(ctx context.Context, sessionId string) (*ConnectionAuthzSummary, error) {
-	tx, err := r.reader.DB()
-	if err != nil {
-		return nil, fmt.Errorf("session summary: unable to get DB: %w", err)
-	}
-	rows, err := tx.QueryContext(ctx, remainingConnectionsCte, sessionId)
+	rows, err := r.reader.Query(ctx, remainingConnectionsCte, []interface{}{sessionId})
 	if err != nil {
 		return nil, fmt.Errorf("session summary: query failed: %w", err)
 	}
@@ -566,7 +557,7 @@ func (r *Repository) ActivateSession(ctx context.Context, sessionId string, sess
 		db.StdRetryCnt,
 		db.ExpBackoff{},
 		func(reader db.Reader, w db.Writer) error {
-			rowsAffected, err := w.Exec(activateStateCte, []interface{}{sessionId, sessionVersion})
+			rowsAffected, err := w.Exec(ctx, activateStateCte, []interface{}{sessionId, sessionVersion})
 			if err != nil {
 				return fmt.Errorf("unable to activate session %s: %w", sessionId, err)
 			}
@@ -662,7 +653,7 @@ func (r *Repository) updateState(ctx context.Context, sessionId string, sessionV
 				updatedSession.CtTofuToken = nil
 			}
 
-			rowsAffected, err = w.Exec(updateSessionState, []interface{}{sessionId, s.String()})
+			rowsAffected, err = w.Exec(ctx, updateSessionState, []interface{}{sessionId, s.String()})
 			if err != nil {
 				return fmt.Errorf("unable to update session %s state to %s: %w", sessionId, s.String(), err)
 			}

internal/target/repository_tcp_target_test.go

@@ -410,7 +410,7 @@ func TestRepository_UpdateTcpTarget(t *testing.T) {
 			foundTarget, _, err := repo.LookupTarget(context.Background(), target.PublicId)
 			assert.NoError(err)
 			assert.True(proto.Equal(targetAfterUpdate.((*TcpTarget)), foundTarget.((*TcpTarget))))
-			dbassert := dbassert.New(t, rw)
+			dbassert := dbassert.New(t, conn.DB())
 			if tt.args.description == "" {
 				assert.Equal(foundTarget.GetDescription(), "")
 				dbassert.IsNull(foundTarget, "description")

internal/target/tcp_target_test.go

@@ -279,7 +279,7 @@ func TestTcpTarget_Update(t *testing.T) {
 			require.NoError(err)
 			assert.True(proto.Equal(updateTarget, foundTarget))
 			if len(tt.args.nullPaths) != 0 {
-				dbassert := dbassert.New(t, rw)
+				dbassert := dbassert.New(t, conn.DB())
 				for _, f := range tt.args.nullPaths {
 					dbassert.IsNull(&foundTarget, f)
 				}