""" Compare Prometheus query_range vs VictoriaMetrics export extraction methods. Loads data from a single region (us-east-1) into separate DuckDB tables for comparison. """ import hashlib import json import logging import os import time from typing import Any, Iterator import requests import yaml 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 # Use a 10-minute window ending 5 minutes ago NOW = int(time.time()) WINDOW_END = NOW - (NOW % 60) - 300 WINDOW_START = WINDOW_END - 600 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() # --------------------------------------------------------------------------- # Method 1: Existing query_range approach (via dlt REST API client) # --------------------------------------------------------------------------- @dlt.resource(name="qr_metrics", write_disposition="merge", primary_key=("metric_labels_hash", "metric_timestamp")) def extract_query_range() -> Iterator[dict]: logger.info(f"[query_range] Querying {WINDOW_START} to {WINDOW_END}, step={STEP}") config: RESTAPIConfig = { "client": { "base_url": TEST_BASE_URL, "paginator": SinglePagePaginator(), }, "resources": [ { "name": "qr_raw", "endpoint": { "path": "query_range", "method": "POST", "params": { "query": TEST_QUERY, "start": WINDOW_START, "end": WINDOW_END, "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"[query_range] Yielded {row_count} rows") # --------------------------------------------------------------------------- # Method 2: VictoriaMetrics /api/v1/export # --------------------------------------------------------------------------- @dlt.resource(name="vm_metrics", write_disposition="merge", primary_key=("metric_labels_hash", "metric_timestamp")) def extract_vm_export(include_zeros: bool = False) -> Iterator[dict]: export_url = TEST_BASE_URL + "export" params = { "match[]": TEST_METRIC, "start": str(WINDOW_START), "end": str(WINDOW_END), } logger.info(f"[vm_export] Fetching {export_url} for {WINDOW_START} to {WINDOW_END}") resp = requests.get(export_url, params=params, timeout=60, stream=True) resp.raise_for_status() row_count = 0 skipped_zero = 0 skipped_null = 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: skipped_null += 1 continue fval = float(val) if not include_zeros and fval <= 0: skipped_zero += 1 continue row = dict(metric) row["metric_timestamp"] = int(ts_ms / 1000) # Convert ms -> seconds row["metric_value"] = fval row["metric_labels_hash"] = lhash row_count += 1 yield row logger.info( f"[vm_export] Yielded {row_count} rows " f"(skipped {skipped_zero} zeros, {skipped_null} nulls)" ) # --------------------------------------------------------------------------- # Method 2b: VictoriaMetrics export, keeping ALL data (no >0 filter) # --------------------------------------------------------------------------- @dlt.resource(name="vm_metrics_all", write_disposition="merge", primary_key=("metric_labels_hash", "metric_timestamp")) def extract_vm_export_all() -> Iterator[dict]: yield from extract_vm_export(include_zeros=True) # --------------------------------------------------------------------------- # Run both extractions # --------------------------------------------------------------------------- @dlt.source(name="query_range_source") def qr_source(): return extract_query_range() @dlt.source(name="vm_export_source") def vm_source(): return extract_vm_export() @dlt.source(name="vm_export_all_source") def vm_all_source(): return extract_vm_export_all() def run_comparison(): logger.info(f"Window: {WINDOW_START} to {WINDOW_END} ({WINDOW_END - WINDOW_START}s)") # Pipeline 1: query_range p1 = dlt.pipeline( pipeline_name="compare_qr", destination="duckdb", dataset_name="query_range_data", pipelines_dir=os.path.join(os.path.dirname(__file__), "..", ".compare_pipelines"), full_refresh=True, ) t0 = time.time() info1 = p1.run(qr_source()) qr_time = time.time() - t0 logger.info(f"[query_range] Load completed in {qr_time:.2f}s") logger.info(f"[query_range] Load info: {info1}") # Pipeline 2: VM export (filtered >0) p2 = dlt.pipeline( pipeline_name="compare_vm", destination="duckdb", dataset_name="vm_export_data", pipelines_dir=os.path.join(os.path.dirname(__file__), "..", ".compare_pipelines"), full_refresh=True, ) t0 = time.time() info2 = p2.run(vm_source()) vm_time = time.time() - t0 logger.info(f"[vm_export] Load completed in {vm_time:.2f}s") logger.info(f"[vm_export] Load info: {info2}") # Pipeline 3: VM export (all data, no filter) p3 = dlt.pipeline( pipeline_name="compare_vm_all", destination="duckdb", dataset_name="vm_export_all_data", pipelines_dir=os.path.join(os.path.dirname(__file__), "..", ".compare_pipelines"), full_refresh=True, ) t0 = time.time() info3 = p3.run(vm_all_source()) vm_all_time = time.time() - t0 logger.info(f"[vm_export_all] Load completed in {vm_all_time:.2f}s") # --- Compare results via DuckDB queries --- print("\n" + "=" * 80) print("COMPARISON RESULTS") print("=" * 80) print(f"Time window: {WINDOW_START} - {WINDOW_END} ({(WINDOW_END - WINDOW_START) // 60} minutes)") print(f"Region: us-east-1") print(f"Metric: {TEST_METRIC}") print() import duckdb for label, pipe, table in [ ("query_range (>0)", p1, "qr_metrics"), ("vm_export (>0 filtered)", p2, "vm_metrics"), ("vm_export (all)", p3, "vm_metrics_all"), ]: db_path = os.path.join(pipe.working_dir, f"{pipe.dataset_name}.duckdb") if not os.path.exists(db_path): db_path = pipe._pipeline_storage.storage_path + f"/{pipe.pipeline_name}.duckdb" con = duckdb.connect(db_path, read_only=True) try: schemas = con.sql("SELECT schema_name FROM information_schema.schemata").fetchall() schema = pipe.dataset_name tables = con.sql(f"SELECT table_name FROM information_schema.tables WHERE table_schema='{schema}'").fetchall() if not tables: print(f" [{label}] No tables found in schema {schema}") con.close() continue actual_table = tables[0][0] full_table = f"{schema}.{actual_table}" row_count = con.sql(f"SELECT COUNT(*) FROM {full_table}").fetchone()[0] distinct_ts = con.sql(f"SELECT COUNT(DISTINCT metric_timestamp) 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] col_count = con.sql(f"SELECT COUNT(*) FROM information_schema.columns WHERE table_schema='{schema}' AND table_name='{actual_table}'").fetchone()[0] sample = con.sql(f"SELECT * FROM {full_table} LIMIT 2").fetchdf() except Exception as e: print(f" [{label}] Error querying: {e}") con.close() continue con.close() print(f"--- {label} ---") print(f" Extraction time: {[qr_time, vm_time, vm_all_time][['query_range (>0)', 'vm_export (>0 filtered)', 'vm_export (all)'].index(label)]:.2f}s") print(f" Total rows: {row_count}") print(f" Distinct series: {distinct_series}") print(f" Distinct ts: {distinct_ts}") print(f" Timestamp range: {min_ts} - {max_ts}") print(f" Columns: {col_count}") print(f" Sample row:") print(f" {sample.iloc[0].to_dict() if len(sample) > 0 else 'empty'}") print() if __name__ == "__main__": run_comparison()