fix(transferegov): download de anexos em lotes e LPG/PNAB primeiro

dags/data_ingest/transferegov_fundo_a_fundo/download_anexos_dag.py

@@ -1,5 +1,6 @@
 import base64
 import logging
+from collections import Counter
 
 import pandas as pd
 import requests
@@ -23,6 +24,13 @@
 # Teto por rodada: com max_active_tis_per_dag=3 isso ja e horas de download.
 _LIMITE_POR_RODADA = 3000
 
+# Anexos por task mapeada. O Airflow recusa expandir acima de `max_map_length`
+# (padrao 1024; o compose local sobe 15000, homologacao nao) e derruba a
+# buscar_ids_pendentes depois do return. Mapear sobre lotes, como a
+# salic_ingestion_trino e a extracao_anexos_dag, nao depende dessa config:
+# 3000 / 50 = 60 tasks.
+_ANEXOS_POR_LOTE = 50
+
 ASSET_ANEXOS = datalakehouse.asset_staging(
     schemas.FONTE_TRANSFEREGOV, schemas.TABELA_ANEXO_RELATORIO
 )
@@ -47,8 +55,9 @@
 def download_anexos_dag() -> None:
 
     @task
-    def buscar_ids_pendentes() -> list:
-        """Anexos de planilha que ainda não estão em ``raw/.../anexos_arquivos/``.
+    def buscar_ids_pendentes() -> list[list[dict]]:
+        """Anexos de planilha que ainda não estão em ``raw/.../anexos_arquivos/``,
+        em lotes de ``_ANEXOS_POR_LOTE``.
 
         A lista sai do staging (anexo -> relatório -> plano de ação, para saber
         o programa e daí a pasta LPG/PNAB) menos o que já foi baixado. Quem
@@ -79,25 +88,37 @@ def staging(entidade: str) -> pd.DataFrame:
             {"id": str(linha.id), "politica": politica_do_programa(linha.id_programa)}
             for linha in anexos.drop_duplicates("id").itertuples()
             if str(linha.id) not in baixados
-        ][:_LIMITE_POR_RODADA]
+        ]
+        # LPG e PNAB antes de `outros`, e so depois o teto: so as duas primeiras
+        # viram tabela na extracao_anexos_dag, entao sao elas que nao podem
+        # esperar outra rodada. `outros` continua sendo baixado, por ultimo. O
+        # sorted e estavel: dentro de cada grupo, a ordem do staging se mantem.
+        pendentes = sorted(pendentes, key=lambda p: p["politica"] == "outros")
+        pendentes = pendentes[:_LIMITE_POR_RODADA]
+
+        lotes = [
+            pendentes[i : i + _ANEXOS_POR_LOTE]
+            for i in range(0, len(pendentes), _ANEXOS_POR_LOTE)
+        ]
 
+        por_politica = Counter(p["politica"] for p in pendentes)
         logging.info(
-            "%d anexos de planilha no staging, %d já baixados, %d nesta rodada",
+            "%d anexos de planilha no staging, %d já baixados, %d nesta rodada "
+            "em %d lotes (lpg %d, pnab %d, outros %d)",
             len(anexos),
             len(baixados),
             len(pendentes),
+            len(lotes),
+            por_politica["lpg"],
+            por_politica["pnab"],
+            por_politica["outros"],
         )
-        return pendentes
+        return lotes
 
-    @task(max_active_tis_per_dag=3)
-    def baixar_e_salvar_anexos(pendente: dict) -> None:
-        """
-        Processa UM único anexo por invocação (Dynamic Task Mapping).
-        Faz o request na API, decodifica o base64 e grava o binário em
+    def _baixar_e_salvar(pendente: dict) -> None:
+        """Faz o request na API, decodifica o base64 e grava o binário em
         ``raw/transferegov/anexos_arquivos/<politica>/`` -- em partes, se passar
         do limite do proxy do MinIO (``datalakehouse.gravar_bytes_em_partes``).
-        A concorrência é controlada por max_active_tis_per_dag=3 para evitar
-        HTTP 429 (rate limit) na API do governo.
         """
         anexo_id = pendente["id"]
 
@@ -120,6 +141,31 @@ def baixar_e_salvar_anexos(pendente: dict) -> None:
         )
         logging.info(f"Anexo {anexo_id} salvo em s3://{datalakehouse.BUCKET}/{key}")
 
+    @task(max_active_tis_per_dag=3)
+    def baixar_e_salvar_anexos(lote: list[dict]) -> None:
+        """
+        Processa um lote de anexos por invocação (Dynamic Task Mapping), um
+        anexo de cada vez. A concorrência é controlada por
+        max_active_tis_per_dag=3 para evitar HTTP 429 (rate limit) na API do
+        governo: são no máximo três requisições ao mesmo tempo.
+
+        Anexo que falha não derruba o lote: fica fora do bucket e volta como
+        pendente na próxima rodada. O lote só falha se nenhum anexo dele
+        passou -- sinal de API fora ou bloqueada, que tem que aparecer
+        vermelho na grade.
+        """
+        falhas = 0
+        for pendente in lote:
+            try:
+                _baixar_e_salvar(pendente)
+            except Exception as exc:
+                falhas += 1
+                logging.warning("Anexo %s falhou: %r", pendente["id"], exc)
+
+        logging.info("Lote: %d de %d anexos com falha", falhas, len(lote))
+        if lote and falhas == len(lote):
+            raise RuntimeError(f"Todos os {len(lote)} anexos do lote falharam")
+
     # Roda mesmo se alguns downloads falharem (ALL_DONE): o que foi baixado
     # já pode ser extraído, e o que falhou volta como pendente na próxima.
     @task(outlets=[ASSET_ANEXOS_ARQUIVOS], trigger_rule=TriggerRule.ALL_DONE)
@@ -128,9 +174,9 @@ def anunciar_downloads() -> None:
             ASSET_ANEXOS_ARQUIVOS, datalakehouse.ler_cadeia(ASSET_ANEXOS)
         )
 
-    # Fluxo da DAG — Dynamic Task Mapping: cada pendente vira uma task individual
-    lista_pendentes = buscar_ids_pendentes()
-    baixar_e_salvar_anexos.expand(pendente=lista_pendentes) >> anunciar_downloads()
+    # Fluxo da DAG — Dynamic Task Mapping: cada lote de pendentes vira uma task
+    lotes_pendentes = buscar_ids_pendentes()
+    baixar_e_salvar_anexos.expand(lote=lotes_pendentes) >> anunciar_downloads()
 
 
 # Instancia a DAG