Add native ADBC FlightSQL driver support for osx-arm64, fix ADBC resources disposal

.github/workflows/pr.yaml

@@ -76,7 +76,6 @@ jobs:
       - name: Build
         run: dotnet build --no-restore --configuration Release
       - name: Test
-        continue-on-error: ${{ startsWith(matrix.os, 'macos') }}
         env:
           SCP_SPICEAI_TPCH_API_KEY: ${{ secrets.SCP_SPICEAI_TPCH_API_KEY }}
         run: dotnet test --no-build --verbosity normal --configuration Release --framework ${{ matrix.dotnet-version == '8.0.x' && 'net8.0' || matrix.dotnet-version == '9.0.x' && 'net9.0' || 'net10.0' }}

Spice/Spice.csproj

@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
 
     <PropertyGroup>
         <PackageId>SpiceAI</PackageId>
@@ -28,6 +28,18 @@
         <None Include="..\README.md" Pack="true" PackagePath="\" />
     </ItemGroup>
 
+    <!-- 
+        Include native ADBC FlightSQL driver for osx-arm64.
+        This platform is not included in the upstream Apache.Arrow.Adbc.Drivers.Interop.FlightSql NuGet package.
+        NuGet/MSBuild automatically copies the correct RID-specific binary to output at build time.
+
+        Native library source: Apache Arrow ADBC 1.9.0
+        SHA-256 (wheel): edee43183625acb2fd2e2b8331f07ac516b6958ef1f27663f6308334db7f7272
+    -->
+    <ItemGroup>
+        <None Include="runtimes\**\*" Pack="true" PackagePath="runtimes" CopyToOutputDirectory="PreserveNewest" Link="runtimes\%(RecursiveDir)%(Filename)%(Extension)" />
+    </ItemGroup>
+
     <ItemGroup>
         <PackageReference Include="Apache.Arrow" Version="22.1.0" />
         <PackageReference Include="Apache.Arrow.Adbc" Version="0.21.0" />
@@ -43,5 +55,4 @@
         </AssemblyAttribute>
     </ItemGroup>
 
-
-</Project>
+</Project>

Spice/src/Adbc/SpiceAdbcClient.cs

@@ -20,9 +20,9 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 */
 
+using System.Runtime.InteropServices;
 using Apache.Arrow;
 using Apache.Arrow.Adbc;
-using Apache.Arrow.Adbc.Drivers.Interop.FlightSql;
 using Apache.Arrow.Ipc;
 using Apache.Arrow.Types;
 using Polly.Retry;
@@ -114,23 +114,36 @@ private void InitializeIfNeeded()
                 return;
             }
 
-            // Format the URI for ADBC FlightSQL Go driver
-            // The Go-based driver handles grpc/grpc+tls schemes natively
+            // Format the URI for ADBC FlightSQL driver
+            // The driver expects grpc:// or grpc+tls:// schemes
             var uri = _flightAddress;
 
-            // Ensure proper scheme if not already present
-            if (!uri.StartsWith("grpc://", StringComparison.OrdinalIgnoreCase) &&
-                !uri.StartsWith("grpc+tls://", StringComparison.OrdinalIgnoreCase) &&
-                !uri.StartsWith("grpc+tcp://", StringComparison.OrdinalIgnoreCase) &&
-                !uri.StartsWith("http://", StringComparison.OrdinalIgnoreCase) &&
-                !uri.StartsWith("https://", StringComparison.OrdinalIgnoreCase))
+            // Convert http/https schemes to grpc/grpc+tls
+            if (uri.StartsWith("https://", StringComparison.OrdinalIgnoreCase))
+            {
+#if NETSTANDARD2_0
+                uri = "grpc+tls://" + uri.Substring(8);
+#else
+                uri = string.Concat("grpc+tls://", uri.AsSpan(8));
+#endif
+            }
+            else if (uri.StartsWith("http://", StringComparison.OrdinalIgnoreCase))
+            {
+#if NETSTANDARD2_0
+                uri = "grpc://" + uri.Substring(7);
+#else
+                uri = string.Concat("grpc://", uri.AsSpan(7));
+#endif
+            }
+            else if (!uri.StartsWith("grpc://", StringComparison.OrdinalIgnoreCase) &&
+                     !uri.StartsWith("grpc+tls://", StringComparison.OrdinalIgnoreCase) &&
+                     !uri.StartsWith("grpc+tcp://", StringComparison.OrdinalIgnoreCase))
             {
                 // No scheme provided - add grpc or grpc+tls based on TLS setting
                 uri = _useTls ? string.Concat("grpc+tls://", uri) : string.Concat("grpc://", uri);
             }
 
-            // Build database parameters for the Go driver
-            // The Go driver uses "uri" as the connection parameter
+            // Build database parameters for the ADBC FlightSQL driver
             var databaseParams = new Dictionary<string, string>
             {
                 { "uri", uri }
@@ -143,11 +156,12 @@ private void InitializeIfNeeded()
                 databaseParams["password"] = _apiKey!;
             }
 
-            // Add user agent header using the Go driver's header prefix
+            // Add user agent header
             databaseParams["adbc.flight.sql.rpc.call_header.user-agent"] = UserAgentHelper.BuildUserAgent(_userAgent);
 
-            // Create the Go-based interop driver and database
-            var driver = FlightSqlDriverLoader.LoadDriver();
+            // Load the ADBC FlightSQL driver from the application base directory
+            var driverPath = ResolveNativeDriverPath();
+            var driver = AdbcDriverLoader.LoadDriver(driverPath, "AdbcDriverFlightsqlInit");
             _database = driver.Open(databaseParams);
             _connection = _database.Connect(new Dictionary<string, string>());
         }
@@ -170,22 +184,34 @@ private void InitializeIfNeeded()
         {
             InitializeIfNeeded();
 
-            using var statement = _connection!.CreateStatement();
-            statement.SqlQuery = sql;
+            var statement = _connection!.CreateStatement();
+            try
+            {
+                statement.SqlQuery = sql;
+
+                // Prepare the statement
+                statement.Prepare();
 
-            // Prepare the statement
-            statement.Prepare();
+                // Bind parameters if provided
+                if (parameters.Length > 0)
+                {
+                    var parameterBatch = CreateParameterBatch(parameters);
+                    statement.Bind(parameterBatch, parameterBatch.Schema);
+                }
 
-            // Bind parameters if provided
-            if (parameters.Length > 0)
+                // Execute the query
+                var result = statement.ExecuteQuery();
+
+                // Wrap the stream to keep the statement alive for the stream's lifetime
+                var wrappedStream = new StatementBoundArrowArrayStream(result.Stream!, statement);
+                return Task.FromResult<IArrowArrayStream?>(wrappedStream);
+            }
+            catch
             {
-                var parameterBatch = CreateParameterBatch(parameters);
-                statement.Bind(parameterBatch, parameterBatch.Schema);
+                // If anything fails before we wrap the stream, dispose the statement
+                statement.Dispose();
+                throw;
             }
-
-            // Execute the query
-            var result = statement.ExecuteQuery();
-            return Task.FromResult(result.Stream);
         });
     }
 
@@ -513,6 +539,61 @@ private static Decimal256Array CreateDecimal256Array(object? value, Decimal256Ty
 
     // ============ Disposal ============
 
+    /// <summary>
+    /// Resolves the path to the native ADBC FlightSQL driver based on the current platform.
+    /// Uses standard .NET conventions: AppContext.BaseDirectory and runtimes/{rid}/native/ layout.
+    /// </summary>
+    private static string ResolveNativeDriverPath()
+    {
+        var rid = GetRuntimeIdentifier();
+        var driverFileName = GetDriverFileName();
+        var driverPath = Path.Combine(AppContext.BaseDirectory, "runtimes", rid, "native", driverFileName);
+
+        if (!File.Exists(driverPath))
+        {
+            throw new FileNotFoundException(
+                $"Could not find native ADBC FlightSQL driver at {driverPath}. " +
+                $"Ensure the SpiceAI NuGet package is properly installed.",
+                driverPath);
+        }
+
+        return driverPath;
+    }
+
+    /// <summary>
+    /// Gets the runtime identifier for the current platform.
+    /// </summary>
+    private static string GetRuntimeIdentifier()
+    {
+        var arch = RuntimeInformation.OSArchitecture.ToString().ToLowerInvariant();
+
+        if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
+            return $"win-{arch}";
+        if (RuntimeInformation.IsOSPlatform(OSPlatform.Linux))
+            return $"linux-{arch}";
+        if (RuntimeInformation.IsOSPlatform(OSPlatform.OSX))
+            return $"osx-{arch}";
+
+        throw new PlatformNotSupportedException(
+            $"Unsupported platform: {RuntimeInformation.OSDescription}");
+    }
+
+    /// <summary>
+    /// Gets the platform-specific driver filename.
+    /// </summary>
+    private static string GetDriverFileName()
+    {
+        if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
+            return "libadbc_driver_flightsql.dll";
+        if (RuntimeInformation.IsOSPlatform(OSPlatform.Linux))
+            return "libadbc_driver_flightsql.so";
+        if (RuntimeInformation.IsOSPlatform(OSPlatform.OSX))
+            return "libadbc_driver_flightsql.dylib";
+
+        throw new PlatformNotSupportedException(
+            $"Unsupported platform: {RuntimeInformation.OSDescription}");
+    }
+
     private bool _disposed;
 
     public void Dispose()

Spice/src/Adbc/StatementBoundArrowArrayStream.cs

@@ -0,0 +1,98 @@
+/*
+Copyright 2024 The Spice.ai OSS Authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*/
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Apache.Arrow;
+using Apache.Arrow.Adbc;
+using Apache.Arrow.Ipc;
+
+namespace Spice.Adbc;
+
+/// <summary>
+/// A wrapper around IArrowArrayStream that keeps the ADBC statement alive
+/// for the lifetime of the stream. When this wrapper is disposed, it disposes
+/// both the underlying stream and the statement in the correct order.
+/// </summary>
+internal sealed class StatementBoundArrowArrayStream : IArrowArrayStream
+{
+    private readonly IArrowArrayStream _innerStream;
+    private readonly AdbcStatement _statement;
+    private bool _disposed;
+
+    /// <summary>
+    /// Creates a new StatementBoundArrowArrayStream that wraps the given stream
+    /// and keeps the statement alive until disposal.
+    /// </summary>
+    /// <param name="innerStream">The underlying Arrow array stream</param>
+    /// <param name="statement">The ADBC statement to keep alive and dispose with the stream</param>
+    public StatementBoundArrowArrayStream(IArrowArrayStream innerStream, AdbcStatement statement)
+    {
+        _innerStream = innerStream ?? throw new ArgumentNullException(nameof(innerStream));
+        _statement = statement ?? throw new ArgumentNullException(nameof(statement));
+    }
+
+    /// <inheritdoc/>
+    public Schema Schema
+    {
+        get
+        {
+#if NET8_0_OR_GREATER
+            ObjectDisposedException.ThrowIf(_disposed, this);
+#else
+            if (_disposed) throw new ObjectDisposedException(GetType().FullName);
+#endif
+            return _innerStream.Schema;
+        }
+    }
+
+    /// <inheritdoc/>
+    public ValueTask<RecordBatch?> ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)
+    {
+#if NET8_0_OR_GREATER
+        ObjectDisposedException.ThrowIf(_disposed, this);
+#else
+        if (_disposed) throw new ObjectDisposedException(GetType().FullName);
+#endif
+        return _innerStream.ReadNextRecordBatchAsync(cancellationToken);
+    }
+
+    /// <inheritdoc/>
+    public void Dispose()
+    {
+        if (_disposed) return;
+        _disposed = true;
+
+        // Dispose in the correct order:
+        // 1. First dispose the stream (finishes reading/closes the stream)
+        // 2. Then dispose the statement (releases server-side resources)
+        try
+        {
+            _innerStream.Dispose();
+        }
+        finally
+        {
+            _statement.Dispose();
+        }
+    }
+}