public async Task TestGetFlightMetadata()
{
var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test");
var expectedBatch1 = CreateTestBatch(0, 100);
var expectedMetadata = ByteString.CopyFromUtf8("test metadata");
var expectedMetadataList = new List <ByteString>()
{
expectedMetadata
};
//Add batch to the in memory store
GivenStoreBatches(flightDescriptor, new RecordBatchWithMetadata(expectedBatch1, expectedMetadata));
//Get the flight info for the ticket
var flightInfo = await _flightClient.GetInfo(flightDescriptor);
Assert.Single(flightInfo.Endpoints);
var endpoint = flightInfo.Endpoints.FirstOrDefault();
var getStream = _flightClient.GetStream(endpoint.Ticket);
List <ByteString> actualMetadata = new List <ByteString>();
while (await getStream.ResponseStream.MoveNext(default))
public override async Task <Protocol.FlightInfo> GetFlightInfo(Protocol.FlightDescriptor request, ServerCallContext context)
{
var flightDescriptor = new FlightDescriptor(request);
var flightInfo = await _flightServer.GetFlightInfo(flightDescriptor, context).ConfigureAwait(false);
return(flightInfo.ToProtocol());
}
private IEnumerable <RecordBatchWithMetadata> GetStoreBatch(FlightDescriptor flightDescriptor)
{
Assert.Contains(flightDescriptor, (IReadOnlyDictionary <FlightDescriptor, FlightHolder>)_flightStore.Flights);
var flightHolder = _flightStore.Flights[flightDescriptor];
return(flightHolder.GetRecordBatches());
}
public override Task <Schema> GetSchema(FlightDescriptor request, ServerCallContext context)
{
if (_flightStore.Flights.TryGetValue(request, out var flightHolder))
{
return(Task.FromResult(flightHolder.GetFlightInfo().Schema));
}
throw new RpcException(new Status(StatusCode.NotFound, "Flight not found"));
}
public override async Task <SchemaResult> GetSchema(Protocol.FlightDescriptor request, ServerCallContext context)
{
var flightDescriptor = new FlightDescriptor(request);
var schema = await _flightServer.GetSchema(flightDescriptor, context).ConfigureAwait(false);
return(new SchemaResult()
{
Schema = SchemaWriter.SerializeSchema(schema)
});
}
public override Task <FlightInfo> GetFlightInfo(FlightDescriptor request, ServerCallContext context)
{
var key = DescriptorAsTicket(request);
if (_flightData.Flights.ContainsKey(key))
{
return(Task.FromResult(_flightData.Flights[key]));
}
else
{
throw new RpcException(new Status(StatusCode.NotFound, "Flight not found."));
}
}
public FlightRecordBatchDuplexStreamingCall StartPut(FlightDescriptor flightDescriptor, Metadata headers = null)
{
var channels = _client.DoPut(headers);
var requestStream = new FlightClientRecordBatchStreamWriter(channels.RequestStream, flightDescriptor);
var readStream = new StreamReader <Protocol.PutResult, FlightPutResult>(channels.ResponseStream, putResult => new FlightPutResult(putResult));
return(new FlightRecordBatchDuplexStreamingCall(
requestStream,
readStream,
channels.ResponseHeadersAsync,
channels.GetStatus,
channels.GetTrailers,
channels.Dispose));
}
private FlightInfo GivenStoreBatches(FlightDescriptor flightDescriptor, params RecordBatchWithMetadata[] batches)
{
var initialBatch = batches.FirstOrDefault();
var flightHolder = new FlightHolder(flightDescriptor, initialBatch.RecordBatch.Schema, _testWebFactory.GetAddress());
foreach (var batch in batches)
{
flightHolder.AddBatch(batch);
}
_flightStore.Flights.Add(flightDescriptor, flightHolder);
return(flightHolder.GetFlightInfo());
}
public override async Task DoGet(FlightTicket ticket, FlightServerRecordBatchStreamWriter responseStream, ServerCallContext context)
{
var flightDescriptor = FlightDescriptor.CreatePathDescriptor(ticket.Ticket.ToStringUtf8());
if (_flightStore.Flights.TryGetValue(flightDescriptor, out var flightHolder))
{
var batches = flightHolder.GetRecordBatches();
foreach (var batch in batches)
{
await responseStream.WriteAsync(batch.RecordBatch, batch.Metadata);
}
}
}
public AsyncUnaryCall <Schema> GetSchema(FlightDescriptor flightDescriptor, Metadata headers = null)
{
var schemaResult = _client.GetSchemaAsync(flightDescriptor.ToProtocol(), headers);
var schema = schemaResult
.ResponseAsync
.ContinueWith(async schema => FlightMessageSerializer.DecodeSchema((await schemaResult.ResponseAsync.ConfigureAwait(false)).Schema.Memory))
.Unwrap();
return(new AsyncUnaryCall <Schema>(
schema,
schemaResult.ResponseHeadersAsync,
schemaResult.GetStatus,
schemaResult.GetTrailers,
schemaResult.Dispose));
}
public AsyncUnaryCall <FlightInfo> GetInfo(FlightDescriptor flightDescriptor, Metadata headers = null)
{
var flightInfoResult = _client.GetFlightInfoAsync(flightDescriptor.ToProtocol(), headers);
var flightInfo = flightInfoResult
.ResponseAsync
.ContinueWith(async flightInfo => new FlightInfo(await flightInfo.ConfigureAwait(false)))
.Unwrap();
return(new AsyncUnaryCall <FlightInfo>(
flightInfo,
flightInfoResult.ResponseHeadersAsync,
flightInfoResult.GetStatus,
flightInfoResult.GetTrailers,
flightInfoResult.Dispose));
}
public async ValueTask <Schema> ReadSchema()
{
if (HasReadSchema)
{
return(Schema);
}
var moveNextResult = await _flightDataStream.MoveNext().ConfigureAwait(false);
if (!moveNextResult)
{
throw new Exception("No records or schema in this flight");
}
//AppMetadata will never be null, but length 0 if empty
//Those are skipped
if (_flightDataStream.Current.AppMetadata.Length > 0)
{
_applicationMetadatas.Add(_flightDataStream.Current.AppMetadata);
}
var header = _flightDataStream.Current.DataHeader.Memory;
Message message = Message.GetRootAsMessage(
ArrowReaderImplementation.CreateByteBuffer(header));
if (_flightDataStream.Current.FlightDescriptor != null)
{
_flightDescriptor = new FlightDescriptor(_flightDataStream.Current.FlightDescriptor);
}
switch (message.HeaderType)
{
case MessageHeader.Schema:
Schema = FlightMessageSerializer.DecodeSchema(message.ByteBuffer);
break;
default:
throw new Exception($"Expected schema as the first message, but got: {message.HeaderType.ToString()}");
}
return(Schema);
}
public async Task TestPutSingleRecordBatch()
{
var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test");
var expectedBatch = CreateTestBatch(0, 100);
var putStream = _flightClient.StartPut(flightDescriptor);
await putStream.RequestStream.WriteAsync(expectedBatch);
await putStream.RequestStream.CompleteAsync();
var putResults = await putStream.ResponseStream.ToListAsync();
Assert.Single(putResults);
var actualBatches = GetStoreBatch(flightDescriptor);
Assert.Single(actualBatches);
ArrowReaderVerifier.CompareBatches(expectedBatch, actualBatches.First().RecordBatch);
}
public async Task TestGetSingleRecordBatch()
{
var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test");
var expectedBatch = CreateTestBatch(0, 100);
//Add batch to the in memory store
GivenStoreBatches(flightDescriptor, new RecordBatchWithMetadata(expectedBatch));
//Get the flight info for the ticket
var flightInfo = await _flightClient.GetInfo(flightDescriptor);
Assert.Single(flightInfo.Endpoints);
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]);
}
private FlightInfo GetFlightInfo(string sql, ServerCallContext context)
{
var partitionsResult = _koraliumTransportService.GetPartitions(IsPartitionsEnabled(context), sql, new Shared.SqlParameters(), context.GetHttpContext()).Result;
var schemaBuilder = new Schema.Builder();
foreach (var column in partitionsResult.Columns)
{
schemaBuilder.Field(new Field(column.Name, TypeConverter.Convert(column), column.IsNullable));
}
var descriptor = FlightDescriptor.CreateCommandDescriptor(sql);
List <FlightEndpoint> endpoints = new List <FlightEndpoint>();
foreach (var partition in partitionsResult.Partitions)
{
List <FlightLocation> locations = new List <FlightLocation>();
foreach (var location in partition.Locations)
{
string uri = null;
if (location.Tls)
{
uri = $"grpc+tls://{location.Host}";
}
else
{
uri = $"grpc+tcp://{location.Host}";
}
locations.Add(new FlightLocation(uri));
}
endpoints.Add(new FlightEndpoint(new FlightTicket(partition.Sql), locations));
}
return(new FlightInfo(schemaBuilder.Build(), descriptor, endpoints));
}
public async Task TestPutTwoRecordBatches()
{
var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test");
var expectedBatch1 = CreateTestBatch(0, 100);
var expectedBatch2 = CreateTestBatch(0, 100);
var putStream = _flightClient.StartPut(flightDescriptor);
await putStream.RequestStream.WriteAsync(expectedBatch1);
await putStream.RequestStream.WriteAsync(expectedBatch2);
await putStream.RequestStream.CompleteAsync();
var putResults = await putStream.ResponseStream.ToListAsync();
Assert.Equal(2, putResults.Count);
var actualBatches = GetStoreBatch(flightDescriptor).ToList();
Assert.Equal(2, actualBatches.Count);
ArrowReaderVerifier.CompareBatches(expectedBatch1, actualBatches[0].RecordBatch);
ArrowReaderVerifier.CompareBatches(expectedBatch2, actualBatches[1].RecordBatch);
}
public virtual Task <FlightInfo> GetFlightInfo(FlightDescriptor request, ServerCallContext context)
{
throw new NotImplementedException();
}
public override Task <Schema> GetSchema(FlightDescriptor request, ServerCallContext context)
{
return(Task.FromResult(GetSchema(request.Command.ToStringUtf8(), context)));
}
public static async Task Main(string[] args)
{
string host = args.Length > 0 ? args[0] : "localhost";
string port = args.Length > 1 ? args[1] : "5000";
// Create client
// (In production systems, you should use https not http)
var address = $"http://{host}:{port}";
Console.WriteLine($"Connecting to: {address}");
var channel = GrpcChannel.ForAddress(address);
var client = new FlightClient(channel);
var recordBatches = new RecordBatch[] {
CreateTestBatch(0, 2000), CreateTestBatch(50, 9000)
};
// Particular flights are identified by a descriptor. This might be a name,
// a SQL query, or a path. Here, just using the name "test".
var descriptor = FlightDescriptor.CreateCommandDescriptor("test");
// Upload data with StartPut
var batchStreamingCall = client.StartPut(descriptor);
foreach (var batch in recordBatches)
{
await batchStreamingCall.RequestStream.WriteAsync(batch);
}
// Signal we are done sending record batches
await batchStreamingCall.RequestStream.CompleteAsync();
// Retrieve final response
await batchStreamingCall.ResponseStream.MoveNext();
Console.WriteLine(batchStreamingCall.ResponseStream.Current.ApplicationMetadata.ToStringUtf8());
Console.WriteLine($"Wrote {recordBatches.Length} batches to server.");
// Request information:
var schema = await client.GetSchema(descriptor).ResponseAsync;
Console.WriteLine($"Schema saved as: \n {schema}");
var info = await client.GetInfo(descriptor).ResponseAsync;
Console.WriteLine($"Info provided: \n {info}");
Console.WriteLine($"Available flights:");
var flights_call = client.ListFlights();
while (await flights_call.ResponseStream.MoveNext())
{
Console.WriteLine(" " + flights_call.ResponseStream.Current.ToString());
}
// Download data
await foreach (var batch in StreamRecordBatches(info))
{
Console.WriteLine($"Read batch from flight server: \n {batch}");
}
// See available comands on this server
var action_stream = client.ListActions();
Console.WriteLine("Actions:");
while (await action_stream.ResponseStream.MoveNext())
{
var action = action_stream.ResponseStream.Current;
Console.WriteLine($" {action.Type}: {action.Description}");
}
// Send clear command to drop all data from the server.
var clear_result = client.DoAction(new FlightAction("clear"));
await clear_result.ResponseStream.MoveNext(default);
public FlightDataStream(IAsyncStreamWriter <FlightData> clientStreamWriter, FlightDescriptor flightDescriptor, Schema schema)
: base(new MemoryStream(), schema)
{
_clientStreamWriter = clientStreamWriter;
_flightDescriptor = flightDescriptor;
}
public override async Task <Schema> GetSchema(FlightDescriptor request, ServerCallContext context)
{
var info = await GetFlightInfo(request, context);
return(info.Schema);
}
private FlightTicket DescriptorAsTicket(FlightDescriptor desc)
{
return(new FlightTicket(desc.Command));
}
internal FlightClientRecordBatchStreamWriter(IClientStreamWriter <FlightData> clientStreamWriter, FlightDescriptor flightDescriptor) : base(clientStreamWriter, flightDescriptor)
{
_clientStreamWriter = clientStreamWriter;
}
public FlightHolder(FlightDescriptor flightDescriptor, Schema schema, string location)
{
_flightDescriptor = flightDescriptor;
_schema = schema;
_location = location;
}
