Skip to content
Merged
Show file tree
Hide file tree
Changes from 7 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
2 changes: 1 addition & 1 deletion src/dremioai/api/dremio/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -185,7 +185,7 @@ async def get_results(
if client is None:
client = AsyncHttpClient()

delay = settings.instance().dremio.api.polling_interval
delay = settings.instance().dremio.api.get("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)
Expand Down
22 changes: 14 additions & 8 deletions src/dremioai/api/transport.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,18 +43,18 @@
class RetryConfig:
def __init__(self):
if settings.instance() and settings.instance().dremio:
self.config = settings.instance().dremio.api.http_retry
self.http_retry = settings.instance().dremio.api.http_retry
else:
self.config = settings.HttpRetry()
self.http_retry = settings.HttpRetry()

@property
def max_retries(self) -> int:
"""Expose max_retries from config for convenience"""
return self.config.max_retries
return self.http_retry.get("max_retries")

def get_config_delay(self, attempt_number: int = 0) -> float:
return self.config.initial_delay * (
self.config.backoff_multiplier**attempt_number
return self.http_retry.get("initial_delay") * (
self.http_retry.get("backoff_multiplier") ** attempt_number
)

def get_delay(
Expand All @@ -72,7 +72,7 @@ def get_delay(
f"Invalid Retry-After header, using exponential backoff - {e}"
)

return min(delay, self.config.max_delay)
return min(delay, self.http_retry.get("max_delay"))


async def retry_middleware(
Expand Down Expand Up @@ -200,7 +200,11 @@ async def post(
async with ClientSession(middlewares=(retry_middleware,)) as session:
self.log_request("POST", endpoint, params)
async with session.post(
f"{self.uri}{endpoint}", params=params, headers=self.headers, json=body, ssl=False
f"{self.uri}{endpoint}",
params=params,
headers=self.headers,
json=body,
ssl=False,
) as response:
return await self.handle_response(
response, deser, file, top_level_list=top_level_list
Expand All @@ -222,5 +226,7 @@ def __init__(self):
pat = dremio.pat

if uri is None or pat is None:
raise RuntimeError("Dremio connection is not properly configured. Both URI and authentication token are required.")
raise RuntimeError(
"Dremio connection is not properly configured. Both URI and authentication token are required."
)
super().__init__(uri, pat)
Loading
Loading