Skip to content

feat(athena): add model-level timeout support - #4

Open
dtaniwaki wants to merge 3 commits into
mainfrom
feat/athena-model-timeout
Open

feat(athena): add model-level timeout support#4
dtaniwaki wants to merge 3 commits into
mainfrom
feat/athena-model-timeout

Conversation

@dtaniwaki

Copy link
Copy Markdown
Owner

Summary

  • Add model_timeout_seconds (model config) to limit total model execution time including batch loops
  • Thread-safe implementation using threading.local() for deadline tracking
  • Check timeout before each batch iteration and in _run_query()
  • Raise AthenaModelTimeoutError when the deadline is exceeded

Motivation

The dbt-athena adapter has no model-level timeout mechanism. When Iceberg force_batch generates a large number of batches, each batch query is fast but the total model execution time becomes unbounded. This makes it impossible to set a reasonable CodeBuild timeout without risking runaway models.

Standardizing timeout handling is being discussed upstream at dbt-labs/dbt-core#8522, but progress has stalled with no concrete implementation plan. This PR provides a practical solution for the Athena adapter in the meantime.

Usage

In dbt_project.yml:

models:
  my_project:
    my_heavy_model:
      +model_timeout_seconds: 3600  # 1 hour

Or in a model file:

{{ config(model_timeout_seconds=3600) }}
SELECT ...

Changes

  • exceptions.py: Add AthenaModelTimeoutError
  • impl.py: Add model_timeout_seconds to AthenaConfig, add set/check/clear_model_timeout() adapter methods with threading.local(), add timeout check in _run_query() and run_operation_with_potential_multiple_runs()
  • table.sql, incremental.sql: Set/clear model timeout at materialization boundaries
  • create_table_as.sql, helpers.sql, merge.sql: Check model timeout before each batch iteration

Test plan

  • Unit tests for model timeout lifecycle (set, check, clear)
  • Thread isolation test (independent deadlines per thread)
  • Edge cases (None, 0, negative values are no-ops)
  • Existing unit tests pass

Recreated from #1, rebased on main for cleaner history.

tauhid621 and others added 3 commits April 1, 2026 19:44
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Add two timeout mechanisms to prevent runaway queries and models:

- query_timeout_seconds: per-query timeout configured in profiles.yml.
  Cancels the Athena query and raises AthenaQueryTimeoutError when a
  single query exceeds the limit.

- model_timeout_seconds: per-model timeout configured in dbt model config.
  Checks deadline before each batch iteration and in _run_query().
  Raises AthenaModelTimeoutError when total model execution exceeds
  the limit. Uses threading.local() for thread safety.

Batch loop timeout checks are added to:
- create_table_as_with_partitions
- batch_incremental_insert
- batch_iceberg_merge
…ut_seconds

Remove per-query timeout (query_timeout_seconds) as only
model-level timeout is needed. The model_timeout_seconds config
covers the batch loop runaway scenario which is the actual problem.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants