nit fix

nit fix

Author: toddmeng-db

csharp/src/Drivers/Databricks/BaseDatabricksReader.cs

@@ -16,8 +16,6 @@
 */
 
 using System;
-using System.Threading.Tasks;
-using System.Threading;
 using Apache.Arrow.Adbc.Tracing;
 
 namespace Apache.Arrow.Adbc.Drivers.Databricks

csharp/src/Drivers/Databricks/CloudFetch/CloudFetchReader.cs

@@ -180,6 +180,7 @@ protected override void Dispose(bool disposing)
                 this.downloadManager.Dispose();
                 this.downloadManager = null;
             }
+            base.Dispose(disposing);
         }
     }
 }

csharp/src/Drivers/Databricks/CloudFetch/CloudFetchResultFetcher.cs

@@ -22,6 +22,7 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Apache.Arrow.Adbc.Drivers.Apache;
 using Apache.Hive.Service.Rpc.Thrift;
 
 namespace Apache.Arrow.Adbc.Drivers.Databricks.CloudFetch
@@ -160,8 +161,11 @@ public async Task StopAsync()
 
                 request.StartRowOffset = offset;
 
+                // Cancelling mid-request breaks the client; Dispose() should not break the underlying client
+                CancellationToken expiringToken = ApacheUtility.GetCancellationToken(DatabricksConstants.DefaultCloudFetchRequestTimeoutSeconds, ApacheUtility.TimeUnit.Seconds);
+
                 // Fetch results
-                TFetchResultsResp response = await _statement.Client.FetchResults(request, cancellationToken);
+                TFetchResultsResp response = await _statement.Client.FetchResults(request, expiringToken);
 
                 // Process the results
                 if (response.Status.StatusCode == TStatusCode.SUCCESS_STATUS &&
@@ -247,10 +251,6 @@ private async Task FetchResultsAsync(CancellationToken cancellationToken)
                 _downloadQueue.Add(EndOfResultsGuard.Instance, cancellationToken);
                 _isCompleted = true;
             }
-            catch (OperationCanceledException)
-            {
-                // downloadQueue could be full, would block indefinitely
-            }
             catch (Exception ex)
             {
                 Debug.WriteLine($"Unhandled error in fetcher: {ex.Message}");
@@ -261,7 +261,7 @@ private async Task FetchResultsAsync(CancellationToken cancellationToken)
                 // Add the end of results guard to the queue even in case of error
                 try
                 {
-                    _downloadQueue.Add(EndOfResultsGuard.Instance, CancellationToken.None);
+                    _downloadQueue.TryAdd(EndOfResultsGuard.Instance, 0);
                 }
                 catch (Exception)
                 {
@@ -286,7 +286,8 @@ private async Task FetchNextResultBatchAsync(long? offset, CancellationToken can
             TFetchResultsResp response;
             try
             {
-                response = await _statement.Client.FetchResults(request, cancellationToken).ConfigureAwait(false);
+                CancellationToken expiringToken = ApacheUtility.GetCancellationToken(DatabricksConstants.DefaultCloudFetchRequestTimeoutSeconds, ApacheUtility.TimeUnit.Seconds);
+                response = await _statement.Client.FetchResults(request, expiringToken).ConfigureAwait(false);
             }
             catch (Exception ex)
             {

csharp/src/Drivers/Databricks/DatabricksCompositeReader.cs

@@ -130,6 +130,7 @@ protected override void Dispose(bool disposing)
                 _activeReader?.Dispose();
                 DisposeOperationStatusPoller();
             }
+            _activeReader = null;
             base.Dispose(disposing);
         }
 

csharp/src/Drivers/Databricks/DatabricksOperationStatusPoller.cs

@@ -101,7 +101,7 @@ public void Dispose()
                 try
                 {
                     if (_operationStatusPollingTask != null)
-                        _operationStatusPollingTask?.ConfigureAwait(false).GetAwaiter().GetResult();
+                        _operationStatusPollingTask?.GetAwaiter().GetResult();
                 }
                 catch (OperationCanceledException)
                 {
@@ -110,6 +110,8 @@ public void Dispose()
 
                 _internalCts.Dispose();
             }
+            _internalCts = null;
+            _operationStatusPollingTask = null;
         }
     }
 }

csharp/src/Drivers/Databricks/DatabricksParameters.cs

@@ -235,6 +235,11 @@ public class DatabricksConstants
         /// </summary>
         public const int DefaultOperationStatusRequestTimeoutSeconds = 30;
 
+        /// <summary>
+        /// Default timeout in seconds for CloudFetch requests
+        /// </summary>
+        public const int DefaultCloudFetchRequestTimeoutSeconds = 30;
+
         /// <summary>
         /// OAuth grant type constants
         /// </summary>

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

@@ -187,10 +187,14 @@ public async Task AllStatementTypesDisposeWithoutErrors(string statementType, st
                 var batch = await queryResult.Stream.ReadNextRecordBatchAsync();
                 // Note: batch might be null for empty results, that's OK
 
+                // test disposing the stream does not throw
+                var streamException = Record.Exception(() => queryResult.Stream.Dispose());
+                Assert.Null(streamException);
+
                 // The critical test: disposal should not throw any exceptions
                 // This specifically tests the fix for the GetColumns bug where _directResults wasn't set
-                var exception = Record.Exception(() => statement.Dispose());
-                Assert.Null(exception);
+                var statementException = Record.Exception(() => statement.Dispose());
+                Assert.Null(statementException);
             }
             catch (Exception ex)
             {