JackieTien97 commented on code in PR #3:
URL: https://github.com/apache/iotdb-mcp-server/pull/3#discussion_r2043350242


##########
src/iotdb_mcp_server/server.py:
##########
@@ -46,85 +49,231 @@
     "user": config.user,
     "password": config.password,
     "database": config.database,
+    "sql_dialect": config.sql_dialect,
 }
 
 logger.info(f"IoTDB Config: {db_config}")
 
-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,
-)
+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;
 
-session_pool = TableSessionPool(session_pool_config)
+        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")
 
-@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.
+    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()
+                # # 将时间戳格式化为易读的时间格式

Review Comment:
   avoid using Chinese



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: dev-unsubscr...@iotdb.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org

Reply via email to