96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
from flask import Flask
|
|
from threading import Lock
|
|
import pika
|
|
import pika.exceptions
|
|
import json
|
|
import time
|
|
|
|
|
|
class MagicAMPQ:
|
|
"""
|
|
This is my pathetic attempt to make RabbitMQ connection in a Flask app reliable and performant.
|
|
"""
|
|
|
|
def __init__(self, app: Flask = None):
|
|
self.app = app
|
|
if app:
|
|
self.init_app(app)
|
|
|
|
self._lock = Lock()
|
|
self._credentials = None
|
|
|
|
def init_app(self, app: Flask):
|
|
self.app = app
|
|
self.app.config.setdefault('FLASK_PIKA_PARAMS', {})
|
|
self.app.config.setdefault('EXCHANGE_NAME', None)
|
|
self.app.config.setdefault('RABBITMQ_QUEUE', None)
|
|
|
|
self._credentials = pika.PlainCredentials(
|
|
app.config['FLASK_PIKA_PARAMS']['username'],
|
|
app.config['FLASK_PIKA_PARAMS']['password']
|
|
)
|
|
|
|
self._reconnect_ampq()
|
|
|
|
def _reconnect_ampq(self):
|
|
self._pika_connection = pika.BlockingConnection(
|
|
pika.ConnectionParameters(
|
|
host=self.app.config['FLASK_PIKA_PARAMS']['host'],
|
|
credentials=self._credentials,
|
|
heartbeat=10,
|
|
socket_timeout=5)
|
|
)
|
|
self._pika_channel = self._pika_connection.channel()
|
|
self._pika_channel.exchange_declare(
|
|
exchange=self.app.config['EXCHANGE_NAME'],
|
|
exchange_type='direct'
|
|
)
|
|
|
|
def loop(self):
|
|
"""
|
|
This method should be called periodically to keep up the connection
|
|
"""
|
|
with self._lock:
|
|
try:
|
|
self._pika_connection.process_data_events(0)
|
|
# We won't attempt retry if this fail
|
|
except pika.exceptions.AMQPConnectionError:
|
|
self._reconnect_ampq()
|
|
|
|
def publish(self, payload=None):
|
|
"""
|
|
Publish a simple json serialized message to the configured queue.
|
|
If the connection is broken, then this call will block until the connection is restored
|
|
"""
|
|
with self._lock:
|
|
tries = 0
|
|
while True:
|
|
try:
|
|
self._pika_channel.basic_publish(
|
|
exchange=self.app.config['EXCHANGE_NAME'],
|
|
routing_key='feature',
|
|
body=json.dumps(payload).encode('UTF-8')
|
|
)
|
|
break # message sent successfully
|
|
except pika.exceptions.AMQPConnectionError:
|
|
|
|
if tries > 30:
|
|
raise # just give up
|
|
|
|
while True:
|
|
try:
|
|
self._reconnect_ampq()
|
|
break
|
|
except pika.exceptions.AMQPConnectionError:
|
|
tries += 1
|
|
|
|
if tries > 30:
|
|
raise # just give up
|
|
|
|
if tries > 10:
|
|
time.sleep(2)
|
|
|
|
|
|
# instance to be used in the flask app
|
|
magic_ampq = MagicAMPQ()
|