diff sat/plugins/plugin_misc_download.py @ 3211:4252176ad993

plugin download: clean unfinished files and re-raise exception in case of download error
author Goffi <goffi@goffi.org>
date Fri, 06 Mar 2020 18:19:03 +0100
parents 2c0628f3927e
children be6d91572633
line wrap: on
line diff
--- a/sat/plugins/plugin_misc_download.py	Fri Mar 06 18:19:03 2020 +0100
+++ b/sat/plugins/plugin_misc_download.py	Fri Mar 06 18:19:03 2020 +0100
@@ -144,6 +144,7 @@
         uri_parsed = urlparse(uri, 'http')
         if dest_path:
             dest_path = Path(dest_path)
+            cache_uid = None
         else:
             filename = Path(unquote(uri_parsed.path)).name.strip() or C.FILE_DEFAULT_NAME
             # we don't use Path.suffixes because we don't want to have more than 2
@@ -151,19 +152,20 @@
             stem, *suffixes = filename.rsplit('.', 2)
             # we hash the URL to have an unique identifier, and avoid double download
             url_hash = hashlib.sha256(uri_parsed.geturl().encode()).hexdigest()
-            uid = f"{stem}_{url_hash}"
-            cache_data = client.cache.getMetadata(uid)
+            cache_uid = f"{stem}_{url_hash}"
+            cache_data = client.cache.getMetadata(cache_uid)
             if cache_data is not None:
                 # file is already in cache, we return it
                 download_d = defer.succeed(cache_data['path'])
                 return '', download_d
             else:
                 # the file is not in cache
-                unique_name = '.'.join([uid] + suffixes)
-                with client.cache.cacheData("DOWNLOAD", uid, filename=unique_name) as f:
+                unique_name = '.'.join([cache_uid] + suffixes)
+                with client.cache.cacheData(
+                    "DOWNLOAD", cache_uid, filename=unique_name) as f:
                     # we close the file and only use its name, the file will be opened
                     # by the registered callback
-                    dest_path = f.name
+                    dest_path = Path(f.name)
 
         # should we check certificates?
         check_certificate = self.host.memory.getParamA(
@@ -178,8 +180,18 @@
         except KeyError:
             raise exceptions.NotFound(f"Can't find any handler for uri {uri}")
         else:
-            progress_id, download_d = await callback(
-                client, uri_parsed, dest_path, options)
+            try:
+                progress_id, download_d = await callback(
+                    client, uri_parsed, dest_path, options)
+            except Exception as e:
+                log.warning(_(
+                    "Can't download URI {uri}: {reason}").format(
+                    uri=uri, reason=e))
+                if cache_uid is not None:
+                    client.cache.removeFromCache(cache_uid)
+                elif dest_path.exists():
+                    dest_path.unlink()
+                raise e
             download_d.addCallback(lambda __: dest_path)
             return progress_id, download_d
 
@@ -230,7 +242,7 @@
         else:
             treq_client = treq
 
-        head_data = await treq_.head(url)
+        head_data = await treq_client.head(url)
         try:
             content_length = int(head_data.headers.getRawHeaders('content-length')[0])
         except (KeyError, TypeError, IndexError):