feat: push directly to prometheus

main.py

@@ -1,10 +1,12 @@
 import os
 import time
 from prometheus_client import start_http_server, Gauge, REGISTRY
+from prometheus_remote_writer import RemoteWriter
 from dotenv import load_dotenv
 
-# 1. Define Metrics with Labels
-# Use labels for domain and feature to avoid creating 100s of unique metric names
+from libs.generate_metrics import aggregate, domain_imports
+from libs.parse_caddyfile import parse
+
 DOMAIN_FEATURE = Gauge(
     "caddy_domain_imports",
     "Binary status of domain features (1=enabled, 0=disabled)",
@@ -15,10 +17,15 @@ GLOBAL_STATS = Gauge(
     "caddy_global_stats", "Aggregated counts of features across all domains", ["type"]
 )
 
-
 load_dotenv()
-
-CADDY_PATH = os.environ["CADDY_PATH"]
+PROMETHEUS_URL: str = os.environ.get("PROMETHEUS_URL", "http://localhost:9090")
+PROMETHEUS_REMOTE_WRITE_URL: str = os.environ.get(
+    "PROMETHEUS_REMOTE_WRITE_URL", "/api/v1/write"
+)
+PROMETHEUS_REMOTE_WRITE_TOKEN: str = os.environ.get(
+    "PROMETHEUS_REMOTE_WRITE_TOKEN", "test"
+)
+CADDY_PATH = os.environ.get("CADDY_PATH", "Caddyfile")
 
 
 def load_caddyfile():
@@ -26,63 +33,38 @@ def load_caddyfile():
         return f.read()
 
 
-def parse(data):
-    from libs.parse_caddyfile import parse as caddy_parse
+class PrometheusWriter:
+    url: str = PROMETHEUS_URL
+    write_path: str = PROMETHEUS_REMOTE_WRITE_URL
+    token: str = PROMETHEUS_REMOTE_WRITE_TOKEN
 
-    return caddy_parse(data)
+    def initialize_writer(self):
+        self.writer = RemoteWriter(
+            url=self.url + self.write_path,
+            headers={f"Authorization": "Bearer {self.token}"},
+        )
+
+    def __init__(self) -> None:
+        self.initialize_writer()
+
+    def write(self, data: dict):
+        self.writer.send([data])
 
 
-def mathing(data: dict):
-    total_domains = len(data.items())
-    # info = dict({"auth": 0, "wan": 0, "noauth": 0, "nowan": 0, "none": 0})
-    info = dict({"no_auth": 0, "no_wan": 0})
-    for domain, imports in data.items():
-        if len(imports) == 0:
-            if "none" not in info:
-                info["none"] = 1
-            info["none"] += 1
-            continue
-        for imp in imports:
-            if imp not in info:
-                info[imp] = 1
-            info[imp] += 1
-        if "auth" not in imports:
-            info["no_auth"] += 1
-        if "wan" not in imports:
-            info["no_wan"] += 1
-    info["total_domains"] = total_domains
-    return info
+def push_metrics(data):
+    for metric in data:
+        print(metric)
+        PrometheusWriter().write(metric)
 
 
-def update_metrics():
-    # Your existing parsing logic
-    try:
-        raw_data = load_caddyfile()
-        hosts = parse(raw_data)
-        info = mathing(hosts)
-
-        # --- Update Global Metrics ---
-        for key, value in info.items():
-            GLOBAL_STATS.labels(type=key).set(value)
-
-        # --- Update Per-Domain Metrics ---
-        tracked_features = ["auth", "wan"]
-
-        for domain, imports in hosts.items():
-            for feature in tracked_features:
-                is_enabled = 1 if feature in imports else 0
-                DOMAIN_FEATURE.labels(domain=domain, feature=feature).set(is_enabled)
-
-        print(f"Metrics updated at {time.ctime()}")
-
-    except Exception as e:
-        print(f"Error updating metrics: {e}")
+def main():
+    caddy = load_caddyfile()
+    domains: dict[str, list[str]] = parse(caddy)
+    aggregated: list[dict] = aggregate(domains)
+    imports = domain_imports(domains)
+    push_metrics(aggregated)
+    push_metrics(imports)
 
 
 if __name__ == "__main__":
-    start_http_server(port=8000, addr="0.0.0.0")
-    print("Prometheus server started on port 8000")
-
-    while True:
-        update_metrics()
-        time.sleep(60)
+    main()