Python asyncio snippet for server streaming RPCs

asyncio_unary_stream_grpc.py

"""
From: https://grpc.io/docs/guides/concepts.html

Server streaming RPCs where the client sends a request to the server and gets a stream to read a sequence of messages back. The client reads from the returned stream until there are no more messages. gRPC guarantees message ordering within an individual RPC call.

rpc LotsOfReplies(HelloRequest) returns (stream HelloResponse){}
"""

import asyncio
import janus
from concurrent.futures.thread import ThreadPoolExecutor

class StreamCompleted:
    pass

def _stream_thread(rpc_method, request, sync_queue):
    try:
        for res in rpc_method(request):
            sync_queue.put(res)
    except:
        print("GRPC Stream Removed")
    finally:
        sync_queue.put(StreamCompleted())
        print("GRPC Stream Finished. Shutting down thread")


async def streaming_response(rpc_method, request):
    event_loop = asyncio.get_event_loop()
    # janus is a thread aware Queue with sync and async api
    thread_safe_queue = janus.Queue(loop=event_loop)
    _ = event_loop.run_in_executor(ThreadPoolExecutor(), _stream_thread, rpc_method, request, thread_safe_queue.sync_q)
    while True:
        response = await thread_safe_queue.async_q.get()
        thread_safe_queue.async_q.task_done()
        if isinstance(response, StreamCompleted):
            return
        else:
            yield response

"""
Example usage:

async for res in streaming_response(client_stub.LotsOfReplies, HelloRequest()):
    print(f"HelloResponse.reply: {res.reply}")
"""