refactor: standardize JSON output and add error handling across tools

- Convert DataFrame results to JSON format using to_json() for consistent output
- Add try/except blocks with descriptive error messages to all query functions
- Update docstrings to reflect JSON-formatted return values
- Add new schema tools: list_sinks() and show_create_sink()
- Import json module in query_tools for enhanced JSON serialization

Author: hahaschool

src/explain.py

@@ -13,7 +13,7 @@ def explain_analyze(query: str) -> str:
         query: The SQL query to analyze streaming job (TABLE, MATERIALIZED VIEW, SINK, INDEX, or ID)
 
     Returns:
-        Query execution plan with actual statistics
+        Query execution plan with detailed runtime statistics
     """
     query_upper = query.strip().upper()
     allowed_keywords = ["TABLE", "MATERIALIZED VIEW", "SINK", "INDEX"]
@@ -26,7 +26,7 @@ def explain_analyze(query: str) -> str:
     try:
         explain_query = f"EXPLAIN ANALYZE {query}"
         result = rw.fetch(explain_query, format=OutputFormat.DATAFRAME)
-        return result
+        return result.to_json()
     except Exception as e:
         return f"Error executing EXPLAIN ANALYZE: {str(e)}"
 
@@ -42,7 +42,7 @@ def explain_query(query: str) -> str:
         query: The SQL query to explain (SELECT, INSERT, UPDATE, DELETE)
 
     Returns:
-        Query execution plan with estimated statistics
+        Query execution plan
     """
     query_upper = query.strip().upper()
     allowed_queries = ['SELECT', 'INSERT', 'UPDATE', 'DELETE', 'WITH']
@@ -59,6 +59,6 @@ def explain_query(query: str) -> str:
     try:
         explain_query = f"EXPLAIN {query}"
         result = rw.fetch(explain_query, format=OutputFormat.DATAFRAME)
-        return result
+        return result.to_json()
     except Exception as e:
         return f"Error executing EXPLAIN: {str(e)}"

src/management_tools.py

@@ -15,8 +15,11 @@ def get_database_version() -> str:
             Database version information
         """
         rw = setup_risingwave_connection()
-        result = rw.fetchone("SELECT version()", format=OutputFormat.DATAFRAME)
-        return result
+        try:
+            result = rw.fetchone("SELECT version()", format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error getting database version: {str(e)}"
 
     @mcp.tool
     def show_running_queries() -> str:
@@ -31,7 +34,7 @@ def show_running_queries() -> str:
             # This may not be available in all RisingWave versions
             result = rw.fetch("SHOW PROCESSLIST",
                               format=OutputFormat.DATAFRAME)
-            return result
+            return result.to_json()
         except Exception as e:
             return f"Show running queries not supported or error occurred: {str(e)}"
 

src/query_tools.py

@@ -1,5 +1,6 @@
 from risingwave import OutputFormat
 from connection import setup_risingwave_connection
+import json
 
 
 def run_select_query(query: str) -> str:
@@ -10,16 +11,20 @@ def run_select_query(query: str) -> str:
         query: The SELECT SQL query to execute (must start with SELECT)
 
     Returns:
-        Query results as a formatted string
+        Query results as a JSON-formatted string
     """
     query_upper = query.strip().upper()
     if not query_upper.startswith('SELECT'):
         raise ValueError(
             "Only SELECT queries are allowed for security reasons")
 
     rw = setup_risingwave_connection()
-    result = rw.fetch(query, format=OutputFormat.DATAFRAME)
-    return result
+    try:
+        result = rw.fetch(query, format=OutputFormat.DATAFRAME)
+        records = result.to_dict(orient='records')
+        return json.dumps(records, default=str, ensure_ascii=False, indent=2)
+    except Exception as e:
+        return f"Error executing SELECT query: {str(e)}"
 
 
 def table_row_count(table_name: str) -> str:
@@ -30,12 +35,16 @@ def table_row_count(table_name: str) -> str:
         table_name: Name of the table
 
     Returns:
-        Row count as a string
+        Row count as a JSON-formatted string
     """
     rw = setup_risingwave_connection()
     query = f"SELECT COUNT(*) as row_count FROM {table_name}"
-    result = rw.fetch(query, format=OutputFormat.DATAFRAME)
-    return result
+    try:
+        result = rw.fetch(query, format=OutputFormat.DATAFRAME)
+        records = result.to_dict(orient='records')
+        return json.dumps(records, default=str, ensure_ascii=False)
+    except Exception as e:
+        return f"Error getting row count for table {table_name}: {str(e)}"
 
 
 def get_table_stats(table_name: str, schema_name: str = "public") -> str:
@@ -61,7 +70,10 @@ def get_table_stats(table_name: str, schema_name: str = "public") -> str:
     FROM information_schema.columns 
     WHERE table_name = '{table_name}' AND table_schema = '{schema_name}'
     """
-    column_info = rw.fetchone(column_query, format=OutputFormat.DATAFRAME)
+    try:
+        column_info = rw.fetchone(column_query, format=OutputFormat.DATAFRAME)
+    except Exception as e:
+        return f"Error getting column info for table {table_name}: {str(e)}"
 
     stats = {
         "table": f"{schema_name}.{table_name}",

src/schema_tools.py

@@ -10,15 +10,21 @@ def register_schema_tools(mcp: FastMCP):
     def show_tables() -> str:
         """List all tables in the database."""
         rw = setup_risingwave_connection()
-        result = rw.fetch("SHOW TABLES", format=OutputFormat.DATAFRAME)
-        return result
+        try:
+            result = rw.fetch("SHOW TABLES", format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error listing tables: {str(e)}"
 
     @mcp.tool
     def list_databases() -> str:
         """List all databases in the RisingWave cluster."""
         rw = setup_risingwave_connection()
-        result = rw.fetch("SHOW DATABASES", format=OutputFormat.DATAFRAME)
-        return result
+        try:
+            result = rw.fetch("SHOW DATABASES", format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error listing databases: {str(e)}" 
 
     @mcp.tool
     def describe_table(table_name: str) -> str:
@@ -33,8 +39,11 @@ def describe_table(table_name: str) -> str:
         """
         rw = setup_risingwave_connection()
         query = f"DESCRIBE {table_name}"
-        result = rw.fetch(query, format=OutputFormat.DATAFRAME)
-        return result
+        try:
+            result = rw.fetch(query, format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error describing table {table_name}: {str(e)}"
 
     @mcp.tool
     def describe_materialized_view(mv_name: str) -> str:
@@ -49,8 +58,11 @@ def describe_materialized_view(mv_name: str) -> str:
         """
         rw = setup_risingwave_connection()
         query = f"DESCRIBE {mv_name}"
-        result = rw.fetch(query, format=OutputFormat.DATAFRAME)
-        return result
+        try:
+            result = rw.fetch(query, format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error describing materialized view {mv_name}: {str(e)}"
 
     @mcp.tool
     def show_create_table(table_name: str) -> str:
@@ -65,8 +77,11 @@ def show_create_table(table_name: str) -> str:
         """
         rw = setup_risingwave_connection()
         query = f"SHOW CREATE TABLE {table_name}"
-        result = rw.fetch(query, format=OutputFormat.DATAFRAME)
-        return result
+        try:
+            result = rw.fetch(query, format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error showing create table for {table_name}: {str(e)}"
 
     @mcp.tool
     def show_create_materialized_view(mv_name: str) -> str:
@@ -81,8 +96,11 @@ def show_create_materialized_view(mv_name: str) -> str:
         """
         rw = setup_risingwave_connection()
         query = f"SHOW CREATE MATERIALIZED VIEW {mv_name}"
-        result = rw.fetch(query, format=OutputFormat.DATAFRAME)
-        return result
+        try:
+            result = rw.fetch(query, format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error showing create materialized view for {mv_name}: {str(e)}"
 
     @mcp.tool
     def check_table_exists(table_name: str, schema_name: str = "public") -> str:
@@ -97,8 +115,11 @@ def check_table_exists(table_name: str, schema_name: str = "public") -> str:
             Boolean result as string indicating if table exists
         """
         rw = setup_risingwave_connection()
-        exists = rw.check_exist(name=table_name, schema_name=schema_name)
-        return f"Table '{table_name}' in schema '{schema_name}' exists: {exists}"
+        try:
+            exists = rw.check_exist(name=table_name, schema_name=schema_name)
+            return f"Table '{table_name}' in schema '{schema_name}' exists: {exists}"
+        except Exception as e:
+            return f"Error checking if table exists: {str(e)}"
 
     @mcp.tool
     def list_schemas() -> str:
@@ -109,9 +130,12 @@ def list_schemas() -> str:
             List of schemas as a formatted string
         """
         rw = setup_risingwave_connection()
-        result = rw.fetch(
-            "SELECT schema_name FROM information_schema.schemata", format=OutputFormat.DATAFRAME)
-        return result
+        try:
+            result = rw.fetch(
+                "SELECT schema_name FROM information_schema.schemata", format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error listing schemas: {str(e)}"
 
     @mcp.tool
     def list_materialized_views() -> str:
@@ -126,8 +150,11 @@ def list_materialized_views() -> str:
         """
         rw = setup_risingwave_connection()
         query = "SHOW MATERIALIZED VIEWS"
-        result = rw.fetch(query, format=OutputFormat.DATAFRAME)
-        return result
+        try:
+            result = rw.fetch(query, format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error listing materialized views: {str(e)}"
 
     @mcp.tool
     def get_table_columns(table_name: str, schema_name: str = "public") -> str:
@@ -148,8 +175,11 @@ def get_table_columns(table_name: str, schema_name: str = "public") -> str:
         WHERE table_name = '{table_name}' AND table_schema = '{schema_name}'
         ORDER BY ordinal_position
         """
-        result = rw.fetch(query, format=OutputFormat.DATAFRAME)
-        return result
+        try:
+            result = rw.fetch(query, format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error getting table columns for {table_name}: {str(e)}"
 
     @mcp.tool
     def list_subscriptions(schema_name: str = "public") -> str:
@@ -166,7 +196,7 @@ def list_subscriptions(schema_name: str = "public") -> str:
         query = f"SHOW SUBSCRIPTIONS FROM {schema_name}"
         try:
             result = rw.fetch(query, format=OutputFormat.DATAFRAME)
-            return result
+            return result.to_json()
         except Exception as e:
             return f"Error listing subscriptions: {str(e)}"
 
@@ -190,6 +220,40 @@ def list_table_privileges(table_name: str, schema_name: str = "public") -> str:
         """
         try:
             result = rw.fetch(query, format=OutputFormat.DATAFRAME)
-            return result
+            return result.to_json()
         except Exception as e:
             return f"Error getting table privileges: {str(e)}"
+        
+    @mcp.tool
+    def list_sinks() -> str:
+        """
+        List all sinks in the RisingWave database.
+
+        Returns:
+            List of sinks as a formatted string
+        """
+        rw = setup_risingwave_connection()
+        query = "SHOW SINKS"
+        try:
+            result = rw.fetch(query, format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error listing sinks: {str(e)}"
+        
+    @mcp.tool
+    def show_create_sink(sink_name: str) -> str:
+        """
+        Show the CREATE SINK statement for a specific sink.
+
+        Args:
+            sink_name: Name of the sink
+        Returns:
+            CREATE SINK statement
+        """
+        rw = setup_risingwave_connection()
+        query = f"SHOW CREATE SINK {sink_name}"
+        try:
+            result = rw.fetch(query, format=OutputFormat.DATAFRAME)
+            return result.to_json()
+        except Exception as e:
+            return f"Error showing create sink: {str(e)}"