Use aiofiles where relevant.

This commit is contained in:
grossmj committed 2019-03-06 23:00:01 +07:00
1 parent b0df7ecabf
commit af80b0bb6e
17 files changed
+90 -234

No files matched your search

-22
View File
@@ -320,28 +320,6 @@ class Compute:
raise aiohttp.web.HTTPNotFound(text="{} not found on compute".format(image))
return response
async def stream_file(self, project, path, timeout=None):
"""
Read file of a project and stream it
:param project: A project object
:param path: The path of the file in the project
:param timeout: timeout
:returns: A file stream
"""
url = self._getUrl("/projects/{}/stream/{}".format(project.id, path))
response = await self._session().request("GET", url, auth=self._auth, timeout=timeout)
if response.status == 404:
raise aiohttp.web.HTTPNotFound(text="file '{}' not found on compute".format(path))
elif response.status == 403:
raise aiohttp.web.HTTPForbidden(text="forbidden to open '{}' on compute".format(path))
elif response.status != 200:
raise aiohttp.web.HTTPInternalServerError(text="Unexpected error {}: {}: while opening {} on compute".format(response.status,
response.reason,
path))
return response
async def http_query(self, method, path, data=None, dont_connect=False, **kwargs):
"""
:param dont_connect: If true do not reconnect if not connected
+30 -31
View File
@@ -19,6 +19,7 @@ import os
import sys
import json
import asyncio
import aiofiles
import aiohttp
import zipfile
import tempfile
@@ -28,6 +29,8 @@ from datetime import datetime
import logging
log = logging.getLogger(__name__)
CHUNK_SIZE = 1024 * 8 # 8KB
async def export_project(zstream, project, temporary_dir, include_images=False, keep_compute_id=False, allow_all_nodes=False, reset_mac_addresses=False):
"""
@@ -36,13 +39,13 @@ async def export_project(zstream, project, temporary_dir, include_images=False,
The file will be read chunk by chunk when you iterate over the zip stream.
Some files like snapshots and packet captures are ignored.
:param zstream: ZipStream object
:param project: Project instance
:param temporary_dir: A temporary dir where to store intermediate data
:param include images: save OS images to the zip file
:param keep_compute_id: If false replace all compute id by local (standard behavior for .gns3project to make it portable)
:param allow_all_nodes: Allow all nodes type to be include in the zip even if not portable
:param reset_mac_addresses: Reset MAC addresses for every nodes.
:returns: ZipStream object
"""
# To avoid issue with data not saved we disallow the export of a running project
@@ -80,28 +83,28 @@ async def export_project(zstream, project, temporary_dir, include_images=False,
zstream.write(path, os.path.relpath(path, project._path))
# Export files from remote computes
downloaded_files = set()
for compute in project.computes:
if compute.id != "local":
compute_files = await compute.list_files(project)
for compute_file in compute_files:
if _is_exportable(compute_file["path"]):
(fd, temp_path) = tempfile.mkstemp(dir=temporary_dir)
f = open(fd, "wb", closefd=True)
log.debug("Downloading file '{}' from compute '{}'".format(compute_file["path"], compute.id))
response = await compute.download_file(project, compute_file["path"])
while True:
try:
data = await response.content.read(1024)
except asyncio.TimeoutError:
raise aiohttp.web.HTTPRequestTimeout(text="Timeout when downloading file '{}' from remote compute {}:{}".format(compute_file["path"], compute.host, compute.port))
if not data:
break
f.write(data)
#if response.status != 200:
# raise aiohttp.web.HTTPConflict(text="Cannot export file from compute '{}'. Compute returned status code {}.".format(compute.id, response.status))
(fd, temp_path) = tempfile.mkstemp(dir=temporary_dir)
async with aiofiles.open(fd, 'wb') as f:
while True:
try:
data = await response.content.read(CHUNK_SIZE)
except asyncio.TimeoutError:
raise aiohttp.web.HTTPRequestTimeout(text="Timeout when downloading file '{}' from remote compute {}:{}".format(compute_file["path"], compute.host, compute.port))
if not data:
break
await f.write(data)
response.close()
f.close()
_patch_mtime(temp_path)
zstream.write(temp_path, arcname=compute_file["path"])
downloaded_files.add(compute_file['path'])
def _patch_mtime(path):
@@ -262,30 +265,26 @@ async def _export_remote_images(project, compute_id, image_type, image, project_
Export specific image from remote compute.
"""
log.info("Downloading image '{}' from compute '{}'".format(image, compute_id))
log.debug("Downloading image '{}' from compute '{}'".format(image, compute_id))
try:
compute = [compute for compute in project.computes if compute.id == compute_id][0]
except IndexError:
raise aiohttp.web.HTTPConflict(text="Cannot export image from '{}' compute. Compute doesn't exist.".format(compute_id))
(fd, temp_path) = tempfile.mkstemp(dir=temporary_dir)
f = open(fd, "wb", closefd=True)
response = await compute.download_image(image_type, image)
if response.status != 200:
raise aiohttp.web.HTTPConflict(text="Cannot export image from '{}' compute. Compute returned status code {}.".format(compute_id, response.status))
raise aiohttp.web.HTTPConflict(text="Cannot export image from compute '{}'. Compute returned status code {}.".format(compute_id, response.status))
while True:
try:
data = await response.content.read(1024)
except asyncio.TimeoutError:
raise aiohttp.web.HTTPRequestTimeout(text="Timeout when downloading image '{}' from remote compute {}:{}".format(image, compute.host, compute.port))
if not data:
break
f.write(data)
(fd, temp_path) = tempfile.mkstemp(dir=temporary_dir)
async with aiofiles.open(fd, 'wb') as f:
while True:
try:
data = await response.content.read(CHUNK_SIZE)
except asyncio.TimeoutError:
raise aiohttp.web.HTTPRequestTimeout(text="Timeout when downloading image '{}' from remote compute {}:{}".format(image, compute.host, compute.port))
if not data:
break
await f.write(data)
response.close()
f.close()
arcname = os.path.join("images", image_type, image)
log.info("Saved {}".format(arcname))
project_zipfile.write(temp_path, arcname=arcname, compress_type=zipfile.ZIP_DEFLATED)
-1
View File
@@ -20,7 +20,6 @@ import sys
import json
import uuid
import shutil
import asyncio
import zipfile
import aiohttp
import itertools
+3 -2
View File
@@ -971,7 +971,8 @@ class Project:
try:
begin = time.time()
with tempfile.TemporaryDirectory() as tmpdir:
with aiozipstream.ZipFile(compression=zipfile.ZIP_DEFLATED) as zstream:
# Do not compress the exported project when duplicating
with aiozipstream.ZipFile(compression=zipfile.ZIP_STORED) as zstream:
await export_project(zstream, self, tmpdir, keep_compute_id=True, allow_all_nodes=True, reset_mac_addresses=True)
# export the project to a temporary location
@@ -985,7 +986,7 @@ class Project:
with open(project_path, "rb") as f:
project = await import_project(self._controller, str(uuid.uuid4()), f, location=location, name=name, keep_compute_id=True)
log.info("Project '{}' duplicated in {:.4f} seconds".format(project.id, time.time() - begin))
log.info("Project '{}' duplicated in {:.4f} seconds".format(project.name, time.time() - begin))
except (ValueError, OSError, UnicodeEncodeError) as e:
raise aiohttp.web.HTTPConflict(text="Cannot duplicate project: {}".format(str(e)))
+1 -1
View File
@@ -101,7 +101,7 @@ class Snapshot:
async with aiofiles.open(self.path, 'wb') as f:
async for chunk in zstream:
await f.write(chunk)
log.info("Snapshot '{}' created in {:.4f} seconds".format(self.path, time.time() - begin))
log.info("Snapshot '{}' created in {:.4f} seconds".format(self.name, time.time() - begin))
except (ValueError, OSError, RuntimeError) as e:
raise aiohttp.web.HTTPConflict(text="Could not create snapshot file '{}': {}".format(self.path, e))