You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by yo...@apache.org on 2026/09/10 03:18:35 UTC
(iotdb-mcp-server) branch main updated: Refactor IoTDB MCP Server and migrate agent enhancements (#18)
This is an automated email from the ASF dual-hosted git repository.
CRZbulabula pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iotdb-mcp-server.git
The following commit(s) were added to refs/heads/main by this push:
new 5851bc4 Refactor IoTDB MCP Server and migrate agent enhancements (#18)
5851bc4 is described below
commit 5851bc4b867b7491a9e001dc3338f796844bce9b
Author: Zeyu Zhang <11...@users.noreply.github.com>
AuthorDate: Thu Sep 10 11:18:30 2026 +0800
Refactor IoTDB MCP Server and migrate agent enhancements (#18)
* Refactor code and update server to current agent style
* Mask password in startup config log output (#19)
* Mask password in startup config log output
The startup log printed the full IoTDB config dict including the password field. Log a copy of the config with the password masked instead, and exclude the password from the Config repr.
* Remove outdated smithery badge links
* Upgrade apache-iotdb to 2.0.10
* fix: fix P1 review problem
* fix: fix P2 review problem
* fix: fix bug in config params
* fix: fix broken TableSession clear bug
* fix: fix broken TableSession discard path
* fix: close broken pipe after `ConnectionError`
---------
Co-authored-by: CritasWang <cr...@outlook.com>
---
Dockerfile | 2 +-
README-zh.md | 110 +-
README.md | 114 +-
pyproject.toml | 3 +-
smithery.yaml | 4 +-
src/iotdb_mcp_server/__init__.py | 113 +-
src/iotdb_mcp_server/config.py | 285 ++++-
src/iotdb_mcp_server/result_store.py | 845 ++++++++++++
src/iotdb_mcp_server/runtime_policy.py | 353 +++++
src/iotdb_mcp_server/server.py | 577 +--------
src/iotdb_mcp_server/services/README.md | 394 ++++++
src/iotdb_mcp_server/services/database.py | 348 +++++
src/iotdb_mcp_server/services/explain.py | 168 +++
src/iotdb_mcp_server/services/json_response.py | 329 +++++
src/iotdb_mcp_server/services/metadata.py | 416 ++++++
src/iotdb_mcp_server/services/model.py | 585 +++++++++
src/iotdb_mcp_server/services/query.py | 434 +++++++
src/iotdb_mcp_server/services/results.py | 226 ++++
src/iotdb_mcp_server/services/runtime_policy.py | 77 ++
src/iotdb_mcp_server/services/sql_driver.py | 1348 ++++++++++++++++++++
src/iotdb_mcp_server/services/table.py | 178 +++
src/iotdb_mcp_server/services/target_selection.py | 492 +++++++
src/iotdb_mcp_server/services/targets.py | 533 ++++++++
src/iotdb_mcp_server/services/timeseries.py | 351 +++++
.../services/tree_sql_guardrails.py | 199 +++
src/iotdb_mcp_server/services/ttl.py | 185 +++
src/iotdb_mcp_server/services/udf.py | 534 ++++++++
src/iotdb_mcp_server/services/write.py | 163 +++
src/iotdb_mcp_server/session_manager.py | 463 +++++++
src/iotdb_mcp_server/target_cli.py | 173 +++
src/iotdb_mcp_server/target_lifecycle.py | 213 ++++
src/iotdb_mcp_server/target_registry.py | 830 ++++++++++++
tests/test_config_fallback.py | 187 +++
tests/test_database_target_routing.py | 175 +++
tests/test_readonly_target_routing.py | 220 ++++
tests/test_result_store.py | 464 +++++++
tests/test_runtime_policy.py | 83 ++
tests/test_session_manager_auth_cache.py | 258 ++++
tests/test_sql_driver_batch.py | 406 ++++++
tests/test_target_cli.py | 186 +++
tests/test_target_lifecycle.py | 335 +++++
tests/test_udf_results.py | 65 +
42 files changed, 12780 insertions(+), 644 deletions(-)
diff --git a/Dockerfile b/Dockerfile
index 92570d5..44cbd62 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,5 +1,5 @@
# Generated by https://smithery.ai. See: https://smithery.ai/docs/config#dockerfile
-FROM python:3.12-alpine
+FROM python:3.11-alpine
# Install dependencies
RUN apk add --no-cache gcc musl-dev linux-headers
diff --git a/README-zh.md b/README-zh.md
index c01ea90..112d2f5 100644
--- a/README-zh.md
+++ b/README-zh.md
@@ -16,10 +16,31 @@ IoTDB MCP Server 是一个基于模型上下文协议（Model Context Protocol,
此服务器不提供任何提示。
+## 权限模型
+
+IoTDB MCP 权限默认是提示层。服务器为 SQL 操作返回所需权限、风险级别和确认参数，
+由宿主 agent system 负责用户审批。执行 DDL/DML 或破坏性 SQL 前，如果可用，应先使用
+`inspect_sql_permission`。
+
+长期托管 agent 可以通过环境变量提供 full permission 默认值，例如
+`IOTDB_SQL_DRIVER_MODE=full` 和 `TIMESEEK_MCP_PERMISSION_ENFORCEMENT=advisory`。
+只有需要 MCP server 自身硬阻断时，才设置
+`TIMESEEK_MCP_PERMISSION_ENFORCEMENT=strict`。
+
### 工具
服务器为 IoTDB 的树模型（Tree Model）和表模型（Table Model）提供了不同的工具。您可以通过设置 "IOTDB_SQL_DIALECT" 配置为 "tree" 或 "table" 来选择使用哪种模型。
+方言相关的标识符语义：
+
+- 树模型：
+ - `FROM` 中使用显式的 `root...` 路径。
+ - `SELECT` 投影通常应使用测点名，而不是完整 `root...` 路径。
+ - `IOTDB_DATABASE` 只作为连接/会话范围提示，树模型 SQL 仍应显式写出 `root...` 路径。
+- 表模型：
+ - `FROM` 中使用当前数据库下的表名。
+ - `SELECT` 投影使用列名。
+
#### 树模型 (Tree Model)
- `metadata_query`
@@ -47,6 +68,32 @@ IoTDB MCP Server 是一个基于模型上下文协议（Model Context Protocol,
- SUM, COUNT, MAX_VALUE, MIN_VALUE, AVG, VARIANCE, MAX_TIME, MIN_TIME 等
- 返回: 查询结果作为对象数组
+- `sql_executor_batch`
+
+ - 并行执行多条只读 SQL，并将每个子结果写入 ResultStore
+ - 输入:
+ - `sqls` (数组): 显式单语句 SQL，或
+ - `sql_template` + `param_sets`: SQL 模板和参数对象列表
+ - `max_concurrency` (整数): 并发语句上限，默认 4
+ - `worker_pool_size` (整数): 线程 worker pool 大小，默认跟随并发度，并受 `IOTDB_SQL_EXECUTOR_BATCH_MAX_WORKER_POOL_SIZE` 限制（默认 16）
+ - `per_item_timeout_ms` (整数): 单条语句等待超时，默认 60000
+ - `batch_timeout_ms` (整数): 整个 batch 等待超时，默认 300000
+ - `max_result_rows_per_item` / `max_result_bytes_per_item`: 单条语句结果 quota，默认 10000 行和 16 MiB
+ - `max_batch_result_rows` / `max_batch_result_bytes`: 整个 batch 结果 quota，默认 100000 行和 64 MiB
+ - 模板占位符:
+ - `{{name}}` 表示 SQL 字面量，`{{name:path}}` 表示 IoTDB 路径，`{{name:identifier}}` 表示 SQL 标识符
+ - 返回: batch 汇总，以及每条 SQL 的 `result_id`、行数、预览行和分页信息
+
+- `read_result_pages`
+
+ - 一次 MCP 调用读取多个 ResultStore page
+ - 输入:
+ - `pages` (数组): page 请求对象，包含 `result_id`，可选 `cursor`、`offset`、`limit`、`owner_session_id`
+ - `default_limit` (整数): 未设置 `limit` 时的默认 page 大小
+ - `max_pages` / `max_total_rows`: 单次调用 quota，默认 32 个 page 和 10000 行
+ - `continue_on_error` (布尔): 是否以 per-item error 返回，默认 true
+ - 返回: 批量读页汇总，以及每个 page 的 rows、cursor 和错误信息
+
- `export_query`
- 执行查询并将结果导出为 CSV 或 Excel 文件
- 输入:
@@ -54,6 +101,37 @@ IoTDB MCP Server 是一个基于模型上下文协议（Model Context Protocol,
- `format` (字符串): 导出格式，可以是 "csv" 或 "excel"（默认: "csv"）
- `filename` (字符串): 导出文件的文件名（可选，如果未提供，将生成唯一文件名）
- 返回: 有关导出文件的信息和数据预览（最多 10 行）
+- `model_inference`
+ - 执行 AINode `CALL INFERENCE(...)` SQL 并返回结果集
+ - 输入:
+ - `inference_sql` (字符串): 以 `CALL INFERENCE` 开头的单条树模型 SQL
+ - 执行前校验模型 ID、带引号的输入 SELECT SQL、显式非通配列，以及支持的参数
+ (`generateTime`, `outputLength`)
+ - 权限信息：模型管理能力通过 MCP advisory policy 层返回。仅在 strict 模式下才硬性执行
+ `IOTDB_ENABLE_MODEL_MANAGEMENT=true` 和 `IOTDB_MODEL_ALLOWED_USERS`。
+- `prepare_model_inference_request`
+ - 通过结构化字段构造并校验 AINode `CALL INFERENCE(...)` SQL
+ - 输入:
+ - `model_id` (字符串): 已注册的 AINode 模型 ID
+ - `input_sql` (字符串): 作为模型输入的有界树模型 SELECT 查询
+ - `output_length` (整数): 预测输出长度（默认: 96）
+ - `generate_time` (布尔): 是否请求 Time 列（默认: false）
+
+#### UDF 工具
+
+- `list_udf_functions`
+ - 对所选 IoTDB target 执行 `SHOW FUNCTIONS`。
+- `prepare_udf_query`
+ - 使用结构化参数生成只读 UDF `SELECT`。
+ - 树模型形式：`SELECT UDF(measurement, "k"="v") FROM root.sg.d1 ...`
+ - 表模型形式：`SELECT UDF(column, "k"="v") FROM table ...`
+- `execute_udf_query`
+ - 执行校验后的 UDF 查询，并返回 ResultStore 支持的预览结果。
+- `export_udf_query`
+ - 执行校验后的 UDF 查询，并导出为 CSV 或 Excel。
+
+UDF 工具会拒绝表达式和过滤条件中的分号、SQL 注释以及 DDL/DML 关键字。
+它们用于只读 UDF 调用，例如数据质量、画像、修复规划和异常评分。
#### 表模型 (Table Model)
@@ -94,11 +172,37 @@ IoTDB MCP Server 支持以下配置选项，可以通过环境变量或命令行
| --host | IOTDB_HOST | 127.0.0.1 | IoTDB 主机地址 |
| --port | IOTDB_PORT | 6667 | IoTDB 端口 |
| --user | IOTDB_USER | root | IoTDB 用户名 |
-| --password | IOTDB_PASSWORD | root | IoTDB 密码 |
-| --database | IOTDB_DATABASE | test | IoTDB 数据库名称 |
+| --password | IOTDB_PASSWORD | 空 | IoTDB 密码 |
+| --database | IOTDB_DATABASE | test | 表模型: 当前数据库名；树模型: 可选的连接/范围提示，查询仍需显式使用 `root...` 路径 |
| --sql-dialect | IOTDB_SQL_DIALECT | table | SQL 方言: tree 或 table |
| --export-path | IOTDB_EXPORT_PATH | /tmp | 查询结果导出路径 |
+target registry 只保存已经成功登录的连接。先用 `prepare_iotdb_target` 创建
+不含凭据、不可执行 SQL 的临时 candidate，再把用户明确提供的账号密码传给
+`connect_iotdb_target`，并且只认证一次。缺少凭据时必须询问用户，禁止探测空
+密码或默认密码；用户明确说明空密码时，空字符串才是有效凭据。
+
+登录成功后，MCP 会原子发布 target，并自动记录 per-target
+`last_known_good_credential`。任何连接层错误都会消费 candidate 或驱逐已发布
+target。只有用户明确要求重连后，才能创建新 candidate，并传入
+`user_confirmed_retry=true`。公开 target 响应会脱敏当前密码和
+last-known-good 密码。
+
+配置 `TIMESEEK_IOTDB_TARGETS_FILE` 后，成功连接默认会持久化。通过
+`iotdb-target-cli` 启动本地 Java CLI，即可复用完全相同的 target：
+
+```bash
+iotdb-target-cli --target-id cloud \
+ --cli /opt/iotdb/sbin/start-cli.sh -- -e "SHOW VERSION"
+```
+
+wrapper 每次启动都会重新加载已验证 target，并注入 host、port、方言、用户名和
+last-known-good 密码。Java CLI 不支持 `-db`，因此 wrapper 不会向
+`start-cli.sh` 传该参数；Table 数据库应通过 SQL `USE <database>` 选择。对于支持
+`-db` 的 `import-data.sh` 和 `import-data.bat`，wrapper 仍会自动注入 Table
+target 的数据库。已验证的空密码会省略 `-pw`，命令预览会脱敏非空密码。直接调用
+`start-cli.sh` 不会读取 target registry。
+
## 性能优化
IoTDB MCP Server 包含以下性能优化特性：
@@ -164,7 +268,7 @@ uv sync
"IOTDB_HOST": "127.0.0.1",
"IOTDB_PORT": "6667",
"IOTDB_USER": "root",
- "IOTDB_PASSWORD": "root",
+ "IOTDB_PASSWORD": "",
"IOTDB_DATABASE": "test",
"IOTDB_SQL_DIALECT": "table",
"IOTDB_EXPORT_PATH": "/path/to/export/folder"
diff --git a/README.md b/README.md
index b94317b..622e785 100644
--- a/README.md
+++ b/README.md
@@ -16,10 +16,33 @@ The server doesn't expose any resources.
The server doesn't provide any prompts.
+## Permission Model
+
+IoTDB MCP permissions are advisory by default. The server reports the required
+permission, risk level, and confirmation parameter for SQL actions, while the
+host agent system owns user approval. Use `inspect_sql_permission` before
+executing DDL/DML or destructive SQL when the tool is available.
+
+Long-running hosted agents can provide full-permission defaults with
+environment variables such as `IOTDB_SQL_DRIVER_MODE=full` and
+`TIMESEEK_MCP_PERMISSION_ENFORCEMENT=advisory`. Set
+`TIMESEEK_MCP_PERMISSION_ENFORCEMENT=strict` only when the MCP server itself
+should hard-block disallowed operations.
+
### Tools
The server offers different tools for IoTDB Tree Model and Table Model. You can choose between them by setting the "IOTDB_SQL_DIALECT" configuration to either "tree" or "table".
+Dialect-specific identifier semantics:
+
+- Tree dialect:
+ - `FROM` targets use explicit `root...` paths.
+ - Projection expressions should usually use measurement names instead of full `root...` paths.
+ - `IOTDB_DATABASE` is only a connection/session hint; tree SQL still needs explicit root paths.
+- Table dialect:
+ - `FROM` targets use table names in the current database.
+ - Projection expressions use column names.
+
#### Tree Model
- `metadata_query`
@@ -43,6 +66,28 @@ The server offers different tools for IoTDB Tree Model and Table Model. You can
- Supported functions:
- SUM, COUNT, MAX_VALUE, MIN_VALUE, AVG, VARIANCE, MAX_TIME, MIN_TIME, etc.
- Returns: Query results as array of objects
+- `sql_executor_batch`
+ - Execute multiple readonly SQL statements in parallel and store each result in ResultStore
+ - Input:
+ - `sqls` (array): Explicit single SQL statements, or
+ - `sql_template` + `param_sets`: Repeated SQL template with parameter objects
+ - `max_concurrency` (integer): Concurrent statement limit, default 4
+ - `worker_pool_size` (integer): Thread worker pool size, default follows concurrency and is capped by `IOTDB_SQL_EXECUTOR_BATCH_MAX_WORKER_POOL_SIZE` (default 16)
+ - `per_item_timeout_ms` (integer): Per-statement wait timeout, default 60000
+ - `batch_timeout_ms` (integer): Whole-batch wait timeout, default 300000
+ - `max_result_rows_per_item` / `max_result_bytes_per_item`: Per-statement result quota, defaults 10000 rows and 16 MiB
+ - `max_batch_result_rows` / `max_batch_result_bytes`: Whole-batch result quota, defaults 100000 rows and 64 MiB
+ - Template placeholders:
+ - `{{name}}` for SQL literals, `{{name:path}}` for IoTDB paths, `{{name:identifier}}` for SQL identifiers
+ - Returns: Batch summary plus per-statement `result_id`, row count, preview rows, and paging metadata
+- `read_result_pages`
+ - Read multiple ResultStore pages in one MCP call
+ - Input:
+ - `pages` (array): Page request objects with `result_id` plus optional `cursor`, `offset`, `limit`, and `owner_session_id`
+ - `default_limit` (integer): Default page size for items without `limit`
+ - `max_pages` / `max_total_rows`: Per-call quotas, defaults 32 pages and 10000 rows
+ - `continue_on_error` (boolean): Return per-item errors instead of aborting, default true
+ - Returns: Batch page summary plus per-page rows, cursors, and errors
- `export_query`
- Execute a query and export the results to a CSV or Excel file
- Input:
@@ -50,6 +95,39 @@ The server offers different tools for IoTDB Tree Model and Table Model. You can
- `format` (string): Export format, either "csv" or "excel" (default: "csv")
- `filename` (string): Optional filename for the exported file. If not provided, a unique filename will be generated.
- Returns: Information about the exported file and a preview of the data (max 10 rows)
+- `model_inference`
+ - Execute AINode `CALL INFERENCE(...)` SQL and return the result set
+ - Input:
+ - `inference_sql` (string): A single Tree-dialect SQL statement starting with `CALL INFERENCE`
+ - Validates model id, quoted input SELECT SQL, explicit non-wildcard columns,
+ and supported parameters (`generateTime`, `outputLength`) before execution
+ - Permission metadata: model management is reported through the advisory MCP
+ policy layer. In strict mode, `IOTDB_ENABLE_MODEL_MANAGEMENT=true` and
+ `IOTDB_MODEL_ALLOWED_USERS` are enforced.
+- `prepare_model_inference_request`
+ - Build and validate AINode `CALL INFERENCE(...)` SQL from structured fields
+ - Input:
+ - `model_id` (string): Registered AINode model id
+ - `input_sql` (string): Bounded Tree-dialect SELECT query used as model input
+ - `output_length` (int): Forecast output length (default: 96)
+ - `generate_time` (bool): Whether to request a Time column (default: false)
+
+#### UDF Tools
+
+- `list_udf_functions`
+ - Execute `SHOW FUNCTIONS` for the selected IoTDB target.
+- `prepare_udf_query`
+ - Build a read-only UDF `SELECT` from structured inputs.
+ - Tree form: `SELECT UDF(measurement, "k"="v") FROM root.sg.d1 ...`
+ - Table form: `SELECT UDF(column, "k"="v") FROM table ...`
+- `execute_udf_query`
+ - Execute the validated UDF query and return a ResultStore-backed preview.
+- `export_udf_query`
+ - Execute the validated UDF query and export the result set to CSV or Excel.
+
+UDF tools reject semicolons, SQL comments, and DDL/DML keywords in expressions
+and filter clauses. They are intended for read-only UDF calls such as data
+quality, profiling, repair planning, and anomaly scoring.
#### Table Model
@@ -93,11 +171,41 @@ IoTDB MCP Server supports the following configuration options, which can be set
| --host | IOTDB_HOST | 127.0.0.1 | IoTDB host address |
| --port | IOTDB_PORT | 6667 | IoTDB port |
| --user | IOTDB_USER | root | IoTDB username |
-| --password | IOTDB_PASSWORD | root | IoTDB password |
-| --database | IOTDB_DATABASE | test | IoTDB database name |
+| --password | IOTDB_PASSWORD | empty | IoTDB password |
+| --database | IOTDB_DATABASE | test | Table dialect: current database name. Tree dialect: optional session/root scope hint; queries still use explicit `root...` paths. |
| --sql-dialect | IOTDB_SQL_DIALECT | table | SQL dialect: tree or table |
| --export-path | IOTDB_EXPORT_PATH | /tmp | Path for exporting query results |
+The target registry contains only connections that have completed a successful
+login. Call `prepare_iotdb_target` with non-secret fields, then pass explicitly
+user-supplied credentials to `connect_iotdb_target` for one authentication
+attempt. If credentials are absent, ask the user; never probe empty or default
+passwords. An explicitly supplied empty password remains valid input.
+
+Successful login atomically publishes the target and records its per-target
+`last_known_good_credential`. Any connection-layer failure consumes the
+candidate or evicts the published target. A retry requires a new candidate and
+`user_confirmed_retry=true` after explicit user instruction. Public target
+responses redact both the active password and last-known-good password.
+
+When `TIMESEEK_IOTDB_TARGETS_FILE` is configured, successful connections are
+persisted by default. The local Java CLI can reuse exactly that target through
+`iotdb-target-cli`:
+
+```bash
+iotdb-target-cli --target-id cloud \
+ --cli /opt/iotdb/sbin/start-cli.sh -- -e "SHOW VERSION"
+```
+
+The wrapper reloads the verified target on every invocation and supplies its
+host, port, dialect, username, and last-known-good password. It does not pass
+`-db` to `start-cli.sh`, because the Java CLI does not support that option;
+select a table database with SQL `USE <database>`. For `import-data.sh` and
+`import-data.bat`, which do support `-db`, a table target's database is injected
+automatically. A verified empty password is represented by omitting `-pw`, and
+command previews redact non-empty passwords. Calling `start-cli.sh` directly
+does not read the target registry.
+
## Performance Optimizations
IoTDB MCP Server includes the following performance optimization features:
@@ -163,7 +271,7 @@ Add the following configuration to Claude Desktop's configuration file:
"IOTDB_HOST": "127.0.0.1",
"IOTDB_PORT": "6667",
"IOTDB_USER": "root",
- "IOTDB_PASSWORD": "root",
+ "IOTDB_PASSWORD": "",
"IOTDB_DATABASE": "test",
"IOTDB_SQL_DIALECT": "table",
"IOTDB_EXPORT_PATH": "/path/to/export/folder"
diff --git a/pyproject.toml b/pyproject.toml
index db52b61..ca73270 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ description = "A Model Context Protocol (MCP) server that enables secure interac
readme = "README.md"
license = "Apache-2.0"
license-file="LICENSE"
-requires-python = ">=3.12"
+requires-python = ">=3.11"
dependencies = [
"fastmcp>=2.8.1",
"apache-iotdb>=2.0.10",
@@ -33,3 +33,4 @@ dev-dependencies = [
[project.scripts]
iotdb-mcp-server = "iotdb_mcp_server:main"
+iotdb-target-cli = "iotdb_mcp_server.target_cli:main"
diff --git a/smithery.yaml b/smithery.yaml
index 9d8a9bb..a820101 100644
--- a/smithery.yaml
+++ b/smithery.yaml
@@ -26,7 +26,7 @@ startCommand:
description: IoTDB username
password:
type: string
- default: root
+ default: ""
description: IoTDB password
database:
type: string
@@ -50,5 +50,5 @@ startCommand:
host: 127.0.0.1
port: 6667
user: root
- password: root
+ password: ""
database: test
diff --git a/src/iotdb_mcp_server/__init__.py b/src/iotdb_mcp_server/__init__.py
index 184f75b..86dc83d 100644
--- a/src/iotdb_mcp_server/__init__.py
+++ b/src/iotdb_mcp_server/__init__.py
@@ -1,73 +1,40 @@
-# 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.
-#
-
-from iotdb_mcp_server.config import Config
-import sys
-import signal
-import asyncio
-import logging
-
-if not "-m" in sys.argv:
- from . import server
-
-# Configure logging for the main module
-logger = logging.getLogger("iotdb_mcp_server")
-
-def main():
- """Main entry point for the package."""
- # Create an event loop with proper exception handling
- loop = asyncio.new_event_loop()
- asyncio.set_event_loop(loop)
-
- # Graceful shutdown handler
- def handle_shutdown(sig, frame):
- logger.info(f"Received shutdown signal {sig}, shutting down gracefully...")
- # Stop the event loop
- if loop.is_running():
- loop.stop()
-
- # Register signal handlers
- signal.signal(signal.SIGINT, handle_shutdown)
- signal.signal(signal.SIGTERM, handle_shutdown)
-
- try:
- logger.info("Starting IoTDB MCP Server...")
- # Run the server in the event loop
- loop.run_until_complete(server.main())
- # Keep the loop running until stopped
- loop.run_forever()
- except Exception as e:
- logger.error(f"Error in IoTDB MCP Server: {e}")
- finally:
- # Clean up tasks
- pending = asyncio.all_tasks(loop)
- if pending:
- logger.info(f"Cancelling {len(pending)} pending tasks...")
- for task in pending:
- task.cancel()
- # Wait for tasks to be cancelled
- loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
- # Close the loop
- loop.close()
- logger.info("IoTDB MCP Server shutdown complete")
-
- return 0
-
-
-# Expose important items at package level
-__all__ = ["main", "server"]
+# 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.
+#
+
+import logging
+
+logger = logging.getLogger("iotdb_mcp_server")
+
+
+def main() -> int:
+ """Console script entrypoint."""
+ from . import server
+
+ try:
+ logger.info("Starting IoTDB MCP Server...")
+ server.main()
+ return 0
+ except KeyboardInterrupt:
+ logger.info("IoTDB MCP Server interrupted")
+ return 0
+ except Exception as e:
+ logger.error(f"Error in IoTDB MCP Server: {e}")
+ return 1
+
+
+__all__ = ["main"]
diff --git a/src/iotdb_mcp_server/config.py b/src/iotdb_mcp_server/config.py
index 7e4ad78..f0b542e 100644
--- a/src/iotdb_mcp_server/config.py
+++ b/src/iotdb_mcp_server/config.py
@@ -19,6 +19,43 @@
import argparse
from dataclasses import dataclass, field
import os
+from pathlib import Path
+from typing import Any
+
+from iotdb_mcp_server.target_registry import IoTDBTarget
+from iotdb_mcp_server.target_registry import IoTDBTargetRegistry
+from iotdb_mcp_server.target_registry import target_from_mapping
+
+
+def env_or_default(name: str, default: str) -> str:
+ value = os.getenv(name)
+ if value is None or value.strip() == "":
+ return default
+ return value
+
+
+def _argument_or_env(value: Any, name: str) -> Any:
+ """Return an explicit CLI value, otherwise the matching environment value."""
+ return value if value is not None else os.getenv(name)
+
+
+def default_export_path() -> str:
+ explicit = os.getenv("IOTDB_EXPORT_PATH")
+ if explicit:
+ return explicit
+
+ workspace = os.getenv("TIMESEEK_WORKSPACE_ROOT") or os.getenv(
+ "TIMESEEK_EXPERIMENT_WORKSPACE_ROOT"
+ )
+ if workspace:
+ return str(
+ Path(workspace).expanduser()
+ / "auto_rule"
+ / ".tmp"
+ / "iotdb_exports"
+ )
+
+ return "/tmp"
@dataclass
@@ -49,19 +86,130 @@ class Config:
database: str
"""
- IoTDB database
+ IoTDB database/session scope.
+ Table dialect: current database name.
+ Tree dialect: optional root scope hint; tree SQL still uses explicit root paths in queries.
"""
sql_dialect: str
"""
SQL dialect: tree or table
"""
+
+ timezone: str
+ """
+ IoTDB session timezone.
+ """
export_path: str
"""
Path for exporting query results
"""
+ target_id: str = "default"
+ display_name: str = "Default IoTDB"
+ target_kind: str = "local"
+ node_urls: tuple[str, ...] = ()
+ iotdb_home: str = ""
+ use_ssl: bool = False
+ ca_certs: str = ""
+ connection_timeout_in_ms: int | None = None
+ enable_redirection: bool = True
+ enable_compression: bool = False
+ fetch_size: int = 1024
+ max_retry: int = 3
+ max_pool_size: int = 100
+ tree_wait_timeout_in_ms: int = 5000
+ table_wait_timeout_in_ms: int = 10000
+ target_fingerprint: str = ""
+ policy: dict[str, str] = field(default_factory=dict)
+ last_known_good_credential: dict[str, object] = field(
+ default_factory=dict,
+ repr=False,
+ compare=False,
+ )
+ verified_at: str = ""
+ target_registry: IoTDBTargetRegistry | None = field(default=None, repr=False, compare=False)
+ session_manager: Any | None = field(default=None, repr=False, compare=False)
+ targets_file: str = ""
+ targets_json: str = field(default="", repr=False, compare=False)
+ targets_file_mtime_ns: int | None = field(default=None, repr=False, compare=False)
+
+ @classmethod
+ def from_target(
+ cls,
+ target: IoTDBTarget,
+ target_registry: IoTDBTargetRegistry | None = None,
+ session_manager: Any | None = None,
+ ) -> "Config":
+ return cls(
+ host=target.host,
+ port=target.port,
+ user=target.user,
+ password=target.password,
+ database=target.database,
+ sql_dialect=target.sql_dialect,
+ timezone=target.timezone,
+ export_path=target.export_path,
+ target_id=target.target_id,
+ display_name=target.display_name,
+ target_kind=target.kind,
+ node_urls=target.node_urls,
+ iotdb_home=target.iotdb_home,
+ use_ssl=target.use_ssl,
+ ca_certs=target.ca_certs,
+ connection_timeout_in_ms=target.connection_timeout_in_ms,
+ enable_redirection=target.enable_redirection,
+ enable_compression=target.enable_compression,
+ fetch_size=target.fetch_size,
+ max_retry=target.max_retry,
+ max_pool_size=target.max_pool_size,
+ tree_wait_timeout_in_ms=target.tree_wait_timeout_in_ms,
+ table_wait_timeout_in_ms=target.table_wait_timeout_in_ms,
+ target_fingerprint=target.fingerprint(),
+ policy={str(key): str(value) for key, value in target.policy.items()},
+ last_known_good_credential=dict(target.last_known_good_credential),
+ verified_at=target.verified_at,
+ target_registry=target_registry,
+ session_manager=session_manager,
+ )
+
+ def to_target(self) -> IoTDBTarget:
+ return IoTDBTarget(
+ target_id=self.target_id,
+ display_name=self.display_name,
+ kind=self.target_kind,
+ host=self.host,
+ port=self.port,
+ node_urls=tuple(self.node_urls) or (f"{self.host}:{self.port}",),
+ user=self.user,
+ password=self.password,
+ database=self.database,
+ sql_dialect=self.sql_dialect,
+ timezone=self.timezone,
+ export_path=self.export_path,
+ iotdb_home=self.iotdb_home,
+ use_ssl=self.use_ssl,
+ ca_certs=self.ca_certs,
+ connection_timeout_in_ms=self.connection_timeout_in_ms,
+ enable_redirection=self.enable_redirection,
+ enable_compression=self.enable_compression,
+ fetch_size=self.fetch_size,
+ max_retry=self.max_retry,
+ max_pool_size=self.max_pool_size,
+ tree_wait_timeout_in_ms=self.tree_wait_timeout_in_ms,
+ table_wait_timeout_in_ms=self.table_wait_timeout_in_ms,
+ policy=dict(self.policy),
+ last_known_good_credential=dict(self.last_known_good_credential),
+ verified_at=self.verified_at,
+ )
+
+ def safe_dict(self) -> dict[str, object]:
+ target = self.to_target()
+ safe = target.as_dict(include_secret=False)
+ safe["password"] = "***" if self.password else ""
+ return safe
+
@staticmethod
def from_env_arguments() -> "Config":
"""
@@ -69,62 +217,161 @@ class Config:
"""
parser = argparse.ArgumentParser(description="IoTDB MCP Server")
+ parser.add_argument("--target-id", type=str, default=None, help="Named IoTDB target id")
parser.add_argument(
- "--host",
+ "--targets-file",
type=str,
- help="IoTDB host",
- default=os.getenv("IOTDB_HOST", "127.0.0.1"),
+ default=None,
+ help="JSON IoTDB target registry file",
)
+ parser.add_argument("--host", type=str, help="IoTDB host", default=None)
parser.add_argument(
"--port",
type=int,
help="IoTDB MySQL protocol port",
- default=os.getenv("IOTDB_PORT", 6667),
+ default=None,
)
parser.add_argument(
"--user",
type=str,
help="IoTDB username",
- default=os.getenv("IOTDB_USER", "root"),
+ default=None,
)
parser.add_argument(
"--password",
type=str,
help="IoTDB password",
- default=os.getenv("IOTDB_PASSWORD", "root"),
+ default=None,
)
parser.add_argument(
"--database",
type=str,
- help="IoTDB connect database name",
- default=os.getenv("IOTDB_DATABASE", "test"),
+ help="IoTDB database/session scope. Table dialect uses database name; tree dialect still requires explicit root paths in SQL.",
+ default=None,
)
parser.add_argument(
"--sql-dialect",
type=str,
help="SQL dialect: tree or table",
- default=os.getenv("IOTDB_SQL_DIALECT", "table"),
+ default=None,
+ )
+
+ parser.add_argument(
+ "--timezone",
+ type=str,
+ help="IoTDB session timezone",
+ default=None,
)
parser.add_argument(
"--export-path",
type=str,
help="Path for exporting query results",
- default=os.getenv("IOTDB_EXPORT_PATH", "/tmp"),
+ default=None,
)
+ parser.add_argument("--node-urls", type=str, default=None, help="Comma-separated IoTDB node URLs")
+ parser.add_argument("--use-ssl", choices=("true", "false"), default=None)
+ parser.add_argument("--ca-certs", type=str, default=None)
+ parser.add_argument("--connection-timeout-ms", type=int, default=None)
+ parser.add_argument("--enable-redirection", choices=("true", "false"), default=None)
+ parser.add_argument("--enable-compression", choices=("true", "false"), default=None)
+ parser.add_argument("--fetch-size", type=int, default=None)
+ parser.add_argument("--max-retry", type=int, default=None)
+ parser.add_argument("--max-pool-size", type=int, default=None)
+ parser.add_argument("--wait-timeout-ms", type=int, default=None)
args = parser.parse_args()
- return Config(
- host=args.host,
- port=args.port,
- user=args.user,
- password=args.password,
- database=args.database,
- sql_dialect=args.sql_dialect,
- export_path=args.export_path,
+ targets_file = args.targets_file or os.getenv("TIMESEEK_IOTDB_TARGETS_FILE", "")
+ targets_json = os.getenv("TIMESEEK_IOTDB_TARGETS_JSON", "")
+ registry = IoTDBTargetRegistry.from_env(
+ targets_file=targets_file or None,
+ targets_json=targets_json or None,
+ )
+ if registry.list_targets(include_secret=False):
+ target = registry.resolve(
+ args.target_id or os.getenv("TIMESEEK_IOTDB_TARGET_ID")
+ )
+ else:
+ wait_timeout = (
+ args.wait_timeout_ms
+ if args.wait_timeout_ms is not None
+ else os.getenv("IOTDB_WAIT_TIMEOUT_MS")
+ )
+ target = target_from_mapping(
+ {
+ "target_id": "unverified-template",
+ "display_name": "Unverified IoTDB template",
+ "host": _argument_or_env(args.host, "IOTDB_HOST"),
+ "port": _argument_or_env(args.port, "IOTDB_PORT"),
+ "user": _argument_or_env(args.user, "IOTDB_USER"),
+ "password": _argument_or_env(args.password, "IOTDB_PASSWORD"),
+ "database": _argument_or_env(args.database, "IOTDB_DATABASE"),
+ "sql_dialect": args.sql_dialect
+ or os.getenv("IOTDB_SQL_DIALECT")
+ or "table",
+ "timezone": _argument_or_env(args.timezone, "IOTDB_TIMEZONE"),
+ "export_path": (
+ args.export_path
+ if args.export_path is not None
+ else default_export_path()
+ ),
+ "node_urls": _argument_or_env(args.node_urls, "IOTDB_NODE_URLS"),
+ "iotdb_home": os.getenv("TIMESEEK_IOTDB_HOME"),
+ "use_ssl": _argument_or_env(args.use_ssl, "IOTDB_USE_SSL"),
+ "ca_certs": _argument_or_env(args.ca_certs, "IOTDB_CA_CERTS"),
+ "connection_timeout_in_ms": _argument_or_env(
+ args.connection_timeout_ms,
+ "IOTDB_CONNECTION_TIMEOUT_MS",
+ ),
+ "enable_redirection": _argument_or_env(
+ args.enable_redirection,
+ "IOTDB_ENABLE_REDIRECTION",
+ ),
+ "enable_compression": _argument_or_env(
+ args.enable_compression,
+ "IOTDB_ENABLE_COMPRESSION",
+ ),
+ "fetch_size": _argument_or_env(args.fetch_size, "IOTDB_FETCH_SIZE"),
+ "max_retry": _argument_or_env(args.max_retry, "IOTDB_MAX_RETRY"),
+ "max_pool_size": _argument_or_env(
+ args.max_pool_size,
+ "IOTDB_MAX_POOL_SIZE",
+ ),
+ "tree_wait_timeout_in_ms": (
+ args.wait_timeout_ms
+ if args.wait_timeout_ms is not None
+ else os.getenv("IOTDB_TREE_WAIT_TIMEOUT_MS")
+ or wait_timeout
+ ),
+ "table_wait_timeout_in_ms": (
+ args.wait_timeout_ms
+ if args.wait_timeout_ms is not None
+ else os.getenv("IOTDB_TABLE_WAIT_TIMEOUT_MS")
+ or wait_timeout
+ ),
+ }
+ )
+ registry = registry.with_target(
+ target,
+ default_target_id=target.target_id,
+ )
+ from iotdb_mcp_server.session_manager import IoTDBSessionManager
+
+ config = Config.from_target(
+ target,
+ target_registry=registry,
+ session_manager=IoTDBSessionManager(registry),
)
+ config.targets_file = targets_file
+ config.targets_json = targets_json
+ if targets_file:
+ try:
+ config.targets_file_mtime_ns = Path(targets_file).expanduser().stat().st_mtime_ns
+ except OSError:
+ config.targets_file_mtime_ns = None
+ return config
diff --git a/src/iotdb_mcp_server/result_store.py b/src/iotdb_mcp_server/result_store.py
new file mode 100644
index 0000000..782318a
--- /dev/null
+++ b/src/iotdb_mcp_server/result_store.py
@@ -0,0 +1,845 @@
+#
+# 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.
+#
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import timedelta
+from datetime import datetime, timezone
+import json
+import os
+from pathlib import Path
+import re
+import shutil
+from typing import Any, Iterable
+import uuid
+
+_RESULT_ID_PATTERN = re.compile(r"^res_[0-9]{8}T[0-9]{6}Z_[a-f0-9]{16}$")
+_OWNER_PATTERN = re.compile(r"^[A-Za-z0-9_.:-]{1,180}$")
+_MANIFEST_NAME = "manifest.json"
+_STANDALONE_OWNER = "standalone"
+
+
+def _env_int(name: str, default: int, minimum: int = 1) -> int:
+ raw = os.getenv(name, "").strip()
+ if not raw:
+ return default
+ try:
+ return max(minimum, int(raw))
+ except ValueError:
+ return default
+
+
+def _bounded_int(value: Any, default: int, minimum: int = 1, maximum: int | None = None) -> int:
+ if value is None or str(value).strip() == "":
+ resolved = default
+ else:
+ resolved = int(value)
+ resolved = max(minimum, resolved)
+ if maximum is not None:
+ resolved = min(maximum, resolved)
+ return resolved
+
+
+def _safe_owner_session_id(value: str | None) -> str:
+ raw = (
+ value
+ or os.getenv("TIMESEEK_MCP_SESSION_ID", "").strip()
+ or os.getenv("TIMESEEK_CODEX_SESSION_ID", "").strip()
+ or _STANDALONE_OWNER
+ )
+ cleaned = re.sub(r"[^A-Za-z0-9_.:-]+", "_", str(raw)).strip("_")
+ if not cleaned:
+ cleaned = _STANDALONE_OWNER
+ cleaned = cleaned[:180]
+ if not _OWNER_PATTERN.fullmatch(cleaned):
+ raise ValueError("Invalid owner_session_id.")
+ return cleaned
+
+
+@dataclass(frozen=True)
+class ResultStoreSettings:
+ preview_rows: int = 50
+ page_size_rows: int = 500
+ max_page_rows: int = 5000
+ shard_rows: int = 10000
+ shard_bytes: int = 8 * 1024 * 1024
+ ttl_seconds: int = 24 * 60 * 60
+ max_session_results: int = 100
+ max_session_bytes: int = 512 * 1024 * 1024
+ max_global_results: int = 1000
+ max_global_bytes: int = 2 * 1024 * 1024 * 1024
+
+ @classmethod
+ def from_env(cls) -> "ResultStoreSettings":
+ return cls(
+ preview_rows=_env_int("TIMESEEK_RESULT_PREVIEW_ROWS", 50, minimum=0),
+ page_size_rows=_env_int("TIMESEEK_RESULT_PAGE_SIZE_ROWS", 500),
+ max_page_rows=_env_int("TIMESEEK_RESULT_MAX_PAGE_ROWS", 5000),
+ shard_rows=_env_int("TIMESEEK_RESULT_SHARD_ROWS", 10000),
+ shard_bytes=_env_int("TIMESEEK_RESULT_SHARD_BYTES", 8 * 1024 * 1024),
+ ttl_seconds=_env_int(
+ "TIMESEEK_RESULT_TTL_SECONDS", 24 * 60 * 60, minimum=0
+ ),
+ max_session_results=_env_int(
+ "TIMESEEK_RESULT_MAX_SESSION_RESULTS", 100, minimum=0
+ ),
+ max_session_bytes=_env_int(
+ "TIMESEEK_RESULT_MAX_SESSION_BYTES", 512 * 1024 * 1024, minimum=0
+ ),
+ max_global_results=_env_int(
+ "TIMESEEK_RESULT_MAX_GLOBAL_RESULTS", 1000, minimum=0
+ ),
+ max_global_bytes=_env_int(
+ "TIMESEEK_RESULT_MAX_GLOBAL_BYTES",
+ 2 * 1024 * 1024 * 1024,
+ minimum=0,
+ ),
+ )
+
+
+@dataclass(frozen=True)
+class StoredCsvResult:
+ result_id: str
+ owner_session_id: str
+ columns: list[str]
+ preview_rows: list[str]
+ row_count: int
+ byte_count: int
+ page_size_rows: int
+ manifest_path: str
+ shard_count: int
+
+
+@dataclass(frozen=True)
+class ResultStoreEntry:
+ result_id: str
+ owner_session_id: str
+ tool: str
+ row_count: int
+ byte_count: int
+ storage_bytes: int
+ shard_count: int
+ created_at: datetime
+ last_accessed_at: datetime
+ manifest_path: Path
+ result_dir: Path
+ legacy: bool = False
+
+ def to_payload(self) -> dict[str, Any]:
+ return {
+ "result_id": self.result_id,
+ "owner_session_id": self.owner_session_id,
+ "tool": self.tool,
+ "row_count": self.row_count,
+ "byte_count": self.byte_count,
+ "storage_bytes": self.storage_bytes,
+ "shard_count": self.shard_count,
+ "created_at": self.created_at.isoformat(),
+ "last_accessed_at": self.last_accessed_at.isoformat(),
+ "manifest_path": str(self.manifest_path),
+ "legacy": self.legacy,
+ }
+
+
+class FileResultStore:
+ """File-backed store for large MCP result sets.
+
+ Rows are stored as JSON lines so callers can fetch a page without parsing a
+ whole CSV file. The public row representation intentionally remains the
+ existing CSV row string used by the IoTDB MCP response schema.
+ """
+
+ def __init__(
+ self,
+ root_dir: str | os.PathLike[str],
+ settings: ResultStoreSettings | None = None,
+ ) -> None:
+ self.root_dir = Path(root_dir).expanduser().resolve()
+ self.settings = settings or ResultStoreSettings.from_env()
+
+ def _ensure_root(self) -> None:
+ self.root_dir.mkdir(parents=True, exist_ok=True)
+
+ def _new_result_id(self) -> str:
+ timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+ return f"res_{timestamp}_{uuid.uuid4().hex[:16]}"
+
+ def _session_dir(self, owner_session_id: str | None = None) -> Path:
+ owner = _safe_owner_session_id(owner_session_id)
+ session_dir = (self.root_dir / "sessions" / owner).resolve()
+ sessions_root = (self.root_dir / "sessions").resolve()
+ if not str(session_dir).startswith(str(sessions_root) + os.sep):
+ raise ValueError("Invalid owner_session_id path.")
+ return session_dir
+
+ def _legacy_result_dir(self, result_id: str) -> Path:
+ if not _RESULT_ID_PATTERN.fullmatch(result_id):
+ raise ValueError("Invalid result_id.")
+ result_dir = (self.root_dir / result_id).resolve()
+ if not str(result_dir).startswith(str(self.root_dir) + os.sep):
+ raise ValueError("Invalid result_id path.")
+ return result_dir
+
+ def _result_dir(self, result_id: str, owner_session_id: str | None = None) -> Path:
+ if not _RESULT_ID_PATTERN.fullmatch(result_id):
+ raise ValueError("Invalid result_id.")
+ session_dir = self._session_dir(owner_session_id)
+ result_dir = (session_dir / result_id).resolve()
+ if not str(result_dir).startswith(str(session_dir) + os.sep):
+ raise ValueError("Invalid result_id path.")
+ return result_dir
+
+ def _existing_result_dir(
+ self, result_id: str, owner_session_id: str | None = None
+ ) -> Path:
+ result_dir = self._result_dir(result_id, owner_session_id)
+ if (result_dir / _MANIFEST_NAME).is_file():
+ return result_dir
+ legacy_dir = self._legacy_result_dir(result_id)
+ if (legacy_dir / _MANIFEST_NAME).is_file():
+ return legacy_dir
+ raise FileNotFoundError(f"Result '{result_id}' was not found.")
+
+ def write_csv_result(
+ self,
+ *,
+ tool: str,
+ columns: list[str],
+ rows: Iterable[str],
+ source: dict[str, Any] | None = None,
+ owner_session_id: str | None = None,
+ preview_rows: int | None = None,
+ page_size_rows: int | None = None,
+ ) -> StoredCsvResult:
+ self._ensure_root()
+ owner = _safe_owner_session_id(owner_session_id)
+ result_id = self._new_result_id()
+ result_dir = self._result_dir(result_id, owner)
+ result_dir.mkdir(parents=True, exist_ok=False)
+
+ preview_limit = (
+ self.settings.preview_rows if preview_rows is None else max(0, preview_rows)
+ )
+ page_size = (
+ self.settings.page_size_rows
+ if page_size_rows is None
+ else max(1, page_size_rows)
+ )
+
+ preview: list[str] = []
+ shards: list[dict[str, Any]] = []
+ row_count = 0
+ byte_count = 0
+ shard_index = -1
+ shard_row_count = 0
+ shard_byte_count = 0
+ shard_start = 0
+ shard_file = None
+ shard_path: Path | None = None
+
+ def open_shard() -> None:
+ nonlocal shard_index, shard_row_count, shard_byte_count
+ nonlocal shard_start, shard_file, shard_path
+ shard_index += 1
+ shard_row_count = 0
+ shard_byte_count = 0
+ shard_start = row_count
+ shard_path = result_dir / f"rows-{shard_index:05d}.jsonl"
+ shard_file = shard_path.open("w", encoding="utf-8")
+
+ def close_shard() -> None:
+ nonlocal shard_file
+ if shard_file is None or shard_path is None:
+ return
+ shard_file.close()
+ shards.append(
+ {
+ "index": shard_index,
+ "file": shard_path.name,
+ "row_start": shard_start,
+ "row_count": shard_row_count,
+ "bytes": shard_byte_count,
+ }
+ )
+ shard_file = None
+
+ try:
+ open_shard()
+ for row in rows:
+ row_text = str(row)
+ encoded = json.dumps(row_text, ensure_ascii=False) + "\n"
+ encoded_bytes = len(encoded.encode("utf-8"))
+ if shard_row_count > 0 and (
+ shard_row_count >= self.settings.shard_rows
+ or shard_byte_count + encoded_bytes > self.settings.shard_bytes
+ ):
+ close_shard()
+ open_shard()
+
+ assert shard_file is not None
+ shard_file.write(encoded)
+ if len(preview) < preview_limit:
+ preview.append(row_text)
+ row_count += 1
+ byte_count += encoded_bytes
+ shard_row_count += 1
+ shard_byte_count += encoded_bytes
+ close_shard()
+ except Exception:
+ if shard_file is not None:
+ try:
+ shard_file.close()
+ except OSError:
+ pass
+ shutil.rmtree(result_dir, ignore_errors=True)
+ raise
+
+ manifest = {
+ "schema_version": "iotdb_result_store_v1",
+ "result_id": result_id,
+ "owner_session_id": owner,
+ "created_at": datetime.now(timezone.utc).isoformat(),
+ "last_accessed_at": datetime.now(timezone.utc).isoformat(),
+ "tool": tool,
+ "format": "csv",
+ "columns": columns,
+ "row_count": row_count,
+ "byte_count": byte_count,
+ "preview_rows": preview,
+ "preview_row_count": len(preview),
+ "page_size_rows": page_size,
+ "source": source or {},
+ "shards": shards,
+ }
+ manifest_path = result_dir / _MANIFEST_NAME
+ try:
+ manifest_path.write_text(
+ json.dumps(manifest, ensure_ascii=False, indent=2),
+ encoding="utf-8",
+ )
+ except Exception:
+ shutil.rmtree(result_dir, ignore_errors=True)
+ raise
+ self.cleanup(owner_session_id=owner)
+ return StoredCsvResult(
+ result_id=result_id,
+ owner_session_id=owner,
+ columns=columns,
+ preview_rows=preview,
+ row_count=row_count,
+ byte_count=byte_count,
+ page_size_rows=page_size,
+ manifest_path=str(manifest_path),
+ shard_count=len(shards),
+ )
+
+ def _parse_time(self, value: Any, fallback: datetime | None = None) -> datetime:
+ if isinstance(value, str) and value.strip():
+ try:
+ parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
+ if parsed.tzinfo is None:
+ return parsed.replace(tzinfo=timezone.utc)
+ return parsed.astimezone(timezone.utc)
+ except ValueError:
+ pass
+ return fallback or datetime.fromtimestamp(0, timezone.utc)
+
+ def _dir_size(self, result_dir: Path) -> int:
+ total = 0
+ try:
+ for path in result_dir.rglob("*"):
+ if path.is_file():
+ total += path.stat().st_size
+ except OSError:
+ return total
+ return total
+
+ def _entry_from_manifest(
+ self,
+ manifest_path: Path,
+ *,
+ legacy: bool = False,
+ ) -> ResultStoreEntry | None:
+ try:
+ manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+ except (OSError, json.JSONDecodeError):
+ return None
+ result_id = str(manifest.get("result_id") or manifest_path.parent.name)
+ if not _RESULT_ID_PATTERN.fullmatch(result_id):
+ return None
+ if legacy:
+ owner = str(manifest.get("owner_session_id") or _STANDALONE_OWNER)
+ else:
+ owner = _safe_owner_session_id(
+ manifest.get("owner_session_id") or manifest_path.parent.parent.name
+ )
+ created_at = self._parse_time(manifest.get("created_at"))
+ last_accessed_at = self._parse_time(
+ manifest.get("last_accessed_at"), fallback=created_at
+ )
+ return ResultStoreEntry(
+ result_id=result_id,
+ owner_session_id=owner,
+ tool=str(manifest.get("tool") or ""),
+ row_count=int(manifest.get("row_count") or 0),
+ byte_count=int(manifest.get("byte_count") or 0),
+ storage_bytes=self._dir_size(manifest_path.parent),
+ shard_count=len(manifest.get("shards") or []),
+ created_at=created_at,
+ last_accessed_at=last_accessed_at,
+ manifest_path=manifest_path,
+ result_dir=manifest_path.parent,
+ legacy=legacy,
+ )
+
+ def _iter_entries(
+ self, owner_session_id: str | None = None
+ ) -> list[ResultStoreEntry]:
+ entries: list[ResultStoreEntry] = []
+ sessions_root = self.root_dir / "sessions"
+ owners: list[str]
+ if owner_session_id is None:
+ owners = []
+ if sessions_root.is_dir():
+ owners = [
+ path.name
+ for path in sessions_root.iterdir()
+ if path.is_dir() and _OWNER_PATTERN.fullmatch(path.name)
+ ]
+ else:
+ owners = [_safe_owner_session_id(owner_session_id)]
+
+ for owner in owners:
+ owner_dir = sessions_root / owner
+ if not owner_dir.is_dir():
+ continue
+ for manifest_path in owner_dir.glob(f"res_*/{_MANIFEST_NAME}"):
+ entry = self._entry_from_manifest(manifest_path)
+ if entry is not None:
+ entries.append(entry)
+
+ if owner_session_id in (None, _STANDALONE_OWNER):
+ for manifest_path in self.root_dir.glob(f"res_*/{_MANIFEST_NAME}"):
+ entry = self._entry_from_manifest(manifest_path, legacy=True)
+ if entry is not None:
+ entries.append(entry)
+
+ return entries
+
+ def list_results(
+ self,
+ owner_session_id: str | None = None,
+ *,
+ all_sessions: bool = False,
+ limit: int = 50,
+ ) -> dict[str, Any]:
+ owner = None if all_sessions else _safe_owner_session_id(owner_session_id)
+ entries = self._iter_entries(owner)
+ entries.sort(key=lambda entry: entry.last_accessed_at, reverse=True)
+ bounded_limit = max(1, min(int(limit), 500))
+ selected = entries[:bounded_limit]
+ return {
+ "owner_session_id": owner,
+ "all_sessions": all_sessions,
+ "result_count": len(entries),
+ "returned_results": len(selected),
+ "storage_bytes": sum(entry.storage_bytes for entry in entries),
+ "results": [entry.to_payload() for entry in selected],
+ }
+
+ def _delete_entry(
+ self,
+ entry: ResultStoreEntry,
+ *,
+ reason: str,
+ deleted_dirs: set[Path],
+ ) -> dict[str, Any] | None:
+ if entry.result_dir in deleted_dirs:
+ return None
+ try:
+ shutil.rmtree(entry.result_dir)
+ except FileNotFoundError:
+ pass
+ deleted_dirs.add(entry.result_dir)
+ return {
+ "result_id": entry.result_id,
+ "owner_session_id": entry.owner_session_id,
+ "storage_bytes": entry.storage_bytes,
+ "reason": reason,
+ }
+
+ def _enforce_lru_quota(
+ self,
+ entries: list[ResultStoreEntry],
+ *,
+ max_results: int,
+ max_bytes: int,
+ reason: str,
+ deleted_dirs: set[Path],
+ ) -> list[dict[str, Any]]:
+ active = [
+ entry
+ for entry in entries
+ if entry.result_dir not in deleted_dirs and entry.manifest_path.is_file()
+ ]
+ active.sort(key=lambda entry: entry.last_accessed_at)
+ total_bytes = sum(entry.storage_bytes for entry in active)
+ deleted: list[dict[str, Any]] = []
+
+ def over_quota() -> bool:
+ return (max_results > 0 and len(active) > max_results) or (
+ max_bytes > 0 and total_bytes > max_bytes
+ )
+
+ while len(active) > 1 and over_quota():
+ entry = active.pop(0)
+ total_bytes -= entry.storage_bytes
+ info = self._delete_entry(entry, reason=reason, deleted_dirs=deleted_dirs)
+ if info is not None:
+ deleted.append(info)
+ return deleted
+
+ def delete_result(
+ self,
+ result_id: str,
+ owner_session_id: str | None = None,
+ *,
+ all_sessions: bool = False,
+ ) -> dict[str, Any]:
+ owner = None if all_sessions else _safe_owner_session_id(owner_session_id)
+ entries = self._iter_entries(owner)
+ for entry in entries:
+ if entry.result_id != result_id:
+ continue
+ deleted_dirs: set[Path] = set()
+ info = self._delete_entry(
+ entry, reason="explicit_delete", deleted_dirs=deleted_dirs
+ )
+ return {
+ "deleted": info is not None,
+ "result": info or entry.to_payload(),
+ }
+ raise FileNotFoundError(f"Result '{result_id}' was not found.")
+
+ def cleanup(
+ self,
+ owner_session_id: str | None = None,
+ *,
+ all_sessions: bool = False,
+ ttl_seconds: int | None = None,
+ max_session_results: int | None = None,
+ max_session_bytes: int | None = None,
+ max_global_results: int | None = None,
+ max_global_bytes: int | None = None,
+ ) -> dict[str, Any]:
+ self._ensure_root()
+ owner = None if all_sessions else _safe_owner_session_id(owner_session_id)
+ deleted_dirs: set[Path] = set()
+ deleted: list[dict[str, Any]] = []
+ now = datetime.now(timezone.utc)
+ effective_ttl = (
+ self.settings.ttl_seconds
+ if ttl_seconds is None
+ else max(0, int(ttl_seconds))
+ )
+
+ ttl_candidates = self._iter_entries(owner)
+ if effective_ttl > 0:
+ cutoff = now - timedelta(seconds=effective_ttl)
+ for entry in ttl_candidates:
+ if entry.last_accessed_at >= cutoff:
+ continue
+ info = self._delete_entry(
+ entry, reason="ttl_expired", deleted_dirs=deleted_dirs
+ )
+ if info is not None:
+ deleted.append(info)
+
+ session_result_limit = (
+ self.settings.max_session_results
+ if max_session_results is None
+ else max(0, int(max_session_results))
+ )
+ session_byte_limit = (
+ self.settings.max_session_bytes
+ if max_session_bytes is None
+ else max(0, int(max_session_bytes))
+ )
+ owners = (
+ sorted({entry.owner_session_id for entry in self._iter_entries(None)})
+ if all_sessions
+ else [owner]
+ )
+ for current_owner in owners:
+ if current_owner is None:
+ continue
+ deleted.extend(
+ self._enforce_lru_quota(
+ self._iter_entries(current_owner),
+ max_results=session_result_limit,
+ max_bytes=session_byte_limit,
+ reason="session_lru_quota",
+ deleted_dirs=deleted_dirs,
+ )
+ )
+
+ global_result_limit = (
+ self.settings.max_global_results
+ if max_global_results is None
+ else max(0, int(max_global_results))
+ )
+ global_byte_limit = (
+ self.settings.max_global_bytes
+ if max_global_bytes is None
+ else max(0, int(max_global_bytes))
+ )
+ deleted.extend(
+ self._enforce_lru_quota(
+ self._iter_entries(None),
+ max_results=global_result_limit,
+ max_bytes=global_byte_limit,
+ reason="global_lru_quota",
+ deleted_dirs=deleted_dirs,
+ )
+ )
+
+ remaining = self._iter_entries(None)
+ return {
+ "owner_session_id": owner,
+ "all_sessions": all_sessions,
+ "ttl_seconds": effective_ttl,
+ "limits": {
+ "max_session_results": session_result_limit,
+ "max_session_bytes": session_byte_limit,
+ "max_global_results": global_result_limit,
+ "max_global_bytes": global_byte_limit,
+ },
+ "deleted_results": len(deleted),
+ "deleted_storage_bytes": sum(item["storage_bytes"] for item in deleted),
+ "remaining_results": len(remaining),
+ "remaining_storage_bytes": sum(entry.storage_bytes for entry in remaining),
+ "deleted": deleted,
+ }
+
+ def read_manifest(
+ self, result_id: str, owner_session_id: str | None = None
+ ) -> dict[str, Any]:
+ manifest_path = (
+ self._existing_result_dir(result_id, owner_session_id) / _MANIFEST_NAME
+ )
+ try:
+ return json.loads(manifest_path.read_text(encoding="utf-8"))
+ except FileNotFoundError as exc:
+ raise FileNotFoundError(f"Result '{result_id}' was not found.") from exc
+
+ def _touch_manifest(self, result_dir: Path, manifest: dict[str, Any]) -> None:
+ manifest["last_accessed_at"] = datetime.now(timezone.utc).isoformat()
+ (result_dir / _MANIFEST_NAME).write_text(
+ json.dumps(manifest, ensure_ascii=False, indent=2),
+ encoding="utf-8",
+ )
+
+ def read_page(
+ self,
+ result_id: str,
+ *,
+ offset: int = 0,
+ limit: int | None = None,
+ owner_session_id: str | None = None,
+ ) -> dict[str, Any]:
+ result_dir = self._existing_result_dir(result_id, owner_session_id)
+ manifest = json.loads((result_dir / _MANIFEST_NAME).read_text(encoding="utf-8"))
+ row_count = int(manifest.get("row_count", 0))
+ page_size = int(manifest.get("page_size_rows") or self.settings.page_size_rows)
+ requested_limit = page_size if limit is None else max(1, int(limit))
+ effective_limit = min(requested_limit, self.settings.max_page_rows)
+ effective_offset = min(max(0, int(offset)), row_count)
+ rows: list[str] = []
+
+ for shard in manifest.get("shards", []):
+ shard_start = int(shard["row_start"])
+ shard_count = int(shard["row_count"])
+ shard_end = shard_start + shard_count
+ if shard_end <= effective_offset:
+ continue
+ if shard_start >= effective_offset + effective_limit:
+ break
+
+ skip = max(0, effective_offset - shard_start)
+ shard_path = result_dir / str(shard["file"])
+ with shard_path.open("r", encoding="utf-8") as handle:
+ for index, line in enumerate(handle):
+ if index < skip:
+ continue
+ if len(rows) >= effective_limit:
+ break
+ rows.append(str(json.loads(line)))
+ if len(rows) >= effective_limit:
+ break
+
+ next_offset = effective_offset + len(rows)
+ has_more = next_offset < row_count
+ self._touch_manifest(result_dir, manifest)
+ return {
+ "format": "csv",
+ "result_id": result_id,
+ "owner_session_id": manifest.get("owner_session_id", _STANDALONE_OWNER),
+ "columns": list(manifest.get("columns", [])),
+ "rows": rows,
+ "row_count": row_count,
+ "offset": effective_offset,
+ "limit": effective_limit,
+ "returned_rows": len(rows),
+ "next_cursor": str(next_offset) if has_more else None,
+ "has_more": has_more,
+ "manifest_path": str(result_dir / _MANIFEST_NAME),
+ }
+
+ def _page_request_offset(self, request: dict[str, Any]) -> int:
+ if request.get("offset") is not None:
+ return max(0, int(request["offset"]))
+ cursor = request.get("cursor")
+ if cursor is None or str(cursor).strip() == "":
+ return 0
+ cleaned = str(cursor).strip()
+ if not cleaned.isdigit():
+ raise ValueError("cursor must be a non-negative integer row offset.")
+ return int(cleaned)
+
+ def read_pages(
+ self,
+ requests: list[dict[str, Any]],
+ *,
+ owner_session_id: str | None = None,
+ default_limit: int | None = None,
+ max_pages: int | None = None,
+ max_total_rows: int | None = None,
+ continue_on_error: bool = True,
+ ) -> dict[str, Any]:
+ if not isinstance(requests, list) or not requests:
+ raise ValueError("pages must be a non-empty list.")
+ effective_max_pages = _bounded_int(
+ max_pages if max_pages is not None else os.getenv("TIMESEEK_RESULT_MAX_BATCH_PAGES"),
+ 32,
+ minimum=1,
+ maximum=256,
+ )
+ if len(requests) > effective_max_pages:
+ raise ValueError(
+ f"read_result_pages accepts at most {effective_max_pages} page requests."
+ )
+ effective_max_total_rows = _bounded_int(
+ max_total_rows
+ if max_total_rows is not None
+ else os.getenv("TIMESEEK_RESULT_MAX_BATCH_READ_ROWS"),
+ 10000,
+ minimum=1,
+ maximum=1_000_000,
+ )
+
+ results: list[dict[str, Any]] = []
+ total_rows = 0
+ failed = 0
+ skipped = 0
+ stop = False
+
+ for index, request in enumerate(requests):
+ if stop:
+ skipped += 1
+ results.append(
+ {
+ "index": index,
+ "ok": False,
+ "status": "skipped",
+ "error_type": "BatchAborted",
+ "error": "Skipped because continue_on_error is false and an earlier page failed.",
+ }
+ )
+ continue
+
+ try:
+ if not isinstance(request, dict):
+ raise ValueError("Each page request must be an object.")
+ result_id = str(request.get("result_id") or "").strip()
+ if not result_id:
+ raise ValueError("Each page request requires result_id.")
+ page_owner = request.get("owner_session_id")
+ if page_owner is not None:
+ page_owner = str(page_owner)
+ else:
+ page_owner = owner_session_id
+ page_limit = (
+ request.get("limit")
+ if request.get("limit") is not None
+ else default_limit
+ )
+ page = self.read_page(
+ result_id,
+ offset=self._page_request_offset(request),
+ limit=None if page_limit is None else int(page_limit),
+ owner_session_id=page_owner,
+ )
+ next_total_rows = total_rows + int(page.get("returned_rows") or 0)
+ if next_total_rows > effective_max_total_rows:
+ raise ValueError(
+ "read_result_pages total row quota exceeded: "
+ f"{next_total_rows}>{effective_max_total_rows}"
+ )
+ total_rows = next_total_rows
+ page["index"] = index
+ page["ok"] = True
+ page["status"] = "success"
+ results.append(page)
+ except Exception as exc:
+ failed += 1
+ if not continue_on_error:
+ stop = True
+ results.append(
+ {
+ "index": index,
+ "ok": False,
+ "status": "error",
+ "result_id": request.get("result_id") if isinstance(request, dict) else None,
+ "error_type": type(exc).__name__,
+ "error": str(exc),
+ }
+ )
+
+ succeeded = sum(1 for item in results if item.get("status") == "success")
+ return {
+ "format": "csv",
+ "page_request_count": len(requests),
+ "returned_pages": succeeded,
+ "succeeded": succeeded,
+ "failed": failed,
+ "skipped": skipped,
+ "continue_on_error": continue_on_error,
+ "limits": {
+ "max_pages": effective_max_pages,
+ "max_total_rows": effective_max_total_rows,
+ "default_limit": default_limit,
+ },
+ "returned_rows": total_rows,
+ "pages": results,
+ }
+
+
+def result_store_from_export_path(export_path: str) -> FileResultStore:
+ return FileResultStore(Path(export_path).expanduser() / "result_store")
diff --git a/src/iotdb_mcp_server/runtime_policy.py b/src/iotdb_mcp_server/runtime_policy.py
new file mode 100644
index 0000000..4b3f418
--- /dev/null
+++ b/src/iotdb_mcp_server/runtime_policy.py
@@ -0,0 +1,353 @@
+# 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.
+#
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+import json
+import os
+from pathlib import Path
+import re
+from threading import RLock
+from typing import Any
+
+
+_DEFAULT_SERVER_NAME = "iotdb"
+_SHELL_DEFAULT_RE = re.compile(r"^\$\{([A-Za-z_][A-Za-z0-9_]*):-([^}]*)\}$")
+_SENSITIVE_KEY_RE = re.compile(r"(PASSWORD|SECRET|TOKEN|API_KEY|ACCESS_KEY)", re.IGNORECASE)
+_WITHHELD_ENV_KEYS = frozenset({"TIMESEEK_IOTDB_TARGETS_JSON"})
+_SESSION_POLICY_LOCK = RLock()
+_SESSION_POLICY: dict[str, str] = {}
+
+_ENFORCEMENT_POLICY_KEYS = frozenset(
+ {
+ "TIMESEEK_MCP_PERMISSION_ENFORCEMENT",
+ "IOTDB_STRICT_PERMISSION_ENFORCEMENT",
+ }
+)
+
+_POLICY_KEYS = frozenset(
+ {
+ "IOTDB_ENABLE_METADATA_QUERY",
+ "IOTDB_METADATA_ALLOWED_USERS",
+ "IOTDB_ENABLE_DATABASE_DDL",
+ "IOTDB_DATABASE_DDL_ALLOWED_USERS",
+ "IOTDB_REQUIRE_DROP_CONFIRM",
+ "IOTDB_ENABLE_TABLE_DDL",
+ "IOTDB_TABLE_DDL_ALLOWED_USERS",
+ "IOTDB_REQUIRE_TABLE_DROP_CONFIRM",
+ "IOTDB_ENABLE_TIMESERIES_DDL",
+ "IOTDB_TIMESERIES_DDL_ALLOWED_USERS",
+ "IOTDB_REQUIRE_TIMESERIES_DROP_CONFIRM",
+ "IOTDB_ENABLE_TTL_SQL",
+ "IOTDB_TTL_ALLOWED_USERS",
+ "IOTDB_REQUIRE_TTL_UNSET_CONFIRM",
+ "IOTDB_ENABLE_WRITE_DML",
+ "IOTDB_WRITE_ALLOWED_USERS",
+ "IOTDB_REQUIRE_DELETE_CONFIRM",
+ "IOTDB_ENABLE_SQL_DRIVER",
+ "IOTDB_SQL_DRIVER_ALLOWED_USERS",
+ "IOTDB_SQL_DRIVER_MODE",
+ "IOTDB_SQL_DRIVER_REQUIRE_DESTRUCTIVE_CONFIRM",
+ "IOTDB_SQL_DRIVER_EXTRA_READONLY_PREFIXES",
+ "IOTDB_SQL_DRIVER_EXTRA_DDL_PREFIXES",
+ "IOTDB_SQL_DRIVER_EXTRA_FULL_PREFIXES",
+ "IOTDB_ENABLE_MODEL_MANAGEMENT",
+ "IOTDB_MODEL_ALLOWED_USERS",
+ "IOTDB_REQUIRE_MODEL_DESTRUCTIVE_CONFIRM",
+ "TIMESEEK_MCP_PERMISSION_ENFORCEMENT",
+ "IOTDB_STRICT_PERMISSION_ENFORCEMENT",
+ }
+)
+
+_FULL_PERMISSION_DEFAULTS = {
+ "IOTDB_ENABLE_METADATA_QUERY": "true",
+ "IOTDB_METADATA_ALLOWED_USERS": "*",
+ "IOTDB_ENABLE_DATABASE_DDL": "true",
+ "IOTDB_DATABASE_DDL_ALLOWED_USERS": "*",
+ "IOTDB_REQUIRE_DROP_CONFIRM": "true",
+ "IOTDB_ENABLE_TABLE_DDL": "true",
+ "IOTDB_TABLE_DDL_ALLOWED_USERS": "*",
+ "IOTDB_REQUIRE_TABLE_DROP_CONFIRM": "true",
+ "IOTDB_ENABLE_TIMESERIES_DDL": "true",
+ "IOTDB_TIMESERIES_DDL_ALLOWED_USERS": "*",
+ "IOTDB_REQUIRE_TIMESERIES_DROP_CONFIRM": "true",
+ "IOTDB_ENABLE_TTL_SQL": "true",
+ "IOTDB_TTL_ALLOWED_USERS": "*",
+ "IOTDB_REQUIRE_TTL_UNSET_CONFIRM": "true",
+ "IOTDB_ENABLE_WRITE_DML": "true",
+ "IOTDB_WRITE_ALLOWED_USERS": "*",
+ "IOTDB_REQUIRE_DELETE_CONFIRM": "true",
+ "IOTDB_ENABLE_SQL_DRIVER": "true",
+ "IOTDB_SQL_DRIVER_ALLOWED_USERS": "*",
+ "IOTDB_SQL_DRIVER_MODE": "full",
+ "IOTDB_SQL_DRIVER_REQUIRE_DESTRUCTIVE_CONFIRM": "true",
+ "IOTDB_ENABLE_MODEL_MANAGEMENT": "true",
+ "IOTDB_MODEL_ALLOWED_USERS": "*",
+ "IOTDB_REQUIRE_MODEL_DESTRUCTIVE_CONFIRM": "true",
+}
+
+_ENFORCEMENT_DEFAULTS = {
+ "TIMESEEK_MCP_PERMISSION_ENFORCEMENT": "advisory",
+ "IOTDB_STRICT_PERMISSION_ENFORCEMENT": "false",
+}
+
+_PRESET_POLICIES = {
+ "full": dict(_FULL_PERMISSION_DEFAULTS),
+ "ddl": {
+ **_FULL_PERMISSION_DEFAULTS,
+ "IOTDB_SQL_DRIVER_MODE": "ddl",
+ "IOTDB_ENABLE_WRITE_DML": "false",
+ },
+ "readonly": {
+ **_FULL_PERMISSION_DEFAULTS,
+ "IOTDB_ENABLE_DATABASE_DDL": "false",
+ "IOTDB_ENABLE_TABLE_DDL": "false",
+ "IOTDB_ENABLE_TIMESERIES_DDL": "false",
+ "IOTDB_ENABLE_TTL_SQL": "false",
+ "IOTDB_ENABLE_WRITE_DML": "false",
+ "IOTDB_SQL_DRIVER_MODE": "readonly",
+ },
+}
+
+
+def _candidate_mcp_config_paths() -> list[Path]:
+ candidates: list[Path] = []
+ explicit = os.getenv("TIMESEEK_MCP_CONFIG_PATH")
+ if explicit:
+ candidates.append(Path(explicit).expanduser())
+
+ timeseek_root = os.getenv("TIMESEEK_ROOT")
+ if timeseek_root:
+ root = Path(timeseek_root).expanduser()
+ candidates.extend(
+ [
+ root / ".mcp.json",
+ root / "timeseek-codex-multi-plugin" / ".mcp.json",
+ root / "timeseek-codex-single-plugin" / ".mcp.json",
+ ]
+ )
+
+ repo_root = os.getenv("REPO_ROOT") or os.getenv("TIMESEEK_REPO_ROOT")
+ if repo_root:
+ root = Path(repo_root).expanduser()
+ candidates.extend(
+ [
+ root / ".mcp.json",
+ root / "timeseek_codex" / "timeseek-codex-multi-plugin" / ".mcp.json",
+ root / "timeseek_codex" / "timeseek-codex-single-plugin" / ".mcp.json",
+ root / "timeseek_cc" / ".mcp.json",
+ ]
+ )
+
+ cwd = Path.cwd()
+ candidates.extend(
+ [
+ cwd / ".mcp.json",
+ cwd / "timeseek_codex" / "timeseek-codex-multi-plugin" / ".mcp.json",
+ cwd / "timeseek_codex" / "timeseek-codex-single-plugin" / ".mcp.json",
+ cwd / "timeseek_cc" / ".mcp.json",
+ ]
+ )
+
+ deduped: list[Path] = []
+ seen: set[str] = set()
+ for candidate in candidates:
+ key = str(candidate)
+ if key not in seen:
+ seen.add(key)
+ deduped.append(candidate)
+ return deduped
+
+
+def _expand_config_value(value: Any) -> str:
+ text = str(value)
+ match = _SHELL_DEFAULT_RE.fullmatch(text)
+ if not match:
+ return text
+ name, default = match.groups()
+ current = os.getenv(name)
+ return current if current not in (None, "") else default
+
+
+def _load_mcp_env() -> tuple[dict[str, str], str]:
+ server_name = os.getenv("TIMESEEK_MCP_SERVER_NAME", _DEFAULT_SERVER_NAME)
+ for path in _candidate_mcp_config_paths():
+ if not path.exists():
+ continue
+ try:
+ data = json.loads(path.read_text(encoding="utf-8"))
+ except Exception:
+ continue
+ server = data.get("mcpServers", {}).get(server_name, {})
+ env = server.get("env", {})
+ if isinstance(env, dict):
+ return {str(key): _expand_config_value(value) for key, value in env.items()}, str(path)
+ return {}, ""
+
+
+def dynamic_getenv(name: str, default: str | None = None) -> str | None:
+ with _SESSION_POLICY_LOCK:
+ value = _SESSION_POLICY.get(name)
+ if value not in (None, ""):
+ return value
+
+ env, _ = _load_mcp_env()
+ value = env.get(name)
+ if value not in (None, ""):
+ return value
+
+ value = os.getenv(name)
+ if value not in (None, ""):
+ return value
+
+ return _FULL_PERMISSION_DEFAULTS.get(
+ name,
+ _ENFORCEMENT_DEFAULTS.get(name, default),
+ )
+
+
+def dynamic_env_bool(name: str, default: bool) -> bool:
+ value = dynamic_getenv(name)
+ if value is None:
+ return default
+ normalized = value.strip().lower()
+ if normalized in {"1", "true", "yes", "on"}:
+ return True
+ if normalized in {"0", "false", "no", "off", ""}:
+ return False
+ return default
+
+
+def permission_enforcement_mode() -> str:
+ if dynamic_env_bool("IOTDB_STRICT_PERMISSION_ENFORCEMENT", False):
+ return "strict"
+ mode = (dynamic_getenv("TIMESEEK_MCP_PERMISSION_ENFORCEMENT", "advisory") or "advisory").strip().lower()
+ if mode not in {"advisory", "strict"}:
+ return "advisory"
+ return mode
+
+
+def strict_permission_enforcement() -> bool:
+ return permission_enforcement_mode() == "strict"
+
+
+def _redact_env(env: dict[str, str]) -> dict[str, str]:
+ return {
+ key: "***"
+ if key in _WITHHELD_ENV_KEYS or _SENSITIVE_KEY_RE.search(key)
+ else value
+ for key, value in env.items()
+ }
+
+
+def _session_policy_copy() -> dict[str, str]:
+ with _SESSION_POLICY_LOCK:
+ return dict(_SESSION_POLICY)
+
+
+def _coerce_policy_value(value: Any) -> str:
+ if isinstance(value, bool):
+ return "true" if value else "false"
+ return str(value)
+
+
+def _validate_policy_keys(policy: Mapping[str, Any]) -> None:
+ unknown = sorted(str(key) for key in policy if str(key) not in _POLICY_KEYS)
+ if unknown:
+ raise ValueError(
+ "Unsupported session policy key(s): "
+ + ", ".join(unknown)
+ + ". Use get_iotdb_session_policy to inspect supported keys."
+ )
+
+
+def session_policy_for_preset(preset: str) -> dict[str, str]:
+ normalized = preset.strip().lower()
+ try:
+ return dict(_PRESET_POLICIES[normalized])
+ except KeyError as exc:
+ raise ValueError("Unsupported policy preset. Expected one of: full, ddl, readonly.") from exc
+
+
+def supported_policy_keys() -> list[str]:
+ return sorted(_POLICY_KEYS)
+
+
+def set_session_policy(
+ policy: Mapping[str, Any] | None = None,
+ *,
+ preset: str | None = None,
+ replace: bool = False,
+) -> dict[str, Any]:
+ updates: dict[str, str] = {}
+ if preset:
+ updates.update(session_policy_for_preset(preset))
+ if policy:
+ _validate_policy_keys(policy)
+ updates.update({str(key): _coerce_policy_value(value) for key, value in policy.items()})
+
+ with _SESSION_POLICY_LOCK:
+ preserved_enforcement = (
+ {
+ key: value
+ for key, value in _SESSION_POLICY.items()
+ if key in _ENFORCEMENT_POLICY_KEYS
+ }
+ if replace and preset
+ else {}
+ )
+ if replace:
+ _SESSION_POLICY.clear()
+ _SESSION_POLICY.update(preserved_enforcement)
+ _SESSION_POLICY.update(updates)
+
+ return dynamic_policy_snapshot()
+
+
+def reset_session_policy(keys: list[str] | None = None) -> dict[str, Any]:
+ if keys is not None:
+ _validate_policy_keys({key: "" for key in keys})
+
+ with _SESSION_POLICY_LOCK:
+ if keys is None:
+ _SESSION_POLICY.clear()
+ else:
+ for key in keys:
+ _SESSION_POLICY.pop(key, None)
+
+ return dynamic_policy_snapshot()
+
+
+def dynamic_policy_snapshot() -> dict[str, Any]:
+ env, path = _load_mcp_env()
+ session_policy = _session_policy_copy()
+ effective: dict[str, str] = {}
+ for key in supported_policy_keys():
+ value = dynamic_getenv(key, "")
+ if value not in (None, ""):
+ effective[key] = value
+ return {
+ "mcp_config_path": path,
+ "session_policy": _redact_env(session_policy),
+ "mcp_env": _redact_env(env),
+ "defaults": {**_FULL_PERMISSION_DEFAULTS, **_ENFORCEMENT_DEFAULTS},
+ "effective_policy": _redact_env(effective),
+ "supported_keys": supported_policy_keys(),
+ "presets": sorted(_PRESET_POLICIES),
+ }
diff --git a/src/iotdb_mcp_server/server.py b/src/iotdb_mcp_server/server.py
index 489687e..b0935e9 100644
--- a/src/iotdb_mcp_server/server.py
+++ b/src/iotdb_mcp_server/server.py
@@ -15,566 +15,59 @@
# specific language governing permissions and limitations
# under the License.
#
-# Security patch for CVE-2026-XXXXX (path traversal vulnerability leading to RCE)
-# Author: Mohammed Tanveer (threatpointer)
-# Date: 2026-01-12
-#
import logging
-import datetime
-import asyncio
-import os
-import uuid
-import re
-import pandas as pd
-from typing import Dict, Any, List, Union
-from iotdb.Session import Session
-from iotdb.SessionPool import SessionPool, PoolConfig
-from iotdb.utils.SessionDataSet import SessionDataSet
-from iotdb.table_session import TableSession
-from iotdb.table_session_pool import TableSessionPool, TableSessionPoolConfig
from fastmcp import FastMCP
-from mcp.types import TextContent
from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.services.database import register_database_tools
+from iotdb_mcp_server.services.explain import register_explain_tools
+from iotdb_mcp_server.services.metadata import register_metadata_tools
+from iotdb_mcp_server.services.model import register_model_tools
+from iotdb_mcp_server.services.query import register_query_tools
+from iotdb_mcp_server.services.results import register_result_store_tools
+from iotdb_mcp_server.services.runtime_policy import register_runtime_policy_tools
+from iotdb_mcp_server.services.sql_driver import register_sql_driver_tools
+from iotdb_mcp_server.services.table import register_table_tools
+from iotdb_mcp_server.services.targets import register_target_tools
+from iotdb_mcp_server.services.timeseries import register_timeseries_tools
+from iotdb_mcp_server.services.ttl import register_ttl_tools
+from iotdb_mcp_server.services.udf import register_udf_tools
+from iotdb_mcp_server.services.write import register_write_tools
-# Initialize FastMCP server
mcp = FastMCP("iotdb_mcp_server")
-# Configure logging
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
-
logger = logging.getLogger("iotdb_mcp_server")
config = Config.from_env_arguments()
-db_config = {
- "host": config.host,
- "port": config.port,
- "user": config.user,
- "password": config.password,
- "database": config.database,
- "sql_dialect": config.sql_dialect,
- "export_path": config.export_path,
-}
-
-max_pool_size = 100 # Increased from 100 for better concurrency
-
-# Never print credentials: log a copy of the config with the password masked.
-logged_config = {**db_config, "password": "***"}
-logger.info(f"IoTDB Config: {logged_config}")
-
-# Ensure export directory exists
-if not os.path.exists(config.export_path):
- try:
- os.makedirs(config.export_path)
- logger.info(f"Created export directory: {config.export_path}")
- except Exception as e:
- logger.warning(
- f"Failed to create export directory {config.export_path}: {str(e)}"
- )
-
-
-def sanitize_filename(filename: str, base_dir: str) -> str:
- """
- Sanitize and validate filename to prevent path traversal attacks.
-
- Security patch for CVE-2026-XXXXX
- Author: Mohammed Tanveer (threatpointer)
- Date: 2026-01-12
-
- Args:
- filename: The user-provided filename
- base_dir: The base directory for exports (must be absolute path)
-
- Returns:
- The sanitized absolute filepath
-
- Raises:
- ValueError: If the filename contains invalid characters or attempts path traversal
-
- Security measures:
- - Rejects any path separators or traversal sequences before processing
- - Validates allowed characters (alphanumeric, underscore, hyphen, dot)
- - Resolves absolute path and verifies it stays within base_dir boundary
- - Prevents directory traversal, symlink attacks, and path manipulation
- """
- if not filename:
- raise ValueError("Filename cannot be empty")
-
- # First, reject any path separators or traversal patterns before processing
- # This catches attacks before os.path.basename can process them
- if "/" in filename or "\\" in filename or ".." in filename:
- raise ValueError(
- "Invalid filename: path separators and directory traversal sequences are not allowed"
- )
-
- # Remove any directory components - only keep the base filename
- # (This is now a safety check since we already blocked separators)
- filename = os.path.basename(filename)
-
- # Validate characters - only allow alphanumeric, underscore, hyphen, and dot
- if not re.match(r"^[a-zA-Z0-9_\-\.]+$", filename):
- raise ValueError(
- "Invalid filename: only alphanumeric characters, underscore, hyphen, and dot are allowed"
- )
-
- # Prevent filenames that are just dots or empty after sanitization
- if not filename or filename in (".", ".."):
- raise ValueError("Invalid filename")
-
- # Prevent filenames that start with a dot followed by a dot (like ..something)
- if filename.startswith(".."):
- raise ValueError("Invalid filename: cannot start with '..'")
-
- # Construct the full path
- filepath = os.path.join(base_dir, filename)
-
- # Resolve to absolute path (resolves symlinks and normalizes path)
- filepath_real = os.path.realpath(filepath)
- basedir_real = os.path.realpath(base_dir)
-
- # Ensure the resolved path is within the base directory boundary
- # This prevents path traversal even with symlinks or complex path manipulation
- if (
- not filepath_real.startswith(basedir_real + os.sep)
- and filepath_real != basedir_real
- ):
- raise ValueError(
- f"Path traversal detected: file must be within export directory"
- )
-
- return filepath_real
-
-
-if config.sql_dialect == "tree":
- # Configure connection pool with optimized settings
- pool_config = PoolConfig(
- node_urls=[str(config.host) + ":" + str(config.port)],
- user_name=config.user,
- password=config.password,
- fetch_size=1024, # Fetch size for queries
- time_zone="UTC+8", # Consistent timezone
- max_retry=3, # Connection retry attempts
- )
- # Optimize pool size based on expected concurrent queries
- wait_timeout_in_ms = 5000 # Increased from 3000 for better reliability
- session_pool = SessionPool(pool_config, max_pool_size, wait_timeout_in_ms)
-
- @mcp.tool()
- async def metadata_query(query_sql: str) -> list[TextContent]:
- """Execute metadata queries on IoTDB to explore database structure and statistics.
-
- Args:
- query_sql: The metadata query to execute. Supported queries:
- - SHOW DATABASES [path]: List all databases or databases under a specific path
- - SHOW TIMESERIES [path]: List all time series or time series under a specific path
- - SHOW CHILD PATHS [path]: List child paths under a specific path
- - SHOW CHILD NODES [path]: List child nodes under a specific path
- - SHOW DEVICES [path]: List all devices or devices under a specific path
- - COUNT TIMESERIES [path]: Count time series under a specific path
- - COUNT NODES [path]: Count nodes under a specific path
- - COUNT DEVICES [path]: Count devices under a specific path
- - if path is not provided, the query will be applied to root.**
-
- Examples:
- SHOW DATABASES root.**
- SHOW TIMESERIES root.ln.**
- SHOW CHILD PATHS root.ln
- SHOW CHILD PATHS root.ln.*.*
- SHOW CHILD NODES root.ln
- SHOW DEVICES root.ln.**
- COUNT TIMESERIES root.ln.**
- COUNT NODES root.ln
- COUNT DEVICES root.ln
- """
- session = None
- try:
- session = session_pool.get_session()
- stmt = query_sql.strip().upper()
-
- # Process SHOW DATABASES
- if (
- stmt.startswith("SHOW DATABASES")
- or stmt.startswith("SHOW TIMESERIES")
- or stmt.startswith("SHOW CHILD PATHS")
- or stmt.startswith("SHOW CHILD NODES")
- or stmt.startswith("SHOW DEVICES")
- or stmt.startswith("COUNT TIMESERIES")
- or stmt.startswith("COUNT NODES")
- or stmt.startswith("COUNT DEVICES")
- ):
- res = session.execute_query_statement(query_sql)
- return prepare_res(res, session)
- else:
- session.close()
- raise ValueError(
- "Unsupported metadata query. Please use one of the supported query types."
- )
- except Exception as e:
- if session:
- session.close()
- logger.error(f"Failed to execute metadata query: {str(e)}")
- raise
-
- @mcp.tool()
- async def select_query(query_sql: str) -> list[TextContent]:
- """Execute a SELECT query on the IoTDB tree SQL dialect.
-
- Args:
- query_sql: The SQL query to execute (using TREE dialect, time using ISO 8601 format, e.g. 2017-11-01T00:08:00.000).
-
- SQL Syntax:
- SELECT [LAST] selectExpr [, selectExpr] ...
- [INTO intoItem [, intoItem] ...]
- FROM prefixPath [, prefixPath] ...
- [WHERE whereCondition]
- [GROUP BY {
- ([startTime, endTime), interval [, slidingStep]) |
- LEVEL = levelNum [, levelNum] ... |
- TAGS(tagKey [, tagKey] ... |
- VARIATION(expression[,delta][,ignoreNull=true/false]) |
- CONDITION(expression,[keep>/>=/=/</<=]threshold[,ignoreNull=true/false]) |
- SESSION(timeInterval) |
- COUNT(expression, size[,ignoreNull=true/false])
- }]
- [HAVING havingCondition]
- [ORDER BY sortKey {ASC | DESC}]
- [FILL ({PREVIOUS | LINEAR | constant}) (, interval=DURATION_LITERAL)?)]
- [SLIMIT seriesLimit] [SOFFSET seriesOffset]
- [LIMIT rowLimit] [OFFSET rowOffset]
- [ALIGN BY {TIME | DEVICE}]
-
- Examples:
- select temperature from root.ln.wf01.wt01 where time < 2017-11-01T00:08:00.000
- select status, temperature from root.ln.wf01.wt01 where (time > 2017-11-01T00:05:00.000 and time < 2017-11-01T00:12:00.000) or (time >= 2017-11-01T16:35:00.000 and time <= 2017-11-01T16:37:00.000)
- select * from root.ln.** where time > 1 order by time desc limit 10;
-
- Supported Aggregate Functions:
- SUM
- COUNT
- MAX_VALUE
- MIN_VALUE
- AVG
- VARIANCE
- MAX_TIME
- MIN_TIME
- ...
- """
- session = None
- try:
- session = session_pool.get_session()
- stmt = query_sql.strip().upper()
-
- # Regular SELECT queries
- if stmt.startswith("SELECT"):
- res = session.execute_query_statement(query_sql)
- return prepare_res(res, session)
- else:
- session.close()
- raise ValueError("Only SELECT queries are allowed for select_query")
- except Exception as e:
- if session:
- session.close()
- logger.error(f"Failed to execute select query: {str(e)}")
- raise
-
- @mcp.tool()
- async def export_query(
- query_sql: str, format: str = "csv", filename: str = None
- ) -> list[TextContent]:
- """Execute a query and export the results to a CSV or Excel file.
-
- Args:
- query_sql: The SQL query to execute (using TREE dialect, time using ISO 8601 format, e.g. 2017-11-01T00:08:00.000)
- format: Export format, either "csv" or "excel" (default: "csv")
- filename: Optional filename for the exported file. If not provided, a unique filename will be generated.
-
- SQL Syntax:
- SELECT ⟨select_list⟩
- FROM ⟨tables⟩
- [WHERE ⟨condition⟩]
- [GROUP BY ⟨groups⟩]
- [HAVING ⟨group_filter⟩]
- [FILL ⟨fill_methods⟩]
- [ORDER BY ⟨order_expression⟩]
- [OFFSET ⟨n⟩]
- [LIMIT ⟨n⟩];
-
- Returns:
- Information about the exported file and a preview of the data (max 10 rows)
- """
- session = None
- try:
- session = session_pool.get_session()
- stmt = query_sql.strip().upper()
-
- if stmt.startswith("SELECT") or stmt.startswith("SHOW"):
- # Execute the query
- res = session.execute_query_statement(query_sql)
-
- # Create a pandas DataFrame
- df = res.todf()
- # Close the session
- session.close()
-
- # Generate unique filename with timestamp
- timestamp = int(datetime.datetime.now().timestamp())
- if filename is None:
- # Generate a unique filename if not provided
- filename = f"dump_{uuid.uuid4().hex[:4]}_{timestamp}"
-
- if format.lower() == "csv":
- if filename.lower().endswith(".csv"):
- filename = filename[:-4]
- # Sanitize filename to prevent path traversal attacks
- filepath = sanitize_filename(f"{filename}.csv", config.export_path)
- df.to_csv(filepath, index=False)
- elif format.lower() == "excel":
- if filename.lower().endswith(".xlsx"):
- filename = filename[:-5]
- # Sanitize filename to prevent path traversal attacks
- filepath = sanitize_filename(f"{filename}.xlsx", config.export_path)
- df.to_excel(filepath, index=False)
- else:
- raise ValueError("Format must be either 'csv' or 'excel'")
-
- # Generate preview (first 10 rows)
- preview_rows = min(10, len(df))
- preview_data = []
- preview_data.append(",".join(df.columns)) # Header
-
- for i in range(preview_rows):
- preview_data.append(",".join(map(str, df.iloc[i])))
-
- # Return information
- return [
- TextContent(
- type="text",
- text=f"Query results exported to {filepath}\n\nPreview (first {preview_rows} rows):\n"
- + "\n".join(preview_data),
- )
- ]
- else:
- raise ValueError("Only SELECT or SHOW queries are allowed for export")
- except Exception as e:
- if session:
- session.close()
- logger.error(f"Failed to export query: {str(e)}")
- raise
-
- def prepare_res(_res: SessionDataSet, _session: Session) -> list[TextContent]:
- columns = _res.get_column_names()
- result = []
- while _res.has_next():
- record = _res.next()
- if columns[0] == "Time":
- timestamp = record.get_timestamp()
- # formatted_time = datetime.datetime.fromtimestamp(timestamp/1000).strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3]
- row = record.get_fields()
- result.append(str(timestamp) + "," + ",".join(map(str, row)))
- else:
- row = record.get_fields()
- result.append(",".join(map(str, row)))
- _session.close()
- return [
- TextContent(
- type="text",
- text="\n".join([",".join(columns)] + result),
- )
- ]
-
-elif config.sql_dialect == "table":
- session_pool_config = TableSessionPoolConfig(
- node_urls=[str(config.host) + ":" + str(config.port)],
- username=config.user,
- password=config.password,
- max_pool_size=max_pool_size, # Increased from 5 for better concurrency
- database=None if len(config.database) == 0 else config.database,
- )
- session_pool = TableSessionPool(session_pool_config)
-
- @mcp.tool()
- async def read_query(query_sql: str) -> list[TextContent]:
- """Execute a SELECT query on the IoTDB. Please use table sql_dialect when generating SQL queries.
-
- Args:
- query_sql: The SQL query to execute (using TABLE dialect, time using ISO 8601 format, e.g. 2017-11-01T00:08:00.000)
- """
- table_session = None
- try:
- table_session = session_pool.get_session()
- stmt = query_sql.strip().upper()
-
- # Regular SELECT queries
- if (
- stmt.startswith("SELECT")
- or stmt.startswith("DESCRIBE")
- or stmt.startswith("SHOW")
- ):
- res = table_session.execute_query_statement(query_sql)
- return prepare_res(res, table_session)
- else:
- table_session.close()
- raise ValueError("Only SELECT queries are allowed for read_query")
- except Exception as e:
- if table_session:
- table_session.close()
- logger.error(f"Failed to execute query: {str(e)}")
- raise
-
- @mcp.tool()
- async def list_tables() -> list[TextContent]:
- """List all tables in the IoTDB database."""
- table_session = None
- try:
- table_session = session_pool.get_session()
- res = table_session.execute_query_statement("SHOW TABLES")
-
- result = ["Tables_in_" + db_config["database"]] # Header
- while res.has_next():
- result.append(str(res.next().get_fields()[0]))
- table_session.close()
- return [TextContent(type="text", text="\n".join(result))]
- except Exception as e:
- if table_session:
- table_session.close()
- logger.error(f"Failed to list tables: {str(e)}")
- raise
-
- @mcp.tool()
- async def describe_table(table_name: str) -> list[TextContent]:
- """Get the schema information for a specific table
- Args:
- table_name: name of the table to describe
- """
- table_session = None
- try:
- table_session = session_pool.get_session()
- res = table_session.execute_query_statement(
- "DESC " + table_name + " details"
- )
- return prepare_res(res, table_session)
- except Exception as e:
- if table_session:
- table_session.close()
- logger.error(f"Failed to describe table {table_name}: {str(e)}")
- raise
-
- @mcp.tool()
- async def export_table_query(
- query_sql: str, format: str = "csv", filename: str = None
- ) -> list[TextContent]:
- """Execute a query and export the results to a CSV or Excel file.
-
- Args:
- query_sql: The SQL query to execute (using TABLE dialect, time using ISO 8601 format, e.g. 2017-11-01T00:08:00.000)
- format: Export format, either "csv" or "excel" (default: "csv")
- filename: Optional filename for the exported file. If not provided, a unique filename will be generated.
-
- SQL Syntax:
- SELECT ⟨select_list⟩
- FROM ⟨tables⟩
- [WHERE ⟨condition⟩]
- [GROUP BY ⟨groups⟩]
- [HAVING ⟨group_filter⟩]
- [FILL ⟨fill_methods⟩]
- [ORDER BY ⟨order_expression⟩]
- [OFFSET ⟨n⟩]
- [LIMIT ⟨n⟩];
-
- Returns:
- Information about the exported file and a preview of the data (max 10 rows)
- """
- table_session = None
- try:
- table_session = session_pool.get_session()
- stmt = query_sql.strip().upper()
-
- if (
- stmt.startswith("SELECT")
- or stmt.startswith("SHOW")
- or stmt.startswith("DESCRIBE")
- or stmt.startswith("DESC")
- ):
- # Execute the query
- res = table_session.execute_query_statement(query_sql)
-
- # Create a pandas DataFrame
- df = res.todf()
-
- # Close the session
- table_session.close()
-
- # Generate unique filename with timestamp
- timestamp = int(datetime.datetime.now().timestamp())
- if filename is None:
- filename = f"dump_{uuid.uuid4().hex[:4]}_{timestamp}"
-
- if format.lower() == "csv":
- if filename.lower().endswith(".csv"):
- filename = filename[:-4]
- # Sanitize filename to prevent path traversal attacks
- filepath = sanitize_filename(f"{filename}.csv", config.export_path)
- df.to_csv(filepath, index=False)
- elif format.lower() == "excel":
- if filename.lower().endswith(".xlsx"):
- filename = filename[:-5]
- # Sanitize filename to prevent path traversal attacks
- filepath = sanitize_filename(f"{filename}.xlsx", config.export_path)
- df.to_excel(filepath, index=False)
- else:
- raise ValueError("Format must be either 'csv' or 'excel'")
-
- # Generate preview (first 10 rows)
- preview_rows = min(10, len(df))
- preview_data = []
- preview_data.append(",".join(df.columns)) # Header
-
- for i in range(preview_rows):
- preview_data.append(",".join(map(str, df.iloc[i])))
-
- # Return information
- return [
- TextContent(
- type="text",
- text=f"Query results exported to {filepath}\n\nPreview (first {preview_rows} rows):\n"
- + "\n".join(preview_data),
- )
- ]
- else:
- raise ValueError(
- "Only SELECT, SHOW or DESCRIBE queries are allowed for export"
- )
- except Exception as e:
- if table_session:
- table_session.close()
- logger.error(f"Failed to export table query: {str(e)}")
- raise
-
- def prepare_res(
- _res: SessionDataSet, _table_session: TableSession
- ) -> list[TextContent]:
- columns = _res.get_column_names()
- result = []
- while _res.has_next():
- row = _res.next().get_fields()
- result.append(",".join(map(str, row)))
- _table_session.close()
- return [
- TextContent(
- type="text",
- text="\n".join([",".join(columns)] + result),
- )
- ]
-
+logger.info(
+ "IoTDB Config: %s",
+ config.safe_dict(),
+)
-def main():
+register_target_tools(mcp, config, logger)
+register_runtime_policy_tools(mcp, config, logger)
+register_result_store_tools(mcp, config, logger)
+register_query_tools(mcp, config, logger)
+register_sql_driver_tools(mcp, config, logger)
+register_metadata_tools(mcp, config, logger)
+register_explain_tools(mcp, config, logger)
+register_database_tools(mcp, config, logger)
+register_timeseries_tools(mcp, config, logger)
+register_ttl_tools(mcp, config, logger)
+register_table_tools(mcp, config, logger)
+register_udf_tools(mcp, config, logger)
+register_write_tools(mcp, config, logger)
+register_model_tools(mcp, config, logger)
+
+
+def main() -> None:
logger.info("iotdb_mcp_server running with stdio transport")
- # Initialize and run the server
mcp.run(transport="stdio")
diff --git a/src/iotdb_mcp_server/services/README.md b/src/iotdb_mcp_server/services/README.md
new file mode 100644
index 0000000..27d1dad
--- /dev/null
+++ b/src/iotdb_mcp_server/services/README.md
@@ -0,0 +1,394 @@
+# IoTDB MCP Services
+
+This directory contains tool modules registered by `src/iotdb_mcp_server/server.py`.
+
+## Registration Overview
+
+Current registration order:
+
+1. `register_target_tools`
+2. `register_runtime_policy_tools`
+3. `register_result_store_tools`
+4. `register_query_tools`
+5. `register_sql_driver_tools`
+6. `register_metadata_tools`
+7. `register_explain_tools`
+8. `register_database_tools`
+9. `register_timeseries_tools`
+10. `register_ttl_tools`
+11. `register_table_tools`
+12. `register_udf_tools`
+13. `register_write_tools`
+14. `register_model_tools`
+
+## Tool Matrix
+
+### `query.py`
+
+- Tree dialect:
+ - `select_query(query_sql, max_inline_rows=None, page_size_rows=None)`
+ - `export_query(query_sql, format="csv", filename=None)`
+- Table dialect:
+ - `read_query(query_sql, max_inline_rows=None, page_size_rows=None)`
+ - `export_table_query(query_sql, format="csv", filename=None)`
+
+Read-query tools stream rows into ResultStore and return only a preview by
+default. Use the returned `result_id` and `next_cursor` with
+`read_result_page`.
+
+### `results.py`
+
+- `read_result_page(result_id, cursor=None, offset=None, limit=None, owner_session_id=None, target_id=None)`
+- `read_result_pages(pages, owner_session_id=None, default_limit=None, max_pages=None, max_total_rows=None, continue_on_error=True, target_id=None)`
+- `list_result_store(owner_session_id=None, all_sessions=False, limit=50, target_id=None)`
+- `cleanup_result_store(owner_session_id=None, all_sessions=False, ttl_seconds=None, target_id=None)`
+- `delete_result(result_id, owner_session_id=None, all_sessions=False, target_id=None)`
+
+Loads a bounded page from a file-backed ResultStore result. Pass the same
+`target_id` as the original query when multiple IoTDB targets are registered.
+When the TimeSeek Codex plugin hook is active, `owner_session_id` is injected
+from the Codex session. Standalone MCP usage can pass `owner_session_id`
+explicitly or set `TIMESEEK_MCP_SESSION_ID`; otherwise the ResultStore owner is
+`standalone`.
+
+ResultStore cleanup is finite and session-aware. Each write runs lightweight
+cleanup for the current owner and then enforces global limits. Entries expire by
+`last_accessed_at`, and quota pressure evicts least-recently-accessed results.
+If a single new result is larger than the byte quota, it is retained so the
+caller can still page through the result that was just returned.
+
+`read_result_pages` batches multiple page requests in one MCP call. Each item
+accepts `result_id` plus optional `cursor`, `offset`, `limit`, and
+`owner_session_id`. The default per-call limits are 32 pages and 10000 returned
+rows; override them with tool arguments or `TIMESEEK_RESULT_MAX_BATCH_PAGES`
+and `TIMESEEK_RESULT_MAX_BATCH_READ_ROWS`.
+
+### `sql_driver.py`
+
+- Both dialects:
+ - `sql_driver_policy()`
+ - `inspect_sql_permission(sql)`
+ - `sql_executor_batch(sqls=None, sql_template=None, param_sets=None, max_concurrency=4, worker_pool_size=None, per_item_timeout_ms=None, batch_timeout_ms=None, max_result_rows_per_item=None, max_result_bytes_per_item=None, max_batch_result_rows=None, max_batch_result_bytes=None, continue_on_error=True, max_inline_rows=5, page_size_rows=None)`
+ - `sql_execute(sql, confirm_destructive=False, max_inline_rows=None, page_size_rows=None)`
+
+Features:
+
+- Statement whitelist by category: `readonly` / `ddl` / `full`
+- Advisory permission metadata via `IOTDB_SQL_DRIVER_MODE`
+- SQL inspection returns `required_permission`, `destructive`, `risk_level`,
+ `approval_required`, and the confirmation parameter without executing SQL
+- Destructive SQL confirmation policy
+- `sql_executor_batch` executes only readonly single statements. It accepts
+ either explicit `sqls=[...]` or a repeated `sql_template` plus `param_sets`.
+ Template placeholders are `{{name}}` for SQL literals, `{{name:path}}` for
+ IoTDB paths, and `{{name:identifier}}` for SQL identifiers.
+- Batch results are per-statement ResultStore entries. The batch response
+ returns status, `result_id`, row counts, a small preview, and paging metadata;
+ use `read_result_page` to fetch full result pages.
+- Batch execution controls include a bounded thread worker pool, per-item and
+ whole-batch wait timeouts, and per-item / whole-batch result row and byte
+ quotas. Defaults: 4 concurrency, 60000 ms per item, 300000 ms per batch,
+ 10000 rows and 16 MiB per item, 100000 rows and 64 MiB per batch.
+- Tree CQ enabled in whitelist:
+ - Readonly: `SHOW CONTINUOUS QUERIES`, `SHOW CQS`
+ - DDL: `CREATE CONTINUOUS QUERY` / `CREATE CQ`, `DROP CONTINUOUS QUERY` / `DROP CQ`
+- Downsampling execution path:
+ - ad-hoc: `SELECT ... GROUP BY ([start, end), interval[, slidingStep])`
+ - scheduled: CQ with `INTO` + optional `RESAMPLE` clause
+
+### `runtime_policy.py`
+
+- Session scoped permission controls:
+ - `get_iotdb_session_policy()`
+ - `set_iotdb_session_policy(policy=None, preset=None, replace=False)`
+ - `reset_iotdb_session_policy(keys=None)`
+- Presets:
+ - `full`: advisory signal for read, DDL, and DML-capable workflows
+ - `ddl`: advisory signal for read and DDL-capable workflows
+ - `readonly`: advisory signal for read-only workflows
+ - In `strict` mode only, these presets become hard execution gates
+- Effective policy priority:
+ - session overlay set by MCP tool
+ - dynamically read `.mcp.json`
+ - process environment
+ - built-in full-permission defaults
+
+### `metadata.py`
+
+- Tree dialect:
+ - `metadata_query(query_sql)`
+ - `list_timeseries(path="root.**")`
+ - `list_devices(path="root.**")`
+ - `list_child_paths(path)`
+ - `list_child_nodes(path)`
+ - `count_timeseries(path="root.**")`
+ - `count_devices(path="root.**")`
+ - `count_nodes(path="root")`
+- Table dialect:
+ - `metadata_query(query_sql)`
+ - `list_tables()`
+ - `describe_table(table_name, details=True)`
+
+Metadata result sets also use ResultStore, which keeps large `SHOW TIMESERIES`
+or `SHOW DEVICES` responses bounded.
+
+### `explain.py`
+
+- Both dialects:
+ - `explain_query(query_sql, analyze=False)`
+
+### `database.py`
+
+- Tree dialect:
+ - `list_databases(details=False)`
+ - `create_database(database)`
+ - `drop_database(database, confirm=False)`
+- Table dialect:
+ - `list_databases(details=False)`
+ - `create_database(database, if_not_exists=True)`
+ - `drop_database(database, if_exists=True, confirm=False)`
+ - `use_database(database)`
+
+### `timeseries.py` (tree only)
+
+- `create_timeseries_ddl(ddl_sql)`
+- `create_timeseries_batch_ddl(ddl_list, continue_on_error=True, skip_if_exists=True, max_statements=1000)`
+- `alter_timeseries_ddl(ddl_sql)`
+- `drop_timeseries_ddl(ddl_sql, confirm=False)`
+
+### `ttl.py` (tree only)
+
+- `ttl_command(ttl_sql, confirm_unset=False)`
+- `ttl_query(ttl_sql)`
+
+### `table.py` (table only)
+
+- `create_table_ddl(ddl_sql)`
+- `alter_table_ddl(ddl_sql)`
+- `drop_table_ddl(ddl_sql, confirm=False)`
+
+### `udf.py`
+
+- Both dialects:
+ - `list_udf_functions()`
+ - `prepare_udf_query(function_name, expressions, table_name=None, from_path=None, attributes=None, where_clause=None, limit=None, alias=None, align_by_device=False)`
+ - `execute_udf_query(function_name, expressions, table_name=None, from_path=None, attributes=None, where_clause=None, limit=None, alias=None, align_by_device=False, max_inline_rows=None, page_size_rows=None)`
+ - `export_udf_query(function_name, expressions, table_name=None, from_path=None, attributes=None, where_clause=None, limit=None, alias=None, align_by_device=False, format="csv", filename=None)`
+
+Features:
+
+- `SHOW FUNCTIONS` discovery through `list_udf_functions`
+- Conservative SQL generation for UDF calls instead of accepting raw SQL
+- Table dialect shape: `SELECT UDF(column, "k"="v") FROM table ...`
+- Tree dialect shape: `SELECT UDF(measurement, "k"="v") FROM root.sg.d1 ...`
+- Read-only execution only; DDL/DML keywords, semicolons, and comments are rejected in expressions and `where_clause`
+- ResultStore-backed previews for inline execution and export support for large UDF outputs
+
+### `write.py`
+
+- Tree dialect:
+ - `write_query(write_sql, confirm_delete=False)` (INSERT/DELETE)
+- Table dialect:
+ - `write_query(write_sql, confirm_delete=False)` (INSERT/UPDATE/DELETE)
+
+### `model.py`
+
+- Both dialects:
+ - `model_query(model_sql)`
+ - `model_command(model_sql, confirm_destructive=False)`
+- Tree dialect:
+ - `prepare_model_inference_request(model_id, input_sql, output_length=96, generate_time=False)`
+ - `model_inference(inference_sql)`
+
+### `json_response.py`
+
+- Shared response wrapper and parser:
+ - `JsonParser.check_format(obj)`
+ - `JsonParser.parse(text)`
+ - `payload_response(tool, payload, message=None)`
+ - `text_payload_response(tool, text, message=None)`
+ - `csv_payload_response(tool, columns, rows, message=None)`
+ - `sql_success_response(tool, sql)`
+
+## Tool Response JSON Envelope
+
+All tools now return `TextContent.text` as a JSON envelope instead of plain text.
+
+Envelope format:
+
+```json
+{
+ "version": "1.0",
+ "tool": "tool_name",
+ "ok": true,
+ "timestamp": "2026-02-26T00:00:00+00:00",
+ "payload": {},
+ "message": "optional"
+}
+```
+
+Payload conventions:
+
+- Query-like tools: `payload.format="csv"` with `columns`, preview `rows`,
+ preview `text`, `row_count`, `inline_row_count`, `inline_truncated`, and
+ ResultStore fields when persisted
+- Text-like tools: `payload.format="text"` with `text`
+- SQL success tools: `payload.sql` + `payload.result="success"`
+
+For ResultStore-backed responses:
+
+```json
+{
+ "format": "csv",
+ "columns": ["Time", "root.sg.d.s0"],
+ "rows": ["0,1.0"],
+ "text": "Time,root.sg.d.s0\n0,1.0",
+ "row_count": 100000,
+ "inline_row_count": 1,
+ "inline_truncated": true,
+ "result_id": "res_20260704T000000Z_0123456789abcdef",
+ "owner_session_id": "ses_abc123",
+ "next_cursor": "1",
+ "result_store": {
+ "type": "file",
+ "page_tool": "read_result_page",
+ "batch_page_tool": "read_result_pages",
+ "owner_session_id": "ses_abc123",
+ "page_size_rows": 500,
+ "shard_count": 10
+ }
+}
+```
+
+Environment knobs:
+
+- `TIMESEEK_RESULT_PREVIEW_ROWS` default `50`
+- `TIMESEEK_RESULT_PAGE_SIZE_ROWS` default `500`
+- `TIMESEEK_RESULT_MAX_PAGE_ROWS` default `5000`
+- `TIMESEEK_RESULT_SHARD_ROWS` default `10000`
+- `TIMESEEK_RESULT_SHARD_BYTES` default `8388608`
+- `TIMESEEK_RESULT_TTL_SECONDS` default `86400`; set `0` to disable TTL expiry
+- `TIMESEEK_RESULT_MAX_SESSION_RESULTS` default `100`; set `0` to disable
+- `TIMESEEK_RESULT_MAX_SESSION_BYTES` default `536870912`; set `0` to disable
+- `TIMESEEK_RESULT_MAX_GLOBAL_RESULTS` default `1000`; set `0` to disable
+- `TIMESEEK_RESULT_MAX_GLOBAL_BYTES` default `2147483648`; set `0` to disable
+- `TIMESEEK_MCP_SESSION_ID` optional standalone ResultStore owner fallback
+- `IOTDB_SQL_EXECUTOR_BATCH_MAX_STATEMENTS` default `64`
+- `IOTDB_SQL_EXECUTOR_BATCH_MAX_CONCURRENCY` default `16`
+
+Parser check:
+
+- Wrapper output is validated by `JsonParser.check_format(...)`
+- Serialized text is parsed again by `JsonParser.parse(...)` before returning
+- Invalid envelope shape raises `ValueError`
+
+## Security and Permission Model
+
+IoTDB MCP permissions default to host-agent approval flow. The MCP server reports
+required permission and risk; Codex, Claude Code, or OpenCode owns user approval
+through its normal question/approval mechanism. This keeps MCP policy aligned
+with the host agent system instead of turning `.mcp.json` into a hard blocker.
+
+Default mode is `advisory`. In advisory mode:
+
+- `IOTDB_SQL_DRIVER_MODE` is a prompt-layer signal, not an execution blocker.
+- `IOTDB_ENABLE_*` and `*_ALLOWED_USERS` are prompt-layer/session settings, not
+ hard gates.
+- Destructive operations still require explicit tool confirmation flags, such
+ as `confirm_destructive=true`, after the host agent has obtained approval.
+
+Strict hard gating is available only when explicitly enabled:
+
+```bash
+TIMESEEK_MCP_PERMISSION_ENFORCEMENT=strict
+# or
+IOTDB_STRICT_PERMISSION_ENFORCEMENT=true
+```
+
+Policy can be changed at runtime without restarting the MCP server:
+
+```json
+{
+ "preset": "readonly",
+ "replace": true
+}
+```
+
+or:
+
+```json
+{
+ "policy": {
+ "IOTDB_SQL_DRIVER_MODE": "full",
+ "IOTDB_ENABLE_WRITE_DML": "true"
+ }
+}
+```
+
+### Per-module advisory keys
+
+- `database.py`
+ - `IOTDB_ENABLE_DATABASE_DDL`
+ - `IOTDB_DATABASE_DDL_ALLOWED_USERS`
+ - `IOTDB_REQUIRE_DROP_CONFIRM`
+- `timeseries.py`
+ - `IOTDB_ENABLE_TIMESERIES_DDL`
+ - `IOTDB_TIMESERIES_DDL_ALLOWED_USERS`
+ - `IOTDB_REQUIRE_TIMESERIES_DROP_CONFIRM`
+- `table.py`
+ - `IOTDB_ENABLE_TABLE_DDL`
+ - `IOTDB_TABLE_DDL_ALLOWED_USERS`
+ - `IOTDB_REQUIRE_TABLE_DROP_CONFIRM`
+- `ttl.py`
+ - `IOTDB_ENABLE_TTL_SQL`
+ - `IOTDB_TTL_ALLOWED_USERS`
+ - `IOTDB_REQUIRE_TTL_UNSET_CONFIRM`
+- `write.py`
+ - `IOTDB_ENABLE_WRITE_DML`
+ - `IOTDB_WRITE_ALLOWED_USERS`
+ - `IOTDB_REQUIRE_DELETE_CONFIRM`
+- `model.py`
+ - `IOTDB_ENABLE_MODEL_MANAGEMENT`
+ - `IOTDB_MODEL_ALLOWED_USERS`
+ - `IOTDB_REQUIRE_MODEL_DESTRUCTIVE_CONFIRM`
+ - `model_inference` is registered only for tree dialect and accepts
+ `CALL INFERENCE(...)` SQL.
+- `metadata.py`
+ - `IOTDB_ENABLE_METADATA_QUERY`
+ - `IOTDB_METADATA_ALLOWED_USERS`
+- `sql_driver.py`
+ - `IOTDB_ENABLE_SQL_DRIVER`
+ - `IOTDB_SQL_DRIVER_ALLOWED_USERS`
+ - `IOTDB_SQL_DRIVER_MODE` (`readonly` / `ddl` / `full`)
+ - `IOTDB_SQL_DRIVER_REQUIRE_DESTRUCTIVE_CONFIRM`
+ - `IOTDB_SQL_DRIVER_EXTRA_READONLY_PREFIXES`
+ - `IOTDB_SQL_DRIVER_EXTRA_DDL_PREFIXES`
+ - `IOTDB_SQL_DRIVER_EXTRA_FULL_PREFIXES`
+
+## Minimal Enablement Example
+
+```bash
+# Dialect
+IOTDB_SQL_DIALECT=tree
+
+# Keep generic driver available in full mode
+IOTDB_ENABLE_SQL_DRIVER=true
+IOTDB_SQL_DRIVER_MODE=full
+IOTDB_SQL_DRIVER_ALLOWED_USERS=*
+
+# Enable specialized tools
+IOTDB_ENABLE_METADATA_QUERY=true
+IOTDB_ENABLE_DATABASE_DDL=true
+IOTDB_ENABLE_TIMESERIES_DDL=true
+IOTDB_ENABLE_TTL_SQL=true
+IOTDB_ENABLE_WRITE_DML=true
+IOTDB_ENABLE_MODEL_MANAGEMENT=true
+```
+
+## Notes
+
+- All tools enforce single-statement execution (`;` separated multi-statement is rejected).
+- Tool availability depends on `sql_dialect` (`tree` or `table`).
+- For production hard-block behavior, set `TIMESEEK_MCP_PERMISSION_ENFORCEMENT=strict`.
+- Keep destructive confirm flags enabled unless an external approval layer
+ provides equivalent per-action confirmation.
diff --git a/src/iotdb_mcp_server/services/database.py b/src/iotdb_mcp_server/services/database.py
new file mode 100644
index 0000000..5578c1e
--- /dev/null
+++ b/src/iotdb_mcp_server/services/database.py
@@ -0,0 +1,348 @@
+#
+# 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.
+#
+
+import logging
+import re
+
+from iotdb.Session import Session
+from iotdb.table_session import TableSession
+from iotdb.utils.SessionDataSet import SessionDataSet
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.runtime_policy import dynamic_env_bool
+from iotdb_mcp_server.runtime_policy import dynamic_getenv
+from iotdb_mcp_server.runtime_policy import strict_permission_enforcement
+from iotdb_mcp_server.services.json_response import (
+ csv_payload_response,
+ sql_success_response,
+)
+from iotdb_mcp_server.services.target_selection import (
+ apply_target_registry,
+ iotdb_target_response_context,
+ select_target_config,
+ table_session_pool,
+ tree_session_pool,
+)
+from iotdb_mcp_server.target_registry import merge_target
+
+_TABLE_DB_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+_TREE_DB_PATTERN = re.compile(r"^root(?:\.[A-Za-z_][A-Za-z0-9_]*)+$")
+
+
+def _env_bool(name: str, default: bool) -> bool:
+ return dynamic_env_bool(name, default)
+
+
+def _csv_set(value: str) -> set[str]:
+ return {item.strip() for item in value.split(",") if item.strip()}
+
+
+def _assert_database_ddl_permission(
+ config: Config, action: str, confirm: bool = False
+) -> None:
+ if strict_permission_enforcement():
+ if not _env_bool("IOTDB_ENABLE_DATABASE_DDL", False):
+ raise PermissionError(
+ "Database DDL tool is disabled by strict MCP policy. "
+ "Set IOTDB_ENABLE_DATABASE_DDL=true to enable."
+ )
+
+ allowed_users = _csv_set(dynamic_getenv("IOTDB_DATABASE_DDL_ALLOWED_USERS", "root") or "root")
+ if "*" not in allowed_users and config.user not in allowed_users:
+ raise PermissionError(
+ f"Current MCP user '{config.user}' is not allowed by IOTDB_DATABASE_DDL_ALLOWED_USERS."
+ )
+
+ if action == "DROP" and _env_bool("IOTDB_REQUIRE_DROP_CONFIRM", True) and not confirm:
+ raise PermissionError(
+ "DROP DATABASE requires confirm=True by server policy."
+ )
+
+
+def _validate_database_name(sql_dialect: str, database: str) -> str:
+ raw = database.strip()
+ if not raw:
+ raise ValueError("Database name/path cannot be empty.")
+
+ if sql_dialect == "tree":
+ if not _TREE_DB_PATTERN.fullmatch(raw):
+ raise ValueError(
+ "Invalid tree database path. Expected pattern like root.sg or root.sg1.dev."
+ )
+ elif sql_dialect == "table":
+ if not _TABLE_DB_PATTERN.fullmatch(raw):
+ raise ValueError(
+ "Invalid table database name. Use letters, numbers, and underscore only."
+ )
+ else:
+ raise ValueError(
+ f"Unsupported sql_dialect '{sql_dialect}'. Expected 'tree' or 'table'."
+ )
+ return raw
+
+
+def _format_result(
+ res: SessionDataSet,
+ session_or_table_session: Session | TableSession,
+ tool_name: str,
+) -> list[TextContent]:
+ columns = res.get_column_names()
+ rows: list[str] = []
+ while res.has_next():
+ row = res.next().get_fields()
+ rows.append(",".join(map(str, row)))
+ session_or_table_session.close()
+ return csv_payload_response(tool_name, columns, rows)
+
+
+def register_database_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register database management tools with per-call IoTDB target selection."""
+ max_pool_size = 100
+
+ @mcp.tool()
+ async def list_databases(
+ details: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """List databases in the selected IoTDB target."""
+ selected_config = select_target_config(config, target_id=target_id, target=target)
+ with iotdb_target_response_context(selected_config):
+ sql = "SHOW DATABASES DETAILS" if details else "SHOW DATABASES"
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="list_databases",
+ )
+ session = None
+ try:
+ session = session_pool.get_session()
+ res = session.execute_query_statement(sql)
+ return _format_result(res, session, "list_databases")
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to list databases: {str(e)}")
+ raise
+
+ _, session_pool = table_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ tool_name="list_databases",
+ )
+ table_session = None
+ try:
+ table_session = session_pool.get_session()
+ res = table_session.execute_query_statement(sql)
+ return _format_result(res, table_session, "list_databases")
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to list databases: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def create_database(
+ database: str,
+ if_not_exists: bool = True,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Create a database/storage group in the selected IoTDB target."""
+ selected_config = select_target_config(config, target_id=target_id, target=target)
+ with iotdb_target_response_context(selected_config):
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="create_database",
+ )
+ session = None
+ try:
+ _assert_database_ddl_permission(selected_config, action="CREATE")
+ database_path = _validate_database_name(
+ selected_config.sql_dialect, database
+ )
+ sql = f"CREATE DATABASE {database_path}"
+ session = session_pool.get_session()
+ session.execute_non_query_statement(sql)
+ session.close()
+ return sql_success_response("create_database", sql)
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to create database: {str(e)}")
+ raise
+
+ _, session_pool = table_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ tool_name="create_database",
+ )
+ table_session = None
+ try:
+ _assert_database_ddl_permission(selected_config, action="CREATE")
+ database_name = _validate_database_name(
+ selected_config.sql_dialect, database
+ )
+ sql = (
+ f"CREATE DATABASE IF NOT EXISTS {database_name}"
+ if if_not_exists
+ else f"CREATE DATABASE {database_name}"
+ )
+ table_session = session_pool.get_session()
+ table_session.execute_non_query_statement(sql)
+ table_session.close()
+ return sql_success_response("create_database", sql)
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to create database: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def drop_database(
+ database: str,
+ if_exists: bool = True,
+ confirm: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Drop a database/storage group in the selected IoTDB target."""
+ selected_config = select_target_config(config, target_id=target_id, target=target)
+ with iotdb_target_response_context(selected_config):
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ selected_config,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="drop_database",
+ )
+ session = None
+ try:
+ _assert_database_ddl_permission(
+ selected_config, action="DROP", confirm=confirm
+ )
+ database_path = _validate_database_name(
+ selected_config.sql_dialect, database
+ )
+ sql = f"DROP DATABASE {database_path}"
+ session = session_pool.get_session()
+ session.execute_non_query_statement(sql)
+ session.close()
+ return sql_success_response("drop_database", sql)
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to drop database: {str(e)}")
+ raise
+
+ _, session_pool = table_session_pool(
+ selected_config,
+ max_pool_size=max_pool_size,
+ tool_name="drop_database",
+ )
+ table_session = None
+ try:
+ _assert_database_ddl_permission(
+ selected_config, action="DROP", confirm=confirm
+ )
+ database_name = _validate_database_name(
+ selected_config.sql_dialect, database
+ )
+ sql = (
+ f"DROP DATABASE IF EXISTS {database_name}"
+ if if_exists
+ else f"DROP DATABASE {database_name}"
+ )
+ table_session = session_pool.get_session()
+ table_session.execute_non_query_statement(sql)
+ table_session.close()
+ return sql_success_response("drop_database", sql)
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to drop database: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def use_database(
+ database: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Switch current database in a table-model IoTDB target."""
+ selected_config, session_pool = table_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=max_pool_size,
+ tool_name="use_database",
+ )
+ table_session = None
+ try:
+ database_name = _validate_database_name(
+ selected_config.sql_dialect, database
+ )
+ sql = f"USE {database_name}"
+ table_session = session_pool.get_session()
+ table_session.execute_non_query_statement(sql)
+ table_session.close()
+ table_session = None
+
+ updated_target = merge_target(
+ selected_config.to_target(),
+ {"database": database_name},
+ )
+ registry = config.target_registry or selected_config.target_registry
+ if registry is not None:
+ updated_registry = registry.with_target(updated_target)
+ apply_target_registry(
+ config,
+ updated_registry,
+ active_target_id=config.target_id or updated_target.target_id,
+ close_pools=False,
+ )
+ if config.session_manager is not None:
+ config.session_manager.close_target_pools(updated_target.target_id)
+ updated_config = select_target_config(
+ config,
+ target_id=updated_target.target_id,
+ )
+ else:
+ selected_config.database = database_name
+ if config.target_id == selected_config.target_id:
+ config.database = database_name
+ updated_config = selected_config
+
+ with iotdb_target_response_context(updated_config):
+ return sql_success_response("use_database", sql)
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to use database: {str(e)}")
+ raise
diff --git a/src/iotdb_mcp_server/services/explain.py b/src/iotdb_mcp_server/services/explain.py
new file mode 100644
index 0000000..3c67e1d
--- /dev/null
+++ b/src/iotdb_mcp_server/services/explain.py
@@ -0,0 +1,168 @@
+#
+# 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.
+#
+
+import logging
+
+from iotdb.Session import Session
+from iotdb.table_session import TableSession
+from iotdb.utils.SessionDataSet import SessionDataSet
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.services.json_response import error_response, payload_response
+from iotdb_mcp_server.services.target_selection import (
+ iotdb_target_response_context,
+ select_target_config,
+ table_session_pool,
+ tree_session_pool,
+)
+from iotdb_mcp_server.services.tree_sql_guardrails import assert_tree_query_shape
+
+
+def _collect_result(
+ res: SessionDataSet, session_or_table_session: Session | TableSession
+) -> tuple[list[str], list[str]]:
+ columns = res.get_column_names()
+ rows: list[str] = []
+ while res.has_next():
+ row = res.next().get_fields()
+ rows.append(",".join(map(str, row)))
+ session_or_table_session.close()
+ return columns, rows
+
+
+def _normalize_and_validate_sql(
+ sql_dialect: str, sql: str, analyze: bool
+) -> tuple[str, str]:
+ raw_sql = sql.strip()
+ if not raw_sql:
+ raise ValueError("SQL cannot be empty")
+
+ explain_sql = raw_sql
+ upper = raw_sql.upper()
+
+ if not upper.startswith("EXPLAIN"):
+ explain_prefix = "EXPLAIN ANALYZE " if analyze else "EXPLAIN "
+ explain_sql = explain_prefix + raw_sql
+ target_sql = raw_sql.strip()
+ else:
+ # Accept user-provided EXPLAIN; do not override existing options.
+ target_sql = raw_sql[len("EXPLAIN") :].strip()
+ if target_sql.upper().startswith("ANALYZE"):
+ target_sql = target_sql[len("ANALYZE") :].strip()
+
+ target_upper = target_sql.upper()
+ allowed_prefixes = (
+ ("SELECT", "SHOW", "COUNT", "WITH")
+ if sql_dialect == "tree"
+ else ("SELECT", "SHOW", "DESC", "DESCRIBE", "WITH")
+ )
+
+ if not target_upper.startswith(allowed_prefixes):
+ raise ValueError(
+ f"explain_query only supports query-like SQL for {sql_dialect} dialect. "
+ f"Allowed prefixes: {', '.join(allowed_prefixes)}"
+ )
+ if sql_dialect == "tree" and target_upper.startswith("SELECT"):
+ assert_tree_query_shape(target_sql)
+
+ return explain_sql, target_sql
+
+
+def register_explain_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register EXPLAIN tool with per-call IoTDB target selection."""
+ max_pool_size = 100
+
+ @mcp.tool()
+ async def explain_query(
+ query_sql: str,
+ analyze: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute EXPLAIN for one statement against the selected IoTDB target."""
+ selected_config = select_target_config(config, target_id=target_id, target=target)
+ with iotdb_target_response_context(selected_config):
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="explain_query",
+ )
+ session = None
+ try:
+ explain_sql, _ = _normalize_and_validate_sql(
+ selected_config.sql_dialect, query_sql, analyze
+ )
+ session = session_pool.get_session()
+ res = session.execute_query_statement(explain_sql)
+ columns, rows = _collect_result(res, session)
+ return payload_response(
+ "explain_query",
+ {
+ "explain_sql": explain_sql,
+ "plan": {
+ "format": "csv",
+ "columns": columns,
+ "rows": rows,
+ "text": "\n".join([",".join(columns)] + rows),
+ },
+ },
+ message="Explain executed.",
+ )
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute explain query: {str(e)}")
+ return error_response("explain_query", str(e))
+
+ _, session_pool = table_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ tool_name="explain_query",
+ )
+ table_session = None
+ try:
+ explain_sql, _ = _normalize_and_validate_sql(
+ selected_config.sql_dialect, query_sql, analyze
+ )
+ table_session = session_pool.get_session()
+ res = table_session.execute_query_statement(explain_sql)
+ columns, rows = _collect_result(res, table_session)
+ return payload_response(
+ "explain_query",
+ {
+ "explain_sql": explain_sql,
+ "plan": {
+ "format": "csv",
+ "columns": columns,
+ "rows": rows,
+ "text": "\n".join([",".join(columns)] + rows),
+ },
+ },
+ message="Explain executed.",
+ )
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to execute explain query: {str(e)}")
+ return error_response("explain_query", str(e))
diff --git a/src/iotdb_mcp_server/services/json_response.py b/src/iotdb_mcp_server/services/json_response.py
new file mode 100644
index 0000000..b4c138d
--- /dev/null
+++ b/src/iotdb_mcp_server/services/json_response.py
@@ -0,0 +1,329 @@
+#
+# 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.
+#
+
+import json
+from contextvars import ContextVar
+from datetime import datetime, timezone
+from typing import Any, Callable, Iterable
+
+from mcp.types import TextContent
+
+from iotdb_mcp_server.result_store import FileResultStore
+
+_RESPONSE_VERSION = "1.0"
+_target_context: ContextVar[dict[str, Any] | None] = ContextVar(
+ "iotdb_target_context", default=None
+)
+
+_TRUE_VALUES = {"1", "true", "yes", "on", "available", "running", "enabled"}
+_FALSE_VALUES = {"0", "false", "no", "off", "unavailable", "stopped", "disabled"}
+
+
+def set_iotdb_target_context(target: dict[str, Any] | None):
+ return _target_context.set(target)
+
+
+def reset_iotdb_target_context(token) -> None:
+ _target_context.reset(token)
+
+
+def _coerce_availability(value: Any) -> str | None:
+ if value is None:
+ return None
+ if isinstance(value, bool):
+ return "yes" if value else "no"
+ normalized = str(value).strip().lower()
+ if normalized in _TRUE_VALUES:
+ return "yes"
+ if normalized in _FALSE_VALUES:
+ return "no"
+ if normalized in {"unknown", "probe_failed", "not_probed"}:
+ return "unknown"
+ return None
+
+
+def _ainode_availability(target: dict[str, Any]) -> tuple[str, Any]:
+ availability = _coerce_availability(target.get("ainode_available")) or "unknown"
+ source = target.get("ainode_availability_source")
+ if source is None:
+ source = {
+ "method": "SHOW AINODES",
+ "ok": False,
+ "error": "not_probed",
+ }
+ return availability, source
+
+
+def _cli_endpoint(target: dict[str, Any]) -> str:
+ node_urls = target.get("node_urls")
+ if isinstance(node_urls, list) and node_urls:
+ first = str(node_urls[0]).strip()
+ if first:
+ return first
+ host = str(target.get("host") or "127.0.0.1").strip()
+ port = str(target.get("port") or "6667").strip()
+ return f"{host}:{port}"
+
+
+def _row_count(payload: Any) -> int:
+ if not isinstance(payload, dict):
+ return 0
+ rows = payload.get("rows")
+ if isinstance(rows, list):
+ return len(rows)
+ for key in ("row_count", "rows_count", "count"):
+ value = payload.get(key)
+ if isinstance(value, int):
+ return value
+ if isinstance(value, str) and value.strip().isdigit():
+ return int(value.strip())
+ return 0
+
+
+def _iotdb_context(tool: str, payload: Any, target: dict[str, Any]) -> dict[str, Any]:
+ ainode_available, ainode_source = _ainode_availability(target)
+ return {
+ "target_id": target.get("target_id"),
+ "cli_endpoint": _cli_endpoint(target),
+ "ainode_available": ainode_available,
+ "ainode_availability_source": ainode_source,
+ "dialect": target.get("sql_dialect"),
+ "database": target.get("database"),
+ "tool": tool,
+ "rows": _row_count(payload),
+ }
+
+
+def _render_iotdb_context(context: dict[str, Any]) -> str:
+ lines = [
+ "TimeSeek IoTDB context:",
+ f"- target_id: {context.get('target_id') or ''}",
+ f"- cli_endpoint: {context.get('cli_endpoint') or ''}",
+ f"- ainode_available: {context.get('ainode_available') or 'unknown'}",
+ f"- dialect: {context.get('dialect') or ''}",
+ f"- database: {context.get('database') or ''}",
+ f"- tool: {context.get('tool') or ''}",
+ f"- rows: {context.get('rows', 0)}",
+ ]
+ return "\n".join(lines)
+
+
+class JsonParser:
+ """Simple parser/checker for MCP tool JSON envelopes."""
+
+ _required_keys = ("version", "tool", "ok", "timestamp", "payload")
+
+ @staticmethod
+ def check_format(obj: Any) -> bool:
+ if not isinstance(obj, dict):
+ return False
+ for key in JsonParser._required_keys:
+ if key not in obj:
+ return False
+ if not isinstance(obj["version"], str):
+ return False
+ if not isinstance(obj["tool"], str) or not obj["tool"].strip():
+ return False
+ if not isinstance(obj["ok"], bool):
+ return False
+ if not isinstance(obj["timestamp"], str) or not obj["timestamp"].strip():
+ return False
+ return True
+
+ @staticmethod
+ def parse(text: str) -> dict[str, Any]:
+ try:
+ obj = json.loads(text)
+ except json.JSONDecodeError as e:
+ raise ValueError(f"Invalid JSON text: {e.msg}") from e
+
+ if not JsonParser.check_format(obj):
+ raise ValueError(
+ "Invalid wrapped response format. Expected keys: "
+ "version, tool, ok, timestamp, payload."
+ )
+ return obj
+
+
+def _envelope(
+ tool: str, payload: Any, ok: bool = True, message: str | None = None
+) -> dict[str, Any]:
+ target = _target_context.get()
+ context = _iotdb_context(tool, payload, target) if target else None
+ if target:
+ if isinstance(payload, dict):
+ payload = dict(payload)
+ payload.setdefault("iotdb_target", target)
+ if context:
+ payload.setdefault("iotdb_context", context)
+ else:
+ payload = {
+ "value": payload,
+ "iotdb_target": target,
+ }
+ if context:
+ payload["iotdb_context"] = context
+ obj: dict[str, Any] = {
+ "version": _RESPONSE_VERSION,
+ "tool": tool,
+ "ok": ok,
+ "timestamp": datetime.now(timezone.utc).isoformat(),
+ "payload": payload,
+ }
+ if context:
+ obj["context"] = {"iotdb": context}
+ if message is not None:
+ obj["message"] = message
+ return obj
+
+
+def _to_text_content(obj: dict[str, Any]) -> list[TextContent]:
+ if not JsonParser.check_format(obj):
+ raise ValueError("JsonParser format check failed before output.")
+ text = json.dumps(obj, ensure_ascii=False)
+ JsonParser.parse(text)
+ context = obj.get("context")
+ iotdb_context = context.get("iotdb") if isinstance(context, dict) else None
+ if isinstance(iotdb_context, dict):
+ return [
+ TextContent(type="text", text=_render_iotdb_context(iotdb_context)),
+ TextContent(type="text", text=text),
+ ]
+ return [TextContent(type="text", text=text)]
+
+
+def payload_response(
+ tool: str, payload: Any, message: str | None = None
+) -> list[TextContent]:
+ return _to_text_content(_envelope(tool=tool, payload=payload, message=message))
+
+
+def error_response(tool: str, error: str) -> list[TextContent]:
+ return _to_text_content(
+ _envelope(
+ tool=tool,
+ payload={"error": error},
+ ok=False,
+ message="Tool execution failed.",
+ )
+ )
+
+
+def text_payload_response(
+ tool: str, text: str, message: str | None = None
+) -> list[TextContent]:
+ return payload_response(
+ tool=tool,
+ payload={"format": "text", "text": text},
+ message=message,
+ )
+
+
+def _csv_text(columns: list[str], rows: list[str]) -> str:
+ return "\n".join([",".join(columns)] + rows)
+
+
+def csv_payload_response(
+ tool: str,
+ columns: list[str],
+ rows: list[str],
+ message: str | None = None,
+ diagnostics: list[dict[str, Any]] | None = None,
+) -> list[TextContent]:
+ csv_text = _csv_text(columns, rows)
+ payload: dict[str, Any] = {
+ "format": "csv",
+ "columns": columns,
+ "rows": rows,
+ "text": csv_text,
+ "row_count": len(rows),
+ "inline_row_count": len(rows),
+ "inline_truncated": False,
+ }
+ if diagnostics:
+ payload["diagnostics"] = diagnostics
+ return payload_response(
+ tool=tool,
+ payload=payload,
+ message=message,
+ )
+
+
+def csv_result_payload_response(
+ tool: str,
+ columns: list[str],
+ rows: Iterable[str],
+ *,
+ result_store: FileResultStore,
+ source: dict[str, Any] | None = None,
+ message: str | None = None,
+ diagnostics: list[dict[str, Any]] | None = None,
+ diagnostics_factory: Callable[[int], list[dict[str, Any]] | None] | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+) -> list[TextContent]:
+ stored = result_store.write_csv_result(
+ tool=tool,
+ columns=columns,
+ rows=rows,
+ source=source,
+ owner_session_id=owner_session_id,
+ preview_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ preview_rows = stored.preview_rows
+ inline_truncated = stored.row_count > len(preview_rows)
+ payload: dict[str, Any] = {
+ "format": "csv",
+ "columns": columns,
+ "rows": preview_rows,
+ "text": _csv_text(columns, preview_rows),
+ "row_count": stored.row_count,
+ "inline_row_count": len(preview_rows),
+ "inline_truncated": inline_truncated,
+ "result_id": stored.result_id,
+ "owner_session_id": stored.owner_session_id,
+ "next_cursor": str(len(preview_rows)) if inline_truncated else None,
+ "result_store": {
+ "type": "file",
+ "result_id": stored.result_id,
+ "owner_session_id": stored.owner_session_id,
+ "page_tool": "read_result_page",
+ "batch_page_tool": "read_result_pages",
+ "page_size_rows": stored.page_size_rows,
+ "manifest_path": stored.manifest_path,
+ "shard_count": stored.shard_count,
+ "byte_count": stored.byte_count,
+ },
+ }
+ combined_diagnostics = list(diagnostics or [])
+ if diagnostics_factory is not None:
+ combined_diagnostics.extend(diagnostics_factory(stored.row_count) or [])
+ if combined_diagnostics:
+ payload["diagnostics"] = combined_diagnostics
+ return payload_response(tool=tool, payload=payload, message=message)
+
+
+def sql_success_response(tool: str, sql: str) -> list[TextContent]:
+ return payload_response(
+ tool=tool,
+ payload={"sql": sql, "result": "success"},
+ message="SQL executed successfully.",
+ )
diff --git a/src/iotdb_mcp_server/services/metadata.py b/src/iotdb_mcp_server/services/metadata.py
new file mode 100644
index 0000000..b7a8398
--- /dev/null
+++ b/src/iotdb_mcp_server/services/metadata.py
@@ -0,0 +1,416 @@
+#
+# 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.
+#
+
+import logging
+import re
+
+from iotdb.Session import Session
+from iotdb.table_session import TableSession
+from iotdb.utils.SessionDataSet import SessionDataSet
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.result_store import result_store_from_export_path
+from iotdb_mcp_server.runtime_policy import dynamic_env_bool
+from iotdb_mcp_server.runtime_policy import dynamic_getenv
+from iotdb_mcp_server.runtime_policy import strict_permission_enforcement
+from iotdb_mcp_server.services.json_response import csv_result_payload_response
+from iotdb_mcp_server.services.target_selection import (
+ iotdb_target_response_context,
+ select_target_config,
+ table_session_pool,
+ tree_session_pool,
+)
+
+_TREE_PATH_PATTERN = re.compile(r"^root(?:\.[A-Za-z_][A-Za-z0-9_]*|\.\*|\.\*\*)*$")
+_TABLE_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+
+def _env_bool(name: str, default: bool) -> bool:
+ return dynamic_env_bool(name, default)
+
+
+def _csv_set(value: str) -> set[str]:
+ return {item.strip() for item in value.split(",") if item.strip()}
+
+
+def _assert_metadata_permission(config: Config) -> None:
+ if strict_permission_enforcement():
+ if not _env_bool("IOTDB_ENABLE_METADATA_QUERY", True):
+ raise PermissionError(
+ "Metadata tools are disabled by strict MCP policy. "
+ "Set IOTDB_ENABLE_METADATA_QUERY=true to enable."
+ )
+
+ allowed_users = _csv_set(
+ dynamic_getenv("IOTDB_METADATA_ALLOWED_USERS", "*") or "*"
+ )
+ if "*" not in allowed_users and config.user not in allowed_users:
+ raise PermissionError(
+ f"Current MCP user '{config.user}' is not allowed by IOTDB_METADATA_ALLOWED_USERS."
+ )
+
+
+def _normalize_sql(sql: str) -> str:
+ cleaned = sql.strip()
+ if not cleaned:
+ raise ValueError("Metadata SQL cannot be empty.")
+ if cleaned.endswith(";"):
+ cleaned = cleaned[:-1].strip()
+ if ";" in cleaned:
+ raise ValueError("Only a single SQL statement is allowed.")
+ return cleaned
+
+
+def _validate_tree_path(path: str) -> str:
+ cleaned = path.strip()
+ if not cleaned:
+ raise ValueError("Tree path cannot be empty.")
+ if not _TREE_PATH_PATTERN.fullmatch(cleaned):
+ raise ValueError(
+ "Invalid tree path. Use path like root.sg, root.sg.dev, root.**, root.sg.*"
+ )
+ return cleaned
+
+
+def _validate_table_identifier(identifier: str) -> str:
+ cleaned = identifier.strip()
+ if not _TABLE_IDENTIFIER_PATTERN.fullmatch(cleaned):
+ raise ValueError(
+ "Invalid table identifier. Use letters, numbers, and underscore only."
+ )
+ return cleaned
+
+
+def _format_result(
+ res: SessionDataSet,
+ session_or_table_session: Session | TableSession,
+ tool_name: str,
+ export_path: str,
+ query_sql: str | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+) -> list[TextContent]:
+ columns = res.get_column_names()
+
+ def rows():
+ while res.has_next():
+ row = res.next().get_fields()
+ yield ",".join(map(str, row))
+
+ try:
+ return csv_result_payload_response(
+ tool_name,
+ columns,
+ rows(),
+ result_store=result_store_from_export_path(export_path),
+ source={"query_sql": query_sql} if query_sql else None,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ finally:
+ session_or_table_session.close()
+
+
+def _ensure_prefix(
+ sql: str, allowed_prefixes: tuple[str, ...], action_name: str
+) -> None:
+ upper = sql.upper()
+ if not upper.startswith(allowed_prefixes):
+ raise ValueError(
+ f"{action_name} only supports SQL starting with: {', '.join(allowed_prefixes)}"
+ )
+
+
+def register_metadata_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register metadata tools with per-call IoTDB target selection."""
+ max_pool_size = 100
+ tree_prefixes = (
+ "SHOW DATABASES",
+ "SHOW TIMESERIES",
+ "SHOW DEVICES",
+ "SHOW CHILD PATHS",
+ "SHOW CHILD NODES",
+ "SHOW FUNCTIONS",
+ "COUNT TIMESERIES",
+ "COUNT NODES",
+ "COUNT DEVICES",
+ )
+ table_prefixes = ("SHOW", "DESC", "DESCRIBE")
+
+ def _require_target_dialect(
+ required_sql_dialect: str,
+ tool_name: str,
+ target_id: str | None,
+ target: dict[str, object] | None,
+ ) -> None:
+ select_target_config(
+ config,
+ target_id=target_id,
+ target=target,
+ required_sql_dialect=required_sql_dialect,
+ tool_name=tool_name,
+ )
+
+ async def _metadata_query_impl(
+ query_sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+ ) -> list[TextContent]:
+ """Execute metadata SQL against the selected IoTDB target."""
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ with iotdb_target_response_context(selected_config):
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="metadata_query",
+ )
+ session = None
+ try:
+ _assert_metadata_permission(selected_config)
+ sql = _normalize_sql(query_sql)
+ _ensure_prefix(sql, tree_prefixes, "metadata_query")
+ session = session_pool.get_session()
+ res = session.execute_query_statement(sql)
+ return _format_result(
+ res,
+ session,
+ "metadata_query",
+ selected_config.export_path,
+ sql,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute metadata_query: {str(e)}")
+ raise
+
+ _, session_pool = table_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ tool_name="metadata_query",
+ )
+ table_session = None
+ try:
+ _assert_metadata_permission(selected_config)
+ sql = _normalize_sql(query_sql)
+ _ensure_prefix(sql, table_prefixes, "metadata_query")
+ table_session = session_pool.get_session()
+ res = table_session.execute_query_statement(sql)
+ return _format_result(
+ res,
+ table_session,
+ "metadata_query",
+ selected_config.export_path,
+ sql,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to execute metadata_query: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def metadata_query(
+ query_sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+ ) -> list[TextContent]:
+ """Execute metadata SQL against the selected IoTDB target."""
+ return await _metadata_query_impl(
+ query_sql,
+ target_id=target_id,
+ target=target,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+
+ @mcp.tool()
+ async def list_timeseries(
+ path: str = "root.**",
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ ) -> list[TextContent]:
+ """List timeseries under a tree path pattern."""
+ _require_target_dialect("tree", "list_timeseries", target_id, target)
+ query_sql = f"SHOW TIMESERIES {_validate_tree_path(path)}"
+ return await _metadata_query_impl(
+ query_sql,
+ target_id=target_id,
+ target=target,
+ owner_session_id=owner_session_id,
+ )
+
+ @mcp.tool()
+ async def list_devices(
+ path: str = "root.**",
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ ) -> list[TextContent]:
+ """List devices under a tree path pattern."""
+ _require_target_dialect("tree", "list_devices", target_id, target)
+ query_sql = f"SHOW DEVICES {_validate_tree_path(path)}"
+ return await _metadata_query_impl(
+ query_sql,
+ target_id=target_id,
+ target=target,
+ owner_session_id=owner_session_id,
+ )
+
+ @mcp.tool()
+ async def list_child_paths(
+ path: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ ) -> list[TextContent]:
+ """List child paths under a tree path."""
+ _require_target_dialect("tree", "list_child_paths", target_id, target)
+ query_sql = f"SHOW CHILD PATHS {_validate_tree_path(path)}"
+ return await _metadata_query_impl(
+ query_sql,
+ target_id=target_id,
+ target=target,
+ owner_session_id=owner_session_id,
+ )
+
+ @mcp.tool()
+ async def list_child_nodes(
+ path: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ ) -> list[TextContent]:
+ """List child nodes under a tree path."""
+ _require_target_dialect("tree", "list_child_nodes", target_id, target)
+ query_sql = f"SHOW CHILD NODES {_validate_tree_path(path)}"
+ return await _metadata_query_impl(
+ query_sql,
+ target_id=target_id,
+ target=target,
+ owner_session_id=owner_session_id,
+ )
+
+ @mcp.tool()
+ async def count_timeseries(
+ path: str = "root.**",
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ ) -> list[TextContent]:
+ """Count timeseries under a tree path pattern."""
+ _require_target_dialect("tree", "count_timeseries", target_id, target)
+ query_sql = f"COUNT TIMESERIES {_validate_tree_path(path)}"
+ return await _metadata_query_impl(
+ query_sql,
+ target_id=target_id,
+ target=target,
+ owner_session_id=owner_session_id,
+ )
+
+ @mcp.tool()
+ async def count_devices(
+ path: str = "root.**",
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ ) -> list[TextContent]:
+ """Count devices under a tree path pattern."""
+ _require_target_dialect("tree", "count_devices", target_id, target)
+ query_sql = f"COUNT DEVICES {_validate_tree_path(path)}"
+ return await _metadata_query_impl(
+ query_sql,
+ target_id=target_id,
+ target=target,
+ owner_session_id=owner_session_id,
+ )
+
+ @mcp.tool()
+ async def count_nodes(
+ path: str = "root",
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ ) -> list[TextContent]:
+ """Count nodes under a tree path."""
+ _require_target_dialect("tree", "count_nodes", target_id, target)
+ query_sql = f"COUNT NODES {_validate_tree_path(path)}"
+ return await _metadata_query_impl(
+ query_sql,
+ target_id=target_id,
+ target=target,
+ owner_session_id=owner_session_id,
+ )
+
+ @mcp.tool()
+ async def list_tables(
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ ) -> list[TextContent]:
+ """List all tables in current table-model database."""
+ _require_target_dialect("table", "list_tables", target_id, target)
+ return await _metadata_query_impl(
+ "SHOW TABLES",
+ target_id=target_id,
+ target=target,
+ owner_session_id=owner_session_id,
+ )
+
+ @mcp.tool()
+ async def describe_table(
+ table_name: str,
+ details: bool = True,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ ) -> list[TextContent]:
+ """Describe schema for a table in current table-model database."""
+ _require_target_dialect("table", "describe_table", target_id, target)
+ safe_table = _validate_table_identifier(table_name)
+ details_suffix = " details" if details else ""
+ return await _metadata_query_impl(
+ f"DESC {safe_table}{details_suffix}",
+ target_id=target_id,
+ target=target,
+ owner_session_id=owner_session_id,
+ )
diff --git a/src/iotdb_mcp_server/services/model.py b/src/iotdb_mcp_server/services/model.py
new file mode 100644
index 0000000..577c172
--- /dev/null
+++ b/src/iotdb_mcp_server/services/model.py
@@ -0,0 +1,585 @@
+#
+# 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.
+#
+
+import logging
+import re
+from typing import Any
+
+from iotdb.Session import Session
+from iotdb.table_session import TableSession
+from iotdb.utils.SessionDataSet import SessionDataSet
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.result_store import result_store_from_export_path
+from iotdb_mcp_server.runtime_policy import dynamic_env_bool
+from iotdb_mcp_server.runtime_policy import dynamic_getenv
+from iotdb_mcp_server.runtime_policy import strict_permission_enforcement
+from iotdb_mcp_server.services.json_response import (
+ csv_result_payload_response,
+ payload_response,
+ sql_success_response,
+)
+from iotdb_mcp_server.services.target_selection import (
+ iotdb_target_response_context,
+ select_target_config,
+ table_session_pool,
+ tree_session_pool,
+)
+
+_MODEL_ID_PATTERN = re.compile(r"^[A-Za-z][A-Za-z0-9_]{1,63}$")
+_INFERENCE_PATTERN = re.compile(
+ r"^\s*CALL\s+INFERENCE\s*\((.*)\)\s*$", re.IGNORECASE | re.DOTALL
+)
+_ALLOWED_INFERENCE_PARAMS = {"generateTime", "outputLength"}
+
+
+def _env_bool(name: str, default: bool) -> bool:
+ return dynamic_env_bool(name, default)
+
+
+def _csv_set(value: str) -> set[str]:
+ return {item.strip() for item in value.split(",") if item.strip()}
+
+
+def _assert_model_permission(
+ config: Config, action: str, confirm_destructive: bool = False
+) -> None:
+ if strict_permission_enforcement():
+ if not _env_bool("IOTDB_ENABLE_MODEL_MANAGEMENT", False):
+ raise PermissionError(
+ "AINode model tools are disabled by strict MCP policy. "
+ "Set IOTDB_ENABLE_MODEL_MANAGEMENT=true to enable."
+ )
+
+ allowed_users = _csv_set(
+ dynamic_getenv("IOTDB_MODEL_ALLOWED_USERS", "root") or "root"
+ )
+ if "*" not in allowed_users and config.user not in allowed_users:
+ raise PermissionError(
+ f"Current MCP user '{config.user}' is not allowed by IOTDB_MODEL_ALLOWED_USERS."
+ )
+
+ if (
+ action == "DESTRUCTIVE"
+ and _env_bool("IOTDB_REQUIRE_MODEL_DESTRUCTIVE_CONFIRM", True)
+ and not confirm_destructive
+ ):
+ raise PermissionError(
+ "Destructive model command requires confirm_destructive=True by server policy."
+ )
+
+
+def _strip_single_trailing_semicolon(sql: str, empty_message: str) -> str:
+ cleaned = sql.strip()
+ if not cleaned:
+ raise ValueError(empty_message)
+
+ quote: str | None = None
+ escaped = False
+ semicolon_indexes: list[int] = []
+ index = 0
+ while index < len(cleaned):
+ char = cleaned[index]
+ if escaped:
+ escaped = False
+ index += 1
+ continue
+ if char == "\\" and quote:
+ escaped = True
+ index += 1
+ continue
+ if quote:
+ if char == quote:
+ if index + 1 < len(cleaned) and cleaned[index + 1] == quote:
+ index += 2
+ continue
+ quote = None
+ index += 1
+ continue
+ if char in {"'", '"'}:
+ quote = char
+ elif char == ";":
+ semicolon_indexes.append(index)
+ index += 1
+
+ if not semicolon_indexes:
+ return cleaned
+ if semicolon_indexes == [len(cleaned) - 1]:
+ return cleaned[:-1].strip()
+ raise ValueError("Only a single SQL statement is allowed.")
+
+
+def _normalize_sql(sql: str) -> str:
+ return _strip_single_trailing_semicolon(sql, "Model SQL cannot be empty.")
+
+
+def _normalize_for_prefix(sql: str) -> str:
+ return re.sub(r"\s+", " ", sql).strip().upper()
+
+
+def _matched_prefix(sql: str, prefixes: tuple[str, ...]) -> str | None:
+ upper = _normalize_for_prefix(sql)
+ for prefix in prefixes:
+ if not upper.startswith(prefix):
+ continue
+ if (
+ len(upper) == len(prefix)
+ or upper[len(prefix)].isspace()
+ or upper[len(prefix)] == "("
+ ):
+ return prefix
+ return None
+
+
+def _assert_prefix(sql: str, prefixes: tuple[str, ...], tool_name: str) -> None:
+ if _matched_prefix(sql, prefixes):
+ return
+ raise ValueError(
+ f"{tool_name} only supports SQL starting with: {', '.join(prefixes)}"
+ )
+
+
+def _split_top_level_args(value: str) -> list[str]:
+ args: list[str] = []
+ start = 0
+ quote: str | None = None
+ escaped = False
+ paren_depth = 0
+
+ for index, char in enumerate(value):
+ if escaped:
+ escaped = False
+ continue
+ if char == "\\" and quote:
+ escaped = True
+ continue
+ if quote:
+ if char == quote:
+ if index + 1 < len(value) and value[index + 1] == quote:
+ continue
+ quote = None
+ continue
+ if char in {"'", '"'}:
+ quote = char
+ elif char == "(":
+ paren_depth += 1
+ elif char == ")":
+ if paren_depth > 0:
+ paren_depth -= 1
+ elif char == "," and paren_depth == 0:
+ args.append(value[start:index].strip())
+ start = index + 1
+
+ if quote:
+ raise ValueError("CALL INFERENCE input has an unterminated string literal.")
+ if paren_depth:
+ raise ValueError("CALL INFERENCE input has unbalanced parentheses.")
+
+ args.append(value[start:].strip())
+ return args
+
+
+def _decode_sql_string(value: str) -> str:
+ if len(value) < 2 or value[0] not in {"'", '"'} or value[-1] != value[0]:
+ raise ValueError(
+ "CALL INFERENCE second argument must be a quoted SELECT SQL string, "
+ 'for example: "SELECT s0 FROM root.AI LIMIT 256".'
+ )
+ quote = value[0]
+ body = value[1:-1]
+ return body.replace(quote * 2, quote)
+
+
+def _quote_sql_string(value: str) -> str:
+ return '"' + value.replace('"', '""') + '"'
+
+
+def _assert_model_id(model_id: str) -> None:
+ if not _MODEL_ID_PATTERN.match(model_id):
+ raise ValueError(
+ "model_id must be 2-64 characters, start with a letter, and contain "
+ "only letters, numbers, or underscores."
+ )
+
+
+def _validate_inference_input_sql(input_sql: str) -> list[dict[str, Any]]:
+ cleaned = _strip_single_trailing_semicolon(
+ input_sql, "CALL INFERENCE input SELECT SQL cannot be empty."
+ )
+ normalized = _normalize_for_prefix(cleaned)
+ diagnostics: list[dict[str, Any]] = []
+
+ if not normalized.startswith("SELECT "):
+ raise ValueError("CALL INFERENCE input SQL must be a SELECT query.")
+ if re.match(r"^\s*SELECT\s+\*", cleaned, re.IGNORECASE):
+ raise ValueError(
+ "CALL INFERENCE input SQL must use an explicit ordered column list; "
+ "do not use SELECT * because wildcard column order is undefined."
+ )
+ if " LIMIT " not in f" {normalized} " and " WHERE " not in f" {normalized} ":
+ diagnostics.append(
+ {
+ "severity": "warning",
+ "code": "unbounded_input_sql",
+ "message": "Prefer LIMIT or a concrete time predicate for inference input SQL.",
+ }
+ )
+ return diagnostics
+
+
+def _parse_inference_params(raw_params: list[str]) -> dict[str, Any]:
+ params: dict[str, Any] = {}
+ for raw in raw_params:
+ if "=" not in raw:
+ raise ValueError(
+ "CALL INFERENCE optional arguments must use name=value syntax. "
+ "Supported names: generateTime, outputLength."
+ )
+ name, raw_value = [part.strip() for part in raw.split("=", 1)]
+ if name not in _ALLOWED_INFERENCE_PARAMS:
+ raise ValueError(
+ f"Unsupported CALL INFERENCE parameter '{name}'. "
+ "Supported parameters: generateTime, outputLength."
+ )
+ if name in params:
+ raise ValueError(f"Duplicate CALL INFERENCE parameter '{name}'.")
+
+ if name == "generateTime":
+ lowered = raw_value.lower()
+ if lowered not in {"true", "false"}:
+ raise ValueError("generateTime must be a boolean: true or false.")
+ params[name] = lowered == "true"
+ elif name == "outputLength":
+ if not re.match(r"^[1-9][0-9]*$", raw_value):
+ raise ValueError("outputLength must be a positive integer.")
+ params[name] = int(raw_value)
+ return params
+
+
+def _validate_inference_sql(sql: str) -> dict[str, Any]:
+ match = _INFERENCE_PATTERN.match(sql)
+ if not match:
+ raise ValueError(
+ "model_inference expects CALL INFERENCE(model_id, "
+ '"SELECT ...", generateTime=true|false, outputLength=<positive int>).'
+ )
+
+ args = _split_top_level_args(match.group(1))
+ if len(args) < 2:
+ raise ValueError(
+ "CALL INFERENCE requires at least model_id and input SELECT SQL."
+ )
+
+ model_id = args[0].strip()
+ _assert_model_id(model_id)
+
+ input_sql = _decode_sql_string(args[1])
+ diagnostics = _validate_inference_input_sql(input_sql)
+ params = _parse_inference_params(args[2:])
+
+ return {
+ "model_id": model_id,
+ "input_sql": input_sql,
+ "parameters": params,
+ "diagnostics": diagnostics,
+ }
+
+
+def _build_inference_sql(
+ model_id: str,
+ input_sql: str,
+ output_length: int = 96,
+ generate_time: bool = False,
+) -> tuple[str, dict[str, Any]]:
+ _assert_model_id(model_id)
+ diagnostics = _validate_inference_input_sql(input_sql)
+ if output_length <= 0:
+ raise ValueError("output_length must be a positive integer.")
+ inference_sql = (
+ f"CALL INFERENCE({model_id}, {_quote_sql_string(input_sql)}, "
+ f"generateTime={str(generate_time).lower()}, outputLength={output_length})"
+ )
+ return inference_sql, {
+ "model_id": model_id,
+ "input_sql": input_sql,
+ "parameters": {
+ "generateTime": generate_time,
+ "outputLength": output_length,
+ },
+ "diagnostics": diagnostics,
+ }
+
+
+def _format_result(
+ res: SessionDataSet,
+ session_or_table_session: Session | TableSession,
+ tool_name: str,
+ export_path: str,
+ sql: str | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+) -> list[TextContent]:
+ columns = res.get_column_names()
+
+ def rows():
+ while res.has_next():
+ row = res.next().get_fields()
+ yield ",".join(map(str, row))
+
+ try:
+ return csv_result_payload_response(
+ tool_name,
+ columns,
+ rows(),
+ result_store=result_store_from_export_path(export_path),
+ source={"sql": sql} if sql else None,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ finally:
+ session_or_table_session.close()
+
+
+def register_model_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register AINode model-management SQL tools with target selection."""
+ max_pool_size = 100
+
+ @mcp.tool()
+ async def model_query(
+ model_sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+ ) -> list[TextContent]:
+ """Execute AINode model-query SQL against the selected IoTDB target."""
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ with iotdb_target_response_context(selected_config):
+ sql = _normalize_sql(model_sql)
+ _assert_prefix(
+ sql,
+ (
+ "SHOW MODELS",
+ "SHOW LOADED MODELS",
+ "SHOW AI_DEVICES",
+ "SHOW AINODES",
+ ),
+ "model_query",
+ )
+ _assert_model_permission(selected_config, action="QUERY")
+
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="model_query",
+ )
+ session = None
+ try:
+ session = session_pool.get_session()
+ res = session.execute_query_statement(sql)
+ return _format_result(
+ res,
+ session,
+ "model_query",
+ selected_config.export_path,
+ sql,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute model_query: {str(e)}")
+ raise
+
+ _, session_pool = table_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ tool_name="model_query",
+ )
+ table_session = None
+ try:
+ table_session = session_pool.get_session()
+ res = table_session.execute_query_statement(sql)
+ return _format_result(
+ res,
+ table_session,
+ "model_query",
+ selected_config.export_path,
+ sql,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to execute model_query: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def prepare_model_inference_request(
+ model_id: str,
+ input_sql: str,
+ output_length: int = 96,
+ generate_time: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Build and validate Tree-dialect AINode CALL INFERENCE SQL."""
+ selected_config = select_target_config(
+ config,
+ target_id=target_id,
+ target=target,
+ required_sql_dialect="tree",
+ tool_name="prepare_model_inference_request",
+ )
+ with iotdb_target_response_context(selected_config):
+ inference_sql, validation = _build_inference_sql(
+ model_id=model_id,
+ input_sql=input_sql,
+ output_length=output_length,
+ generate_time=generate_time,
+ )
+ return payload_response(
+ "prepare_model_inference_request",
+ {
+ "inference_sql": inference_sql,
+ "validation": validation,
+ "recommended_next_tool": "model_inference",
+ },
+ )
+
+ @mcp.tool()
+ async def model_inference(
+ inference_sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+ ) -> list[TextContent]:
+ """Execute Tree-dialect AINode inference SQL."""
+ selected_config, session_pool = tree_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="model_inference",
+ )
+ with iotdb_target_response_context(selected_config):
+ session = None
+ try:
+ _assert_model_permission(selected_config, action="QUERY")
+ sql = _normalize_sql(inference_sql)
+ _assert_prefix(sql, ("CALL INFERENCE",), "model_inference")
+ _validate_inference_sql(sql)
+ session = session_pool.get_session()
+ res = session.execute_query_statement(sql)
+ return _format_result(
+ res,
+ session,
+ "model_inference",
+ selected_config.export_path,
+ sql,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute model_inference: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def model_command(
+ model_sql: str,
+ confirm_destructive: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute AINode model-management command SQL."""
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ with iotdb_target_response_context(selected_config):
+ sql = _normalize_sql(model_sql)
+
+ if _matched_prefix(sql, ("DROP MODEL", "UNLOAD MODEL", "REMOVE AINODE")):
+ _assert_model_permission(
+ selected_config,
+ action="DESTRUCTIVE",
+ confirm_destructive=confirm_destructive,
+ )
+ elif _matched_prefix(sql, ("CREATE MODEL", "LOAD MODEL")):
+ _assert_model_permission(selected_config, action="MANAGE")
+ else:
+ raise ValueError(
+ "model_command only supports SQL starting with: "
+ "CREATE MODEL, DROP MODEL, LOAD MODEL, UNLOAD MODEL, REMOVE AINODE"
+ )
+
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ selected_config,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="model_command",
+ )
+ session = None
+ try:
+ session = session_pool.get_session()
+ session.execute_non_query_statement(sql)
+ session.close()
+ return sql_success_response("model_command", sql)
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute model_command: {str(e)}")
+ raise
+
+ _, session_pool = table_session_pool(
+ selected_config,
+ max_pool_size=max_pool_size,
+ tool_name="model_command",
+ )
+ table_session = None
+ try:
+ table_session = session_pool.get_session()
+ table_session.execute_non_query_statement(sql)
+ table_session.close()
+ return sql_success_response("model_command", sql)
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to execute model_command: {str(e)}")
+ raise
diff --git a/src/iotdb_mcp_server/services/query.py b/src/iotdb_mcp_server/services/query.py
new file mode 100644
index 0000000..096a36b
--- /dev/null
+++ b/src/iotdb_mcp_server/services/query.py
@@ -0,0 +1,434 @@
+import datetime
+import logging
+import os
+import re
+import uuid
+
+from iotdb.Session import Session
+from iotdb.table_session import TableSession
+from iotdb.utils.SessionDataSet import SessionDataSet
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.result_store import result_store_from_export_path
+from iotdb_mcp_server.services.json_response import (
+ csv_result_payload_response,
+ text_payload_response,
+)
+from iotdb_mcp_server.services.target_selection import (
+ iotdb_target_response_context,
+ table_session_pool,
+ tree_session_pool,
+)
+from iotdb_mcp_server.services.tree_sql_guardrails import assert_tree_query_shape
+from iotdb_mcp_server.services.tree_sql_guardrails import (
+ tree_from_wildcard_runtime_hint,
+)
+
+
+def sanitize_filename(filename: str, base_dir: str) -> str:
+ """
+ Sanitize and validate filename to prevent path traversal attacks.
+
+ Security patch for CVE-2026-XXXXX
+ Author: Mohammed Tanveer (threatpointer)
+ Date: 2026-01-12
+
+ Args:
+ filename: The user-provided filename
+ base_dir: The base directory for exports (must be absolute path)
+
+ Returns:
+ The sanitized absolute filepath
+
+ Raises:
+ ValueError: If the filename contains invalid characters or attempts path traversal
+
+ Security measures:
+ - Accepts path-like inputs but strips directories and uses only the basename
+ under base_dir
+ - Rejects directory traversal segments before processing
+ - Validates allowed characters (alphanumeric, underscore, hyphen, dot)
+ - Resolves absolute path and verifies it stays within base_dir boundary
+ - Prevents directory traversal, symlink attacks, and path manipulation
+ """
+ if not filename:
+ raise ValueError("Filename cannot be empty")
+
+ normalized = filename.replace("\\", "/").strip()
+ parts = [part for part in normalized.split("/") if part]
+ if any(part == ".." for part in parts):
+ raise ValueError(
+ "Invalid filename: directory traversal sequences are not allowed"
+ )
+
+ filename = os.path.basename(parts[-1] if parts else normalized)
+
+ if not re.match(r"^[a-zA-Z0-9_\-\.]+$", filename):
+ raise ValueError(
+ "Invalid filename: only alphanumeric characters, underscore, hyphen, and dot are allowed"
+ )
+
+ if not filename or filename in (".", ".."):
+ raise ValueError("Invalid filename")
+
+ if filename.startswith(".."):
+ raise ValueError("Invalid filename: cannot start with '..'")
+
+ filepath = os.path.join(base_dir, filename)
+ filepath_real = os.path.realpath(filepath)
+ basedir_real = os.path.realpath(base_dir)
+
+ if (
+ not filepath_real.startswith(basedir_real + os.sep)
+ and filepath_real != basedir_real
+ ):
+ raise ValueError(
+ "Path traversal detected: file must be within export directory"
+ )
+
+ return filepath_real
+
+
+def _ensure_export_directory(export_path: str, logger: logging.Logger) -> None:
+ if os.path.exists(export_path):
+ return
+ try:
+ os.makedirs(export_path)
+ logger.info(f"Created export directory: {export_path}")
+ except Exception as e:
+ logger.warning(f"Failed to create export directory {export_path}: {str(e)}")
+
+
+def _prepare_tree_res(
+ _res: SessionDataSet,
+ _session: Session,
+ tool_name: str,
+ export_path: str,
+ query_sql: str | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+) -> list[TextContent]:
+ columns = _res.get_column_names()
+
+ def rows():
+ while _res.has_next():
+ record = _res.next()
+ if columns and columns[0] == "Time":
+ timestamp = record.get_timestamp()
+ row = record.get_fields()
+ yield str(timestamp) + "," + ",".join(map(str, row))
+ else:
+ row = record.get_fields()
+ yield ",".join(map(str, row))
+
+ def diagnostics_factory(row_count: int):
+ diagnostics = []
+ if query_sql:
+ hint = tree_from_wildcard_runtime_hint(query_sql, row_count=row_count)
+ if hint:
+ diagnostics.append(hint)
+ return diagnostics
+
+ try:
+ return csv_result_payload_response(
+ tool_name,
+ columns,
+ rows(),
+ result_store=result_store_from_export_path(export_path),
+ source={"query_sql": query_sql} if query_sql else None,
+ diagnostics_factory=diagnostics_factory,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ finally:
+ _session.close()
+
+
+def _prepare_table_res(
+ _res: SessionDataSet,
+ _table_session: TableSession,
+ tool_name: str,
+ export_path: str,
+ query_sql: str | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+) -> list[TextContent]:
+ columns = _res.get_column_names()
+
+ def rows():
+ while _res.has_next():
+ row = _res.next().get_fields()
+ yield ",".join(map(str, row))
+
+ try:
+ return csv_result_payload_response(
+ tool_name,
+ columns,
+ rows(),
+ result_store=result_store_from_export_path(export_path),
+ source={"query_sql": query_sql} if query_sql else None,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ finally:
+ _table_session.close()
+
+
+def register_query_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register query tools with per-call IoTDB target selection."""
+ _ensure_export_directory(config.export_path, logger)
+ max_pool_size = 100
+
+ @mcp.tool()
+ async def select_query(
+ query_sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+ ) -> list[TextContent]:
+ """Execute one tree-dialect read query and return rows.
+
+ Use target_id or target selector fields to choose a registered IoTDB target.
+ """
+ selected_config, session_pool = tree_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="select_query",
+ )
+ with iotdb_target_response_context(selected_config):
+ session = None
+ try:
+ stmt = query_sql.strip().upper()
+ if stmt.startswith("SELECT"):
+ assert_tree_query_shape(query_sql)
+ session = session_pool.get_session()
+ res = session.execute_query_statement(query_sql)
+ return _prepare_tree_res(
+ res,
+ session,
+ "select_query",
+ selected_config.export_path,
+ query_sql,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ raise ValueError("Only SELECT queries are allowed for select_query")
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute select query: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def export_query(
+ query_sql: str,
+ format: str = "csv",
+ filename: str = None,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute one tree-dialect read query and export its result set."""
+ selected_config, session_pool = tree_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="export_query",
+ )
+ with iotdb_target_response_context(selected_config):
+ _ensure_export_directory(selected_config.export_path, logger)
+ session = None
+ try:
+ stmt = query_sql.strip().upper()
+ if not (stmt.startswith("SELECT") or stmt.startswith("SHOW")):
+ raise ValueError(
+ "Only SELECT or SHOW queries are allowed for export"
+ )
+ if stmt.startswith("SELECT"):
+ assert_tree_query_shape(query_sql)
+
+ session = session_pool.get_session()
+ res = session.execute_query_statement(query_sql)
+ df = res.todf()
+ session.close()
+
+ timestamp = int(datetime.datetime.now().timestamp())
+ if filename is None:
+ filename = f"dump_{uuid.uuid4().hex[:4]}_{timestamp}"
+
+ if format.lower() == "csv":
+ if filename.lower().endswith(".csv"):
+ filename = filename[:-4]
+ filepath = sanitize_filename(
+ f"{filename}.csv", selected_config.export_path
+ )
+ df.to_csv(filepath, index=False)
+ elif format.lower() == "excel":
+ if filename.lower().endswith(".xlsx"):
+ filename = filename[:-5]
+ filepath = sanitize_filename(
+ f"{filename}.xlsx", selected_config.export_path
+ )
+ df.to_excel(filepath, index=False)
+ else:
+ raise ValueError("Format must be either 'csv' or 'excel'")
+
+ preview_rows = min(10, len(df))
+ preview_data = [",".join(df.columns)]
+ for i in range(preview_rows):
+ preview_data.append(",".join(map(str, df.iloc[i])))
+
+ wildcard_hint = tree_from_wildcard_runtime_hint(
+ query_sql, row_count=len(df)
+ )
+ diagnostic_text = ""
+ if wildcard_hint:
+ diagnostic_text = (
+ "\n\nDiagnostic: "
+ + wildcard_hint["message"]
+ + " "
+ + wildcard_hint["rewrite_hint"]
+ )
+
+ return text_payload_response(
+ "export_query",
+ f"Query results exported to {filepath}\n\nPreview (first {preview_rows} rows):\n"
+ + "\n".join(preview_data)
+ + diagnostic_text,
+ )
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to export query: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def read_query(
+ query_sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+ ) -> list[TextContent]:
+ """Execute one table-dialect read query and return rows."""
+ selected_config, session_pool = table_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=max_pool_size,
+ tool_name="read_query",
+ )
+ with iotdb_target_response_context(selected_config):
+ table_session = None
+ try:
+ table_session = session_pool.get_session()
+ stmt = query_sql.strip().upper()
+ if (
+ stmt.startswith("SELECT")
+ or stmt.startswith("DESCRIBE")
+ or stmt.startswith("SHOW")
+ ):
+ res = table_session.execute_query_statement(query_sql)
+ return _prepare_table_res(
+ res,
+ table_session,
+ "read_query",
+ selected_config.export_path,
+ query_sql,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ table_session.close()
+ raise ValueError("Only SELECT queries are allowed for read_query")
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to execute query: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def export_table_query(
+ query_sql: str,
+ format: str = "csv",
+ filename: str = None,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute one table-dialect read query and export its result set."""
+ selected_config, session_pool = table_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=max_pool_size,
+ tool_name="export_table_query",
+ )
+ with iotdb_target_response_context(selected_config):
+ _ensure_export_directory(selected_config.export_path, logger)
+ table_session = None
+ try:
+ table_session = session_pool.get_session()
+ stmt = query_sql.strip().upper()
+ if not (
+ stmt.startswith("SELECT")
+ or stmt.startswith("SHOW")
+ or stmt.startswith("DESCRIBE")
+ or stmt.startswith("DESC")
+ ):
+ raise ValueError(
+ "Only SELECT, SHOW or DESCRIBE queries are allowed for export"
+ )
+
+ res = table_session.execute_query_statement(query_sql)
+ df = res.todf()
+ table_session.close()
+
+ timestamp = int(datetime.datetime.now().timestamp())
+ if filename is None:
+ filename = f"dump_{uuid.uuid4().hex[:4]}_{timestamp}"
+
+ if format.lower() == "csv":
+ if filename.lower().endswith(".csv"):
+ filename = filename[:-4]
+ filepath = sanitize_filename(
+ f"{filename}.csv", selected_config.export_path
+ )
+ df.to_csv(filepath, index=False)
+ elif format.lower() == "excel":
+ if filename.lower().endswith(".xlsx"):
+ filename = filename[:-5]
+ filepath = sanitize_filename(
+ f"{filename}.xlsx", selected_config.export_path
+ )
+ df.to_excel(filepath, index=False)
+ else:
+ raise ValueError("Format must be either 'csv' or 'excel'")
+
+ preview_rows = min(10, len(df))
+ preview_data = [",".join(df.columns)]
+ for i in range(preview_rows):
+ preview_data.append(",".join(map(str, df.iloc[i])))
+
+ return text_payload_response(
+ "export_table_query",
+ f"Query results exported to {filepath}\n\nPreview (first {preview_rows} rows):\n"
+ + "\n".join(preview_data),
+ )
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to export table query: {str(e)}")
+ raise
diff --git a/src/iotdb_mcp_server/services/results.py b/src/iotdb_mcp_server/services/results.py
new file mode 100644
index 0000000..7d64f52
--- /dev/null
+++ b/src/iotdb_mcp_server/services/results.py
@@ -0,0 +1,226 @@
+#
+# 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.
+#
+
+from __future__ import annotations
+
+import logging
+
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.result_store import result_store_from_export_path
+from iotdb_mcp_server.services.json_response import payload_response
+from iotdb_mcp_server.services.target_selection import select_target_config
+
+
+def _cursor_to_offset(cursor: str | None, offset: int | None) -> int:
+ if offset is not None:
+ return max(0, int(offset))
+ if cursor is None or str(cursor).strip() == "":
+ return 0
+ cleaned = str(cursor).strip()
+ if not cleaned.isdigit():
+ raise ValueError("cursor must be a non-negative integer row offset.")
+ return int(cleaned)
+
+
+def _csv_text(columns: list[str], rows: list[str]) -> str:
+ return "\n".join([",".join(columns)] + rows)
+
+
+def register_result_store_tools(
+ mcp,
+ config: Config,
+ logger: logging.Logger,
+) -> None:
+ """Register tools for paging file-backed MCP result sets."""
+
+ @mcp.tool()
+ async def read_result_page(
+ result_id: str,
+ cursor: str | None = None,
+ offset: int | None = None,
+ limit: int | None = None,
+ owner_session_id: str | None = None,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Read one page from a previously returned ResultStore result_id.
+
+ Pass the same target_id used by the original query when multiple IoTDB
+ targets are registered. cursor is the next_cursor returned by a prior
+ page; offset can be used for random access and takes precedence.
+ """
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ store = result_store_from_export_path(selected_config.export_path)
+ try:
+ payload = store.read_page(
+ result_id,
+ offset=_cursor_to_offset(cursor, offset),
+ limit=limit,
+ owner_session_id=owner_session_id,
+ )
+ payload["target_id"] = selected_config.target_id
+ payload["text"] = _csv_text(payload["columns"], payload["rows"])
+ payload["inline_row_count"] = payload["returned_rows"]
+ payload["inline_truncated"] = payload["has_more"]
+ return payload_response(
+ "read_result_page",
+ payload,
+ message="Result page loaded from ResultStore.",
+ )
+ except Exception as e:
+ logger.error(f"Failed to read ResultStore page: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def read_result_pages(
+ pages: list[dict[str, object]],
+ owner_session_id: str | None = None,
+ default_limit: int | None = None,
+ max_pages: int | None = None,
+ max_total_rows: int | None = None,
+ continue_on_error: bool = True,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Read multiple ResultStore pages in one MCP call.
+
+ Each page request is an object with result_id plus optional cursor,
+ offset, limit, and owner_session_id. offset takes precedence over
+ cursor. owner_session_id on an item takes precedence over the shared
+ owner_session_id argument.
+ """
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ store = result_store_from_export_path(selected_config.export_path)
+ try:
+ payload = store.read_pages(
+ pages,
+ owner_session_id=owner_session_id,
+ default_limit=default_limit,
+ max_pages=max_pages,
+ max_total_rows=max_total_rows,
+ continue_on_error=continue_on_error,
+ )
+ payload["target_id"] = selected_config.target_id
+ for page in payload["pages"]:
+ if page.get("status") != "success":
+ continue
+ page["text"] = _csv_text(page["columns"], page["rows"])
+ page["inline_row_count"] = page["returned_rows"]
+ page["inline_truncated"] = page["has_more"]
+ return payload_response(
+ "read_result_pages",
+ payload,
+ message="Result pages loaded from ResultStore.",
+ )
+ except Exception as e:
+ logger.error(f"Failed to read ResultStore pages: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def list_result_store(
+ owner_session_id: str | None = None,
+ all_sessions: bool = False,
+ limit: int = 50,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """List ResultStore entries for the current session by default."""
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ store = result_store_from_export_path(selected_config.export_path)
+ try:
+ payload = store.list_results(
+ owner_session_id=owner_session_id,
+ all_sessions=all_sessions,
+ limit=limit,
+ )
+ payload["target_id"] = selected_config.target_id
+ return payload_response(
+ "list_result_store",
+ payload,
+ message="ResultStore entries listed.",
+ )
+ except Exception as e:
+ logger.error(f"Failed to list ResultStore entries: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def cleanup_result_store(
+ owner_session_id: str | None = None,
+ all_sessions: bool = False,
+ ttl_seconds: int | None = None,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Run TTL and quota cleanup for ResultStore entries."""
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ store = result_store_from_export_path(selected_config.export_path)
+ try:
+ payload = store.cleanup(
+ owner_session_id=owner_session_id,
+ all_sessions=all_sessions,
+ ttl_seconds=ttl_seconds,
+ )
+ payload["target_id"] = selected_config.target_id
+ return payload_response(
+ "cleanup_result_store",
+ payload,
+ message="ResultStore cleanup completed.",
+ )
+ except Exception as e:
+ logger.error(f"Failed to cleanup ResultStore: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def delete_result(
+ result_id: str,
+ owner_session_id: str | None = None,
+ all_sessions: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Delete one cached ResultStore result."""
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ store = result_store_from_export_path(selected_config.export_path)
+ try:
+ payload = store.delete_result(
+ result_id,
+ owner_session_id=owner_session_id,
+ all_sessions=all_sessions,
+ )
+ payload["target_id"] = selected_config.target_id
+ return payload_response(
+ "delete_result",
+ payload,
+ message="ResultStore result deleted.",
+ )
+ except Exception as e:
+ logger.error(f"Failed to delete ResultStore result: {str(e)}")
+ raise
diff --git a/src/iotdb_mcp_server/services/runtime_policy.py b/src/iotdb_mcp_server/services/runtime_policy.py
new file mode 100644
index 0000000..c768389
--- /dev/null
+++ b/src/iotdb_mcp_server/services/runtime_policy.py
@@ -0,0 +1,77 @@
+# 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.
+#
+
+from __future__ import annotations
+
+import logging
+
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.runtime_policy import dynamic_policy_snapshot
+from iotdb_mcp_server.runtime_policy import reset_session_policy
+from iotdb_mcp_server.runtime_policy import set_session_policy
+from iotdb_mcp_server.services.json_response import payload_response
+
+
+def register_runtime_policy_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register session-scoped policy controls for this MCP process."""
+
+ @mcp.tool()
+ async def get_iotdb_session_policy() -> list[TextContent]:
+ """Inspect effective IoTDB MCP permissions for this running session."""
+ return payload_response(
+ "get_iotdb_session_policy",
+ dynamic_policy_snapshot(),
+ message="IoTDB session policy snapshot.",
+ )
+
+ @mcp.tool()
+ async def set_iotdb_session_policy(
+ policy: dict[str, object] | None = None,
+ preset: str | None = None,
+ replace: bool = False,
+ ) -> list[TextContent]:
+ """
+ Set session-scoped IoTDB MCP permissions without restarting the MCP server.
+
+ preset accepts: full, ddl, readonly. Explicit policy values override preset values.
+ """
+ snapshot = set_session_policy(policy=policy, preset=preset, replace=replace)
+ logger.info(
+ "Updated IoTDB session policy preset=%s replace=%s keys=%s",
+ preset,
+ replace,
+ sorted((policy or {}).keys()),
+ )
+ return payload_response(
+ "set_iotdb_session_policy",
+ snapshot,
+ message="IoTDB session policy updated.",
+ )
+
+ @mcp.tool()
+ async def reset_iotdb_session_policy(keys: list[str] | None = None) -> list[TextContent]:
+ """Clear all or selected session-scoped IoTDB MCP permission overrides."""
+ snapshot = reset_session_policy(keys=keys)
+ logger.info("Reset IoTDB session policy keys=%s", keys or "*")
+ return payload_response(
+ "reset_iotdb_session_policy",
+ snapshot,
+ message="IoTDB session policy reset.",
+ )
diff --git a/src/iotdb_mcp_server/services/sql_driver.py b/src/iotdb_mcp_server/services/sql_driver.py
new file mode 100644
index 0000000..055c3e5
--- /dev/null
+++ b/src/iotdb_mcp_server/services/sql_driver.py
@@ -0,0 +1,1348 @@
+#
+# 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.
+#
+
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
+from functools import partial
+import logging
+import os
+import re
+import threading
+import time
+import uuid
+from datetime import datetime, timezone
+from typing import Any
+
+from iotdb.Session import Session
+from iotdb.table_session import TableSession
+from iotdb.utils.SessionDataSet import SessionDataSet
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.result_store import result_store_from_export_path
+from iotdb_mcp_server.result_store import StoredCsvResult
+from iotdb_mcp_server.runtime_policy import dynamic_env_bool
+from iotdb_mcp_server.runtime_policy import dynamic_getenv
+from iotdb_mcp_server.runtime_policy import dynamic_policy_snapshot
+from iotdb_mcp_server.runtime_policy import permission_enforcement_mode
+from iotdb_mcp_server.runtime_policy import strict_permission_enforcement
+from iotdb_mcp_server.services.json_response import (
+ csv_result_payload_response,
+ payload_response,
+)
+from iotdb_mcp_server.services.target_selection import (
+ iotdb_target_response_context,
+ select_target_config,
+ table_session_pool,
+ tree_session_pool,
+)
+from iotdb_mcp_server.services.tree_sql_guardrails import assert_tree_query_shape
+from iotdb_mcp_server.services.tree_sql_guardrails import (
+ tree_from_wildcard_runtime_hint,
+)
+
+_READONLY_PREFIXES_COMMON = (
+ "SELECT",
+ "SHOW",
+ "EXPLAIN",
+ "DESC",
+ "DESCRIBE",
+)
+_READONLY_PREFIXES_TREE = (
+ "COUNT TIMESERIES",
+ "COUNT NODES",
+ "COUNT DEVICES",
+ "CALL INFERENCE",
+ "SHOW CONTINUOUS QUERIES",
+ "SHOW CQS",
+)
+_READONLY_PREFIXES_TABLE: tuple[str, ...] = ()
+
+_DDL_PREFIXES_COMMON = (
+ "CREATE DATABASE",
+ "ALTER DATABASE",
+ "DROP DATABASE",
+ "CREATE TABLE",
+ "ALTER TABLE",
+ "DROP TABLE",
+ "USE",
+ "SET TTL",
+ "UNSET TTL",
+ "CREATE MODEL",
+ "DROP MODEL",
+ "LOAD MODEL",
+ "UNLOAD MODEL",
+ "REMOVE AINODE",
+)
+_DDL_PREFIXES_TREE = (
+ "CREATE TIMESERIES",
+ "CREATE ALIGNED TIMESERIES",
+ "ALTER TIMESERIES",
+ "DROP TIMESERIES",
+ "DELETE TIMESERIES",
+ "CREATE CONTINUOUS QUERY",
+ "CREATE CQ",
+ "DROP CONTINUOUS QUERY",
+ "DROP CQ",
+)
+_DDL_PREFIXES_TABLE: tuple[str, ...] = ()
+
+_FULL_PREFIXES_COMMON = (
+ "INSERT INTO",
+ "UPDATE",
+ "DELETE FROM",
+ "DELETE DEVICES",
+)
+_FULL_PREFIXES_TREE: tuple[str, ...] = ()
+_FULL_PREFIXES_TABLE: tuple[str, ...] = ()
+
+_DESTRUCTIVE_PREFIXES = (
+ "DROP ",
+ "DELETE ",
+ "UNSET TTL",
+ "REMOVE AINODE",
+ "UNLOAD MODEL",
+)
+
+_PERMISSION_RANK = {"readonly": 0, "ddl": 1, "full": 2}
+_TEMPLATE_FIELD_PATTERN = re.compile(
+ r"\{\{\s*([A-Za-z_][A-Za-z0-9_]*)(?::(literal|path|identifier))?\s*\}\}"
+)
+_PATH_SEGMENT_PATTERN = re.compile(r"^([A-Za-z0-9_]+|\*|\*\*)$")
+_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+_DEFAULT_BATCH_PER_ITEM_TIMEOUT_MS = 60_000
+_DEFAULT_BATCH_TIMEOUT_MS = 300_000
+_DEFAULT_BATCH_MAX_ROWS_PER_ITEM = 10_000
+_DEFAULT_BATCH_MAX_BYTES_PER_ITEM = 16 * 1024 * 1024
+_DEFAULT_BATCH_MAX_TOTAL_ROWS = 100_000
+_DEFAULT_BATCH_MAX_TOTAL_BYTES = 64 * 1024 * 1024
+_DEFAULT_BATCH_WORKER_POOL_SIZE = 4
+
+
+def _env_bool(name: str, default: bool) -> bool:
+ return dynamic_env_bool(name, default)
+
+
+def _env_int(
+ name: str, default: int, minimum: int = 1, maximum: int | None = None
+) -> int:
+ raw = (dynamic_getenv(name, "") or os.getenv(name, "") or "").strip()
+ if not raw:
+ value = default
+ else:
+ try:
+ value = int(raw)
+ except ValueError:
+ value = default
+ value = max(minimum, value)
+ if maximum is not None:
+ value = min(maximum, value)
+ return value
+
+
+def _positive_int(
+ value: int | None,
+ *,
+ env_name: str,
+ default: int,
+ minimum: int = 1,
+ maximum: int | None = None,
+) -> int:
+ if value is None:
+ return _env_int(env_name, default, minimum=minimum, maximum=maximum)
+ resolved = max(minimum, int(value))
+ if maximum is not None:
+ resolved = min(maximum, resolved)
+ return resolved
+
+
+def _csv_set(value: str) -> set[str]:
+ return {item.strip() for item in value.split(",") if item.strip()}
+
+
+def _parse_mode() -> str:
+ mode = (
+ (dynamic_getenv("IOTDB_SQL_DRIVER_MODE", "readonly") or "readonly")
+ .strip()
+ .lower()
+ )
+ if mode not in ("readonly", "ddl", "full"):
+ raise ValueError(
+ "Invalid IOTDB_SQL_DRIVER_MODE. Expected one of: readonly, ddl, full."
+ )
+ return mode
+
+
+def _strip_single_trailing_semicolon(sql: str, empty_message: str) -> str:
+ cleaned = sql.strip()
+ if not cleaned:
+ raise ValueError(empty_message)
+
+ quote: str | None = None
+ escaped = False
+ semicolon_indexes: list[int] = []
+ index = 0
+ while index < len(cleaned):
+ char = cleaned[index]
+ if escaped:
+ escaped = False
+ index += 1
+ continue
+ if char == "\\" and quote:
+ escaped = True
+ index += 1
+ continue
+ if quote:
+ if char == quote:
+ if index + 1 < len(cleaned) and cleaned[index + 1] == quote:
+ index += 2
+ continue
+ quote = None
+ index += 1
+ continue
+ if char in {"'", '"'}:
+ quote = char
+ elif char == ";":
+ semicolon_indexes.append(index)
+ index += 1
+
+ if not semicolon_indexes:
+ return cleaned
+ if semicolon_indexes == [len(cleaned) - 1]:
+ return cleaned[:-1].strip()
+ raise ValueError("Only a single SQL statement is allowed.")
+
+
+def _normalize_sql(sql: str) -> str:
+ return _strip_single_trailing_semicolon(sql, "SQL cannot be empty.")
+
+
+def _normalize_for_prefix(sql: str) -> str:
+ return re.sub(r"\s+", " ", sql).strip().upper()
+
+
+def _top_level_sql_keywords(sql: str) -> list[str]:
+ """Return unquoted SQL words outside parentheses and comments."""
+ keywords: list[str] = []
+ token: list[str] = []
+ quote: str | None = None
+ depth = 0
+ index = 0
+
+ def flush_token() -> None:
+ if token:
+ keywords.append("".join(token).upper())
+ token.clear()
+
+ while index < len(sql):
+ char = sql[index]
+ next_char = sql[index + 1] if index + 1 < len(sql) else ""
+
+ if quote is not None:
+ if char == quote:
+ if next_char == quote:
+ index += 2
+ continue
+ quote = None
+ elif char == "\\":
+ index += 2
+ continue
+ index += 1
+ continue
+
+ if char in {"'", '"', "`"}:
+ flush_token()
+ quote = char
+ index += 1
+ continue
+ if char == "-" and next_char == "-":
+ flush_token()
+ newline = sql.find("\n", index + 2)
+ index = len(sql) if newline < 0 else newline + 1
+ continue
+ if char == "/" and next_char == "*":
+ flush_token()
+ comment_end = sql.find("*/", index + 2)
+ index = len(sql) if comment_end < 0 else comment_end + 2
+ continue
+ if char == "(":
+ flush_token()
+ depth += 1
+ index += 1
+ continue
+ if char == ")":
+ flush_token()
+ depth = max(0, depth - 1)
+ index += 1
+ continue
+ if depth == 0 and (char.isalnum() or char == "_"):
+ token.append(char)
+ else:
+ flush_token()
+ index += 1
+
+ flush_token()
+ return keywords
+
+
+def _select_writes_results(sql: str) -> bool:
+ keywords = _top_level_sql_keywords(sql)
+ return bool(keywords and keywords[0] == "SELECT" and "INTO" in keywords[1:])
+
+
+def _matches_prefix(normalized_upper: str, prefix: str) -> bool:
+ if not normalized_upper.startswith(prefix):
+ return False
+ return (
+ len(normalized_upper) == len(prefix)
+ or normalized_upper[len(prefix)].isspace()
+ or normalized_upper[len(prefix)] == "("
+ )
+
+
+def _merge_prefixes(base: tuple[str, ...], env_var: str) -> tuple[str, ...]:
+ extra = tuple(
+ item.strip().upper()
+ for item in (dynamic_getenv(env_var, "") or "").split(",")
+ if item.strip()
+ )
+ all_prefixes = tuple(dict.fromkeys([*base, *extra]))
+ return tuple(sorted(all_prefixes, key=len, reverse=True))
+
+
+def _resolve_whitelists(sql_dialect: str) -> dict[str, tuple[str, ...]]:
+ readonly = _READONLY_PREFIXES_COMMON + (
+ _READONLY_PREFIXES_TREE if sql_dialect == "tree" else _READONLY_PREFIXES_TABLE
+ )
+ ddl = _DDL_PREFIXES_COMMON + (
+ _DDL_PREFIXES_TREE if sql_dialect == "tree" else _DDL_PREFIXES_TABLE
+ )
+ full = _FULL_PREFIXES_COMMON + (
+ _FULL_PREFIXES_TREE if sql_dialect == "tree" else _FULL_PREFIXES_TABLE
+ )
+
+ return {
+ "readonly": _merge_prefixes(
+ readonly, "IOTDB_SQL_DRIVER_EXTRA_READONLY_PREFIXES"
+ ),
+ "ddl": _merge_prefixes(ddl, "IOTDB_SQL_DRIVER_EXTRA_DDL_PREFIXES"),
+ "full": _merge_prefixes(full, "IOTDB_SQL_DRIVER_EXTRA_FULL_PREFIXES"),
+ }
+
+
+def _classify_sql(sql: str, whitelists: dict[str, tuple[str, ...]]) -> tuple[str, str]:
+ normalized_upper = _normalize_for_prefix(sql)
+
+ if _select_writes_results(sql):
+ return "full", "SELECT INTO"
+
+ for category in ("readonly", "ddl", "full"):
+ for prefix in whitelists[category]:
+ if _matches_prefix(normalized_upper, prefix):
+ return category, prefix
+
+ raise ValueError(
+ "SQL is not in whitelist. Allowed categories and prefixes can be inspected "
+ "with sql_driver_policy tool."
+ )
+
+
+def _is_destructive_prefix(prefix: str) -> bool:
+ upper_prefix = prefix.upper()
+ return upper_prefix.startswith(_DESTRUCTIVE_PREFIXES)
+
+
+def _sql_literal(value: Any) -> str:
+ if value is None:
+ return "NULL"
+ if isinstance(value, bool):
+ return "true" if value else "false"
+ if isinstance(value, int):
+ return str(value)
+ if isinstance(value, float):
+ if value != value or value in (float("inf"), float("-inf")):
+ raise ValueError("Non-finite float values are not valid SQL literals.")
+ return repr(value)
+ return "'" + str(value).replace("'", "''") + "'"
+
+
+def _validate_iotdb_path(value: Any) -> str:
+ text = str(value).strip()
+ if not text:
+ raise ValueError("Path placeholder value cannot be empty.")
+ if any(char in text for char in (";", "'", '"', "`", "\\", "\n", "\r", "\t", " ")):
+ raise ValueError(f"Unsafe IoTDB path placeholder value: {text!r}")
+ parts = text.split(".")
+ if not parts or any(not _PATH_SEGMENT_PATTERN.fullmatch(part) for part in parts):
+ raise ValueError(f"Invalid IoTDB path placeholder value: {text!r}")
+ return text
+
+
+def _validate_identifier(value: Any) -> str:
+ text = str(value).strip()
+ if not _IDENTIFIER_PATTERN.fullmatch(text):
+ raise ValueError(f"Invalid SQL identifier placeholder value: {text!r}")
+ return text
+
+
+def _render_sql_template(sql_template: str, params: dict[str, Any]) -> str:
+ template = str(sql_template or "").strip()
+ if not template:
+ raise ValueError("sql_template cannot be empty.")
+ seen: set[str] = set()
+
+ def replace(match: re.Match[str]) -> str:
+ name = match.group(1)
+ kind = match.group(2) or "literal"
+ if name not in params:
+ raise ValueError(f"Missing SQL template parameter: {name}")
+ seen.add(name)
+ value = params[name]
+ if kind == "literal":
+ return _sql_literal(value)
+ if kind == "path":
+ return _validate_iotdb_path(value)
+ if kind == "identifier":
+ return _validate_identifier(value)
+ raise ValueError(f"Unsupported SQL template placeholder kind: {kind}")
+
+ rendered = _TEMPLATE_FIELD_PATTERN.sub(replace, template)
+ if "{{" in rendered or "}}" in rendered:
+ raise ValueError(
+ "Invalid sql_template placeholder. Use {{name}}, {{name:path}}, or {{name:identifier}}."
+ )
+ if not seen:
+ raise ValueError(
+ "sql_template must contain at least one placeholder such as {{device:path}}."
+ )
+ return _normalize_sql(rendered)
+
+
+def _resolve_batch_sqls(
+ sqls: list[str] | None,
+ sql_template: str | None,
+ param_sets: list[dict[str, Any]] | None,
+) -> list[dict[str, Any]]:
+ direct_sqls = [str(sql) for sql in (sqls or []) if str(sql).strip()]
+ has_template = bool(sql_template and str(sql_template).strip())
+ has_param_sets = bool(param_sets)
+
+ if direct_sqls and (has_template or has_param_sets):
+ raise ValueError("Use either sqls or sql_template+param_sets, not both.")
+ if direct_sqls:
+ return [
+ {"index": index, "sql": _normalize_sql(sql), "params": None}
+ for index, sql in enumerate(direct_sqls)
+ ]
+
+ if has_template or has_param_sets:
+ if not has_template:
+ raise ValueError("sql_template is required when param_sets is provided.")
+ if not has_param_sets:
+ raise ValueError("param_sets must contain at least one parameter set.")
+ statements: list[dict[str, Any]] = []
+ for index, params in enumerate(param_sets or []):
+ if not isinstance(params, dict):
+ raise ValueError("Each param_sets item must be an object.")
+ statements.append(
+ {
+ "index": index,
+ "sql": _render_sql_template(str(sql_template), params),
+ "params": dict(params),
+ }
+ )
+ return statements
+
+ raise ValueError("Provide either sqls or sql_template+param_sets.")
+
+
+def _mode_satisfies(required_category: str, mode: str) -> bool:
+ return _PERMISSION_RANK[mode] >= _PERMISSION_RANK[required_category]
+
+
+def _sql_permission_payload(
+ sql: str,
+ selected_config: Config,
+ whitelists: dict[str, tuple[str, ...]],
+ mode: str,
+) -> dict[str, object]:
+ normalized_sql = _normalize_sql(sql)
+ category, matched_prefix = _classify_sql(normalized_sql, whitelists)
+ destructive = _is_destructive_prefix(matched_prefix)
+ destructive_confirmation_required = destructive and _env_bool(
+ "IOTDB_SQL_DRIVER_REQUIRE_DESTRUCTIVE_CONFIRM", True
+ )
+ current_satisfies = _mode_satisfies(category, mode)
+ approval_required = destructive_confirmation_required or not current_satisfies
+ approval_reasons: list[str] = []
+ if not current_satisfies:
+ approval_reasons.append(
+ f"current advisory permission is '{mode}', SQL requires '{category}'"
+ )
+ if destructive_confirmation_required:
+ approval_reasons.append("SQL is destructive and requires explicit confirmation")
+
+ return {
+ "sql": normalized_sql,
+ "sql_dialect": selected_config.sql_dialect,
+ "matched_prefix": matched_prefix,
+ "required_permission": category,
+ "current_advisory_permission": mode,
+ "current_permission_satisfies": current_satisfies,
+ "permission_enforcement": permission_enforcement_mode(),
+ "strict_enforcement": strict_permission_enforcement(),
+ "destructive": destructive,
+ "risk_level": (
+ "high" if destructive else ("medium" if category != "readonly" else "low")
+ ),
+ "confirmation_parameter": (
+ "confirm_destructive" if destructive_confirmation_required else None
+ ),
+ "approval_required": approval_required,
+ "approval_reasons": approval_reasons,
+ "agent_approval_protocol": {
+ "owner": "host agent system",
+ "before_execution": (
+ "If approval_required is true, ask the user through the host agent "
+ "question/approval mechanism before executing."
+ ),
+ "after_user_approval": (
+ "Call sql_execute with confirm_destructive=true when the "
+ "confirmation_parameter is confirm_destructive."
+ ),
+ "do_not_use_environment_as_primary_gate": True,
+ },
+ }
+
+
+def _assert_sql_driver_permission(
+ config: Config,
+ required_category: str,
+ mode: str,
+ confirm_destructive: bool,
+ matched_prefix: str,
+) -> None:
+ if strict_permission_enforcement():
+ if not _env_bool("IOTDB_ENABLE_SQL_DRIVER", False):
+ raise PermissionError(
+ "SQL driver tool is disabled by strict MCP policy. "
+ "Set IOTDB_ENABLE_SQL_DRIVER=true to enable."
+ )
+
+ allowed_users = _csv_set(
+ dynamic_getenv("IOTDB_SQL_DRIVER_ALLOWED_USERS", "root") or "root"
+ )
+ if "*" not in allowed_users and config.user not in allowed_users:
+ raise PermissionError(
+ f"Current MCP user '{config.user}' is not allowed by IOTDB_SQL_DRIVER_ALLOWED_USERS."
+ )
+
+ if not _mode_satisfies(required_category, mode):
+ raise PermissionError(
+ f"Current sql driver mode '{mode}' does not allow SQL requiring "
+ f"'{required_category}' permission ('{matched_prefix}')."
+ )
+
+ if (
+ _is_destructive_prefix(matched_prefix)
+ and _env_bool("IOTDB_SQL_DRIVER_REQUIRE_DESTRUCTIVE_CONFIRM", True)
+ and not confirm_destructive
+ ):
+ raise PermissionError(
+ "Destructive SQL requires confirm_destructive=True by server policy."
+ )
+
+
+def _format_tree_result(
+ res: SessionDataSet,
+ session: Session,
+ tool_name: str,
+ export_path: str,
+ sql: str | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+) -> list[TextContent]:
+ columns = res.get_column_names()
+
+ def rows():
+ while res.has_next():
+ record = res.next()
+ if columns and columns[0] == "Time":
+ timestamp = record.get_timestamp()
+ row = record.get_fields()
+ yield str(timestamp) + "," + ",".join(map(str, row))
+ else:
+ yield ",".join(map(str, record.get_fields()))
+
+ def diagnostics_factory(row_count: int):
+ diagnostics = []
+ if sql:
+ hint = tree_from_wildcard_runtime_hint(sql, row_count=row_count)
+ if hint:
+ diagnostics.append(hint)
+ return diagnostics
+
+ try:
+ return csv_result_payload_response(
+ tool_name,
+ columns,
+ rows(),
+ result_store=result_store_from_export_path(export_path),
+ source={"sql": sql} if sql else None,
+ diagnostics_factory=diagnostics_factory,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ finally:
+ session.close()
+
+
+def _format_table_result(
+ res: SessionDataSet,
+ table_session: TableSession,
+ tool_name: str,
+ export_path: str,
+ sql: str | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+) -> list[TextContent]:
+ columns = res.get_column_names()
+
+ def rows():
+ while res.has_next():
+ yield ",".join(map(str, res.next().get_fields()))
+
+ try:
+ return csv_result_payload_response(
+ tool_name,
+ columns,
+ rows(),
+ result_store=result_store_from_export_path(export_path),
+ source={"sql": sql} if sql else None,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ finally:
+ table_session.close()
+
+
+def _tree_csv_rows(res: SessionDataSet, columns: list[str]):
+ while res.has_next():
+ record = res.next()
+ if columns and columns[0] == "Time":
+ timestamp = record.get_timestamp()
+ row = record.get_fields()
+ yield str(timestamp) + "," + ",".join(map(str, row))
+ else:
+ yield ",".join(map(str, record.get_fields()))
+
+
+def _table_csv_rows(res: SessionDataSet):
+ while res.has_next():
+ yield ",".join(map(str, res.next().get_fields()))
+
+
+class BatchQuotaExceeded(RuntimeError):
+ pass
+
+
+class BatchItemTimeout(TimeoutError):
+ pass
+
+
+class BatchTimeout(TimeoutError):
+ pass
+
+
+@dataclass
+class _ResultQuota:
+ label: str
+ max_rows: int
+ max_bytes: int
+ rows: int = 0
+ bytes: int = 0
+ lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
+
+ def consume(self, row: str) -> None:
+ row_bytes = len(row.encode("utf-8")) + 1
+ with self.lock:
+ next_rows = self.rows + 1
+ next_bytes = self.bytes + row_bytes
+ if next_rows > self.max_rows:
+ raise BatchQuotaExceeded(
+ f"{self.label} row quota exceeded: {next_rows}>{self.max_rows}"
+ )
+ if next_bytes > self.max_bytes:
+ raise BatchQuotaExceeded(
+ f"{self.label} byte quota exceeded: {next_bytes}>{self.max_bytes}"
+ )
+ self.rows = next_rows
+ self.bytes = next_bytes
+
+
+def _quota_csv_rows(rows, *quotas: _ResultQuota):
+ for row in rows:
+ for quota in quotas:
+ quota.consume(row)
+ yield row
+
+
+def _stored_result_payload(
+ *,
+ index: int,
+ sql: str,
+ stored: StoredCsvResult,
+ columns: list[str],
+ duration_ms: int,
+ diagnostics: list[dict[str, Any]] | None = None,
+) -> dict[str, Any]:
+ preview_rows = stored.preview_rows
+ inline_truncated = stored.row_count > len(preview_rows)
+ payload: dict[str, Any] = {
+ "index": index,
+ "ok": True,
+ "status": "success",
+ "sql": sql,
+ "format": "csv",
+ "columns": columns,
+ "rows": preview_rows,
+ "row_count": stored.row_count,
+ "inline_row_count": len(preview_rows),
+ "inline_truncated": inline_truncated,
+ "result_id": stored.result_id,
+ "owner_session_id": stored.owner_session_id,
+ "next_cursor": str(len(preview_rows)) if inline_truncated else None,
+ "duration_ms": duration_ms,
+ "result_store": {
+ "type": "file",
+ "result_id": stored.result_id,
+ "owner_session_id": stored.owner_session_id,
+ "page_tool": "read_result_page",
+ "batch_page_tool": "read_result_pages",
+ "page_size_rows": stored.page_size_rows,
+ "manifest_path": stored.manifest_path,
+ "shard_count": stored.shard_count,
+ "byte_count": stored.byte_count,
+ },
+ }
+ if diagnostics:
+ payload["diagnostics"] = diagnostics
+ return payload
+
+
+def _batch_error_payload(index: int, sql: str, exc: BaseException) -> dict[str, Any]:
+ return {
+ "index": index,
+ "ok": False,
+ "status": "error",
+ "sql": sql,
+ "error_type": type(exc).__name__,
+ "error": str(exc),
+ }
+
+
+def _batch_skipped_payload(index: int, sql: str) -> dict[str, Any]:
+ return {
+ "index": index,
+ "ok": False,
+ "status": "skipped",
+ "sql": sql,
+ "error_type": "BatchAborted",
+ "error": "Skipped because continue_on_error is false and an earlier item failed.",
+ }
+
+
+def _store_tree_batch_item(
+ *,
+ session_pool: Any,
+ selected_config: Config,
+ index: int,
+ sql: str,
+ batch_id: str,
+ owner_session_id: str | None,
+ max_inline_rows: int | None,
+ page_size_rows: int | None,
+ item_quota: _ResultQuota,
+ batch_quota: _ResultQuota,
+) -> dict[str, Any]:
+ started = time.monotonic()
+ session = None
+ try:
+ session = session_pool.get_session()
+ res = session.execute_query_statement(sql)
+ columns = res.get_column_names()
+ store = result_store_from_export_path(selected_config.export_path)
+ stored = store.write_csv_result(
+ tool="sql_executor_batch",
+ columns=columns,
+ rows=_quota_csv_rows(_tree_csv_rows(res, columns), item_quota, batch_quota),
+ source={"sql": sql, "batch_id": batch_id, "index": index},
+ owner_session_id=owner_session_id,
+ preview_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ diagnostics = []
+ hint = tree_from_wildcard_runtime_hint(sql, row_count=stored.row_count)
+ if hint:
+ diagnostics.append(hint)
+ return _stored_result_payload(
+ index=index,
+ sql=sql,
+ stored=stored,
+ columns=columns,
+ duration_ms=int((time.monotonic() - started) * 1000),
+ diagnostics=diagnostics,
+ )
+ finally:
+ if session is not None:
+ session.close()
+
+
+def _store_table_batch_item(
+ *,
+ session_pool: Any,
+ selected_config: Config,
+ index: int,
+ sql: str,
+ batch_id: str,
+ owner_session_id: str | None,
+ max_inline_rows: int | None,
+ page_size_rows: int | None,
+ item_quota: _ResultQuota,
+ batch_quota: _ResultQuota,
+) -> dict[str, Any]:
+ started = time.monotonic()
+ table_session = None
+ try:
+ table_session = session_pool.get_session()
+ res = table_session.execute_query_statement(sql)
+ columns = res.get_column_names()
+ store = result_store_from_export_path(selected_config.export_path)
+ stored = store.write_csv_result(
+ tool="sql_executor_batch",
+ columns=columns,
+ rows=_quota_csv_rows(_table_csv_rows(res), item_quota, batch_quota),
+ source={"sql": sql, "batch_id": batch_id, "index": index},
+ owner_session_id=owner_session_id,
+ preview_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ return _stored_result_payload(
+ index=index,
+ sql=sql,
+ stored=stored,
+ columns=columns,
+ duration_ms=int((time.monotonic() - started) * 1000),
+ )
+ finally:
+ if table_session is not None:
+ table_session.close()
+
+
+def register_sql_driver_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register generic SQL execution tools with advisory permission metadata."""
+ logger.info("SQL driver registered; mode is resolved dynamically per tool call")
+
+ @mcp.tool()
+ async def sql_driver_policy(
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Show current sql_driver policy and statement whitelists."""
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ mode = _parse_mode()
+ whitelists = _resolve_whitelists(selected_config.sql_dialect)
+ with iotdb_target_response_context(selected_config):
+ return payload_response(
+ "sql_driver_policy",
+ {
+ "sql_dialect": selected_config.sql_dialect,
+ "mode": mode,
+ "advisory_permission": mode,
+ "enable_sql_driver": _env_bool("IOTDB_ENABLE_SQL_DRIVER", True),
+ "permission_enforcement": permission_enforcement_mode(),
+ "strict_enforcement": strict_permission_enforcement(),
+ "permission_model": {
+ "readonly": "read-only SQL such as SELECT, SHOW, DESC, EXPLAIN",
+ "ddl": "schema/model/TTL management SQL",
+ "full": "DML or data-changing SQL such as INSERT, UPDATE, DELETE",
+ "host_agent_approval": (
+ "MCP reports required_permission and risk. The host agent "
+ "asks the user through its question/approval mechanism before "
+ "executing risky SQL."
+ ),
+ },
+ "require_destructive_confirm": _env_bool(
+ "IOTDB_SQL_DRIVER_REQUIRE_DESTRUCTIVE_CONFIRM", True
+ ),
+ "runtime_policy": dynamic_policy_snapshot(),
+ "whitelists": whitelists,
+ },
+ message="SQL driver policy snapshot.",
+ )
+
+ @mcp.tool()
+ async def inspect_sql_permission(
+ sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Inspect SQL permission, risk, and approval requirements without executing it."""
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ mode = _parse_mode()
+ whitelists = _resolve_whitelists(selected_config.sql_dialect)
+ with iotdb_target_response_context(selected_config):
+ return payload_response(
+ "inspect_sql_permission",
+ _sql_permission_payload(sql, selected_config, whitelists, mode),
+ message="SQL permission inspection completed.",
+ )
+
+ @mcp.tool()
+ async def sql_executor_batch(
+ sqls: list[str] | None = None,
+ sql_template: str | None = None,
+ param_sets: list[dict[str, object]] | None = None,
+ max_concurrency: int = 4,
+ worker_pool_size: int | None = None,
+ per_item_timeout_ms: int | None = None,
+ batch_timeout_ms: int | None = None,
+ max_result_rows_per_item: int | None = None,
+ max_result_bytes_per_item: int | None = None,
+ max_batch_result_rows: int | None = None,
+ max_batch_result_bytes: int | None = None,
+ continue_on_error: bool = True,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = 5,
+ page_size_rows: int | None = None,
+ ) -> list[TextContent]:
+ """Execute many readonly SQL statements in parallel and store each result.
+
+ Use either `sqls=[...]` for explicit single statements or
+ `sql_template` with `param_sets` to generate repeated SELECT/SHOW
+ statements without asking the model to spell out each SQL. Template
+ placeholders are `{{name}}` for SQL literals, `{{name:path}}` for IoTDB
+ paths, and `{{name:identifier}}` for SQL identifiers.
+ """
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ mode = _parse_mode()
+ whitelists = _resolve_whitelists(selected_config.sql_dialect)
+ max_statements = _env_int(
+ "IOTDB_SQL_EXECUTOR_BATCH_MAX_STATEMENTS", 64, minimum=1, maximum=1000
+ )
+ max_allowed_concurrency = _env_int(
+ "IOTDB_SQL_EXECUTOR_BATCH_MAX_CONCURRENCY", 16, minimum=1, maximum=128
+ )
+ max_allowed_workers = _env_int(
+ "IOTDB_SQL_EXECUTOR_BATCH_MAX_WORKER_POOL_SIZE",
+ 16,
+ minimum=1,
+ maximum=128,
+ )
+ per_item_timeout_ms = _positive_int(
+ per_item_timeout_ms,
+ env_name="IOTDB_SQL_EXECUTOR_BATCH_PER_ITEM_TIMEOUT_MS",
+ default=_DEFAULT_BATCH_PER_ITEM_TIMEOUT_MS,
+ maximum=3_600_000,
+ )
+ batch_timeout_ms = _positive_int(
+ batch_timeout_ms,
+ env_name="IOTDB_SQL_EXECUTOR_BATCH_TIMEOUT_MS",
+ default=_DEFAULT_BATCH_TIMEOUT_MS,
+ maximum=24 * 3_600_000,
+ )
+ max_result_rows_per_item = _positive_int(
+ max_result_rows_per_item,
+ env_name="IOTDB_SQL_EXECUTOR_BATCH_MAX_RESULT_ROWS_PER_ITEM",
+ default=_DEFAULT_BATCH_MAX_ROWS_PER_ITEM,
+ maximum=10_000_000,
+ )
+ max_result_bytes_per_item = _positive_int(
+ max_result_bytes_per_item,
+ env_name="IOTDB_SQL_EXECUTOR_BATCH_MAX_RESULT_BYTES_PER_ITEM",
+ default=_DEFAULT_BATCH_MAX_BYTES_PER_ITEM,
+ maximum=10 * 1024 * 1024 * 1024,
+ )
+ max_batch_result_rows = _positive_int(
+ max_batch_result_rows,
+ env_name="IOTDB_SQL_EXECUTOR_BATCH_MAX_TOTAL_RESULT_ROWS",
+ default=_DEFAULT_BATCH_MAX_TOTAL_ROWS,
+ maximum=100_000_000,
+ )
+ max_batch_result_bytes = _positive_int(
+ max_batch_result_bytes,
+ env_name="IOTDB_SQL_EXECUTOR_BATCH_MAX_TOTAL_RESULT_BYTES",
+ default=_DEFAULT_BATCH_MAX_TOTAL_BYTES,
+ maximum=100 * 1024 * 1024 * 1024,
+ )
+ statements = _resolve_batch_sqls(
+ sqls,
+ sql_template,
+ param_sets,
+ )
+ if len(statements) > max_statements:
+ raise ValueError(
+ f"sql_executor_batch accepts at most {max_statements} statements."
+ )
+
+ for item in statements:
+ normalized_sql = item["sql"]
+ category, matched_prefix = _classify_sql(normalized_sql, whitelists)
+ if category != "readonly":
+ raise PermissionError(
+ "sql_executor_batch only supports readonly SQL. "
+ f"Statement {item['index']} requires '{category}' permission."
+ )
+ _assert_sql_driver_permission(
+ selected_config,
+ required_category=category,
+ mode=mode,
+ confirm_destructive=False,
+ matched_prefix=matched_prefix,
+ )
+ if (
+ selected_config.sql_dialect == "tree"
+ and normalized_sql.upper().startswith("SELECT")
+ ):
+ assert_tree_query_shape(normalized_sql)
+ item["matched_prefix"] = matched_prefix
+
+ batch_id = (
+ "batch_"
+ + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+ + "_"
+ + uuid.uuid4().hex[:12]
+ )
+ concurrency = min(
+ max(1, int(max_concurrency or 1)),
+ max_allowed_concurrency,
+ len(statements),
+ )
+ workers = min(
+ _positive_int(
+ worker_pool_size,
+ env_name="IOTDB_SQL_EXECUTOR_BATCH_WORKER_POOL_SIZE",
+ default=max(concurrency, _DEFAULT_BATCH_WORKER_POOL_SIZE),
+ maximum=max_allowed_workers,
+ ),
+ concurrency,
+ len(statements),
+ )
+ started = time.monotonic()
+ batch_deadline = started + (batch_timeout_ms / 1000.0)
+ stop_event = asyncio.Event()
+ semaphore = asyncio.Semaphore(concurrency)
+ batch_quota = _ResultQuota(
+ "batch",
+ max_rows=max_batch_result_rows,
+ max_bytes=max_batch_result_bytes,
+ )
+
+ with iotdb_target_response_context(selected_config):
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=concurrency,
+ wait_timeout_in_ms=selected_config.tree_wait_timeout_in_ms,
+ tool_name="sql_executor_batch",
+ )
+ store_item = _store_tree_batch_item
+ else:
+ _, session_pool = table_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=concurrency,
+ wait_timeout_in_ms=selected_config.table_wait_timeout_in_ms,
+ tool_name="sql_executor_batch",
+ )
+ store_item = _store_table_batch_item
+
+ loop = asyncio.get_running_loop()
+ executor = ThreadPoolExecutor(
+ max_workers=workers,
+ thread_name_prefix="iotdb-sql-batch",
+ )
+
+ async def execute_item(item: dict[str, Any]) -> dict[str, Any]:
+ if stop_event.is_set():
+ return _batch_skipped_payload(item["index"], item["sql"])
+ async with semaphore:
+ if stop_event.is_set():
+ return _batch_skipped_payload(item["index"], item["sql"])
+ try:
+ remaining_seconds = batch_deadline - time.monotonic()
+ if remaining_seconds <= 0:
+ raise BatchTimeout(
+ f"Batch timeout exceeded after {batch_timeout_ms} ms."
+ )
+ timeout_seconds = min(
+ per_item_timeout_ms / 1000.0,
+ remaining_seconds,
+ )
+ item_quota = _ResultQuota(
+ f"statement {item['index']}",
+ max_rows=max_result_rows_per_item,
+ max_bytes=max_result_bytes_per_item,
+ )
+ call = partial(
+ store_item,
+ session_pool=session_pool,
+ selected_config=selected_config,
+ index=item["index"],
+ sql=item["sql"],
+ batch_id=batch_id,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ item_quota=item_quota,
+ batch_quota=batch_quota,
+ )
+ return await asyncio.wait_for(
+ loop.run_in_executor(executor, call),
+ timeout=timeout_seconds,
+ )
+ except asyncio.TimeoutError as exc:
+ if not continue_on_error:
+ stop_event.set()
+ return _batch_error_payload(
+ item["index"],
+ item["sql"],
+ BatchItemTimeout(
+ f"Statement timed out after {per_item_timeout_ms} ms."
+ ),
+ )
+ except Exception as exc:
+ if not continue_on_error:
+ stop_event.set()
+ logger.error(
+ "sql_executor_batch item %s failed: %s",
+ item["index"],
+ str(exc),
+ )
+ return _batch_error_payload(item["index"], item["sql"], exc)
+
+ tasks = {
+ asyncio.create_task(execute_item(item)): item for item in statements
+ }
+ done, pending = await asyncio.wait(
+ tasks,
+ timeout=batch_timeout_ms / 1000.0,
+ )
+ results = []
+ for task in done:
+ try:
+ results.append(task.result())
+ except asyncio.CancelledError:
+ item = tasks[task]
+ results.append(
+ _batch_error_payload(
+ item["index"],
+ item["sql"],
+ BatchTimeout(
+ f"Batch timeout exceeded after {batch_timeout_ms} ms."
+ ),
+ )
+ )
+ if pending:
+ stop_event.set()
+ for task in pending:
+ item = tasks[task]
+ task.cancel()
+ results.append(
+ _batch_error_payload(
+ item["index"],
+ item["sql"],
+ BatchTimeout(
+ f"Batch timeout exceeded after {batch_timeout_ms} ms."
+ ),
+ )
+ )
+ executor.shutdown(wait=False, cancel_futures=True)
+ results.sort(key=lambda item: int(item.get("index", 0)))
+ succeeded = sum(1 for item in results if item.get("status") == "success")
+ failed = sum(1 for item in results if item.get("status") == "error")
+ skipped = sum(1 for item in results if item.get("status") == "skipped")
+ payload = {
+ "batch_id": batch_id,
+ "sql_dialect": selected_config.sql_dialect,
+ "target_id": selected_config.target_id,
+ "statement_count": len(statements),
+ "succeeded": succeeded,
+ "failed": failed,
+ "skipped": skipped,
+ "continue_on_error": continue_on_error,
+ "max_concurrency": concurrency,
+ "worker_pool_size": workers,
+ "duration_ms": int((time.monotonic() - started) * 1000),
+ "generated_from_template": bool(sql_template),
+ "execution_controls": {
+ "per_item_timeout_ms": per_item_timeout_ms,
+ "batch_timeout_ms": batch_timeout_ms,
+ "timeout_cancellation": "best_effort",
+ "max_result_rows_per_item": max_result_rows_per_item,
+ "max_result_bytes_per_item": max_result_bytes_per_item,
+ "max_batch_result_rows": max_batch_result_rows,
+ "max_batch_result_bytes": max_batch_result_bytes,
+ "batch_rows_consumed": batch_quota.rows,
+ "batch_bytes_consumed": batch_quota.bytes,
+ },
+ "result_store": {
+ "type": "file",
+ "page_tool": "read_result_page",
+ "batch_page_tool": "read_result_pages",
+ "owner_session_id": owner_session_id,
+ },
+ "results": results,
+ }
+ return payload_response(
+ "sql_executor_batch",
+ payload,
+ message="Batch SQL execution completed.",
+ )
+
+ @mcp.tool()
+ async def sql_execute(
+ sql: str,
+ confirm_destructive: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+ ) -> list[TextContent]:
+ """Execute exactly one SQL statement on the selected IoTDB target.
+
+ The selected target's dialect determines the whitelist and session pool.
+ """
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ mode = _parse_mode()
+ whitelists = _resolve_whitelists(selected_config.sql_dialect)
+ with iotdb_target_response_context(selected_config):
+ normalized_sql = _normalize_sql(sql)
+ category, matched_prefix = _classify_sql(normalized_sql, whitelists)
+ _assert_sql_driver_permission(
+ selected_config,
+ required_category=category,
+ mode=mode,
+ confirm_destructive=confirm_destructive,
+ matched_prefix=matched_prefix,
+ )
+
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=100,
+ wait_timeout_in_ms=5000,
+ tool_name="sql_execute",
+ )
+ session = None
+ try:
+ if category == "readonly":
+ if normalized_sql.upper().startswith("SELECT"):
+ assert_tree_query_shape(normalized_sql)
+ session = session_pool.get_session()
+ res = session.execute_query_statement(normalized_sql)
+ return _format_tree_result(
+ res,
+ session,
+ "sql_execute",
+ selected_config.export_path,
+ normalized_sql,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+
+ session = session_pool.get_session()
+ session.execute_non_query_statement(normalized_sql)
+ session.close()
+ return payload_response(
+ "sql_execute",
+ {
+ "sql": normalized_sql,
+ "category": category,
+ "mode": mode,
+ "matched_prefix": matched_prefix,
+ "result": "success",
+ },
+ message="SQL executed successfully.",
+ )
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute sql_execute (tree): {str(e)}")
+ raise
+
+ _, session_pool = table_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=100,
+ tool_name="sql_execute",
+ )
+ table_session = None
+ try:
+ table_session = session_pool.get_session()
+ if category == "readonly":
+ res = table_session.execute_query_statement(normalized_sql)
+ return _format_table_result(
+ res,
+ table_session,
+ "sql_execute",
+ selected_config.export_path,
+ normalized_sql,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+
+ table_session.execute_non_query_statement(normalized_sql)
+ table_session.close()
+ return payload_response(
+ "sql_execute",
+ {
+ "sql": normalized_sql,
+ "category": category,
+ "mode": mode,
+ "matched_prefix": matched_prefix,
+ "result": "success",
+ },
+ message="SQL executed successfully.",
+ )
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to execute sql_execute (table): {str(e)}")
+ raise
diff --git a/src/iotdb_mcp_server/services/table.py b/src/iotdb_mcp_server/services/table.py
new file mode 100644
index 0000000..958a05f
--- /dev/null
+++ b/src/iotdb_mcp_server/services/table.py
@@ -0,0 +1,178 @@
+#
+# 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.
+#
+
+import logging
+
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.runtime_policy import dynamic_env_bool
+from iotdb_mcp_server.runtime_policy import dynamic_getenv
+from iotdb_mcp_server.runtime_policy import strict_permission_enforcement
+from iotdb_mcp_server.services.json_response import sql_success_response
+from iotdb_mcp_server.services.target_selection import (
+ iotdb_target_response_context,
+ table_session_pool,
+)
+
+
+def _env_bool(name: str, default: bool) -> bool:
+ return dynamic_env_bool(name, default)
+
+
+def _csv_set(value: str) -> set[str]:
+ return {item.strip() for item in value.split(",") if item.strip()}
+
+
+def _assert_table_ddl_permission(
+ config: Config, action: str, confirm: bool = False
+) -> None:
+ if strict_permission_enforcement():
+ if not _env_bool("IOTDB_ENABLE_TABLE_DDL", False):
+ raise PermissionError(
+ "Table DDL tools are disabled by strict MCP policy. "
+ "Set IOTDB_ENABLE_TABLE_DDL=true to enable."
+ )
+
+ allowed_users = _csv_set(dynamic_getenv("IOTDB_TABLE_DDL_ALLOWED_USERS", "root") or "root")
+ if "*" not in allowed_users and config.user not in allowed_users:
+ raise PermissionError(
+ f"Current MCP user '{config.user}' is not allowed by IOTDB_TABLE_DDL_ALLOWED_USERS."
+ )
+
+ if (
+ action == "DROP"
+ and _env_bool("IOTDB_REQUIRE_TABLE_DROP_CONFIRM", True)
+ and not confirm
+ ):
+ raise PermissionError("DROP TABLE requires confirm=True by server policy.")
+
+
+def _normalize_sql(sql: str) -> str:
+ cleaned = sql.strip()
+ if not cleaned:
+ raise ValueError("DDL SQL cannot be empty.")
+ if cleaned.endswith(";"):
+ cleaned = cleaned[:-1].strip()
+ if ";" in cleaned:
+ raise ValueError("Only a single SQL statement is allowed.")
+ return cleaned
+
+
+def _ensure_prefix(sql: str, allowed_prefixes: tuple[str, ...], action_name: str) -> None:
+ upper = sql.upper()
+ if not upper.startswith(allowed_prefixes):
+ raise ValueError(
+ f"{action_name} only supports SQL starting with: {', '.join(allowed_prefixes)}"
+ )
+
+
+def register_table_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register table-model table-level DDL tools."""
+
+ @mcp.tool()
+ async def create_table_ddl(
+ ddl_sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute table-model CREATE TABLE DDL."""
+ selected_config, session_pool = table_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=100,
+ tool_name="create_table_ddl",
+ )
+ with iotdb_target_response_context(selected_config):
+ table_session = None
+ try:
+ _assert_table_ddl_permission(selected_config, action="CREATE")
+ sql = _normalize_sql(ddl_sql)
+ _ensure_prefix(sql, ("CREATE TABLE",), "create_table_ddl")
+ table_session = session_pool.get_session()
+ table_session.execute_non_query_statement(sql)
+ table_session.close()
+ return sql_success_response("create_table_ddl", sql)
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to execute create_table_ddl: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def alter_table_ddl(
+ ddl_sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute table-model ALTER TABLE DDL."""
+ selected_config, session_pool = table_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=100,
+ tool_name="alter_table_ddl",
+ )
+ with iotdb_target_response_context(selected_config):
+ table_session = None
+ try:
+ _assert_table_ddl_permission(selected_config, action="ALTER")
+ sql = _normalize_sql(ddl_sql)
+ _ensure_prefix(sql, ("ALTER TABLE",), "alter_table_ddl")
+ table_session = session_pool.get_session()
+ table_session.execute_non_query_statement(sql)
+ table_session.close()
+ return sql_success_response("alter_table_ddl", sql)
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to execute alter_table_ddl: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def drop_table_ddl(
+ ddl_sql: str,
+ confirm: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute table-model DROP TABLE DDL (destructive)."""
+ selected_config, session_pool = table_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=100,
+ tool_name="drop_table_ddl",
+ )
+ with iotdb_target_response_context(selected_config):
+ table_session = None
+ try:
+ _assert_table_ddl_permission(selected_config, action="DROP", confirm=confirm)
+ sql = _normalize_sql(ddl_sql)
+ _ensure_prefix(sql, ("DROP TABLE",), "drop_table_ddl")
+ table_session = session_pool.get_session()
+ table_session.execute_non_query_statement(sql)
+ table_session.close()
+ return sql_success_response("drop_table_ddl", sql)
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to execute drop_table_ddl: {str(e)}")
+ raise
diff --git a/src/iotdb_mcp_server/services/target_selection.py b/src/iotdb_mcp_server/services/target_selection.py
new file mode 100644
index 0000000..e07aadd
--- /dev/null
+++ b/src/iotdb_mcp_server/services/target_selection.py
@@ -0,0 +1,492 @@
+from __future__ import annotations
+
+import copy
+from contextlib import contextmanager
+import os
+from pathlib import Path
+import time
+from typing import Any, Iterator, Mapping
+
+from iotdb.Session import Session
+from iotdb.table_session import TableSession, TableSessionConfig
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.session_manager import (
+ create_table_session_pool,
+ create_tree_session_pool,
+)
+from iotdb_mcp_server.target_registry import IoTDBTargetRegistry
+from iotdb_mcp_server.services.json_response import (
+ reset_iotdb_target_context,
+ set_iotdb_target_context,
+)
+
+
+TargetSelector = str | Mapping[str, Any] | None
+_AINODE_PROBE_CACHE: dict[str, dict[str, Any]] = {}
+_DEFAULT_AINODE_PROBE_TTL_SECONDS = 10.0
+
+
+def _copy_target_fields(config: Config, source: Config) -> None:
+ for field_name in (
+ "host",
+ "port",
+ "user",
+ "password",
+ "database",
+ "sql_dialect",
+ "timezone",
+ "export_path",
+ "target_id",
+ "display_name",
+ "target_kind",
+ "node_urls",
+ "iotdb_home",
+ "use_ssl",
+ "ca_certs",
+ "connection_timeout_in_ms",
+ "enable_redirection",
+ "enable_compression",
+ "fetch_size",
+ "max_retry",
+ "max_pool_size",
+ "tree_wait_timeout_in_ms",
+ "table_wait_timeout_in_ms",
+ "target_fingerprint",
+ "policy",
+ "last_known_good_credential",
+ "verified_at",
+ ):
+ setattr(config, field_name, getattr(source, field_name))
+
+
+def apply_target_registry(
+ config: Config,
+ registry: IoTDBTargetRegistry,
+ *,
+ active_target_id: str | None = None,
+ targets_file_mtime_ns: int | None = None,
+ close_pools: bool = True,
+) -> Config:
+ if not registry.list_targets(include_secret=False):
+ config.target_registry = registry
+ config.target_id = ""
+ config.display_name = ""
+ config.target_kind = ""
+ config.host = ""
+ config.port = 0
+ config.node_urls = ()
+ config.user = ""
+ config.password = ""
+ config.database = ""
+ config.target_fingerprint = ""
+ config.policy = {}
+ config.last_known_good_credential = {}
+ config.verified_at = ""
+ if config.session_manager is not None:
+ config.session_manager.update_registry(registry, close_pools=close_pools)
+ return config
+ active_target = registry.resolve(active_target_id or registry.default_target_id)
+ selected = Config.from_target(
+ active_target,
+ target_registry=registry,
+ session_manager=config.session_manager,
+ )
+ _copy_target_fields(config, selected)
+ config.target_registry = registry
+ if targets_file_mtime_ns is not None:
+ config.targets_file_mtime_ns = targets_file_mtime_ns
+ if config.session_manager is not None:
+ update = getattr(config.session_manager, "update_registry", None)
+ if update is not None:
+ update(registry, close_pools=close_pools)
+ else:
+ config.session_manager.registry = registry
+ return config
+
+
+def maybe_reload_target_registry(config: Config) -> bool:
+ targets_file = (config.targets_file or os.getenv("TIMESEEK_IOTDB_TARGETS_FILE", "")).strip()
+ if not targets_file:
+ return False
+ path = Path(targets_file).expanduser()
+ try:
+ mtime_ns = path.stat().st_mtime_ns
+ except OSError:
+ return False
+ if config.targets_file_mtime_ns == mtime_ns:
+ return False
+ registry = IoTDBTargetRegistry.from_env(
+ targets_file=str(path),
+ targets_json=config.targets_json or os.getenv("TIMESEEK_IOTDB_TARGETS_JSON", "") or None,
+ )
+ apply_target_registry(
+ config,
+ registry,
+ targets_file_mtime_ns=mtime_ns,
+ close_pools=True,
+ )
+ return True
+
+
+def _string_fields(row: Any) -> list[str]:
+ fields = row.get_fields() if hasattr(row, "get_fields") else row
+ if isinstance(fields, (list, tuple)):
+ return [str(field) for field in fields]
+ return [str(fields)]
+
+
+def _row_dict(columns: list[str], row: list[str]) -> dict[str, str] | dict[str, list[str]]:
+ if len(columns) == len(row):
+ return dict(zip(columns, row))
+ return {"values": row}
+
+
+def _has_running_ainode(columns: list[str], rows: list[list[str]]) -> bool:
+ status_index = next(
+ (index for index, column in enumerate(columns) if column.strip().lower() == "status"),
+ None,
+ )
+ if status_index is not None:
+ return any(
+ len(row) > status_index and row[status_index].strip().lower() == "running"
+ for row in rows
+ )
+ return any(any(value.strip().lower() == "running" for value in row) for row in rows)
+
+
+def _read_dataset(res: Any) -> tuple[list[str], list[list[str]]]:
+ columns = [str(column) for column in res.get_column_names()]
+ rows: list[list[str]] = []
+ while res.has_next():
+ rows.append(_string_fields(res.next()))
+ return columns, rows
+
+
+def _ainode_probe_ttl_seconds() -> float:
+ raw = os.getenv("TIMESEEK_AINODE_PROBE_TTL_SECONDS", "").strip()
+ if not raw:
+ return _DEFAULT_AINODE_PROBE_TTL_SECONDS
+ try:
+ return max(0.0, float(raw))
+ except ValueError:
+ return _DEFAULT_AINODE_PROBE_TTL_SECONDS
+
+
+def _ainode_probe_cache_key(config: Config) -> str:
+ try:
+ return config.to_target().auth_fingerprint()
+ except Exception:
+ return config.target_fingerprint or f"{config.target_id}:{config.host}:{config.port}"
+
+
+def _cached_ainode_probe_result(
+ entry: dict[str, Any], age_seconds: float, ttl_seconds: float
+) -> dict[str, Any]:
+ result = copy.deepcopy(entry["result"])
+ source = result.get("ainode_availability_source")
+ if isinstance(source, dict):
+ source["cached"] = True
+ source["cache_age_ms"] = int(max(0.0, age_seconds) * 1000)
+ source["cache_ttl_seconds"] = ttl_seconds
+ return result
+
+
+def _last_success_for_stale_result(
+ cache_key: str, error_source: dict[str, Any]
+) -> dict[str, Any] | None:
+ cached = _AINODE_PROBE_CACHE.get(cache_key)
+ if not cached:
+ return None
+ age_seconds = time.monotonic() - float(cached["monotonic"])
+ cached_result = cached["result"]
+ source = {
+ **error_source,
+ "stale_from_last_success": True,
+ "cache_age_ms": int(max(0.0, age_seconds) * 1000),
+ "last_success": copy.deepcopy(
+ cached_result.get("ainode_availability_source")
+ ),
+ }
+ return {
+ "ainode_available": cached_result.get("ainode_available", "unknown"),
+ "ainode_availability_source": source,
+ }
+
+
+def _open_ainode_probe_session(config: Config):
+ """Open a short-lived direct session for AINode discovery.
+
+ This intentionally avoids the shared SessionPool used by business tools.
+ The IoTDB Python SessionPool requires `put_back()`, while many call sites
+ close raw sessions; a one-session probe pool can therefore time out under
+ repeated response annotation. A direct session keeps probe health isolated.
+ """
+ if config.sql_dialect == "tree":
+ session = Session.init_from_node_urls(
+ list(config.node_urls),
+ config.user,
+ config.password,
+ config.fetch_size,
+ config.timezone,
+ enable_redirection=config.enable_redirection,
+ use_ssl=config.use_ssl,
+ ca_certs=config.ca_certs or None,
+ connection_timeout_in_ms=config.connection_timeout_in_ms,
+ )
+ session.sql_dialect = "tree"
+ session.database = config.database
+ session.open(config.enable_compression)
+ return session
+
+ return TableSession(
+ TableSessionConfig(
+ node_urls=list(config.node_urls),
+ username=config.user,
+ password=config.password,
+ database=config.database or None,
+ fetch_size=config.fetch_size,
+ time_zone=config.timezone,
+ enable_redirection=config.enable_redirection,
+ enable_compression=config.enable_compression,
+ use_ssl=config.use_ssl,
+ ca_certs=config.ca_certs or None,
+ connection_timeout_in_ms=config.connection_timeout_in_ms,
+ )
+ )
+
+
+def _probe_ainode_availability(config: Config) -> dict[str, Any]:
+ session = None
+ cache_key = _ainode_probe_cache_key(config)
+ ttl_seconds = _ainode_probe_ttl_seconds()
+ cached = _AINODE_PROBE_CACHE.get(cache_key)
+ if cached and ttl_seconds > 0:
+ age_seconds = time.monotonic() - float(cached["monotonic"])
+ if age_seconds <= ttl_seconds:
+ return _cached_ainode_probe_result(cached, age_seconds, ttl_seconds)
+
+ try:
+ session = _open_ainode_probe_session(config)
+ res = session.execute_query_statement("SHOW AINODES")
+ columns, raw_rows = _read_dataset(res)
+ running_count = sum(
+ 1
+ for row in raw_rows
+ if _has_running_ainode(columns, [row])
+ )
+ result = {
+ "ainode_available": "yes" if running_count > 0 else "no",
+ "ainode_availability_source": {
+ "method": "SHOW AINODES",
+ "ok": True,
+ "columns": columns,
+ "rows": [_row_dict(columns, row) for row in raw_rows],
+ "row_count": len(raw_rows),
+ "running_count": running_count,
+ },
+ }
+ _AINODE_PROBE_CACHE[cache_key] = {
+ "monotonic": time.monotonic(),
+ "result": copy.deepcopy(result),
+ }
+ return result
+ except Exception as exc:
+ error_source = {
+ "method": "SHOW AINODES",
+ "ok": False,
+ "error_type": type(exc).__name__,
+ "error": str(exc),
+ }
+ stale = _last_success_for_stale_result(cache_key, error_source)
+ if stale:
+ return stale
+ return {
+ "ainode_available": "unknown",
+ "ainode_availability_source": error_source,
+ }
+ finally:
+ if session is not None:
+ close = getattr(session, "close", None)
+ if close is not None:
+ close()
+
+
+def _target_payload(config: Config) -> dict[str, Any]:
+ payload = {
+ "target_id": config.target_id,
+ "display_name": config.display_name,
+ "kind": config.target_kind,
+ "host": config.host,
+ "port": config.port,
+ "node_urls": list(config.node_urls),
+ "database": config.database,
+ "sql_dialect": config.sql_dialect,
+ "timezone": config.timezone,
+ "policy": dict(config.policy),
+ "fingerprint": config.target_fingerprint,
+ "ainode_available": "unknown",
+ "ainode_availability_source": {"method": "not_probed", "ok": None},
+ }
+ return payload
+
+
+def select_target_config(
+ config: Config,
+ target_id: str | None = None,
+ target: TargetSelector = None,
+ required_sql_dialect: str | None = None,
+ tool_name: str | None = None,
+) -> Config:
+ maybe_reload_target_registry(config)
+ registry = config.target_registry
+ if registry is None:
+ if target_id and target_id != config.target_id:
+ raise KeyError(
+ f"IoTDB target '{target_id}' is not registered in this MCP server."
+ )
+ if isinstance(target, str) and target != config.target_id:
+ raise KeyError(
+ f"IoTDB target '{target}' is not registered in this MCP server."
+ )
+ if isinstance(target, Mapping):
+ requested_id = target.get("target_id") or target.get("id") or target.get("name")
+ if requested_id and str(requested_id) != config.target_id:
+ raise KeyError(
+ f"IoTDB target '{requested_id}' is not registered in this MCP server."
+ )
+ candidate = _target_payload(config)
+ candidate["target_kind"] = candidate["kind"]
+ candidate["dialect"] = candidate["sql_dialect"]
+ for key in (
+ "display_name",
+ "kind",
+ "target_kind",
+ "host",
+ "port",
+ "database",
+ "sql_dialect",
+ "dialect",
+ ):
+ expected = target.get(key)
+ if expected is not None and str(candidate[key]) != str(expected):
+ raise KeyError(
+ "IoTDB target selector does not match the active target."
+ )
+ endpoint = target.get("node_url") or target.get("endpoint")
+ if endpoint and str(endpoint).strip() not in config.node_urls:
+ raise KeyError(
+ "IoTDB target selector does not match the active target."
+ )
+ raw_node_urls = (
+ target.get("node_urls")
+ or target.get("endpoints")
+ or target.get("nodes")
+ )
+ if raw_node_urls is not None:
+ if isinstance(raw_node_urls, (list, tuple)):
+ requested_node_urls = [
+ str(item).strip() for item in raw_node_urls if str(item).strip()
+ ]
+ else:
+ requested_node_urls = [
+ item.strip()
+ for item in str(raw_node_urls).split(",")
+ if item.strip()
+ ]
+ if not all(item in config.node_urls for item in requested_node_urls):
+ raise KeyError(
+ "IoTDB target selector does not match the active target."
+ )
+ selected = config
+ else:
+ selected_target = registry.resolve_selector(
+ target,
+ target_id=target_id or (None if target is not None else config.target_id),
+ )
+ selected = Config.from_target(
+ selected_target,
+ target_registry=config.target_registry,
+ session_manager=config.session_manager,
+ )
+ # Nested pool helpers select again. Preserve the reload watermark so a
+ # derived explicit-target config is not reset to the registry default.
+ selected.targets_file = config.targets_file
+ selected.targets_json = config.targets_json
+ selected.targets_file_mtime_ns = config.targets_file_mtime_ns
+
+ if required_sql_dialect and selected.sql_dialect != required_sql_dialect:
+ owner = f"{tool_name} " if tool_name else ""
+ raise ValueError(
+ f"{owner}requires {required_sql_dialect} SQL dialect, but IoTDB target "
+ f"'{selected.target_id}' uses {selected.sql_dialect}."
+ )
+ return selected
+
+
+def tree_session_pool(
+ config: Config,
+ target_id: str | None = None,
+ target: TargetSelector = None,
+ max_pool_size: int | None = None,
+ wait_timeout_in_ms: int | None = None,
+ tool_name: str | None = None,
+):
+ selected = select_target_config(
+ config,
+ target_id=target_id,
+ target=target,
+ required_sql_dialect="tree",
+ tool_name=tool_name,
+ )
+ if selected.session_manager is not None:
+ return selected, selected.session_manager.tree_pool(
+ selected.target_id,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=wait_timeout_in_ms,
+ )
+ return selected, create_tree_session_pool(
+ selected,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=wait_timeout_in_ms,
+ )
+
+
+def table_session_pool(
+ config: Config,
+ target_id: str | None = None,
+ target: TargetSelector = None,
+ max_pool_size: int | None = None,
+ database: str | None = None,
+ wait_timeout_in_ms: int | None = None,
+ tool_name: str | None = None,
+):
+ selected = select_target_config(
+ config,
+ target_id=target_id,
+ target=target,
+ required_sql_dialect="table",
+ tool_name=tool_name,
+ )
+ if selected.session_manager is not None:
+ return selected, selected.session_manager.table_pool(
+ selected.target_id,
+ max_pool_size=max_pool_size,
+ database=database,
+ wait_timeout_in_ms=wait_timeout_in_ms,
+ )
+ return selected, create_table_session_pool(
+ selected,
+ max_pool_size=max_pool_size,
+ database=database,
+ wait_timeout_in_ms=wait_timeout_in_ms,
+ )
+
+
+@contextmanager
+def iotdb_target_response_context(config: Config) -> Iterator[None]:
+ token = set_iotdb_target_context(_target_payload(config))
+ try:
+ yield
+ finally:
+ reset_iotdb_target_context(token)
diff --git a/src/iotdb_mcp_server/services/targets.py b/src/iotdb_mcp_server/services/targets.py
new file mode 100644
index 0000000..3e51e52
--- /dev/null
+++ b/src/iotdb_mcp_server/services/targets.py
@@ -0,0 +1,533 @@
+#
+# 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.
+#
+
+import logging
+import json
+import os
+from pathlib import Path
+import tempfile
+
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.services.json_response import payload_response
+from iotdb_mcp_server.services.target_selection import (
+ apply_target_registry,
+ iotdb_target_response_context,
+ maybe_reload_target_registry,
+ select_target_config,
+)
+from iotdb_mcp_server.target_registry import (
+ IoTDBTargetRegistry,
+ mark_target_last_known_good,
+ target_from_mapping,
+)
+from iotdb_mcp_server.target_lifecycle import (
+ IoTDBTargetCandidateStore,
+ candidate_target,
+ same_target_identity,
+ verify_target_once,
+)
+
+
+def _registry(config: Config) -> IoTDBTargetRegistry:
+ maybe_reload_target_registry(config)
+ if config.target_registry is not None:
+ return config.target_registry
+ registry = IoTDBTargetRegistry({})
+ config.target_registry = registry
+ return registry
+
+
+def _targets_file(config: Config, targets_file: str | None = None) -> Path:
+ resolved = (
+ targets_file
+ or config.targets_file
+ or os.getenv("TIMESEEK_IOTDB_TARGETS_FILE", "")
+ ).strip()
+ if not resolved:
+ raise ValueError(
+ "No IoTDB targets file is configured. Pass targets_file or set "
+ "TIMESEEK_IOTDB_TARGETS_FILE to persist dynamic target changes."
+ )
+ return Path(resolved).expanduser()
+
+
+def _persist_registry(
+ config: Config, registry: IoTDBTargetRegistry, targets_file: str | None
+) -> None:
+ path = _targets_file(config, targets_file)
+ path.parent.mkdir(parents=True, exist_ok=True)
+ payload = (
+ json.dumps(registry.as_dict(include_secret=True), ensure_ascii=False, indent=2)
+ + "\n"
+ )
+ temporary_path: Path | None = None
+ try:
+ with tempfile.NamedTemporaryFile(
+ mode="w",
+ encoding="utf-8",
+ dir=path.parent,
+ prefix=f".{path.name}.",
+ suffix=".tmp",
+ delete=False,
+ ) as handle:
+ handle.write(payload)
+ handle.flush()
+ os.fsync(handle.fileno())
+ temporary_path = Path(handle.name)
+ temporary_path.chmod(0o600)
+ os.replace(temporary_path, path)
+ finally:
+ if temporary_path is not None and temporary_path.exists():
+ temporary_path.unlink()
+ try:
+ path.chmod(0o600)
+ except OSError:
+ pass
+ config.targets_file = str(path)
+ try:
+ config.targets_file_mtime_ns = path.stat().st_mtime_ns
+ except OSError:
+ config.targets_file_mtime_ns = None
+
+
+def _targets_payload(config: Config) -> dict[str, object]:
+ registry = _registry(config)
+ return {
+ "default_target_id": registry.default_target_id,
+ "active_target_id": config.target_id or None,
+ "targets": registry.list_targets(include_secret=False),
+ "selector_fields": [
+ "target_id",
+ "id",
+ "display_name",
+ "kind",
+ "host",
+ "port",
+ "database",
+ "sql_dialect",
+ "user",
+ "node_url",
+ "node_urls",
+ ],
+ "dynamic_tools": [
+ "prepare_iotdb_target",
+ "connect_iotdb_target",
+ "register_iotdb_target",
+ "remove_iotdb_target",
+ "set_default_iotdb_target",
+ "reload_iotdb_targets",
+ ],
+ "targets_file": config.targets_file
+ or os.getenv("TIMESEEK_IOTDB_TARGETS_FILE", ""),
+ }
+
+
+def register_target_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register IoTDB target discovery tools."""
+ candidate_store = IoTDBTargetCandidateStore()
+
+ def publish_target(
+ target,
+ *,
+ set_default: bool,
+ persist: bool | None,
+ targets_file: str | None,
+ ):
+ registry = _registry(config)
+ updated = registry
+ duplicate_ids = []
+ for item in registry.list_targets(include_secret=True):
+ existing = target_from_mapping(item)
+ if existing.target_id != target.target_id and same_target_identity(
+ existing, target
+ ):
+ duplicate_ids.append(existing.target_id)
+ for duplicate_id in duplicate_ids:
+ updated = updated.without_target(duplicate_id)
+ make_default = (
+ set_default
+ or updated.default_target_id is None
+ or registry.default_target_id in duplicate_ids
+ )
+ updated = updated.with_target(
+ target,
+ default_target_id=(
+ target.target_id if make_default else updated.default_target_id
+ ),
+ )
+ persist_resolved = (
+ persist
+ if persist is not None
+ else bool(
+ targets_file
+ or config.targets_file
+ or os.getenv("TIMESEEK_IOTDB_TARGETS_FILE", "")
+ )
+ )
+ if persist_resolved:
+ _persist_registry(config, updated, targets_file)
+ apply_target_registry(
+ config,
+ updated,
+ active_target_id=(
+ target.target_id if make_default else updated.default_target_id
+ ),
+ close_pools=True,
+ )
+ return updated, duplicate_ids, persist_resolved
+
+ def evict_on_connection_error(
+ target_id: str,
+ error: BaseException,
+ updated: IoTDBTargetRegistry,
+ ) -> None:
+ apply_target_registry(
+ config,
+ updated,
+ active_target_id=updated.default_target_id,
+ close_pools=False,
+ )
+ if config.targets_file or os.getenv("TIMESEEK_IOTDB_TARGETS_FILE", ""):
+ try:
+ _persist_registry(config, updated, None)
+ except Exception as persist_error:
+ logger.error(
+ "Removed failed IoTDB target %s in memory but could not persist: %s",
+ target_id,
+ persist_error,
+ )
+ logger.error(
+ "Removed IoTDB target %s after connection error %s: %s",
+ target_id,
+ type(error).__name__,
+ error,
+ )
+
+ if config.session_manager is not None:
+ config.session_manager.set_connection_error_callback(evict_on_connection_error)
+
+ @mcp.tool()
+ async def list_iotdb_targets() -> list[TextContent]:
+ """List registered IoTDB targets with redacted credentials."""
+ return payload_response(
+ "list_iotdb_targets",
+ _targets_payload(config),
+ message="Registered IoTDB targets.",
+ )
+
+ @mcp.tool()
+ async def describe_iotdb_target(
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Describe one IoTDB target selected by id or selector fields."""
+ selected_config = select_target_config(
+ config, target_id=target_id, target=target
+ )
+ with iotdb_target_response_context(selected_config):
+ return payload_response(
+ "describe_iotdb_target",
+ selected_config.to_target().as_dict(include_secret=False),
+ message="IoTDB target description.",
+ )
+
+ @mcp.tool()
+ async def prepare_iotdb_target(
+ target_template: dict[str, object],
+ user_confirmed_retry: bool = False,
+ ) -> list[TextContent]:
+ """Create a credential-free, non-routable target candidate.
+
+ A previous failed endpoint requires an explicit user-requested retry.
+ Templates are memory-only and never appear in list_iotdb_targets.
+ """
+ candidate = candidate_store.prepare(
+ target_template,
+ user_confirmed_retry=user_confirmed_retry,
+ )
+ return payload_response(
+ "prepare_iotdb_target",
+ candidate.as_dict(),
+ message="IoTDB target candidate prepared; credentials are still required.",
+ )
+
+ @mcp.tool()
+ async def connect_iotdb_target(
+ candidate_id: str,
+ username: str,
+ password: str,
+ set_default: bool = True,
+ persist: bool | None = None,
+ targets_file: str | None = None,
+ ) -> list[TextContent]:
+ """Authenticate once and publish only after success.
+
+ When a targets file is configured, the verified credential is persisted
+ by default so MCP and iotdb-target-cli read the same target state.
+ """
+ candidate = candidate_store.consume(candidate_id)
+ target = candidate_target(candidate, user=username, password=password)
+ try:
+ verified = verify_target_once(target)
+ except Exception as exc:
+ candidate_store.record_failure(candidate)
+ logger.error(
+ "IoTDB candidate %s was deleted after connection failure: %s",
+ candidate_id,
+ exc,
+ )
+ raise ConnectionError(
+ "IoTDB connection failed; the one-time candidate was deleted. "
+ "Do not retry until the user explicitly requests reconnection. "
+ f"Cause: {type(exc).__name__}: {exc}"
+ ) from None
+ updated, duplicate_ids, persisted = publish_target(
+ verified,
+ set_default=set_default,
+ persist=persist,
+ targets_file=targets_file,
+ )
+ return payload_response(
+ "connect_iotdb_target",
+ {
+ "registered_target": verified.as_dict(include_secret=False),
+ "replaced_target_ids": duplicate_ids,
+ "registry": _targets_payload(config),
+ "persisted": persisted,
+ },
+ message="IoTDB target authenticated and registered.",
+ )
+
+ @mcp.tool()
+ async def register_iotdb_target(
+ target: dict[str, object],
+ set_default: bool = False,
+ persist: bool | None = None,
+ targets_file: str | None = None,
+ user_confirmed_retry: bool = False,
+ ) -> list[TextContent]:
+ """Compatibility entry point that verifies credentials before registration."""
+ if not isinstance(target, dict):
+ raise ValueError("target must be an object.")
+ username = target.get("user", target.get("username"))
+ if username is None or "password" not in target:
+ raise ValueError(
+ "register_iotdb_target requires user and an explicitly supplied password. "
+ "When credentials are unavailable, call prepare_iotdb_target and ask the user."
+ )
+ template = {
+ key: value
+ for key, value in target.items()
+ if key
+ not in {
+ "user",
+ "username",
+ "password",
+ "last_known_good_credential",
+ "verified_at",
+ }
+ }
+ candidate = candidate_store.prepare(
+ template,
+ user_confirmed_retry=user_confirmed_retry,
+ )
+ consumed = candidate_store.consume(candidate.candidate_id)
+ parsed_target = candidate_target(
+ consumed,
+ user=str(username),
+ password="" if target["password"] is None else str(target["password"]),
+ )
+ try:
+ verified = verify_target_once(parsed_target)
+ except Exception as exc:
+ candidate_store.record_failure(consumed)
+ raise ConnectionError(
+ "IoTDB connection failed; no target was registered. Do not retry "
+ "until the user explicitly requests reconnection. "
+ f"Cause: {type(exc).__name__}: {exc}"
+ ) from None
+ updated, duplicate_ids, persisted = publish_target(
+ verified,
+ set_default=set_default,
+ persist=persist,
+ targets_file=targets_file,
+ )
+ logger.info("Authenticated and registered IoTDB target %s", verified.target_id)
+ return payload_response(
+ "register_iotdb_target",
+ {
+ "registered_target": verified.as_dict(include_secret=False),
+ "replaced_target_ids": duplicate_ids,
+ "registry": _targets_payload(config),
+ "persisted": persisted,
+ },
+ message="IoTDB target authenticated and registered.",
+ )
+
+ @mcp.tool()
+ async def remember_iotdb_target_credential(
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ persist: bool = False,
+ targets_file: str | None = None,
+ ) -> list[TextContent]:
+ """Record the selected target's current credential as last-known-good.
+
+ Call this only after a connection/query has succeeded with the selected
+ target. This tool does not test login and does not try fallback
+ passwords.
+ """
+ registry = _registry(config)
+ selected_target = registry.resolve_selector(
+ target,
+ target_id=target_id or (None if target is not None else config.target_id),
+ )
+ updated_target = mark_target_last_known_good(selected_target)
+ updated = registry.with_target(updated_target)
+ apply_target_registry(
+ config,
+ updated,
+ active_target_id=(
+ config.target_id
+ if config.target_id
+ in {
+ item["target_id"]
+ for item in updated.list_targets(include_secret=False)
+ }
+ else updated.default_target_id
+ ),
+ close_pools=False,
+ )
+ if persist:
+ _persist_registry(config, updated, targets_file)
+ logger.info(
+ "Remembered IoTDB target credential for %s",
+ updated_target.target_id,
+ )
+ return payload_response(
+ "remember_iotdb_target_credential",
+ {
+ "remembered_target": updated_target.as_dict(include_secret=False),
+ "registry": _targets_payload(config),
+ "persisted": persist,
+ },
+ message="IoTDB target last-known-good credential recorded.",
+ )
+
+ @mcp.tool()
+ async def remove_iotdb_target(
+ target_id: str,
+ persist: bool = False,
+ targets_file: str | None = None,
+ ) -> list[TextContent]:
+ """Remove an IoTDB target from this running MCP process."""
+ registry = _registry(config)
+ updated = registry.without_target(target_id)
+ apply_target_registry(
+ config,
+ updated,
+ active_target_id=updated.default_target_id,
+ close_pools=True,
+ )
+ if persist:
+ _persist_registry(config, updated, targets_file)
+ logger.info("Removed IoTDB target %s", target_id)
+ return payload_response(
+ "remove_iotdb_target",
+ {
+ "removed_target_id": target_id,
+ "registry": _targets_payload(config),
+ "persisted": persist,
+ },
+ message="IoTDB target removed.",
+ )
+
+ @mcp.tool()
+ async def set_default_iotdb_target(
+ target_id: str,
+ persist: bool = False,
+ targets_file: str | None = None,
+ ) -> list[TextContent]:
+ """Set the default IoTDB target for this running MCP process."""
+ registry = _registry(config)
+ updated = registry.with_default(target_id)
+ apply_target_registry(
+ config,
+ updated,
+ active_target_id=target_id,
+ close_pools=False,
+ )
+ if persist:
+ _persist_registry(config, updated, targets_file)
+ logger.info("Set default IoTDB target %s", target_id)
+ return payload_response(
+ "set_default_iotdb_target",
+ {
+ "default_target_id": target_id,
+ "registry": _targets_payload(config),
+ "persisted": persist,
+ },
+ message="Default IoTDB target updated.",
+ )
+
+ @mcp.tool()
+ async def reload_iotdb_targets(
+ targets_file: str | None = None,
+ targets_json: str | None = None,
+ ) -> list[TextContent]:
+ """Reload IoTDB targets from env, a targets file, or an explicit JSON string."""
+ resolved_file = (
+ targets_file
+ or config.targets_file
+ or os.getenv("TIMESEEK_IOTDB_TARGETS_FILE", "")
+ or None
+ )
+ resolved_json = (
+ targets_json
+ or config.targets_json
+ or os.getenv("TIMESEEK_IOTDB_TARGETS_JSON", "")
+ or None
+ )
+ registry = IoTDBTargetRegistry.from_env(
+ targets_file=resolved_file,
+ targets_json=resolved_json,
+ )
+ mtime_ns = None
+ if resolved_file:
+ try:
+ mtime_ns = Path(resolved_file).expanduser().stat().st_mtime_ns
+ except OSError:
+ mtime_ns = None
+ config.targets_file = resolved_file
+ if resolved_json:
+ config.targets_json = resolved_json
+ apply_target_registry(
+ config,
+ registry,
+ targets_file_mtime_ns=mtime_ns,
+ close_pools=True,
+ )
+ logger.info("Reloaded IoTDB targets")
+ return payload_response(
+ "reload_iotdb_targets",
+ _targets_payload(config),
+ message="IoTDB targets reloaded.",
+ )
diff --git a/src/iotdb_mcp_server/services/timeseries.py b/src/iotdb_mcp_server/services/timeseries.py
new file mode 100644
index 0000000..89b034e
--- /dev/null
+++ b/src/iotdb_mcp_server/services/timeseries.py
@@ -0,0 +1,351 @@
+#
+# 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.
+#
+
+import logging
+import time
+from typing import Any
+
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.runtime_policy import dynamic_env_bool
+from iotdb_mcp_server.runtime_policy import dynamic_getenv
+from iotdb_mcp_server.runtime_policy import strict_permission_enforcement
+from iotdb_mcp_server.services.json_response import payload_response, sql_success_response
+from iotdb_mcp_server.services.target_selection import (
+ iotdb_target_response_context,
+ tree_session_pool,
+)
+
+
+def _env_bool(name: str, default: bool) -> bool:
+ return dynamic_env_bool(name, default)
+
+
+def _csv_set(value: str) -> set[str]:
+ return {item.strip() for item in value.split(",") if item.strip()}
+
+
+def _assert_timeseries_ddl_permission(
+ config: Config, action: str, confirm: bool = False
+) -> None:
+ if strict_permission_enforcement():
+ if not _env_bool("IOTDB_ENABLE_TIMESERIES_DDL", False):
+ raise PermissionError(
+ "Timeseries DDL tools are disabled by strict MCP policy. "
+ "Set IOTDB_ENABLE_TIMESERIES_DDL=true to enable."
+ )
+
+ allowed_users = _csv_set(dynamic_getenv("IOTDB_TIMESERIES_DDL_ALLOWED_USERS", "root") or "root")
+ if "*" not in allowed_users and config.user not in allowed_users:
+ raise PermissionError(
+ f"Current MCP user '{config.user}' is not allowed by IOTDB_TIMESERIES_DDL_ALLOWED_USERS."
+ )
+
+ if (
+ action == "DROP"
+ and _env_bool("IOTDB_REQUIRE_TIMESERIES_DROP_CONFIRM", True)
+ and not confirm
+ ):
+ raise PermissionError(
+ "DROP/DELETE TIMESERIES requires confirm=True by server policy."
+ )
+
+
+def _normalize_sql(sql: str) -> str:
+ cleaned = sql.strip()
+ if not cleaned:
+ raise ValueError("DDL SQL cannot be empty.")
+ if cleaned.endswith(";"):
+ cleaned = cleaned[:-1].strip()
+ if ";" in cleaned:
+ raise ValueError("Only a single SQL statement is allowed.")
+ return cleaned
+
+
+def _ensure_prefix(sql: str, allowed_prefixes: tuple[str, ...], action_name: str) -> None:
+ upper = sql.upper()
+ if not upper.startswith(allowed_prefixes):
+ raise ValueError(
+ f"{action_name} only supports SQL starting with: {', '.join(allowed_prefixes)}"
+ )
+
+
+def _is_already_exists_error(message: str) -> bool:
+ lower = message.lower()
+ patterns = (
+ "already exist",
+ "already_exists",
+ "path already exists",
+ "timeseries already exists",
+ )
+ return any(p in lower for p in patterns)
+
+
+def register_timeseries_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register tree-model timeseries-level DDL tools."""
+
+ @mcp.tool()
+ async def create_timeseries_ddl(
+ ddl_sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute tree-model CREATE TIMESERIES DDL.
+
+ Supported prefixes:
+ - CREATE TIMESERIES
+ - CREATE ALIGNED TIMESERIES
+ """
+ selected_config, session_pool = tree_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=100,
+ wait_timeout_in_ms=5000,
+ tool_name="create_timeseries_ddl",
+ )
+ with iotdb_target_response_context(selected_config):
+ session = None
+ try:
+ _assert_timeseries_ddl_permission(selected_config, action="CREATE")
+ sql = _normalize_sql(ddl_sql)
+ _ensure_prefix(
+ sql,
+ ("CREATE TIMESERIES", "CREATE ALIGNED TIMESERIES"),
+ "create_timeseries_ddl",
+ )
+ session = session_pool.get_session()
+ session.execute_non_query_statement(sql)
+ session.close()
+ return sql_success_response("create_timeseries_ddl", sql)
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute create_timeseries_ddl: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def create_timeseries_batch_ddl(
+ ddl_list: list[str],
+ continue_on_error: bool = True,
+ skip_if_exists: bool = True,
+ max_statements: int = 1000,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute a batch of tree-model CREATE TIMESERIES DDL in one session pipeline.
+
+ Supported prefixes in each statement:
+ - CREATE TIMESERIES
+ - CREATE ALIGNED TIMESERIES
+ """
+ selected_config, session_pool = tree_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=100,
+ wait_timeout_in_ms=5000,
+ tool_name="create_timeseries_batch_ddl",
+ )
+ with iotdb_target_response_context(selected_config):
+ session = None
+ started_at = time.perf_counter()
+ try:
+ _assert_timeseries_ddl_permission(selected_config, action="CREATE")
+
+ if not isinstance(ddl_list, list) or len(ddl_list) == 0:
+ raise ValueError("ddl_list must be a non-empty list of SQL statements.")
+ if max_statements <= 0:
+ raise ValueError("max_statements must be greater than 0.")
+ if len(ddl_list) > max_statements:
+ raise ValueError(
+ f"ddl_list length {len(ddl_list)} exceeds max_statements={max_statements}."
+ )
+
+ validated_sqls: list[dict[str, Any]] = []
+ per_statement: list[dict[str, Any]] = []
+ failed_count = 0
+
+ for idx, raw_sql in enumerate(ddl_list, start=1):
+ statement_started = time.perf_counter()
+ try:
+ sql = _normalize_sql(str(raw_sql))
+ _ensure_prefix(
+ sql,
+ ("CREATE TIMESERIES", "CREATE ALIGNED TIMESERIES"),
+ "create_timeseries_batch_ddl",
+ )
+ validated_sqls.append({"index": idx, "sql": sql})
+ except Exception as e:
+ failed_count += 1
+ per_statement.append(
+ {
+ "index": idx,
+ "sql": str(raw_sql),
+ "status": "invalid_sql",
+ "elapsed_ms": round(
+ (time.perf_counter() - statement_started) * 1000, 3
+ ),
+ "error": str(e),
+ }
+ )
+ if not continue_on_error:
+ break
+
+ applied_count = 0
+ skipped_count = 0
+ execution_aborted = False
+
+ if validated_sqls and (failed_count == 0 or continue_on_error):
+ session = session_pool.get_session()
+ for item in validated_sqls:
+ idx = int(item["index"])
+ sql = str(item["sql"])
+ statement_started = time.perf_counter()
+ status = "applied"
+ error_msg = ""
+ try:
+ session.execute_non_query_statement(sql)
+ applied_count += 1
+ except Exception as e:
+ if skip_if_exists and _is_already_exists_error(str(e)):
+ status = "skipped_existing"
+ skipped_count += 1
+ else:
+ status = "error"
+ error_msg = str(e)
+ failed_count += 1
+ if not continue_on_error:
+ execution_aborted = True
+ per_statement.append(
+ {
+ "index": idx,
+ "sql": sql,
+ "status": status,
+ "elapsed_ms": round(
+ (time.perf_counter() - statement_started) * 1000, 3
+ ),
+ "error": error_msg,
+ }
+ )
+ if execution_aborted:
+ break
+
+ if session:
+ session.close()
+
+ per_statement.sort(key=lambda x: x.get("index", 0))
+ total_elapsed_ms = round((time.perf_counter() - started_at) * 1000, 3)
+ executed_count = applied_count + skipped_count + (
+ len([x for x in per_statement if x.get("status") == "error"])
+ )
+
+ return payload_response(
+ "create_timeseries_batch_ddl",
+ {
+ "result": "success" if failed_count == 0 else "partial_success",
+ "execution_mode": "single_session_pipeline",
+ "continue_on_error": continue_on_error,
+ "skip_if_exists": skip_if_exists,
+ "total_input": len(ddl_list),
+ "validated_count": len(validated_sqls),
+ "executed_count": executed_count,
+ "applied_count": applied_count,
+ "skipped_count": skipped_count,
+ "failed_count": failed_count,
+ "total_elapsed_ms": total_elapsed_ms,
+ "statements": per_statement,
+ },
+ message="Batch CREATE TIMESERIES execution completed.",
+ )
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute create_timeseries_batch_ddl: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def alter_timeseries_ddl(
+ ddl_sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute tree-model ALTER TIMESERIES DDL."""
+ selected_config, session_pool = tree_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=100,
+ wait_timeout_in_ms=5000,
+ tool_name="alter_timeseries_ddl",
+ )
+ with iotdb_target_response_context(selected_config):
+ session = None
+ try:
+ _assert_timeseries_ddl_permission(selected_config, action="ALTER")
+ sql = _normalize_sql(ddl_sql)
+ _ensure_prefix(sql, ("ALTER TIMESERIES",), "alter_timeseries_ddl")
+ session = session_pool.get_session()
+ session.execute_non_query_statement(sql)
+ session.close()
+ return sql_success_response("alter_timeseries_ddl", sql)
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute alter_timeseries_ddl: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def drop_timeseries_ddl(
+ ddl_sql: str,
+ confirm: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute tree-model DROP/DELETE TIMESERIES DDL (destructive)."""
+ selected_config, session_pool = tree_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=100,
+ wait_timeout_in_ms=5000,
+ tool_name="drop_timeseries_ddl",
+ )
+ with iotdb_target_response_context(selected_config):
+ session = None
+ try:
+ _assert_timeseries_ddl_permission(
+ selected_config, action="DROP", confirm=confirm
+ )
+ sql = _normalize_sql(ddl_sql)
+ _ensure_prefix(
+ sql,
+ ("DROP TIMESERIES", "DELETE TIMESERIES"),
+ "drop_timeseries_ddl",
+ )
+ session = session_pool.get_session()
+ session.execute_non_query_statement(sql)
+ session.close()
+ return sql_success_response("drop_timeseries_ddl", sql)
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute drop_timeseries_ddl: {str(e)}")
+ raise
diff --git a/src/iotdb_mcp_server/services/tree_sql_guardrails.py b/src/iotdb_mcp_server/services/tree_sql_guardrails.py
new file mode 100644
index 0000000..7d2a87b
--- /dev/null
+++ b/src/iotdb_mcp_server/services/tree_sql_guardrails.py
@@ -0,0 +1,199 @@
+import re
+
+
+AGGREGATE_FUNCTIONS = (
+ "COUNT",
+ "SUM",
+ "AVG",
+ "MAX_VALUE",
+ "MIN_VALUE",
+ "FIRST_VALUE",
+ "LAST_VALUE",
+ "MAX",
+ "MIN",
+)
+
+
+def _split_top_level_csv(expr: str) -> list[str]:
+ parts: list[str] = []
+ cur: list[str] = []
+ depth = 0
+ in_single_quote = False
+
+ idx = 0
+ while idx < len(expr):
+ ch = expr[idx]
+ if ch == "'":
+ in_single_quote = not in_single_quote
+ cur.append(ch)
+ elif not in_single_quote and ch == "(":
+ depth += 1
+ cur.append(ch)
+ elif not in_single_quote and ch == ")":
+ depth = max(0, depth - 1)
+ cur.append(ch)
+ elif not in_single_quote and ch == "," and depth == 0:
+ part = "".join(cur).strip()
+ if part:
+ parts.append(part)
+ cur = []
+ else:
+ cur.append(ch)
+ idx += 1
+
+ tail = "".join(cur).strip()
+ if tail:
+ parts.append(tail)
+ return parts
+
+
+def _strip_alias(expr: str) -> str:
+ out = re.sub(r"\s+AS\s+[A-Z_][A-Z0-9_]*\s*$", "", expr, flags=re.IGNORECASE)
+ m = re.match(r"^(.*\S)\s+([A-Z_][A-Z0-9_]*)\s*$", out, flags=re.IGNORECASE)
+ if m and "(" in m.group(1):
+ return m.group(1)
+ return out
+
+
+def _strip_identifier_quotes(identifier: str) -> str:
+ value = identifier.strip()
+ if len(value) >= 2 and value[0] == value[-1] and value[0] in {"`", '"'}:
+ return value[1:-1].strip()
+ return value
+
+
+def _extract_projection_expr(sql: str) -> str:
+ match = re.search(
+ r"\bSELECT\b\s+(?:LAST\s+)?(.+?)\s+\bFROM\b",
+ sql,
+ flags=re.IGNORECASE | re.DOTALL,
+ )
+ if not match:
+ return ""
+ return match.group(1).strip()
+
+
+def _extract_from_identifier(sql: str) -> str:
+ match = re.search(r"\bFROM\s+([^\s,;]+)", sql, flags=re.IGNORECASE)
+ if not match:
+ return ""
+ identifier = match.group(1).strip()
+ if not identifier or identifier.startswith("("):
+ return ""
+ return _strip_identifier_quotes(identifier.rstrip(",;"))
+
+
+def _projection_measurement_names(sql: str) -> set[str]:
+ projection = _extract_projection_expr(sql)
+ if not projection:
+ return set()
+
+ aggregate_re = re.compile(
+ rf"^(?:{'|'.join(AGGREGATE_FUNCTIONS)})\s*\((.*)\)$",
+ flags=re.IGNORECASE | re.DOTALL,
+ )
+ simple_identifier_re = re.compile(
+ r"^[`\"]?[A-Za-z_][A-Za-z0-9_]*[`\"]?"
+ r"(?:\.[`\"]?[A-Za-z_][A-Za-z0-9_]*[`\"]?)*$"
+ )
+
+ measurements: set[str] = set()
+ for raw_item in _split_top_level_csv(projection):
+ item = _strip_alias(raw_item.strip())
+ aggregate = aggregate_re.match(item)
+ if aggregate:
+ item = aggregate.group(1).strip()
+ if item == "*" or not simple_identifier_re.fullmatch(item):
+ continue
+ measurement = _strip_identifier_quotes(item.rsplit(".", 1)[-1])
+ if measurement:
+ measurements.add(measurement.lower())
+ return measurements
+
+
+def validate_tree_query_shape(sql: str) -> list[dict[str, object]]:
+ """Return high-confidence Tree SQL shape issues before runtime execution.
+
+ This intentionally covers only patterns that are almost certainly wrong in
+ ordinary Tree reads. It does not try to infer schema or rewrite SQL.
+ """
+ from_identifier = _extract_from_identifier(sql)
+ if not from_identifier.lower().startswith("root.") or "*" in from_identifier:
+ return []
+
+ path_segments = [_strip_identifier_quotes(part) for part in from_identifier.split(".") if part]
+ if len(path_segments) < 4:
+ return []
+
+ from_tail = path_segments[-1].lower()
+ projection_measurements = _projection_measurement_names(sql)
+ if from_tail not in projection_measurements:
+ return []
+
+ device_path = ".".join(path_segments[:-1])
+ measurement = path_segments[-1]
+ return [
+ {
+ "code": "tree_dialect_from_must_use_device_path",
+ "message": (
+ "Tree SQL should use a device/path-pattern in FROM and project "
+ "measurement names in SELECT. The FROM path appears to include "
+ f"measurement `{measurement}`."
+ ),
+ "rewrite_hint": (
+ f"Use FROM {device_path} and keep `{measurement}` in the SELECT "
+ "projection or aggregate expression."
+ ),
+ "from_path": from_identifier,
+ "suggested_from_path": device_path,
+ "measurement": measurement,
+ }
+ ]
+
+
+def tree_from_wildcard_runtime_hint(sql: str, *, row_count: int | None = None) -> dict[str, object] | None:
+ """Return a non-blocking hint for device-path wildcard FROM shapes.
+
+ IoTDB Tree accepts prefix paths such as ``root.sg.d1.*`` syntactically, but
+ agents often build that shape by appending ``.*`` to a device path copied
+ from schema metadata. For ordinary measurement projections, the canonical
+ query shape is ``SELECT measurement FROM root.sg.d1``. This helper is a
+ runtime diagnostic only; it does not reject broad scans such as ``root.**``.
+ """
+ from_identifier = _extract_from_identifier(sql)
+ if not from_identifier.lower().startswith("root."):
+ return None
+ if not from_identifier.endswith(".*"):
+ return None
+ suggested_from_path = from_identifier[:-2]
+ if not suggested_from_path or suggested_from_path.endswith("."):
+ return None
+ projection_measurements = _projection_measurement_names(sql)
+ if not projection_measurements:
+ return None
+
+ qualifier = " If this query returned zero rows," if row_count == 0 else ""
+ return {
+ "code": "tree_dialect_from_path_wildcard_noncanonical",
+ "message": (
+ f"Tree SQL `FROM {from_identifier}` is a path-pattern form, not the "
+ "canonical device-scoped form for ordinary measurement reads."
+ f"{qualifier} retry with the device path in FROM and keep measurement names in SELECT."
+ ),
+ "rewrite_hint": (
+ f"Use FROM {suggested_from_path}; metadata `paths`/`timeseries` values are full "
+ "timeseries paths and should not be copied as FROM path.* clauses."
+ ),
+ "from_path": from_identifier,
+ "suggested_from_path": suggested_from_path,
+ }
+
+
+def assert_tree_query_shape(sql: str) -> None:
+ issues = validate_tree_query_shape(sql)
+ if not issues:
+ return
+ issue = issues[0]
+ raise ValueError(
+ f"{issue['code']}: {issue['message']} {issue['rewrite_hint']}"
+ )
diff --git a/src/iotdb_mcp_server/services/ttl.py b/src/iotdb_mcp_server/services/ttl.py
new file mode 100644
index 0000000..b22786e
--- /dev/null
+++ b/src/iotdb_mcp_server/services/ttl.py
@@ -0,0 +1,185 @@
+#
+# 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.
+#
+
+import logging
+
+from iotdb.Session import Session
+from iotdb.utils.SessionDataSet import SessionDataSet
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.runtime_policy import dynamic_env_bool
+from iotdb_mcp_server.runtime_policy import dynamic_getenv
+from iotdb_mcp_server.runtime_policy import strict_permission_enforcement
+from iotdb_mcp_server.services.json_response import (
+ csv_payload_response,
+ sql_success_response,
+)
+from iotdb_mcp_server.services.target_selection import (
+ iotdb_target_response_context,
+ tree_session_pool,
+)
+
+
+def _env_bool(name: str, default: bool) -> bool:
+ return dynamic_env_bool(name, default)
+
+
+def _csv_set(value: str) -> set[str]:
+ return {item.strip() for item in value.split(",") if item.strip()}
+
+
+def _assert_ttl_permission(
+ config: Config, action: str, confirm_unset: bool = False
+) -> None:
+ if strict_permission_enforcement():
+ if not _env_bool("IOTDB_ENABLE_TTL_SQL", False):
+ raise PermissionError(
+ "TTL tools are disabled by strict MCP policy. "
+ "Set IOTDB_ENABLE_TTL_SQL=true to enable."
+ )
+
+ allowed_users = _csv_set(dynamic_getenv("IOTDB_TTL_ALLOWED_USERS", "root") or "root")
+ if "*" not in allowed_users and config.user not in allowed_users:
+ raise PermissionError(
+ f"Current MCP user '{config.user}' is not allowed by IOTDB_TTL_ALLOWED_USERS."
+ )
+
+ if (
+ action == "UNSET"
+ and _env_bool("IOTDB_REQUIRE_TTL_UNSET_CONFIRM", True)
+ and not confirm_unset
+ ):
+ raise PermissionError(
+ "UNSET TTL requires confirm_unset=True by server policy."
+ )
+
+
+def _normalize_sql(sql: str) -> str:
+ cleaned = sql.strip()
+ if not cleaned:
+ raise ValueError("TTL SQL cannot be empty.")
+ if cleaned.endswith(";"):
+ cleaned = cleaned[:-1].strip()
+ if ";" in cleaned:
+ raise ValueError("Only a single SQL statement is allowed.")
+ return cleaned
+
+
+def _format_result(res: SessionDataSet, session: Session) -> list[TextContent]:
+ columns = res.get_column_names()
+ rows: list[str] = []
+ while res.has_next():
+ row = res.next().get_fields()
+ rows.append(",".join(map(str, row)))
+ session.close()
+ return csv_payload_response("ttl_query", columns, rows)
+
+
+def register_ttl_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register TTL SQL tools.
+
+ Note: TTL grammar is tree-model SQL.
+ """
+ @mcp.tool()
+ async def ttl_command(
+ ttl_sql: str,
+ confirm_unset: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute TTL command SQL.
+
+ Supported prefixes:
+ - SET TTL
+ - UNSET TTL
+ """
+ selected_config, session_pool = tree_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=100,
+ wait_timeout_in_ms=5000,
+ tool_name="ttl_command",
+ )
+ with iotdb_target_response_context(selected_config):
+ session = None
+ try:
+ sql = _normalize_sql(ttl_sql)
+ upper = sql.upper()
+ if upper.startswith("SET TTL"):
+ _assert_ttl_permission(selected_config, action="SET")
+ elif upper.startswith("UNSET TTL"):
+ _assert_ttl_permission(
+ selected_config, action="UNSET", confirm_unset=confirm_unset
+ )
+ else:
+ raise ValueError(
+ "ttl_command only supports SQL starting with: SET TTL, UNSET TTL"
+ )
+
+ session = session_pool.get_session()
+ session.execute_non_query_statement(sql)
+ session.close()
+ return sql_success_response("ttl_command", sql)
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute ttl_command: {str(e)}")
+ raise
+
+ @mcp.tool()
+ async def ttl_query(
+ ttl_sql: str,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute TTL query SQL.
+
+ Supported prefixes:
+ - SHOW TTL ON
+ - SHOW ALL TTL
+ """
+ selected_config, session_pool = tree_session_pool(
+ config,
+ target_id=target_id,
+ target=target,
+ max_pool_size=100,
+ wait_timeout_in_ms=5000,
+ tool_name="ttl_query",
+ )
+ with iotdb_target_response_context(selected_config):
+ session = None
+ try:
+ _assert_ttl_permission(selected_config, action="SHOW")
+ sql = _normalize_sql(ttl_sql)
+ upper = sql.upper()
+ if not upper.startswith(("SHOW TTL ON", "SHOW ALL TTL")):
+ raise ValueError(
+ "ttl_query only supports SQL starting with: SHOW TTL ON, SHOW ALL TTL"
+ )
+
+ session = session_pool.get_session()
+ res = session.execute_query_statement(sql)
+ return _format_result(res, session)
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute ttl_query: {str(e)}")
+ raise
diff --git a/src/iotdb_mcp_server/services/udf.py b/src/iotdb_mcp_server/services/udf.py
new file mode 100644
index 0000000..e720c35
--- /dev/null
+++ b/src/iotdb_mcp_server/services/udf.py
@@ -0,0 +1,534 @@
+from __future__ import annotations
+
+import datetime
+import logging
+import re
+import uuid
+from typing import Any
+
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.result_store import result_store_from_export_path
+from iotdb_mcp_server.services.json_response import (
+ csv_result_payload_response,
+ payload_response,
+ text_payload_response,
+)
+from iotdb_mcp_server.services.query import (
+ _ensure_export_directory,
+ _prepare_table_res,
+ _prepare_tree_res,
+ sanitize_filename,
+)
+from iotdb_mcp_server.services.target_selection import (
+ iotdb_target_response_context,
+ select_target_config,
+ table_session_pool,
+ tree_session_pool,
+)
+from iotdb_mcp_server.services.tree_sql_guardrails import assert_tree_query_shape
+
+_IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+_TREE_PATH_RE = re.compile(r"^root(?:\.[A-Za-z_][A-Za-z0-9_]*|\.\*|\.\*\*)*$")
+_TREE_EXPR_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*$")
+_DANGEROUS_SQL_RE = re.compile(
+ r";|--|/\*|\*/|\b(INSERT|UPDATE|DELETE|CREATE|DROP|ALTER|LOAD|SET|GRANT|REVOKE|TRUNCATE)\b",
+ re.IGNORECASE,
+)
+
+
+def _validate_identifier(value: str, label: str) -> str:
+ cleaned = str(value or "").strip()
+ if not _IDENTIFIER_RE.fullmatch(cleaned):
+ raise ValueError(f"Invalid {label}. Use letters, numbers, and underscore only.")
+ return cleaned
+
+
+def _validate_function_name(function_name: str) -> str:
+ return _validate_identifier(function_name, "function_name")
+
+
+def _validate_tree_path(path: str) -> str:
+ cleaned = str(path or "").strip()
+ if not _TREE_PATH_RE.fullmatch(cleaned):
+ raise ValueError("Invalid tree from_path. Use path like root.sg.d1 or root.sg.*.")
+ return cleaned
+
+
+def _validate_expression(expression: str, *, dialect: str) -> str:
+ cleaned = str(expression or "").strip()
+ if not cleaned:
+ raise ValueError("UDF expression cannot be empty.")
+ if _DANGEROUS_SQL_RE.search(cleaned):
+ raise ValueError("UDF expression contains unsupported SQL syntax.")
+ if dialect == "table":
+ return _validate_identifier(cleaned, "table expression")
+ if not _TREE_EXPR_RE.fullmatch(cleaned):
+ raise ValueError("Invalid tree expression. Use measurement names or relative paths.")
+ return cleaned
+
+
+def _validate_where_clause(where_clause: str | None) -> str | None:
+ if where_clause is None:
+ return None
+ cleaned = str(where_clause).strip()
+ if not cleaned:
+ return None
+ if _DANGEROUS_SQL_RE.search(cleaned):
+ raise ValueError("where_clause contains unsupported or non-readonly SQL syntax.")
+ if cleaned.upper().startswith("WHERE "):
+ cleaned = cleaned[6:].strip()
+ if not cleaned:
+ return None
+ return cleaned
+
+
+def _render_attribute_value(value: Any) -> str:
+ text = str(value)
+ escaped = text.replace("\\", "\\\\").replace('"', '\\"')
+ return f'"{escaped}"'
+
+
+def _render_attributes(attributes: dict[str, Any] | None) -> list[str]:
+ if not attributes:
+ return []
+ rendered = []
+ for key, value in attributes.items():
+ safe_key = _validate_identifier(str(key), "attribute key")
+ rendered.append(f'"{safe_key}"={_render_attribute_value(value)}')
+ return rendered
+
+
+def _normalize_expressions(expressions: list[str] | str | None, *, dialect: str) -> list[str]:
+ if expressions is None:
+ raise ValueError("expressions is required.")
+ raw_values = [expressions] if isinstance(expressions, str) else list(expressions)
+ if not raw_values:
+ raise ValueError("expressions cannot be empty.")
+ return [_validate_expression(str(value), dialect=dialect) for value in raw_values]
+
+
+def _build_udf_query_plan(
+ *,
+ dialect: str,
+ function_name: str,
+ expressions: list[str] | str,
+ table_name: str | None = None,
+ from_path: str | None = None,
+ attributes: dict[str, Any] | None = None,
+ where_clause: str | None = None,
+ limit: int | None = None,
+ alias: str | None = None,
+ align_by_device: bool = False,
+) -> dict[str, Any]:
+ safe_dialect = str(dialect or "").strip().lower()
+ if safe_dialect not in {"tree", "table"}:
+ raise ValueError("dialect must be tree or table.")
+ safe_function = _validate_function_name(function_name)
+ safe_expressions = _normalize_expressions(expressions, dialect=safe_dialect)
+ function_args = safe_expressions + _render_attributes(attributes)
+ projection = f"{safe_function}({', '.join(function_args)})"
+ if alias:
+ projection += f" AS {_validate_identifier(alias, 'alias')}"
+
+ where = _validate_where_clause(where_clause)
+ clauses: list[str]
+ if safe_dialect == "table":
+ source = _validate_identifier(str(table_name or ""), "table_name")
+ clauses = [f"SELECT {projection}", f"FROM {source}"]
+ else:
+ source = _validate_tree_path(str(from_path or ""))
+ clauses = [f"SELECT {projection}", f"FROM {source}"]
+
+ if where:
+ clauses.append(f"WHERE {where}")
+ if limit is not None:
+ safe_limit = int(limit)
+ if safe_limit <= 0:
+ raise ValueError("limit must be positive.")
+ clauses.append(f"LIMIT {safe_limit}")
+ if safe_dialect == "tree" and align_by_device:
+ clauses.append("ALIGN BY DEVICE")
+
+ sql = " ".join(clauses)
+ return {
+ "version": "v0",
+ "kind": "iotdb_udf_query_plan_ir",
+ "sql": sql,
+ "dialect": safe_dialect,
+ "function_name": safe_function,
+ "expressions": safe_expressions,
+ "source": {"table_name": source} if safe_dialect == "table" else {"from_path": source},
+ "attributes": attributes or {},
+ "where_clause": where,
+ "limit": limit,
+ "alias": alias,
+ "align_by_device": bool(align_by_device),
+ "readonly": True,
+ }
+
+
+def _export_dataset(
+ *,
+ res: Any,
+ session: Any,
+ selected_config: Config,
+ tool_name: str,
+ query_sql: str,
+ filename: str | None,
+ fmt: str,
+ logger: logging.Logger,
+) -> list[TextContent]:
+ try:
+ df = res.todf()
+ finally:
+ session.close()
+
+ timestamp = int(datetime.datetime.now().timestamp())
+ if filename is None:
+ filename = f"udf_{uuid.uuid4().hex[:4]}_{timestamp}"
+
+ fmt_lower = fmt.lower()
+ if fmt_lower == "csv":
+ if filename.lower().endswith(".csv"):
+ filename = filename[:-4]
+ filepath = sanitize_filename(f"{filename}.csv", selected_config.export_path)
+ df.to_csv(filepath, index=False)
+ elif fmt_lower == "excel":
+ if filename.lower().endswith(".xlsx"):
+ filename = filename[:-5]
+ filepath = sanitize_filename(f"{filename}.xlsx", selected_config.export_path)
+ df.to_excel(filepath, index=False)
+ else:
+ raise ValueError("format must be either 'csv' or 'excel'.")
+
+ preview_rows = min(10, len(df))
+ preview_data = [",".join(df.columns)]
+ for index in range(preview_rows):
+ preview_data.append(",".join(map(str, df.iloc[index])))
+ logger.info("Exported UDF query result to %s", filepath)
+ return text_payload_response(
+ tool_name,
+ f"Query results exported to {filepath}\n\nSQL:\n{query_sql}\n\nPreview (first {preview_rows} rows):\n"
+ + "\n".join(preview_data),
+ )
+
+
+def _prepare_udf_res(
+ res: Any,
+ session: Any,
+ tool_name: str,
+ export_path: str,
+ plan: dict[str, Any],
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+) -> list[TextContent]:
+ columns = res.get_column_names()
+
+ def rows():
+ while res.has_next():
+ record = res.next()
+ fields = record.get_fields()
+ if plan["dialect"] == "tree" and columns and columns[0] == "Time":
+ yield str(record.get_timestamp()) + "," + ",".join(
+ map(str, fields)
+ )
+ else:
+ yield ",".join(map(str, fields))
+
+ try:
+ return csv_result_payload_response(
+ tool_name,
+ columns,
+ rows(),
+ result_store=result_store_from_export_path(export_path),
+ source={"query_sql": plan["sql"], "udf_query_plan": plan},
+ diagnostics=[
+ {
+ "code": "iotdb_udf_query_plan",
+ "query_sql": plan["sql"],
+ "function_name": plan["function_name"],
+ "dialect": plan["dialect"],
+ "readonly": True,
+ }
+ ],
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ finally:
+ session.close()
+
+
+def register_udf_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register read-only IoTDB UDF discovery and execution tools."""
+ _ensure_export_directory(config.export_path, logger)
+ max_pool_size = 100
+
+ @mcp.tool()
+ async def list_udf_functions(
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+ ) -> list[TextContent]:
+ """List IoTDB functions/UDFs visible to the selected target via SHOW FUNCTIONS."""
+ selected_config = select_target_config(config, target_id=target_id, target=target, tool_name="list_udf_functions")
+ with iotdb_target_response_context(selected_config):
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="list_udf_functions",
+ )
+ session = None
+ try:
+ session = session_pool.get_session()
+ res = session.execute_query_statement("SHOW FUNCTIONS")
+ return _prepare_tree_res(
+ res,
+ session,
+ "list_udf_functions",
+ selected_config.export_path,
+ "SHOW FUNCTIONS",
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ except Exception:
+ if session:
+ session.close()
+ raise
+
+ _, session_pool = table_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ tool_name="list_udf_functions",
+ )
+ table_session = None
+ try:
+ table_session = session_pool.get_session()
+ res = table_session.execute_query_statement("SHOW FUNCTIONS")
+ return _prepare_table_res(
+ res,
+ table_session,
+ "list_udf_functions",
+ selected_config.export_path,
+ "SHOW FUNCTIONS",
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ except Exception:
+ if table_session:
+ table_session.close()
+ raise
+
+ @mcp.tool()
+ async def prepare_udf_query(
+ function_name: str,
+ expressions: list[str] | str,
+ table_name: str | None = None,
+ from_path: str | None = None,
+ attributes: dict[str, object] | None = None,
+ where_clause: str | None = None,
+ limit: int | None = None,
+ alias: str | None = None,
+ align_by_device: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Build a conservative read-only SQL plan for invoking an IoTDB UDF."""
+ selected_config = select_target_config(config, target_id=target_id, target=target, tool_name="prepare_udf_query")
+ with iotdb_target_response_context(selected_config):
+ payload = _build_udf_query_plan(
+ dialect=selected_config.sql_dialect,
+ function_name=function_name,
+ expressions=expressions,
+ table_name=table_name,
+ from_path=from_path,
+ attributes=attributes,
+ where_clause=where_clause,
+ limit=limit,
+ alias=alias,
+ align_by_device=align_by_device,
+ )
+ return payload_response("prepare_udf_query", payload, message="UDF query plan prepared.")
+
+ @mcp.tool()
+ async def execute_udf_query(
+ function_name: str,
+ expressions: list[str] | str,
+ table_name: str | None = None,
+ from_path: str | None = None,
+ attributes: dict[str, object] | None = None,
+ where_clause: str | None = None,
+ limit: int | None = None,
+ alias: str | None = None,
+ align_by_device: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ owner_session_id: str | None = None,
+ max_inline_rows: int | None = None,
+ page_size_rows: int | None = None,
+ ) -> list[TextContent]:
+ """Execute one read-only IoTDB UDF SELECT built from validated inputs."""
+ selected_config = select_target_config(config, target_id=target_id, target=target, tool_name="execute_udf_query")
+ with iotdb_target_response_context(selected_config):
+ plan = _build_udf_query_plan(
+ dialect=selected_config.sql_dialect,
+ function_name=function_name,
+ expressions=expressions,
+ table_name=table_name,
+ from_path=from_path,
+ attributes=attributes,
+ where_clause=where_clause,
+ limit=limit,
+ alias=alias,
+ align_by_device=align_by_device,
+ )
+ sql = plan["sql"]
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="execute_udf_query",
+ )
+ session = None
+ try:
+ assert_tree_query_shape(sql)
+ session = session_pool.get_session()
+ res = session.execute_query_statement(sql)
+ return _prepare_udf_res(
+ res,
+ session,
+ "execute_udf_query",
+ selected_config.export_path,
+ plan,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ except Exception:
+ if session:
+ session.close()
+ raise
+
+ _, session_pool = table_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ tool_name="execute_udf_query",
+ )
+ table_session = None
+ try:
+ table_session = session_pool.get_session()
+ res = table_session.execute_query_statement(sql)
+ return _prepare_udf_res(
+ res,
+ table_session,
+ "execute_udf_query",
+ selected_config.export_path,
+ plan,
+ owner_session_id=owner_session_id,
+ max_inline_rows=max_inline_rows,
+ page_size_rows=page_size_rows,
+ )
+ except Exception:
+ if table_session:
+ table_session.close()
+ raise
+
+ @mcp.tool()
+ async def export_udf_query(
+ function_name: str,
+ expressions: list[str] | str,
+ table_name: str | None = None,
+ from_path: str | None = None,
+ attributes: dict[str, object] | None = None,
+ where_clause: str | None = None,
+ limit: int | None = None,
+ alias: str | None = None,
+ align_by_device: bool = False,
+ format: str = "csv",
+ filename: str | None = None,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute one read-only IoTDB UDF SELECT and export the result set."""
+ selected_config = select_target_config(config, target_id=target_id, target=target, tool_name="export_udf_query")
+ with iotdb_target_response_context(selected_config):
+ _ensure_export_directory(selected_config.export_path, logger)
+ plan = _build_udf_query_plan(
+ dialect=selected_config.sql_dialect,
+ function_name=function_name,
+ expressions=expressions,
+ table_name=table_name,
+ from_path=from_path,
+ attributes=attributes,
+ where_clause=where_clause,
+ limit=limit,
+ alias=alias,
+ align_by_device=align_by_device,
+ )
+ sql = plan["sql"]
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="export_udf_query",
+ )
+ session = None
+ try:
+ assert_tree_query_shape(sql)
+ session = session_pool.get_session()
+ res = session.execute_query_statement(sql)
+ return _export_dataset(
+ res=res,
+ session=session,
+ selected_config=selected_config,
+ tool_name="export_udf_query",
+ query_sql=sql,
+ filename=filename,
+ fmt=format,
+ logger=logger,
+ )
+ except Exception:
+ if session:
+ session.close()
+ raise
+
+ _, session_pool = table_session_pool(
+ config,
+ target_id=selected_config.target_id,
+ max_pool_size=max_pool_size,
+ tool_name="export_udf_query",
+ )
+ table_session = None
+ try:
+ table_session = session_pool.get_session()
+ res = table_session.execute_query_statement(sql)
+ return _export_dataset(
+ res=res,
+ session=table_session,
+ selected_config=selected_config,
+ tool_name="export_udf_query",
+ query_sql=sql,
+ filename=filename,
+ fmt=format,
+ logger=logger,
+ )
+ except Exception:
+ if table_session:
+ table_session.close()
+ raise
diff --git a/src/iotdb_mcp_server/services/write.py b/src/iotdb_mcp_server/services/write.py
new file mode 100644
index 0000000..5250112
--- /dev/null
+++ b/src/iotdb_mcp_server/services/write.py
@@ -0,0 +1,163 @@
+#
+# 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.
+#
+
+import logging
+
+from mcp.types import TextContent
+
+from iotdb_mcp_server.config import Config
+from iotdb_mcp_server.runtime_policy import dynamic_env_bool
+from iotdb_mcp_server.runtime_policy import dynamic_getenv
+from iotdb_mcp_server.runtime_policy import strict_permission_enforcement
+from iotdb_mcp_server.services.json_response import sql_success_response
+from iotdb_mcp_server.services.target_selection import (
+ iotdb_target_response_context,
+ select_target_config,
+ table_session_pool,
+ tree_session_pool,
+)
+
+
+def _env_bool(name: str, default: bool) -> bool:
+ return dynamic_env_bool(name, default)
+
+
+def _csv_set(value: str) -> set[str]:
+ return {item.strip() for item in value.split(",") if item.strip()}
+
+
+def _assert_write_permission(
+ config: Config, action: str, confirm_delete: bool = False
+) -> None:
+ if strict_permission_enforcement():
+ if not _env_bool("IOTDB_ENABLE_WRITE_DML", False):
+ raise PermissionError(
+ "Write DML tools are disabled by strict MCP policy. "
+ "Set IOTDB_ENABLE_WRITE_DML=true to enable."
+ )
+
+ allowed_users = _csv_set(dynamic_getenv("IOTDB_WRITE_ALLOWED_USERS", "root") or "root")
+ if "*" not in allowed_users and config.user not in allowed_users:
+ raise PermissionError(
+ f"Current MCP user '{config.user}' is not allowed by IOTDB_WRITE_ALLOWED_USERS."
+ )
+
+ if (
+ action == "DELETE"
+ and _env_bool("IOTDB_REQUIRE_DELETE_CONFIRM", True)
+ and not confirm_delete
+ ):
+ raise PermissionError(
+ "DELETE operation requires confirm_delete=True by server policy."
+ )
+
+
+def _normalize_sql(sql: str) -> str:
+ cleaned = sql.strip()
+ if not cleaned:
+ raise ValueError("Write SQL cannot be empty.")
+ if cleaned.endswith(";"):
+ cleaned = cleaned[:-1].strip()
+ if ";" in cleaned:
+ raise ValueError("Only a single SQL statement is allowed.")
+ return cleaned
+
+
+def register_write_tools(mcp, config: Config, logger: logging.Logger) -> None:
+ """Register write-related DML tools with per-call IoTDB target selection."""
+ max_pool_size = 100
+
+ @mcp.tool()
+ async def write_query(
+ write_sql: str,
+ confirm_delete: bool = False,
+ target_id: str | None = None,
+ target: dict[str, object] | None = None,
+ ) -> list[TextContent]:
+ """Execute exactly one write statement on the selected IoTDB target."""
+ selected_config = select_target_config(config, target_id=target_id, target=target)
+ with iotdb_target_response_context(selected_config):
+ sql = _normalize_sql(write_sql)
+ upper = sql.upper()
+
+ if selected_config.sql_dialect == "tree":
+ _, session_pool = tree_session_pool(
+ selected_config,
+ max_pool_size=max_pool_size,
+ wait_timeout_in_ms=5000,
+ tool_name="write_query",
+ )
+ session = None
+ try:
+ if upper.startswith("INSERT INTO"):
+ _assert_write_permission(selected_config, action="INSERT")
+ elif upper.startswith("DELETE FROM"):
+ _assert_write_permission(
+ selected_config,
+ action="DELETE",
+ confirm_delete=confirm_delete,
+ )
+ else:
+ raise ValueError(
+ "tree write_query only supports SQL starting with: "
+ "INSERT INTO, DELETE FROM"
+ )
+
+ session = session_pool.get_session()
+ session.execute_non_query_statement(sql)
+ session.close()
+ return sql_success_response("write_query", sql)
+ except Exception as e:
+ if session:
+ session.close()
+ logger.error(f"Failed to execute tree write_query: {str(e)}")
+ raise
+
+ _, session_pool = table_session_pool(
+ selected_config,
+ max_pool_size=max_pool_size,
+ tool_name="write_query",
+ )
+ table_session = None
+ try:
+ if upper.startswith("INSERT INTO"):
+ _assert_write_permission(selected_config, action="INSERT")
+ elif upper.startswith("UPDATE"):
+ _assert_write_permission(selected_config, action="UPDATE")
+ elif upper.startswith("DELETE FROM") or upper.startswith("DELETE DEVICES"):
+ _assert_write_permission(
+ selected_config,
+ action="DELETE",
+ confirm_delete=confirm_delete,
+ )
+ else:
+ raise ValueError(
+ "table write_query only supports SQL starting with: "
+ "INSERT INTO, UPDATE, DELETE FROM, DELETE DEVICES"
+ )
+
+ table_session = session_pool.get_session()
+ table_session.execute_non_query_statement(sql)
+ table_session.close()
+ return sql_success_response("write_query", sql)
+ except Exception as e:
+ if table_session:
+ table_session.close()
+ logger.error(f"Failed to execute table write_query: {str(e)}")
+ raise
diff --git a/src/iotdb_mcp_server/session_manager.py b/src/iotdb_mcp_server/session_manager.py
new file mode 100644
index 0000000..3321440
--- /dev/null
+++ b/src/iotdb_mcp_server/session_manager.py
@@ -0,0 +1,463 @@
+# 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.
+#
+
+from __future__ import annotations
+
+from typing import Any, Callable, Mapping
+
+from iotdb.SessionPool import PoolConfig, SessionPool
+from iotdb.table_session import TableSession
+from iotdb.table_session_pool import TableSessionPoolConfig
+from iotdb.utils.exception import IoTDBConnectionException
+
+from iotdb_mcp_server.target_registry import (
+ IoTDBTarget,
+ IoTDBTargetRegistry,
+ merge_target,
+ target_from_mapping,
+)
+
+
+def target_from_config(config: Any) -> IoTDBTarget:
+ if isinstance(config, IoTDBTarget):
+ return config
+ if hasattr(config, "to_target"):
+ return config.to_target()
+
+ keys = (
+ "target_id",
+ "display_name",
+ "target_kind",
+ "kind",
+ "host",
+ "port",
+ "node_urls",
+ "user",
+ "password",
+ "database",
+ "sql_dialect",
+ "timezone",
+ "export_path",
+ "iotdb_home",
+ "use_ssl",
+ "ca_certs",
+ "connection_timeout_in_ms",
+ "enable_redirection",
+ "enable_compression",
+ "fetch_size",
+ "max_retry",
+ "max_pool_size",
+ "tree_wait_timeout_in_ms",
+ "table_wait_timeout_in_ms",
+ "policy",
+ "last_known_good_credential",
+ "verified_at",
+ )
+ payload: dict[str, Any] = {}
+ for key in keys:
+ if hasattr(config, key):
+ payload[key] = getattr(config, key)
+ if "target_kind" in payload and "kind" not in payload:
+ payload["kind"] = payload["target_kind"]
+ return target_from_mapping(payload)
+
+
+def _target_with_overrides(
+ config_or_target: Any,
+ overrides: Mapping[str, Any] | None = None,
+) -> IoTDBTarget:
+ target = target_from_config(config_or_target)
+ clean_overrides = {
+ key: value for key, value in (overrides or {}).items() if value is not None
+ }
+ if not clean_overrides:
+ return target
+ return merge_target(target, clean_overrides)
+
+
+def make_tree_pool_config(target: IoTDBTarget) -> PoolConfig:
+ return PoolConfig(
+ host=target.host,
+ port=target.port,
+ node_urls=None,
+ user_name=target.user,
+ password=target.password,
+ fetch_size=target.fetch_size,
+ time_zone=target.timezone,
+ max_retry=target.max_retry,
+ enable_compression=target.enable_compression,
+ enable_redirection=target.enable_redirection,
+ use_ssl=target.use_ssl,
+ ca_certs=target.ca_certs or None,
+ connection_timeout_in_ms=target.connection_timeout_in_ms,
+ )
+
+
+def make_table_pool_config(
+ target: IoTDBTarget, database: str | None = None
+) -> TableSessionPoolConfig:
+ resolved_database = database if database is not None else target.database
+ return TableSessionPoolConfig(
+ node_urls=list(target.node_urls),
+ username=target.user,
+ password=target.password,
+ max_pool_size=target.max_pool_size,
+ database=resolved_database or None,
+ fetch_size=target.fetch_size,
+ time_zone=target.timezone,
+ enable_redirection=target.enable_redirection,
+ enable_compression=target.enable_compression,
+ wait_timeout_in_ms=target.table_wait_timeout_in_ms,
+ max_retry=target.max_retry,
+ use_ssl=target.use_ssl,
+ ca_certs=target.ca_certs or None,
+ connection_timeout_in_ms=target.connection_timeout_in_ms,
+ )
+
+
+def create_tree_session_pool(
+ config_or_target: Any,
+ max_pool_size: int | None = None,
+ wait_timeout_in_ms: int | None = None,
+) -> SessionPool:
+ target = _target_with_overrides(
+ config_or_target,
+ {
+ "max_pool_size": max_pool_size,
+ "tree_wait_timeout_in_ms": wait_timeout_in_ms,
+ },
+ )
+ return SessionPool(
+ make_tree_pool_config(target),
+ target.max_pool_size,
+ target.tree_wait_timeout_in_ms,
+ )
+
+
+def create_table_session_pool(
+ config_or_target: Any,
+ max_pool_size: int | None = None,
+ database: str | None = None,
+ wait_timeout_in_ms: int | None = None,
+) -> Any:
+ target = _target_with_overrides(
+ config_or_target,
+ {
+ "max_pool_size": max_pool_size,
+ "database": database,
+ "table_wait_timeout_in_ms": wait_timeout_in_ms,
+ },
+ )
+ return _SingleEndpointTableSessionPool(target)
+
+
+class _SingleEndpointTableSessionPool:
+ """TableSessionPool equivalent that opens only the target's primary endpoint."""
+
+ def __init__(self, target: IoTDBTarget) -> None:
+ self.database = target.database or None
+ self._session_pool = SessionPool(
+ make_tree_pool_config(target),
+ target.max_pool_size,
+ target.table_wait_timeout_in_ms,
+ )
+ self._session_pool.sql_dialect = "table"
+ self._session_pool.database = self.database
+
+ def get_session(self) -> TableSession:
+ return TableSession(None, session_pool=self._session_pool)
+
+ @staticmethod
... 3951 lines suppressed ...