Foundry MCP Server

Instructions

Implementation Reference

• src/mcp_server_foundry/server.py:97-118 (handler)

  The main handler function for the 'query_dataset' MCP tool. It is registered via the @mcp.tool() decorator. Takes a SQL query string, executes it via FoundryClient SQL queries (preview mode), polls until succeeded, fetches results as Arrow stream, converts to Pandas DataFrame, and returns as JSON.

  @mcp.tool()
  def query_dataset(ctx: Context, query: str) -> dict:
      """ Query a dataset using Spark SQL dialect e.g. "SELECT COUNT(*) FROM `dataset_rid`" """
      foundry_client: FoundryClient = ctx.request_context.lifespan_context.foundry_client
      query_id = foundry_client.sql_queries.Query.execute(
          query=query,
          preview=True
      ).query_id
  
      succeeded = False
  
      while not succeeded:
          status = foundry_client.sql_queries.Query.get_status(query_id, preview=True)
          succeeded = status.type == "succeeded"
          if status.type == "failed" or status.type == "cancelled":
              raise Exception("Query failed or cancelled")
          sleep(2)
  
      results = foundry_client.sql_queries.Query.get_results(query_id, preview=True)
  
      return pa.ipc.open_stream(results).read_all().to_pandas().to_json()