Skip to content

Allow AthenaOperator queries without a database argument#69846

Merged
eladkal merged 3 commits into
apache:mainfrom
te-horie:make_database_optional_aws_athena
Jul 14, 2026
Merged

Allow AthenaOperator queries without a database argument#69846
eladkal merged 3 commits into
apache:mainfrom
te-horie:make_database_optional_aws_athena

Conversation

@te-horie

Copy link
Copy Markdown
Contributor

Summary

This PR makes the database argument optional in AthenaOperator, allowing queries that do not require a default database.
This is useful when all referenced table names are fully qualified, such as in queries across multiple databases.

When database is not provided, Database from query_execution_context is used as the fallback database for OpenLineage dataset extraction.

Testing

  • breeze run pytest providers/amazon/tests/unit/amazon/aws/operators/test_athena.py passed (16 test cases).
  • prek run --from-ref upstream/main passed.
  • Successfully ran the system test against Amazon Athena
    1. Locally removed the database argument from the read_table task in example_athena.py like this:
      diff --git a/providers/amazon/tests/system/amazon/aws/example_athena.py b/providers/amazon/tests/system/amazon/aws/example_athena.py
      index 36a37a8e1c..3c72e9b309 100644
      --- a/providers/amazon/tests/system/amazon/aws/example_athena.py
      +++ b/providers/amazon/tests/system/amazon/aws/example_athena.py
      @@ -129,7 +129,6 @@ with DAG(
           read_table = AthenaOperator(
               task_id="read_table",
               query=query_read_table,
      -        database=athena_database,
               output_location=f"s3://{s3_bucket}/",
           )
           # [END howto_operator_athena]
    2. Ran the system test with the following command; it passed ✅:
      $ SYSTEM_TESTS_ENV_ID=airflowtest$(date +%s) breeze run --forward-credentials env AWS_PROFILE=airflow-test pytest --system providers/amazon/tests/system/amazon/aws/example_athena.py
      ...
      providers/amazon/tests/system/amazon/aws/example_athena.py::test_run <- devel-common/src/tests_common/test_utils/system_tests.py PASSED [100%]
      
      ================================================= 1 passed, 1 warning in 56.41s ==================================================
    3. Confirmed that the query succeeded with only Catalog in QueryExecutionContext:
      $ aws athena get-query-execution --query-execution-id <<<REDACTED>>> --profile airflow-test
      QueryExecution:
        ...
        Query: SELECT * from airflowtest1784003174_default.airflowtest1784003174_test_table
        QueryExecutionContext:
          Catalog: awsdatacatalog
        QueryExecutionId: <<<REDACTED>>>
        ...
        Status:
          CompletionDateTime: '2026-07-14T13:27:19.237000+09:00'
          State: SUCCEEDED
          SubmissionDateTime: '2026-07-14T13:27:17.950000+09:00'
        SubstatementType: SELECT
        WorkGroup: primary

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Generated-by: OpenAI Codex (GPT-5.5) following the guidelines


Drafted-by: OpenAI Codex; reviewed by @te-horie before posting

@boring-cyborg boring-cyborg Bot added area:providers provider:amazon AWS/Amazon - related issues labels Jul 14, 2026
@boring-cyborg

boring-cyborg Bot commented Jul 14, 2026

Copy link
Copy Markdown

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide
Here are some useful points:

  • Pay attention to the quality of your code (ruff, mypy and type annotations). Our prek-hooks will help you with that.
  • In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example Dag that shows how users should use it.
  • Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
  • Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
  • Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
  • Be sure to read the Airflow Coding style.
  • Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
    Apache Airflow is a community-driven project and together we are making it better 🚀.
    In case of doubts contact the developers at:
    Mailing List: dev@airflow.apache.org
    Slack: https://s.apache.org/airflow-slack

te-horie added 2 commits July 14, 2026 14:00
Athena does not require the `Database` field in `QueryExecutionContext`.
Queries with fully qualified table names can therefore run without a
default database.
https://docs.aws.amazon.com/athena/latest/APIReference/API_StartQueryExecution.html#API_StartQueryExecution_RequestSyntax

Requiring the argument forces users to provide a default database even
when the query does not need one.
When the `database` argument of `AthenaOperator` is `None`, the
`Database` field in `QueryExecutionContext` can still provide the
default database for unqualified table names in a query.
@te-horie te-horie force-pushed the make_database_optional_aws_athena branch from b49829e to d7d2c25 Compare July 14, 2026 05:00
@te-horie te-horie marked this pull request as ready for review July 14, 2026 05:03
@te-horie te-horie requested a review from o-nikolas as a code owner July 14, 2026 05:03
@Vamsi-klu

Copy link
Copy Markdown
Contributor

The if self.database: guard in execute() looks right, it skips writing Database into query_execution_context when the arg is unset instead of forcing a None in there, and defaulting the signature to database: str | None = None leaves existing keyword callers alone. Nice that fallback_database = self.database or self.query_execution_context.get("Database") keeps the OpenLineage schema resolution working when the database only arrives through the execution context. One subtle behavior maybe worth a docstring line: with database=None and a query_execution_context that already sets Database, execute() now preserves that value instead of overwriting it, which looks like the right call but is easy to miss. The two new tests cover the no database run and the context fallback, so this looks good to me.

@eladkal eladkal requested a review from vincbeck July 14, 2026 10:34
Suggested in apache#69846 (comment)

It is better to explain to users who omit `database` or set it to `None`
that `Database` in `query_execution_context` remains effective.
@te-horie

te-horie commented Jul 14, 2026

Copy link
Copy Markdown
Contributor Author

One subtle behavior maybe worth a docstring line: with database=None and a query_execution_context that already sets Database, execute() now preserves that value instead of overwriting it, which looks like the right call but is easy to miss.

Thank you for your suggestion!
I added a note to the database parameter docstring in f72d787 explaining that any Database value in query_execution_context is used when database is omitted or set to None.

@eladkal eladkal merged commit 26c2c12 into apache:main Jul 14, 2026
83 checks passed
@boring-cyborg

boring-cyborg Bot commented Jul 14, 2026

Copy link
Copy Markdown

Awesome work, congrats on your first merged pull request! You are invited to check our Issue Tracker for additional contributions.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:providers provider:amazon AWS/Amazon - related issues

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants