Skip to content

Commit f5ce4ae

Browse files
Support configurable delay for SQL REST API polling interval (#71)
* Support configurable delay for SQL REST API polling interval * fix test --------- Co-authored-by: eduardo-munoz-dremio <eduardo.munoz@dremio.com>
1 parent 732a488 commit f5ce4ae

4 files changed

Lines changed: 18 additions & 8 deletions

File tree

src/dremioai/api/dremio/sql.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -185,10 +185,12 @@ async def get_results(
185185
if client is None:
186186
client = AsyncHttpClient()
187187

188+
delay = settings.instance().dremio.api.polling_interval
189+
188190
endpoint = f"/v0/projects/{project_id}" if project_id else "/api/v3"
189191
job: Job = await client.get(f"{endpoint}/job/{qs.id}", deser=Job)
190192
while not job.done:
191-
await asyncio.sleep(0.5)
193+
await asyncio.sleep(delay)
192194
job = await client.get(f"{endpoint}/job/{qs.id}", deser=Job)
193195

194196
if not job.succeeded:

src/dremioai/api/transport.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@
4343
class RetryConfig:
4444
def __init__(self):
4545
if settings.instance() and settings.instance().dremio:
46-
self.config = settings.instance().dremio.http_retry
46+
self.config = settings.instance().dremio.api.http_retry
4747
else:
4848
self.config = settings.HttpRetry()
4949

src/dremioai/config/settings.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,7 @@ class HttpRetry(BaseModel):
147147
"""Configuration for HTTP retry behavior with exponential backoff"""
148148

149149
max_retries: Optional[int] = Field(
150-
default=3,
150+
default=20,
151151
description="Maximum number of retry attempts for rate-limited requests",
152152
)
153153
initial_delay: Optional[float] = Field(
@@ -162,6 +162,15 @@ class HttpRetry(BaseModel):
162162
model_config = ConfigDict(validate_assignment=True)
163163

164164

165+
class ApiSettings(BaseModel):
166+
# HTTP retry configuration
167+
http_retry: Optional[HttpRetry] = Field(default_factory=HttpRetry)
168+
polling_interval: Optional[float] = Field(
169+
default=1, description="Polling interval for REST api in seconds"
170+
)
171+
model_config = ConfigDict(validate_assignment=True)
172+
173+
165174
class Dremio(BaseModel):
166175
uri: Annotated[
167176
Union[str, HttpUrl, DremioCloudUri], AfterValidator(_resolve_dremio_uri)
@@ -177,10 +186,8 @@ class Dremio(BaseModel):
177186
allow_dml: Optional[bool] = False
178187
auth_issuer_uri_override: Optional[str] = None
179188
wlm: Optional[Wlm] = None
180-
# Metrics server configuration
189+
api: Optional[ApiSettings] = Field(default_factory=ApiSettings)
181190
metrics: Optional[Metrics] = None
182-
# HTTP retry configuration
183-
http_retry: Optional[HttpRetry] = Field(default_factory=HttpRetry)
184191
model_config = ConfigDict(validate_assignment=True)
185192

186193
@field_serializer("raw_pat")

tests/api/test_transport_retry.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ def test_retry_config_with_default_settings(self):
3737
retry_config = RetryConfig()
3838

3939
# Verify default values from HttpRetry model
40-
assert retry_config.config.max_retries == 3
40+
assert retry_config.config.max_retries == 20
4141
assert retry_config.config.initial_delay == 1.0
4242
assert retry_config.config.max_delay == 60.0
4343
assert retry_config.config.backoff_multiplier == 2.0
@@ -300,7 +300,8 @@ def mock_settings_instance():
300300
)
301301

302302
mock_dremio = MagicMock()
303-
mock_dremio.http_retry = http_retry_config
303+
mock_dremio.api = MagicMock()
304+
mock_dremio.api.http_retry = http_retry_config
304305

305306
# Create mock settings object
306307
mock_settings = MagicMock()

0 commit comments

Comments
 (0)