error handling for operation status poller

Author: toddmeng-db

csharp/src/Drivers/Databricks/BaseDatabricksReader.cs

@@ -16,7 +16,8 @@
 */
 
 using System;
-using Apache.Arrow.Adbc.Drivers.Apache;
+using System.Threading.Tasks;
+using System.Threading;
 using Apache.Arrow.Adbc.Tracing;
 
 namespace Apache.Arrow.Adbc.Drivers.Databricks
@@ -94,5 +95,35 @@ protected void ThrowIfDisposed()
         public override string AssemblyName => DatabricksConnection.s_assemblyName;
 
         public override string AssemblyVersion => DatabricksConnection.s_assemblyVersion;
+
+        public override async ValueTask<RecordBatch?> ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)
+        {
+            try
+            {
+                var result = await ReadNextRecordBatchInternalAsync(cancellationToken);
+                
+                // Stop the poller when we've reached the end of results
+                if (result == null)
+                {
+                    StopOperationStatusPoller();
+                }
+                
+                return result;
+            }
+            catch
+            {
+                // Stop the poller immediately on any exception to prevent unnecessary polling
+                StopOperationStatusPoller();
+                throw;
+            }
+        }
+
+        /// <summary>
+        /// Internal method that derived classes implement to read the next record batch.
+        /// This method is called by ReadNextRecordBatchAsync which provides exception handling.
+        /// </summary>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>The next record batch, or null if there are no more batches.</returns>
+        protected abstract ValueTask<RecordBatch?> ReadNextRecordBatchInternalAsync(CancellationToken cancellationToken);
     }
 }

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

@@ -81,7 +81,7 @@ public CloudFetchReader(DatabricksStatement statement, Schema schema, TFetchResu
         /// </summary>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns>The next record batch, or null if there are no more batches.</returns>
-        public override async ValueTask<RecordBatch?> ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)
+        protected override async ValueTask<RecordBatch?> ReadNextRecordBatchInternalAsync(CancellationToken cancellationToken = default)
         {
             return await this.TraceActivityAsync(async _ =>
             {
@@ -155,7 +155,6 @@ public CloudFetchReader(DatabricksStatement statement, Schema schema, TFetchResu
                         }
                     }
 
-                    StopOperationStatusPoller();
                     // If we get here, there are no more files
                     return null;
                 }

csharp/src/Drivers/Databricks/DatabricksReader.cs

@@ -47,7 +47,7 @@ public DatabricksReader(DatabricksStatement statement, Schema schema, TFetchResu
             }
         }
 
-        public override async ValueTask<RecordBatch?> ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)
+        protected override async ValueTask<RecordBatch?> ReadNextRecordBatchInternalAsync(CancellationToken cancellationToken = default)
         {
             return await this.TraceActivity(async activity =>
             {
@@ -77,7 +77,6 @@ public DatabricksReader(DatabricksStatement statement, Schema schema, TFetchResu
 
                     if (this.hasNoMoreRows)
                     {
-                        StopOperationStatusPoller();
                         return null;
                     }