File size: 2,899 Bytes
c7f05a9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
# -*- coding: utf-8 -*-
"""
O STORAGE DO LADO DO WORKER.

O vídeo pronto tem dezenas de MB, e o caminho dele até a tela decide tudo:
com S3 configurado, o worker sobe o arquivo e devolve SÓ a URL — o backend e
o Firestore nunca carregam os bytes. Sem S3, volta base64 (funciona, mas
esbarra no teto de 1 MB do documento; o job store do backend já sabe marcar
`resultado_grande`). O DEPLOYMENT.md manda configurar S3 para vídeo, e este
comentário é o porquê.

Cópia deliberadamente PEQUENA do armazenamento do backend, e não um import
dele: os dois serviços têm deploys separados, e um worker que importa código
do backend acopla os ciclos de vida — a primeira mudança lá derruba cá.
O contrato compartilhado é o FORMATO da resposta ({url} ou {base64}), que é
minúsculo e está testado dos dois lados.
"""
from __future__ import annotations

import base64
import logging
import os
import secrets
import time

logger = logging.getLogger("worker.armazenamento")

_EXTENSAO = {"image/png": ".png", "image/webp": ".webp", "image/jpeg": ".jpg",
             "audio/wav": ".wav", "audio/mpeg": ".mp3", "video/mp4": ".mp4"}


def _s3_configurado() -> bool:
    return bool(os.getenv("STORAGE_ENDPOINT", "").strip()
                and os.getenv("STORAGE_BUCKET", "").strip())


def como_url_ou_base64(conteudo: bytes, mime: str) -> dict:
    if _s3_configurado():
        try:
            import boto3  # type: ignore
            nome = (f"ger/{time.strftime('%Y%m')}/"
                    f"{secrets.token_urlsafe(16)}{_EXTENSAO.get(mime, '.bin')}")
            cliente = boto3.client(
                "s3", endpoint_url=os.getenv("STORAGE_ENDPOINT"),
                aws_access_key_id=os.getenv("STORAGE_ACCESS_KEY", ""),
                aws_secret_access_key=os.getenv("STORAGE_SECRET_KEY", ""),
                region_name=os.getenv("STORAGE_REGION", "auto"))
            cliente.put_object(Bucket=os.getenv("STORAGE_BUCKET"), Key=nome,
                               Body=conteudo, ContentType=mime)
            base_url = os.getenv("STORAGE_URL_PUBLICA", "").strip().rstrip("/")
            if base_url:
                return {"url": f"{base_url}/{nome}", "mime": mime}
            url = cliente.generate_presigned_url(
                "get_object",
                Params={"Bucket": os.getenv("STORAGE_BUCKET"), "Key": nome},
                ExpiresIn=int(os.getenv("STORAGE_LINK_VALIDADE", "604800")))
            return {"url": url, "mime": mime}
        except Exception as e:
            # Storage caído não joga fora uma geração que custou GPU: cai
            # para o base64 e registra — a criação chega, só que pesada.
            logger.warning("S3 falhou (%s); devolvendo embutido", str(e)[:120])
    return {"base64": "data:%s;base64,%s" % (mime,
                                             base64.b64encode(conteudo).decode()),
            "mime": mime}