diff --git a/mcp_clickhouse/mcp_server.py b/mcp_clickhouse/mcp_server.py index b6aafc8..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,6 +70,11 @@ 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}`" @@ -63,6 +83,7 @@ def get_table_info(table): return { "database": database, "name": table, + "comment": table_comments.get(table), "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" },