How do you run TestGen from Airflow?

TestGen runs from Airflow: an Airflow task calls the TestGen REST API, and TestGen runs the tests as SQL inside your database. No connector exists, and none is needed. Airflow never sees a row. TestGen tests any of the 14 databases it supports, plus files.

This page is the how-to. To choose a tool in the first place, read the Airflow tool comparison. To decide what a failed check should do to the pipeline, read stop, quarantine, retry, or publish.

Get a token

The API authentication page offers two kinds. An OAuth 2.1 client works in open source: register once, store the refresh token as an Airflow Variable or secret, and trade it for a one-hour access token on every run. The refresh token lasts 30 days, and getting a new one means signing in through a browser again, so put the renewal on a calendar. TestGen Enterprise adds personal access tokens, a single static token that suits a quick setup. Every request acts with the role of the user who owns the token.

The base URL is the TestGen API server (http://<host>:8530 on a Docker or pip install), not the UI.

The task

The Run Profiling and Tests page gives the pattern: submit, poll, fetch. Submission returns at once with a job ID. Three statuses are terminal: completed, error, and canceled. Give the poll a timeout.

import time, requests
from airflow.sdk import task, Variable   # Airflow 3; on Airflow 2: airflow.decorators, airflow.models
from airflow.exceptions import AirflowFailException

def access_token(base):
    r = requests.post(f"{base}/oauth/token",
        auth=(Variable.get("tg_client_id"), Variable.get("tg_client_secret")),
        data={"grant_type": "refresh_token",
              "refresh_token": Variable.get("tg_refresh_token")})
    r.raise_for_status()
    return r.json()["access_token"]

@task
def testgen_gate(test_suite_id: str, **context):
    TG = Variable.get("testgen_api_url")    # e.g. http://testgen-host:8530
    h = {"Authorization": f"Bearer {access_token(TG)}"}
    job = requests.post(f"{TG}/api/v1/test-suites/{test_suite_id}/test-runs",
        headers=h,
        json={"metadata": {"dag_id": context["dag"].dag_id,
                           "run_id": context["run_id"]}}).json()
    deadline = time.time() + 3600
    while True:
        status = requests.get(f"{TG}/api/v1/jobs/{job['id']}", headers=h).json()["status"]
        if status in ("completed", "error", "canceled"):
            break
        if time.time() > deadline:
            raise TimeoutError(f"TestGen job {job['id']} still {status}")
        time.sleep(30)
    if status != "completed":
        raise RuntimeError(f"TestGen job {job['id']} ended {status}")  # retries apply
    counts = requests.get(f"{TG}/api/v1/test-runs/{job['id']}",
                          headers=h).json()["result"]["result_counts"]
    if counts["failed"] + counts["error"] > 0:
        raise AirflowFailException(f"TestGen: {counts['failed']} failed, {counts['error']} errored")
    return counts

Put testgen_gate between the load and whatever reads from it: load >> testgen_gate("<suite id>") >> publish. When it raises, publish never runs and the dashboards keep yesterday’s numbers.

Blocking and warning results

A test run summary counts passed, failed, warning, error, log, and dismissed. Block on failed and on error, which means a test couldn’t run. Let warning through and send it to a person. That split is the tripwire pattern DataKitchen built in Databricks: a gate task that exits non-zero, so the orchestrator’s own dependency graph stops the next step.

Run profiling on its own schedule, not on every load. The scheduling docs recommend biweekly to monthly.

Where the results land

The run shows up on TestGen’s Test Results page with your metadata under Run details. The same summary comes back from the API. Notification rules send failures to email, Slack, Microsoft Teams, or any HTTPS webhook, including a Jira automation rule.

For one view across the pipeline, DataOps Observability has a documented Airflow agent. It polls the Airflow REST API with a service account, every 5 seconds by default. Connect TestGen to Observability and test results arrive as Test Outcomes events on a dataset component, beside the DAG runs. DataKitchen’s post on the four points in a medallion architecture where testing matters explains why you watch the tools as well as the tables: an Airflow job that hangs for three hours fails no data test.

The same calls work from a CI runner; see running data tests in CI/CD.

Documentation: TestGen API, API authentication, OAuth clients, Run profiling and tests, Scheduling, Notifications, Connect to Observability, Observability Airflow agent.