diff --git a/mcp_clickhouse/mcp_server.py b/mcp_clickhouse/mcp_server.py index 244a6cd..c2a0fdf 100644 --- a/mcp_clickhouse/mcp_server.py +++ b/mcp_clickhouse/mcp_server.py @@ -137,15 +137,8 @@ def execute_query(query: str): try: read_only = get_readonly_setting(client) res = client.query(query, settings={"readonly": read_only}) - column_names = res.column_names - rows = [] - for row in res.result_rows: - row_dict = {} - for i, col_name in enumerate(column_names): - row_dict[col_name] = row[i] - rows.append(row_dict) - logger.info(f"Query returned {len(rows)} rows") - return rows + logger.info(f"Query returned {len(res.result_rows)} rows") + return {"columns": res.column_names, "rows": res.result_rows} except Exception as err: logger.error(f"Error executing query: {err}") # Return a structured dictionary rather than a string to ensure proper serialization diff --git a/tests/test_tool.py b/tests/test_tool.py index e931c3f..94899ea 100644 --- a/tests/test_tool.py +++ b/tests/test_tool.py @@ -62,10 +62,10 @@ def test_run_select_query_success(self): """Test running a SELECT query successfully.""" query = f"SELECT * FROM {self.test_db}.{self.test_table}" result = run_select_query(query) - self.assertIsInstance(result, list) + self.assertIsInstance(result, dict) self.assertEqual(len(result), 2) - self.assertEqual(result[0]["id"], 1) - self.assertEqual(result[0]["name"], "Alice") + self.assertEqual(result["rows"][0][0], 1) + self.assertEqual(result["rows"][0][1], "Alice") def test_run_select_query_failure(self): """Test running a SELECT query with an error."""