add export tools (#7)
* add export tools
* change format
* todf using SessionDataSet
* update readme add readme-zh
update to fastmcp 2
* 更新 readme
* change notifications settings
diff --git a/.asf.yaml b/.asf.yaml
index bbc8e7f..723b370 100644
--- a/.asf.yaml
+++ b/.asf.yaml
@@ -17,6 +17,10 @@
# under the License.
#
+notifications:
+ issues: reviews@iotdb.apache.org
+ pullrequests: reviews@iotdb.apache.org
+
github:
description: "Apache IoTDB MCP Server"
homepage: https://iotdb.apache.org/
@@ -33,8 +37,8 @@
issues: true
enabled_merge_buttons:
# enable squash button:
- squash: true
+ squash: true
# enable merge button:
- merge: false
+ merge: false
# disable rebase button:
- rebase: true
+ rebase: true
diff --git a/README-zh.md b/README-zh.md
new file mode 100644
index 0000000..fcbd6c4
--- /dev/null
+++ b/README-zh.md
@@ -0,0 +1,200 @@
+# IoTDB MCP Server
+
+[](https://smithery.ai/server/@apache/iotdb-mcp-server)
+
+[English](README.md) | 中文
+
+## 概述
+
+IoTDB MCP Server 是一个基于模型上下文协议（Model Context Protocol, MCP）的服务器实现，通过 IoTDB 提供数据库交互和商业智能能力。该服务器支持执行 SQL 查询，并可以通过不同的 SQL 方言（树模型和表模型）与 IoTDB 进行交互。
+
+## 组件
+
+### 资源
+
+此服务器不暴露任何资源。
+
+### 提示
+
+此服务器不提供任何提示。
+
+### 工具
+
+服务器为 IoTDB 的树模型（Tree Model）和表模型（Table Model）提供了不同的工具。您可以通过设置 "IOTDB_SQL_DIALECT" 配置为 "tree" 或 "table" 来选择使用哪种模型。
+
+#### 树模型 (Tree Model)
+
+- `metadata_query`
+
+ - 执行 SHOW/COUNT 查询以从数据库读取元数据
+ - 输入:
+ - `query_sql` (字符串): 要执行的 SHOW/COUNT SQL 查询
+ - 支持的查询类型:
+ - SHOW DATABASES [path]
+ - SHOW TIMESERIES [path]
+ - SHOW CHILD PATHS [path]
+ - SHOW CHILD NODES [path]
+ - SHOW DEVICES [path]
+ - COUNT TIMESERIES [path]
+ - COUNT NODES [path]
+ - COUNT DEVICES [path]
+ - 返回: 查询结果作为对象数组
+
+- `select_query`
+
+ - 执行 SELECT 查询以从数据库读取数据
+ - 输入:
+ - `query_sql` (字符串): 要执行的 SELECT SQL 查询（使用树模型方言，时间使用 ISO 8601 格式，例如 2017-11-01T00:08:00.000）
+ - 支持的函数:
+ - SUM, COUNT, MAX_VALUE, MIN_VALUE, AVG, VARIANCE, MAX_TIME, MIN_TIME 等
+ - 返回: 查询结果作为对象数组
+
+- `export_query`
+ - 执行查询并将结果导出为 CSV 或 Excel 文件
+ - 输入:
+ - `query_sql` (字符串): 要执行的 SQL 查询（使用树模型方言）
+ - `format` (字符串): 导出格式，可以是 "csv" 或 "excel"（默认: "csv"）
+ - `filename` (字符串): 导出文件的文件名（可选，如果未提供，将生成唯一文件名）
+ - 返回: 有关导出文件的信息和数据预览（最多 10 行）
+
+#### 表模型 (Table Model)
+
+- `read_query`
+
+ - 执行 SELECT 查询以从数据库读取数据
+ - 输入:
+ - `query_sql` (字符串): 要执行的 SELECT SQL 查询（使用表模型方言，时间使用 ISO 8601 格式，例如 `2017-11-01T00:08:00.000`）
+ - 返回: 查询结果作为对象数组
+
+- `list_tables`
+
+ - 获取数据库中所有表的列表
+ - 无需输入参数
+ - 返回: 表名数组
+
+- `describe_table`
+
+ - 查看特定表的模式信息
+ - 输入:
+ - `table_name` (字符串): 要描述的表名
+ - 返回: 包含列名和类型的列定义数组
+
+- `export_table_query`
+ - 执行查询并将结果导出为 CSV 或 Excel 文件
+ - 输入:
+ - `query_sql` (字符串): 要执行的 SQL 查询（使用表模型方言）
+ - `format` (字符串): 导出格式，可以是 "csv" 或 "excel"（默认: "csv"）
+ - `filename` (字符串): 导出文件的文件名（可选，如果未提供，将生成唯一文件名）
+ - 返回: 有关导出文件的信息和数据预览（最多 10 行）
+
+## 配置选项
+
+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 数据库名称 |
+| --sql-dialect | IOTDB_SQL_DIALECT | table | SQL 方言: tree 或 table |
+| --export-path | IOTDB_EXPORT_PATH | /tmp | 查询结果导出路径 |
+
+## 性能优化
+
+IoTDB MCP Server 包含以下性能优化特性：
+
+1. **会话池管理**：使用优化的会话池配置，可支持最多 100 个并发会话
+2. **优化的获取大小**：对于查询，设置了 1024 的获取大小
+3. **连接重试**：配置了连接失败时的自动重试机制
+4. **超时管理**：会话等待超时设置为 5000 毫秒，提高可靠性
+5. **导出功能**：支持将查询结果导出为 CSV 或 Excel 格式
+
+## 前提条件
+
+- Python 环境
+- `uv` 包管理器
+- IoTDB 安装
+- MCP 服务器依赖项
+
+## 开发
+
+```bash
+# 克隆仓库
+git clone https://github.com/apache/iotdb-mcp-server.git
+cd iotdb-mcp-server
+
+# 创建虚拟环境
+uv venv
+source venv/bin/activate # 或在 Windows 上使用 `venv\Scripts\activate`
+
+# 安装开发依赖
+uv sync
+```
+
+## 在 Claude Desktop 中配置
+
+在 Claude Desktop 的配置文件中设置 MCP 服务器：
+
+#### macOS
+
+位置: `~/Library/Application Support/Claude/claude_desktop_config.json`
+
+#### Windows
+
+位置: `%APPDATA%/Claude/claude_desktop_config.json`
+
+**你可能需要在命令字段中放入 uv 可执行文件的完整路径。你可以通过在 macOS/Linux 上运行 `which uv` 或在 Windows 上运行 `where uv` 来获取这个路径。**
+
+### Claude Desktop 配置示例
+
+将以下配置添加到 Claude Desktop 的配置文件中：
+
+```json
+{
+ "mcpServers": {
+ "iotdb": {
+ "command": "uv",
+ "args": [
+ "--directory",
+ "/Users/your_username/iotdb-mcp-server/src/iotdb_mcp_server",
+ "run",
+ "server.py"
+ ],
+ "env": {
+ "IOTDB_HOST": "127.0.0.1",
+ "IOTDB_PORT": "6667",
+ "IOTDB_USER": "root",
+ "IOTDB_PASSWORD": "root",
+ "IOTDB_DATABASE": "test",
+ "IOTDB_SQL_DIALECT": "table",
+ "IOTDB_EXPORT_PATH": "/path/to/export/folder"
+ }
+ }
+ }
+}
+```
+
+> **注意**：请确保将 `--directory` 参数后的路径替换为你实际克隆仓库的路径。
+
+## 错误处理与日志
+
+IoTDB MCP Server 包含全面的错误处理和日志记录功能：
+
+1. **日志级别**：日志记录级别设置为 INFO，可以在控制台查看服务器的运行状态
+2. **异常处理**：所有的数据库操作都包含了异常处理，确保在出现错误时能够优雅地处理并返回有意义的错误消息
+3. **会话管理**：自动关闭已使用的会话，防止资源泄露
+4. **参数验证**：对用户输入的 SQL 查询进行基本验证，确保只有允许的查询类型被执行
+
+## Docker 支持
+
+您可以使用项目根目录下的 `Dockerfile` 构建 IoTDB MCP Server 的容器镜像：
+
+```bash
+# 构建 Docker 镜像
+docker build -t iotdb-mcp-server .
+
+# 运行容器
+docker run -e IOTDB_HOST=<your-iotdb-host> -e IOTDB_PORT=<your-iotdb-port> -e IOTDB_USER=<your-iotdb-user> -e IOTDB_PASSWORD=<your-iotdb-password> iotdb-mcp-server
+```
diff --git a/README.md b/README.md
index 457021f..5da5738 100644
--- a/README.md
+++ b/README.md
@@ -2,67 +2,127 @@
[](https://smithery.ai/server/@apache/iotdb-mcp-server)
+English | [中文](README-zh.md)
+
## Overview
-A Model Context Protocol (MCP) server implementation that provides database interaction and business intelligence capabilities through IoTDB. This server enables running SQL queries.
+
+A Model Context Protocol (MCP) server implementation that provides database interaction and business intelligence capabilities through IoTDB. This server enables running SQL queries and interacting with IoTDB using different SQL dialects (Tree Model and Table Model).
## Components
### Resources
+
The server doesn't expose any resources.
### Prompts
+
The server doesn't provide any prompts.
### 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".
#### Tree Model
+
- `metadata_query`
- - Execute SHOW/COUNT queries to read metadata from the database
- - Input:
- - `query_sql` (string): The SHOW/COUNT SQL query to execute
- - Returns: Query results as array of objects
+ - Execute SHOW/COUNT queries to read metadata from the database
+ - Input:
+ - `query_sql` (string): The SHOW/COUNT SQL query to execute
+ - Supported query types:
+ - SHOW DATABASES [path]
+ - SHOW TIMESERIES [path]
+ - SHOW CHILD PATHS [path]
+ - SHOW CHILD NODES [path]
+ - SHOW DEVICES [path]
+ - COUNT TIMESERIES [path]
+ - COUNT NODES [path]
+ - COUNT DEVICES [path]
+ - Returns: Query results as array of objects
- `select_query`
- - Execute SELECT queries to read data from the database
- - Input:
- - `query_sql` (string): The SELECT SQL query to execute
- - Returns: Query results as array of objects
+ - Execute SELECT queries to read data from the database
+ - Input:
+ - `query_sql` (string): The SELECT SQL query to execute (using TREE dialect, time using ISO 8601 format, e.g. 2017-11-01T00:08:00.000)
+ - Supported functions:
+ - SUM, COUNT, MAX_VALUE, MIN_VALUE, AVG, VARIANCE, MAX_TIME, MIN_TIME, etc.
+ - Returns: Query results as array of objects
+- `export_query`
+ - Execute a query and export the results to a CSV or Excel file
+ - Input:
+ - `query_sql` (string): The SQL query to execute (using TREE dialect)
+ - `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)
#### Table Model
##### Query Tools
+
- `read_query`
- - Execute SELECT queries to read data from the database
- - Input:
- - `query` (string): The SELECT SQL query to execute
- - Returns: Query results as array of objects
+ - Execute SELECT queries to read data from the database
+ - Input:
+ - `query_sql` (string): The SELECT SQL query to execute (using TABLE dialect, time using ISO 8601 format, e.g. 2017-11-01T00:08:00.000)
+ - Returns: Query results as array of objects
##### Schema Tools
+
- `list_tables`
- - Get a list of all tables in the database
- - No input required
- - Returns: Array of table names
-- `describe-table`
- - View schema information for a specific table
- - Input:
- - `table_name` (string): Name of table to describe
- - Returns: Array of column definitions with names and types
+ - Get a list of all tables in the database
+ - No input required
+ - Returns: Array of table names
+- `describe_table`
-## Claude Desktop Integration
+ - View schema information for a specific table
+ - Input:
+ - `table_name` (string): Name of table to describe
+ - Returns: Array of column definitions with names and types
+
+- `export_table_query`
+ - Execute a query and export the results to a CSV or Excel file
+ - Input:
+ - `query_sql` (string): The SQL query to execute (using TABLE dialect)
+ - `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)
+
+## Configuration Options
+
+IoTDB MCP Server supports the following configuration options, which can be set via environment variables or command-line arguments:
+
+| Option | Environment Variable | Default Value | Description |
+| ------------- | -------------------- | ------------- | -------------------------------- |
+| --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 |
+| --sql-dialect | IOTDB_SQL_DIALECT | table | SQL dialect: tree or table |
+| --export-path | IOTDB_EXPORT_PATH | /tmp | Path for exporting query results |
+
+## Performance Optimizations
+
+IoTDB MCP Server includes the following performance optimization features:
+
+1. **Session Pool Management**: Uses optimized session pool configurations, supporting up to 100 concurrent sessions
+2. **Optimized Fetch Size**: For queries, a fetch size of 1024 is set
+3. **Connection Retry**: Configured automatic retry mechanism for connection failures
+4. **Timeout Management**: Session wait timeout set to 5000 milliseconds for improved reliability
+5. **Export Functionality**: Support for exporting query results to CSV or Excel formats
## Prerequisites
-- Python with `uv` package manager
+
+- Python environment
+- `uv` package manager
- IoTDB installation
- MCP server dependencies
## Development
-```
+```bash
# Clone the repository
git clone https://github.com/apache/iotdb-mcp-server.git
-cd iotdb_mcp_server
+cd iotdb-mcp-server
# Create virtual environment
uv venv
@@ -72,11 +132,11 @@
uv sync
```
-
+## Claude Desktop Integration
Configure the MCP server in Claude Desktop's configuration file:
-#### MacOS
+#### macOS
Location: `~/Library/Application Support/Claude/claude_desktop_config.json`
@@ -86,6 +146,10 @@
**You may need to put the full path to the uv executable in the command field. You can get this by running `which uv` on MacOS/Linux or `where uv` on Windows.**
+### Claude Desktop Configuration Example
+
+Add the following configuration to Claude Desktop's configuration file:
+
```json
{
"mcpServers": {
@@ -93,7 +157,7 @@
"command": "uv",
"args": [
"--directory",
- "YOUR_REPO_PATH/src/iotdb_mcp_server",
+ "/Users/your_username/iotdb-mcp-server/src/iotdb_mcp_server",
"run",
"server.py"
],
@@ -103,9 +167,33 @@
"IOTDB_USER": "root",
"IOTDB_PASSWORD": "root",
"IOTDB_DATABASE": "test",
- "IOTDB_SQL_DIALECT": "table"
+ "IOTDB_SQL_DIALECT": "table",
+ "IOTDB_EXPORT_PATH": "/path/to/export/folder"
}
}
}
}
```
+
+> **Note**: Make sure to replace the `--directory` parameter's path with your actual repository clone path.
+
+## Error Handling and Logging
+
+IoTDB MCP Server includes comprehensive error handling and logging capabilities:
+
+1. **Log Level**: Logging level is set to INFO, allowing you to view server status in the console
+2. **Exception Handling**: All database operations include exception handling to ensure graceful handling and meaningful error messages when errors occur
+3. **Session Management**: Automatic closure of used sessions to prevent resource leaks
+4. **Parameter Validation**: Basic validation of user-input SQL queries to ensure only allowed query types are executed
+
+## Docker Support
+
+You can build a container image for the IoTDB MCP Server using the `Dockerfile` in the project root:
+
+```bash
+# Build Docker image
+docker build -t iotdb-mcp-server .
+
+# Run container
+docker run -e IOTDB_HOST=<your-iotdb-host> -e IOTDB_PORT=<your-iotdb-port> -e IOTDB_USER=<your-iotdb-user> -e IOTDB_PASSWORD=<your-iotdb-password> iotdb-mcp-server
+```
diff --git a/pyproject.toml b/pyproject.toml
index 2a7f0b3..b7a38ac 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,15 +1,16 @@
[project]
name = "iotdb-mcp-server"
packages = [{ include = "src/iotdb_mcp_server" }]
-version = "0.1.0"
+version = "0.2.0"
description = "A Model Context Protocol (MCP) server that enables secure interaction with IoTDB databases. This server allows AI assistants to list tables, read data, and execute SQL queries through a controlled interface, making database exploration and analysis safer and more structured."
readme = "README.md"
license = "Apache-2.0"
license-file="LICENSE"
requires-python = ">=3.12"
dependencies = [
- "mcp>=1.0.0",
- "apache-iotdb>=2.0.1b0"
+ "fastmcp>=2.8.1",
+ "apache-iotdb>=2.0.3",
+ "openpyxl>=3.1.5"
]
[build-system]
diff --git a/src/iotdb_mcp_server/__init__.py b/src/iotdb_mcp_server/__init__.py
index a2a10c4..184f75b 100644
--- a/src/iotdb_mcp_server/__init__.py
+++ b/src/iotdb_mcp_server/__init__.py
@@ -18,15 +18,55 @@
from iotdb_mcp_server.config import Config
import sys
+import signal
+import asyncio
+import logging
if not "-m" in sys.argv:
from . import server
-import asyncio
+# Configure logging for the main module
+logger = logging.getLogger("iotdb_mcp_server")
def main():
"""Main entry point for the package."""
- asyncio.run(server.main())
+ # 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
diff --git a/src/iotdb_mcp_server/config.py b/src/iotdb_mcp_server/config.py
index c571e2b..a646891 100644
--- a/src/iotdb_mcp_server/config.py
+++ b/src/iotdb_mcp_server/config.py
@@ -1,117 +1,130 @@
-# 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 argparse
-from dataclasses import dataclass
-import os
-
-
-@dataclass
-class Config:
- """
- Configuration for the IoTDB mcp server.
- """
-
- host: str
- """
- IoTDB host
- """
-
- port: int
- """
- IoTDB port
- """
-
- user: str
- """
- IoTDB username
- """
-
- password: str
- """
- IoTDB password
- """
-
- database: str
- """
- IoTDB database
- """
-
- sql_dialect: str
- """
- SQL dialect: tree or table
- """
-
- @staticmethod
- def from_env_arguments() -> "Config":
- """
- Parse command line arguments.
- """
- parser = argparse.ArgumentParser(description="IoTDB MCP Server")
-
- parser.add_argument(
- "--host",
- type=str,
- help="IoTDB host",
- default=os.getenv("IOTDB_HOST", "127.0.0.1"),
- )
-
- parser.add_argument(
- "--port",
- type=int,
- help="IoTDB MySQL protocol port",
- default=os.getenv("IOTDB_PORT", 6667),
- )
-
- parser.add_argument(
- "--user",
- type=str,
- help="IoTDB username",
- default=os.getenv("IOTDB_USER", "root"),
- )
-
- parser.add_argument(
- "--password",
- type=str,
- help="IoTDB password",
- default=os.getenv("IOTDB_PASSWORD", "root"),
- )
-
- parser.add_argument(
- "--database",
- type=str,
- help="IoTDB connect database name",
- default=os.getenv("IOTDB_DATABASE", "test"),
- )
-
- parser.add_argument(
- "--sql-dialect",
- type=str,
- help="SQL dialect: tree or table",
- default=os.getenv("IOTDB_SQL_DIALECT", "table"),
- )
-
- 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,
- )
+# 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 argparse
+from dataclasses import dataclass
+import os
+
+
+@dataclass
+class Config:
+ """
+ Configuration for the IoTDB mcp server.
+ """
+
+ host: str
+ """
+ IoTDB host
+ """
+
+ port: int
+ """
+ IoTDB port
+ """
+
+ user: str
+ """
+ IoTDB username
+ """
+
+ password: str
+ """
+ IoTDB password
+ """
+
+ database: str
+ """
+ IoTDB database
+ """
+
+ sql_dialect: str
+ """
+ SQL dialect: tree or table
+ """
+
+ export_path: str
+ """
+ Path for exporting query results
+ """
+
+ @staticmethod
+ def from_env_arguments() -> "Config":
+ """
+ Parse command line arguments.
+ """
+ parser = argparse.ArgumentParser(description="IoTDB MCP Server")
+
+ parser.add_argument(
+ "--host",
+ type=str,
+ help="IoTDB host",
+ default=os.getenv("IOTDB_HOST", "127.0.0.1"),
+ )
+
+ parser.add_argument(
+ "--port",
+ type=int,
+ help="IoTDB MySQL protocol port",
+ default=os.getenv("IOTDB_PORT", 6667),
+ )
+
+ parser.add_argument(
+ "--user",
+ type=str,
+ help="IoTDB username",
+ default=os.getenv("IOTDB_USER", "root"),
+ )
+
+ parser.add_argument(
+ "--password",
+ type=str,
+ help="IoTDB password",
+ default=os.getenv("IOTDB_PASSWORD", "root"),
+ )
+
+ parser.add_argument(
+ "--database",
+ type=str,
+ help="IoTDB connect database name",
+ default=os.getenv("IOTDB_DATABASE", "test"),
+ )
+
+ parser.add_argument(
+ "--sql-dialect",
+ type=str,
+ help="SQL dialect: tree or table",
+ default=os.getenv("IOTDB_SQL_DIALECT", "table"),
+ )
+
+ parser.add_argument(
+ "--export-path",
+ type=str,
+ help="Path for exporting query results",
+ default=os.getenv("IOTDB_EXPORT_PATH", "/tmp"),
+ )
+
+ 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,
+ )
diff --git a/src/iotdb_mcp_server/server.py b/src/iotdb_mcp_server/server.py
index 53e1fb4..6e970bf 100644
--- a/src/iotdb_mcp_server/server.py
+++ b/src/iotdb_mcp_server/server.py
@@ -1,284 +1,485 @@
-# 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 datetime
-
-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 mcp.server.fastmcp import FastMCP
-from mcp.types import (
- TextContent,
-)
-
-from iotdb_mcp_server.config import Config
-
-# 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,
-}
-
-logger.info(f"IoTDB Config: {db_config}")
-
-if config.sql_dialect == "tree":
-
- pool_config = PoolConfig(
- node_urls=[str(config.host) + ":" + str(config.port)],
- user_name=config.user,
- password=config.password,
- fetch_size=1024,
- time_zone="UTC+8",
- max_retry=3
- )
- max_pool_size = 5
- wait_timeout_in_ms = 3000
- 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 = session_pool.get_session()
- try:
- stmt = query_sql.strip().upper()
-
- # 处理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:
- raise ValueError("Unsupported metadata query. Please use one of the supported query types.")
-
- except Exception as e:
- session.close()
- raise e
-
- @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)
-
- 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 = session_pool.get_session()
- res = session.execute_query_statement(query_sql)
-
- stmt = query_sql.strip().upper()
- # Regular SELECT queries
- if (
- stmt.startswith("SELECT")
- ):
- return prepare_res(res, session)
- # Non-SELECT queries
- else:
- raise ValueError("Only SELECT queries are allowed for read_query")
-
- 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,
- 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)
- """
- table_session = session_pool.get_session()
- res = table_session.execute_query_statement(query_sql)
-
- stmt = query_sql.strip().upper()
- # Regular SELECT queries
- if (
- stmt.startswith("SELECT")
- or stmt.startswith("DESCRIBE")
- or stmt.startswith("SHOW")
- ):
- return prepare_res(res, table_session)
- # Non-SELECT queries
- else:
- raise ValueError("Only SELECT queries are allowed for read_query")
-
-
- @mcp.tool()
- async def list_tables() -> list[TextContent]:
- """List all tables in the IoTDB database."""
- 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))]
-
-
- @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 = session_pool.get_session()
- res = table_session.execute_query_statement("DESC " + table_name)
-
- return prepare_res(res, table_session)
-
-
- 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),
- )
- ]
-
-
-def main():
- logger.info("iotdb_mcp_server running with stdio transport")
- # Initialize and run the server
- mcp.run(transport="stdio")
-
-
-if __name__ == "__main__":
- main()
+# 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 datetime
+import asyncio
+import os
+import uuid
+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 fastmcp.utilities.types import TextContent
+
+from iotdb_mcp_server.config import Config
+
+# 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
+
+logger.info(f"IoTDB Config: {db_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)}")
+
+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}"
+ filepath = ""
+
+ if format.lower() == "csv":
+ if(filename.lower().endswith(".csv")):
+ filename = filename[:-4]
+ filepath = f"{config.export_path}/{filename}.csv"
+ df.to_csv(filepath, index=False)
+ elif format.lower() == "excel":
+ if(filename.lower().endswith(".xlsx")):
+ filename = filename[:-5]
+ filepath = f"{config.export_path}/{filename}.xlsx"
+ 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}"
+ filepath = ""
+
+ if format.lower() == "csv":
+ if(filename.lower().endswith(".csv")):
+ filename = filename[:-4]
+ filepath = f"{config.export_path}/{filename}.csv"
+ df.to_csv(filepath, index=False)
+ elif format.lower() == "excel":
+ if(filename.lower().endswith(".xlsx")):
+ filename = filename[:-5]
+ filepath = f"{config.export_path}/{filename}.xlsx"
+ 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),
+ )
+ ]
+
+
+def main():
+ logger.info("iotdb_mcp_server running with stdio transport")
+ # Initialize and run the server
+ mcp.run(transport="stdio")
+
+
+if __name__ == "__main__":
+ main()
