chore: cleanup and CI adjustments

- Removed unused `reduce` import from Athena module.
- Applied ruff formatting to `start_query_executions`.
- Fixed static check issues to pass CI.
- Added ruff check on Athena tests file.

Author: ggiallo28

awswrangler/athena/_executions.py

@@ -3,22 +3,21 @@
 from __future__ import annotations
 
 import logging
+import os
 import time
+from concurrent.futures import ThreadPoolExecutor
 from typing import (
     Any,
     Dict,
     cast,
 )
 
-import os
 import boto3
 import botocore
 from typing_extensions import Literal
 
-from concurrent.futures import ThreadPoolExecutor
 from awswrangler import _utils, exceptions, typing
 from awswrangler._config import apply_configs
-from functools import reduce
 
 from ._cache import _CacheInfo, _check_for_cached_results
 from ._utils import (
@@ -34,6 +33,7 @@
 
 _DEFAULT_MAX_WORKERS = max(4, os.cpu_count() or 4)
 
+
 @apply_configs
 def start_query_execution(
     sql: str,
@@ -172,6 +172,7 @@ def start_query_execution(
 
     return query_execution_id
 
+
 @apply_configs
 def start_query_executions(
     sqls: list[str],

awswrangler/athena/_executions.pyi

@@ -58,6 +58,71 @@ def start_query_execution(
     data_source: str | None = ...,
     wait: bool,
 ) -> str | dict[str, Any]: ...
+@overload
+def start_query_executions(
+    sqls: list[str],
+    database: str | None = ...,
+    s3_output: str | None = ...,
+    workgroup: str = ...,
+    encryption: str | None = ...,
+    kms_key: str | None = ...,
+    params: dict[str, Any] | list[str] | None = ...,
+    paramstyle: Literal["qmark", "named"] = ...,
+    boto3_session: boto3.Session | None = ...,
+    client_request_token: str | list[str] | None = ...,
+    athena_cache_settings: typing.AthenaCacheSettings | None = ...,
+    athena_query_wait_polling_delay: float = ...,
+    data_source: str | None = ...,
+    wait: Literal[False] = ...,
+    check_workgroup: bool = ...,
+    enforce_workgroup: bool = ...,
+    as_iterator: bool = ...,
+    use_threads: bool | int = ...,
+) -> list[str]: ...
+@overload
+def start_query_executions(
+    sqls: list[str],
+    *,
+    database: str | None = ...,
+    s3_output: str | None = ...,
+    workgroup: str = ...,
+    encryption: str | None = ...,
+    kms_key: str | None = ...,
+    params: dict[str, Any] | list[str] | None = ...,
+    paramstyle: Literal["qmark", "named"] = ...,
+    boto3_session: boto3.Session | None = ...,
+    client_request_token: str | list[str] | None = ...,
+    athena_cache_settings: typing.AthenaCacheSettings | None = ...,
+    athena_query_wait_polling_delay: float = ...,
+    data_source: str | None = ...,
+    wait: Literal[True],
+    check_workgroup: bool = ...,
+    enforce_workgroup: bool = ...,
+    as_iterator: bool = ...,
+    use_threads: bool | int = ...,
+) -> list[dict[str, Any]]: ...
+@overload
+def start_query_executions(
+    sqls: list[str],
+    *,
+    database: str | None = ...,
+    s3_output: str | None = ...,
+    workgroup: str = ...,
+    encryption: str | None = ...,
+    kms_key: str | None = ...,
+    params: dict[str, Any] | list[str] | None = ...,
+    paramstyle: Literal["qmark", "named"] = ...,
+    boto3_session: boto3.Session | None = ...,
+    client_request_token: str | list[str] | None = ...,
+    athena_cache_settings: typing.AthenaCacheSettings | None = ...,
+    athena_query_wait_polling_delay: float = ...,
+    data_source: str | None = ...,
+    wait: bool,
+    check_workgroup: bool = ...,
+    enforce_workgroup: bool = ...,
+    as_iterator: bool = ...,
+    use_threads: bool | int = ...,
+) -> list[str] | list[dict[str, Any]]: ...
 def stop_query_execution(query_execution_id: str, boto3_session: boto3.Session | None = ...) -> None: ...
 def wait_query(
     query_execution_id: str,

tests/unit/test_athena.py

@@ -1708,3 +1708,62 @@ def test_athena_date_recovery(path, glue_database, glue_table):
         ctas_approach=False,
     )
     assert pandas_equals(df, df2)
+
+
+def test_start_query_executions_ids_and_results(path, glue_database, glue_table):
+    # Prepare table
+    wr.s3.to_parquet(
+        df=get_df(),
+        path=path,
+        index=True,
+        dataset=True,
+        mode="overwrite",
+        database=glue_database,
+        table=glue_table,
+        partition_cols=["par0", "par1"],
+    )
+
+    sqls = [
+        f"SELECT * FROM {glue_table} LIMIT 1",
+        f"SELECT COUNT(*) FROM {glue_table}",
+    ]
+
+    # Case 1: Sequential, return query IDs
+    qids = wr.athena.start_query_executions(sqls=sqls, database=glue_database, wait=False, use_threads=False)
+    assert isinstance(qids, list)
+    assert all(isinstance(qid, str) for qid in qids)
+    assert len(qids) == len(sqls)
+
+    # Case 2: Sequential, wait for results
+    results = wr.athena.start_query_executions(sqls=sqls, database=glue_database, wait=True, use_threads=False)
+    assert isinstance(results, list)
+    assert all(isinstance(r, dict) for r in results)
+    assert all("Status" in r for r in results)
+
+    # Case 3: Parallel execution with threads
+    results_parallel = wr.athena.start_query_executions(sqls=sqls, database=glue_database, wait=True, use_threads=True)
+    assert isinstance(results_parallel, list)
+    assert all(isinstance(r, dict) for r in results_parallel)
+
+
+def test_start_query_executions_as_iterator(path, glue_database, glue_table):
+    # Prepare table
+    wr.s3.to_parquet(
+        df=get_df(),
+        path=path,
+        index=True,
+        dataset=True,
+        mode="overwrite",
+        database=glue_database,
+        table=glue_table,
+        partition_cols=["par0", "par1"],
+    )
+
+    sqls = [f"SELECT * FROM {glue_table} LIMIT 1"]
+
+    # Case: as_iterator=True should return a generator-like object
+    qids_iter = wr.athena.start_query_executions(sqls=sqls, database=glue_database, wait=False, as_iterator=True)
+    assert not isinstance(qids_iter, list)
+    qids = list(qids_iter)
+    assert len(qids) == 1
+    assert isinstance(qids[0], str)