""" Test incremental loading behavior with VictoriaMetrics export. Runs the pipeline 3 times to verify: 1. First run loads data for window [start, start+window) 2. Second run advances state and loads next window 3. Third run with no new data available yields nothing """ import hashlib import json import logging import os import time from typing import Iterator import duckdb import requests import dlt 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" WINDOW_SECONDS = 300 # 5-minute windows for faster testing 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_vm_resource( metric_match: str, base_url: str, resource_name: str, default_start: int, 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 # 2 min buffer if end > latest_end: logger.info(f"[{resource_name}] Caught up -- end {end} > latest {latest_end}, skipping") return logger.info(f"[{resource_name}] Loading window {start} -> {end}") export_url = base_url + "export" params = { "match[]": metric_match, "start": str(start), "end": str(end), } resp = requests.get(export_url, params=params, timeout=60, stream=True) resp.raise_for_status() row_count = 0 for line in resp.iter_lines(decode_unicode=True): if not line: continue series = json.loads(line) metric = series.get("metric", {}) values = series.get("values", []) timestamps = series.get("timestamps", []) lhash = _labels_hash(metric) for ts_ms, val in zip(timestamps, values): if val is None or float(val) <= 0: continue row = dict(metric) row["metric_timestamp"] = int(ts_ms / 1000) row["metric_value"] = float(val) row["metric_labels_hash"] = lhash row_count += 1 yield row logger.info(f"[{resource_name}] Yielded {row_count} rows") state["last_end"] = end return _resource @dlt.source(name="vm_incremental_source") def vm_incremental_source(default_start: int): return make_vm_resource( metric_match=TEST_METRIC, base_url=TEST_BASE_URL, resource_name="vm_incremental", default_start=default_start, window_seconds=WINDOW_SECONDS, ) def test_incremental(): # Start 20 minutes ago so we have room for 3 x 5-min windows now = int(time.time()) initial_start = now - 1200 # 20 min ago initial_start = initial_start - (initial_start % 60) # align to minute pipeline_dir = os.path.join(os.path.dirname(__file__), "..", ".test_incremental") # Clean up any previous test import shutil if os.path.exists(pipeline_dir): shutil.rmtree(pipeline_dir) db_file = "test_incremental.duckdb" if os.path.exists(db_file): os.remove(db_file) p = dlt.pipeline( pipeline_name="test_incremental", destination="duckdb", dataset_name="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(vm_incremental_source(default_start=initial_start)) elapsed = time.time() - t0 # Query the current state 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") 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() # Verify incremental behavior if len(results) >= 2: r1, r2 = results[0], results[1] if r2["total_rows"] > r1["total_rows"]: print("OK: Run 2 added new rows (incremental state advanced)") elif r2["total_rows"] == r1["total_rows"]: print("NOTE: Run 2 same row count -- may have merged duplicates or caught up") if r2["max_ts"] and r1["max_ts"] and r2["max_ts"] > r1["max_ts"]: print("OK: Run 2 has later timestamps (window advanced)") if len(results) >= 3: r2, r3 = results[1], results[2] if r3["total_rows"] == r2["total_rows"]: print("OK: Run 3 added no new rows (caught up or same data merged)") elif r3["total_rows"] > r2["total_rows"]: print("OK: Run 3 added more rows (still catching up)") # Check for duplicates 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!") for d in dupes[:3]: print(f" hash={d[0]}, ts={d[1]}, count={d[2]}") else: print("\nOK: No duplicate (hash, timestamp) pairs found -- deduplication works") con.close() if __name__ == "__main__": test_incremental()