YDB MCP

Instructions

Implementation Reference

• ydb_mcp/server.py:431-464 (handler)

  Core handler function that executes the ydb_query tool: runs SQL against YDB, handles optional parameters, processes multiple result sets, formats output as JSON TextContent.

  async def query(self, sql: str, params: Optional[Dict[str, Any]] = None) -> List[TextContent]:
      """Run a SQL query against YDB.
  
      Args:
          sql: SQL query to execute
          params: Optional query parameters
  
      Returns:
          List of TextContent objects with JSON-formatted results
      """
      # 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
          result_sets = await pool.execute_with_retries(sql, ydb_params)
          all_results = []
          for result_set in result_sets:
              processed = self._process_result_set(result_set)
              all_results.append(processed)
          # Convert all dict keys to strings for JSON serialization
          safe_result = self._stringify_dict_keys({"result_sets": all_results})
          return [TextContent(type="text", text=json.dumps(safe_result, indent=2, cls=CustomJSONEncoder))]
      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:589-598 (registration)

  Tool registration specification for ydb_query, defining name, description, handler (self.query), and input parameters schema.

  {
      "name": "ydb_query",
      "description": "Run a SQL query against YDB database",
      "handler": self.query,  # Use real handler
      "parameters": {
          "properties": {"sql": {"type": "string", "title": "Sql"}},
          "required": ["sql"],
          "type": "object",
      },
  },

• ydb_mcp/server.py:593-597 (schema)

  JSON schema defining the input parameters for the ydb_query tool: requires 'sql' string.

  "parameters": {
      "properties": {"sql": {"type": "string", "title": "Sql"}},
      "required": ["sql"],
      "type": "object",
  },

• ydb_mcp/server.py:466-502 (helper)

  Helper function used by the query handler to process YDB result sets into structured dictionaries with columns and rows.

  def _process_result_set(self, result_set):
      """Process YDB result set into a dictionary format.
  
      Args:
          result_set: YDB result set object
  
      Returns:
          Processed result set as a dictionary
      """
      try:
          # Extract columns
          columns = []
          try:
              # Get column names from the columns attribute
              columns_attr = getattr(result_set, "columns")
              columns = [col.name for col in columns_attr]
          except Exception as e:
              logger.exception(f"Error getting columns: {e}")
              return {"error": str(e), "columns": [], "rows": []}
  
          # Extract rows
          rows = []
          try:
              rows_attr = getattr(result_set, "rows")
              for row in rows_attr:
                  row_values = []
                  for i in range(len(columns)):
                      row_values.append(row[i])
                  rows.append(row_values)
          except Exception as e:
              logger.exception(f"Error getting rows: {e}")
              return {"error": str(e), "columns": columns, "rows": []}
  
          return {"columns": columns, "rows": rows}
      except Exception as e:
          logger.exception(f"Error processing result set: {e}")
          return {"error": str(e), "columns": [], "rows": []}

• ydb_mcp/server.py:1060-1062 (handler)

  Dispatch logic in call_tool method that invokes the ydb_query handler when the tool is called.

  if tool_name == "ydb_query" and "sql" in params:
      result = await self.query(sql=params["sql"])
  elif tool_name == "ydb_query_with_params" and "sql" in params and "params" in params: