def ready(self):
dramatiq.set_encoder(self.select_encoder())
result_backend_settings = self.result_backend_settings()
if result_backend_settings:
result_backend_path = result_backend_settings.get(
"BACKEND", "dramatiq.results.backends.StubBackend")
result_backend_class = load_class(result_backend_path)
result_backend_options = result_backend_settings.get(
"BACKEND_OPTIONS", {})
result_backend = result_backend_class(**result_backend_options)
results_middleware_options = result_backend_settings.get(
"MIDDLEWARE_OPTIONS", {})
results_middleware = Results(backend=result_backend,
**results_middleware_options)
else:
result_backend = None
results_middleware = None
broker_settings = self.broker_settings()
broker_path = broker_settings["BROKER"]
broker_class = load_class(broker_path)
broker_options = broker_settings.get("OPTIONS", {})
middleware = [
load_middleware(path)
for path in broker_settings.get("MIDDLEWARE", [])
]
if result_backend is not None:
middleware.append(results_middleware)
broker = broker_class(middleware=middleware, **broker_options)
dramatiq.set_broker(broker)
def set_default(self):
"""Configure this broker instance to be the global broker instance.
Calls :func:`dramatiq.set_broker` internally.
"""
dramatiq.set_broker(self)
def stub_broker():
broker = StubBroker()
broker.emit_after("process_boot")
dramatiq.set_broker(broker)
yield broker
broker.flush_all()
broker.close()
def configure():
result_backend = RedisBackend(port=63791)
rabbitmq_broker = RabbitmqBroker(port=56721)
rabbitmq_broker.add_middleware(Results(backend=result_backend))
dramatiq.set_broker(rabbitmq_broker)
def schedule(task_file, debug, rabbitmq, redis_url, expire):
try:
config = yaml.safe_load(task_file)
except yaml.YAMLError as e:
raise ClickException(f"Yaml task file is invalid: {e}")
if rabbitmq:
rabbitmq_broker = RabbitmqBroker(url=rabbitmq)
set_broker(rabbitmq_broker)
if debug:
logging.getLogger().setLevel(logging.DEBUG)
logging.getLogger("pika").setLevel(logging.CRITICAL)
scheduler = BlockingScheduler(timezone=utc)
try:
add_all_jobs(scheduler, config["jobs"])
except KeyError as e:
raise ClickException(
f"Config file missing required parameter: {e.args}")
if not redis_url:
scheduler.start()
return
conn = redis.Redis.from_url(redis_url)
with redis_lock.Lock(conn,
LOCK_NAME,
id=PROCESS_KEY,
expire=expire,
auto_renewal=True):
scheduler.start()
def sqlite_broker_in_tests(sqlite_in_tests):
"""Reset dramatiq broker to a new instance pointing at sqlite DB for duration of
tests.
This is required because the dramatiq design is such that a broker needs to be
installed at import-time, and actors declared using the decorator will be pointing
at that broker. Since that happens too early for us to set up our test fixtures,
we need to reinstall another broker pointing at our sqlite DB after we've set
that up.
"""
old_broker = dramatiq.get_broker()
new_broker = Broker()
dramatiq.set_broker(new_broker)
# All actors are currently pointing at old_broker which is not using sqlite.
# This will break a call to .send() on those actors.
# Point them towards new_broker instead which will allow them to work.
actors = []
for actor in old_broker.actors.values():
actors.append(actor)
actor.broker = new_broker
new_broker.declare_actor(actor)
# Everything now points at the sqlite-enabled broker, so proceed with test
yield
# Now roll back our changes
for actor in actors:
actor.broker = old_broker
dramatiq.set_broker(old_broker)
def stub_broker(r):
result_backend = RedisBackend()
result_backend.client = r
broker.add_middleware(Results(backend=result_backend))
broker.flush_all()
dramatiq.set_broker(stub_broker)
return broker
def rabbitmq_broker():
broker = RabbitmqBroker()
check_rabbitmq(broker)
broker.emit_after("process_boot")
dramatiq.set_broker(broker)
yield broker
broker.flush_all()
broker.close()
def rabbitmq_broker():
broker = RabbitmqBroker(host="127.0.0.1", max_priority=10)
check_rabbitmq(broker)
broker.emit_after("process_boot")
dramatiq.set_broker(broker)
yield broker
broker.flush_all()
broker.close()
def broker(sentry_init):
sentry_init(integrations=[DramatiqIntegration()])
broker = StubBroker()
broker.emit_after("process_boot")
dramatiq.set_broker(broker)
yield broker
broker.flush_all()
broker.close()
def redis_broker():
broker = RedisBroker()
broker.client.flushall()
broker.emit_after("process_boot")
dramatiq.set_broker(broker)
yield broker
broker.client.flushall()
broker.close()
def setup_broker(settings):
"""Setup all brokers."""
from dramatiq.brokers.redis import RedisBroker
dramatiq.set_broker(RedisBroker(url=settings['app.broker.url']))
dramatiq.get_broker().add_middleware(ConfigMiddleware())
dramatiq.get_broker().add_middleware(DBSessionMiddleware())
dramatiq.get_broker().add_middleware(ElasticsearchMiddleware())
dramatiq.get_broker().add_middleware(RateLimiterMiddleware())
def set_broker(self, broker):
import dramatiq
self.broker = broker
dramatiq.set_broker(broker)
self.func_to_actor = {}
for func, actor_kwargs in self.func_to_actor_kwargs.items():
self.func_to_actor[func] = dramatiq.actor(**actor_kwargs)(func)
def initialize_broker(broker, results_backend=None):
dramatiq.set_broker(broker)
dramatiq.set_encoder(middleware.TypedEncoder())
if results_backend:
results_middleware = dramatiq.results.Results(backend=results_backend)
broker.add_middleware(results_middleware)
progress_middleware = middleware.ProgressMiddleware(cache.client)
broker.add_middleware(progress_middleware)
def execute():
"""Initialises MQ broker & connects dramatiq library.
"""
# Configure broker.
broker = get_broker()
# Configure dramatiq.
dramatiq.set_broker(broker)
def __init__(self):
# create RabbitMQ broker
self.broker = RabbitmqBroker(host=settings.MESSAGE_QUEUE_HOST,
port=settings.MESSAGE_QUEUE_PORT,
credentials=PlainCredentials(username=settings.MESSAGE_QUEUE_USERNAME,
password=settings.MESSAGE_QUEUE_PASSWORD))
# set the broker as default
set_broker(self.broker)
def setup_worker():
"""Set up Dramatiq with RabbitMQ.
Dramatiq manages the message passing to background workers which run
long tasks to avoid stalling the application responses for too long.
"""
logger.info('creating tasks queue')
broker = RabbitmqBroker(url=f'{settings.RABBITMQ_URL}')
dramatiq.set_broker(broker)
def rabbitmq_broker():
broker = RabbitmqBroker(
host="127.0.0.1",
max_priority=10,
credentials=RABBITMQ_CREDENTIALS,
)
check_rabbitmq(broker)
broker.emit_after("process_boot")
dramatiq.set_broker(broker)
yield broker
broker.flush_all()
broker.close()
def ready(self):
broker_settings = self.broker_settings()
broker_path = broker_settings["BROKER"]
broker_class = load_class(broker_path)
broker_options = broker_settings.get("OPTIONS", {})
middleware = [
load_middleware(path)
for path in broker_settings.get("MIDDLEWARE", [])
]
broker = broker_class(middleware=middleware, **broker_options)
dramatiq.set_broker(broker)
dramatiq.set_encoder(self.select_encoder())
def initialize_context(app: Flask, event_handler: Callable):
dramatiq.set_broker(broker)
broker.add_middleware(AppContextMiddleware(app))
dramatiq_actor = dramatiq.actor(fn=lambda *args, **kwargs: event_handler,
broker=broker,
max_retries=3)
builtins.DramatiqEvent = dramatiq_actor
return dramatiq_actor
def redis_broker(r):
broker = RedisBroker()
broker.client = r
broker.client.flushall()
result_backend = RedisBackend()
result_backend.client = r
broker.add_middleware(Results(backend=result_backend))
broker.emit_after("process_boot")
# monkeypatch.setattr("yatsm.jobs.redis_broker", broker)
dramatiq.set_broker(broker)
yield broker
broker.client.flushall()
broker.close()
def setup_worker():
"""Set up Dramatiq with Redis as a queue.
Dramatiq manages the message passing to background workers which run
long tasks to avoid stalling the application responses for too long.
Redis is used as a message queue for simplicity sake. A more robust
infrastructure could configure a RabbitMQ here, for example.
"""
logger.info('creating tasks queue')
broker = RedisBroker(url=f'{settings.REDIS_URL}')
dramatiq.set_broker(broker)
def setup_redis_broker() -> None:
_connection_pool: redis.BlockingConnectionPool = redis.BlockingConnectionPool(
host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, password=REDIS_PASSWORD
)
_redis_db: redis.StrictRedis = redis.StrictRedis(connection_pool=_connection_pool)
_result_backend = RedisBackend(encoder=PickleEncoder(), client=_redis_db)
_result_middleware = Results(backend=_result_backend)
broker: Broker = RedisBroker(
connection_pool=_connection_pool,
middleware=[_result_middleware],
namespace="lrw",
)
dramatiq.set_broker(broker)
dramatiq.set_encoder(dramatiq.PickleEncoder())
def ready(self):
# this line should be first, to set up global encoder before any middleware is created
dramatiq.set_encoder(self.select_encoder())
broker_settings = self.broker_settings()
broker_path = broker_settings["BROKER"]
broker_class = load_class(broker_path)
broker_options = broker_settings.get("OPTIONS", {})
middleware = [
load_middleware(path)
for path in broker_settings.get("MIDDLEWARE", [])
]
broker = broker_class(middleware=middleware, **broker_options)
self.load_results_backend(broker, broker_settings)
dramatiq.set_broker(broker)
def broker():
broker = SQSBroker(
namespace="dramatiq_sqs_tests",
middleware=[
AgeLimit(),
TimeLimit(),
Callbacks(),
Pipelines(),
Retries(min_backoff=1000, max_backoff=900000, max_retries=96),
],
)
dramatiq.set_broker(broker)
yield broker
for queue in broker.queues.values():
queue.delete()
def confirm_delivery_rabbitmq_broker():
broker = RabbitmqBroker(
host="127.0.0.1",
max_priority=10,
credentials=RABBITMQ_CREDENTIALS,
confirm_delivery=True,
)
check_rabbitmq(broker)
broker.emit_after("process_boot")
dramatiq.set_broker(broker)
yield broker
try:
broker.flush_all()
except Exception:
pass
broker.close()
def execute():
"""Initialises MQ broker & connects dramatiq library.
"""
# Instantiate a broker.
broker = get_broker()
dramatiq.set_broker(broker)
# Inject middleware.
from stests.core.mq.middleware import get_middleware
for mware in get_middleware():
broker.add_middleware(mware)
# Inject encoder.
encoder.initialise()
dramatiq.set_encoder(encoder)
def initialize(cls):
global RATE_LIMITER_BACKEND
dramatiq.set_encoder(cls.select_encoder())
result_backend_settings = cls.result_backend_settings()
if result_backend_settings:
result_backend_path = result_backend_settings.get(
"BACKEND", "dramatiq.results.backends.StubBackend")
result_backend_class = import_string(result_backend_path)
result_backend_options = result_backend_settings.get(
"BACKEND_OPTIONS", {})
result_backend = result_backend_class(**result_backend_options)
results_middleware_options = result_backend_settings.get(
"MIDDLEWARE_OPTIONS", {})
results_middleware = Results(backend=result_backend,
**results_middleware_options)
else:
result_backend = None
results_middleware = None
rate_limiter_backend_settings = cls.rate_limiter_backend_settings()
if rate_limiter_backend_settings:
rate_limiter_backend_path = rate_limiter_backend_settings.get(
"BACKEND", "dramatiq.rate_limits.backends.stub.StubBackend")
rate_limiter_backend_class = import_string(
rate_limiter_backend_path)
rate_limiter_backend_options = rate_limiter_backend_settings.get(
"BACKEND_OPTIONS", {})
RATE_LIMITER_BACKEND = rate_limiter_backend_class(
**rate_limiter_backend_options)
broker_settings = cls.broker_settings()
broker_path = broker_settings["BROKER"]
broker_class = import_string(broker_path)
broker_options = broker_settings.get("OPTIONS", {})
middleware = [
load_middleware(path)
for path in broker_settings.get("MIDDLEWARE", [])
]
if result_backend is not None:
middleware.append(results_middleware)
broker = broker_class(middleware=middleware, **broker_options)
dramatiq.set_broker(broker)
def execute():
"""Initialises MQ broker & connects dramatiq library.
"""
# JIT import to avoid circularity - TODO remove.
from stests.core.mq.middleware import get_middleware
# Configure broker.
broker = get_broker()
for mware in get_middleware():
broker.add_middleware(mware)
# Configure dramatiq.
dramatiq.set_broker(broker)
dramatiq.set_encoder(encoder)
log_event(EventType.CORE_BROKER_CONNECTION_ESTABLISHED, None)
def execute(mode: BrokerMode = BrokerMode.ACTORS):
"""Initialises dramatiq library.
:param mode: Mode in which MQ package is being used.
"""
# Instantiate a broker.
broker = get_broker()
dramatiq.set_broker(broker)
# Inject middleware.
from stests.core.mq.middleware import get_middleware
for mware in get_middleware(mode):
broker.add_middleware(mware)
# Simply broker & encoder.
dramatiq.set_encoder(encoder)