chore(ci): fix static checks to pass CI

Author: ggiallo28

awswrangler/athena/_executions.py

@@ -3,19 +3,19 @@
 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
 

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,65 @@ 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)