From 450738f3cef0578ce12b2df47c7a53e225a4ad2f Mon Sep 17 00:00:00 2001 From: Kaushik Iska Date: Thu, 20 Feb 2025 13:51:04 -0600 Subject: [PATCH 1/2] Gather comments from CH tables as well --- mcp_clickhouse/mcp_server.py | 15 +++++++++++++++ tests/test_tool.py | 29 ++++++++++++++++++++++++++--- uv.lock | 2 +- 3 files changed, 42 insertions(+), 4 deletions(-) diff --git a/mcp_clickhouse/mcp_server.py b/mcp_clickhouse/mcp_server.py index b6aafc8..824e1cb 100644 --- a/mcp_clickhouse/mcp_server.py +++ b/mcp_clickhouse/mcp_server.py @@ -60,9 +60,24 @@ def get_table_info(table): create_table_query = f"SHOW CREATE TABLE {database}.`{table}`" create_table_result = client.command(create_table_query) + # Get table comment + table_comment_query = f"SELECT comment FROM system.tables WHERE database = '{database}' AND name = '{table}'" + table_comment_result = client.query(table_comment_query) + table_comment = table_comment_result.result_rows[0][0] if table_comment_result.result_rows else None + + # Get column comments + column_comments_query = f"SELECT name, comment FROM system.columns WHERE database = '{database}' AND table = '{table}'" + column_comments_result = client.query(column_comments_query) + column_comments = {row[0]: row[1] for row in column_comments_result.result_rows} + + # Add comments to columns + for column in columns: + column['comment'] = column_comments.get(column['name']) + return { "database": database, "name": table, + "comment": table_comment, "columns": columns, "create_table_query": create_table_result, } diff --git a/tests/test_tool.py b/tests/test_tool.py index 2de3a20..d8e71e3 100644 --- a/tests/test_tool.py +++ b/tests/test_tool.py @@ -17,12 +17,18 @@ def setUpClass(cls): cls.test_db = "test_tool_db" cls.test_table = "test_table" cls.client.command(f"CREATE DATABASE IF NOT EXISTS {cls.test_db}") + + # Drop table if exists to ensure clean state + cls.client.command(f"DROP TABLE IF EXISTS {cls.test_db}.{cls.test_table}") + + # Create table with comments cls.client.command(f""" - CREATE TABLE IF NOT EXISTS {cls.test_db}.{cls.test_table} ( - id UInt32, - name String + CREATE TABLE {cls.test_db}.{cls.test_table} ( + id UInt32 COMMENT 'Primary identifier', + name String COMMENT 'User name field' ) ENGINE = MergeTree() ORDER BY id + COMMENT 'Test table for unit testing' """) cls.client.command(f""" INSERT INTO {cls.test_db}.{cls.test_table} (id, name) VALUES (1, 'Alice'), (2, 'Bob') @@ -68,6 +74,23 @@ def test_run_select_query_failure(self): self.assertIsInstance(result, str) self.assertIn("error running query", result) + def test_table_and_column_comments(self): + """Test that table and column comments are correctly retrieved.""" + result = list_tables(self.test_db) + self.assertIsInstance(result, list) + self.assertEqual(len(result), 1) + + table_info = result[0] + # Verify table comment + self.assertEqual(table_info["comment"], "Test table for unit testing") + + # Get columns by name for easier testing + columns = {col["name"]: col for col in table_info["columns"]} + + # Verify column comments + self.assertEqual(columns["id"]["comment"], "Primary identifier") + self.assertEqual(columns["name"]["comment"], "User name field") + if __name__ == "__main__": unittest.main() diff --git a/uv.lock b/uv.lock index c37b231..7537024 100644 --- a/uv.lock +++ b/uv.lock @@ -218,7 +218,7 @@ wheels = [ [[package]] name = "mcp-clickhouse" -version = "0.1.0" +version = "0.1.1" source = { editable = "." } dependencies = [ { name = "clickhouse-connect" }, From 67098bda8bbd8620d21c3a59868ad33eb1bfb717 Mon Sep 17 00:00:00 2001 From: Kaushik Iska Date: Thu, 20 Feb 2025 13:59:56 -0600 Subject: [PATCH 2/2] 1 query --- mcp_clickhouse/mcp_server.py | 36 +++++++++++++++++++++--------------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/mcp_clickhouse/mcp_server.py b/mcp_clickhouse/mcp_server.py index 824e1cb..ecff216 100644 --- a/mcp_clickhouse/mcp_server.py +++ b/mcp_clickhouse/mcp_server.py @@ -44,6 +44,21 @@ def list_tables(database: str, like: str = None): query += f" LIKE '{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_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_result = client.query(column_comments_query) + column_comments = {} + for row in column_comments_result.result_rows: + table, col_name, comment = row + if table not in column_comments: + column_comments[table] = {} + column_comments[table][col_name] = comment + def get_table_info(table): logger.info(f"Getting schema info for table {database}.{table}") schema_query = f"DESCRIBE TABLE {database}.`{table}`" @@ -55,29 +70,20 @@ def get_table_info(table): column_dict = {} for i, col_name in enumerate(column_names): column_dict[col_name] = row[i] + # Add comment from our pre-fetched comments + if table in column_comments and column_dict['name'] in column_comments[table]: + column_dict['comment'] = column_comments[table][column_dict['name']] + else: + column_dict['comment'] = None columns.append(column_dict) create_table_query = f"SHOW CREATE TABLE {database}.`{table}`" create_table_result = client.command(create_table_query) - # Get table comment - table_comment_query = f"SELECT comment FROM system.tables WHERE database = '{database}' AND name = '{table}'" - table_comment_result = client.query(table_comment_query) - table_comment = table_comment_result.result_rows[0][0] if table_comment_result.result_rows else None - - # Get column comments - column_comments_query = f"SELECT name, comment FROM system.columns WHERE database = '{database}' AND table = '{table}'" - column_comments_result = client.query(column_comments_query) - column_comments = {row[0]: row[1] for row in column_comments_result.result_rows} - - # Add comments to columns - for column in columns: - column['comment'] = column_comments.get(column['name']) - return { "database": database, "name": table, - "comment": table_comment, + "comment": table_comments.get(table), "columns": columns, "create_table_query": create_table_result, }