From 3837f968aafa02ab4dd13c9ac2a219387e7a4fdb Mon Sep 17 00:00:00 2001
From: Francia Csaba <cf@gerhardt.io>
Date: Fri, 16 Feb 2024 14:56:42 +0100
Subject: [PATCH] GTFS file size metric

---
 amarillo/plugins/metrics/metrics.py | 20 ++++++++++++++++++++
 pyproject.toml                      |  2 +-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/amarillo/plugins/metrics/metrics.py b/amarillo/plugins/metrics/metrics.py
index 445c7a9..6f106a1 100644
--- a/amarillo/plugins/metrics/metrics.py
+++ b/amarillo/plugins/metrics/metrics.py
@@ -30,6 +30,24 @@ def amarillo_trips_number_total() -> Callable[[Info], None]:
 
     return instrumentation
 
+def amarillo_gtfs_file_size() -> Callable[[Info], None]:
+    METRIC = Gauge("amarillo_gtfs_file_size_MB", "Total file size of GTFS data.")
+
+    def instrumentation(info: Info) -> None:
+        total_size = sum(os.path.getsize(os.path.join(path, file)) for path, dirnames, files in os.walk("./data/gtfs") for file in files)
+        METRIC.set(total_size/1_000_000)
+
+    return instrumentation
+
+def amarillo_grfs_file_size() -> Callable[[Info], None]:
+    METRIC = Gauge("amarillo_grfs_file_size_MB", "Total file size of GRFS data.")
+
+    def instrumentation(info: Info) -> None:
+        total_size = sum(os.path.getsize(os.path.join(path, file)) for path, dirnames, files in os.walk("./data/grfs") for file in files)
+        METRIC.set(total_size/1_000_000)
+
+    return instrumentation
+
 
 router = APIRouter(
     prefix="/metrics",
@@ -59,6 +77,8 @@ def setup(app: FastAPI):
     instrumentator = Instrumentator().instrument(app)
     instrumentator.add(pfi_metrics.default())
     instrumentator.add(amarillo_trips_number_total())
+    instrumentator.add(amarillo_gtfs_file_size())
+    instrumentator.add(amarillo_grfs_file_size())
 
 
     instrumentator.instrument(app)
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 01d6682..e6eadb4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "amarillo-metrics"
-version = "0.0.4"
+version = "0.0.5"
 description = "Prometheus metrics for Amarillo"
 readme = "README.md"
 license = {file = "LICENSE"}