FastAPI websockets not working when using Redis pubsub functionality - python

currently I'm using websockets to pass through data that I receive from a Redis queue (pub/sub). But for some reason the websocket doesn't send messages when using this redis queue.
What my code looks like
My code works as folllow:
I accept the socket connection
I connect to the redis queue
For each message that I receive from the subscription, i sent a message through the socket. (at the moment only text for testing)
#check_route.websocket_route("/check")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
redis = Redis(host='::1', port=6379, db=1)
subscribe = redis.pubsub()
subscribe.subscribe('websocket_queue')
try:
for result in subscribe.listen():
await websocket.send_text('test')
print('test send')
except Exception as e:
await websocket.close()
raise e
The issue with the code
When I'm using this code it's just not sending the message through the socket. But when I accept the websocket within the subscribe.listen() loop it does work but it reconnects every time (see code below).
#check_route.websocket_route("/check")
async def websocket_endpoint(websocket: WebSocket):
redis = Redis(host='::1', port=6379, db=1)
subscribe = redis.pubsub()
subscribe.subscribe('websocket_queue')
try:
for result in subscribe.listen():
await websocket.accept()
await websocket.send_text('test')
print('test send')
except Exception as e:
await websocket.close()
raise e
I think that the subscribe.listen() causes some problems that make the websocket do nothing when websocket.accept() is outside the for loop.
I hope someone knows whats wrong with this.

I'm not sure if this will work, but you could try this:
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
redis = Redis(host='::1', port=6379, db=1)
subscribe = redis.pubsub()
subscribe.subscribe('websocket_queue')
try:
results = await subscribe.listen()
for result in results:
await websocket.send_text('test')
print('test send')
except Exception as e:
await websocket.close()
raise e

After a few days more research I found a solution for this issue. I solved it by using aioredis. This solution is based on the following GitHub Gist.
import json
import aioredis
from fastapi import APIRouter, WebSocket
from app.service.config_service import load_config
check_route = APIRouter()
#check_route.websocket("/check")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
# ---------------------------- REDIS REQUIREMENTS ---------------------------- #
config = load_config()
redis_uri: str = f"redis://{config.redis.host}:{config.redis.port}"
redis_channel = config.redis.redis_socket_queue.channel
redis = await aioredis.create_redis_pool(redis_uri)
# ------------------ SEND SUBSCRIBE RESULT THROUGH WEBSOCKET ----------------- #
(channel,) = await redis.subscribe(redis_channel)
assert isinstance(channel, aioredis.Channel)
try:
while True:
response_raw = await channel.get()
response_str = response_raw.decode("utf-8")
response = json.loads(response_str)
if response:
await websocket.send_json({
"event": 'NEW_CHECK_RESULT',
"data": response
})
except Exception as e:
raise e

Related

Telegram client operating on multiple thread

I'm trying to attach to Telegram APIs to interact through such platform.
I need to send messages on multiple different threads, but I've not found any solution so far allowing me to do that.
Here you can find two sketch codes of my attempts to manage messages on different threads, using both telethon and pyrogram libraries.
So far, I've not found a solution yet.
Telethon case
import time
import threading
from telethon import TelegramClient, events
telegram_client = None
async def telethon_telegram_init(api_id, api_hash):
global telegram_client
client = TelegramClient('my_account', api_id, api_hash)
await client.start()
try: assert await client.connect()
except: pass
if not await client.is_user_authorized():
client.send_code_request(phone_number)
me = client.sign_in(phone_number, input('Enter code: '))
client.parse_mode = 'html'
telegram_client = client
#telegram_client.on(events.NewMessage(incoming=True))
async def telethon_telegram_receive(event):
if event.is_private: await event.respond('Thank you for your message')
def telethon_telegram_send(to, message):
global telegram_client
while True:
telegram_client.send_message(to, message)
time.sleep(60)
# initiating the client object
telegram_client = telethon_telegram_init(api_id, api_hash)
# running the thread aimed to send messages
threading.Thread(target=telethon_telegram_send, args=('me', 'Hello')).start()
Results:
initialization: ok
message reception: not working (no errors are returned)
message sending: not working (AttributeError: 'coroutine' object has no attribute 'send_message' returned)
Pyrogram case
import time
import threading
from pyrogram import Client, filters
def pyrogram_telegram_init(api_id, api_hash):
client = Client('my_account', api_id, api_hash)
#client.on_message(filters.private)
async def pyrogram_telegram_receive(event):
await event.reply_text(f'Thank you for your message')
client.start()
return client
def pyrogram_telegram_send(client, to, message):
while True:
client.send_message(to, message)
time.sleep(60)
# initiating the client object
telegram_client = pyrogram_telegram_init(api_id, api_hash)
# running the thread aimed to send messages
threading.Thread(target=pyrogram_telegram_send, args=(telegram_client, 'me', 'Hello')).start()
Results:
initialization: ok
message reception: not working (no errors are returned)
message sending: not working (no errors are returned)
Update:
Attempts with asyncio
I've also tried to use asyncio, as suggested by #Lonami. I've focused on the telethon library.
Nevertheless, please note that in my case, new Telegram messages are generated by a TCP server listening on the host, in a way similar to the simplified sketch code below.
Here is the code reporting my attempts.
import time
import asyncio
from telethon import TelegramClient, events
RECIPIENT = 'me'
api_id = '...'
api_hash = '...'
telegram_client = None
# server management part
from socketserver import ThreadingMixIn, TCPServer, StreamRequestHandler
class ThreadingTCPServer(ThreadingMixIn, TCPServer): pass
class ServerSample(StreamRequestHandler):
def handle(self):
message = 'New connection from %s:%s' % self.client_address
print(message)
self.server.loop.call_soon_threadsafe(self.server.queue.put_nowait, message)
# also tried with the following:
#self.server.queue.put_nowait(message)
async def initialize_server(loop, queue):
with ThreadingTCPServer(('127.0.0.1', 8080), ServerSample) as server:
server.loop = loop
server.queue = queue
server.serve_forever()
async def telethon_telegram_init(api_id, api_hash):
global telegram_client
client = TelegramClient('my_account', api_id, api_hash)
await client.start()
try: assert await client.connect()
except Exception as e: print(str(e))
if not await client.is_user_authorized():
client.send_code_request(phone_number)
me = client.sign_in(phone_number, input('Enter code: '))
client.parse_mode = 'html'
telegram_client = client
#telegram_client.on(events.NewMessage(incoming=True))
async def telethon_telegram_receive(event):
if event.is_private: await event.respond('Thank you for your message')
return telegram_client
async def telethon_telegram_generate(loop, queue, t):
loop.call_soon_threadsafe(queue.put_nowait, 'Hello from direct {} call'.format(t))
async def main():
global telegram_client
loop = asyncio.get_running_loop()
queue = asyncio.Queue()
# initiating the Telegram client object
await telethon_telegram_init(api_id, api_hash)
# initializing the server (non blocking)
loop.create_task(initialize_server(loop, queue))
# creating a new message (non blocking)
loop.create_task(telethon_telegram_generate(loop, queue, 'non blocking'))
# creating a new message (blocking)
loop.run_in_executor(None, telethon_telegram_generate, loop, queue, 'blocking')
# managing the queue
while True:
message = await queue.get()
print("Sending '{}'...".format(message))
await telegram_client.send_message(RECIPIENT, message)
time.sleep(1)
asyncio.run(main())
In this case, to trigger the generation of a new Telegram message from the TCP server itself, it is needed to launch a curl command like the following one:
curl http://localhost:8080
Nevertheless, results are not promising.
Results:
initialization: ok for telethon_telegram_init and initialize_server, apparently not working for both telethon_telegram_generate method calls
message reception: not working (no errors are returned)
message sending: not working (no errors are returned)
additional notes: after the curl command is executed, a New connection from 127.0.0.1:<random_port> message is printed, but no Telegram messages are sent

How to make Postgres LISTEN async (non blocking) with FastAPI websocket?

I'm trying to make the Postgress LISTEN async with FastAPI so that the WebSocket connection is not blocking while waiting for Postgres table updates.
What I got so far:
router = APIRouter()
#router.websocket("/pg_notify")
async def get_notifications(websocket: WebSocket):
await websocket.accept()
conn = psycopg2.connect("*****")
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
curs = conn.cursor()
curs.execute("LISTEN channel;")
while True:
try:
conn.poll()
while conn.notifies:
print("Waiting for notification...")
notify = await conn.notifies.pop(0)
print(notify.payload)
except Exception as e:
print("exception triggered: ", str(e))
await websocket.close()
This way an exception is raised on the conn.notifies,pop(0):
object psycopg2.extensions.Notify can't be used in the 'await' expression
Finally the below works for me.
while True:
try:
async with aiopg.create_pool(dsn) as pool:
async with pool.acquire() as conn:
listener = listen(conn, "channel")
task_pg_listen = asyncio.create_task(listener)
result = await task_pg_listen

FastAPI with Redis pubsub

I am looking for a solution for using FastAPI with Redis pubsub without using threads. I have looked around for solutions using aioredis, redis-py and web sockets, but I can't seem to make them work.
Would love to have a simple example of a unicorn server running, while subscribed to a Redis channel, printing the message when receiving it. I am using Python3.6.12 so python3.6 solutions are preferred.
Thanks!
Below is the most recent try on aioredis without FASTapi, can't get the print message when publishing on the Redis in another terminal.
import asyncio
import aioredis
import async_timeout
STOPWORD = "STOP"
async def pubsub():
redis = aioredis.Redis.from_url(
"redis://:password#localhost:6379", max_connections=10, decode_responses=True
)
await redis.publish("channel:1", "abc")
psub = redis.pubsub()
async def reader(channel: aioredis.client.PubSub):
while True:
try:
async with async_timeout.timeout(1):
# print("trying to get message")
message = await channel.get_message(ignore_subscribe_messages=True)
if message is not None:
print(f"(Reader) Message Received: {message}")
if message["data"] == STOPWORD:
print("(Reader) STOP")
break
await asyncio.sleep(0.01)
except asyncio.TimeoutError:
pass
async with psub as p:
await p.subscribe("channel:1")
await reader(p) # wait for reader to complete
await p.unsubscribe("channel:1")
# closing all open connections
await psub.close()
if __name__ == '__main__':
loop = asyncio.get_event_loop()
result = loop.run_until_complete(pubsub())
print('finished')

How to properly retry websockets connection in python?

This code tries to connect for 20 seconds. If the connection is established and lost then it tries to reconnect for 10 seconds. If the client can't connect or reconnect then it prints "Failure" and exits. If the server responds with "exit" then the client exits.
import asyncio
import websockets
class Client:
async def connect_with_retries(self, uri: str):
while True:
try:
return await websockets.connect(uri)
except (ConnectionError, websockets.exceptions.WebSocketException):
await asyncio.sleep(5)
async def loop(self, uri: str):
connection_timeout = 20
try:
while True:
websocket = await asyncio.wait_for(self.connect_with_retries(uri), connection_timeout)
connection_timeout = 10
try:
while True:
await websocket.send("req")
resp = await websocket.recv()
if resp == "exit":
return
await asyncio.sleep(1)
except (ConnectionError, websockets.exceptions.WebSocketException) as exc:
pass
finally:
await websocket.close()
except asyncio.TimeoutError:
print("Failure")
client = Client()
asyncio.run(client.loop("ws://localhost:8000"))
I don't like the explicit websocket.close(). How to use asyncio.wait_for with the context manager of the websocket?

Make websocket callback asynchronous with asyncio

I am trying to implement a basic websocket client using asyncio and websockets with Python 3.5.2.
Basically, I want connect_to_dealer to be a blocking call, but wait for the websocket message on a different thread.
After reading some docs (I have very little exp with Python), I concluded that asyncio.ensure_future() passing a coroutine (listen_for_message) was the way to go.
Now, I get to run listen_for_message on a different thread, but from within the coroutine I can't seem to use await or any other mechanism to make the calls synchronous. If I do it, the execution waits forever (it hangs) even for a simple sleep.
I'd like to know what I'm doing wrong.
async def listen_for_message(self, future, websocket):
while (True):
try:
await asyncio.sleep(1) # It hangs here
print('Listening for a message...')
message = await websocket.recv() # If I remove the sleep, hangs here
print("< {}".format(message))
future.set_result(message)
future.done()
except websockets.ConnectionClosed as cc:
print('Connection closed')
except Exception as e:
print('Something happened')
def handle_connect_message(self, future):
# We must first remove the websocket-specific payload because we're only interested in the connect protocol msg
print(future.result)
async def connect_to_dealer(self):
print('connect to dealer')
websocket = await websockets.connect('wss://mywebsocket'))
hello_message = await websocket.recv()
print("< {}".format(hello_message))
# We need to parse the connection ID out of the message
connection_id = hello_message['connectionId']
print('Got connection id {}'.format(connection_id))
sub_response = requests.put('https://subscribetotraffic{user_id}?connection={connection_id}'.format(user_id='username', connection_id=connection_id), headers=headers)
if sub_response.status_code == 200:
print('Now we\'re observing traffic')
else:
print('Oops request failed with {code}'.format(code=sub_response.status_code))
# Now we need to handle messages but continue with the regular execution
try:
future = asyncio.get_event_loop().create_future()
future.add_done_callback(self.handle_connect_message)
asyncio.ensure_future(self.listen_for_message(future, websocket))
except Exception as e:
print(e)
Is there a specific reason you need to work with explicit futures?
With asyncio you can use a combination of coroutines and Tasks to achieve most purposes. Tasks are essentially wrapped coroutines that go about cranking themselves over in the background, independently of other async code, so you don't have to explicitly manage their flow or juggle them with other bits of code.
I am not entirely sure of your end goal, but perhaps the approach elaborated below gives you something to work with:
import asyncio
async def listen_for_message(websocket):
while True:
await asyncio.sleep(0)
try:
print('Listening for a message...')
message = await websocket.recv()
print("< {}".format(message))
except websockets.ConnectionClosed as cc:
print('Connection closed')
except Exception as e:
print('Something happened')
async def connect_to_dealer():
print('connect to dealer')
websocket = await websockets.connect('wss://mywebsocket')
hello_message = await websocket.recv()
print("< {}".format(hello_message))
# We need to parse the connection ID out of the message
connection_id = hello_message['connectionId']
print('Got connection id {}'.format(connection_id))
sub_response = requests.put('https://subscribetotraffic{user_id}?connection={connection_id}'.format(
user_id='username', connection_id=connection_id), headers=headers)
if sub_response.status_code == 200:
print('Now we\'re observing traffic')
else:
print('Oops request failed with {code}'.format(code=sub_response.status_code))
async def my_app():
# this will block until connect_to_dealer() returns
websocket = await connect_to_dealer()
# start listen_for_message() in its own task wrapper, so doing it continues in the background
asyncio.ensure_future(listen_for_message(websocket))
# you can continue with other code here that can now coexist with listen_for_message()
if __name__ == '__main__':
loop = asyncio.get_event_loop()
loop.run_until_complete(my_app())
loop.run_forever()

Categories

Resources