Upgrade to kit v0.3.0

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/accessreview/access_entry_service.go

@@ -54,7 +54,7 @@ func (s AccessEntryService) Get(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			return entry.LoadByID(ctx, conn, s.scope, entryID)
 		},
 	)
@@ -83,7 +83,7 @@ func (s AccessEntryService) RecordDecision(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := entry.LoadByID(ctx, conn, s.scope, req.EntryID); err != nil {
 				return fmt.Errorf("cannot load access entry: %w", err)
 			}
@@ -173,7 +173,7 @@ func (s AccessEntryService) RecordDecisions(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			// Track verified campaigns to avoid repeated loads within the
 			// same transaction.
 			verifiedCampaigns := make(map[gid.GID]bool)
@@ -259,7 +259,7 @@ func (s AccessEntryService) FlagEntry(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := entry.LoadByID(ctx, conn, s.scope, req.EntryID); err != nil {
 				return fmt.Errorf("cannot load access entry: %w", err)
 			}
@@ -304,7 +304,7 @@ func (s AccessEntryService) ListForCampaignID(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			return entries.LoadByCampaignID(ctx, conn, s.scope, campaignID, cursor, filter)
 		},
 	)
@@ -326,7 +326,7 @@ func (s AccessEntryService) ListForCampaignIDAndSourceID(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			return entries.LoadByCampaignIDAndSourceID(ctx, conn, s.scope, campaignID, sourceID, cursor, filter)
 		},
 	)
@@ -346,7 +346,7 @@ func (s AccessEntryService) CountForCampaignID(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) (err error) {
+		func(ctx context.Context, conn pg.Querier) (err error) {
 			entries := coredata.AccessEntries{}
 			count, err = entries.CountByCampaignID(ctx, conn, s.scope, campaignID, filter)
 			if err != nil {
@@ -372,7 +372,7 @@ func (s AccessEntryService) CountForCampaignIDAndSourceID(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) (err error) {
+		func(ctx context.Context, conn pg.Querier) (err error) {
 			entries := coredata.AccessEntries{}
 			count, err = entries.CountByCampaignIDAndSourceID(ctx, conn, s.scope, campaignID, sourceID, filter)
 			if err != nil {
@@ -396,7 +396,7 @@ func (s AccessEntryService) CountPendingForCampaignID(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) (err error) {
+		func(ctx context.Context, conn pg.Querier) (err error) {
 			entries := coredata.AccessEntries{}
 			count, err = entries.CountPendingByCampaignID(ctx, conn, s.scope, campaignID)
 			if err != nil {
@@ -420,7 +420,7 @@ func (s AccessEntryService) DecisionHistory(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			return histories.LoadByEntryID(ctx, conn, s.scope, entryID)
 		},
 	)
@@ -439,7 +439,7 @@ func (s AccessEntryService) Statistics(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			return stats.LoadByCampaignID(ctx, conn, s.scope, campaignID)
 		},
 	)
@@ -459,7 +459,7 @@ func (s AccessEntryService) StatisticsForSource(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			return stats.LoadByCampaignIDAndSourceID(ctx, conn, s.scope, campaignID, sourceID)
 		},
 	)

pkg/accessreview/access_source_service.go

@@ -114,7 +114,7 @@ func (s AccessSourceService) Create(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			// Validate connector exists if provided
 			if req.ConnectorID != nil {
 				connector := &coredata.Connector{}
@@ -145,7 +145,7 @@ func (s AccessSourceService) Get(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			return source.LoadByID(ctx, conn, s.scope, accessSourceID)
 		},
 	)
@@ -168,7 +168,7 @@ func (s AccessSourceService) Update(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := source.LoadByID(ctx, conn, s.scope, req.AccessSourceID); err != nil {
 				return fmt.Errorf("cannot load access source: %w", err)
 			}
@@ -219,7 +219,7 @@ func (s AccessSourceService) Delete(
 
 	return s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			return source.Delete(ctx, conn, s.scope)
 		},
 	)
@@ -234,7 +234,7 @@ func (s AccessSourceService) ListForOrganizationID(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			return sources.LoadByOrganizationID(ctx, conn, s.scope, organizationID, cursor)
 		},
 	)
@@ -253,7 +253,7 @@ func (s AccessSourceService) CountForOrganizationID(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) (err error) {
+		func(ctx context.Context, conn pg.Querier) (err error) {
 			sources := coredata.AccessSources{}
 			count, err = sources.CountByOrganizationID(ctx, conn, s.scope, organizationID)
 			return err
@@ -274,7 +274,7 @@ func (s AccessSourceService) ListScopeSourcesForCampaignID(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			return sources.LoadScopeSourcesByCampaignID(ctx, conn, s.scope, campaignID)
 		},
 	)
@@ -296,7 +296,7 @@ func (s AccessSourceService) ConnectorHTTPClient(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			if err := dbConnector.LoadByID(ctx, conn, s.scope, connectorID, s.encryptionKey); err != nil {
 				return fmt.Errorf("cannot load connector: %w", err)
 			}
@@ -336,10 +336,10 @@ func (s AccessSourceService) ConnectorHTTPClient(
 	// Persist refreshed token if it changed.
 	if isOAuth2 && oauth2Conn.AccessToken != tokenBefore {
 		dbConnector.UpdatedAt = time.Now()
-		if err := s.pg.WithConn(
+		if err := s.pg.WithTx(
 			ctx,
-			func(conn pg.Conn) error {
-				return dbConnector.Update(ctx, conn, s.scope, s.encryptionKey)
+			func(ctx context.Context, tx pg.Tx) error {
+				return dbConnector.Update(ctx, tx, s.scope, s.encryptionKey)
 			},
 		); err != nil {
 			return nil, nil, fmt.Errorf("cannot persist refreshed token: %w", err)
@@ -361,7 +361,7 @@ func (s AccessSourceService) ConfigureAccessSource(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := source.LoadByID(ctx, conn, s.scope, req.AccessSourceID); err != nil {
 				return fmt.Errorf("cannot load access source: %w", err)
 			}

pkg/accessreview/campaign_service.go

@@ -59,7 +59,7 @@ func (s *CampaignService) Create(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := campaign.Insert(ctx, conn, s.scope); err != nil {
 				return fmt.Errorf("cannot insert access review campaign: %w", err)
 			}
@@ -101,7 +101,7 @@ func (s *CampaignService) Get(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			if err := campaign.LoadByID(ctx, conn, s.scope, campaignID); err != nil {
 				return fmt.Errorf("cannot load campaign: %w", err)
 			}
@@ -127,7 +127,7 @@ func (s *CampaignService) Update(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := lockCampaignForUpdate(ctx, conn, s.scope, req.CampaignID); err != nil {
 				return fmt.Errorf("cannot lock campaign: %w", err)
 			}
@@ -174,7 +174,7 @@ func (s *CampaignService) Delete(
 ) error {
 	return s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := lockCampaignForUpdate(ctx, conn, s.scope, campaignID); err != nil {
 				return fmt.Errorf("cannot lock campaign: %w", err)
 			}
@@ -205,7 +205,7 @@ func (s *CampaignService) AddScopeSource(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := lockCampaignForUpdate(ctx, conn, s.scope, req.CampaignID); err != nil {
 				return fmt.Errorf("cannot lock campaign: %w", err)
 			}
@@ -253,7 +253,7 @@ func (s *CampaignService) RemoveScopeSource(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := lockCampaignForUpdate(ctx, conn, s.scope, req.CampaignID); err != nil {
 				return fmt.Errorf("cannot lock campaign: %w", err)
 			}
@@ -292,7 +292,7 @@ func (s *CampaignService) Start(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := lockCampaignForUpdate(ctx, conn, s.scope, campaignID); err != nil {
 				return fmt.Errorf("cannot lock campaign: %w", err)
 			}
@@ -346,7 +346,7 @@ func (s *CampaignService) Close(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := lockCampaignForUpdate(ctx, conn, s.scope, campaignID); err != nil {
 				return fmt.Errorf("cannot lock campaign: %w", err)
 			}
@@ -388,9 +388,9 @@ func (s *CampaignService) Close(
 	return campaign, nil
 }
 
-func lockCampaignForUpdate(ctx context.Context, conn pg.Conn, scope coredata.Scoper, campaignID gid.GID) error {
+func lockCampaignForUpdate(ctx context.Context, tx pg.Tx, scope coredata.Scoper, campaignID gid.GID) error {
 	c := &coredata.AccessReviewCampaign{ID: campaignID}
-	if err := c.LockForUpdate(ctx, conn, scope); err != nil {
+	if err := c.LockForUpdate(ctx, tx, scope); err != nil {
 		return fmt.Errorf("cannot lock campaign for update: %w", err)
 	}
 	return nil
@@ -398,7 +398,7 @@ func lockCampaignForUpdate(ctx context.Context, conn pg.Conn, scope coredata.Sco
 
 func (s *CampaignService) enqueueSourceFetches(
 	ctx context.Context,
-	conn pg.Conn,
+	tx pg.Tx,
 	campaignID gid.GID,
 	sources coredata.AccessSources,
 ) error {
@@ -408,7 +408,7 @@ func (s *CampaignService) enqueueSourceFetches(
 			AccessReviewCampaignID: campaignID,
 			AccessSourceID:         source.ID,
 		}
-		if err := fetch.UpsertQueued(ctx, conn, s.scope, now); err != nil {
+		if err := fetch.UpsertQueued(ctx, tx, s.scope, now); err != nil {
 			return fmt.Errorf("cannot queue source fetch %s: %w", source.ID, err)
 		}
 	}
@@ -424,7 +424,7 @@ func (s *CampaignService) Cancel(
 
 	err := s.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			if err := lockCampaignForUpdate(ctx, conn, s.scope, campaignID); err != nil {
 				return fmt.Errorf("cannot lock campaign: %w", err)
 			}
@@ -466,7 +466,7 @@ func (s *CampaignService) ListForOrganizationID(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			if err := campaigns.LoadByOrganizationID(ctx, conn, s.scope, organizationID, cursor); err != nil {
 				return fmt.Errorf("cannot load campaigns by organization: %w", err)
 			}
@@ -488,7 +488,7 @@ func (s *CampaignService) ListSourceFetches(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			if err := fetches.LoadByCampaignID(ctx, conn, s.scope, campaignID); err != nil {
 				return fmt.Errorf("cannot load source fetches by campaign: %w", err)
 			}
@@ -510,7 +510,7 @@ func (s *CampaignService) CountForOrganizationID(
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) (err error) {
+		func(ctx context.Context, conn pg.Querier) (err error) {
 			campaigns := coredata.AccessReviewCampaigns{}
 			count, err = campaigns.CountByOrganizationID(ctx, conn, s.scope, organizationID)
 			if err != nil {

pkg/accessreview/drivers/probo_memberships.go

@@ -49,7 +49,7 @@ func (d *ProboMembershipsDriver) ListAccounts(ctx context.Context) ([]AccountRec
 
 	err := d.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			accounts, err := coredata.LoadMembershipAccountsByOrganizationID(
 				ctx,
 				conn,

pkg/accessreview/review_engine.go

@@ -73,11 +73,11 @@ func (e *ReviewEngine) FetchSource(
 		baseline []coredata.BaselineAccountEntry
 	)
 
-	err := e.pg.WithConn(
+	err := e.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, tx pg.Tx) error {
 			source = &coredata.AccessSource{}
-			if err := source.LoadByID(ctx, conn, e.scope, sourceID); err != nil {
+			if err := source.LoadByID(ctx, tx, e.scope, sourceID); err != nil {
 				return fmt.Errorf("cannot load access source %s: %w", sourceID, err)
 			}
 			if source.OrganizationID != campaign.OrganizationID {
@@ -85,19 +85,19 @@ func (e *ReviewEngine) FetchSource(
 			}
 
 			var err error
-			driver, err = e.resolveDriver(ctx, conn, source)
+			driver, err = e.resolveDriver(ctx, tx, source)
 			if err != nil {
 				return fmt.Errorf("cannot resolve driver for source %s: %w", source.Name, err)
 			}
 
 			lastCompletedCampaign := &coredata.AccessReviewCampaign{}
-			if err := lastCompletedCampaign.LoadLastCompletedByOrganizationID(ctx, conn, e.scope, campaign.OrganizationID); err != nil {
+			if err := lastCompletedCampaign.LoadLastCompletedByOrganizationID(ctx, tx, e.scope, campaign.OrganizationID); err != nil {
 				if !errors.Is(err, coredata.ErrResourceNotFound) {
 					return fmt.Errorf("cannot load last completed campaign: %w", err)
 				}
 			} else {
 				entries := &coredata.AccessEntries{}
-				baseline, err = entries.LoadBaselineBySourceID(ctx, conn, e.scope, lastCompletedCampaign.ID, sourceID)
+				baseline, err = entries.LoadBaselineBySourceID(ctx, tx, e.scope, lastCompletedCampaign.ID, sourceID)
 				if err != nil {
 					return fmt.Errorf("cannot load baseline entries by source: %w", err)
 				}
@@ -125,7 +125,7 @@ func (e *ReviewEngine) FetchSource(
 
 	err = e.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Tx) error {
 			now := time.Now()
 			seenAccountKeys := make(map[string]struct{}, len(accounts))
 
@@ -252,7 +252,7 @@ func (e *ReviewEngine) connectorHTTPClient(
 // connector_id (null = built-in, set = connector-backed).
 func (e *ReviewEngine) resolveDriver(
 	ctx context.Context,
-	conn pg.Conn,
+	tx pg.Tx,
 	source *coredata.AccessSource,
 ) (drivers.Driver, error) {
 	if source.ConnectorID == nil {
@@ -267,7 +267,7 @@ func (e *ReviewEngine) resolveDriver(
 
 	// Connector-backed: look up the connector and resolve driver by provider
 	dbConnector := &coredata.Connector{}
-	if err := dbConnector.LoadByID(ctx, conn, e.scope, *source.ConnectorID, e.encryptionKey); err != nil {
+	if err := dbConnector.LoadByID(ctx, tx, e.scope, *source.ConnectorID, e.encryptionKey); err != nil {
 		return nil, fmt.Errorf("cannot load connector %s: %w", *source.ConnectorID, err)
 	}
 
@@ -291,7 +291,7 @@ func (e *ReviewEngine) resolveDriver(
 	if oauth2Conn, ok := dbConnector.Connection.(*connector.OAuth2Connection); ok {
 		if oauth2Conn.AccessToken != tokenBefore {
 			dbConnector.UpdatedAt = time.Now()
-			if err := dbConnector.Update(ctx, conn, e.scope, e.encryptionKey); err != nil {
+			if err := dbConnector.Update(ctx, tx, e.scope, e.encryptionKey); err != nil {
 				return nil, fmt.Errorf("cannot persist refreshed token for connector %s: %w", *source.ConnectorID, err)
 			}
 		}

pkg/accessreview/service.go

@@ -116,7 +116,7 @@ func (s *Service) ResolveEntryOrganizationID(ctx context.Context, entryID gid.GI
 
 	err := s.pg.WithConn(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, conn pg.Querier) error {
 			var err error
 			entry := &coredata.AccessEntry{}
 			organizationID, err = entry.LoadOrganizationID(ctx, conn, entryID)

pkg/accessreview/source_name_worker.go

@@ -83,7 +83,7 @@ func (w *SourceNameWorker) processNext(ctx context.Context) error {
 
 	err := w.pg.WithTx(
 		ctx,
-		func(tx pg.Conn) error {
+		func(ctx context.Context, tx pg.Tx) error {
 			return source.LoadNextUnsyncedNameForUpdateSkipLocked(ctx, tx)
 		},
 	)
@@ -101,15 +101,15 @@ func (w *SourceNameWorker) processNext(ctx context.Context) error {
 		resolver    drivers.NameResolver
 	)
 
-	err = w.pg.WithConn(
+	err = w.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
+		func(ctx context.Context, tx pg.Tx) error {
 			scope := coredata.NewScopeFromObjectID(source.ID)
 			if source.ConnectorID == nil {
 				return fmt.Errorf("source %s has no connector", source.ID)
 			}
 
-			if err := dbConnector.LoadByID(ctx, conn, scope, *source.ConnectorID, w.encryptionKey); err != nil {
+			if err := dbConnector.LoadByID(ctx, tx, scope, *source.ConnectorID, w.encryptionKey); err != nil {
 				return fmt.Errorf("cannot load connector %s: %w", *source.ConnectorID, err)
 			}
 
@@ -126,7 +126,7 @@ func (w *SourceNameWorker) processNext(ctx context.Context) error {
 			if oauth2Conn, ok := dbConnector.Connection.(*connector.OAuth2Connection); ok {
 				if oauth2Conn.AccessToken != tokenBefore {
 					dbConnector.UpdatedAt = time.Now()
-					if err := dbConnector.Update(ctx, conn, scope, w.encryptionKey); err != nil {
+					if err := dbConnector.Update(ctx, tx, scope, w.encryptionKey); err != nil {
 						return fmt.Errorf("cannot persist refreshed token for connector %s: %w", *source.ConnectorID, err)
 					}
 				}
@@ -192,7 +192,7 @@ func (w *SourceNameWorker) markNameSynced(
 ) error {
 	return w.pg.WithTx(
 		ctx,
-		func(tx pg.Conn) error {
+		func(ctx context.Context, tx pg.Tx) error {
 			scope := coredata.NewScopeFromObjectID(source.ID)
 			now := time.Now()
 

pkg/accessreview/worker.go

@@ -127,7 +127,7 @@ func (w *SourceFetchWorker) processNext(
 
 	if err := w.pg.WithTx(
 		nonCancelableCtx,
-		func(tx pg.Conn) error {
+		func(ctx context.Context, tx pg.Tx) error {
 			if err := sourceFetch.LoadNextQueuedForUpdateSkipLocked(nonCancelableCtx, tx); err != nil {
 				return err // sentinel errors checked by caller
 			}
@@ -212,7 +212,7 @@ func (w *SourceFetchWorker) recoverStaleRows(ctx context.Context) {
 
 	err := w.pg.WithTx(
 		ctx,
-		func(tx pg.Conn) error {
+		func(ctx context.Context, tx pg.Tx) error {
 			var fetches coredata.AccessReviewCampaignSourceFetches
 			count, err := fetches.RecoverStale(ctx, tx, staleThreshold, now)
 			if err != nil {
@@ -251,10 +251,10 @@ func (w *SourceFetchWorker) commitFailedSourceFetch(
 	sourceFetch.CompletedAt = new(now)
 	sourceFetch.UpdatedAt = now
 
-	return w.pg.WithConn(
+	return w.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
-			return sourceFetch.Update(ctx, conn, scope)
+		func(ctx context.Context, tx pg.Tx) error {
+			return sourceFetch.Update(ctx, tx, scope)
 		},
 	)
 }
@@ -275,10 +275,10 @@ func (w *SourceFetchWorker) commitSuccessfulSourceFetch(
 	sourceFetch.CompletedAt = new(now)
 	sourceFetch.UpdatedAt = now
 
-	return w.pg.WithConn(
+	return w.pg.WithTx(
 		ctx,
-		func(conn pg.Conn) error {
-			return sourceFetch.Update(ctx, conn, scope)
+		func(ctx context.Context, tx pg.Tx) error {
+			return sourceFetch.Update(ctx, tx, scope)
 		},
 	)
 }
@@ -292,7 +292,7 @@ func (w *SourceFetchWorker) finalizeCampaignFetchLifecycle(
 
 	return w.pg.WithTx(
 		ctx,
-		func(tx pg.Conn) error {
+		func(ctx context.Context, tx pg.Tx) error {
 			if err := lockCampaignForUpdate(ctx, tx, scope, campaignID); err != nil {
 				return fmt.Errorf("cannot lock campaign: %w", err)
 			}