fix: use the db transaction reader and writer

ddebko · ddebko · commit 65aa0400f0e4 · 2025-02-07T10:59:47.000-08:00
diff --git a/internal/auth/oidc/repository_auth_method.go b/internal/auth/oidc/repository_auth_method.go
@@ -179,7 +179,7 @@ func (r *Repository) upsertAccount(ctx context.Context, am *AuthMethod, IdTokenC
 			var rowCnt int
 			for rows.Next() {
 				rowCnt += 1
-				err = r.reader.ScanRows(ctx, rows, &result)
+				err = reader.ScanRows(ctx, rows, &result)
 				if err != nil {
 					return errors.Wrap(ctx, err, op, errors.WithMsg("unable to scan rows for account"))
 				}
diff --git a/internal/auth/repository_auth_method.go b/internal/auth/repository_auth_method.go
@@ -147,7 +147,7 @@ func (amr *AuthMethodRepository) ListDeletedIds(ctx context.Context, since time.
 	var deletedAuthMethodIDs []string
 	var transactionTimestamp time.Time
 	if _, err := amr.writer.DoTx(ctx, db.StdRetryCnt, db.ExpBackoff{}, func(r db.Reader, w db.Writer) error {
-		rows, err := amr.writer.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})
+		rows, err := w.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})
 		if err != nil {
 			return errors.Wrap(ctx, err, op)
 		}
diff --git a/internal/credential/repository_store.go b/internal/credential/repository_store.go
@@ -118,7 +118,7 @@ func (s *StoreRepository) ListDeletedIds(ctx context.Context, since time.Time) (
 	var deletedStoreIDs []string
 	var transactionTimestamp time.Time
 	if _, err := s.writer.DoTx(ctx, db.StdRetryCnt, db.ExpBackoff{}, func(r db.Reader, w db.Writer) error {
-		rows, err := s.writer.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})
+		rows, err := w.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})
 		if err != nil {
 			return errors.Wrap(ctx, err, op)
 		}
diff --git a/internal/host/repository_catalog.go b/internal/host/repository_catalog.go
@@ -119,7 +119,7 @@ func (s *CatalogRepository) ListDeletedIds(ctx context.Context, since time.Time)
 	var deletedCatalogIDs []string
 	var transactionTimestamp time.Time
 	if _, err := s.writer.DoTx(ctx, db.StdRetryCnt, db.ExpBackoff{}, func(r db.Reader, w db.Writer) error {
-		rows, err := s.writer.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})
+		rows, err := w.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})
 		if err != nil {
 			return errors.Wrap(ctx, err, op)
 		}
diff --git a/internal/host/static/repository_host.go b/internal/host/static/repository_host.go
@@ -168,7 +168,7 @@ func (r *Repository) UpdateHost(ctx context.Context, projectId string, h *Host,
 	var rowsUpdated int
 	var returnedHost *Host
 	_, err = r.writer.DoTx(ctx, db.StdRetryCnt, db.ExpBackoff{},
-		func(_ db.Reader, w db.Writer) error {
+		func(r db.Reader, w db.Writer) error {
 			returnedHost = h.clone()
 			var err error
 			rowsUpdated, err = w.Update(ctx, returnedHost, dbMask, nullFields,
@@ -183,7 +183,7 @@ func (r *Repository) UpdateHost(ctx context.Context, projectId string, h *Host,
 			ha := &hostAgg{
 				PublicId: h.PublicId,
 			}
-			if err := r.reader.LookupByPublicId(ctx, ha); err != nil {
+			if err := r.LookupByPublicId(ctx, ha); err != nil {
 				return errors.Wrap(ctx, err, op, errors.WithMsg("failed to lookup host after update"))
 			}
 			returnedHost.SetIds = ha.getSetIds()
diff --git a/internal/iam/repository_role.go b/internal/iam/repository_role.go
@@ -326,7 +326,7 @@ func (r *Repository) queryRoles(ctx context.Context, whereClause string, args []
 			for _, retRole := range retRoles {
 				roleIds = append(roleIds, retRole.PublicId)
 			}
-			retRoleGrantScopes, err = r.ListRoleGrantScopes(ctx, roleIds)
+			retRoleGrantScopes, err = r.ListRoleGrantScopes(ctx, roleIds, WithReaderWriter(rd, w))
 			if err != nil {
 				return errors.Wrap(ctx, err, op, errors.WithMsg("failed to query role grant scopes"))
 			}
diff --git a/internal/iam/repository_role_grant.go b/internal/iam/repository_role_grant.go
@@ -359,7 +359,7 @@ func (r *Repository) SetRoleGrants(ctx context.Context, roleId string, roleVersi
 				return errors.Wrap(ctx, err, op, errors.WithMsg("unable to write oplog"))
 			}
 
-			currentRoleGrants, err = r.ListRoleGrants(ctx, roleId)
+			currentRoleGrants, err = r.ListRoleGrants(ctx, roleId, WithReaderWriter(reader, w))
 			if err != nil {
 				return errors.Wrap(ctx, err, op, errors.WithMsg("unable to retrieve current role grants after set"))
 			}
diff --git a/internal/server/options.go b/internal/server/options.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/hashicorp/boundary/internal/db"
 	"github.com/hashicorp/boundary/version"
 	"github.com/hashicorp/nodeenrollment/types"
 )
@@ -54,6 +55,8 @@ type options struct {
 	withFeature                            version.Feature
 	withDirectlyConnected                  bool
 	withWorkerPool                         []string
+	WithReader                             db.Reader
+	WithWriter                             db.Writer
 }
 
 func getDefaultOptions() options {
@@ -276,3 +279,12 @@ func WithLocalStorageState(state string) Option {
 		o.withLocalStorageState = state
 	}
 }
+
+// WithReaderWriter is used to share the same database reader
+// and writer when executing sql within a transaction.
+func WithReaderWriter(r db.Reader, w db.Writer) Option {
+	return func(o *options) {
+		o.WithReader = r
+		o.WithWriter = w
+	}
+}
diff --git a/internal/server/options_test.go b/internal/server/options_test.go
@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/hashicorp/boundary/internal/db"
 	"github.com/hashicorp/boundary/version"
 	"github.com/stretchr/testify/assert"
 )
@@ -251,4 +252,19 @@ func Test_GetOpts(t *testing.T) {
 		testOpts.withNewIdFunc = nil
 		assert.Equal(t, opts, testOpts)
 	})
+	t.Run("WithReaderWriter", func(t *testing.T) {
+		reader := &db.Db{}
+		writer := &db.Db{}
+		testOpts := getDefaultOptions()
+		assert.Nil(t, testOpts.WithReader)
+		assert.Nil(t, testOpts.WithWriter)
+		testOpts.WithReader = reader
+		testOpts.WithWriter = writer
+		opts := GetOpts(WithReaderWriter(reader, writer))
+		opts.withNewIdFunc = nil
+		testOpts.withNewIdFunc = nil
+		assert.Equal(t, reader, opts.WithReader)
+		assert.Equal(t, writer, opts.WithWriter)
+		assert.Equal(t, opts, testOpts)
+	})
 }
diff --git a/internal/server/repository_worker.go b/internal/server/repository_worker.go
@@ -569,7 +569,7 @@ func (r *Repository) UpdateWorker(ctx context.Context, worker *Worker, version u
 			if ret, err = wAgg.toWorker(ctx); err != nil {
 				return err
 			}
-			ret.RemoteStorageStates, err = r.ListWorkerStorageBucketCredentialState(ctx, ret.GetPublicId())
+			ret.RemoteStorageStates, err = r.ListWorkerStorageBucketCredentialState(ctx, ret.GetPublicId(), WithReaderWriter(reader, w))
 			if err != nil {
 				return err
 			}
diff --git a/internal/server/repository_worker_storage_bucket_credential_state.go b/internal/server/repository_worker_storage_bucket_credential_state.go
@@ -9,23 +9,29 @@ import (
 
 	"github.com/hashicorp/boundary/internal/db/timestamp"
 	"github.com/hashicorp/boundary/internal/errors"
+	"github.com/hashicorp/boundary/internal/util"
 	plgpb "github.com/hashicorp/boundary/sdk/pbs/plugin"
 )
 
 // ListWorkerStorageBucketCredentialState returns a list of storage bucket credential states for the given worker.
-func (r *Repository) ListWorkerStorageBucketCredentialState(ctx context.Context, workerId string) (map[string]*plgpb.StorageBucketCredentialState, error) {
+func (r *Repository) ListWorkerStorageBucketCredentialState(ctx context.Context, workerId string, opts ...Option) (map[string]*plgpb.StorageBucketCredentialState, error) {
 	const op = "server.(Repository).ListWorkerStorageBucketCredentialState"
 	if workerId == "" {
 		return nil, errors.New(ctx, errors.InvalidParameter, op, "empty worker id")
 	}
+	opt := GetOpts(opts...)
+	reader := r.reader
+	if !util.IsNil(opt.WithReader) {
+		reader = opt.WithReader
+	}
 	type remoteStorageState struct {
 		StorageBucketId string
 		PermissionType  string
 		State           string
 		CheckedAt       *timestamp.Timestamp
 		ErrorDetails    string
 	}
-	rows, err := r.reader.Query(ctx, getStorageBucketCredentialStatesByWorkerId, []any{sql.Named("worker_id", workerId)})
+	rows, err := reader.Query(ctx, getStorageBucketCredentialStatesByWorkerId, []any{sql.Named("worker_id", workerId)})
 	if err != nil && !errors.Match(errors.T(errors.RecordNotFound), err) {
 		return nil, errors.Wrap(ctx, err, op)
 	}
@@ -36,7 +42,7 @@ func (r *Repository) ListWorkerStorageBucketCredentialState(ctx context.Context,
 			return nil, errors.Wrap(ctx, err, op)
 		}
 		var row remoteStorageState
-		if err := r.reader.ScanRows(ctx, rows, &row); err != nil {
+		if err := reader.ScanRows(ctx, rows, &row); err != nil {
 			return nil, errors.Wrap(ctx, err, op, errors.WithMsg("failed to fetch remote storage state"))
 		}
 		s, ok := remoteStorageStates[row.StorageBucketId]

Original file line number	Diff line number	Diff line change
`@@ -179,7 +179,7 @@ func (r Repository) upsertAccount(ctx context.Context, am AuthMethod, IdTokenC`
`179`	`179`	`var rowCnt int`
`180`	`180`	`for rows.Next() {`
`181`	`181`	`rowCnt += 1`
`182`		`- err = r.reader.ScanRows(ctx, rows, &result)`
	`182`	`+ err = reader.ScanRows(ctx, rows, &result)`
`183`	`183`	`if err != nil {`
`184`	`184`	`return errors.Wrap(ctx, err, op, errors.WithMsg("unable to scan rows for account"))`
`185`	`185`	`}`
Original file line number	Diff line number	Diff line change
`@@ -147,7 +147,7 @@ func (amr *AuthMethodRepository) ListDeletedIds(ctx context.Context, since time.`
`147`	`147`	`var deletedAuthMethodIDs []string`
`148`	`148`	`var transactionTimestamp time.Time`
`149`	`149`	`if _, err := amr.writer.DoTx(ctx, db.StdRetryCnt, db.ExpBackoff{}, func(r db.Reader, w db.Writer) error {`
`150`		`- rows, err := amr.writer.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})`
	`150`	`+ rows, err := w.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})`
`151`	`151`	`if err != nil {`
`152`	`152`	`return errors.Wrap(ctx, err, op)`
`153`	`153`	`}`
Original file line number	Diff line number	Diff line change
`@@ -118,7 +118,7 @@ func (s *StoreRepository) ListDeletedIds(ctx context.Context, since time.Time) (`
`118`	`118`	`var deletedStoreIDs []string`
`119`	`119`	`var transactionTimestamp time.Time`
`120`	`120`	`if _, err := s.writer.DoTx(ctx, db.StdRetryCnt, db.ExpBackoff{}, func(r db.Reader, w db.Writer) error {`
`121`		`- rows, err := s.writer.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})`
	`121`	`+ rows, err := w.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})`
`122`	`122`	`if err != nil {`
`123`	`123`	`return errors.Wrap(ctx, err, op)`
`124`	`124`	`}`
Original file line number	Diff line number	Diff line change
`@@ -119,7 +119,7 @@ func (s *CatalogRepository) ListDeletedIds(ctx context.Context, since time.Time)`
`119`	`119`	`var deletedCatalogIDs []string`
`120`	`120`	`var transactionTimestamp time.Time`
`121`	`121`	`if _, err := s.writer.DoTx(ctx, db.StdRetryCnt, db.ExpBackoff{}, func(r db.Reader, w db.Writer) error {`
`122`		`- rows, err := s.writer.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})`
	`122`	`+ rows, err := w.Query(ctx, listDeletedIdsQuery, []any{sql.Named("since", since)})`
`123`	`123`	`if err != nil {`
`124`	`124`	`return errors.Wrap(ctx, err, op)`
`125`	`125`	`}`
Original file line number	Diff line number	Diff line change
`@@ -326,7 +326,7 @@ func (r *Repository) queryRoles(ctx context.Context, whereClause string, args []`
`326`	`326`	`for _, retRole := range retRoles {`
`327`	`327`	`roleIds = append(roleIds, retRole.PublicId)`
`328`	`328`	`}`
`329`		`- retRoleGrantScopes, err = r.ListRoleGrantScopes(ctx, roleIds)`
	`329`	`+ retRoleGrantScopes, err = r.ListRoleGrantScopes(ctx, roleIds, WithReaderWriter(rd, w))`
`330`	`330`	`if err != nil {`
`331`	`331`	`return errors.Wrap(ctx, err, op, errors.WithMsg("failed to query role grant scopes"))`
`332`	`332`	`}`
Original file line number	Diff line number	Diff line change
`@@ -359,7 +359,7 @@ func (r *Repository) SetRoleGrants(ctx context.Context, roleId string, roleVersi`
`359`	`359`	`return errors.Wrap(ctx, err, op, errors.WithMsg("unable to write oplog"))`
`360`	`360`	`}`
`361`	`361`
`362`		`- currentRoleGrants, err = r.ListRoleGrants(ctx, roleId)`
	`362`	`+ currentRoleGrants, err = r.ListRoleGrants(ctx, roleId, WithReaderWriter(reader, w))`
`363`	`363`	`if err != nil {`
`364`	`364`	`return errors.Wrap(ctx, err, op, errors.WithMsg("unable to retrieve current role grants after set"))`
`365`	`365`	`}`
Original file line number	Diff line number	Diff line change
`@@ -569,7 +569,7 @@ func (r Repository) UpdateWorker(ctx context.Context, worker Worker, version u`
`569`	`569`	`if ret, err = wAgg.toWorker(ctx); err != nil {`
`570`	`570`	`return err`
`571`	`571`	`}`
`572`		`- ret.RemoteStorageStates, err = r.ListWorkerStorageBucketCredentialState(ctx, ret.GetPublicId())`
	`572`	`+ ret.RemoteStorageStates, err = r.ListWorkerStorageBucketCredentialState(ctx, ret.GetPublicId(), WithReaderWriter(reader, w))`
`573`	`573`	`if err != nil {`
`574`	`574`	`return err`
`575`	`575`	`}`