MCP Server for Apache Airflow

Instructions

Implementation Reference

• src/airflow/dataset.py:66-70 (handler)

  The handler function implementing the 'get_dataset' tool. It takes a dataset URI, calls the Airflow DatasetApi.get_dataset, converts the response to dict string, and returns it wrapped in MCP TextContent.

  async def get_dataset(
      uri: str,
  ) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
      response = dataset_api.get_dataset(uri=uri)
      return [types.TextContent(type="text", text=str(response.to_dict()))]

• src/airflow/dataset.py:15-15 (registration)

  Registration entry for the 'get_dataset' tool in the list returned by get_all_functions, specifying the handler, tool name, description, and read-only flag.

  (get_dataset, "get_dataset", "Get a dataset by URI", True),

• src/airflow/dataset.py:11-39 (helper)

  Helper function that returns the list of dataset-related tools for registration, including 'get_dataset'.

  def get_all_functions() -> list[tuple[Callable, str, str, bool]]:
      """Return list of (function, name, description, is_read_only) tuples for registration."""
      return [
          (get_datasets, "get_datasets", "List datasets", True),
          (get_dataset, "get_dataset", "Get a dataset by URI", True),
          (get_dataset_events, "get_dataset_events", "Get dataset events", True),
          (create_dataset_event, "create_dataset_event", "Create dataset event", False),
          (get_dag_dataset_queued_event, "get_dag_dataset_queued_event", "Get a queued Dataset event for a DAG", True),
          (get_dag_dataset_queued_events, "get_dag_dataset_queued_events", "Get queued Dataset events for a DAG", True),
          (
              delete_dag_dataset_queued_event,
              "delete_dag_dataset_queued_event",
              "Delete a queued Dataset event for a DAG",
              False,
          ),
          (
              delete_dag_dataset_queued_events,
              "delete_dag_dataset_queued_events",
              "Delete queued Dataset events for a DAG",
              False,
          ),
          (get_dataset_queued_events, "get_dataset_queued_events", "Get queued Dataset events for a Dataset", True),
          (
              delete_dataset_queued_events,
              "delete_dataset_queued_events",
              "Delete queued Dataset events for a Dataset",
              False,
          ),
      ]

• src/main.py:24-40 (registration)

  Mapping of API types to their get_all_functions, including DATASET to get_dataset_functions, used to load tools for registration.

  APITYPE_TO_FUNCTIONS = {
      APIType.CONFIG: get_config_functions,
      APIType.CONNECTION: get_connection_functions,
      APIType.DAG: get_dag_functions,
      APIType.DAGRUN: get_dagrun_functions,
      APIType.DAGSTATS: get_dagstats_functions,
      APIType.DATASET: get_dataset_functions,
      APIType.EVENTLOG: get_eventlog_functions,
      APIType.IMPORTERROR: get_importerror_functions,
      APIType.MONITORING: get_monitoring_functions,
      APIType.PLUGIN: get_plugin_functions,
      APIType.POOL: get_pool_functions,
      APIType.PROVIDER: get_provider_functions,
      APIType.TASKINSTANCE: get_taskinstance_functions,
      APIType.VARIABLE: get_variable_functions,
      APIType.XCOM: get_xcom_functions,
  }

• src/main.py:95-97 (registration)

  Code that iterates over the functions (including get_dataset) and registers them as MCP tools using app.add_tool.

  for func, name, description, *_ in functions:
      app.add_tool(Tool.from_function(func, name=name, description=description))