diff --git a/scripts/deploy_docs_to_oss.py b/scripts/deploy_docs_to_oss.py index b518b7e..3a6eae8 100644 --- a/scripts/deploy_docs_to_oss.py +++ b/scripts/deploy_docs_to_oss.py @@ -186,20 +186,32 @@ def signed_get_sample(self, key: str) -> int: return len(self.request("GET", key)[:256]) def list_prefix(self, prefix: str) -> list[str]: - query = urllib.parse.urlencode({"prefix": prefix, "max-keys": "1000"}) - date, authorization = self.sign("GET", "") - request = urllib.request.Request( - f"https://{self.host}/?{query}", - headers={"Date": date, "Authorization": authorization}, - method="GET", - ) - with urllib.request.urlopen(request, timeout=30) as response: - root = ET.fromstring(response.read()) - return [ - node.text - for node in root.iter() - if node.tag.endswith("Key") and node.text - ] + keys: list[str] = [] + params = {"prefix": prefix, "max-keys": "1000"} + while True: + query = urllib.parse.urlencode(params) + date, authorization = self.sign("GET", "") + request = urllib.request.Request( + f"https://{self.host}/?{query}", + headers={"Date": date, "Authorization": authorization}, + method="GET", + ) + with urllib.request.urlopen(request, timeout=30) as response: + root = ET.fromstring(response.read()) + keys.extend( + node.text + for node in root.findall("{*}Contents/{*}Key") + if node.text + ) + truncated = root.findtext("{*}IsTruncated") + if truncated == "false": + return keys + if truncated != "true": + raise ValueError("OSS listing has an invalid IsTruncated value") + marker = root.findtext("{*}NextMarker") + if not marker or marker <= params.get("marker", ""): + raise ValueError("OSS listing is missing an advancing NextMarker") + params["marker"] = marker def public_head_status(self, key: str) -> int: request = urllib.request.Request( diff --git a/scripts/test_deploy_docs_to_oss.py b/scripts/test_deploy_docs_to_oss.py index f0e7746..b3552a9 100644 --- a/scripts/test_deploy_docs_to_oss.py +++ b/scripts/test_deploy_docs_to_oss.py @@ -1,8 +1,12 @@ +import io import tempfile import unittest from pathlib import Path +from unittest.mock import patch +from urllib.parse import parse_qs, urlsplit from scripts.deploy_docs_to_oss import ( + OssClient, cache_control_for, collect_files, collect_uploads, @@ -14,6 +18,75 @@ class DeployDocsToOssTests(unittest.TestCase): + def test_list_prefix_reads_beyond_1000_objects(self) -> None: + keys = [f"site/file-{index:04d}" for index in range(1001)] + for namespace in ( + "", + ' xmlns="http://doc.oss-cn-hangzhou.aliyuncs.com"', + ): + with self.subTest(namespace=namespace): + first_page = ( + f"true" + f"{keys[999]}" + + "".join( + f"{key}" + for key in keys[:1000] + ) + + "" + ) + last_page = ( + f"false" + f"{keys[1000]}" + ) + with patch( + "scripts.deploy_docs_to_oss.urllib.request.urlopen", + side_effect=[ + io.BytesIO(first_page.encode()), + io.BytesIO(last_page.encode()), + ], + ) as urlopen: + client = OssClient("test", "test", "bucket", "example.com") + actual = client.list_prefix("site/") + + self.assertEqual(actual, keys) + queries = [ + parse_qs(urlsplit(call.args[0].full_url).query) + for call in urlopen.call_args_list + ] + self.assertEqual( + queries, + [ + {"prefix": ["site/"], "max-keys": ["1000"]}, + { + "prefix": ["site/"], + "max-keys": ["1000"], + "marker": [keys[999]], + }, + ], + ) + + def test_list_prefix_rejects_invalid_pagination(self) -> None: + first_page = ( + b"true" + b"site/b" + ) + for pagination in ( + "true", + "truesite/b", + "truesite/a", + "", + ): + with self.subTest(pagination=pagination): + last_page = ( + f"{pagination}".encode() + ) + with patch( + "scripts.deploy_docs_to_oss.urllib.request.urlopen", + side_effect=[io.BytesIO(first_page), io.BytesIO(last_page)], + ), self.assertRaisesRegex(ValueError, "OSS listing"): + client = OssClient("test", "test", "bucket", "example.com") + client.list_prefix("site/") + def test_parse_bucket(self) -> None: self.assertEqual(parse_bucket("oss://docs-bucket/"), "docs-bucket") self.assertEqual(parse_bucket("oss://docs-bucket"), "docs-bucket")