Gather comments from CH tables as well (#13)

Author: iskakaushik

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,
         }

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()