Django management command to push whitenoise static files to S3, respecting headers

push_static_to_s3.py

import logging
from concurrent import futures
from typing import List, Tuple

import boto3
from django.conf import settings
from django.core.management import BaseCommand
from whitenoise.middleware import WhiteNoiseMiddleware

log = logging.getLogger(__name__)


class Command(BaseCommand):
    help = (
        "Push static assets to S3. "
        "Should be done during release phase (same time as migrations)"
    )

    def add_arguments(self, parser):
        parser.add_argument(
            "-t",
            "--threads",
            type=int,
            default=10,
            dest="threads",
            help="Simultaneous threads to spawn for upload",
        )

    def handle(self, **options):
        """Upload static files to S3 maintaining WhiteNoise headers"""
        static_files = WhiteNoiseMiddleware().files
        log.info("Uploading %s files to S3...", len(static_files))
        bucket = settings.AWS_STORAGE_BUCKET_NAME
        s3 = boto3.client("s3")

        def upload_file(key: str, filename: str, headers: List[Tuple[str, str]]):
            key = key.lstrip("/")
            extra_args = {
                k.replace("-", ""): v
                for k, v in headers
                if k in ["Cache-Control", "Content-Type"]
            }
            s3.upload_file(
                Filename=filename, Bucket=bucket, Key=key, ExtraArgs=extra_args
            )
            log.info("Uploaded s3://%s/%s", bucket, key)

        with futures.ThreadPoolExecutor(max_workers=options["threads"]) as executor:
            futures.wait(
                [
                    executor.submit(upload_file, path, *sf.get_path_and_headers({}))
                    for path, sf in static_files.items()
                ],
                return_when=futures.FIRST_EXCEPTION,
            )