Prefect MCP Server

from typing import Any, Callable, Dict, List, Optional, Union from uuid import UUID import mcp.types as types from prefect import get_client def get_all_functions() -> list[tuple[Callable, str, str]]: return [ (get_work_pools, "get_work_pools", "Get all work pools"), (get_work_pool, "get_work_pool", "Get a work pool by name"), (create_work_pool, "create_work_pool", "Create a work pool"), (update_work_pool, "update_work_pool", "Update a work pool"), (delete_work_pool, "delete_work_pool", "Delete a work pool"), (get_scheduled_flow_runs, "get_scheduled_flow_runs", "Get scheduled flow runs for a work pool"), (get_work_pool_queues, "get_work_pool_queues", "Get work queues for a work pool"), (get_work_pool_queue, "get_work_pool_queue", "Get a specific work queue in a work pool"), (create_work_pool_queue, "create_work_pool_queue", "Create a work queue in a work pool"), (update_work_pool_queue, "update_work_pool_queue", "Update a work queue in a work pool"), (delete_work_pool_queue, "delete_work_pool_queue", "Delete a work queue from a work pool"), ] async def get_work_pools( limit: Optional[int] = None, offset: Optional[int] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Get all work pools with optional pagination. Args: limit: Maximum number of work pools to return offset: Number of work pools to skip Returns: A list of work pools """ async with get_client() as client: try: # The method name might be different based on Prefect's API work_pools = await client.read_work_pools( limit=limit, offset=offset ) work_pools_result = { "work_pools": [work_pool.dict() for work_pool in work_pools] } return [types.TextContent(type="text", text=str(work_pools_result))] except Exception as e: error_message = f"Error fetching work pools: {str(e)}" return [types.TextContent(type="text", text=error_message)] async def get_work_pool( name: str, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Get a work pool by name. Args: name: The work pool name Returns: Work pool details """ async with get_client() as client: try: work_pool = await client.read_work_pool(name) return [types.TextContent(type="text", text=str(work_pool.dict()))] except Exception as e: error_message = f"Error fetching work pool '{name}': {str(e)}" return [types.TextContent(type="text", text=error_message)] async def create_work_pool( name: str, type: str, description: Optional[str] = None, base_job_template: Optional[Dict] = None, is_paused: Optional[bool] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Create a work pool. Args: name: The name for the work pool type: The type of work pool (e.g., 'kubernetes', 'process', etc.) description: Optional description base_job_template: Optional base job template as JSON is_paused: Whether the work pool should be paused Returns: Details of the created work pool """ async with get_client() as client: try: work_pool = await client.create_work_pool( name=name, work_pool_type=type, description=description, base_job_template=base_job_template or {}, is_paused=is_paused ) return [types.TextContent(type="text", text=str(work_pool.dict()))] except Exception as e: error_message = f"Error creating work pool: {str(e)}" return [types.TextContent(type="text", text=error_message)] async def update_work_pool( name: str, description: Optional[str] = None, base_job_template: Optional[Dict] = None, is_paused: Optional[bool] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Update a work pool. Args: name: The work pool name description: New description base_job_template: New base job template as JSON is_paused: New paused status Returns: Details of the updated work pool """ async with get_client() as client: try: # Prepare update data update_data = {} if description is not None: update_data["description"] = description if base_job_template is not None: update_data["base_job_template"] = base_job_template if is_paused is not None: update_data["is_paused"] = is_paused work_pool = await client.update_work_pool( name=name, **update_data ) return [types.TextContent(type="text", text=str(work_pool.dict()))] except Exception as e: error_message = f"Error updating work pool '{name}': {str(e)}" return [types.TextContent(type="text", text=error_message)] async def delete_work_pool( name: str, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Delete a work pool by name. Args: name: The work pool name Returns: Confirmation message """ async with get_client() as client: try: await client.delete_work_pool(name) return [types.TextContent(type="text", text=f"Work pool '{name}' deleted successfully.")] except Exception as e: error_message = f"Error deleting work pool '{name}': {str(e)}" return [types.TextContent(type="text", text=error_message)] async def get_scheduled_flow_runs( work_pool_name: str, limit: Optional[int] = None, scheduled_before: Optional[str] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Get scheduled flow runs for a work pool. Args: work_pool_name: The work pool name limit: Maximum number of flow runs to return scheduled_before: ISO formatted timestamp Returns: A list of scheduled flow runs """ async with get_client() as client: try: flow_runs = await client.get_scheduled_flow_runs( work_pool_name=work_pool_name, limit=limit, scheduled_before=scheduled_before ) flow_runs_result = { "scheduled_flow_runs": [flow_run.dict() for flow_run in flow_runs] } return [types.TextContent(type="text", text=str(flow_runs_result))] except Exception as e: error_message = f"Error fetching scheduled flow runs for work pool '{work_pool_name}': {str(e)}" return [types.TextContent(type="text", text=error_message)] async def get_work_pool_queues( work_pool_name: str, limit: Optional[int] = None, offset: Optional[int] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Get work queues for a work pool. Args: work_pool_name: The work pool name limit: Maximum number of queues to return offset: Number of queues to skip Returns: A list of work queues in the work pool """ async with get_client() as client: try: queues = await client.read_work_pool_queues( work_pool_name=work_pool_name, limit=limit, offset=offset ) queues_result = { "work_pool_queues": [queue.dict() for queue in queues] } return [types.TextContent(type="text", text=str(queues_result))] except Exception as e: error_message = f"Error fetching work queues for work pool '{work_pool_name}': {str(e)}" return [types.TextContent(type="text", text=error_message)] async def get_work_pool_queue( work_pool_name: str, queue_name: str, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Get a specific work queue in a work pool. Args: work_pool_name: The work pool name queue_name: The work queue name Returns: Work queue details """ async with get_client() as client: try: queue = await client.read_work_pool_queue( work_pool_name=work_pool_name, queue_name=queue_name ) return [types.TextContent(type="text", text=str(queue.dict()))] except Exception as e: error_message = f"Error fetching work queue '{queue_name}' in work pool '{work_pool_name}': {str(e)}" return [types.TextContent(type="text", text=error_message)] async def create_work_pool_queue( work_pool_name: str, name: str, description: Optional[str] = None, is_paused: Optional[bool] = None, concurrency_limit: Optional[int] = None, priority: Optional[int] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Create a work queue in a work pool. Args: work_pool_name: The work pool name name: The name for the work queue description: Optional description is_paused: Whether the queue is paused concurrency_limit: Optional concurrency limit priority: Optional priority Returns: Details of the created work queue """ async with get_client() as client: try: queue = await client.create_work_pool_queue( work_pool_name=work_pool_name, name=name, description=description, is_paused=is_paused, concurrency_limit=concurrency_limit, priority=priority ) return [types.TextContent(type="text", text=str(queue.dict()))] except Exception as e: error_message = f"Error creating work queue in work pool '{work_pool_name}': {str(e)}" return [types.TextContent(type="text", text=error_message)] async def update_work_pool_queue( work_pool_name: str, queue_name: str, description: Optional[str] = None, is_paused: Optional[bool] = None, concurrency_limit: Optional[int] = None, priority: Optional[int] = None, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Update a work queue in a work pool. Args: work_pool_name: The work pool name queue_name: The work queue name description: New description is_paused: New paused status concurrency_limit: New concurrency limit priority: New priority Returns: Details of the updated work queue """ async with get_client() as client: try: # Prepare update data update_data = {} if description is not None: update_data["description"] = description if is_paused is not None: update_data["is_paused"] = is_paused if concurrency_limit is not None: update_data["concurrency_limit"] = concurrency_limit if priority is not None: update_data["priority"] = priority queue = await client.update_work_pool_queue( work_pool_name=work_pool_name, queue_name=queue_name, **update_data ) return [types.TextContent(type="text", text=str(queue.dict()))] except Exception as e: error_message = f"Error updating work queue '{queue_name}' in work pool '{work_pool_name}': {str(e)}" return [types.TextContent(type="text", text=error_message)] async def delete_work_pool_queue( work_pool_name: str, queue_name: str, ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]: """ Delete a work queue from a work pool. Args: work_pool_name: The work pool name queue_name: The work queue name Returns: Confirmation message """ async with get_client() as client: try: await client.delete_work_pool_queue( work_pool_name=work_pool_name, queue_name=queue_name ) return [types.TextContent(type="text", text=f"Work queue '{queue_name}' deleted from work pool '{work_pool_name}' successfully.")] except Exception as e: error_message = f"Error deleting work queue '{queue_name}' from work pool '{work_pool_name}': {str(e)}" return [types.TextContent(type="text", text=error_message)]