Chore/csrf support (#40)

* Trigger release of prometheus exporter automatically

* Ref

* CSRF Support

* Comment

* Unneeded code

* Use requests and add csrf token method to class

* Expiration token validation

* Use alternative logger

* Boolean for csrf_enabled

* Add env var to avoid unnecessary prefect logging alerts

* Remove API attempts entirely

* remove unneeded logging

* Use var

Dockerfile

@@ -6,7 +6,9 @@ ENV PYTHONDONTWRITEBYTECODE=1 \
 WORKDIR /app
 
 COPY ./ ./
-
+ENV PREFECT_HOME=/app
+ENV PREFECT_LOGGING_EXTRA_LOGGERS=prometheus-prefect-exporter
+ENV PREFECT_LOGGING_TO_API_ENABLED=False
 RUN pip install --upgrade pip
 RUN pip install \
       --disable-pip-version-check \

main.py

@@ -1,13 +1,13 @@
 import os
 import logging
 import time
+import uuid
 
 from metrics.metrics import PrefectMetrics
 from metrics.healthz import PrefectHealthz
 from prometheus_client import start_http_server, REGISTRY
 
-
-if __name__ == "__main__":
+def metrics():
     """
     Main entry point for the PrefectMetrics exporter.
     """
@@ -19,6 +19,7 @@
     offset_minutes = int(os.getenv("OFFSET_MINUTES", "5"))
     url = str(os.getenv("PREFECT_API_URL", "http://localhost:4200/api"))
     api_key = str(os.getenv("PREFECT_API_KEY", ""))
+    csrf_client_id = str(uuid.uuid4())
 
     # Configure logging
     logging.basicConfig(
@@ -43,6 +44,8 @@
         headers=headers,
         offset_minutes=offset_minutes,
         max_retries=max_retries,
+        client_id=csrf_client_id,
+        csrf_enabled=str(os.getenv("PREFECT_CSRF_ENABLED", "False")) == "True",
         logger=logger,
     )
 
@@ -57,3 +60,6 @@
     # Run the loop to collect Prefect metrics
     while True:
         time.sleep(5)
+
+if __name__ == "__main__":
+    metrics()

metrics/metrics.py

@@ -3,6 +3,10 @@
 from metrics.flows import PrefectFlows
 from metrics.work_pools import PrefectWorkPools
 from metrics.work_queues import PrefectWorkQueues
+import requests
+import time
+from prefect.client.schemas.objects import CsrfToken
+import pendulum
 from prometheus_client.core import GaugeMetricFamily, CounterMetricFamily
 
 
@@ -11,7 +15,7 @@ class PrefectMetrics(object):
     PrefectMetrics class for collecting and exposing Prometheus metrics related to Prefect.
     """
 
-    def __init__(self, url, headers, offset_minutes, max_retries, logger) -> None:
+    def __init__(self, url, headers, offset_minutes, max_retries, csrf_enabled, client_id, logger) -> None:
         """
         Initialize the PrefectMetrics instance.
 
@@ -28,13 +32,28 @@ def __init__(self, url, headers, offset_minutes, max_retries, logger) -> None:
         self.url = url
         self.max_retries = max_retries
         self.logger = logger
+        self.client_id = client_id
+        self.csrf_enabled = csrf_enabled
+        self.csrf_token = None
+        self.csrf_token_expiration = None
 
     def collect(self):
         """
         Get and set Prefect work queues metrics.
 
         """
         ##
+        # PREFECT GET CSRF TOKEN IF ENABLED
+        #
+        if self.csrf_enabled:
+            if not self.csrf_token or pendulum.now("UTC") > self.csrf_token_expiration:
+                self.logger.info("CSRF Token is expired or has not been generated yet. Fetching new CSRF Token...")
+                token_information = self.get_csrf_token()
+                self.csrf_token = token_information.token
+                self.csrf_token_expiration = token_information.expiration
+            self.headers["Prefect-Csrf-Token"] = self.csrf_token
+            self.headers["Prefect-Csrf-Client"] = self.client_id
+        ##
         # PREFECT GET RESOURCES
         #
         deployments = PrefectDeployments(
@@ -377,3 +396,20 @@ def collect(self):
             )
 
         yield prefect_info_work_queues
+
+    def get_csrf_token(self) -> CsrfToken:
+        """
+        Pull CSRF Token from CSRF Endpoint.
+
+        """
+        for retry in range(self.max_retries):
+            try:
+                csrf_token = requests.get(f"{self.url}/csrf-token?client={self.client_id}")
+            except requests.exceptions.HTTPError as err:
+                self.logger.error(err)
+                if retry >= self.max_retries - 1:
+                    time.sleep(1)
+                    raise SystemExit(err)
+            else:
+                break
+        return CsrfToken.parse_obj(csrf_token.json())

requirements.txt

@@ -1,2 +1,4 @@
 requests==2.31.0
 prometheus_client==0.20.0
+pendulum==2.1.2
+prefect==2.20.1