diff --git a/src/Apache.Arrow.Flight/Client/FlightClient.cs b/src/Apache.Arrow.Flight/Client/FlightClient.cs
index d0942669..c8dc82f3 100644
--- a/src/Apache.Arrow.Flight/Client/FlightClient.cs
+++ b/src/Apache.Arrow.Flight/Client/FlightClient.cs
@@ -27,15 +27,42 @@ public class FlightClient
internal static readonly Empty EmptyInstance = new Empty();
private readonly FlightService.FlightServiceClient _client;
+ private readonly ArrowContext _context;
public FlightClient(ChannelBase grpcChannel)
+ : this(grpcChannel, null)
+ {
+ }
+
+ public FlightClient(ChannelBase grpcChannel, ArrowContext context)
{
_client = new FlightService.FlightServiceClient(grpcChannel);
+ _context = context;
}
public FlightClient(CallInvoker callInvoker)
+ : this(callInvoker, null)
+ {
+ }
+
+ private FlightClient(CallInvoker callInvoker, ArrowContext context)
{
_client = new FlightService.FlightServiceClient(callInvoker);
+ _context = context;
+ }
+
+ ///
+ /// Creates a from a with an .
+ ///
+ ///
+ /// A factory method is used instead of a constructor overload so that dependency injection
+ /// (e.g. Grpc.Net.ClientFactory) can still resolve the single-argument constructor unambiguously.
+ ///
+ /// The to use for gRPC calls
+ /// The carrying the compression codec factory and other reader configuration
+ public static FlightClient Create(CallInvoker callInvoker, ArrowContext context)
+ {
+ return new FlightClient(callInvoker, context);
}
public AsyncServerStreamingCall ListFlights(FlightCriteria criteria = null, Metadata headers = null)
@@ -77,7 +104,7 @@ public FlightRecordBatchStreamingCall GetStream(FlightTicket ticket, Metadata he
public FlightRecordBatchStreamingCall GetStream(FlightTicket ticket, Metadata headers, System.DateTime? deadline, CancellationToken cancellationToken = default)
{
var stream = _client.DoGet(ticket.ToProtocol(), headers, deadline, cancellationToken);
- var responseStream = new FlightClientRecordBatchStreamReader(stream.ResponseStream);
+ var responseStream = new FlightClientRecordBatchStreamReader(stream.ResponseStream, _context);
return new FlightRecordBatchStreamingCall(responseStream, stream.ResponseHeadersAsync, stream.GetStatus, stream.GetTrailers, stream.Dispose);
}
@@ -207,7 +234,7 @@ public FlightRecordBatchExchangeCall DoExchange(FlightDescriptor flightDescripto
{
var channel = _client.DoExchange(headers, deadline, cancellationToken);
var requestStream = new FlightClientRecordBatchStreamWriter(channel.RequestStream, flightDescriptor);
- var responseStream = new FlightClientRecordBatchStreamReader(channel.ResponseStream);
+ var responseStream = new FlightClientRecordBatchStreamReader(channel.ResponseStream, _context);
var call = new FlightRecordBatchExchangeCall(
requestStream,
responseStream,
diff --git a/src/Apache.Arrow.Flight/Client/FlightClientRecordBatchStreamReader.cs b/src/Apache.Arrow.Flight/Client/FlightClientRecordBatchStreamReader.cs
index 65eb6b65..fc0bd898 100644
--- a/src/Apache.Arrow.Flight/Client/FlightClientRecordBatchStreamReader.cs
+++ b/src/Apache.Arrow.Flight/Client/FlightClientRecordBatchStreamReader.cs
@@ -21,7 +21,7 @@ namespace Apache.Arrow.Flight.Client
{
public class FlightClientRecordBatchStreamReader : FlightRecordBatchStreamReader
{
- internal FlightClientRecordBatchStreamReader(IAsyncStreamReader flightDataStream) : base(flightDataStream)
+ internal FlightClientRecordBatchStreamReader(IAsyncStreamReader flightDataStream, ArrowContext context = null) : base(flightDataStream, context)
{
}
}
diff --git a/src/Apache.Arrow.Flight/FlightRecordBatchStreamReader.cs b/src/Apache.Arrow.Flight/FlightRecordBatchStreamReader.cs
index 7400ec15..7336b853 100644
--- a/src/Apache.Arrow.Flight/FlightRecordBatchStreamReader.cs
+++ b/src/Apache.Arrow.Flight/FlightRecordBatchStreamReader.cs
@@ -40,9 +40,9 @@ public abstract class FlightRecordBatchStreamReader : IAsyncStreamReader flightDataStream)
+ private protected FlightRecordBatchStreamReader(IAsyncStreamReader flightDataStream, ArrowContext context = null)
{
- _arrowReaderImplementation = new RecordBatchReaderImplementation(flightDataStream);
+ _arrowReaderImplementation = new RecordBatchReaderImplementation(flightDataStream, context);
}
public ValueTask Schema => _arrowReaderImplementation.GetSchemaAsync();
diff --git a/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs b/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs
index cb64867e..a320f7b3 100644
--- a/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs
+++ b/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs
@@ -31,7 +31,8 @@ internal class RecordBatchReaderImplementation : ArrowReaderImplementation
private FlightDescriptor _flightDescriptor;
private readonly List _applicationMetadatas;
- public RecordBatchReaderImplementation(IAsyncStreamReader streamReader)
+ public RecordBatchReaderImplementation(IAsyncStreamReader streamReader, ArrowContext context = null)
+ : base(context?.Allocator, context?.CompressionCodecFactory, context?.ExtensionRegistry)
{
_flightDataStream = streamReader;
_applicationMetadatas = new List();
diff --git a/src/Apache.Arrow.Flight/Server/FlightServerRecordBatchStreamReader.cs b/src/Apache.Arrow.Flight/Server/FlightServerRecordBatchStreamReader.cs
index c52b761a..ca93610f 100644
--- a/src/Apache.Arrow.Flight/Server/FlightServerRecordBatchStreamReader.cs
+++ b/src/Apache.Arrow.Flight/Server/FlightServerRecordBatchStreamReader.cs
@@ -21,11 +21,16 @@ namespace Apache.Arrow.Flight.Server
{
public class FlightServerRecordBatchStreamReader : FlightRecordBatchStreamReader
{
- public FlightServerRecordBatchStreamReader(IAsyncStreamReader flightDataStream) : base(new StreamReader(flightDataStream, data => data.ToProtocol()))
+ public FlightServerRecordBatchStreamReader(IAsyncStreamReader flightDataStream)
+ : this(flightDataStream, null)
{
}
- internal FlightServerRecordBatchStreamReader(IAsyncStreamReader flightDataStream) : base(flightDataStream)
+ public FlightServerRecordBatchStreamReader(IAsyncStreamReader flightDataStream, ArrowContext context) : base(new StreamReader(flightDataStream, data => data.ToProtocol()), context)
+ {
+ }
+
+ internal FlightServerRecordBatchStreamReader(IAsyncStreamReader flightDataStream, ArrowContext context = null) : base(flightDataStream, context)
{
}
diff --git a/test/Apache.Arrow.Flight.Tests/FlightTests.cs b/test/Apache.Arrow.Flight.Tests/FlightTests.cs
index 8c692115..46662cce 100644
--- a/test/Apache.Arrow.Flight.Tests/FlightTests.cs
+++ b/test/Apache.Arrow.Flight.Tests/FlightTests.cs
@@ -577,5 +577,110 @@ public async Task TestIntegrationWithGrpcNetClientFactory()
SchemaComparer.Compare(expectedSchema, actualSchema);
}
+
+ [Fact]
+ public async Task TestIntegrationWithGrpcNetClientFactoryAndArrowContext()
+ {
+ IServiceCollection services = new ServiceCollection();
+
+ services.AddGrpcClient(grpc => grpc.Address = new Uri(_testWebFactory.GetAddress()))
+ .ConfigureGrpcClientCreator(invoker =>
+ {
+ return FlightClient.Create(invoker, new ArrowContext());
+ });
+
+ IServiceProvider provider = services.BuildServiceProvider();
+
+ // Test that a FlightClient carrying an ArrowContext can be created via the
+ // Grpc.Net.ClientFactory library using the FlightClient.Create factory method.
+ FlightClient flightClient = provider.GetRequiredService();
+
+ // Test that the resolved client is functional.
+ var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test-factory-context");
+ var expectedBatch = CreateTestBatch(0, 100);
+ var expectedSchema = expectedBatch.Schema;
+
+ GivenStoreBatches(flightDescriptor, new RecordBatchWithMetadata(expectedBatch));
+
+ var actualSchema = await flightClient.GetSchema(flightDescriptor);
+
+ SchemaComparer.Compare(expectedSchema, actualSchema);
+ }
+
+ [Fact]
+ public async Task TestGetWithArrowContext()
+ {
+ // Verify that FlightClient works when constructed with an ArrowContext
+ var context = new ArrowContext();
+ var flightClient = new FlightClient(_testWebFactory.GetChannel(), context);
+
+ var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test-context");
+ var expectedBatch = CreateTestBatch(0, 100);
+
+ GivenStoreBatches(flightDescriptor, new RecordBatchWithMetadata(expectedBatch));
+
+ var flightInfo = await flightClient.GetInfo(flightDescriptor);
+ var endpoint = flightInfo.Endpoints.FirstOrDefault();
+
+ var getStream = flightClient.GetStream(endpoint.Ticket);
+ var resultList = await getStream.ResponseStream.ToListAsync();
+
+ Assert.Single(resultList);
+ ArrowReaderVerifier.CompareBatches(expectedBatch, resultList[0]);
+ }
+
+ [Fact]
+ public async Task TestGetWithNullArrowContext()
+ {
+ // Verify that FlightClient works with null ArrowContext (backward compat)
+ var flightClient = new FlightClient(_testWebFactory.GetChannel(), null);
+
+ var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test-null-context");
+ var expectedBatch = CreateTestBatch(0, 100);
+
+ GivenStoreBatches(flightDescriptor, new RecordBatchWithMetadata(expectedBatch));
+
+ var flightInfo = await flightClient.GetInfo(flightDescriptor);
+ var endpoint = flightInfo.Endpoints.FirstOrDefault();
+
+ var getStream = flightClient.GetStream(endpoint.Ticket);
+ var resultList = await getStream.ResponseStream.ToListAsync();
+
+ Assert.Single(resultList);
+ ArrowReaderVerifier.CompareBatches(expectedBatch, resultList[0]);
+ }
+
+ [Fact]
+ public async Task TestPutAndGetWithArrowContext()
+ {
+ // Verify put + get round-trip works with ArrowContext
+ var context = new ArrowContext();
+ var flightClient = new FlightClient(_testWebFactory.GetChannel(), context);
+
+ var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test-context-roundtrip");
+ var expectedBatch = CreateTestBatch(0, 100);
+
+ var putStream = await flightClient.StartPut(flightDescriptor, expectedBatch.Schema);
+ await putStream.RequestStream.WriteAsync(expectedBatch);
+ await putStream.RequestStream.CompleteAsync();
+ await putStream.ResponseStream.ToListAsync();
+
+ var flightInfo = await flightClient.GetInfo(flightDescriptor);
+ var endpoint = flightInfo.Endpoints.FirstOrDefault();
+
+ var getStream = flightClient.GetStream(endpoint.Ticket);
+ var resultList = await getStream.ResponseStream.ToListAsync();
+
+ Assert.Single(resultList);
+ ArrowReaderVerifier.CompareBatches(expectedBatch, resultList[0]);
+ }
+
+ [Fact]
+ public void TestFlightClientDefaultConstructorStillWorks()
+ {
+ // Verify the original constructor without ArrowContext still works
+ var client = new FlightClient(_testWebFactory.GetChannel());
+ Assert.NotNull(client);
+ }
}
}