Merge pull request #6 from rated-network/feat/promql-client

Fix URL construction and add JSON Content-Type header in requests.

rated_exporter_sdk/providers/prometheus/auth.py

@@ -96,7 +96,10 @@ def get_auth_headers(self) -> Dict[str, Any]:
         Raises:
             AuthenticationError: If token file cannot be read
         """
-        headers: Dict[str, Union[bool, Tuple[str, str], str]] = {}
+        headers: Dict[str, Union[bool, Tuple[str, str], str]] = {
+            "Content-Type": "application/json"
+        }
+
         if self.token:
             headers["Authorization"] = f"Bearer {self.token}"
         elif isinstance(self.token_file, Path):

rated_exporter_sdk/providers/prometheus/client.py

@@ -2,7 +2,6 @@
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timedelta
 from typing import Any, Dict, Generator, List, Optional
-from urllib.parse import urljoin
 
 import requests
 import structlog
@@ -145,7 +144,11 @@ def _make_request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]:
             PrometheusAPIError: If API returns an error
             PrometheusTimeoutError: If request times out
         """
-        url = urljoin(self.base_url, endpoint)
+
+        if endpoint.startswith("/api/v1"):
+            endpoint = endpoint[7:]
+
+        url = f"{self.base_url.rstrip('/')}/api/v1{endpoint}"
         timeout = kwargs.pop("timeout", self.timeout)
 
         try:

rated_exporter_sdk/providers/prometheus/managed/gcloud_auth.py

@@ -0,0 +1,85 @@
+from datetime import datetime
+from typing import Optional
+
+import google.auth
+from google.auth import impersonated_credentials
+from google.auth.credentials import Credentials
+from google.auth.transport.requests import Request
+from google.oauth2 import service_account
+from rated_exporter_sdk.core.exceptions import AuthenticationError
+from rated_exporter_sdk.providers.prometheus.auth import PrometheusAuth
+
+
+class GCPPrometheusAuth(PrometheusAuth):
+    """Extended PrometheusAuth with GCP token refresh capabilities"""
+
+    def __init__(
+        self,
+        credentials: Optional[Credentials] = None,
+        service_account_file: Optional[str] = None,
+        target_principal: Optional[str] = None,
+        *args,
+        **kwargs
+    ):
+        self.credentials: Optional[Credentials] = credentials
+        self.service_account_file = service_account_file
+        self.target_principal = target_principal
+        self.scopes = (["https://www.googleapis.com/auth/monitoring.read"],)
+        self._setup_credentials()
+
+        if self.credentials is None:
+            raise AuthenticationError(
+                "Failed to initialize Prometheus client with GCP cloud credentials"
+            )
+
+        current_token = self._get_fresh_token()
+        super().__init__(token=current_token, *args, **kwargs)  # type: ignore
+
+    def _setup_credentials(self):
+        """Set up GCP credentials"""
+        if not self.credentials:
+            if self.service_account_file:
+                # Load from service account file
+                base_credentials = (
+                    service_account.Credentials.from_service_account_file(
+                        self.service_account_file, scopes=self.scopes
+                    )
+                )
+            else:
+                # Try to load default credentials
+                base_credentials, _ = google.auth.default(scopes=self.scopes)
+
+            if self.target_principal:
+                # Set up service account impersonation if needed
+                self.credentials = impersonated_credentials.Credentials(
+                    source_credentials=base_credentials,
+                    target_principal=self.target_principal,
+                    target_scopes=self.scopes,
+                    lifetime=3600,  # 1 hour
+                )
+            else:
+                self.credentials = base_credentials
+
+    def _get_fresh_token(self) -> str:
+        """Get a fresh token, refreshing if necessary"""
+        if self.credentials is None:
+            raise AuthenticationError("Credentials not initialized")
+
+        if not self.credentials.valid:
+            self.credentials.refresh(Request())
+        elif (
+            self.credentials.expiry
+            and (self.credentials.expiry - datetime.now()).total_seconds() < 300
+        ):
+            # Proactively refresh if less than 5 minutes until expiry
+            self.credentials.refresh(Request())
+
+        if not self.credentials.token:
+            raise AuthenticationError("No token available after refresh")
+
+        return self.credentials.token
+
+    def get_auth_headers(self) -> dict:
+        """Override to ensure fresh token"""
+        self.token = self._get_fresh_token()
+        return super().get_auth_headers()

requirements.txt

@@ -1,3 +1,8 @@
 pydantic==2.8.2
 urllib3==1.26.18
 structlog==24.4.0
+google==3.0.0
+google-auth==2.35.0
+google-auth-httplib2==0.2.0
+google-api-python-client==2.151.0
+oauth2client==4.1.3