Port/1.1.2: load comments (#4)

* 0.1.1 (#12)

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

* 1.1.2

* Update test_tool.py

---------

Co-authored-by: Philip Dubé <serprex@users.noreply.github.com>
Co-authored-by: Kaushik Iska <iska.kaushik@gmail.com>
Co-authored-by: Philip Dubé <philip@peerdb.io>

Author: 3 people

.editorconfig

@@ -0,0 +1,3 @@
+[*]
+trim_trailing_whitespace = true
+end_of_line = lf

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,
         }
@@ -108,7 +129,8 @@ def create_timeplus_client():
         f"as {client_config['username']} "
         f"(secure={client_config['secure']}, verify={client_config['verify']}, "
         f"connect_timeout={client_config['connect_timeout']}s, "
-        f"send_receive_timeout={client_config['send_receive_timeout']}s)")
+        f"send_receive_timeout={client_config['send_receive_timeout']}s)"
+    )
 
     try:
         client = timeplus_connect.get_client(**client_config)

pyproject.toml

@@ -1,17 +1,17 @@
 [project]
 name = "mcp-timeplus"
-version = "0.1.0"
+version = "0.1.2"
 description = "An MCP server for Timeplus."
 readme = "README.md"
 license = "Apache-2.0"
 license-files = ["LICENSE"]
 requires-python = ">=3.13"
 dependencies = [
-     "mcp>=1.0.0",
-     "python-dotenv>=1.0.1",
-     "fastmcp>=0.4.0",
-     "uvicorn>=0.34.0",
-     "timeplus-connect>=0.8.14",
+    "mcp>=1.0.0",
+    "python-dotenv>=1.0.1",
+    "fastmcp>=0.4.0",
+    "uvicorn>=0.34.0",
+    "timeplus-connect>=0.8.14",
 ]
 
 [project.scripts]
@@ -21,10 +21,7 @@ mcp-timeplus = "mcp_timeplus.main:main"
 Home = "https://github.com/jovezhong/mcp-timeplus"
 
 [project.optional-dependencies]
-dev = [
-    "ruff",
-    "pytest"
-]
+dev = ["ruff", "pytest"]
 
 [tool.hatch.build.targets.wheel]
 packages = ["mcp_timeplus"]

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