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); + } } }