feat(csharp/src/Drivers/Databricks): Implement CloudFetchUrlManager to handle presigned URL expiration in CloudFetch (#2855)

### Problem
The Databricks driver's CloudFetch functionality was not properly
handling expired cloud file URLs, which could lead to failed downloads
and errors during query execution. The system needed a way to track,
cache, and refresh presigned URLs before they expire.

### Solution
- Improve `CloudFetchResultFetcher` class that:
  - Manages a cache of cloud file URLs with their expiration times
  - Proactively refreshes URLs that are about to expire
  - Provides thread-safe access to URL information
- Added an `IClock` interface and implementations to facilitate testing
with controlled time
- Extended the `IDownloadResult` interface to support URL refreshing and
expiration checking
- Updated namespace from
`Apache.Arrow.Adbc.Drivers.Apache.Databricks.CloudFetch` to
`Apache.Arrow.Adbc.Drivers.Databricks.CloudFetch` for better
organization

Author: jadewang-db

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

@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+using System;
+
+namespace Apache.Arrow.Adbc.Drivers.Databricks.CloudFetch
+{
+    /// <summary>
+    /// Abstraction for time operations to enable testing with controlled time.
+    /// </summary>
+    internal interface IClock
+    {
+        /// <summary>
+        /// Gets the current UTC time.
+        /// </summary>
+        DateTime UtcNow { get; }
+    }
+
+    /// <summary>
+    /// Default implementation that uses system time.
+    /// </summary>
+    internal class SystemClock : IClock
+    {
+        public DateTime UtcNow => DateTime.UtcNow;
+    }
+}

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

@@ -22,9 +22,8 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Apache.Arrow.Adbc.Drivers.Apache.Hive2;
-using Apache.Arrow.Adbc.Drivers.Databricks;
 
-namespace Apache.Arrow.Adbc.Drivers.Apache.Databricks.CloudFetch
+namespace Apache.Arrow.Adbc.Drivers.Databricks.CloudFetch
 {
     /// <summary>
     /// Manages the CloudFetch download pipeline.
@@ -38,6 +37,8 @@ internal sealed class CloudFetchDownloadManager : ICloudFetchDownloadManager
         private const bool DefaultPrefetchEnabled = true;
         private const int DefaultFetchBatchSize = 2000000;
         private const int DefaultTimeoutMinutes = 5;
+        private const int DefaultMaxUrlRefreshAttempts = 3;
+        private const int DefaultUrlExpirationBufferSeconds = 60;
 
         private readonly DatabricksStatement _statement;
         private readonly Schema _schema;
@@ -151,6 +152,34 @@ public CloudFetchDownloadManager(DatabricksStatement statement, Schema schema, b
                 }
             }
 
+            // Parse URL expiration buffer seconds
+            int urlExpirationBufferSeconds = DefaultUrlExpirationBufferSeconds;
+            if (connectionProps.TryGetValue(DatabricksParameters.CloudFetchUrlExpirationBufferSeconds, out string? urlExpirationBufferStr))
+            {
+                if (int.TryParse(urlExpirationBufferStr, out int parsedUrlExpirationBuffer) && parsedUrlExpirationBuffer > 0)
+                {
+                    urlExpirationBufferSeconds = parsedUrlExpirationBuffer;
+                }
+                else
+                {
+                    throw new ArgumentException($"Invalid value for {DatabricksParameters.CloudFetchUrlExpirationBufferSeconds}: {urlExpirationBufferStr}. Expected a positive integer.");
+                }
+            }
+
+            // Parse max URL refresh attempts
+            int maxUrlRefreshAttempts = DefaultMaxUrlRefreshAttempts;
+            if (connectionProps.TryGetValue(DatabricksParameters.CloudFetchMaxUrlRefreshAttempts, out string? maxUrlRefreshAttemptsStr))
+            {
+                if (int.TryParse(maxUrlRefreshAttemptsStr, out int parsedMaxUrlRefreshAttempts) && parsedMaxUrlRefreshAttempts > 0)
+                {
+                    maxUrlRefreshAttempts = parsedMaxUrlRefreshAttempts;
+                }
+                else
+                {
+                    throw new ArgumentException($"Invalid value for {DatabricksParameters.CloudFetchMaxUrlRefreshAttempts}: {maxUrlRefreshAttemptsStr}. Expected a positive integer.");
+                }
+            }
+
             // Initialize the memory manager
             _memoryManager = new CloudFetchMemoryBufferManager(memoryBufferSizeMB);
 
@@ -161,23 +190,27 @@ public CloudFetchDownloadManager(DatabricksStatement statement, Schema schema, b
             _httpClient = httpClient;
             _httpClient.Timeout = TimeSpan.FromMinutes(timeoutMinutes);
 
-            // Initialize the result fetcher
+            // Initialize the result fetcher with URL management capabilities
             _resultFetcher = new CloudFetchResultFetcher(
                 _statement,
                 _memoryManager,
                 _downloadQueue,
-                DefaultFetchBatchSize);
+                DefaultFetchBatchSize,
+                urlExpirationBufferSeconds);
 
             // Initialize the downloader
             _downloader = new CloudFetchDownloader(
                 _downloadQueue,
                 _resultQueue,
                 _memoryManager,
                 _httpClient,
+                _resultFetcher,
                 parallelDownloads,
                 _isLz4Compressed,
                 maxRetries,
-                retryDelayMs);
+                retryDelayMs,
+                maxUrlRefreshAttempts,
+                urlExpirationBufferSeconds);
         }
 
         /// <summary>

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

@@ -24,7 +24,7 @@
 using System.Threading.Tasks;
 using K4os.Compression.LZ4.Streams;
 
-namespace Apache.Arrow.Adbc.Drivers.Apache.Databricks.CloudFetch
+namespace Apache.Arrow.Adbc.Drivers.Databricks.CloudFetch
 {
     /// <summary>
     /// Downloads files from URLs.
@@ -35,10 +35,13 @@ internal sealed class CloudFetchDownloader : ICloudFetchDownloader
         private readonly BlockingCollection<IDownloadResult> _resultQueue;
         private readonly ICloudFetchMemoryBufferManager _memoryManager;
         private readonly HttpClient _httpClient;
+        private readonly ICloudFetchResultFetcher _resultFetcher;
         private readonly int _maxParallelDownloads;
         private readonly bool _isLz4Compressed;
         private readonly int _maxRetries;
         private readonly int _retryDelayMs;
+        private readonly int _maxUrlRefreshAttempts;
+        private readonly int _urlExpirationBufferSeconds;
         private readonly SemaphoreSlim _downloadSemaphore;
         private Task? _downloadTask;
         private CancellationTokenSource? _cancellationTokenSource;
@@ -53,29 +56,37 @@ internal sealed class CloudFetchDownloader : ICloudFetchDownloader
         /// <param name="resultQueue">The queue to add completed downloads to.</param>
         /// <param name="memoryManager">The memory buffer manager.</param>
         /// <param name="httpClient">The HTTP client to use for downloads.</param>
+        /// <param name="resultFetcher">The result fetcher that manages URLs.</param>
         /// <param name="maxParallelDownloads">The maximum number of parallel downloads.</param>
         /// <param name="isLz4Compressed">Whether the results are LZ4 compressed.</param>
-        /// <param name="logger">The logger instance.</param>
         /// <param name="maxRetries">The maximum number of retry attempts.</param>
         /// <param name="retryDelayMs">The delay between retry attempts in milliseconds.</param>
+        /// <param name="maxUrlRefreshAttempts">The maximum number of URL refresh attempts.</param>
+        /// <param name="urlExpirationBufferSeconds">Buffer time in seconds before URL expiration to trigger refresh.</param>
         public CloudFetchDownloader(
             BlockingCollection<IDownloadResult> downloadQueue,
             BlockingCollection<IDownloadResult> resultQueue,
             ICloudFetchMemoryBufferManager memoryManager,
             HttpClient httpClient,
+            ICloudFetchResultFetcher resultFetcher,
             int maxParallelDownloads,
             bool isLz4Compressed,
             int maxRetries = 3,
-            int retryDelayMs = 500)
+            int retryDelayMs = 500,
+            int maxUrlRefreshAttempts = 3,
+            int urlExpirationBufferSeconds = 60)
         {
             _downloadQueue = downloadQueue ?? throw new ArgumentNullException(nameof(downloadQueue));
             _resultQueue = resultQueue ?? throw new ArgumentNullException(nameof(resultQueue));
             _memoryManager = memoryManager ?? throw new ArgumentNullException(nameof(memoryManager));
             _httpClient = httpClient ?? throw new ArgumentNullException(nameof(httpClient));
+            _resultFetcher = resultFetcher ?? throw new ArgumentNullException(nameof(resultFetcher));
             _maxParallelDownloads = maxParallelDownloads > 0 ? maxParallelDownloads : throw new ArgumentOutOfRangeException(nameof(maxParallelDownloads));
             _isLz4Compressed = isLz4Compressed;
             _maxRetries = maxRetries > 0 ? maxRetries : throw new ArgumentOutOfRangeException(nameof(maxRetries));
             _retryDelayMs = retryDelayMs > 0 ? retryDelayMs : throw new ArgumentOutOfRangeException(nameof(retryDelayMs));
+            _maxUrlRefreshAttempts = maxUrlRefreshAttempts > 0 ? maxUrlRefreshAttempts : throw new ArgumentOutOfRangeException(nameof(maxUrlRefreshAttempts));
+            _urlExpirationBufferSeconds = urlExpirationBufferSeconds > 0 ? urlExpirationBufferSeconds : throw new ArgumentOutOfRangeException(nameof(urlExpirationBufferSeconds));
             _downloadSemaphore = new SemaphoreSlim(_maxParallelDownloads, _maxParallelDownloads);
             _isCompleted = false;
         }
@@ -237,6 +248,19 @@ private async Task DownloadFilesAsync(CancellationToken cancellationToken)
                         break;
                     }
 
+                    // Check if the URL is expired or about to expire
+                    if (downloadResult.IsExpiredOrExpiringSoon(_urlExpirationBufferSeconds))
+                    {
+                        // Get a refreshed URL before starting the download
+                        var refreshedLink = await _resultFetcher.GetUrlAsync(downloadResult.Link.StartRowOffset, cancellationToken);
+                        if (refreshedLink != null)
+                        {
+                            // Update the download result with the refreshed link
+                            downloadResult.UpdateWithRefreshedLink(refreshedLink);
+                            Trace.TraceInformation($"Updated URL for file at offset {refreshedLink.StartRowOffset} before download");
+                        }
+                    }
+
                     // Acquire a download slot
                     await _downloadSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
 
@@ -341,6 +365,37 @@ private async Task DownloadFileAsync(IDownloadResult downloadResult, Cancellatio
                         HttpCompletionOption.ResponseHeadersRead,
                         cancellationToken).ConfigureAwait(false);
 
+                    // Check if the response indicates an expired URL (typically 403 or 401)
+                    if (response.StatusCode == System.Net.HttpStatusCode.Forbidden ||
+                        response.StatusCode == System.Net.HttpStatusCode.Unauthorized)
+                    {
+                        // If we've already tried refreshing too many times, fail
+                        if (downloadResult.RefreshAttempts >= _maxUrlRefreshAttempts)
+                        {
+                            throw new InvalidOperationException($"Failed to download file after {downloadResult.RefreshAttempts} URL refresh attempts.");
+                        }
+
+                        // Try to refresh the URL
+                        var refreshedLink = await _resultFetcher.GetUrlAsync(downloadResult.Link.StartRowOffset, cancellationToken);
+                        if (refreshedLink != null)
+                        {
+                            // Update the download result with the refreshed link
+                            downloadResult.UpdateWithRefreshedLink(refreshedLink);
+                            url = refreshedLink.FileLink;
+                            sanitizedUrl = SanitizeUrl(url);
+
+                            Trace.TraceInformation($"URL for file at offset {refreshedLink.StartRowOffset} was refreshed after expired URL response");
+
+                            // Continue to the next retry attempt with the refreshed URL
+                            continue;
+                        }
+                        else
+                        {
+                            // If refresh failed, throw an exception
+                            throw new InvalidOperationException("Failed to refresh expired URL.");
+                        }
+                    }
+
                     response.EnsureSuccessStatusCode();
 
                     // Log the download size if available from response headers

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

@@ -19,7 +19,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 
-namespace Apache.Arrow.Adbc.Drivers.Apache.Databricks.CloudFetch
+namespace Apache.Arrow.Adbc.Drivers.Databricks.CloudFetch
 {
     /// <summary>
     /// Manages memory allocation for prefetched files.

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

@@ -21,8 +21,6 @@
 using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
-using Apache.Arrow.Adbc.Drivers.Apache.Databricks.CloudFetch;
-using Apache.Arrow.Adbc.Drivers.Databricks;
 using Apache.Arrow.Ipc;
 using Apache.Hive.Service.Rpc.Thrift;