diff --git a/docker-compose.yml b/docker-compose.yml index 1e36f3d86..dd2261c3a 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -54,6 +54,20 @@ services: command: './querybook/scripts/runservice scheduler --pidfile="/opt/celerybeat.pid"' volumes: *querybook-volumes depends_on: *querybook-depends-on + mcp: + container_name: querybook_mcp + image: *querybook-image + tty: true + stdin_open: true + command: "./querybook/scripts/runservice mcp" + ports: + - "${MCP_PORT:-8771}:${MCP_PORT:-8771}" + expose: + - "${MCP_PORT:-8771}" + environment: + MCP_PORT: "${MCP_PORT:-8771}" + volumes: *querybook-volumes + depends_on: *querybook-depends-on redis: container_name: querybook_redis image: redis:5.0.9 diff --git a/querybook/config/querybook_default_config.yaml b/querybook/config/querybook_default_config.yaml index 9680bdaca..03504df31 100644 --- a/querybook/config/querybook_default_config.yaml +++ b/querybook/config/querybook_default_config.yaml @@ -114,6 +114,9 @@ GITHUB_REPO_NAME: ~ GITHUB_BRANCH: 'main' GITHUB_CRYPTO_SECRET: '' +# --------------- MCP Server --------------- +MCP_PORT: 8771 + # --------------- Cache Control --------------- # Cache control settings for HTTP responses on static assets # Maximum age in seconds for static assets (CSS, JS, images, fonts) diff --git a/querybook/migrations/versions/320ccafaa979_add_mcp_event_type.py b/querybook/migrations/versions/320ccafaa979_add_mcp_event_type.py new file mode 100644 index 000000000..24ade3839 --- /dev/null +++ b/querybook/migrations/versions/320ccafaa979_add_mcp_event_type.py @@ -0,0 +1,63 @@ +"""add MCP event type + +Revision ID: 320ccafaa979 +Revises: a1b2c3d4e5f6 +Create Date: 2026-03-04 14:00:00.000000 + +""" + +from alembic import op +import sqlalchemy as sa + +# revision identifiers, used by Alembic. +revision = '320ccafaa979' +down_revision = 'a1b2c3d4e5f6' +branch_labels = None +depends_on = None + + +# Define the old and new EventType enum types +old_event_type_enum = sa.Enum("API", "WEBSOCKET", "VIEW", "CLICK", name="eventtype") +new_event_type_enum = sa.Enum("API", "WEBSOCKET", "VIEW", "CLICK", "MCP", name="eventtype") + + +def upgrade(): + """Add MCP to EventType enum""" + conn = op.get_bind() + dialect = conn.dialect.name + + if dialect == "postgresql": + # PostgreSQL: Add new enum value to existing 'eventtype' enum + op.execute("ALTER TYPE eventtype ADD VALUE 'MCP'") + else: + # Other Databases (e.g., MySQL, SQLite): Alter 'event_log.event_type' column to use the new enum + op.alter_column( + "event_log", + "event_type", + existing_type=old_event_type_enum, + type_=new_event_type_enum, + ) + + +def downgrade(): + """Remove MCP from EventType enum""" + bind = op.get_bind() + dialect = bind.dialect.name + + if dialect == "postgresql": + # PostgreSQL: does not support removing enum values directly + # We need to create a new enum without 'MCP', rename the old one, and update the column + op.execute("ALTER TYPE eventtype RENAME TO eventtype_old") + old_event_type_enum.create(bind, checkfirst=True) + op.execute( + "ALTER TABLE event_log ALTER COLUMN event_type TYPE eventtype USING event_type::text::eventtype" + ) + op.execute("DROP TYPE eventtype_old") + else: + # Other Databases (e.g., MySQL, SQLite): Revert 'event_log.event_type' column to the old enum + op.alter_column( + "event_log", + "event_type", + existing_type=new_event_type_enum, + type_=old_event_type_enum, + ) diff --git a/querybook/scripts/runservice b/querybook/scripts/runservice index f80b70d74..20729d256 100755 --- a/querybook/scripts/runservice +++ b/querybook/scripts/runservice @@ -41,6 +41,12 @@ case "$SERVICE_NAME" in "flower") COMMAND="celery -A tasks.all_tasks flower" ;; +"mcp") + COMMAND="watchmedo auto-restart -d querybook -p '*.py' -R -- python3 querybook/server/run_mcp.py" + ;; +"prod_mcp") + COMMAND="python3 querybook/server/run_mcp.py" + ;; esac if [ -z "$COMMAND" ]; then diff --git a/querybook/server/const/event_log.py b/querybook/server/const/event_log.py index bae653dc8..1bbfdb7a5 100644 --- a/querybook/server/const/event_log.py +++ b/querybook/server/const/event_log.py @@ -11,6 +11,8 @@ class EventType(Enum): VIEW = "VIEW" # a UI element gets clicked CLICK = "CLICK" + # MCP (Model Context Protocol) operation + MCP = "MCP" class FrontendEvent(TypedDict): diff --git a/querybook/server/env.py b/querybook/server/env.py index 4c634e3bf..f98f4d333 100644 --- a/querybook/server/env.py +++ b/querybook/server/env.py @@ -170,6 +170,9 @@ class QuerybookSettings(object): GITHUB_BRANCH = get_env_config("GITHUB_BRANCH") GITHUB_CRYPTO_SECRET = get_env_config("GITHUB_CRYPTO_SECRET") + # MCP Server + MCP_PORT = int(get_env_config("MCP_PORT") or 8771) + # Cache Control CACHE_CONTROL_MAX_AGE = int( get_env_config("CACHE_CONTROL_MAX_AGE") or "604800" diff --git a/querybook/server/lib/mcp/auth.py b/querybook/server/lib/mcp/auth.py new file mode 100644 index 000000000..ee3ed249e --- /dev/null +++ b/querybook/server/lib/mcp/auth.py @@ -0,0 +1,29 @@ +import hashlib + +from fastmcp.server.auth import AccessToken, TokenVerifier + +from app.db import DBSession +from models.admin import APIAccessToken + + +class QuerybookTokenVerifier(TokenVerifier): + """Validate Querybook API keys (SHA-512 hashed) against the database.""" + + async def verify_token(self, token: str) -> AccessToken | None: + token_hash = hashlib.sha512(token.encode("utf-8")).hexdigest() + with DBSession() as session: + api_token = ( + session.query(APIAccessToken) + .filter(APIAccessToken.token == token_hash) + .filter(APIAccessToken.enabled.is_(True)) + .first() + ) + if api_token is None: + return None + return AccessToken( + token=token, + client_id=str(api_token.creator_uid), + scopes=[], + expires_at=None, + claims={"creator_uid": api_token.creator_uid}, + ) diff --git a/querybook/server/lib/mcp/lib/comments.py b/querybook/server/lib/mcp/lib/comments.py new file mode 100644 index 000000000..cf062031b --- /dev/null +++ b/querybook/server/lib/mcp/lib/comments.py @@ -0,0 +1,204 @@ +"""Comment utility functions for MCP tools and resources.""" + +from collections import defaultdict + +from logic.comment import get_comment_by_id, get_comments_by_data_cell_id +from logic.datadoc import get_data_cell_by_id +from logic.datadoc_permission import user_can_read, DocDoesNotExist +from models.comment import Comment, CommentReaction, DataCellComment +from models.datadoc import DataDocDataCell + + +def serialize_reaction(reaction) -> dict: + """Serialize a CommentReaction model to dict.""" + return { + "id": reaction.id, + "reaction": reaction.reaction, + "created_by": reaction.created_by, + "created_by_resource_uri": f"querybook://user/{reaction.created_by}", + } + + +def serialize_comment(comment, reactions: list = None) -> dict: + """Serialize a Comment model to dict. + + Args: + comment: Comment model object + reactions: Optional list of CommentReaction objects for this comment + """ + return { + "id": comment.id, + "text": "" if comment.archived else comment.text, + "created_by": comment.created_by, + "created_by_resource_uri": f"querybook://user/{comment.created_by}", + "created_at": comment.created_at.isoformat() if comment.created_at else None, + "updated_at": comment.updated_at.isoformat() if comment.updated_at else None, + "archived": comment.archived, + "parent_comment_id": comment.parent_comment_id, + "reactions": [serialize_reaction(r) for r in reactions] if reactions else [], + "resource_uri": f"querybook://comment/{comment.id}", + } + + +def serialize_comments( + comments: list[Comment], + session, + include_threads: bool = True, +) -> list[dict]: + """Serialize comments with optional thread replies and reactions. + + Args: + comments: List of Comment objects to serialize + session: Database session + include_threads: If True, include thread replies for each comment + + Returns: + List of serialized comment dicts with reactions and optional thread replies + """ + if not comments: + return [] + + # Collect all comment IDs (top-level) + all_comment_ids = [c.id for c in comments] + replies_by_parent = defaultdict(list) + + if include_threads: + # Fetch thread replies for all comments + comment_ids = [c.id for c in comments] + thread_replies = ( + session.query(Comment) + .filter(Comment.parent_comment_id.in_(comment_ids)) + .order_by(Comment.created_at) + .all() + ) + + # Group thread replies by parent_comment_id + for reply in thread_replies: + replies_by_parent[reply.parent_comment_id].append(reply) + all_comment_ids.append(reply.id) + + # Batch fetch all reactions for all comments (both top-level and replies) + reactions_by_comment = defaultdict(list) + if all_comment_ids: + reactions = ( + session.query(CommentReaction) + .filter(CommentReaction.comment_id.in_(all_comment_ids)) + .all() + ) + for reaction in reactions: + reactions_by_comment[reaction.comment_id].append(reaction) + + # Serialize all comments + result = [] + for comment in comments: + comment_dict = serialize_comment( + comment, reactions_by_comment.get(comment.id, []) + ) + + # Add thread replies if they exist + if include_threads and comment.id in replies_by_parent: + comment_dict["replies"] = [ + serialize_comment(reply, reactions_by_comment.get(reply.id, [])) + for reply in replies_by_parent[comment.id] + ] + + result.append(comment_dict) + + return result + + +def get_comment_data(comment_id: int, uid: int, session) -> dict: + """Get comment data with thread replies and reactions. + + Args: + comment_id: Comment ID + uid: User ID for permission checking + session: Database session + + Returns: + Serialized comment dict with replies and reactions + + Raises: + ValueError: If comment not found or user lacks permission + """ + comment = get_comment_by_id(comment_id, session=session) + if not comment: + raise ValueError(f"Comment {comment_id} not found.") + + # Walk up to the root comment if this is a thread reply + root_comment = comment + if root_comment.parent_comment_id: + root_comment = get_comment_by_id( + root_comment.parent_comment_id, session=session + ) + if not root_comment: + raise ValueError("Parent comment not found.") + + # Find the DataDoc cell this comment belongs to + doc_cell_comment = ( + session.query(DataCellComment) + .filter(DataCellComment.comment_id == root_comment.id) + .first() + ) + if doc_cell_comment: + doc_cell = ( + session.query(DataDocDataCell) + .filter(DataDocDataCell.data_cell_id == doc_cell_comment.data_cell_id) + .first() + ) + if doc_cell: + try: + if not user_can_read(doc_cell.data_doc_id, uid, session=session): + raise ValueError( + "You do not have access to this comment's DataDoc." + ) + except DocDoesNotExist: + raise ValueError("The DataDoc for this comment was not found.") + + # Serialize with threads and reactions + result = serialize_comments([comment], session, include_threads=True) + return result[0] if result else {} + + +def get_datadoc_cell_comments_data( + cell_id: int, uid: int, include_threads: bool, session +) -> list[dict]: + """Get comments for a DataDoc cell with permission checking. + + Args: + cell_id: DataDoc cell ID + uid: User ID for permission checking + include_threads: Include thread replies for each comment + session: Database session + + Returns: + List of serialized comment dicts + + Raises: + ValueError: If cell not found or user lacks permission + """ + # Permission check + cell = get_data_cell_by_id(cell_id, session=session) + if not cell: + raise ValueError(f"DataDoc cell {cell_id} not found.") + + doc_cell = ( + session.query(DataDocDataCell) + .filter(DataDocDataCell.data_cell_id == cell_id) + .first() + ) + if not doc_cell: + raise ValueError(f"DataDoc cell {cell_id} is not associated with a DataDoc.") + datadoc_id = doc_cell.data_doc_id + + try: + if not user_can_read(datadoc_id, uid, session=session): + raise ValueError("You do not have access to this DataDoc.") + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + # Get comments + comments = get_comments_by_data_cell_id(cell_id, session=session) + + # Serialize with threads and reactions + return serialize_comments(comments, session, include_threads=include_threads) diff --git a/querybook/server/lib/mcp/lib/datadocs.py b/querybook/server/lib/mcp/lib/datadocs.py new file mode 100644 index 000000000..5b61e30e6 --- /dev/null +++ b/querybook/server/lib/mcp/lib/datadocs.py @@ -0,0 +1,326 @@ +"""DataDoc utility functions for MCP tools.""" + +from collections import defaultdict + +from lib.mcp.lib.comments import serialize_comments +from lib.mcp.lib.query_executions import serialize_query_execution +from lib.mcp.lib.schedules import serialize_schedule +from lib.mcp.utils import build_querybook_url +from logic.datadoc import ( + get_data_cell_by_id, + get_data_doc_by_id, + get_data_doc_editors_by_doc_id, + get_data_cells_executions, +) +from logic.datadoc_permission import DocDoesNotExist, user_can_read +from logic.schedule import get_task_schedule_by_name, get_data_doc_schedule_name +from models.comment import Comment, DataCellComment +from models.datadoc import DataDocDataCell, FavoriteDataDoc +from models.environment import Environment + + +def serialize_datadoc_editor(editor) -> dict: + """Serialize a DataDoc editor model with user resource URI. + + Args: + editor: DataDocEditor model object + + Returns: + Dict with all editor fields and user_resource_uri + """ + editor_dict = editor.to_dict() + editor_dict["user_resource_uri"] = f"querybook://user/{editor.uid}" + return editor_dict + + +def serialize_datadoc_cell( + cell_dict: dict, + session, + with_latest_execution: bool = False, + datadoc_id: int | None = None, +) -> dict: + """Serialize a DataDoc cell dict with resource URIs and optional enrichments. + + Args: + cell_dict: Cell dictionary (from cell.to_dict()) + session: Database session + with_latest_execution: If True, add latest_execution for query cells + datadoc_id: If provided, add datadoc_id and datadoc_resource_uri + + Returns: + Enriched cell dict with resource URIs and optional data + """ + cell_id = cell_dict["id"] + cell_dict["resource_uri"] = f"querybook://datadoc-cell/{cell_id}" + cell_dict["comments_resource_uri"] = ( + f"querybook://datadoc-cell/{cell_id}/comments?include_threads=true" + ) + + # Add executions_resource_uri for query cells + if cell_dict.get("cell_type") == "query": + cell_dict["executions_resource_uri"] = ( + f"querybook://datadoc-cell/{cell_id}/executions?limit=&offset=" + ) + + # Add latest execution if requested + if with_latest_execution: + cells_executions = get_data_cells_executions([cell_id], session=session) + if cells_executions and cells_executions[0][1]: + latest_execution = cells_executions[0][1][0] # First execution in list + cell_dict["latest_execution"] = serialize_query_execution( + latest_execution + ) + else: + cell_dict["latest_execution"] = None + + # Add parent datadoc reference if provided + if datadoc_id is not None: + cell_dict["datadoc_id"] = datadoc_id + cell_dict["datadoc_resource_uri"] = f"querybook://datadoc/{datadoc_id}" + + return cell_dict + + +def serialize_datadoc( + doc, + session, + uid: int | None = None, + with_cells: bool = False, + with_editors: bool = False, + with_schedule: bool = False, + with_comments: bool = False, +) -> dict: + """Serialize a DataDoc model to dict with all fields, resource_uri, and URL. + + Args: + doc: DataDoc model object + session: Database session + uid: User ID to check if this DataDoc is favorited by this user + with_cells: If True, include cells + with_editors: If True, include editors + with_schedule: If True, include schedule details + with_comments: If True, include comments for each cell (requires with_cells=True) + + Returns: + Dict with all datadoc fields, resource_uri, url, favorite status, and optional related data + """ + result = doc.to_dict(with_cells=with_cells) + result["owner_resource_uri"] = f"querybook://user/{doc.owner_uid}" + + # Serialize each cell with resource URIs and latest executions + if with_cells and "cells" in result: + # Batch-fetch latest executions for all query cells (single DB query) + query_cell_ids = [ + cell["id"] for cell in result["cells"] if cell.get("cell_type") == "query" + ] + latest_executions_by_cell = {} + if query_cell_ids: + cells_executions = get_data_cells_executions( + query_cell_ids, session=session + ) + for cell_id, executions in cells_executions: + if executions: + latest_executions_by_cell[cell_id] = executions[0] + + for cell in result["cells"]: + serialize_datadoc_cell(cell, session) + # Attach pre-fetched latest execution for query cells + if cell.get("cell_type") == "query": + execution = latest_executions_by_cell.get(cell["id"]) + cell["latest_execution"] = ( + serialize_query_execution(execution) if execution else None + ) + + # Check if this DataDoc is favorited by the current user + if uid is not None: + favorite = ( + session.query(FavoriteDataDoc) + .filter_by(data_doc_id=doc.id, uid=uid) + .first() + ) + result["favorite"] = favorite is not None + else: + result["favorite"] = False + + if with_editors: + editors = get_data_doc_editors_by_doc_id(doc.id, session=session) + result["editors"] = [serialize_datadoc_editor(editor) for editor in editors] + + if with_schedule: + schedule_name = get_data_doc_schedule_name(doc.id) + schedule = get_task_schedule_by_name(schedule_name, session=session) + result["schedule"] = serialize_schedule(schedule) if schedule else None + + if with_comments and with_cells and "cells" in result: + # Get all cell IDs + cell_ids = [cell["id"] for cell in result["cells"]] + + # Group comments by cell_id + comments_by_cell = defaultdict(list) + + if cell_ids: + # Batch fetch all comments for all cells + cell_comments = ( + session.query(DataCellComment) + .filter(DataCellComment.data_cell_id.in_(cell_ids)) + .all() + ) + + # Get all comment IDs + comment_ids = [cc.comment_id for cc in cell_comments] + + # Batch fetch all comments + if comment_ids: + comments = ( + session.query(Comment).filter(Comment.id.in_(comment_ids)).all() + ) + + # Serialize all comments with threads and reactions + serialized_comments = serialize_comments( + comments, session, include_threads=True + ) + + # Create a map of comment_id to serialized comment + comments_dict = {c["id"]: c for c in serialized_comments} + + # Group by cell_id + for cc in cell_comments: + if cc.comment_id in comments_dict: + comments_by_cell[cc.data_cell_id].append( + comments_dict[cc.comment_id] + ) + + # Add comments to each cell + for cell in result["cells"]: + cell["comments"] = comments_by_cell.get(cell["id"], []) + + # Add resource_uri and URL + result["resource_uri"] = f"querybook://datadoc/{doc.id}" + + environment = session.query(Environment).get(doc.environment_id) + if environment: + url = build_querybook_url(environment.name, f"datadoc/{doc.id}") + if url: + result["url"] = url + + return result + + +def get_datadoc_data(datadoc_id: int, uid: int, session) -> dict: + """Get datadoc data with permission checking. + + Args: + datadoc_id: DataDoc ID + uid: User ID for permission checking + session: Database session + + Returns: + Serialized datadoc dict + + Raises: + ValueError: If datadoc not found or user lacks permission + """ + try: + if not user_can_read(datadoc_id, uid, session=session): + raise ValueError("You do not have access to this DataDoc.") + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + doc = get_data_doc_by_id(id=datadoc_id, session=session) + return serialize_datadoc( + doc, + session, + uid=uid, + with_cells=True, + with_editors=True, + with_schedule=True, + with_comments=True, + ) + + +def get_datadoc_cell_data(cell_id: int, uid: int, session) -> dict: + """Get datadoc cell data with permission checking. + + Args: + cell_id: DataDoc cell ID + uid: User ID for permission checking + session: Database session + + Returns: + Serialized cell dict with resource_uri and comments_resource_uri + + Raises: + ValueError: If cell not found or user lacks permission + """ + cell = get_data_cell_by_id(cell_id, session=session) + if not cell: + raise ValueError(f"DataDoc cell {cell_id} not found.") + + # Get parent datadoc for permission check + doc_cell = session.query(DataDocDataCell).filter_by(data_cell_id=cell_id).first() + if not doc_cell: + raise ValueError(f"DataDoc cell {cell_id} not associated with any datadoc.") + + # Check permission on parent datadoc + try: + if not user_can_read(doc_cell.data_doc_id, uid, session=session): + raise ValueError("You do not have access to this DataDoc cell.") + except DocDoesNotExist: + raise ValueError("Parent DataDoc not found.") + + # Serialize cell with all enrichments + result = cell.to_dict() + serialize_datadoc_cell( + result, + session, + with_latest_execution=True, + datadoc_id=doc_cell.data_doc_id, + ) + return result + + +def get_datadoc_cell_executions_data( + cell_id: int, uid: int, session, limit: int = 20, offset: int = 0 +) -> list[dict]: + """Get execution history for a DataDoc cell with permission checking. + + Args: + cell_id: DataDoc cell ID + uid: User ID for permission checking + session: Database session + limit: Maximum number of results to return + offset: Number of results to skip + + Returns: + List of serialized execution dicts + + Raises: + ValueError: If cell not found or user lacks permission + """ + + cell = get_data_cell_by_id(cell_id, session=session) + if not cell: + raise ValueError(f"DataDoc cell {cell_id} not found.") + + # Get parent datadoc for permission check + doc_cell = session.query(DataDocDataCell).filter_by(data_cell_id=cell_id).first() + if not doc_cell: + raise ValueError(f"DataDoc cell {cell_id} not associated with any datadoc.") + + # Check permission on parent datadoc + try: + if not user_can_read(doc_cell.data_doc_id, uid, session=session): + raise ValueError("You do not have access to this DataDoc cell.") + except DocDoesNotExist: + raise ValueError("Parent DataDoc not found.") + + # Get executions + cells_executions = get_data_cells_executions([cell_id], session=session) + if not cells_executions or not cells_executions[0][1]: + return [] + + executions = cells_executions[0][1] + + # Slice first, then serialize (avoid serializing executions we won't return) + page = executions[offset : offset + limit] + return [serialize_query_execution(execution) for execution in page] diff --git a/querybook/server/lib/mcp/lib/environments.py b/querybook/server/lib/mcp/lib/environments.py new file mode 100644 index 000000000..97c395168 --- /dev/null +++ b/querybook/server/lib/mcp/lib/environments.py @@ -0,0 +1,64 @@ +"""Environment utility functions for MCP tools.""" + +from lib.mcp.lib.query_engines import serialize_query_engine +from logic import admin as admin_logic +from logic.environment import get_environment_by_id, get_all_visible_environments_by_uid + + +def serialize_environment( + env, uid: int | None = None, session=None, with_query_engines: bool = False +) -> dict: + """Serialize an Environment model to dict with all fields and resource_uri. + + Args: + env: Environment model object + uid: User ID for filtering accessible query engines (required if with_query_engines=True) + session: Database session (required if with_query_engines=True) + with_query_engines: If True, include list of accessible query engines with resource URIs + + Returns: + Dict with all environment fields and resource_uri + """ + # Handle both dict and model object + if hasattr(env, "to_dict"): + env_dict = env.to_dict() + else: + env_dict = {"id": env.id, "name": env.name} + + env_dict["resource_uri"] = f"querybook://environment/{env_dict['id']}" + + if with_query_engines and uid is not None and session is not None: + engines = admin_logic.get_query_engines_by_environment( + env_dict["id"], ordered=True, session=session + ) + env_dict["query_engines"] = [serialize_query_engine(qe) for qe in engines] + + return env_dict + + +def get_environment_data(environment_id: int, uid: int, session) -> dict: + """Get environment data with permission checking. + + Args: + environment_id: Environment ID + uid: User ID for permission checking + session: Database session + + Returns: + Serialized environment dict + + Raises: + ValueError: If environment not found or user lacks permission + """ + environment = get_environment_by_id(environment_id, session=session) + if not environment: + raise ValueError(f"Environment {environment_id} not found.") + + # Check if user has access to this environment + visible_envs = get_all_visible_environments_by_uid(uid, session=session) + if environment.id not in [e.id for e in visible_envs]: + raise ValueError("You do not have access to this environment.") + + return serialize_environment( + environment, uid=uid, session=session, with_query_engines=True + ) diff --git a/querybook/server/lib/mcp/lib/lists.py b/querybook/server/lib/mcp/lib/lists.py new file mode 100644 index 000000000..73005e041 --- /dev/null +++ b/querybook/server/lib/mcp/lib/lists.py @@ -0,0 +1,134 @@ +"""List utility functions for MCP tools.""" + +from lib.mcp.utils import build_querybook_url +from logic.board import get_board_editors_by_board_id +from logic.board_permission import BoardDoesNotExist, user_can_read +from models.board import Board +from models.environment import Environment + + +def serialize_list_editor(editor) -> dict: + """Serialize a List editor model with user resource URI. + + Args: + editor: BoardEditor model object + + Returns: + Dict with all editor fields and user_resource_uri + """ + editor_dict = editor.to_dict() + editor_dict["user_resource_uri"] = f"querybook://user/{editor.uid}" + return editor_dict + + +def serialize_list( + list_obj, + environment_id: int, + session, + include_items: bool = False, + with_editors: bool = False, +) -> dict: + """Serialize a List model to dict with all fields, resource_uri, and URL. + + Args: + list_obj: Board model object + environment_id: Environment ID for URL construction + session: Database session + include_items: If True, include serialized list items + with_editors: If True, include list of editors with permissions + + Returns: + Dict with all list fields, resource_uri, url, and optionally items and editors + """ + result = { + "id": list_obj.id, + "name": list_obj.name, + "description": list_obj.description, + "public": list_obj.public, + "owner_uid": list_obj.owner_uid, + "owner_resource_uri": f"querybook://user/{list_obj.owner_uid}", + "environment_id": list_obj.environment_id, + "list_type": list_obj.board_type, + "created_at": ( + list_obj.created_at.isoformat() if list_obj.created_at else None + ), + "updated_at": ( + list_obj.updated_at.isoformat() if list_obj.updated_at else None + ), + "resource_uri": f"querybook://list/{list_obj.id}", + } + + if include_items: + result["items"] = [serialize_list_item(item) for item in list_obj.items] + + if with_editors: + editors = get_board_editors_by_board_id(list_obj.id, session=session) + result["editors"] = [serialize_list_editor(editor) for editor in editors] + + environment = session.query(Environment).get(environment_id) + if environment: + url = build_querybook_url(environment.name, f"list/{list_obj.id}") + if url: + result["url"] = url + + return result + + +def serialize_list_item(item) -> dict: + """Serialize a ListItem model to a dictionary.""" + result = { + "id": item.id, + "data_doc_id": item.data_doc_id, + "table_id": item.table_id, + "list_id": item.board_id, + "query_execution_id": item.query_execution_id, + "item_order": item.item_order, + "description": item.description, + "created_at": (item.created_at.isoformat() if item.created_at else None), + } + + # Determine and add item_type (translate "board" to "list" for external API) + if item.data_doc_id is not None: + result["item_type"] = "data_doc" + result["datadoc_resource_uri"] = f"querybook://datadoc/{item.data_doc_id}" + elif item.table_id is not None: + result["item_type"] = "table" + elif item.board_id is not None: + result["item_type"] = ( + "list" # Translate from internal "board" to external "list" + ) + result["list_resource_uri"] = f"querybook://list/{item.board_id}" + elif item.query_execution_id is not None: + result["item_type"] = "query" + + return result + + +def get_list_data(list_id: int, uid: int, session) -> dict: + """Get list data with permission checking. + + Args: + list_id: List ID + uid: User ID for permission checking + session: Database session + + Returns: + Serialized list dict with items and editors + + Raises: + ValueError: If list not found or user lacks permission + """ + try: + if not user_can_read(list_id, uid, session=session): + raise ValueError("You do not have access to this list.") + except BoardDoesNotExist: + raise ValueError(f"List {list_id} not found.") + + list_obj = Board.get(id=list_id, session=session) + return serialize_list( + list_obj, + list_obj.environment_id, + session, + include_items=True, + with_editors=True, + ) diff --git a/querybook/server/lib/mcp/lib/query_engines.py b/querybook/server/lib/mcp/lib/query_engines.py new file mode 100644 index 000000000..ee5cbd302 --- /dev/null +++ b/querybook/server/lib/mcp/lib/query_engines.py @@ -0,0 +1,50 @@ +"""Query engine utility functions for MCP tools.""" + +from lib.mcp.utils import verify_query_engine_access +from logic import admin as admin_logic + + +def serialize_query_engine(engine, with_environments: bool = False) -> dict: + """Serialize query engine model with resource_uri. + + Args: + engine: QueryEngine model object + with_environments: If True, include list of environments with resource URIs + + Returns: + Dict with all query engine fields and resource_uri + """ + engine_dict = engine.to_dict() + engine_dict["resource_uri"] = f"querybook://query-engine/{engine.id}" + + if with_environments and hasattr(engine, "environments"): + # Late import to avoid circular dependency with environments.py + from lib.mcp.lib.environments import serialize_environment + + engine_dict["environments"] = [ + serialize_environment(env) for env in engine.environments + ] + + return engine_dict + + +def get_query_engine_data(engine_id: int, uid: int, session) -> dict: + """Get query engine data with permission checking. + + Args: + engine_id: Query engine ID + uid: User ID for permission checking + session: Database session + + Returns: + Serialized query engine dict + + Raises: + ValueError: If engine not found or user lacks permission + """ + # Verify user has access to this engine (also validates it exists) + verify_query_engine_access(engine_id, uid, session) + + engine = admin_logic.get_query_engine_by_id(engine_id, session=session) + + return serialize_query_engine(engine, with_environments=True) diff --git a/querybook/server/lib/mcp/lib/query_executions.py b/querybook/server/lib/mcp/lib/query_executions.py new file mode 100644 index 000000000..299144ef1 --- /dev/null +++ b/querybook/server/lib/mcp/lib/query_executions.py @@ -0,0 +1,143 @@ +"""Query execution utility functions for MCP tools.""" + +import json + +from const.query_execution import ( + QueryExecutionErrorType, + QueryExecutionStatus, + StatementExecutionStatus, +) + + +def serialize_query_execution_summary(execution) -> dict: + """Serialize query execution model for list results (lightweight, no statements).""" + execution_dict = execution.to_dict(with_statement=False) + execution_dict["query_execution_resource_uri"] = ( + f"querybook://query-execution/{execution.id}" + ) + if "status" in execution_dict: + try: + status_enum = QueryExecutionStatus(execution_dict["status"]) + execution_dict["status_name"] = status_enum.name + except (ValueError, KeyError): + pass + return execution_dict + + +def serialize_query_execution(execution) -> dict: + """Serialize query execution model with human-readable status names and resource URIs. + + This function converts a QueryExecution model to a dictionary with human-readable + status names and resource URIs, while removing internal implementation details. + + Adds: + - status_name for QueryExecutionStatus (e.g., "DONE", "ERROR") + - status_name for each StatementExecutionStatus + - results_resource_uri for each statement execution (MCP resource) + - download_url for each statement execution: + * S3/GCS stores: Pre-signed URL (no authentication needed, 24hr expiration) + * File/DB stores: Flask endpoint URL (requires api-access-token header) + + Removes: + - result_path (internal field) + - log_path (internal field) + + Args: + execution: QueryExecution model object + + Returns: + Serialized execution dict with status_name fields, resource URIs, + and internal fields removed. + """ + # Convert model to dict with statements + execution_dict = execution.to_dict(with_statement=True) + + # Add query execution resource URI + execution_dict["query_execution_resource_uri"] = ( + f"querybook://query-execution/{execution.id}" + ) + + # Add query execution status name + if "status" in execution_dict: + try: + status_enum = QueryExecutionStatus(execution_dict["status"]) + execution_dict["status_name"] = status_enum.name + except (ValueError, KeyError): + pass + + # Add statement execution status names and resource URIs + if "statement_executions" in execution_dict: + for stmt in execution_dict["statement_executions"]: + if "status" in stmt: + try: + status_enum = StatementExecutionStatus(stmt["status"]) + stmt["status_name"] = status_enum.name + except (ValueError, KeyError): + pass + + # Add resource URIs for results + if "id" in stmt: + stmt["results_resource_uri"] = ( + f"querybook://statement-execution/{stmt['id']}/results" + ) + # Add download URL for efficient large dataset downloads + # (save result_path before removing it) + result_path = stmt.get("result_path") + if result_path: + from env import QuerybookSettings + from lib.result_store import GenericReader + + try: + # For S3/GCS: Generate pre-signed URL (no auth needed) + reader = GenericReader(result_path) + if reader.has_download_url: + download_file_name = ( + f"result_{execution.id}_{stmt['id']}.csv" + ) + stmt["results_download_url"] = reader.get_download_url( + custom_name=download_file_name + ) + # For file/db stores: Use Flask endpoint (requires api-access-token) + elif QuerybookSettings.PUBLIC_URL: + stmt["results_download_url"] = ( + f"{QuerybookSettings.PUBLIC_URL}/ds/statement_execution/{stmt['id']}/result/download/" + ) + except Exception: + # If download URL generation fails, try Flask endpoint as fallback + if QuerybookSettings.PUBLIC_URL: + stmt["results_download_url"] = ( + f"{QuerybookSettings.PUBLIC_URL}/ds/statement_execution/{stmt['id']}/result/download/" + ) + + # Remove internal result_path and log_path fields + stmt.pop("result_path", None) + stmt.pop("log_path", None) + + # Include error details for failed executions + if execution.error: + error = execution.error + + # Add human-readable error type name + error_type_name = None + if error.error_type is not None: + try: + error_type_name = QueryExecutionErrorType(error.error_type).name + except (ValueError, KeyError): + pass + + # Try to parse JSON-encoded error messages (e.g. SYNTAX errors) + error_message = error.error_message + if error_message: + try: + error_message = json.loads(error_message) + except (json.JSONDecodeError, TypeError): + pass + + execution_dict["error"] = { + "error_type": error.error_type, + "error_type_name": error_type_name, + "error_message_extracted": error.error_message_extracted, + "error_message": error_message, + } + + return execution_dict diff --git a/querybook/server/lib/mcp/lib/schedules.py b/querybook/server/lib/mcp/lib/schedules.py new file mode 100644 index 000000000..a40f5f14b --- /dev/null +++ b/querybook/server/lib/mcp/lib/schedules.py @@ -0,0 +1,284 @@ +"""Schedule-related utilities for MCP.""" + +from const.schedule import TaskRunStatus +from logic.datadoc_permission import user_can_read, DocDoesNotExist +from logic.schedule import get_task_schedule_by_id, get_task_run_record_run_by_name + + +def _add_run_status_name(run_dict: dict) -> dict: + """Add human-readable status name to run dict (MCP enhancement over API). + + Args: + run_dict: Run dictionary with integer status field + + Returns: + Run dict with status_name field added + """ + if "status" in run_dict: + try: + status_enum = TaskRunStatus(run_dict["status"]) + run_dict["status_name"] = status_enum.name + except (ValueError, KeyError): + pass + return run_dict + + +def get_datadoc_id_from_schedule(schedule) -> int | None: + """Extract datadoc ID from schedule kwargs (safer than parsing name). + + Args: + schedule: TaskSchedule object + + Returns: + DataDoc ID if found in kwargs, None otherwise + """ + if schedule.kwargs and isinstance(schedule.kwargs, dict): + return schedule.kwargs.get("doc_id") + return None + + +def humanize_notifications(notifications: list[dict]) -> list[dict]: + """Convert notification 'on' values from integers to human-readable strings. + + Args: + notifications: List of notification dicts with 'on' as integer (0=all, 1=failure, 2=success) + + Returns: + List of notification dicts with 'on' as string + """ + on_map = {0: "all", 1: "failure", 2: "success"} + humanized = [] + for notif in notifications: + notif_copy = dict(notif) + if "on" in notif_copy: + notif_copy["on"] = on_map.get(notif_copy["on"], notif_copy["on"]) + # Also flatten config for easier reading + if "config" in notif_copy: + config = notif_copy["config"] + notif_copy["recipients"] = config.get("to", []) + notif_copy["user_ids"] = config.get("to_user", []) + del notif_copy["config"] + # Rename 'with' to 'notifier' for clarity + if "with" in notif_copy: + notif_copy["notifier"] = notif_copy.pop("with") + humanized.append(notif_copy) + return humanized + + +def notifications_to_kwargs(notifications) -> list[dict]: + """Convert NotificationConfig list to internal kwargs format.""" + return [ + { + "with": notif.notifier, + "on": {"all": 0, "failure": 1, "success": 2}[notif.on], + "config": { + "to": notif.recipients or [], + "to_user": notif.user_ids or [], + }, + } + for notif in notifications + ] + + +def retry_to_kwargs(retry) -> dict: + """Convert RetryConfig to internal kwargs format.""" + return { + "enabled": retry.enabled, + "max_retries": retry.max_retries, + "delay_sec": retry.delay_sec, + } + + +def exports_to_kwargs(exports) -> list[dict]: + """Convert ExportConfig list to internal kwargs format.""" + return [ + { + "exporter_cell_id": export.cell_id, + "exporter_name": export.exporter_name, + "exporter_params": export.exporter_params, + } + for export in exports + ] + + +DEFAULT_RETRY = {"enabled": False, "max_retries": 1, "delay_sec": 60} + + +def serialize_schedule(schedule) -> dict: + """Convert TaskSchedule object to dict with all fields and resource_uri. + + Args: + schedule: TaskSchedule object + + Returns: + Dict with all schedule fields, datadoc_id, resource_uri, and humanized kwargs + """ + datadoc_id = get_datadoc_id_from_schedule(schedule) + + result = { + "id": schedule.id, + "name": schedule.name, + "task": schedule.task, + "cron": schedule.cron, + "args": schedule.args, + "kwargs": schedule.kwargs, # Keep raw kwargs for backward compatibility + "options": schedule.options, + "last_run_at": ( + schedule.last_run_at.isoformat() if schedule.last_run_at else None + ), + "total_run_count": schedule.total_run_count, + "enabled": schedule.enabled, + "datadoc_id": datadoc_id, + "resource_uri": f"querybook://schedule/{schedule.id}", + "runs_resource_uri": f"querybook://schedule/{schedule.id}/runs?limit=&offset=&hide_successful=", + } + + # Add datadoc resource_uri if datadoc_id exists + if datadoc_id is not None: + result["datadoc_resource_uri"] = f"querybook://datadoc/{datadoc_id}" + + # Add humanized versions of kwargs at top level (matching Pydantic model format) + # Raw internal format is still available in kwargs for compatibility + if schedule.kwargs and isinstance(schedule.kwargs, dict): + kwargs = schedule.kwargs + + if "notifications" in kwargs and kwargs["notifications"]: + result["notifications"] = humanize_notifications(kwargs["notifications"]) + else: + result["notifications"] = [] + + if "retry" in kwargs: + result["retry"] = kwargs["retry"] + else: + result["retry"] = dict(DEFAULT_RETRY) + + if "exports" in kwargs and kwargs["exports"]: + result["exports"] = kwargs["exports"] + else: + result["exports"] = [] + + if "disable_if_running_doc" in kwargs: + result["disable_if_running"] = kwargs["disable_if_running_doc"] + else: + result["disable_if_running"] = False + + return result + + +def _get_schedule_with_read_permission(schedule_id: int, uid: int, session): + """Fetch a schedule and verify the user has read access to its DataDoc. + + Args: + schedule_id: Schedule ID + uid: User ID for permission checking + session: Database session + + Returns: + TaskSchedule object + + Raises: + ValueError: If schedule not found, not a DataDoc schedule, or user lacks permission + """ + schedule = get_task_schedule_by_id(schedule_id, session=session) + if not schedule: + raise ValueError(f"Schedule {schedule_id} not found.") + + datadoc_id = get_datadoc_id_from_schedule(schedule) + if datadoc_id is None: + raise ValueError("Schedule is not a DataDoc schedule.") + + try: + if not user_can_read(datadoc_id, uid, session=session): + raise ValueError("You do not have access to this schedule.") + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + return schedule + + +def get_schedule_runs_data( + schedule_id: int, + uid: int, + session, + limit: int = 20, + offset: int = 0, + hide_successful: bool = False, +) -> list[dict]: + """Get execution history for a schedule with permission checking. + + Args: + schedule_id: Schedule ID + uid: User ID for permission checking + session: Database session + limit: Maximum number of results + offset: Pagination offset + hide_successful: Hide successful runs, show only failures + + Returns: + List of serialized run dicts + + Raises: + ValueError: If schedule not found, not a DataDoc schedule, or user lacks permission + """ + schedule = _get_schedule_with_read_permission(schedule_id, uid, session) + + runs, count = get_task_run_record_run_by_name( + schedule.name, + limit=limit, + offset=offset, + hide_successful_jobs=hide_successful, + session=session, + ) + + return [ + _add_run_status_name( + { + "id": run.id, + "schedule_name": run.name, + "status": run.status.value, # Convert enum to integer + "error_message": run.error_message, + "created_at": (run.created_at.isoformat() if run.created_at else None), + "updated_at": (run.updated_at.isoformat() if run.updated_at else None), + } + ) + for run in runs + ] + + +def get_schedule_data(schedule_id: int, uid: int, session) -> dict: + """Get schedule data with permission checking and recent runs. + + Args: + schedule_id: Schedule ID + uid: User ID for permission checking + session: Database session + + Returns: + Serialized schedule dict with recent runs + + Raises: + ValueError: If schedule not found, not a DataDoc schedule, or user lacks permission + """ + schedule = _get_schedule_with_read_permission(schedule_id, uid, session) + + # Get recent runs + runs, count = get_task_run_record_run_by_name( + schedule.name, limit=5, session=session + ) + + # Serialize schedule and add recent runs + result = serialize_schedule(schedule) + result["recent_runs"] = [ + _add_run_status_name( + { + "id": run.id, + "status": run.status.value, # Convert enum to integer + "error_message": run.error_message, + "created_at": (run.created_at.isoformat() if run.created_at else None), + "updated_at": (run.updated_at.isoformat() if run.updated_at else None), + } + ) + for run in runs + ] + + return result diff --git a/querybook/server/lib/mcp/lib/users.py b/querybook/server/lib/mcp/lib/users.py new file mode 100644 index 000000000..6862809c8 --- /dev/null +++ b/querybook/server/lib/mcp/lib/users.py @@ -0,0 +1,57 @@ +"""User utility functions for MCP tools.""" + +from logic.user import get_user_by_id + + +def serialize_user(user) -> dict: + """Serialize a User model to dict with all fields and resource_uri. + + Args: + user: User model object + + Returns: + Dict with all user fields and resource_uri + """ + user_dict = { + "id": user.id, + "username": user.username, + "fullname": user.fullname, + "email": user.email, + "deleted": user.deleted, + "is_group": user.is_group, + "properties": user.properties.get("public_info", {}), + "resource_uri": f"querybook://user/{user.id}", + } + return user_dict + + +def get_user_data(user_id: int | str, uid: int, session) -> dict: + """Get user data. + + Args: + user_id: User ID or 'me' for current user + uid: Current user ID (for resolving 'me') + session: Database session + + Returns: + Serialized user dict + + Raises: + ValueError: If user_id is invalid or user not found + """ + # Handle 'me' as current user + if user_id == "me": + target_uid = uid + elif isinstance(user_id, int): + target_uid = user_id + else: + try: + target_uid = int(user_id) + except ValueError: + raise ValueError(f"Invalid user_id: {user_id}. Must be an integer or 'me'.") + + user = get_user_by_id(target_uid, session=session) + if not user: + raise ValueError(f"User {target_uid} not found.") + + return serialize_user(user) diff --git a/querybook/server/lib/mcp/middleware.py b/querybook/server/lib/mcp/middleware.py new file mode 100644 index 000000000..944003400 --- /dev/null +++ b/querybook/server/lib/mcp/middleware.py @@ -0,0 +1,271 @@ +"""MCP Event Logging Middleware and Decorators + +Tracks all MCP tool calls and resource reads using Querybook's existing EventLog system. +Uses a dedicated EventType.MCP for clean semantic separation from REST API events. + +NOTE: FastMCP middleware only supports tool hooks. Resource logging is handled via +a decorator pattern since on_read_resource() is not supported. +""" + +import functools +import time + +from fastmcp.server.dependencies import get_access_token +from fastmcp.server.middleware import Middleware, MiddlewareContext + +from const.event_log import EventType +from lib.event_logger import event_logger +from lib.logger import get_logger + +LOG = get_logger(__file__) + +# Maximum length for string parameters (same as BaseEventLogger) +MAX_STR_PARAM_LENGTH = 128 + + +class MCPEventLoggingMiddleware(Middleware): + """Middleware that logs MCP tool calls and resource reads to Querybook's event log.""" + + async def on_call_tool(self, context: MiddlewareContext, call_next): + """Log MCP tool calls with timing and parameters. + + Args: + context: FastMCP middleware context containing tool name and arguments + call_next: Next middleware/handler in the chain + + Returns: + Tool execution result + + Raises: + Exception: Re-raises any exceptions after logging them + """ + start_time = time.perf_counter() + tool_name = context.message.name + user_id = self._get_user_id(context) + + try: + result = await call_next(context) + duration_ms = (time.perf_counter() - start_time) * 1000 + + # Log successful tool execution + _log_mcp_event( + user_id=user_id, + event_data={ + "operation_type": "tool", + "tool": tool_name, + "status": "success", + "duration_ms": round(duration_ms, 2), + "parameters": self._sanitize_params(context.message.arguments), + }, + ) + return result + + except Exception as e: + duration_ms = (time.perf_counter() - start_time) * 1000 + + # Log failed tool execution + _log_mcp_event( + user_id=user_id, + event_data={ + "operation_type": "tool", + "tool": tool_name, + "status": "error", + "error": str(e)[:MAX_STR_PARAM_LENGTH], + "duration_ms": round(duration_ms, 2), + "parameters": self._sanitize_params(context.message.arguments), + }, + ) + raise + + def _get_user_id(self, context: MiddlewareContext) -> int: + """Extract user ID from the current request's access token. + + Uses FastMCP's get_access_token() which reads from the HTTP request + scope or SDK context var, matching how CurrentAccessToken() resolves + in tool/resource handlers. + + Args: + context: FastMCP middleware context + + Returns: + User ID from access token claims, or 0 if not available + """ + try: + token = get_access_token() + if token: + return token.claims.get("creator_uid", 0) + except Exception as e: + LOG.warning(f"Failed to extract user ID from MCP context: {e}") + + return 0 + + def _sanitize_params(self, params: dict) -> dict: + """Sanitize parameters by trimming long strings. + + Recursively processes nested dictionaries and trims strings longer than + MAX_STR_PARAM_LENGTH to prevent logging large payloads. + + Args: + params: Dictionary of parameters to sanitize + + Returns: + New dictionary with sanitized parameters + """ + if not isinstance(params, dict): + return params + + sanitized = {} + for key, value in params.items(): + if isinstance(value, str) and len(value) > MAX_STR_PARAM_LENGTH: + sanitized[key] = value[:MAX_STR_PARAM_LENGTH] + "..." + elif isinstance(value, dict): + sanitized[key] = self._sanitize_params(value) + elif isinstance(value, list): + sanitized[key] = self._sanitize_list(value) + else: + sanitized[key] = value + + return sanitized + + def _sanitize_list(self, items: list) -> list: + """Sanitize list items recursively. + + Args: + items: List to sanitize + + Returns: + New list with sanitized items + """ + sanitized = [] + for item in items: + if isinstance(item, str) and len(item) > MAX_STR_PARAM_LENGTH: + sanitized.append(item[:MAX_STR_PARAM_LENGTH] + "...") + elif isinstance(item, dict): + sanitized.append(self._sanitize_params(item)) + elif isinstance(item, list): + sanitized.append(self._sanitize_list(item)) + else: + sanitized.append(item) + + return sanitized + + +def wrap_mcp_resources(mcp): + """Wrap FastMCP's resource decorator to add logging. + + FastMCP middleware doesn't support resource hooks, so we monkey-patch + the resource decorator to automatically add logging to all resources. + + Args: + mcp: FastMCP instance to wrap + + Returns: + The same FastMCP instance with wrapped resource decorator + """ + original_resource = mcp.resource + + def logging_resource(*decorator_args, **decorator_kwargs): + """Wrapped resource decorator that adds logging.""" + + def decorator(func): + """Actual decorator applied to resource functions.""" + + @functools.wraps(func) + def wrapper(*args, **kwargs): + start_time = time.perf_counter() + + # Extract user ID from token + token = None + for arg in args: + if ( + hasattr(arg, "__class__") + and arg.__class__.__name__ == "AccessToken" + ): + token = arg + break + if not token: + token = kwargs.get("token") + + user_id = 0 + if token and hasattr(token, "claims"): + user_id = token.claims.get("creator_uid", 0) + + # Extract resource URI from decorator kwargs if available + resource_uri = decorator_kwargs.get( + "uri", f"{func.__module__}.{func.__name__}" + ) + + # Substitute parameters in URI template (e.g., {datadoc_id}) + for key, value in kwargs.items(): + placeholder = f"{{{key}}}" + if placeholder in resource_uri: + resource_uri = resource_uri.replace(placeholder, str(value)) + + # Remove query string template from URI + resource_uri = resource_uri.split("{?")[0] + + try: + result = func(*args, **kwargs) + duration_ms = (time.perf_counter() - start_time) * 1000 + + # Log successful resource read + _log_mcp_event( + user_id=user_id, + event_data={ + "operation_type": "resource", + "resource_uri": resource_uri, + "status": "success", + "duration_ms": round(duration_ms, 2), + }, + ) + return result + + except Exception as e: + duration_ms = (time.perf_counter() - start_time) * 1000 + + # Log failed resource read + _log_mcp_event( + user_id=user_id, + event_data={ + "operation_type": "resource", + "resource_uri": resource_uri, + "status": "error", + "error": str(e)[:MAX_STR_PARAM_LENGTH], + "duration_ms": round(duration_ms, 2), + }, + ) + raise + + # Apply original resource decorator to wrapped function + return original_resource(*decorator_args, **decorator_kwargs)(wrapper) + + return decorator + + # Replace mcp.resource with our wrapped version + mcp.resource = logging_resource + return mcp + + +def _log_mcp_event(user_id: int, event_data: dict) -> None: + """Log MCP event using EventLogger singleton. + + Wraps event_logger.log() with additional error handling to prevent + logging failures from interrupting MCP operations. + + Args: + user_id: ID of the user performing the action + event_data: Event data dictionary + """ + try: + # Note: event_logger.log() normally gets user_id from current_user (Flask context) + # but in MCP we don't have Flask context, so we need to pass uid directly + # to the underlying logger. Since event_logger.log() doesn't accept uid parameter, + # we call the underlying logger directly. + event_logger.logger.log( + uid=user_id, + event_type=EventType.MCP, + event_data=event_data, + ) + except Exception as e: + # Log error but don't interrupt MCP operation + LOG.error(f"Failed to log MCP event: {e}", exc_info=True) diff --git a/querybook/server/lib/mcp/resources/comments.py b/querybook/server/lib/mcp/resources/comments.py new file mode 100644 index 000000000..30b979b0b --- /dev/null +++ b/querybook/server/lib/mcp/resources/comments.py @@ -0,0 +1,51 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.resources import ResourceContent +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.comments import get_comment_data, get_datadoc_cell_comments_data +from lib.mcp.utils import RESOURCE_ANNOTATIONS + + +def register(mcp: FastMCP) -> None: + """Register comment resources on the given MCP server.""" + + @mcp.resource( + uri="querybook://comment/{comment_id}", + name="Comment Content", + description="Get a comment with its thread replies and reactions", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_comment_resource( + comment_id: Annotated[int, "Comment ID"], + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get a comment with its thread replies and reactions.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + result = get_comment_data(comment_id, uid, session) + return [ResourceContent(result)] + + @mcp.resource( + uri="querybook://datadoc-cell/{cell_id}/comments{?include_threads}", + name="DataDoc Cell Comments", + description="Get all comments for a DataDoc cell with optional thread replies", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_datadoc_cell_comments_resource( + cell_id: Annotated[int, "DataDoc cell ID"], + include_threads: Annotated[bool, "Include thread replies"] = True, + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get comments for a DataDoc cell.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + result = get_datadoc_cell_comments_data( + cell_id, uid, include_threads, session + ) + return [ResourceContent({"comments": result})] diff --git a/querybook/server/lib/mcp/resources/datadocs.py b/querybook/server/lib/mcp/resources/datadocs.py new file mode 100644 index 000000000..2ea5f801b --- /dev/null +++ b/querybook/server/lib/mcp/resources/datadocs.py @@ -0,0 +1,73 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.resources import ResourceContent +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.datadocs import ( + get_datadoc_data, + get_datadoc_cell_data, + get_datadoc_cell_executions_data, +) +from lib.mcp.utils import RESOURCE_ANNOTATIONS + + +def register(mcp: FastMCP) -> None: + """Register datadoc resources on the given MCP server.""" + + @mcp.resource( + uri="querybook://datadoc/{datadoc_id}", + name="DataDoc Content", + description="Get a single DataDoc with its cells and editors by ID", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_datadoc_resource( + datadoc_id: Annotated[int, "DataDoc ID"], + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get a single DataDoc with its cells and editors.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + result = get_datadoc_data(datadoc_id, uid, session) + return [ResourceContent(result)] + + @mcp.resource( + uri="querybook://datadoc-cell/{cell_id}", + name="DataDoc Cell", + description="Get a single DataDoc cell with its content and metadata", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_datadoc_cell_resource( + cell_id: Annotated[int, "DataDoc cell ID"], + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get a single DataDoc cell.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + result = get_datadoc_cell_data(cell_id, uid, session) + return [ResourceContent(result)] + + @mcp.resource( + uri="querybook://datadoc-cell/{cell_id}/executions{?limit,offset}", + name="DataDoc Cell Executions", + description="Get query executions for a DataDoc cell with pagination", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_datadoc_cell_executions_resource( + cell_id: Annotated[int, "DataDoc cell ID"], + limit: Annotated[int, "Maximum number of results"] = 20, + offset: Annotated[int, "Pagination offset"] = 0, + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get execution history for a DataDoc cell.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + result = get_datadoc_cell_executions_data( + cell_id, uid, session, limit, offset + ) + return [ResourceContent({"executions": result})] diff --git a/querybook/server/lib/mcp/resources/environments.py b/querybook/server/lib/mcp/resources/environments.py new file mode 100644 index 000000000..c2ae1c50e --- /dev/null +++ b/querybook/server/lib/mcp/resources/environments.py @@ -0,0 +1,31 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.resources import ResourceContent +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.environments import get_environment_data +from lib.mcp.utils import RESOURCE_ANNOTATIONS + + +def register(mcp: FastMCP) -> None: + """Register environment resources on the given MCP server.""" + + @mcp.resource( + uri="querybook://environment/{environment_id}", + name="Environment", + description="Get a single environment's details by ID", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_environment_resource( + environment_id: Annotated[int, "Environment ID"], + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get a single environment's details.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + result = get_environment_data(environment_id, uid, session) + return [ResourceContent(result)] diff --git a/querybook/server/lib/mcp/resources/guide.py b/querybook/server/lib/mcp/resources/guide.py new file mode 100644 index 000000000..cab261eb1 --- /dev/null +++ b/querybook/server/lib/mcp/resources/guide.py @@ -0,0 +1,93 @@ +from fastmcp import FastMCP +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from lib.mcp.utils import RESOURCE_ANNOTATIONS + +RESOURCE_GUIDE = """\ +# Querybook MCP Resource Guide + +This server exposes resource templates that provide read-only access to +Querybook data. Use these URIs to fetch structured JSON for any entity +by substituting the appropriate ID. + +## Resource Templates + +### DataDocs +| URI | Description | +|-----|-------------| +| `querybook://datadoc/{datadoc_id}` | DataDoc with its cells and editors | +| `querybook://datadoc-cell/{cell_id}` | Single DataDoc cell content and metadata | +| `querybook://datadoc-cell/{cell_id}/executions{?limit,offset}` | Query execution history for a cell (default limit=20, offset=0) | + +### Query Execution +| URI | Description | +|-----|-------------| +| `querybook://query-execution/{query_execution_id}` | Query execution status and statement details | +| `querybook://statement-execution/{statement_execution_id}/results{?limit}` | Actual result data rows for a statement execution (limit defaults to server-configured size) | + +### Comments +| URI | Description | +|-----|-------------| +| `querybook://comment/{comment_id}` | Comment with thread replies and reactions | +| `querybook://datadoc-cell/{cell_id}/comments{?include_threads}` | All comments for a cell (include_threads defaults to true) | + +### Schedules +| URI | Description | +|-----|-------------| +| `querybook://schedule/{schedule_id}` | Schedule configuration and recent run history | +| `querybook://schedule/{schedule_id}/runs{?limit,offset,hide_successful}` | Schedule run history with filtering and pagination (default limit=20, offset=0, hide_successful=false) | + +### Lists +| URI | Description | +|-----|-------------| +| `querybook://list/{list_id}` | List with its items and editors | + +### Environments & Engines +| URI | Description | +|-----|-------------| +| `querybook://environment/{environment_id}` | Environment details | +| `querybook://query-engine/{engine_id}` | Query engine configuration | + +### Users +| URI | Description | +|-----|-------------| +| `querybook://user/{user_id}` | User profile (use `me` for current user) | + +## Query Parameters + +Some resources accept optional query parameters using RFC 6570 syntax: + +- **`limit`** / **`offset`** — Pagination (e.g., `querybook://datadoc-cell/42/executions?limit=10&offset=20`) +- **`include_threads`** — Boolean to include/exclude thread replies on comments +- **`hide_successful`** — Boolean to filter schedule runs to only failures + +## Tips + +- Many tools return a `resource_uri` field in their response — use it to + fetch the full resource representation. +- Use `list_environments` and `list_query_engines` tools to discover valid + environment and engine IDs. +- Use `search_datadocs` or `list_datadocs` tools to find DataDoc IDs. +- The `querybook://user/me` shorthand resolves to your authenticated user. +""" + + +def register(mcp: FastMCP) -> None: + """Register the resource guide on the given MCP server.""" + + @mcp.resource( + uri="querybook://resource-guide", + name="Resource Guide", + description=( + "Markdown guide listing all available Querybook resource templates, " + "their URI patterns, query parameters, and usage tips" + ), + mime_type="text/markdown", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_resource_guide( + token: AccessToken = CurrentAccessToken(), + ) -> str: + """Returns a Markdown document describing all available resource templates.""" + return RESOURCE_GUIDE diff --git a/querybook/server/lib/mcp/resources/lists.py b/querybook/server/lib/mcp/resources/lists.py new file mode 100644 index 000000000..f94e2de8b --- /dev/null +++ b/querybook/server/lib/mcp/resources/lists.py @@ -0,0 +1,31 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.resources import ResourceContent +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.lists import get_list_data +from lib.mcp.utils import RESOURCE_ANNOTATIONS + + +def register(mcp: FastMCP) -> None: + """Register list resources on the given MCP server.""" + + @mcp.resource( + uri="querybook://list/{list_id}", + name="List Content", + description="Get a single list with its items and editors by ID", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_list_resource( + list_id: Annotated[int, "List ID"], + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get a single list with its items and editors.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + result = get_list_data(list_id, uid, session) + return [ResourceContent(result)] diff --git a/querybook/server/lib/mcp/resources/query_engines.py b/querybook/server/lib/mcp/resources/query_engines.py new file mode 100644 index 000000000..66369e563 --- /dev/null +++ b/querybook/server/lib/mcp/resources/query_engines.py @@ -0,0 +1,31 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.resources import ResourceContent +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.query_engines import get_query_engine_data +from lib.mcp.utils import RESOURCE_ANNOTATIONS + + +def register(mcp: FastMCP) -> None: + """Register query engine resources on the given MCP server.""" + + @mcp.resource( + uri="querybook://query-engine/{engine_id}", + name="Query Engine", + description="Get a single query engine's configuration by ID", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_query_engine_resource( + engine_id: Annotated[int, "Query engine ID"], + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get a single query engine's configuration.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + result = get_query_engine_data(engine_id, uid, session) + return [ResourceContent(result)] diff --git a/querybook/server/lib/mcp/resources/query_executions.py b/querybook/server/lib/mcp/resources/query_executions.py new file mode 100644 index 000000000..59b6a5748 --- /dev/null +++ b/querybook/server/lib/mcp/resources/query_executions.py @@ -0,0 +1,52 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.resources import ResourceContent +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.query_executions import serialize_query_execution +from lib.mcp.utils import RESOURCE_ANNOTATIONS +from logic import query_execution as logic +from logic.query_execution_permission import user_can_access_query_execution + + +def register(mcp: FastMCP) -> None: + """Register query execution resources on the given MCP server.""" + + @mcp.resource( + uri="querybook://query-execution/{query_execution_id}", + name="Query Execution", + description="Get query execution status and statement details", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_query_execution_resource( + query_execution_id: Annotated[int, "Query execution ID"], + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get query execution status and statement details.""" + uid = token.claims["creator_uid"] + + with DBSession() as session: + execution = logic.get_query_execution_by_id( + query_execution_id, session=session + ) + + if not execution: + raise ValueError(f"Query execution {query_execution_id} not found.") + + # Check permissions + has_execution_access = user_can_access_query_execution( + uid=uid, + execution_id=query_execution_id, + session=session, + ) + + if not has_execution_access: + raise ValueError( + "You do not have permission to access this query execution." + ) + + return [ResourceContent(serialize_query_execution(execution))] diff --git a/querybook/server/lib/mcp/resources/schedules.py b/querybook/server/lib/mcp/resources/schedules.py new file mode 100644 index 000000000..a75a7da76 --- /dev/null +++ b/querybook/server/lib/mcp/resources/schedules.py @@ -0,0 +1,55 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.resources import ResourceContent +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.schedules import get_schedule_data, get_schedule_runs_data +from lib.mcp.utils import RESOURCE_ANNOTATIONS + + +def register(mcp: FastMCP) -> None: + """Register schedule resources on the given MCP server.""" + + @mcp.resource( + uri="querybook://schedule/{schedule_id}", + name="Schedule Details", + description="Get a single schedule with its run history", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_schedule_resource( + schedule_id: Annotated[int, "Schedule ID"], + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get schedule details with recent run history.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + result = get_schedule_data(schedule_id, uid, session) + return [ResourceContent(result)] + + @mcp.resource( + uri="querybook://schedule/{schedule_id}/runs{?limit,offset,hide_successful}", + name="Schedule Run History", + description="Get execution history for a schedule with optional filtering and pagination", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_schedule_runs_resource( + schedule_id: Annotated[int, "Schedule ID"], + limit: Annotated[int, "Maximum number of results"] = 20, + offset: Annotated[int, "Pagination offset"] = 0, + hide_successful: Annotated[ + bool, "Hide successful runs, show only failures" + ] = False, + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get execution history for a schedule.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + result = get_schedule_runs_data( + schedule_id, uid, session, limit, offset, hide_successful + ) + return [ResourceContent({"runs": result})] diff --git a/querybook/server/lib/mcp/resources/statement_executions.py b/querybook/server/lib/mcp/resources/statement_executions.py new file mode 100644 index 000000000..d504683fe --- /dev/null +++ b/querybook/server/lib/mcp/resources/statement_executions.py @@ -0,0 +1,100 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.resources import ResourceContent +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.config import get_config_value +from lib.mcp.utils import RESOURCE_ANNOTATIONS +from lib.result_store import GenericReader +from logic import query_execution as logic +from logic.query_execution_permission import ( + get_user_environments_by_execution_id, + user_can_access_query_execution, +) + +QUERY_RESULT_LIMIT_CONFIG = get_config_value("query_result_limit") + + +def register(mcp: FastMCP) -> None: + """Register statement execution resources on the given MCP server.""" + + @mcp.resource( + uri="querybook://statement-execution/{statement_execution_id}/results{?limit}", + name="Statement Execution Results", + description="Get the actual result data rows for a statement execution", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_statement_execution_results_resource( + statement_execution_id: Annotated[int, "Statement execution ID"], + limit: Annotated[ + int | None, + f"Maximum number of rows to return (default: {QUERY_RESULT_LIMIT_CONFIG['default_query_result_size']})", + ] = None, + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get the actual result data rows for a statement execution.""" + uid = token.claims["creator_uid"] + + # Apply limits + if limit is None: + limit = QUERY_RESULT_LIMIT_CONFIG["default_query_result_size"] + + max_limit = QUERY_RESULT_LIMIT_CONFIG["query_result_size_options"][-1] + if limit > max_limit: + raise ValueError(f"Too many rows requested. Maximum is {max_limit}") + + with DBSession() as session: + statement_execution = logic.get_statement_execution_by_id( + statement_execution_id, session=session + ) + + if statement_execution is None: + raise ValueError( + f"Statement execution {statement_execution_id} not found." + ) + + # Check permissions on the parent query execution + query_execution_id = statement_execution.query_execution_id + user_envs = get_user_environments_by_execution_id( + query_execution_id, uid, session=session + ) + + # User must either be in a shareable environment or have explicit access + has_env_access = len(user_envs) > 0 and any(e.shareable for e in user_envs) + has_execution_access = user_can_access_query_execution( + uid=uid, + execution_id=query_execution_id, + session=session, + ) + + if not (has_env_access or has_execution_access): + raise ValueError( + "You do not have permission to access this query execution." + ) + + # Read result data + # Check if result_path exists (can be None for failed queries) + if statement_execution.result_path is None: + return [ResourceContent({"columns": [], "data": []})] + + try: + with GenericReader(statement_execution.result_path) as reader: + rows = reader.read_csv(number_of_lines=limit + 1) + + if not rows: + return [ResourceContent({"columns": [], "data": []})] + + return [ + ResourceContent( + { + "columns": rows[0] if len(rows) > 0 else [], + "data": rows[1:] if len(rows) > 1 else [], + } + ) + ] + except Exception as e: + raise ValueError(f"Failed to read result data: {str(e)}") diff --git a/querybook/server/lib/mcp/resources/users.py b/querybook/server/lib/mcp/resources/users.py new file mode 100644 index 000000000..650441185 --- /dev/null +++ b/querybook/server/lib/mcp/resources/users.py @@ -0,0 +1,31 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.resources import ResourceContent +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.users import get_user_data +from lib.mcp.utils import RESOURCE_ANNOTATIONS + + +def register(mcp: FastMCP) -> None: + """Register user resources on the given MCP server.""" + + @mcp.resource( + uri="querybook://user/{user_id}", + name="User Profile", + description="Get a single user's profile information by ID or 'me' for current user", + mime_type="application/json", + annotations=RESOURCE_ANNOTATIONS, + ) + def get_user_resource( + user_id: Annotated[int | str, "User ID or 'me' for current user"], + token: AccessToken = CurrentAccessToken(), + ) -> list[ResourceContent]: + """Get a single user's profile information.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + result = get_user_data(user_id, uid, session) + return [ResourceContent(result)] diff --git a/querybook/server/lib/mcp/tools/comments.py b/querybook/server/lib/mcp/tools/comments.py new file mode 100644 index 000000000..69497f8ff --- /dev/null +++ b/querybook/server/lib/mcp/tools/comments.py @@ -0,0 +1,210 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.comments import serialize_comment, serialize_reaction +from lib.mcp.utils import CREATE_ANNOTATIONS, DELETE_ANNOTATIONS, WRITE_ANNOTATIONS +from logic.comment import ( + add_comment_to_data_cell, + add_thread_comment, + get_comment_by_id, + edit_comment, + add_reaction, + remove_reaction, +) +from logic.datadoc import get_data_cell_by_id +from logic.datadoc_permission import user_can_read, DocDoesNotExist +from models.comment import CommentReaction, DataCellComment +from models.datadoc import DataDocDataCell + + +def register(mcp: FastMCP) -> None: + """Register comment tools on the given MCP server.""" + + @mcp.tool( + title="Add DataDoc Cell Comment", + annotations=CREATE_ANNOTATIONS, + ) + def add_datadoc_cell_comment( + cell_id: Annotated[int, "ID of the DataDoc cell to comment on"], + text: Annotated[str, "Comment text"], + parent_comment_id: Annotated[ + int | None, "Parent comment ID for thread replies" + ] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Add a comment to a DataDoc cell, optionally as a reply to another comment.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + # Permission check + cell = get_data_cell_by_id(cell_id, session=session) + if not cell: + raise ValueError(f"DataDoc cell {cell_id} not found.") + + # Get datadoc_id from cell's data_doc_cells relationship + doc_cell = ( + session.query(DataDocDataCell) + .filter(DataDocDataCell.data_cell_id == cell_id) + .first() + ) + if not doc_cell: + raise ValueError( + f"DataDoc cell {cell_id} is not associated with a DataDoc." + ) + datadoc_id = doc_cell.data_doc_id + + try: + if not user_can_read(datadoc_id, uid, session=session): + raise ValueError("You do not have access to this DataDoc.") + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + # Create comment + if parent_comment_id is not None: + comment = add_thread_comment( + parent_comment_id, uid, text, session=session + ) + else: + comment = add_comment_to_data_cell(cell_id, uid, text, session=session) + + # New comments have no reactions yet + return serialize_comment(comment, []) + + @mcp.tool( + title="Update DataDoc Cell Comment", + annotations=WRITE_ANNOTATIONS, + ) + def update_datadoc_cell_comment( + comment_id: Annotated[int, "DataDoc cell comment ID"], + text: Annotated[str | None, "New comment text"] = None, + archived: Annotated[bool | None, "Archive/unarchive the comment"] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Update a DataDoc cell comment's text or archive status. Only non-null fields are updated.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + comment = get_comment_by_id(comment_id, session=session) + if not comment: + raise ValueError(f"Comment {comment_id} not found.") + + if comment.created_by != uid: + raise ValueError("You can only edit your own comments.") + + fields = {} + if text is not None: + fields["text"] = text + if archived is not None: + fields["archived"] = archived + + edit_comment(comment_id, session=session, **fields) + + # Return updated comment with reactions + comment = get_comment_by_id(comment_id, session=session) + reactions = ( + session.query(CommentReaction) + .filter(CommentReaction.comment_id == comment_id) + .all() + ) + return serialize_comment(comment, reactions) + + @mcp.tool( + title="Delete DataDoc Cell Comment", + annotations=DELETE_ANNOTATIONS, + ) + def delete_datadoc_cell_comment( + comment_id: Annotated[int, "DataDoc cell comment ID"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Delete (archive) a DataDoc cell comment. This is a soft delete.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + comment = get_comment_by_id(comment_id, session=session) + if not comment: + raise ValueError(f"Comment {comment_id} not found.") + + if comment.created_by != uid: + raise ValueError("You can only delete your own comments.") + + edit_comment(comment_id, archived=True, session=session) + return {"deleted": comment_id, "archived": True} + + @mcp.tool( + title="Add Comment Reaction", + annotations=CREATE_ANNOTATIONS, + ) + def add_comment_reaction( + comment_id: Annotated[int, "Comment ID to react to"], + reaction: Annotated[str, "Reaction emoji or text"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Add a reaction to a comment.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + # Check if comment exists + comment = get_comment_by_id(comment_id, session=session) + if not comment: + raise ValueError(f"Comment {comment_id} not found.") + + # Check permission: can read the datadoc containing this comment + cell_comment = ( + session.query(DataCellComment) + .filter(DataCellComment.comment_id == comment_id) + .first() + ) + + if cell_comment: + # This is a DataDoc cell comment - verify read permission + cell = get_data_cell_by_id(cell_comment.data_cell_id, session=session) + if not cell: + raise ValueError("DataDoc cell not found.") + + # Get datadoc_id from cell's data_doc_cells relationship + doc_cell = ( + session.query(DataDocDataCell) + .filter(DataDocDataCell.data_cell_id == cell.id) + .first() + ) + if not doc_cell: + raise ValueError("DataDoc cell is not associated with a DataDoc.") + + datadoc_id = doc_cell.data_doc_id + try: + if not user_can_read(datadoc_id, uid, session=session): + raise ValueError("You do not have access to this comment.") + except DocDoesNotExist: + raise ValueError("DataDoc not found.") + # If it's not a DataDoc cell comment, it's a table comment - no permission check needed + + # Add the reaction + reaction_obj = add_reaction( + comment_id=comment_id, reaction=reaction, uid=uid, session=session + ) + + return serialize_reaction(reaction_obj) + + @mcp.tool( + title="Remove Comment Reaction", + annotations=DELETE_ANNOTATIONS, + ) + def remove_comment_reaction( + reaction_id: Annotated[int, "Reaction ID to remove"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Remove a reaction from a comment. You can only remove your own reactions.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + # Check if reaction exists + reaction = session.query(CommentReaction).get(reaction_id) + if not reaction: + raise ValueError(f"Reaction {reaction_id} not found.") + + # Check if user owns this reaction + if reaction.created_by != uid: + raise ValueError("You can only remove your own reactions.") + + # Remove the reaction + remove_reaction(reaction_id, session=session) + return {"deleted": reaction_id} diff --git a/querybook/server/lib/mcp/tools/datadocs.py b/querybook/server/lib/mcp/tools/datadocs.py new file mode 100644 index 000000000..08426aef2 --- /dev/null +++ b/querybook/server/lib/mcp/tools/datadocs.py @@ -0,0 +1,597 @@ +from typing import Annotated, Literal + +from fastmcp import FastMCP +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.datadocs import ( + serialize_datadoc, + serialize_datadoc_cell, + serialize_datadoc_editor, + get_datadoc_data, +) +from lib.mcp.utils import ( + READ_ONLY_ANNOTATIONS, + WRITE_ANNOTATIONS, + CREATE_ANNOTATIONS, + DELETE_ANNOTATIONS, + build_querybook_url, +) +from logic.datadoc import ( + create_data_doc, + get_data_doc_by_id, + update_data_doc, + favorite_data_doc, + unfavorite_data_doc, + create_data_cell, + get_data_cell_by_id, + insert_data_doc_cell, + update_data_cell, + delete_data_doc_cell, + create_data_doc_editor, + get_data_doc_editor_by_id, + update_data_doc_editor, + delete_data_doc_editor, +) +from logic.datadoc_permission import ( + DocDoesNotExist, + user_can_write, +) +from models.datadoc import DataDocDataCell +from models.environment import Environment + + +VALID_VARIABLE_TYPES = {"string", "number", "boolean"} +VARIABLE_TYPE_CHECKS = { + "string": lambda v: isinstance(v, str), + "number": lambda v: isinstance(v, (int, float)) and not isinstance(v, bool), + "boolean": lambda v: isinstance(v, bool), +} + + +def validate_variables(variables: list[dict]) -> None: + """Validate template variables before passing to the data layer.""" + seen_names = set() + for i, var in enumerate(variables): + name = var.get("name") + var_type = var.get("type") + + if not name or not name.strip(): + raise ValueError(f"Variable at index {i} has an empty or missing name.") + + if var_type not in VALID_VARIABLE_TYPES: + raise ValueError( + f"Variable '{name}' has invalid type '{var_type}'. " + f"Must be one of: {', '.join(sorted(VALID_VARIABLE_TYPES))}." + ) + + if not VARIABLE_TYPE_CHECKS[var_type](var.get("value")): + raise ValueError( + f"Variable '{name}' has type '{var_type}' but value is not a {var_type}." + ) + + if name in seen_names: + raise ValueError(f"Duplicate variable name '{name}'.") + seen_names.add(name) + + +def register(mcp: FastMCP) -> None: + """Register datadoc tools on the given MCP server.""" + + @mcp.tool( + title="Get DataDoc", + annotations=READ_ONLY_ANNOTATIONS, + ) + def get_datadoc( + datadoc_id: Annotated[int, "DataDoc ID"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Get a single DataDoc with its cells and editors by ID.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + return get_datadoc_data(datadoc_id, uid, session) + + @mcp.tool( + title="List DataDocs", + annotations=READ_ONLY_ANNOTATIONS, + ) + def list_datadocs( + environment_id: Annotated[int, "Environment ID from list_environments"], + scope: Annotated[ + Literal["mine", "favorite", "recent"], + "Scope: 'mine' (owned by you), 'favorite' (favorited by you), 'recent' (recently viewed by you)", + ] = "mine", + offset: Annotated[int, "Pagination offset"] = 0, + limit: Annotated[int, "Maximum number of results"] = 100, + token: AccessToken = CurrentAccessToken(), + ) -> list[dict]: + """List DataDocs in the given environment, filtered by scope.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + docs = [] + + if scope == "mine": + from logic.datadoc import get_data_doc_by_user + + docs = get_data_doc_by_user( + uid, + environment_id=environment_id, + offset=offset, + limit=limit, + session=session, + ) + elif scope == "favorite": + from logic.datadoc import get_user_favorite_data_docs + + all_docs = get_user_favorite_data_docs( + uid, environment_id=environment_id, session=session + ) + docs = all_docs[offset : offset + limit] + elif scope == "recent": + from logic.datadoc import get_user_recent_data_docs + + all_docs = get_user_recent_data_docs( + uid, environment_id=environment_id, session=session + ) + docs = all_docs[offset : offset + limit] + + return [serialize_datadoc(doc, session, uid=uid) for doc in docs] + + @mcp.tool( + title="Search DataDocs", + annotations=READ_ONLY_ANNOTATIONS, + ) + def search_datadocs( + environment_id: Annotated[int, "Environment ID from list_environments"], + keywords: Annotated[str, "Search keywords"], + filters: Annotated[ + list[list] | None, "Filters as list of [field, value] pairs" + ] = None, + limit: Annotated[int, "Maximum number of results"] = 100, + offset: Annotated[int, "Pagination offset"] = 0, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Search DataDocs using keywords and filters.""" + if limit > 100: + raise ValueError("Limit cannot exceed 100") + + uid = token.claims["creator_uid"] + + from lib.elasticsearch.search_datadoc import construct_datadoc_query + from lib.elasticsearch.search_utils import get_matching_objects, ES_CONFIG + + with DBSession() as session: + # Add environment filter + filters_with_env = (filters or []) + [["environment_id", environment_id]] + + query = construct_datadoc_query( + uid=uid, + keywords=keywords.strip(), + filters=filters_with_env, + fields=[], + limit=limit, + offset=offset, + sort_key=None, + sort_order=None, + ) + + results, count = get_matching_objects( + query, ES_CONFIG["datadocs"]["index_name"], True + ) + + # Add resource URIs and URLs to search results + environment = session.query(Environment).get(environment_id) + for result in results: + doc_id = result.get("_source", {}).get("id") or result.get("id") + if doc_id: + result["resource_uri"] = f"querybook://datadoc/{doc_id}" + if environment: + url = build_querybook_url(environment.name, f"datadoc/{doc_id}") + if url: + result["url"] = url + + return {"count": count, "results": results} + + @mcp.tool( + title="Create DataDoc", + annotations=CREATE_ANNOTATIONS, + ) + def create_datadoc( + environment_id: Annotated[int, "Environment ID from list_environments"], + title: Annotated[str, "DataDoc title"] = "Untitled", + public: Annotated[bool, "Whether the DataDoc is public"] = False, + favorite: Annotated[ + bool, "Whether to favorite this DataDoc immediately" + ] = False, + variables: Annotated[ + list[dict] | None, + "Template variables as list of {name, type, value} objects. " + "type must be 'string', 'number', or 'boolean'. " + "value must match the declared type.", + ] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Create a new empty DataDoc.""" + owner_uid = token.claims["creator_uid"] + if variables is not None: + validate_variables(variables) + meta = {"variables": variables} if variables else {} + + with DBSession() as session: + data_doc = create_data_doc( + environment_id=environment_id, + owner_uid=owner_uid, + cells=[], + title=title, + meta=meta, + public=public, + archived=False, + session=session, + ) + + # Favorite if requested + if favorite: + favorite_data_doc( + data_doc_id=data_doc.id, uid=owner_uid, session=session + ) + + return serialize_datadoc(data_doc, session, uid=owner_uid, with_cells=True) + + @mcp.tool( + title="Update DataDoc", + annotations=WRITE_ANNOTATIONS, + ) + def update_datadoc( + datadoc_id: Annotated[int, "DataDoc ID"], + title: Annotated[str | None, "New title for the DataDoc"] = None, + public: Annotated[bool | None, "Whether the DataDoc is public"] = None, + archived: Annotated[bool | None, "Whether the DataDoc is archived"] = None, + environment_id: Annotated[int | None, "Move to different environment"] = None, + favorite: Annotated[ + bool | None, "Favorite (true) or unfavorite (false) this DataDoc" + ] = None, + variables: Annotated[ + list[dict] | None, + "Template variables as list of {name, type, value} objects. " + "type must be 'string', 'number', or 'boolean'. " + "value must match the declared type. Replaces all existing variables.", + ] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Update DataDoc properties. Only non-null fields are updated.""" + if variables is not None: + validate_variables(variables) + uid = token.claims["creator_uid"] + with DBSession() as session: + try: + if not user_can_write(datadoc_id, uid, session=session): + raise ValueError("You do not have permission to edit this DataDoc.") + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + # Build fields dict with only non-null values + fields = {} + if title is not None: + fields["title"] = title + if public is not None: + fields["public"] = public + if archived is not None: + fields["archived"] = archived + if environment_id is not None: + fields["environment_id"] = environment_id + if variables is not None: + fields["meta"] = {"variables": variables} + + update_data_doc(id=datadoc_id, commit=True, session=session, **fields) + + # Update favorite status if specified + if favorite is not None: + if favorite: + favorite_data_doc(data_doc_id=datadoc_id, uid=uid, session=session) + else: + unfavorite_data_doc( + data_doc_id=datadoc_id, uid=uid, session=session + ) + + # Return updated doc + doc = get_data_doc_by_id(datadoc_id, session=session) + return serialize_datadoc(doc, session, uid=uid, with_cells=True) + + @mcp.tool( + title="Delete DataDoc", + annotations=DELETE_ANNOTATIONS, + ) + def delete_datadoc( + datadoc_id: Annotated[int, "DataDoc ID"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Delete (archive) a DataDoc. This is a soft delete - the DataDoc is archived, not permanently deleted.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + try: + if not user_can_write(datadoc_id, uid, session=session): + raise ValueError( + "You do not have permission to delete this DataDoc." + ) + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + update_data_doc(id=datadoc_id, archived=True, commit=True, session=session) + return {"deleted": datadoc_id, "archived": True} + + @mcp.tool( + title="Add DataDoc Cell", + annotations=CREATE_ANNOTATIONS, + ) + def add_datadoc_cell( + datadoc_id: Annotated[int, "DataDoc ID"], + cell_type: Annotated[ + Literal["query", "text", "chart", "python"], + "Type of cell: 'query', 'text', 'chart', or 'python'", + ], + context: Annotated[str, "Cell SQL, code, or rich text content"] = "", + title: Annotated[str | None, "Cell title"] = None, + engine_id: Annotated[ + int | None, + "Query engine ID from list_query_engines (required for query cells)", + ] = None, + meta: Annotated[dict | None, "Additional cell metadata"] = None, + index: Annotated[int | None, "Position to insert cell (default: end)"] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Add a cell to a DataDoc.""" + uid = token.claims["creator_uid"] + + with DBSession() as session: + try: + if not user_can_write(datadoc_id, uid, session=session): + raise ValueError("You do not have permission to edit this DataDoc.") + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + doc = get_data_doc_by_id(datadoc_id, session=session) + if index is None: + index = len(doc.cells) + + # Build meta dict with top-level fields + cell_meta = {**(meta or {})} + if title is not None: + cell_meta["title"] = title + if engine_id is not None: + cell_meta["engine"] = engine_id + + data_cell = create_data_cell( + cell_type=cell_type, + context=context, + meta=cell_meta, + commit=False, + session=session, + ) + + insert_data_doc_cell( + data_doc_id=datadoc_id, + cell_id=data_cell.id, + index=index, + commit=True, + session=session, + ) + + return serialize_datadoc_cell( + data_cell.to_dict(), session, datadoc_id=datadoc_id + ) + + @mcp.tool( + title="Update DataDoc Cell", + annotations=WRITE_ANNOTATIONS, + ) + def update_datadoc_cell( + cell_id: Annotated[int, "Cell ID not the cell's index"], + context: Annotated[str | None, "Cell content/code"] = None, + title: Annotated[str | None, "Cell title"] = None, + engine_id: Annotated[ + int | None, "Query engine ID from list_query_engines (for query cells)" + ] = None, + meta: Annotated[dict | None, "Additional cell metadata"] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Update a DataDoc cell's context or metadata. Only updates non-null fields.""" + uid = token.claims["creator_uid"] + + with DBSession() as session: + # Load the cell first to validate it exists + cell = get_data_cell_by_id(cell_id, session=session) + if not cell: + raise ValueError(f"DataDoc cell {cell_id} not found.") + + # Get datadoc_id from cell's data_doc_cells relationship + doc_cell = ( + session.query(DataDocDataCell) + .filter(DataDocDataCell.data_cell_id == cell_id) + .first() + ) + if not doc_cell: + raise ValueError( + f"DataDoc cell {cell_id} is not associated with a DataDoc." + ) + datadoc_id = doc_cell.data_doc_id + + # Check permission against the actual parent DataDoc + try: + if not user_can_write(datadoc_id, uid, session=session): + raise ValueError("You do not have permission to edit this DataDoc.") + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + fields = {} + if context is not None: + fields["context"] = context + + # Build updated meta if any top-level fields or meta provided + if title is not None or engine_id is not None or meta is not None: + # Merge with existing meta + cell_meta = {**(cell.meta or {})} + + # Merge provided meta + if meta is not None: + cell_meta.update(meta) + + # Override with top-level fields + if title is not None: + cell_meta["title"] = title + if engine_id is not None: + cell_meta["engine"] = engine_id + + fields["meta"] = cell_meta + + update_data_cell(cell_id, session=session, **fields) + + session.refresh(cell) + return serialize_datadoc_cell( + cell.to_dict(), + session, + with_latest_execution=True, + datadoc_id=datadoc_id, + ) + + @mcp.tool( + title="Delete DataDoc Cell", + annotations=DELETE_ANNOTATIONS, + ) + def delete_datadoc_cell( + cell_id: Annotated[int, "Cell ID not the cell's index"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Delete a cell from a DataDoc.""" + uid = token.claims["creator_uid"] + + with DBSession() as session: + # Load the cell first to validate it exists + cell = get_data_cell_by_id(cell_id, session=session) + if not cell: + raise ValueError(f"DataDoc cell {cell_id} not found.") + + # Get datadoc_id from cell's data_doc_cells relationship + doc_cell = ( + session.query(DataDocDataCell) + .filter(DataDocDataCell.data_cell_id == cell_id) + .first() + ) + if not doc_cell: + raise ValueError( + f"DataDoc cell {cell_id} is not associated with a DataDoc." + ) + datadoc_id = doc_cell.data_doc_id + + # Check permission against the actual parent DataDoc + try: + if not user_can_write(datadoc_id, uid, session=session): + raise ValueError("You do not have permission to edit this DataDoc.") + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + delete_data_doc_cell( + data_doc_id=datadoc_id, data_cell_id=cell_id, session=session + ) + + return { + "deleted": cell_id, + "datadoc_id": datadoc_id, + "datadoc_resource_uri": f"querybook://datadoc/{datadoc_id}", + } + + @mcp.tool( + title="Add DataDoc Editor", + annotations=WRITE_ANNOTATIONS, + ) + def add_datadoc_editor( + datadoc_id: Annotated[int, "DataDoc ID"], + user_id: Annotated[int, "User ID to add as editor, from search_users"], + read: Annotated[bool, "Grant read permission"] = False, + write: Annotated[bool, "Grant write permission"] = False, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Add an editor to a DataDoc with specified permissions.""" + caller_uid = token.claims["creator_uid"] + + with DBSession() as session: + try: + if not user_can_write(datadoc_id, caller_uid, session=session): + raise ValueError("You do not have permission to edit this DataDoc.") + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + editor = create_data_doc_editor( + data_doc_id=datadoc_id, + uid=user_id, + read=read, + write=write, + session=session, + ) + + return serialize_datadoc_editor(editor) + + @mcp.tool( + title="Update DataDoc Editor", + annotations=WRITE_ANNOTATIONS, + ) + def update_datadoc_editor( + editor_id: Annotated[int, "Editor ID not the user_id"], + read: Annotated[bool | None, "New read permission (null = no change)"] = None, + write: Annotated[bool | None, "New write permission (null = no change)"] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Update permissions for an editor on a DataDoc.""" + uid = token.claims["creator_uid"] + + with DBSession() as session: + editor = get_data_doc_editor_by_id(editor_id, session=session) + if not editor: + raise ValueError(f"Editor {editor_id} not found.") + + try: + if not user_can_write(editor.data_doc_id, uid, session=session): + raise ValueError("You do not have permission to edit this DataDoc.") + except DocDoesNotExist: + raise ValueError(f"DataDoc {editor.data_doc_id} not found.") + + updated_editor = update_data_doc_editor( + id=editor_id, + read=read, + write=write, + session=session, + ) + + return serialize_datadoc_editor(updated_editor) if updated_editor else {} + + @mcp.tool( + title="Remove DataDoc Editor", + annotations=DELETE_ANNOTATIONS, + ) + def remove_datadoc_editor( + editor_id: Annotated[int, "Editor ID not the user_id"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Remove an editor from a DataDoc.""" + uid = token.claims["creator_uid"] + + with DBSession() as session: + editor = get_data_doc_editor_by_id(editor_id, session=session) + if not editor: + raise ValueError(f"Editor {editor_id} not found.") + + doc_id = editor.data_doc_id + + try: + if not user_can_write(doc_id, uid, session=session): + raise ValueError("You do not have permission to edit this DataDoc.") + except DocDoesNotExist: + raise ValueError(f"DataDoc {doc_id} not found.") + + delete_data_doc_editor( + id=editor_id, + doc_id=doc_id, + session=session, + ) + + return {"removed": editor_id} diff --git a/querybook/server/lib/mcp/tools/environments.py b/querybook/server/lib/mcp/tools/environments.py new file mode 100644 index 000000000..72a65de63 --- /dev/null +++ b/querybook/server/lib/mcp/tools/environments.py @@ -0,0 +1,30 @@ +from fastmcp import FastMCP +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.environments import serialize_environment +from lib.mcp.utils import READ_ONLY_ANNOTATIONS +from logic.environment import get_all_visible_environments_by_uid + + +def register(mcp: FastMCP) -> None: + """Register environment tools on the given MCP server.""" + + @mcp.tool( + title="List Environments", + annotations=READ_ONLY_ANNOTATIONS, + ) + def list_environments( + token: AccessToken = CurrentAccessToken(), + ) -> list[dict]: + """List all environments visible to the authenticated user.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + environments = get_all_visible_environments_by_uid(uid, session=session) + return [ + serialize_environment( + env, uid=uid, session=session, with_query_engines=True + ) + for env in environments + ] diff --git a/querybook/server/lib/mcp/tools/lists.py b/querybook/server/lib/mcp/tools/lists.py new file mode 100644 index 000000000..e1fde3afe --- /dev/null +++ b/querybook/server/lib/mcp/tools/lists.py @@ -0,0 +1,433 @@ +from typing import Annotated, Literal + +from fastmcp import FastMCP +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.lists import ( + serialize_list, + serialize_list_editor, + serialize_list_item, + get_list_data, +) +from lib.mcp.utils import ( + READ_ONLY_ANNOTATIONS, + WRITE_ANNOTATIONS, + CREATE_ANNOTATIONS, + DELETE_ANNOTATIONS, + build_querybook_url, +) +from logic.board import ( + get_user_boards, + get_all_public_boards, + create_board as create_board_logic, + update_board as update_board_logic, + add_item_to_board, + update_board_item as update_board_item_logic, + remove_item_from_board, + create_board_editor, + get_board_editor_by_id, + update_board_editor as update_board_editor_logic, + delete_board_editor as delete_board_editor_logic, +) +from logic.board_permission import BoardDoesNotExist, user_can_edit +from models.board import Board, BoardItem +from models.environment import Environment + + +def register(mcp: FastMCP) -> None: + """Register list tools on the given MCP server.""" + + @mcp.tool( + title="Get List", + annotations=READ_ONLY_ANNOTATIONS, + ) + def get_list( + list_id: Annotated[int, "List ID"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Get a single list with its items and editors by ID.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + return get_list_data(list_id, uid, session) + + @mcp.tool( + title="List Lists", + annotations=READ_ONLY_ANNOTATIONS, + ) + def list_lists( + environment_id: Annotated[int, "Environment ID from list_environments"], + scope: Annotated[ + Literal["mine", "public"], + "Which lists to return: 'mine' (owned by you), 'public' (all public lists).", + ] = "mine", + token: AccessToken = CurrentAccessToken(), + ) -> list[dict]: + """Enumerates lists in the given environment, filtered by scope.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + if scope == "mine": + lists = get_user_boards(uid, environment_id, session=session) + elif scope == "public": + lists = get_all_public_boards(environment_id, session=session) + else: + raise ValueError(f"Invalid scope: {scope}") + + return [ + serialize_list(list_obj, environment_id, session) for list_obj in lists + ] + + @mcp.tool( + title="Search Lists", + annotations=READ_ONLY_ANNOTATIONS, + ) + def search_lists( + environment_id: Annotated[int, "Environment ID from list_environments"], + keywords: Annotated[str, "Search keywords"], + filters: Annotated[ + list[list] | None, "Filters as list of [field, value] pairs" + ] = None, + limit: Annotated[int, "Maximum number of results"] = 100, + offset: Annotated[int, "Pagination offset"] = 0, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Search lists using keywords and filters.""" + if limit > 100: + raise ValueError("Limit cannot exceed 100") + + uid = token.claims["creator_uid"] + + from lib.elasticsearch.search_board import construct_board_query + from lib.elasticsearch.search_utils import get_matching_objects, ES_CONFIG + + with DBSession() as session: + # Add environment filter + filters_with_env = (filters or []) + [["environment_id", environment_id]] + + query = construct_board_query( + uid=uid, + keywords=keywords.strip(), + filters=filters_with_env, + fields=[], + limit=limit, + offset=offset, + sort_key=None, + sort_order=None, + ) + + results, count = get_matching_objects( + query, ES_CONFIG["boards"]["index_name"], True + ) + + # Add resource URIs and URLs to search results, translate field names + environment = session.query(Environment).get(environment_id) + for result in results: + # Translate board_type to list_type in result + source = result.get("_source", {}) + if "board_type" in source: + source["list_type"] = source.pop("board_type") + if "board_type" in result: + result["list_type"] = result.pop("board_type") + + list_id = source.get("id") or result.get("id") + if list_id: + result["resource_uri"] = f"querybook://list/{list_id}" + if environment: + url = build_querybook_url(environment.name, f"list/{list_id}") + if url: + result["url"] = url + + return {"count": count, "results": results} + + @mcp.tool( + title="Create List", + annotations=CREATE_ANNOTATIONS, + ) + def create_list( + environment_id: Annotated[int, "Environment ID from list_environments"], + name: Annotated[str, "List name"], + description: Annotated[str | None, "List description"] = None, + public: Annotated[bool, "Whether the list is public"] = False, + list_type: Annotated[str, "List type"] = "", + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Create a new list.""" + owner_uid = token.claims["creator_uid"] + with DBSession() as session: + list_obj = create_board_logic( + name=name, + environment_id=environment_id, + owner_uid=owner_uid, + description=description, + public=public, + board_type=list_type, + commit=True, + session=session, + ) + return serialize_list(list_obj, environment_id, session) + + @mcp.tool( + title="Update List", + annotations=WRITE_ANNOTATIONS, + ) + def update_list( + list_id: Annotated[int, "List ID from list_lists or search_lists"], + name: Annotated[str | None, "New name for the list"] = None, + description: Annotated[str | None, "New description for the list"] = None, + public: Annotated[bool | None, "Whether the list is public"] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Update list properties. Only non-null fields are updated.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + try: + if not user_can_edit(list_id, uid, session=session): + raise ValueError("You do not have permission to edit this list.") + except BoardDoesNotExist: + raise ValueError(f"List {list_id} not found.") + + fields = {} + if name is not None: + fields["name"] = name + if description is not None: + fields["description"] = description + if public is not None: + fields["public"] = public + + update_board_logic(id=list_id, commit=True, session=session, **fields) + + list_obj = Board.get(id=list_id, session=session) + return serialize_list(list_obj, list_obj.environment_id, session) + + @mcp.tool( + title="Delete List", + annotations=DELETE_ANNOTATIONS, + ) + def delete_list( + list_id: Annotated[int, "List ID from list_lists or search_lists"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Delete a list. Cannot delete favorite lists.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + try: + if not user_can_edit(list_id, uid, session=session): + raise ValueError("You do not have permission to delete this list.") + except BoardDoesNotExist: + raise ValueError(f"List {list_id} not found.") + + list_obj = Board.get(id=list_id, session=session) + if list_obj.board_type == "favorite": + raise ValueError("Cannot delete favorite lists.") + + Board.delete(list_obj.id, session=session) + return {"deleted": list_id} + + @mcp.tool( + title="Add List Item", + annotations=WRITE_ANNOTATIONS, + ) + def add_list_item( + list_id: Annotated[int, "List ID from list_lists or search_lists"], + item_id: Annotated[ + int, "ID of the item to add (datadoc_id, table_id, etc. based on item_type)" + ], + item_type: Annotated[ + Literal["data_doc", "table", "list", "query"], + "Type of item: 'data_doc', 'table', 'list', or 'query'", + ], + description: Annotated[ + str | None, "Optional description for the list item" + ] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Add an item to a list.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + try: + if not user_can_edit(list_id, uid, session=session): + raise ValueError("You do not have permission to edit this list.") + except BoardDoesNotExist: + raise ValueError(f"List {list_id} not found.") + + # Translate "list" to "board" for internal database API + internal_item_type = "board" if item_type == "list" else item_type + + list_item = add_item_to_board( + list_id, item_id, internal_item_type, session=session + ) + + # Optionally set description if provided + if description is not None: + list_item = update_board_item_logic( + list_item.id, description=description, session=session + ) + + return serialize_list_item(list_item) + + @mcp.tool( + title="Update List Item", + annotations=WRITE_ANNOTATIONS, + ) + def update_list_item( + list_item_id: Annotated[ + int, "List item ID from get_list, not the datadoc_id/table_id/etc" + ], + description: Annotated[str | None, "Item description"] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Update a list item's description.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + # Load the list item first to get its parent list + list_item = BoardItem.get(id=list_item_id, session=session) + if not list_item: + raise ValueError(f"List item {list_item_id} not found.") + + # Check permission against the actual parent list + try: + if not user_can_edit(list_item.parent_board_id, uid, session=session): + raise ValueError("You do not have permission to edit this list.") + except BoardDoesNotExist: + raise ValueError(f"List {list_item.parent_board_id} not found.") + + fields = {} + if description is not None: + fields["description"] = description + + if fields: + updated_item = update_board_item_logic( + list_item_id, session=session, **fields + ) + if not updated_item: + raise ValueError(f"Failed to update list item {list_item_id}.") + list_item = updated_item + + return serialize_list_item(list_item) + + @mcp.tool( + title="Remove List Item", + annotations=DELETE_ANNOTATIONS, + ) + def remove_list_item( + list_id: Annotated[int, "List ID from list_lists or search_lists"], + item_id: Annotated[ + int, + "ID of the item to remove: datadoc_id, table_id, etc. based on item_type", + ], + item_type: Annotated[ + Literal["data_doc", "table", "list", "query"], + "Type of item: 'data_doc', 'table', 'list', or 'query'", + ], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Remove an item from a list.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + try: + if not user_can_edit(list_id, uid, session=session): + raise ValueError("You do not have permission to edit this list.") + except BoardDoesNotExist: + raise ValueError(f"List {list_id} not found.") + + # Translate "list" to "board" for internal database API + internal_item_type = "board" if item_type == "list" else item_type + + remove_item_from_board( + list_id, item_id, internal_item_type, session=session + ) + return {"removed": item_id, "list_id": list_id, "item_type": item_type} + + @mcp.tool( + title="Add List Editor", + annotations=WRITE_ANNOTATIONS, + ) + def add_list_editor( + list_id: Annotated[int, "List ID"], + user_id: Annotated[int, "User ID to add as editor, from search_users"], + read: Annotated[bool, "Grant read permission"] = False, + write: Annotated[bool, "Grant write permission"] = False, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Add an editor to a list with specified permissions.""" + caller_uid = token.claims["creator_uid"] + + with DBSession() as session: + try: + if not user_can_edit(list_id, caller_uid, session=session): + raise ValueError("You do not have permission to edit this list.") + except BoardDoesNotExist: + raise ValueError(f"List {list_id} not found.") + + editor = create_board_editor( + board_id=list_id, + uid=user_id, + read=read, + write=write, + commit=True, + session=session, + ) + + return serialize_list_editor(editor) + + @mcp.tool( + title="Update List Editor", + annotations=WRITE_ANNOTATIONS, + ) + def update_list_editor( + editor_id: Annotated[int, "Editor ID not the user_id"], + read: Annotated[bool | None, "New read permission (null = no change)"] = None, + write: Annotated[bool | None, "New write permission (null = no change)"] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Update permissions for an editor on a list.""" + uid = token.claims["creator_uid"] + + with DBSession() as session: + editor = get_board_editor_by_id(editor_id, session=session) + if not editor: + raise ValueError(f"Editor {editor_id} not found.") + + try: + if not user_can_edit(editor.board_id, uid, session=session): + raise ValueError("You do not have permission to edit this list.") + except BoardDoesNotExist: + raise ValueError(f"List {editor.board_id} not found.") + + updated_editor = update_board_editor_logic( + id=editor_id, + read=read, + write=write, + session=session, + ) + + return serialize_list_editor(updated_editor) if updated_editor else {} + + @mcp.tool( + title="Remove List Editor", + annotations=DELETE_ANNOTATIONS, + ) + def remove_list_editor( + editor_id: Annotated[int, "Editor ID not the user_id"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Remove an editor from a list.""" + uid = token.claims["creator_uid"] + + with DBSession() as session: + editor = get_board_editor_by_id(editor_id, session=session) + if not editor: + raise ValueError(f"Editor {editor_id} not found.") + + try: + if not user_can_edit(editor.board_id, uid, session=session): + raise ValueError("You do not have permission to edit this list.") + except BoardDoesNotExist: + raise ValueError(f"List {editor.board_id} not found.") + + delete_board_editor_logic( + id=editor_id, board_id=editor.board_id, session=session + ) + + return {"removed": editor_id, "list_id": editor.board_id} diff --git a/querybook/server/lib/mcp/tools/query_engines.py b/querybook/server/lib/mcp/tools/query_engines.py new file mode 100644 index 000000000..f4b025d0a --- /dev/null +++ b/querybook/server/lib/mcp/tools/query_engines.py @@ -0,0 +1,41 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.query_engines import serialize_query_engine +from lib.mcp.utils import READ_ONLY_ANNOTATIONS +from logic import admin as admin_logic +from logic.environment import get_all_accessible_environment_ids_by_uid + + +def register(mcp: FastMCP) -> None: + """Register query engine tools on the given MCP server.""" + + @mcp.tool( + title="List Query Engines", + annotations=READ_ONLY_ANNOTATIONS, + ) + def list_query_engines( + environment_id: Annotated[int, "Environment ID from list_environments"], + token: AccessToken = CurrentAccessToken(), + ) -> list[dict]: + """List query engines accessible to the user in the given environment.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + accessible_env_ids = get_all_accessible_environment_ids_by_uid( + uid, session=session + ) + if environment_id not in accessible_env_ids: + raise ValueError( + f"You do not have access to environment {environment_id}." + ) + engines = admin_logic.get_query_engines_by_environment( + environment_id, ordered=True, session=session + ) + return [ + serialize_query_engine(engine, with_environments=True) + for engine in engines + ] diff --git a/querybook/server/lib/mcp/tools/query_executions.py b/querybook/server/lib/mcp/tools/query_executions.py new file mode 100644 index 000000000..2e9d7ad2b --- /dev/null +++ b/querybook/server/lib/mcp/tools/query_executions.py @@ -0,0 +1,257 @@ +from typing import Annotated, Literal + +from fastmcp import FastMCP +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from const.query_execution import QueryExecutionStatus, QueryExecutionType +from lib.mcp.lib.query_executions import ( + serialize_query_execution, + serialize_query_execution_summary, +) +from lib.mcp.utils import ( + CREATE_ANNOTATIONS, + READ_ONLY_ANNOTATIONS, + verify_query_engine_access, +) +from lib.query_analysis.templating import render_templated_query +from logic import query_execution as logic +from logic import datadoc as datadoc_logic +from logic.datadoc_permission import user_can_write, DocDoesNotExist + + +def _build_mcp_metadata(user_metadata: dict | None) -> dict: + """Build execution metadata with MCP source marker.""" + metadata = dict(user_metadata) if user_metadata else {} + metadata["source"] = "mcp" + return metadata + + +def register(mcp: FastMCP) -> None: + """Register query execution tools on the given MCP server.""" + + @mcp.tool( + title="List Query Executions", + annotations=READ_ONLY_ANNOTATIONS, + ) + def list_query_executions( + environment_id: Annotated[int, "Environment ID from list_environments"], + engine_id: Annotated[ + int | None, + "Filter by query engine ID from list_query_engines", + ] = None, + status: Annotated[ + Literal[ + "INITIALIZED", + "DELIVERED", + "RUNNING", + "DONE", + "ERROR", + "CANCEL", + "PENDING_REVIEW", + "REJECTED", + ] + | None, + "Filter by execution status", + ] = None, + running: Annotated[ + bool | None, + "If true, show only active executions (INITIALIZED, DELIVERED, RUNNING)", + ] = None, + limit: Annotated[int, "Maximum results, max 100"] = 20, + offset: Annotated[int, "Pagination offset"] = 0, + token: AccessToken = CurrentAccessToken(), + ) -> list[dict]: + """List query executions for the current user. Returns summaries only.""" + if limit > 100: + raise ValueError("limit must be 100 or less") + + uid = token.claims["creator_uid"] + + filters = {"user": uid} + if engine_id is not None: + filters["engine"] = engine_id + if status is not None: + filters["status"] = QueryExecutionStatus[status].value + if running: + filters["running"] = True + + with DBSession() as session: + executions = logic.search_query_execution( + environment_id=environment_id, + filters=filters, + orderBy="created_at", + limit=limit, + offset=offset, + session=session, + ) + return [serialize_query_execution_summary(e) for e in executions] + + @mcp.tool( + title="Run DataDoc Cell", + annotations=CREATE_ANNOTATIONS, + ) + def run_datadoc_cell( + cell_id: Annotated[int, "Cell ID from get_datadoc"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Execute a query cell from a DataDoc. Returns the query execution ID.""" + uid = token.claims["creator_uid"] + + with DBSession() as session: + # Get the cell and verify it's a query cell + cell = datadoc_logic.get_data_cell_by_id(cell_id, session=session) + if not cell: + raise ValueError(f"Cell {cell_id} not found.") + + if cell.cell_type.name != "query": + raise ValueError( + f"Cell {cell_id} is not a query cell (type: {cell.cell_type.name})." + ) + + # Check write permission on the parent datadoc (required to execute) + data_doc = cell.doc + try: + if not user_can_write(data_doc.id, uid, session=session): + raise ValueError( + "You do not have permission to execute this DataDoc." + ) + except DocDoesNotExist: + raise ValueError(f"DataDoc {data_doc.id} not found.") + + # Get engine_id from cell metadata + engine_id = cell.meta.get("engine") + if not engine_id: + raise ValueError("Cell does not have an engine specified.") + + # Verify engine permission (matches REST API's verify_query_engine_permission) + verify_query_engine_access(engine_id, uid, session) + + # Render query with datadoc variables + query = render_templated_query( + cell.context, + data_doc.meta_variables, + engine_id, + session=session, + ) + + # Create and run the query execution + query_execution = logic.create_query_execution( + query=query, + engine_id=engine_id, + uid=uid, + status=QueryExecutionStatus.INITIALIZED, + session=session, + ) + + # Mark execution as originating from MCP + logic.create_query_execution_metadata( + query_execution.id, _build_mcp_metadata(None), session=session + ) + + # Associate with the data cell + datadoc_logic.append_query_executions_to_data_cell( + cell_id, [query_execution.id], session=session + ) + + # Initiate execution (this queues it for processing) + from datasources.query_execution import initiate_query_execution + + initiate_query_execution( + query_execution=query_execution, + uid=uid, + peer_review_params=None, + session=session, + ) + + return serialize_query_execution(query_execution) + + @mcp.tool( + title="Run DataDoc", + annotations=CREATE_ANNOTATIONS, + ) + def run_datadoc( + datadoc_id: Annotated[int, "DataDoc ID"], + start_index: Annotated[int, "Index of first cell to execute (0-based)"] = 0, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Execute all query cells in a DataDoc starting from start_index, runs asynchronously via Celery.""" + uid = token.claims["creator_uid"] + + with DBSession() as session: + # Check write permission (required to execute) + try: + if not user_can_write(datadoc_id, uid, session=session): + raise ValueError( + "You do not have permission to execute this DataDoc." + ) + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + # Send Celery task to run the datadoc + from app.flask_app import celery + + celery.send_task( + "tasks.run_datadoc.run_datadoc", + args=[], + kwargs={ + "doc_id": datadoc_id, + "start_index": start_index, + "user_id": uid, + "execution_type": QueryExecutionType.ADHOC.value, + "notifications": [], + "metadata": {"source": "mcp"}, + }, + ) + + return { + "datadoc_id": datadoc_id, + "datadoc_resource_uri": f"querybook://datadoc/{datadoc_id}", + "start_index": start_index, + "message": "DataDoc execution queued. Use get_datadoc_cell_executions to check progress.", + } + + @mcp.tool( + title="Execute Ad-Hoc Query", + annotations=CREATE_ANNOTATIONS, + ) + def execute_ad_hoc_query( + query: Annotated[str, "SQL query to execute"], + engine_id: Annotated[int, "Query engine ID from list_query_engines"], + metadata: Annotated[dict | None, "Optional metadata for the execution"] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Execute an ad-hoc SQL query without creating a DataDoc. Returns execution details with query_execution_resource_uri for polling and results_resource_uri for each statement.""" + uid = token.claims["creator_uid"] + + with DBSession() as session: + # Verify engine permission + verify_query_engine_access(engine_id, uid, session) + + # Create query execution + query_execution = logic.create_query_execution( + query=query, + engine_id=engine_id, + uid=uid, + status=QueryExecutionStatus.INITIALIZED, + session=session, + ) + + # Always create metadata with MCP source marker + mcp_metadata = _build_mcp_metadata(metadata) + logic.create_query_execution_metadata( + query_execution.id, mcp_metadata, session=session + ) + + # Initiate execution (queues to Celery) + from datasources.query_execution import initiate_query_execution + + initiate_query_execution( + query_execution=query_execution, + uid=uid, + peer_review_params=None, + session=session, + ) + + return serialize_query_execution(query_execution) diff --git a/querybook/server/lib/mcp/tools/schedules.py b/querybook/server/lib/mcp/tools/schedules.py new file mode 100644 index 000000000..1a8f77421 --- /dev/null +++ b/querybook/server/lib/mcp/tools/schedules.py @@ -0,0 +1,289 @@ +from typing import Annotated, Literal +from pydantic import BaseModel, Field + +from fastmcp import FastMCP +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.mcp.lib.schedules import ( + DEFAULT_RETRY, + exports_to_kwargs, + get_datadoc_id_from_schedule, + notifications_to_kwargs, + retry_to_kwargs, + serialize_schedule, +) +from lib.mcp.utils import ( + READ_ONLY_ANNOTATIONS, + WRITE_ANNOTATIONS, + CREATE_ANNOTATIONS, + DELETE_ANNOTATIONS, +) +from logic.datadoc_permission import user_can_write, DocDoesNotExist +from logic.schedule import ( + create_task_schedule, + get_task_schedule_by_id, + get_scheduled_data_docs_by_user, + update_task_schedule, + delete_task_schedule, + get_data_doc_schedule_name, +) + + +class NotificationConfig(BaseModel): + """Configuration for schedule notifications.""" + + notifier: Annotated[ + str, + Field( + description="Notifier to use: 'email', 'slack', 'teams', etc. Available notifiers depend on server configuration." + ), + ] + on: Annotated[ + Literal["all", "failure", "success"], + Field( + description="When to notify: 'all' (every run), 'failure' (only on failure), 'success' (only on success)" + ), + ] = "all" + recipients: Annotated[ + list[str] | None, + Field( + description="Recipients to notify (e.g., email addresses, Slack channel IDs, etc.)" + ), + ] = None + user_ids: Annotated[ + list[int] | None, + Field(description="User IDs to notify via their preferred notifier"), + ] = None + + +class RetryConfig(BaseModel): + """Configuration for automatic retry on failure.""" + + enabled: Annotated[bool, Field(description="Enable automatic retry on failure")] = ( + False + ) + max_retries: Annotated[ + int, Field(description="Maximum number of retry attempts") + ] = 1 + delay_sec: Annotated[ + int, Field(description="Delay in seconds between retry attempts") + ] = 60 + + +class ExportConfig(BaseModel): + """Configuration for exporting query results after execution.""" + + cell_id: Annotated[int, Field(description="Cell ID to export results from")] + exporter_name: Annotated[ + str, Field(description="Exporter to use (e.g., 's3', 'gcs', 'email')") + ] + exporter_params: Annotated[ + dict, Field(description="Exporter-specific parameters") + ] = {} + + +def register(mcp: FastMCP) -> None: + """Register schedule tools on the given MCP server.""" + + @mcp.tool( + title="Create DataDoc Schedule", + annotations=CREATE_ANNOTATIONS, + ) + def create_datadoc_schedule( + datadoc_id: Annotated[int, "DataDoc ID to schedule"], + cron: Annotated[ + str, + "Cron expression in UTC timezone, e.g., '0 9 * * *' for daily at 9am UTC", + ], + enabled: Annotated[bool, "Whether the schedule is enabled"] = True, + notifications: Annotated[ + list[NotificationConfig] | None, + "Notification settings for schedule completion", + ] = None, + retry: Annotated[ + RetryConfig | None, "Automatic retry configuration on failure" + ] = None, + exports: Annotated[ + list[ExportConfig] | None, "Export configurations for query results" + ] = None, + disable_if_running: Annotated[ + bool, "Skip scheduled run if DataDoc is already running" + ] = False, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Schedule a DataDoc to run on a cron schedule with optional notifications, retries, and exports.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + try: + if not user_can_write(datadoc_id, uid, session=session): + raise ValueError( + "You do not have permission to schedule this DataDoc." + ) + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + schedule_name = get_data_doc_schedule_name(datadoc_id) + + # Build kwargs from structured parameters + kwargs = { + "user_id": uid, + "doc_id": datadoc_id, + "disable_if_running_doc": disable_if_running, + } + + # Convert structured configs to internal format + kwargs["notifications"] = ( + notifications_to_kwargs(notifications) if notifications else [] + ) + kwargs["retry"] = retry_to_kwargs(retry) if retry else dict(DEFAULT_RETRY) + kwargs["exports"] = exports_to_kwargs(exports) if exports else [] + + schedule = create_task_schedule( + name=schedule_name, + task="tasks.run_datadoc.run_datadoc", + cron=cron, + kwargs=kwargs, + enabled=enabled, + commit=True, + session=session, + ) + + return serialize_schedule(schedule) + + @mcp.tool( + title="List DataDoc Schedules", + annotations=READ_ONLY_ANNOTATIONS, + ) + def list_datadoc_schedules( + environment_id: Annotated[int, "Environment ID from list_environments"], + enabled: Annotated[bool | None, "Filter by enabled status"] = None, + offset: Annotated[int, "Pagination offset"] = 0, + limit: Annotated[int, "Maximum number of results"] = 100, + token: AccessToken = CurrentAccessToken(), + ) -> list[dict]: + """List all DataDoc schedules for the current user in an environment.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + filters = {} + if enabled is not None: + filters["status"] = enabled + + docs_with_schedules, count = get_scheduled_data_docs_by_user( + uid, environment_id, offset, limit, filters=filters, session=session + ) + + result = [] + for item in docs_with_schedules: + schedule = item.get("schedule") + if schedule: + result.append(serialize_schedule(schedule)) + + return result + + @mcp.tool( + title="Update DataDoc Schedule", + annotations=WRITE_ANNOTATIONS, + ) + def update_datadoc_schedule( + schedule_id: Annotated[int, "Schedule ID"], + cron: Annotated[ + str | None, + "Cron expression in UTC timezone, e.g., '0 9 * * *' for daily at 9am UTC", + ] = None, + enabled: Annotated[bool | None, "Enable or disable the schedule"] = None, + notifications: Annotated[ + list[NotificationConfig] | None, + "New notification settings (replaces existing)", + ] = None, + retry: Annotated[ + RetryConfig | None, "New retry configuration (replaces existing)" + ] = None, + exports: Annotated[ + list[ExportConfig] | None, "New export configurations (replaces existing)" + ] = None, + disable_if_running: Annotated[ + bool | None, "Skip scheduled run if DataDoc is already running" + ] = None, + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Update a DataDoc schedule's settings. Only non-null fields are updated.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + schedule = get_task_schedule_by_id(schedule_id, session=session) + if not schedule: + raise ValueError(f"Schedule {schedule_id} not found.") + + datadoc_id = get_datadoc_id_from_schedule(schedule) + if datadoc_id is None: + raise ValueError("Schedule is not a DataDoc schedule.") + + try: + if not user_can_write(datadoc_id, uid, session=session): + raise ValueError( + "You do not have permission to update this schedule." + ) + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + fields = {} + if cron is not None: + fields["cron"] = cron + if enabled is not None: + fields["enabled"] = enabled + + # Update kwargs if any advanced settings changed + if any( + x is not None + for x in [notifications, retry, exports, disable_if_running] + ): + kwargs = dict(schedule.kwargs) # Copy existing kwargs + + if notifications is not None: + kwargs["notifications"] = notifications_to_kwargs(notifications) + if retry is not None: + kwargs["retry"] = retry_to_kwargs(retry) + if exports is not None: + kwargs["exports"] = exports_to_kwargs(exports) + + if disable_if_running is not None: + kwargs["disable_if_running_doc"] = disable_if_running + + fields["kwargs"] = kwargs + + update_task_schedule(schedule_id, commit=True, session=session, **fields) + + # Return updated schedule + schedule = get_task_schedule_by_id(schedule_id, session=session) + return serialize_schedule(schedule) + + @mcp.tool( + title="Delete DataDoc Schedule", + annotations=DELETE_ANNOTATIONS, + ) + def delete_datadoc_schedule( + schedule_id: Annotated[int, "Schedule ID"], + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Delete a DataDoc schedule.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + schedule = get_task_schedule_by_id(schedule_id, session=session) + if not schedule: + raise ValueError(f"Schedule {schedule_id} not found.") + + datadoc_id = get_datadoc_id_from_schedule(schedule) + if datadoc_id is None: + raise ValueError("Schedule is not a DataDoc schedule.") + + try: + if not user_can_write(datadoc_id, uid, session=session): + raise ValueError( + "You do not have permission to delete this schedule." + ) + except DocDoesNotExist: + raise ValueError(f"DataDoc {datadoc_id} not found.") + + delete_task_schedule(schedule_id, commit=True, session=session) + return {"deleted": schedule_id} diff --git a/querybook/server/lib/mcp/tools/users.py b/querybook/server/lib/mcp/tools/users.py new file mode 100644 index 000000000..ef6e2cf78 --- /dev/null +++ b/querybook/server/lib/mcp/tools/users.py @@ -0,0 +1,73 @@ +from typing import Annotated + +from fastmcp import FastMCP +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import CurrentAccessToken + +from app.db import DBSession +from lib.elasticsearch.search_utils import get_matching_suggestions, ES_CONFIG +from lib.elasticsearch.suggest_user import construct_suggest_user_query +from lib.mcp.lib.users import serialize_user, get_user_data +from lib.mcp.utils import READ_ONLY_ANNOTATIONS +from logic.user import get_users_by_ids + + +def register(mcp: FastMCP) -> None: + """Register user tools on the given MCP server.""" + + @mcp.tool( + title="Get Current User", + annotations=READ_ONLY_ANNOTATIONS, + ) + def get_current_user( + token: AccessToken = CurrentAccessToken(), + ) -> dict: + """Get the current authenticated user's profile information.""" + uid = token.claims["creator_uid"] + with DBSession() as session: + return get_user_data("me", uid, session) + + @mcp.tool( + title="Get Users", + annotations=READ_ONLY_ANNOTATIONS, + ) + def get_users( + user_ids: Annotated[list[int], "List of user IDs to retrieve, max 100"], + token: AccessToken = CurrentAccessToken(), + ) -> list[dict]: + """Get user information for multiple user IDs in batch.""" + if len(user_ids) > 100: + raise ValueError("Requesting too many users (max 100)") + with DBSession() as session: + users = get_users_by_ids(user_ids, session=session) + return [serialize_user(user) for user in users] + + @mcp.tool( + title="Search Users", + annotations=READ_ONLY_ANNOTATIONS, + ) + def search_users( + query: Annotated[str, "Search query to match against username or fullname"], + limit: Annotated[int, "Maximum number of results"] = 20, + token: AccessToken = CurrentAccessToken(), + ) -> list[dict]: + """Search for users by username or fullname with case-insensitive prefix match.""" + if limit > 100: + raise ValueError("Requesting too many users (max 100)") + + es_query = construct_suggest_user_query(prefix=query, limit=limit) + options = get_matching_suggestions(es_query, ES_CONFIG["users"]["index_name"]) + + users = [] + for option in options: + user_id = option.get("_source", {}).get("id") + if user_id: + users.append( + { + "id": user_id, + "username": option.get("_source", {}).get("username"), + "fullname": option.get("_source", {}).get("fullname"), + "resource_uri": f"querybook://user/{user_id}", + } + ) + return users diff --git a/querybook/server/lib/mcp/utils.py b/querybook/server/lib/mcp/utils.py new file mode 100644 index 000000000..5a117c828 --- /dev/null +++ b/querybook/server/lib/mcp/utils.py @@ -0,0 +1,90 @@ +"""Shared utilities for MCP tools and resources.""" + +from logic import admin as admin_logic +from logic.environment import get_all_accessible_environment_ids_by_uid +from models.admin import QueryEngineEnvironment +from env import QuerybookSettings + + +def build_querybook_url(environment_name: str, path: str) -> str | None: + """Build a Querybook web UI URL. + + Args: + environment_name: Environment name for the URL path + path: Resource path, e.g. "datadoc/123" or "list/456" + + Returns: + Full URL string, or None if PUBLIC_URL is not configured + """ + if not QuerybookSettings.PUBLIC_URL: + return None + return f"{QuerybookSettings.PUBLIC_URL}/{environment_name}/{path}/" + + +def verify_query_engine_access(engine_id: int, uid: int, session) -> None: + """Verify user has access to query engine through environment permissions. + + This mimics the upstream permission.py verify_query_engine_permission logic. + + Args: + engine_id: Query engine ID to check + uid: User ID making the request + session: Database session + + Raises: + ValueError: If engine doesn't exist or user lacks access + """ + # Get the query engine + engine = admin_logic.get_query_engine_by_id(engine_id, session=session) + if not engine: + raise ValueError(f"Query engine {engine_id} not found.") + + # Get environments this engine belongs to + engine_env_ids = [ + eid + for eid, in session.query(QueryEngineEnvironment.environment_id).filter( + QueryEngineEnvironment.query_engine_id == engine_id + ) + ] + + # Get environments the user has access to + accessible_env_ids = get_all_accessible_environment_ids_by_uid(uid, session=session) + + # Check if user has access to at least one environment that has this engine + if not any(eid in accessible_env_ids for eid in engine_env_ids): + raise ValueError(f"You do not have access to query engine {engine_id}.") + + +# Annotation constants for MCP tools and resources +READ_ONLY_ANNOTATIONS = { + "destructiveHint": False, + "idempotentHint": True, + "openWorldHint": False, + "readOnlyHint": True, +} + +WRITE_ANNOTATIONS = { + "destructiveHint": False, + "idempotentHint": True, + "openWorldHint": False, + "readOnlyHint": False, +} + +CREATE_ANNOTATIONS = { + "destructiveHint": False, + "idempotentHint": False, + "openWorldHint": False, + "readOnlyHint": False, +} + +DELETE_ANNOTATIONS = { + "destructiveHint": True, + "idempotentHint": True, + "openWorldHint": False, + "readOnlyHint": False, +} + +RESOURCE_ANNOTATIONS = { + "readOnlyHint": True, + "idempotentHint": True, +} diff --git a/querybook/server/run_mcp.py b/querybook/server/run_mcp.py new file mode 100644 index 000000000..c4290d1f3 --- /dev/null +++ b/querybook/server/run_mcp.py @@ -0,0 +1,66 @@ +from fastmcp import FastMCP + +from env import QuerybookSettings +from lib.mcp.auth import QuerybookTokenVerifier +from lib.mcp.middleware import MCPEventLoggingMiddleware, wrap_mcp_resources +from lib.mcp.resources import ( + comments as comments_resources, + datadocs as datadocs_resources, + environments as environments_resources, + guide as guide_resources, + lists as lists_resources, + query_engines as query_engines_resources, + query_executions as query_executions_resources, + schedules as schedules_resources, + statement_executions as statement_executions_resources, + users as users_resources, +) +from lib.mcp.tools import ( + comments, + datadocs, + environments, + lists, + query_engines, + query_executions, + schedules, + users, +) + +# Querybook MCP server +mcp = FastMCP("Querybook MCP", auth=QuerybookTokenVerifier()) + +# Add event logging middleware for tools (must be registered before tools) +mcp.add_middleware(MCPEventLoggingMiddleware()) + +# Wrap resource decorator to add logging (FastMCP middleware doesn't support resource hooks) +wrap_mcp_resources(mcp) + +comments.register(mcp) +datadocs.register(mcp) +environments.register(mcp) +lists.register(mcp) +query_engines.register(mcp) +query_executions.register(mcp) +schedules.register(mcp) +users.register(mcp) + +# Register resources +guide_resources.register(mcp) +comments_resources.register(mcp) +datadocs_resources.register(mcp) +environments_resources.register(mcp) +lists_resources.register(mcp) +query_engines_resources.register(mcp) +query_executions_resources.register(mcp) +schedules_resources.register(mcp) +statement_executions_resources.register(mcp) +users_resources.register(mcp) + +if __name__ == "__main__": + mcp.run( + transport="http", + host="0.0.0.0", + port=QuerybookSettings.MCP_PORT, + # Stateless mode used to enable horizontally-scaled deployments. + stateless_http=True, + ) diff --git a/requirements/base.txt b/requirements/base.txt index 58383357e..420a1bd76 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -40,9 +40,12 @@ markdown2 # Utils pandas==1.5.3 -typing-extensions==4.9.0 +typing-extensions==4.15.0 setuptools>=65.5.1 # not directly required, pinned by Snyk to avoid a vulnerability numpy>=1.22.2,<2.0.0 # not directly required, pinned by Snyk to avoid a vulnerability # Query engine - PostgreSQL psycopg2==2.9.5 + +# MCP server +fastmcp==3.0.2