handler.py

1import argparse2import json3import os4import random5import shutil6import time7from datetime import datetime, timezone8from typing import Any, Dict9from config import DATA_DIR, INDEX, POSTS, PREFIX, files10from main import expired, make1112# PLACE1314LAMBDA = "AWS_LAMBDA_FUNCTION_NAME"15TMP = "/tmp/feed"16IMMUTABLE = "public, max-age=31536000, immutable"17NO_CACHE = "no-cache"18TYPES = {".mp4": "video/mp4", ".webp": "image/webp", ".json": "application/json"}1920SITE_FUNCTION = "carlomitchener-site"2122def in_lambda() -> bool:23    return bool(os.environ.get(LAMBDA))2425def dry() -> bool:26    return bool(os.environ.get("DRY"))2728def root() -> str:29    return TMP if in_lambda() else DATA_DIR3031def content_type(name: str) -> str:32    return TYPES.get(os.path.splitext(name)[1], "application/octet-stream")3334# S33536def wake_site(name: str) -> None:37    import boto338    try:39        boto3.client("lambda").invoke(FunctionName=SITE_FUNCTION, InvocationType="Event", Payload=json.dumps({"source": "manual", "reason": f"post {name}"}).encode())40        print(f"woke {SITE_FUNCTION}")41    except Exception as error:42        print(f"could not wake {SITE_FUNCTION}: {error}")4344def upload(base: str, name: str, row: Dict[str, Any]) -> int:45    import boto346    from botocore.exceptions import ClientError47    s3 = boto3.client("s3")48    bucket = os.environ["CARLOMITCHENER_BUCKET"]49    count = 050    for file in files(name):51        s3.upload_file(52            os.path.join(base, POSTS, name, file), bucket, f"{PREFIX}/{POSTS}/{name}/{file}",53            ExtraArgs={"ContentType": content_type(file), "CacheControl": IMMUTABLE},54        )55        count += 156    rows = []57    try:58        body = s3.get_object(Bucket=bucket, Key=f"{PREFIX}/{INDEX}")["Body"].read()59        rows = [item for item in json.loads(body) if item.get("name") != row["name"]]60    except ClientError as error:61        if error.response["Error"]["Code"] not in ("NoSuchKey", "404"):62            raise63    for item in expired(rows, datetime.now(timezone.utc)):64        keys = [{"Key": f"{PREFIX}/{POSTS}/{item['name']}/{file}"} for file in files(item["name"])]65        s3.delete_objects(Bucket=bucket, Delete={"Objects": keys})66        rows.remove(item)67        print(f"reap {item['name']}")68    rows.insert(0, row)69    s3.put_object(70        Bucket=bucket, Key=f"{PREFIX}/{INDEX}", Body=json.dumps(rows).encode(),71        ContentType="application/json", CacheControl=NO_CACHE,72    )73    print(f"upload {count + 1} keys to s3://{bucket}/{PREFIX}/ ({len(rows)} posts)")74    return count + 17576# RUN7778def run(seed: int = None) -> Dict[str, Any]:79    started = time.time()80    seed = random.getrandbits(32) if seed is None else int(seed)81    base = root()82    if in_lambda():83        shutil.rmtree(base, ignore_errors=True)84    os.makedirs(base, exist_ok=True)85    print(f"start seed {seed} root {base}")86    record = make(seed, base)87    manifest = record["manifest"]88    if dry():89        print("dry, nothing uploaded")90    else:91        upload(base, manifest["name"], record["row"])92        wake_site(manifest["name"])93    answer = {94        "name": manifest["name"],95        "seed": seed,96        "duration": manifest["duration"],97        "sizes": manifest["sizes"],98        "ms": dict(record["ms"], total=int((time.time() - started) * 1000)),99    }100    print(f"done {json.dumps(answer)}")101    return answer102103def handler(event, context):104    event = event or {}105    return run(event.get("seed"))106107if __name__ == "__main__":108    parser = argparse.ArgumentParser()109    parser.add_argument("--seed", type=int, default=None)110    run(parser.parse_args().seed)