Simplify v2 downloader implementation

Author: dkirov-dd

datadog_checks_downloader/datadog_checks/downloader/cli.py

@@ -8,11 +8,14 @@
 import os
 import re
 import sys
+import urllib.error
 
 # 2nd party.
+from tuf.api.exceptions import DownloadError
+
 from .download import DEFAULT_ROOT_LAYOUT_TYPE, REPOSITORY_URL_PREFIX, ROOT_LAYOUTS, TUFDownloader
 from .download_v2 import V2_REPOSITORY_URL, TUFPointerDownloader
-from .exceptions import NonCanonicalVersion, NonDatadogPackage
+from .exceptions import NonCanonicalVersion, NonDatadogPackage, TargetNotFoundError
 
 # Private module functions.
 
@@ -26,6 +29,14 @@ def __is_canonical(version):
     return re.match(P, version) is not None
 
 
+def _v2_failure_category(exc: Exception) -> str:
+    if isinstance(exc, TargetNotFoundError):
+        return 'target version not found'
+    if isinstance(exc, (DownloadError, TimeoutError, urllib.error.URLError)):
+        return 'network error'
+    return 'other'
+
+
 def __find_shipped_integrations():
     # Recurse up from site-packages until we find the Agent root directory.
     # The relative path differs between operating systems.
@@ -158,7 +169,12 @@ def download():
         except Exception as exc:
             import logging
 
-            logging.getLogger(__name__).info('v2 download failed (%s: %s), falling back to v1', type(exc).__name__, exc)
+            logging.getLogger(__name__).info(
+                'v2 download failed (%s, %s: %s), falling back to v1',
+                _v2_failure_category(exc),
+                type(exc).__name__,
+                exc,
+            )
             tuf_downloader, standard_distribution_name, version, ignore_python_version = instantiate_downloader()
             run_downloader(tuf_downloader, standard_distribution_name, version, ignore_python_version)
 
@@ -187,18 +203,14 @@ def _download_v2():
     # v1 flags that are not applicable in v2: accept and warn so that callers
     # upgrading from v1 get a clear message instead of an argument error.
     parser.add_argument('--type', type=str, default=None, dest='_type_ignored')
-    parser.add_argument('--force', action='store_true', dest='_force_ignored')
     parser.add_argument('--ignore-python-version', action='store_true', dest='_ignore_python_version')
 
     args = parser.parse_args()
 
     if not args.standard_distribution_name.startswith('datadog-'):
         raise NonDatadogPackage(args.standard_distribution_name)
 
-    if args.version and not re.match(
-        r'^([1-9]\d*!)?(0|[1-9]\d*)(\.(0|[1-9]\d*))*((a|b|rc)(0|[1-9]\d*))?(\.post(0|[1-9]\d*))?(\.dev(0|[1-9]\d*))?$',
-        args.version,
-    ):
+    if args.version and not __is_canonical(args.version):
         raise NonCanonicalVersion(args.version)
 
     if args._type_ignored is not None:

…/downloader/data/v2/metadata/1.root.json …ks/downloader/data/v2/metadata/root.jsondatadog_checks_downloader/datadog_checks/downloader/data/v2/metadata/1.root.json renamed to datadog_checks_downloader/datadog_checks/downloader/data/v2/metadata/root.json datadog_checks_downloader/datadog_checks/downloader/data/v2/metadata/1.root.json renamed to datadog_checks_downloader/datadog_checks/downloader/data/v2/metadata/root.json

@@ -2,32 +2,10 @@
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
 
-"""TUF pointer-file downloader (v2 repository format).
+"""TUF pointer-file downloader for the v2 repository format.
 
-The v2 format stores a JSON pointer file as a TUF target at:
-
-    targets/<project>/<version>.json   (versioned pointer)
-    targets/<project>/latest.json      (copy of latest stable version pointer)
-
-Target files are content-addressed on S3 (consistent-snapshot format):
-the actual file at ``targets/<project>/<name>.json`` is stored as
-``targets/<project>/{sha256}.{name}.json``.  The TUF Updater resolves
-the hash-prefixed path automatically via TUF metadata.
-
-Each pointer file contains:
-
-    {
-      "digest":           "<sha256 hex>",
-      "length":           <int>,
-      "version":          "<semver>",
-      "repository":       "https://<bucket>.s3.amazonaws.com",
-      "wheel_path":       "/wheels/<project>/<wheel>.whl",
-      "attestation_path": "/attestations/<project>/<version>.sigstore.json"
-    }
-
-The downloader TUF-verifies the pointer file, then fetches the wheel from
-``repository + wheel_path`` and verifies the sha256 digest and byte length
-before writing to disk.
+The downloader TUF-verifies a JSON pointer target, downloads the referenced
+wheel, and verifies the wheel digest and byte length before writing it to disk.
 """
 
 import hashlib
@@ -82,7 +60,7 @@ def _bootstrap_metadata_dir(self, metadata_dir: Path) -> None:
         """Seed *metadata_dir* with the bundled initial root.json trust anchor."""
         dest = metadata_dir / 'root.json'
         metadata = importlib.resources.files('datadog_checks.downloader') / 'data' / 'v2' / 'metadata'
-        dest.write_bytes((metadata / '1.root.json').read_bytes())
+        dest.write_bytes((metadata / 'root.json').read_bytes())
 
     def _make_updater(self, metadata_dir: Path, target_dir: Path) -> Updater:
         return Updater(

datadog_checks_downloader/datadog_checks/downloader/download_v2.py

@@ -56,7 +56,7 @@ include = [
     "/datadog_checks/downloader",
 ]
 artifacts = [
-    "/datadog_checks/downloader/data/v2/metadata/1.root.json",
+    "/datadog_checks/downloader/data/v2/metadata/root.json",
 ]
 dev-mode-dirs = [
     ".",

datadog_checks_downloader/pyproject.toml

@@ -1,7 +1,23 @@
 # (C) Datadog, Inc. 2023-present
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
+import urllib.error
+
+from datadog_checks.downloader.cli import _v2_failure_category
 from datadog_checks.downloader.download import TUFDownloader
+from datadog_checks.downloader.exceptions import TargetNotFoundError
+
+
+def test_v2_target_not_found_failure_category():
+    assert _v2_failure_category(TargetNotFoundError('missing')) == 'target version not found'
+
+
+def test_v2_network_failure_category():
+    assert _v2_failure_category(urllib.error.URLError('timeout')) == 'network error'
+
+
+def test_v2_other_failure_category():
+    assert _v2_failure_category(ValueError('bad pointer')) == 'other'
 
 
 def test_non_official_wheel_filter(mocker):

datadog_checks_downloader/tests/test_unit.py

@@ -2,11 +2,7 @@
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
 
-"""Unit tests for TUFPointerDownloader (v2 repository format).
-
-All tests are offline: the TUF Updater and HTTP calls are mocked so that no
-network traffic is needed.
-"""
+"""Unit tests for TUFPointerDownloader (v2 repository format)."""
 
 import hashlib
 import json
@@ -18,30 +14,25 @@
 from datadog_checks.downloader.download_v2 import TUFPointerDownloader
 from datadog_checks.downloader.exceptions import DigestMismatch, TargetNotFoundError
 
-# ---------------------------------------------------------------------------
-# Shared fixtures
-# ---------------------------------------------------------------------------
-
 _PROJECT = 'datadog-postgres'
 _VERSION = '14.0.0'
 _WHEEL_CONTENT = b'fake wheel bytes for testing'
 _WHEEL_DIGEST = hashlib.sha256(_WHEEL_CONTENT).hexdigest()
 _WHEEL_LENGTH = len(_WHEEL_CONTENT)
+_REPO_URL = 'https://agent-integration-wheels-staging.s3.amazonaws.com'
 
 _POINTER = {
     'digest': _WHEEL_DIGEST,
     'length': _WHEEL_LENGTH,
     'version': _VERSION,
-    'repository': 'https://agent-integration-wheels-staging.s3.amazonaws.com',
+    'repository': _REPO_URL,
     'wheel_path': f'/wheels/{_PROJECT}/datadog_postgres-{_VERSION}-py3-none-any.whl',
     'attestation_path': f'/attestations/{_PROJECT}/{_VERSION}.sigstore.json',
 }
 
-_REPO_URL = 'https://agent-integration-wheels-staging.s3.amazonaws.com'
-
 
 def _mock_tuf_updater(pointer: dict) -> MagicMock:
-    """Return a mock Updater that yields *pointer* as the target content."""
+    """Return an Updater mock that writes *pointer* as target content."""
     pointer_bytes = json.dumps(pointer).encode()
     mock_updater = MagicMock()
     mock_updater.get_targetinfo.return_value = MagicMock()
@@ -54,9 +45,26 @@ def fake_download_target(_target_info, dest_path):
     return mock_updater
 
 
-# ---------------------------------------------------------------------------
-# target-path unit tests
-# ---------------------------------------------------------------------------
+def _mock_response(content: bytes) -> MagicMock:
+    response = MagicMock()
+    response.__enter__ = lambda s: s
+    response.__exit__ = MagicMock(return_value=False)
+    response.read.return_value = content
+    return response
+
+
+@pytest.fixture
+def mock_urlopen():
+    with patch('datadog_checks.downloader.download_v2.urllib.request.urlopen') as mock:
+        mock.return_value = _mock_response(_WHEEL_CONTENT)
+        yield mock
+
+
+@pytest.fixture
+def mock_updater_cls():
+    with patch('datadog_checks.downloader.download_v2.Updater') as mock:
+        mock.return_value = _mock_tuf_updater(_POINTER)
+        yield mock
 
 
 class TestTargetPath:
@@ -67,85 +75,38 @@ def test_missing_version_uses_latest_pointer(self):
         assert TUFPointerDownloader._target_path(_PROJECT, None) == f'{_PROJECT}/latest.json'
 
 
-# ---------------------------------------------------------------------------
-# Happy-path integration tests
-# ---------------------------------------------------------------------------
-
-
 @pytest.mark.offline
 class TestHappyPath:
-    @patch('datadog_checks.downloader.download_v2.Updater')
-    @patch('datadog_checks.downloader.download_v2.urllib.request.urlopen')
     def test_download_returns_wheel_path(self, mock_urlopen, mock_updater_cls, tmp_path):
-        mock_updater_cls.return_value = _mock_tuf_updater(_POINTER)
-
-        mock_urlopen.return_value.__enter__ = lambda s: s
-        mock_urlopen.return_value.__exit__ = MagicMock(return_value=False)
-        mock_urlopen.return_value.read.return_value = _WHEEL_CONTENT
-
         downloader = TUFPointerDownloader(repository_url=_REPO_URL)
         wheel_path = downloader.download(_PROJECT, version=_VERSION, dest_dir=tmp_path)
 
         assert wheel_path.exists()
         assert wheel_path.read_bytes() == _WHEEL_CONTENT
         assert wheel_path.name == f'datadog_postgres-{_VERSION}-py3-none-any.whl'
 
-    @patch('datadog_checks.downloader.download_v2.Updater')
-    @patch('datadog_checks.downloader.download_v2.urllib.request.urlopen')
     def test_repository_flag_overrides_pointer_repository(self, mock_urlopen, mock_updater_cls, tmp_path):
-        """--repository supersedes the repository field baked into the pointer."""
-        staging_url = 'https://agent-integration-wheels-staging.s3.amazonaws.com'
         prod_pointer = {**_POINTER, 'repository': 'https://agent-integration-wheels-prod.s3.amazonaws.com'}
         mock_updater_cls.return_value = _mock_tuf_updater(prod_pointer)
 
-        captured_urls: list[str] = []
-
-        def fake_urlopen(url):
-            captured_urls.append(url)
-            mock_resp = MagicMock()
-            mock_resp.__enter__ = lambda s: s
-            mock_resp.__exit__ = MagicMock(return_value=False)
-            mock_resp.read.return_value = _WHEEL_CONTENT
-            return mock_resp
-
-        mock_urlopen.side_effect = fake_urlopen
-
-        downloader = TUFPointerDownloader(repository_url=staging_url)
+        downloader = TUFPointerDownloader(repository_url=_REPO_URL)
         downloader.download(_PROJECT, version=_VERSION, dest_dir=tmp_path)
 
-        wheel_fetch_url = captured_urls[-1]
-        assert wheel_fetch_url.startswith(staging_url), (
-            f'Expected wheel fetch from {staging_url!r}, got {wheel_fetch_url!r}'
+        mock_urlopen.assert_called_once_with(
+            f'{_REPO_URL}/wheels/{_PROJECT}/datadog_postgres-{_VERSION}-py3-none-any.whl'
         )
 
-    @patch('datadog_checks.downloader.download_v2.Updater')
-    @patch('datadog_checks.downloader.download_v2.urllib.request.urlopen')
     def test_version_none_fetches_latest_pointer(self, mock_urlopen, mock_updater_cls, tmp_path):
-        """version=None fetches <project>/latest.json."""
-        mock_updater = _mock_tuf_updater(_POINTER)
-        mock_updater_cls.return_value = mock_updater
-
-        mock_urlopen.return_value.__enter__ = lambda s: s
-        mock_urlopen.return_value.__exit__ = MagicMock(return_value=False)
-        mock_urlopen.return_value.read.return_value = _WHEEL_CONTENT
-
         downloader = TUFPointerDownloader(repository_url=_REPO_URL)
         downloader.download(_PROJECT, dest_dir=tmp_path)
 
-        call_args = mock_updater.get_targetinfo.call_args[0][0]
-        assert call_args == f'{_PROJECT}/latest.json'
-
-
-# ---------------------------------------------------------------------------
-# Error-path tests
-# ---------------------------------------------------------------------------
+        mock_updater = mock_updater_cls.return_value
+        assert mock_updater.get_targetinfo.call_args[0][0] == f'{_PROJECT}/latest.json'
 
 
 @pytest.mark.offline
 class TestTargetNotFound:
-    @patch('datadog_checks.downloader.download_v2.Updater')
-    @patch('datadog_checks.downloader.download_v2.urllib.request.urlopen')
-    def test_raises_when_tuf_target_absent(self, mock_urlopen, mock_updater_cls, tmp_path):
+    def test_raises_when_tuf_target_absent(self, mock_urlopen, mock_updater_cls):
         mock_updater = MagicMock()
         mock_updater.get_targetinfo.return_value = None
         mock_updater_cls.return_value = mock_updater
@@ -157,66 +118,37 @@ def test_raises_when_tuf_target_absent(self, mock_urlopen, mock_updater_cls, tmp
 
 @pytest.mark.offline
 class TestDigestMismatch:
-    @patch('datadog_checks.downloader.download_v2.Updater')
-    @patch('datadog_checks.downloader.download_v2.urllib.request.urlopen')
     def test_raises_on_corrupted_wheel(self, mock_urlopen, mock_updater_cls, tmp_path):
-        mock_updater_cls.return_value = _mock_tuf_updater(_POINTER)
-
-        mock_urlopen.return_value.__enter__ = lambda s: s
-        mock_urlopen.return_value.__exit__ = MagicMock(return_value=False)
-        mock_urlopen.return_value.read.return_value = b'tampered bytes'
+        mock_urlopen.return_value = _mock_response(b'tampered bytes')
 
         downloader = TUFPointerDownloader(repository_url=_REPO_URL)
         with pytest.raises(DigestMismatch, match=_PROJECT):
             downloader.download(_PROJECT, version=_VERSION, dest_dir=tmp_path)
 
-    @patch('datadog_checks.downloader.download_v2.Updater')
-    @patch('datadog_checks.downloader.download_v2.urllib.request.urlopen')
     def test_raises_on_length_mismatch(self, mock_urlopen, mock_updater_cls, tmp_path):
         bad_pointer = {**_POINTER, 'length': _WHEEL_LENGTH + 1}
         mock_updater_cls.return_value = _mock_tuf_updater(bad_pointer)
 
-        mock_urlopen.return_value.__enter__ = lambda s: s
-        mock_urlopen.return_value.__exit__ = MagicMock(return_value=False)
-        mock_urlopen.return_value.read.return_value = _WHEEL_CONTENT
-
         downloader = TUFPointerDownloader(repository_url=_REPO_URL)
         with pytest.raises(DigestMismatch, match='length'):
             downloader.download(_PROJECT, version=_VERSION, dest_dir=tmp_path)
 
 
-# ---------------------------------------------------------------------------
-# Disable-verification mode
-# ---------------------------------------------------------------------------
-
-
 @pytest.mark.offline
 class TestDisableVerification:
-    @patch('datadog_checks.downloader.download_v2.urllib.request.urlopen')
     def test_directly_downloads_wheel_without_tuf_or_digest_checks(self, mock_urlopen, tmp_path):
-        """disable_verification bypasses TUF/pointers and downloads the wheel by
-        convention from /wheels/<project>/<wheel>.whl."""
-        captured_urls: list[str] = []
-
-        def fake_urlopen(url):
-            captured_urls.append(url)
-            mock_resp = MagicMock()
-            mock_resp.__enter__ = lambda s: s
-            mock_resp.__exit__ = MagicMock(return_value=False)
-            mock_resp.read.return_value = b'bytes not matching any signed pointer'
-            return mock_resp
-
-        mock_urlopen.side_effect = fake_urlopen
+        content = b'bytes not matching any signed pointer'
+        mock_urlopen.return_value = _mock_response(content)
 
         with patch('datadog_checks.downloader.download_v2.Updater') as mock_updater_cls:
             downloader = TUFPointerDownloader(repository_url=_REPO_URL, disable_verification=True)
             wheel_path = downloader.download(_PROJECT, version=_VERSION, dest_dir=tmp_path)
 
-        assert captured_urls == [
-            f'{_REPO_URL}/wheels/{_PROJECT}/datadog_postgres-{_VERSION}-py3-none-any.whl',
-        ]
+        mock_urlopen.assert_called_once_with(
+            f'{_REPO_URL}/wheels/{_PROJECT}/datadog_postgres-{_VERSION}-py3-none-any.whl'
+        )
         assert wheel_path.name == f'datadog_postgres-{_VERSION}-py3-none-any.whl'
-        assert wheel_path.read_bytes() == b'bytes not matching any signed pointer'
+        assert wheel_path.read_bytes() == content
         mock_updater_cls.assert_not_called()
 
     def test_direct_download_requires_explicit_version(self, tmp_path):