airbyte_ops_mcp.mcp.data_ops

  1# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
  2"""MCP tools for Airbyte data team operations (the data warehouse, `airbytehq/airbyte-warehouse`).
  3
  4This module exposes tools for running on-demand dbt workflows in the warehouse repository.
  5These tools target the analytics warehouse, not Airbyte Cloud's operational database (`prod_db_ops`).
  6
  7## MCP reference
  8
  9.. include:: ../../../docs/mcp-generated/data_ops.md
 10    :start-line: 2
 11"""
 12
 13# NOTE: We intentionally do NOT use `from __future__ import annotations` here.
 14# FastMCP has issues resolving forward references when PEP 563 deferred
 15# annotations are used. See: https://github.com/jlowin/fastmcp/issues/905
 16
 17__all__: list[str] = []
 18
 19from typing import Annotated
 20
 21from fastmcp import FastMCP
 22from fastmcp_extensions import mcp_tool, register_mcp_tools
 23from pydantic import Field
 24
 25from airbyte_ops_mcp.mcp.github_ops import (
 26    TriggerCIWorkflowResult,
 27    trigger_ci_workflow,
 28)
 29
 30WAREHOUSE_DBT_ON_DEMAND_WORKFLOW_FILE = "dbt-run-command.yml"
 31
 32
 33@mcp_tool(read_only=False, idempotent=False, open_world=True)
 34def run_bigquery_dbt_job(
 35    select: Annotated[
 36        str | None,
 37        Field(description="dbt `--select` value; `None` runs the whole project."),
 38    ] = None,
 39    exclude: Annotated[
 40        str | None,
 41        Field(description="dbt `--exclude` value."),
 42    ] = None,
 43    seed: Annotated[
 44        bool,
 45        Field(description="Run `dbt seed` for all seeds; selectors are ignored."),
 46    ] = False,
 47    run: Annotated[bool, Field(description="Run selected dbt models.")] = True,
 48    snapshot: Annotated[
 49        bool,
 50        Field(description="Run selected dbt snapshots."),
 51    ] = False,
 52    test: Annotated[
 53        bool,
 54        Field(description="Run selected dbt tests."),
 55    ] = False,
 56    full_refresh: Annotated[
 57        bool,
 58        Field(
 59            description="Pass `--full-refresh` to seed and run, rebuilding "
 60            "incremental models from scratch (e.g. after adding a column)."
 61        ),
 62    ] = False,
 63    pr: Annotated[
 64        int | None,
 65        Field(
 66            description="Optional airbyte-warehouse PR number. When set, runs that "
 67            "PR's dbt code into the non-prod `airbyte-data-dev.dbt_run_pr_<N>_*` "
 68            "datasets (unselected upstream models are deferred to prod). When "
 69            "omitted, runs `main` against prod."
 70        ),
 71    ] = None,
 72    wait_for_completion: Annotated[
 73        bool,
 74        Field(
 75            description="Wait for the workflow to complete before returning. "
 76            "Set `False` to return after dispatch."
 77        ),
 78    ] = True,
 79    max_wait_seconds: Annotated[
 80        int,
 81        Field(description="Maximum time to wait for workflow completion in seconds."),
 82    ] = 600,
 83) -> TriggerCIWorkflowResult:
 84    """Run dbt on demand in the `airbytehq/airbyte-warehouse` repository.
 85
 86    Uses the `dbt Run (On Demand)` GitHub Actions workflow on `main`. Without `pr`,
 87    it runs `main` against production. With `pr`, it checks out that PR's dbt code
 88    and runs into non-prod `airbyte-data-dev.dbt_run_pr_<N>_*` datasets, deferring
 89    unselected upstream models to production. Steps run in order: seed → run →
 90    snapshot → test. `seed` runs all seeds; `select` and `exclude` are ignored for
 91    seeds.
 92
 93    `full_refresh` adds `--full-refresh` to seed and run.
 94
 95    Requires a token with `actions:write` on `airbyte-warehouse`.
 96    """
 97    if pr is not None and pr <= 0:
 98        raise ValueError("PR number must be a positive integer.")
 99
100    if not (seed or run or snapshot or test):
101        raise ValueError(
102            "At least one of seed, run, snapshot, or test must be enabled."
103        )
104
105    inputs = {
106        "seed": str(seed).lower(),
107        "run": str(run).lower(),
108        "snapshot": str(snapshot).lower(),
109        "test": str(test).lower(),
110        "full_refresh": str(full_refresh).lower(),
111    }
112    if pr is not None:
113        inputs["pr"] = str(pr)
114    if select and select.strip():
115        inputs["select"] = select
116    if exclude and exclude.strip():
117        inputs["exclude"] = exclude
118
119    return trigger_ci_workflow(
120        owner="airbytehq",
121        repo="airbyte-warehouse",
122        workflow_file=WAREHOUSE_DBT_ON_DEMAND_WORKFLOW_FILE,
123        workflow_definition_ref="main",
124        inputs=inputs,
125        wait_for_completion=wait_for_completion,
126        max_wait_seconds=max_wait_seconds,
127    )
128
129
130def register_data_ops_tools(app: FastMCP) -> None:
131    """Register data_ops tools with the FastMCP app."""
132    register_mcp_tools(app, mcp_module=__name__)