From 6c4e5e429212da60a2eae655458c87ffcdd4c25b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philip=20Dub=C3=A9?= Date: Thu, 20 Feb 2025 20:32:36 +0000 Subject: [PATCH] escape strings/identifiers in generated queries (#14) --- mcp_timeplus/mcp_server.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/mcp_timeplus/mcp_server.py b/mcp_timeplus/mcp_server.py index 290f31f..e8615a0 100644 --- a/mcp_timeplus/mcp_server.py +++ b/mcp_timeplus/mcp_server.py @@ -2,6 +2,7 @@ from typing import Sequence import timeplus_connect +from timeplus_connect.driver.binding import quote_identifier, format_query_value from dotenv import load_dotenv from fastmcp import FastMCP @@ -38,18 +39,18 @@ def list_databases(): def list_tables(database: str, like: str = None): logger.info(f"Listing tables in database '{database}'") client = create_timeplus_client() - query = f"SHOW STREAMS FROM {database}" + query = f"SHOW STREAMS FROM {quote_identifier(database)}" if like: - query += f" LIKE '{like}'" + query += f" LIKE {format_query_value(like)}" result = client.command(query) # Get all table comments in one query - table_comments_query = f"SELECT name, comment FROM system.tables WHERE database = '{database}'" + table_comments_query = f"SELECT name, comment FROM system.tables WHERE database = {format_query_value(database)}" table_comments_result = client.query(table_comments_query) table_comments = {row[0]: row[1] for row in table_comments_result.result_rows} # Get all column comments in one query - column_comments_query = f"SELECT table, name, comment FROM system.columns WHERE database = '{database}'" + column_comments_query = f"SELECT table, name, comment FROM system.columns WHERE database = {format_query_value(database)}" column_comments_result = client.query(column_comments_query) column_comments = {} for row in column_comments_result.result_rows: @@ -60,7 +61,7 @@ def list_tables(database: str, like: str = None): def get_table_info(table): logger.info(f"Getting schema info for table {database}.{table}") - schema_query = f"DESCRIBE STREAM {database}.`{table}`" + schema_query = f"DESCRIBE STREAM {quote_identifier(database)}.{quote_identifier(table)}" schema_result = client.query(schema_query) columns = []