Coverage for adrs/data/datasource.py: 59%
203 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-29 15:31 +0800
1import httpx
2import asyncio
3import logging
4import pandas as pd
5import polars as pl
6import clickhouse_connect
8from typing import cast
9from abc import abstractmethod
10from datetime import datetime, timedelta, timezone
12from adrs.types import Data, Topic, SortedDataList
13from adrs.data.progress import inner_task
15logger = logging.getLogger(__name__)
17DEFAULT_API_URL = "https://api.datasource.cybotrade.rs"
20def ms_to_dt(ms: int) -> datetime:
21 return datetime.fromtimestamp(ms / 1000, tz=timezone.utc)
24class Datasource:
25 @abstractmethod
26 async def query_paginated(
27 self,
28 topic: Topic | str,
29 start_time: datetime | None = None,
30 end_time: datetime | None = None,
31 limit: int | None = None,
32 flatten: bool = False,
33 ) -> pl.DataFrame:
34 """Paginate across multiple query() calls, returning all rows as one DataFrame.
36 Resolves missing args from topic.params before dispatching.
37 - Range mode (start_time + end_time): pages forward until end_time.
38 - Limit mode (limit only): pages backward from last closed candle.
40 Returns DataFrame with start_time as Datetime("ms", "UTC").
41 Raises ValueError if neither (start_time, end_time) nor limit can be resolved,
42 or if topic has no interval.
43 """
44 raise NotImplementedError()
47class CybotradeDatasource(Datasource):
48 def __init__(
49 self,
50 api_key: str,
51 base_url: str = DEFAULT_API_URL,
52 max_limit: int = 100_000,
53 max_retries: int = 3,
54 backoff_base: float = 2.0,
55 ):
56 self.api_key = api_key
57 self.base_url = base_url.rstrip("/")
58 self.max_limit = max_limit
59 self.max_retries = max_retries
60 self.backoff_base = backoff_base
61 self._client = httpx.AsyncClient(timeout=120.0)
63 async def query(
64 self,
65 topic: Topic | str,
66 start_time: datetime | None = None,
67 end_time: datetime | None = None,
68 limit: int | None = None,
69 flatten: bool = False,
70 ) -> list[Data]:
71 if isinstance(topic, str):
72 topic = Topic.from_str(topic)
73 params: dict[str, str] = dict(topic.params)
74 if start_time is not None:
75 params["start_time"] = str(int(start_time.timestamp() * 1000))
76 if end_time is not None:
77 params["end_time"] = str(int(end_time.timestamp() * 1000))
78 if limit is not None:
79 params["limit"] = str(limit)
80 if flatten:
81 params["flatten"] = "true"
83 url = f"{self.base_url}/{topic.provider}/{topic.endpoint}"
84 last_exc: Exception | None = None
85 for attempt in range(self.max_retries):
86 try:
87 resp = await self._client.get(
88 url, params=params, headers={"X-API-KEY": self.api_key}
89 )
90 resp.raise_for_status()
91 batch = resp.json()["data"]
92 return cast(
93 list[Data],
94 [
95 {**item, "start_time": ms_to_dt(item["start_time"])}
96 for item in batch
97 ],
98 )
99 except (httpx.TimeoutException, httpx.ConnectError) as e:
100 last_exc = e
101 wait = self.backoff_base**attempt
102 logger.warning(
103 "[%s] attempt %d/%d failed (%s), retrying in %.1fs",
104 topic,
105 attempt + 1,
106 self.max_retries,
107 type(e).__name__,
108 wait,
109 )
110 await asyncio.sleep(wait)
111 raise last_exc # type: ignore[misc]
113 async def query_paginated(
114 self,
115 topic: Topic | str,
116 start_time: datetime | None = None,
117 end_time: datetime | None = None,
118 limit: int | None = None,
119 flatten: bool = False,
120 ) -> pl.DataFrame:
121 if isinstance(topic, str):
122 topic = Topic.from_str(topic)
124 if start_time is None:
125 if (st := topic.params.get("start_time")) is not None:
126 start_time = ms_to_dt(int(st))
127 if end_time is None:
128 if (st := topic.params.get("end_time")) is not None:
129 end_time = ms_to_dt(int(st))
130 if limit is None:
131 if (lim := topic.params.get("limit")) is not None:
132 limit = int(lim)
134 if (start_time is None or end_time is None) and limit is None:
135 raise ValueError(
136 "either (start_time, end_time) or (limit) must be specified"
137 )
138 interval = topic.interval()
139 if interval is None:
140 raise ValueError(f"{topic} does not include an interval")
141 interval_ms = int(interval.total_seconds() * 1000)
142 # Pages are kept as separate frames and stitched once at the end. Folding
143 # each page into a running accumulator instead would re-materialise the
144 # whole history per page — quadratic, and the reason a 1m topic used to
145 # spike memory during a long backfill.
146 pages: list[pl.DataFrame] = []
148 if start_time and end_time:
149 end_ms = (int(end_time.timestamp() * 1000) // interval_ms) * interval_ms
150 inner_total: int | None = (
151 None
152 if topic.is_block()
153 else max(
154 0,
155 (end_ms - int(start_time.timestamp() * 1000)) // interval_ms,
156 )
157 )
158 else:
159 inner_total = None if topic.is_block() else limit
161 async with inner_task(str(topic), inner_total) as _bar:
162 if start_time and end_time:
163 # pyrefly: ignore [unbound-name]
164 end_time = ms_to_dt(end_ms)
166 current_start = start_time
167 total = (end_ms - int(start_time.timestamp() * 1000)) // interval_ms
168 iter_n = 1
170 while current_start < end_time:
171 current_limit = (
172 self.max_limit
173 if topic.is_block()
174 else min(self.max_limit, total)
175 )
176 current_end_check = current_start + timedelta(
177 milliseconds=current_limit * interval_ms
178 )
180 logger.debug(
181 "[query %d] fetching %s (start_time: %s, limit: %d)",
182 iter_n,
183 topic,
184 current_start,
185 current_limit,
186 )
187 resp = await self.query(
188 topic,
189 start_time=current_start,
190 limit=current_limit,
191 flatten=flatten,
192 )
193 if len(resp) == 0: # stop pagination if no new data
194 break
196 if "start_time" not in resp[0].keys():
197 raise ValueError(
198 f"{type(self).__name__}.query() does not return 'start_time' in its records"
199 )
200 if not isinstance(resp[0]["start_time"], datetime):
201 raise TypeError(
202 f"'start_time' must be tz-aware datetime, got {type(resp[0]['start_time'])}"
203 )
204 num = len(resp)
205 _bar.advance(num)
206 logger.debug("[query %d] %s got %d datapoints", iter_n, topic, num)
208 page_end = max(r["start_time"] for r in resp)
209 pages.append(pl.DataFrame(resp, infer_schema_length=None))
211 if topic.is_block():
212 current_start = page_end + timedelta(seconds=1)
213 else:
214 total = max(0, total - num)
215 current_start = page_end + interval
216 if total == 0:
217 break
219 iter_n += 1
221 if (
222 not topic.is_block()
223 and current_end_check >= end_time
224 and num < current_limit
225 ):
226 break
227 elif topic.is_block() and num < (current_limit - 1):
228 break
230 await asyncio.sleep(0.1)
232 if topic.is_block() and pages:
233 block_window = pl.col("start_time").is_between(
234 start_time, end_time, closed="left"
235 )
236 pages = [page.filter(block_window) for page in pages]
237 else:
238 # paginate backwards from last_closed_time
239 current_end = topic.last_closed_time(is_collect=False)
240 remaining = limit
241 iter_n = 1
243 while remaining > 0:
244 current_limit = min(self.max_limit, remaining)
245 logger.debug(
246 "[query %d] fetching %s (end_time: %s, limit: %d)",
247 iter_n,
248 topic,
249 current_end,
250 current_limit,
251 )
252 resp = await self.query(
253 topic,
254 end_time=current_end,
255 limit=current_limit,
256 flatten=flatten,
257 )
258 num = len(resp)
259 _bar.advance(num)
260 logger.debug("[query %d] %s got %d datapoints", iter_n, topic, num)
262 page_start = min(r["start_time"] for r in resp)
263 remaining -= num
264 current_end = (
265 page_start - timedelta(seconds=1)
266 if topic.is_block()
267 else page_start - interval
268 )
269 pages.append(pl.DataFrame(resp, infer_schema_length=None))
270 iter_n += 1
272 if num < current_limit:
273 break
274 await asyncio.sleep(0.1)
276 pages = [page for page in pages if page.height]
277 if not pages:
278 return pl.DataFrame()
280 # Pages are appended in fetch order, so keeping the last row for a
281 # start_time preserves the old per-page merge rule: later wins.
282 # maintain_order pins which duplicate that is rather than leaving it to
283 # whichever thread finishes first.
284 stitched = (
285 pl.concat(pages, how="diagonal_relaxed")
286 .unique(subset="start_time", keep="last", maintain_order=True)
287 .sort("start_time")
288 )
289 # Round-trip through the store so start_time lands as Datetime("ms", "UTC").
290 return SortedDataList.from_df(stitched).to_df()
293class ClickhouseDatasource(Datasource):
294 def __init__(self, **kwargs):
295 self.connect_params = kwargs
296 self.ch = None
298 async def _client(self) -> clickhouse_connect.driver.AsyncClient:
299 if self.ch is None:
300 self.ch = await clickhouse_connect.get_async_client(**self.connect_params)
301 return self.ch
303 def _make_query_params_filter(self, topic: Topic):
304 if topic.provider == "bybit" and topic.endpoint == "double-win":
305 return f"query_params LIKE '%interval={topic.params.get('interval')}&product={topic.params.get('currency')}%'"
306 else:
307 return f"query_params = '{topic.query_params_str()}'"
309 def _make_query(self, topic: Topic, start_time: datetime, end_time: datetime):
310 return f"""
311 SELECT timestamp AS start_time, data FROM `{topic.provider}`.`{topic.endpoint}` FINAL
312 WHERE {self._make_query_params_filter(topic)}
313 AND timestamp >= '{start_time.strftime("%Y-%m-%d %H:%M:%S")}'
314 AND timestamp < '{end_time.strftime("%Y-%m-%d %H:%M:%S")}'
315 ORDER BY timestamp ASC
316 """
318 # Empirically-set chunk size (not a config knob elsewhere): a single
319 # unpaginated multi-year query against a multi-record-per-timestamp
320 # topic (full option chains, GEX snapshots -- ~800 nested structs per
321 # row) overloaded clickhouse-connect's streaming parser
322 # (StreamFailureError: "unrecognized data found in stream", ~5min in,
323 # against a 7-year range). A single month (~744 rows * ~800 structs)
324 # completed reliably in the same test. Chunking unconditionally (not
325 # just for detected multi-record topics) keeps this simple and safe for
326 # any topic, at the cost of a few extra round-trips for small
327 # single-record ranges that were already fine unchunked.
328 _MAX_CHUNK_DAYS = 31
330 async def query_paginated(
331 self,
332 topic: Topic | str,
333 start_time: datetime | None = None,
334 end_time: datetime | None = None,
335 limit: int | None = None,
336 flatten: bool = False,
337 ) -> pl.DataFrame:
338 if isinstance(topic, str):
339 topic = Topic.from_str(topic)
340 if start_time is None or end_time is None:
341 raise ValueError("start_time and end_time must be provided")
343 if end_time - start_time <= timedelta(days=self._MAX_CHUNK_DAYS):
344 return await self._query_one_range(topic, start_time, end_time)
346 frames = []
347 cur = start_time
348 while cur < end_time:
349 nxt = min(cur + timedelta(days=self._MAX_CHUNK_DAYS), end_time)
350 chunk = await self._query_one_range(topic, cur, nxt)
351 if len(chunk):
352 frames.append(chunk)
353 cur = nxt
354 return pl.concat(frames, how="vertical") if frames else pl.DataFrame()
356 async def _query_one_range(
357 self, topic: Topic, start_time: datetime, end_time: datetime
358 ) -> pl.DataFrame:
359 ch = await self._client()
361 # clickhouse-connect for pyarrow doesn't support JSON column yet,
362 # see: https://github.com/ClickHouse/clickhouse-connect/issues/398
363 # df = await ch.query_df_arrow(
364 # self._make_query(topic, start_time, end_time), dataframe_library="polars"
365 # )
367 pandas_df = await ch.query_df(self._make_query(topic, start_time, end_time))
368 if pandas_df.empty:
369 return pl.DataFrame()
371 # Multi-record-per-timestamp topics (full option chains, GEX
372 # snapshots): "data" is a LIST of per-instrument dicts, not one
373 # flat dict. json_normalize on a list-per-cell Series indexes
374 # positionally ("0".."N" columns) -- meaningless here since column
375 # identity isn't stable across snapshots (chain composition
376 # changes every hour), and it drops the "data" column name that
377 # DataColumn(src="data", ...) needs downstream. Keep the list
378 # intact; normalize numeric fields to float so polars can infer
379 # one consistent Struct schema across all records instead of
380 # choking on mixed int/float typing for the same field across
381 # instruments.
382 #
383 # ClickHouse's raw storage preserves the source API's full
384 # per-timestamp envelope, which for multi-record topics is itself
385 # `{"data": [...]}` -- one extra wrapper level versus what the
386 # cybotrade REST proxy client hands back (it unwraps `resp.json()
387 # ["data"]` itself before this code ever sees it). Verified directly
388 # via a raw SQL query against the real flow.`options/gamma-exposure`
389 # table: cells are dicts of the form {"data": [<840 instrument
390 # dicts>]}, not the list directly. Unwrap that one level, if present,
391 # before checking for the list shape -- a no-op for any topic where
392 # "data" is already the list itself.
393 def _unwrap(v):
394 if (
395 isinstance(v, dict)
396 and set(v.keys()) == {"data"}
397 and isinstance(v["data"], list)
398 ):
399 return v["data"]
400 return v
402 first_non_null = _unwrap(
403 next((v for v in pandas_df["data"] if v is not None), None)
404 )
405 is_multi_record = isinstance(first_non_null, list)
407 if is_multi_record:
409 def _normalize_record(rec: dict) -> dict:
410 # Two independent sources of mixed typing for the SAME
411 # struct field across different records, both verified
412 # directly against real GEX data, not assumed:
413 # 1. isinstance(v, (int, float)) doesn't catch numpy scalar
414 # types (np.int64, np.float64) that clickhouse-connect's
415 # driver returns for numeric JSON fields.
416 # 2. The underlying JSON itself is inconsistently typed --
417 # the SAME field (e.g. dealerTotalInventory) arrives as a
418 # JSON string in some records and a JSON number in
419 # others. A blanket "leave strings alone" rule preserves
420 # that inconsistency instead of fixing it.
421 # Try a numeric cast on everything except None/bool; genuine
422 # text fields (currency, exchange, putCall,
423 # instrumentNormalized) simply fail the cast and keep their
424 # original string value.
425 out = {}
426 for k, v in rec.items():
427 if v is None or isinstance(v, bool):
428 out[k] = v
429 continue
430 try:
431 out[k] = float(v)
432 except (TypeError, ValueError):
433 # Some numeric fields serialize missing values as
434 # the literal JSON string "null" (verified: real
435 # data, not a true JSON null) instead of an actual
436 # null -- treat those as None rather than leaving a
437 # stray string in an otherwise-numeric field.
438 if isinstance(v, str) and v.strip().lower() in (
439 "null",
440 "none",
441 "nan",
442 ):
443 out[k] = None
444 else:
445 out[k] = v
446 return out
448 normalized = pandas_df["data"].map(
449 lambda cell: [_normalize_record(r) for r in _unwrap(cell)]
450 if cell is not None
451 else cell
452 )
453 pandas_df = pandas_df.drop(columns=["data"])
454 else:
455 pandas_df = pandas_df.drop(columns=["data"]).join(
456 # pyrefly: ignore [bad-argument-type]
457 pd.json_normalize(pandas_df["data"])
458 )
460 for col in pandas_df.columns:
461 # Fix dynamic 'object' columns containing mixed PyArrow-unfriendly types
462 if pandas_df[col].dtype == "object":
463 # Check if it's primarily numeric despite being labeled an 'object'
464 converted_numeric = pd.to_numeric(pandas_df[col], errors="coerce")
465 if not converted_numeric.isna().all():
466 pandas_df[col] = converted_numeric
467 else:
468 # If it's mixed text/strings/None, force everything to string safely
469 pandas_df[col] = (
470 pandas_df[col]
471 .astype(str)
472 .replace("None", None)
473 .replace("nan", None)
474 )
476 df = pl.from_pandas(pandas_df)
477 if is_multi_record:
478 # Build "data" via pl.Series directly from the normalized
479 # Python list rather than pl.from_pandas' generic object-column
480 # handling -- polars' own type inference over a list of dicts
481 # reliably produces List(Struct(...)) once values are
482 # type-consistent (verified against real GEX data); pandas'
483 # pyarrow-conversion path for nested Python objects is not.
484 df = df.with_columns(pl.Series("data", normalized.tolist()))
486 # Enforce the Datasource contract (start_time as Datetime("ms",
487 # "UTC")) the same way CybotradeDatasource does via SortedDataList.
488 # ClickHouse's plain `DateTime` column type carries no explicit tz,
489 # so clickhouse-connect -> pandas -> pl.from_pandas produces a
490 # tz-naive start_time here -- without this, cache.py's pl.concat
491 # of old (Cybotrade-sourced, tz-aware) and new (Clickhouse-sourced,
492 # tz-naive) cached shards for the same series raises SchemaError.
493 df = df.with_columns(
494 pl.col("start_time")
495 .dt.replace_time_zone(time_zone="UTC")
496 .dt.cast_time_unit(time_unit="ms")
497 )
499 return df