#!/usr/bin/env python3
"""
Test script to verify AWS Access Keys can query AWS Managed Prometheus.
No IRSA/terraform required - just access keys!
"""
import os
import sys
from datetime import datetime
from typing import Optional
from urllib.parse import urlencode
import boto3
import httpx
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
def test_amp_query():
"""Test querying AMP with access keys."""
# Get config from environment
workspace_id = os.environ.get("PROMETHEUS_WORKSPACE_ID")
region = os.environ.get("AWS_REGION", "us-east-1")
access_key = os.environ.get("AWS_ACCESS_KEY_ID")
secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
# Validate
missing = []
if not workspace_id:
missing.append("PROMETHEUS_WORKSPACE_ID")
if not access_key:
missing.append("AWS_ACCESS_KEY_ID")
if not secret_key:
missing.append("AWS_SECRET_ACCESS_KEY")
if missing:
print(f"ā Missing environment variables: {', '.join(missing)}")
print("\nSet them like this:")
print(' export AWS_ACCESS_KEY_ID="AKIA..."')
print(' export AWS_SECRET_ACCESS_KEY="..."')
print(' export PROMETHEUS_WORKSPACE_ID="ws-..."')
sys.exit(1)
print(f"š§ Configuration:")
print(f" Region: {region}")
print(f" Workspace ID: {workspace_id}")
print(f" Access Key ID: {access_key[:8]}...{access_key[-4:]}")
print()
# Build URL
base_url = f"https://aps-workspaces.{region}.amazonaws.com/workspaces/{workspace_id}"
# Test 1: Simple query
print("š Test 1: Instant query (up metric)")
test_query(base_url, region, "up", "query")
# Test 2: Get labels
print("\nš·ļø Test 2: Get label names")
test_query(base_url, region, None, "labels")
# Test 3: Get series
print("\nš Test 3: Get series count")
test_query(base_url, region, "up", "series")
print("\nā
All tests passed! Access keys are working correctly.")
print("\nš” You can now use these credentials in your satellite deployment:")
print("""
kubectl create secret generic prometheus-mcp-credentials \\
--namespace deeptrace \\
--from-literal=aws_access_key_id="$AWS_ACCESS_KEY_ID" \\
--from-literal=aws_secret_access_key="$AWS_SECRET_ACCESS_KEY"
""")
def test_query(base_url: str, region: str, query: Optional[str], endpoint: str):
"""Execute a signed request to AMP."""
# Create boto3 session (uses env vars automatically)
session = boto3.Session(region_name=region)
credentials = session.get_credentials()
if credentials is None:
print(" ā Failed to get AWS credentials")
sys.exit(1)
# Build request
if endpoint == "query":
url = f"{base_url}/api/v1/query"
method = "POST"
body = urlencode({"query": query}).encode()
elif endpoint == "labels":
url = f"{base_url}/api/v1/labels"
method = "GET"
body = None
elif endpoint == "series":
url = f"{base_url}/api/v1/series"
method = "POST"
body = urlencode({"match[]": query}).encode()
else:
raise ValueError(f"Unknown endpoint: {endpoint}")
headers = {
"Content-Type": "application/x-www-form-urlencoded",
"Host": f"aps-workspaces.{region}.amazonaws.com",
"X-Amz-Date": datetime.utcnow().strftime("%Y%m%dT%H%M%SZ"),
}
# Sign request with SigV4
aws_request = AWSRequest(method=method, url=url, headers=headers, data=body or b"")
SigV4Auth(credentials, "aps", region).add_auth(aws_request)
signed_headers = dict(aws_request.headers)
# Make request
try:
with httpx.Client(timeout=30.0) as client:
response = client.request(
method=method,
url=url,
headers=signed_headers,
content=body,
)
if response.status_code == 200:
data = response.json()
if endpoint == "query":
result_type = data.get("data", {}).get("resultType", "unknown")
result_count = len(data.get("data", {}).get("result", []))
print(f" ā
Success! Result type: {result_type}, results: {result_count}")
elif endpoint == "labels":
labels = data.get("data", [])
print(f" ā
Success! Found {len(labels)} labels")
if labels:
print(f" š Sample labels: {labels[:5]}")
elif endpoint == "series":
series = data.get("data", [])
print(f" ā
Success! Found {len(series)} series")
else:
print(f" ā Failed with status {response.status_code}")
print(f" Response: {response.text[:500]}")
sys.exit(1)
except Exception as e:
print(f" ā Request failed: {e}")
sys.exit(1)
if __name__ == "__main__":
test_amp_query()
