From 887ae918ec6816b8fa76e82e1da51a7bb430a010 Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Wed, 7 Aug 2024 09:57:26 +0300 Subject: [PATCH 1/6] Refactor things to ArtContext --- art/command.py | 41 ++++++++++++++++++++++++----------------- art/context.py | 8 ++++++++ art/s3.py | 12 ++++++++++-- art/write.py | 30 +++++++++++++++++++++--------- art_tests/test_s3.py | 3 ++- art_tests/test_write.py | 8 +++++--- 6 files changed, 70 insertions(+), 32 deletions(-) create mode 100644 art/context.py diff --git a/art/command.py b/art/command.py index 7c89c89..ccf32d5 100644 --- a/art/command.py +++ b/art/command.py @@ -5,10 +5,11 @@ import os import shutil import tempfile -from typing import Any, Dict, List, Optional +from typing import List, Optional from art.config import ArtConfig, FileMapEntry from art.consts import DEFAULT_CONFIG_FILENAME +from art.context import ArtContext from art.excs import Problem from art.git import git_clone from art.manifest import Manifest @@ -95,33 +96,35 @@ def run_command(argv: Optional[List[str]] = None) -> None: args = Args(**vars(ap.parse_args(argv))) logging.basicConfig(level=(args.log_level or logging.INFO)) - config_args: Dict[str, Any] = {"dests": list(args.dests), "name": ""} - is_git = False if args.git_source: - config_args.update( + work_dir = tempfile.mkdtemp(prefix="art-git-") + atexit.register(shutil.rmtree, work_dir) + config = ArtConfig( + dests=list(args.dests), + name="", repo_url=args.git_source, ref=args.git_ref, - work_dir=tempfile.mkdtemp(prefix="art-git-"), + work_dir=work_dir, ) - is_git = True + git_clone(config) elif args.local_source: work_dir = os.path.abspath(args.local_source) - config_args.update( + config = ArtConfig( + dests=list(args.dests), + name="", repo_url=work_dir, work_dir=work_dir, ) else: ap.error("Either a git source or a local source must be defined") - - config = ArtConfig(**config_args) - - if is_git: - git_clone(config) - atexit.register(shutil.rmtree, config.work_dir) + return + context = ArtContext( + dry_run=bool(args.dry_run), + ) for forked_config in fork_configs_from_work_dir(config, filename=args.config_file): try: - process_config_postfork(args, forked_config) + process_config_postfork(context, args, forked_config) except Problem as p: ap.error(f"config {forked_config.name}: {p}") @@ -132,7 +135,11 @@ def clean_dest(dest: str) -> str: return dest -def process_config_postfork(args: Args, config: ArtConfig) -> None: +def process_config_postfork( + context: ArtContext, + args: Args, + config: ArtConfig, +) -> None: if not config.dests: raise Problem("No destination(s) specified (on command line or in config in source)") config.dests = [clean_dest(dest) for dest in config.dests] @@ -152,12 +159,12 @@ def process_config_postfork(args: Args, config: ArtConfig) -> None: for dest in config.dests: for suffix in suffixes: write( - config, + context=context, + config=config, dest=dest, path_suffix=suffix, manifest=manifest, wrap_filename=wrap_temp, - dry_run=args.dry_run, ) if wrap_temp: os.unlink(wrap_temp) diff --git a/art/context.py b/art/context.py new file mode 100644 index 0000000..36de930 --- /dev/null +++ b/art/context.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +import dataclasses + + +@dataclasses.dataclass(frozen=True) +class ArtContext: + dry_run: bool = False diff --git a/art/s3.py b/art/s3.py index c6624c9..b38c34e 100644 --- a/art/s3.py +++ b/art/s3.py @@ -2,6 +2,8 @@ from typing import IO, Any, Dict from urllib.parse import urlparse +from art.context import ArtContext + _s3_client = None log = logging.getLogger(__name__) @@ -15,7 +17,13 @@ def get_s3_client() -> Any: return _s3_client -def s3_write(url: str, source_fp: IO[bytes], *, options: Dict[str, Any], dry_run: bool) -> None: +def s3_write( + url: str, + source_fp: IO[bytes], + *, + options: Dict[str, Any], + context: ArtContext, +) -> None: purl = urlparse(url) s3_client = get_s3_client() assert purl.scheme == "s3" @@ -27,7 +35,7 @@ def s3_write(url: str, source_fp: IO[bytes], *, options: Dict[str, Any], dry_run if acl: kwargs["ACL"] = acl - if dry_run: + if context.dry_run: log.info("Dry-run: would write to S3 (ACL %s): %s", acl, url) return s3_client.put_object(**kwargs) diff --git a/art/write.py b/art/write.py index 2149eb0..f503a99 100644 --- a/art/write.py +++ b/art/write.py @@ -7,6 +7,7 @@ from urllib.parse import parse_qsl from art.config import ArtConfig +from art.context import ArtContext from art.manifest import Manifest from art.s3 import s3_write @@ -17,13 +18,13 @@ def _write_file( dest: str, source_fp: IO[bytes], *, + context: ArtContext, options: Optional[Dict[str, Any]] = None, - dry_run: bool = False, ) -> None: if options is None: options = {} writer = _get_writer_for_dest(dest) - writer(dest, source_fp, options=options, dry_run=dry_run) + writer(dest, source_fp, options=options, context=context) def _get_writer_for_dest(dest: str) -> Callable: # type: ignore[type-arg] @@ -34,8 +35,14 @@ def _get_writer_for_dest(dest: str) -> Callable: # type: ignore[type-arg] raise ValueError(f"Invalid destination: {dest}") -def local_write(dest: str, source_fp: IO[bytes], *, options: Dict[str, Any], dry_run: bool) -> None: - if dry_run: +def local_write( + dest: str, + source_fp: IO[bytes], + *, + context: ArtContext, + options: Dict[str, Any], +) -> None: + if context.dry_run: log.info("Dry-run: Would have written local file %s", dest) return os.makedirs(os.path.dirname(dest), exist_ok=True) @@ -45,12 +52,12 @@ def local_write(dest: str, source_fp: IO[bytes], *, options: Dict[str, Any], dry def write( - config: ArtConfig, *, + context: ArtContext, + config: ArtConfig, dest: str, path_suffix: str, manifest: Manifest, - dry_run: bool, wrap_filename: Optional[str] = None, ) -> None: options = {} @@ -63,13 +70,18 @@ def write( dest_path = posixpath.join(dest, dest_filename) local_path = os.path.join(config.work_dir, fileinfo["path"]) with open(local_path, "rb") as infp: - _write_file(dest_path, infp, options=options, dry_run=dry_run) + _write_file( + dest_path, + infp, + context=context, + options=options, + ) _write_file( dest=posixpath.join(dest, ".manifest.json"), source_fp=io.BytesIO(manifest.as_json_bytes()), + context=context, options=options, - dry_run=dry_run, ) if config.wrap and wrap_filename: @@ -77,6 +89,6 @@ def write( _write_file( dest=posixpath.join(dest, config.wrap), source_fp=infp, + context=context, options=options, - dry_run=dry_run, ) diff --git a/art_tests/test_s3.py b/art_tests/test_s3.py index daedb90..ec5d76d 100644 --- a/art_tests/test_s3.py +++ b/art_tests/test_s3.py @@ -1,5 +1,6 @@ import io +from art.context import ArtContext from art.s3 import get_s3_client from art.write import _write_file @@ -9,5 +10,5 @@ def test_s3_acl(mocker): cli.put_object = cli.put_object # avoid magic mocker.patch.object(cli, "put_object") body = io.BytesIO(b"test") - _write_file("s3://bukkit/key", body, options={"acl": "public-read"}) + _write_file("s3://bukkit/key", body, options={"acl": "public-read"}, context=ArtContext()) cli.put_object.assert_called_with(Bucket="bukkit", Key="key", ACL="public-read", Body=body) diff --git a/art_tests/test_write.py b/art_tests/test_write.py index 2a104ad..72fbc8a 100644 --- a/art_tests/test_write.py +++ b/art_tests/test_write.py @@ -1,5 +1,6 @@ import art.write from art.config import ArtConfig +from art.context import ArtContext from art.manifest import Manifest @@ -7,12 +8,13 @@ def test_dest_options(mocker, tmpdir): cfg = ArtConfig(work_dir=str(tmpdir), dests=[str(tmpdir)], name="", repo_url=str(tmpdir)) mf = Manifest(files={}) wf = mocker.patch("art.write._write_file") + context = ArtContext(dry_run=False) art.write.write( - cfg, + config=cfg, + context=context, dest="derp://foo/bar/?acl=quux", - path_suffix="blag", manifest=mf, - dry_run=False, + path_suffix="blag", ) call_kwargs = wf.call_args[1] assert call_kwargs["options"] == {"acl": "quux"} From a1ecbae00171a8783b9a010dce802d130ec43dd1 Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Wed, 7 Aug 2024 10:00:49 +0300 Subject: [PATCH 2/6] Use `@cache` for S3 client caching --- art/s3.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/art/s3.py b/art/s3.py index b38c34e..f5f9b2f 100644 --- a/art/s3.py +++ b/art/s3.py @@ -1,20 +1,18 @@ import logging +from functools import cache from typing import IO, Any, Dict from urllib.parse import urlparse from art.context import ArtContext -_s3_client = None log = logging.getLogger(__name__) +@cache def get_s3_client() -> Any: - global _s3_client - if not _s3_client: - import boto3 + import boto3 - _s3_client = boto3.client("s3") - return _s3_client + return boto3.client("s3") def s3_write( From 7bc62a0cfa157b00e7c41c4bfa808e70b9cf6102 Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Wed, 7 Aug 2024 10:23:22 +0300 Subject: [PATCH 3/6] Drop pytest-mock in favor of standard library --- art_tests/test_s3.py | 6 ++++-- art_tests/test_write.py | 7 +++++-- pyproject.toml | 1 - 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/art_tests/test_s3.py b/art_tests/test_s3.py index ec5d76d..968811b 100644 --- a/art_tests/test_s3.py +++ b/art_tests/test_s3.py @@ -1,14 +1,16 @@ import io +from unittest.mock import Mock from art.context import ArtContext from art.s3 import get_s3_client from art.write import _write_file -def test_s3_acl(mocker): +def test_s3_acl(monkeypatch): cli = get_s3_client() cli.put_object = cli.put_object # avoid magic - mocker.patch.object(cli, "put_object") + put_object = Mock() + monkeypatch.setattr(cli, "put_object", put_object) body = io.BytesIO(b"test") _write_file("s3://bukkit/key", body, options={"acl": "public-read"}, context=ArtContext()) cli.put_object.assert_called_with(Bucket="bukkit", Key="key", ACL="public-read", Body=body) diff --git a/art_tests/test_write.py b/art_tests/test_write.py index 72fbc8a..3beaa86 100644 --- a/art_tests/test_write.py +++ b/art_tests/test_write.py @@ -1,13 +1,16 @@ +import unittest.mock + import art.write from art.config import ArtConfig from art.context import ArtContext from art.manifest import Manifest -def test_dest_options(mocker, tmpdir): +def test_dest_options(monkeypatch, tmpdir): cfg = ArtConfig(work_dir=str(tmpdir), dests=[str(tmpdir)], name="", repo_url=str(tmpdir)) mf = Manifest(files={}) - wf = mocker.patch("art.write._write_file") + wf = unittest.mock.MagicMock() + monkeypatch.setattr(art.write, "_write_file", wf) context = ArtContext(dry_run=False) art.write.write( config=cfg, diff --git a/pyproject.toml b/pyproject.toml index 8a81d75..ba4bb80 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,6 @@ mypy = [ test = [ "build", "pytest-cov", - "pytest-mock~=3.12", "pytest~=8.1", ] From 6128e6a62dda01ca4690d65fbdecbba51e78e1f0 Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Wed, 7 Aug 2024 10:35:26 +0300 Subject: [PATCH 4/6] Add support for CloudFront invalidations --- art/cloudfront.py | 30 ++++++++++++++++++++++++++++++ art/command.py | 1 + art/context.py | 11 +++++++++++ art/s3.py | 4 ++++ art_tests/test_s3.py | 36 ++++++++++++++++++++++++++++++++++++ 5 files changed, 82 insertions(+) create mode 100644 art/cloudfront.py diff --git a/art/cloudfront.py b/art/cloudfront.py new file mode 100644 index 0000000..f42929a --- /dev/null +++ b/art/cloudfront.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import logging +from typing import Any + +log = logging.getLogger(__name__) + + +# Separated for testing purposes +def get_cloudfront_client() -> Any: + import boto3 + + return boto3.client("cloudfront") + + +def execute_cloudfront_invalidations(invalidations: dict[str, set[str]]) -> None: + cf_client = get_cloudfront_client() + for dist_id, paths in invalidations.items(): + log.info("Creating CloudFront invalidation for %s: %d paths", dist_id, len(paths)) + inv = cf_client.create_invalidation( + DistributionId=dist_id, + InvalidationBatch={ + "Paths": { + "Quantity": len(paths), + "Items": sorted(paths), + }, + "CallerReference": "art", + }, + ) + log.info("Created CloudFront invalidation: %s", inv["Invalidation"]["Id"]) diff --git a/art/command.py b/art/command.py index ccf32d5..70836a2 100644 --- a/art/command.py +++ b/art/command.py @@ -127,6 +127,7 @@ def run_command(argv: Optional[List[str]] = None) -> None: process_config_postfork(context, args, forked_config) except Problem as p: ap.error(f"config {forked_config.name}: {p}") + context.execute_post_run_tasks() def clean_dest(dest: str) -> str: diff --git a/art/context.py b/art/context.py index 36de930..7854fdf 100644 --- a/art/context.py +++ b/art/context.py @@ -2,7 +2,18 @@ import dataclasses +from art.cloudfront import execute_cloudfront_invalidations + @dataclasses.dataclass(frozen=True) class ArtContext: dry_run: bool = False + _cloudfront_invalidations: dict[str, set[str]] = dataclasses.field(default_factory=dict) + + def add_cloudfront_invalidation(self, dist_id: str, path: str) -> None: + self._cloudfront_invalidations.setdefault(dist_id, set()).add(path) + + def execute_post_run_tasks(self) -> None: + if self._cloudfront_invalidations: + execute_cloudfront_invalidations(self._cloudfront_invalidations) + self._cloudfront_invalidations.clear() diff --git a/art/s3.py b/art/s3.py index f5f9b2f..75474af 100644 --- a/art/s3.py +++ b/art/s3.py @@ -38,3 +38,7 @@ def s3_write( return s3_client.put_object(**kwargs) log.info("Wrote to S3 (ACL %s): %s", acl, url) + + cf_distribution_id = options.get("cf-distribution-id") + if cf_distribution_id: + context.add_cloudfront_invalidation(cf_distribution_id, purl.path) diff --git a/art_tests/test_s3.py b/art_tests/test_s3.py index 968811b..d37e5f3 100644 --- a/art_tests/test_s3.py +++ b/art_tests/test_s3.py @@ -1,11 +1,21 @@ import io from unittest.mock import Mock +import pytest +from boto3 import _get_default_session + +from art import cloudfront from art.context import ArtContext from art.s3 import get_s3_client from art.write import _write_file +@pytest.fixture(autouse=True) +def aws_fake_credentials(monkeypatch): + # Makes sure we don't accidentally use real AWS credentials. + monkeypatch.setattr(_get_default_session()._session, "_credentials", Mock()) + + def test_s3_acl(monkeypatch): cli = get_s3_client() cli.put_object = cli.put_object # avoid magic @@ -14,3 +24,29 @@ def test_s3_acl(monkeypatch): body = io.BytesIO(b"test") _write_file("s3://bukkit/key", body, options={"acl": "public-read"}, context=ArtContext()) cli.put_object.assert_called_with(Bucket="bukkit", Key="key", ACL="public-read", Body=body) + + +def test_s3_invalidate_cloudfront(monkeypatch): + cli = get_s3_client() + cli.put_object = cli.put_object # avoid magic + put_object = Mock() + monkeypatch.setattr(cli, "put_object", put_object) + body = io.BytesIO(b"test") + options = {"acl": "public-read", "cf-distribution-id": "UWUWU"} + context = ArtContext() + _write_file("s3://bukkit/key/foo/bar", body, options=options, context=context) + _write_file("s3://bukkit/key/baz/quux", body, options=options, context=context) + _write_file("s3://bukkit/key/baz/barple", body, options=options, context=context) + cf_client = Mock() + cf_client.create_invalidation.return_value = {"Invalidation": {"Id": "AAAAA"}} + monkeypatch.setattr(cloudfront, "get_cloudfront_client", Mock(return_value=cf_client)) + context.execute_post_run_tasks() + # Assert the 3 files get a single invalidation + cf_client.create_invalidation.assert_called_once() + call_kwargs = cf_client.create_invalidation.call_args.kwargs + assert call_kwargs["DistributionId"] == "UWUWU" + assert set(call_kwargs["InvalidationBatch"]["Paths"]["Items"]) == { + "/key/baz/barple", + "/key/baz/quux", + "/key/foo/bar", + } From 548f0ecd0d783db5a7a6ab9532c6eaf3b440dfd2 Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Wed, 7 Aug 2024 13:19:37 +0300 Subject: [PATCH 5/6] Upgrade lint tools --- .pre-commit-config.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cc405d5..937675f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,13 +1,13 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.3.3 + rev: v0.5.6 hooks: - id: ruff args: - --fix - id: ruff-format - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.5.0 + rev: v4.6.0 hooks: - id: check-merge-conflict - id: check-yaml @@ -18,6 +18,6 @@ repos: args: - --fix=lf - repo: https://github.com/crate-ci/typos - rev: v1.19.0 + rev: v1.23.6 hooks: - id: typos From 1704eff3fe198f13de705125902ba6ae9ab309d8 Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Wed, 7 Aug 2024 14:03:44 +0300 Subject: [PATCH 6/6] Make Cloudfront caller references unique --- art/cloudfront.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/art/cloudfront.py b/art/cloudfront.py index f42929a..08ec143 100644 --- a/art/cloudfront.py +++ b/art/cloudfront.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import time from typing import Any log = logging.getLogger(__name__) @@ -15,8 +16,10 @@ def get_cloudfront_client() -> Any: def execute_cloudfront_invalidations(invalidations: dict[str, set[str]]) -> None: cf_client = get_cloudfront_client() + ts = int(time.time()) for dist_id, paths in invalidations.items(): log.info("Creating CloudFront invalidation for %s: %d paths", dist_id, len(paths)) + caller_reference = f"art-{dist_id}-{ts}" inv = cf_client.create_invalidation( DistributionId=dist_id, InvalidationBatch={ @@ -24,7 +27,11 @@ def execute_cloudfront_invalidations(invalidations: dict[str, set[str]]) -> None "Quantity": len(paths), "Items": sorted(paths), }, - "CallerReference": "art", + "CallerReference": caller_reference, }, ) - log.info("Created CloudFront invalidation: %s", inv["Invalidation"]["Id"]) + log.info( + "Created CloudFront invalidation with caller reference %s: %s", + caller_reference, + inv["Invalidation"]["Id"], + )