如何在Arrow Flight RPC服务端/客户端代码中调用带参函数？

完全可以基于Arrow Flight RPC实现带参数的多查询接口集成，核心思路是利用Flight Ticket传递查询参数，在服务端解析参数后调用你的SQLAlchemy查询函数，再将查询结果转为Arrow格式返回给客户端。下面是具体实现方案：

1. 服务端改造（集成SQLAlchemy异步查询） #

修改FlightServerBase的do_get方法，解析Ticket中的参数，调用你的get_drives函数，再将SQLAlchemy模型结果转为Arrow Table：

import json
import pyarrow as pa
import pyarrow.flight as flight
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from your_models import Drives  # 替换为你的模型导入路径

# 初始化SQLAlchemy异步引擎和会话工厂
DATABASE_URL = "postgresql+asyncpg://postgres:{passwd}@{host}:5432/local_db"
async_engine = create_async_engine(DATABASE_URL)
AsyncSessionLocal = sessionmaker(async_engine, class_=AsyncSession, expire_on_commit=False)

async def get_drives(id: str, mb_box: str, session: AsyncSession):
    query = session.query(Drives)
    if id:
        result = await query.filter(Drives.id == id).first()
        return [result] if result else []
    if mb_box:
        result = await query.filter(Drives.box == mb_box).all()
        return result
    return await query.all()

class FlightQueryServer(flight.FlightServerBase):
    async def do_get(self, context, ticket):
        # 解析Ticket中的参数（JSON序列化的字节）
        try:
            params = json.loads(ticket.ticket.decode("utf-8"))
        except json.JSONDecodeError:
            raise flight.FlightError("无效的参数格式，需为JSON序列化字节")
        
        # 创建异步会话
        async with AsyncSessionLocal() as session:
            # 调用查询函数
            results = await get_drives(
                id=params.get("id"),
                mb_box=params.get("mb_box"),
                session=session
            )
        
        # 将SQLAlchemy模型转为Arrow Table
        if not results:
            # 空结果返回空表
            schema = pa.schema([
                ("id", pa.string()),
                ("box", pa.string()),
                # 补充你的Drives模型其他字段的类型
            ])
            table = pa.Table.from_arrays([[] for _ in schema.names], schema=schema)
        else:
            # 提取模型字段转为字典列表
            data = [
                {col.name: getattr(obj, col.name) for col in Drives.__table__.columns}
                for obj in results
            ]
            table = pa.Table.from_pylist(data)
        
        return flight.RecordBatchStream(table)

if __name__ == '__main__':
    location = flight.Location.for_grpc_tcp("localhost", 50051)
    server = FlightQueryServer(location)
    print("服务启动于 localhost:50051")
    server.serve()

2. 客户端改造（传递参数并接收结果） #

客户端将查询参数序列化为JSON字节，放入Flight Ticket中发送，然后接收Arrow数据并转为Python对象：

import json
import asyncio
import pyarrow as pa
import pyarrow.flight as flight

async def main():
    # 定义查询参数
    query_params = {
        "mb_box": "box_123"  # 可替换为id参数，比如{"id": "drive_456"}
    }
    
    # 将参数转为JSON字节，创建Ticket
    ticket_bytes = json.dumps(query_params).encode("utf-8")
    ticket = flight.Ticket(ticket_bytes)
    
    # 连接服务端并发送请求
    location = flight.Location.for_grpc_tcp("localhost", 50051)
    async with flight.FlightClient(location) as client:
        reader = await client.do_get(ticket)
        table = await reader.read_all()
    
    # 将Arrow Table转为Python字典列表
    results = table.to_pylist()
    print("查询结果:", results)

if __name__ == '__main__':
    asyncio.run(main())

关键注意事项 #

• 参数序列化：这里用JSON作为参数序列化方式，也可使用MessagePack等其他格式，只要服务端和客户端统一即可。
• 异步兼容：因get_drives是异步函数，服务端需使用SQLAlchemy的异步引擎和会话，同时Flight的do_get支持异步重写。
• 结果转换：SQLAlchemy模型对象需转为Arrow可识别的格式（如字典列表），确保字段类型匹配。
• 异常处理：建议在服务端增加参数校验、数据库异常捕获，返回清晰的错误信息。

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