jenkins-bot has submitted this change. (
https://gerrit.wikimedia.org/r/c/pywikibot/core/+/1326394?usp=email )
Change subject: download_dump: Harden streamed downloads
......................................................................
download_dump: Harden streamed downloads
Streamed responses were never closed, including when the server
returned an error status. Missing Content-Length headers also raised
KeyError before the intended warning and progress-free fallback could
run.
Manage each response with its context manager and treat a missing
length as unknown so the complete stream is still written.
Change-Id: Ie229f207cb0a9e59268dca4764b98e481e3b988c
---
M scripts/download_dump.py
M tests/__init__.py
A tests/download_dump_tests.py
3 files changed, 99 insertions(+), 42 deletions(-)
Approvals:
Xqt: Looks good to me, approved
jenkins-bot: Verified
diff --git a/scripts/download_dump.py b/scripts/download_dump.py
index 800a9c0..8721a24 100755
--- a/scripts/download_dump.py
+++ b/scripts/download_dump.py
@@ -111,53 +111,56 @@
url = (f'https://dumps.wikimedia.org/{self.opt.wikiname}/'
f'{self.opt.dumpdate}/{download_filename}')
pywikibot.info('Downloading file from ' + url)
- response = fetch(url, stream=True)
+ with fetch(url, stream=True) as response:
+ if response.status_code != HTTPStatus.OK:
+ if response.status_code == HTTPStatus.NOT_FOUND:
+ pywikibot.info(
+ 'File with name {filename!r}, from '
+ 'dumpdate {dumpdate!r}, and wiki '
+ '{wikiname!r} ({url}) '
+ "isn't available in the Wikimedia Dumps"
+ .format(url=url, **self.opt))
+ else:
+ pywikibot.info(HTTPStatus(
+ response.status_code).description)
+ return
- if response.status_code != HTTPStatus.OK:
- if response.status_code == HTTPStatus.NOT_FOUND:
- pywikibot.info(
- 'File with name {filename!r}, from dumpdate '
- '{dumpdate!r}, and wiki {wikiname!r} ({url}) '
- "isn't available in the Wikimedia Dumps"
- .format(url=url, **self.opt))
- else:
- pywikibot.info(
- HTTPStatus(response.status_code).description)
- return
+ with open(file_current_storepath,
+ 'wb') as result_file:
+ total = int(response.headers.get(
+ 'content-length', -1))
+ if total == -1:
+ pywikibot.warning(
+ "'content-length' missing in response "
+ 'headers')
+ downloaded = 0
+ parts = 50
+ display_string = ''
- with open(file_current_storepath, 'wb') as result_file:
- total = int(response.headers['content-length'])
- if total == -1:
- pywikibot.warning("'content-length' missing in "
- 'response headers')
- downloaded = 0
- parts = 50
- display_string = ''
+ pywikibot.info()
+ for data in response.iter_content(100 * 1024):
+ result_file.write(data)
- pywikibot.info()
- for data in response.iter_content(100 * 1024):
- result_file.write(data)
+ if total <= 0:
+ continue
- if total <= 0:
- continue
+ downloaded += len(data)
+ done = int(parts * downloaded / total)
+ display = map(convert_from_bytes,
+ (downloaded, total))
+ prior_display = display_string
+ display_string = '\r|{}{}|{}{}/{}'.format(
+ '=' * done,
+ '-' * (parts - done),
+ ' ' * 5,
+ *display)
+ # Add whitespace to cover up prior bar
+ display_string += ' ' * (
+ len(prior_display.rstrip())
+ - len(display_string.rstrip()))
- downloaded += len(data)
- done = int(parts * downloaded / total)
- display = map(convert_from_bytes,
- (downloaded, total))
- prior_display = display_string
- display_string = '\r|{}{}|{}{}/{}'.format(
- '=' * done,
- '-' * (parts - done),
- ' ' * 5,
- *display)
- # Add whitespace to cover up prior bar
- display_string += ' ' * (
- len(prior_display.rstrip())
- - len(display_string.rstrip()))
-
- pywikibot.info(display_string, newline=False)
- pywikibot.info()
+ pywikibot.info(display_string, newline=False)
+ pywikibot.info()
# Rename the temporary file to the target file
# if the download completes successfully
diff --git a/tests/__init__.py b/tests/__init__.py
index b3fe091..7c0145f 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -183,6 +183,7 @@
'commonscat',
'data_ingestion',
'deletionbot',
+ 'download_dump',
'fixing_redirects',
'generate_family_file',
'generate_user_files',
diff --git a/tests/download_dump_tests.py b/tests/download_dump_tests.py
new file mode 100755
index 0000000..7d090be
--- /dev/null
+++ b/tests/download_dump_tests.py
@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+#
+# (C) Pywikibot team, 2026
+#
+# Distributed under the terms of the MIT license.
+#
+"""Tests for the download_dump script."""
+from __future__ import annotations
+
+import unittest
+from http import HTTPStatus
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from unittest.mock import MagicMock, patch
+
+from scripts import download_dump
+from tests.aspects import TestCase
+
+
+class DownloadDumpBotTestCase(TestCase):
+
+ """Test :class:`download_dump.DownloadDumpBot`."""
+
+ net = False
+
+ @patch.object(download_dump.DownloadDumpBot, 'get_dump_name',
+ return_value=None)
+ @patch.object(download_dump, 'fetch')
+ def test_download_without_content_length(
+ self, fetch_mock, get_dump_name_mock,
+ ) -> None:
+ """Test downloading a response without a content length."""
+ response = MagicMock()
+ response.__enter__.return_value = response
+ response.status_code = HTTPStatus.OK
+ response.headers = {}
+ response.iter_content.return_value = [b'first', b'second']
+ fetch_mock.return_value = response
+
+ with TemporaryDirectory() as directory:
+ bot = download_dump.DownloadDumpBot(
+ wikiname='enwiki', filename='pages.xml.bz2',
+ storepath=directory, dumpdate='latest')
+ bot.run()
+
+ path = Path(directory) / 'enwiki-latest-pages.xml.bz2'
+ self.assertEqual(path.read_bytes(), b'firstsecond')
+
+ response.__exit__.assert_called_once()
+
+
+if __name__ == '__main__':
+ unittest.main()
--
To view, visit
https://gerrit.wikimedia.org/r/c/pywikibot/core/+/1326394?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: Ie229f207cb0a9e59268dca4764b98e481e3b988c
Gerrit-Change-Number: 1326394
Gerrit-PatchSet: 2
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]
