def process_add_file(db, ch, method, properties, body):
info = json.loads(body)
hash = info['hash']
file_ref = info['file']
treename = info['treename']
# Update state. Just silently ignore if this thing isn't there.
state = db.lookup_results(hash)
if state is not None:
new_files = list(state.files)
new_files.append((file_ref, treename))
new_files = list(set(new_files))
new_done = len(new_files) >= state.jobs
new_state = ADLRequestInfo(done=new_done,
files=new_files,
jobs=state.jobs,
phase=state.phase,
hash=state.hash if not new_done else 'done',
log=state.log,
message=state.message)
db.save_results(hash, new_state)
else:
print(
f'Unable to find an entry for hash {hash}. Ignoring adding file {file_ref}.'
)
ch.basic_ack(delivery_tag=method.delivery_tag)
def process_message(db, ch, method, properties, body):
'Process the incoming message'
# Get the AST out of the body of the message.
a = pickle.loads(body)
if a is None or not isinstance(a, ast.AST):
logging.warning(f"Body of message wasn't of type AST: {a}")
ch.basic_ack(delivery_tag=method.delivery_tag)
return
# Great. Next - see if we know about this already.
status = db.lookup_results(a)
# If we know nothing about this, then fire off a new task
if status is None:
status = db.save_results(
a,
ADLRequestInfo(done=False,
files=[],
jobs=-1,
phase='waiting_for_data',
hash='',
log=None,
message=None))
finder_message = {
'hash': status.hash,
'ast': base64.b64encode(pickle.dumps(a)).decode(),
}
logging.info(f'Running new request: {status.hash}')
ch.basic_publish(exchange='',
routing_key='find_did',
body=json.dumps(finder_message))
else:
logging.info(
f'Request already running: {status.hash} Phase: {status.phase} Files: {status.files}'
)
# Next, we have to let everyone know the thing is off and going (or done, or whatever).
ch.basic_publish(exchange='',
routing_key=properties.reply_to,
properties=pika.BasicProperties(
correlation_id=properties.correlation_id),
body=json.dumps({
'files': status.files,
'phase': status.phase,
'done': status.done,
'jobs': status.jobs,
'log': status.log,
'message': status.message
}))
# Done!
ch.basic_ack(delivery_tag=method.delivery_tag)
def test_hash_update(empty_db):
db = FuncADLDBAccess(empty_db)
db.save_results(
'bogus1',
ADLRequestInfo(done=False,
files=['file://root1.root'],
jobs=5,
phase='running',
hash=''))
db.save_results(
'bogus1',
ADLRequestInfo(done=False,
files=['file://root1.root'],
jobs=1,
phase='done',
hash=''))
r = db.lookup_results('bogus1')
assert r is not None
assert r.done == False
assert len(r.files) == 1
assert r.files[0] == 'file://root1.root'
assert r.jobs == 1
assert r.phase == 'done'
def test_ast_is_there(empty_db):
a = ast.Add()
db = FuncADLDBAccess(empty_db)
db.save_results(
a,
ADLRequestInfo(done=False,
files=['file://root1.root'],
jobs=5,
phase='downloading',
hash=''))
r = db.lookup_results(a)
assert r is not None
assert r.done == False
assert len(r.files) == 1
assert r.files[0] == 'file://root1.root'
assert r.jobs == 5
assert r.phase == 'downloading'
def process_number_jobs(db, ch, method, properties, body):
info = json.loads(body)
hash = info['hash']
new_n_jobs = info['njobs']
# Update state. Just silently ignore if this thing isn't there.
state = db.lookup_results(hash)
if state is not None:
new_state = ADLRequestInfo(done=state.done,
files=state.files,
jobs=new_n_jobs,
phase=state.phase,
hash=state.hash,
log=state.log,
message=state.message)
db.save_results(hash, new_state)
else:
print(
f'Unable to find an entry for hash {hash} to update it to state {new_phase}.'
)
ch.basic_ack(delivery_tag=method.delivery_tag)