From f74806029d58a25cf38fe1fc28753f2249c430e8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philip=20Dub=C3=A9?= Date: Thu, 20 Feb 2025 20:26:12 +0000 Subject: [PATCH] escape strings/identifiers in generated queries --- mcp_clickhouse/mcp_server.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/mcp_clickhouse/mcp_server.py b/mcp_clickhouse/mcp_server.py index ecff216..9293873 100644 --- a/mcp_clickhouse/mcp_server.py +++ b/mcp_clickhouse/mcp_server.py @@ -2,6 +2,7 @@ from typing import Sequence import clickhouse_connect +from clickhouse_connect.driver.binding import quote_identifier, format_query_value from dotenv import load_dotenv from fastmcp import FastMCP @@ -39,18 +40,18 @@ def list_databases(): def list_tables(database: str, like: str = None): logger.info(f"Listing tables in database '{database}'") client = create_clickhouse_client() - query = f"SHOW TABLES FROM {database}" + query = f"SHOW TABLES 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: @@ -61,7 +62,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 TABLE {database}.`{table}`" + schema_query = f"DESCRIBE TABLE {quote_identifier(database)}.{quote_identifier(table)}" schema_result = client.query(schema_query) columns = []