Skip to content
Draft
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
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