refactor: Modularize MCP server with command batching
- Split monolithic mcp/server.py (1874 lines) into haproxy_mcp/ package:
- config.py: Configuration constants and environment variables
- exceptions.py: Custom exception classes
- validation.py: Input validation functions
- haproxy_client.py: HAProxy Runtime API client with batch support
- file_ops.py: Atomic file operations with locking
- utils.py: CSV parsing utilities
- tools/: MCP tools organized by function
- domains.py: Domain management (3 tools)
- servers.py: Server management (7 tools)
- health.py: Health checks (3 tools)
- monitoring.py: Monitoring (4 tools)
- configuration.py: Config management (4 tools)
- Add haproxy_cmd_batch() for sending multiple commands in single TCP connection
- Optimize server operations: 1 connection instead of 2 per server
- Optimize startup restore: All servers in 1 connection (was 2×N)
- Update type hints to Python 3.9+ style (built-in generics)
- Remove unused imports and functions
- Update CLAUDE.md with new structure and performance notes
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
kaffa
223
haproxy_mcp/tools/domains.py
Normal file
223
haproxy_mcp/tools/domains.py
Normal file
@@ -0,0 +1,223 @@
|
||||
"""Domain management tools for HAProxy MCP Server."""
|
||||
|
||||
import fcntl
|
||||
from typing import Annotated
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from ..config import (
|
||||
MAP_FILE,
|
||||
MAP_FILE_CONTAINER,
|
||||
POOL_COUNT,
|
||||
MAX_SLOTS,
|
||||
StateField,
|
||||
STATE_MIN_COLUMNS,
|
||||
logger,
|
||||
)
|
||||
from ..exceptions import HaproxyError
|
||||
from ..validation import validate_domain, validate_ip
|
||||
from ..haproxy_client import haproxy_cmd
|
||||
from ..file_ops import (
|
||||
get_map_contents,
|
||||
save_map_file,
|
||||
get_domain_backend,
|
||||
is_legacy_backend,
|
||||
add_server_to_config,
|
||||
remove_server_from_config,
|
||||
remove_domain_from_config,
|
||||
)
|
||||
|
||||
|
||||
def register_domain_tools(mcp):
|
||||
"""Register domain management tools with MCP server."""
|
||||
|
||||
@mcp.tool()
|
||||
def haproxy_list_domains(
|
||||
include_wildcards: Annotated[bool, Field(default=False, description="Include wildcard entries (.example.com). Default: False")]
|
||||
) -> str:
|
||||
"""List all configured domains with their backend servers."""
|
||||
try:
|
||||
domains = []
|
||||
state = haproxy_cmd("show servers state")
|
||||
|
||||
# Build server map from HAProxy state
|
||||
server_map: dict[str, list] = {}
|
||||
for line in state.split("\n"):
|
||||
parts = line.split()
|
||||
if len(parts) >= STATE_MIN_COLUMNS and parts[StateField.SRV_ADDR] != "0.0.0.0":
|
||||
backend = parts[StateField.BE_NAME]
|
||||
if backend not in server_map:
|
||||
server_map[backend] = []
|
||||
server_map[backend].append(
|
||||
f"{parts[StateField.SRV_NAME]}={parts[StateField.SRV_ADDR]}:{parts[StateField.SRV_PORT]}"
|
||||
)
|
||||
|
||||
# Read from domains.map
|
||||
seen_domains: set[str] = set()
|
||||
for domain, backend in get_map_contents():
|
||||
# Skip wildcard entries unless explicitly requested
|
||||
if domain.startswith(".") and not include_wildcards:
|
||||
continue
|
||||
if domain in seen_domains:
|
||||
continue
|
||||
seen_domains.add(domain)
|
||||
servers = server_map.get(backend, ["(none)"])
|
||||
if domain.startswith("."):
|
||||
backend_type = "wildcard"
|
||||
elif backend.startswith("pool_"):
|
||||
backend_type = "pool"
|
||||
else:
|
||||
backend_type = "static"
|
||||
domains.append(f"• {domain} -> {backend} ({backend_type}): {', '.join(servers)}")
|
||||
|
||||
return "\n".join(domains) if domains else "No domains configured"
|
||||
except HaproxyError as e:
|
||||
return f"Error: {e}"
|
||||
|
||||
@mcp.tool()
|
||||
def haproxy_add_domain(
|
||||
domain: Annotated[str, Field(description="Domain name to add (e.g., api.example.com, example.com)")],
|
||||
ip: Annotated[str, Field(default="", description="Optional: Initial server IP. If provided, adds server to slot 1")],
|
||||
http_port: Annotated[int, Field(default=80, description="HTTP port for backend server (default: 80)")]
|
||||
) -> str:
|
||||
"""Add a new domain to HAProxy (no reload required).
|
||||
|
||||
Creates domain→pool mapping. Use haproxy_add_server to add more servers later.
|
||||
|
||||
Example: haproxy_add_domain("api.example.com", ip="10.0.0.1", http_port=8080)
|
||||
"""
|
||||
# Validate inputs
|
||||
if domain.startswith("."):
|
||||
return "Error: Domain cannot start with '.' (wildcard entries are added automatically)"
|
||||
if not validate_domain(domain):
|
||||
return "Error: Invalid domain format"
|
||||
if not validate_ip(ip, allow_empty=True):
|
||||
return "Error: Invalid IP address format"
|
||||
if not (1 <= http_port <= 65535):
|
||||
return "Error: Port must be between 1 and 65535"
|
||||
|
||||
# Use file locking for the entire pool allocation operation
|
||||
lock_path = f"{MAP_FILE}.lock"
|
||||
with open(lock_path, 'w') as lock_file:
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
|
||||
try:
|
||||
# Read map contents once for both existence check and pool lookup
|
||||
entries = get_map_contents()
|
||||
|
||||
# Check if domain already exists (using cached entries)
|
||||
for domain_entry, backend in entries:
|
||||
if domain_entry == domain:
|
||||
return f"Error: Domain {domain} already exists (mapped to {backend})"
|
||||
|
||||
# Find available pool (using cached entries)
|
||||
used_pools: set[str] = set()
|
||||
for _, backend in entries:
|
||||
if backend.startswith("pool_"):
|
||||
used_pools.add(backend)
|
||||
|
||||
pool = None
|
||||
for i in range(1, POOL_COUNT + 1):
|
||||
pool_name = f"pool_{i}"
|
||||
if pool_name not in used_pools:
|
||||
pool = pool_name
|
||||
break
|
||||
if not pool:
|
||||
return f"Error: All {POOL_COUNT} pool backends are in use"
|
||||
|
||||
try:
|
||||
# Save to disk first (atomic write for persistence)
|
||||
entries.append((domain, pool))
|
||||
entries.append((f".{domain}", pool))
|
||||
try:
|
||||
save_map_file(entries)
|
||||
except IOError as e:
|
||||
return f"Error: Failed to save map file: {e}"
|
||||
|
||||
# Then update HAProxy map via Runtime API
|
||||
try:
|
||||
haproxy_cmd(f"add map {MAP_FILE_CONTAINER} {domain} {pool}")
|
||||
haproxy_cmd(f"add map {MAP_FILE_CONTAINER} .{domain} {pool}")
|
||||
except HaproxyError as e:
|
||||
# Rollback: remove the domain we just added from entries and re-save
|
||||
rollback_entries = [(d, b) for d, b in entries if d != domain and d != f".{domain}"]
|
||||
try:
|
||||
save_map_file(rollback_entries)
|
||||
except IOError:
|
||||
logger.error("Failed to rollback map file after HAProxy error")
|
||||
return f"Error: Failed to update HAProxy map: {e}"
|
||||
|
||||
# If IP provided, add server to slot 1
|
||||
if ip:
|
||||
# Save server config to disk first
|
||||
add_server_to_config(domain, 1, ip, http_port)
|
||||
|
||||
try:
|
||||
server = f"{pool}_1"
|
||||
haproxy_cmd(f"set server {pool}/{server} addr {ip} port {http_port}")
|
||||
haproxy_cmd(f"set server {pool}/{server} state ready")
|
||||
except HaproxyError as e:
|
||||
# Rollback server config on failure
|
||||
remove_server_from_config(domain, 1)
|
||||
return f"Domain {domain} added to {pool} but server config failed: {e}"
|
||||
|
||||
return f"Domain {domain} added to {pool} with server {ip}:{http_port}"
|
||||
|
||||
return f"Domain {domain} added to {pool} (no servers configured)"
|
||||
|
||||
except HaproxyError as e:
|
||||
return f"Error: {e}"
|
||||
finally:
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
|
||||
|
||||
@mcp.tool()
|
||||
def haproxy_remove_domain(
|
||||
domain: Annotated[str, Field(description="Domain name to remove (e.g., api.example.com)")]
|
||||
) -> str:
|
||||
"""Remove a domain from HAProxy (no reload required)."""
|
||||
if not validate_domain(domain):
|
||||
return "Error: Invalid domain format"
|
||||
|
||||
# Look up the domain in the map
|
||||
backend = get_domain_backend(domain)
|
||||
if not backend:
|
||||
return f"Error: Domain {domain} not found"
|
||||
|
||||
# Check if this is a legacy backend (not a pool)
|
||||
if is_legacy_backend(backend):
|
||||
return f"Error: Cannot remove legacy domain {domain} (uses static backend {backend})"
|
||||
|
||||
try:
|
||||
# Save to disk first (atomic write for persistence)
|
||||
entries = get_map_contents()
|
||||
new_entries = [(d, b) for d, b in entries if d != domain and d != f".{domain}"]
|
||||
save_map_file(new_entries)
|
||||
|
||||
# Remove from persistent server config
|
||||
remove_domain_from_config(domain)
|
||||
|
||||
# Clear map entries via Runtime API (immediate effect)
|
||||
haproxy_cmd(f"del map {MAP_FILE_CONTAINER} {domain}")
|
||||
try:
|
||||
haproxy_cmd(f"del map {MAP_FILE_CONTAINER} .{domain}")
|
||||
except HaproxyError as e:
|
||||
logger.warning("Failed to remove wildcard entry for %s: %s", domain, e)
|
||||
|
||||
# Disable all servers in the pool (reset to 0.0.0.0:0)
|
||||
for slot in range(1, MAX_SLOTS + 1):
|
||||
server = f"{backend}_{slot}"
|
||||
try:
|
||||
haproxy_cmd(f"set server {backend}/{server} state maint")
|
||||
haproxy_cmd(f"set server {backend}/{server} addr 0.0.0.0 port 0")
|
||||
except HaproxyError as e:
|
||||
logger.warning(
|
||||
"Failed to clear server %s/%s for domain %s: %s",
|
||||
backend, server, domain, e
|
||||
)
|
||||
# Continue with remaining cleanup
|
||||
|
||||
return f"Domain {domain} removed from {backend}"
|
||||
|
||||
except IOError as e:
|
||||
return f"Error: Failed to update map file: {e}"
|
||||
except HaproxyError as e:
|
||||
return f"Error: {e}"
|
||||
Reference in New Issue
Block a user