Support Athena managed query result storage#665
Open
laughingman7743 wants to merge 2 commits intomasterfrom
Open
Support Athena managed query result storage#665laughingman7743 wants to merge 2 commits intomasterfrom
laughingman7743 wants to merge 2 commits intomasterfrom
Conversation
When a workgroup has managed query result storage enabled, GetQueryExecution does not return ResultConfiguration.OutputLocation. File-based cursors (Pandas, Arrow, Polars, S3FS) checked output_location and returned empty results even though data was available via API. Fall back to GetQueryResults API when output_location is None but the query succeeded. S3 file reading remains the primary path when output_location is available. - Add _fetch_all_rows() to AthenaResultSet base class for paginated API fetching with DefaultTypeConverter - Add _parse_result_rows() to share response parsing between _process_rows (normal path) and _fetch_all_rows (fallback path) - Add _rows_to_columnar() for shared row-to-columnar conversion - Extract __get_query_results() from __fetch() for reuse - Add _as_*_from_api() fallback methods to each result set subclass - Remove empty ResultConfiguration from StartQueryExecution request when no S3 staging dir or encryption is configured - Emit warning log when falling back to API Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…with ProgrammingError - Add test_fetch_all_rows to all cursor test files (Cursor, Pandas, Arrow, Polars, S3FS) parameterized by workgroup (default vs managed) using fixture indirect - Support s3_staging_dir="" to explicitly disable env var fallback, required for managed workgroups where ResultConfiguration conflicts with ManagedQueryResultsConfiguration - Replace all assert statements with ProgrammingError exceptions across connection.py, cursor files, result_set files (11 occurrences) - Add AWS_ATHENA_MANAGED_WORKGROUP env var to tests/Env, GitHub Actions, docs - Add ManagedWorkGroup resource to CloudFormation template Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #664
When a workgroup has managed query result storage enabled,
GetQueryExecutiondoes not returnResultConfiguration.OutputLocation. This caused file-based cursors (Pandas, Arrow, Polars, S3FS) to return empty results even though data was available via theGetQueryResultsAPI.Changes
API fallback for managed storage (
result_set.py,pandas/result_set.py,arrow/result_set.py,polars/result_set.py,s3fs/result_set.py)output_locationisNonebut the query succeeded, fall back toGetQueryResultsAPI to fetch dataoutput_locationis available_fetch_all_rows(),_parse_result_rows(),_rows_to_columnar()s3_staging_dir=""support (connection.py)AWS_ATHENA_S3_STAGING_DIRenv var fallbackResultConfigurationconflicts withManagedQueryResultsConfigurationassert→ProgrammingError(11 occurrences across 8 files)assertstatements used for input validation with properProgrammingErrorexceptionsassertcan be stripped bypython -O, making validation unreliableOmit empty
ResultConfiguration(common.py)ResultConfigurationdict inStartQueryExecutionwhen no S3 staging dir or encryption is configuredTests (5 test files +
tests/__init__.py)test_fetch_all_rowsadded to all cursor types, parameterized by workgroup (default vs managed) using fixtureindirectAWS_ATHENA_MANAGED_WORKGROUPenv var controls managed workgroup name; tests skip when not setCI & Docs
AWS_ATHENA_MANAGED_WORKGROUPenv varManagedWorkGroupresourcedocs/testing.md: documented managed workgroup setupTest plan
test_fetch_all_rowstests pass (5 cursor types × 2 workgroup modes)make fmt/make chk(ruff + mypy) passpyathena-managed)🤖 Generated with Claude Code