query_and_plotly_chart

Instructions

Implementation Reference

• src/mcp_server_starrocks/server.py:324-403 (handler)

  The main handler function for the 'query_and_plotly_chart' tool. Executes an SQL query, creates a DataFrame, evaluates a Plotly expression to generate a chart, and returns the result as either JSON or an image (PNG/JPEG). Registered via @mcp.tool decorator.

  @mcp.tool(description="using sql `query` to extract data from database, then using python `plotly_expr` to generate a chart for UI to display" + description_suffix)
  def query_and_plotly_chart(
          query: Annotated[str, Field(description="SQL query to execute")],
          plotly_expr: Annotated[
              str, Field(description="a one function call expression, with 2 vars binded: `px` as `import plotly.express as px`, and `df` as dataframe generated by query `plotly_expr` example: `px.scatter(df, x=\"sepal_width\", y=\"sepal_length\", color=\"species\", marginal_y=\"violin\", marginal_x=\"box\", trendline=\"ols\", template=\"simple_white\")`")],
          format: Annotated[str, Field(description="chart output format, json|png|jpeg")] = "jpeg",
          db: Annotated[str|None, Field(description="database")] = None
  ) -> ToolResult:
      """
      Executes an SQL query, creates a Pandas DataFrame, generates a Plotly chart
      using the provided expression, encodes the chart as a base64 PNG image,
      and returns it along with optional text.
  
      Args:
          query: The SQL query string to execute.
          plotly_expr: A Python string expression using 'px' (plotly.express)
                       and 'df' (the DataFrame from the query) to generate a figure.
                       Example: "px.scatter(df, x='col1', y='col2')"
          format: chat output format, json|png|jpeg, default is jpeg
          db: Optional database name to execute the query in.
  
      Returns:
          A list containing types.TextContent and types.ImageContent,
          or just types.TextContent in case of an error or no data.
      """
      try:
          logger.info(f'query_and_plotly_chart query:{one_line_summary(query)}, plotly:{one_line_summary(plotly_expr)} format:{format}, db:{db}')
          result = db_client.execute(query, db=db, return_format="pandas")
          errmsg = None
          if not result.success:
              errmsg = result.error_message
          elif result.pandas is None:
              errmsg = 'Query did not return data suitable for plotting.'
          else:
              df = result.pandas
              if df.empty:
                  errmsg = 'Query returned no data to plot.'
          if errmsg:
              logger.warning(f"Query or data issue: {errmsg}")
              return ToolResult(
                  content=[TextContent(type='text', text=f'Error: {errmsg}')],
                  structured_content={'success': False, 'error_message': errmsg},
              )
          # Validate and evaluate the plotly expression using px and df
          local_vars = {'df': df}
          validate_plotly_expr(plotly_expr)
          fig : plotly.graph_objs.Figure = eval(plotly_expr, {"px": px}, local_vars)
          if format == 'json':
              # return json representation of the figure for front-end rendering
              plot_json = json.loads(fig.to_json())
              structured_content = result.to_dict()
              structured_content['data'] = plot_json['data']
              structured_content['layout'] = plot_json['layout']
              summary = result.to_string()
              return ToolResult(
                  content=[
                      TextContent(type='text', text=f'{summary}\nChart Generated for UI rendering'),
                  ],
                  structured_content=structured_content,
              )
          else:
              if not hasattr(fig, 'to_image'):
                  raise ToolError(f"The evaluated expression did not return a Plotly figure object. Result type: {type(fig)}")
              if format == 'jpg':
                  format = 'jpeg'
              img_bytes = fig.to_image(format=format, width=960, height=720)
              structured_content = result.to_dict()
              structured_content['img_bytes_base64'] = base64.b64encode(img_bytes)
              return ToolResult(
                  content=[
                     TextContent(type='text', text=f'dataframe data:\n{df}\nChart generated but for UI only'),
                     Image(data=img_bytes, format="jpeg").to_image_content()
                  ],
                  structured_content=structured_content
              )
      except Exception as err:
          return ToolResult(
              content=[TextContent(type='text', text=f'Error: {err}')],
              structured_content={'success': False, 'error_message': str(err)},
          )

• src/mcp_server_starrocks/server.py:325-331 (schema)

  Input schema/type definitions for the tool using Pydantic Field annotations. Parameters: query (SQL string), plotly_expr (Plotly expression string), format (output format: json|png|jpeg), db (optional database name).

  def query_and_plotly_chart(
          query: Annotated[str, Field(description="SQL query to execute")],
          plotly_expr: Annotated[
              str, Field(description="a one function call expression, with 2 vars binded: `px` as `import plotly.express as px`, and `df` as dataframe generated by query `plotly_expr` example: `px.scatter(df, x=\"sepal_width\", y=\"sepal_length\", color=\"species\", marginal_y=\"violin\", marginal_x=\"box\", trendline=\"ols\", template=\"simple_white\")`")],
          format: Annotated[str, Field(description="chart output format, json|png|jpeg")] = "jpeg",
          db: Annotated[str|None, Field(description="database")] = None
  ) -> ToolResult:

• src/mcp_server_starrocks/server.py:324-324 (registration)

  Tool registration via the @mcp.tool() decorator at line 324. The decorator registers 'query_and_plotly_chart' as an MCP tool with a description.

  @mcp.tool(description="using sql `query` to extract data from database, then using python `plotly_expr` to generate a chart for UI to display" + description_suffix)

• src/mcp_server_starrocks/server.py:256-313 (helper)

  Helper function that validates the plotly expression for security. Ensures it's a single function call on the 'px' object without nested calls or dangerous constructs.

  def validate_plotly_expr(expr: str):
      """
      Validates a string to ensure it represents a single call to a method
      of the 'px' object, without containing other statements or imports,
      and ensures its arguments do not contain nested function calls.
  
      Args:
          expr: The string expression to validate.
  
      Raises:
          ValueError: If the expression does not meet the security criteria.
          SyntaxError: If the expression is not valid Python syntax.
      """
      # 1. Check for valid Python syntax
      try:
          tree = ast.parse(expr)
      except SyntaxError as e:
          raise SyntaxError(f"Invalid Python syntax in expression: {e}") from e
  
      # 2. Check that the tree contains exactly one top-level node (statement/expression)
      if len(tree.body) != 1:
          raise ValueError("Expression must be a single statement or expression.")
  
      node = tree.body[0]
  
      # 3. Check that the single node is an expression
      if not isinstance(node, ast.Expr):
          raise ValueError(
              "Expression must be a single expression, not a statement (like assignment, function definition, import, etc.).")
  
      # 4. Get the actual value of the expression and check it's a function call
      expr_value = node.value
      if not isinstance(expr_value, ast.Call):
          raise ValueError("Expression must be a function call.")
  
      # 5. Check that the function being called is an attribute lookup (like px.scatter)
      if not isinstance(expr_value.func, ast.Attribute):
          raise ValueError("Function call must be on an object attribute (e.g., px.scatter).")
  
      # 6. Check that the attribute is being accessed on a simple variable name
      if not isinstance(expr_value.func.value, ast.Name):
          raise ValueError("Function call must be on a simple variable name (e.g., px.scatter, not obj.px.scatter).")
  
      # 7. Check that the simple variable name is 'px'
      if expr_value.func.value.id != 'px':
          raise ValueError("Function call must be on the 'px' object.")
  
      # Check positional arguments
      for i, arg_node in enumerate(expr_value.args):
          for sub_node in ast.walk(arg_node):
              if isinstance(sub_node, ast.Call):
                  raise ValueError(f"Positional argument at index {i} contains a disallowed nested function call.")
      # Check keyword arguments
      for kw in expr_value.keywords:
          for sub_node in ast.walk(kw.value):
              if isinstance(sub_node, ast.Call):
                  keyword_name = kw.arg if kw.arg else '<unknown>'
                  raise ValueError(f"Keyword argument '{keyword_name}' contains a disallowed nested function call.")

• src/mcp_server_starrocks/server.py:316-321 (helper)

  Helper function that truncates text to a one-line summary for logging purposes, used within the query_and_plotly_chart handler.

  def one_line_summary(text: str, limit:int=100) -> str:
      """Generate a one-line summary of the given text, truncated to the specified limit."""
      single_line = ' '.join(text.split())
      if len(single_line) > limit:
          return single_line[:limit-3] + '...'
      return single_line