Skip to content

Redesign the PyAthena test fixtures: define once, generate data, set up once per job #848

Closed

Description

@laughingman7743

Use case

The PyAthena suite's shared fixtures have barely changed since 2022 (#295). They are costly to set up, and they are hard to change.

Setup cost

  • tests/__init__.py gives every process its own random schema name (pyathena_test_<random>).
  • pytest_sessionstart in tests/pyathena/conftest.py runs in every process. It uploads tests/resources/rows/* to S3, creates the database, and runs create_table.sql.jinja2: 34 statements, namely 16 DROP TABLE IF EXISTS (redundant in a fresh schema), 16 CREATE EXTERNAL TABLE, and 2 CREATE OR REPLACE VIEW.
  • With pytest -n 8, one job sets up the same fixtures 9 times: 8 workers plus the pytest-xdist controller, which runs no tests. A full run on 5 Python versions creates 45 copies. In the sample on Reduce the AWS cost of integration tests without reducing real-service coverage #834, CREATE TABLE/DROP TABLE statements were 15.4% of all CI queries.

Maintainability

  • The data is files in tests/resources/rows/, prepared by hand in advance. The maintainer has confirmed that this approach does not need to be kept:
    • one_row_complex.gz is a single gzipped row in Hive text format, with \002/\003 separators for arrays, maps, and structs. It cannot practically be edited by hand.
    • many_rows.tsv has 10,000 lines.
  • one_row_complex (16 columns) is referenced 83 times in 9 files. The cursor tests (default, pandas, arrow, polars, s3fs) spell out the full row or DataFrame for each cursor type, and the SQLAlchemy reflection tests assert its column list. Adding a column means editing the Hive text file and every one of those expectations.
  • 9 of the 16 tables are execute_many* tables: mutable, one per cursor type, created for every session whether used or not.
  • tests/resources/queries/insert_into_table.sql.jinja2 is not referenced anywhere.

Proposed change

Redesign the shared fixtures so that they are defined once, generated, and set up once per job.

  1. Single definition in Python.
    • Describe each shared table in code: name, columns with Athena types, storage format, partitions, and rows as Python values.
    • Generate the DDL and the data files from this definition, and upload them. No queries are needed beyond CREATE EXTERNAL TABLE/CREATE VIEW.
    • The data format is free to choose; TSV does not need to be kept. Parquet written with pyarrow (already a dev dependency) is a good default: it represents binary, decimal, microsecond timestamps, and nested array/map/struct values natively, with no delimiter escaping.
    • Keep text-format tables only where a test depends on text-format behavior, such as NA handling.
  2. Expectations derived from the definition.
    • Tests that read a whole row get their expected values from the definition, with per-cursor conversion helpers (Python objects, pandas, Arrow, Polars).
    • Other tests select explicit columns instead of SELECT *.
    • Adding a column or a type-coverage table then means adding one entry.
  3. Set up once per job.
    • Skip the setup in the xdist controller and drop the redundant DROP TABLE IF EXISTS statements (a small first step).
    • Then create the read-only fixtures once, in one schema shared by all workers. For example, the controller creates it and passes the name through pytest-xdist workerinput.
    • Keep cleanup working after interrupted runs (the scheduled database sweep covers leftovers).
  4. Mutable tables per test. Replace the session-wide execute_many* tables with per-test tables, like the existing executemany_table fixture, so the shared schema stays read-only.
  5. Remove the unused insert_into_table.sql.jinja2.

This can land in steps. Step 3's first part is independent and cheap; the definition and the expectation helpers can move one table at a time, starting with one_row_complex.

Validation plan (if implementing)

  • just test pyathena against AWS passes with -n 8 and with -n 1, with the same number of tests and no weakened assertions. Compare the collected test IDs and the pass/skip counts before and after.
  • Count the setup statements and S3 uploads per job before and after.
  • Show that adding a column to a shared table is a one-place change, by adding one in the PR that introduces the definition, or in a follow-up.
  • Confirm that the schema and the S3 Tables namespace are cleaned up after normal and interrupted runs.

Activity

  1. changed the title [-]Set up the PyAthena test fixtures once per job instead of in every xdist process[/-] [+]Redesign the PyAthena test fixtures: define once, generate data, set up once per job[/+] on Sep 26, 2026
  2. added this to the 4.0.0 milestone on Sep 26, 2026
  3. laughingman7743 commented on Sep 27, 2026

    @laughingman7743
    MemberAuthor

    A related teardown failure, from the SQLAlchemy compliance suite rather than tests/pyathena/, may be worth covering in the fixture redesign.

    What happened

    In run 36252040653 (PR #858), job test-sqla-async / run (3.11), attempt 1:

    • All 577 tests passed.
    • The job then failed with INTERNALERROR in pytest_testnodedown. The follower-database hook _awsathena_drop_db in tests/sqlalchemy/conftest.py ran DROP DATABASE test_b4798d487ea4 CASCADE, which failed with SemanticException [Error 10072]: Database does not exist.
    • A rerun passed.

    The database was dropped

    • Only one such query ran: query 4f17347c-..., the only DROP DATABASE for this name among StartQueryExecution calls between 15:41:55 and 15:42:25 UTC.
    • CloudTrail showed no throttled or duplicate client calls.
    • Glue CloudTrail events for that query, all UTC:
    Time Glue call Result
    15:42:01 GetDatabase ok
    15:42:01 GetTables ok
    15:42:02 DeleteDatabase ok
    15:42:03 GetDatabase (invoked by Athena) EntityNotFoundException

    So Athena deleted the database, looked it up again, and reported the query as FAILED.
    The error came from the service side, after a successful delete.

    Effect on the current hook

    • The hook drops {ident}, {ident}_test_schema, and {ident}_test_schema_2 in sequence without IF EXISTS.
    • The first failure aborted the teardown, so the two _test_schema databases were left behind until the 7-day sweep.
    • The PyAthena suite's drop_database.sql.jinja2 already uses DROP DATABASE IF EXISTS ... CASCADE.

    Suggestion for the redesign

    • Make the follower-database teardown tolerate a database that is already gone. For example, use DROP DATABASE IF EXISTS ... CASCADE and attempt each of the three drops even if an earlier one fails.
    • I have not verified that IF EXISTS avoids this particular post-delete failure, because it did not reproduce on rerun.
    • Treating Database does not exist from a DROP as success would also keep a passing test run from failing in teardown.
  4. laughingman7743 commented on Sep 27, 2026

    @laughingman7743
    MemberAuthor

    Proposal for the fixture work: add IF EXISTS to the follower-database drops in tests/sqlalchemy/conftest.py.

    The PyAthena suite's teardown already uses it (tests/resources/queries/drop_database.sql.jinja2: DROP DATABASE IF EXISTS {{ schema }} CASCADE).
    The SQLAlchemy suite's _awsathena_drop_db does not:

    @drop_db.for_db("awsathena")
    def _awsathena_drop_db(cfg, eng, ident):
        with eng.begin() as conn:
            conn.exec_driver_sql(f"DROP DATABASE {ident} CASCADE")
            conn.exec_driver_sql(f"DROP DATABASE {ident}_test_schema CASCADE")
            conn.exec_driver_sql(f"DROP DATABASE {ident}_test_schema_2 CASCADE")

    Suggested change:

    @drop_db.for_db("awsathena")
    def _awsathena_drop_db(cfg, eng, ident):
        with eng.begin() as conn:
            for name in (ident, f"{ident}_test_schema", f"{ident}_test_schema_2"):
                conn.exec_driver_sql(f"DROP DATABASE IF EXISTS {name} CASCADE")

    Why:

    • As described in the comment above, Athena can report Database does not exist for a DROP DATABASE ... CASCADE that already deleted the database.
    • Without IF EXISTS, that turns a passing run into an INTERNALERROR in pytest_testnodedown, and it aborts the remaining two drops, leaving the _test_schema databases for the 7-day sweep.
    • IF EXISTS also matches the PyAthena suite's teardown.
    • _awsathena_create_db calls this hook inside contextlib.suppress(Exception) to clear leftovers before creating, so IF EXISTS does not change that path.

    Limit: I could not reproduce the failure, so it is unverified whether IF EXISTS suppresses Athena's post-delete GetDatabase failure.
    If it does not, attempting each drop independently and treating Database does not exist as success would still let the three drops complete.

  5. laughingman7743 commented on Sep 28, 2026

    @laughingman7743
    MemberAuthor

    Status of the proposed changes:

    Step PR State
    1. Single definition in Python (generated DDL and data) #874 Merged
    2. Expectations derived from the definition #875, #876, #877 Closed without merging (see below)
    3. Set up once per job: first part #866 Merged
    3. Set up once per job: shared fixture schema #878 Open
    4. Mutable tables per test #873 Merged
    5. Remove insert_into_table.sql.jinja2 #866 Merged

    Step 2 was dropped by the maintainer's decision. To derive the expected results of every cursor from the table definitions, tests/pyathena/expected.py had to re-model how each cursor returns each type:

    • Athena's text rendering of arrays, maps, and structs, and how the converters parse it back
    • JSON casts
    • the pandas, Arrow, and Polars value and type conversions, for both the CSV and UNLOAD paths

    It reached about 800 lines, enough that it would need tests of its own. The hand-written expectations are more verbose, but they can be read as they are.

    Adding a column to one_row_complex therefore means changing two things:

    • the definition in tests/pyathena/tables.py
    • the hand-written expectations of the tests that assert the whole row

    Notes from the work that stay useful:

    • many_rows and one_row remain text tables. For SELECT ... LIMIT without ORDER BY, Athena returned a 10,000-row Parquet file out of file order. The reflection tests assert the text SerDe of one_row.
    • The sync and aio test_spark_sql tests are always skipped. Their pytest.mark.dependency(depends="test_spark_dataframe") passes a string instead of a list. This existing problem is unrelated to this issue.

    This issue can be closed once #878 is merged.

  6. laughingman7743 commented on Sep 28, 2026

    @laughingman7743
    MemberAuthor

    Closing. Final state of the proposed changes:

    Step PR State
    1. Single definition in Python (generated DDL and data) #874 Merged
    2. Expectations derived from the definition #875, #876, #877 Not done (see the previous comment)
    3. Set up once per job: first part (controller skip, redundant DROPs) #866 Merged
    3. Set up once per job: shared fixture schema #878 Not done
    4. Mutable tables per test #873 Merged
    5. Remove insert_into_table.sql.jinja2 #866 Merged

    Step 3's shared fixture schema was dropped by the maintainer's decision. #866 and #873 already brought the session setup from 333 to 88 statements per pytest -n 8 job. The shared schema would have removed 61 more.

    That saving is small next to the few thousand queries in a job. Athena does not charge for DDL, so it comes to under a tenth of a cent per job, and the wall time does not change.

    To get there, the session lifecycle would have needed the following:

    • pytest-xdist hooks and workerinput
    • a session-start wrapper with cleanup fallbacks
    • a guard on the fixture schema
    • a new default schema, which every test would have had to keep separate from ENV.schema

    The saving is not worth that complexity.

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

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions