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
4 changes: 3 additions & 1 deletion src/dremioai/api/dremio/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -185,10 +185,12 @@ async def get_results(
if client is None:
client = AsyncHttpClient()

delay = settings.instance().dremio.api.polling_interval

endpoint = f"/v0/projects/{project_id}" if project_id else "/api/v3"
job: Job = await client.get(f"{endpoint}/job/{qs.id}", deser=Job)
while not job.done:
await asyncio.sleep(0.5)
await asyncio.sleep(delay)
job = await client.get(f"{endpoint}/job/{qs.id}", deser=Job)

if not job.succeeded:
Expand Down
2 changes: 1 addition & 1 deletion src/dremioai/api/transport.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@
class RetryConfig:
def __init__(self):
if settings.instance() and settings.instance().dremio:
self.config = settings.instance().dremio.http_retry
self.config = settings.instance().dremio.api.http_retry
else:
self.config = settings.HttpRetry()

Expand Down
15 changes: 11 additions & 4 deletions src/dremioai/config/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,7 @@ class HttpRetry(BaseModel):
"""Configuration for HTTP retry behavior with exponential backoff"""

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


class ApiSettings(BaseModel):
# HTTP retry configuration
http_retry: Optional[HttpRetry] = Field(default_factory=HttpRetry)
polling_interval: Optional[float] = Field(
default=1, description="Polling interval for REST api in seconds"
)
model_config = ConfigDict(validate_assignment=True)


class Dremio(BaseModel):
uri: Annotated[
Union[str, HttpUrl, DremioCloudUri], AfterValidator(_resolve_dremio_uri)
Expand All @@ -177,10 +186,8 @@ class Dremio(BaseModel):
allow_dml: Optional[bool] = False
auth_issuer_uri_override: Optional[str] = None
wlm: Optional[Wlm] = None
# Metrics server configuration
api: Optional[ApiSettings] = Field(default_factory=ApiSettings)
metrics: Optional[Metrics] = None
# HTTP retry configuration
http_retry: Optional[HttpRetry] = Field(default_factory=HttpRetry)
model_config = ConfigDict(validate_assignment=True)

@field_serializer("raw_pat")
Expand Down
5 changes: 3 additions & 2 deletions tests/api/test_transport_retry.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ def test_retry_config_with_default_settings(self):
retry_config = RetryConfig()

# Verify default values from HttpRetry model
assert retry_config.config.max_retries == 3
assert retry_config.config.max_retries == 20
assert retry_config.config.initial_delay == 1.0
assert retry_config.config.max_delay == 60.0
assert retry_config.config.backoff_multiplier == 2.0
Expand Down Expand Up @@ -300,7 +300,8 @@ def mock_settings_instance():
)

mock_dremio = MagicMock()
mock_dremio.http_retry = http_retry_config
mock_dremio.api = MagicMock()
mock_dremio.api.http_retry = http_retry_config

# Create mock settings object
mock_settings = MagicMock()
Expand Down