Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 59 additions & 65 deletions tests/pyathena/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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]
Expand All @@ -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"),
[
Expand All @@ -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",
Expand All @@ -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):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Round one, deferred with reasons:

  • pytest.raises(TypeError) without match="positional": the message is CPython's, which docs/testing.md does not count as a contract. Mutations that accept a positional setting in Connection, Cursor, AsyncCursor, DictCursor, and Cursor.execute() fail the tests.
  • test_cursor_execute_keyword_only opens a cursor per class (about 3 s offline). A real call is kept as the evidence: the positional call reaches AWS or a future in the mutated case.
  • test_binary_parameter loses its seven test IDs: Remove and rewrite unnecessary tests in the top-level cursor, connection, converter, and utility tests #1243 rewrites it to one query.

factory("s3://bucket/path/")


class RecordingCursor(Cursor):
Expand Down Expand Up @@ -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")

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. P2: no remaining synchronous test calls execute(operation=...). Rejected: docs/usage.md documents the operation and parameters as positional and the other settings as keyword-only; it does not document the keyword form, and Remove and rewrite unnecessary tests in the top-level cursor, connection, converter, and utility tests #1243 rewrites the test to the rejection of a positional setting.

Also repaired in this commit, from the #1245 work: test_connect_keyword_only passes one positional argument. With two, a factory that accepted one positional setting still raised TypeError for the second.



@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'"):
Expand All @@ -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),

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. P2: the dictionary stub answered GetQueryExecution for ever, so a cursor that kept polling would hang. Repaired: answered_requests() in tests/pyathena/util.py answers each operation once; a repeat raises AssertionError, and the block raises it again at exit when the cursor caught it. Treating FAILED as not terminal fails test_cursor_accepts_connection_callbacks[Cursor] in about 1 s.

_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(
Expand Down Expand Up @@ -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)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. P2: removing the lock still passed if one thread finished the build before the others entered. Repaired: the first build waits up to 1 s for a second build, so without the lock a second thread builds meanwhile; with the lock no second build happens and the wait times out. Removing the S3 or the Glue lock fails the tests.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up (relayed): FINDINGS. Without the lock, the 1 s wait for a second build still passes if every other released thread stays descheduled until the first client is cached. Deferred with reasons: without the private lock, the test cannot observe contention inside the property, and #1243 removes the private _s3_client_lock.locked() / glue._lock.locked() checks. The test now misses a missing lock only if all seven released threads stay descheduled for a full second; removing either lock failed it in each offline run.

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(_):
Expand All @@ -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"

Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading