This is an automated email from the ASF dual-hosted git repository.
FreeOnePlus pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris-mcp-server.git
The following commit(s) were added to refs/heads/master by this push:
new ce116c2 feat: support multi-instance and multi-FE routing (#163)
ce116c2 is described below
commit ce116c2ec47ae8594c1041738128bc7f1b609675
Author: Yijia Su <[email protected]>
AuthorDate: Thu Jul 30 21:59:48 2026 +0800
feat: support multi-instance and multi-FE routing (#163)
Add ordered FE failover for global and token-bound Doris routes,
token-aware HTTP endpoint routing, and explicit fail-closed ADBC behavior for
token-bound clusters.
Closes #73
---
.env.example | 5 +
README.md | 36 +-
README.zh-CN.md | 30 ++
docs/doris-fine-grained-access-control.md | 18 +-
doris_mcp_server/auth/token_handlers.py | 84 ++++-
doris_mcp_server/auth/token_manager.py | 91 ++++-
doris_mcp_server/main.py | 6 +
doris_mcp_server/utils/adbc_query_tools.py | 38 ++-
doris_mcp_server/utils/analysis_tools.py | 30 +-
doris_mcp_server/utils/config.py | 61 ++++
doris_mcp_server/utils/db.py | 494 +++++++++++++++++-----------
doris_mcp_server/utils/doris_http_client.py | 94 +++++-
doris_mcp_server/utils/monitoring_tools.py | 41 ++-
test/auth/test_token_handlers_boundaries.py | 16 +-
test/protocol/test_multiworker_config.py | 11 +
test/security/test_token_digest_storage.py | 58 +++-
test/utils/test_multi_fe_routing.py | 354 ++++++++++++++++++++
17 files changed, 1224 insertions(+), 243 deletions(-)
diff --git a/.env.example b/.env.example
index 84b4686..ab9b64f 100644
--- a/.env.example
+++ b/.env.example
@@ -25,6 +25,9 @@
# Doris FE (Frontend) connection settings
DORIS_HOST=localhost
+# Optional ordered FE MySQL endpoints for one Doris cluster. DORIS_HOST is
+# prepended when both settings are present.
+DORIS_HOSTS=
DORIS_PORT=9030
DORIS_USER=root
DORIS_PASSWORD=
@@ -33,6 +36,8 @@ DORIS_DATABASE=information_schema
# Doris FE HTTP API endpoint (for Profile and other HTTP APIs).
# Leave the host empty to reuse DORIS_HOST for backward compatibility.
DORIS_FE_HTTP_HOST=
+# Optional ordered FE HTTP endpoints for the same Doris cluster.
+DORIS_FE_HTTP_HOSTS=
DORIS_FE_HTTP_PORT=8030
# Doris BE (Backend) HTTP allowlist for monitoring metrics.
diff --git a/README.md b/README.md
index a8a88b8..4629b3a 100644
--- a/README.md
+++ b/README.md
@@ -296,6 +296,7 @@ cp .env.example .env
* **Database Connection**:
* `DORIS_HOST`: Database hostname (default: localhost)
+ * `DORIS_HOSTS`: Ordered FE MySQL failover hosts for one Doris cluster
(comma-separated; `DORIS_HOST` is prepended when both are set)
* `DORIS_PORT`: Database port (default: 9030)
* `DORIS_USER`: Database username (default: root)
* `DORIS_PASSWORD`: Database password
@@ -303,6 +304,7 @@ cp .env.example .env
* `DORIS_MIN_CONNECTIONS`: Minimum connection pool size (default: 5)
* `DORIS_MAX_CONNECTIONS`: Maximum connection pool size (default: 20)
* `DORIS_FE_HTTP_HOST`: Independent FE HTTP host for profile,
table-size, and monitoring tools (default: empty, falling back to `DORIS_HOST`)
+ * `DORIS_FE_HTTP_HOSTS`: Ordered FE HTTP failover hosts for the same
Doris cluster (comma-separated)
* `DORIS_FE_HTTP_PORT`: Independent FE HTTP API port (default: 8030)
* `DORIS_BE_HOSTS`: Explicit BE HTTP allowlist for monitoring
(comma-separated; BE HTTP metrics are disabled when empty)
* `DORIS_BE_WEBSERVER_PORT`: BE webserver port for monitoring tools
(default: 8040)
@@ -1731,10 +1733,12 @@ proxy, tunnel, or split network exposes them at
different addresses:
```bash
# SQL/MySQL protocol endpoint
DORIS_HOST=sql-gateway.internal
+DORIS_HOSTS=sql-gateway.internal,fe-2.internal,fe-3.internal
DORIS_PORT=9030
# FE HTTP endpoint; omit DORIS_FE_HTTP_HOST to reuse DORIS_HOST
DORIS_FE_HTTP_HOST=fe-http-proxy.internal
+DORIS_FE_HTTP_HOSTS=fe-http-proxy.internal,fe-2.internal,fe-3.internal
DORIS_FE_HTTP_PORT=8030
# Explicit BE HTTP allowlist
@@ -1753,6 +1757,17 @@ connection/read/total timeouts plus a response byte
limit. Private and loopback
addresses remain available for normal internal Doris deployments and SSH
tunnels.
+`DORIS_HOSTS` and `DORIS_FE_HTTP_HOSTS` are ordered failover lists, not
+load-balancing or cluster-discovery settings. Every host in one list must
+belong to the same Doris cluster and use the configured shared port and
+credentials. The server probes candidates in order during global/static-token
+pool creation and recovery; FE HTTP requests move to the next configured
+endpoint only on a transport error or `502`/`503`/`504`. Doris-backed OAuth
+tries the candidates during sign-in, but an established per-user pool cannot
+be reconstructed after failure because the server intentionally does not
+retain the user's raw password; the user must sign in again. A stable load
+balancer or SQL gateway is still recommended for large production deployments.
+
### Q: How to use SQL Explain/Profile files with LLM for optimization?
**A:** The tools provide both truncated content and complete files for LLM
analysis:
@@ -2076,11 +2091,20 @@ cat logs/doris_mcp_server_critical.log
"is_active": true,
"database_config": {
"host": "tenant-alpha-db.company.com",
+ "hosts": [
+ "tenant-alpha-fe-1.company.com",
+ "tenant-alpha-fe-2.company.com"
+ ],
"port": 9030,
"user": "alpha_user",
"password": "secure_password",
"database": "alpha_analytics",
- "charset": "UTF8"
+ "charset": "UTF8",
+ "fe_http_hosts": [
+ "tenant-alpha-fe-1.company.com",
+ "tenant-alpha-fe-2.company.com"
+ ],
+ "fe_http_port": 8030
}
}
]
@@ -2109,6 +2133,16 @@ cat logs/doris_mcp_server_critical.log
curl -H "Authorization: Bearer $TOKEN_TENANT_BETA" http://localhost:3000/mcp
```
+ Each token may bind a different Doris cluster. Within one token binding,
+ `hosts` and `fe_http_hosts` are ordered FE candidates for that same cluster.
+ The authenticated token fixes the route; MCP tool arguments cannot select
+ or override another cluster. This multi-instance mode requires the HTTP
+ transport with static-token authentication. A stdio process has one global
+ database route, so run separate stdio processes when clients need separate
+ clusters. `exec_adbc_query` is intentionally fail-closed on token-bound
+ routes because the current Arrow Flight client is process-global; use
+ `exec_query` or a separate MCP process for that cluster.
+
### Q: How is Doris-backed OAuth different from external OAuth/OIDC?
**A:** External OAuth/OIDC delegates identity to an external provider such as
Google, Azure AD, GitHub, GitLab, or Keycloak. Doris-backed OAuth is issued by
this MCP server after the user signs in with Doris credentials. The server
validates the Doris username/password, creates a per-user Doris connection
pool, issues `doa_` access and refresh tokens, and lets Doris RBAC decide what
data and metadata that user can access.
diff --git a/README.zh-CN.md b/README.zh-CN.md
index ab05c1b..24478f2 100644
--- a/README.zh-CN.md
+++ b/README.zh-CN.md
@@ -151,6 +151,7 @@ curl --fail http://127.0.0.1:3000/ready
| 环境变量 | 用途 | 默认值 |
|:---------|:-----|:-------|
| `DORIS_HOST` | Doris FE MySQL 主机 | `localhost` |
+| `DORIS_HOSTS` | 同一集群的有序 FE MySQL 故障切换主机，逗号分隔 | 空 |
| `DORIS_PORT` | Doris FE MySQL 端口 | `9030` |
| `DORIS_USER` | Doris 用户 | `root` |
| `DORIS_PASSWORD` | Doris 密码 | 空 |
@@ -158,6 +159,7 @@ curl --fail http://127.0.0.1:3000/ready
| `DORIS_MIN_CONNECTIONS` | 最小连接数 | `5` |
| `DORIS_MAX_CONNECTIONS` | 最大连接数 | `20` |
| `DORIS_FE_HTTP_HOST` | FE HTTP 工具使用的主机 | 回退到 `DORIS_HOST` |
+| `DORIS_FE_HTTP_HOSTS` | 同一集群的有序 FE HTTP 故障切换主机，逗号分隔 | 空 |
| `DORIS_FE_HTTP_PORT` | FE HTTP 端口 | `8030` |
| `DORIS_BE_HOSTS` | BE HTTP 主机白名单 | 空，表示禁用 BE HTTP 指标 |
| `MAX_RESULT_ROWS` | 查询返回行数上限 | `10000` |
@@ -262,10 +264,38 @@ curl \
请求通过该 Token 鉴权后，数据库操作使用它自己的连接池，不回退到全局服务账号。
这适合多租户或按客户隔离 Doris 身份的部署。
+一个 HTTP 进程可以为不同 Token 绑定不同 Doris 集群；Token 是固定路由边界，
+工具参数不能临时切换集群。每套绑定还可以通过 `hosts` 和 `fe_http_hosts`
+配置同一集群内的多个 FE。stdio 进程只有一套全局数据库路由，需要连接不同
+集群时应分别启动多个进程。当前 Arrow Flight 客户端仍是进程级状态，因此
+`exec_adbc_query` 在 Token 绑定路由上会拒绝执行；应使用 `exec_query` 或为该
+集群单独启动 MCP 进程。
+
`tokens.json` v2 只保存 Bearer Token 的摘要，不保存可恢复的明文。托管创建接口
只返回一次明文，调用方必须立即存入客户端密钥存储。手工配置示例、迁移规则和
文件格式见[英文 Token 绑定说明](README.md#token-bound-database-configuration-new-in-v060)。
+### 多 FE 故障切换
+
+同一 Doris 集群可以配置多个 FE：
+
+```bash
+DORIS_HOST=fe-1.internal
+DORIS_HOSTS=fe-1.internal,fe-2.internal,fe-3.internal
+DORIS_PORT=9030
+
+DORIS_FE_HTTP_HOST=fe-1.internal
+DORIS_FE_HTTP_HOSTS=fe-1.internal,fe-2.internal,fe-3.internal
+DORIS_FE_HTTP_PORT=8030
+```
+
+列表是有序故障切换，不是负载均衡或集群发现。全局连接池和静态 Token 连接池
+会在创建、恢复时依次尝试 FE；Doris OAuth 会在登录时依次尝试 FE，但服务端
+不会保留用户明文密码，已建立的用户连接池失效后不能自动重建，用户需要重新
+登录。Profile、Trace、表容量和 FE 监控只在网络错误或 `502`/`503`/`504` 时
+尝试下一节点。列表中的节点必须属于同一集群并共用端口和凭据。大型生产部署
+仍建议在前面使用稳定的负载均衡或 SQL 网关。
+
### 外部 OAuth/OIDC
外部 OAuth 使用受信任的 RFC 7662 Introspection 端点校验访问 Token，并检查
diff --git a/docs/doris-fine-grained-access-control.md
b/docs/doris-fine-grained-access-control.md
index 2218c3b..e522cc6 100644
--- a/docs/doris-fine-grained-access-control.md
+++ b/docs/doris-fine-grained-access-control.md
@@ -200,12 +200,20 @@ Doris user:
```json
{
"database_config": {
- "host": "doris-fe.internal.example",
+ "host": "doris-fe-1.internal.example",
+ "hosts": [
+ "doris-fe-1.internal.example",
+ "doris-fe-2.internal.example"
+ ],
"port": 9030,
"user": "mcp_orders_east",
"password": "<secret-store-value>",
"database": "sales",
"charset": "UTF8",
+ "fe_http_hosts": [
+ "doris-fe-1.internal.example",
+ "doris-fe-2.internal.example"
+ ],
"fe_http_port": 8030
}
}
@@ -216,6 +224,10 @@ pool for this route. It must not fall back to the global
Doris account. Store
only a token digest in `tokens.json`, keep the file at mode `0600`, and protect
the Doris password as a secret.
+All entries in `hosts` and `fe_http_hosts` must be FE nodes of the same Doris
+cluster. Use separate static-token bindings for separate clusters; the bearer
+token selects the route and a tool call cannot override it.
+
Do not bind a fine-grained token to `root`, `admin`, or a service account that
has global/table-level privileges.
@@ -237,6 +249,10 @@ DORIS_OAUTH_EXPLAIN_TOOLS_ENABLED=true
`exec_query`, reviewed metadata tools, and SQL explain then reach Doris as the
signed-in user. Doris-backed OAuth must fail closed if its per-user pool is
missing; it must never use the global service account as a fallback.
+When multiple FE candidates are configured, sign-in tries them in order.
+If an established per-user pool later fails, the user must sign in again:
+the server intentionally does not retain the raw Doris password needed to
+reconstruct that pool.
The current Doris-backed OAuth implementation is single-process and requires
`WORKERS=1`. See the main README for its complete operational limits.
diff --git a/doris_mcp_server/auth/token_handlers.py
b/doris_mcp_server/auth/token_handlers.py
index c855640..6ea38a5 100644
--- a/doris_mcp_server/auth/token_handlers.py
+++ b/doris_mcp_server/auth/token_handlers.py
@@ -37,6 +37,25 @@ if TYPE_CHECKING:
from ..utils.security import DorisSecurityManager
+def _parse_hosts(value: object) -> list[str]:
+ if value in (None, ""):
+ return []
+ values = value.split(",") if isinstance(value, str) else value
+ if not isinstance(values, list):
+ raise ValueError("host list must be a list or comma-separated string")
+ if len(values) > 16:
+ raise ValueError("host list cannot contain more than 16 entries")
+ hosts = [host.strip() for host in values if isinstance(host, str) and
host.strip()]
+ if len(hosts) != len(values):
+ raise ValueError("host list contains an invalid entry")
+ if any(
+ "://" in host or any(character in host for character in "/\\?#@%")
+ for host in hosts
+ ):
+ raise ValueError("host list contains a malformed entry")
+ return list(dict.fromkeys(hosts))
+
+
class TokenHandlers:
"""Token Authentication HTTP Handlers"""
@@ -85,14 +104,24 @@ class TokenHandlers:
custom_token = query_params.get("custom_token")
# Database configuration from query params
db_config = None
- if query_params.get("db_host"):
+ if query_params.get("db_host") or query_params.get("db_hosts"):
+ db_hosts = _parse_hosts(query_params.get("db_hosts"))
db_config = DatabaseConfig(
- host=query_params.get("db_host", "localhost"),
+ host=query_params.get("db_host") or db_hosts[0],
+ hosts=db_hosts,
port=int(query_params.get("db_port", "9030")),
user=query_params.get("db_user", "root"),
password=query_params.get("db_password", ""),
database=query_params.get("db_database",
"information_schema"),
+ fe_http_host=query_params.get("db_fe_http_host", ""),
+ fe_http_hosts=_parse_hosts(
+ query_params.get("db_fe_http_hosts")
+ ),
fe_http_port=int(query_params.get("db_fe_http_port",
"8030")),
+ be_hosts=_parse_hosts(query_params.get("db_be_hosts")),
+ be_webserver_port=int(
+ query_params.get("db_be_webserver_port", "8040")
+ ),
)
else:
# POST request with JSON body
@@ -110,13 +139,40 @@ class TokenHandlers:
if body.get("database_config"):
db_data = body["database_config"]
try:
+ db_hosts = _parse_hosts(db_data.get("hosts"))
db_config = DatabaseConfig(
- host=db_data.get("host", "localhost"),
+ host=db_data.get("host") or (
+ db_hosts[0] if db_hosts else "localhost"
+ ),
+ hosts=db_hosts,
port=int(db_data.get("port", 9030)),
user=db_data.get("user", "root"),
password=db_data.get("password", ""),
database=db_data.get("database",
"information_schema"),
+ fe_http_host=db_data.get("fe_http_host", ""),
+ fe_http_hosts=_parse_hosts(
+ db_data.get("fe_http_hosts")
+ ),
fe_http_port=int(db_data.get("fe_http_port",
8030)),
+ be_hosts=_parse_hosts(db_data.get("be_hosts")),
+ be_webserver_port=int(
+ db_data.get("be_webserver_port", 8040)
+ ),
+ http_connect_timeout_seconds=float(
+ db_data.get("http_connect_timeout_seconds",
3.0)
+ ),
+ http_read_timeout_seconds=float(
+ db_data.get("http_read_timeout_seconds", 15.0)
+ ),
+ http_total_timeout_seconds=float(
+ db_data.get("http_total_timeout_seconds", 30.0)
+ ),
+ http_max_response_bytes=int(
+ db_data.get(
+ "http_max_response_bytes",
+ 4 * 1024 * 1024,
+ )
+ ),
)
except (ValueError, TypeError):
return JSONResponse(
@@ -502,6 +558,10 @@ class TokenHandlers:
<label for="db_host">Host:</label>
<input type="text" id="db_host"
name="db_host" placeholder="localhost">
</div>
+ <div class="form-group">
+ <label for="db_hosts">FE MySQL
Hosts:</label>
+ <input type="text" id="db_hosts"
name="db_hosts" placeholder="fe-1,fe-2,fe-3">
+ </div>
<div class="form-group">
<label for="db_port">Port:</label>
<input type="number" id="db_port"
name="db_port" placeholder="9030">
@@ -518,6 +578,10 @@ class TokenHandlers:
<label
for="db_database">Database:</label>
<input type="text" id="db_database"
name="db_database" placeholder="information_schema">
</div>
+ <div class="form-group">
+ <label for="db_fe_http_hosts">FE HTTP
Hosts:</label>
+ <input type="text"
id="db_fe_http_hosts" name="db_fe_http_hosts" placeholder="fe-1,fe-2,fe-3">
+ </div>
<div class="form-group">
<label for="db_fe_http_port">FE HTTP
Port:</label>
<input type="number"
id="db_fe_http_port" name="db_fe_http_port" placeholder="8030">
@@ -589,23 +653,33 @@ class TokenHandlers:
if (!data.custom_token) delete data.custom_token;
// Handle database configuration
- if (data.db_host) {{
+ if (data.db_host || data.db_hosts) {{
+ const sqlHosts = data.db_hosts
+ ? data.db_hosts.split(',').map(host =>
host.trim()).filter(Boolean)
+ : [];
+ const httpHosts = data.db_fe_http_hosts
+ ? data.db_fe_http_hosts.split(',').map(host =>
host.trim()).filter(Boolean)
+ : [];
data.database_config = {{
- host: data.db_host,
+ host: data.db_host || sqlHosts[0],
+ hosts: sqlHosts,
port: data.db_port ? parseInt(data.db_port) :
9030,
user: data.db_user || 'root',
password: data.db_password || '',
database: data.db_database ||
'information_schema',
+ fe_http_hosts: httpHosts,
fe_http_port: data.db_fe_http_port ?
parseInt(data.db_fe_http_port) : 8030
}};
}}
// Remove individual database fields from data
delete data.db_host;
+ delete data.db_hosts;
delete data.db_port;
delete data.db_user;
delete data.db_password;
delete data.db_database;
+ delete data.db_fe_http_hosts;
delete data.db_fe_http_port;
try {{
diff --git a/doris_mcp_server/auth/token_manager.py
b/doris_mcp_server/auth/token_manager.py
index ac256ac..45abd09 100644
--- a/doris_mcp_server/auth/token_manager.py
+++ b/doris_mcp_server/auth/token_manager.py
@@ -53,12 +53,21 @@ class DatabaseConfig:
"""Database connection configuration for token binding"""
host: str
+ hosts: list[str] = field(default_factory=list)
port: int = 9030
user: str = ""
password: str = ""
database: str = "information_schema"
charset: str = "UTF8"
+ fe_http_host: str = ""
+ fe_http_hosts: list[str] = field(default_factory=list)
fe_http_port: int = 8030
+ be_hosts: list[str] = field(default_factory=list)
+ be_webserver_port: int = 8040
+ http_connect_timeout_seconds: float = 3.0
+ http_read_timeout_seconds: float = 15.0
+ http_total_timeout_seconds: float = 30.0
+ http_max_response_bytes: int = 4 * 1024 * 1024
@dataclass
@@ -236,6 +245,33 @@ class TokenManager:
raise ValueError(f"{setting} must be an RFC 3339 timestamp") from
exc
return as_utc(parsed)
+ @staticmethod
+ def _parse_hosts(value: Any, *, setting: str) -> list[str]:
+ """Parse an ordered hostname/IP list from JSON or a CSV query value."""
+ if value in (None, ""):
+ return []
+ values = value.split(",") if isinstance(value, str) else value
+ if not isinstance(values, list):
+ raise ValueError(f"{setting} must be a list or comma-separated
string")
+ if len(values) > 16:
+ raise ValueError(f"{setting} cannot contain more than 16 entries")
+ hosts: list[str] = []
+ for raw_host in values:
+ if not isinstance(raw_host, str):
+ raise ValueError(f"{setting} entries must be strings")
+ host = raw_host.strip()
+ if (
+ not host
+ or "://" in host
+ or any(character in host for character in "/\\?#@%")
+ ):
+ raise ValueError(
+ f"{setting} entries must be hostname or IP address values"
+ )
+ if host not in hosts:
+ hosts.append(host)
+ return hosts
+
def _add_token_from_config(
self,
token_config: dict[str, Any],
@@ -274,14 +310,46 @@ class TokenManager:
database_config = None
if 'database_config' in token_config:
db_config = token_config['database_config']
+ hosts = self._parse_hosts(
+ db_config.get('hosts'),
+ setting=f"static token '{token_id}' database hosts",
+ )
+ fe_http_hosts = self._parse_hosts(
+ db_config.get('fe_http_hosts'),
+ setting=f"static token '{token_id}' FE HTTP hosts",
+ )
database_config = DatabaseConfig(
- host=db_config.get('host', 'localhost'),
+ host=db_config.get('host') or (hosts[0] if hosts else
'localhost'),
+ hosts=hosts,
port=db_config.get('port', 9030),
user=db_config.get('user', 'root'),
password=db_config.get('password', ''),
database=db_config.get('database', 'information_schema'),
charset=db_config.get('charset', 'UTF8'),
- fe_http_port=db_config.get('fe_http_port', 8030)
+ fe_http_host=db_config.get('fe_http_host', ''),
+ fe_http_hosts=fe_http_hosts,
+ fe_http_port=db_config.get('fe_http_port', 8030),
+ be_hosts=self._parse_hosts(
+ db_config.get('be_hosts'),
+ setting=f"static token '{token_id}' BE HTTP hosts",
+ ),
+ be_webserver_port=db_config.get('be_webserver_port', 8040),
+ http_connect_timeout_seconds=db_config.get(
+ 'http_connect_timeout_seconds',
+ 3.0,
+ ),
+ http_read_timeout_seconds=db_config.get(
+ 'http_read_timeout_seconds',
+ 15.0,
+ ),
+ http_total_timeout_seconds=db_config.get(
+ 'http_total_timeout_seconds',
+ 30.0,
+ ),
+ http_max_response_bytes=db_config.get(
+ 'http_max_response_bytes',
+ 4 * 1024 * 1024,
+ ),
)
# Create token info
@@ -938,12 +1006,29 @@ class TokenManager:
if token_info.database_config:
token_config["database_config"] = {
"host": token_info.database_config.host,
+ "hosts": token_info.database_config.hosts,
"port": token_info.database_config.port,
"user": token_info.database_config.user,
"password": token_info.database_config.password,
"database": token_info.database_config.database,
"charset": token_info.database_config.charset,
- "fe_http_port": token_info.database_config.fe_http_port
+ "fe_http_host": token_info.database_config.fe_http_host,
+ "fe_http_hosts": token_info.database_config.fe_http_hosts,
+ "fe_http_port": token_info.database_config.fe_http_port,
+ "be_hosts": token_info.database_config.be_hosts,
+ "be_webserver_port":
token_info.database_config.be_webserver_port,
+ "http_connect_timeout_seconds": (
+ token_info.database_config.http_connect_timeout_seconds
+ ),
+ "http_read_timeout_seconds": (
+ token_info.database_config.http_read_timeout_seconds
+ ),
+ "http_total_timeout_seconds": (
+ token_info.database_config.http_total_timeout_seconds
+ ),
+ "http_max_response_bytes": (
+ token_info.database_config.http_max_response_bytes
+ ),
}
return token_config
diff --git a/doris_mcp_server/main.py b/doris_mcp_server/main.py
index 40017fa..ad703ae 100644
--- a/doris_mcp_server/main.py
+++ b/doris_mcp_server/main.py
@@ -72,10 +72,16 @@ def _multiworker_environment(
"""Serialize resolved parent settings inherited by Uvicorn workers."""
return {
"DORIS_HOST": config.database.host,
+ "DORIS_HOSTS": ",".join(config.database.hosts),
"DORIS_PORT": str(config.database.port),
"DORIS_USER": config.database.user,
"DORIS_PASSWORD": config.database.password,
"DORIS_DATABASE": config.database.database,
+ "DORIS_FE_HTTP_HOST": config.database.fe_http_host,
+ "DORIS_FE_HTTP_HOSTS": ",".join(config.database.fe_http_hosts),
+ "DORIS_FE_HTTP_PORT": str(config.database.fe_http_port),
+ "DORIS_BE_HOSTS": ",".join(config.database.be_hosts),
+ "DORIS_BE_WEBSERVER_PORT": str(config.database.be_webserver_port),
"SERVER_HOST": host,
"SERVER_PORT": str(port),
"MCP_ALLOWED_HOSTS": ",".join(config.mcp_allowed_hosts),
diff --git a/doris_mcp_server/utils/adbc_query_tools.py
b/doris_mcp_server/utils/adbc_query_tools.py
index 7de8b44..58c37cb 100644
--- a/doris_mcp_server/utils/adbc_query_tools.py
+++ b/doris_mcp_server/utils/adbc_query_tools.py
@@ -34,6 +34,7 @@ from ..result_limits import (
resolve_result_limits,
)
from ..utils.db import DorisConnectionManager
+from ..utils.doris_http_client import database_config_for_request
from ..utils.logger import get_logger
from ..utils.security import AuthContext
from ..utils.sql_security_utils import get_auth_context
@@ -78,6 +79,26 @@ class DorisADBCQueryTools:
# ADBC connections and cursors are not safe to replace concurrently.
self._query_lock = asyncio.Lock()
+ def _token_bound_route_error(self) -> dict[str, Any] | None:
+ global_database_config = getattr(
+ self.connection_manager.config,
+ "database",
+ None,
+ )
+ if global_database_config is None:
+ return None
+ selected = database_config_for_request(self.connection_manager)
+ if selected is global_database_config:
+ return None
+ return {
+ "success": False,
+ "error": (
+ "ADBC is unavailable for token-bound database routes; use "
+ "exec_query or a separately configured MCP process"
+ ),
+ "error_type": "token_bound_adbc_unsupported",
+ }
+
async def exec_adbc_query(
self,
sql: str,
@@ -100,6 +121,9 @@ class DorisADBCQueryTools:
Query results in specified format with metadata
"""
try:
+ route_error = self._token_bound_route_error()
+ if route_error is not None:
+ return route_error
# Use configuration defaults if parameters not specified
adbc_config = self.connection_manager.config.adbc
try:
@@ -218,8 +242,7 @@ class DorisADBCQueryTools:
}
# Get host address
- db_config = self.connection_manager.config.database
- fe_host = db_config.host
+ fe_host = self.connection_manager.host
# Check FE Arrow Flight SQL port availability
fe_available = self._check_port_connectivity(fe_host, fe_port)
@@ -393,7 +416,7 @@ class DorisADBCQueryTools:
fe_port = int(fe_port_raw)
# Build connection URI
- uri = f"grpc://{db_config.host}:{fe_port}"
+ uri = f"grpc://{self.connection_manager.host}:{fe_port}"
# Create database connection parameters
db_kwargs = {
@@ -623,6 +646,13 @@ class DorisADBCQueryTools:
async def get_adbc_connection_info(self) -> dict[str, Any]:
"""Get ADBC connection information and status"""
try:
+ route_error = self._token_bound_route_error()
+ if route_error is not None:
+ return {
+ **route_error,
+ "status": "not_ready",
+ "timestamp": datetime.now().isoformat(),
+ }
# Check port status
port_status = await self._check_arrow_flight_ports()
@@ -638,7 +668,7 @@ class DorisADBCQueryTools:
"adbc_available": module_status["success"],
"ports_available": port_status["success"],
"configuration": {
- "fe_host": db_config.host,
+ "fe_host": self.connection_manager.host,
"fe_arrow_flight_port": fe_port,
"be_arrow_flight_port": be_port,
"user": db_config.user
diff --git a/doris_mcp_server/utils/analysis_tools.py
b/doris_mcp_server/utils/analysis_tools.py
index 08fc8b4..5fd0011 100644
--- a/doris_mcp_server/utils/analysis_tools.py
+++ b/doris_mcp_server/utils/analysis_tools.py
@@ -27,7 +27,11 @@ from typing import Any
from urllib.parse import quote
from .db import DorisConnectionManager
-from .doris_http_client import DorisHTTPClient
+from .doris_http_client import (
+ DorisHTTPClient,
+ configured_fe_http_hosts,
+ database_config_for_request,
+)
from .logger import get_logger
from .sql_security_utils import (
SQLSecurityError,
@@ -924,12 +928,12 @@ class SQLAnalyzer:
Query ID string or None if not found
"""
try:
- db_config = self.connection_manager.config.database
- fe_http_host = getattr(db_config, "fe_http_host", "") or
db_config.host
+ db_config = database_config_for_request(self.connection_manager)
+ fe_http_hosts = configured_fe_http_hosts(db_config)
http_client = DorisHTTPClient.from_database_config(db_config)
- response = await http_client.get(
+ response = await http_client.get_first_available(
role="fe",
- host=fe_http_host,
+ hosts=fe_http_hosts,
port=db_config.fe_http_port,
path=(f"/rest/v2/manager/query/trace_id/{quote(trace_id,
safe='')}"),
headers={"Accept": "application/json"},
@@ -1001,8 +1005,8 @@ class SQLAnalyzer:
Profile data dict or None if failed
"""
try:
- db_config = self.connection_manager.config.database
- fe_http_host = getattr(db_config, "fe_http_host", "") or
db_config.host
+ db_config = database_config_for_request(self.connection_manager)
+ fe_http_hosts = configured_fe_http_hosts(db_config)
endpoints = [
(
f"/rest/v2/manager/query/profile/text/{quote(query_id,
safe='')}",
@@ -1013,9 +1017,9 @@ class SQLAnalyzer:
http_client = DorisHTTPClient.from_database_config(db_config)
for i, (path, params) in enumerate(endpoints):
logger.info("Requesting profile from configured FE endpoint
%s", i + 1)
- response = await http_client.get(
+ response = await http_client.get_first_available(
role="fe",
- host=fe_http_host,
+ hosts=fe_http_hosts,
port=db_config.fe_http_port,
path=path,
params=params,
@@ -1091,7 +1095,7 @@ class SQLAnalyzer:
Dict containing table data size information
"""
try:
- db_config = self.connection_manager.config.database
+ db_config = database_config_for_request(self.connection_manager)
params = {}
if db_name:
params["db"] = db_name
@@ -1103,11 +1107,11 @@ class SQLAnalyzer:
"Requesting table data size from configured FE endpoint with
params: %s",
params,
)
- fe_http_host = getattr(db_config, "fe_http_host", "") or
db_config.host
+ fe_http_hosts = configured_fe_http_hosts(db_config)
http_client = DorisHTTPClient.from_database_config(db_config)
- response = await http_client.get(
+ response = await http_client.get_first_available(
role="fe",
- host=fe_http_host,
+ hosts=fe_http_hosts,
port=db_config.fe_http_port,
path="/api/show_table_data",
params=params,
diff --git a/doris_mcp_server/utils/config.py b/doris_mcp_server/utils/config.py
index 42f237d..1714d03 100644
--- a/doris_mcp_server/utils/config.py
+++ b/doris_mcp_server/utils/config.py
@@ -359,6 +359,10 @@ class DatabaseConfig:
"""Database connection configuration"""
host: str = "localhost"
+ # Ordered FE MySQL endpoints for one Doris cluster. When populated, the
+ # connection manager tries these hosts in order and keeps ``host`` as the
+ # backward-compatible primary endpoint.
+ hosts: list[str] = field(default_factory=list)
port: int = 9030
user: str = "root"
password: str = ""
@@ -368,6 +372,7 @@ class DatabaseConfig:
# FE HTTP API endpoint for profile and other HTTP APIs. An empty host keeps
# backward compatibility by falling back to the SQL host.
fe_http_host: str = ""
+ fe_http_hosts: list[str] = field(default_factory=list)
fe_http_port: int = 8030
# BE HTTP nodes must be configured explicitly. SQL metadata is not trusted
@@ -936,6 +941,19 @@ class DorisConfig:
doris_host = os.getenv("DORIS_HOST", "").strip()
config.database.host = doris_host if doris_host else
config.database.host
+ doris_hosts_env = os.getenv("DORIS_HOSTS", "")
+ if doris_hosts_env:
+ doris_hosts = [
+ host.strip()
+ for host in doris_hosts_env.split(",")
+ if host.strip()
+ ]
+ if doris_host:
+ doris_hosts = list(dict.fromkeys([doris_host, *doris_hosts]))
+ if doris_hosts:
+ config.database.hosts = doris_hosts
+ config.database.host = doris_hosts[0]
+
doris_port = os.getenv("DORIS_PORT", "").strip()
if doris_port and doris_port.isdigit():
config.database.port = int(doris_port)
@@ -952,6 +970,21 @@ class DorisConfig:
doris_fe_http_host = os.getenv("DORIS_FE_HTTP_HOST", "").strip()
config.database.fe_http_host = doris_fe_http_host
+ doris_fe_http_hosts_env = os.getenv("DORIS_FE_HTTP_HOSTS", "")
+ if doris_fe_http_hosts_env:
+ doris_fe_http_hosts = [
+ host.strip()
+ for host in doris_fe_http_hosts_env.split(",")
+ if host.strip()
+ ]
+ if doris_fe_http_host:
+ doris_fe_http_hosts = list(
+ dict.fromkeys([doris_fe_http_host, *doris_fe_http_hosts])
+ )
+ if doris_fe_http_hosts:
+ config.database.fe_http_hosts = doris_fe_http_hosts
+ config.database.fe_http_host = doris_fe_http_hosts[0]
+
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)
@@ -1519,6 +1552,10 @@ class DorisConfig:
for key, value in db_config.items():
if hasattr(config.database, key):
setattr(config.database, key, value)
+ if config.database.hosts and "host" not in db_config:
+ config.database.host = config.database.hosts[0]
+ if config.database.fe_http_hosts and "fe_http_host" not in
db_config:
+ config.database.fe_http_host = config.database.fe_http_hosts[0]
# Update security configuration
if "security" in config_data:
@@ -1585,12 +1622,14 @@ class DorisConfig:
"temp_files_dir": self.temp_files_dir,
"database": {
"host": self.database.host,
+ "hosts": self.database.hosts,
"port": self.database.port,
"user": self.database.user,
"password": "***", # Hide password
"database": self.database.database,
"charset": self.database.charset,
"fe_http_host": self.database.fe_http_host,
+ "fe_http_hosts": self.database.fe_http_hosts,
"fe_http_port": self.database.fe_http_port,
"be_hosts": self.database.be_hosts,
"be_webserver_port": self.database.be_webserver_port,
@@ -1745,6 +1784,28 @@ class DorisConfig:
if not self.database.host:
errors.append("Database host address cannot be empty")
+ database_host_lists: tuple[tuple[str, Any], ...] = (
+ ("Doris FE SQL hosts", self.database.hosts),
+ ("Doris FE HTTP hosts", self.database.fe_http_hosts),
+ )
+ for field_name, hosts in database_host_lists:
+ if not isinstance(hosts, list):
+ errors.append(f"{field_name} must be a list")
+ continue
+ if len(hosts) > 16:
+ errors.append(f"{field_name} cannot contain more than 16
entries")
+ if any(
+ not isinstance(host, str)
+ or host != host.strip()
+ or not host
+ or "://" in host
+ or any(character in host for character in "/\\?#@%")
+ for host in hosts
+ ):
+ errors.append(
+ f"{field_name} must contain only hostname or IP address
values"
+ )
+
if not (1 <= self.database.port <= 65535):
errors.append("Database port must be in the range 1-65535")
diff --git a/doris_mcp_server/utils/db.py b/doris_mcp_server/utils/db.py
index b596525..0f2c217 100644
--- a/doris_mcp_server/utils/db.py
+++ b/doris_mcp_server/utils/db.py
@@ -36,7 +36,7 @@ from collections.abc import AsyncIterator, Mapping
from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import datetime
-from typing import TYPE_CHECKING, Any, TypedDict
+from typing import TYPE_CHECKING, Any, NotRequired, TypedDict
import aiomysql
from aiomysql import Connection, Pool
@@ -96,6 +96,7 @@ class DatabasePoolConfig(TypedDict):
"""Connection parameters shared by global and token-owned pools."""
host: str
+ hosts: NotRequired[list[str]]
port: int
user: str
password: str
@@ -533,8 +534,13 @@ class DorisConnectionManager:
)
# Store original database config for fallback
+ configured_hosts = self._ordered_hosts(
+ config.database.host,
+ getattr(config.database, "hosts", []),
+ )
self.original_db_config: DatabasePoolConfig = {
- "host": config.database.host,
+ "host": configured_hosts[0],
+ "hosts": configured_hosts,
"port": config.database.port,
"user": config.database.user,
"password": config.database.password,
@@ -596,6 +602,25 @@ class DorisConnectionManager:
db_config["charset"].upper(), db_config["charset"].lower()
)
+ @staticmethod
+ def _ordered_hosts(primary: str, configured: object) -> list[str]:
+ hosts = (
+ [host.strip() for host in configured if isinstance(host, str) and
host.strip()]
+ if isinstance(configured, list)
+ else []
+ )
+ return list(dict.fromkeys(([primary] if primary else []) + hosts))
+
+ def _host_candidates(self, db_config: Mapping[str, Any]) -> list[str]:
+ return self._ordered_hosts(
+ str(db_config.get("host", "")),
+ db_config.get("hosts", []),
+ )
+
+ def _set_active_global_host(self, host: str) -> None:
+ self.host = host
+ self.active_db_config["host"] = host
+
def _is_config_empty(self, config_value: object) -> bool:
"""Check if a config value is empty (None, empty string, or 'null')"""
return (
@@ -629,7 +654,7 @@ class DorisConnectionManager:
token_db_config =
self.token_manager.get_database_config_by_token(token)
if token_db_config:
- return {
+ db_config: DatabasePoolConfig = {
"host": token_db_config.host,
"port": token_db_config.port,
"user": token_db_config.user,
@@ -637,6 +662,13 @@ class DorisConnectionManager:
"database": token_db_config.database,
"charset": token_db_config.charset,
}
+ token_hosts = getattr(token_db_config, "hosts", [])
+ if token_hosts:
+ db_config["hosts"] = self._ordered_hosts(
+ token_db_config.host,
+ token_hosts,
+ )
+ return db_config
return None
def _config_changed(
@@ -649,7 +681,7 @@ class DorisConnectionManager:
return old_config != new_config
# Compare key fields
- for key in ["host", "port", "user", "password", "database"]:
+ for key in ["host", "hosts", "port", "user", "password", "database"]:
if old_config.get(key) != new_config.get(key):
return True
return False
@@ -705,7 +737,8 @@ class DorisConnectionManager:
def _build_doris_user_db_config(self, user: str, password: str) -> dict:
return {
- "host": self.original_db_config["host"],
+ "host": self.host,
+ "hosts": self._host_candidates(self.original_db_config),
"port": self.original_db_config["port"],
"user": user,
"password": password,
@@ -803,35 +836,41 @@ class DorisConnectionManager:
retries = self._validate_doris_auth_retries(max_retries)
last_error: Exception | None = None
+ candidate_config = self.original_db_config.copy()
+ candidate_config["host"] = self.host
+ hosts = self._host_candidates(candidate_config)
for attempt in range(1, retries + 1):
- conn = None
- try:
- conn = await aiomysql.connect(
- host=self.original_db_config["host"],
- port=self.original_db_config["port"],
- user=normalized_user,
- password=password,
- db="information_schema",
- charset=self.charset,
- connect_timeout=self.connect_timeout,
- autocommit=True,
- )
- return
- except DorisUserAuthenticationError:
- raise
- except Exception as e:
- last_error = e
- self.logger.warning(
- "Doris user authentication failed for %s@%s:%s on attempt
%s/%s: %s",
- normalized_user,
- self.original_db_config["host"],
- self.original_db_config["port"],
- attempt,
- retries,
- type(e).__name__,
- )
- finally:
- await self._close_auth_connection(conn)
+ for host in hosts:
+ conn = None
+ try:
+ conn = await aiomysql.connect(
+ host=host,
+ port=self.original_db_config["port"],
+ user=normalized_user,
+ password=password,
+ db="information_schema",
+ charset=self.charset,
+ connect_timeout=self.connect_timeout,
+ autocommit=True,
+ )
+ self._set_active_global_host(host)
+ return
+ except DorisUserAuthenticationError:
+ raise
+ except Exception as exc:
+ last_error = exc
+ self.logger.warning(
+ "Doris user authentication failed for %s@%s:%s "
+ "on attempt %s/%s (%s)",
+ normalized_user,
+ host,
+ self.original_db_config["port"],
+ attempt,
+ retries,
+ type(exc).__name__,
+ )
+ finally:
+ await self._close_auth_connection(conn)
raise DorisUserAuthenticationError(
f"Doris user authentication failed for {normalized_user}:
{type(last_error).__name__}"
@@ -1125,6 +1164,12 @@ class DorisConnectionManager:
"database": token_db_config.database,
"charset": token_db_config.charset,
}
+ token_hosts = getattr(token_db_config, "hosts", [])
+ if token_hosts:
+ db_config["hosts"] = self._ordered_hosts(
+ token_db_config.host,
+ token_hosts,
+ )
config_source = "token-bound"
# Fallback to global config if token has no specific config
@@ -1169,57 +1214,105 @@ class DorisConnectionManager:
self,
db_config: Mapping[str, Any],
) -> Pool:
- """Create a connection pool with specified configuration
+ """Create a pool, failing over when multiple FE hosts are
configured."""
+ pool, _ = await self._create_pool_with_candidates(db_config)
+ return pool
- Args:
- db_config: Database configuration dictionary
+ async def _probe_pool(self, pool: Pool) -> bool:
+ raw_connection = None
+ try:
+ raw_connection = await asyncio.wait_for(
+ pool.acquire(),
+ timeout=self.connect_timeout,
+ )
+ async with raw_connection.cursor() as cursor:
+ await cursor.execute("SELECT 1")
+ result = await cursor.fetchone()
+ return bool(result and result[0] == 1)
+ except Exception as exc:
+ self.logger.debug(
+ "Doris FE candidate health probe failed: %s",
+ type(exc).__name__,
+ )
+ return False
+ finally:
+ if raw_connection is not None:
+ try:
+ pool.release(raw_connection)
+ except Exception:
+ await self._force_close_raw_connection(
+ raw_connection,
+ "candidate health probe release failure",
+ )
- Returns:
- Created connection pool
- """
+ async def _create_pool_with_candidates(
+ self,
+ db_config: Mapping[str, Any],
+ *,
+ minsize: int = 0,
+ maxsize: int | None = None,
+ timeout: float | None = None,
+ require_health: bool | None = None,
+ ) -> tuple[Pool, str]:
+ """Create a pool against the first reachable FE in an ordered list."""
# Convert charset to aiomysql compatible format
charset_map = {"UTF8": "utf8", "UTF8MB4": "utf8mb4"}
charset = charset_map.get(
db_config["charset"].upper(), db_config["charset"].lower()
)
+ candidates = self._host_candidates(db_config)
+ should_probe = len(candidates) > 1 if require_health is None else
require_health
+ pool_timeout = timeout or (self.connect_timeout + 5)
+ last_error: Exception | None = None
- self.logger.debug(
- f"Creating pool for
{db_config['user']}@{db_config['host']}:{db_config['port']}/{db_config['database']}"
- )
-
- try:
- pool = await asyncio.wait_for(
- aiomysql.create_pool(
- host=db_config["host"],
- port=db_config["port"],
- user=db_config["user"],
- password=db_config["password"],
- db=db_config["database"],
- charset=charset,
- minsize=0, # Don't pre-create connections
- maxsize=db_config.get("maxsize", self.maxsize),
- connect_timeout=self.connect_timeout,
- autocommit=True,
- pool_recycle=self.pool_recycle,
- ),
- timeout=self.connect_timeout + 5, # Give extra time for pool
creation
- )
- self.logger.info(
- f"Successfully created pool for
{db_config['user']}@{db_config['host']}:{db_config['port']}"
- )
- return pool
- except TimeoutError:
- self.logger.error(
- f"Timeout creating pool for
{db_config['user']}@{db_config['host']}:{db_config['port']}"
- )
- raise RuntimeError(
- f"Timeout creating connection pool for
{db_config['user']}@{db_config['host']}:{db_config['port']}"
- )
- except Exception as e:
- self.logger.error(
- f"Failed to create pool for
{db_config['user']}@{db_config['host']}:{db_config['port']}:
{type(e).__name__}: {e}"
+ for host in candidates:
+ pool: Pool | None = None
+ self.logger.debug(
+ "Creating pool for %s@%s:%s/%s",
+ db_config["user"],
+ host,
+ db_config["port"],
+ db_config["database"],
)
- raise
+ try:
+ pool = await asyncio.wait_for(
+ aiomysql.create_pool(
+ host=host,
+ port=db_config["port"],
+ user=db_config["user"],
+ password=db_config["password"],
+ db=db_config["database"],
+ charset=charset,
+ minsize=minsize,
+ maxsize=maxsize or db_config.get("maxsize",
self.maxsize),
+ connect_timeout=self.connect_timeout,
+ autocommit=True,
+ pool_recycle=self.pool_recycle,
+ ),
+ timeout=pool_timeout,
+ )
+ if should_probe and not await self._probe_pool(pool):
+ raise RuntimeError("Doris FE candidate failed its health
probe")
+ self.logger.info(
+ "Successfully created pool for %s@%s:%s",
+ db_config["user"],
+ host,
+ db_config["port"],
+ )
+ return pool, host
+ except Exception as exc:
+ last_error = exc
+ await self._close_pool_safely(pool, f"failed FE candidate
{host}")
+ self.logger.warning(
+ "Doris FE candidate %s:%s failed during pool creation
(%s)",
+ host,
+ db_config["port"],
+ type(exc).__name__,
+ )
+
+ raise RuntimeError(
+ "Unable to create a Doris connection pool from the configured FE
hosts"
+ ) from last_error
async def get_connection_for_token(
self,
@@ -1235,41 +1328,73 @@ class DorisConnectionManager:
Returns:
DorisConnection wrapper
"""
- pool, db_config = await self.get_pool_for_token(token)
+ token_hash = self._get_token_hash(token)
+ last_error: Exception | None = None
+ for attempt in range(2):
+ pool, db_config = await self.get_pool_for_token(token)
+ try:
+ connection = await asyncio.wait_for(
+ pool.acquire(), timeout=self.connect_timeout
+ )
+ if getattr(connection, "closed", False):
+ pool.release(connection)
+ raise RuntimeError("Token pool returned a closed
connection")
- try:
- connection = await asyncio.wait_for(
- pool.acquire(), timeout=self.connect_timeout
- )
+ self.logger.debug(
+ f"Session {session_id}: Acquired connection from token
pool "
+ f"(user: {db_config['user']}@{db_config['host']})"
+ )
- self.logger.debug(
- f"Session {session_id}: Acquired connection from token pool "
- f"(user: {db_config['user']}@{db_config['host']})"
- )
+ return DorisConnection(
+ connection,
+ session_id,
+ self.security_manager,
+ pool_kind="static_token",
+ route_key=f"static_token:{token_hash}",
+ owner_id=self._token_pool_owner_ids.get(
+ token_hash, f"static_token:{token_hash}:0"
+ ),
+ generation=self._token_pool_generations.get(token_hash, 0),
+ owner_pool=pool,
+ )
+ except Exception as exc:
+ last_error = exc
+ has_failover = len(self._host_candidates(db_config)) > 1
+ if attempt == 0 and has_failover:
+ self.logger.warning(
+ "Session %s: token pool acquisition failed; "
+ "recreating it from configured FE candidates (%s)",
+ session_id,
+ type(exc).__name__,
+ )
+ await self._evict_token_pool_for_recovery(token_hash, pool)
+ continue
+ self.logger.error(
+ "Session %s: Failed to acquire connection from token pool
(%s)",
+ session_id,
+ type(exc).__name__,
+ )
+ raise
- token_hash = self._get_token_hash(token)
- return DorisConnection(
- connection,
- session_id,
- self.security_manager,
- pool_kind="static_token",
- route_key=f"static_token:{token_hash}",
- owner_id=self._token_pool_owner_ids.get(
- token_hash, f"static_token:{token_hash}:0"
- ),
- generation=self._token_pool_generations.get(token_hash, 0),
- owner_pool=pool,
- )
+ raise RuntimeError("Token pool recovery failed") from last_error
- except Exception as e:
- self.logger.error(
- "Session %s: Failed to acquire connection from token pool "
- "(%s): %s",
- session_id,
- type(e).__name__,
- str(e) or "<no message>",
+ async def _evict_token_pool_for_recovery(
+ self,
+ token_hash: str,
+ expected_pool: Pool,
+ ) -> None:
+ """Remove a failed token pool without evicting a concurrent
replacement."""
+ pool_to_close: Pool | None = None
+ async with self._token_pools_lock:
+ if self.token_pools.get(token_hash) is expected_pool:
+ pool_to_close = self.token_pools.pop(token_hash)
+ self.token_configs.pop(token_hash, None)
+ self._token_pool_owner_ids.pop(token_hash, None)
+ if pool_to_close is not None:
+ await self._close_pool_safely(
+ pool_to_close,
+ f"failed static token pool {token_hash[:8]}",
)
- raise
async def release_connection_for_token(
self,
@@ -1469,20 +1594,7 @@ class DorisConnectionManager:
async def _create_pool_with_current_config(self) -> None:
"""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,
- )
- self._mark_global_pool_created()
+ await self._create_global_pool()
# Store the current config for comparison later
self._last_pool_config = {
@@ -1517,6 +1629,21 @@ class DorisConnectionManager:
self.logger.error(f"Failed to create connection pool: {e}")
raise
+ async def _create_global_pool(self, *, timeout: float | None = None) ->
None:
+ """Create the global pool against the first healthy configured FE."""
+ candidate_config = self.original_db_config.copy()
+ candidate_config["host"] = self.host
+ pool, selected_host = await self._create_pool_with_candidates(
+ candidate_config,
+ minsize=self.minsize,
+ maxsize=self.maxsize,
+ timeout=timeout,
+ require_health=True,
+ )
+ self.pool = pool
+ self._set_active_global_host(selected_host)
+ self._mark_global_pool_created()
+
async def _recreate_pool(self) -> None:
"""Recreate connection pool with current database configuration"""
try:
@@ -1628,21 +1755,8 @@ class DorisConnectionManager:
)
return
- # Create connection pool
- 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,
- )
- self._mark_global_pool_created()
+ # Create connection pool against the first healthy configured FE.
+ await self._create_global_pool()
# Test initial connection
if not await self._test_pool_health():
@@ -1875,30 +1989,43 @@ class DorisConnectionManager:
"""
import aiomysql
- conn = None
- try:
- conn = await aiomysql.connect(
- host=self.host,
- port=self.port,
- user=self.user,
- password=self.password,
- db=self.database,
- charset=self.charset,
- connect_timeout=self.connect_timeout,
- autocommit=True,
- )
+ last_error: Exception | None = None
+ for host in self._host_candidates(self.original_db_config):
+ conn = None
+ try:
+ conn = await aiomysql.connect(
+ host=host,
+ port=self.port,
+ user=self.user,
+ password=self.password,
+ db=self.database,
+ charset=self.charset,
+ connect_timeout=self.connect_timeout,
+ autocommit=True,
+ )
- async with conn.cursor() as cursor:
- await cursor.execute("SELECT 1")
- result = await cursor.fetchone()
- if not result or result[0] != 1:
- raise RuntimeError("Database connectivity test query
failed")
+ async with conn.cursor() as cursor:
+ await cursor.execute("SELECT 1")
+ result = await cursor.fetchone()
+ if not result or result[0] != 1:
+ raise RuntimeError("Database connectivity test query
failed")
+ self._set_active_global_host(host)
+ return
+ except Exception as exc:
+ last_error = exc
+ self.logger.warning(
+ "Doris FE candidate %s:%s failed connectivity check (%s)",
+ host,
+ self.port,
+ type(exc).__name__,
+ )
+ finally:
+ if conn:
+ conn.close()
- except Exception as e:
- raise RuntimeError(f"Database connectivity test failed: {e}")
- finally:
- if conn:
- conn.close()
+ raise RuntimeError(
+ "Database connectivity test failed for all configured Doris FE
hosts"
+ ) from last_error
async def _create_connection_pool(self) -> None:
"""
@@ -1907,29 +2034,7 @@ class DorisConnectionManager:
Raises:
Exception: If pool creation fails
"""
- 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,
- )
- self._mark_global_pool_created()
-
- # Test pool health
- if not await self._test_pool_health():
- # Clean up the pool if health test fails
- if self.pool:
- self.pool.close()
- await self.pool.wait_closed()
- self.pool = None
- raise RuntimeError("Connection pool health check failed")
+ await self._create_global_pool()
async def _test_pool_health(self) -> bool:
"""Test connection pool health"""
@@ -2144,23 +2249,7 @@ class DorisConnectionManager:
# Recreate pool with timeout
self.logger.debug("Creating new connection pool...")
- self.pool = await asyncio.wait_for(
- 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,
- ),
- timeout=10.0,
- )
- self._mark_global_pool_created()
+ await self._create_global_pool(timeout=10.0)
# Test recovered pool with timeout
if await asyncio.wait_for(
@@ -2242,6 +2331,19 @@ class DorisConnectionManager:
self.logger.debug(f"Could not get auth_context: {e}")
return None
+ def get_database_config_for_auth_context(
+ self,
+ auth_context: AuthContext | None = None,
+ ) -> Any:
+ """Return the endpoint config for the current authenticated route."""
+ effective_context = self._get_effective_auth_context(auth_context)
+ token = getattr(effective_context, "token", "") if effective_context
else ""
+ if token and self.token_manager:
+ token_config =
self.token_manager.get_database_config_by_token(token)
+ if token_config is not None:
+ return token_config
+ return self.config.database
+
async def _get_connection_for_auth_context(
self,
session_id: str,
diff --git a/doris_mcp_server/utils/doris_http_client.py
b/doris_mcp_server/utils/doris_http_client.py
index 4e47eb9..798eadb 100644
--- a/doris_mcp_server/utils/doris_http_client.py
+++ b/doris_mcp_server/utils/doris_http_client.py
@@ -34,6 +34,7 @@ DEFAULT_TOTAL_TIMEOUT_SECONDS = 30.0
DEFAULT_MAX_RESPONSE_BYTES = 4 * 1024 * 1024
MAX_TIMEOUT_SECONDS = 60.0
MAX_RESPONSE_BYTES = 16 * 1024 * 1024
+MAX_CONFIGURED_HOSTS = 16
_METADATA_HOSTS = frozenset(
{
@@ -182,6 +183,53 @@ def _bounded_response_limit(value: Any) -> int:
return min(parsed, MAX_RESPONSE_BYTES)
+def configured_fe_http_hosts(database_config: Any) -> tuple[str, ...]:
+ """Return ordered FE HTTP candidates without widening the allowlist."""
+ configured_hosts = getattr(database_config, "fe_http_hosts", []) or []
+ if not isinstance(configured_hosts, list) or any(
+ not isinstance(host, str) for host in configured_hosts
+ ):
+ raise DorisHTTPPolicyError("Configured Doris FE HTTP hosts are
invalid")
+ if len(configured_hosts) > MAX_CONFIGURED_HOSTS:
+ raise DorisHTTPPolicyError("Too many Doris FE HTTP hosts are
configured")
+
+ configured_host = getattr(database_config, "fe_http_host", "")
+ if configured_hosts:
+ candidates = (
+ [configured_host, *configured_hosts]
+ if configured_host
+ else configured_hosts
+ )
+ else:
+ if configured_host:
+ candidates = [configured_host]
+ else:
+ sql_hosts = getattr(database_config, "hosts", []) or []
+ if not isinstance(sql_hosts, list) or any(
+ not isinstance(host, str) for host in sql_hosts
+ ):
+ raise DorisHTTPPolicyError("Configured Doris FE SQL hosts are
invalid")
+ if len(sql_hosts) > MAX_CONFIGURED_HOSTS:
+ raise DorisHTTPPolicyError(
+ "Too many Doris FE SQL hosts are configured"
+ )
+ candidates = sql_hosts or [database_config.host]
+
+ return tuple(dict.fromkeys(_normalize_host(host) for host in candidates))
+
+
+def database_config_for_request(connection_manager: Any) -> Any:
+ """Resolve a token-bound endpoint config with legacy-manager fallback."""
+ resolver = getattr(
+ connection_manager,
+ "get_database_config_for_auth_context",
+ None,
+ )
+ if callable(resolver):
+ return resolver()
+ return connection_manager.config.database
+
+
class DorisHTTPClient:
"""Fetch only configured Doris HTTP endpoints with SSRF controls."""
@@ -221,9 +269,7 @@ class DorisHTTPClient:
@classmethod
def from_database_config(cls, database_config: Any) -> DorisHTTPClient:
- fe_host = _normalize_host(
- getattr(database_config, "fe_http_host", "") or
database_config.host
- )
+ fe_hosts = configured_fe_http_hosts(database_config)
fe_port = _validate_port(database_config.fe_http_port)
be_port = _validate_port(getattr(database_config, "be_webserver_port",
8040))
be_hosts = getattr(database_config, "be_hosts", []) or []
@@ -235,7 +281,7 @@ class DorisHTTPClient:
user=str(database_config.user),
password=str(database_config.password),
allowed_endpoints={
- "fe": {(fe_host, fe_port)},
+ "fe": {(host, fe_port) for host in fe_hosts},
"be": {(_normalize_host(host), be_port) for host in be_hosts},
},
connect_timeout_seconds=getattr(
@@ -260,6 +306,46 @@ class DorisHTTPClient:
),
)
+ async def get_first_available(
+ self,
+ *,
+ role: Literal["fe", "be"],
+ hosts: list[str] | tuple[str, ...],
+ port: int,
+ path: str,
+ params: Mapping[str, str] | None = None,
+ headers: Mapping[str, str] | None = None,
+ ) -> DorisHTTPResponse:
+ """Try configured endpoints in order on transport or gateway
failure."""
+ candidates = tuple(dict.fromkeys(_normalize_host(host) for host in
hosts))
+ if not candidates:
+ raise DorisHTTPPolicyError("No Doris HTTP endpoint is configured")
+
+ last_error: DorisHTTPRequestError | None = None
+ last_response: DorisHTTPResponse | None = None
+ for host in candidates:
+ try:
+ response = await self.get(
+ role=role,
+ host=host,
+ port=port,
+ path=path,
+ params=params,
+ headers=headers,
+ )
+ except DorisHTTPRequestError as exc:
+ last_error = exc
+ continue
+ if response.status not in {502, 503, 504}:
+ return response
+ last_response = response
+
+ if last_response is not None:
+ return last_response
+ if last_error is not None:
+ raise last_error
+ raise DorisHTTPRequestError("All configured Doris HTTP endpoints
failed")
+
async def get(
self,
*,
diff --git a/doris_mcp_server/utils/monitoring_tools.py
b/doris_mcp_server/utils/monitoring_tools.py
index 30767ab..5ac22d3 100644
--- a/doris_mcp_server/utils/monitoring_tools.py
+++ b/doris_mcp_server/utils/monitoring_tools.py
@@ -28,6 +28,8 @@ from .doris_http_client import (
DorisHTTPPolicyError,
DorisHTTPRequestError,
DorisHTTPResponseTooLarge,
+ configured_fe_http_hosts,
+ database_config_for_request,
)
from .logger import get_logger
@@ -686,7 +688,7 @@ class DorisMonitoringTools:
async def get_be_nodes(self) -> list[dict[str, Any]]:
"""Return only BE HTTP nodes from the explicit configuration
allowlist."""
try:
- db_config = self.connection_manager.config.database
+ db_config = database_config_for_request(self.connection_manager)
be_hosts = getattr(db_config, "be_hosts", []) or []
if be_hosts:
logger.info(f"Using configured BE hosts: {be_hosts}")
@@ -726,16 +728,25 @@ class DorisMonitoringTools:
"""Fetch metrics from an explicitly configured FE or BE node."""
endpoint = f"{node_info.get('host')}:{node_info.get('port') or
node_info.get('http_port')}"
try:
- db_config = self.connection_manager.config.database
+ db_config = database_config_for_request(self.connection_manager)
port_key = "port" if node_type == "fe" else "http_port"
http_client = DorisHTTPClient.from_database_config(db_config)
- response = await http_client.get(
- role=node_type,
- host=node_info["host"],
- port=int(node_info[port_key]),
- path="/metrics",
- headers={"Accept": "text/plain"},
- )
+ if node_type == "fe":
+ response = await http_client.get_first_available(
+ role="fe",
+ hosts=node_info.get("hosts") or [node_info["host"]],
+ port=int(node_info[port_key]),
+ path="/metrics",
+ headers={"Accept": "text/plain"},
+ )
+ else:
+ response = await http_client.get(
+ role="be",
+ host=node_info["host"],
+ port=int(node_info[port_key]),
+ path="/metrics",
+ headers={"Accept": "text/plain"},
+ )
if response.status == 200:
metrics_data = self._parse_prometheus_metrics(response.text())
return {
@@ -970,13 +981,15 @@ class DorisMonitoringTools:
async def _get_fe_metrics(self, monitor_type: str, priority: str,
format_type: str, include_raw_metrics: bool) -> dict[str, Any]:
"""Get FE monitoring metrics"""
try:
- db_config = self.connection_manager.config.database
- fe_http_host = (
- getattr(db_config, "fe_http_host", "") or db_config.host
- )
+ db_config = database_config_for_request(self.connection_manager)
+ fe_http_hosts = configured_fe_http_hosts(db_config)
fe_result = await self.fetch_metrics_from_node(
"fe",
- {"host": fe_http_host, "port": db_config.fe_http_port}
+ {
+ "host": fe_http_hosts[0],
+ "hosts": list(fe_http_hosts),
+ "port": db_config.fe_http_port,
+ },
)
if not fe_result.get("success"):
return fe_result
diff --git a/test/auth/test_token_handlers_boundaries.py
b/test/auth/test_token_handlers_boundaries.py
index 12edf71..58d4684 100644
--- a/test/auth/test_token_handlers_boundaries.py
+++ b/test/auth/test_token_handlers_boundaries.py
@@ -89,10 +89,15 @@ async def
test_create_token_from_get_query_includes_database_config():
"description": "read only",
"custom_token": "chosen-token",
"db_host": "doris.internal",
+ "db_hosts": "doris.internal,doris-2.internal",
"db_port": "9031",
"db_user": "readonly",
"db_password": "secret",
"db_database": "analytics",
+ "db_fe_http_host": "doris-http.internal",
+ "db_fe_http_hosts": (
+ "doris-http.internal,doris-http-2.internal"
+ ),
"db_fe_http_port": "8031",
},
)
@@ -105,10 +110,16 @@ async def
test_create_token_from_get_query_includes_database_config():
assert call.kwargs["custom_token"] == "chosen-token"
database = call.kwargs["database_config"]
assert database.host == "doris.internal"
+ assert database.hosts == ["doris.internal", "doris-2.internal"]
assert database.port == 9031
assert database.user == "readonly"
assert database.password == "secret"
assert database.database == "analytics"
+ assert database.fe_http_host == "doris-http.internal"
+ assert database.fe_http_hosts == [
+ "doris-http.internal",
+ "doris-http-2.internal",
+ ]
assert database.fe_http_port == 8031
@@ -121,13 +132,16 @@ async def
test_create_token_from_post_json_uses_defaults():
"/token/create",
json={
"token_id": "writer",
- "database_config": {"host": "doris.internal"},
+ "database_config": {
+ "hosts": ["doris.internal", "doris-2.internal"]
+ },
},
)
assert response.status_code == 200
database = manager.create_token.await_args.kwargs["database_config"]
assert database.host == "doris.internal"
+ assert database.hosts == ["doris.internal", "doris-2.internal"]
assert database.port == 9030
assert database.user == "root"
assert database.database == "information_schema"
diff --git a/test/protocol/test_multiworker_config.py
b/test/protocol/test_multiworker_config.py
index 25996b2..73f4357 100644
--- a/test/protocol/test_multiworker_config.py
+++ b/test/protocol/test_multiworker_config.py
@@ -100,10 +100,14 @@ def
test_state_handle_secret_and_ttl_are_configurable_without_serializing_secret
def test_multiworker_environment_preserves_resolved_parent_config(monkeypatch):
config = DorisConfig()
config.database.host = "127.0.0.1"
+ config.database.hosts = ["127.0.0.1", "127.0.0.2"]
config.database.port = 19030
config.database.user = "loader"
config.database.password = "test-password"
config.database.database = "hhm_dt_sim"
+ config.database.fe_http_host = "127.0.0.10"
+ config.database.fe_http_hosts = ["127.0.0.10", "127.0.0.11"]
+ config.database.be_hosts = ["127.0.0.20", "127.0.0.21"]
config.server_name = "doris-mcp-server"
config.mcp_allowed_hosts = ["mcp.example.test", "mcp.example.test:*"]
config.mcp_allowed_origins = ["https://client.example.test";]
@@ -126,10 +130,17 @@ def
test_multiworker_environment_preserves_resolved_parent_config(monkeypatch):
child_config = DorisConfig.from_env()
assert child_config.database.host == "127.0.0.1"
+ assert child_config.database.hosts == ["127.0.0.1", "127.0.0.2"]
assert child_config.database.port == 19030
assert child_config.database.user == "loader"
assert child_config.database.password == "test-password"
assert child_config.database.database == "hhm_dt_sim"
+ assert child_config.database.fe_http_host == "127.0.0.10"
+ assert child_config.database.fe_http_hosts == [
+ "127.0.0.10",
+ "127.0.0.11",
+ ]
+ assert child_config.database.be_hosts == ["127.0.0.20", "127.0.0.21"]
assert child_config.server_host == "127.0.0.1"
assert child_config.server_port == 31133
assert child_config.mcp_allowed_hosts == [
diff --git a/test/security/test_token_digest_storage.py
b/test/security/test_token_digest_storage.py
index 22ef3b3..6ca46a3 100644
--- a/test/security/test_token_digest_storage.py
+++ b/test/security/test_token_digest_storage.py
@@ -23,7 +23,7 @@ from pathlib import Path
import pytest
-from doris_mcp_server.auth.token_manager import TokenManager
+from doris_mcp_server.auth.token_manager import DatabaseConfig, TokenManager
from doris_mcp_server.utils.config import (
AuthConfigError,
DorisConfig,
@@ -94,6 +94,62 @@ async def
test_created_token_is_returned_once_and_persisted_as_digest(
reloaded.stop_hot_reload()
[email protected]
+async def test_token_database_fe_candidates_round_trip(
+ tmp_path,
+ monkeypatch,
+):
+ _clear_static_token_environment(monkeypatch)
+ config = _config(tmp_path)
+ manager = TokenManager(config)
+ try:
+ raw_token = await manager.create_token(
+ "multi-fe",
+ database_config=DatabaseConfig(
+ host="fe-1.internal",
+ hosts=["fe-1.internal", "fe-2.internal"],
+ user="tenant_reader",
+ password="tenant-password",
+ database="analytics",
+ fe_http_hosts=["fe-http-1.internal", "fe-http-2.internal"],
+ be_hosts=["be-1.internal", "be-2.internal"],
+ ),
+ )
+ stored = json.loads(
+ Path(config.security.token_file_path).read_text(encoding="utf-8")
+ )["tokens"][0]["database_config"]
+
+ assert stored["hosts"] == ["fe-1.internal", "fe-2.internal"]
+ assert stored["fe_http_hosts"] == [
+ "fe-http-1.internal",
+ "fe-http-2.internal",
+ ]
+ assert stored["be_hosts"] == ["be-1.internal", "be-2.internal"]
+
+ selected = manager.get_database_config_by_token(raw_token)
+ assert selected is not None
+ assert selected.hosts == ["fe-1.internal", "fe-2.internal"]
+ assert selected.fe_http_hosts == [
+ "fe-http-1.internal",
+ "fe-http-2.internal",
+ ]
+ finally:
+ manager.stop_hot_reload()
+
+ reloaded = TokenManager(config)
+ try:
+ result = await reloaded.validate_token(raw_token)
+ assert result.is_valid is True
+ assert result.token_info.token_id == "multi-fe"
+ assert result.token_info.database_config is not None
+ assert result.token_info.database_config.hosts == [
+ "fe-1.internal",
+ "fe-2.internal",
+ ]
+ finally:
+ reloaded.stop_hot_reload()
+
+
@pytest.mark.asyncio
async def
test_legacy_plaintext_file_is_atomically_migrated_without_expiry_extension(
tmp_path,
diff --git a/test/utils/test_multi_fe_routing.py
b/test/utils/test_multi_fe_routing.py
new file mode 100644
index 0000000..7c0a7b1
--- /dev/null
+++ b/test/utils/test_multi_fe_routing.py
@@ -0,0 +1,354 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Regression tests for multi-instance routing and multi-FE failover."""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from doris_mcp_server.auth.token_manager import DatabaseConfig as
TokenDatabaseConfig
+from doris_mcp_server.utils import db as db_module
+from doris_mcp_server.utils.adbc_query_tools import DorisADBCQueryTools
+from doris_mcp_server.utils.config import DorisConfig
+from doris_mcp_server.utils.db import DorisConnectionManager
+from doris_mcp_server.utils.doris_http_client import (
+ DorisHTTPClient,
+ DorisHTTPRequestError,
+ DorisHTTPResponse,
+ configured_fe_http_hosts,
+)
+from doris_mcp_server.utils.security import (
+ AuthContext,
+ reset_auth_context,
+ set_current_auth_context,
+)
+
+
+class _ProbeCursor:
+ def __init__(self, healthy: bool):
+ self.healthy = healthy
+
+ async def __aenter__(self):
+ return self
+
+ async def __aexit__(self, exc_type, exc, traceback):
+ return False
+
+ async def execute(self, sql: str) -> None:
+ assert sql == "SELECT 1"
+
+ async def fetchone(self):
+ return (1,) if self.healthy else None
+
+
+class _ProbeConnection:
+ def __init__(self, healthy: bool):
+ self.healthy = healthy
+ self.closed = False
+
+ def cursor(self):
+ return _ProbeCursor(self.healthy)
+
+
+class _ProbePool:
+ def __init__(self, healthy: bool):
+ self.healthy = healthy
+ self.closed = False
+ self.released = []
+ self.close_calls = 0
+
+ async def acquire(self):
+ return _ProbeConnection(self.healthy)
+
+ def release(self, connection):
+ self.released.append(connection)
+
+ def close(self):
+ self.closed = True
+ self.close_calls += 1
+
+ async def wait_closed(self):
+ return None
+
+
+class _BrokenAcquirePool(_ProbePool):
+ async def acquire(self):
+ raise OSError("active FE is unavailable")
+
+
+def _config(*hosts: str) -> DorisConfig:
+ config = DorisConfig()
+ config.database.host = hosts[0]
+ config.database.hosts = list(hosts)
+ config.database.user = "root"
+ config.database.password = "secret"
+ config.database.connection_timeout = 1
+ config.database.max_connections = 2
+ return config
+
+
+def test_environment_loads_ordered_sql_and_http_fe_hosts(monkeypatch):
+ monkeypatch.setenv("DORIS_HOST", "sql-primary")
+ monkeypatch.setenv("DORIS_HOSTS", "sql-secondary,sql-primary,sql-third")
+ monkeypatch.setenv("DORIS_FE_HTTP_HOST", "http-primary")
+ monkeypatch.setenv(
+ "DORIS_FE_HTTP_HOSTS",
+ "http-secondary,http-primary,http-third",
+ )
+
+ config = DorisConfig.from_env()
+
+ assert config.database.host == "sql-primary"
+ assert config.database.hosts == [
+ "sql-primary",
+ "sql-secondary",
+ "sql-third",
+ ]
+ assert config.database.fe_http_host == "http-primary"
+ assert config.database.fe_http_hosts == [
+ "http-primary",
+ "http-secondary",
+ "http-third",
+ ]
+ serialized = config.to_dict()["database"]
+ assert serialized["hosts"] == config.database.hosts
+ assert serialized["fe_http_hosts"] == config.database.fe_http_hosts
+
+
+def test_config_file_host_lists_set_backward_compatible_primary():
+ config = DorisConfig._from_dict(
+ {
+ "database": {
+ "hosts": ["fe-a", "fe-b"],
+ "fe_http_hosts": ["fe-http-a", "fe-http-b"],
+ }
+ }
+ )
+
+ assert config.database.host == "fe-a"
+ assert config.database.fe_http_host == "fe-http-a"
+
+
[email protected]
+async def test_global_pool_fails_over_to_second_fe(monkeypatch):
+ manager = DorisConnectionManager(_config("fe-down", "fe-up"))
+ failed_pool = _ProbePool(healthy=False)
+ healthy_pool = _ProbePool(healthy=True)
+ create_pool = AsyncMock(side_effect=[failed_pool, healthy_pool])
+ monkeypatch.setattr(db_module.aiomysql, "create_pool", create_pool)
+
+ await manager._create_global_pool()
+
+ assert [call.kwargs["host"] for call in create_pool.await_args_list] == [
+ "fe-down",
+ "fe-up",
+ ]
+ assert failed_pool.closed is True
+ assert manager.pool is healthy_pool
+ assert manager.host == "fe-up"
+ assert manager.active_db_config["host"] == "fe-up"
+
+
[email protected]
+async def test_doris_user_authentication_uses_same_fe_candidates(monkeypatch):
+ manager = DorisConnectionManager(_config("fe-down", "fe-up"))
+ good_connection = SimpleNamespace(close=lambda: None)
+ connect = AsyncMock(
+ side_effect=[OSError("unreachable"), good_connection],
+ )
+ monkeypatch.setattr(db_module.aiomysql, "connect", connect)
+
+ await manager.authenticate_doris_user("alice", "correct-password")
+
+ assert [call.kwargs["host"] for call in connect.await_args_list] == [
+ "fe-down",
+ "fe-up",
+ ]
+ assert manager.host == "fe-up"
+
+
+def test_token_selects_an_independent_cluster_and_http_allowlist():
+ manager = DorisConnectionManager(_config("global-fe"))
+ tenant_config = TokenDatabaseConfig(
+ host="tenant-a-fe-1",
+ hosts=["tenant-a-fe-1", "tenant-a-fe-2"],
+ user="tenant_a",
+ password="tenant-secret",
+ fe_http_hosts=["tenant-a-http-1", "tenant-a-http-2"],
+ )
+ manager.token_manager = SimpleNamespace(
+ get_database_config_by_token=lambda token: (
+ tenant_config if token == "tenant-a-token" else None
+ )
+ )
+ auth_context = AuthContext(
+ user_id="tenant-a",
+ auth_method="token",
+ token="tenant-a-token",
+ )
+
+ selected = manager.get_database_config_for_auth_context(auth_context)
+
+ assert selected is tenant_config
+ assert configured_fe_http_hosts(selected) == (
+ "tenant-a-http-1",
+ "tenant-a-http-2",
+ )
+ client = DorisHTTPClient.from_database_config(selected)
+ assert client.allowed_endpoints["fe"] == {
+ ("tenant-a-http-1", 8030),
+ ("tenant-a-http-2", 8030),
+ }
+
+
[email protected]
+async def test_token_bound_pool_fails_over_inside_its_own_cluster(monkeypatch):
+ manager = DorisConnectionManager(_config("global-fe"))
+ tenant_config = TokenDatabaseConfig(
+ host="tenant-fe-down",
+ hosts=["tenant-fe-down", "tenant-fe-up"],
+ user="tenant_user",
+ password="tenant-secret",
+ database="tenant_db",
+ )
+ manager.token_manager = SimpleNamespace(
+ get_database_config_by_token=lambda token: (
+ tenant_config if token == "tenant-token" else None
+ )
+ )
+ failed_pool = _ProbePool(healthy=False)
+ healthy_pool = _ProbePool(healthy=True)
+ create_pool = AsyncMock(side_effect=[failed_pool, healthy_pool])
+ monkeypatch.setattr(db_module.aiomysql, "create_pool", create_pool)
+
+ pool, selected_config = await manager.get_pool_for_token("tenant-token")
+
+ assert pool is healthy_pool
+ assert selected_config["hosts"] == ["tenant-fe-down", "tenant-fe-up"]
+ assert [call.kwargs["host"] for call in create_pool.await_args_list] == [
+ "tenant-fe-down",
+ "tenant-fe-up",
+ ]
+ assert failed_pool.closed is True
+
+
[email protected]
+async def test_token_bound_pool_recreates_after_active_fe_failure(monkeypatch):
+ manager = DorisConnectionManager(_config("global-fe"))
+ tenant_config = TokenDatabaseConfig(
+ host="tenant-fe-1",
+ hosts=["tenant-fe-1", "tenant-fe-2"],
+ user="tenant_user",
+ password="tenant-secret",
+ database="tenant_db",
+ )
+ manager.token_manager = SimpleNamespace(
+ get_database_config_by_token=lambda token: (
+ tenant_config if token == "tenant-token" else None
+ )
+ )
+ token_hash = manager._get_token_hash("tenant-token")
+ broken_pool = _BrokenAcquirePool(healthy=False)
+ recovered_pool = _ProbePool(healthy=True)
+ manager.token_pools[token_hash] = broken_pool
+ manager.token_configs[token_hash] = {
+ "host": "tenant-fe-1",
+ "hosts": ["tenant-fe-1", "tenant-fe-2"],
+ "port": 9030,
+ "user": "tenant_user",
+ "password": "tenant-secret",
+ "database": "tenant_db",
+ "charset": "UTF8",
+ }
+ manager._token_pool_owner_ids[token_hash] = "static_token:old"
+ manager._token_pool_generations[token_hash] = 3
+ create_pool = AsyncMock(return_value=recovered_pool)
+ monkeypatch.setattr(manager, "_create_pool_with_config", create_pool)
+
+ connection = await manager.get_connection_for_token(
+ "tenant-token",
+ "recovering-session",
+ )
+
+ assert broken_pool.closed is True
+ assert connection.owner_pool is recovered_pool
+ assert manager.token_pools[token_hash] is recovered_pool
+ assert connection.generation == 4
+ create_pool.assert_awaited_once()
+
+
[email protected]
+async def test_adbc_fails_closed_for_token_bound_cluster():
+ manager = DorisConnectionManager(_config("global-fe"))
+ tenant_config = TokenDatabaseConfig(
+ host="tenant-fe",
+ user="tenant_user",
+ password="tenant-secret",
+ )
+ manager.token_manager = SimpleNamespace(
+ get_database_config_by_token=lambda token: (
+ tenant_config if token == "tenant-token" else None
+ )
+ )
+ context_token = set_current_auth_context(
+ AuthContext(
+ user_id="tenant",
+ auth_method="token",
+ token="tenant-token",
+ )
+ )
+ try:
+ result = await DorisADBCQueryTools(manager).exec_adbc_query("SELECT 1")
+ finally:
+ reset_auth_context(context_token)
+
+ assert result["success"] is False
+ assert result["error_type"] == "token_bound_adbc_unsupported"
+
+
[email protected]
+async def test_http_request_fails_over_only_within_allowlisted_hosts():
+ client = DorisHTTPClient(
+ user="root",
+ password="secret",
+ allowed_endpoints={
+ "fe": {("fe-a", 8030), ("fe-b", 8030)},
+ "be": set(),
+ },
+ )
+ expected = DorisHTTPResponse(
+ status=200,
+ headers={},
+ body=b"ok",
+ url="http://fe-b:8030/metrics";,
+ )
+ client.get = AsyncMock( # type: ignore[method-assign]
+ side_effect=[DorisHTTPRequestError("down"), expected]
+ )
+
+ response = await client.get_first_available(
+ role="fe",
+ hosts=["fe-a", "fe-b"],
+ port=8030,
+ path="/metrics",
+ )
+
+ assert response is expected
+ assert [call.kwargs["host"] for call in client.get.await_args_list] == [
+ "fe-a",
+ "fe-b",
+ ]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
