Added connection params of rateLimitRetry and rateLimitRetryTimeout

Author: msrathore-db

csharp/src/Drivers/Databricks/DatabricksConnection.cs

@@ -73,7 +73,9 @@ internal class DatabricksConnection : SparkHttpConnection
         private const long DefaultMaxBytesPerFetchRequest = 400 * 1024 * 1024; // 400MB
         private long _maxBytesPerFetchRequest = DefaultMaxBytesPerFetchRequest;
         private const bool DefaultRetryOnUnavailable = true;
+        private const bool DefaultRateLimitRetry = true;
         private const int DefaultTemporarilyUnavailableRetryTimeout = 900;
+        private const int DefaultRateLimitRetryTimeout = 120;
         private bool _useDescTableExtended = false;
 
         // Trace propagation configuration
@@ -540,15 +542,25 @@ protected internal override bool TrySetGetDirectResults(IRequest request)
         public bool RunAsyncInThrift => _runAsyncInThrift;
 
         /// <summary>
-        /// Gets a value indicating whether to retry requests that receive a 503 response with a Retry-After header.
+        /// Gets a value indicating whether to retry requests that receive retryable responses (408, 502, 503, 504) .
         /// </summary>
         protected bool TemporarilyUnavailableRetry { get; private set; } = DefaultRetryOnUnavailable;
 
         /// <summary>
-        /// Gets the maximum total time in seconds to retry 503 responses before failing.
+        /// Gets the maximum total time in seconds to retry retryable responses (408, 502, 503, 504) before failing.
         /// </summary>
         protected int TemporarilyUnavailableRetryTimeout { get; private set; } = DefaultTemporarilyUnavailableRetryTimeout;
 
+        /// <summary>
+        /// Gets a value indicating whether to retry requests that receive HTTP 429 responses.
+        /// </summary>
+        protected bool RateLimitRetry { get; private set; } = DefaultRateLimitRetry;
+
+        /// <summary>
+        /// Gets the number of seconds to wait before stopping an attempt to retry HTTP 429 responses.
+        /// </summary>
+        protected int RateLimitRetryTimeout { get; private set; } = DefaultRateLimitRetryTimeout;
+
         protected override HttpMessageHandler CreateHttpHandler()
         {
             HttpMessageHandler baseHandler = base.CreateHttpHandler();
@@ -565,7 +577,7 @@ protected override HttpMessageHandler CreateHttpHandler()
             // Current chain order (outermost to innermost):
             // 1. OAuth handlers (OAuthDelegatingHandler, etc.) - only on baseHandler for API requests
             // 2. ThriftErrorMessageHandler - extracts x-thriftserver-error-message and throws descriptive exceptions
-            // 3. RetryHttpHandler - retries 408, 502, 503, 504 with Retry-After support
+            // 3. RetryHttpHandler - retries 408, 429, 502, 503, 504 with Retry-After support
             // 4. TracingDelegatingHandler - propagates W3C trace context
             // 5. Base HTTP handler - actual network communication
             //
@@ -586,16 +598,16 @@ protected override HttpMessageHandler CreateHttpHandler()
                 baseAuthHandler = new TracingDelegatingHandler(baseAuthHandler, this, _traceParentHeaderName, _traceStateEnabled);
             }
 
-            if (TemporarilyUnavailableRetry)
+            if (TemporarilyUnavailableRetry || RateLimitRetry)
             {
-                // Add retry handler for 408, 502, 503, 504 responses with Retry-After support
+                // Add retry handler for 408, 429, 502, 503, 504 responses with Retry-After support
                 // This must be INSIDE ThriftErrorMessageHandler so retries happen before exceptions are thrown
-                baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout);
-                baseAuthHandler = new RetryHttpHandler(baseAuthHandler, TemporarilyUnavailableRetryTimeout);
+                baseHandler = new RetryHttpHandler(baseHandler, TemporarilyUnavailableRetryTimeout, RateLimitRetryTimeout, TemporarilyUnavailableRetry, RateLimitRetry);
+                baseAuthHandler = new RetryHttpHandler(baseAuthHandler, TemporarilyUnavailableRetryTimeout, RateLimitRetryTimeout, TemporarilyUnavailableRetry, RateLimitRetry);
             }
 
             // Add Thrift error message handler AFTER retry handler (OUTSIDE in the chain)
-            // This ensures retryable status codes (408, 502, 503, 504) are retried by RetryHttpHandler
+            // This ensures retryable status codes (408, 429, 502, 503, 504) are retried by RetryHttpHandler
             // before ThriftErrorMessageHandler throws exceptions with Thrift error messages
             baseHandler = new ThriftErrorMessageHandler(baseHandler);
             baseAuthHandler = new ThriftErrorMessageHandler(baseAuthHandler);
@@ -900,6 +912,16 @@ protected override void ValidateOptions()
                 TemporarilyUnavailableRetry = tempUnavailableRetryValue;
             }
 
+            if (Properties.TryGetValue(DatabricksParameters.RateLimitRetry, out string? rateLimitRetryStr))
+            {
+                if (!bool.TryParse(rateLimitRetryStr, out bool rateLimitRetryValue))
+                {
+                    throw new ArgumentOutOfRangeException(DatabricksParameters.RateLimitRetry, rateLimitRetryStr,
+                        $"must be a value of false (disabled) or true (enabled). Default is true.");
+                }
+
+                RateLimitRetry = rateLimitRetryValue;
+            }
 
             if (Properties.TryGetValue(DatabricksParameters.TemporarilyUnavailableRetryTimeout, out string? tempUnavailableRetryTimeoutStr))
             {
@@ -912,6 +934,17 @@ protected override void ValidateOptions()
                 TemporarilyUnavailableRetryTimeout = tempUnavailableRetryTimeoutValue;
             }
 
+            if (Properties.TryGetValue(DatabricksParameters.RateLimitRetryTimeout, out string? rateLimitRetryTimeoutStr))
+            {
+                if (!int.TryParse(rateLimitRetryTimeoutStr, out int rateLimitRetryTimeoutValue) ||
+                    rateLimitRetryTimeoutValue < 0)
+                {
+                    throw new ArgumentOutOfRangeException(DatabricksParameters.RateLimitRetryTimeout, rateLimitRetryTimeoutStr,
+                        $"must be a value of 0 (retry indefinitely) or a positive integer representing seconds. Default is 120 seconds (2 minutes).");
+                }
+                RateLimitRetryTimeout = rateLimitRetryTimeoutValue;
+            }
+
             // When TemporarilyUnavailableRetry is enabled, we need to make sure connection timeout (which is used to cancel the HttpConnection) is equal
             // or greater than TemporarilyUnavailableRetryTimeout so that it won't timeout before server startup timeout (TemporarilyUnavailableRetryTimeout)
             if (TemporarilyUnavailableRetry && TemporarilyUnavailableRetryTimeout * 1000 > ConnectTimeoutMilliseconds)

csharp/src/Drivers/Databricks/DatabricksParameters.cs

@@ -99,17 +99,31 @@ public class DatabricksParameters : SparkParameters
         /// </summary>
         public const string ServerSidePropertyPrefix = "adbc.databricks.ssp_";
 
-        /// Controls whether to retry requests that receive a 503 response with a Retry-After header.
+        /// <summary>
+        /// Controls whether to retry requests that receive retryable responses (408, 502, 503, 504).
         /// Default value is true (enabled). Set to false to disable retry behavior.
         /// </summary>
         public const string TemporarilyUnavailableRetry = "adbc.spark.temporarily_unavailable_retry";
 
         /// <summary>
-        /// Maximum total time in seconds to retry 503 responses before failing.
+        /// Maximum total time in seconds to retry retryable responses (408, 502, 503, 504) before failing.
         /// Default value is 900 seconds (15 minutes). Set to 0 to retry indefinitely.
         /// </summary>
         public const string TemporarilyUnavailableRetryTimeout = "adbc.spark.temporarily_unavailable_retry_timeout";
 
+        /// <summary>
+        /// Controls whether to retry requests that receive HTTP 429 (TooManyRequests) response.
+        /// Default value is true. Set to false to disable rate limit retry behavior.
+        /// </summary>
+        public const string RateLimitRetry = "adbc.databricks.rate_limit_retry";
+
+        /// <summary>
+        /// The number of seconds that the connector waits before stopping an attempt to retry an operation
+        /// when the operation receives an HTTP 429 response.
+        /// Default value is 120 seconds. Set to 0 to retry indefinitely.
+        /// </summary>
+        public const string RateLimitRetryTimeout = "adbc.databricks.rate_limit_retry_timeout";
+
         /// <summary>
         /// Maximum number of parallel downloads for CloudFetch operations.
         /// Default value is 3 if not specified.

csharp/src/Drivers/Databricks/RetryHttpHandler.cs

@@ -25,24 +25,44 @@
 namespace Apache.Arrow.Adbc.Drivers.Databricks
 {
     /// <summary>
-    /// HTTP handler that implements retry behavior for 408, 502, 503, and 504 responses.
+    /// HTTP handler that implements retry behavior for 408, 429, 502, 503, and 504 responses.
     /// Uses Retry-After header if present, otherwise uses exponential backoff.
     /// </summary>
     internal class RetryHttpHandler : DelegatingHandler
     {
         private readonly int _retryTimeoutSeconds;
+        private readonly int _rateLimitRetryTimeoutSeconds;
+        private readonly bool _retryTemporarilyUnavailableEnabled;
+        private readonly bool _rateLimitRetryEnabled;
         private readonly int _initialBackoffSeconds = 1;
         private readonly int _maxBackoffSeconds = 32;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="RetryHttpHandler"/> class.
         /// </summary>
         /// <param name="innerHandler">The inner handler to delegate to.</param>
-        /// <param name="retryTimeoutSeconds">Maximum total time in seconds to retry before failing.</param>
-        public RetryHttpHandler(HttpMessageHandler innerHandler, int retryTimeoutSeconds)
+        /// <param name="retryTimeoutSeconds">Maximum total time in seconds to retry retryable responses (408, 502, 503, 504) before failing.</param>
+        /// <param name="rateLimitRetryTimeoutSeconds">Maximum total time in seconds to retry HTTP 429 responses before failing.</param>
+        public RetryHttpHandler(HttpMessageHandler innerHandler, int retryTimeoutSeconds, int rateLimitRetryTimeoutSeconds)
+            : this(innerHandler, retryTimeoutSeconds, rateLimitRetryTimeoutSeconds, true, true)
+        {
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="RetryHttpHandler"/> class.
+        /// </summary>
+        /// <param name="innerHandler">The inner handler to delegate to.</param>
+        /// <param name="retryTimeoutSeconds">Maximum total time in seconds to retry retryable responses (408, 502, 503, 504) before failing.</param>
+        /// <param name="rateLimitRetryTimeoutSeconds">Maximum total time in seconds to retry 429 (rate limit) responses before failing.</param>
+        /// <param name="retryTemporarilyUnavailableEnabled">Whether to retry temporarily unavailable (408, 502, 503, 504) responses.</param>
+        /// <param name="rateLimitRetryEnabled">Whether to retry HTTP 429 responses.</param>
+        public RetryHttpHandler(HttpMessageHandler innerHandler, int retryTimeoutSeconds, int rateLimitRetryTimeoutSeconds, bool retryTemporarilyUnavailableEnabled, bool rateLimitRetryEnabled)
             : base(innerHandler)
         {
             _retryTimeoutSeconds = retryTimeoutSeconds;
+            _rateLimitRetryTimeoutSeconds = rateLimitRetryTimeoutSeconds;
+            _retryTemporarilyUnavailableEnabled = retryTemporarilyUnavailableEnabled;
+            _rateLimitRetryEnabled = rateLimitRetryEnabled;
         }
 
         /// <summary>
@@ -59,10 +79,10 @@ protected override async Task<HttpResponseMessage> SendAsync(
 
             HttpResponseMessage response;
             string? lastErrorMessage = null;
-            DateTime startTime = DateTime.UtcNow;
             int attemptCount = 0;
             int currentBackoffSeconds = _initialBackoffSeconds;
-            int totalRetrySeconds = 0;
+            int totalServiceUnavailableRetrySeconds = 0;
+            int totalTooManyRequestsRetrySeconds = 0;
 
             do
             {
@@ -82,14 +102,6 @@ protected override async Task<HttpResponseMessage> SendAsync(
 
                 attemptCount++;
 
-                // Check if we've exceeded the timeout
-                TimeSpan elapsedTime = DateTime.UtcNow - startTime;
-                if (_retryTimeoutSeconds > 0 && elapsedTime.TotalSeconds > _retryTimeoutSeconds)
-                {
-                    // We've exceeded the timeout, so break out of the loop
-                    break;
-                }
-
                 int waitSeconds;
 
                 // Check for Retry-After header
@@ -123,12 +135,27 @@ protected override async Task<HttpResponseMessage> SendAsync(
                 // Reset the request content for the next attempt
                 request.Content = null;
 
-                // Update total retry time
-                totalRetrySeconds += waitSeconds;
-                if (_retryTimeoutSeconds > 0 && totalRetrySeconds > _retryTimeoutSeconds)
+                // Check if we would exceed the timeout after waiting, based on error type
+                bool isTooManyRequests = response.StatusCode == HttpStatusCode.TooManyRequests;
+                if (isTooManyRequests)
+                {
+                    // Check 429 rate limit timeout
+                    if (_rateLimitRetryTimeoutSeconds > 0 && totalTooManyRequestsRetrySeconds + waitSeconds > _rateLimitRetryTimeoutSeconds)
+                    {
+                        // We've exceeded the rate limit retry timeout, so break out of the loop
+                        break;
+                    }
+                    totalTooManyRequestsRetrySeconds += waitSeconds;
+                }
+                else
                 {
-                    // We've exceeded the timeout, so break out of the loop
-                    break;
+                    // Check service unavailable timeout for other retryable errors (408, 502, 503, 504)
+                    if (_retryTimeoutSeconds > 0 && totalServiceUnavailableRetrySeconds + waitSeconds > _retryTimeoutSeconds)
+                    {
+                        // We've exceeded the retry timeout, so break out of the loop
+                        break;
+                    }
+                    totalServiceUnavailableRetrySeconds += waitSeconds;
                 }
 
                 // Wait for the calculated time
@@ -153,10 +180,18 @@ protected override async Task<HttpResponseMessage> SendAsync(
         /// </summary>
         private bool IsRetryableStatusCode(HttpStatusCode statusCode)
         {
-            return statusCode == HttpStatusCode.RequestTimeout ||      // 408
-                   statusCode == HttpStatusCode.BadGateway ||          // 502
-                   statusCode == HttpStatusCode.ServiceUnavailable ||  // 503
-                   statusCode == HttpStatusCode.GatewayTimeout;        // 504
+            // Check too many requests separately
+            if (statusCode == HttpStatusCode.TooManyRequests)         // 429
+                return _rateLimitRetryEnabled;
+
+            // Check other retryable codes
+            if (statusCode == HttpStatusCode.RequestTimeout ||        // 408
+                statusCode == HttpStatusCode.BadGateway ||            // 502
+                statusCode == HttpStatusCode.ServiceUnavailable ||    // 503
+                statusCode == HttpStatusCode.GatewayTimeout)          // 504
+                return _retryTemporarilyUnavailableEnabled;
+
+            return false;
         }
 
         /// <summary>