chore: respect transaction in error index reporter

rudderlabs · Oct 19, 2023 · b90288b · b90288b
1 parent 4db35e9
commit b90288b
Show file tree

Hide file tree

Showing 27 changed files with 107 additions and 65 deletions.
diff --git a/enterprise/reporting/error_index_reporting.go b/enterprise/reporting/error_index_reporting.go
@@ -2,7 +2,6 @@ package reporting
 
 import (
  "context"
- "database/sql"
  "encoding/json"
  "fmt"
  "time"
@@ -13,6 +12,7 @@ import (
  "github.com/rudderlabs/rudder-go-kit/logger"
  "github.com/rudderlabs/rudder-server/jobsdb"
  "github.com/rudderlabs/rudder-server/utils/misc"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
  "github.com/rudderlabs/rudder-server/utils/types"
 )
 
@@ -79,7 +79,7 @@ func NewErrorIndexReporter(
 }
 
 // Report reports the metrics to the errorIndex JobsDB
-func (eir *ErrorIndexReporter) Report(metrics []*types.PUReportedMetric, _ *sql.Tx) error {
+func (eir *ErrorIndexReporter) Report(metrics []*types.PUReportedMetric, tx *Tx) error {
  failedAt := eir.now()
 
  var jobs []*jobsdb.JobT
@@ -134,7 +134,9 @@ func (eir *ErrorIndexReporter) Report(metrics []*types.PUReportedMetric, _ *sql.
  return nil
  }
 
- if err := eir.errIndexDB.Store(eir.ctx, jobs); err != nil {
+ if err := eir.errIndexDB.WithStoreSafeTxFromTx(eir.ctx, tx, func(tx jobsdb.StoreSafeTx) error {
+ return eir.errIndexDB.StoreInTx(eir.ctx, tx, jobs)
+ }); err != nil {
  return fmt.Errorf("failed to store jobs: %v", err)
  }
 

diff --git a/enterprise/reporting/error_index_reporting_test.go b/enterprise/reporting/error_index_reporting_test.go
@@ -19,6 +19,7 @@ import (
  "github.com/rudderlabs/rudder-server/jobsdb"
  mocksBackendConfig "github.com/rudderlabs/rudder-server/mocks/backend-config"
  "github.com/rudderlabs/rudder-server/utils/pubsub"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
  "github.com/rudderlabs/rudder-server/utils/types"
 )
 
@@ -263,10 +264,12 @@ func TestErrorIndexReporter(t *testing.T) {
  defer func() {
  eir.errIndexDB.TearDown()
  }()
-
- err = eir.Report(tc.reports, nil)
+ sqltx, err := postgresContainer.DB.Begin()
  require.NoError(t, err)
-
+ tx := &Tx{Tx: sqltx}
+ err = eir.Report(tc.reports, tx)
+ require.NoError(t, err)
+ require.NoError(t, tx.Commit())
  jr, err := eir.errIndexDB.GetUnprocessed(ctx, jobsdb.GetQueryParams{
  JobsLimit: 100,
  })
@@ -328,9 +331,12 @@ func TestErrorIndexReporter(t *testing.T) {
 
  eir := NewErrorIndexReporter(ctx, c, logger.NOP, cs)
  defer eir.errIndexDB.TearDown()
-
- err = eir.Report([]*types.PUReportedMetric{}, nil)
+ sqltx, err := postgresContainer.DB.Begin()
+ require.NoError(t, err)
+ tx := &Tx{Tx: sqltx}
+ err = eir.Report([]*types.PUReportedMetric{}, tx)
  require.NoError(t, err)
+ require.NoError(t, tx.Commit())
 
  syncerDone := make(chan struct{})
  go func() {

diff --git a/enterprise/reporting/error_reporting.go b/enterprise/reporting/error_reporting.go
@@ -25,6 +25,7 @@ import (
  migrator "github.com/rudderlabs/rudder-server/services/sql-migrator"
  "github.com/rudderlabs/rudder-server/utils/httputil"
  "github.com/rudderlabs/rudder-server/utils/misc"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
  "github.com/rudderlabs/rudder-server/utils/types"
 )
 
@@ -149,7 +150,7 @@ func (edr *ErrorDetailReporter) GetSyncer(syncerKey string) *types.SyncSource {
  return edr.syncers[syncerKey]
 }
 
-func (edr *ErrorDetailReporter) Report(metrics []*types.PUReportedMetric, txn *sql.Tx) error {
+func (edr *ErrorDetailReporter) Report(metrics []*types.PUReportedMetric, txn *Tx) error {
  edr.log.Debug("[ErrorDetailReport] Report method called\n")
  if len(metrics) == 0 {
  return nil

diff --git a/enterprise/reporting/mediator.go b/enterprise/reporting/mediator.go
@@ -2,13 +2,13 @@ package reporting
 
 import (
  "context"
- "database/sql"
 
  "golang.org/x/sync/errgroup"
 
  "github.com/rudderlabs/rudder-go-kit/config"
  "github.com/rudderlabs/rudder-go-kit/logger"
  backendconfig "github.com/rudderlabs/rudder-server/backend-config"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
  "github.com/rudderlabs/rudder-server/utils/types"
 )
 
@@ -56,7 +56,7 @@ func NewReportingMediator(ctx context.Context, log logger.Logger, enterpriseToke
  return rm
 }
 
-func (rm *Mediator) Report(metrics []*types.PUReportedMetric, txn *sql.Tx) error {
+func (rm *Mediator) Report(metrics []*types.PUReportedMetric, txn *Tx) error {
  for _, reporter := range rm.reporters {
  if err := reporter.Report(metrics, txn); err != nil {
  return err

diff --git a/enterprise/reporting/noop.go b/enterprise/reporting/noop.go
@@ -1,15 +1,14 @@
 package reporting
 
 import (
- "database/sql"
-
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
  "github.com/rudderlabs/rudder-server/utils/types"
 )
 
 // NOOP reporting implementation that does nothing
 type NOOP struct{}
 
-func (*NOOP) Report(_ []*types.PUReportedMetric, _ *sql.Tx) error {
+func (*NOOP) Report(_ []*types.PUReportedMetric, _ *Tx) error {
  return nil
 }
 

diff --git a/enterprise/reporting/reporting.go b/enterprise/reporting/reporting.go
@@ -26,6 +26,7 @@ import (
  migrator "github.com/rudderlabs/rudder-server/services/sql-migrator"
  "github.com/rudderlabs/rudder-server/utils/httputil"
  "github.com/rudderlabs/rudder-server/utils/misc"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
  "github.com/rudderlabs/rudder-server/utils/types"
 )
 
@@ -512,7 +513,7 @@ func transformMetricForPII(metric types.PUReportedMetric, piiColumns []string) t
  return metric
 }
 
-func (r *DefaultReporter) Report(metrics []*types.PUReportedMetric, txn *sql.Tx) error {
+func (r *DefaultReporter) Report(metrics []*types.PUReportedMetric, txn *Tx) error {
  if len(metrics) == 0 {
  return nil
  }

diff --git a/jobsdb/backup.go b/jobsdb/backup.go
@@ -19,6 +19,7 @@ import (
  "github.com/rudderlabs/rudder-go-kit/stats"
  "github.com/rudderlabs/rudder-server/services/fileuploader"
  "github.com/rudderlabs/rudder-server/utils/misc"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
 )
 
 func (jd *Handle) isBackupEnabled() bool {

diff --git a/jobsdb/integration_test.go b/jobsdb/integration_test.go
@@ -16,6 +16,7 @@ import (
  "github.com/rudderlabs/rudder-go-kit/testhelper/rand"
  "github.com/rudderlabs/rudder-server/jobsdb/prebackup"
  fileuploader "github.com/rudderlabs/rudder-server/services/fileuploader"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
 )
 
 const (

diff --git a/jobsdb/jobsdb.go b/jobsdb/jobsdb.go
@@ -46,6 +46,7 @@ import (
  "github.com/rudderlabs/rudder-server/jobsdb/internal/cache"
  "github.com/rudderlabs/rudder-server/jobsdb/internal/lock"
  "github.com/rudderlabs/rudder-server/jobsdb/prebackup"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
 
  "github.com/rudderlabs/rudder-go-kit/config"
  "github.com/rudderlabs/rudder-go-kit/stats"
@@ -101,29 +102,6 @@ type GetQueryParams struct {
  PayloadSizeLimit int64
 }
 
-// Tx is a wrapper around sql.Tx that supports registering and executing
-// post-commit actions, a.k.a. success listeners.
-type Tx struct {
- *sql.Tx
- successListeners []func()
-}
-
-// AddSuccessListener registers a listener to be executed after the transaction has been committed successfully.
-func (tx *Tx) AddSuccessListener(listener func()) {
- tx.successListeners = append(tx.successListeners, listener)
-}
-
-// Commit commits the transaction and executes all listeners.
-func (tx *Tx) Commit() error {
- err := tx.Tx.Commit()
- if err == nil {
- for _, successListener := range tx.successListeners {
- successListener()
- }
- }
- return err
-}
-
 // StoreSafeTx sealed interface
 type StoreSafeTx interface {
  Tx() *Tx
@@ -237,6 +215,9 @@ type JobsDB interface {
  // that can be used by the provided function.
  WithStoreSafeTx(context.Context, func(tx StoreSafeTx) error) error
 
+ // WithStoreSafeTxFromTx prepares a store-safe environment for an existing transaction.
+ WithStoreSafeTxFromTx(ctx context.Context, tx *Tx, f func(tx StoreSafeTx) error) error
+
  // Store stores the provided jobs to the database
  Store(ctx context.Context, jobList []*JobT) error
 
@@ -1740,6 +1721,12 @@ func (jd *Handle) WithStoreSafeTx(ctx context.Context, f func(tx StoreSafeTx) er
  })
 }
 
+func (jd *Handle) WithStoreSafeTxFromTx(ctx context.Context, tx *Tx, f func(tx StoreSafeTx) error) error {
+ return jd.inStoreSafeCtx(ctx, func() error {
+ return f(&storeSafeTx{tx: tx, identity: jd.tablePrefix})
+ })
+}
+
 func (jd *Handle) inStoreSafeCtx(ctx context.Context, f func() error) error {
  // Only locks the list
  op := func() error {

diff --git a/jobsdb/jobsdb_backup_test.go b/jobsdb/jobsdb_backup_test.go
@@ -28,6 +28,7 @@ import (
  "github.com/rudderlabs/rudder-server/services/fileuploader"
  "github.com/rudderlabs/rudder-server/testhelper"
  "github.com/rudderlabs/rudder-server/testhelper/destination"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
 )
 
 func TestBackupTable(t *testing.T) {

diff --git a/jobsdb/jobsdb_renameDs_test.go b/jobsdb/jobsdb_renameDs_test.go
@@ -10,6 +10,7 @@ import (
 
  rsRand "github.com/rudderlabs/rudder-go-kit/testhelper/rand"
  "github.com/rudderlabs/rudder-server/jobsdb/prebackup"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
 )
 
 func Test_mustRenameDS(t *testing.T) {

diff --git a/jobsdb/jobsdb_test.go b/jobsdb/jobsdb_test.go
@@ -33,6 +33,7 @@ import (
  "github.com/rudderlabs/rudder-server/jobsdb/prebackup"
  fileuploader "github.com/rudderlabs/rudder-server/services/fileuploader"
  "github.com/rudderlabs/rudder-server/utils/misc"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
 )
 
 var _ = Describe("Calculate newDSIdx for internal migrations", Ordered, func() {

diff --git a/jobsdb/migration.go b/jobsdb/migration.go
@@ -16,6 +16,7 @@ import (
  "github.com/rudderlabs/rudder-server/jobsdb/internal/dsindex"
  "github.com/rudderlabs/rudder-server/jobsdb/internal/lock"
  "github.com/rudderlabs/rudder-server/utils/misc"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
 )
 
 // startMigrateDSLoop migrates jobs from src dataset (srcDS) to destination dataset (dest_ds)

diff --git a/mocks/jobsdb/mock_jobsdb.go b/mocks/jobsdb/mock_jobsdb.go
diff --git a/mocks/utils/types/mock_types.go b/mocks/utils/types/mock_types.go
diff --git a/processor/processor.go b/processor/processor.go
@@ -47,6 +47,7 @@ import (
  "github.com/rudderlabs/rudder-server/services/transientsource"
  "github.com/rudderlabs/rudder-server/utils/httputil"
  "github.com/rudderlabs/rudder-server/utils/misc"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
  "github.com/rudderlabs/rudder-server/utils/types"
  "github.com/rudderlabs/rudder-server/utils/workerpool"
 )
@@ -2131,7 +2132,7 @@ func (proc *Handle) Store(partition string, in *storeMessage) {
  }
 
  if proc.isReportingEnabled() {
- if err = proc.reporting.Report(in.reportMetrics, tx.SqlTx()); err != nil {
+ if err = proc.reporting.Report(in.reportMetrics, tx.Tx()); err != nil {
  return fmt.Errorf("reporting metrics: %w", err)
  }
  }
@@ -2573,7 +2574,7 @@ func (proc *Handle) transformSrcDest(
  }
 }
 
-func (proc *Handle) saveDroppedJobs(failedJobs []*jobsdb.JobT, tx *jobsdb.Tx) error {
+func (proc *Handle) saveDroppedJobs(failedJobs []*jobsdb.JobT, tx *Tx) error {
  if len(failedJobs) > 0 {
  rsourcesStats := rsources.NewDroppedJobsCollector(proc.rsourcesService)
  rsourcesStats.JobsDropped(failedJobs)

diff --git a/processor/processor_test.go b/processor/processor_test.go
@@ -45,6 +45,7 @@ import (
  "github.com/rudderlabs/rudder-server/utils/misc"
  "github.com/rudderlabs/rudder-server/utils/pubsub"
  testutils "github.com/rudderlabs/rudder-server/utils/tests"
+ . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
  "github.com/rudderlabs/rudder-server/utils/types"
 )
 
@@ -2024,8 +2025,8 @@ var _ = Describe("Processor", Ordered, func() {
  StoreInTx(gomock.Any(), gomock.Any(), gomock.Any()).
  AnyTimes()
 
- c.mockWriteProcErrorsDB.EXPECT().WithTx(gomock.Any()).Do(func(f func(tx *jobsdb.Tx) error) {
- _ = f(&jobsdb.Tx{})
+ c.mockWriteProcErrorsDB.EXPECT().WithTx(gomock.Any()).Do(func(f func(tx *Tx) error) {
+ _ = f(&Tx{})
  }).Return(nil).Times(0)
 
  // One Store call is expected for all events

diff --git a/router/batchrouter/handle.go b/router/batchrouter/handle.go
@@ -762,7 +762,7 @@ func (brt *Handle) updateJobStatus(batchJobs *BatchedJobs, isWarehouse bool, err
  }
 
  if brt.reporting != nil && brt.reportingEnabled {
- if err = brt.reporting.Report(reportMetrics, tx.SqlTx()); err != nil {
+ if err = brt.reporting.Report(reportMetrics, tx.Tx()); err != nil {
  return fmt.Errorf("reporting metrics: %w", err)
  }
  }

diff --git a/router/batchrouter/handle_async.go b/router/batchrouter/handle_async.go
@@ -58,7 +58,7 @@ func (brt *Handle) updateJobStatuses(ctx context.Context, destinationID string,
  }
 
  if brt.reporting != nil && brt.reportingEnabled {
- if err = brt.reporting.Report(reportMetrics, tx.SqlTx()); err != nil {
+ if err = brt.reporting.Report(reportMetrics, tx.Tx()); err != nil {
  return fmt.Errorf("reporting metrics: %w", err)
  }
  }
@@ -688,7 +688,7 @@ func (brt *Handle) setMultipleJobStatus(asyncOutput common.AsyncUploadOutput, at
  }
 
  if brt.reporting != nil && brt.reportingEnabled {
- if err = brt.reporting.Report(reportMetrics, tx.SqlTx()); err != nil {
+ if err = brt.reporting.Report(reportMetrics, tx.Tx()); err != nil {
  return fmt.Errorf("reporting metrics: %w", err)
  }
  }

diff --git a/router/batchrouter/worker.go b/router/batchrouter/worker.go
@@ -142,7 +142,7 @@ func (w *worker) processJobAsync(jobsWg *sync.WaitGroup, destinationJobs *Destin
  return fmt.Errorf("marking %s job statuses as aborted: %w", brt.destType, err)
  }
  if brt.reporting != nil && brt.reportingEnabled {
- if err = brt.reporting.Report(reportMetrics, tx.SqlTx()); err != nil {
+ if err = brt.reporting.Report(reportMetrics, tx.Tx()); err != nil {
  return fmt.Errorf("reporting metrics: %w", err)
  }
  }