使用PyArrow Flight连接Dremio时出现Could not finish writing before closing错误

使用基于PyArrow Flight的Python代码连接Dremio执行查询，此前运行正常，近期调用client.authenticate_basic_token()时出现错误：

E pyarrow._flight.FlightInternalError: Could not finish writing before closing

相关代码如下：

class DremioConnector:
    env: str
    auth_token: str

    def __init__(self, env: str,  auth_token: str):
        self.env = env
        self.auth_token = auth_token

    def get_dremio_client(self):
        dremio_url = get_url_from_template(DREMIO_URL, self.env)

        return FlightClient(f"grpc+tls://{dremio_url}:*****", tls_root_certs=get_cert(),
                            middleware=[DremioClientAuthMiddlewareFactory(), CookieMiddlewareFactory()], **{}, )

    def get_dremio_credentials(self):
       
        ...
        ...

        return username, password

    def __create_flight_call_options(self, username: str, password: str, client: FlightClient) -> FlightCallOptions:
        headers: list[Any] = []

        bearer_token = client.authenticate_basic_token(username, password, FlightCallOptions(headers=headers))
        headers.append(bearer_token)
        return FlightCallOptions(headers=headers)

    def run_query(self, query: str, username: str, password: str, client: FlightClient) -> FlightStreamReader:
        flight_desc = FlightDescriptor.for_command(query)
        flight_info = client.get_flight_info(flight_desc, self.__create_flight_call_options(username, password, client))
        reader = client.do_get(flight_info.endpoints[0].ticket, self.__create_flight_call_options(username, password, client))
        return reader


connector = DremioConnector('env', "auth_token")

dremio_client = connector.get_dremio_client()

• 认证中间件冲突：创建FlightClient时同时添加了DremioClientAuthMiddlewareFactory和CookieMiddlewareFactory，又手动调用authenticate_basic_token做认证，双重认证流程会打乱gRPC流的写入/关闭逻辑，触发该错误。
• 重复认证导致流异常：__create_flight_call_options每次被调用都会重新执行认证，重复的认证请求可能导致连接池中的连接被异常关闭，gRPC请求流未正确完成。
• 版本兼容性问题：近期PyArrow或Dremio服务端版本升级后，Flight协议的认证流程发生变化，旧逻辑不再适配新的协议实现。
• 网络或超时问题：网络波动、服务端超时设置变更，导致认证请求的gRPC流还未完成写入就被关闭。

1. 移除冲突的认证中间件
   如果手动使用authenticate_basic_token做认证，无需保留DremioClientAuthMiddlewareFactory，仅保留CookieMiddlewareFactory即可：

def get_dremio_client(self):
    dremio_url = get_url_from_template(DREMIO_URL, self.env)
    return FlightClient(f"grpc+tls://{dremio_url}:*****", 
                        tls_root_certs=get_cert(),
                        middleware=[CookieMiddlewareFactory()])

2. 缓存认证Token
   避免每次调用都重新认证，将Token缓存到类实例中，仅在首次调用或过期时重新获取：

class DremioConnector:
    env: str
    auth_token: str
    _bearer_token = None  # 缓存Token

    # ... 其他方法不变 ...

    def __create_flight_call_options(self, username: str, password: str, client: FlightClient) -> FlightCallOptions:
        if self._bearer_token is None:
            headers: list[Any] = []
            self._bearer_token = client.authenticate_basic_token(username, password, FlightCallOptions(headers=headers))
        return FlightCallOptions(headers=[self._bearer_token])

3. 检查版本兼容性
   确认PyArrow版本与Dremio服务端版本匹配，参考官方兼容矩阵调整版本，必要时升级或降级PyArrow。

4. 添加超时处理
   在创建FlightCallOptions时设置超时，避免因请求超时导致流异常关闭：

from pyarrow.flight import FlightCallOptions, Deadline

def __create_flight_call_options(self, username: str, password: str, client: FlightClient) -> FlightCallOptions:
    if self._bearer_token is None:
        headers: list[Any] = []
        # 设置10秒超时
        call_options = FlightCallOptions(headers=headers, deadline=Deadline.from_timeout(10))
        self._bearer_token = client.authenticate_basic_token(username, password, call_options)
    return FlightCallOptions(headers=[self._bearer_token], deadline=Deadline.from_timeout(10))

内容的提问来源于stack exchange，提问作者dom159