feat(mc2mc): expose retry config (#93)

feat: expose retry config

Author: deryrahman

mc2mc/internal/client/client.go

@@ -18,6 +18,7 @@ type OdpsClient interface {
 	SetDefaultProject(project string)
 	SetLogViewRetentionInDays(days int)
 	SetDryRun(dryRun bool)
+	SetRetry(max int, backoffMs int)
 }
 
 type Client struct {

mc2mc/internal/client/odps.go

@@ -19,6 +19,7 @@ type odpsClient struct {
 
 	logViewRetentionInDays int
 	isDryRun               bool
+	retry                  func(f func() error) error
 }
 
 // NewODPSClient creates a new odpsClient instance
@@ -27,6 +28,9 @@ func NewODPSClient(logger *slog.Logger, client *odps.Odps) *odpsClient {
 		logger:                 logger,
 		client:                 client,
 		logViewRetentionInDays: 2,
+		retry: func(f func() error) error {
+			return retry(logger, 3, 1000, f)
+		},
 	}
 }
 
@@ -49,7 +53,7 @@ func (c *odpsClient) ExecSQL(ctx context.Context, query string, additionalHints
 	// generate log view
 	url, err := c.generateLogView(taskIns)
 	if err != nil {
-		err = e.Join(err, taskIns.Terminate())
+		err = e.Join(err, c.terminate(taskIns))
 		return errors.WithStack(err)
 	}
 	c.logger.Info(fmt.Sprintf("taskId: %s, log view: %s , hints: (%s)", taskIns.Id(), url, getHintsString(hints)))
@@ -83,6 +87,13 @@ func (c *odpsClient) SetDefaultProject(project string) {
 	c.client.SetDefaultProjectName(project)
 }
 
+// SetRetry sets the retry configuration for the odps client
+func (c *odpsClient) SetRetry(max int, backoffMs int) {
+	c.retry = func(f func() error) error {
+		return retry(c.logger, max, int64(backoffMs), f)
+	}
+}
+
 // GetPartitionNames returns the partition names of the given table
 // by querying the table schema.
 func (c *odpsClient) GetPartitionNames(_ context.Context, tableID string) ([]string, error) {
@@ -158,11 +169,6 @@ func (c *odpsClient) wait(taskIns *odps.Instance) <-chan error {
 	return errChan
 }
 
-// retry retries the given function with exponential backoff
-func (c *odpsClient) retry(f func() error) error {
-	return retry(c.logger, 3, 1000, f)
-}
-
 func (c *odpsClient) terminate(instance *odps.Instance) error {
 	if instance == nil {
 		return nil

mc2mc/internal/client/setup.go

@@ -80,3 +80,10 @@ func SetupOTelSDK(collectorGRPCEndpoint string, otelAttributes string) SetupFn {
 		return nil
 	}
 }
+
+func SetupRetry(max int, backoffMs int) SetupFn {
+	return func(c *Client) error {
+		c.OdpsClient.SetRetry(max, backoffMs)
+		return nil
+	}
+}

mc2mc/internal/config/config.go

@@ -24,6 +24,8 @@ type ConfigEnv struct {
 	LogViewRetentionInDays      int               `env:"LOG_VIEW_RETENTION_IN_DAYS" envDefault:"2"`
 	DisableMultiQueryGeneration bool              `env:"DISABLE_MULTI_QUERY_GENERATION" envDefault:"false"`
 	DryRun                      bool              `env:"DRY_RUN" envDefault:"false"`
+	RetryMax                    int               `env:"RETRY_MAX" envDefault:"3"`
+	RetryBackoffMs              int               `env:"RETRY_BACKOFF_MS" envDefault:"1000"`
 	// TODO: delete this
 	DevEnablePartitionValue string `env:"DEV__ENABLE_PARTITION_VALUE" envDefault:"false"`
 	DevEnableAutoPartition  string `env:"DEV__ENABLE_AUTO_PARTITION" envDefault:"false"`

mc2mc/mc2mc.go

@@ -45,6 +45,7 @@ func mc2mc(envs []string) error {
 		client.SetupDefaultProject(cfg.ExecutionProject),
 		client.SetUpLogViewRetentionInDays(cfg.LogViewRetentionInDays),
 		client.SetupDryRun(cfg.DryRun),
+		client.SetupRetry(cfg.RetryMax, cfg.RetryBackoffMs),
 	)
 	if err != nil {
 		return errors.WithStack(err)