AI_DATA/doris-mcp-server
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Performance]Add a controllable MCP Server DB Pool permission authentication system (#53)

Browse Source 
* 0.5.1 Version

* fix 0.5.1 schema async bug

* fix security bug

* fix security bug

* Add complete Token, JWT, OAuth authentication system

* Add complete Token, JWT, OAuth authentication system

* Add complete Token, JWT, OAuth authentication system

* Add complete Token, JWT, OAuth authentication system

* Add a controllable MCP Server DB Pool permission authentication system, connect it with the Doris permission system, and provide it to enterprise-level applications concurrently with the multi-Worker mode.
This commit is contained in:
Yijia Su
2025-09-02 18:40:48 +08:00
committed by GitHub
parent c3d487ccdd
commit f99399c6c7
10 changed files with 636 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
doris_mcp_server/utils/config.py
View File

@@ -372,19 +372,34 @@ class DorisConfig:
config = cls()
# Database configuration
config.database.host = os.getenv("DORIS_HOST", config.database.host)
config.database.port = int(os.getenv("DORIS_PORT", str(config.database.port)))
config.database.user = os.getenv("DORIS_USER", config.database.user)
config.database.password = os.getenv("DORIS_PASSWORD", config.database.password)
config.database.database = os.getenv("DORIS_DATABASE", config.database.database)
config.database.fe_http_port = int(os.getenv("DORIS_FE_HTTP_PORT", str(config.database.fe_http_port)))
# Database configuration - handle empty strings properly
doris_host = os.getenv("DORIS_HOST", "").strip()
config.database.host = doris_host if doris_host else config.database.host
doris_port = os.getenv("DORIS_PORT", "").strip()
if doris_port and doris_port.isdigit():
config.database.port = int(doris_port)
doris_user = os.getenv("DORIS_USER", "").strip()
config.database.user = doris_user if doris_user else config.database.user
doris_password = os.getenv("DORIS_PASSWORD", "")
config.database.password = doris_password if doris_password else config.database.password
doris_database = os.getenv("DORIS_DATABASE", "").strip()
config.database.database = doris_database if doris_database else config.database.database
doris_fe_http_port = os.getenv("DORIS_FE_HTTP_PORT", "").strip()
if doris_fe_http_port and doris_fe_http_port.isdigit():
config.database.fe_http_port = int(doris_fe_http_port)
# BE nodes configuration
be_hosts_env = os.getenv("DORIS_BE_HOSTS", "")
if be_hosts_env:
config.database.be_hosts = [host.strip() for host in be_hosts_env.split(",") if host.strip()]
config.database.be_webserver_port = int(os.getenv("DORIS_BE_WEBSERVER_PORT", str(config.database.be_webserver_port)))
be_webserver_port = os.getenv("DORIS_BE_WEBSERVER_PORT", "").strip()
if be_webserver_port and be_webserver_port.isdigit():
config.database.be_webserver_port = int(be_webserver_port)
# Arrow Flight SQL Configuration
fe_arrow_port_env = os.getenv("FE_ARROW_FLIGHT_SQL_PORT")
@@ -557,7 +572,9 @@ class DorisConfig:
# Server configuration
config.server_name = os.getenv("SERVER_NAME", config.server_name)
config.server_version = os.getenv("SERVER_VERSION", config.server_version)
config.server_port = int(os.getenv("SERVER_PORT", str(config.server_port)))
server_port = os.getenv("SERVER_PORT", "").strip()
if server_port and server_port.isdigit():
config.server_port = int(server_port)
config.temp_files_dir = os.getenv("TEMP_FILES_DIR", config.temp_files_dir)
return config

336
doris_mcp_server/utils/db.py
View File

@@ -240,15 +240,30 @@ class DorisConnectionManager:
Uses direct connection pool management with proper synchronization
Implements connection pool health monitoring and proactive cleanup
Supports token-bound database configurations for multi-tenant access
"""
def __init__(self, config, security_manager=None):
def __init__(self, config, security_manager=None, token_manager=None):
self.config = config
self.pool: Pool | None = None
self.logger = get_logger(__name__)
self.security_manager = security_manager
self.token_manager = token_manager # Token manager for token-bound DB config
self.session_cache = DorisSessionCache(self)
# Store original database config for fallback
self.original_db_config = {
'host': config.database.host,
'port': config.database.port,
'user': config.database.user,
'password': config.database.password,
'database': config.database.database,
'charset': config.database.charset
}
# Current active database config (may be overridden by token-bound config)
self.active_db_config = self.original_db_config.copy()
# Connection pool state management
self.pool_recovering = False
@@ -267,14 +282,7 @@ class DorisConnectionManager:
# Database connection parameters from config.database
self.pool_recovery_lock = self._recovery_lock # Compatibility alias
self.host = config.database.host
self.port = config.database.port
self.user = config.database.user
self.password = config.database.password
self.database = config.database.database
# Convert charset to aiomysql compatible format
charset_map = {"UTF8": "utf8", "UTF8MB4": "utf8mb4"}
self.charset = charset_map.get(config.database.charset.upper(), config.database.charset.lower())
self._update_db_params_from_config(self.active_db_config)
self.connect_timeout = config.database.connection_timeout
# Connection pool parameters - more conservative settings
@@ -285,12 +293,307 @@ class DorisConnectionManager:
# 🔧 FIX: Add missing monitoring parameters that were removed during refactoring
self.health_check_interval = 30 # seconds
self.pool_warmup_size = 3 # connections to maintain
def _update_db_params_from_config(self, db_config: dict):
"""Update database connection parameters from config dictionary"""
self.host = db_config['host']
self.port = db_config['port']
self.user = db_config['user']
self.password = db_config['password']
self.database = db_config['database']
# Convert charset to aiomysql compatible format
charset_map = {"UTF8": "utf8", "UTF8MB4": "utf8mb4"}
self.charset = charset_map.get(db_config['charset'].upper(), db_config['charset'].lower())
def _is_config_empty(self, config_value) -> bool:
"""Check if a config value is empty (None, empty string, or 'null')"""
return config_value is None or config_value == '' or str(config_value).lower() == 'null'
def _has_valid_global_config(self) -> bool:
"""Check if global database configuration is valid and non-empty"""
return (not self._is_config_empty(self.original_db_config['host']) and
not self._is_config_empty(self.original_db_config['user']))
def _find_available_token_with_db_config(self) -> str:
"""Find the first available token with database configuration
Returns:
Raw token string if found, empty string if not found
"""
if not self.token_manager:
return ""
try:
for token_hash, token_info in self.token_manager._tokens.items():
if (token_info.database_config and
token_info.is_active and
not self._is_config_empty(token_info.database_config.host) and
not self._is_config_empty(token_info.database_config.user)):
# We need to find the raw token from the hash
# This is a bit tricky since we only store hashes
# We'll need to use the admin token from tokens.json if it has db config
if token_info.token_id == 'admin-token':
# Try the known admin token
return 'doris_admin_token_123456'
elif 'tenant' in token_info.token_id:
# For tenant tokens, we'll need a different approach
# For now, skip these as we don't know the raw token
continue
return ""
except Exception as e:
self.logger.error(f"Error finding available token: {e}")
return ""
async def configure_for_token(self, token: str) -> tuple[bool, str]:
"""Configure connection manager for token with new priority logic
Priority: Token-bound DB config > .env config > error
Args:
token: Authentication token to get database config for
Returns:
(success: bool, config_source: str): Result and which config was used
Raises:
RuntimeError: If no valid database configuration is available
"""
try:
# Priority 1: Try token-bound database config first
if self.token_manager:
db_config = self.token_manager.get_database_config_by_token(token)
if db_config:
# Convert DatabaseConfig to dictionary
token_db_config = {
'host': db_config.host,
'port': db_config.port,
'user': db_config.user,
'password': db_config.password,
'database': db_config.database,
'charset': db_config.charset
}
# Check if token-bound config is valid
if (not self._is_config_empty(token_db_config['host']) and
not self._is_config_empty(token_db_config['user'])):
self.logger.info(f"Using token-bound database configuration for host: {token_db_config['host']}")
self.active_db_config = token_db_config
self._update_db_params_from_config(self.active_db_config)
# Create/recreate connection pool with token-bound config
await self._ensure_pool_with_current_config()
return True, "token-bound"
# Priority 2: Use global .env config if available
if self._has_valid_global_config():
self.logger.info("Using global .env database configuration")
self.active_db_config = self.original_db_config.copy()
self._update_db_params_from_config(self.active_db_config)
# Create/recreate connection pool with global config
await self._ensure_pool_with_current_config()
return True, "global-env"
# Priority 3: No valid configuration available
error_msg = (
"No valid database configuration available for this token. "
"Please contact administrator to:\n"
"1. Add database configuration to tokens.json for this token, OR\n"
"2. Configure valid global database settings in .env file\n"
"Required fields: DB_HOST, DB_USER"
)
self.logger.error(error_msg)
raise RuntimeError(error_msg)
except Exception as e:
self.logger.error(f"Failed to configure database for token: {e}")
raise
async def _ensure_pool_with_current_config(self):
"""Ensure connection pool exists with current configuration"""
try:
# If pool exists with different config, need to recreate it
# If no pool exists, create one with current config
if self.pool and not self.pool.closed:
# Since we can't reliably check pool config attributes,
# we'll recreate the pool if we detect a potential config change
# by checking if current config differs from what we stored
pool_needs_recreation = False
# Compare current config with what we might have used before
if hasattr(self, '_last_pool_config'):
current_config = {
'host': self.host,
'port': self.port,
'user': self.user,
'database': self.database
}
if current_config != self._last_pool_config:
pool_needs_recreation = True
if pool_needs_recreation:
self.logger.info("Database configuration changed, recreating connection pool")
await self._recreate_pool()
elif not self.pool:
self.logger.info("Creating connection pool with current configuration")
await self._create_pool_with_current_config()
# Test the connection immediately
if not await self._test_pool_health():
raise RuntimeError(f"Database connection test failed for {self.host}:{self.port}")
except Exception as e:
self.logger.error(f"Failed to ensure connection pool: {e}")
raise
async def _create_pool_with_current_config(self):
"""Create connection pool with current database configuration"""
try:
self.pool = await aiomysql.create_pool(
host=self.host,
port=self.port,
user=self.user,
password=self.password,
db=self.database,
charset=self.charset,
minsize=self.minsize,
maxsize=self.maxsize,
pool_recycle=self.pool_recycle,
connect_timeout=self.connect_timeout,
autocommit=True
)
# Store the current config for comparison later
self._last_pool_config = {
'host': self.host,
'port': self.port,
'user': self.user,
'database': self.database
}
# Test initial connection
if not await self._test_pool_health():
raise RuntimeError("Connection pool health check failed")
# Start background monitoring tasks if not already running
if not self.pool_health_check_task or self.pool_health_check_task.done():
self.pool_health_check_task = asyncio.create_task(self._pool_health_monitor())
if not self.pool_cleanup_task or self.pool_cleanup_task.done():
self.pool_cleanup_task = asyncio.create_task(self._pool_cleanup_monitor())
# Perform initial pool warmup
await self._warmup_pool()
self.logger.info(f"Connection pool created successfully with {self.host}:{self.port}")
except Exception as e:
self.logger.error(f"Failed to create connection pool: {e}")
raise
async def _recreate_pool(self):
"""Recreate connection pool with current database configuration"""
try:
# Close existing pool
if self.pool and not self.pool.closed:
self.pool.close()
await self.pool.wait_closed()
self.pool = None
# Create new pool with current config
await self._create_pool_with_current_config()
except Exception as e:
self.logger.error(f"Failed to recreate connection pool: {e}")
raise
def validate_database_configuration(self) -> tuple[bool, str]:
"""Validate database configuration completeness
Returns:
(is_valid, error_message): Configuration validation result
"""
# Check if Token authentication is enabled
token_auth_enabled = getattr(self.config.security, 'enable_token_auth', False)
# Check if tokens.json exists and has valid tokens with database configs
tokens_file_available = False
token_bound_configs_available = False
if self.token_manager:
try:
# Check if tokens.json file exists
import os
tokens_file_path = getattr(self.token_manager, 'token_file_path', 'tokens.json')
tokens_file_available = os.path.exists(tokens_file_path)
# Check if any tokens have database configurations
if tokens_file_available or self.token_manager._tokens:
for token_hash, token_info in self.token_manager._tokens.items():
if token_info.database_config:
token_bound_configs_available = True
break
except Exception:
pass
# Validate .env database configuration
env_config_valid = self._has_valid_global_config()
# Decision logic
if token_auth_enabled:
if tokens_file_available:
# tokens.json exists - either .env OR token-bound config must be valid
if env_config_valid or token_bound_configs_available:
return True, "Configuration valid"
else:
return False, (
"Token authentication is enabled and tokens.json exists, but no valid database "
"configuration found. Please provide either:\n"
"1. Valid database configuration in .env file (DB_HOST, DB_USER, etc.)\n"
"2. Database configuration in tokens.json for at least one token"
)
else:
# tokens.json does not exist - must have valid .env config
if env_config_valid:
return True, "Configuration valid"
else:
return False, (
"Token authentication is enabled but tokens.json file not found. "
"Either:\n"
"1. Create tokens.json file with token configurations\n"
"2. Provide valid database configuration in .env file (DB_HOST, DB_USER, etc.)"
)
else:
# Token auth is disabled, must have valid .env config
if env_config_valid:
return True, "Configuration valid"
else:
return False, (
"Token authentication is disabled. Valid database configuration is required "
"in .env file (DB_HOST, DB_USER, etc.)"
)
async def initialize(self):
"""Initialize connection pool with health monitoring"""
try:
# First validate configuration
is_valid, error_message = self.validate_database_configuration()
if not is_valid:
self.logger.error(f"Database configuration validation failed: {error_message}")
raise RuntimeError(f"Database configuration validation failed:\n{error_message}")
self.logger.info(f"Database configuration validated successfully")
self.logger.info(f"Initializing connection pool to {self.host}:{self.port}")
# Only create connection pool if we have valid global config
# Token-bound configs will be handled dynamically during requests
if not self._has_valid_global_config():
self.logger.info("No valid global database config, pool will be created dynamically for token-bound configs")
return
# Create connection pool
self.pool = await aiomysql.create_pool(
host=self.host,
@@ -592,7 +895,20 @@ class DorisConnectionManager:
# Check if pool is available
if not self.pool:
self.logger.warning("Connection pool is not available, attempting recovery...")
await self._recover_pool_with_lock()
# Try to use token-bound configuration if available
if self.token_manager and not self._has_valid_global_config():
available_token = self._find_available_token_with_db_config()
if available_token:
self.logger.info(f"Using token-bound configuration for pool creation: {available_token}")
try:
await self.configure_for_token(available_token)
except Exception as e:
self.logger.error(f"Failed to configure with token-bound config: {e}")
# Fallback to recovery
if not self.pool:
await self._recover_pool_with_lock()
if not self.pool:
raise RuntimeError("Connection pool is not available and recovery failed")

4
doris_mcp_server/utils/query_executor.py
View File

@@ -426,6 +426,10 @@ class DorisQueryExecutor:
self, query_request: QueryRequest, auth_context
) -> QueryResult:
"""Internal query execution"""
# Database configuration should already be handled during authentication
# No need to configure again during query execution
# Optimize query
optimized_sql = await self.query_optimizer.optimize_query(
query_request.sql, {"user_roles": getattr(auth_context, 'roles', [])}

61
doris_mcp_server/utils/security.py
View File

@@ -47,11 +47,16 @@ class SecurityLevel(Enum):
class AuthContext:
"""Authentication context for audit and session tracking"""
token_id: str # Token identifier for audit logging
token_id: str = "" # Token identifier for audit logging
user_id: str = "" # User identifier
roles: list[str] = field(default_factory=list) # User roles
permissions: list[str] = field(default_factory=list) # User permissions
security_level: 'SecurityLevel' = field(default_factory=lambda: SecurityLevel.INTERNAL) # Security level
client_ip: str = "unknown" # Client IP address
session_id: str = "" # Session identifier
login_time: datetime = field(default_factory=datetime.utcnow)
last_activity: datetime | None = None
token: str = "" # Raw token for token-bound database configuration
@dataclass
@@ -84,12 +89,13 @@ class DorisSecurityManager:
Provides complete security control functionality, including authentication, authorization, SQL security validation and data masking
"""
def __init__(self, config):
def __init__(self, config, connection_manager=None):
self.config = config
self.logger = get_logger(__name__)
self.connection_manager = connection_manager
# Initialize security components
self.auth_provider = AuthenticationProvider(config)
self.auth_provider = AuthenticationProvider(config, self)
self.authz_provider = AuthorizationProvider(config)
self.sql_validator = SQLSecurityValidator(config)
self.masking_processor = DataMaskingProcessor(config)
@@ -226,6 +232,10 @@ class DorisSecurityManager:
# Return anonymous context when no authentication is enabled
return AuthContext(
token_id="anonymous",
user_id="anonymous",
roles=["anonymous"],
permissions=["read"],
security_level=SecurityLevel.PUBLIC,
client_ip=auth_info.get("client_ip", "unknown"),
session_id="anonymous_session"
)
@@ -392,18 +402,50 @@ class DorisSecurityManager:
return {"error": "Token manager not initialized"}
return self.auth_provider.token_manager.get_token_stats()
async def _validate_token_database_config(self, token: str, token_info) -> None:
"""Validate database configuration for token immediately during authentication
This ensures database connectivity issues are caught at authentication time,
not during query execution, providing better user experience.
Args:
token: Raw authentication token
token_info: TokenInfo object from token validation
Raises:
ValueError: If database configuration is invalid or connection fails
"""
try:
if not self.connection_manager:
self.logger.warning("Connection manager not available for immediate database validation")
return
# Configure and test database connection for this token
success, config_source = await self.connection_manager.configure_for_token(token)
if success:
self.logger.info(f"Database configuration validated successfully for token {token_info.token_id} (source: {config_source})")
else:
raise ValueError("Database configuration validation failed")
except Exception as e:
error_msg = f"Database configuration validation failed for token {token_info.token_id}: {str(e)}"
self.logger.error(error_msg)
raise ValueError(error_msg)
class AuthenticationProvider:
"""Authentication provider"""
def __init__(self, config):
def __init__(self, config, security_manager=None):
self.config = config
self.logger = get_logger(__name__)
self.session_cache = {}
self.jwt_manager = None
self.oauth_provider = None
self.token_manager = None
self.security_manager = security_manager
# Initialize authentication providers based on individual switches
auth_methods_enabled = []
@@ -583,12 +625,21 @@ class AuthenticationProvider:
token_info = validation_result.token_info
# Immediately validate database configuration for this token
if self.security_manager:
await self.security_manager._validate_token_database_config(token, token_info)
return AuthContext(
token_id=token_info.token_id,
user_id=token_info.token_id, # Use token_id as user_id for token auth
roles=["token_user"], # Default role for token users
permissions=["read", "write"], # Default permissions for token users
security_level=SecurityLevel.INTERNAL,
client_ip=auth_info.get("client_ip", "unknown"),
session_id=auth_info.get("session_id", f"session_{token_info.token_id}"),
login_time=datetime.utcnow(),
last_activity=token_info.last_used
last_activity=token_info.last_used,
token=token # Store raw token for token-bound database configuration
)
except Exception as e: