Add adaptive query utilities and parallel scan

Introduce slicks.query_utils with error classification, adaptive_query (recursive split/fallback) and run_chunks_parallel (threaded chunk execution with client_factory). Refactor discovery.discover_sensors and scanner._fetch_bins_adaptive to use these utilities, add progress bar handling, per-thread InfluxDBClient3 creation/closing, and surface permanent errors as RuntimeError. Add comprehensive unit tests for discovery, query_utils and update scanner tests to match parallel client usage.

Author: haoruizhou

src/slicks/discovery.py

@@ -1,55 +1,116 @@
-from datetime import timedelta
+"""
+Sensor discovery module.
+
+Scans the database for all unique sensor names within a time range.
+Uses adaptive chunking with parallel execution.
+"""
+
+from __future__ import annotations
+
+import threading
+from datetime import datetime, timedelta
+from typing import List, Optional
+
+from influxdb_client_3 import InfluxDBClient3
+from tqdm.auto import tqdm
+
 from . import config
-from .fetcher import get_influx_client
+from .query_utils import adaptive_query, run_chunks_parallel, PermanentQueryError
 
-def discover_sensors(start_time, end_time, chunk_size_days=1, client=None):
+
+def discover_sensors(
+    start_time: datetime,
+    end_time: datetime,
+    chunk_size_days: int = 7,
+    client: Optional[InfluxDBClient3] = None,
+    show_progress: bool = True,
+) -> List[str]:
     """
-    Scans the database for ALL unique sensor names within the time range.
-    Uses recursive splitting to handle server resource limits.
+    Scan the database for ALL unique sensor names within the time range.
+
+    Uses adaptive chunking with parallel execution to handle server
+    resource limits efficiently.
+
+    Args:
+        start_time: Start of scan range.
+        end_time: End of scan range.
+        chunk_size_days: Days per chunk (default 7).
+        client: Ignored (kept for backward compatibility).
+        show_progress: Show progress bar (default True).
+
+    Returns:
+        Sorted list of unique sensor name strings.
     """
-    if client is None:
-        client = get_influx_client()
-
-    unique_sensors = set()
-    
-    def _scan_recursive(start, end, depth=0):
-        # Stop recursion if interval is too small (< 10 seconds) or depth too high
-        if (end - start).total_seconds() < 10 or depth > 5:
-            # print(f"  Skipping small/deep chunk: {start} to {end}")
-            return
-
-        query = f"""
+
+    def _make_client() -> InfluxDBClient3:
+        return InfluxDBClient3(
+            host=config.INFLUX_URL,
+            token=config.INFLUX_TOKEN,
+            database=config.INFLUX_DB,
+        )
+
+    def _query_distinct(
+        client: InfluxDBClient3, t0: datetime, t1: datetime,
+    ) -> List[str]:
+        sql = f"""
         SELECT DISTINCT "signalName"
         FROM "iox"."{config.INFLUX_DB}"
-        WHERE time >= '{start.isoformat()}Z'
-        AND time < '{end.isoformat()}Z'
+        WHERE time >= '{t0.isoformat()}Z'
+        AND time < '{t1.isoformat()}Z'
         """
-        
-        try:
-            # print(f"Scanning {start} -> {end} (Depth {depth})...")
-            table = client.query(query=query, mode="all")
-            df = table.to_pandas()
-            
-            if not df.empty and "signalName" in df.columns:
-                batch_sensors = set(df["signalName"].unique())
-                unique_sensors.update(batch_sensors)
-                
-        except Exception as e:
-            # print(f"  Chunk failed ({e}). Splitting...")
-            mid_point = start + (end - start) / 2
-            _scan_recursive(start, mid_point, depth + 1)
-            _scan_recursive(mid_point, end, depth + 1)
-
-    print(f"Discovering sensors from {start_time} to {end_time}...")
-    current = start_time
-    while current < end_time:
-        next_step = min(current + timedelta(days=chunk_size_days), end_time)
-        if next_step <= current: break
-        
-        # Start recursion for this chunk
-        _scan_recursive(current, next_step)
-        current = next_step
-
-    sorted_sensors = sorted(list(unique_sensors))
-    print(f"Discovery Complete. Found {len(sorted_sensors)} unique sensors.")
-    return sorted_sensors
+        table = client.query(query=sql)
+        if table.num_rows == 0:
+            return []
+        col = table.column("signalName")
+        return [v.as_py() for v in col if v.as_py() is not None]
+
+    def _process_chunk(
+        client: InfluxDBClient3, t0: datetime, t1: datetime,
+    ) -> List[str]:
+        return adaptive_query(
+            client=client,
+            t0=t0,
+            t1=t1,
+            primary_fn=_query_distinct,
+            fallback_fn=None,
+            min_span=timedelta(seconds=10),
+            max_depth=5,
+        )
+
+    # Build chunk list
+    chunks = []
+    cur = start_time
+    while cur < end_time:
+        nxt = min(cur + timedelta(days=chunk_size_days), end_time)
+        if nxt <= cur:
+            break
+        chunks.append((cur, nxt))
+        cur = nxt
+
+    pbar = tqdm(
+        total=len(chunks),
+        desc="Discovering sensors",
+        unit="chunk",
+        disable=not show_progress,
+    )
+    pbar_lock = threading.Lock()
+
+    def on_chunk_done(idx: int) -> None:
+        with pbar_lock:
+            pbar.update(1)
+
+    try:
+        all_names = run_chunks_parallel(
+            client_factory=_make_client,
+            chunks=chunks,
+            query_fn=_process_chunk,
+            max_workers=4,
+            on_chunk_done=on_chunk_done,
+        )
+    except PermanentQueryError as e:
+        raise RuntimeError(f"Sensor discovery aborted: {e}") from e
+    finally:
+        pbar.close()
+
+    unique = sorted(set(all_names))
+    return unique

src/slicks/query_utils.py

@@ -0,0 +1,165 @@
+"""
+Shared utilities for adaptive chunked querying against InfluxDB 3.x (IOx).
+
+Provides:
+- Error classification (recoverable vs permanent)
+- Parallel chunk execution via ThreadPoolExecutor
+- Adaptive recursive splitting on resource-limit failures
+"""
+
+from __future__ import annotations
+
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime, timedelta
+from typing import Callable, List, Optional, Sequence, Tuple, TypeVar
+
+from influxdb_client_3 import InfluxDBClient3
+
+T = TypeVar("T")
+
+# ---------------------------------------------------------------------------
+# Error classification
+# ---------------------------------------------------------------------------
+
+_PERMANENT_ERROR_PATTERNS = (
+    "table not found",
+    "not found",
+    "unauthorized",
+    "unauthenticated",
+    "permission denied",
+    "invalid token",
+    "database not found",
+    "bucket not found",
+    "syntax error",
+)
+
+
+class PermanentQueryError(Exception):
+    """An error that will not resolve by splitting the time range."""
+
+
+def is_permanent_error(exc: Exception) -> bool:
+    """Classify an exception as permanent (non-retryable) vs recoverable."""
+    msg = str(exc).lower()
+    return any(pattern in msg for pattern in _PERMANENT_ERROR_PATTERNS)
+
+
+# ---------------------------------------------------------------------------
+# Adaptive recursive query
+# ---------------------------------------------------------------------------
+
+def adaptive_query(
+    client: InfluxDBClient3,
+    t0: datetime,
+    t1: datetime,
+    primary_fn: Callable[[InfluxDBClient3, datetime, datetime], List[T]],
+    fallback_fn: Optional[Callable[[InfluxDBClient3, datetime, datetime], List[T]]] = None,
+    min_span: Optional[timedelta] = None,
+    max_depth: int = 10,
+    _depth: int = 0,
+) -> List[T]:
+    """
+    Execute *primary_fn* on [t0, t1).  On a recoverable failure the range is
+    split in half and each half is retried recursively.
+
+    When the remaining span is smaller than *min_span* (or *max_depth* is
+    reached) *fallback_fn* is used instead — if provided — otherwise an empty
+    list is returned.
+
+    Raises ``PermanentQueryError`` immediately for non-retryable errors such
+    as authentication failures or missing tables.
+    """
+    if min_span and (t1 - t0) <= min_span:
+        if fallback_fn:
+            return fallback_fn(client, t0, t1)
+        return []
+
+    if _depth > max_depth:
+        if fallback_fn:
+            return fallback_fn(client, t0, t1)
+        return []
+
+    try:
+        return primary_fn(client, t0, t1)
+    except Exception as exc:
+        if is_permanent_error(exc):
+            raise PermanentQueryError(str(exc)) from exc
+
+        mid = t0 + (t1 - t0) / 2
+        if mid <= t0 or mid >= t1:
+            if fallback_fn:
+                return fallback_fn(client, t0, t1)
+            return []
+
+        left = adaptive_query(
+            client, t0, mid, primary_fn, fallback_fn,
+            min_span, max_depth, _depth + 1,
+        )
+        right = adaptive_query(
+            client, mid, t1, primary_fn, fallback_fn,
+            min_span, max_depth, _depth + 1,
+        )
+        return left + right
+
+
+# ---------------------------------------------------------------------------
+# Parallel chunk execution
+# ---------------------------------------------------------------------------
+
+def run_chunks_parallel(
+    client_factory: Callable[[], InfluxDBClient3],
+    chunks: Sequence[Tuple[datetime, datetime]],
+    query_fn: Callable[[InfluxDBClient3, datetime, datetime], List[T]],
+    max_workers: int = 4,
+    on_chunk_done: Optional[Callable[[int], None]] = None,
+) -> List[T]:
+    """
+    Execute *query_fn* across time-range *chunks* in parallel.
+
+    Each worker thread receives its own ``InfluxDBClient3`` instance
+    (via *client_factory*) because the client is not guaranteed thread-safe.
+
+    Results are returned in chunk order regardless of completion order.
+
+    Raises ``PermanentQueryError`` immediately, cancelling remaining work.
+    """
+    if not chunks:
+        return []
+
+    results: dict[int, List[T]] = {}
+    lock = threading.Lock()
+
+    with ThreadPoolExecutor(max_workers=max_workers) as pool:
+        future_to_idx: dict = {}
+        clients: list[InfluxDBClient3] = []
+
+        for idx, (t0, t1) in enumerate(chunks):
+            client = client_factory()
+            clients.append(client)
+            future = pool.submit(query_fn, client, t0, t1)
+            future_to_idx[future] = idx
+
+        try:
+            for future in as_completed(future_to_idx):
+                idx = future_to_idx[future]
+                result = future.result()  # raises on exception
+                with lock:
+                    results[idx] = result
+                if on_chunk_done:
+                    on_chunk_done(idx)
+        except PermanentQueryError:
+            for f in future_to_idx:
+                f.cancel()
+            raise
+        finally:
+            for c in clients:
+                try:
+                    c.close()
+                except Exception:
+                    pass
+
+    ordered: List[T] = []
+    for idx in sorted(results.keys()):
+        ordered.extend(results[idx])
+    return ordered