import json
import logging
from typing import Any, Dict, Literal, Optional
import requests
from airflow.exceptions import AirflowException
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from config import Config
[docs]
class GrapiBaseOperator(BaseOperator):
"""Base operator for making GraphQL requests to the Grapi API.
This operator handles authentication, request preparation, and response processing.
Subclasses must implement the `get_request_payload` and `post_process_response` methods.
:param tenant: The tenant identifier for the request.
:param stage: The API stage to use ("main" or "stage").
:param timeout: Request timeout in seconds.
:param retry_count: Number of retry attempts for the request.
:param retry_delay: Delay between retries in seconds.
"""
template_fields = ["tenant"]
ui_color = "#e8f5e9"
iv_length = 12 # Length of the IV for AES-GCM
tag_length = 16 # Length of the authentication tag for AES-GCM
@apply_defaults
def __init__(
self,
tenant: str,
stage: Literal["main", "stage"] = "main",
timeout: int = 30,
retry_count: int = 3,
retry_delay: int = 5,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.tenant = tenant
self.api_url = (
Config.GRAPHQL_API_URL if stage == "main" else Config.GRAPHQL_API_URL_STAGE
)
self.grapi_secret_name = (
Config.GRAPI_API_SECRET_NAME
if stage == "main"
else Config.GRAPI_API_SECRET_NAME_STAGE
)
self.aws_session_token = None
self.timeout = timeout
self.retry_count = retry_count
self.retry_delay = retry_delay
[docs]
def get_request_payload(self) -> Dict[str, Any]:
"""Prepare the GraphQL request payload."""
raise NotImplementedError("Subclasses must implement this method.")
[docs]
def post_process_response(
self, response: requests.Response, context: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
raise NotImplementedError("Subclasses must implement this method.")
def _get_secret(self, secret_name: str) -> str:
"""Retrieve the secret from AWS Secrets Manager."""
import boto3
region_name = "us-east-1"
# Create a Secrets Manager client
session = boto3.session.Session()
client = session.client(service_name="secretsmanager", region_name=region_name)
try:
get_secret_value_response = client.get_secret_value(SecretId=secret_name)
secret = get_secret_value_response["SecretString"].encode("utf-8")
return secret
except Exception as e:
# For a list of exceptions thrown, see
# https://docs.aws.amazon.com/secretsmanager/latest/apireference/API_GetSecretValue.html
raise e
def _get_aws_session_token_boto3(self) -> None:
"""
Get AWS session credentials using boto3.
This automatically handles IMDSv1/v2.
Returns:
dict: Contains credentials
None: If no credentials available
"""
try:
import boto3
session = boto3.Session()
credentials = session.get_credentials()
if credentials and credentials.token:
frozen_creds = credentials.get_frozen_credentials()
self.aws_session_token = frozen_creds.token
return frozen_creds
else:
raise AirflowException(
"No valid AWS session credentials found. "
"Ensure your environment is configured with AWS credentials."
)
except Exception as e:
raise AirflowException(
f"Failed to retrieve AWS session credentials: {str(e)}"
)
def _get_headers(self) -> Dict[str, str]:
"""Prepare request headers."""
import time
import uuid
if not self.tenant:
raise AirflowException("Missing required field: tenant")
headers = {"Content-Type": "application/json", "x-tenant": self.tenant}
# Get the auth key from secret manager
auth_key = self._derive_auth_key()
aws_token = self._get_aws_session_token_boto3()
timestamp_ms = int(time.time() * 1000)
object_to_encrypt = {
"credentials": {
"accessKeyId": aws_token.access_key,
"secretAccessKey": aws_token.secret_key,
"sessionToken": aws_token.token,
},
"nonce": str(uuid.uuid4()),
"timestamp": timestamp_ms,
}
auth_key = self._derive_auth_key()
encrypted_aws_session_token = self._encrypt(object_to_encrypt, auth_key)
headers["x-amz-security-token"] = encrypted_aws_session_token
return headers
def _make_request_with_retry(
self, url: str, data: Dict[str, Any]
) -> requests.Response:
"""Make HTTP request with retry logic."""
import time
last_exception = None
for attempt in range(self.retry_count):
try:
headers = self._get_headers()
self.log.info(f"Request headers: {headers}")
self.log.info(
f"Making GraphQL request to {url} (attempt {attempt + 1}/{self.retry_count})"
)
response = requests.post(
url, headers=headers, data=json.dumps(data), timeout=self.timeout
)
if response.ok:
return response
self.log.warning(
f"Request returned status {response.status_code}: {response.text}"
)
# Don't retry on 4xx errors (client errors)
if 400 <= response.status_code < 500:
raise AirflowException(
f"Client error (status {response.status_code}): {response.text}"
)
except Exception as e:
last_exception = e
self.log.warning(f"Request attempt {attempt + 1} failed: {str(e)}")
if attempt < self.retry_count - 1:
time.sleep(1)
# All retries exhausted
raise AirflowException(
f"Request failed after {self.retry_count} attempts: {str(last_exception)}"
)
[docs]
def execute(self, context):
"""Execute the GraphQL mutation."""
data = self.get_request_payload()
self.log.info(f"Making GraphQL request to {self.api_url}")
self.log.info(f"Request data: {data}")
response = self._make_request_with_retry(self.api_url, data)
self.log.info(f"Response status code: {response.status_code}")
self.log.info(f"Response content: {response.text}")
result = self.post_process_response(response, context)
self.log.info(f"Response result: {result}")
return result
def _derive_auth_key(self) -> bytes:
"""Drive a key from a secret manager
:return: _description_
:rtype: bytes
"""
import base64
raw_key = self._get_secret(secret_name=self.grapi_secret_name)
return base64.b64decode(raw_key)
def _encrypt(self, object_to_encript: dict, auth_key: str) -> str:
"""Encrypt an object using AES-256-GCM.
:param object_to_encript: The object to encrypt.
:param auth_key: The AES key to use for encryption.
:return: Base64 encoded encrypted string.
:rtype: str
"""
import json
from cryptography.hazmat.backends import default_backend
from uuid import uuid4
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
import base64
if len(auth_key) != 32:
raise ValueError(
f"Secret key must be 32 bytes for AES-256-GCM, got {len(auth_key)} bytes"
)
nonce = uuid4().bytes[: self.iv_length]
data = json.dumps(object_to_encript).encode("utf-8")
cipher = Cipher(
algorithms.AES(auth_key), modes.GCM(nonce), backend=default_backend()
)
encryptor = cipher.encryptor()
ciphertext = encryptor.update(data) + encryptor.finalize()
tag = encryptor.tag
encrypted = nonce + ciphertext + tag
return base64.b64encode(encrypted).decode("utf-8")
[docs]
@staticmethod
def get_grapi_api_connections():
"""
Get grapi API connections
"""
return (
{
"stage": "STAGE grapi",
"main": "PROD grapi",
}
if Config.PRODUCTION
else {
"main": "DEV grapi",
}
)