3232)
3333from urllib .parse import urlencode , urlparse , parse_qs , urlunparse
3434
35+ import aiofiles
3536import aiohttp_jinja2
3637import jinja2
3738import prometheus_client
4445 MultipartReader ,
4546 ClientSession ,
4647 TCPConnector ,
48+ BodyPartReader ,
4749)
4850from aiohttp .abc import AbstractStreamWriter
4951from aiohttp .hdrs import METH_ANY
@@ -1380,6 +1382,23 @@ def line_to_js(line: ParsedCommandLine) -> Json:
13801382 @timed ("api" , "execute" )
13811383 async def execute (self , request : Request , deps : TenantDependencies ) -> StreamResponse :
13821384 temp_dir : Optional [str ] = None
1385+
1386+ async def write_files (mpr : MultipartReader , tmp_dir : str ) -> Dict [str , str ]:
1387+ files : Dict [str , str ] = {}
1388+ async for part in mpr :
1389+ if isinstance (part , MultipartReader ):
1390+ files .update (await write_files (part , tmp_dir ))
1391+ elif isinstance (part , BodyPartReader ):
1392+ name = part .filename
1393+ if not name :
1394+ raise AttributeError ("Multipart request: content disposition name is required!" )
1395+ path = os .path .join (tmp_dir , rnd_str ()) # use random local path to avoid clashes
1396+ files [name ] = path
1397+ async with aiofiles .open (path , "wb" ) as writer :
1398+ while not part .at_eof ():
1399+ await writer .write (await part .read_chunk ())
1400+ return files
1401+
13831402 try :
13841403 ctx = self .cli_context_from_request (request )
13851404 if request .content_type .startswith ("text" ):
@@ -1388,18 +1407,9 @@ async def execute(self, request: Request, deps: TenantDependencies) -> StreamRes
13881407 command = request .headers ["Fix-Shell-Command" ].strip ()
13891408 temp = tempfile .mkdtemp ()
13901409 temp_dir = temp
1391- files = {}
13921410 # for now, we assume that all multi-parts are file uploads
1393- async for part in MultipartReader (request .headers , request .content ):
1394- name = part .name
1395- if not name :
1396- raise AttributeError ("Multipart request: content disposition name is required!" )
1397- path = os .path .join (temp , rnd_str ()) # use random local path to avoid clashes
1398- files [name ] = path
1399- with open (path , "wb" ) as writer :
1400- while not part .at_eof ():
1401- writer .write (await part .read_chunk ())
1402- ctx = evolve (ctx , uploaded_files = files )
1411+ uploaded = await write_files (MultipartReader (request .headers , request .content ), temp )
1412+ ctx = evolve (ctx , uploaded_files = uploaded )
14031413 else :
14041414 raise AttributeError (f"Not able to handle: { request .content_type } " )
14051415
0 commit comments