Orchestrate Your Data via LLMs: Meet the Dagster MCP Server

from mcp.server.fastmcp import FastMCP
import httpx
import json
from typing import Dict, Optional, Any, List, Union

class DagsterGraphqlClient:
    """Client for interacting with Dagster GraphQL API"""
    
    def __init__(self, graphql_url: str = "http://localhost:3000/graphql"):
        self.graphql_url = graphql_url
        
    async def execute_graphql(self, query: str, variables: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Execute a GraphQL query against the Dagster API"""
        async with httpx.AsyncClient() as client:
            response = await client.post(
                self.graphql_url,
                json={"query": query, "variables": variables or {}}
            )
            return response.json()
    
    async def list_repositories(self) -> Dict[str, Union[List[Dict[str, str]], str]]:
        """List all available Dagster repositories"""
        query = """
        query {
          repositoriesOrError {
            ... on RepositoryConnection {
              nodes {
                name
                location {
                  name
                }
              }
            }
          }
        }
        """
        result = await self.execute_graphql(query)
        if "errors" in result:
            return {"error": result["errors"][0]["message"]}
        
        repos = result.get("data", {}).get("repositoriesOrError", {}).get("nodes", [])
        return {"repositories": [{"name": r["name"], "location": r["location"]["name"]} for r in repos]}
    
    # See full code on github https://github.com/kyryl-opens-ml/mcp-server-dagster/blob/main/src/mcp_dagster/server.py

    async def get_asset_info(self, asset_key: str) -> Dict[str, Any]:
        """Get detailed information about a specific asset"""
        query = """
        query AssetQuery($assetKey: AssetKeyInput!) {
          assetOrError(assetKey: $assetKey) {
            ... on Asset {
              key {
                path
              }
              definition {
                description
              }
            }
            ... on AssetNotFoundError {
              message
            }
          }
        }
        """
        
        asset_key_parts = asset_key.split("/")
        variables = {
            "assetKey": {"path": asset_key_parts}
        }
        
        result = await self.execute_graphql(query, variables)
        if "errors" in result:
            return {"error": result["errors"][0]["message"]}
        
        asset_data = result.get("data", {}).get("assetOrError", {})
        
        if "message" in asset_data:
            return {"error": asset_data["message"]}
        
        formatted_asset = {
            "asset_key": "/".join(asset_data.get("key", {}).get("path", [])),
            "description": asset_data.get("definition", {}).get("description")
        }
        jobs_query = """
        query AssetJobsQuery($assetKey: AssetKeyInput!) {
          assetOrError(assetKey: $assetKey) {
            ... on Asset {
              definition {
                jobNames
              }
            }
          }
        }
        """
        
        jobs_result = await self.execute_graphql(jobs_query, variables)
        if "errors" not in jobs_result:
            jobs_data = jobs_result.get("data", {}).get("assetOrError", {}).get("definition", {})
            if "jobNames" in jobs_data:
                formatted_asset["materialized_by_jobs"] = jobs_data["jobNames"]
        
        return formatted_asset


mcp = FastMCP("Dagster Explorer", description="A simple MCP server for Dagster", dependencies=["httpx"])
dagster_client = DagsterGraphqlClient()

@mcp.tool()
async def list_repositories() -> dict:
    """List all available Dagster repositories"""
    return await dagster_client.list_repositories()

@mcp.tool()
async def list_jobs(repository_location: str, repository_name: str) -> dict:
    """List all jobs in a specific repository"""
    return await dagster_client.list_jobs(repository_location, repository_name)

@mcp.tool()
async def list_assets(repository_location: str, repository_name: str) -> dict:
    """List all assets in a specific repository"""
    return await dagster_client.list_assets(repository_location, repository_name)

@mcp.tool()
async def recent_runs(limit: int = 10) -> dict:
    """Get recent Dagster runs"""
    return await dagster_client.get_recent_runs(limit)

@mcp.tool()
async def get_run_info(run_id: str) -> dict:
    """Get detailed information about a specific run"""
    return await dagster_client.get_run_info(run_id)

@mcp.tool()
async def launch_run(repository_location: str, repository_name: str, job_name: str, run_config: str = "{}") -> dict:
    """Launch a Dagster job run"""
    return await dagster_client.launch_run(repository_location, repository_name, job_name, run_config)

@mcp.tool()
async def materialize_asset(asset_key: str, repository_location: str, repository_name: str) -> dict:
    """Materialize a specific Dagster asset"""
    return await dagster_client.materialize_asset(asset_key, repository_location, repository_name)

@mcp.tool()
async def terminate_run(run_id: str) -> dict:
    """Terminate an in-progress Dagster run"""
    return await dagster_client.terminate_run(run_id)

@mcp.tool()
async def get_asset_info(asset_key: str) -> dict:
    """Get detailed information about a specific asset"""
    return await dagster_client.get_asset_info(asset_key)

if __name__ == "__main__":
    # Run with SSE transport
    mcp.run(transport="sse")

import asyncio

from agents import Agent, Runner, ModelSettings
from agents.mcp import MCPServerSse


async def main():
    async with MCPServerSse(
        name="Dagster MCP Server",
        params={
            "url": "http://0.0.0.0:8000/sse",
        },
    ) as server:
        tools = await server.list_tools()
        print(f"Available tools: {[x.name for x in tools]}")
    
        agent = Agent(
            name="Dagster Explorer",
            instructions="""You are a Data Engineering Assistant specializing in Dagster workflows.
You help users explore their Dagster instance, monitor pipelines, and troubleshoot issues.
When interacting with the Dagster instance:
1. Start by exploring available repositories, jobs, and assets
2. Provide clear, concise explanations of pipeline components
3. Help users materialize assets and monitor runs
4. Summarize complex Dagster information in a user-friendly way
Always confirm actions before materializing assets or launching jobs.""",
            mcp_servers=[server],
            model_settings=ModelSettings(tool_choice="auto", parallel_tool_calls=True),
        )

        while True:
            user_input = input("\nEnter your question (or 'exit' to quit): ")
            if user_input.lower() == 'exit':
                break
                
            print("\nProcessing your request...")
            result = await Runner.run(starting_agent=agent, input=user_input)
            print(f"\nResponse: {result.final_output}")


if __name__ == "__main__":
    asyncio.run(main())

from dagster import asset
from pandas import DataFrame, read_html, get_dummies
from sklearn.linear_model import LinearRegression

@asset
def country_populations() -> DataFrame:
    df = read_html("https://tinyurl.com/mry64ebh")[0]
    df.columns = ["country", "pop2022", "pop2023", "change", "continent", "region"]
    df["change"] = df["change"].str.rstrip("%").str.replace("−", "-").astype("float")
    return df

@asset
def continent_change_model(country_populations: DataFrame) -> LinearRegression:
    data = country_populations.dropna(subset=["change"])
    return LinearRegression().fit(get_dummies(data[["continent"]]), data["change"])

@asset
def continent_stats(country_populations: DataFrame, continent_change_model: LinearRegression) -> DataFrame:
    result = country_populations.groupby("continent").sum()
    result["pop_change_factor"] = continent_change_model.coef_
    return result