Amazon Security Lake MCP Server

import asyncio import time from typing import Any, Dict, List, Optional import boto3 import structlog from botocore.exceptions import ClientError, NoCredentialsError from ..config.settings import Settings logger = structlog.get_logger(__name__) class AthenaClient: def __init__(self, settings: Settings): self.settings = settings self._client: Optional[boto3.client] = None self._session: Optional[boto3.Session] = None def _get_session(self) -> boto3.Session: if self._session is None: if self.settings.aws_profile: self._session = boto3.Session( profile_name=self.settings.aws_profile, region_name=self.settings.aws_region ) else: self._session = boto3.Session(region_name=self.settings.aws_region) return self._session def _get_client(self) -> boto3.client: if self._client is None: session = self._get_session() self._client = session.client("athena") return self._client async def execute_query( self, query: str, parameters: Optional[List[str]] = None ) -> List[Dict[str, Any]]: try: client = self._get_client() execution_params = { "QueryString": query, "QueryExecutionContext": { "Database": self.settings.security_lake_database }, "ResultConfiguration": { "OutputLocation": self.settings.athena_output_location }, "WorkGroup": self.settings.athena_workgroup, } if parameters: execution_params["ExecutionParameters"] = parameters logger.info( "Starting Athena query execution", query=query[:200] + "..." if len(query) > 200 else query, database=self.settings.security_lake_database, workgroup=self.settings.athena_workgroup ) response = client.start_query_execution(**execution_params) execution_id = response["QueryExecutionId"] logger.info("Query execution started", execution_id=execution_id) execution_status = await self._wait_for_completion(execution_id) if execution_status != "SUCCEEDED": error_msg = await self._get_execution_error(execution_id) raise RuntimeError(f"Query failed: {error_msg}") results = await self._get_query_results(execution_id) logger.info( "Query completed successfully", execution_id=execution_id, result_count=len(results) ) return results except NoCredentialsError: logger.error("AWS credentials not found") raise RuntimeError( "AWS credentials not configured. Please set up AWS credentials " "or specify an AWS profile." ) except ClientError as e: error_code = e.response.get("Error", {}).get("Code", "Unknown") error_message = e.response.get("Error", {}).get("Message", str(e)) logger.error( "AWS client error", error_code=error_code, error_message=error_message ) raise RuntimeError(f"AWS error ({error_code}): {error_message}") except Exception as e: logger.error("Unexpected error during query execution", error=str(e)) raise async def _wait_for_completion(self, execution_id: str) -> str: client = self._get_client() start_time = time.time() while True: if time.time() - start_time > self.settings.query_timeout_seconds: raise RuntimeError(f"Query timeout after {self.settings.query_timeout_seconds} seconds") response = client.get_query_execution(QueryExecutionId=execution_id) status = response["QueryExecution"]["Status"]["State"] if status in ["SUCCEEDED", "FAILED", "CANCELLED"]: return status await asyncio.sleep(1) async def _get_execution_error(self, execution_id: str) -> str: client = self._get_client() response = client.get_query_execution(QueryExecutionId=execution_id) status_info = response["QueryExecution"]["Status"] return status_info.get("StateChangeReason", "Unknown error") async def _get_query_results(self, execution_id: str) -> List[Dict[str, Any]]: client = self._get_client() results = [] next_token = None while True: params = {"QueryExecutionId": execution_id} if next_token: params["NextToken"] = next_token response = client.get_query_results(**params) # Extract column names from the first row (header) if not results and response["ResultSet"]["Rows"]: columns = [ col["VarCharValue"] for col in response["ResultSet"]["Rows"][0]["Data"] ] # Skip the header row for data processing data_rows = response["ResultSet"]["Rows"][1:] else: data_rows = response["ResultSet"]["Rows"] columns = [] # Process data rows for row in data_rows: if len(results) >= self.settings.max_query_results: logger.warning( "Query result limit reached", limit=self.settings.max_query_results ) return results row_data = {} for i, cell in enumerate(row["Data"]): if i < len(columns): value = cell.get("VarCharValue") row_data[columns[i]] = value if row_data: # Only add non-empty rows results.append(row_data) next_token = response.get("NextToken") if not next_token: break return results async def list_databases(self) -> List[str]: try: client = self._get_client() response = client.list_databases(CatalogName="AwsDataCatalog") return [db["Name"] for db in response["DatabaseList"]] except Exception as e: logger.error("Failed to list databases", error=str(e)) raise async def list_tables(self, database: str) -> List[Dict[str, Any]]: try: client = self._get_client() response = client.list_table_metadata( CatalogName="AwsDataCatalog", DatabaseName=database ) return [ { "name": table["Name"], "type": table.get("TableType", "UNKNOWN"), "columns": len(table.get("Columns", [])), "last_accessed": table.get("LastAccessTime"), "location": table.get("Parameters", {}).get("location") } for table in response["TableMetadataList"] ] except Exception as e: logger.error("Failed to list tables", database=database, error=str(e)) raise async def test_connection(self) -> bool: try: await self.list_databases() logger.info("Athena connection test successful") return True except Exception as e: logger.error("Athena connection test failed", error=str(e)) return False