add scraping capabilities

Author: agavra

Cargo.lock

@@ -12,20 +12,24 @@ edition = "2024"
 
 [workspace.dependencies]
 async-trait = "0.1"
+axum = "0.7"
 bytes = "1.0"
 chrono = "0.4"
+clap = { version = "4", features = ["derive", "env"] }
+config = "0.14"
 dashmap = "6.1.0"
 fail = "0.5"
 moka = { version = "0.12", features = ["future"] }
 opendata-common = { path = "./opendata-common" }
+reqwest = { version = "0.12", default-features = false, features = ["rustls-tls"] }
 rstest = "0.19"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
+serde_yaml = "0.9"
 slatedb = "0.9.1"
 thiserror = "2.0"
 tokio = { version = "1.0", features = ["full"] }
+tower-http = { version = "0.5", features = ["cors"] }
 tracing = "0.1"
 tracing-subscriber = "0.3"
-uuid = "1.14.0"
-axum = "0.7"
-tower-http = { version = "0.5", features = ["cors"] }
+uuid = "1.14.0"

Cargo.toml

@@ -5,21 +5,25 @@ edition.workspace = true
 
 [dependencies]
 async-trait.workspace = true
+axum.workspace = true
 bytes.workspace = true
 chrono.workspace = true
+clap.workspace = true
+config.workspace = true
 dashmap.workspace = true
 moka.workspace = true
 opendata-common.workspace = true
+reqwest.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+serde_yaml.workspace = true
 slatedb.workspace = true
+thiserror.workspace = true
 tokio.workspace = true
+tower-http.workspace = true
 tracing.workspace = true
 tracing-subscriber.workspace = true
-serde.workspace = true
-serde_json.workspace = true
-thiserror.workspace = true
 uuid.workspace = true
-axum.workspace = true
-tower-http.workspace = true
 
 blake3 = "1.8.2"
 promql-parser = "0.6"

open-tsdb/Cargo.toml

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+"""
+Mock metrics server that exposes Prometheus-format metrics for testing.
+
+Usage:
+    python mock_metrics_server.py [--port PORT]
+
+The server exposes metrics at /metrics endpoint.
+"""
+
+import argparse
+import random
+import time
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from threading import Thread
+
+# Simulated metric values
+start_time = time.time()
+request_count = 0
+error_count = 0
+
+
+class MetricsHandler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        global request_count, error_count
+
+        if self.path == "/metrics":
+            request_count += 1
+            # Randomly increment error count
+            if random.random() < 0.1:
+                error_count += 1
+
+            metrics = self._generate_metrics()
+            self.send_response(200)
+            self.send_header("Content-Type", "text/plain; version=0.0.4")
+            self.end_headers()
+            self.wfile.write(metrics.encode("utf-8"))
+        elif self.path == "/health":
+            self.send_response(200)
+            self.send_header("Content-Type", "text/plain")
+            self.end_headers()
+            self.wfile.write(b"OK")
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+    def _generate_metrics(self):
+        uptime = time.time() - start_time
+
+        # Generate some realistic-looking metrics
+        metrics = []
+
+        # Counter: total requests
+        metrics.append("# HELP mock_requests_total Total number of requests received")
+        metrics.append("# TYPE mock_requests_total counter")
+        metrics.append(f"mock_requests_total {request_count}")
+
+        # Counter: errors
+        metrics.append("# HELP mock_errors_total Total number of errors")
+        metrics.append("# TYPE mock_errors_total counter")
+        metrics.append(f"mock_errors_total {error_count}")
+
+        # Gauge: uptime
+        metrics.append("# HELP mock_uptime_seconds Time since server started")
+        metrics.append("# TYPE mock_uptime_seconds gauge")
+        metrics.append(f"mock_uptime_seconds {uptime:.2f}")
+
+        # Gauge: random CPU usage simulation
+        metrics.append("# HELP mock_cpu_usage_percent Simulated CPU usage")
+        metrics.append("# TYPE mock_cpu_usage_percent gauge")
+        metrics.append(f"mock_cpu_usage_percent {random.uniform(10, 90):.2f}")
+
+        # Gauge: random memory usage simulation
+        metrics.append("# HELP mock_memory_usage_bytes Simulated memory usage")
+        metrics.append("# TYPE mock_memory_usage_bytes gauge")
+        metrics.append(f"mock_memory_usage_bytes {random.randint(100000000, 500000000)}")
+
+        # Histogram: simulated request duration
+        metrics.append("# HELP mock_request_duration_seconds Request duration histogram")
+        metrics.append("# TYPE mock_request_duration_seconds histogram")
+        # Simulate bucket counts
+        le_buckets = [0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0]
+        cumulative = 0
+        for le in le_buckets:
+            cumulative += random.randint(0, 10)
+            metrics.append(f'mock_request_duration_seconds_bucket{{le="{le}"}} {cumulative}')
+        metrics.append(f'mock_request_duration_seconds_bucket{{le="+Inf"}} {cumulative + random.randint(0, 5)}')
+        metrics.append(f"mock_request_duration_seconds_sum {random.uniform(0.1, 10.0):.4f}")
+        metrics.append(f"mock_request_duration_seconds_count {cumulative}")
+
+        # Metrics with labels
+        metrics.append("# HELP mock_http_requests_by_status HTTP requests by status code")
+        metrics.append("# TYPE mock_http_requests_by_status counter")
+        metrics.append(f'mock_http_requests_by_status{{status="200"}} {request_count - error_count}')
+        metrics.append(f'mock_http_requests_by_status{{status="500"}} {error_count}')
+
+        # Multi-label metric
+        metrics.append("# HELP mock_api_calls API calls by endpoint and method")
+        metrics.append("# TYPE mock_api_calls counter")
+        metrics.append(f'mock_api_calls{{endpoint="/users",method="GET"}} {random.randint(100, 1000)}')
+        metrics.append(f'mock_api_calls{{endpoint="/users",method="POST"}} {random.randint(10, 100)}')
+        metrics.append(f'mock_api_calls{{endpoint="/orders",method="GET"}} {random.randint(50, 500)}')
+        metrics.append(f'mock_api_calls{{endpoint="/orders",method="POST"}} {random.randint(20, 200)}')
+
+        return "\n".join(metrics) + "\n"
+
+    def log_message(self, format, *args):
+        # Suppress default logging, or customize it
+        print(f"[{self.log_date_time_string()}] {args[0]}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Mock Prometheus metrics server")
+    parser.add_argument("--port", type=int, default=8080, help="Port to listen on (default: 8080)")
+    args = parser.parse_args()
+
+    server = HTTPServer(("0.0.0.0", args.port), MetricsHandler)
+    print(f"Mock metrics server running on http://0.0.0.0:{args.port}")
+    print(f"Metrics available at http://0.0.0.0:{args.port}/metrics")
+    print("Press Ctrl+C to stop")
+
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        print("\nShutting down...")
+        server.shutdown()
+
+
+if __name__ == "__main__":
+    main()

open-tsdb/etc/mock_metrics_server.py

@@ -0,0 +1,23 @@
+# Prometheus configuration for testing with mock metrics server
+#
+# Usage:
+#   1. Start the mock server: python mock_metrics_server.py --port 8080
+#   2. Start Prometheus: prometheus --config.file=prometheus.yaml
+
+global:
+  scrape_interval: 15s
+  evaluation_interval: 15s
+
+scrape_configs:
+  # Scrape the mock metrics server
+  - job_name: "mock-metrics"
+    static_configs:
+      - targets: ["localhost:8080"]
+        labels:
+          env: "test"
+          service: "mock"
+
+  # Optionally scrape Prometheus itself
+  - job_name: "prometheus"
+    static_configs:
+      - targets: ["localhost:9090"]

open-tsdb/etc/prometheus.yaml

@@ -15,8 +15,10 @@ mod util;
 
 use std::sync::Arc;
 
+use clap::Parser;
 use opendata_common::storage::in_memory::InMemoryStorage;
 
+use promql::config::{CliArgs, PrometheusConfig, load_config};
 use promql::server::{PromqlServer, ServerConfig};
 use storage::merge_operator::OpenTsdbMergeOperator;
 use tsdb::Tsdb;
@@ -26,6 +28,26 @@ async fn main() {
     // Initialize tracing
     tracing_subscriber::fmt::init();
 
+    // Parse CLI arguments
+    let args = CliArgs::parse();
+
+    // Load Prometheus configuration if provided
+    let prometheus_config = if let Some(config_path) = &args.config {
+        match load_config(config_path) {
+            Ok(config) => {
+                tracing::info!("Loaded configuration from {}", config_path);
+                config
+            }
+            Err(e) => {
+                tracing::error!("Failed to load configuration: {}", e);
+                std::process::exit(1);
+            }
+        }
+    } else {
+        tracing::info!("No configuration file provided, using defaults");
+        PrometheusConfig::default()
+    };
+
     // Create in-memory storage with merge operator
     let storage = Arc::new(InMemoryStorage::with_merge_operator(Arc::new(
         OpenTsdbMergeOperator,
@@ -34,10 +56,18 @@ async fn main() {
     // Create Tsdb
     let tsdb = Arc::new(Tsdb::new(storage));
 
-    // Create and run server on default port (9090)
-    let config = ServerConfig::default();
+    // Create server configuration
+    let config = ServerConfig {
+        port: args.port,
+        prometheus_config,
+    };
+
+    // Create and run server
     let server = PromqlServer::new(tsdb, config);
 
-    println!("Starting open-tsdb Prometheus-compatible server on port 9090...");
+    tracing::info!(
+        "Starting open-tsdb Prometheus-compatible server on port {}...",
+        args.port
+    );
     server.run().await;
 }