From e0c9784242eb5224775de8a80180c27dcd58c190 Mon Sep 17 00:00:00 2001
From: Pedro Marques <pedrovitor.marques@gmail.com>
Date: Fri, 2 Aug 2024 15:13:23 -0300
Subject: [PATCH] Injecting Paylod CNES in Datalake Models

---
 app/datalake/models.py      | 4 ++++
 app/routers/entities_raw.py | 6 ++++++
 2 files changed, 10 insertions(+)

diff --git a/app/datalake/models.py b/app/datalake/models.py
index a97b6aa..3cbc053 100644
--- a/app/datalake/models.py
+++ b/app/datalake/models.py
@@ -44,6 +44,7 @@ class SMSRioPaciente(BaseModel):
     data__timestamp: Optional[str]
     data__cns_provisorio: Optional[str]
     data__telefones: Optional[str]
+    payload_cnes: str
 
     class Config:
         dataset_id = "brutos_plataforma_smsrio"
@@ -136,6 +137,7 @@ class VitacarePaciente(BaseModel):
     data__dataAtualizacaoVinculoEquipe: Optional[str]
     data__familiaBeneficiariaAuxilioBrasil: Optional[str]
     data__criancaMatriculadaCrechePreEscola: Optional[str]
+    payload_cnes: str
 
     class Config:
         dataset_id = "brutos_prontuario_vitacare"
@@ -189,6 +191,7 @@ class VitacarePacienteHistorico(BaseModel):
     data__NOME_DA_PESSOA_CADASTRADA: Optional[str]
     data__N_CNS_DA_PESSOA_CADASTRADA: Optional[str]
     data__NOME_DA_MAE_PESSOA_CADASTRADA: Optional[str]
+    payload_cnes: str
 
     class Config:
         dataset_id = "brutos_prontuario_vitacare"
@@ -232,6 +235,7 @@ class VitacareAtendimento(BaseModel):
     data__alergias_anamnese: str
     data__indicadores: str
     data__encaminhamentos: str
+    payload_cnes: str
 
     class Config:
         dataset_id = "brutos_prontuario_vitacare"
diff --git a/app/routers/entities_raw.py b/app/routers/entities_raw.py
index d3ee1c0..7935d77 100644
--- a/app/routers/entities_raw.py
+++ b/app/routers/entities_raw.py
@@ -107,11 +107,17 @@ async def create_raw_data(
     # ====================
     uploaded_to_datalake = False
 
+    # Inject CNES in records
+    for record in records:
+        record["payload_cnes"] = data_source.cnes
+
+    # Get Formatter
     formatter = get_formatter(
         system=data_source.system.value,
         entity=entity_name
     )
 
+    # Format and Upload
     try:
         if upload_to_datalake and formatter:
             uploader = DatalakeUploader(