YDB MCP

Instructions

Implementation Reference

• ydb_mcp/server.py:397-430 (handler)

  The handler function for the 'ydb_explain_query_with_params' tool. It parses the JSON params string, converts to YDB parameters, and delegates to the explain_query method to get the execution plan.

  async def explain_query_with_params(self, sql: str, params: str) -> List[TextContent]:
      """Explain a SQL query against YDB
  
      Args:
          sql: SQL query to execute
          params: Optional query parameters
  
      Returns:
          Execution plan of the query as TextContent object with JSON-formatted execution plan
      """
      """Run a parameterized SQL query with JSON parameters.
  
      Args:
          sql: SQL query to execute
          params: Parameters as a JSON string
  
      Returns:
          Query results as a list of TextContent objects or a dictionary
      """
      # Handle authentication errors
      if self.auth_error:
          logger.error(f"Authentication error: {self.auth_error}")
          safe_error = self._stringify_dict_keys({"error": f"Authentication error: {self.auth_error}"})
          return [TextContent(type="text", text=json.dumps(safe_error, indent=2))]
  
      try:
          ydb_params = self._parse_str_to_ydb_params(params)
      except json.JSONDecodeError as e:
          logger.error(f"Error parsing JSON parameters: {str(e)}")
          safe_error = self._stringify_dict_keys({"error": f"Error parsing JSON parameters: {str(e)}"})
          return [TextContent(type="text", text=json.dumps(safe_error, indent=2))]
  
      return await self.explain_query(sql, ydb_params)

• ydb_mcp/server.py:641-658 (registration)

  The registration loop in register_tools() that registers the ydb_explain_query_with_params tool (along with others) using FastMCP.add_tool and ToolManager.register_tool, linking the name to the handler function.

  for spec in tool_specs:
      self.add_tool(
          spec["handler"],
          name=spec["name"],
          description=spec["description"],
  
          # Structured output is temporarily disabled until proper schema definitions are implemented.
          # See https://github.com/ydb-platform/ydb-mcp/issues/12 for details.
          structured_output=False,
      )
  
      # Also register with our tool manager
      self.tool_manager.register_tool(
          name=spec["name"],
          handler=spec["handler"],
          description=spec["description"],
          parameters=spec.get("parameters"),
      )

• ydb_mcp/server.py:575-587 (schema)

  The tool specification dictionary defining the schema (input parameters: sql string, params string) and description for ydb_explain_query_with_params.

  {
      "name": "ydb_explain_query_with_params",
      "description": "Explain a parametrized SQL query with JSON parameters",
      "handler": self.explain_query_with_params,  # Use real handler
      "parameters": {
          "properties": {
              "sql": {"type": "string", "title": "Sql"},
              "params": {"type": "string", "title": "Params"},
          },
          "required": ["sql", "params"],
          "type": "object",
      },

• ydb_mcp/server.py:359-396 (helper)

  The core explain_query helper method called by the tool handler, which performs the actual YDB explain operation using the session pool.

  async def explain_query(self, sql: str, params: Optional[Dict[str, Any]] = None) -> List[TextContent]:
      """Explain a SQL query against YDB
  
      Args:
          sql: SQL query to execute
          params: Optional query parameters
  
      Returns:
          Execution plan of the query as TextContent object with JSON-formatted execution plan
      """
      # Check if there's an authentication error
      if self.auth_error:
          return [TextContent(type="text", text=json.dumps({"error": self.auth_error}, indent=2))]
  
      try:
          pool = await self.get_pool()
          ydb_params = None
          if params:
              ydb_params = {}
              for key, value in params.items():
                  param_key = key if key.startswith("$") else f"${key}"
                  ydb_params[param_key] = value
  
          structured_plan = await pool.explain_with_retries(
              query=sql,
              parameters=ydb_params,
              result_format=ydb.QueryExplainResultFormat.DICT,
          )
  
          safe_plan = self._stringify_dict_keys(structured_plan)
          formatted_plan = json.dumps(safe_plan, indent=2, cls=CustomJSONEncoder)
          logger.info(f"Query plan: {formatted_plan}")
          return [TextContent(type="text", text=formatted_plan)]
      except Exception as e:
          error_message = str(e)
          safe_error = self._stringify_dict_keys({"error": error_message})
          return [TextContent(type="text", text=json.dumps(safe_error, indent=2))]

• ydb_mcp/server.py:533-553 (helper)

  Helper function to parse the JSON string parameters into YDB-compatible parameter dictionary, used by the tool handler.

  def _parse_str_to_ydb_params(self, params: str) -> Dict:
      parsed_params = {}
      if params and params.strip():
          parsed_params = json.loads(params)
  
      # Convert [value, type] to YDB type if needed
      ydb_params = {}
      for key, value in parsed_params.items():
          param_key = key if key.startswith("$") else f"${key}"
          if isinstance(value, (list, tuple)) and len(value) == 2:
              param_value, type_name = value
              if isinstance(type_name, str) and hasattr(ydb.PrimitiveType, type_name):
                  ydb_type = getattr(ydb.PrimitiveType, type_name)
                  ydb_params[param_key] = (param_value, ydb_type)
              else:
                  ydb_params[param_key] = param_value
          else:
              ydb_params[param_key] = value
  
      return ydb_params