jenkins-bot has submitted this change. (
https://gerrit.wikimedia.org/r/c/pywikibot/core/+/1332197?usp=email )
Change subject: upload: Clean up temporary downloads
......................................................................
upload: Clean up temporary downloads
Remove partial files when URL downloads fail and close streamed responses after
each attempt.
Delete owned download files after upload completion or failure.
Preserve caller-provided local files.
Change-Id: Ia52d9fbb0209f23e7b52c96e2376e7339e043baa
---
M pywikibot/specialbots/_upload.py
M tests/uploadbot_tests.py
2 files changed, 263 insertions(+), 88 deletions(-)
Approvals:
jenkins-bot: Verified
Xqt: Looks good to me, approved
diff --git a/pywikibot/specialbots/_upload.py b/pywikibot/specialbots/_upload.py
index 06715f9..2415662 100644
--- a/pywikibot/specialbots/_upload.py
+++ b/pywikibot/specialbots/_upload.py
@@ -138,7 +138,16 @@
else:
self.target_site = pywikibot.Site()
- def read_file_content(self, file_url: str):
+ @staticmethod
+ def _remove_temp_file(path: Path) -> None:
+ """Remove a temporary file without masking an earlier error."""
+ try:
+ path.unlink(missing_ok=True)
+ except OSError as e:
+ pywikibot.warning(
+ f'Unable to remove temporary file {path}: {e}')
+
+ def read_file_content(self, file_url: str) -> str:
"""Return name of temp file in which remote file is saved."""
pywikibot.info('Reading file ' + file_url)
@@ -146,61 +155,78 @@
os.close(temp_fd)
path = Path(tempname)
size = 0
+ complete = False
- dt_gen = (el for el in (15, 30, 45, 60, 120, 180, 240, 300))
- while True:
- file_len = path.stat().st_size
- if file_len:
- pywikibot.info('Download resumed.')
- headers = {'Range': f'bytes={file_len}-'}
- else:
- headers = {}
+ try:
+ dt_gen = (el for el in (15, 30, 45, 60, 120, 180, 240, 300))
+ while True:
+ file_len = path.stat().st_size
+ if file_len:
+ pywikibot.info('Download resumed.')
+ headers = {'Range': f'bytes={file_len}-'}
+ else:
+ headers = {}
- with path.open('ab') as fd:
+ with path.open('ab') as fd:
+ try:
+ with http.fetch(file_url, stream=True,
+ headers=headers) as response:
+ try:
+ response.raise_for_status()
+ except requests.HTTPError as e:
+ # exit criteria if size is not available
+ # error on last iteration is OK, we're
+ # requesting {'Range': 'bytes=file_len-'}
+ err = (
+ HTTPStatus.REQUESTED_RANGE_NOT_SATISFIABLE
+ )
+ if response.status_code == err \
+ and path.stat().st_size:
+ break
+ raise FatalServerError(str(e)) from e
+
+ # get download info, if available
+ # Note: this is not enough to exclude pages
+ # e.g. 'application/json' is also not a media
+ content_type = response.headers.get(
+ 'Content-Type', '')
+ if 'text/' in content_type:
+ raise FatalServerError(
+ 'The requested URL was not found on '
+ 'server.')
+ size = max(
+ size,
+ int(response.headers.get('Content-Length', 0))
+ )
+
+ # stream content to temp file (in chunks of 1Mb)
+ for chunk in response.iter_content(
+ chunk_size=1024 * 1024):
+ fd.write(chunk)
+
+ # raised from connection lost during iter_content()
+ except requests.ConnectionError:
+ fd.flush()
+ pywikibot.info(
+ 'Connection closed at byte '
+ f'{path.stat().st_size}')
+
+ if size and size == path.stat().st_size:
+ break
try:
- response = http.fetch(file_url, stream=True,
- headers=headers)
- response.raise_for_status()
+ dt = next(dt_gen)
+ pywikibot.info(f'Sleeping for {dt} seconds ...')
+ pywikibot.sleep(dt)
+ except StopIteration:
+ raise FatalServerError(
+ 'Download failed, too many retries!')
- # get download info, if available
- # Note: this is not enough to exclude pages
- # e.g. 'application/json' is also not a media
- if 'text/' in response.headers['Content-Type']:
- raise FatalServerError('The requested URL was not '
- 'found on server.')
- size = max(size,
- int(response.headers.get('Content-Length', 0)))
-
- # stream content to temp file (in chunks of 1Mb)
- for chunk in response.iter_content(chunk_size=1024 * 1024):
- fd.write(chunk)
-
- # raised from connection lost during response.iter_content()
- except requests.ConnectionError:
- fd.flush()
- pywikibot.info(
- f'Connection closed at byte {path.stat().st_size}')
- # raised from response.raise_for_status()
- except requests.HTTPError as e:
- # exit criteria if size is not available
- # error on last iteration is OK, we're requesting
- # {'Range': 'bytes=file_len-'}
- err = HTTPStatus.REQUESTED_RANGE_NOT_SATISFIABLE
- if response.status_code == err and path.stat().st_size:
- break
- raise FatalServerError(str(e)) from e
-
- if size and size == path.stat().st_size:
- break
- try:
- dt = next(dt_gen)
- pywikibot.info(f'Sleeping for {dt} seconds ...')
- pywikibot.sleep(dt)
- except StopIteration:
- raise FatalServerError('Download failed, too many retries!')
-
- pywikibot.info(f'Downloaded {path.stat().st_size} bytes')
- return tempname
+ pywikibot.info(f'Downloaded {path.stat().st_size} bytes')
+ complete = True
+ return tempname
+ finally:
+ if not complete:
+ self._remove_temp_file(path)
def _handle_warning(self, warning: str) -> bool | None:
"""Return whether the warning cause an abort or be ignored.
@@ -414,45 +440,53 @@
ignore_warnings = self.ignore_warning is True or self._handle_warnings
download = False
- while True:
- if '://' in file_url \
- and (not site.has_right('upload_by_url') or download):
+ temp_filename: str | None = None
+ try:
+ while True:
+ if '://' in file_url \
+ and (not site.has_right('upload_by_url') or download):
+ try:
+ temp_filename = self.read_file_content(file_url)
+ file_url = temp_filename
+ except FatalServerError as e:
+ pywikibot.error(e)
+ return None
+
try:
- file_url = self.read_file_content(file_url)
- except FatalServerError as e:
- pywikibot.error(e)
- return None
-
- try:
- success = imagepage.upload(file_url,
- ignore_warnings=ignore_warnings,
- chunk_size=self.chunk_size,
- asynchronous=self.asynchronous,
- comment=self.summary)
- except APIError as error:
- if error.code == 'uploaddisabled':
- pywikibot.error(f'Upload error: Local file uploads are '
- f'disabled on {site}.')
- elif error.code == 'copyuploadbaddomain' and not download \
- and '://' in file_url:
- pywikibot.error(error)
- pywikibot.info('Downloading the file and retry...')
- download = True
- continue
- else:
+ success = imagepage.upload(
+ file_url,
+ ignore_warnings=ignore_warnings,
+ chunk_size=self.chunk_size,
+ asynchronous=self.asynchronous,
+ comment=self.summary)
+ except APIError as error:
+ if error.code == 'uploaddisabled':
+ pywikibot.error(
+ f'Upload error: Local file uploads are disabled '
+ f'on {site}.')
+ elif error.code == 'copyuploadbaddomain' and not download \
+ and '://' in file_url:
+ pywikibot.error(error)
+ pywikibot.info('Downloading the file and retry...')
+ download = True
+ continue
+ else:
+ pywikibot.exception('Upload error: ')
+ except Exception:
pywikibot.exception('Upload error: ')
- except Exception:
- pywikibot.exception('Upload error: ')
- else:
- if success:
- # No warning, upload complete.
- pywikibot.info(f'Upload of {filename} successful.')
- self.counter['upload'] += 1
- return filename # data['filename']
- pywikibot.info('Upload aborted.')
- break
+ else:
+ if success:
+ # No warning, upload complete.
+ pywikibot.info(f'Upload of {filename} successful.')
+ self.counter['upload'] += 1
+ return filename # data['filename']
+ pywikibot.info('Upload aborted.')
+ break
- return None
+ return None
+ finally:
+ if temp_filename is not None:
+ self._remove_temp_file(Path(temp_filename))
def skip_run(self) -> bool:
"""Check whether processing is to be skipped."""
diff --git a/tests/uploadbot_tests.py b/tests/uploadbot_tests.py
index 1c84614..1945210 100755
--- a/tests/uploadbot_tests.py
+++ b/tests/uploadbot_tests.py
@@ -11,10 +11,15 @@
from __future__ import annotations
import os
+import tempfile
import unittest
from contextlib import suppress
+from pathlib import Path
from unittest import mock
+import requests
+
+from pywikibot.exceptions import APIError
from pywikibot.specialbots import UploadRobot
from tests import join_images_path
from tests.aspects import DefaultSiteTestCase, TestCase
@@ -85,6 +90,142 @@
self.assertIsNone(bot.post_processor)
+class TestUploadbotTempFiles(TestCase):
+
+ """Dry tests for UploadRobot temporary files."""
+
+ net = False
+
+ def test_download_failure_removes_partial_file(self) -> None:
+ """Test that a failed download removes its partial file."""
+ response = mock.MagicMock()
+ response.__enter__.return_value = response
+ response.headers = {
+ 'Content-Type': 'image/png',
+ 'Content-Length': '10',
+ }
+
+ def chunks():
+ yield b'partial'
+ raise requests.Timeout('Download timed out')
+
+ response.iter_content.return_value = chunks()
+ bot = UploadRobot(
+ url=['mahveo.png'], target_site=mock.Mock(), always=False)
+
+ with tempfile.TemporaryDirectory() as directory:
+ temp_fd, tempname = tempfile.mkstemp(dir=directory)
+ with (
+ mock.patch('pywikibot.specialbots._upload.tempfile.mkstemp',
+ return_value=(temp_fd, tempname)),
+ mock.patch('pywikibot.specialbots._upload.http.fetch',
+ return_value=response),
+ self.assertRaises(requests.Timeout),
+ ):
+ bot.read_file_content('https://yo.wikipedia.org/mahveo.png')
+
+ self.assertFalse(Path(tempname).exists())
+
+ response.__exit__.assert_called_once()
+
+ def test_downloaded_file_cleanup(self) -> None:
+ """Test that downloaded files are removed after an upload attempt."""
+ file_url = 'https://yo.wikipedia.org/mahveo.png'
+ cases = (
+ (True, 'mahveo.png'),
+ (False, None),
+ (APIError('mahveo', 'Upload failed'), None),
+ )
+
+ for outcome, expected in cases:
+ with self.subTest(outcome=outcome), \
+ tempfile.TemporaryDirectory() as directory:
+ temp_path = Path(directory) / 'download'
+ temp_path.write_bytes(b'content')
+ site = mock.Mock()
+ site.has_right.return_value = False
+ imagepage = mock.Mock()
+
+ def upload(source, *, expected_path=temp_path,
+ upload_outcome=outcome, **kwargs):
+ self.assertEqual(source, str(expected_path))
+ self.assertTrue(expected_path.exists())
+ if isinstance(upload_outcome, Exception):
+ raise upload_outcome
+ return upload_outcome
+
+ imagepage.upload.side_effect = upload
+ bot = UploadRobot(
+ url=[file_url], target_site=site, always=False)
+ with (
+ mock.patch.object(bot, 'process_filename',
+ return_value='mahveo.png'),
+ mock.patch.object(bot, 'read_file_content',
+ return_value=str(temp_path)),
+ mock.patch('pywikibot.FilePage', return_value=imagepage),
+ mock.patch('pywikibot.exception'),
+ ):
+ result = bot.upload_file(file_url)
+
+ self.assertEqual(result, expected)
+ self.assertFalse(temp_path.exists())
+
+ def test_local_file_is_preserved(self) -> None:
+ """Test that a caller-owned local file is not removed."""
+ site = mock.Mock()
+ imagepage = mock.Mock()
+ imagepage.upload.return_value = True
+
+ with tempfile.TemporaryDirectory() as directory:
+ path = Path(directory) / 'mahveo.png'
+ path.write_bytes(b'content')
+ bot = UploadRobot(
+ url=[str(path)], target_site=site, always=False)
+ with (
+ mock.patch.object(bot, 'process_filename',
+ return_value='mahveo.png'),
+ mock.patch('pywikibot.FilePage', return_value=imagepage),
+ ):
+ result = bot.upload_file(str(path))
+
+ self.assertEqual(result, 'mahveo.png')
+ self.assertTrue(path.exists())
+
+ def test_downloaded_file_cleanup_after_retry(self) -> None:
+ """Test cleanup after falling back from URL to file upload."""
+ file_url = 'https://yo.wikipedia.org/mahveo.png'
+ site = mock.Mock()
+ site.has_right.return_value = True
+ imagepage = mock.Mock()
+
+ with tempfile.TemporaryDirectory() as directory:
+ temp_path = Path(directory) / 'download'
+ temp_path.write_bytes(b'content')
+
+ def upload(source, **kwargs):
+ if source == file_url:
+ raise APIError('copyuploadbaddomain', 'Bad domain')
+ self.assertEqual(source, str(temp_path))
+ self.assertTrue(temp_path.exists())
+ return True
+
+ imagepage.upload.side_effect = upload
+ bot = UploadRobot(
+ url=[file_url], target_site=site, always=False)
+ with (
+ mock.patch.object(bot, 'process_filename',
+ return_value='mahveo.png'),
+ mock.patch.object(bot, 'read_file_content',
+ return_value=str(temp_path)) as download,
+ mock.patch('pywikibot.FilePage', return_value=imagepage),
+ ):
+ result = bot.upload_file(file_url)
+
+ self.assertEqual(result, 'mahveo.png')
+ download.assert_called_once_with(file_url)
+ self.assertFalse(temp_path.exists())
+
+
class TestUploadbotCounter(TestCase):
"""Dry tests for UploadRobot counters."""
--
To view, visit
https://gerrit.wikimedia.org/r/c/pywikibot/core/+/1332197?usp=email
To unsubscribe, or for help writing mail filters, visit
https://gerrit.wikimedia.org/r/settings?usp=email
Gerrit-MessageType: merged
Gerrit-Project: pywikibot/core
Gerrit-Branch: master
Gerrit-Change-Id: Ia52d9fbb0209f23e7b52c96e2376e7339e043baa
Gerrit-Change-Number: 1332197
Gerrit-PatchSet: 5
Gerrit-Owner: Mahveotm <[email protected]>
Gerrit-Reviewer: Xqt <[email protected]>
Gerrit-Reviewer: jenkins-bot
_______________________________________________
Pywikibot-commits mailing list -- [email protected]
To unsubscribe send an email to [email protected]
