diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 6df0b1686..8d25856ef 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -8,7 +8,7 @@ import os import threading from concurrent.futures import ThreadPoolExecutor -from inspect import Parameter, signature +from inspect import signature from typing import Any from unittest.mock import MagicMock, patch @@ -19,11 +19,10 @@ from pyathena.arrow.async_cursor import AsyncArrowCursor from pyathena.arrow.cursor import ArrowCursor from pyathena.async_cursor import AsyncCursor, AsyncDictCursor -from pyathena.common import BaseCursor from pyathena.connection import Connection from pyathena.converter import DefaultTypeConverter from pyathena.cursor import Cursor, DictCursor -from pyathena.error import ProgrammingError +from pyathena.error import OperationalError, ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.pandas.async_cursor import AsyncPandasCursor from pyathena.pandas.cursor import PandasCursor @@ -36,6 +35,7 @@ from pyathena.spark.common import SparkBaseCursor from pyathena.spark.cursor import SparkCursor from pyathena.util import RetryConfig +from tests.pyathena.util import EVENT_TIMEOUT, answered_requests # Cursors whose arraysize is the GetQueryResults page size, capped at 1000. PAGED_CURSORS = [Cursor, DictCursor, AsyncCursor, AsyncDictCursor] @@ -51,11 +51,12 @@ AsyncS3FSCursor, ] CURSOR_CLASSES = PAGED_CURSORS + UNCAPPED_CURSORS + [SparkCursor, AsyncSparkCursor] +# Cursors that run on_start_query_execution. +START_CALLBACK_CURSORS = [Cursor, DictCursor, ArrowCursor, PandasCursor, PolarsCursor, S3FSCursor] # Cursors that pass execute() keyword arguments to a DataFrame reader. READER_CURSOR_CLASSES = [PandasCursor, PolarsCursor, AsyncPandasCursor, AsyncPolarsCursor] -@pytest.mark.parametrize("factory", [Connection, pyathena.connect]) @pytest.mark.parametrize( ("setting", "kwargs", "expected"), [ @@ -69,10 +70,10 @@ ("work_group", {"work_group": "explicit-workgroup"}, "explicit-workgroup"), ], ) -def test_environment_fallback(factory, setting, kwargs, expected, monkeypatch): +def test_environment_fallback(setting, kwargs, expected, monkeypatch): monkeypatch.setenv("AWS_ATHENA_S3_STAGING_DIR", "s3://environment/path/") monkeypatch.setenv("AWS_ATHENA_WORK_GROUP", "environment-workgroup") - conn = factory( + conn = Connection( region_name="us-east-1", aws_access_key_id="access_key", aws_secret_access_key="secret_key", @@ -84,60 +85,17 @@ def test_environment_fallback(factory, setting, kwargs, expected, monkeypatch): conn.close() -@pytest.mark.parametrize("factory", [Connection, pyathena.connect]) -def test_empty_settings_disable_both_environment_fallbacks(factory, monkeypatch): +def test_empty_settings_disable_both_environment_fallbacks(monkeypatch): monkeypatch.setenv("AWS_ATHENA_S3_STAGING_DIR", "s3://environment/path/") monkeypatch.setenv("AWS_ATHENA_WORK_GROUP", "environment-workgroup") with pytest.raises(ProgrammingError, match="Required argument"): - factory(region_name="us-east-1", s3_staging_dir="", work_group="") + Connection(region_name="us-east-1", s3_staging_dir="", work_group="") @pytest.mark.parametrize("factory", [pyathena.connect, Connection]) def test_connect_keyword_only(factory): - assert all( - parameter.kind in (Parameter.KEYWORD_ONLY, Parameter.VAR_KEYWORD) - for parameter in signature(factory).parameters.values() - ) - with patch("pyathena.connection.Session") as session: - with pytest.raises(TypeError, match="positional"): - factory("s3://bucket/path/", "us-east-1") - session.assert_not_called() - - -@pytest.mark.parametrize( - "cursor_class", [*CURSOR_CLASSES, BaseCursor, SparkBaseCursor, WithResultSet] -) -def test_cursor_constructor_keyword_only(cursor_class): - constructor = cursor_class.__init__ - assert [ - parameter.name - for parameter in signature(constructor).parameters.values() - if parameter.kind in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD) - ] == ["self"] - with pytest.raises(TypeError, match="positional"): - constructor(object(), "positional-setting") - - -@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) -def test_cursor_execute_keyword_only(cursor_class): - execute_signature = signature(cursor_class.execute) - assert [ - parameter.name - for parameter in execute_signature.parameters.values() - if parameter.kind in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD) - ] == ["self", "operation", "parameters"] - parameters = {"value": 1} - cursor = object.__new__(cursor_class) - bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters) - assert bound.arguments["parameters"] is parameters - if cursor_class in (SparkCursor, AsyncSparkCursor): - assert "work_group" not in execute_signature.parameters - else: - bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters, work_group="group") - assert bound.arguments["work_group"] == "group" - execute_signature.bind(cursor, operation="SELECT 1") - with pytest.raises(TypeError, match="positional"): - cursor.execute("SELECT %(value)s", parameters, "group") + with pytest.raises(TypeError): + factory("s3://bucket/path/") class RecordingCursor(Cursor): @@ -196,6 +154,21 @@ def _open_cursor(conn: Connection[Any], cursor_class: type[Any], **kwargs: Any) return conn.cursor(cursor_class, **kwargs) +@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) +def test_cursor_constructor_keyword_only(cursor_class): + # Binds the arguments as a call does before the constructor runs. A call + # through Connection.cursor() passes every setting by keyword, so a positional + # setting would fail there as a duplicate keyword argument too. + with pytest.raises(TypeError): + signature(cursor_class).bind("positional-setting") + + +@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) +def test_cursor_execute_keyword_only(offline_connection, cursor_class): + with _open_cursor(offline_connection, cursor_class) as cursor, pytest.raises(TypeError): + cursor.execute("SELECT %(value)s", {"value": 1}, "group") + + @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): @@ -207,13 +180,25 @@ def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class): - on_start_query_execution, on_poll = MagicMock(), MagicMock() - offline_connection.on_start_query_execution = on_start_query_execution - offline_connection.on_poll = on_poll + started, polled = [], [] + offline_connection.on_start_query_execution = started.append + offline_connection.on_poll = polled.append - with _open_cursor(offline_connection, cursor_class) as cursor: - assert cursor._on_start_query_execution is on_start_query_execution - assert cursor._on_poll is on_poll + with ( + answered_requests(offline_connection.client), + _open_cursor(offline_connection, cursor_class) as cursor, + ): + try: + result = cursor.execute("SELECT 1") + except OperationalError: + pass + else: + # The Async cursors run the execution in a future. + result[1].result(timeout=EVENT_TIMEOUT) + + assert [execution.state for execution in polled] == ["FAILED"] + # The Async and Spark cursors do not run on_start_query_execution. + assert started == (["query-id"] if cursor_class in START_CALLBACK_CURSORS else []) @pytest.mark.parametrize( @@ -335,11 +320,18 @@ def test_s3_client_built_once_across_threads(self): created = [] session_client = conn.session.client + second_build = threading.Event() + def contended_client(*args, **kwargs): - created.append((args, conn._s3_client_lock.locked())) + created.append(args) + # Without the lock, another thread would build a client meanwhile. + if len(created) == 1: + second_build.wait(1) + else: + second_build.set() return session_client(*args, **kwargs) - conn._session.client = contended_client + conn.session.client = contended_client barrier = threading.Barrier(8) def get_client(_): @@ -349,7 +341,7 @@ def get_client(_): with ThreadPoolExecutor(max_workers=8) as executor: clients = list(executor.map(get_client, range(8))) - assert created == [(("s3",), True)] + assert created == [("s3",)] assert all(client is clients[0] for client in clients) assert clients[0].meta.service_model.service_name == "s3" @@ -401,13 +393,15 @@ def test_s3_config_merged_over_config(self): def test_close_closes_built_clients(self): conn = _connection() - with patch.object(conn.client, "close") as athena_close: + with ( + patch.object(conn.session, "client") as build, + patch.object(conn.client, "close") as athena_close, + ): conn.close() athena_close.assert_called_once_with() # Clients that were not used are not built to be closed. - assert conn._glue._client is None - assert conn._s3_client is None + build.assert_not_called() with ( patch.object(conn.client, "close") as athena_close, diff --git a/tests/pyathena/test_converter.py b/tests/pyathena/test_converter.py index bac5a8e0c..5f8507c7f 100644 --- a/tests/pyathena/test_converter.py +++ b/tests/pyathena/test_converter.py @@ -123,18 +123,15 @@ def test_to_struct_athena_nested_formats(input_value, expected): @pytest.mark.parametrize( - "input_value", + ("input_value", "expected"), [ - "{formula=x=y+1, status=active}", - '{json={"key": "value"}, name=test}', - '{message=He said "hello", name=John}', + ("{formula=x=y+1, status=active}", {"formula": "x=y+1", "status": "active"}), + ('{json={"key": "value"}, name=test}', {"json": {"key": "value"}, "name": "test"}), + ('{message=He said "hello", name=John}', {"message": 'He said "hello"', "name": "John"}), ], ) -def test_to_struct_athena_complex_cases(input_value): - result = _to_struct(input_value) - # Values too complex to parse fall back to the original string - # instead of being silently dropped - assert result == input_value or isinstance(result, dict) +def test_to_struct_athena_complex_cases(input_value, expected): + assert _to_struct(input_value) == expected @pytest.mark.parametrize( @@ -332,17 +329,6 @@ def test_to_array_invalid_formats(input_value): assert _to_array(input_value) is None -@pytest.mark.parametrize( - ("value", "type_hint", "expected"), - [ - ('[""]', "array(json)", [""]), - ('{"k": ""}', "map(varchar,json)", {"k": ""}), - ], -) -def test_nested_json_empty_string(value, type_hint, expected): - assert DefaultTypeConverter().convert(type_hint.split("(")[0], value, type_hint) == expected - - class TestDefaultTypeConverter: @pytest.mark.parametrize( ("value", "expected"), @@ -351,77 +337,11 @@ class TestDefaultTypeConverter: def test_boolean_conversion(self, value, expected): assert DefaultTypeConverter().convert("boolean", value) is expected - @pytest.mark.parametrize( - ("input_value", "expected"), - [ - ('{"name": "Alice", "age": 25}', {"name": "Alice", "age": 25}), - (None, None), - ("", None), - ("invalid json", None), - ("{a=1, b=2}", {"a": "1", "b": "2"}), - ], - ) - def test_struct_conversion(self, input_value, expected): - converter = DefaultTypeConverter() - assert converter.convert("row", input_value) == expected - - @pytest.mark.parametrize( - ("input_value", "expected"), - [ - ("[1, 2, 3]", [1, 2, 3]), - ('["a", "b", "c"]', ["a", "b", "c"]), - (None, None), - ("", None), - ("invalid json", None), - ("[apple, banana]", ["apple", "banana"]), - ("[]", []), - ], - ) - def test_array_conversion(self, input_value, expected): - converter = DefaultTypeConverter() - assert converter.convert("array", input_value) == expected - - def test_array_varchar_keeps_strings(self): - converter = DefaultTypeConverter() - result = converter.convert("array", "[1234, 5678]", type_hint="array(varchar)") - assert result == ["1234", "5678"] - - def test_array_integer_converts_to_int(self): - converter = DefaultTypeConverter() - result = converter.convert("array", "[1, 2, 3]", type_hint="array(integer)") - assert result == [1, 2, 3] - - def test_array_boolean_converts(self): - converter = DefaultTypeConverter() - result = converter.convert("array", "[true, false]", type_hint="array(boolean)") - assert result == [True, False] - - def test_array_with_null(self): - converter = DefaultTypeConverter() - result = converter.convert("array", "[1, null, 3]", type_hint="array(integer)") - assert result == [1, None, 3] - - def test_map_varchar_integer(self): - converter = DefaultTypeConverter() - result = converter.convert( - "map", '{"key1": 1, "key2": 2}', type_hint="map(varchar, integer)" - ) - assert result == {"key1": 1, "key2": 2} - def test_map_native_format_with_hints(self): converter = DefaultTypeConverter() result = converter.convert("map", "{a=1, b=2}", type_hint="map(varchar, integer)") assert result == {"a": 1, "b": 2} - def test_row_type_hint(self): - converter = DefaultTypeConverter() - result = converter.convert( - "row", - '{"name": "Alice", "age": 25}', - type_hint="row(name varchar, age integer)", - ) - assert result == {"name": "Alice", "age": 25} - def test_row_native_format_with_hints(self): converter = DefaultTypeConverter() result = converter.convert( @@ -443,30 +363,18 @@ def test_nested_array_of_row(self): {"name": "Bob", "age": 30}, ] - def test_array_varchar_prevents_number_inference(self): + @pytest.mark.parametrize( + ("input_value", "expected"), + [("[1234, 5678]", ["1234", "5678"]), ("[1234, 5678, hello]", ["1234", "5678", "hello"])], + ) + def test_array_varchar_prevents_number_inference(self, input_value, expected): converter = DefaultTypeConverter() - result = converter.convert( - "array", - "[1234, 5678, hello]", - type_hint="array(varchar)", - ) - assert result == ["1234", "5678", "hello"] + assert converter.convert("array", input_value, type_hint="array(varchar)") == expected def test_none_value_with_type_hint(self): converter = DefaultTypeConverter() assert converter.convert("array", None, type_hint="array(varchar)") is None - def test_simple_type_hint(self): - converter = DefaultTypeConverter() - assert converter.convert("varchar", "hello", type_hint="varchar") == "hello" - - def test_type_hint_caching(self): - converter = DefaultTypeConverter() - converter.convert("array", "[1, 2]", type_hint="array(integer)") - assert "array(integer)" in converter._parsed_hints - converter.convert("array", "[3, 4]", type_hint="array(integer)") - assert len(converter._parsed_hints) == 1 - def test_empty_array_with_type_hint(self): converter = DefaultTypeConverter() assert converter.convert("array", "[]", type_hint="array(varchar)") == [] @@ -485,71 +393,13 @@ def test_row_with_nested_struct(self): ) assert result == {"header": {"seq": 123, "stamp": "2024"}, "x": 4.5} - def test_fallback_on_malformed_value(self): - """When typed conversion fails (returns None), fall back to untyped conversion.""" - converter = DefaultTypeConverter() - # "not-an-array" doesn't look like an array — typed converter returns None. - # Untyped _to_array also returns None for this input, which is correct. - result = converter.convert("array", "not-an-array", type_hint="array(integer)") - assert result is None - - def test_fallback_preserves_struct_value(self): - """Malformed struct with type_hint still falls back to untyped parsing.""" - converter = DefaultTypeConverter() - # Struct with no closing brace — typed converter returns None. - # Untyped _to_struct also returns None here. - result = converter.convert("row", "{unclosed", type_hint="row(a integer)") - assert result is None - def test_fallback_returns_untyped_result(self): - """When typed conversion returns None, untyped conversion is used.""" + """A value that the type hint does not describe is converted without the hint.""" converter = DefaultTypeConverter() - # The typed converter returns None for a struct that doesn't start with "{". - # The untyped _to_struct also returns None for non-struct input. - # Use an array example where typed converter returns None (not a bracket-wrapped - # value), but untyped _to_array can still parse it via JSON. - result = converter.convert( - "row", - '{"a": 1}', - type_hint="row(a varchar)", - ) - # Typed conversion succeeds here — "a" is varchar so "1" stays a string - assert result == {"a": "1"} - - def test_hive_syntax_through_converter(self): - """Hive-style syntax works end-to-end through DefaultTypeConverter.""" - converter = DefaultTypeConverter() - result = converter.convert("array", "[1, 2, 3]", type_hint="array") - assert result == [1, 2, 3] - - def test_hive_syntax_struct_through_converter(self): - """Hive struct syntax works end-to-end.""" - converter = DefaultTypeConverter() - result = converter.convert( - "row", - "{name=Alice, age=25}", - type_hint="struct", - ) - assert result == {"name": "Alice", "age": 25} - - def test_hive_syntax_caching(self): - """Hive syntax is normalized before cache lookup.""" - converter = DefaultTypeConverter() - converter.convert("array", "[1]", type_hint="array") - converter.convert("array", "[2]", type_hint="array(integer)") - # Both should normalize to "array(integer)" in the cache - assert "array(integer)" in converter._parsed_hints - assert len(converter._parsed_hints) == 1 - - def test_normalize_hive_syntax_noop(self): - """Trino-style input passes through unchanged.""" - assert DefaultTypeConverter._normalize_hive_syntax("array(integer)") == "array(integer)" - - def test_normalize_hive_syntax_replaces(self): - assert ( - DefaultTypeConverter._normalize_hive_syntax("array>") - == "array(struct(a int))" - ) + assert converter.convert("row", "{a=1, b=2}", type_hint="array(integer)") == { + "a": "1", + "b": "2", + } def test_normalize_hive_syntax_struct(self): converter = DefaultTypeConverter() @@ -571,11 +421,8 @@ def test_normalize_hive_syntax_nested(self): def test_normalize_hive_syntax_map(self): converter = DefaultTypeConverter() - result = converter.convert( - "map", - '{"x": 1, "y": 2}', - type_hint="map", - ) + # The native values are strings without the hint. + result = converter.convert("map", "{x=1, y=2}", type_hint="map") assert result == {"x": 1, "y": 2} def test_normalize_hive_syntax_mixed(self): diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index b5604e854..6401c07f6 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -1,6 +1,4 @@ import contextlib -import json -import logging import random import re import string @@ -13,7 +11,7 @@ from datetime import UTC, date, datetime, timedelta from decimal import Decimal from random import randint -from unittest.mock import MagicMock, call, patch +from unittest.mock import MagicMock, patch import pytest from botocore.exceptions import ClientError @@ -31,7 +29,6 @@ ExecuteOptions, ) from pyathena.async_cursor import AsyncCursor -from pyathena.converter import _to_array, _to_map, _to_struct from pyathena.cursor import Cursor, DictCursor from pyathena.error import ( DatabaseError, @@ -52,15 +49,15 @@ CONVERTED_VALUES_QUERY, CONVERTED_VALUES_ROW, EVENT_TIMEOUT, + QueryExecutionHistory, ThrottledOperation, interrupt_start_waits, succeeded_query_execution, + succeeded_query_execution_data, throttle_metadata_api, unreachable_glue, ) -_logger = logging.getLogger(__name__) - def _offline_cursor(kill_on_interrupt, cursor_class=Cursor): """A cursor whose query requests go to a mocked Athena client. @@ -94,6 +91,40 @@ def _offline_cursor(kill_on_interrupt, cursor_class=Cursor): return cursor, cancel +def _cache_search_cursor(client): + """A cursor whose result cache search sends its requests to a client stand-in. + + Args: + client: The Athena client stand-in, such as ``QueryExecutionHistory``. + + Returns: + The cursor, for the schema ``this_schema``. + """ + cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls + cursor._connection = MagicMock(_client=client) + cursor._retry_config = RetryConfig(attempt=1) + cursor._work_group = None + cursor._schema_name = "this_schema" + cursor._catalog_name = None + return cursor + + +def _types(value): + """The type of a value and of each value it contains, for comparing them. + + Args: + value: A fetched value, such as a row of lists and dictionaries. + + Returns: + The types, in the shape of the value. + """ + if isinstance(value, dict): + return {k: _types(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_types(v) for v in value] + return type(value) + + def _block_start(cursor, response=None): """Make the cursor's StartQueryExecution request block until released. @@ -394,46 +425,50 @@ def execution(catalog): == "query_id_awsdatacatalog" ) - def test_cache_search_stops_at_expired_execution(self, caplog): + def test_cache_search_stops_at_expired_execution(self): query = "SELECT * FROM one_row" now = datetime.now(UTC) - cursor = Cursor.__new__(Cursor) - cursor._schema_name = "this_schema" - cursor._catalog_name = None - page = [ - succeeded_query_execution("expired", query, now - timedelta(hours=2)), - succeeded_query_execution("other", "SELECT 1", now), + page_size = Cursor.LIST_QUERY_EXECUTIONS_MAX_RESULTS + executions = [ + succeeded_query_execution_data("other", "SELECT 1", now), + succeeded_query_execution_data("expired", query, now - timedelta(hours=2)), + *( + succeeded_query_execution_data(f"older_{i}", "SELECT 1", now - timedelta(hours=3)) + for i in range(page_size - 2) + ), + # A search that went on to the second page would return this match. + succeeded_query_execution_data("next_page", query, now), ] + history = QueryExecutionHistory(executions) + cursor = _cache_search_cursor(history) + + # Without cache_size, cache_expiration_time alone bounds the search. + assert cursor._find_previous_query_id(query, None, cache_expiration_time=3600) is None + assert len(history.list_requests) == 1 + # A failed search would also return None, but the same search reuses the + # execution once it is within cache_expiration_time. + assert ( + cursor._find_previous_query_id(query, None, cache_expiration_time=3 * 3600) == "expired" + ) - # A page read after the expired execution would return this match. - next_page = (None, [succeeded_query_execution("next_page", query, now)]) - - with patch.object( - Cursor, "_list_query_executions", side_effect=[("next_token", page), next_page] - ) as list_mock: - # Without cache_size, cache_expiration_time alone bounds the search. - assert cursor._find_previous_query_id(query, None, cache_expiration_time=3600) is None - list_mock.assert_called_once_with(None, next_token=None, max_results=50) - # A failed search also returns None; the expiry must stop it without an error. - assert "Failed to check the cache" not in caplog.text - - def test_cache_search_reads_pages_up_to_cache_size(self): + @pytest.mark.parametrize( + ("position", "expected"), + [(59, "match"), (60, None)], + ids=["last_searched", "after_cache_size"], + ) + def test_cache_search_reads_pages_up_to_cache_size(self, position, expected): + """With cache_size=60, the search reads the 60 latest executions over two pages.""" query = "SELECT * FROM one_row" now = datetime.now(UTC) - cursor = Cursor.__new__(Cursor) - cursor._schema_name = "this_schema" - cursor._catalog_name = None - pages = [ - ("next_token", [succeeded_query_execution("other", "SELECT 1", now)]), - ("last_token", [succeeded_query_execution("match", query, now)]), + executions = [ + succeeded_query_execution_data(f"other_{i}", "SELECT 1", now) for i in range(70) ] + executions[position] = succeeded_query_execution_data("match", query, now) + history = QueryExecutionHistory(executions) + cursor = _cache_search_cursor(history) - with patch.object(Cursor, "_list_query_executions", side_effect=pages) as list_mock: - assert cursor._find_previous_query_id(query, "wg", cache_size=60) == "match" - assert list_mock.call_args_list == [ - call("wg", next_token=None, max_results=50), - call("wg", next_token="next_token", max_results=10), - ] + assert cursor._find_previous_query_id(query, "wg", cache_size=60) == expected + assert [request["WorkGroup"] for request in history.list_requests] == ["wg", "wg"] def test_cache_search_prefers_latest_execution(self): query = "SELECT * FROM one_row" @@ -582,9 +617,9 @@ def test_null_param(self, cursor): cursor.execute("SELECT %(param)s FROM one_row", {"param": None}) assert cursor.fetchall() == [(None,)] - @pytest.mark.parametrize( - ("value", "expected"), - [ + def test_binary_parameter(self, cursor): + # Empty, special, all bytes, bytearray, memoryview, DB API Binary, and NULL. + values = [ (b"", b""), (b"\x00\xff'\\%", b"\x00\xff'\\%"), (bytes(range(256)), bytes(range(256))), @@ -592,25 +627,13 @@ def test_null_param(self, cursor): (memoryview(b"\x00\xff"), b"\x00\xff"), (Binary(bytearray(b"abc")), b"abc"), (None, None), - ], - ids=["empty", "special", "all_bytes", "bytearray", "memoryview", "dbapi_binary", "null"], - ) - def test_binary_parameter(self, cursor, value, expected): - cursor.execute("SELECT CAST(%(value)s AS VARBINARY)", {"value": value}) - assert cursor.fetchone() == (expected,) - assert cursor.description[0][1] == BINARY - - def test_timestamp_result_precision(self, cursor): + ] cursor.execute( - """ - SELECT TIMESTAMP '2020-01-01 00:00:00', - CAST(TIMESTAMP '2020-01-01 00:00:00.123456' AS TIMESTAMP(9)) - """ - ) - assert cursor.fetchone() == ( - datetime(2020, 1, 1), - datetime(2020, 1, 1, 0, 0, 0, 123456), + "SELECT " + ", ".join(f"CAST(%(v{i})s AS VARBINARY)" for i in range(len(values))), + {f"v{i}": value for i, (value, _) in enumerate(values)}, ) + assert cursor.fetchone() == tuple(expected for _, expected in values) + assert [d[1] for d in cursor.description] == [BINARY] * len(values) def test_datetime_parameter_microseconds(self, cursor): value = datetime(2017, 1, 1, 12, 0, 0, 789012) @@ -880,76 +903,48 @@ def test_complex(self, cursor): ] def test_complex_with_type_hints(self, cursor): - # 1. Basic complex columns from one_row_complex + # Without the hints, the numbers in the structs and the map are strings. cursor.execute( """ - SELECT col_array, col_map, col_struct + SELECT + col_array + ,col_map + ,col_struct + ,CAST( + ROW(ROW('2024-01-01', 123), 4.736, 0.583) + AS ROW(header ROW(stamp VARCHAR, seq INTEGER), x DOUBLE, y DOUBLE) + ) AS positions + ,CAST( + ARRAY[ + ROW(ROW(1, 2), ROW(CAST(0.5 AS DOUBLE))), + ROW(ROW(3, 4), ROW(CAST(1.5 AS DOUBLE))) + ] + AS ARRAY(ROW(pos ROW(x INTEGER, y INTEGER), vel ROW(x DOUBLE))) + ) AS data FROM one_row_complex """, result_set_type_hints={ "col_array": "array(integer)", "col_map": "map(integer, integer)", "col_struct": "row(a integer, b integer)", - }, - ) - row = cursor.fetchone() - assert row[0] == [1, 2] - assert row[1] == {"1": 2, "3": 4} - assert row[2] == {"a": 1, "b": 2} - # Without type hints col_struct values are strings; with hints they are ints - assert isinstance(row[2]["a"], int) - assert isinstance(row[2]["b"], int) - # Map values should also be ints - assert isinstance(row[1]["1"], int) - - # 2. Nested struct - cursor.execute( - """ - SELECT CAST( - ROW(ROW('2024-01-01', 123), 4.736, 0.583) - AS ROW(header ROW(stamp VARCHAR, seq INTEGER), x DOUBLE, y DOUBLE) - ) AS positions - """, - result_set_type_hints={ "positions": "row(header row(stamp varchar, seq integer), x double, y double)", - }, - ) - row = cursor.fetchone() - positions = row[0] - assert positions["header"]["stamp"] == "2024-01-01" - assert positions["header"]["seq"] == 123 - assert isinstance(positions["header"]["seq"], int) - assert positions["x"] == 4.736 - assert isinstance(positions["x"], float) - assert positions["y"] == 0.583 - assert isinstance(positions["y"], float) - - # 3. Array of nested structs - cursor.execute( - """ - SELECT CAST( - ARRAY[ - ROW(ROW(1, 2), ROW(CAST(0.5 AS DOUBLE))), - ROW(ROW(3, 4), ROW(CAST(1.5 AS DOUBLE))) - ] - AS ARRAY(ROW(pos ROW(x INTEGER, y INTEGER), vel ROW(x DOUBLE))) - ) AS data - """, - result_set_type_hints={ "data": "array(row(pos row(x integer, y integer), vel row(x double)))", }, ) + expected = ( + [1, 2], + {"1": 2, "3": 4}, + {"a": 1, "b": 2}, + {"header": {"stamp": "2024-01-01", "seq": 123}, "x": 4.736, "y": 0.583}, + [ + {"pos": {"x": 1, "y": 2}, "vel": {"x": 0.5}}, + {"pos": {"x": 3, "y": 4}, "vel": {"x": 1.5}}, + ], + ) row = cursor.fetchone() - data = row[0] - assert len(data) == 2 - assert data[0]["pos"]["x"] == 1 - assert isinstance(data[0]["pos"]["x"], int) - assert data[0]["pos"]["y"] == 2 - assert isinstance(data[0]["pos"]["y"], int) - assert data[0]["vel"]["x"] == 0.5 - assert isinstance(data[0]["vel"]["x"], float) - assert data[1]["pos"]["x"] == 3 - assert data[1]["vel"]["x"] == 1.5 + assert row == expected + # == does not tell 123 from 123.0. + assert _types(row) == _types(expected) def test_cancel(self, cursor): def cancel(c): @@ -1044,21 +1039,16 @@ def test_executemany(self, cursor, empty_table): cursor.execute(f"SELECT * FROM {empty_table}") assert sorted(cursor.fetchall()) == list(rows) + # The sum of the update counts does not depend on the statement; test_iceberg_table + # checks the update count of INSERT, UPDATE, and DELETE. @pytest.mark.parametrize( ("operation", "expected"), [ - pytest.param( - "INSERT INTO {table} SELECT id+10, group_id+10, value " - "FROM {table} WHERE group_id=%(group_id)d", - [(1, 10), (2, 20), (3, 30), (11, 10), (12, 20), (13, 30)], - id="insert", - ), pytest.param( "UPDATE {table} SET value=value+1 WHERE group_id=%(group_id)d", [(1, 11), (2, 21), (3, 31)], id="update", ), - pytest.param("DELETE FROM {table} WHERE group_id=%(group_id)d", [], id="delete"), ], ) def test_executemany_rowcount(self, cursor, executemany_table, operation, expected): @@ -1341,36 +1331,6 @@ def test_execute_internal_options_passthrough(self): cache_expiration_time=100, ) - def test_execute_qmark_parameters_skip_cache(self): - """A qmark query with parameters never searches the cache (no AWS, #941).""" - cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls - cursor._connection = MagicMock() - cursor._connection.client.start_query_execution.return_value = { - "QueryExecutionId": "test_query_id" - } - cursor._retry_config = RetryConfig() - cursor._kill_on_interrupt = True - - with ( - patch.object( - Cursor, - "_build_start_query_execution_request", - return_value={"ExecutionParameters": ["'1'"]}, - ) as request_mock, - patch.object(Cursor, "_find_previous_query_id", return_value="cached") as cache_mock, - ): - query_id = cursor._execute( - "SELECT ?", - ["'1'"], - options=ExecuteOptions( - paramstyle="qmark", cache_size=10, cache_expiration_time=100 - ), - ) - - assert query_id == "test_query_id" - assert request_mock.call_args.kwargs["execution_parameters"] == ["'1'"] - cache_mock.assert_not_called() - def test_connection_level_callback(self): """Test connection-level default callback.""" callback_results = [] @@ -1417,78 +1377,6 @@ def execute_callback(query_id): assert connection_callbacks[0] == execute_callbacks[0] # Same query_id assert connection_callbacks[0] == cursor.query_id - def test_callback_thread_safety(self, cursor): - """Test that callback can be used for query ID access from another thread.""" - query_started = threading.Event() - stored_query_id = [] - - def on_start(query_id): - stored_query_id.append(query_id) - query_started.set() - - def verify_query_id(): - # Wait for query to start - query_started.wait(timeout=10) - - # Verify query_id is accessible - if stored_query_id: - assert stored_query_id[0] == cursor.query_id - - # Use ThreadPoolExecutor for proper thread management - with ThreadPoolExecutor(max_workers=1) as executor: - # Start verification task - verify_future = executor.submit(verify_query_id) - - # Execute query with callback - cursor.execute("SELECT 5", on_start_query_execution=on_start) - - # Wait for verification to complete - verify_future.result(timeout=5) - - # Verify events occurred - assert query_started.is_set() - assert len(stored_query_id) == 1 - - # Verify query completed successfully - row = cursor.fetchone() - assert row == (5,) - - def test_on_poll_invoked_each_iteration(self): - """on_poll fires once per poll iteration with the current execution (no AWS).""" - states = [ - AthenaQueryExecution.STATE_QUEUED, - AthenaQueryExecution.STATE_RUNNING, - AthenaQueryExecution.STATE_SUCCEEDED, - ] - executions = [MagicMock(state=state) for state in states] - received = [] - - cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls - cursor._poll_interval = 0 - cursor._kill_on_interrupt = False - cursor._on_poll = received.append - - with patch.object(Cursor, "_get_query_execution", side_effect=executions): - result = cursor._poll("query_id") - - # Callback received every state in order, including the terminal one - assert [execution.state for execution in received] == states - assert result is executions[-1] - - def test_on_poll_none_is_noop(self): - """A None on_poll callback does not affect polling (no AWS).""" - execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) - - cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls - cursor._poll_interval = 0 - cursor._kill_on_interrupt = False - cursor._on_poll = None - - with patch.object(Cursor, "_get_query_execution", return_value=execution): - result = cursor._poll("query_id") - - assert result is execution - @pytest.mark.parametrize( "final_state", [AthenaQueryExecution.STATE_CANCELLED, AthenaQueryExecution.STATE_SUCCEEDED], @@ -1760,9 +1648,18 @@ def test_fetch_complex_values(self, cursor): ARRAY['x', '', 'y'] AS col_empty ,ARRAY['x', ' ', 'a,b'] AS col_space_comma ,ARRAY[json_parse('1234567890'), json_parse('"a,b"')] AS col_json + ,ARRAY[DATE '2023-01-01', DATE '2023-12-31'] AS col_date + ,ARRAY[CAST('2023-01-01 12:00:00.123' AS TIMESTAMP(3))] AS col_timestamp """ cursor.execute(query) - expected = (["x", "", "y"], ["x", " ", "a,b"], [1234567890, "a,b"]) + # Without a hint, the dates and timestamps in an array stay strings. + expected = ( + ["x", "", "y"], + ["x", " ", "a,b"], + [1234567890, "a,b"], + ["2023-01-01", "2023-12-31"], + ["2023-01-01 12:00:00.123"], + ) assert cursor.fetchall() == [expected] cursor.execute( query, @@ -1786,41 +1683,27 @@ def _metadata_view(metadata): ) def test_throttled_metadata_reads_glue(self, cursor, monkeypatch): + tables = ("one_row", "parquet_with_compression", "partition_table", "view_one_row") + def read(): return ( - [ - self._metadata_view(cursor.get_table_metadata(t)) - for t in ("one_row", "parquet_with_compression", "partition_table") - ], + [self._metadata_view(cursor.get_table_metadata(t)) for t in tables], sorted(self._metadata_view(m) for m in cursor.list_table_metadata()), + sorted(m.name for m in cursor.list_table_metadata(expression="one_row")), ENV.schema in [d.name for d in cursor.list_databases("AwsDataCatalog")], ) expected = read() + # The listing filtered by the expression is not empty. + assert "one_row" in expected[2] calls = throttle_metadata_api(cursor.connection.client, monkeypatch) # Glue reports each table as Athena does, and nothing is retried first. assert read() == expected assert sorted(calls) == sorted( - ["get_table_metadata"] * 3 + ["list_table_metadata", "list_databases"] + ["get_table_metadata"] * len(tables) + ["list_table_metadata"] * 2 + ["list_databases"] ) - def test_wrapped_glue_throttling_reads_glue(self, cursor, monkeypatch): - # Athena reports Glue's own throttling inside a MetadataException, and - # the first attempt does not retry it either. - calls = throttle_metadata_api( - cursor.connection.client, - monkeypatch, - code="MetadataException", - message=( - "Rate exceeded (Service: AmazonDataCatalog; Status Code: 400; " - "Error Code: ThrottlingException; Request ID: example; Proxy: null)" - ), - ) - - assert cursor.get_table_metadata("one_row").name == "one_row" - assert calls == ["get_table_metadata"] - @staticmethod def _unreachable_glue(connection, monkeypatch): glue = unreachable_glue(connection) @@ -2061,393 +1944,3 @@ def test_null_vs_empty_string(self, dict_cursor): # Row 5: "NULL" string literal should be preserved as-is assert rows[4]["value"] == "NULL" - - -class TestComplexDataTypes: - """Test complex data types (STRUCT, ARRAY, MAP) with actual Athena queries.""" - - @pytest.mark.parametrize( - ("query", "description"), - [ - ("SELECT ROW('John', 30) AS simple_struct", "simple_struct"), - ( - "SELECT CAST(ROW('Alice', 25) AS ROW(name VARCHAR, age INTEGER)) AS named_struct", - "named_struct", - ), - ("SELECT ROW('Hello, world', 'x=y+1') AS special_chars_struct", "special_chars_struct"), - ("SELECT ROW('He said \"hello\"', 'It''s working') AS quotes_struct", "quotes_struct"), - ("SELECT ROW('Alice', NULL, 'active') AS null_struct", "null_struct"), - ( - "SELECT ROW(ROW('John', 30), ROW('Engineer', 'Tech')) AS nested_struct", - "nested_struct", - ), - ("SELECT CAST(ROW('Alice', 25, 'Hello, world') AS JSON) AS json_struct", "json_struct"), - ], - ) - def test_struct_types(self, cursor, query, description): - """Test various STRUCT type scenarios to understand Athena's behavior.""" - _logger.info("=== STRUCT Type Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - struct_value = result[0] - _logger.info("%s: %r (type: %s)", description, struct_value, type(struct_value).__name__) - - # Validate struct value and converter behavior - assert struct_value is not None, f"STRUCT value should not be None for {description}" - - # Test struct conversion behavior - if isinstance(struct_value, str): - converted = _to_struct(struct_value) - _logger.info("%s: Converted %r -> %r", description, struct_value, converted) - # Conversion either fully parses to a dict or falls back to the - # original string for values too complex to parse losslessly - if converted is not None: - assert isinstance(converted, dict) or converted == struct_value, ( - f"Converted struct should be dict or the original string for {description}" - ) - elif isinstance(struct_value, dict): - # Already converted by the cursor converter - _logger.info("%s: Already converted to dict: %r", description, struct_value) - else: - # Log unexpected types for debugging but don't fail - _logger.warning( - "%s: Unexpected type %s: %r", - description, - type(struct_value).__name__, - struct_value, - ) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ("SELECT ARRAY[1, 2, 3, 4, 5] AS simple_array", "simple_array"), - ("SELECT ARRAY['apple', 'banana', 'cherry'] AS string_array", "string_array"), - ( - "SELECT ARRAY['Hello, world', 'x=y+1', 'It''s working'] AS special_array", - "special_array", - ), - ("SELECT ARRAY[ROW('Alice', 25), ROW('Bob', 30)] AS struct_array", "struct_array"), - ("SELECT ARRAY[ARRAY[1, 2], ARRAY[3, 4]] AS nested_array", "nested_array"), - ( - "SELECT CAST(MAP(ARRAY['data'], ARRAY[ARRAY['Alice', 'Bob']]) AS JSON) " - "AS array_json", - "array_json", - ), - ], - ) - def test_array_types(self, cursor, query, description): - """Test various ARRAY type scenarios.""" - _logger.info("=== ARRAY Type Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - array_value = result[0] - _logger.info("%s: %r (type: %s)", description, array_value, type(array_value).__name__) - - # Validate array value - assert array_value is not None, f"ARRAY value should not be None for {description}" - _logger.info("%s: Array value type %s", description, type(array_value).__name__) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ( - "SELECT MAP(ARRAY[1, 2, 3], ARRAY['one', 'two', 'three']) AS simple_map", - "simple_map", - ), - ( - "SELECT MAP(ARRAY['name', 'age', 'city'], ARRAY['John', '30', 'Tokyo']) " - "AS string_map", - "string_map", - ), - ( - "SELECT MAP(ARRAY['msg', 'formula'], ARRAY['Hello, world', 'x=y+1']) " - "AS special_map", - "special_map", - ), - ( - "SELECT CAST(MAP(ARRAY['person1', 'person2'], " - "ARRAY[ROW('Alice', 25), ROW('Bob', 30)]) AS JSON) AS struct_value_map", - "struct_value_map", - ), - ( - "SELECT CAST(MAP(ARRAY['name', 'age'], ARRAY['Alice', '25']) AS JSON) AS json_map", - "json_map", - ), - ], - ) - def test_map_types(self, cursor, query, description): - """Test various MAP type scenarios.""" - _logger.info("=== MAP Type Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - map_value = result[0] - _logger.info("%s: %r (type: %s)", description, map_value, type(map_value).__name__) - - # Validate map value and converter behavior - assert map_value is not None, f"MAP value should not be None for {description}" - - # Test map conversion behavior - if isinstance(map_value, str): - # For complex MAP structures, string is expected (JSON or native format) - if "ROW(" in map_value or "ARRAY[" in map_value: - # Complex structure, expect string format - _logger.info("%s: Complex MAP kept as string: %r", description, map_value) - else: - # Simple MAP, try conversion - converted = _to_map(map_value) - _logger.info("%s: Converted %r -> %r", description, map_value, converted) - # Conversion either fully parses to a dict or falls back to the - # original string for values too complex to parse losslessly - if converted is not None: - assert isinstance(converted, dict) or converted == map_value, ( - f"Converted map should be dict or the original string for {description}" - ) - elif isinstance(map_value, dict): - # Already converted by the cursor converter - _logger.info("%s: Already converted to dict: %r", description, map_value) - else: - # Log unexpected types for debugging but don't fail - _logger.warning( - "%s: Unexpected type %s: %r", - description, - type(map_value).__name__, - map_value, - ) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ( - "SELECT CAST(ROW(ARRAY[1, 2, 3], MAP(ARRAY['a', 'b'], ARRAY[1, 2])) AS JSON) " - "AS struct_with_collections", - "struct_with_collections", - ), - ( - "SELECT CAST(ARRAY[MAP(ARRAY['name'], ARRAY['Alice']), " - "MAP(ARRAY['name'], ARRAY['Bob'])] AS JSON) AS array_of_maps", - "array_of_maps", - ), - ( - "SELECT CAST(MAP(ARRAY['numbers', 'letters'], " - "ARRAY[ARRAY['1', '2', '3'], ARRAY['a', 'b', 'c']]) AS JSON) AS map_with_arrays", - "map_with_arrays", - ), - ], - ) - def test_complex_combinations(self, cursor, query, description): - """Test complex combinations of data types.""" - _logger.info("=== Complex Combination Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - complex_value = result[0] - _logger.info("%s: %r (type: %s)", description, complex_value, type(complex_value).__name__) - - # For JSON cast results, expect string values that can be parsed as JSON - if isinstance(complex_value, str): - try: - # Test that the JSON string can be parsed - parsed = json.loads(complex_value) - _logger.info(" Parsed JSON: %r", parsed) - assert parsed is not None, f"Parsed JSON should not be None for {description}" - except json.JSONDecodeError as e: - raise AssertionError(f"JSON parsing failed for {description}: {e}") from e - else: - # If it's not a string, it should still be a valid value (not None) - assert complex_value is not None, f"Complex value should not be None for {description}" - _logger.info("%s: Complex value type %s", description, type(complex_value).__name__) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ("SELECT ARRAY[1, 2, 3, 4, 5] AS simple_array", "simple_array"), - ("SELECT ARRAY['apple', 'banana', 'cherry'] AS string_array", "string_array"), - ("SELECT ARRAY[true, false, null] AS boolean_array", "boolean_array"), - ("SELECT ARRAY[1.5, 2.7, 3.14] AS float_array", "float_array"), - ("SELECT ARRAY[] AS empty_array", "empty_array"), - ("SELECT ARRAY[1, null, 3, null, 5] AS null_elements_array", "null_elements_array"), - ( - "SELECT ARRAY[CAST(1 AS VARCHAR), 'mixed', 'true', 'null', '2.5'] AS mixed_array", - "mixed_array", - ), - ( - "SELECT ARRAY[CAST(1.23 AS DECIMAL(10,2)), CAST(4.56 AS DECIMAL(10,2))] " - "AS decimal_array", - "decimal_array", - ), - ( - "SELECT ARRAY[DATE '2023-01-01', DATE '2023-12-31'] AS date_array", - "date_array", - ), - ( - "SELECT ARRAY[TIMESTAMP '2023-01-01 12:00:00'] AS timestamp_array", - "timestamp_array", - ), - ], - ) - def test_array_types_basic(self, cursor, query, description): - """Test basic ARRAY type scenarios.""" - _logger.info("=== ARRAY Type Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - array_value = result[0] - _logger.info("%s: %r (type: %s)", description, array_value, type(array_value).__name__) - - # Validate array value - assert array_value is not None or description == "empty_array", ( - f"ARRAY value should not be None for {description}" - ) - if description == "empty_array": - assert array_value == [], f"Empty array should be [] for {description}" - else: - assert isinstance(array_value, list), f"ARRAY value should be list for {description}" - _logger.info("%s: Array value type %s", description, type(array_value).__name__) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ( - "SELECT ARRAY[ROW(1, 'Alice'), ROW(2, 'Bob'), ROW(3, 'Charlie')] AS struct_array", - "struct_array", - ), - ( - "SELECT ARRAY[ROW('name', 'John', 25), ROW('name', 'Jane', 30)] " - "AS unnamed_struct_array", - "unnamed_struct_array", - ), - ( - "SELECT ARRAY[CAST(ROW('John', 25) AS ROW(name VARCHAR, age INT)), " - "CAST(ROW('Jane', 30) AS ROW(name VARCHAR, age INT))] AS named_struct_array", - "named_struct_array", - ), - ], - ) - def test_array_types_with_structs(self, cursor, query, description): - """Test ARRAY types containing STRUCT elements.""" - _logger.info("=== ARRAY with STRUCT Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - array_value = result[0] - _logger.info("%s: %r (type: %s)", description, array_value, type(array_value).__name__) - - # Validate array value - assert array_value is not None, f"ARRAY value should not be None for {description}" - assert isinstance(array_value, list), f"ARRAY value should be list for {description}" - assert len(array_value) > 0, f"ARRAY should not be empty for {description}" - - # Check first element is a dict (converted struct) - first_element = array_value[0] - assert isinstance(first_element, dict), ( - f"First array element should be dict (struct) for {description}" - ) - _logger.info("%s: First element: %r", description, first_element) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ( - "SELECT CAST(ARRAY[1, 2, 3] AS JSON) AS arr_json", - "arr_json", - ), - ( - "SELECT CAST(ARRAY['a', 'b', 'c'] AS JSON) AS str_arr_json", - "str_arr_json", - ), - ( - "SELECT CAST(ARRAY[ROW(1, 'Alice'), ROW(2, 'Bob')] AS JSON) AS struct_arr_json", - "struct_arr_json", - ), - ], - ) - def test_array_types_json_cast(self, cursor, query, description): - """Test ARRAY types with JSON casting.""" - _logger.info("=== ARRAY JSON Cast Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - array_value = result[0] - _logger.info("%s: %r (type: %s)", description, array_value, type(array_value).__name__) - - # Validate array value - assert array_value is not None, f"ARRAY value should not be None for {description}" - assert isinstance(array_value, list), ( - f"JSON cast ARRAY value should be list for {description}" - ) - _logger.info("%s: JSON cast array type %s", description, type(array_value).__name__) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ("SELECT CARDINALITY(ARRAY[1, 2, 3, 4, 5]) AS array_size", "array_size"), - ("SELECT ARRAY[10, 20, 30, 40][2] AS array_element", "array_element"), - ("SELECT ARRAY[1, 2] || ARRAY[3, 4] AS array_concat", "array_concat"), - ("SELECT CONTAINS(ARRAY[1, 2, 3], 2) AS array_contains", "array_contains"), - ], - ) - def test_array_operations(self, cursor, query, description): - """Test ARRAY operations and functions.""" - _logger.info("=== ARRAY Operation Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - operation_result = result[0] - _logger.info( - "%s: %r (type: %s)", - description, - operation_result, - type(operation_result).__name__, - ) - - # Validate operation result - assert operation_result is not None, ( - f"ARRAY operation result should not be None for {description}" - ) - - # Type-specific validations - if description == "array_size": - assert operation_result == 5, f"Array size should be 5 for {description}" - elif description == "array_element": - assert operation_result == 20, f"Array element [2] should be 20 for {description}" - elif description == "array_concat": - assert operation_result == [1, 2, 3, 4], ( - f"Array concat should be [1,2,3,4] for {description}" - ) - elif description == "array_contains": - assert operation_result is True, f"Array contains should be True for {description}" - - def test_array_converter_behavior(self, cursor): - """Test ARRAY converter behavior with different formats.""" - _logger.info("=== ARRAY Converter Behavior Test ===") - - # Test simple array conversion - cursor.execute("SELECT ARRAY[1, 2, 3] AS simple") - result = cursor.fetchone() - simple_array = result[0] - _logger.info("Simple array: %r", simple_array) - assert simple_array == [1, 2, 3] - - # Test array with struct conversion - cursor.execute( - "SELECT ARRAY[CAST(ROW(1, 2) AS ROW(a INT, b INT)), " - "CAST(ROW(3, 4) AS ROW(a INT, b INT))] AS struct_array" - ) - result = cursor.fetchone() - struct_array = result[0] - _logger.info("Struct array: %r", struct_array) - assert isinstance(struct_array, list) - assert len(struct_array) == 2 - assert isinstance(struct_array[0], dict) - - # Test converter function directly - test_cases = [ - ("[1, 2, 3]", [1, 2, 3]), - ('["a", "b", "c"]', ["a", "b", "c"]), - ("[{a=1, b=2}]", [{"a": "1", "b": "2"}]), - ("[]", []), - (None, None), - ("invalid", None), - ] - - for test_input, expected in test_cases: - result = _to_array(test_input) - _logger.info("Converter test: %r -> %r", test_input, result) - assert result == expected, ( - f"Converter failed for {test_input}: expected {expected}, got {result}" - ) diff --git a/tests/pyathena/test_formatter.py b/tests/pyathena/test_formatter.py index 8290ee62d..2e40f4df4 100644 --- a/tests/pyathena/test_formatter.py +++ b/tests/pyathena/test_formatter.py @@ -12,7 +12,6 @@ import pytest from pyathena.error import ProgrammingError -from pyathena.formatter import _escape_trino # A value whose single quote must not be allowed to terminate the string # literal. Trino and Athena do not treat a backslash as an escape character @@ -599,7 +598,6 @@ def test_table_name_starting_with_as_is_still_hive_ddl(self, formatter): ) -def test_escape_trino_doubles_single_quotes(): - assert _escape_trino("a'b") == "'a''b'" +def test_backslash_before_quote_stays_literal_in_trino_statement(formatter): # A backslash is not an escape character in Trino/Presto; it stays literal. - assert _escape_trino("a\\'b") == "'a\\''b'" + assert formatter.format("SELECT %(v)s", {"v": "a\\'b"}) == "SELECT 'a\\''b'" diff --git a/tests/pyathena/test_glue.py b/tests/pyathena/test_glue.py index eecd3443b..b9700e024 100644 --- a/tests/pyathena/test_glue.py +++ b/tests/pyathena/test_glue.py @@ -9,13 +9,11 @@ from concurrent.futures import ThreadPoolExecutor import pytest -from botocore.exceptions import ClientError, ParamValidationError -from botocore.exceptions import ConnectionError as BotoConnectionError +from botocore.exceptions import ParamValidationError from pyathena.glue import GlueMetadataClient from tests import ENV from tests.pyathena.conftest import connect -from tests.pyathena.util import unreachable_glue class TestGlueMetadataClient: @@ -43,20 +41,6 @@ def _view(metadata): metadata.parameters, ) - def test_reads_what_athena_reports(self, cursor): - glue = cursor.connection._glue - catalog = cursor.connection.catalog_name - - for table in ("one_row", "parquet_with_compression", "partition_table", "view_one_row"): - assert self._view(glue.get_table(catalog, ENV.schema, table)) == self._view( - cursor.get_table_metadata(table) - ) - assert sorted(self._view(m) for m in glue.list_tables(catalog, ENV.schema)) == sorted( - self._view(m) for m in cursor.list_table_metadata() - ) - assert [m.name for m in glue.list_tables(catalog, ENV.schema, "one_row")] == ["one_row"] - assert ENV.schema in [d.name for d in glue.list_databases(catalog)] - @pytest.mark.skipif( not ENV.s3tables_catalog, reason="AWS_ATHENA_S3_TABLES_CATALOG is not configured", @@ -91,28 +75,6 @@ def test_reads_s3_tables_catalog(self): finally: cursor.execute(f"DROP TABLE IF EXISTS {schema}.{table}") - def test_reports_a_missing_table(self, cursor): - with pytest.raises(ClientError) as caught: - cursor.connection._glue.get_table( - cursor.connection.catalog_name, ENV.schema, "no_such_table_786" - ) - - assert caught.value.response["Error"]["Code"] == "EntityNotFoundException" - assert cursor.connection._glue.reachable - - def test_rejects_an_unsupported_catalog(self, cursor): - with pytest.raises(ValueError, match="federated_catalog"): - cursor.connection._glue.get_table("federated_catalog", ENV.schema, "one_row") - - def test_stops_after_it_cannot_reach_glue(self, cursor): - glue = unreachable_glue(cursor.connection) - - with pytest.raises(BotoConnectionError): - glue.get_table("AwsDataCatalog", ENV.schema, "one_row") - - assert not glue.reachable - assert not glue.usable_for("AwsDataCatalog") - def test_keeps_going_after_a_request_it_rejects(self, cursor): # botocore rejects the request before sending it; Glue stays in use. glue = cursor.connection._glue @@ -133,11 +95,18 @@ def test_client_leaves_out_athena_endpoint(self): created = [] session_client = conn.session.client + second_build = threading.Event() + def contended_client(*args, **kwargs): - created.append((args, glue._lock.locked())) + created.append(args) + # Without the lock, another thread would build a client meanwhile. + if len(created) == 1: + second_build.wait(1) + else: + second_build.set() return session_client(*args, **kwargs) - conn._session.client = contended_client + conn.session.client = contended_client barrier = threading.Barrier(8) def get_client(_): @@ -147,9 +116,10 @@ def get_client(_): with ThreadPoolExecutor(max_workers=8) as executor: clients = list(executor.map(get_client, range(8))) - assert created == [(("glue",), True)] + assert created == [("glue",)] assert all(client is clients[0] for client in clients) assert clients[0].meta.service_model.service_name == "glue" + assert clients[0].meta.service_model.api_version == "2017-03-31" assert clients[0].meta.endpoint_url == f"https://glue.{ENV.region_name}.amazonaws.com" def test_table_metadata_leaves_out_columns_iceberg_no_longer_has(self): diff --git a/tests/pyathena/test_options.py b/tests/pyathena/test_options.py index 2ea323c09..bfd3e9624 100644 --- a/tests/pyathena/test_options.py +++ b/tests/pyathena/test_options.py @@ -30,45 +30,18 @@ def test_is_frozen(self): with pytest.raises(dataclasses.FrozenInstanceError): options.work_group = "primary" # type: ignore[misc] - def test_merge_ignores_none_overrides(self): - options = ExecuteOptions(work_group="primary", cache_size=100) - merged = options.merge(work_group=None, cache_size=None, s3_staging_dir=None) - assert merged.work_group == "primary" - assert merged.cache_size == 100 - assert merged.s3_staging_dir is None - - def test_merge_overrides_take_precedence(self): - options = ExecuteOptions( - work_group="primary", - cache_size=100, - result_reuse_enable=True, - paramstyle="pyformat", - ) - merged = options.merge( - work_group="adhoc", - cache_size=0, - result_reuse_enable=False, - paramstyle="qmark", - ) - assert merged.work_group == "adhoc" - assert merged.cache_size == 0 - assert merged.result_reuse_enable is False - assert merged.paramstyle == "qmark" - def test_merge_does_not_mutate_original(self): options = ExecuteOptions(work_group="primary") merged = options.merge(work_group="adhoc") assert options.work_group == "primary" assert merged is not options - def test_merge_without_applied_overrides_is_equal(self): - options = ExecuteOptions(work_group="primary") - assert options.merge() == options - assert options.merge(work_group=None) == options - - def test_merge_raises_on_unknown_field(self): - with pytest.raises(TypeError, match="unknown_field"): - ExecuteOptions().merge(unknown_field="value") + def test_merge_ignores_none_overrides(self): + options = ExecuteOptions(work_group="primary", cache_size=100) + merged = options.merge(work_group=None, cache_size=None, s3_staging_dir=None) + assert merged.work_group == "primary" + assert merged.cache_size == 100 + assert merged.s3_staging_dir is None @pytest.mark.parametrize( ("base", "overrides", "expected"), @@ -100,10 +73,6 @@ def test_merge_precedence(self, base, overrides, expected): for name, value in expected.items(): assert getattr(merged, name) == value - def test_resolve_with_none_returns_defaults_with_overrides(self): - resolved = ExecuteOptions.resolve(None, work_group="wg") - assert resolved == ExecuteOptions(work_group="wg") - def test_resolve_applies_overrides_to_given_options(self): options = ExecuteOptions(work_group="primary", cache_size=100) resolved = ExecuteOptions.resolve(options, work_group="adhoc", cache_size=None) @@ -127,39 +96,11 @@ def test_resolve_without_base_preserves_override_values(self, overrides, expecte assert resolved.cache_expiration_time == 0 assert resolved.result_reuse_minutes is None - def test_resolve_returns_distinct_default_instances(self): - first = ExecuteOptions.resolve(None) - second = ExecuteOptions.resolve(None) - assert first == second - assert first is not second - - def test_resolve_preserves_callback_and_hint_references(self): - def callback(query_id): - pass - - first_hints = {"tags": "array(varchar)"} - second_hints = {"tags": "map(varchar, integer)"} - first = ExecuteOptions.resolve( - None, on_start_query_execution=callback, result_set_type_hints=first_hints - ) - second = ExecuteOptions.resolve(None, result_set_type_hints=second_hints) - assert first.on_start_query_execution is callback - assert first.result_set_type_hints is first_hints - assert second.result_set_type_hints is second_hints - first_hints["tags"] = "array(bigint)" - assert second.result_set_type_hints == {"tags": "map(varchar, integer)"} - @pytest.mark.parametrize("options", [None, ExecuteOptions(work_group="primary")]) def test_resolve_rejects_unknown_non_none_fields(self, options): with pytest.raises(TypeError, match="unknown_field"): ExecuteOptions.resolve(options, unknown_field="value") - def test_resolve_given_options_preserves_new_instance_behavior(self): - options = ExecuteOptions(work_group="primary") - resolved = ExecuteOptions.resolve(options, work_group=None) - assert resolved == options - assert resolved is not options - @pytest.mark.parametrize("provided", [False, True]) def test_resolve_retains_subclass_merge_dispatch(self, provided): @dataclasses.dataclass(frozen=True) diff --git a/tests/pyathena/test_parser.py b/tests/pyathena/test_parser.py index dd3b827f1..e1c776423 100644 --- a/tests/pyathena/test_parser.py +++ b/tests/pyathena/test_parser.py @@ -10,7 +10,6 @@ from pyathena.converter import _DEFAULT_CONVERTERS, _to_default, _to_struct from pyathena.parser import ( TypedValueConverter, - TypeNode, TypeSignatureParser, ) @@ -47,12 +46,6 @@ def test_array_type(self): assert len(node.children) == 1 assert node.children[0].type_name == "varchar" - def test_array_of_integer(self): - parser = TypeSignatureParser() - node = parser.parse("array(integer)") - assert node.type_name == "array" - assert node.children[0].type_name == "integer" - def test_map_type(self): parser = TypeSignatureParser() node = parser.parse("map(varchar, integer)") @@ -106,16 +99,6 @@ def test_deeply_nested(self): assert row_node.children[0].field_names == ["x", "y"] assert row_node.children[1].type_name == "varchar" - def test_parameterized_type(self): - parser = TypeSignatureParser() - node = parser.parse("decimal(10, 2)") - assert node.type_name == "decimal" - - def test_varchar_with_length(self): - parser = TypeSignatureParser() - node = parser.parse("varchar(255)") - assert node.type_name == "varchar" - def test_type_alias_int(self): parser = TypeSignatureParser() node = parser.parse("int") @@ -127,13 +110,6 @@ def test_type_alias_in_complex_type(self): assert node.type_name == "array" assert node.children[0].type_name == "integer" - def test_trailing_modifier_after_paren(self): - """Type with content after closing paren should not break parsing.""" - parser = TypeSignatureParser() - # Simulates a hypothetical "timestamp(3) with time zone" style input - node = parser.parse("decimal(10, 2) extra") - assert node.type_name == "decimal" - class TestTypedValueConverter: @pytest.fixture @@ -144,40 +120,11 @@ def converter(self): struct_parser=_to_struct, ) - def test_simple_varchar(self, converter): - node = TypeNode("varchar") - assert converter.convert("hello", node) == "hello" - - def test_simple_integer(self, converter): - node = TypeNode("integer") - assert converter.convert("42", node) == 42 - - def test_array_of_varchar(self, converter): - parser = TypeSignatureParser() - node = parser.parse("array(varchar)") - assert converter.convert("[1234, 5678]", node) == ["1234", "5678"] - - def test_array_of_integer(self, converter): - parser = TypeSignatureParser() - node = parser.parse("array(integer)") - assert converter.convert("[1, 2, 3]", node) == [1, 2, 3] - def test_map_varchar_integer(self, converter): parser = TypeSignatureParser() node = parser.parse("map(varchar, integer)") assert converter.convert('{"a": 1, "b": 2}', node) == {"a": 1, "b": 2} - def test_row_named_fields(self, converter): - parser = TypeSignatureParser() - node = parser.parse("row(name varchar, age integer)") - assert converter.convert("{name=Alice, age=25}", node) == {"name": "Alice", "age": 25} - - def test_nested_row(self, converter): - parser = TypeSignatureParser() - node = parser.parse("row(header row(seq integer, stamp varchar), x double)") - result = converter.convert("{header={seq=123, stamp=2024}, x=4.5}", node) - assert result == {"header": {"seq": 123, "stamp": "2024"}, "x": 4.5} - def test_array_of_row_json(self, converter): """JSON path: array(row(...)) with nested dict elements.""" parser = TypeSignatureParser() diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index 0612754f1..d901cdae5 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -9,12 +9,10 @@ import pytest -from pyathena.aio.common import AioBaseCursor, WithAsyncFetch -from pyathena.common import BaseCursor, CursorIterator from pyathena.converter import DefaultTypeConverter -from pyathena.error import OperationalError, ProgrammingError +from pyathena.error import OperationalError from pyathena.model import AthenaQueryExecution -from pyathena.result_set import AthenaResultSet, WithFetch, WithResultSet +from pyathena.result_set import AthenaResultSet from pyathena.util import RetryConfig from tests.pyathena.util import assert_fetches_fail, get_query_results_responses @@ -107,12 +105,6 @@ def test_fetch_after_close(self): assert result_set.fetchall() == [] assert result_set.rownumber is None - def test_get_query_results_after_close(self): - result_set = _fetching_result_set(("1", "2"), ("3",)) - result_set.close() - with pytest.raises(ProgrammingError, match=r"AthenaResultSet is closed\."): - result_set._get_query_results(2, "token1") - def test_fetchone_across_pages(self): result_set = _fetching_result_set(("1", "2"), ("3",)) assert [result_set.fetchone() for _ in range(4)] == [(1,), (2,), (3,), None] @@ -159,16 +151,3 @@ def test_fetch_all_rows_as_csv(self): with patch.object(AthenaResultSet, "_get_query_results", side_effect=pages): data = result_set._fetch_all_rows_as_csv() assert data == b'"v","w"""\n"a,""b""\nc",\n"v","w"""\n"","1"\n' - - -class TestWithResultSet: - def test_is_mixin(self): - assert WithResultSet.__bases__ == (object,) - - @pytest.mark.parametrize( - ("cursor_base", "base"), - [(WithFetch, BaseCursor), (WithAsyncFetch, AioBaseCursor)], - ) - def test_precedes_cursor_bases(self, cursor_base, base): - # Listed first, so that its members take precedence over the cursor bases'. - assert cursor_base.__bases__ == (WithResultSet, base, CursorIterator) diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index f3c93759e..6021212ba 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -1,6 +1,5 @@ import threading import time -from typing import Any import pytest import tenacity.nap @@ -79,41 +78,12 @@ def test_strtobool_invalid(value): strtobool(value) -class _WithCodeError(Exception): - def __init__(self, code: int) -> None: - super().__init__(f"error:{code}") - self.response = {"Error": {"Code": code}} - - class _NoResponseError(Exception): def __init__(self) -> None: super().__init__("error") self.response = None -def _test_retry(ex: Exception) -> None: - calls = {"n": 0} - - def fn() -> Any: - calls["n"] += 1 - raise ex - - cfg = RetryConfig(attempt=1, max_delay=1) - - with pytest.raises(type(ex)): - retry_api_call(fn, config=cfg) - - assert calls["n"] == 1 - - -def test_retry_api_call(): - _test_retry(_WithCodeError(500)) - - -def test_retry_api_call_with_none_error(): - _test_retry(_NoResponseError()) - - @pytest.mark.parametrize( ("code", "message", "exceptions", "expected_calls"), [ @@ -251,23 +221,12 @@ def call(): retry_api_call(call, config) assert calls == 7 assert len(sleeps) == len(expected_bases) - for waited, base in zip(sleeps, expected_bases, strict=True): - assert base <= waited <= base + multiplier - - -def test_retry_api_call_jitter_varies(monkeypatch): - sleeps = [] - monkeypatch.setattr(tenacity.nap.time, "sleep", sleeps.append) - error = throttling_error("GetTableMetadata") - - def call(): - raise error - - config = RetryConfig(attempt=20, multiplier=1, max_delay=1) - with pytest.raises(ClientError): - retry_api_call(call, config) - assert all(1 <= waited <= 2 for waited in sleeps) - assert len(set(sleeps)) > 1 + jitters = [waited - base for waited, base in zip(sleeps, expected_bases, strict=True)] + assert all(0 <= jitter <= multiplier for jitter in jitters) + if multiplier: + # Random, so the six jitters spread over more than the float residuals of a + # constant one would. + assert max(jitters) - min(jitters) > multiplier / 100 @pytest.mark.parametrize( @@ -292,7 +251,11 @@ def call(): def test_is_retryable_error(code, message, expected): error = ClientError({"Error": {"Code": code, "Message": message}}, "GetTableMetadata") assert is_retryable_error(error, RetryConfig()) is expected - assert is_retryable_error(ValueError("no response"), RetryConfig()) is False + + +@pytest.mark.parametrize("error", [ValueError("no response"), _NoResponseError()]) +def test_is_retryable_error_without_response(error): + assert is_retryable_error(error, RetryConfig()) is False @pytest.mark.parametrize( diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 8d3a08412..733416c23 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -10,6 +10,7 @@ import threading import time from concurrent.futures import Future, wait +from contextlib import contextmanager from datetime import UTC, datetime, timedelta, timezone from pathlib import Path from unittest.mock import patch @@ -266,6 +267,30 @@ def fail(**kwargs): return calls +def succeeded_query_execution_data(query_id, query, completion_date_time, schema="this_schema"): + """Build a succeeded DML query execution as ``BatchGetQueryExecution`` returns it. + + Args: + query_id: The query execution ID. + query: The query string. + completion_date_time: The completion time. + schema: The database the query ran against. + + Returns: + The ``QueryExecution`` dictionary. + """ + return { + "QueryExecutionId": query_id, + "Query": query, + "StatementType": AthenaQueryExecution.STATEMENT_TYPE_DML, + "QueryExecutionContext": {"Database": schema}, + "Status": { + "State": AthenaQueryExecution.STATE_SUCCEEDED, + "CompletionDateTime": completion_date_time, + }, + } + + def succeeded_query_execution(query_id, query, completion_date_time, schema="this_schema"): """Build a succeeded DML query execution, as the result cache search lists it. @@ -280,20 +305,129 @@ def succeeded_query_execution(query_id, query, completion_date_time, schema="thi """ return AthenaQueryExecution( { - "QueryExecution": { - "QueryExecutionId": query_id, - "Query": query, - "StatementType": AthenaQueryExecution.STATEMENT_TYPE_DML, - "QueryExecutionContext": {"Database": schema}, - "Status": { - "State": AthenaQueryExecution.STATE_SUCCEEDED, - "CompletionDateTime": completion_date_time, - }, - } + "QueryExecution": succeeded_query_execution_data( + query_id, query, completion_date_time, schema + ) } ) +# Athena's answers to a query or calculation that fails at its first status request. +FAILED_EXECUTION_RESPONSES = { + "StartQueryExecution": {"QueryExecutionId": "query-id"}, + "GetQueryExecution": { + "QueryExecution": { + "QueryExecutionId": "query-id", + "Query": "SELECT 1", + "StatementType": "DML", + "Status": {"State": "FAILED", "StateChangeReason": "failed"}, + } + }, + "StartCalculationExecution": {"CalculationExecutionId": "calculation-id", "State": "CREATING"}, + "GetCalculationExecutionStatus": { + "Status": {"State": "FAILED"}, + "Statistics": {"DpuExecutionInMillis": 0}, + }, + "GetCalculationExecution": { + "CalculationExecutionId": "calculation-id", + "SessionId": "session", + "Status": {"State": "FAILED", "StateChangeReason": "failed"}, + "Statistics": {"DpuExecutionInMillis": 0}, + "Result": {}, + }, +} + + +@contextmanager +def answered_requests(client, responses=FAILED_EXECUTION_RESPONSES): + """Answer each request of a boto3 client once, without AWS, and record them. + + A repeated request raises ``AssertionError``, so that a cursor that keeps + polling stops instead of waiting for ever, and the block then raises it too, + also when the cursor caught it. + + Args: + client: The boto3 client. + responses: The response to each operation name. + + Yields: + The operation name and parameters of each request, in order. + """ + requests = [] + repeated = [] + + def answer(operation, parameters): + if any(operation == requested for requested, _ in requests): + repeated.append(operation) + raise AssertionError(f"repeated {operation} request") + requests.append((operation, parameters)) + return responses[operation] + + with patch.object(client, "_make_api_call", side_effect=answer): + yield requests + assert not repeated, f"repeated requests: {repeated}" + + +class QueryExecutionHistory: + """An Athena client stand-in that lists the query executions of a work group. + + ``ListQueryExecutions`` returns up to ``MaxResults`` IDs per page in the given + order, latest first as Athena lists them, and ``BatchGetQueryExecution`` + returns their executions. + + Attributes: + executions: The ``QueryExecution`` dictionaries, latest first. + list_requests: The parameters of each ``ListQueryExecutions`` request. + """ + + def __init__(self, executions): + """Initialize the history. + + Args: + executions: The ``QueryExecution`` dictionaries, latest first. + """ + self.executions = executions + self.list_requests = [] + + def list_query_executions(self, **request): + """List a page of query execution IDs. + + Args: + **request: The request parameters: ``MaxResults``, the ``NextToken`` + of a later page, and others such as ``WorkGroup``. + + Returns: + The response, with a ``NextToken`` unless the page is the last. + """ + self.list_requests.append(request) + # ListQueryExecutions rejects a MaxResults above 50. + if request["MaxResults"] > 50: + raise ClientError( + {"Error": {"Code": "InvalidRequestException", "Message": "MaxResults"}}, + "ListQueryExecutions", + ) + start = int(request.get("NextToken", 0)) + end = start + request["MaxResults"] + response = { + "QueryExecutionIds": [e["QueryExecutionId"] for e in self.executions[start:end]] + } + if end < len(self.executions): + response["NextToken"] = str(end) + return response + + def batch_get_query_execution(self, **request): + """Get the executions of query execution IDs. + + Args: + **request: The request parameters, with ``QueryExecutionIds``. + + Returns: + The response. + """ + executions = {e["QueryExecutionId"]: e for e in self.executions} + return {"QueryExecutions": [executions[i] for i in request["QueryExecutionIds"]]} + + def unreachable_glue(connection): """A Glue client for the connection that sends requests to a closed proxy port.""" return GlueMetadataClient(