392 lines
16 KiB
Python
392 lines
16 KiB
Python
"""Runtime helpers for mapped custom data sources."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.target_schema_registry import TARGET_SCHEMAS
|
|
from app.db.session import async_session_factory
|
|
from app.models.datasource_config import DataSourceConfig
|
|
from app.models.datasource_mapping import DataSourceMappingTemplate
|
|
from app.services.datasource_mapping import (
|
|
MappingError,
|
|
execute_mapping,
|
|
extract_path,
|
|
persist_mapped_records,
|
|
)
|
|
|
|
DEFAULT_MAPPING_TEMPLATES: dict[str, dict[str, Any]] = {
|
|
"vessel_ais": {
|
|
"source": {"items_path": "$"},
|
|
"fields": {
|
|
"mmsi": {"path": "$.mmsi", "type": "integer"},
|
|
"name": {"path": "$.name", "type": "string", "default": None},
|
|
"lat": {"path": "$.lat", "type": "float"},
|
|
"lon": {"path": "$.lon", "type": "float"},
|
|
"sog": {"path": "$.sog", "type": "float", "default": None},
|
|
"cog": {"path": "$.cog", "type": "float", "default": None},
|
|
"heading": {"path": "$.heading", "type": "integer", "default": None},
|
|
"nav_status": {"path": "$.nav_status", "type": "integer", "default": None},
|
|
"callsign": {"path": "$.callsign", "type": "string", "default": None},
|
|
"vessel_type": {"path": "$.vessel_type", "type": "string", "default": None},
|
|
"vessel_type_name": {"path": "$.vessel_type_name", "type": "string", "default": None},
|
|
"received_at": {"path": "$.received_at", "type": "datetime", "default": None},
|
|
},
|
|
"meta": {"generated_by": "default_template", "requires_review": False},
|
|
},
|
|
}
|
|
|
|
RUNNING_CUSTOM_STREAM_TASKS: dict[int, asyncio.Task[Any]] = {}
|
|
|
|
|
|
class CustomDatasourceRuntimeError(RuntimeError):
|
|
"""Raised when a custom datasource cannot run."""
|
|
|
|
|
|
def build_request_headers(auth_type: str, auth_config: dict, headers: dict) -> dict[str, str]:
|
|
request_headers = {str(key): str(value) for key, value in (headers or {}).items()}
|
|
auth_type = str(auth_type or "none").lower()
|
|
auth_config = auth_config or {}
|
|
|
|
if auth_type == "bearer" and auth_config.get("token"):
|
|
request_headers["Authorization"] = f"Bearer {auth_config['token']}"
|
|
elif auth_type == "api_key" and auth_config.get("api_key"):
|
|
location = str(auth_config.get("in") or auth_config.get("location") or "header").lower()
|
|
if location != "query":
|
|
key_name = auth_config.get("key_name", "X-API-Key")
|
|
request_headers[str(key_name)] = str(auth_config["api_key"])
|
|
elif auth_type == "basic":
|
|
username = auth_config.get("username", "")
|
|
password = auth_config.get("password", "")
|
|
credentials = f"{username}:{password}"
|
|
encoded = base64.b64encode(credentials.encode()).decode()
|
|
request_headers["Authorization"] = f"Basic {encoded}"
|
|
return request_headers
|
|
|
|
|
|
def build_query_params(auth_type: str, auth_config: dict, config: dict) -> dict[str, Any]:
|
|
params: dict[str, Any] = {}
|
|
candidate = (config or {}).get("params") or (config or {}).get("query_params")
|
|
if isinstance(candidate, dict):
|
|
params.update(candidate)
|
|
|
|
auth_type = str(auth_type or "none").lower()
|
|
auth_config = auth_config or {}
|
|
if auth_type == "api_key" and auth_config.get("api_key"):
|
|
location = str(auth_config.get("in") or auth_config.get("location") or "header").lower()
|
|
if location == "query":
|
|
key_name = auth_config.get("key_name") or auth_config.get("param_name") or "api_key"
|
|
params[str(key_name)] = auth_config["api_key"]
|
|
return params
|
|
|
|
|
|
async def load_active_mapping(
|
|
db: AsyncSession,
|
|
datasource_config_id: int,
|
|
) -> DataSourceMappingTemplate:
|
|
result = await db.execute(
|
|
select(DataSourceMappingTemplate)
|
|
.where(DataSourceMappingTemplate.datasource_config_id == datasource_config_id)
|
|
.where(DataSourceMappingTemplate.is_active.is_(True))
|
|
.order_by(DataSourceMappingTemplate.version.desc())
|
|
.limit(1)
|
|
)
|
|
mapping = result.scalar_one_or_none()
|
|
if mapping is not None:
|
|
return mapping
|
|
|
|
datasource = await db.get(DataSourceConfig, datasource_config_id)
|
|
if datasource is None:
|
|
raise CustomDatasourceRuntimeError("Configuration not found")
|
|
target_schema = (datasource.config or {}).get("target_schema")
|
|
template_body = DEFAULT_MAPPING_TEMPLATES.get(str(target_schema or "")) if target_schema else None
|
|
if not template_body or target_schema not in TARGET_SCHEMAS:
|
|
raise CustomDatasourceRuntimeError(
|
|
"No active mapping template found and no default template available for this target schema"
|
|
)
|
|
|
|
mapping = DataSourceMappingTemplate(
|
|
datasource_config_id=datasource_config_id,
|
|
target_schema=str(target_schema),
|
|
mapping_json=template_body,
|
|
sample_payload_hash=None,
|
|
validation_status="valid",
|
|
version=1,
|
|
is_active=True,
|
|
)
|
|
db.add(mapping)
|
|
await db.commit()
|
|
await db.refresh(mapping)
|
|
return mapping
|
|
|
|
|
|
async def fetch_rest_payload(config: DataSourceConfig, limit_bytes: int) -> Any:
|
|
request_config = config.config or {}
|
|
method = str(request_config.get("method") or request_config.get("request_method") or "GET").upper()
|
|
if method not in {"GET", "POST"}:
|
|
raise CustomDatasourceRuntimeError("Only GET and POST sample requests are supported.")
|
|
|
|
headers = build_request_headers(config.auth_type, config.auth_config or {}, config.headers or {})
|
|
params = build_query_params(config.auth_type, config.auth_config or {}, request_config)
|
|
timeout = float(request_config.get("timeout", 30))
|
|
json_body = request_config.get("json_body")
|
|
if json_body is None and str(request_config.get("body_type") or "").lower() in {"json", ""}:
|
|
candidate = request_config.get("body")
|
|
if isinstance(candidate, (dict, list)):
|
|
json_body = candidate
|
|
|
|
async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
|
|
response = await client.request(
|
|
method,
|
|
config.endpoint,
|
|
headers=headers,
|
|
params=params or None,
|
|
json=json_body,
|
|
)
|
|
response.raise_for_status()
|
|
content = response.content[:limit_bytes]
|
|
if "application/json" in response.headers.get("content-type", ""):
|
|
return json.loads(content.decode(response.encoding or "utf-8"))
|
|
return {"text": content.decode(response.encoding or "utf-8", errors="replace")}
|
|
|
|
|
|
async def run_mapped_rest_config(
|
|
db: AsyncSession,
|
|
datasource: DataSourceConfig,
|
|
) -> dict[str, Any]:
|
|
mapping = await load_active_mapping(db, datasource.id)
|
|
sample = await fetch_rest_payload(datasource, 5_000_000)
|
|
mapped = execute_mapping(sample, mapping.mapping_json, mapping.target_schema)
|
|
if mapped["failed_count"] > 0:
|
|
return {
|
|
"status": "failed",
|
|
"datasource_config_id": datasource.id,
|
|
"mapping_id": mapping.id,
|
|
"mapping_version": mapping.version,
|
|
"target_schema": mapping.target_schema,
|
|
"mapped_count": mapped["mapped_count"],
|
|
"failed_count": mapped["failed_count"],
|
|
"errors": mapped["errors"][:20],
|
|
}
|
|
|
|
request_config = datasource.config or {}
|
|
written_count = await persist_mapped_records(
|
|
db,
|
|
datasource_name=datasource.name,
|
|
datasource_config_id=datasource.id,
|
|
target_schema=mapping.target_schema,
|
|
records=mapped["records"],
|
|
mapping_version=mapping.version,
|
|
delivery_mode=request_config.get("delivery_mode") or "polling",
|
|
transport="http",
|
|
)
|
|
return {
|
|
"status": "success",
|
|
"datasource_config_id": datasource.id,
|
|
"mapping_id": mapping.id,
|
|
"mapping_version": mapping.version,
|
|
"target_schema": mapping.target_schema,
|
|
"fetched_count": mapped["total_items"],
|
|
"mapped_count": mapped["mapped_count"],
|
|
"written_count": written_count,
|
|
}
|
|
|
|
|
|
def _items_from_ws_message(payload: Any, config: dict) -> Any:
|
|
message_path = config.get("ws_message_path")
|
|
items_path = config.get("ws_items_path")
|
|
value = extract_path(payload, message_path) if message_path else payload
|
|
return extract_path(value, items_path) if items_path else value
|
|
|
|
|
|
async def _connect_websocket(endpoint: str, headers: dict[str, str]):
|
|
import websockets
|
|
|
|
try:
|
|
return await websockets.connect(endpoint, additional_headers=headers or None)
|
|
except TypeError:
|
|
return await websockets.connect(endpoint, extra_headers=headers or None)
|
|
|
|
|
|
async def test_websocket_config(config: DataSourceConfig) -> dict[str, Any]:
|
|
if not str(config.endpoint or "").startswith(("ws://", "wss://")):
|
|
raise CustomDatasourceRuntimeError("WebSocket datasource endpoint must start with ws:// or wss://")
|
|
|
|
runtime_config = config.config or {}
|
|
headers = build_request_headers(config.auth_type, config.auth_config or {}, config.headers or {})
|
|
receive_timeout = float(runtime_config.get("receive_timeout_seconds") or runtime_config.get("timeout") or 10)
|
|
async with await _connect_websocket(config.endpoint, headers) as websocket:
|
|
subscribe_message = runtime_config.get("ws_subscribe_message")
|
|
if isinstance(subscribe_message, (dict, list)):
|
|
await websocket.send(json.dumps(subscribe_message))
|
|
elif isinstance(subscribe_message, str) and subscribe_message.strip():
|
|
await websocket.send(subscribe_message)
|
|
raw_message = await asyncio.wait_for(websocket.recv(), timeout=receive_timeout)
|
|
return {
|
|
"success": True,
|
|
"message_preview": raw_message[:1000] if isinstance(raw_message, str) else str(raw_message)[:1000],
|
|
}
|
|
|
|
|
|
async def run_mapped_websocket_config(
|
|
db: AsyncSession,
|
|
datasource: DataSourceConfig,
|
|
*,
|
|
debug_max_messages: int | None = None,
|
|
use_config_debug_max_messages: bool = True,
|
|
) -> dict[str, Any]:
|
|
if not str(datasource.endpoint or "").startswith(("ws://", "wss://")):
|
|
raise CustomDatasourceRuntimeError("WebSocket datasource endpoint must start with ws:// or wss://")
|
|
|
|
mapping = await load_active_mapping(db, datasource.id)
|
|
runtime_config = datasource.config or {}
|
|
max_messages = debug_max_messages
|
|
if max_messages is None and use_config_debug_max_messages:
|
|
max_messages = runtime_config.get("debug_max_messages")
|
|
max_messages = int(max_messages) if max_messages else None
|
|
receive_timeout = float(runtime_config.get("receive_timeout_seconds") or runtime_config.get("timeout") or 30)
|
|
reconnect = bool(runtime_config.get("ws_reconnect", True))
|
|
reconnect_delay = float(runtime_config.get("reconnect_delay_seconds") or 3)
|
|
headers = build_request_headers(datasource.auth_type, datasource.auth_config or {}, datasource.headers or {})
|
|
|
|
messages_seen = 0
|
|
mapped_count = 0
|
|
failed_count = 0
|
|
written_count = 0
|
|
errors: list[dict[str, Any]] = []
|
|
started_at = datetime.now(UTC)
|
|
|
|
while True:
|
|
try:
|
|
async with await _connect_websocket(datasource.endpoint, headers) as websocket:
|
|
subscribe_message = runtime_config.get("ws_subscribe_message")
|
|
if isinstance(subscribe_message, (dict, list)):
|
|
await websocket.send(json.dumps(subscribe_message))
|
|
elif isinstance(subscribe_message, str) and subscribe_message.strip():
|
|
await websocket.send(subscribe_message)
|
|
|
|
while True:
|
|
raw_message = await asyncio.wait_for(websocket.recv(), timeout=receive_timeout)
|
|
messages_seen += 1
|
|
try:
|
|
payload = json.loads(raw_message)
|
|
except json.JSONDecodeError as exc:
|
|
failed_count += 1
|
|
errors.append({"message": "invalid_json", "error": str(exc)})
|
|
continue
|
|
|
|
extracted = _items_from_ws_message(payload, runtime_config)
|
|
try:
|
|
mapped = execute_mapping(extracted, mapping.mapping_json, mapping.target_schema)
|
|
except (MappingError, ValueError) as exc:
|
|
failed_count += 1
|
|
errors.append({"message": "mapping_failed", "error": str(exc)})
|
|
continue
|
|
|
|
mapped_count += mapped["mapped_count"]
|
|
failed_count += mapped["failed_count"]
|
|
if mapped["errors"]:
|
|
errors.extend(mapped["errors"][:5])
|
|
if mapped["records"]:
|
|
written_count += await persist_mapped_records(
|
|
db,
|
|
datasource_name=datasource.name,
|
|
datasource_config_id=datasource.id,
|
|
target_schema=mapping.target_schema,
|
|
records=mapped["records"],
|
|
mapping_version=mapping.version,
|
|
delivery_mode=runtime_config.get("delivery_mode") or "realtime_stream",
|
|
transport="websocket",
|
|
)
|
|
|
|
if max_messages and messages_seen >= max_messages:
|
|
return {
|
|
"status": "success",
|
|
"datasource_config_id": datasource.id,
|
|
"mapping_id": mapping.id,
|
|
"mapping_version": mapping.version,
|
|
"target_schema": mapping.target_schema,
|
|
"messages_seen": messages_seen,
|
|
"mapped_count": mapped_count,
|
|
"failed_count": failed_count,
|
|
"written_count": written_count,
|
|
"errors": errors[:20],
|
|
"execution_time_seconds": (datetime.now(UTC) - started_at).total_seconds(),
|
|
}
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
failed_count += 1
|
|
errors.append({"message": "websocket_error", "error": f"{exc.__class__.__name__}: {exc}"})
|
|
if not reconnect or max_messages:
|
|
return {
|
|
"status": "failed" if written_count == 0 else "partial",
|
|
"datasource_config_id": datasource.id,
|
|
"mapping_id": mapping.id,
|
|
"mapping_version": mapping.version,
|
|
"target_schema": mapping.target_schema,
|
|
"messages_seen": messages_seen,
|
|
"mapped_count": mapped_count,
|
|
"failed_count": failed_count,
|
|
"written_count": written_count,
|
|
"errors": errors[:20],
|
|
}
|
|
await asyncio.sleep(reconnect_delay)
|
|
|
|
|
|
async def run_custom_stream_by_id(config_id: int) -> dict[str, Any]:
|
|
async with async_session_factory() as db:
|
|
datasource = await db.get(DataSourceConfig, config_id)
|
|
if not datasource:
|
|
raise CustomDatasourceRuntimeError("Configuration not found")
|
|
return await run_mapped_websocket_config(
|
|
db,
|
|
datasource,
|
|
use_config_debug_max_messages=False,
|
|
)
|
|
|
|
|
|
def start_custom_stream(config_id: int) -> bool:
|
|
existing = RUNNING_CUSTOM_STREAM_TASKS.get(config_id)
|
|
if existing is not None and not existing.done():
|
|
return False
|
|
task = asyncio.create_task(run_custom_stream_by_id(config_id), name=f"custom-stream:{config_id}")
|
|
RUNNING_CUSTOM_STREAM_TASKS[config_id] = task
|
|
|
|
def _cleanup(done_task: asyncio.Task[Any]) -> None:
|
|
if RUNNING_CUSTOM_STREAM_TASKS.get(config_id) is done_task:
|
|
RUNNING_CUSTOM_STREAM_TASKS.pop(config_id, None)
|
|
|
|
task.add_done_callback(_cleanup)
|
|
return True
|
|
|
|
|
|
async def stop_custom_stream(config_id: int) -> bool:
|
|
task = RUNNING_CUSTOM_STREAM_TASKS.get(config_id)
|
|
if task is None or task.done():
|
|
RUNNING_CUSTOM_STREAM_TASKS.pop(config_id, None)
|
|
return False
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
return True
|
|
return task.cancelled()
|
|
|
|
|
|
def get_custom_stream_status(config_id: int) -> dict[str, Any]:
|
|
task = RUNNING_CUSTOM_STREAM_TASKS.get(config_id)
|
|
return {
|
|
"config_id": config_id,
|
|
"running": bool(task and not task.done()),
|
|
"done": bool(task and task.done()),
|
|
}
|