have _chunk_actions always return data and actions

This commit is contained in:
Honza Král
2017-07-18 14:18:09 -04:00
parent c8231cad0e
commit f9f4f0c61a
2 changed files with 32 additions and 44 deletions
+20 -39
View File
@@ -47,13 +47,12 @@ def expand_action(data):
return action, data.get('_source', data)
def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer, include_data=False):
def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer):
"""
Split actions into chunks by number or size, serialize them into strings in
the process.
"""
bulk_actions = []
bulk_data = []
bulk_actions, bulk_data = [], []
size, action_count = 0, 0
for action, data in actions:
action = serializer.dumps(action)
@@ -65,34 +64,24 @@ def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer, include_dat
# full chunk, send it and start a new one
if bulk_actions and (size + cur_size > max_chunk_bytes or action_count == chunk_size):
if include_data:
yield bulk_data, bulk_actions
else:
yield bulk_actions
bulk_actions = []
bulk_data = []
yield bulk_data, bulk_actions
bulk_actions, bulk_data = [], []
size, action_count = 0, 0
bulk_actions.append(action)
if data is not None:
bulk_actions.append(data)
if include_data:
if data is not None:
bulk_data.append((action, data))
else:
bulk_data.append((action, ))
bulk_data.append((action, data))
else:
bulk_data.append((action, ))
size += cur_size
action_count += 1
if bulk_actions:
if include_data:
yield bulk_data, bulk_actions
else:
yield bulk_actions
yield bulk_data, bulk_actions
def _process_bulk_chunk(client, bulk_actions, raise_on_exception=True, raise_on_error=True, **kwargs):
def _process_bulk_chunk(client, bulk_actions, bulk_data, raise_on_exception=True, raise_on_error=True, **kwargs):
"""
Send a bulk request to elasticsearch and process the output.
"""
@@ -111,22 +100,14 @@ def _process_bulk_chunk(client, bulk_actions, raise_on_exception=True, raise_on_
err_message = str(e)
exc_errors = []
# deserialize the data back, thisis expensive but only run on
# errors if raise_on_exception is false, so shouldn't be a real
# issue
bulk_data = map(client.transport.serializer.loads, bulk_actions)
while True:
try:
# collect all the information about failed actions
action = next(bulk_data)
op_type, action = action.popitem()
info = {"error": err_message, "status": e.status_code, "exception": e}
if op_type != 'delete':
info['data'] = next(bulk_data)
info.update(action)
exc_errors.append({op_type: info})
except StopIteration:
break
for data in bulk_data:
# collect all the information about failed actions
op_type, action = data[0].popitem()
info = {"error": err_message, "status": e.status_code, "exception": e}
if op_type != 'delete':
info['data'] = data[1]
info.update(action)
exc_errors.append({op_type: info})
# emulate standard behavior for failed actions
if raise_on_error:
@@ -175,8 +156,8 @@ def streaming_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 *
"""
actions = map(expand_action_callback, actions)
for bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer):
for result in _process_bulk_chunk(client, bulk_actions, raise_on_exception, raise_on_error, **kwargs):
for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer):
for result in _process_bulk_chunk(client, bulk_actions, bulk_data, raise_on_exception, raise_on_error, **kwargs):
yield result
def bulk(client, actions, stats_only=False, **kwargs):
@@ -256,7 +237,7 @@ def parallel_bulk(client, actions, thread_count=4, chunk_size=500,
try:
for result in pool.imap(
lambda chunk: list(_process_bulk_chunk(client, chunk, **kwargs)),
lambda bulk_data, bulk_actions: list(_process_bulk_chunk(client, bulk_actions, bulk_data, **kwargs)),
_chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer)
):
for item in result:
+12 -5
View File
@@ -30,16 +30,24 @@ def backoff_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 * 1
"""
actions = map(expand_action_callback, actions)
for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer, include_data=True):
for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size,
max_chunk_bytes,
client.transport.serializer):
retry = 0
while max_retries == -1 or retry <= max_retries:
to_retry = []
to_retry_data = []
to_retry, to_retry_data = [], []
if retry:
time.sleep(min(max_backoff, initial_backoff * 2**(retry-1)))
try:
for data, (ok, info) in zip(bulk_data, _process_bulk_chunk(client, bulk_actions, raise_on_exception=True, raise_on_error=False, **kwargs)):
for data, (ok, info) in zip(bulk_data,
_process_bulk_chunk(client,
bulk_actions,
bulk_data,
raise_on_exception=True,
raise_on_error=False,
**kwargs)):
if not ok:
action, info = info.popitem()
if info['status'] == 429 and (retry+1) <= max_retries:
@@ -52,7 +60,6 @@ def backoff_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 * 1
except TransportError as e:
if e.status_code != 429:
raise
retry += 1
else:
if not to_retry: