feat(csharp/src/Drivers/Databricks): Added support for user-configurable Fetch heartbeat interval param (apache#3472)

## Description
Added support for the Fetch heartbeat interval connection param. The
default values are kept same as before. Users can now configure this
param.

[PECO-2736](https://databricks.atlassian.net/browse/PECO-2736)

Author: msrathore-db

csharp/src/Drivers/Databricks/DatabricksConnection.cs

@@ -83,6 +83,12 @@ internal class DatabricksConnection : SparkHttpConnection
         // Identity federation client ID for token exchange
         private string? _identityFederationClientId;
 
+        // Heartbeat interval configuration
+        private int _fetchHeartbeatIntervalSeconds = DatabricksConstants.DefaultOperationStatusPollingIntervalSeconds;
+
+        // Request timeout configuration
+        private int _operationStatusRequestTimeoutSeconds = DatabricksConstants.DefaultOperationStatusRequestTimeoutSeconds;
+
         // Default namespace
         private TNamespace? _defaultNamespace;
 
@@ -386,6 +392,40 @@ private void ValidateProperties()
             {
                 _identityFederationClientId = identityFederationClientId;
             }
+
+            if (Properties.TryGetValue(DatabricksParameters.FetchHeartbeatInterval, out string? fetchHeartbeatIntervalStr))
+            {
+                if (!int.TryParse(fetchHeartbeatIntervalStr, out int fetchHeartbeatIntervalValue))
+                {
+                    throw new ArgumentException($"Parameter '{DatabricksParameters.FetchHeartbeatInterval}' value '{fetchHeartbeatIntervalStr}' could not be parsed. Valid values are positive integers.");
+                }
+
+                if (fetchHeartbeatIntervalValue <= 0)
+                {
+                    throw new ArgumentOutOfRangeException(
+                        nameof(Properties),
+                        fetchHeartbeatIntervalValue,
+                        $"Parameter '{DatabricksParameters.FetchHeartbeatInterval}' value must be a positive integer.");
+                }
+                _fetchHeartbeatIntervalSeconds = fetchHeartbeatIntervalValue;
+            }
+
+            if (Properties.TryGetValue(DatabricksParameters.OperationStatusRequestTimeout, out string? operationStatusRequestTimeoutStr))
+            {
+                if (!int.TryParse(operationStatusRequestTimeoutStr, out int operationStatusRequestTimeoutValue))
+                {
+                    throw new ArgumentException($"Parameter '{DatabricksParameters.OperationStatusRequestTimeout}' value '{operationStatusRequestTimeoutStr}' could not be parsed. Valid values are positive integers.");
+                }
+
+                if (operationStatusRequestTimeoutValue <= 0)
+                {
+                    throw new ArgumentOutOfRangeException(
+                        nameof(Properties),
+                        operationStatusRequestTimeoutValue,
+                        $"Parameter '{DatabricksParameters.OperationStatusRequestTimeout}' value must be a positive integer.");
+                }
+                _operationStatusRequestTimeoutSeconds = operationStatusRequestTimeoutValue;
+            }
         }
 
         /// <summary>
@@ -428,6 +468,16 @@ protected internal override bool TrySetGetDirectResults(IRequest request)
         /// </summary>
         internal TNamespace? DefaultNamespace => _defaultNamespace;
 
+        /// <summary>
+        /// Gets the heartbeat interval in seconds for long-running operations.
+        /// </summary>
+        internal int FetchHeartbeatIntervalSeconds => _fetchHeartbeatIntervalSeconds;
+
+        /// <summary>
+        /// Gets the request timeout in seconds for operation status polling requests.
+        /// </summary>
+        internal int OperationStatusRequestTimeoutSeconds => _operationStatusRequestTimeoutSeconds;
+
         /// <summary>
         /// Gets whether multiple catalog is supported
         /// </summary>

csharp/src/Drivers/Databricks/DatabricksParameters.cs

@@ -228,6 +228,22 @@ public class DatabricksParameters : SparkParameters
         /// Default value is false if not specified.
         /// </summary>
         public const string DriverConfigTakePrecedence = "adbc.databricks.driver_config_take_precedence";
+
+        /// <summary>
+        /// The interval in seconds for heartbeat polling during long-running operations.
+        /// This prevents queries from timing out by periodically checking operation status.
+        /// Default value is 60 seconds if not specified.
+        /// Must be a positive integer value.
+        /// </summary>
+        public const string FetchHeartbeatInterval = "adbc.databricks.fetch_heartbeat_interval";
+
+        /// <summary>
+        /// The timeout in seconds for operation status polling requests.
+        /// This controls how long to wait for each individual polling request to complete.
+        /// Default value is 30 seconds if not specified.
+        /// Must be a positive integer value.
+        /// </summary>
+        public const string OperationStatusRequestTimeout = "adbc.databricks.operation_status_request_timeout";
     }
 
     /// <summary>
@@ -236,12 +252,12 @@ public class DatabricksParameters : SparkParameters
     public class DatabricksConstants
     {
         /// <summary>
-        /// Default heartbeat interval in seconds for long-running operations. TODO: make this user-configurable
+        /// Default heartbeat interval in seconds for long-running operations.
         /// </summary>
         public const int DefaultOperationStatusPollingIntervalSeconds = 60;
 
         /// <summary>
-        /// Default timeout in seconds for operation status polling requests. TODO: make this user-configurable
+        /// Default timeout in seconds for operation status polling requests.
         /// </summary>
         public const int DefaultOperationStatusRequestTimeoutSeconds = 30;
 

csharp/src/Drivers/Databricks/Reader/DatabricksCompositeReader.cs

@@ -80,7 +80,7 @@ internal DatabricksCompositeReader(
             }
             if (_response.DirectResults?.ResultSet?.HasMoreRows ?? true)
             {
-                operationStatusPoller = operationPoller ?? new DatabricksOperationStatusPoller(_statement, response);
+                operationStatusPoller = operationPoller ?? new DatabricksOperationStatusPoller(_statement, response, GetHeartbeatIntervalFromConnection(), GetRequestTimeoutFromConnection());
                 operationStatusPoller.Start();
             }
         }
@@ -197,5 +197,41 @@ private void StopOperationStatusPoller()
             operationStatusPoller?.Dispose();
             operationStatusPoller = null;
         }
+
+        /// <summary>
+        /// Gets the heartbeat interval from the statement's connection.
+        /// </summary>
+        /// <returns>The heartbeat interval in seconds, or default if not available.</returns>
+        private int GetHeartbeatIntervalFromConnection()
+        {
+            if (_statement is DatabricksStatement databricksStatement)
+            {
+                var connection = databricksStatement.Connection;
+                if (connection is DatabricksConnection databricksConnection)
+                {
+                    return databricksConnection.FetchHeartbeatIntervalSeconds;
+                }
+            }
+
+            return DatabricksConstants.DefaultOperationStatusPollingIntervalSeconds;
+        }
+
+        /// <summary>
+        /// Gets the request timeout from the statement's connection.
+        /// </summary>
+        /// <returns>The request timeout in seconds, or default if not available.</returns>
+        private int GetRequestTimeoutFromConnection()
+        {
+            if (_statement is DatabricksStatement databricksStatement)
+            {
+                var connection = databricksStatement.Connection;
+                if (connection is DatabricksConnection databricksConnection)
+                {
+                    return databricksConnection.OperationStatusRequestTimeoutSeconds;
+                }
+            }
+
+            return DatabricksConstants.DefaultOperationStatusRequestTimeoutSeconds;
+        }
     }
 }

csharp/test/Drivers/Databricks/E2E/DatabricksConnectionTest.cs

@@ -325,6 +325,16 @@ public InvalidConnectionParametersTestData()
                 Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.TracePropagationEnabled] = "notabool" }, typeof(ArgumentException)));
                 Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.TraceParentHeaderName] = "" }, typeof(ArgumentException)));
                 Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.TraceStateEnabled] = "notabool" }, typeof(ArgumentException)));
+
+                // Tests for fetch heartbeat interval parameter
+                Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.FetchHeartbeatInterval] = "notanumber" }, typeof(ArgumentException)));
+                Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.FetchHeartbeatInterval] = "0" }, typeof(ArgumentOutOfRangeException)));
+                Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.FetchHeartbeatInterval] = "-1" }, typeof(ArgumentOutOfRangeException)));
+
+                // Tests for operation status request timeout parameter
+                Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.OperationStatusRequestTimeout] = "notanumber" }, typeof(ArgumentException)));
+                Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.OperationStatusRequestTimeout] = "0" }, typeof(ArgumentOutOfRangeException)));
+                Add(new(new() { [SparkParameters.Type] = SparkServerTypeConstants.Http, [SparkParameters.HostName] = "valid.server.com", [AdbcOptions.Username] = "user", [AdbcOptions.Password] = "myPassword", [DatabricksParameters.OperationStatusRequestTimeout] = "-1" }, typeof(ArgumentOutOfRangeException)));
             }
         }
 

csharp/test/Drivers/Databricks/E2E/DatabricksTestConfiguration.cs

@@ -55,6 +55,12 @@ public class DatabricksTestConfiguration : SparkTestConfiguration
         [JsonPropertyName("enableRunAsyncInThriftOp"), JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)]
         public string EnableRunAsyncInThriftOp { get; set; } = string.Empty;
 
+        [JsonPropertyName("fetchHeartbeatInterval"), JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)]
+        public string FetchHeartbeatInterval { get; set; } = string.Empty;
+
+        [JsonPropertyName("operationStatusRequestTimeout"), JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)]
+        public string OperationStatusRequestTimeout { get; set; } = string.Empty;
+
         [JsonPropertyName("enableDirectResults"), JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)]
         public string EnableDirectResults { get; set; } = string.Empty;
     }

csharp/test/Drivers/Databricks/E2E/DatabricksTestEnvironment.cs

@@ -141,6 +141,14 @@ public override Dictionary<string, string> GetDriverParameters(DatabricksTestCon
             {
                 parameters.Add(ApacheParameters.QueryTimeoutSeconds, testConfiguration.QueryTimeoutSeconds!);
             }
+            if (!string.IsNullOrEmpty(testConfiguration.FetchHeartbeatInterval))
+            {
+                parameters.Add(DatabricksParameters.FetchHeartbeatInterval, testConfiguration.FetchHeartbeatInterval!);
+            }
+            if (!string.IsNullOrEmpty(testConfiguration.OperationStatusRequestTimeout))
+            {
+                parameters.Add(DatabricksParameters.OperationStatusRequestTimeout, testConfiguration.OperationStatusRequestTimeout!);
+            }
             if (!string.IsNullOrEmpty(testConfiguration.EnableMultipleCatalogSupport))
             {
                 parameters.Add(DatabricksParameters.EnableMultipleCatalogSupport, testConfiguration.EnableMultipleCatalogSupport!);

csharp/test/Drivers/Databricks/Unit/DatabricksOperationStatusPollerTests.cs

@@ -190,5 +190,47 @@ public async Task StopsPollingOnException()
             {
             }
         }
+
+        [Fact]
+        public async Task UsesCustomHeartbeatInterval()
+        {
+            // Arrange
+            int customHeartbeatInterval = 2; // 2 seconds
+            using var poller = new DatabricksOperationStatusPoller(_mockStatement.Object, _mockResponse.Object, customHeartbeatInterval);
+            var pollCount = 0;
+            _mockClient.Setup(c => c.GetOperationStatus(It.IsAny<TGetOperationStatusReq>(), It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new TGetOperationStatusResp())
+                .Callback(() => pollCount++);
+
+            // Act
+            poller.Start();
+            await Task.Delay(TimeSpan.FromSeconds(customHeartbeatInterval * 3)); // Wait for 3 intervals
+
+            // Assert
+            Assert.True(pollCount > 0, "Should have polled at least once");
+            _mockClient.Verify(c => c.GetOperationStatus(It.IsAny<TGetOperationStatusReq>(), It.IsAny<CancellationToken>()), Times.AtLeastOnce);
+        }
+
+        [Fact]
+        public async Task UsesCustomRequestTimeout()
+        {
+            // Arrange
+            int customRequestTimeout = 5;
+            using var poller = new DatabricksOperationStatusPoller(_mockStatement.Object, _mockResponse.Object, _heartbeatIntervalSeconds, customRequestTimeout);
+            var pollCount = 0;
+
+            _mockClient.Setup(c => c.GetOperationStatus(It.IsAny<TGetOperationStatusReq>(), It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new TGetOperationStatusResp())
+                .Callback(() => pollCount++);
+
+            // Act
+            poller.Start();
+            await Task.Delay(TimeSpan.FromSeconds(_heartbeatIntervalSeconds * 2));
+
+            // Assert: This test verifies that the poller can be instantiated with a custom request timeout
+
+            Assert.True(pollCount > 0, "Should have polled at least once");
+            _mockClient.Verify(c => c.GetOperationStatus(It.IsAny<TGetOperationStatusReq>(), It.IsAny<CancellationToken>()), Times.AtLeastOnce);
+        }
     }
 }