%pip install great_expectations

%pip install -U "typing_extensions>=4.12"

dbutils.library.restartPython()

import great_expectations as gx
from great_expectations.exceptions import DataContextError

# Cria o File Data Context na pasta atual
context = gx.get_context(mode="file", context_root_dir="/Workspace/Repos/arthurfr23@gmail.com/learn_databricks/src/great_expectations/gx")

data_source_name = "databricks.sql_saturday.pandas"
try:
    data_source = context.data_sources.get(data_source_name)
except Exception:
    data_source = context.data_sources.add_pandas(name=data_source_name)

GOLD_TABLE_CONFIGS = [
    {"asset_name": "sql_saturday.gold.dim_parlamentares", "batch_def_name": "whole_df", "expectation_suite_name": "dim_parlamentares_suite"},
    {"asset_name": "sql_saturday.gold.dim_calendario", "batch_def_name": "whole_df", "expectation_suite_name": "dim_calendario_suite"},
    {"asset_name": "sql_saturday.gold.dim_fornecedores", "batch_def_name": "whole_df", "expectation_suite_name": "dim_fornecedores_suite"},
    {"asset_name": "sql_saturday.gold.fato_reembolso", "batch_def_name": "whole_df", "expectation_suite_name": "fato_reembolso_suite"},
]

def build_expectation_suite(suite_name: str, asset_name: str):
    """Constrói a ExpectationSuite com as expectativas específicas de cada tabela gold."""
    suite = gx.ExpectationSuite(name=suite_name)

    if "dim_parlamentares" in asset_name:
        # dim_parlamentares: not null, unicidade Id_Parlamentar, CPF 11 dígitos
        for col in ["Id_Parlamentar", "Nome_Parlamentar", "cpf", "UF", "Partido"]:
            suite.add_expectation(gx.expectations.ExpectColumnValuesToNotBeNull(column=col))
        suite.add_expectation(gx.expectations.ExpectColumnValuesToBeUnique(column="Id_Parlamentar"))
        suite.add_expectation(
            gx.expectations.ExpectColumnValueLengthsToBeBetween(column="cpf", min_value=11, max_value=11, strict_min=True, strict_max=True)
        )

    elif "dim_calendario" in asset_name:
        # dim_calendario: not null, unicidade, domínio Mes/Dia/Ano
        suite.add_expectation(gx.expectations.ExpectColumnValuesToBeUnique(column="Id_Calendario")
        )

    elif "dim_fornecedores" in asset_name:
        # dim_fornecedores: not null, unicidade Id e CNPJ_CPF
        for col in ["Id_Fornecedor", "Fornecedor", "CNPJ_CPF"]:
            suite.add_expectation(gx.expectations.ExpectColumnValuesToNotBeNull(column=col))
        suite.add_expectation(gx.expectations.ExpectColumnValuesToBeUnique(column="Id_Fornecedor"))
        suite.add_expectation(gx.expectations.ExpectColumnValuesToBeUnique(column="CNPJ_CPF"))

    elif "fato_reembolso" in asset_name:
        # fato_reembolso: not null FKs e chave, unicidade Id_Reembolso, valores >= 0
        for col in ["Id_Reembolso", "Id_Parlamentar", "Id_Calendario", "Id_Fornecedor"]:
            suite.add_expectation(gx.expectations.ExpectColumnValuesToNotBeNull(column=col))
        suite.add_expectation(gx.expectations.ExpectColumnValuesToBeUnique(column="Id_Reembolso"))
        for col in ["Valor_Documento", "Valor_Glosa", "Valor_Liquido", "Valor_Ressarcimento"]:
            suite.add_expectation(
                gx.expectations.ExpectColumnValuesToBeBetween(column=col, min_value=0, max_value=None, strict_min=True)
            )

    return suite

# Para cada tabela: get/create asset, get/create batch definition, get/create suíte (evita erro ao rodar de novo)
for config in GOLD_TABLE_CONFIGS:
    asset_name = config["asset_name"]
    batch_def_name = config["batch_def_name"]
    expectation_suite_name = config["expectation_suite_name"]

    try:
        data_asset = data_source.get_asset(asset_name)
    except Exception:
        data_asset = data_source.add_dataframe_asset(name=asset_name)

    try:
        data_asset.get_batch_definition(batch_def_name)
    except Exception:
        data_asset.add_batch_definition_whole_dataframe(batch_def_name)

    try:
        context.suites.get(name=expectation_suite_name)
        print(f"Suite '{expectation_suite_name}' já existe, pulando.")
    except Exception:
        suite = build_expectation_suite(expectation_suite_name, asset_name)
        context.suites.add(suite)
        print(f"Suite '{expectation_suite_name}' registrada para asset '{asset_name}'.")

for config in GOLD_TABLE_CONFIGS:
    validation_definition_name = config["expectation_suite_name"].replace("_suite", "_validation")
    batch_definition = (
        data_source.get_asset(config["asset_name"]).get_batch_definition(config["batch_def_name"])
    )
    suite = context.suites.get(name=config["expectation_suite_name"])
    vd = gx.ValidationDefinition(
        data=batch_definition,
        suite=suite,
        name=validation_definition_name,
    )
    try:
        context.validation_definitions.add(vd)
        print(f"Validation Definition '{validation_definition_name}' criada.")
    except DataContextError as e:
        if "already exists" in str(e):
            print(f"Validation Definition '{validation_definition_name}' já existe, pulando.")
        else:
            raise

import great_expectations as gx
import great_expectations.expectations as gxe
from great_expectations.validator.validator import Validator
from great_expectations.checkpoint import UpdateDataDocsAction
from great_expectations.exceptions import DataContextError
from datetime import datetime, timezone
from pyspark.sql import Row

class GreatExpectationsExecute:

    def __init__ (self, 
                  data_source_name, 
                  suite_name, 
                  asset_name, 
                  batch_def_name, 
                  expectation_suite_name, 
                  table_name,
                  context_mode: str = "file",
                  context_root_dir: str = "/Workspace/Repos/arthurfr23@gmail.com/learn_databricks/src/great_expectations/gx"):
        
        self.data_source_name = data_source_name
        self.suite_name = suite_name
        self.asset_name = asset_name
        self.batch_def_name = batch_def_name
        self.expectation_suite_name = expectation_suite_name
        self.table_name = table_name
        self.context = gx.get_context(mode=context_mode, context_root_dir=context_root_dir)

    def data_source_name_check(self):
        try:
            self.data_source = self.context.data_sources.get(self.data_source_name)
        except Exception:
            print("Data Source não encontrado")
            raise

    def suite_name_check(self):
        try:
            self.suite = self.context.suites.get(self.suite_name)
        except Exception:
            print("Suite não encontrada")
            raise

    def asset_name_check(self):
        try:
            self.asset = self.data_source.get_asset(self.asset_name)
        except Exception:
            print("Asset name não encontrado")
            raise

    def batch_def_name_check(self):
        try:
            self.asset.get_batch_definition(self.batch_def_name)
        except Exception:
            print("Batch não definido")
            raise

    def expectation_suite_name_check(self):
        try:
            self.suite = self.context.suites.get(self.expectation_suite_name)
        except Exception:
            print("Expectation suite não encontrada")
            raise

    def table_name_check(self):
        try:
            self.table = spark.table(self.table_name)
        except Exception:
            print("Tabela não encontrada")
            raise

    def run(self, build_data_docs: bool = False, results_table: str = None):
        self.data_source_name_check()
        self.suite_name_check()
        self.asset_name_check()
        self.batch_def_name_check()
        self.expectation_suite_name_check()
        self.table_name_check()
        try:
            self._execute(results_table=results_table)
        finally:
            if build_data_docs:
                self._geracao_html()

    def _execute(self, results_table: str = None):
        validation_definition_name = self.expectation_suite_name.replace("_suite", "_validation")
        vd = self.context.validation_definitions.get(validation_definition_name)
        df_spark = spark.read.table(self.asset_name)
        df = df_spark.toPandas()
        result = vd.run(batch_parameters={"dataframe": df})
        success = result.get("success", False) if isinstance(result, dict) else getattr(result, "success", False)
#--------------------------------------------------------------------------------
        status = "OK" if success else "FALHOU"
        print(f"\n{'='*60}")

#--------------------------------------------------------------------------------
        print(f"Asset: {self.asset_name} → {status}")
        print(f"{'='*60}")

        results_list = (
            result.get("results", []) if isinstance(result, dict)
            else getattr(result, "results", [])
        )
        for r in results_list:
            r_success = r.get("success", True) if isinstance(r, dict) else getattr(r, "success", True)
            expectation_type = (
                r.get("expectation_config", {}).get("type", "")
                if isinstance(r, dict)
                else getattr(getattr(r, "expectation_config", None), "type", "")
            )
            kwargs = (
                r.get("expectation_config", {}).get("kwargs", {})
                if isinstance(r, dict)
                else getattr(getattr(r, "expectation_config", None), "kwargs", {})
            )
            result_detail = (
                r.get("result", {})
                if isinstance(r, dict)
                else getattr(r, "result", {})
            )
            icon = "✓" if r_success else "✗"
            print(f"  {icon} {expectation_type} | kwargs={kwargs} | result={result_detail}")

        if results_table:
            self._save_results(results_list, success, results_table)

        if not success:
            raise ValueError(
                f"Validação GE falhou para {self.asset_name}. "
                f"Ver resultado completo em result / Data Docs."
            )
        return status

    def _save_results(self, results_list: list, overall_success: bool, results_table: str):
        run_ts = datetime.now(timezone.utc)
        rows = []
        for r in results_list:
            r_success = r.get("success", True) if isinstance(r, dict) else getattr(r, "success", True)
            cfg = (
                r.get("expectation_config", {}) if isinstance(r, dict)
                else getattr(r, "expectation_config", {})
            )
            expectation_type = cfg.get("type", "") if isinstance(cfg, dict) else getattr(cfg, "type", "")
            kwargs = cfg.get("kwargs", {}) if isinstance(cfg, dict) else getattr(cfg, "kwargs", {})
            result_detail = (
                r.get("result", {}) if isinstance(r, dict) else getattr(r, "result", {})
            )
            rows.append(Row(
                run_timestamp=run_ts,
                asset_name=self.asset_name,
                expectation_suite=self.expectation_suite_name,
                overall_success=overall_success,
                expectation_type=expectation_type,
                column=kwargs.get("column", None) if isinstance(kwargs, dict) else None,
                success=bool(r_success),
                unexpected_count=int(result_detail.get("unexpected_count", 0)) if isinstance(result_detail, dict) else None,
                unexpected_percent=float(result_detail.get("unexpected_percent", 0.0)) if isinstance(result_detail, dict) else None,
                kwargs_str=str(kwargs),
                result_str=str(result_detail),
            ))

        df_results = spark.createDataFrame(rows)
        df_results.write.format("delta").mode("append").saveAsTable(results_table)
        print(f"\nResultados salvos em: {results_table}")

    def _geracao_html(self):
        docs = self.context.build_data_docs()
        if docs:
            for site_name, path in docs.items():
                print(f"Data Docs ({site_name}): {path}")
                print(f"Abra no navegador: file://{path}")
        else:
            print("Nenhum site de Data Docs configurado. Verifique great_expectations.yml.")

runner = GreatExpectationsExecute(
    data_source_name="databricks.sql_saturday.pandas",
    suite_name="dim_parlamentares_suite",
    asset_name="sql_saturday.gold.dim_parlamentares",
    batch_def_name="whole_df",
    expectation_suite_name="dim_parlamentares_suite",
    table_name="sql_saturday.gold.dim_parlamentares"
)

runner.run(build_data_docs=True, results_table="sql_saturday.monitoring.ge_results")

%sql
SELECT *
FROM sql_saturday.monitoring.ge_results

{
  "expectations": [
    {
      "id": "8c635dcf-a459-4865-9a3d-7944153a3266",
      "kwargs": {
        "column": "Id_Fornecedor"
      },
      "meta": {},
      "severity": "critical",
      "type": "expect_column_values_to_not_be_null"
    },
    {
      "id": "43771305-4b0d-4061-9139-61f841eef3a6",
      "kwargs": {
        "column": "Fornecedor"
      },
      "meta": {},
      "severity": "critical",
      "type": "expect_column_values_to_not_be_null"
    },
    {
      "id": "248948ee-9ee8-490d-8f0a-6dae1e78b8fc",
      "kwargs": {
        "column": "CNPJ_CPF"
      },
      "meta": {},
      "severity": "critical",
      "type": "expect_column_values_to_not_be_null"
    },
    {
      "id": "79c9b68c-618f-4678-bd52-3fc7dbb5553f",
      "kwargs": {
        "column": "Id_Fornecedor"
      },
      "meta": {},
      "severity": "critical",
      "type": "expect_column_values_to_be_unique"
    },
    {
      "id": "0faa514d-d1d8-4e64-9b61-6bb8eb983523",
      "kwargs": {
        "column": "CNPJ_CPF"
      },
      "meta": {},
      "severity": "critical",
      "type": "expect_column_values_to_be_unique"
    }
  ],
  "id": "bf7e47cf-ae1d-4896-be49-85fbec12c54f",
  "meta": {
    "great_expectations_version": "1.16.0"
  },
  "name": "dim_fornecedores_suite",
  "notes": null
}

from pyspark.sql import functions as F
from datetime import date

def check_nulls(table_name, col, criticality):

    test_date = date.today()

    test_type = 'check null'

    df = spark.read.table(f'sql_saturday.gold.{table_name}')
        
    nulos = df.select(F.count(F.when(F.col(col).isNull(), 1)).alias('num_nulos'))

    num_nulos = nulos.collect()[0][0]

    spark.sql(f"""
        INSERT INTO sql_saturday.custom_tests.tests_results (test_date, table_name, column_name, test_type, test_result, criticality)
        VALUES ('{test_date}', '{table_name}', '{col}', '{test_type}', {num_nulos}, '{criticality}')
    """)

def check_unique(table_name, col, criticality):

    test_date = date.today()

    test_type = 'check unique'

    df = spark.read.table(f'sql_saturday.gold.{table_name}')

    df_duplicate = df.groupBy(col).count().filter("count > 1")
        
    total_duplicates = df_duplicate.agg(F.sum(col)).collect()[0][0]

    if total_duplicates is None:
        total_duplicates = 0

    spark.sql(f"""
        INSERT INTO sql_saturday.custom_tests.tests_results (test_date, table_name, column_name, test_type, test_result, criticality)
        VALUES ('{test_date}', '{table_name}', '{col}', '{test_type}', {total_duplicates}, '{criticality}')
    """)

def check_lenght(table_name, col, criticality):

    test_date = date.today()

    test_type = 'check lenght'

    df = spark.read.table(f'sql_saturday.gold.{table_name}')

    df2 = df.withColumn('length_rows', F.length(df[col])).select('length_rows').filter('length_rows != 11').agg(F.sum('length_rows'))

    total_lenght = df2.collect()[0][0]

    if total_lenght is None:
        total_lenght = 0

    spark.sql(f"""
        INSERT INTO sql_saturday.custom_tests.tests_results (test_date, table_name, column_name, test_type, test_result, criticality)
        VALUES ('{test_date}', '{table_name}', '{col}', '{test_type}', {total_lenght}, '{criticality}')
    """)

def check_value_positive(table_name, col, criticality):

    test_date = date.today()

    test_type = 'check value positive'

    df = spark.read.table(f'sql_saturday.gold.{table_name}')

    total = df.select(F.count(F.when(F.col(col) <= 0, 1)).alias('non_positive')).collect()[0][0]

    if total is None:
        total = 0

    spark.sql(f"""
        INSERT INTO sql_saturday.custom_tests.tests_results (test_date, table_name, column_name, test_type, test_result, criticality)
        VALUES ('{test_date}', '{table_name}', '{col}', '{test_type}', {total}, '{criticality}')
    """)

check_unique('dim_calendario', 'Id_Calendario', 'critical')

%sql
SELECT *
FROM sql_saturday.custom_tests.tests_results

%pip install databricks-labs-dqx

%pip install 'databricks-labs-dqx[llm]'

%pip install 'databricks-labs-dqx[datacontract]'

dbutils.library.restartPython()

from databricks.labs.dqx import check_funcs
from databricks.labs.dqx.rule import DQDatasetRule
from databricks.labs.dqx.engine import DQEngine
from databricks.sdk import WorkspaceClient

ws = WorkspaceClient()
dq_engine = DQEngine(ws)

# Criação da expectation
checks = [
  DQDatasetRule(
    name="id_calendar_month_is_unique",
    criticality="warn",
    check_func=check_funcs.is_unique,
    columns=["Id_Calendario, Mes"]
    )
  ]

# Execução do check
input_df = spark.read.table('sql_saturday.gold.dim_calendario')
result_df = dq_engine.apply_checks(input_df, checks)
display(result_df)

from databricks.labs.dqx.rule import DQRowRule
from databricks.sdk import WorkspaceClient
from databricks.labs.dqx.engine import DQEngine
from databricks.labs.dqx.profiler.generator import DQGenerator
from databricks.labs.dqx.profiler.profiler import DQProfiler
from databricks.labs.dqx.config import WorkspaceFileChecksStorageConfig
import logging
import json
import yaml

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

ws = WorkspaceClient()
dq_engine = DQEngine(ws)
generator = DQGenerator(ws)
profiler = DQProfiler(ws)

input_df = spark.read.table("sql_saturday.gold.dim_calendario")

# Databricks faz a checagem e cria as regras, usando uma llm para detectar a PK
ws = WorkspaceClient()
profiler = DQProfiler(ws)
summary_stats, profiles = profiler.profile(input_df, options={'llm_primary_key_detection': False})

generator = DQGenerator(ws)
checks = generator.generate_dq_rules(profiles)

# Salva o profile e o check no workspace
dq_engine.save_checks(checks, config=WorkspaceFileChecksStorageConfig(location="/Workspace/Repos/arthurfr23@gmail.com/learn_databricks/src/dqx_tests/dqx_apresentacao/dim_calendar_profile"))

checks = generator.generate_dq_rules(profiles)

# Carrega lista de checks a partir do arquivo YAML
with open("/Workspace/Repos/arthurfr23@gmail.com/learn_databricks/src/dqx_tests/dqx_apresentacao/dim_calendar_profile", "r") as f:
    checks = yaml.safe_load(f)

# Aplica os checks carregados aos dados
result_df = dq_engine.apply_checks_by_metadata(input_df, checks)
display(result_df)

# Carrega lista de checks a partir do arquivo YAML
with open("/Workspace/Repos/arthurfr23@gmail.com/learn_databricks/src/dqx_tests/dqx_apresentacao/dim_calendar_profile", "r") as f:
    checks = yaml.safe_load(f)

# Aplica os checks carregados aos dados
result_df = dq_engine.apply_checks_by_metadata(input_df, checks)
display(result_df)

from databricks.labs.dqx.profiler.generator import DQGenerator
from databricks.labs.dqx.config import InputConfig
from databricks.sdk import WorkspaceClient

ws = WorkspaceClient()
generator = DQGenerator(workspace_client=ws, spark=spark)

user_input = """
Id_calendario needs to be unique
Ano need to be between 2000 and 2023
Id_calendario needs to be between 1 and 365
"""

checks = generator.generate_dq_rules_ai_assisted(
  user_input=user_input,
  input_config=InputConfig(location="sql_saturday.gold.dim_calendario")
)

print(checks)

from databricks.labs.dqx.profiler.generator import DQGenerator
from databricks.sdk import WorkspaceClient
from databricks.labs.dqx.engine import DQEngine

ws = WorkspaceClient()
engine = DQEngine(workspace_client=ws)
generator = DQGenerator(workspace_client=ws, spark=spark)

rules = generator.generate_rules_from_contract(
    contract_file="/Workspace/Repos/arthurfr23@gmail.com/learn_databricks/src/dqx_tests/dqx_apresentacao/dim_calendar_data_contract.yml"
)

input_df = spark.read.table("sql_saturday.gold.dim_calendario")

result_df = engine.apply_checks_by_metadata(
    input_df,
    rules
)

display(result_df)

from databricks.labs.dqx.engine import DQEngine
from databricks.labs.dqx.metrics_observer import DQMetricsObserver
from databricks.sdk import WorkspaceClient
from databricks.labs.dqx.profiler.generator import DQGenerator
from databricks.labs.dqx import check_funcs
from databricks.labs.dqx.rule import DQDatasetRule
from databricks.labs.dqx.config import InputConfig, OutputConfig

ws = WorkspaceClient()

observer = DQMetricsObserver(name="dq_metrics")

engine = DQEngine(WorkspaceClient(), observer=observer)

generator = DQGenerator(workspace_client=ws, spark=spark)

# Lê o df e o Data Contract

df = spark.read.table('sql_saturday.gold.dim_calendario')

checks = generator.generate_rules_from_contract(
    contract_file="/Workspace/Repos/arthurfr23@gmail.com/learn_databricks/src/dqx_tests/dqx_apresentacao/dim_calendar_data_contract.yml"
)

# Gera as métricas
checked_df, observation = engine.apply_checks_by_metadata(df, checks)

row_count = checked_df.count()


metrics = observation.get
print(f"Input row count: {metrics['input_row_count']}")
print(f"Error row count: {metrics['error_row_count']}")
print(f"Warning row count: {metrics['warning_row_count']}")
print(f"Valid row count: {metrics['valid_row_count']}")

# Cria tabelas de validação, métricas e quarentenas para serem analisadas posteriormente
input_config = InputConfig("sql_saturday.gold.dim_calendario")

output_config = OutputConfig("sql_saturday.tests.valid_data")
quarantine_config = OutputConfig("sql_saturday.tests.quarantine_data") 
metrics_config = OutputConfig("sql_saturday.tests.metrics_data") 

engine.apply_checks_by_metadata_and_save_in_table(
    checks=checks,
    input_config=input_config,
    output_config=output_config,
    quarantine_config=quarantine_config,
    metrics_config=metrics_config
)

%sql
SELECT *
FROM sql_saturday.tests.valid_data

%sql
SELECT *
FROM sql_saturday.tests.metrics_data

- check:
    arguments:
      column: Id_Calendario
    function: is_not_null
  criticality: error
  name: Id_Calendario_is_null
- check:
    arguments:
      column: Id_Calendario
      max_limit: 9364
      min_limit: 1
    function: is_in_range
  criticality: error
  name: Id_Calendario_isnt_in_range
- check:
    arguments:
      column: Ano
    function: is_not_null
  criticality: error
  name: Ano_is_null
- check:
    arguments:
      column: Ano
      max_limit: 2035
      min_limit: 2010
    function: is_in_range
  criticality: error
  name: Ano_isnt_in_range
- check:
    arguments:
      column: Mes
    function: is_not_null
  criticality: error
  name: Mes_is_null
- check:
    arguments:
      column: Mes
      max_limit: 12
      min_limit: 1
    function: is_in_range
  criticality: error
  name: Mes_isnt_in_range
- check:
    arguments:
      column: Dia
    function: is_not_null
  criticality: error
  name: Dia_is_null
- check:
    arguments:
      column: Dia
      max_limit: 31
      min_limit: 1
    function: is_in_range
  criticality: error
  name: Dia_isnt_in_range
- check:
    arguments:
      column: Data_Completa
    function: is_not_null
  criticality: error
  name: Data_Completa_is_null
- check:
    arguments:
      column: Data_Completa
      max_limit: '2035-08-21'
      min_limit: '2010-01-01'
    function: is_in_range
  criticality: error
  name: Data_Completa_isnt_in_range

kind: DataContract
apiVersion: v3.0.2
id: urn:datacontract:sql_saturday:gold_tables
name: Regras das Tabelas Gold
version: 2.0.0
status: active
domain: SQL Saturday
dataProduct: sql_saturday_gold_rules

schema:
  - name: dim_calendario
    physicalType: table
    properties:
      - name: Id_Calendario
        logicalType: string
        unique: true
      - name: Ano
        logicalType: string
        unique: true