Monthly report from a survey CSV¶
A published CSV, a stakeholder who wants the same question answered from it every month, and nobody who wants to hand-write the SQL or learn from the report that a column was renamed. This Dag downloads the file, checks its schema against a reference, has the model write the SQL, runs it with Apache DataFusion, and emails the rows. Airflow supplies the monthly schedule and a record of every schema change.
The example uses the Airflow community survey CSV, which is public and needs no credentials.
What this demonstrates¶
Natural language to SQL: LLMSQLQueryOperator –
LLMSQLQueryOperatorturns a question into SQL against a described schema, without executing it.Detect schema drift: LLMSchemaCompareOperator –
LLMSchemaCompareOperatorrecords how the downloaded file differs from a reference before any SQL is generated. It reports; it does not block (see Adapting it).AnalyticsOperatorfrom thecommon.sqlprovider – runs the generated SQL over a local file with DataFusion, no database needed.HttpOperatorandSmtpHookfrom thehttpandsmtpproviders do the download and delivery.
Run it¶
Install the provider with the SQL extra, plus the HTTP and SMTP providers:
pip install "apache-airflow-providers-common-ai[openai,sql]" \ "apache-airflow-providers-common-sql[datafusion]" \ apache-airflow-providers-http apache-airflow-providers-smtp
Create an HTTP connection named
airflow_websitewith hosthttps://airflow.apache.organd no auth.Optionally set
SMTP_CONN_IDandNOTIFY_EMAILin the environment. Without them the result goes to the task log.Trigger the Dag:
airflow dags test example_llm_survey_scheduled
The check_schema XCom holds the comparison, generate_sql holds the query the model
wrote, and send_result logs or mails the rows.
The Dag¶
@dag(schedule="@monthly", start_date=datetime.datetime(2025, 1, 1), catchup=False, tags=["example"])
def example_llm_survey_scheduled():
"""
Download, validate, query, and report on the survey CSV on a schedule.
Task graph::
download_survey (HttpOperator)
→ prepare_csv (@task)
→ check_schema (LLMSchemaCompareOperator)
→ generate_sql (LLMSQLQueryOperator)
→ run_query (AnalyticsOperator)
→ extract_data (@task)
→ send_result (@task)
No human review steps -- suitable for recurring reporting or dashboards.
Change ``schedule`` to any cron expression or Airflow timetable to adjust
the run frequency.
Prerequisites:
- HTTP connection ``airflow_website`` pointing at ``https://airflow.apache.org``.
- Set ``SMTP_CONN_ID`` and ``NOTIFY_EMAIL`` environment variables to enable
email delivery of results; otherwise results are logged to the task log.
"""
# ------------------------------------------------------------------
# Step 1: Download the survey CSV from the Airflow website.
# ------------------------------------------------------------------
download_survey = HttpOperator(
task_id="download_survey",
http_conn_id=AIRFLOW_WEBSITE_CONN_ID,
endpoint=SURVEY_CSV_ENDPOINT,
method="GET",
response_filter=lambda r: r.text,
log_response=False,
)
# ------------------------------------------------------------------
# Step 2: Write the downloaded CSV to disk and generate a reference
# schema file for the schema comparison step.
# ------------------------------------------------------------------
@task
def prepare_csv(csv_text: str) -> None:
os.makedirs(os.path.dirname(SURVEY_CSV_PATH), exist_ok=True)
with open(SURVEY_CSV_PATH, "w", encoding="utf-8") as f:
f.write(csv_text)
# Write a single-row reference CSV from the schema context so
# LLMSchemaCompareOperator has a structured baseline to compare against.
os.makedirs(os.path.dirname(REFERENCE_CSV_PATH), exist_ok=True)
columns = [line.split('"')[1] for line in SURVEY_SCHEMA.strip().splitlines() if '"' in line]
with open(REFERENCE_CSV_PATH, "w", newline="", encoding="utf-8") as ref:
csv_mod.writer(ref).writerow(columns)
csv_ready = prepare_csv(download_survey.output)
# ------------------------------------------------------------------
# Step 3: Compare the downloaded CSV schema against the reference.
# Reports what changed; add a @task.branch on ``compatible`` to block on drift.
# ------------------------------------------------------------------
check_schema = LLMSchemaCompareOperator(
task_id="check_schema",
prompt="""\
Compare the survey CSV schema against the reference schema.
Flag any missing or renamed columns that would break the downstream SQL queries.""",
llm_conn_id=LLM_CONN_ID,
data_sources=[survey_datasource, reference_datasource],
context_strategy="basic",
)
csv_ready >> check_schema
# ------------------------------------------------------------------
# Step 4: SQL generation -- LLM translates the fixed question.
# ------------------------------------------------------------------
generate_sql = LLMSQLQueryOperator(
task_id="generate_sql",
prompt=SCHEDULED_PROMPT,
llm_conn_id=LLM_CONN_ID,
datasource_config=survey_datasource,
schema_context=SURVEY_SCHEMA,
)
check_schema >> generate_sql
# ------------------------------------------------------------------
# Step 5: SQL execution via Apache DataFusion.
# ------------------------------------------------------------------
run_query = AnalyticsOperator(
task_id="run_query",
datasource_configs=[survey_datasource],
queries=["{{ ti.xcom_pull(task_ids='generate_sql') }}"],
result_output_format="json",
)
# ------------------------------------------------------------------
# Step 6: Extract data rows from the JSON result.
# AnalyticsOperator returns [{"query": "...", "data": [...]}, ...]
# ------------------------------------------------------------------
@task
def extract_data(raw: str) -> str:
results = json.loads(raw)
data = [row for item in results for row in item["data"]]
return json.dumps(data, indent=2)
result_data = extract_data(run_query.output)
# ------------------------------------------------------------------
# Step 7: Send result via email if SMTP is configured, otherwise log.
# Set the SMTP_CONN_ID and NOTIFY_EMAIL environment variables to enable
# email delivery.
# ------------------------------------------------------------------
@task
def send_result(data: str) -> None:
if SMTP_CONN_ID and NOTIFY_EMAIL:
from airflow.providers.smtp.hooks.smtp import SmtpHook
with SmtpHook(smtp_conn_id=SMTP_CONN_ID) as hook:
hook.send_email_smtp(
to=NOTIFY_EMAIL,
subject=f"Airflow Survey Analysis: {SCHEDULED_PROMPT}",
html_content=f"<pre>{data}</pre>",
)
else:
print(f"Survey analysis result:\n{data}")
generate_sql >> run_query >> result_data >> send_result(result_data)
Variants¶
example_llm_survey_interactive in the same file takes ad hoc questions: a
HITLEntryOperator edits the question, an ApprovalOperator reviews the rows. No
download, no schedule; the CSV is assumed in place.
@dag(tags=["example"])
def example_llm_survey_interactive():
"""
Ask a natural language question about the survey with human review at each end.
Task graph::
prompt_confirmation (HITLEntryOperator)
→ generate_sql (LLMSQLQueryOperator)
→ run_query (AnalyticsOperator)
→ extract_data (@task)
→ result_confirmation (ApprovalOperator)
The first HITL step lets the analyst review and optionally reword the
question before it reaches the LLM. The final HITL step presents the
query result for approval or rejection.
"""
# ------------------------------------------------------------------
# Step 1: Prompt confirmation -- review or edit the question.
# ------------------------------------------------------------------
prompt_confirmation = HITLEntryOperator(
task_id="prompt_confirmation",
subject="Review the survey analysis question",
params={
"prompt": Param(
INTERACTIVE_PROMPT,
type="string",
description="The natural language question to answer via SQL",
)
},
response_timeout=datetime.timedelta(hours=1),
)
# ------------------------------------------------------------------
# Step 2: SQL generation -- LLM translates the confirmed question.
# ------------------------------------------------------------------
generate_sql = LLMSQLQueryOperator(
task_id="generate_sql",
prompt="{{ ti.xcom_pull(task_ids='prompt_confirmation')['params_input']['prompt'] }}",
llm_conn_id=LLM_CONN_ID,
datasource_config=survey_datasource,
schema_context=SURVEY_SCHEMA,
)
# ------------------------------------------------------------------
# Step 3: SQL execution via Apache DataFusion.
# ------------------------------------------------------------------
run_query = AnalyticsOperator(
task_id="run_query",
datasource_configs=[survey_datasource],
queries=["{{ ti.xcom_pull(task_ids='generate_sql') }}"],
result_output_format="json",
)
# ------------------------------------------------------------------
# Step 4: Extract data rows from the JSON result.
# AnalyticsOperator returns [{"query": "...", "data": [...]}, ...]
# This step strips the query field so only the rows reach the reviewer.
# ------------------------------------------------------------------
@task
def extract_data(raw: str) -> str:
results = json.loads(raw)
data = [row for item in results for row in item["data"]]
return json.dumps(data, indent=2)
result_data = extract_data(run_query.output)
# ------------------------------------------------------------------
# Step 5: Result confirmation -- approve or reject the query result.
# ------------------------------------------------------------------
result_confirmation = ApprovalOperator( # noqa: F841
task_id="result_confirmation",
subject="Review the survey query result",
body=result_data,
response_timeout=datetime.timedelta(hours=1),
)
prompt_confirmation >> generate_sql >> run_query
When one question is not enough, example_llm_survey_agentic splits a research question
into sub-questions, maps SQL generation and execution over them, and synthesizes the
results behind an approval gate.
Adapting it¶
Make the schema check block: add a
@task.branchoncompatiblebetweencheck_schemaandgenerate_sql, as Block a load when the schema drifts does.Point
SURVEY_CSV_ENDPOINTand the HTTP connection at your own published file, and replaceSURVEY_SCHEMAwith a description of its columns; that description is what the model writes SQL against.Change
SCHEDULED_PROMPTto the question the report answers.Swap the local-file
DataSourceConfigfor a warehouse connection when the data lives in a database rather than a file;LLMSQLQueryOperatoraccepts either.