Skip to content

Commit 7268255

Browse files
committed
fix(athena): allow spark_connect_max_retries=0 at profile load
The retry-semantics fix turned `max_retries=0` into a valid setting ("no retries, single attempt"), but the profile-load validator at `AthenaCredentials.__post_init__` still rejected anything below 1, so the new value crashed before the runtime could see it. Generalize the validator to a per-field minimum: max_retries allows non-negative; counts / sizes / timeouts stay positive.
1 parent c4dff84 commit 7268255

2 files changed

Lines changed: 84 additions & 10 deletions

File tree

dbt-athena/src/dbt/adapters/athena/connections.py

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,9 @@
5151
from dbt.adapters.sql import SQLConnectionManager
5252

5353

54-
def _is_positive_int(value: Any) -> bool:
54+
def _is_int_at_least(value: Any, minimum: int) -> bool:
5555
try:
56-
return int(value) >= 1
56+
return int(value) >= minimum
5757
except (TypeError, ValueError):
5858
return False
5959

@@ -102,19 +102,22 @@ class AthenaCredentials(Credentials):
102102
def __post_init__(self) -> None:
103103
# Validate Spark Connect integer fields at profile load so a typo
104104
# cannot wait until a python model is submitted to surface.
105-
for field_name in (
106-
"spark_connect_max_sessions",
107-
"spark_connect_session_concurrency",
108-
"spark_connect_dpu_budget",
109-
"spark_connect_pool_acquire_timeout",
110-
"spark_connect_max_retries",
105+
# max_retries allows 0 (= no retries, single attempt); other knobs
106+
# are counts/sizes/timeouts where 0 has no meaning.
107+
for field_name, minimum in (
108+
("spark_connect_max_sessions", 1),
109+
("spark_connect_session_concurrency", 1),
110+
("spark_connect_dpu_budget", 1),
111+
("spark_connect_pool_acquire_timeout", 1),
112+
("spark_connect_max_retries", 0),
111113
):
112114
raw = getattr(self, field_name)
113115
if raw is None:
114116
continue
115-
if not _is_positive_int(raw):
117+
if not _is_int_at_least(raw, minimum):
118+
bound = "non-negative" if minimum == 0 else "positive"
116119
raise DbtRuntimeError(
117-
f"{field_name} must be a positive integer (got {raw!r}). "
120+
f"{field_name} must be a {bound} integer (got {raw!r}). "
118121
"Omit the field to use the default."
119122
)
120123
setattr(self, field_name, int(raw))
Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
import pytest
2+
3+
from dbt.adapters.athena.connections import AthenaCredentials
4+
from dbt_common.exceptions import DbtRuntimeError
5+
6+
from tests.unit import constants
7+
8+
9+
def _make(**overrides):
10+
base = dict(
11+
database=constants.DATA_CATALOG_NAME,
12+
schema=constants.DATABASE_NAME,
13+
s3_staging_dir=constants.S3_STAGING_DIR,
14+
region_name=constants.AWS_REGION,
15+
work_group=constants.ATHENA_WORKGROUP,
16+
spark_work_group=constants.SPARK_WORKGROUP,
17+
)
18+
base.update(overrides)
19+
return AthenaCredentials(**base)
20+
21+
22+
class TestSparkConnectIntegerValidation:
23+
"""Profile-load validation guards against typos in Spark Connect int fields."""
24+
25+
def test_max_retries_zero_is_accepted(self):
26+
# 0 retries = single attempt; allowed semantic since rebuild 15 fix.
27+
c = _make(spark_connect_max_retries=0)
28+
assert c.spark_connect_max_retries == 0
29+
30+
def test_max_retries_positive_is_accepted(self):
31+
c = _make(spark_connect_max_retries=3)
32+
assert c.spark_connect_max_retries == 3
33+
34+
def test_max_retries_negative_is_rejected(self):
35+
with pytest.raises(
36+
DbtRuntimeError, match="spark_connect_max_retries must be a non-negative integer"
37+
):
38+
_make(spark_connect_max_retries=-1)
39+
40+
@pytest.mark.parametrize(
41+
"field_name",
42+
[
43+
"spark_connect_max_sessions",
44+
"spark_connect_session_concurrency",
45+
"spark_connect_dpu_budget",
46+
"spark_connect_pool_acquire_timeout",
47+
],
48+
)
49+
def test_count_fields_reject_zero(self, field_name):
50+
# Counts / sizes / timeouts have no meaning at 0 — keep the strict guard.
51+
with pytest.raises(DbtRuntimeError, match=f"{field_name} must be a positive integer"):
52+
_make(**{field_name: 0})
53+
54+
@pytest.mark.parametrize(
55+
"field_name",
56+
[
57+
"spark_connect_max_sessions",
58+
"spark_connect_session_concurrency",
59+
"spark_connect_dpu_budget",
60+
"spark_connect_pool_acquire_timeout",
61+
],
62+
)
63+
def test_count_fields_accept_one(self, field_name):
64+
c = _make(**{field_name: 1})
65+
assert getattr(c, field_name) == 1
66+
67+
def test_none_is_passed_through(self):
68+
# Explicitly omitted fields stay None so the runtime falls back to defaults.
69+
c = _make()
70+
assert c.spark_connect_max_retries is None
71+
assert c.spark_connect_dpu_budget is None

0 commit comments

Comments
 (0)