import ast

import base64

import itertools

import os

import pathlib

import signal

import struct

import subprocess

import sys

import tempfile

import threading

import time

import traceback

import json

from datetime import datetime

try:

import numpy as np

except ImportError:

np = None

import pytest

import pyarrow as pa

from pyarrow.lib import IpcReadOptions, ReadStats, is_opentelemetry_enabled, tobytes

from pyarrow.util import find_free_port

from pyarrow.tests import util

try:

from pyarrow import flight

from pyarrow.flight import (

FlightClient, FlightServerBase,

ServerAuthHandler, ClientAuthHandler,

ServerMiddleware, ServerMiddlewareFactory,

ClientMiddleware, ClientMiddlewareFactory,

FlightCallOptions,

)

except ImportError:

flight = None

FlightClient, FlightServerBase = object, object

ServerAuthHandler, ClientAuthHandler = object, object

ServerMiddleware, ServerMiddlewareFactory = object, object

ClientMiddleware, ClientMiddlewareFactory = object, object

FlightCallOptions = object

pytestmark = pytest.mark.flight

def test_import():

# So we see the ImportError somewhere

import pyarrow.flight # noqa

def resource_root():

return pathlib.Path(os.environ["ARROW_TEST_DATA"]) / "flight"

def read_flight_resource(path):

)

def example_tls_certs():

}

def simple_ints_table():

return pa.Table.from_arrays(data, names=['some_ints'])

def simple_dicts_table():

return pa.Table.from_arrays(data, names=['some_dicts'])

def multiple_column_table():

return pa.Table.from_arrays([pa.array(['foo', 'bar', 'baz', 'qux']),

pa.array([1, 2, 3, 4])],

names=['a', 'b'])

class ConstantFlightServer(FlightServerBase):

return flight.RecordBatchStream(table, options=self.options)

class MetadataFlightServer(FlightServerBase):

yield batch, buf

class EchoFlightServer(FlightServerBase):

pass

class EchoStreamFlightServer(EchoFlightServer):

raise NotImplementedError

class EchoTableStreamFlightServer(EchoFlightServer):

raise NotImplementedError

class EchoRecordBatchReaderStreamFlightServer(EchoFlightServer):

raise NotImplementedError

class GetInfoFlightServer(FlightServerBase):

return flight.SchemaResult(info.schema)

class ListActionsFlightServer(FlightServerBase):

yield from self.expected_actions()

class ListActionsErrorFlightServer(FlightServerBase):

yield "foo"

class CheckTicketFlightServer(FlightServerBase):

self.last_message = reader.read_all()

class InvalidStreamFlightServer(FlightServerBase):

return flight.GeneratorStream(self.schema, [table1, table2])

class NeverSendsDataFlightServer(FlightServerBase):

self.schema, itertools.repeat(self.schema.empty_table()))

class SlowFlightServer(FlightServerBase):

yield pa.Table.from_arrays(data1, names=['a'])

class ErrorFlightServer(FlightServerBase):

raise flight.FlightUnauthorizedError("foo", err_msg)

class ExchangeFlightServer(FlightServerBase):

writer.write_table(result)

class HttpBasicServerAuthHandler(ServerAuthHandler):

return token

class HttpBasicClientAuthHandler(ClientAuthHandler):

return self.token

class TokenServerAuthHandler(ServerAuthHandler):

return token[7:]

class TokenClientAuthHandler(ClientAuthHandler):

return self.token

class NoopAuthHandler(ServerAuthHandler):

return ""

def case_insensitive_header_lookup(headers, lookup_key):

return headers.get(key)

class ClientHeaderAuthMiddlewareFactory(ClientMiddlewareFactory):

self.call_credential = call_credential

class ClientHeaderAuthMiddleware(ClientMiddleware):

auth_header[0].encode("utf-8")])

class HeaderAuthServerMiddlewareFactory(ServerMiddlewareFactory):

return HeaderAuthServerMiddleware(token)

class HeaderAuthServerMiddleware(ServerMiddleware):

return {'authorization': 'Bearer ' + self.token}

class HeaderAuthFlightServer(FlightServerBase):

'No token auth middleware found.')

class ArbitraryHeadersServerMiddlewareFactory(ServerMiddlewareFactory):

return ArbitraryHeadersServerMiddleware(headers)

class ArbitraryHeadersServerMiddleware(ServerMiddleware):

return self.incoming

class ArbitraryHeadersFlightServer(FlightServerBase):

raise flight.FlightServerError("No headers middleware found")

class HeaderServerMiddleware(ServerMiddleware):

self.special_value = special_value

class HeaderServerMiddlewareFactory(ServerMiddlewareFactory):

return HeaderServerMiddleware("right value")

class HeaderFlightServer(FlightServerBase):

return [b""]

class MultiHeaderFlightServer(FlightServerBase):

return [headers]

class SelectiveAuthServerMiddlewareFactory(ServerMiddlewareFactory):

return HeaderServerMiddleware(token)

class SelectiveAuthClientMiddlewareFactory(ClientMiddlewareFactory):

def start_call(self, info):

return SelectiveAuthClientMiddleware()

class SelectiveAuthClientMiddleware(ClientMiddleware):

}

class RecordingServerMiddlewareFactory(ServerMiddlewareFactory):

return None

class RecordingClientMiddlewareFactory(ClientMiddlewareFactory):

return None

class MultiHeaderClientMiddlewareFactory(ClientMiddlewareFactory):

return MultiHeaderClientMiddleware(self)

class MultiHeaderClientMiddleware(ClientMiddleware):

self.factory.last_headers.update(headers)

class MultiHeaderServerMiddlewareFactory(ServerMiddlewareFactory):

return MultiHeaderServerMiddleware(headers)

class MultiHeaderServerMiddleware(ServerMiddleware):

return MultiHeaderClientMiddleware.EXPECTED

class LargeMetadataFlightServer(FlightServerBase):

writer.write_metadata(self._metadata)

def test_repr():

flight.FlightEndpoint("foo", [], app_metadata=object())