feat: implement retries with exponential backoff (#55)

* feat: implement retries with exponential backoff

* refactor: restructure query runners and optimize backoff

* style: remove redundant comment line

* Apply suggestion from @Copilot

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* refactor: extract retry invocation logic

* fix: handle overflow when calculating delay

* test: add case for Mgmt error

* doc: adjust log statements, comments

* test: fix expected error

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: chrischangcode

README.md

@@ -39,7 +39,23 @@ Where `./testdata/commands.kql` contains Kusto management commands, e.g.:
 - `--auth-azcli` or other authentication options (see below)
 - `--kusto-endpoint` (required)
 - `--kusto-database` (required)
+- `--max-retries=3` (optional)
+- `--max-timeout=60` (optional)
 
+#### Retry Options
+
+Both file ingestion and management commands support automatic retry with exponential backoff for transient errors. 
+Default values can be overridden by specifying the arguments as below:
+
+```
+$ kusto-ingest file ./testdata/logs.multijson \
+    --max-retries=5 \
+    --max-timeout=120 \
+    # ... other options
+```
+
+- `--max-retries=3` - Maximum number of retry attempts (default: 3)
+- `--max-timeout=60` - Maximum total time in seconds for all retries (default: 60)
 
 ### Authentication
 

internal/kusto/file.go

@@ -46,6 +46,8 @@ func (f FileIngestOptions) Run(cli cli.Provider) error {
 		"target.table", f.KustoTarget.Table,
 		"auth.tenant", f.Auth.TenantID,
 		"auth.clientID", f.Auth.ClientID,
+		"maxRetries", f.MaxRetries,
+		"maxTimeout", f.MaxTimeout,
 	)
 
 	fileOptions, err := f.FileOptions()
@@ -62,17 +64,27 @@ func (f FileIngestOptions) Run(cli cli.Provider) error {
 	ctx, cancel := cli.Context()
 	defer cancel()
 
+	invokeIngest := func() error {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		_, err := ingestor.FromFile(ctx, f.SourceFile, fileOptions...)
+		return err
+	}
+
 	cli.Logger().Info("file ingestion started")
 	start := time.Now()
-	_, err = ingestor.FromFile(
-		ctx,
-		f.SourceFile,
-		fileOptions...,
+	err = invokeWithRetries(
+		invokeIngest,
+		f.MaxRetries,
+		f.MaxTimeout,
+		cli.Logger(),
 	)
 	if err != nil {
-		return fmt.Errorf("ingest from file %q: %w", f.SourceFile, err)
+		cli.Logger().Error("failed to ingest file", "error", err, "file", f.SourceFile)
+		return err
 	}
-	cli.Logger().Info("file ingestion completed", "duration", time.Since(start))
 
+	cli.Logger().Info("file ingestion completed successfully", "duration", time.Since(start))
 	return nil
 }

internal/kusto/file_test.go

@@ -25,11 +25,11 @@ func writeToTestFile(t testing.TB, fileName string, content []byte) string {
 }
 
 func Test_FileIngestOptions_FileOptions(t *testing.T) {
-	t.Run("with mapping", func (t *testing.T)  {
+	t.Run("with mapping", func(t *testing.T) {
 		mappingFile := writeToTestFile(t, "test-mapping.json", []byte(`[]`))
 
 		options := FileIngestOptions{
-			Format: "csv",
+			Format:       "csv",
 			MappingsFile: mappingFile,
 		}
 
@@ -40,15 +40,15 @@ func Test_FileIngestOptions_FileOptions(t *testing.T) {
 
 	t.Run("with invalid mapping file", func(t *testing.T) {
 		options := FileIngestOptions{
-			Format: "csv",
+			Format:       "csv",
 			MappingsFile: "some-random-file",
 		}
 
 		_, err := options.FileOptions()
 		assert.Error(t, err)
 	})
 
-	t.Run("without mapping", func (t *testing.T)  {
+	t.Run("without mapping", func(t *testing.T) {
 		options := FileIngestOptions{
 			Format: "csv",
 		}
@@ -61,7 +61,7 @@ func Test_FileIngestOptions_FileOptions(t *testing.T) {
 
 func Test_FileIngestOptions_Run_IngestFile_NoMapping(t *testing.T) {
 	sourceFile := writeToTestFile(t, "logs.json", []byte("{}"))
-	
+
 	cli := testingcli.New()
 
 	ingestor := testingkusto.New(func(ing *testingkusto.Ingestor) {
@@ -74,9 +74,9 @@ func Test_FileIngestOptions_Run_IngestFile_NoMapping(t *testing.T) {
 	})
 
 	opts := FileIngestOptions{
-		SourceFile: sourceFile,
-		Format: "multijson",
-		Auth: newTestAuth(),
+		SourceFile:  sourceFile,
+		Format:      "multijson",
+		Auth:        newTestAuth(),
 		KustoTarget: newTestKustoTarget(),
 
 		ingestorBuildSettings: ingestorBuildSettings{
@@ -93,7 +93,7 @@ func Test_FileIngestOptions_Run_IngestFile_NoMapping(t *testing.T) {
 func Test_FileIngestOptions_Run_IngestFile_WithMapping(t *testing.T) {
 	sourceFile := writeToTestFile(t, "logs.json", []byte("{}"))
 	sourceFileMapping := writeToTestFile(t, "logs-mapping.json", []byte("[]"))
-	
+
 	cli := testingcli.New()
 
 	ingestor := testingkusto.New(func(ing *testingkusto.Ingestor) {
@@ -106,11 +106,11 @@ func Test_FileIngestOptions_Run_IngestFile_WithMapping(t *testing.T) {
 	})
 
 	opts := FileIngestOptions{
-		SourceFile: sourceFile,
-		Format: "multijson",
+		SourceFile:   sourceFile,
+		Format:       "multijson",
 		MappingsFile: sourceFileMapping,
-		Auth: newTestAuth(),
-		KustoTarget: newTestKustoTarget(),
+		Auth:         newTestAuth(),
+		KustoTarget:  newTestKustoTarget(),
 
 		ingestorBuildSettings: ingestorBuildSettings{
 			CreateIngestor: func(target KustoTargetOptions, auth AuthOptions) (ingest.Ingestor, error) {
@@ -121,4 +121,25 @@ func Test_FileIngestOptions_Run_IngestFile_WithMapping(t *testing.T) {
 
 	err := opts.Run(cli)
 	assert.NoError(t, err)
-}
+}
+
+func Test_FileIngestOptions_Run_CreateIngestorError(t *testing.T) {
+	sourceFile := writeToTestFile(t, "logs.json", []byte("{}"))
+	cli := testingcli.New()
+
+	opts := FileIngestOptions{
+		SourceFile:  sourceFile,
+		Format:      "multijson",
+		Auth:        newTestAuth(),
+		KustoTarget: newTestKustoTarget(),
+		ingestorBuildSettings: ingestorBuildSettings{
+			CreateIngestor: func(target KustoTargetOptions, auth AuthOptions) (ingest.Ingestor, error) {
+				return nil, assert.AnError
+			},
+		},
+	}
+
+	err := opts.Run(cli)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "create Kusto ingestor")
+}

internal/kusto/management.go

@@ -15,6 +15,8 @@ func (m ManagementOptions) Run(cli cli.Provider) error {
 		"target.database", m.KustoTarget.Database,
 		"auth.tenant", m.Auth.TenantID,
 		"auth.clientID", m.Auth.ClientID,
+		"maxRetries", m.MaxRetries,
+		"maxTimeout", m.MaxTimeout,
 	)
 
 	queryer, err := m.createQueryClient(m.KustoTarget, m.Auth)
@@ -26,14 +28,29 @@ func (m ManagementOptions) Run(cli cli.Provider) error {
 	ctx, cancel := cli.Context()
 	defer cancel()
 
-	cli.Logger().Info("executing management commands")
-	start := time.Now()
 	stmt := kql.New("").AddUnsafe(string(m.Source))
-	_, err = queryer.Mgmt(ctx, m.KustoTarget.Database, stmt)
+	invokeQuery := func() error {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		_, err := queryer.Mgmt(ctx, m.KustoTarget.Database, stmt)
+		return err
+	}
+
+	cli.Logger().Info("executing management command")
+
+	start := time.Now()
+	err = invokeWithRetries(
+		invokeQuery,
+		m.MaxRetries,
+		m.MaxTimeout,
+		cli.Logger(),
+	)
 	if err != nil {
-		return fmt.Errorf("execute management commands: %w", err)
+		cli.Logger().Error("failed to execute management command", "error", err)
+		return err
 	}
-	cli.Logger().Info("management commands executed", "duration", time.Since(start))
 
+	cli.Logger().Info("management command executed successfully", "duration", time.Since(start))
 	return nil
 }

internal/kusto/management_test.go

@@ -64,7 +64,8 @@ func Test_ManagementOptions_Run_MgmtError(t *testing.T) {
 
 	q := testingkusto.NewQueryClient(func(qc *testingkusto.QueryClient) {
 		qc.MgmtFn = func(_ context.Context, db string, stmt kusto.Statement, _ ...kusto.QueryOption) (*kusto.RowIterator, error) {
-			return nil, errors.New("mgmt failed")
+			// Return a non-retryable error
+			return nil, errors.New("mgmt command failed")
 		}
 	})
 
@@ -81,5 +82,6 @@ func Test_ManagementOptions_Run_MgmtError(t *testing.T) {
 
 	err := opts.Run(cli)
 	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "execute management commands")
+	// The error will be wrapped by invokeWithRetries as "non-retryable error"
+	assert.Contains(t, err.Error(), "non-retryable kusto error")
 }

internal/kusto/options.go

@@ -28,6 +28,10 @@ type FileIngestOptions struct {
 	Auth        AuthOptions        `embed:"" prefix:"auth-"`
 	KustoTarget KustoTargetOptions `embed:"" prefix:"kusto-"`
 
+	// Retry and timeout configuration
+	MaxRetries int `optional:"" default:"3" help:"Maximum number of retries for transient errors (default: 3)."`
+	MaxTimeout int `optional:"" default:"60" help:"Maximum timeout in seconds for all retries (default: 60)."`
+
 	// for unit test
 	ingestorBuildSettings `kong:"-"`
 }
@@ -39,6 +43,10 @@ type ManagementOptions struct {
 	Auth        AuthOptions        `embed:"" prefix:"auth-"`
 	KustoTarget KustoTargetOptions `embed:"" prefix:"kusto-"`
 
+	// Retry and timeout configuration
+	MaxRetries int `optional:"" default:"3" help:"Maximum number of retries for transient errors (default: 3)."`
+	MaxTimeout int `optional:"" default:"60" help:"Maximum timeout in seconds for all retries (default: 60)."`
+
 	// for unit test
 	ingestorBuildSettings `kong:"-"`
 }

internal/kusto/retry.go

@@ -0,0 +1,119 @@
+package kusto
+
+import (
+	"fmt"
+	"math"
+	"math/rand/v2"
+	"time"
+
+	"github.com/Azure/azure-kusto-go/kusto/data/errors"
+	"github.com/charmbracelet/log"
+)
+
+const (
+	retryBaseDelay = 1 * time.Second
+)
+
+/*
+invokeWithRetries executes the provided function with retry logic for transient Kusto errors.
+It uses exponential backoff with jitter and respects both maxRetries and maxTimeout constraints.
+
+Parameters:
+  - invoke: The function to execute with retries
+  - maxRetries: Maximum number of retry attempts (e.g., 3 means 4 total attempts: initial + 3 retries)
+  - maxTimeout: Maximum total duration in seconds before timing out
+  - logger: Logger for logging retry attempts (optional, can be nil)
+
+Returns an error if:
+  - A non-retryable error is encountered
+  - Maximum timeout is reached
+  - Maximum retries are exhausted
+*/
+func invokeWithRetries(
+	invoke func() error,
+	maxRetries int,
+	maxTimeout int,
+	logger *log.Logger,
+) error {
+	var err error
+	maxTimeoutDuration := time.Duration(maxTimeout) * time.Second
+	deadline := time.Now().Add(maxTimeoutDuration)
+
+	for attempt := 0; attempt <= maxRetries; attempt++ {
+		// Check if we've exceeded the deadline before attempting
+		if attempt > 0 && time.Now().After(deadline) {
+			return fmt.Errorf("max timeout reached after %d retries: %w", attempt-1, err)
+		}
+
+		err = invoke()
+		if err == nil {
+			return nil
+		}
+
+		// Check error type for retry logic using azure-kusto-go SDK's Retry function
+		if !errors.Retry(err) {
+			return fmt.Errorf("non-retryable kusto error: %w", err)
+		}
+
+		// Calculate next backoff duration
+		backoffDelay := calculateDelay(attempt, retryBaseDelay)
+
+		if time.Now().Add(backoffDelay).After(deadline) {
+			return fmt.Errorf("max timeout reached after %d retries: %w", attempt, err)
+		}
+
+		if logger != nil {
+			logger.Warn("transient kusto error, will retry", "error", err, "attempt", attempt+1, "backoff", backoffDelay)
+		}
+
+		time.Sleep(backoffDelay)
+	}
+
+	return fmt.Errorf("exhausted max retries (%d): %w", maxRetries, err)
+}
+
+/*
+calculateDelay computes the next retry delay duration.
+Uses exponential backoff with jitter to prevent thundering herd issues.
+Formula: baseDelay * (2^attempt) * (1 + jitter)
+Where jitter is a random value between 0 and 0.1 (10% jitter)
+*/
+func calculateDelay(attempt int, baseDelay time.Duration) time.Duration {
+	if attempt < 0 {
+		return baseDelay
+	}
+
+	// Check for overflow before multiplying by checking if we would exceed math.MaxInt64
+	maxDelay := time.Duration(math.MaxInt64)
+
+	// Calculate 2^attempt using bit shifting, checking for overflow
+	// If attempt >= 63, shifting would overflow (2^63 > MaxInt64)
+	if attempt >= 63 {
+		return maxDelay
+	}
+
+	multiplier := int64(1) << attempt // 2^attempt
+
+	// Check if baseDelay * multiplier would overflow
+	if baseDelay > maxDelay/time.Duration(multiplier) {
+		return maxDelay
+	}
+
+	delay := baseDelay * time.Duration(multiplier)
+
+	// Add jitter: random value between 0% and 10% of the delay
+	jitter := rand.Float64() * 0.1 // 0% to 10% jitter
+	jitteredDelay := time.Duration(float64(delay) * (1.0 + jitter))
+
+	// Check for overflow after jitter (float64 conversion could overflow)
+	if jitteredDelay < 0 || jitteredDelay > maxDelay {
+		return maxDelay
+	}
+
+	// Apply max cap again after jitter
+	if jitteredDelay > maxDelay {
+		return maxDelay
+	}
+
+	return jitteredDelay
+}