query_dataset
Execute Spark SQL queries on Foundry datasets to retrieve and analyze data directly through natural language commands.
Instructions
Query a dataset using Spark SQL dialect e.g. "SELECT COUNT(*) FROM dataset_rid"
Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| query | Yes |
Implementation Reference
- src/mcp_server_foundry/server.py:98-117 (handler)The handler function for the query_dataset tool. It executes a SQL query using FoundryClient's sql_queries.Query.execute, polls the status until completion, fetches results, converts Arrow IPC stream to Pandas DataFrame, and returns as JSON.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()
- src/mcp_server_foundry/server.py:97-97 (registration)The @mcp.tool() decorator registers the query_dataset function as a tool in the FastMCP server.@mcp.tool()
- Input schema: ctx (Context), query (str). Output: dict. Docstring describes usage.def query_dataset(ctx: Context, query: str) -> dict: """ Query a dataset using Spark SQL dialect e.g. "SELECT COUNT(*) FROM `dataset_rid`" """