Updates
This commit is contained in:
Iliyan Angelov
@@ -0,0 +1,416 @@
|
||||
"""Message migration tools (Broker <-> Broker)."""
|
||||
import socket
|
||||
from functools import partial
|
||||
from itertools import cycle, islice
|
||||
|
||||
from kombu import Queue, eventloop
|
||||
from kombu.common import maybe_declare
|
||||
from kombu.utils.encoding import ensure_bytes
|
||||
|
||||
from celery.app import app_or_default
|
||||
from celery.utils.nodenames import worker_direct
|
||||
from celery.utils.text import str_to_list
|
||||
|
||||
__all__ = (
|
||||
'StopFiltering', 'State', 'republish', 'migrate_task',
|
||||
'migrate_tasks', 'move', 'task_id_eq', 'task_id_in',
|
||||
'start_filter', 'move_task_by_id', 'move_by_idmap',
|
||||
'move_by_taskmap', 'move_direct', 'move_direct_by_id',
|
||||
)
|
||||
|
||||
MOVING_PROGRESS_FMT = """\
|
||||
Moving task {state.filtered}/{state.strtotal}: \
|
||||
{body[task]}[{body[id]}]\
|
||||
"""
|
||||
|
||||
|
||||
class StopFiltering(Exception):
|
||||
"""Semi-predicate used to signal filter stop."""
|
||||
|
||||
|
||||
class State:
|
||||
"""Migration progress state."""
|
||||
|
||||
count = 0
|
||||
filtered = 0
|
||||
total_apx = 0
|
||||
|
||||
@property
|
||||
def strtotal(self):
|
||||
if not self.total_apx:
|
||||
return '?'
|
||||
return str(self.total_apx)
|
||||
|
||||
def __repr__(self):
|
||||
if self.filtered:
|
||||
return f'^{self.filtered}'
|
||||
return f'{self.count}/{self.strtotal}'
|
||||
|
||||
|
||||
def republish(producer, message, exchange=None, routing_key=None,
|
||||
remove_props=None):
|
||||
"""Republish message."""
|
||||
if not remove_props:
|
||||
remove_props = ['application_headers', 'content_type',
|
||||
'content_encoding', 'headers']
|
||||
body = ensure_bytes(message.body) # use raw message body.
|
||||
info, headers, props = (message.delivery_info,
|
||||
message.headers, message.properties)
|
||||
exchange = info['exchange'] if exchange is None else exchange
|
||||
routing_key = info['routing_key'] if routing_key is None else routing_key
|
||||
ctype, enc = message.content_type, message.content_encoding
|
||||
# remove compression header, as this will be inserted again
|
||||
# when the message is recompressed.
|
||||
compression = headers.pop('compression', None)
|
||||
|
||||
expiration = props.pop('expiration', None)
|
||||
# ensure expiration is a float
|
||||
expiration = float(expiration) if expiration is not None else None
|
||||
|
||||
for key in remove_props:
|
||||
props.pop(key, None)
|
||||
|
||||
producer.publish(ensure_bytes(body), exchange=exchange,
|
||||
routing_key=routing_key, compression=compression,
|
||||
headers=headers, content_type=ctype,
|
||||
content_encoding=enc, expiration=expiration,
|
||||
**props)
|
||||
|
||||
|
||||
def migrate_task(producer, body_, message, queues=None):
|
||||
"""Migrate single task message."""
|
||||
info = message.delivery_info
|
||||
queues = {} if queues is None else queues
|
||||
republish(producer, message,
|
||||
exchange=queues.get(info['exchange']),
|
||||
routing_key=queues.get(info['routing_key']))
|
||||
|
||||
|
||||
def filter_callback(callback, tasks):
|
||||
|
||||
def filtered(body, message):
|
||||
if tasks and body['task'] not in tasks:
|
||||
return
|
||||
|
||||
return callback(body, message)
|
||||
return filtered
|
||||
|
||||
|
||||
def migrate_tasks(source, dest, migrate=migrate_task, app=None,
|
||||
queues=None, **kwargs):
|
||||
"""Migrate tasks from one broker to another."""
|
||||
app = app_or_default(app)
|
||||
queues = prepare_queues(queues)
|
||||
producer = app.amqp.Producer(dest, auto_declare=False)
|
||||
migrate = partial(migrate, producer, queues=queues)
|
||||
|
||||
def on_declare_queue(queue):
|
||||
new_queue = queue(producer.channel)
|
||||
new_queue.name = queues.get(queue.name, queue.name)
|
||||
if new_queue.routing_key == queue.name:
|
||||
new_queue.routing_key = queues.get(queue.name,
|
||||
new_queue.routing_key)
|
||||
if new_queue.exchange.name == queue.name:
|
||||
new_queue.exchange.name = queues.get(queue.name, queue.name)
|
||||
new_queue.declare()
|
||||
|
||||
return start_filter(app, source, migrate, queues=queues,
|
||||
on_declare_queue=on_declare_queue, **kwargs)
|
||||
|
||||
|
||||
def _maybe_queue(app, q):
|
||||
if isinstance(q, str):
|
||||
return app.amqp.queues[q]
|
||||
return q
|
||||
|
||||
|
||||
def move(predicate, connection=None, exchange=None, routing_key=None,
|
||||
source=None, app=None, callback=None, limit=None, transform=None,
|
||||
**kwargs):
|
||||
"""Find tasks by filtering them and move the tasks to a new queue.
|
||||
|
||||
Arguments:
|
||||
predicate (Callable): Filter function used to decide the messages
|
||||
to move. Must accept the standard signature of ``(body, message)``
|
||||
used by Kombu consumer callbacks. If the predicate wants the
|
||||
message to be moved it must return either:
|
||||
|
||||
1) a tuple of ``(exchange, routing_key)``, or
|
||||
|
||||
2) a :class:`~kombu.entity.Queue` instance, or
|
||||
|
||||
3) any other true value means the specified
|
||||
``exchange`` and ``routing_key`` arguments will be used.
|
||||
connection (kombu.Connection): Custom connection to use.
|
||||
source: List[Union[str, kombu.Queue]]: Optional list of source
|
||||
queues to use instead of the default (queues
|
||||
in :setting:`task_queues`). This list can also contain
|
||||
:class:`~kombu.entity.Queue` instances.
|
||||
exchange (str, kombu.Exchange): Default destination exchange.
|
||||
routing_key (str): Default destination routing key.
|
||||
limit (int): Limit number of messages to filter.
|
||||
callback (Callable): Callback called after message moved,
|
||||
with signature ``(state, body, message)``.
|
||||
transform (Callable): Optional function to transform the return
|
||||
value (destination) of the filter function.
|
||||
|
||||
Also supports the same keyword arguments as :func:`start_filter`.
|
||||
|
||||
To demonstrate, the :func:`move_task_by_id` operation can be implemented
|
||||
like this:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def is_wanted_task(body, message):
|
||||
if body['id'] == wanted_id:
|
||||
return Queue('foo', exchange=Exchange('foo'),
|
||||
routing_key='foo')
|
||||
|
||||
move(is_wanted_task)
|
||||
|
||||
or with a transform:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def transform(value):
|
||||
if isinstance(value, str):
|
||||
return Queue(value, Exchange(value), value)
|
||||
return value
|
||||
|
||||
move(is_wanted_task, transform=transform)
|
||||
|
||||
Note:
|
||||
The predicate may also return a tuple of ``(exchange, routing_key)``
|
||||
to specify the destination to where the task should be moved,
|
||||
or a :class:`~kombu.entity.Queue` instance.
|
||||
Any other true value means that the task will be moved to the
|
||||
default exchange/routing_key.
|
||||
"""
|
||||
app = app_or_default(app)
|
||||
queues = [_maybe_queue(app, queue) for queue in source or []] or None
|
||||
with app.connection_or_acquire(connection, pool=False) as conn:
|
||||
producer = app.amqp.Producer(conn)
|
||||
state = State()
|
||||
|
||||
def on_task(body, message):
|
||||
ret = predicate(body, message)
|
||||
if ret:
|
||||
if transform:
|
||||
ret = transform(ret)
|
||||
if isinstance(ret, Queue):
|
||||
maybe_declare(ret, conn.default_channel)
|
||||
ex, rk = ret.exchange.name, ret.routing_key
|
||||
else:
|
||||
ex, rk = expand_dest(ret, exchange, routing_key)
|
||||
republish(producer, message,
|
||||
exchange=ex, routing_key=rk)
|
||||
message.ack()
|
||||
|
||||
state.filtered += 1
|
||||
if callback:
|
||||
callback(state, body, message)
|
||||
if limit and state.filtered >= limit:
|
||||
raise StopFiltering()
|
||||
|
||||
return start_filter(app, conn, on_task, consume_from=queues, **kwargs)
|
||||
|
||||
|
||||
def expand_dest(ret, exchange, routing_key):
|
||||
try:
|
||||
ex, rk = ret
|
||||
except (TypeError, ValueError):
|
||||
ex, rk = exchange, routing_key
|
||||
return ex, rk
|
||||
|
||||
|
||||
def task_id_eq(task_id, body, message):
|
||||
"""Return true if task id equals task_id'."""
|
||||
return body['id'] == task_id
|
||||
|
||||
|
||||
def task_id_in(ids, body, message):
|
||||
"""Return true if task id is member of set ids'."""
|
||||
return body['id'] in ids
|
||||
|
||||
|
||||
def prepare_queues(queues):
|
||||
if isinstance(queues, str):
|
||||
queues = queues.split(',')
|
||||
if isinstance(queues, list):
|
||||
queues = dict(tuple(islice(cycle(q.split(':')), None, 2))
|
||||
for q in queues)
|
||||
if queues is None:
|
||||
queues = {}
|
||||
return queues
|
||||
|
||||
|
||||
class Filterer:
|
||||
|
||||
def __init__(self, app, conn, filter,
|
||||
limit=None, timeout=1.0,
|
||||
ack_messages=False, tasks=None, queues=None,
|
||||
callback=None, forever=False, on_declare_queue=None,
|
||||
consume_from=None, state=None, accept=None, **kwargs):
|
||||
self.app = app
|
||||
self.conn = conn
|
||||
self.filter = filter
|
||||
self.limit = limit
|
||||
self.timeout = timeout
|
||||
self.ack_messages = ack_messages
|
||||
self.tasks = set(str_to_list(tasks) or [])
|
||||
self.queues = prepare_queues(queues)
|
||||
self.callback = callback
|
||||
self.forever = forever
|
||||
self.on_declare_queue = on_declare_queue
|
||||
self.consume_from = [
|
||||
_maybe_queue(self.app, q)
|
||||
for q in consume_from or list(self.queues)
|
||||
]
|
||||
self.state = state or State()
|
||||
self.accept = accept
|
||||
|
||||
def start(self):
|
||||
# start migrating messages.
|
||||
with self.prepare_consumer(self.create_consumer()):
|
||||
try:
|
||||
for _ in eventloop(self.conn, # pragma: no cover
|
||||
timeout=self.timeout,
|
||||
ignore_timeouts=self.forever):
|
||||
pass
|
||||
except socket.timeout:
|
||||
pass
|
||||
except StopFiltering:
|
||||
pass
|
||||
return self.state
|
||||
|
||||
def update_state(self, body, message):
|
||||
self.state.count += 1
|
||||
if self.limit and self.state.count >= self.limit:
|
||||
raise StopFiltering()
|
||||
|
||||
def ack_message(self, body, message):
|
||||
message.ack()
|
||||
|
||||
def create_consumer(self):
|
||||
return self.app.amqp.TaskConsumer(
|
||||
self.conn,
|
||||
queues=self.consume_from,
|
||||
accept=self.accept,
|
||||
)
|
||||
|
||||
def prepare_consumer(self, consumer):
|
||||
filter = self.filter
|
||||
update_state = self.update_state
|
||||
ack_message = self.ack_message
|
||||
if self.tasks:
|
||||
filter = filter_callback(filter, self.tasks)
|
||||
update_state = filter_callback(update_state, self.tasks)
|
||||
ack_message = filter_callback(ack_message, self.tasks)
|
||||
consumer.register_callback(filter)
|
||||
consumer.register_callback(update_state)
|
||||
if self.ack_messages:
|
||||
consumer.register_callback(self.ack_message)
|
||||
if self.callback is not None:
|
||||
callback = partial(self.callback, self.state)
|
||||
if self.tasks:
|
||||
callback = filter_callback(callback, self.tasks)
|
||||
consumer.register_callback(callback)
|
||||
self.declare_queues(consumer)
|
||||
return consumer
|
||||
|
||||
def declare_queues(self, consumer):
|
||||
# declare all queues on the new broker.
|
||||
for queue in consumer.queues:
|
||||
if self.queues and queue.name not in self.queues:
|
||||
continue
|
||||
if self.on_declare_queue is not None:
|
||||
self.on_declare_queue(queue)
|
||||
try:
|
||||
_, mcount, _ = queue(
|
||||
consumer.channel).queue_declare(passive=True)
|
||||
if mcount:
|
||||
self.state.total_apx += mcount
|
||||
except self.conn.channel_errors:
|
||||
pass
|
||||
|
||||
|
||||
def start_filter(app, conn, filter, limit=None, timeout=1.0,
|
||||
ack_messages=False, tasks=None, queues=None,
|
||||
callback=None, forever=False, on_declare_queue=None,
|
||||
consume_from=None, state=None, accept=None, **kwargs):
|
||||
"""Filter tasks."""
|
||||
return Filterer(
|
||||
app, conn, filter,
|
||||
limit=limit,
|
||||
timeout=timeout,
|
||||
ack_messages=ack_messages,
|
||||
tasks=tasks,
|
||||
queues=queues,
|
||||
callback=callback,
|
||||
forever=forever,
|
||||
on_declare_queue=on_declare_queue,
|
||||
consume_from=consume_from,
|
||||
state=state,
|
||||
accept=accept,
|
||||
**kwargs).start()
|
||||
|
||||
|
||||
def move_task_by_id(task_id, dest, **kwargs):
|
||||
"""Find a task by id and move it to another queue.
|
||||
|
||||
Arguments:
|
||||
task_id (str): Id of task to find and move.
|
||||
dest: (str, kombu.Queue): Destination queue.
|
||||
transform (Callable): Optional function to transform the return
|
||||
value (destination) of the filter function.
|
||||
**kwargs (Any): Also supports the same keyword
|
||||
arguments as :func:`move`.
|
||||
"""
|
||||
return move_by_idmap({task_id: dest}, **kwargs)
|
||||
|
||||
|
||||
def move_by_idmap(map, **kwargs):
|
||||
"""Move tasks by matching from a ``task_id: queue`` mapping.
|
||||
|
||||
Where ``queue`` is a queue to move the task to.
|
||||
|
||||
Example:
|
||||
>>> move_by_idmap({
|
||||
... '5bee6e82-f4ac-468e-bd3d-13e8600250bc': Queue('name'),
|
||||
... 'ada8652d-aef3-466b-abd2-becdaf1b82b3': Queue('name'),
|
||||
... '3a2b140d-7db1-41ba-ac90-c36a0ef4ab1f': Queue('name')},
|
||||
... queues=['hipri'])
|
||||
"""
|
||||
def task_id_in_map(body, message):
|
||||
return map.get(message.properties['correlation_id'])
|
||||
|
||||
# adding the limit means that we don't have to consume any more
|
||||
# when we've found everything.
|
||||
return move(task_id_in_map, limit=len(map), **kwargs)
|
||||
|
||||
|
||||
def move_by_taskmap(map, **kwargs):
|
||||
"""Move tasks by matching from a ``task_name: queue`` mapping.
|
||||
|
||||
``queue`` is the queue to move the task to.
|
||||
|
||||
Example:
|
||||
>>> move_by_taskmap({
|
||||
... 'tasks.add': Queue('name'),
|
||||
... 'tasks.mul': Queue('name'),
|
||||
... })
|
||||
"""
|
||||
def task_name_in_map(body, message):
|
||||
return map.get(body['task']) # <- name of task
|
||||
|
||||
return move(task_name_in_map, **kwargs)
|
||||
|
||||
|
||||
def filter_status(state, body, message, **kwargs):
|
||||
print(MOVING_PROGRESS_FMT.format(state=state, body=body, **kwargs))
|
||||
|
||||
|
||||
move_direct = partial(move, transform=worker_direct)
|
||||
move_direct_by_id = partial(move_task_by_id, transform=worker_direct)
|
||||
move_direct_by_idmap = partial(move_by_idmap, transform=worker_direct)
|
||||
move_direct_by_taskmap = partial(move_by_taskmap, transform=worker_direct)
|
||||
Reference in New Issue
Block a user