| # 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. |
| |
| # cython: language_level = 3 |
| |
| from __future__ import absolute_import |
| |
| import collections |
| import enum |
| |
| import six |
| |
| from cython.operator cimport dereference as deref |
| |
| from pyarrow.compat import frombytes, tobytes |
| from pyarrow.lib cimport * |
| from pyarrow.lib import as_buffer |
| from pyarrow.includes.libarrow_flight cimport * |
| from pyarrow.ipc import _ReadPandasOption |
| import pyarrow.lib as lib |
| |
| |
| cdef CFlightCallOptions DEFAULT_CALL_OPTIONS |
| |
| |
| cdef class FlightCallOptions: |
| """RPC-layer options for a Flight call.""" |
| |
| cdef: |
| CFlightCallOptions options |
| |
| def __init__(self, timeout=None): |
| """Create call options. |
| |
| Parameters |
| ---------- |
| timeout : float or None |
| A timeout for the call, in seconds. None means that the |
| timeout defaults to an implementation-specific value. |
| |
| """ |
| if timeout is not None: |
| self.options.timeout = CTimeoutDuration(timeout) |
| |
| @staticmethod |
| cdef CFlightCallOptions* unwrap(obj): |
| if not obj: |
| return &DEFAULT_CALL_OPTIONS |
| elif isinstance(obj, FlightCallOptions): |
| return &((<FlightCallOptions> obj).options) |
| raise TypeError("Expected a FlightCallOptions object, not " |
| "'{}'".format(type(obj))) |
| |
| |
| _CertKeyPair = collections.namedtuple('_CertKeyPair', ['cert', 'key']) |
| |
| |
| class CertKeyPair(_CertKeyPair): |
| """A TLS certificate and key for use in Flight.""" |
| |
| |
| cdef class Action: |
| """An action executable on a Flight service.""" |
| cdef: |
| CAction action |
| |
| def __init__(self, action_type, buf): |
| """Create an action from a type and a buffer. |
| |
| Parameters |
| ---------- |
| action_type : bytes or str |
| buf : Buffer or bytes-like object |
| """ |
| self.action.type = tobytes(action_type) |
| self.action.body = pyarrow_unwrap_buffer(as_buffer(buf)) |
| |
| @property |
| def type(self): |
| """The action type.""" |
| return frombytes(self.action.type) |
| |
| @property |
| def body(self): |
| """The action body (arguments for the action).""" |
| return pyarrow_wrap_buffer(self.action.body) |
| |
| @staticmethod |
| cdef CAction unwrap(action) except *: |
| if not isinstance(action, Action): |
| raise TypeError("Must provide Action, not '{}'".format( |
| type(action))) |
| return (<Action> action).action |
| |
| |
| _ActionType = collections.namedtuple('_ActionType', ['type', 'description']) |
| |
| |
| class ActionType(_ActionType): |
| """A type of action that is executable on a Flight service.""" |
| |
| def make_action(self, buf): |
| """Create an Action with this type. |
| |
| Parameters |
| ---------- |
| buf : obj |
| An Arrow buffer or Python bytes or bytes-like object. |
| """ |
| return Action(self.type, buf) |
| |
| |
| cdef class Result: |
| """A result from executing an Action.""" |
| cdef: |
| unique_ptr[CFlightResult] result |
| |
| def __init__(self, buf): |
| """Create a new result. |
| |
| Parameters |
| ---------- |
| buf : Buffer or bytes-like object |
| """ |
| self.result.reset(new CFlightResult()) |
| self.result.get().body = pyarrow_unwrap_buffer(as_buffer(buf)) |
| |
| @property |
| def body(self): |
| """Get the Buffer containing the result.""" |
| return pyarrow_wrap_buffer(self.result.get().body) |
| |
| |
| class DescriptorType(enum.Enum): |
| """ |
| The type of a FlightDescriptor. |
| |
| Attributes |
| ---------- |
| |
| UNKNOWN |
| An unknown descriptor type. |
| |
| PATH |
| A Flight stream represented by a path. |
| |
| CMD |
| A Flight stream represented by an application-defined command. |
| |
| """ |
| |
| UNKNOWN = 0 |
| PATH = 1 |
| CMD = 2 |
| |
| |
| cdef class FlightDescriptor: |
| """A description of a data stream available from a Flight service.""" |
| cdef: |
| CFlightDescriptor descriptor |
| |
| def __init__(self): |
| raise TypeError("Do not call {}'s constructor directly, use " |
| "`pyarrow.flight.FlightDescriptor.for_{path,command}` " |
| "function instead." |
| .format(self.__class__.__name__)) |
| |
| @staticmethod |
| def for_path(*path): |
| """Create a FlightDescriptor for a resource path.""" |
| cdef FlightDescriptor result = \ |
| FlightDescriptor.__new__(FlightDescriptor) |
| result.descriptor.type = CDescriptorTypePath |
| result.descriptor.path = [tobytes(p) for p in path] |
| return result |
| |
| @staticmethod |
| def for_command(command): |
| """Create a FlightDescriptor for an opaque command.""" |
| cdef FlightDescriptor result = \ |
| FlightDescriptor.__new__(FlightDescriptor) |
| result.descriptor.type = CDescriptorTypeCmd |
| result.descriptor.cmd = tobytes(command) |
| return result |
| |
| @property |
| def descriptor_type(self): |
| """Get the type of this descriptor.""" |
| if self.descriptor.type == CDescriptorTypeUnknown: |
| return DescriptorType.UNKNOWN |
| elif self.descriptor.type == CDescriptorTypePath: |
| return DescriptorType.PATH |
| elif self.descriptor.type == CDescriptorTypeCmd: |
| return DescriptorType.CMD |
| raise RuntimeError("Invalid descriptor type!") |
| |
| @property |
| def command(self): |
| """Get the command for this descriptor.""" |
| if self.descriptor_type != DescriptorType.CMD: |
| return None |
| return self.descriptor.cmd |
| |
| @property |
| def path(self): |
| """Get the path for this descriptor.""" |
| if self.descriptor_type != DescriptorType.PATH: |
| return None |
| return self.descriptor.path |
| |
| def __repr__(self): |
| return "<FlightDescriptor type: {!r}>".format(self.descriptor_type) |
| |
| @staticmethod |
| cdef CFlightDescriptor unwrap(descriptor) except *: |
| if not isinstance(descriptor, FlightDescriptor): |
| raise TypeError("Must provide a FlightDescriptor, not '{}'".format( |
| type(descriptor))) |
| return (<FlightDescriptor> descriptor).descriptor |
| |
| |
| cdef class Ticket: |
| """A ticket for requesting a Flight stream.""" |
| |
| cdef: |
| CTicket ticket |
| |
| def __init__(self, ticket): |
| self.ticket.ticket = tobytes(ticket) |
| |
| @property |
| def ticket(self): |
| return self.ticket.ticket |
| |
| def __repr__(self): |
| return '<Ticket {}>'.format(self.ticket.ticket) |
| |
| |
| cdef class Location: |
| """The location of a Flight service.""" |
| cdef: |
| CLocation location |
| |
| def __init__(self, uri): |
| check_status(CLocation.Parse(tobytes(uri), &self.location)) |
| |
| def __repr__(self): |
| return '<Location {}>'.format(self.location.ToString()) |
| |
| @property |
| def uri(self): |
| return self.location.ToString() |
| |
| def equals(self, Location other): |
| return self == other |
| |
| def __eq__(self, other): |
| if not isinstance(other, Location): |
| return NotImplemented |
| return self.location.Equals((<Location> other).location) |
| |
| @staticmethod |
| def for_grpc_tcp(host, port): |
| """Create a Location for a TCP-based gRPC service.""" |
| cdef: |
| c_string c_host = tobytes(host) |
| int c_port = port |
| Location result = Location.__new__(Location) |
| check_status(CLocation.ForGrpcTcp(c_host, c_port, &result.location)) |
| return result |
| |
| @staticmethod |
| def for_grpc_tls(host, port): |
| """Create a Location for a TLS-based gRPC service.""" |
| cdef: |
| c_string c_host = tobytes(host) |
| int c_port = port |
| Location result = Location.__new__(Location) |
| check_status(CLocation.ForGrpcTls(c_host, c_port, &result.location)) |
| return result |
| |
| @staticmethod |
| def for_grpc_unix(path): |
| """Create a Location for a domain socket-based gRPC service.""" |
| cdef: |
| c_string c_path = tobytes(path) |
| Location result = Location.__new__(Location) |
| check_status(CLocation.ForGrpcUnix(c_path, &result.location)) |
| return result |
| |
| @staticmethod |
| cdef Location wrap(CLocation location): |
| cdef Location result = Location.__new__(Location) |
| result.location = location |
| return result |
| |
| @staticmethod |
| cdef CLocation unwrap(object location) except *: |
| cdef CLocation c_location |
| if isinstance(location, six.text_type): |
| check_status(CLocation.Parse(tobytes(location), &c_location)) |
| return c_location |
| elif not isinstance(location, Location): |
| raise TypeError("Must provide a Location, not '{}'".format( |
| type(location))) |
| return (<Location> location).location |
| |
| |
| cdef class FlightEndpoint: |
| """A Flight stream, along with the ticket and locations to access it.""" |
| cdef: |
| CFlightEndpoint endpoint |
| |
| def __init__(self, ticket, locations): |
| """Create a FlightEndpoint from a ticket and list of locations. |
| |
| Parameters |
| ---------- |
| ticket : Ticket or bytes |
| the ticket needed to access this flight |
| locations : list of string URIs |
| locations where this flight is available |
| |
| Raises |
| ------ |
| ArrowException |
| If one of the location URIs is not a valid URI. |
| """ |
| cdef: |
| CLocation c_location |
| |
| if isinstance(ticket, Ticket): |
| self.endpoint.ticket.ticket = tobytes(ticket.ticket) |
| else: |
| self.endpoint.ticket.ticket = tobytes(ticket) |
| |
| for location in locations: |
| if isinstance(location, Location): |
| c_location = (<Location> location).location |
| else: |
| c_location = CLocation() |
| check_status(CLocation.Parse(tobytes(location), &c_location)) |
| self.endpoint.locations.push_back(c_location) |
| |
| @property |
| def ticket(self): |
| """Get the ticket in this endpoint.""" |
| return Ticket(self.endpoint.ticket.ticket) |
| |
| @property |
| def locations(self): |
| return [Location.wrap(location) |
| for location in self.endpoint.locations] |
| |
| |
| cdef class FlightInfo: |
| """A description of a Flight stream.""" |
| cdef: |
| unique_ptr[CFlightInfo] info |
| |
| def __init__(self, Schema schema, FlightDescriptor descriptor, endpoints, |
| total_records, total_bytes): |
| """Create a FlightInfo object from a schema, descriptor, and endpoints. |
| |
| Parameters |
| ---------- |
| schema : Schema |
| the schema of the data in this flight. |
| descriptor : FlightDescriptor |
| the descriptor for this flight. |
| endpoints : list of FlightEndpoint |
| a list of endpoints where this flight is available. |
| total_records : int |
| the total records in this flight, or -1 if unknown |
| total_bytes : int |
| the total bytes in this flight, or -1 if unknown |
| """ |
| cdef: |
| shared_ptr[CSchema] c_schema = pyarrow_unwrap_schema(schema) |
| vector[CFlightEndpoint] c_endpoints |
| |
| for endpoint in endpoints: |
| if isinstance(endpoint, FlightEndpoint): |
| c_endpoints.push_back((<FlightEndpoint> endpoint).endpoint) |
| else: |
| raise TypeError('Endpoint {} is not instance of' |
| ' FlightEndpoint'.format(endpoint)) |
| |
| check_status(CreateFlightInfo(c_schema, |
| descriptor.descriptor, |
| c_endpoints, |
| total_records, |
| total_bytes, &self.info)) |
| |
| @property |
| def total_records(self): |
| """The total record count of this flight, or -1 if unknown.""" |
| return self.info.get().total_records() |
| |
| @property |
| def total_bytes(self): |
| """The size in bytes of the data in this flight, or -1 if unknown.""" |
| return self.info.get().total_bytes() |
| |
| @property |
| def schema(self): |
| """The schema of the data in this flight.""" |
| cdef: |
| shared_ptr[CSchema] schema |
| CDictionaryMemo dummy_memo |
| |
| check_status(self.info.get().GetSchema(&dummy_memo, &schema)) |
| return pyarrow_wrap_schema(schema) |
| |
| @property |
| def descriptor(self): |
| """The descriptor of the data in this flight.""" |
| cdef FlightDescriptor result = \ |
| FlightDescriptor.__new__(FlightDescriptor) |
| result.descriptor = self.info.get().descriptor() |
| return result |
| |
| @property |
| def endpoints(self): |
| """The endpoints where this flight is available.""" |
| # TODO: get Cython to iterate over reference directly |
| cdef: |
| vector[CFlightEndpoint] endpoints = self.info.get().endpoints() |
| FlightEndpoint py_endpoint |
| |
| result = [] |
| for endpoint in endpoints: |
| py_endpoint = FlightEndpoint.__new__(FlightEndpoint) |
| py_endpoint.endpoint = endpoint |
| result.append(py_endpoint) |
| return result |
| |
| |
| cdef class FlightStreamChunk: |
| """A RecordBatch with application metadata on the side.""" |
| cdef: |
| CFlightStreamChunk chunk |
| |
| @property |
| def data(self): |
| if self.chunk.data == NULL: |
| return None |
| return pyarrow_wrap_batch(self.chunk.data) |
| |
| @property |
| def app_metadata(self): |
| if self.chunk.app_metadata == NULL: |
| return None |
| return pyarrow_wrap_buffer(self.chunk.app_metadata) |
| |
| def __iter__(self): |
| return iter((self.data, self.app_metadata)) |
| |
| |
| cdef class _MetadataRecordBatchReader: |
| """A reader for Flight streams.""" |
| |
| # Needs to be separate class so the "real" class can subclass the |
| # pure-Python mixin class |
| |
| cdef dict __dict__ |
| cdef shared_ptr[CMetadataRecordBatchReader] reader |
| |
| cdef readonly: |
| Schema schema |
| |
| |
| cdef class MetadataRecordBatchReader(_MetadataRecordBatchReader, |
| _ReadPandasOption): |
| """A reader for Flight streams.""" |
| |
| def __iter__(self): |
| while True: |
| yield self.read_chunk() |
| |
| def read_all(self): |
| """Read the entire contents of the stream as a Table.""" |
| cdef: |
| shared_ptr[CTable] c_table |
| with nogil: |
| check_status(self.reader.get().ReadAll(&c_table)) |
| return pyarrow_wrap_table(c_table) |
| |
| def read_chunk(self): |
| """Read the next RecordBatch along with any metadata. |
| |
| Returns |
| ------- |
| data : RecordBatch |
| The next RecordBatch in the stream. |
| app_metadata : Buffer or None |
| Application-specific metadata for the batch as defined by |
| Flight. |
| |
| Raises |
| ------ |
| StopIteration |
| when the stream is finished |
| """ |
| cdef: |
| FlightStreamChunk chunk = FlightStreamChunk() |
| |
| with nogil: |
| check_status(self.reader.get().Next(&chunk.chunk)) |
| |
| if chunk.chunk.data == NULL: |
| raise StopIteration |
| |
| return chunk |
| |
| |
| cdef class FlightStreamReader(MetadataRecordBatchReader): |
| """A reader that can also be canceled.""" |
| |
| def cancel(self): |
| """Cancel the read operation.""" |
| with nogil: |
| (<CFlightStreamReader*> self.reader.get()).Cancel() |
| |
| |
| cdef class FlightStreamWriter(_CRecordBatchWriter): |
| """A RecordBatchWriter that also allows writing application metadata.""" |
| |
| def write_with_metadata(self, RecordBatch batch, buf): |
| """Write a RecordBatch along with Flight metadata. |
| |
| Parameters |
| ---------- |
| batch : RecordBatch |
| The next RecordBatch in the stream. |
| buf : Buffer |
| Application-specific metadata for the batch as defined by |
| Flight. |
| """ |
| cdef shared_ptr[CBuffer] c_buf = pyarrow_unwrap_buffer(as_buffer(buf)) |
| with nogil: |
| check_status( |
| (<CFlightStreamWriter*> self.writer.get()) |
| .WriteWithMetadata(deref(batch.batch), |
| c_buf, |
| 1)) |
| |
| |
| cdef class FlightMetadataReader: |
| """A reader for Flight metadata messages sent during a DoPut.""" |
| |
| cdef: |
| unique_ptr[CFlightMetadataReader] reader |
| |
| def read(self): |
| """Read the next metadata message.""" |
| cdef shared_ptr[CBuffer] buf |
| with nogil: |
| check_status(self.reader.get().ReadMetadata(&buf)) |
| if buf == NULL: |
| return None |
| return pyarrow_wrap_buffer(buf) |
| |
| |
| cdef class FlightMetadataWriter: |
| """A sender for Flight metadata messages during a DoPut.""" |
| |
| cdef: |
| unique_ptr[CFlightMetadataWriter] writer |
| |
| def write(self, message): |
| """Write the next metadata message. |
| |
| Parameters |
| ---------- |
| message : Buffer |
| """ |
| cdef shared_ptr[CBuffer] buf = \ |
| pyarrow_unwrap_buffer(as_buffer(message)) |
| with nogil: |
| check_status(self.writer.get().WriteMetadata(deref(buf))) |
| |
| |
| cdef class FlightClient: |
| """A client to a Flight service.""" |
| cdef: |
| unique_ptr[CFlightClient] client |
| |
| def __init__(self): |
| raise TypeError("Do not call {}'s constructor directly, use " |
| "`pyarrow.flight.FlightClient.connect` instead." |
| .format(self.__class__.__name__)) |
| |
| @staticmethod |
| def connect(location, tls_root_certs=None, override_hostname=None): |
| """ |
| Connect to a Flight service on the given host and port. |
| |
| Parameters |
| ---------- |
| location : Location |
| location to connect to |
| |
| tls_root_certs : bytes or None |
| PEM-encoded |
| unsafe_override_hostname : str or None |
| Override the hostname checked by TLS. Insecure, use with caution. |
| """ |
| cdef: |
| FlightClient result = FlightClient.__new__(FlightClient) |
| int c_port = 0 |
| CLocation c_location = Location.unwrap(location) |
| CFlightClientOptions c_options |
| |
| if tls_root_certs: |
| c_options.tls_root_certs = tobytes(tls_root_certs) |
| if override_hostname: |
| c_options.override_hostname = tobytes(override_hostname) |
| |
| with nogil: |
| check_status(CFlightClient.Connect(c_location, c_options, |
| &result.client)) |
| |
| return result |
| |
| def authenticate(self, auth_handler, options: FlightCallOptions = None): |
| """Authenticate to the server. |
| |
| Parameters |
| ---------- |
| auth_handler : ClientAuthHandler |
| The authentication mechanism to use. |
| options : FlightCallOptions |
| Options for this call. |
| """ |
| cdef: |
| unique_ptr[CClientAuthHandler] handler |
| CFlightCallOptions* c_options = FlightCallOptions.unwrap(options) |
| |
| if not isinstance(auth_handler, ClientAuthHandler): |
| raise TypeError( |
| "FlightClient.authenticate takes a ClientAuthHandler, " |
| "not '{}'".format(type(auth_handler))) |
| handler.reset((<ClientAuthHandler> auth_handler).to_handler()) |
| with nogil: |
| check_status(self.client.get().Authenticate(deref(c_options), |
| move(handler))) |
| |
| def list_actions(self, options: FlightCallOptions = None): |
| """List the actions available on a service.""" |
| cdef: |
| vector[CActionType] results |
| CFlightCallOptions* c_options = FlightCallOptions.unwrap(options) |
| |
| with nogil: |
| check_status( |
| self.client.get().ListActions(deref(c_options), &results)) |
| |
| result = [] |
| for action_type in results: |
| py_action = ActionType(frombytes(action_type.type), |
| frombytes(action_type.description)) |
| result.append(py_action) |
| |
| return result |
| |
| def do_action(self, action: Action, options: FlightCallOptions = None): |
| """Execute an action on a service.""" |
| cdef: |
| unique_ptr[CResultStream] results |
| Result result |
| CAction c_action = Action.unwrap(action) |
| CFlightCallOptions* c_options = FlightCallOptions.unwrap(options) |
| with nogil: |
| check_status( |
| self.client.get().DoAction(deref(c_options), c_action, |
| &results)) |
| |
| while True: |
| result = Result.__new__(Result) |
| with nogil: |
| check_status(results.get().Next(&result.result)) |
| if result.result == NULL: |
| break |
| yield result |
| |
| def list_flights(self, options: FlightCallOptions = None): |
| """List the flights available on a service.""" |
| cdef: |
| unique_ptr[CFlightListing] listing |
| FlightInfo result |
| CFlightCallOptions* c_options = FlightCallOptions.unwrap(options) |
| CCriteria c_criteria |
| |
| with nogil: |
| check_status(self.client.get().ListFlights(deref(c_options), |
| c_criteria, &listing)) |
| |
| while True: |
| result = FlightInfo.__new__(FlightInfo) |
| with nogil: |
| check_status(listing.get().Next(&result.info)) |
| if result.info == NULL: |
| break |
| yield result |
| |
| def get_flight_info(self, descriptor: FlightDescriptor, |
| options: FlightCallOptions = None): |
| """Request information about an available flight.""" |
| cdef: |
| FlightInfo result = FlightInfo.__new__(FlightInfo) |
| CFlightCallOptions* c_options = FlightCallOptions.unwrap(options) |
| CFlightDescriptor c_descriptor = \ |
| FlightDescriptor.unwrap(descriptor) |
| |
| with nogil: |
| check_status(self.client.get().GetFlightInfo( |
| deref(c_options), c_descriptor, &result.info)) |
| |
| return result |
| |
| def do_get(self, ticket: Ticket, options: FlightCallOptions = None): |
| """Request the data for a flight. |
| |
| Returns |
| ------- |
| reader : FlightStreamReader |
| """ |
| cdef: |
| unique_ptr[CFlightStreamReader] reader |
| CFlightCallOptions* c_options = FlightCallOptions.unwrap(options) |
| |
| with nogil: |
| check_status( |
| self.client.get().DoGet( |
| deref(c_options), ticket.ticket, &reader)) |
| result = FlightStreamReader() |
| result.reader.reset(reader.release()) |
| result.schema = pyarrow_wrap_schema(result.reader.get().schema()) |
| return result |
| |
| def do_put(self, descriptor: FlightDescriptor, schema: Schema, |
| options: FlightCallOptions = None): |
| """Upload data to a flight. |
| |
| Returns |
| ------- |
| writer : FlightStreamWriter |
| reader : FlightMetadataReader |
| """ |
| cdef: |
| shared_ptr[CSchema] c_schema = pyarrow_unwrap_schema(schema) |
| unique_ptr[CFlightStreamWriter] writer |
| unique_ptr[CFlightMetadataReader] metadata_reader |
| CFlightCallOptions* c_options = FlightCallOptions.unwrap(options) |
| CFlightDescriptor c_descriptor = \ |
| FlightDescriptor.unwrap(descriptor) |
| FlightMetadataReader reader = FlightMetadataReader() |
| |
| with nogil: |
| check_status(self.client.get().DoPut( |
| deref(c_options), |
| c_descriptor, |
| c_schema, |
| &writer, |
| &reader.reader)) |
| result = FlightStreamWriter() |
| result.writer.reset(writer.release()) |
| return result, reader |
| |
| |
| cdef class FlightDataStream: |
| """Abstract base class for Flight data streams.""" |
| |
| cdef CFlightDataStream* to_stream(self) except *: |
| """Create the C++ data stream for the backing Python object. |
| |
| We don't expose the C++ object to Python, so we can manage its |
| lifetime from the Cython/C++ side. |
| """ |
| raise NotImplementedError |
| |
| |
| cdef class RecordBatchStream(FlightDataStream): |
| """A Flight data stream backed by RecordBatches.""" |
| cdef: |
| object data_source |
| |
| def __init__(self, data_source): |
| """Create a RecordBatchStream from a data source. |
| |
| Parameters |
| ---------- |
| data_source : RecordBatchReader or Table |
| """ |
| if (not isinstance(data_source, _CRecordBatchReader) and |
| not isinstance(data_source, lib.Table)): |
| raise TypeError("Expected RecordBatchReader or Table, " |
| "but got: {}".format(type(data_source))) |
| self.data_source = data_source |
| |
| cdef CFlightDataStream* to_stream(self) except *: |
| cdef: |
| shared_ptr[CRecordBatchReader] reader |
| if isinstance(self.data_source, _CRecordBatchReader): |
| reader = (<_CRecordBatchReader> self.data_source).reader |
| elif isinstance(self.data_source, lib.Table): |
| table = (<Table> self.data_source).table |
| reader.reset(new TableBatchReader(deref(table))) |
| else: |
| raise RuntimeError("Can't construct RecordBatchStream " |
| "from type {}".format(type(self.data_source))) |
| return new CRecordBatchStream(reader) |
| |
| |
| cdef class GeneratorStream(FlightDataStream): |
| """A Flight data stream backed by a Python generator.""" |
| cdef: |
| shared_ptr[CSchema] schema |
| object generator |
| # A substream currently being consumed by the client, if |
| # present. Produced by the generator. |
| unique_ptr[CFlightDataStream] current_stream |
| |
| def __init__(self, schema, generator): |
| """Create a GeneratorStream from a Python generator. |
| |
| Parameters |
| ---------- |
| schema : Schema |
| The schema for the data to be returned. |
| |
| generator : iterator or iterable |
| The generator should yield other FlightDataStream objects, |
| Tables, RecordBatches, or RecordBatchReaders. |
| """ |
| self.schema = pyarrow_unwrap_schema(schema) |
| self.generator = iter(generator) |
| |
| cdef CFlightDataStream* to_stream(self) except *: |
| cdef: |
| function[cb_data_stream_next] callback = &_data_stream_next |
| return new CPyGeneratorFlightDataStream(self, self.schema, callback) |
| |
| |
| cdef class ServerCallContext: |
| """Per-call state/context.""" |
| cdef: |
| const CServerCallContext* context |
| |
| def peer_identity(self): |
| """Get the identity of the authenticated peer. |
| |
| May be the empty string. |
| """ |
| return tobytes(self.context.peer_identity()) |
| |
| @staticmethod |
| cdef ServerCallContext wrap(const CServerCallContext& context): |
| cdef ServerCallContext result = \ |
| ServerCallContext.__new__(ServerCallContext) |
| result.context = &context |
| return result |
| |
| |
| cdef class ServerAuthReader: |
| """A reader for messages from the client during an auth handshake.""" |
| cdef: |
| CServerAuthReader* reader |
| |
| def read(self): |
| cdef c_string token |
| if not self.reader: |
| raise ValueError("Cannot use ServerAuthReader outside " |
| "ServerAuthHandler.authenticate") |
| with nogil: |
| check_status(self.reader.Read(&token)) |
| return token |
| |
| cdef void poison(self): |
| """Prevent further usage of this object. |
| |
| This object is constructed by taking a pointer to a reference, |
| so we want to make sure Python users do not access this after |
| the reference goes away. |
| """ |
| self.reader = NULL |
| |
| @staticmethod |
| cdef ServerAuthReader wrap(CServerAuthReader* reader): |
| cdef ServerAuthReader result = \ |
| ServerAuthReader.__new__(ServerAuthReader) |
| result.reader = reader |
| return result |
| |
| |
| cdef class ServerAuthSender: |
| """A writer for messages to the client during an auth handshake.""" |
| cdef: |
| CServerAuthSender* sender |
| |
| def write(self, message): |
| cdef c_string c_message = tobytes(message) |
| if not self.sender: |
| raise ValueError("Cannot use ServerAuthSender outside " |
| "ServerAuthHandler.authenticate") |
| with nogil: |
| check_status(self.sender.Write(c_message)) |
| |
| cdef void poison(self): |
| """Prevent further usage of this object. |
| |
| This object is constructed by taking a pointer to a reference, |
| so we want to make sure Python users do not access this after |
| the reference goes away. |
| """ |
| self.sender = NULL |
| |
| @staticmethod |
| cdef ServerAuthSender wrap(CServerAuthSender* sender): |
| cdef ServerAuthSender result = \ |
| ServerAuthSender.__new__(ServerAuthSender) |
| result.sender = sender |
| return result |
| |
| |
| cdef class ClientAuthReader: |
| """A reader for messages from the server during an auth handshake.""" |
| cdef: |
| CClientAuthReader* reader |
| |
| def read(self): |
| cdef c_string token |
| if not self.reader: |
| raise ValueError("Cannot use ClientAuthReader outside " |
| "ClientAuthHandler.authenticate") |
| with nogil: |
| check_status(self.reader.Read(&token)) |
| return token |
| |
| cdef void poison(self): |
| """Prevent further usage of this object. |
| |
| This object is constructed by taking a pointer to a reference, |
| so we want to make sure Python users do not access this after |
| the reference goes away. |
| """ |
| self.reader = NULL |
| |
| @staticmethod |
| cdef ClientAuthReader wrap(CClientAuthReader* reader): |
| cdef ClientAuthReader result = \ |
| ClientAuthReader.__new__(ClientAuthReader) |
| result.reader = reader |
| return result |
| |
| |
| cdef class ClientAuthSender: |
| """A writer for messages to the server during an auth handshake.""" |
| cdef: |
| CClientAuthSender* sender |
| |
| def write(self, message): |
| cdef c_string c_message = tobytes(message) |
| if not self.sender: |
| raise ValueError("Cannot use ClientAuthSender outside " |
| "ClientAuthHandler.authenticate") |
| with nogil: |
| check_status(self.sender.Write(c_message)) |
| |
| cdef void poison(self): |
| """Prevent further usage of this object. |
| |
| This object is constructed by taking a pointer to a reference, |
| so we want to make sure Python users do not access this after |
| the reference goes away. |
| """ |
| self.sender = NULL |
| |
| @staticmethod |
| cdef ClientAuthSender wrap(CClientAuthSender* sender): |
| cdef ClientAuthSender result = \ |
| ClientAuthSender.__new__(ClientAuthSender) |
| result.sender = sender |
| return result |
| |
| |
| cdef void _data_stream_next(void* self, CFlightPayload* payload) except *: |
| """Callback for implementing FlightDataStream in Python.""" |
| cdef: |
| unique_ptr[CFlightDataStream] data_stream |
| |
| py_stream = <object> self |
| if not isinstance(py_stream, GeneratorStream): |
| raise RuntimeError("self object in callback is not GeneratorStream") |
| stream = <GeneratorStream> py_stream |
| |
| if stream.current_stream != nullptr: |
| check_status(stream.current_stream.get().Next(payload)) |
| # If the stream ended, see if there's another stream from the |
| # generator |
| if payload.ipc_message.metadata != nullptr: |
| return |
| stream.current_stream.reset(nullptr) |
| |
| try: |
| result = next(stream.generator) |
| except StopIteration: |
| payload.ipc_message.metadata.reset(<CBuffer*> nullptr) |
| return |
| |
| if isinstance(result, (list, tuple)): |
| result, metadata = result |
| else: |
| result, metadata = result, None |
| |
| if isinstance(result, (Table, _CRecordBatchReader)): |
| if metadata: |
| raise ValueError("Can only return metadata alongside a " |
| "RecordBatch.") |
| result = RecordBatchStream(result) |
| |
| stream_schema = pyarrow_wrap_schema(stream.schema) |
| if isinstance(result, FlightDataStream): |
| if metadata: |
| raise ValueError("Can only return metadata alongside a " |
| "RecordBatch.") |
| data_stream = unique_ptr[CFlightDataStream]( |
| (<FlightDataStream> result).to_stream()) |
| substream_schema = pyarrow_wrap_schema(data_stream.get().schema()) |
| if substream_schema != stream_schema: |
| raise ValueError("Got a FlightDataStream whose schema does not " |
| "match the declared schema of this " |
| "GeneratorStream. " |
| "Got: {}\nExpected: {}".format(substream_schema, |
| stream_schema)) |
| stream.current_stream.reset( |
| new CPyFlightDataStream(result, move(data_stream))) |
| _data_stream_next(self, payload) |
| elif isinstance(result, RecordBatch): |
| batch = <RecordBatch> result |
| if batch.schema != stream_schema: |
| raise ValueError("Got a RecordBatch whose schema does not " |
| "match the declared schema of this " |
| "GeneratorStream. " |
| "Got: {}\nExpected: {}".format(batch.schema, |
| stream_schema)) |
| check_status(_GetRecordBatchPayload( |
| deref(batch.batch), |
| c_default_memory_pool(), |
| &payload.ipc_message)) |
| if metadata: |
| payload.app_metadata = pyarrow_unwrap_buffer(as_buffer(metadata)) |
| else: |
| raise TypeError("GeneratorStream must be initialized with " |
| "an iterator of FlightDataStream, Table, " |
| "RecordBatch, or RecordBatchStreamReader objects, " |
| "not {}.".format(type(result))) |
| |
| |
| cdef void _list_flights(void* self, const CServerCallContext& context, |
| const CCriteria* c_criteria, |
| unique_ptr[CFlightListing]* listing) except *: |
| """Callback for implementing ListFlights in Python.""" |
| cdef: |
| vector[CFlightInfo] flights |
| result = (<object> self).list_flights(ServerCallContext.wrap(context), |
| c_criteria.expression) |
| for info in result: |
| if not isinstance(info, FlightInfo): |
| raise TypeError("FlightServerBase.list_flights must return " |
| "FlightInfo instances, but got {}".format( |
| type(info))) |
| flights.push_back(deref((<FlightInfo> info).info.get())) |
| listing.reset(new CSimpleFlightListing(flights)) |
| |
| |
| cdef void _get_flight_info(void* self, const CServerCallContext& context, |
| CFlightDescriptor c_descriptor, |
| unique_ptr[CFlightInfo]* info) except *: |
| """Callback for implementing Flight servers in Python.""" |
| cdef: |
| FlightDescriptor py_descriptor = \ |
| FlightDescriptor.__new__(FlightDescriptor) |
| py_descriptor.descriptor = c_descriptor |
| result = (<object> self).get_flight_info(ServerCallContext.wrap(context), |
| py_descriptor) |
| if not isinstance(result, FlightInfo): |
| raise TypeError("FlightServerBase.get_flight_info must return " |
| "a FlightInfo instance, but got {}".format( |
| type(result))) |
| info.reset(new CFlightInfo(deref((<FlightInfo> result).info.get()))) |
| |
| |
| cdef void _do_put(void* self, const CServerCallContext& context, |
| unique_ptr[CFlightMessageReader] reader, |
| unique_ptr[CFlightMetadataWriter] writer) except *: |
| """Callback for implementing Flight servers in Python.""" |
| cdef: |
| MetadataRecordBatchReader py_reader = MetadataRecordBatchReader() |
| FlightMetadataWriter py_writer = FlightMetadataWriter() |
| FlightDescriptor descriptor = \ |
| FlightDescriptor.__new__(FlightDescriptor) |
| |
| descriptor.descriptor = reader.get().descriptor() |
| py_reader.reader.reset(reader.release()) |
| py_reader.schema = pyarrow_wrap_schema( |
| py_reader.reader.get().schema()) |
| py_writer.writer.reset(writer.release()) |
| (<object> self).do_put(ServerCallContext.wrap(context), descriptor, |
| py_reader, py_writer) |
| |
| |
| cdef void _do_get(void* self, const CServerCallContext& context, |
| CTicket ticket, |
| unique_ptr[CFlightDataStream]* stream) except *: |
| """Callback for implementing Flight servers in Python.""" |
| cdef: |
| unique_ptr[CFlightDataStream] data_stream |
| |
| py_ticket = Ticket(ticket.ticket) |
| result = (<object> self).do_get(ServerCallContext.wrap(context), |
| py_ticket) |
| if not isinstance(result, FlightDataStream): |
| raise TypeError("FlightServerBase.do_get must return " |
| "a FlightDataStream") |
| data_stream = unique_ptr[CFlightDataStream]( |
| (<FlightDataStream> result).to_stream()) |
| stream[0] = unique_ptr[CFlightDataStream]( |
| new CPyFlightDataStream(result, move(data_stream))) |
| |
| |
| cdef void _do_action_result_next(void* self, |
| unique_ptr[CFlightResult]* result) except *: |
| """Callback for implementing Flight servers in Python.""" |
| cdef: |
| CFlightResult* c_result |
| |
| try: |
| action_result = next(<object> self) |
| if not isinstance(action_result, Result): |
| raise TypeError("Result of FlightServerBase.do_action must " |
| "return an iterator of Result objects") |
| c_result = (<Result> action_result).result.get() |
| result.reset(new CFlightResult(deref(c_result))) |
| except StopIteration: |
| result.reset(nullptr) |
| |
| |
| cdef void _do_action(void* self, const CServerCallContext& context, |
| const CAction& action, |
| unique_ptr[CResultStream]* result) except *: |
| """Callback for implementing Flight servers in Python.""" |
| cdef: |
| function[cb_result_next] ptr = &_do_action_result_next |
| py_action = Action(action.type, pyarrow_wrap_buffer(action.body)) |
| responses = (<object> self).do_action(ServerCallContext.wrap(context), |
| py_action) |
| result.reset(new CPyFlightResultStream(responses, ptr)) |
| |
| |
| cdef void _list_actions(void* self, const CServerCallContext& context, |
| vector[CActionType]* actions) except *: |
| """Callback for implementing Flight servers in Python.""" |
| cdef: |
| CActionType action_type |
| # Method should return a list of ActionTypes or similar tuple |
| result = (<object> self).list_actions(ServerCallContext.wrap(context)) |
| for action in result: |
| action_type.type = tobytes(action[0]) |
| action_type.description = tobytes(action[1]) |
| actions.push_back(action_type) |
| |
| |
| cdef void _server_authenticate(void* self, CServerAuthSender* outgoing, |
| CServerAuthReader* incoming) except *: |
| """Callback for implementing authentication in Python.""" |
| sender = ServerAuthSender.wrap(outgoing) |
| reader = ServerAuthReader.wrap(incoming) |
| try: |
| (<object> self).authenticate(sender, reader) |
| finally: |
| sender.poison() |
| reader.poison() |
| |
| |
| cdef void _is_valid(void* self, const c_string& token, |
| c_string* peer_identity) except *: |
| """Callback for implementing authentication in Python.""" |
| cdef c_string c_result |
| c_result = tobytes((<object> self).is_valid(token)) |
| peer_identity[0] = c_result |
| |
| |
| cdef void _client_authenticate(void* self, CClientAuthSender* outgoing, |
| CClientAuthReader* incoming) except *: |
| """Callback for implementing authentication in Python.""" |
| sender = ClientAuthSender.wrap(outgoing) |
| reader = ClientAuthReader.wrap(incoming) |
| try: |
| (<object> self).authenticate(sender, reader) |
| finally: |
| sender.poison() |
| reader.poison() |
| |
| |
| cdef void _get_token(void* self, c_string* token) except *: |
| """Callback for implementing authentication in Python.""" |
| cdef c_string c_result |
| c_result = tobytes((<object> self).get_token()) |
| token[0] = c_result |
| |
| |
| cdef class ServerAuthHandler: |
| """Authentication middleware for a server. |
| |
| To implement an authentication mechanism, subclass this class and |
| override its methods. |
| |
| """ |
| |
| def authenticate(self, outgoing, incoming): |
| """Conduct the handshake with the client. |
| |
| May raise an error if the client cannot authenticate. |
| |
| Parameters |
| ---------- |
| outgoing : ServerAuthSender |
| A channel to send messages to the client. |
| incoming : ServerAuthReader |
| A channel to read messages from the client. |
| """ |
| raise NotImplementedError |
| |
| def is_valid(self, token): |
| """Validate a client token, returning their identity. |
| |
| May return an empty string (if the auth mechanism does not |
| name the peer) or raise an exception (if the token is |
| invalid). |
| |
| Parameters |
| ---------- |
| token : bytes |
| The authentication token from the client. |
| |
| """ |
| raise NotImplementedError |
| |
| cdef PyServerAuthHandler* to_handler(self): |
| cdef PyServerAuthHandlerVtable vtable |
| vtable.authenticate = _server_authenticate |
| vtable.is_valid = _is_valid |
| return new PyServerAuthHandler(self, vtable) |
| |
| |
| cdef class ClientAuthHandler: |
| """Authentication plugin for a client.""" |
| |
| def authenticate(self, outgoing, incoming): |
| """Conduct the handshake with the server. |
| |
| Parameters |
| ---------- |
| outgoing : ClientAuthSender |
| A channel to send messages to the server. |
| incoming : ClientAuthReader |
| A channel to read messages from the server. |
| """ |
| raise NotImplementedError |
| |
| def get_token(self): |
| """Get the auth token for a call.""" |
| raise NotImplementedError |
| |
| cdef PyClientAuthHandler* to_handler(self): |
| cdef PyClientAuthHandlerVtable vtable |
| vtable.authenticate = _client_authenticate |
| vtable.get_token = _get_token |
| return new PyClientAuthHandler(self, vtable) |
| |
| |
| cdef class FlightServerBase: |
| """A Flight service definition. |
| |
| Override methods to define your Flight service. |
| |
| """ |
| |
| cdef: |
| unique_ptr[PyFlightServer] server |
| |
| def run(self, location, auth_handler=None, tls_certificates=None): |
| """Start this server. |
| |
| Parameters |
| ---------- |
| location : Location |
| auth_handler : ServerAuthHandler |
| An authentication mechanism to use. May be None. |
| tls_certificates : list |
| A list of (certificate, key) pairs. |
| """ |
| cdef: |
| PyFlightServerVtable vtable = PyFlightServerVtable() |
| PyFlightServer* c_server |
| unique_ptr[CFlightServerOptions] c_options |
| CCertKeyPair c_cert |
| |
| c_options.reset(new CFlightServerOptions(Location.unwrap(location))) |
| |
| if auth_handler: |
| if not isinstance(auth_handler, ServerAuthHandler): |
| raise TypeError("auth_handler must be a ServerAuthHandler, " |
| "not a '{}'".format(type(auth_handler))) |
| c_options.get().auth_handler.reset( |
| (<ServerAuthHandler> auth_handler).to_handler()) |
| |
| if tls_certificates: |
| for cert, key in tls_certificates: |
| c_cert.pem_cert = tobytes(cert) |
| c_cert.pem_key = tobytes(key) |
| c_options.get().tls_certificates.push_back(c_cert) |
| |
| vtable.list_flights = &_list_flights |
| vtable.get_flight_info = &_get_flight_info |
| vtable.do_put = &_do_put |
| vtable.do_get = &_do_get |
| vtable.list_actions = &_list_actions |
| vtable.do_action = &_do_action |
| |
| c_server = new PyFlightServer(self, vtable) |
| self.server.reset(c_server) |
| with nogil: |
| check_status(c_server.Init(deref(c_options))) |
| check_status(c_server.ServeWithSignals()) |
| |
| def list_flights(self, context, criteria): |
| raise NotImplementedError |
| |
| def get_flight_info(self, context, descriptor): |
| raise NotImplementedError |
| |
| def do_put(self, context, descriptor, reader, |
| writer: FlightMetadataWriter): |
| raise NotImplementedError |
| |
| def do_get(self, context, ticket): |
| raise NotImplementedError |
| |
| def list_actions(self, context): |
| raise NotImplementedError |
| |
| def do_action(self, context, action): |
| raise NotImplementedError |
| |
| def shutdown(self): |
| """Shut down the server, blocking until current requests finish. |
| |
| Do not call this directly from the implementation of a Flight |
| method, as then the server will block forever waiting for that |
| request to finish. Instead, call this method from a background |
| thread. |
| """ |
| # Must not hold the GIL: shutdown waits for pending RPCs to |
| # complete. Holding the GIL means Python-implemented Flight |
| # methods will never get to run, so this will hang |
| # indefinitely. |
| with nogil: |
| if self.server.get() != NULL: |
| self.server.get().Shutdown() |
