def __init__(self,
host=None,
port=None,
socket_timeout=None,
socket_connect_timeout=None,
socket_keepalive=None,
socket_keepalive_options=None,
connection=None,
**kwargs):
if not connection:
self.socket_timeout = socket_timeout
kwargs2 = {
'host': host,
'port': int(port),
'socket_connect_timeout': socket_connect_timeout,
'socket_keepalive': socket_keepalive,
'socket_keepalive_options': socket_keepalive_options,
'socket_timeout': socket_timeout,
}
connection = Connection(**kwargs2)
self.conn_queue = LifoQueue()
self.conn_queue.put_nowait(connection)
self._connection = connection
self.response_callbacks = self.__class__.RESPONSE_CALLBACKS.copy()
self.expected_ok = self.__class__.EXPECTED_OK.copy()
self.expected_err = self.__class__.EXPECTED_ERR.copy()
class Beanstalk(object):
RESPONSE_CALLBACKS = dict_merge({
'list-tubes': parse_yaml,
'peek': parse_body,
'peek-buried': parse_body,
'peek-delayed': parse_body,
'peek-ready': parse_body,
'reserve': parse_body,
'reserve-with-timeout': parse_body,
'stats': parse_yaml,
'stats-tube': parse_yaml
})
EXPECTED_OK = dict_merge({
'bury': ['BURIED'],
'delete': ['DELETED'],
'list-tubes': ['OK'],
'kick': ['KICKED'],
'kick-job': ['KICKED'],
'peek': ['FOUND'],
'peek-buried': ['FOUND'],
'peek-delayed': ['FOUND'],
'peek-ready': ['FOUND'],
'put': ['INSERTED'],
'release': ['RELEASED'],
'reserve': ['RESERVED'],
'reserve-with-timeout': ['RESERVED'],
'stats': ['OK'],
'stats-tube': ['OK'],
'use': ['USING'],
'watch': ['WATCHING'],
})
EXPECTED_ERR = dict_merge({
'bury': ['NOT_FOUND', 'OUT_OF_MEMORY'],
'delete': ['NOT_FOUND'],
'list-tubes': [],
'kick': ['OUT_OF_MEMORY'],
'kick-job': ['NOT_FOUND', 'OUT_OF_MEMORY'],
'peek': ['NOT_FOUND'],
'peek-buried': ['NOT_FOUND'],
'peek-delayed': ['NOT_FOUND'],
'peek-ready': ['NOT_FOUND'],
'put': ['JOB_TOO_BIG', 'BURIED', 'DRAINING', 'OUT_OF_MEMORY'],
'reserve': ['DEADLINE_SOON', 'TIMED_OUT'],
'reserve-with-timeout': ['DEADLINE_SOON', 'TIMED_OUT'],
'release': ['BURIED', 'NOT_FOUND', 'OUT_OF_MEMORY'],
'stats': [],
'stats-tube': ['NOT_FOUND'],
'use': [],
'watch': [],
})
@classmethod
def from_url(cls, url, **kwargs):
if url is None or not url:
raise ConnectionError('Empty URL')
if not url.startswith('beanstalk://'):
import warnings
warnings.warn(
'Invalid URL scheme, expecting beanstalk',
DeprecationWarning)
connection = Connection.from_url(url, **kwargs)
return cls(connection=connection)
def __init__(self, host=None, port=None, socket_timeout=None,
socket_connect_timeout=None, socket_keepalive=None,
socket_keepalive_options=None, connection=None,
**kwargs):
if not connection:
self.socket_timeout = socket_timeout
kwargs2 = {
'host': host,
'port': int(port),
'socket_connect_timeout': socket_connect_timeout,
'socket_keepalive': socket_keepalive,
'socket_keepalive_options': socket_keepalive_options,
'socket_timeout': socket_timeout,
}
connection = Connection(**kwargs2)
self.conn_queue = LifoQueue()
self.conn_queue.put_nowait(connection)
self._connection = connection
self.response_callbacks = self.__class__.RESPONSE_CALLBACKS.copy()
self.expected_ok = self.__class__.EXPECTED_OK.copy()
self.expected_err = self.__class__.EXPECTED_ERR.copy()
def _get_connection(self):
try:
connection = self.conn_queue.get(block=True, timeout=None)
except Empty:
raise ConnectionError("No connection available")
return connection
def _release_connection(self, connection):
self.conn_queue.put_nowait(connection)
def execute_command(self, *args, **kwargs):
connection = self._get_connection()
command_name = args[0]
try:
connection.send_command(*args, **kwargs)
return self.parse_response(connection, command_name, **kwargs)
except (ConnectionError, TimeoutError):
connection.disconnect()
raise
finally:
self._release_connection(connection)
def parse_response(self, connection, command_name, **kwargs):
response = connection.read_response()
status, results = response
if status in self.expected_ok[command_name]:
if command_name in self.response_callbacks:
return self.response_callbacks[command_name](
connection, response, **kwargs)
return response
elif status in self.expected_err[command_name]:
raise ResponseError(command_name, status, results)
else:
raise InvalidResponse(command_name, status, results)
return response
def put(self, body, priority=DEFAULT_PRIORITY, delay=0, ttr=DEFAULT_TTR):
assert isinstance(body, str), 'body must be str'
job_id = self.execute_command('put', priority, delay, ttr, body=body)
return job_id
def use(self, tube):
self._connection.use(tube)
def watch(self, tube):
self._connection.watch(tube)
def reserve(self, timeout=None):
if timeout is not None:
return self.execute_command('reserve-with-timeout', timeout)
else:
return self.execute_command('reserve')
def bury(self, job_id, priority=DEFAULT_PRIORITY):
self.execute_command('bury', job_id, priority)
def release(self, job_id, priority=DEFAULT_PRIORITY, delay=0):
self.execute_command('release', job_id, priority, delay)
def delete(self, job_id):
self.execute_command('delete', job_id)
def _drain(self, fetch_func):
try:
job_id = True
while job_id is not None:
job_id, _ = fetch_func()
self.delete(job_id)
except ResponseError:
pass
def drain_buried(self, tube):
self.use(tube)
return self._drain(self.peek_buried)
def drain_tube(self, tube):
"""Delete all jobs from the specified tube."""
self.watch(tube)
from functools import partial
return self._drain(partial(self.reserve, timeout=0))
def kick_job(self, job_id):
"""
Variant of` kick` that operates with a single job.
:param job_id: the job id to kick
:type job_id: `str`
"""
self.execute_command('kick-job', job_id)
def kick(self, bound=1000):
"""
Move jobs into the ready queue.
If there are any buried jobs, it will only kick buried jobs.
Otherwise it will kick delayed jobs.
:param bound: upper bound on the number of jobs to kick
:type bound: `int`
"""
kicked = int(self.execute_command('kick', str(bound))[1][0])
return kicked
def _peek_generic(self, command_suffix=''):
command = 'peek' + command_suffix
try:
return self.execute_command(command)
except ResponseError as err:
if err.args[0] == command and err.args[1] == 'NOT_FOUND':
return None, None
else:
raise
def peek_buried(self):
"""
Read the next buried job without kicking it.
"""
return self._peek_generic('-buried')
def peek_ready(self):
"""
read the next ready job without reserving it.
"""
return self._peek_generic('-ready')
def wait_until_empty(self, tube, timeout=float('inf'), poll_interval=0.2,
initial_delay=0.0):
"""
Wait until the the specified tube is empty, or the timeout expires.
"""
# TODO(FVE): check tube stats to ensure some jobs have passed through
# and then get rid of the initial_delay
# peek-ready requires "use", not "watch"
self.use(tube)
if initial_delay > 0.0:
time.sleep(initial_delay)
job_id, _ = self.peek_ready()
deadline = time.time() + timeout
while job_id is not None and time.time() < deadline:
time.sleep(poll_interval)
job_id, _ = self.peek_ready()
def wait_for_ready_job(self, tube, timeout=float('inf'),
poll_interval=0.2):
"""
Wait until the the specified tube has a ready job,
or the timeout expires.
"""
self.use(tube)
job_id, data = self.peek_ready()
deadline = time.time() + timeout
while job_id is None and time.time() < deadline:
time.sleep(poll_interval)
job_id, data = self.peek_ready()
return job_id, data
def stats(self):
return self.execute_command('stats')
def stats_tube(self, tube):
return self.execute_command('stats-tube', tube)
def tubes(self):
return self.execute_command('list-tubes')
def close(self):
if self._connection:
self._connection.disconnect()
self._connection = None