Merge branch 'main' into staging/fix-fn-is-up-to-date-br_cnj

aspeddro · web-flow · commit 78a9fc4c7869 · 2024-04-02T10:00:25.000-03:00
diff --git a/pipelines/datasets/br_bcb_agencia/flows.py b/pipelines/datasets/br_bcb_agencia/flows.py
@@ -103,6 +103,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/datasets/br_bcb_taxa_cambio/flows.py b/pipelines/datasets/br_bcb_taxa_cambio/flows.py
@@ -77,6 +77,7 @@
                 "mode": materialization_mode,
                 "dbt_alias": dbt_alias,
                 "dbt_command": "run/test",
+                "disable_elementary": False,
             },
             labels=current_flow_labels,
             run_name=f"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/datasets/br_bd_metadados/flows.py b/pipelines/datasets/br_bd_metadados/flows.py
@@ -71,6 +71,7 @@
                 "mode": materialization_mode,
                 "dbt_alias": dbt_alias,
                 "dbt_command": "run/test",
+                "disable_elementary": False,
             },
             labels=current_flow_labels,
             run_name=f"Materialize {dataset_id}.{table_id}",
@@ -140,6 +141,7 @@
                 "mode": materialization_mode,
                 "dbt_alias": dbt_alias,
                 "dbt_command": "run/test",
+                "disable_elementary": False,
             },
             labels=current_flow_labels,
             run_name=f"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/datasets/br_cgu_beneficios_cidadao/flows.py b/pipelines/datasets/br_cgu_beneficios_cidadao/flows.py
@@ -121,6 +121,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -148,7 +149,7 @@
                         "month": "mes_competencia",
                     },
                     date_format="%Y-%m",
-                    coverage_type="all_bdpro",
+                    coverage_type="part_bdpro",
                     time_delta={"months": 6},
                     prefect_mode=materialization_mode,
                     bq_project="basedosdados",
@@ -251,6 +252,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -375,6 +377,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/datasets/br_cnj_improbidade_administrativa/flows.py b/pipelines/datasets/br_cnj_improbidade_administrativa/flows.py
@@ -80,6 +80,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=r"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/datasets/br_cvm_fi/flows.py b/pipelines/datasets/br_cvm_fi/flows.py
@@ -602,6 +602,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/datasets/br_denatran_frota/flows.py b/pipelines/datasets/br_denatran_frota/flows.py
@@ -221,6 +221,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/datasets/br_me_cnpj/flows.py b/pipelines/datasets/br_me_cnpj/flows.py
@@ -88,6 +88,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -180,6 +181,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -274,6 +276,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/datasets/br_ons_avaliacao_operacao/flows.py b/pipelines/datasets/br_ons_avaliacao_operacao/flows.py
@@ -96,6 +96,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -193,6 +194,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -293,6 +295,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -394,6 +397,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -496,6 +500,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -596,6 +601,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/datasets/br_ons_estimativa_custos/flows.py b/pipelines/datasets/br_ons_estimativa_custos/flows.py
@@ -97,6 +97,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -199,6 +200,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -302,6 +304,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -405,6 +408,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
@@ -508,6 +512,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/datasets/mundo_transfermarkt_competicoes_internacionais/flows.py b/pipelines/datasets/mundo_transfermarkt_competicoes_internacionais/flows.py
@@ -96,6 +96,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=r"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/utils/crawler_camara_dados_abertos/flows.py b/pipelines/utils/crawler_camara_dados_abertos/flows.py
@@ -81,6 +81,7 @@
                     "mode": materialization_mode,
                     "dbt_alias": dbt_alias,
                     "dbt_command": "run/test",
+                    "disable_elementary": False,
                 },
                 labels=current_flow_labels,
                 run_name=f"Materialize {dataset_id}.{table_id}",
diff --git a/pipelines/utils/metadata/flows.py b/pipelines/utils/metadata/flows.py
@@ -3,17 +3,17 @@
 Flows for temporal_coverage_updater
 """
 
+
+
 from prefect import Parameter
 from prefect.run_configs import KubernetesRun
 from prefect.storage import GCS
 
 from pipelines.constants import constants
 
-# from pipelines.datasets.temporal_coverage_updater.schedules import every_two_weeks
 from pipelines.utils.decorators import Flow
-from pipelines.utils.metadata.tasks import update_django_metadata
-
-# from pipelines.utils.utils import log
+from pipelines.utils.metadata.tasks import update_django_metadata, create_update_quality_checks, query_tests_results
+from pipelines.utils.metadata.schedules import every_day_quality_checks
 
 with Flow(
     name="update_temporal_coverage_teste",
@@ -43,3 +43,22 @@
     image=constants.DOCKER_IMAGE.value
 )
 # flow.schedule = every_two_weeks
+
+
+with Flow(
+    name="create_update_quality_checks",
+    code_owners=[
+        "arthurfg",
+    ],
+) as quality_checks_updater:
+
+    tests_results = query_tests_results()
+    results = create_update_quality_checks(tests_results = tests_results, upstream_tasks=[tests_results])
+
+
+quality_checks_updater.storage = GCS(constants.GCS_FLOWS_BUCKET.value)
+quality_checks_updater.run_config = KubernetesRun(
+    image=constants.DOCKER_IMAGE.value
+)
+quality_checks_updater.schedule = every_day_quality_checks
+
diff --git a/pipelines/utils/metadata/schedules.py b/pipelines/utils/metadata/schedules.py
@@ -0,0 +1,23 @@
+# -*- coding: utf-8 -*-
+"""
+Schedules for metadata flows
+"""
+
+
+from datetime import datetime
+from prefect.schedules import Schedule
+from prefect.schedules.clocks import CronClock
+from pipelines.constants import constants
+
+
+every_day_quality_checks = Schedule(
+    clocks=[
+        CronClock(
+            cron="0 21 * * *",  # At 23:00 on every day-of-week from Monday through Friday.
+            start_date=datetime(2024, 3, 20, 12, 26),
+            labels=[
+                constants.BASEDOSDADOS_PROD_AGENT_LABEL.value,
+            ],
+        ),
+    ],
+)
diff --git a/pipelines/utils/metadata/tasks.py b/pipelines/utils/metadata/tasks.py
@@ -2,8 +2,9 @@
 """
 Tasks for metadata
 """
-
+import asyncio
 from datetime import datetime
+import basedosdados as bd
 
 import pandas as pd
 from prefect import task
@@ -17,11 +18,12 @@
     get_coverage_parameters,
     get_credentials_utils,
     get_ids,
+    get_id,
     get_table_status,
     update_row_access_policy,
 )
 from pipelines.utils.utils import log
-
+from pipelines.utils.metadata.utils_async import create_update_quality_checks_async
 
 @task
 def get_today_date():
@@ -200,3 +202,69 @@ def task_get_api_most_recent_date(dataset_id, table_id, date_format):
     return get_api_most_recent_date(
         dataset_id=dataset_id, table_id=table_id, date_format=date_format
     )
+
+#####             #####
+## Quality check's   ##
+#####             #####
+
+
+@task
+def query_tests_results() -> pd.DataFrame:
+    """
+    Task to query recent test results from basedosdados.
+
+    Returns:
+    - pd.DataFrame: A pandas DataFrame containing recent test results with the following columns:
+        - 'name': The name of the test.
+        - 'description': The description of the test result (typically the column name).
+        - 'status': Indicates whether the test passed or failed.
+        - 'dataset_id': The ID of the dataset containing the tested table.
+        - 'table_id': The ID of the tested table.
+    """
+
+    billing_project_id = get_billing_project_id(mode="prod")
+    query_bd = f"""
+    with tests_order as (
+    select
+        test_short_name as name,
+        column_name as description,
+        status,
+        schema_name as dataset_id,
+        table_name as table_id,
+        row_number() over (partition by test_short_name, schema_name, table_name order by created_at desc) as position,
+        created_at
+    from
+        `basedosdados.elementary.elementary_test_results`
+    where
+        date(created_at) >= date_sub(current_date(), interval 7 DAY))
+    select
+        name,
+        description,
+        status,
+        dataset_id,
+        table_id
+    from tests_order
+    where position = 1
+    """
+
+    t = bd.read_sql(
+        query=query_bd,
+        billing_project_id=billing_project_id,
+        from_file=True,
+    )
+
+    return t
+
+@task
+def create_update_quality_checks(tests_results: pd.DataFrame) -> None:
+    """
+    Task to create or update multiple quality checks based on test results asynchronously.
+
+    Parameters:
+    - tests_results (pd.DataFrame): A pandas DataFrame containing test results.
+
+    Returns:
+    - None
+
+    """
+    asyncio.run(create_update_quality_checks_async(tests_results=tests_results))
diff --git a/pipelines/utils/metadata/utils.py b/pipelines/utils/metadata/utils.py
diff --git a/pipelines/utils/metadata/utils_async.py b/pipelines/utils/metadata/utils_async.py
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
