Uma COPY lambda sempre lê os dados, né? Não é bem assim

Quando comecei a escrever uma lambda que dispara um COPY no Redshift eu achava que por padrão, a lambda leria os dados que estava copiando, e não é o que acontece (mas poderia), portanto, você pode estar copiando dados duplicados sem nunca perceber, ou queimando retries com dados envenenados. Esse post vai tentar explicar a diferença entre control planes e data planes, e como uma COPY lambda disparada pela criação de mensagens no SQS se encaixa nisso.

Control planes e data planes

Para deixar claro, o termo plane se refere a uma camada da arquitetura (originalmente uma camada de rede, e só conceitualmente) onde algum tipo específico de tarefa é executado, e o significado dele é definido pela fronteira que você está desenhando.
Então a diferença entre control e data planes de forma geral é a seguinte:

Control plane: é o tomador de decisão, o que decide para onde os dados devem ir, e faz o trabalho administrativo, o setup, as políticas de segurança.

Data plane: esse aqui executa essas decisões, move os dados da origem para o destino, e segue as regras determinadas pelo control plane. É o que de fato toca os dados.

A AWS tem um conceito diferente, porque para eles control e data plane podem ser diferenciados dentro do mesmo serviço. Por exemplo, o control plane do Lambda é o CreateFunction, e o data plane dele é o Invoke. No meu contexto, esses conceitos envolvem o pipeline de ingestão e não o serviço em si. A fronteira que estou desenhando é basicamente por qual memória os bytes passam. Em outras palavras, pela definição da AWS essa lambda chama sim uma API de data plane, redshift-data:ExecuteStatement, já que o Redshift está fazendo o trabalho principal dele ali e não uma chamada de gerenciamento.

No pipeline de ingestão, a COPY lambda parece no começo ser um componente de data plane, mas quando olhamos mais de perto a implementação percebemos que os dados não são lidos em momento nenhum, porque a lambda só fala para o Redshift movê-los. Ela entrega um ponteiro (os manifests JSON) e os bytes dos dados reais vão para o cluster do Redshift sem nem tocar na memória da lambda.

O que o código realmente faz

Para fazer uma operação de COPY que carrega os dados em batch para um cluster do Redshift não precisamos ler os dados em momento algum, e ler eles só daria mais pontos de falha para a lambda como um todo. Aqui está uma implementação simples disso, com parte do tratamento de erro e do carregamento de configuração de fora:

import json
import time

import boto3

s3 = boto3.client("s3")
redshift = boto3.client("redshift-data")

MAX_POLL_ATTEMPTS = 60
POLL_INTERVAL_SECONDS = 1


def handler(event, context):
    urls = []
    failures = []

    for record in event["Records"]:  # mensagens do SQS anunciando arquivos, e não os arquivos
        try:
            message = json.loads(record["body"])
            urls.extend(message["urls"])
        except (json.JSONDecodeError, KeyError):
            failures.append({"itemIdentifier": record["messageId"]})

    manifest = {"entries": [{"url": url, "mandatory": True} for url in urls]}
    key = f"manifests/{time.time_ns()}.json"
    s3.put_object(Bucket=BUCKET, Key=key, Body=json.dumps(manifest))

    statement = redshift.execute_statement(
        ClusterIdentifier=CLUSTER,
        Database=DATABASE,
        Sql=f"COPY raw_events FROM 's3://{BUCKET}/{key}' IAM_ROLE '{IAM_ROLE}' FORMAT AS JSON 'auto' MANIFEST",
    )

    status = "SUBMITTED"
    for _ in range(MAX_POLL_ATTEMPTS):
        status = redshift.describe_statement(Id=statement["Id"])["Status"]
        if status not in ("SUBMITTED", "PICKED", "STARTED"):
            break
        time.sleep(POLL_INTERVAL_SECONDS)

    if status != "FINISHED":  # um único COPY, então o batch inteiro volta para a fila
        return {"batchItemFailures": [{"itemIdentifier": r["messageId"]} for r in event["Records"]]}

    return {"batchItemFailures": failures}

Se olharmos os dois clients criados no topo conseguimos ver o ponto principal desse post, porque essa função só conversa com o S3 através do put_object e com o Redshift através do execute_statement, e não existe nenhum get_object em lugar nenhum e nenhum terceiro client. As urls que vêm dentro da mensagem são copiadas para o manifest como strings e nunca são abertas, então a lambda só anota onde os dados estão e pede para o Redshift ir lá buscar.

Como podemos ver, produzimos um único manifest a partir de todos os records do evento do SQS, mas se por algum motivo (como aquele que vou explicar na próxima seção) alguma das mensagens for re-enfileirada no SQS podemos ter problemas de duplicação, já que por natureza o statement de COPY não é idempotente. E como o COPY é feito sobre um manifest, quando a mensagem volta o manifest inteiro é copiado de novo, então um retry duplica um batch e não só um record.

O que pode quebrar

O código acima nunca valida os dados. Ele valida a mensagem que aponta para eles, o que não é a mesma coisa. Os records são basicamente referências para os dados reais, então se existirem erros lá, o COPY vai falhar do lado do Redshift e a lambda não tem como saber disso de antemão. Isso significa que o statement de COPY é a primeira leitura que impõe o schema, e nada checa esses bytes contra o schema de destino até o COPY fazer isso.

Além disso, sem um tratamento adequado disso em outra lambda ou uma configuração adequada da DLQ (dead-letter queue), a mesma mensagem envenenada pode voltar para a fila do SQS de novo e de novo, resultando no mesmo erro, até a mensagem simplesmente expirar. Vamos imaginar um cenário onde temos um visibility timeout de 15 minutos e uma redrive policy com um maxReceiveCount de 5. O visibility timeout precisa ser tão grande porque ele precisa cobrir a invocação inteira, incluindo o tempo que a lambda pode gastar fazendo polling, senão o SQS entregaria a mesma mensagem para uma segunda invocação enquanto a primeira ainda está esperando o COPY dela terminar.

Com esses números a história do retry fica assim. A mensagem envenenada é recebida, o COPY falha, e 15 minutos depois a mensagem fica visível de novo, e isso se repete até o quinto recebimento, quando o SQS finalmente move ela para a DLQ, onde vai esperar o período de retenção (até 14 dias) até alguém olhar para ela. Até esses limites serem atingidos, a mensagem vai voltar para essa lambda e vamos replicar os mesmos erros.

Também vale dizer que retornar batchItemFailures é o que torna isso suportável. Com o ReportBatchItemFailures habilitado no event source mapping, só as mensagens listadas nesse retorno voltam para a fila, e não o batch inteiro. A parte interessante é que só conseguimos ser precisos com metade disso, porque uma mensagem que não é um JSON válido ou que não tem a chave urls é algo que a lambda tem em memória e consegue checar, então conseguimos mandar de volta exatamente esses ids. No contexto do COPY, a história é outra, já que existe um único statement sobre um único manifest, e a menor coisa que pode falhar ali é o batch inteiro, então um arquivo ruim dentro dele manda todas as mensagens de volta, e no retry todo arquivo bom daquele batch é carregado de novo. Em outras palavras, a lambda consegue ser específica sobre as mensagens que ela lê, e só consegue ser genérica sobre os dados que ela nunca toca.

O loop de poll tem um problema parecido. Quando o MAX_POLL_ATTEMPTS acaba paramos de observar o statement, mas parar de observar não para o COPY, e ele continua rodando no cluster. As mensagens voltam para a fila, a próxima invocação escreve um novo manifest com as mesmas urls, e acabamos com dois statements carregando os mesmos arquivos. Poderíamos chamar o cancel_statement nesse branch, mas isso só ajudaria enquanto o COPY ainda estivesse rodando, e quando o poll desiste ele já poderia ter carregado tudo, então a duplicação aconteceria de qualquer jeito.

Outro cenário é mandar a mensagem sem pensar em idempotência. O Redshift não impõe primary keys nem unique constraints, então se a mesma mensagem for entregue duas vezes (e com uma fila SQS standard, a entrega at-least-once significa que ela pode ser), poderíamos ter duas mensagens corretas e entradas duplicadas no Redshift.

Possíveis correções

Para corrigir os erros descritos acima podemos desenvolver alguns tipos de solução. Por exemplo, para evitar mensagens envenenadas, deveríamos ter alguma outra lambda que vai validar elas antes de chegarem no SQS. Outra opção é validar elas na própria COPY lambda, o que para mim não é uma boa solução, já que implicaria baixar os dados reais da origem, e arrastaria a lambda para o caminho dos dados e a desviaria do objetivo original dela.

Para o problema de idempotência, de novo podemos ter outra lambda antes dessa para ser responsável por criar algumas chaves de idempotência e passar essa informação para essa lambda, ou também podemos fazer esse trabalho aqui, mas de novo vai ser mais um caso de um componente com múltiplas responsabilidades (data plane e control plane ao mesmo tempo), o que não é o comportamento desejado. Para esse caso a melhor solução é provavelmente deixar um statement do Redshift ser responsável por garantir a idempotência descartando as duplicatas. A query abaixo, por exemplo, pode rodar como um model em uma ferramenta de analytics como o dbt, e tudo que está downstream leria esse model em vez de ler a tabela raw:

with ranked as (
    select
        event_id,
        payload,
        loaded_at,
        row_number() over (
            partition by md5(payload)
            order by loaded_at
        ) as duplicate_rank
    from raw_events
)

select
    event_id,
    payload,
    loaded_at
from ranked
where duplicate_rank = 1

A chave de deduplicação aqui é o hash do payload inteiro, que já carrega o identificador do evento dentro dele, o row_number() numera as linhas dentro de cada grupo de linhas idênticas, e manter só o rank 1 colapsa cada grupo em uma única linha. É importante notar que essa versão mantém uma linha por grupo, porque a variante que filtra por count(*) over (partition by ...) = 1 parece muito parecida mas joga fora o grupo duplicado inteiro, perdendo o registro em vez de deduplicar ele. O order by só decide qual cópia sobrevive, e como as linhas são idênticas não importa qual seja.

Os leitores mais atentos provavelmente já notaram que o COPY em si continua não sendo idempotente, mas como estamos pensando no contexto de uma implementação geral do pipeline, o objetivo final é garantir que o que quer que leia os dados downstream não veja linhas duplicadas. Como sempre existe a necessidade de rodar o statement acima podemos dizer que esse pipeline é meio que eventualmente idempotente, mas a chamada da lambda em si não é. A palavra eventualmente é importante aqui, porque as linhas duplicadas continuam fisicamente na tabela raw, e qualquer coisa que consultar ela antes da deduplicação rodar ainda vai ver as duas cópias. Então a janela onde os dados estão errados é o intervalo entre o COPY e a próxima execução daquele model, o que significa que essa janela não é definida pela lambda e sim pelo agendamento em que a ferramenta de deduplicação está rodando. E qualquer coisa que decidir ler a tabela raw diretamente em vez do model não está coberta por isso de jeito nenhum.

Existe também uma segunda coisa que esse design não garante. Como a chave é um hash do payload, dois eventos que são realmente distintos mas que por acaso são idênticos são indistinguíveis de uma reentrega, então a deduplicação vai descartar um deles. Eu aceitei isso porque nesse pipeline uma duplicata idêntica é sempre um retry, mas isso é uma suposição sobre os dados e não uma propriedade do código, então se um dia isso deixar de ser verdade o statement vai descartar linhas reais silenciosamente.

Conclusão

Essa COPY lambda foi um exemplo para mim de que até uma função lambda simples responsável só por criar manifests e copiar dados pode levar a suposições erradas se você não olhar mais os detalhes, e também pode introduzir um tipo de erro que um olhar simples no happy path consegue esconder. Além disso, a mesma lambda pode cuidar da validação de dados malformados ou da deduplicação, mas isso vai contradizer o princípio da responsabilidade única (virando control e data plane ao mesmo tempo). A validação se encaixa melhor na função que escreve o arquivo, porque aquela já tem os bytes em memória e checar eles ali não custa nada a mais, enquanto fazer isso aqui significaria baixar os mesmos dados uma segunda vez só para olhar para eles. Então ter uma visão clara de cada passo da função e uma solução para cada cenário de falha pode aumentar as chances da implementação geral não virar um fracasso completo.