id: bigquery-agent-query-spend-gate
namespace: company.analytics
inputs:
- id: sql
type: STRING
description: The BigQuery Standard SQL SELECT statement to run.
- id: reason
type: STRING
defaults: ad hoc analysis
description: What the result is for. The optimizer uses it to judge which
filters keep the answer useful.
variables:
project: your-gcp-project
budget_gib: 10
usd_per_tib: 6.25
model: claude-haiku-5-5
dry_run_url: "https://bigquery.googleapis.com/bigquery/v2/projects/{{
vars.project }}/queries"
checks:
- message: Provide the SQL to run.
when: "{{ (inputs.sql | trim | length) > 0 }}"
style: ERROR
behavior: BLOCK_EXECUTION
tasks:
- id: token
type: io.kestra.plugin.gcp.auth.OauthAccessToken
timeout: PT1M
serviceAccount: "{{ secret('GCP_SERVICE_ACCOUNT') }}"
projectId: "{{ vars.project }}"
scopes:
- https://www.googleapis.com/auth/bigquery
- id: price_original
type: io.kestra.plugin.core.http.Request
description: Dry-run the query through the BigQuery API. Nothing executes and
nothing is billed.
uri: "{{ render(vars.dry_run_url) }}"
method: POST
contentType: application/json
headers:
Authorization: "Bearer {{ outputs.token.accessToken.tokenValue }}"
body: "{{ {'query': inputs.sql, 'useLegacySql': false, 'dryRun': true} | toJson
}}"
timeout: PT1M
# A dry run has no side effects, so a failed call is safe to retry.
retry:
type: exponential
maxAttempts: 3
interval: PT2S
maxInterval: PT20S
- id: only_select
type: io.kestra.plugin.core.flow.If
condition: "{{ fromJson(outputs.price_original.body).statementType != 'SELECT' }}"
then:
- id: block_non_select
type: io.kestra.plugin.core.execution.Fail
errorMessage: "Only SELECT statements may run through this gate, got {{
fromJson(outputs.price_original.body).statementType }}."
- id: within_budget
type: io.kestra.plugin.core.flow.If
condition: "{{ (fromJson(outputs.price_original.body).totalBytesProcessed |
number) <= vars.budget_gib * 1073741824 }}"
then:
- id: run_query
type: io.kestra.plugin.gcp.bigquery.Query
description: Run the original query, capped slightly above its own estimate so a
bad estimate cannot overspend.
serviceAccount: "{{ secret('GCP_SERVICE_ACCOUNT') }}"
projectId: "{{ vars.project }}"
sql: "{{ inputs.sql }}"
maximumBytesBilled: "{{
(fromJson(outputs.price_original.body).totalBytesProcessed | number) *
12 / 10 + 10485760 }}"
fetchType: STORE
timeout: PT10M
else:
- id: optimize
type: io.kestra.plugin.ai.agent.AIAgent
description: Rewrite the query to fit the budget, checking each attempt with a
real BigQuery dry run.
timeout: PT5M
maxSequentialToolsInvocations: 6
provider:
type: io.kestra.plugin.ai.provider.Anthropic
apiKey: "{{ secret('ANTHROPIC_API_KEY') }}"
modelName: "{{ vars.model }}"
systemMessage: |
You reduce BigQuery on-demand cost. Rewrite the user's SELECT so it scans at most the budget.
Prefer the changes that keep the answer useful for the stated reason: select only needed
columns instead of SELECT *, add a filter on the table's partitioning column, and avoid
scanning tables twice. Keep it a single SELECT statement.
Check every candidate with the kestra_task_bigquery_dry_run tool before answering. Its body
must be a JSON string exactly like {"query": "<your SQL>", "useLegacySql": false, "dryRun": true}.
The response field totalBytesProcessed is the scan size. A response with an error message means
the SQL is invalid: fix it and dry-run again. Quote column names that are reserved words with
backticks, for example `hash`. Stop at the first candidate under budget.
Answer with the final SQL, its totalBytesProcessed from your last dry run, and one or two
sentences on what the rewrite drops or narrows compared to the original.
prompt: |
Budget: {{ vars.budget_gib }} GiB ({{ vars.budget_gib * 1073741824 }} bytes).
Original estimate: {{ fromJson(outputs.price_original.body).totalBytesProcessed }} bytes.
Reason for the query: {{ inputs.reason }}
Original SQL:
{{ inputs.sql }}
tools:
- type: io.kestra.plugin.ai.tool.KestraTask
tasks:
- id: bigquery_dry_run
type: io.kestra.plugin.core.http.Request
uri: "{{ render(vars.dry_run_url) }}"
method: POST
contentType: application/json
headers:
Authorization: "Bearer {{ outputs.token.accessToken.tokenValue }}"
options:
allowFailed: true
body: "..."
- id: proposal
type: io.kestra.plugin.ai.completion.JSONStructuredExtraction
description: Turn the agent's answer into fields the flow can check.
schemaName: query_rewrite
jsonFields:
- rewritten_sql
- estimated_bytes
- changes
prompt: |
Extract the final rewritten SQL exactly as written (no code fences), the bytes estimate, and the
explanation of what changed from this answer:
{{ outputs.optimize.textOutput }}
provider:
type: io.kestra.plugin.ai.provider.Anthropic
apiKey: "{{ secret('ANTHROPIC_API_KEY') }}"
modelName: "{{ vars.model }}"
timeout: PT2M
- id: price_rewrite
type: io.kestra.plugin.core.http.Request
description: Dry-run the proposed SQL independently, so the gate trusts
BigQuery's number and not the agent's.
uri: "{{ render(vars.dry_run_url) }}"
method: POST
contentType: application/json
headers:
Authorization: "Bearer {{ outputs.token.accessToken.tokenValue }}"
body: "{{ {'query': fromJson(outputs.proposal.extractedJson).rewritten_sql,
'useLegacySql': false, 'dryRun': true} | toJson }}"
options:
allowFailed: true
timeout: PT1M
retry:
type: exponential
maxAttempts: 3
interval: PT2S
maxInterval: PT20S
- id: rewrite_fits
type: io.kestra.plugin.core.flow.If
condition: "{{ outputs.price_rewrite.code == 200 and
fromJson(outputs.price_rewrite.body).statementType == 'SELECT' and
(fromJson(outputs.price_rewrite.body).totalBytesProcessed | number) <=
vars.budget_gib * 1073741824 }}"
else:
- id: block_over_budget
type: io.kestra.plugin.core.execution.Fail
errorMessage: "Query scans {{
((fromJson(outputs.price_original.body).totalBytesProcessed |
number) / 1073741824.0) | numberFormat('#,##0.00') }} GiB and no
verified rewrite fits the {{ vars.budget_gib }} GiB budget."
then:
- id: review_summary
type: io.kestra.plugin.core.log.Log
message: |
Original: {{ ((fromJson(outputs.price_original.body).totalBytesProcessed | number) / 1073741824.0) | numberFormat('#,##0.00') }} GiB, about ${{ ((fromJson(outputs.price_original.body).totalBytesProcessed | number) / 1099511627776.0 * vars.usd_per_tib) | numberFormat('#,##0.00') }}
Rewrite: {{ ((fromJson(outputs.price_rewrite.body).totalBytesProcessed | number) / 1073741824.0) | numberFormat('#,##0.00') }} GiB, about ${{ ((fromJson(outputs.price_rewrite.body).totalBytesProcessed | number) / 1099511627776.0 * vars.usd_per_tib) | numberFormat('#,##0.00') }} (verified by a second dry run)
What changes: {{ fromJson(outputs.proposal.extractedJson).changes }}
Rewritten SQL:
{{ fromJson(outputs.proposal.extractedJson).rewritten_sql }}
- id: approve_rewrite
type: io.kestra.plugin.core.flow.Pause
description: The rewrite can change the answer, so a person decides whether it
still serves the reason.
pauseDuration: PT30M
behavior: FAIL
onResume:
- id: decision
type: SELECT
values:
- run_rewrite
- reject
defaults: reject
description: Run the verified cheaper rewrite, or reject and run nothing.
- id: approved
type: io.kestra.plugin.core.flow.If
condition: "{{ outputs.approve_rewrite.onResume.decision == 'run_rewrite' }}"
then:
- id: run_rewrite
type: io.kestra.plugin.gcp.bigquery.Query
serviceAccount: "{{ secret('GCP_SERVICE_ACCOUNT') }}"
projectId: "{{ vars.project }}"
sql: "{{ fromJson(outputs.proposal.extractedJson).rewritten_sql }}"
maximumBytesBilled: "{{
(fromJson(outputs.price_rewrite.body).totalBytesProcessed |
number) * 12 / 10 + 10485760 }}"
fetchType: STORE
timeout: PT10M
else:
- id: rejected
type: io.kestra.plugin.core.execution.Fail
errorMessage: The reviewer rejected the rewrite, so nothing ran.
outputs:
- id: result
type: STRING
description: Internal storage URI of the query result (ION rows).
value: "{{ outputs.run_query.uri ?? outputs.run_rewrite.uri ?? '' }}"
- id: executed_sql
type: STRING
description: The SQL that actually ran, empty when the gate blocked the request.
value: "{{ outputs.run_rewrite is defined ?
fromJson(outputs.proposal.extractedJson).rewritten_sql :
(outputs.run_query is defined ? inputs.sql : '') }}"
- id: scanned_gib_estimate
type: STRING
description: Dry-run estimate of the query that ran, in GiB.
value: "{{ outputs.run_rewrite is defined ?
(((fromJson(outputs.price_rewrite.body).totalBytesProcessed | number) /
1073741824.0) | numberFormat('#,##0.000')) : (outputs.run_query is defined
? (((fromJson(outputs.price_original.body).totalBytesProcessed | number) /
1073741824.0) | numberFormat('#,##0.000')) : '0') }}"
triggers:
- id: agent_tool
type: io.kestra.plugin.core.trigger.McpToolTrigger
toolName: bigquery_query_with_budget
title: Run a BigQuery query within the team budget
toolDescription: Use this instead of querying BigQuery directly. It prices the
SELECT with a dry run, runs it when it fits the budget, and otherwise
proposes a cheaper rewrite that a person must approve, so the call can
wait for a human.
annotations:
readOnly: true
openWorld: true
- id: http
type: io.kestra.plugin.core.trigger.Webhook
key: bigquery-spend-gate
inputs:
sql: "{{ trigger.body.sql }}"
reason: "{{ trigger.body.reason ?? 'webhook request' }}"
