athena: methods for adding and cancelling multiple upload requests (#23366)
* multiple upload cancel * multiple uploads athena method * cleanup * cleanup * more cleanup * isnt used * fix test * actually fix testpull/23390/head
parent
aad7ebdc2a
commit
47bb62b875
|
@ -263,20 +263,29 @@ def reboot():
|
|||
|
||||
@dispatcher.add_method
|
||||
def uploadFileToUrl(fn, url, headers):
|
||||
if len(fn) == 0 or fn[0] == '/' or '..' in fn:
|
||||
return 500
|
||||
path = os.path.join(ROOT, fn)
|
||||
if not os.path.exists(path):
|
||||
return 404
|
||||
return uploadFilesToUrls([[fn, url, headers]])
|
||||
|
||||
item = UploadItem(path=path, url=url, headers=headers, created_at=int(time.time() * 1000), id=None)
|
||||
upload_id = hashlib.sha1(str(item).encode()).hexdigest()
|
||||
item = item._replace(id=upload_id)
|
||||
|
||||
upload_queue.put_nowait(item)
|
||||
@dispatcher.add_method
|
||||
def uploadFilesToUrls(files_data):
|
||||
items = []
|
||||
for fn, url, headers in files_data:
|
||||
if len(fn) == 0 or fn[0] == '/' or '..' in fn:
|
||||
return 500
|
||||
path = os.path.join(ROOT, fn)
|
||||
if not os.path.exists(path):
|
||||
return 404
|
||||
|
||||
item = UploadItem(path=path, url=url, headers=headers, created_at=int(time.time() * 1000), id=None)
|
||||
upload_id = hashlib.sha1(str(item).encode()).hexdigest()
|
||||
items.append(item._replace(id=upload_id))
|
||||
|
||||
for item in items:
|
||||
upload_queue.put_nowait(item)
|
||||
|
||||
UploadQueueCache.cache(upload_queue)
|
||||
|
||||
return {"enqueued": 1, "item": item._asdict()}
|
||||
return {"enqueued": len(items), "items": [i._asdict() for i in items]}
|
||||
|
||||
|
||||
@dispatcher.add_method
|
||||
|
@ -287,11 +296,15 @@ def listUploadQueue():
|
|||
|
||||
@dispatcher.add_method
|
||||
def cancelUpload(upload_id):
|
||||
upload_ids = {item.id for item in list(upload_queue.queue)}
|
||||
if upload_id not in upload_ids:
|
||||
if not isinstance(upload_id, list):
|
||||
upload_id = [upload_id]
|
||||
|
||||
uploading_ids = {item.id for item in list(upload_queue.queue)}
|
||||
cancelled_ids = uploading_ids.intersection(upload_id)
|
||||
if len(cancelled_ids) == 0:
|
||||
return 404
|
||||
|
||||
cancelled_uploads.add(upload_id)
|
||||
cancelled_uploads.update(cancelled_ids)
|
||||
return {"success": 1}
|
||||
|
||||
|
||||
|
|
|
@ -141,8 +141,8 @@ class TestAthenadMethods(unittest.TestCase):
|
|||
|
||||
resp = dispatcher["uploadFileToUrl"]("qlog.bz2", f"{host}/qlog.bz2", {})
|
||||
self.assertEqual(resp['enqueued'], 1)
|
||||
self.assertDictContainsSubset({"path": fn, "url": f"{host}/qlog.bz2", "headers": {}}, resp['item'])
|
||||
self.assertIsNotNone(resp['item'].get('id'))
|
||||
self.assertDictContainsSubset({"path": fn, "url": f"{host}/qlog.bz2", "headers": {}}, resp['items'][0])
|
||||
self.assertIsNotNone(resp['items'][0].get('id'))
|
||||
self.assertEqual(athenad.upload_queue.qsize(), 1)
|
||||
|
||||
@with_http_server
|
||||
|
|
Loading…
Reference in New Issue