Async Query Client
Asked at Datadog· Onsite
Reported by the problem's creator, not verified
You are given the class QueryClient below. Its API methods are already implemented and scripted by the test; you only implement run.
submit_query(sql)starts a query and returns{"query_id": ...}.get_query_status(query_id)returns the next scripted response:{"status": "RUNNING"}or{"status": "QUEUED"}while the query is in progress;{"status": "SUCCEEDED", "rows": [...], "next_token": ...}with the first page of rows and the token of the next page (next_tokenisNonewhen there is none);{"status": "FAILED", "error": message}; or{"status": "CANCELED"}.fetch_result_page(query_id, token)returns{"rows": [...], "next_token": ...}for one later page.stats()returns[queries_submitted, status_calls, page_fetches].
The tests build the client with the status responses to return in order — the last one repeats once the script runs out — and a dict of pages keyed by token, call run, then call stats.
class QueryClient:
def __init__(self, polls: list[dict], pages: dict[str, dict], max_polls: int) -> None:
self.max_polls = max_polls
self._polls = polls
self._pages = pages
self._submits = 0
self._status_calls = 0
self._page_calls = 0
def submit_query(self, sql: str) -> dict:
"""Starts a query and returns {"query_id": ...}."""
self._submits += 1
return {"query_id": f"q{self._submits}"}
def get_query_status(self, query_id: str) -> dict:
"""The next scripted response; the last one repeats once the script runs out."""
index = min(self._status_calls, len(self._polls) - 1)
self._status_calls += 1
return self._polls[index]
def fetch_result_page(self, query_id: str, token: str) -> dict:
"""The page stored under `token`: {"rows": [...], "next_token": ... or None}."""
self._page_calls += 1
return self._pages[token]
def stats(self) -> list[int]:
"""[queries submitted, status calls made, page fetches made]."""
return [self._submits, self._status_calls, self._page_calls]
def run(self, sql: str) -> list:
raise NotImplementedErrorExamples
Example 1
Input: init =
[[{"status":"RUNNING"},{"rows":[[1,"alice"],[2,"bob"]],"status":"SUCCEEDED","next_token":null}],{},3], operations =[["run",["SELECT id, name FROM users"]],["stats",[]]]Output:
[[[1,"alice"],[2,"bob"]],[1,2,0]]Explanation: The first poll returns RUNNING, so
runpolls again. The second poll returns SUCCEEDED with both rows and no next token, so the rows are returned and no page is fetched.Example 2
Input: init =
[[{"rows":[[1]],"status":"SUCCEEDED","next_token":"p2"}],{"p2":{"rows":[[2],[3]],"next_token":"p3"},"p3":{"rows":[[4]],"next_token":null}},5], operations =[["run",["SELECT id, name FROM users"]],["stats",[]]]Output:
[[[1],[2],[3],[4]],[1,1,2]]Explanation: The SUCCEEDED response carries the first page
[[1]]and tokenp2.runfetchesp2for[[2], [3]]andp3for[[4]], whose next token isNone, and returns the rows in order. The first page is never fetched, so there are 2 page fetches.