""" Test incremental loading with query_range (existing method) for comparison. """ import hashlib import json import logging import os import shutil import time from typing import Iterator import duckdb import dlt from dlt.sources.rest_api import RESTAPIConfig, rest_api_resources from dlt.sources.helpers.rest_client.paginators import SinglePagePaginator logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger(__name__) TEST_BASE_URL = "http://www.ds.us-east-1.aws.observability.tidbcloud.com/internal/metrics/f502cc26-7a0e-4a42-a92e-8816c323ccc7/api/v1/" TEST_METRIC = "tidb_vector_search_table_nums_with_vector_index" TEST_QUERY = f"{TEST_METRIC}>0" STEP = 30 WINDOW_SECONDS = 300 def _labels_hash(metric: dict) -> str: labels = {k: v for k, v in sorted(metric.items()) if k != "__name__"} return hashlib.md5(json.dumps(labels, sort_keys=True).encode()).hexdigest() def make_qr_resource( query: str, base_url: str, resource_name: str, default_start: int, step: int = 30, window_seconds: int = 300, ): @dlt.resource( name=resource_name, write_disposition="merge", primary_key=("metric_labels_hash", "metric_timestamp"), ) def _resource() -> Iterator[dict]: state = dlt.current.resource_state() start = state.setdefault("last_end", default_start) end = start + window_seconds current_time = int(time.time()) latest_end = current_time - 120 if end > latest_end: logger.info(f"[{resource_name}] Caught up, skipping") return logger.info(f"[{resource_name}] Loading window {start} -> {end}") config: RESTAPIConfig = { "client": { "base_url": base_url, "paginator": SinglePagePaginator(), }, "resources": [ { "name": "qr_raw", "endpoint": { "path": "query_range", "method": "POST", "params": { "query": query, "start": start, "end": end - 1, "step": step, }, }, } ], } row_count = 0 for resource in rest_api_resources(config): for item in resource: data = getattr(item, "data", item) if not isinstance(data, list): data = [data] for subitem in data: metric = subitem.get("metric", {}) for ts, value in subitem.get("values", []): row = dict(metric) row["metric_timestamp"] = int(ts) row["metric_value"] = float(value) row["metric_labels_hash"] = _labels_hash(metric) row_count += 1 yield row logger.info(f"[{resource_name}] Yielded {row_count} rows") state["last_end"] = end return _resource @dlt.source(name="qr_incremental_source") def qr_incremental_source(default_start: int): return make_qr_resource( query=TEST_QUERY, base_url=TEST_BASE_URL, resource_name="qr_incremental", default_start=default_start, step=STEP, window_seconds=WINDOW_SECONDS, ) def test_incremental(): now = int(time.time()) initial_start = now - 1200 initial_start = initial_start - (initial_start % 60) pipeline_dir = os.path.join(os.path.dirname(__file__), "..", ".test_incremental_qr") if os.path.exists(pipeline_dir): shutil.rmtree(pipeline_dir) db_file = "test_incremental_qr.duckdb" if os.path.exists(db_file): os.remove(db_file) p = dlt.pipeline( pipeline_name="test_incremental_qr", destination="duckdb", dataset_name="qr_incremental_test", pipelines_dir=pipeline_dir, ) results = [] for run_num in range(1, 4): logger.info(f"\n{'='*60}\nRUN {run_num}\n{'='*60}") t0 = time.time() info = p.run(qr_incremental_source(default_start=initial_start)) elapsed = time.time() - t0 con = duckdb.connect(db_file, read_only=True) schemas = [r[0] for r in con.sql("SELECT DISTINCT table_schema FROM information_schema.tables").fetchall() if r[0] not in ('information_schema', 'main', 'pg_catalog')] if schemas: schema = schemas[0] tables = [r[0] for r in con.sql(f"SELECT table_name FROM information_schema.tables WHERE table_schema='{schema}'").fetchall() if not r[0].startswith('_dlt')] if tables: full_table = f'"{schema}"."{tables[0]}"' row_count = con.sql(f'SELECT COUNT(*) FROM {full_table}').fetchone()[0] distinct_series = con.sql(f'SELECT COUNT(DISTINCT metric_labels_hash) FROM {full_table}').fetchone()[0] min_ts = con.sql(f'SELECT MIN(metric_timestamp) FROM {full_table}').fetchone()[0] max_ts = con.sql(f'SELECT MAX(metric_timestamp) FROM {full_table}').fetchone()[0] else: row_count = distinct_series = 0 min_ts = max_ts = None else: row_count = distinct_series = 0 min_ts = max_ts = None con.close() results.append({ "run": run_num, "elapsed": elapsed, "total_rows": row_count, "distinct_series": distinct_series, "min_ts": min_ts, "max_ts": max_ts, }) logger.info(f"Run {run_num}: {row_count} total rows, {distinct_series} series, " f"ts range {min_ts}-{max_ts}, took {elapsed:.2f}s") print("\n" + "=" * 80) print("INCREMENTAL LOADING TEST RESULTS (query_range)") print("=" * 80) print(f"Window size: {WINDOW_SECONDS}s ({WINDOW_SECONDS // 60} minutes)") print(f"Initial start: {initial_start}") print() for r in results: print(f"Run {r['run']}:") print(f" Total rows in DB: {r['total_rows']}") print(f" Distinct series: {r['distinct_series']}") print(f" Timestamp range: {r['min_ts']} - {r['max_ts']}") print(f" Time taken: {r['elapsed']:.2f}s") print() if len(results) >= 2 and results[1]["total_rows"] > results[0]["total_rows"]: print("OK: Run 2 added new rows (incremental state advanced)") if len(results) >= 2 and results[1]["max_ts"] and results[0]["max_ts"] and results[1]["max_ts"] > results[0]["max_ts"]: print("OK: Run 2 has later timestamps (window advanced)") con = duckdb.connect(db_file, read_only=True) schemas = [r[0] for r in con.sql("SELECT DISTINCT table_schema FROM information_schema.tables").fetchall() if r[0] not in ('information_schema', 'main', 'pg_catalog')] if schemas: schema = schemas[0] tables = [r[0] for r in con.sql(f"SELECT table_name FROM information_schema.tables WHERE table_schema='{schema}'").fetchall() if not r[0].startswith('_dlt')] if tables: full_table = f'"{schema}"."{tables[0]}"' dupes = con.sql(f""" SELECT metric_labels_hash, metric_timestamp, COUNT(*) as cnt FROM {full_table} GROUP BY metric_labels_hash, metric_timestamp HAVING COUNT(*) > 1 LIMIT 5 """).fetchall() if dupes: print(f"\nWARNING: Found {len(dupes)}+ duplicate (hash, timestamp) pairs!") else: print("\nOK: No duplicate (hash, timestamp) pairs found -- deduplication works") con.close() if __name__ == "__main__": test_incremental()