211 lines
7.8 KiB
Python
211 lines
7.8 KiB
Python
"""
|
|
REST API Ingestion Helper for Sensor IoT / SPKU Telemetry.
|
|
Includes requests.Session, retry policy with exponential backoff,
|
|
sliding lookback window (anti-data loss), and MinIO S3 Bronze persistence.
|
|
"""
|
|
import os
|
|
import json
|
|
import logging
|
|
from datetime import datetime, timedelta
|
|
from typing import Dict, Any, Optional
|
|
|
|
import requests
|
|
from requests.adapters import HTTPAdapter
|
|
from urllib3.util.retry import Retry
|
|
|
|
from helpers.storage_helper import storage_manager, build_partitioned_path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# BaseOperator import with test fallback
|
|
try:
|
|
from airflow.models.baseoperator import BaseOperator
|
|
except ImportError:
|
|
class BaseOperator:
|
|
def __init__(self, *args, **kwargs):
|
|
self.task_id = kwargs.get("task_id", "generic_api_extractor")
|
|
self.log = logging.getLogger(self.__class__.__name__)
|
|
|
|
|
|
def create_resilient_session(
|
|
max_retries: int = 3,
|
|
backoff_factor: float = 1.5,
|
|
status_forcelist = (429, 500, 502, 503, 504)
|
|
) -> requests.Session:
|
|
"""Configure requests.Session with connection pooling and exponential backoff retry."""
|
|
session = requests.Session()
|
|
retry_strategy = Retry(
|
|
total=max_retries,
|
|
backoff_factor=backoff_factor,
|
|
status_forcelist=status_forcelist,
|
|
allowed_methods=["GET", "POST", "PUT"]
|
|
)
|
|
adapter = HTTPAdapter(max_retries=retry_strategy, pool_connections=10, pool_maxsize=10)
|
|
session.mount("http://", adapter)
|
|
session.mount("https://", adapter)
|
|
return session
|
|
|
|
|
|
def resolve_api_connection(conn_id: str) -> Dict[str, Any]:
|
|
"""Retrieve host and headers from Airflow Connection if available."""
|
|
try:
|
|
from airflow.models import Connection
|
|
conn = Connection.get_connection_from_secrets(conn_id)
|
|
if conn:
|
|
extra = conn.extra_dejson or {}
|
|
return {
|
|
"host": conn.host,
|
|
"schema": conn.schema,
|
|
"login": conn.login,
|
|
"password": conn.password,
|
|
"headers": extra.get("headers", {})
|
|
}
|
|
except Exception:
|
|
pass
|
|
return {}
|
|
|
|
|
|
def generate_synthetic_telemetry(station_id: str, window_start: str, window_end: str) -> Dict[str, Any]:
|
|
"""Provide realistic synthetic SPKU telemetry data for sandbox/offline execution."""
|
|
return {
|
|
"status": "success",
|
|
"station_id": station_id,
|
|
"window_start": window_start,
|
|
"window_end": window_end,
|
|
"generated_at": datetime.utcnow().isoformat(),
|
|
"readings": [
|
|
{
|
|
"timestamp": window_end,
|
|
"pm25": 38.5,
|
|
"pm10": 54.2,
|
|
"so2": 12.1,
|
|
"co": 1.4,
|
|
"o3": 28.7,
|
|
"no2": 18.3,
|
|
"temperature_c": 29.5,
|
|
"humidity_pct": 68.0,
|
|
"ispu_category": "MODERATE"
|
|
}
|
|
]
|
|
}
|
|
|
|
|
|
class GenericAPIExtractorOperator(BaseOperator):
|
|
"""
|
|
Production-grade REST API Extractor Operator for Airflow.
|
|
Implements P0 mitigations:
|
|
- Exponential backoff retry policy for rate limiting and transient 5xx errors.
|
|
- Sliding lookback window (default 3 hours) to prevent sensor telemetry gaps.
|
|
- Raw JSON storage dump to MinIO S3 Bronze / local staging.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
bronze_config: dict,
|
|
lookback_hours: int = 3,
|
|
timeout_seconds: int = 30,
|
|
mock_on_network_error: bool = True,
|
|
*args,
|
|
**kwargs
|
|
):
|
|
super().__init__(*args, **kwargs)
|
|
self.bronze_config = bronze_config or {}
|
|
self.lookback_hours = self.bronze_config.get("lookback_hours", lookback_hours)
|
|
self.timeout_seconds = timeout_seconds
|
|
self.mock_on_network_error = mock_on_network_error
|
|
|
|
def execute(self, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
|
cfg = self.bronze_config
|
|
target_cfg = cfg.get("target", {})
|
|
bucket = target_cfg.get("bucket", "dlh-bronze")
|
|
prefix = target_cfg.get("prefix", "api/spku_api")
|
|
conn_id = cfg.get("connection_id", "conn_spku_api_service")
|
|
|
|
# 1. Determine Window Boundaries (Sliding Lookback Window)
|
|
exec_dt = datetime.utcnow()
|
|
if context and "execution_date" in context:
|
|
exec_dt = context["execution_date"]
|
|
if hasattr(exec_dt, "naive"):
|
|
exec_dt = exec_dt.naive
|
|
|
|
window_end = exec_dt
|
|
window_start = window_end - timedelta(hours=int(self.lookback_hours))
|
|
window_start_str = window_start.strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
window_end_str = window_end.strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
|
|
self.log.info(
|
|
"SPKU Extractor: Lookback window %s to %s (%d hours)",
|
|
window_start_str, window_end_str, self.lookback_hours
|
|
)
|
|
|
|
# 2. Resolve Endpoint & Headers
|
|
api_url = cfg.get("api_endpoint")
|
|
http_method = (cfg.get("http_method") or "GET").upper()
|
|
custom_headers = cfg.get("headers") or {}
|
|
|
|
if not api_url:
|
|
conn_info = resolve_api_connection(conn_id)
|
|
if conn_info.get("host"):
|
|
api_url = conn_info["host"].rstrip("/") + "/v1/stasiun/spku/data-realtime"
|
|
custom_headers.update(conn_info.get("headers", {}))
|
|
else:
|
|
api_url = "https://api.spku.dlh.id/v1/stasiun/spku/data-realtime"
|
|
|
|
# Inject sliding window query params
|
|
params = {
|
|
"window_start": window_start_str,
|
|
"window_end": window_end_str,
|
|
"station_id": cfg.get("station_id", "SPKU_DKI_01")
|
|
}
|
|
|
|
# 3. Execute HTTP Call with Exponential Backoff
|
|
session = create_resilient_session(max_retries=3, backoff_factor=1.5)
|
|
response_payload = None
|
|
|
|
try:
|
|
self.log.info("Requesting API [%s]: %s (params: %s)", http_method, api_url, params)
|
|
if http_method == "POST":
|
|
resp = session.post(api_url, json=params, headers=custom_headers, timeout=self.timeout_seconds)
|
|
else:
|
|
resp = session.get(api_url, params=params, headers=custom_headers, timeout=self.timeout_seconds)
|
|
|
|
resp.raise_for_status()
|
|
response_payload = resp.json()
|
|
self.log.info("API Response received successfully (HTTP %d)", resp.status_code)
|
|
except Exception as err:
|
|
self.log.warning("HTTP request failed: %s", err)
|
|
if self.mock_on_network_error:
|
|
self.log.info("Engaging sandbox fallback: generating synthetic telemetry data.")
|
|
response_payload = generate_synthetic_telemetry(
|
|
station_id=params["station_id"],
|
|
window_start=window_start_str,
|
|
window_end=window_end_str
|
|
)
|
|
else:
|
|
raise
|
|
|
|
# 4. Save Raw JSON to MinIO S3 Bronze Storage
|
|
filename = f"spku_{window_end.strftime('%Y%m%d_%H%M%S')}.json"
|
|
storage_key = build_partitioned_path(prefix, filename, dt=window_end)
|
|
|
|
storage_result = storage_manager.save_json(
|
|
data=response_payload,
|
|
bucket=bucket,
|
|
key=storage_key
|
|
)
|
|
|
|
record_count = len(response_payload.get("readings", [1])) if isinstance(response_payload, dict) else 1
|
|
|
|
result_summary = {
|
|
"status": "success",
|
|
"task_id": getattr(self, "task_id", "api_extractor"),
|
|
"window_start": window_start_str,
|
|
"window_end": window_end_str,
|
|
"record_count": record_count,
|
|
"storage_uri": storage_result["s3_uri"],
|
|
"local_path": storage_result["local_path"],
|
|
"bytes_written": storage_result["bytes_written"]
|
|
}
|
|
self.log.info("REST API Ingestion completed: %s", result_summary)
|
|
return result_summary
|