Use Data Quality Gates Locally in Python

import os

from dotenv import find_dotenv, load_dotenv

from ataccama_one import Client, OpenIdConnectAuth


def get_env_var(name: str) -> str:
    value = os.getenv(name)
    if value is None:
        raise ValueError(f"Environment variable {name} is required but not set.")
    return value


load_dotenv(find_dotenv(".env"))

client = Client(
    platform_version=get_env_var("ATACCAMA_PLATFORM_VERSION"),
    url=get_env_var("ATACCAMA_INSTANCE_URL"),
    auth=OpenIdConnectAuth(
        client_id=get_env_var("ATACCAMA_CLIENT_ID"),
        client_secret=get_env_var("ATACCAMA_CLIENT_SECRET"),
        keycloak_host=get_env_var("ATACCAMA_KEYCLOAK_HOST"),
        keycloak_realm=get_env_var("ATACCAMA_KEYCLOAK_REALM"),
    ),
)

firewall_id = get_env_var("YOUR_DQ_FIREWALL_ID")
client.export_firewall_definition(firewall_id, save_to="firewall_definition.zip")

# Adjust `firewall_filter` to limit which firewalls are fetched from Ataccama ONE
firewall_filter = None

for firewall in client.get_firewall_definitions(filter=firewall_filter):
    client.export_firewall_definition(
        firewall.firewall_id,
        save_to=f"{firewall.name}.zip",
    )

# Return the value to use as each record's identifier.
# Adapt this to your data: here, the first value in each record is used.
def get_record_id(record, index):
    return str(records[index][0])

results = firewall.validate_records(records, record_identifier=get_record_id)

DqFirewallResult(
    record_results=[
        DqFirewallRecordResult(
            record_id='893794994',
            overall_result='FAILED',
            score=1000,
            rule_instance_results=[
                DqFirewallRuleInstanceResult(
                    rule_instance_label='Not empty (String)',
                    overall_result='PASSED',
                    dimension_result='COMPLETE',
                    explanation='OTHER',
                ),
                DqFirewallRuleInstanceResult(
                    rule_instance_label='customername evaluation',
                    overall_result='FAILED',
                    dimension_result='INVALID',
                    explanation='INVALID_FORMAT',
                ),
                DqFirewallRuleInstanceResult(
                    rule_instance_label='accuracy prefix',
                    overall_result='FAILED',
                    dimension_result='INVALID',
                    explanation='WRONG_REGIONAL_PREFIX',
                ),
            ],
        ),
    ],
)

for record_result in results.record_results:
    print(record_result.record_id, record_result.overall_result, record_result.score)

    for rule_result in record_result.rule_instance_results:
        print(" ", rule_result.rule_instance_label, rule_result.overall_result)

import pandas as pd

flat_data = []


def mark_failed(input: str) -> str:
    return f"🔴 {input}" if "FAILED" in input else f"🟢 {input}"


for record_result in results.record_results:
    flat_data.append(
        {
            "record_id": record_result.record_id,
            "overall_result": mark_failed(record_result.overall_result),
            "score": record_result.score,
            **{
                f"{rule.rule_instance_label}": f"{mark_failed(rule.overall_result)}  "
                f"(dimension: {rule.dimension_result} - {rule.explanation})"
                for rule in record_result.rule_instance_results
            },
        }
    )

results_frame = pd.DataFrame(flat_data)
results_frame

combined = []
for source_record, record_result in zip(records, results.record_results):
    combined.append({
        "record": source_record,
        "record_id": record_result.record_id,
        "overall_result": record_result.overall_result,
        "score": record_result.score,
        "rule_instance_results": record_result.rule_instance_results,
    })

combined_frame = pd.concat([input_frame, results_frame], axis=1)

valid_records = []
quarantined_records = []

# Choose a score threshold that fits your quality requirements.
SCORE_THRESHOLD = 1000

for source_record, record_result in zip(records, results.record_results):
    if record_result.overall_result == "PASSED":
        valid_records.append(source_record)
    elif record_result.score < SCORE_THRESHOLD:
        failed_reasons = [
            (rule.rule_instance_label, rule.explanation)
            for rule in record_result.rule_instance_results
            if rule.overall_result == "FAILED"
        ]
        quarantined_records.append((source_record, record_result.score, failed_reasons))

# Pass valid_records on to the next layer.
# Write quarantined_records, with their explanations, to your quarantine target.

# Only the DQ results, without the source data.
results_only = [
    {
        "record_id": record_result.record_id,
        "overall_result": record_result.overall_result,
        "score": record_result.score,
    }
    for record_result in results.record_results
]

# Only the records that failed, from the combined list.
failed_only = [row for row in combined if row["overall_result"] == "FAILED"]

total = len(results.record_results)
passed = sum(1 for r in results.record_results if r.overall_result == "PASSED")
pass_rate = (passed / total) * 100 if total else 0

# Set the minimum acceptable pass rate for your pipeline.
PASS_RATE_THRESHOLD = 95
if pass_rate < PASS_RATE_THRESHOLD:
    raise RuntimeError(f"Data quality below threshold: {pass_rate:.1f}% passed")

# For monitoring, record pass_rate and the passed/total counts to your metrics store.