Gather comments from CH tables as well (ClickHouse#13)

Author: iskakaushik

mcp_timeplus/mcp_server.py

@@ -43,6 +43,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 STREAM {database}.`{table}`"
@@ -54,6 +69,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 STREAM {database}.`{table}`"
@@ -62,6 +82,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,11 +17,17 @@ def setUpClass(cls):
         cls.test_db = "default"
         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 STREAM IF NOT EXISTS {cls.test_db}.{cls.test_table} (
                 id uint32,
                 name string
             )
+            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')
@@ -67,6 +73,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()