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 test
pull/23390/head
Joost Wooning 2022-01-04 17:01:33 +01:00 committed by GitHub
parent aad7ebdc2a
commit 47bb62b875
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 28 additions and 15 deletions

View File

@ -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}

View File

@ -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