diff --git a/dags/data_ingest/transferegov_fundo_a_fundo/download_anexos_dag.py b/dags/data_ingest/transferegov_fundo_a_fundo/download_anexos_dag.py index 2ddbcc71..28736d20 100644 --- a/dags/data_ingest/transferegov_fundo_a_fundo/download_anexos_dag.py +++ b/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//`` -- 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