feat(mcp): MCP service implementation (PRs 3-9 consolidated) (#35877)
This commit is contained in:
Amin Ghadersohi
@@ -3,14 +3,3 @@
|
||||
For complete documentation on using GitHub Codespaces with Apache Superset, please see:
|
||||
|
||||
**[Setting up a Development Environment - GitHub Codespaces](https://superset.apache.org/docs/contributing/development#github-codespaces-cloud-development)**
|
||||
|
||||
## Pre-installed Development Environment
|
||||
|
||||
When you create a new Codespace from this repository, it automatically:
|
||||
|
||||
1. **Creates a Python virtual environment** using `uv venv`
|
||||
2. **Installs all development dependencies** via `uv pip install -r requirements/development.txt`
|
||||
3. **Sets up pre-commit hooks** with `pre-commit install`
|
||||
4. **Activates the virtual environment** automatically in all terminals
|
||||
|
||||
The virtual environment is located at `/workspaces/{repository-name}/.venv` and is automatically activated through environment variables set in the devcontainer configuration.
|
||||
|
||||
19
.devcontainer/default/devcontainer.json
Normal file
19
.devcontainer/default/devcontainer.json
Normal file
@@ -0,0 +1,19 @@
|
||||
{
|
||||
// Extend the base configuration
|
||||
"extends": "../devcontainer-base.json",
|
||||
|
||||
"name": "Apache Superset Development (Default)",
|
||||
|
||||
// Forward ports for development
|
||||
"forwardPorts": [9001],
|
||||
"portsAttributes": {
|
||||
"9001": {
|
||||
"label": "Superset (via Webpack Dev Server)",
|
||||
"onAutoForward": "notify",
|
||||
"visibility": "public"
|
||||
}
|
||||
},
|
||||
|
||||
// Auto-start Superset on Codespace resume
|
||||
"postStartCommand": ".devcontainer/start-superset.sh"
|
||||
}
|
||||
39
.devcontainer/devcontainer-base.json
Normal file
39
.devcontainer/devcontainer-base.json
Normal file
@@ -0,0 +1,39 @@
|
||||
{
|
||||
"name": "Apache Superset Development",
|
||||
// Keep this in sync with the base image in Dockerfile (ARG PY_VER)
|
||||
// Using the same base as Dockerfile, but non-slim for dev tools
|
||||
"image": "python:3.11.13-bookworm",
|
||||
|
||||
"features": {
|
||||
"ghcr.io/devcontainers/features/docker-in-docker:2": {
|
||||
"moby": true,
|
||||
"dockerDashComposeVersion": "v2"
|
||||
},
|
||||
"ghcr.io/devcontainers/features/node:1": {
|
||||
"version": "20"
|
||||
},
|
||||
"ghcr.io/devcontainers/features/git:1": {},
|
||||
"ghcr.io/devcontainers/features/common-utils:2": {
|
||||
"configureZshAsDefaultShell": true
|
||||
},
|
||||
"ghcr.io/devcontainers/features/sshd:1": {
|
||||
"version": "latest"
|
||||
}
|
||||
},
|
||||
|
||||
// Run commands after container is created
|
||||
"postCreateCommand": "chmod +x .devcontainer/setup-dev.sh && .devcontainer/setup-dev.sh",
|
||||
|
||||
// VS Code customizations
|
||||
"customizations": {
|
||||
"vscode": {
|
||||
"extensions": [
|
||||
"ms-python.python",
|
||||
"ms-python.vscode-pylance",
|
||||
"charliermarsh.ruff",
|
||||
"dbaeumer.vscode-eslint",
|
||||
"esbenp.prettier-vscode"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3,76 +3,30 @@
|
||||
|
||||
echo "🔧 Setting up Superset development environment..."
|
||||
|
||||
# System dependencies and uv are now pre-installed in the Docker image
|
||||
# This speeds up Codespace creation significantly!
|
||||
# The universal image has most tools, just need Superset-specific libs
|
||||
echo "📦 Installing Superset-specific dependencies..."
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y \
|
||||
libsasl2-dev \
|
||||
libldap2-dev \
|
||||
libpq-dev \
|
||||
tmux \
|
||||
gh
|
||||
|
||||
# Create virtual environment using uv
|
||||
echo "🐍 Creating Python virtual environment..."
|
||||
if ! uv venv; then
|
||||
echo "❌ Failed to create virtual environment"
|
||||
exit 1
|
||||
fi
|
||||
# Install uv for fast Python package management
|
||||
echo "📦 Installing uv..."
|
||||
curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
|
||||
# Install Python dependencies
|
||||
echo "📦 Installing Python dependencies..."
|
||||
if ! uv pip install -r requirements/development.txt; then
|
||||
echo "❌ Failed to install Python dependencies"
|
||||
echo "💡 You may need to run this manually after the Codespace starts"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Install pre-commit hooks
|
||||
echo "🪝 Installing pre-commit hooks..."
|
||||
if source .venv/bin/activate && pre-commit install; then
|
||||
echo "✅ Pre-commit hooks installed"
|
||||
else
|
||||
echo "⚠️ Pre-commit hooks installation failed (non-critical)"
|
||||
fi
|
||||
# Add cargo/bin to PATH for uv
|
||||
echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> ~/.bashrc
|
||||
echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> ~/.zshrc
|
||||
|
||||
# Install Claude Code CLI via npm
|
||||
echo "🤖 Installing Claude Code..."
|
||||
if npm install -g @anthropic-ai/claude-code; then
|
||||
echo "✅ Claude Code installed"
|
||||
else
|
||||
echo "⚠️ Claude Code installation failed (non-critical)"
|
||||
fi
|
||||
npm install -g @anthropic-ai/claude-code
|
||||
|
||||
# Make the start script executable
|
||||
chmod +x .devcontainer/start-superset.sh
|
||||
|
||||
# Add bashrc additions for automatic venv activation
|
||||
echo "🔧 Setting up automatic environment activation..."
|
||||
if [ -f ~/.bashrc ]; then
|
||||
# Check if we've already added our additions
|
||||
if ! grep -q "Superset Codespaces environment setup" ~/.bashrc; then
|
||||
echo "" >> ~/.bashrc
|
||||
cat .devcontainer/bashrc-additions >> ~/.bashrc
|
||||
echo "✅ Added automatic venv activation to ~/.bashrc"
|
||||
else
|
||||
echo "✅ Bashrc additions already present"
|
||||
fi
|
||||
else
|
||||
# Create bashrc if it doesn't exist
|
||||
cat .devcontainer/bashrc-additions > ~/.bashrc
|
||||
echo "✅ Created ~/.bashrc with automatic venv activation"
|
||||
fi
|
||||
|
||||
# Also add to zshrc since that's the default shell
|
||||
if [ -f ~/.zshrc ] || [ -n "$ZSH_VERSION" ]; then
|
||||
if ! grep -q "Superset Codespaces environment setup" ~/.zshrc; then
|
||||
echo "" >> ~/.zshrc
|
||||
cat .devcontainer/bashrc-additions >> ~/.zshrc
|
||||
echo "✅ Added automatic venv activation to ~/.zshrc"
|
||||
fi
|
||||
fi
|
||||
|
||||
echo "✅ Development environment setup complete!"
|
||||
echo ""
|
||||
echo "📝 The virtual environment will be automatically activated in new terminals"
|
||||
echo ""
|
||||
echo "🔄 To activate in this terminal, run:"
|
||||
echo " source ~/.bashrc"
|
||||
echo ""
|
||||
echo "🚀 To start Superset:"
|
||||
echo " start-superset"
|
||||
echo ""
|
||||
echo "🚀 Run '.devcontainer/start-superset.sh' to start Superset"
|
||||
|
||||
@@ -1,14 +1,14 @@
|
||||
#!/bin/bash
|
||||
# Startup script for Superset in Codespaces
|
||||
|
||||
# Log to a file for debugging
|
||||
LOG_FILE="/tmp/superset-startup.log"
|
||||
echo "[$(date)] Starting Superset startup script" >> "$LOG_FILE"
|
||||
echo "[$(date)] User: $(whoami), PWD: $(pwd)" >> "$LOG_FILE"
|
||||
|
||||
echo "🚀 Starting Superset in Codespaces..."
|
||||
echo "🌐 Frontend will be available at port 9001"
|
||||
|
||||
# Check if MCP is enabled
|
||||
if [ "$ENABLE_MCP" = "true" ]; then
|
||||
echo "🤖 MCP Service will be available at port 5008"
|
||||
fi
|
||||
|
||||
# Find the workspace directory (Codespaces clones as 'superset', not 'superset-2')
|
||||
WORKSPACE_DIR=$(find /workspaces -maxdepth 1 -name "superset*" -type d | head -1)
|
||||
if [ -n "$WORKSPACE_DIR" ]; then
|
||||
@@ -18,71 +18,32 @@ else
|
||||
echo "📁 Using current directory: $(pwd)"
|
||||
fi
|
||||
|
||||
# Wait for Docker to be available
|
||||
echo "⏳ Waiting for Docker to start..."
|
||||
echo "[$(date)] Waiting for Docker..." >> "$LOG_FILE"
|
||||
max_attempts=30
|
||||
attempt=0
|
||||
while ! docker info > /dev/null 2>&1; do
|
||||
if [ $attempt -eq $max_attempts ]; then
|
||||
echo "❌ Docker failed to start after $max_attempts attempts"
|
||||
echo "[$(date)] Docker failed to start after $max_attempts attempts" >> "$LOG_FILE"
|
||||
echo "🔄 Please restart the Codespace or run this script manually later"
|
||||
exit 1
|
||||
fi
|
||||
echo " Attempt $((attempt + 1))/$max_attempts..."
|
||||
echo "[$(date)] Docker check attempt $((attempt + 1))/$max_attempts" >> "$LOG_FILE"
|
||||
sleep 2
|
||||
attempt=$((attempt + 1))
|
||||
done
|
||||
echo "✅ Docker is ready!"
|
||||
echo "[$(date)] Docker is ready" >> "$LOG_FILE"
|
||||
|
||||
# Check if Superset containers are already running
|
||||
if docker ps | grep -q "superset"; then
|
||||
echo "✅ Superset containers are already running!"
|
||||
echo ""
|
||||
echo "🌐 To access Superset:"
|
||||
echo " 1. Click the 'Ports' tab at the bottom of VS Code"
|
||||
echo " 2. Find port 9001 and click the globe icon to open"
|
||||
echo " 3. Wait 10-20 minutes for initial startup"
|
||||
echo ""
|
||||
echo "📝 Login credentials: admin/admin"
|
||||
exit 0
|
||||
# Check if docker is running
|
||||
if ! docker info > /dev/null 2>&1; then
|
||||
echo "⏳ Waiting for Docker to start..."
|
||||
sleep 5
|
||||
fi
|
||||
|
||||
# Clean up any existing containers
|
||||
echo "🧹 Cleaning up existing containers..."
|
||||
docker-compose -f docker-compose-light.yml down
|
||||
docker-compose -f docker-compose-light.yml --profile mcp down
|
||||
|
||||
# Start services
|
||||
echo "🏗️ Starting Superset in background (daemon mode)..."
|
||||
echo "🏗️ Building and starting services..."
|
||||
echo ""
|
||||
echo "📝 Once started, login with:"
|
||||
echo " Username: admin"
|
||||
echo " Password: admin"
|
||||
echo ""
|
||||
echo "📋 Running in foreground with live logs (Ctrl+C to stop)..."
|
||||
|
||||
# Start in detached mode
|
||||
docker-compose -f docker-compose-light.yml up -d
|
||||
|
||||
echo ""
|
||||
echo "✅ Docker Compose started successfully!"
|
||||
echo ""
|
||||
echo "📋 Important information:"
|
||||
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
|
||||
echo "⏱️ Initial startup takes 10-20 minutes"
|
||||
echo "🌐 Check the 'Ports' tab for your Superset URL (port 9001)"
|
||||
echo "👤 Login: admin / admin"
|
||||
echo ""
|
||||
echo "📊 Useful commands:"
|
||||
echo " docker-compose -f docker-compose-light.yml logs -f # Follow logs"
|
||||
echo " docker-compose -f docker-compose-light.yml ps # Check status"
|
||||
echo " docker-compose -f docker-compose-light.yml down # Stop services"
|
||||
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
|
||||
echo ""
|
||||
echo "💤 Keeping terminal open for 60 seconds to test persistence..."
|
||||
sleep 60
|
||||
echo "✅ Test complete - check if this terminal is still visible!"
|
||||
|
||||
# Show final status
|
||||
docker-compose -f docker-compose-light.yml ps
|
||||
# Run docker-compose and capture exit code
|
||||
if [ "$ENABLE_MCP" = "true" ]; then
|
||||
echo "🤖 Starting with MCP Service enabled..."
|
||||
docker-compose -f docker-compose-light.yml --profile mcp up
|
||||
else
|
||||
docker-compose -f docker-compose-light.yml up
|
||||
fi
|
||||
EXIT_CODE=$?
|
||||
|
||||
# If it failed, provide helpful instructions
|
||||
|
||||
29
.devcontainer/with-mcp/devcontainer.json
Normal file
29
.devcontainer/with-mcp/devcontainer.json
Normal file
@@ -0,0 +1,29 @@
|
||||
{
|
||||
// Extend the base configuration
|
||||
"extends": "../devcontainer-base.json",
|
||||
|
||||
"name": "Apache Superset Development with MCP",
|
||||
|
||||
// Forward ports for development
|
||||
"forwardPorts": [9001, 5008],
|
||||
"portsAttributes": {
|
||||
"9001": {
|
||||
"label": "Superset (via Webpack Dev Server)",
|
||||
"onAutoForward": "notify",
|
||||
"visibility": "public"
|
||||
},
|
||||
"5008": {
|
||||
"label": "MCP Service (Model Context Protocol)",
|
||||
"onAutoForward": "notify",
|
||||
"visibility": "private"
|
||||
}
|
||||
},
|
||||
|
||||
// Auto-start Superset with MCP on Codespace resume
|
||||
"postStartCommand": "ENABLE_MCP=true .devcontainer/start-superset.sh",
|
||||
|
||||
// Environment variables
|
||||
"containerEnv": {
|
||||
"ENABLE_MCP": "true"
|
||||
}
|
||||
}
|
||||
@@ -86,6 +86,10 @@ case "${1}" in
|
||||
echo "Starting web app..."
|
||||
/usr/bin/run-server.sh
|
||||
;;
|
||||
mcp)
|
||||
echo "Starting MCP service..."
|
||||
superset mcp run --host 0.0.0.0 --port ${MCP_PORT:-5008} --debug
|
||||
;;
|
||||
*)
|
||||
echo "Unknown Operation!!!"
|
||||
;;
|
||||
|
||||
523
superset/mcp_service/CLAUDE.md
Normal file
523
superset/mcp_service/CLAUDE.md
Normal file
@@ -0,0 +1,523 @@
|
||||
# MCP Service - LLM Agent Guide
|
||||
|
||||
This guide helps LLM agents understand the Superset MCP (Model Context Protocol) service architecture and development conventions.
|
||||
|
||||
## ⚠️ CRITICAL: Apache License Headers
|
||||
|
||||
**EVERY Python file in the MCP service MUST have the Apache Software Foundation license header.**
|
||||
|
||||
This includes:
|
||||
- All `.py` files (tool files, schemas, __init__.py files, etc.)
|
||||
- **NEVER remove existing license headers during refactoring or edits**
|
||||
- **ALWAYS add license headers when creating new files**
|
||||
- **ALWAYS verify license headers are present after editing files**
|
||||
|
||||
If you see a file without a license header, ADD IT IMMEDIATELY. If you accidentally remove one during editing, ADD IT BACK.
|
||||
|
||||
Use this exact template at the top of EVERY Python file:
|
||||
|
||||
```python
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
```
|
||||
|
||||
**Note**: LLM instruction files like `CLAUDE.md`, `AGENTS.md`, etc. are excluded from this requirement (listed in `.rat-excludes`) to avoid token overhead, but ALL other Python files require it.
|
||||
|
||||
## Architecture Overview
|
||||
|
||||
The MCP service provides programmatic access to Superset via the Model Context Protocol, allowing AI assistants to interact with dashboards, charts, datasets, SQL Lab, and instance metadata.
|
||||
|
||||
### Key Components
|
||||
|
||||
```
|
||||
superset/mcp_service/
|
||||
├── app.py # FastMCP app factory and tool registration
|
||||
├── auth.py # Authentication and authorization
|
||||
├── mcp_config.py # Default configuration
|
||||
├── mcp_core.py # Reusable core classes for tools
|
||||
├── flask_singleton.py # Flask app singleton for MCP context
|
||||
├── chart/ # Chart-related tools
|
||||
│ ├── schemas.py # Pydantic schemas for chart responses
|
||||
│ └── tool/ # Chart tool implementations
|
||||
│ ├── __init__.py # Tool exports
|
||||
│ ├── list_charts.py
|
||||
│ └── get_chart_info.py
|
||||
├── dashboard/ # Dashboard-related tools
|
||||
│ ├── schemas.py
|
||||
│ └── tool/
|
||||
├── dataset/ # Dataset-related tools
|
||||
│ ├── schemas.py
|
||||
│ └── tool/
|
||||
└── system/ # System/instance tools
|
||||
├── schemas.py
|
||||
└── tool/
|
||||
```
|
||||
|
||||
## Critical Convention: Tool, Prompt, and Resource Registration
|
||||
|
||||
**IMPORTANT**: When creating new MCP tools, prompts, or resources, you MUST add their imports to `app.py` for auto-registration. Do NOT add them to `server.py` - that approach doesn't work properly.
|
||||
|
||||
### How to Add a New Tool
|
||||
|
||||
1. **Create the tool file** in the appropriate directory (e.g., `chart/tool/my_new_tool.py`)
|
||||
2. **Decorate with `@mcp.tool`** to register it with FastMCP
|
||||
3. **Add import to `app.py`** at the bottom of the file where other tools are imported (around line 210-242)
|
||||
|
||||
**Example**:
|
||||
```python
|
||||
# superset/mcp_service/chart/tool/my_new_tool.py
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
def my_new_tool(param: str) -> dict:
|
||||
"""Tool description for LLMs."""
|
||||
return {"result": "success"}
|
||||
```
|
||||
|
||||
**Then add to app.py**:
|
||||
```python
|
||||
# superset/mcp_service/app.py (at the bottom, around line 207-224)
|
||||
from superset.mcp_service.chart.tool import ( # noqa: F401, E402
|
||||
get_chart_info,
|
||||
list_charts,
|
||||
my_new_tool, # ADD YOUR TOOL HERE
|
||||
)
|
||||
```
|
||||
|
||||
**Why this matters**: Tools use `@mcp.tool` decorators and register automatically on import. The import MUST be in `app.py` at the bottom of the file (after the `mcp` instance is created). If you don't import the tool in `app.py`, it won't be available to MCP clients. DO NOT add imports to `server.py` - that file is for running the server only.
|
||||
|
||||
### How to Add a New Prompt
|
||||
|
||||
1. **Create the prompt file** in the appropriate directory (e.g., `chart/prompts/my_new_prompt.py`)
|
||||
2. **Decorate with `@mcp.prompt`** to register it with FastMCP
|
||||
3. **Add import to module's `__init__.py`** (e.g., `chart/prompts/__init__.py`)
|
||||
4. **Ensure module is imported in `app.py`** (around line 244-253)
|
||||
|
||||
**Example**:
|
||||
```python
|
||||
# superset/mcp_service/chart/prompts/my_new_prompt.py
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
|
||||
@mcp.prompt("my_new_prompt")
|
||||
@mcp_auth_hook
|
||||
async def my_new_prompt_handler(ctx: Context) -> str:
|
||||
"""Interactive prompt for doing something."""
|
||||
return "Prompt instructions here..."
|
||||
```
|
||||
|
||||
**Then add to `chart/prompts/__init__.py`**:
|
||||
```python
|
||||
# superset/mcp_service/chart/prompts/__init__.py
|
||||
from . import create_chart_guided # existing
|
||||
from . import my_new_prompt # ADD YOUR PROMPT HERE
|
||||
```
|
||||
|
||||
**Verify module import exists in `app.py`** (around line 248):
|
||||
```python
|
||||
# superset/mcp_service/app.py
|
||||
from superset.mcp_service.chart import prompts as chart_prompts # This imports all prompts
|
||||
```
|
||||
|
||||
### How to Add a New Resource
|
||||
|
||||
1. **Create the resource file** in the appropriate directory (e.g., `chart/resources/my_new_resource.py`)
|
||||
2. **Decorate with `@mcp.resource`** to register it with FastMCP
|
||||
3. **Add import to module's `__init__.py`** (e.g., `chart/resources/__init__.py`)
|
||||
4. **Ensure module is imported in `app.py`** (around line 244-253)
|
||||
|
||||
**Example**:
|
||||
```python
|
||||
# superset/mcp_service/chart/resources/my_new_resource.py
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
|
||||
@mcp.resource("superset://chart/my_resource")
|
||||
@mcp_auth_hook
|
||||
def get_my_resource() -> str:
|
||||
"""Resource description for LLMs."""
|
||||
return "Resource data here..."
|
||||
```
|
||||
|
||||
**Then add to `chart/resources/__init__.py`**:
|
||||
```python
|
||||
# superset/mcp_service/chart/resources/__init__.py
|
||||
from . import chart_configs # existing
|
||||
from . import my_new_resource # ADD YOUR RESOURCE HERE
|
||||
```
|
||||
|
||||
**Verify module import exists in `app.py`** (around line 249):
|
||||
```python
|
||||
# superset/mcp_service/app.py
|
||||
from superset.mcp_service.chart import resources as chart_resources # This imports all resources
|
||||
```
|
||||
|
||||
**Why this matters**: Prompts and resources work similarly to tools - they use decorators and register on import. The module-level imports (`chart/prompts/__init__.py`, `chart/resources/__init__.py`) ensure individual files are imported when the module is imported. The `app.py` imports ensure the modules are loaded when the MCP service starts.
|
||||
|
||||
## Tool Development Patterns
|
||||
|
||||
### 1. Use Core Classes for Reusability
|
||||
|
||||
The `mcp_core.py` module provides reusable patterns:
|
||||
|
||||
- **`ModelListCore`**: For listing resources (dashboards, charts, datasets)
|
||||
- **`ModelGetInfoCore`**: For getting resource details by ID/UUID
|
||||
- **`ModelGetAvailableFiltersCore`**: For retrieving filterable columns
|
||||
|
||||
**Example**:
|
||||
```python
|
||||
from superset.mcp_service.mcp_core import ModelListCore
|
||||
from superset.daos.dashboard import DashboardDAO
|
||||
from superset.mcp_service.dashboard.schemas import DashboardList
|
||||
|
||||
list_core = ModelListCore(
|
||||
dao_class=DashboardDAO,
|
||||
output_schema=DashboardList,
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
def list_dashboards(filters: List[DashboardFilter], page: int = 1) -> DashboardList:
|
||||
return list_core.run_tool(filters=filters, page=page, page_size=10)
|
||||
```
|
||||
|
||||
### 2. Always Use Authentication
|
||||
|
||||
**Every tool must use `@mcp_auth_hook`** to ensure:
|
||||
- User authentication from JWT or configured admin user
|
||||
- Permission checking via JWT scopes
|
||||
- Audit logging of tool access
|
||||
|
||||
```python
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook # REQUIRED
|
||||
def my_tool() -> dict:
|
||||
# g.user is set by mcp_auth_hook
|
||||
return {"user": g.user.username}
|
||||
```
|
||||
|
||||
### 3. Use Pydantic Schemas
|
||||
|
||||
**All tool inputs and outputs must be Pydantic models** for:
|
||||
- Automatic validation
|
||||
- LLM-friendly schema generation
|
||||
- Type safety
|
||||
|
||||
**Convention**: Place schemas in `{module}/schemas.py`
|
||||
|
||||
```python
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
class MyToolRequest(BaseModel):
|
||||
param: str = Field(..., description="Parameter description for LLMs")
|
||||
|
||||
class MyToolResponse(BaseModel):
|
||||
result: str = Field(..., description="Result description")
|
||||
timestamp: datetime = Field(
|
||||
default_factory=lambda: datetime.now(timezone.utc),
|
||||
description="Response timestamp"
|
||||
)
|
||||
```
|
||||
|
||||
### 4. Follow the DAO Pattern
|
||||
|
||||
**Use Superset's DAO (Data Access Object) layer** instead of direct database queries:
|
||||
|
||||
```python
|
||||
from superset.daos.dashboard import DashboardDAO
|
||||
|
||||
# GOOD: Use DAO
|
||||
dashboard = DashboardDAO.find_by_id(dashboard_id)
|
||||
|
||||
# BAD: Don't query directly
|
||||
dashboard = db.session.query(Dashboard).filter_by(id=dashboard_id).first()
|
||||
```
|
||||
|
||||
### 5. Python Type Hints (Python 3.10+ Style)
|
||||
|
||||
**CRITICAL**: Always use modern Python 3.10+ union syntax for type hints.
|
||||
|
||||
```python
|
||||
# GOOD - Modern Python 3.10+ syntax
|
||||
from typing import List, Dict, Any
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
class MySchema(BaseModel):
|
||||
name: str | None = Field(None, description="Optional name")
|
||||
tags: List[str] = Field(default_factory=list)
|
||||
metadata: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
def my_function(
|
||||
id: int,
|
||||
filters: List[str] | None = None,
|
||||
options: Dict[str, Any] | None = None
|
||||
) -> MySchema | None:
|
||||
pass
|
||||
|
||||
# BAD - Old-style Optional (DO NOT USE)
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
class MySchema(BaseModel):
|
||||
name: Optional[str] = Field(None, description="Optional name") # Wrong!
|
||||
|
||||
def my_function(
|
||||
id: int,
|
||||
filters: Optional[List[str]] = None, # Wrong!
|
||||
options: Optional[Dict[str, Any]] = None # Wrong!
|
||||
) -> Optional[MySchema]: # Wrong!
|
||||
pass
|
||||
```
|
||||
|
||||
**Key rules:**
|
||||
- Use `T | None` instead of `Optional[T]`
|
||||
- Do NOT import `Optional` from typing
|
||||
- Still import `List`, `Dict`, `Any`, etc. from typing (for now)
|
||||
- All new code must follow this pattern
|
||||
|
||||
### 6. Error Handling
|
||||
|
||||
**Use consistent error schemas**:
|
||||
|
||||
```python
|
||||
class MyError(BaseModel):
|
||||
error: str = Field(..., description="Error message")
|
||||
error_type: str = Field(..., description="Type of error")
|
||||
timestamp: datetime = Field(
|
||||
default_factory=lambda: datetime.now(timezone.utc),
|
||||
description="Error timestamp"
|
||||
)
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
def my_tool(id: int) -> MyResponse:
|
||||
try:
|
||||
result = process_data(id)
|
||||
return MyResponse(data=result)
|
||||
except NotFound:
|
||||
raise ValueError(f"Resource {id} not found")
|
||||
```
|
||||
|
||||
## Testing Conventions
|
||||
|
||||
### Unit Tests
|
||||
|
||||
Place unit tests in `tests/unit_tests/mcp_service/{module}/tool/test_{tool_name}.py`
|
||||
|
||||
**Test structure**:
|
||||
```python
|
||||
from unittest.mock import MagicMock, patch
|
||||
import pytest
|
||||
|
||||
class TestMyTool:
|
||||
@pytest.fixture
|
||||
def mock_dao(self):
|
||||
"""Create mock DAO for testing."""
|
||||
dao = MagicMock()
|
||||
dao.find_by_id.return_value = create_mock_object()
|
||||
return dao
|
||||
|
||||
@patch("superset.mcp_service.chart.tool.my_tool.ChartDAO")
|
||||
def test_my_tool_success(self, mock_dao_class, mock_dao):
|
||||
"""Test successful tool execution."""
|
||||
mock_dao_class.return_value = mock_dao
|
||||
|
||||
result = my_tool(id=1)
|
||||
|
||||
assert result.data is not None
|
||||
mock_dao.find_by_id.assert_called_once_with(1)
|
||||
```
|
||||
|
||||
### Integration Tests
|
||||
|
||||
Use Flask test client for integration tests:
|
||||
|
||||
```python
|
||||
def test_tool_with_flask_context(app):
|
||||
"""Test tool with full Flask app context."""
|
||||
with app.app_context():
|
||||
result = my_tool(id=1)
|
||||
assert result is not None
|
||||
```
|
||||
|
||||
## Common Pitfalls to Avoid
|
||||
|
||||
### 1. ❌ Forgetting Tool Import in app.py
|
||||
**Problem**: Tool exists but isn't available to MCP clients.
|
||||
**Solution**: Always add tool import to `app.py` (at the bottom) after creating it. Never add to `server.py`.
|
||||
|
||||
### 2. ❌ Adding Tool Imports to server.py
|
||||
**Problem**: Tools won't register properly, causing runtime errors.
|
||||
**Solution**: Tool imports must be in `app.py` at the bottom of the file, not in `server.py`. The `server.py` file is only for running the server.
|
||||
|
||||
### 3. ❌ Missing @mcp_auth_hook Decorator
|
||||
**Problem**: Tool bypasses authentication and authorization.
|
||||
**Solution**: Always use `@mcp_auth_hook` on every tool.
|
||||
|
||||
### 4. ❌ Using `Optional` Instead of Union Syntax
|
||||
**Problem**: Old-style Optional[T] is not Python 3.10+ style.
|
||||
**Solution**: Use `T | None` instead of `Optional[T]` for all type hints.
|
||||
```python
|
||||
# GOOD - Modern Python 3.10+ syntax
|
||||
def my_function(param: str | None = None) -> int | None:
|
||||
pass
|
||||
|
||||
# BAD - Old-style Optional
|
||||
from typing import Optional
|
||||
def my_function(param: Optional[str] = None) -> Optional[int]:
|
||||
pass
|
||||
```
|
||||
|
||||
### 5. ❌ Using `any` Types in Schemas
|
||||
**Problem**: Violates TypeScript modernization goals, no validation.
|
||||
**Solution**: Use proper Pydantic types with Field descriptions.
|
||||
|
||||
### 6. ❌ Direct Database Queries
|
||||
**Problem**: Bypasses Superset's security and caching layers.
|
||||
**Solution**: Use DAO classes (ChartDAO, DashboardDAO, etc.).
|
||||
|
||||
### 7. ❌ Not Using Core Classes
|
||||
**Problem**: Duplicating list/get_info/filter logic across tools.
|
||||
**Solution**: Use ModelListCore, ModelGetInfoCore, ModelGetAvailableFiltersCore.
|
||||
|
||||
### 8. ❌ Missing Apache License Headers
|
||||
**Problem**: CI fails on license check.
|
||||
**Solution**: Add Apache license header to all new .py files. Use this exact template at the top of every new Python file:
|
||||
|
||||
```python
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
```
|
||||
|
||||
**Note**: LLM instruction files like `CLAUDE.md`, `AGENTS.md`, etc. are excluded from this requirement (listed in `.rat-excludes`) to avoid token overhead.
|
||||
|
||||
### 9. ❌ Using `@mcp.tool()` with Empty Parentheses
|
||||
**Problem**: Inconsistent decorator style.
|
||||
**Solution**: Use `@mcp.tool` without parentheses unless passing arguments.
|
||||
```python
|
||||
# GOOD
|
||||
@mcp.tool
|
||||
def my_tool():
|
||||
pass
|
||||
|
||||
# BAD
|
||||
@mcp.tool()
|
||||
def my_tool():
|
||||
pass
|
||||
```
|
||||
|
||||
### 10. ❌ Circular Imports
|
||||
**Problem**: Importing `mcp` from `app.py` creates circular dependency.
|
||||
**Solution**: Import `mcp` at module level in tool files:
|
||||
```python
|
||||
# GOOD
|
||||
from superset.mcp_service.app import mcp
|
||||
|
||||
@mcp.tool
|
||||
def my_tool():
|
||||
pass
|
||||
|
||||
# BAD - causes circular import
|
||||
from superset.mcp_service.app import mcp, some_other_function
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
Default configuration is in `mcp_config.py`. Users can override in `superset_config.py`:
|
||||
|
||||
```python
|
||||
# superset_config.py
|
||||
MCP_ADMIN_USERNAME = "your_admin"
|
||||
MCP_AUTH_ENABLED = True
|
||||
MCP_JWT_PUBLIC_KEY = "your_public_key"
|
||||
```
|
||||
|
||||
## Tool Discovery
|
||||
|
||||
MCP clients discover tools via:
|
||||
1. **Tool listing**: All tools with `@mcp.tool` are automatically listed
|
||||
2. **Schema introspection**: Pydantic schemas generate JSON Schema for LLMs
|
||||
3. **Instructions**: `DEFAULT_INSTRUCTIONS` in `app.py` documents available tools
|
||||
|
||||
## Resources for Learning
|
||||
|
||||
- **MCP Specification**: https://modelcontextprotocol.io/
|
||||
- **FastMCP Documentation**: https://github.com/jlowin/fastmcp
|
||||
- **Superset DAO Patterns**: See `superset/daos/` for examples
|
||||
- **Pydantic Documentation**: https://docs.pydantic.dev/
|
||||
|
||||
## Quick Checklist for New Tools
|
||||
|
||||
- [ ] Created tool file in `{module}/tool/{tool_name}.py`
|
||||
- [ ] Added `@mcp.tool` decorator
|
||||
- [ ] Added `@mcp_auth_hook` decorator
|
||||
- [ ] Created Pydantic request/response schemas in `{module}/schemas.py`
|
||||
- [ ] Used DAO classes instead of direct queries
|
||||
- [ ] Added tool import to `app.py` (around line 210-242)
|
||||
- [ ] Added Apache license header to new files
|
||||
- [ ] Created unit tests in `tests/unit_tests/mcp_service/{module}/tool/test_{tool_name}.py`
|
||||
- [ ] Updated `DEFAULT_INSTRUCTIONS` in `app.py` if adding new capability
|
||||
- [ ] Tested locally with MCP client (e.g., Claude Desktop)
|
||||
|
||||
## Quick Checklist for New Prompts
|
||||
|
||||
- [ ] Created prompt file in `{module}/prompts/{prompt_name}.py`
|
||||
- [ ] Added `@mcp.prompt("prompt_name")` decorator
|
||||
- [ ] Added `@mcp_auth_hook` decorator
|
||||
- [ ] Made function async: `async def prompt_handler(ctx: Context) -> str`
|
||||
- [ ] Added import to `{module}/prompts/__init__.py`
|
||||
- [ ] Verified module import exists in `app.py` (around line 244-253)
|
||||
- [ ] Added Apache license header to new file
|
||||
- [ ] Updated `DEFAULT_INSTRUCTIONS` in `app.py` to list the new prompt
|
||||
- [ ] Tested locally with MCP client (e.g., Claude Desktop)
|
||||
|
||||
## Quick Checklist for New Resources
|
||||
|
||||
- [ ] Created resource file in `{module}/resources/{resource_name}.py`
|
||||
- [ ] Added `@mcp.resource("superset://{path}")` decorator with unique URI
|
||||
- [ ] Added `@mcp_auth_hook` decorator
|
||||
- [ ] Implemented resource data retrieval logic
|
||||
- [ ] Added import to `{module}/resources/__init__.py`
|
||||
- [ ] Verified module import exists in `app.py` (around line 244-253)
|
||||
- [ ] Added Apache license header to new file
|
||||
- [ ] Updated `DEFAULT_INSTRUCTIONS` in `app.py` to list the new resource
|
||||
- [ ] Tested locally with MCP client (e.g., Claude Desktop)
|
||||
|
||||
## Getting Help
|
||||
|
||||
- Check existing tool implementations for patterns (chart/tool/, dashboard/tool/)
|
||||
- Review core classes in `mcp_core.py` for reusable functionality
|
||||
- See `CLAUDE.md` in project root for general Superset development guidelines
|
||||
- Consult Superset documentation: https://superset.apache.org/docs/
|
||||
@@ -77,6 +77,40 @@ Available Prompts:
|
||||
- superset_quickstart: Interactive guide for getting started with the MCP service
|
||||
- create_chart_guided: Step-by-step chart creation wizard
|
||||
|
||||
Common Chart Types (viz_type) and Behaviors:
|
||||
|
||||
Interactive Charts (support sorting, filtering, drill-down):
|
||||
- table: Standard table view with sorting and filtering
|
||||
- pivot_table_v2: Pivot table with grouping and aggregations
|
||||
- echarts_timeseries_line: Time series line chart
|
||||
- echarts_timeseries_bar: Time series bar chart
|
||||
- echarts_timeseries_area: Time series area chart
|
||||
- echarts_timeseries_scatter: Time series scatter plot
|
||||
- mixed_timeseries: Combined line/bar time series
|
||||
|
||||
Common Visualization Types:
|
||||
- big_number: Single metric display
|
||||
- big_number_total: Total value display
|
||||
- pie: Pie chart for proportions
|
||||
- echarts_timeseries: Generic time series chart
|
||||
- funnel: Funnel chart for conversion analysis
|
||||
- gauge_chart: Gauge/speedometer visualization
|
||||
- heatmap_v2: Heat map for correlation analysis
|
||||
- sankey_v2: Sankey diagram for flow visualization
|
||||
- sunburst_v2: Sunburst chart for hierarchical data
|
||||
- treemap_v2: Tree map for hierarchical proportions
|
||||
- word_cloud: Word cloud visualization
|
||||
- world_map: Geographic world map
|
||||
- box_plot: Box plot for distribution analysis
|
||||
- bubble: Bubble chart for 3-dimensional data
|
||||
|
||||
Query Examples:
|
||||
- List all interactive tables:
|
||||
filters=[{"col": "viz_type", "opr": "in", "value": ["table", "pivot_table_v2"]}]
|
||||
- List time series charts:
|
||||
filters=[{"col": "viz_type", "opr": "sw", "value": "echarts_timeseries"}]
|
||||
- Search by name: search="sales"
|
||||
|
||||
General usage tips:
|
||||
- All listing tools use 1-based pagination (first page is 1)
|
||||
- Use 'filters' parameter for advanced queries (see *_available_filters tools)
|
||||
@@ -202,13 +236,56 @@ def create_mcp_app(
|
||||
|
||||
# Create default MCP instance for backward compatibility
|
||||
# Tool modules can import this and use @mcp.tool decorators
|
||||
mcp = create_mcp_app()
|
||||
mcp = create_mcp_app(stateless_http=True)
|
||||
|
||||
from superset.mcp_service.chart.tool import ( # noqa: F401, E402
|
||||
get_chart_info,
|
||||
list_charts,
|
||||
# Import all MCP tools to register them with the mcp instance
|
||||
# NOTE: Always add new tool imports here when creating new MCP tools.
|
||||
# Tools use @mcp.tool decorators and register automatically on import.
|
||||
# Import prompts and resources to register them with the mcp instance
|
||||
# NOTE: Always add new prompt/resource imports here when creating new prompts/resources.
|
||||
# Prompts use @mcp.prompt decorators and resources use @mcp.resource decorators.
|
||||
# They register automatically on import, similar to tools.
|
||||
from superset.mcp_service.chart import ( # noqa: F401, E402
|
||||
prompts as chart_prompts,
|
||||
resources as chart_resources,
|
||||
)
|
||||
from superset.mcp_service.chart.tool import ( # noqa: F401, E402
|
||||
generate_chart,
|
||||
get_chart_available_filters,
|
||||
get_chart_data,
|
||||
get_chart_info,
|
||||
get_chart_preview,
|
||||
list_charts,
|
||||
update_chart,
|
||||
update_chart_preview,
|
||||
)
|
||||
from superset.mcp_service.dashboard.tool import ( # noqa: F401, E402
|
||||
add_chart_to_existing_dashboard,
|
||||
generate_dashboard,
|
||||
get_dashboard_available_filters,
|
||||
get_dashboard_info,
|
||||
list_dashboards,
|
||||
)
|
||||
from superset.mcp_service.dataset.tool import ( # noqa: F401, E402
|
||||
get_dataset_available_filters,
|
||||
get_dataset_info,
|
||||
list_datasets,
|
||||
)
|
||||
from superset.mcp_service.explore.tool import ( # noqa: F401, E402
|
||||
generate_explore_link,
|
||||
)
|
||||
from superset.mcp_service.sql_lab.tool import ( # noqa: F401, E402
|
||||
execute_sql,
|
||||
open_sql_lab_with_context,
|
||||
)
|
||||
from superset.mcp_service.system import ( # noqa: F401, E402
|
||||
prompts as system_prompts,
|
||||
resources as system_resources,
|
||||
)
|
||||
from superset.mcp_service.system.tool import ( # noqa: F401, E402
|
||||
get_superset_instance_info,
|
||||
health_check,
|
||||
)
|
||||
from superset.mcp_service.system.tool import health_check # noqa: F401, E402
|
||||
|
||||
|
||||
def init_fastmcp_server(
|
||||
|
||||
@@ -28,11 +28,14 @@ Future enhancements (to be added in separate PRs):
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Callable, TypeVar
|
||||
from typing import Any, Callable, TYPE_CHECKING, TypeVar
|
||||
|
||||
from flask import g
|
||||
from flask_appbuilder.security.sqla.models import User
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from superset.connectors.sqla.models import SqlaTable
|
||||
|
||||
# Type variable for decorated functions
|
||||
F = TypeVar("F", bound=Callable[..., Any])
|
||||
|
||||
@@ -45,15 +48,14 @@ def get_user_from_request() -> User:
|
||||
|
||||
TODO (future PR): Add JWT token extraction and validation.
|
||||
TODO (future PR): Add user impersonation support.
|
||||
TODO (future PR): Add fallback user configuration.
|
||||
|
||||
For now, this uses MCP_DEV_USERNAME from configuration for development.
|
||||
|
||||
Raises:
|
||||
ValueError: If MCP_DEV_USERNAME is not configured or user doesn't exist
|
||||
For now, this returns the admin user for development.
|
||||
"""
|
||||
from flask import current_app
|
||||
from sqlalchemy.orm import joinedload
|
||||
|
||||
from superset import security_manager
|
||||
from superset.extensions import db
|
||||
|
||||
# TODO: Extract from JWT token once authentication is implemented
|
||||
# For now, use MCP_DEV_USERNAME from configuration
|
||||
@@ -62,14 +64,57 @@ def get_user_from_request() -> User:
|
||||
if not username:
|
||||
raise ValueError("Username not configured")
|
||||
|
||||
user = security_manager.find_user(username)
|
||||
# Query user directly with eager loading to ensure fresh session-bound object
|
||||
# Do NOT use security_manager.find_user() as it may return cached/detached user
|
||||
user = (
|
||||
db.session.query(User)
|
||||
.options(joinedload(User.roles), joinedload(User.groups))
|
||||
.filter(User.username == username)
|
||||
.first()
|
||||
)
|
||||
|
||||
if not user:
|
||||
raise ValueError(f"User '{username}' not found")
|
||||
raise ValueError(
|
||||
f"User '{username}' not found. "
|
||||
f"Please create admin user with: superset fab create-admin"
|
||||
)
|
||||
|
||||
return user
|
||||
|
||||
|
||||
def has_dataset_access(dataset: "SqlaTable") -> bool:
|
||||
"""
|
||||
Validate user has access to the dataset.
|
||||
|
||||
This function checks if the current user (from Flask g.user context)
|
||||
has permission to access the given dataset using Superset's security manager.
|
||||
|
||||
Args:
|
||||
dataset: The SqlaTable dataset to check access for
|
||||
|
||||
Returns:
|
||||
True if user has access, False otherwise
|
||||
|
||||
Security Note:
|
||||
This should be called after mcp_auth_hook has set g.user.
|
||||
Returns False on any error to fail securely.
|
||||
"""
|
||||
try:
|
||||
from superset import security_manager
|
||||
|
||||
# Check if user has read access to the dataset
|
||||
if hasattr(g, "user") and g.user:
|
||||
# Use Superset's security manager to check dataset access
|
||||
return security_manager.can_access_datasource(datasource=dataset)
|
||||
|
||||
# If no user context, deny access
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("Error checking dataset access: %s", e)
|
||||
return False # Deny access on error
|
||||
|
||||
|
||||
def mcp_auth_hook(tool_func: F) -> F:
|
||||
"""
|
||||
Authentication and authorization decorator for MCP tools.
|
||||
@@ -87,17 +132,48 @@ def mcp_auth_hook(tool_func: F) -> F:
|
||||
|
||||
@functools.wraps(tool_func)
|
||||
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
# Get user and set Flask context
|
||||
from superset.extensions import db
|
||||
|
||||
# Get user and set Flask context OUTSIDE try block
|
||||
user = get_user_from_request()
|
||||
|
||||
# Force load relationships NOW while session is definitely active
|
||||
_ = user.roles
|
||||
if hasattr(user, "groups"):
|
||||
_ = user.groups
|
||||
|
||||
g.user = user
|
||||
|
||||
# TODO: Add permission checks here in future PR
|
||||
# TODO: Add audit logging here in future PR
|
||||
try:
|
||||
# TODO: Add permission checks here in future PR
|
||||
# TODO: Add audit logging here in future PR
|
||||
|
||||
logger.debug(
|
||||
"MCP tool call: user=%s, tool=%s", user.username, tool_func.__name__
|
||||
)
|
||||
logger.debug(
|
||||
"MCP tool call: user=%s, tool=%s", user.username, tool_func.__name__
|
||||
)
|
||||
|
||||
return tool_func(*args, **kwargs)
|
||||
result = tool_func(*args, **kwargs)
|
||||
|
||||
return result
|
||||
|
||||
except Exception:
|
||||
# On error, rollback and cleanup session
|
||||
# pylint: disable=consider-using-transaction
|
||||
try:
|
||||
db.session.rollback()
|
||||
db.session.remove()
|
||||
except Exception as e:
|
||||
logger.warning("Error cleaning up session after exception: %s", e)
|
||||
raise
|
||||
|
||||
finally:
|
||||
# Only rollback if session is still active (no exception occurred)
|
||||
# Do NOT call remove() on success to avoid detaching user
|
||||
try:
|
||||
if db.session.is_active:
|
||||
# pylint: disable=consider-using-transaction
|
||||
db.session.rollback()
|
||||
except Exception as e:
|
||||
logger.warning("Error in finally block: %s", e)
|
||||
|
||||
return wrapper # type: ignore[return-value]
|
||||
|
||||
484
superset/mcp_service/chart/chart_utils.py
Normal file
484
superset/mcp_service/chart/chart_utils.py
Normal file
@@ -0,0 +1,484 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Shared chart utilities for MCP tools
|
||||
|
||||
This module contains shared logic for chart configuration mapping and explore link
|
||||
generation that can be used by both generate_chart and generate_explore_link tools.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
ChartCapabilities,
|
||||
ChartSemantics,
|
||||
ColumnRef,
|
||||
TableChartConfig,
|
||||
XYChartConfig,
|
||||
)
|
||||
from superset.mcp_service.utils.url_utils import get_superset_base_url
|
||||
from superset.utils import json
|
||||
|
||||
|
||||
def generate_explore_link(dataset_id: int | str, form_data: Dict[str, Any]) -> str:
|
||||
"""Generate an explore link for the given dataset and form data."""
|
||||
base_url = get_superset_base_url()
|
||||
numeric_dataset_id = None
|
||||
|
||||
try:
|
||||
from superset.commands.explore.form_data.parameters import CommandParameters
|
||||
|
||||
# Find the dataset to get its numeric ID
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
from superset.mcp_service.commands.create_form_data import (
|
||||
MCPCreateFormDataCommand,
|
||||
)
|
||||
from superset.utils.core import DatasourceType
|
||||
|
||||
dataset = None
|
||||
|
||||
if isinstance(dataset_id, int) or (
|
||||
isinstance(dataset_id, str) and dataset_id.isdigit()
|
||||
):
|
||||
numeric_dataset_id = (
|
||||
int(dataset_id) if isinstance(dataset_id, str) else dataset_id
|
||||
)
|
||||
dataset = DatasetDAO.find_by_id(numeric_dataset_id)
|
||||
else:
|
||||
# Try UUID lookup using DAO flexible method
|
||||
dataset = DatasetDAO.find_by_id(dataset_id, id_column="uuid")
|
||||
if dataset:
|
||||
numeric_dataset_id = dataset.id
|
||||
|
||||
if not dataset or numeric_dataset_id is None:
|
||||
# Fallback to basic explore URL
|
||||
return (
|
||||
f"{base_url}/explore/?datasource_type=table&datasource_id={dataset_id}"
|
||||
)
|
||||
|
||||
# Add datasource to form_data
|
||||
form_data_with_datasource = {
|
||||
**form_data,
|
||||
"datasource": f"{numeric_dataset_id}__table",
|
||||
}
|
||||
|
||||
# Try to create form_data in cache using MCP-specific CreateFormDataCommand
|
||||
cmd_params = CommandParameters(
|
||||
datasource_type=DatasourceType.TABLE,
|
||||
datasource_id=numeric_dataset_id,
|
||||
chart_id=0, # 0 for new charts
|
||||
tab_id=None,
|
||||
form_data=json.dumps(form_data_with_datasource),
|
||||
)
|
||||
|
||||
# Create the form_data cache entry and get the key
|
||||
form_data_key = MCPCreateFormDataCommand(cmd_params).run()
|
||||
|
||||
# Return URL with just the form_data_key
|
||||
return f"{base_url}/explore/?form_data_key={form_data_key}"
|
||||
|
||||
except Exception:
|
||||
# Fallback to basic explore URL with numeric ID if available
|
||||
if numeric_dataset_id is not None:
|
||||
return (
|
||||
f"{base_url}/explore/?datasource_type=table"
|
||||
f"&datasource_id={numeric_dataset_id}"
|
||||
)
|
||||
else:
|
||||
return (
|
||||
f"{base_url}/explore/?datasource_type=table&datasource_id={dataset_id}"
|
||||
)
|
||||
|
||||
|
||||
def map_config_to_form_data(
|
||||
config: TableChartConfig | XYChartConfig,
|
||||
) -> Dict[str, Any]:
|
||||
"""Map chart config to Superset form_data."""
|
||||
if isinstance(config, TableChartConfig):
|
||||
return map_table_config(config)
|
||||
elif isinstance(config, XYChartConfig):
|
||||
return map_xy_config(config)
|
||||
else:
|
||||
raise ValueError(f"Unsupported config type: {type(config)}")
|
||||
|
||||
|
||||
def map_table_config(config: TableChartConfig) -> Dict[str, Any]:
|
||||
"""Map table chart config to form_data with defensive validation."""
|
||||
# Early validation to prevent empty charts
|
||||
if not config.columns:
|
||||
raise ValueError("Table chart must have at least one column")
|
||||
|
||||
# Separate columns with aggregates from raw columns
|
||||
raw_columns = []
|
||||
aggregated_metrics = []
|
||||
|
||||
for col in config.columns:
|
||||
if col.aggregate:
|
||||
# Column has aggregation - treat as metric
|
||||
aggregated_metrics.append(create_metric_object(col))
|
||||
else:
|
||||
# No aggregation - treat as raw column
|
||||
raw_columns.append(col.name)
|
||||
|
||||
# Final validation - ensure we have some data to display
|
||||
if not raw_columns and not aggregated_metrics:
|
||||
raise ValueError("Table chart configuration resulted in no displayable columns")
|
||||
|
||||
form_data: Dict[str, Any] = {
|
||||
"viz_type": "table",
|
||||
}
|
||||
|
||||
# Handle raw columns (no aggregation)
|
||||
if raw_columns and not aggregated_metrics:
|
||||
# Pure raw columns - show individual rows
|
||||
form_data.update(
|
||||
{
|
||||
"all_columns": raw_columns,
|
||||
"query_mode": "raw",
|
||||
"include_time": False,
|
||||
"order_desc": True,
|
||||
"row_limit": 1000, # Reasonable limit for raw data
|
||||
}
|
||||
)
|
||||
|
||||
# Handle aggregated columns only
|
||||
elif aggregated_metrics and not raw_columns:
|
||||
# Pure aggregation - show totals
|
||||
form_data.update(
|
||||
{
|
||||
"metrics": aggregated_metrics,
|
||||
"query_mode": "aggregate",
|
||||
}
|
||||
)
|
||||
|
||||
# Handle mixed columns (raw + aggregated)
|
||||
elif raw_columns and aggregated_metrics:
|
||||
# Mixed mode - group by raw columns, aggregate metrics
|
||||
form_data.update(
|
||||
{
|
||||
"all_columns": raw_columns,
|
||||
"metrics": aggregated_metrics,
|
||||
"groupby": raw_columns,
|
||||
"query_mode": "aggregate",
|
||||
}
|
||||
)
|
||||
|
||||
if config.filters:
|
||||
form_data["adhoc_filters"] = [
|
||||
{
|
||||
"clause": "WHERE",
|
||||
"expressionType": "SIMPLE",
|
||||
"subject": filter_config.column,
|
||||
"operator": map_filter_operator(filter_config.op),
|
||||
"comparator": filter_config.value,
|
||||
}
|
||||
for filter_config in config.filters
|
||||
if filter_config is not None
|
||||
]
|
||||
|
||||
if config.sort_by:
|
||||
form_data["order_by_cols"] = config.sort_by
|
||||
|
||||
return form_data
|
||||
|
||||
|
||||
def create_metric_object(col: ColumnRef) -> Dict[str, Any]:
|
||||
"""Create a metric object for a column with enhanced validation."""
|
||||
# Ensure aggregate is valid - default to SUM if not specified or invalid
|
||||
valid_aggregates = {
|
||||
"SUM",
|
||||
"COUNT",
|
||||
"AVG",
|
||||
"MIN",
|
||||
"MAX",
|
||||
"COUNT_DISTINCT",
|
||||
"STDDEV",
|
||||
"VAR",
|
||||
"MEDIAN",
|
||||
"PERCENTILE",
|
||||
}
|
||||
aggregate = col.aggregate or "SUM"
|
||||
|
||||
# Validate aggregate function (final safety check)
|
||||
if aggregate.upper() not in valid_aggregates:
|
||||
aggregate = "SUM" # Safe fallback
|
||||
|
||||
return {
|
||||
"aggregate": aggregate.upper(),
|
||||
"column": {
|
||||
"column_name": col.name,
|
||||
},
|
||||
"expressionType": "SIMPLE",
|
||||
"label": col.label or f"{aggregate.upper()}({col.name})",
|
||||
"optionName": f"metric_{col.name}",
|
||||
"sqlExpression": None,
|
||||
"hasCustomLabel": bool(col.label),
|
||||
"datasourceWarning": False,
|
||||
}
|
||||
|
||||
|
||||
def add_axis_config(form_data: Dict[str, Any], config: XYChartConfig) -> None:
|
||||
"""Add axis configurations to form_data."""
|
||||
if config.x_axis:
|
||||
if config.x_axis.title:
|
||||
form_data["x_axis_title"] = config.x_axis.title
|
||||
if config.x_axis.format:
|
||||
form_data["x_axis_format"] = config.x_axis.format
|
||||
|
||||
if config.y_axis:
|
||||
if config.y_axis.title:
|
||||
form_data["y_axis_title"] = config.y_axis.title
|
||||
if config.y_axis.format:
|
||||
form_data["y_axis_format"] = config.y_axis.format
|
||||
if config.y_axis.scale == "log":
|
||||
form_data["y_axis_scale"] = "log"
|
||||
|
||||
|
||||
def add_legend_config(form_data: Dict[str, Any], config: XYChartConfig) -> None:
|
||||
"""Add legend configuration to form_data."""
|
||||
if config.legend:
|
||||
if not config.legend.show:
|
||||
form_data["show_legend"] = False
|
||||
if config.legend.position:
|
||||
form_data["legend_orientation"] = config.legend.position
|
||||
|
||||
|
||||
def map_xy_config(config: XYChartConfig) -> Dict[str, Any]:
|
||||
"""Map XY chart config to form_data with defensive validation."""
|
||||
# Early validation to prevent empty charts
|
||||
if not config.y:
|
||||
raise ValueError("XY chart must have at least one Y-axis metric")
|
||||
|
||||
# Map chart kind to viz_type
|
||||
viz_type_map = {
|
||||
"line": "echarts_timeseries_line",
|
||||
"bar": "echarts_timeseries_bar",
|
||||
"area": "echarts_area",
|
||||
"scatter": "echarts_timeseries_scatter",
|
||||
}
|
||||
|
||||
# Convert Y columns to metrics with validation
|
||||
metrics = []
|
||||
for col in config.y:
|
||||
if not col.name.strip(): # Validate column name is not empty
|
||||
raise ValueError("Y-axis column name cannot be empty")
|
||||
metrics.append(create_metric_object(col))
|
||||
|
||||
# Final validation - ensure we have metrics to display
|
||||
if not metrics:
|
||||
raise ValueError("XY chart configuration resulted in no displayable metrics")
|
||||
|
||||
form_data: Dict[str, Any] = {
|
||||
"viz_type": viz_type_map.get(config.kind, "echarts_timeseries_line"),
|
||||
"x_axis": config.x.name,
|
||||
"metrics": metrics,
|
||||
}
|
||||
|
||||
# CRITICAL FIX: For time series charts, handle groupby carefully to avoid duplicates
|
||||
# The x_axis field already tells Superset which column to use for time grouping
|
||||
groupby_columns = []
|
||||
|
||||
# Only add groupby columns if there's an explicit group_by specified
|
||||
# The x_axis column should NOT be duplicated in groupby as it causes
|
||||
# "Duplicate column/metric labels" errors in Superset
|
||||
# Only add group_by column if it's specified AND different from x_axis
|
||||
# NEVER add the x_axis column to groupby as it creates duplicate labels
|
||||
if config.group_by and config.group_by.name != config.x.name:
|
||||
groupby_columns.append(config.group_by.name)
|
||||
|
||||
# Set the groupby in form_data only if we have valid columns
|
||||
# Don't set empty groupby - let Superset handle x_axis grouping automatically
|
||||
if groupby_columns:
|
||||
form_data["groupby"] = groupby_columns
|
||||
|
||||
# Add filters if specified
|
||||
if config.filters:
|
||||
form_data["adhoc_filters"] = [
|
||||
{
|
||||
"clause": "WHERE",
|
||||
"expressionType": "SIMPLE",
|
||||
"subject": filter_config.column,
|
||||
"operator": map_filter_operator(filter_config.op),
|
||||
"comparator": filter_config.value,
|
||||
}
|
||||
for filter_config in config.filters
|
||||
if filter_config is not None
|
||||
]
|
||||
|
||||
# Add configurations
|
||||
add_axis_config(form_data, config)
|
||||
add_legend_config(form_data, config)
|
||||
|
||||
return form_data
|
||||
|
||||
|
||||
def map_filter_operator(op: str) -> str:
|
||||
"""Map filter operator to Superset format."""
|
||||
operator_map = {
|
||||
"=": "==",
|
||||
">": ">",
|
||||
"<": "<",
|
||||
">=": ">=",
|
||||
"<=": "<=",
|
||||
"!=": "!=",
|
||||
}
|
||||
return operator_map.get(op, op)
|
||||
|
||||
|
||||
def generate_chart_name(config: TableChartConfig | XYChartConfig) -> str:
|
||||
"""Generate a chart name based on the configuration."""
|
||||
if isinstance(config, TableChartConfig):
|
||||
return f"Table Chart - {', '.join(col.name for col in config.columns)}"
|
||||
elif isinstance(config, XYChartConfig):
|
||||
chart_type = config.kind.capitalize()
|
||||
x_col = config.x.name
|
||||
y_cols = ", ".join(col.name for col in config.y)
|
||||
return f"{chart_type} Chart - {x_col} vs {y_cols}"
|
||||
else:
|
||||
return "Chart"
|
||||
|
||||
|
||||
def analyze_chart_capabilities(chart: Any | None, config: Any) -> ChartCapabilities:
|
||||
"""Analyze chart capabilities based on type and configuration."""
|
||||
if chart:
|
||||
viz_type = getattr(chart, "viz_type", "unknown")
|
||||
else:
|
||||
# Map config chart_type to viz_type
|
||||
chart_type = getattr(config, "chart_type", "unknown")
|
||||
if chart_type == "xy":
|
||||
kind = getattr(config, "kind", "line")
|
||||
viz_type_map = {
|
||||
"line": "echarts_timeseries_line",
|
||||
"bar": "echarts_timeseries_bar",
|
||||
"area": "echarts_area",
|
||||
"scatter": "echarts_timeseries_scatter",
|
||||
}
|
||||
viz_type = viz_type_map.get(kind, "echarts_timeseries_line")
|
||||
elif chart_type == "table":
|
||||
viz_type = "table"
|
||||
else:
|
||||
viz_type = "unknown"
|
||||
|
||||
# Determine interaction capabilities based on chart type
|
||||
interactive_types = [
|
||||
"echarts_timeseries_line",
|
||||
"echarts_timeseries_bar",
|
||||
"echarts_area",
|
||||
"echarts_timeseries_scatter",
|
||||
"deck_scatter",
|
||||
"deck_hex",
|
||||
]
|
||||
|
||||
supports_interaction = viz_type in interactive_types
|
||||
supports_drill_down = viz_type in ["table", "pivot_table_v2"]
|
||||
supports_real_time = viz_type in [
|
||||
"echarts_timeseries_line",
|
||||
"echarts_timeseries_bar",
|
||||
]
|
||||
|
||||
# Determine optimal formats
|
||||
optimal_formats = ["url"] # Always include static image
|
||||
if supports_interaction:
|
||||
optimal_formats.extend(["interactive", "vega_lite"])
|
||||
optimal_formats.extend(["ascii", "table"])
|
||||
|
||||
# Classify data types
|
||||
data_types = []
|
||||
if hasattr(config, "x") and config.x:
|
||||
data_types.append("categorical" if not config.x.aggregate else "metric")
|
||||
if hasattr(config, "y") and config.y:
|
||||
data_types.extend(["metric"] * len(config.y))
|
||||
if "time" in viz_type or "timeseries" in viz_type:
|
||||
data_types.append("time_series")
|
||||
|
||||
return ChartCapabilities(
|
||||
supports_interaction=supports_interaction,
|
||||
supports_real_time=supports_real_time,
|
||||
supports_drill_down=supports_drill_down,
|
||||
supports_export=True, # All charts can be exported
|
||||
optimal_formats=optimal_formats,
|
||||
data_types=list(set(data_types)),
|
||||
)
|
||||
|
||||
|
||||
def analyze_chart_semantics(chart: Any | None, config: Any) -> ChartSemantics:
|
||||
"""Generate semantic understanding of the chart."""
|
||||
if chart:
|
||||
viz_type = getattr(chart, "viz_type", "unknown")
|
||||
else:
|
||||
# Map config chart_type to viz_type
|
||||
chart_type = getattr(config, "chart_type", "unknown")
|
||||
if chart_type == "xy":
|
||||
kind = getattr(config, "kind", "line")
|
||||
viz_type_map = {
|
||||
"line": "echarts_timeseries_line",
|
||||
"bar": "echarts_timeseries_bar",
|
||||
"area": "echarts_area",
|
||||
"scatter": "echarts_timeseries_scatter",
|
||||
}
|
||||
viz_type = viz_type_map.get(kind, "echarts_timeseries_line")
|
||||
elif chart_type == "table":
|
||||
viz_type = "table"
|
||||
else:
|
||||
viz_type = "unknown"
|
||||
|
||||
# Generate primary insight based on chart type
|
||||
insights_map = {
|
||||
"echarts_timeseries_line": "Shows trends and changes over time",
|
||||
"echarts_timeseries_bar": "Compares values across categories or time periods",
|
||||
"table": "Displays detailed data in tabular format",
|
||||
"pie": "Shows proportional relationships within a dataset",
|
||||
"echarts_area": "Emphasizes cumulative totals and part-to-whole relationships",
|
||||
}
|
||||
|
||||
primary_insight = insights_map.get(
|
||||
viz_type, f"Visualizes data using {viz_type} format"
|
||||
)
|
||||
|
||||
# Generate data story
|
||||
columns = []
|
||||
if hasattr(config, "x") and config.x:
|
||||
columns.append(config.x.name)
|
||||
if hasattr(config, "y") and config.y:
|
||||
columns.extend([col.name for col in config.y])
|
||||
|
||||
if columns:
|
||||
ellipsis = "..." if len(columns) > 3 else ""
|
||||
data_story = (
|
||||
f"This {viz_type} chart analyzes {', '.join(columns[:3])}{ellipsis}"
|
||||
)
|
||||
else:
|
||||
data_story = "This chart provides insights into the selected dataset"
|
||||
|
||||
# Generate recommended actions
|
||||
recommended_actions = [
|
||||
"Review data patterns and trends",
|
||||
"Consider filtering or drilling down for more detail",
|
||||
"Export chart for reporting or sharing",
|
||||
]
|
||||
|
||||
if viz_type in ["echarts_timeseries_line", "echarts_timeseries_bar"]:
|
||||
recommended_actions.append("Analyze seasonal patterns or cyclical trends")
|
||||
|
||||
return ChartSemantics(
|
||||
primary_insight=primary_insight,
|
||||
data_story=data_story,
|
||||
recommended_actions=recommended_actions,
|
||||
anomalies=[], # Would need actual data analysis to populate
|
||||
statistical_summary={}, # Would need actual data analysis to populate
|
||||
)
|
||||
561
superset/mcp_service/chart/preview_utils.py
Normal file
561
superset/mcp_service/chart/preview_utils.py
Normal file
@@ -0,0 +1,561 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Preview utilities for chart generation without saving.
|
||||
|
||||
This module provides utilities for generating chart previews
|
||||
from form data without requiring a saved chart object.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from superset.commands.chart.data.get_data_command import ChartDataCommand
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
ASCIIPreview,
|
||||
ChartError,
|
||||
TablePreview,
|
||||
VegaLitePreview,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def generate_preview_from_form_data(
|
||||
form_data: Dict[str, Any], dataset_id: int, preview_format: str
|
||||
) -> Any:
|
||||
"""
|
||||
Generate preview from form data without a saved chart.
|
||||
|
||||
Args:
|
||||
form_data: Chart configuration form data
|
||||
dataset_id: Dataset ID
|
||||
preview_format: Preview format (ascii, table, etc.)
|
||||
|
||||
Returns:
|
||||
Preview object or ChartError
|
||||
"""
|
||||
try:
|
||||
# Execute query to get data
|
||||
from superset.connectors.sqla.models import SqlaTable
|
||||
from superset.extensions import db
|
||||
|
||||
dataset = db.session.query(SqlaTable).get(dataset_id)
|
||||
if not dataset:
|
||||
return ChartError(
|
||||
error=f"Dataset {dataset_id} not found", error_type="DatasetNotFound"
|
||||
)
|
||||
|
||||
# Create query context from form data using factory
|
||||
from superset.common.query_context_factory import QueryContextFactory
|
||||
|
||||
factory = QueryContextFactory()
|
||||
query_context_obj = factory.create(
|
||||
datasource={"id": dataset_id, "type": "table"},
|
||||
queries=[
|
||||
{
|
||||
"columns": form_data.get("columns", []),
|
||||
"metrics": form_data.get("metrics", []),
|
||||
"orderby": form_data.get("orderby", []),
|
||||
"row_limit": form_data.get("row_limit", 100),
|
||||
"filters": form_data.get("adhoc_filters", []),
|
||||
"time_range": form_data.get("time_range", "No filter"),
|
||||
}
|
||||
],
|
||||
form_data=form_data,
|
||||
)
|
||||
|
||||
# Execute query
|
||||
command = ChartDataCommand(query_context_obj)
|
||||
result = command.run()
|
||||
|
||||
if not result or not result.get("queries"):
|
||||
return ChartError(
|
||||
error="No data returned from query", error_type="EmptyResult"
|
||||
)
|
||||
|
||||
query_result = result["queries"][0]
|
||||
data = query_result.get("data", [])
|
||||
|
||||
# Generate preview based on format
|
||||
if preview_format == "ascii":
|
||||
return _generate_ascii_preview_from_data(data, form_data)
|
||||
elif preview_format == "table":
|
||||
return _generate_table_preview_from_data(data, form_data)
|
||||
elif preview_format == "vega_lite":
|
||||
return _generate_vega_lite_preview_from_data(data, form_data)
|
||||
else:
|
||||
return ChartError(
|
||||
error=f"Unsupported preview format: {preview_format}",
|
||||
error_type="UnsupportedFormat",
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Preview generation from form data failed: %s", e)
|
||||
return ChartError(
|
||||
error=f"Failed to generate preview: {str(e)}", error_type="PreviewError"
|
||||
)
|
||||
|
||||
|
||||
def _generate_ascii_preview_from_data(
|
||||
data: List[Dict[str, Any]], form_data: Dict[str, Any]
|
||||
) -> ASCIIPreview:
|
||||
"""Generate ASCII preview from raw data."""
|
||||
viz_type = form_data.get("viz_type", "table")
|
||||
|
||||
# Handle different chart types
|
||||
if viz_type in ["bar", "dist_bar", "column"]:
|
||||
content = _generate_safe_ascii_bar_chart(data)
|
||||
elif viz_type in ["line", "area"]:
|
||||
content = _generate_safe_ascii_line_chart(data)
|
||||
elif viz_type == "pie":
|
||||
content = _generate_safe_ascii_pie_chart(data)
|
||||
else:
|
||||
content = _generate_safe_ascii_table(data)
|
||||
|
||||
return ASCIIPreview(
|
||||
ascii_content=content, width=80, height=20, supports_color=False
|
||||
)
|
||||
|
||||
|
||||
def _calculate_column_widths(
|
||||
display_columns: List[str], data: List[Dict[str, Any]]
|
||||
) -> Dict[str, int]:
|
||||
"""Calculate optimal width for each column."""
|
||||
column_widths = {}
|
||||
for col in display_columns:
|
||||
# Start with column name length
|
||||
max_width = len(str(col))
|
||||
|
||||
# Check data values to determine width
|
||||
for row in data[:20]: # Sample first 20 rows
|
||||
val = row.get(col, "")
|
||||
if isinstance(val, float):
|
||||
val_str = f"{val:.2f}"
|
||||
elif isinstance(val, int):
|
||||
val_str = str(val)
|
||||
else:
|
||||
val_str = str(val)
|
||||
max_width = max(max_width, len(val_str))
|
||||
|
||||
# Set reasonable bounds
|
||||
column_widths[col] = min(max(max_width, 8), 25)
|
||||
return column_widths
|
||||
|
||||
|
||||
def _format_value(val: Any, width: int) -> str:
|
||||
"""Format a value based on its type."""
|
||||
if isinstance(val, float):
|
||||
if abs(val) >= 1000000:
|
||||
val_str = f"{val:.2e}" # Scientific notation for large numbers
|
||||
elif abs(val) >= 1000:
|
||||
val_str = f"{val:,.2f}" # Thousands separator
|
||||
else:
|
||||
val_str = f"{val:.2f}"
|
||||
elif isinstance(val, int):
|
||||
if abs(val) >= 1000:
|
||||
val_str = f"{val:,}" # Thousands separator
|
||||
else:
|
||||
val_str = str(val)
|
||||
elif val is None:
|
||||
val_str = "NULL"
|
||||
else:
|
||||
val_str = str(val)
|
||||
|
||||
# Truncate if too long
|
||||
if len(val_str) > width:
|
||||
val_str = val_str[: width - 2] + ".."
|
||||
return val_str
|
||||
|
||||
|
||||
def _generate_table_preview_from_data(
|
||||
data: List[Dict[str, Any]], form_data: Dict[str, Any]
|
||||
) -> TablePreview:
|
||||
"""Generate table preview from raw data with improved formatting."""
|
||||
if not data:
|
||||
return TablePreview(
|
||||
table_data="No data available", row_count=0, supports_sorting=False
|
||||
)
|
||||
|
||||
# Get columns
|
||||
columns = list(data[0].keys()) if data else []
|
||||
|
||||
# Determine optimal column widths and how many columns to show
|
||||
max_columns = 8 # Show more columns than before
|
||||
display_columns = columns[:max_columns]
|
||||
|
||||
# Calculate optimal width for each column
|
||||
column_widths = _calculate_column_widths(display_columns, data)
|
||||
|
||||
# Format table with proper alignment
|
||||
lines = ["Table Preview", "=" * 80]
|
||||
|
||||
# Header with dynamic width
|
||||
header_parts = []
|
||||
separator_parts = []
|
||||
for col in display_columns:
|
||||
width = column_widths[col]
|
||||
col_name = str(col)
|
||||
if len(col_name) > width:
|
||||
col_name = col_name[: width - 2] + ".."
|
||||
header_parts.append(f"{col_name:<{width}}")
|
||||
separator_parts.append("-" * width)
|
||||
|
||||
lines.append(" | ".join(header_parts))
|
||||
lines.append("-+-".join(separator_parts))
|
||||
|
||||
# Data rows with proper formatting
|
||||
rows_shown = min(len(data), 15) # Show more rows
|
||||
for row in data[:rows_shown]:
|
||||
row_parts = []
|
||||
for col in display_columns:
|
||||
width = column_widths[col]
|
||||
val = row.get(col, "")
|
||||
val_str = _format_value(val, width)
|
||||
row_parts.append(f"{val_str:<{width}}")
|
||||
lines.append(" | ".join(row_parts))
|
||||
|
||||
# Summary information
|
||||
if len(data) > rows_shown:
|
||||
lines.append(f"... and {len(data) - rows_shown} more rows")
|
||||
|
||||
if len(columns) > max_columns:
|
||||
lines.append(f"... and {len(columns) - max_columns} more columns")
|
||||
|
||||
lines.append("")
|
||||
lines.append(f"Total: {len(data)} rows × {len(columns)} columns")
|
||||
|
||||
return TablePreview(
|
||||
table_data="\n".join(lines), row_count=len(data), supports_sorting=True
|
||||
)
|
||||
|
||||
|
||||
def _generate_safe_ascii_bar_chart(data: List[Dict[str, Any]]) -> str:
|
||||
"""Generate ASCII bar chart with proper error handling."""
|
||||
if not data:
|
||||
return "No data available for bar chart"
|
||||
|
||||
lines = ["ASCII Bar Chart", "=" * 50]
|
||||
|
||||
# Extract values safely
|
||||
values = []
|
||||
labels = []
|
||||
|
||||
for row in data[:10]:
|
||||
label = None
|
||||
value = None
|
||||
|
||||
for _, val in row.items():
|
||||
if isinstance(val, (int, float)) and not _is_nan(val) and value is None:
|
||||
value = val
|
||||
elif isinstance(val, str) and label is None:
|
||||
label = val
|
||||
|
||||
if value is not None:
|
||||
values.append(value)
|
||||
labels.append(label or f"Item {len(values)}")
|
||||
|
||||
if not values:
|
||||
return "No numeric data found for bar chart"
|
||||
|
||||
# Generate bars
|
||||
max_val = max(values)
|
||||
if max_val == 0:
|
||||
return "All values are zero"
|
||||
|
||||
for label, value in zip(labels, values, strict=False):
|
||||
bar_length = int((value / max_val) * 30)
|
||||
bar = "█" * bar_length
|
||||
lines.append(f"{label[:10]:>10} |{bar:<30} {value:.2f}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _generate_safe_ascii_line_chart(data: List[Dict[str, Any]]) -> str:
|
||||
"""Generate ASCII line chart with proper NaN handling."""
|
||||
if not data:
|
||||
return "No data available for line chart"
|
||||
|
||||
lines = ["ASCII Line Chart", "=" * 50]
|
||||
values = _extract_numeric_values_safe(data)
|
||||
|
||||
if not values:
|
||||
return "No valid numeric data found for line chart"
|
||||
|
||||
range_str = _format_range_display(values)
|
||||
lines.append(range_str)
|
||||
|
||||
sparkline = _generate_sparkline_safe(values)
|
||||
lines.append(sparkline)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _extract_numeric_values_safe(data: List[Dict[str, Any]]) -> List[float]:
|
||||
"""Extract numeric values safely from data."""
|
||||
values = []
|
||||
for row in data[:20]:
|
||||
for _, val in row.items():
|
||||
if isinstance(val, (int, float)) and not _is_nan(val):
|
||||
values.append(val)
|
||||
break
|
||||
return values
|
||||
|
||||
|
||||
def _format_range_display(values: List[float]) -> str:
|
||||
"""Format range display safely."""
|
||||
min_val = min(values)
|
||||
max_val = max(values)
|
||||
|
||||
if _is_nan(min_val) or _is_nan(max_val):
|
||||
return "Range: Unable to calculate"
|
||||
else:
|
||||
return f"Range: {min_val:.2f} to {max_val:.2f}"
|
||||
|
||||
|
||||
def _generate_sparkline_safe(values: List[float]) -> str:
|
||||
"""Generate sparkline from values."""
|
||||
if not values:
|
||||
return ""
|
||||
|
||||
min_val = min(values)
|
||||
|
||||
if (max_val := max(values)) != min_val:
|
||||
sparkline = ""
|
||||
for val in values:
|
||||
normalized = (val - min_val) / (max_val - min_val)
|
||||
if normalized < 0.2:
|
||||
sparkline += "▁"
|
||||
elif normalized < 0.4:
|
||||
sparkline += "▂"
|
||||
elif normalized < 0.6:
|
||||
sparkline += "▄"
|
||||
elif normalized < 0.8:
|
||||
sparkline += "▆"
|
||||
else:
|
||||
sparkline += "█"
|
||||
return sparkline
|
||||
else:
|
||||
return "─" * len(values) # Flat line if all values are same
|
||||
|
||||
|
||||
def _generate_safe_ascii_pie_chart(data: List[Dict[str, Any]]) -> str:
|
||||
"""Generate ASCII pie chart representation."""
|
||||
if not data:
|
||||
return "No data available for pie chart"
|
||||
|
||||
lines = ["ASCII Pie Chart", "=" * 50]
|
||||
|
||||
# Extract values and labels
|
||||
values = []
|
||||
labels = []
|
||||
|
||||
for row in data[:8]: # Limit slices
|
||||
label = None
|
||||
value = None
|
||||
|
||||
for _, val in row.items():
|
||||
if isinstance(val, (int, float)) and not _is_nan(val) and value is None:
|
||||
value = val
|
||||
elif isinstance(val, str) and label is None:
|
||||
label = val
|
||||
|
||||
if value is not None and value > 0:
|
||||
values.append(value)
|
||||
labels.append(label or f"Slice {len(values)}")
|
||||
|
||||
if not values:
|
||||
return "No valid data for pie chart"
|
||||
|
||||
# Calculate percentages
|
||||
total = sum(values)
|
||||
if total == 0:
|
||||
return "Total is zero"
|
||||
|
||||
for label, value in zip(labels, values, strict=False):
|
||||
percentage = (value / total) * 100
|
||||
bar_length = int(percentage / 3) # Scale to fit
|
||||
bar = "●" * bar_length
|
||||
lines.append(f"{label[:15]:>15}: {bar} {percentage:.1f}%")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _generate_safe_ascii_table(data: List[Dict[str, Any]]) -> str:
|
||||
"""Generate ASCII table with safe formatting."""
|
||||
if not data:
|
||||
return "No data available"
|
||||
|
||||
lines = ["Data Table", "=" * 50]
|
||||
|
||||
# Get columns
|
||||
columns = list(data[0].keys()) if data else []
|
||||
|
||||
# Format header
|
||||
header = " | ".join(str(col)[:10] for col in columns[:5])
|
||||
lines.append(header)
|
||||
lines.append("-" * len(header))
|
||||
|
||||
# Format rows
|
||||
for row in data[:10]:
|
||||
row_str = " | ".join(str(row.get(col, ""))[:10] for col in columns[:5])
|
||||
lines.append(row_str)
|
||||
|
||||
if len(data) > 10:
|
||||
lines.append(f"... {len(data) - 10} more rows")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _is_nan(value: Any) -> bool:
|
||||
"""Check if a value is NaN."""
|
||||
try:
|
||||
import math
|
||||
|
||||
return math.isnan(float(value))
|
||||
except (ValueError, TypeError):
|
||||
return False
|
||||
|
||||
|
||||
def _generate_vega_lite_preview_from_data( # noqa: C901
|
||||
data: List[Dict[str, Any]], form_data: Dict[str, Any]
|
||||
) -> VegaLitePreview:
|
||||
"""Generate Vega-Lite preview from raw data and form_data."""
|
||||
viz_type = form_data.get("viz_type", "table")
|
||||
|
||||
# Map Superset viz types to Vega-Lite marks
|
||||
viz_to_mark = {
|
||||
"echarts_timeseries_line": "line",
|
||||
"echarts_timeseries_bar": "bar",
|
||||
"echarts_area": "area",
|
||||
"echarts_timeseries_scatter": "point",
|
||||
"bar": "bar",
|
||||
"line": "line",
|
||||
"area": "area",
|
||||
"scatter": "point",
|
||||
"pie": "arc",
|
||||
"table": "text",
|
||||
}
|
||||
|
||||
mark = viz_to_mark.get(viz_type, "bar")
|
||||
|
||||
# Basic Vega-Lite spec
|
||||
spec = {
|
||||
"$schema": "https://vega.github.io/schema/vega-lite/v5.json",
|
||||
"data": {"values": data},
|
||||
"mark": mark,
|
||||
}
|
||||
|
||||
# Get x_axis and metrics from form_data
|
||||
x_axis = form_data.get("x_axis")
|
||||
metrics = form_data.get("metrics", [])
|
||||
groupby = form_data.get("groupby", [])
|
||||
|
||||
# Build encoding based on available fields
|
||||
encoding = {}
|
||||
|
||||
# Handle X-axis
|
||||
if x_axis and x_axis in (data[0] if data else {}):
|
||||
# Detect field type from data
|
||||
field_type = "nominal" # default
|
||||
if data and len(data) > 0:
|
||||
sample_val = data[0].get(x_axis)
|
||||
if isinstance(sample_val, str):
|
||||
# Check if it's a date/time
|
||||
if any(char in str(sample_val) for char in ["-", "/", ":"]):
|
||||
field_type = "temporal"
|
||||
else:
|
||||
field_type = "nominal"
|
||||
elif isinstance(sample_val, (int, float)):
|
||||
field_type = "quantitative"
|
||||
|
||||
encoding["x"] = {
|
||||
"field": x_axis,
|
||||
"type": field_type,
|
||||
"title": x_axis,
|
||||
}
|
||||
|
||||
# Handle Y-axis (metrics)
|
||||
if metrics and data:
|
||||
# Find the first metric column in the data
|
||||
metric_col = None
|
||||
for col in data[0].keys():
|
||||
# Check if this is a metric column (usually has aggregation in name)
|
||||
if any(
|
||||
agg in str(col).upper()
|
||||
for agg in ["SUM", "AVG", "COUNT", "MIN", "MAX", "TOTAL"]
|
||||
):
|
||||
metric_col = col
|
||||
break
|
||||
# Or check if it's numeric
|
||||
elif isinstance(data[0].get(col), (int, float)):
|
||||
metric_col = col
|
||||
break
|
||||
|
||||
if metric_col:
|
||||
encoding["y"] = {
|
||||
"field": metric_col,
|
||||
"type": "quantitative",
|
||||
"title": metric_col,
|
||||
}
|
||||
|
||||
# Handle color encoding for groupby
|
||||
if groupby and len(groupby) > 0 and groupby[0] in (data[0] if data else {}):
|
||||
encoding["color"] = {
|
||||
"field": groupby[0],
|
||||
"type": "nominal",
|
||||
"title": groupby[0],
|
||||
}
|
||||
|
||||
# Special handling for pie charts
|
||||
if mark == "arc" and data:
|
||||
# For pie charts, we need theta encoding
|
||||
if "y" in encoding:
|
||||
encoding["theta"] = encoding.pop("y")
|
||||
encoding["theta"]["stack"] = True
|
||||
if "x" in encoding:
|
||||
# Use x as color for pie
|
||||
encoding["color"] = {
|
||||
"field": encoding["x"]["field"],
|
||||
"type": "nominal",
|
||||
}
|
||||
del encoding["x"]
|
||||
|
||||
# Add encoding to spec
|
||||
if encoding:
|
||||
spec["encoding"] = encoding
|
||||
|
||||
# Add responsive sizing - Vega-Lite supports "container" as a special width value
|
||||
spec["width"] = "container"
|
||||
spec["height"] = 400 # type: ignore
|
||||
|
||||
# Add interactivity
|
||||
if mark in ["line", "point", "bar", "area"]:
|
||||
spec["selection"] = {
|
||||
"highlight": {
|
||||
"type": "single",
|
||||
"on": "mouseover",
|
||||
"empty": "none",
|
||||
}
|
||||
}
|
||||
|
||||
return VegaLitePreview(
|
||||
specification=spec,
|
||||
data_url=None,
|
||||
supports_streaming=False,
|
||||
)
|
||||
21
superset/mcp_service/chart/prompts/__init__.py
Normal file
21
superset/mcp_service/chart/prompts/__init__.py
Normal file
@@ -0,0 +1,21 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Chart prompts for Superset MCP service"""
|
||||
|
||||
# Import to register prompts when module is imported
|
||||
from . import create_chart_guided # noqa: F401
|
||||
195
superset/mcp_service/chart/prompts/create_chart_guided.py
Normal file
195
superset/mcp_service/chart/prompts/create_chart_guided.py
Normal file
@@ -0,0 +1,195 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Chart prompts for visualization guidance
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.prompt("create_chart_guided")
|
||||
@mcp_auth_hook
|
||||
async def create_chart_guided_prompt(
|
||||
chart_type: str = "auto", business_goal: str = "exploration"
|
||||
) -> str:
|
||||
"""
|
||||
AI-powered chart creation guide following Anthropic's agent design principles.
|
||||
|
||||
This prompt implements:
|
||||
- Transparency: Clear reasoning at each step
|
||||
- Proactive Intelligence: Suggests insights before being asked
|
||||
- Context Awareness: Maintains conversational flow
|
||||
- Business Focus: Translates data into actionable insights
|
||||
- Validation: Verifies choices before proceeding
|
||||
- Natural Interaction: Conversational, not configuration-driven
|
||||
|
||||
Args:
|
||||
chart_type: Preferred chart type (auto, line, bar, pie, table, scatter, area)
|
||||
business_goal: Purpose (exploration, reporting, monitoring, presentation)
|
||||
"""
|
||||
|
||||
# Enhanced chart intelligence with business context
|
||||
chart_intelligence = {
|
||||
"line": {
|
||||
"description": "Time series visualization for trend analysis",
|
||||
"best_for": "Tracking performance over time, identifying patterns",
|
||||
"business_value": "Reveals growth trends, seasonality, and patterns",
|
||||
"data_requirements": "Temporal column + continuous metrics",
|
||||
},
|
||||
"bar": {
|
||||
"description": "Category comparison visualization",
|
||||
"best_for": "Ranking, comparisons, and performance by category",
|
||||
"business_value": "Identifies top performers, bottlenecks, and gaps",
|
||||
"data_requirements": "Categorical dimensions + aggregatable metrics",
|
||||
},
|
||||
"scatter": {
|
||||
"description": "Correlation and relationship analysis",
|
||||
"best_for": "Finding relationships, outlier detection, clustering",
|
||||
"business_value": "Uncovers hidden correlations and identifies anomalies",
|
||||
"data_requirements": "Two continuous variables, optional grouping",
|
||||
},
|
||||
"table": {
|
||||
"description": "Detailed data exploration and exact values",
|
||||
"best_for": "Detailed analysis, data validation, precise values",
|
||||
"business_value": "Provides granular insights and detailed reporting",
|
||||
"data_requirements": "Any combination of dimensions and metrics",
|
||||
},
|
||||
"area": {
|
||||
"description": "Volume and composition over time",
|
||||
"best_for": "Showing cumulative effects, stacked comparisons",
|
||||
"business_value": "Visualizes contribution and total volume trends",
|
||||
"data_requirements": "Temporal dimension + stackable metrics",
|
||||
},
|
||||
"auto": {
|
||||
"description": "AI-powered visualization recommendation",
|
||||
"best_for": "When you're not sure what chart type to use",
|
||||
"business_value": "Optimizes chart choice based on data characteristics",
|
||||
"data_requirements": "I'll analyze your data and recommend the best type",
|
||||
},
|
||||
}
|
||||
|
||||
# Business context intelligence
|
||||
goal_intelligence = {
|
||||
"exploration": {
|
||||
"approach": "Interactive discovery and pattern finding",
|
||||
"features": "Filters, drill-downs, multiple perspectives",
|
||||
"outcome": "Uncover hidden insights and generate hypotheses",
|
||||
},
|
||||
"reporting": {
|
||||
"approach": "Clear, professional, and consistent presentation",
|
||||
"features": "Clean design, appropriate aggregation, clear labels",
|
||||
"outcome": "Reliable, repeatable business reporting",
|
||||
},
|
||||
"monitoring": {
|
||||
"approach": "Real-time tracking with clear thresholds",
|
||||
"features": "Alert conditions, trend indicators, key metrics",
|
||||
"outcome": "Proactive issue detection and performance tracking",
|
||||
},
|
||||
"presentation": {
|
||||
"approach": "Compelling visual storytelling",
|
||||
"features": "Engaging colors, clear messaging, audience-appropriate detail",
|
||||
"outcome": "Persuasive data-driven presentations for stakeholders",
|
||||
},
|
||||
}
|
||||
|
||||
selected_chart = chart_intelligence.get(chart_type, chart_intelligence["auto"])
|
||||
selected_goal = goal_intelligence.get(
|
||||
business_goal, goal_intelligence["exploration"]
|
||||
)
|
||||
|
||||
return f"""🎯 **AI-Powered Chart Creation Assistant**
|
||||
|
||||
I'm your intelligent data visualization partner! Let me help you create charts.
|
||||
|
||||
**Your Visualization Goal:**
|
||||
📊 **Chart Focus**: {chart_type.title()} - {selected_chart["description"]}
|
||||
🎯 **Business Purpose**: {business_goal.title()} - {selected_goal["approach"]}
|
||||
💡 **Expected Value**: {selected_chart["business_value"]}
|
||||
|
||||
---
|
||||
|
||||
## 🚀 My Intelligent Approach
|
||||
|
||||
### **Phase 1: Data Intelligence** 📊
|
||||
I'll automatically analyze your dataset to understand:
|
||||
- **Data characteristics** (types, distributions, quality)
|
||||
- **Business relationships** (correlations, hierarchies, trends)
|
||||
- **Visualization opportunities** (what stories your data can tell)
|
||||
- **Performance considerations** (size, complexity, aggregation needs)
|
||||
|
||||
*Why this matters: The right chart depends on your data's unique characteristics*
|
||||
|
||||
### **Phase 2: Smart Recommendations** 🧠
|
||||
Based on your data analysis, I'll:
|
||||
- **Recommend optimal chart types** with confidence scores and reasoning
|
||||
- **Suggest meaningful metrics** that align with your business goal
|
||||
- **Identify interesting patterns** you might want to highlight
|
||||
- **Propose filters** to focus on what matters most
|
||||
|
||||
*Why this matters: I'll spot opportunities you might miss and save you time*
|
||||
|
||||
### **Phase 3: Intelligent Configuration** ⚙️
|
||||
I'll configure your chart with:
|
||||
- **Business-appropriate aggregations** (daily, weekly, monthly for time series)
|
||||
- **Meaningful labels and formatting** (currency, percentages, readable names)
|
||||
- **Performance optimizations** (appropriate limits, caching strategies)
|
||||
- **Visual best practices** (colors, scales, legends that enhance understanding)
|
||||
|
||||
*Why this matters: Proper configuration makes charts both beautiful and actionable*
|
||||
|
||||
### **Phase 4: Validation & Refinement** 🎯
|
||||
Before finalizing, I'll:
|
||||
- **Verify the chart answers your business question**
|
||||
- **Check data quality and completeness**
|
||||
- **Suggest improvements** based on visualization best practices
|
||||
- **Provide preview** so you can see exactly what you're getting
|
||||
|
||||
*Why this matters: Great charts require iteration and validation*
|
||||
|
||||
---
|
||||
|
||||
## 🎬 Let's Begin Your Data Story
|
||||
|
||||
I'm ready to be your proactive data exploration partner. Here's how we can start:
|
||||
|
||||
**Option 1: Quick Start** ⚡
|
||||
Tell me: *"What business question are you trying to answer?"*
|
||||
(e.g., "How are our sales trending?" or "Which products perform best?")
|
||||
|
||||
**Option 2: Dataset Exploration** 🔍
|
||||
I can show you available datasets: `list_datasets`
|
||||
Or explore a specific one: `get_dataset_info [dataset_id]`
|
||||
|
||||
**Option 3: Visual Inspiration** 🎨
|
||||
Browse pre-built chart configurations: `superset://chart/configs` resource
|
||||
Perfect for when you want to see examples of great charts!
|
||||
|
||||
**Option 4: Autonomous Discovery** 🤖
|
||||
Just point me to a dataset and say *"Find something interesting"*
|
||||
I'll explore autonomously and surface the most compelling insights!
|
||||
|
||||
---
|
||||
|
||||
💡 **Pro Tip**: Great charts combine business intuition with data analysis!
|
||||
|
||||
**What's your data challenge today?** 🚀"""
|
||||
21
superset/mcp_service/chart/resources/__init__.py
Normal file
21
superset/mcp_service/chart/resources/__init__.py
Normal file
@@ -0,0 +1,21 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Chart resources for Superset MCP service"""
|
||||
|
||||
# Import to register resources when module is imported
|
||||
from . import chart_configs # noqa: F401
|
||||
362
superset/mcp_service/chart/resources/chart_configs.py
Normal file
362
superset/mcp_service/chart/resources/chart_configs.py
Normal file
@@ -0,0 +1,362 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Chart resources for valid ChartConfig examples and templates
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.resource("superset://chart/configs")
|
||||
@mcp_auth_hook
|
||||
def get_chart_configs_resource() -> str:
|
||||
"""
|
||||
Provide valid ChartConfig examples that match the exact schema.
|
||||
|
||||
This resource gives LLMs access to:
|
||||
- Valid ChartConfig examples for XYChartConfig and TableChartConfig
|
||||
- Working configurations that pass schema validation
|
||||
- Examples with proper ColumnRef, FilterConfig, AxisConfig, and LegendConfig
|
||||
- Best practices for each chart type configuration
|
||||
"""
|
||||
|
||||
# Valid XYChartConfig examples - these match the exact schema
|
||||
xy_chart_configs = {
|
||||
"line_chart": {
|
||||
"description": "Basic line chart for time series analysis",
|
||||
"config": {
|
||||
"chart_type": "xy",
|
||||
"kind": "line",
|
||||
"x": {"name": "created_on", "label": "Date Created"},
|
||||
"y": [
|
||||
{
|
||||
"name": "count_metric",
|
||||
"aggregate": "COUNT",
|
||||
"label": "Total Count",
|
||||
}
|
||||
],
|
||||
},
|
||||
"use_cases": [
|
||||
"Time series trends",
|
||||
"Historical analysis",
|
||||
"Growth tracking",
|
||||
],
|
||||
},
|
||||
"bar_chart": {
|
||||
"description": "Bar chart for category comparison",
|
||||
"config": {
|
||||
"chart_type": "xy",
|
||||
"kind": "bar",
|
||||
"x": {"name": "category", "label": "Category"},
|
||||
"y": [{"name": "sales", "aggregate": "SUM", "label": "Total Sales"}],
|
||||
"x_axis": {"title": "Product Categories", "scale": "linear"},
|
||||
"y_axis": {
|
||||
"title": "Revenue ($)",
|
||||
"format": "$,.0f",
|
||||
"scale": "linear",
|
||||
},
|
||||
},
|
||||
"use_cases": ["Category comparison", "Rankings", "Performance metrics"],
|
||||
},
|
||||
"multi_metric_line": {
|
||||
"description": "Multi-metric line chart with grouping",
|
||||
"config": {
|
||||
"chart_type": "xy",
|
||||
"kind": "line",
|
||||
"x": {"name": "date_column", "label": "Date"},
|
||||
"y": [
|
||||
{"name": "revenue", "aggregate": "SUM", "label": "Revenue"},
|
||||
{
|
||||
"name": "users",
|
||||
"aggregate": "COUNT_DISTINCT",
|
||||
"label": "Unique Users",
|
||||
},
|
||||
],
|
||||
"group_by": {"name": "region", "label": "Region"},
|
||||
"legend": {"show": True, "position": "right"},
|
||||
"filters": [{"column": "status", "op": "=", "value": "active"}],
|
||||
},
|
||||
"use_cases": [
|
||||
"Multi-dimensional analysis",
|
||||
"Regional comparisons",
|
||||
"KPI tracking",
|
||||
],
|
||||
},
|
||||
"scatter_plot": {
|
||||
"description": "Scatter plot for correlation analysis",
|
||||
"config": {
|
||||
"chart_type": "xy",
|
||||
"kind": "scatter",
|
||||
"x": {
|
||||
"name": "advertising_spend",
|
||||
"aggregate": "AVG",
|
||||
"label": "Avg Ad Spend",
|
||||
},
|
||||
"y": [
|
||||
{
|
||||
"name": "conversion_rate",
|
||||
"aggregate": "AVG",
|
||||
"label": "Avg Conversion Rate",
|
||||
}
|
||||
],
|
||||
"group_by": {"name": "campaign_type", "label": "Campaign Type"},
|
||||
"x_axis": {"title": "Average Advertising Spend", "format": "$,.0f"},
|
||||
"y_axis": {"title": "Conversion Rate", "format": ".2%"},
|
||||
},
|
||||
"use_cases": [
|
||||
"Correlation analysis",
|
||||
"Outlier detection",
|
||||
"Performance relationships",
|
||||
],
|
||||
},
|
||||
"area_chart": {
|
||||
"description": "Area chart for volume visualization",
|
||||
"config": {
|
||||
"chart_type": "xy",
|
||||
"kind": "area",
|
||||
"x": {"name": "month", "label": "Month"},
|
||||
"y": [
|
||||
{"name": "signups", "aggregate": "SUM", "label": "Monthly Signups"}
|
||||
],
|
||||
"filters": [
|
||||
{"column": "year", "op": ">=", "value": 2023},
|
||||
{"column": "active", "op": "=", "value": True},
|
||||
],
|
||||
},
|
||||
"use_cases": ["Volume trends", "Cumulative metrics", "Stacked comparisons"],
|
||||
},
|
||||
}
|
||||
|
||||
# Valid TableChartConfig examples - these match the exact schema
|
||||
table_chart_configs = {
|
||||
"basic_table": {
|
||||
"description": "Basic data table with multiple columns",
|
||||
"config": {
|
||||
"chart_type": "table",
|
||||
"columns": [
|
||||
{"name": "name", "label": "Customer Name"},
|
||||
{"name": "email", "label": "Email Address"},
|
||||
{"name": "orders", "aggregate": "COUNT", "label": "Total Orders"},
|
||||
{"name": "revenue", "aggregate": "SUM", "label": "Total Revenue"},
|
||||
],
|
||||
"sort_by": ["Total Revenue"],
|
||||
},
|
||||
"use_cases": [
|
||||
"Detailed data views",
|
||||
"Customer lists",
|
||||
"Transaction records",
|
||||
],
|
||||
},
|
||||
"aggregated_table": {
|
||||
"description": "Table with aggregated metrics and filters",
|
||||
"config": {
|
||||
"chart_type": "table",
|
||||
"columns": [
|
||||
{"name": "region", "label": "Sales Region"},
|
||||
{
|
||||
"name": "sales_amount",
|
||||
"aggregate": "SUM",
|
||||
"label": "Total Sales",
|
||||
},
|
||||
{
|
||||
"name": "sales_amount",
|
||||
"aggregate": "AVG",
|
||||
"label": "Average Sale",
|
||||
},
|
||||
{
|
||||
"name": "customer_id",
|
||||
"aggregate": "COUNT_DISTINCT",
|
||||
"label": "Unique Customers",
|
||||
},
|
||||
],
|
||||
"filters": [
|
||||
{"column": "sale_date", "op": ">=", "value": "2024-01-01"},
|
||||
{"column": "status", "op": "!=", "value": "cancelled"},
|
||||
],
|
||||
"sort_by": ["Total Sales", "Sales Region"],
|
||||
},
|
||||
"use_cases": ["Summary reports", "Regional analysis", "Performance tables"],
|
||||
},
|
||||
}
|
||||
|
||||
# Schema reference for developers
|
||||
schema_reference = {
|
||||
"ChartConfig": {
|
||||
"description": "Union type - XYChartConfig or TableChartConfig by type",
|
||||
"discriminator": "chart_type",
|
||||
"types": ["xy", "table"],
|
||||
},
|
||||
"XYChartConfig": {
|
||||
"required_fields": ["chart_type", "x", "y"],
|
||||
"optional_fields": [
|
||||
"kind",
|
||||
"group_by",
|
||||
"x_axis",
|
||||
"y_axis",
|
||||
"legend",
|
||||
"filters",
|
||||
],
|
||||
"chart_type": "xy",
|
||||
"kind_options": ["line", "bar", "area", "scatter"],
|
||||
"validation_rules": [
|
||||
"All column labels must be unique across x, y, and group_by",
|
||||
"Y-axis must have at least one column",
|
||||
"Column names must match pattern: ^[a-zA-Z0-9_][a-zA-Z0-9_\\s\\-\\.]*$",
|
||||
],
|
||||
},
|
||||
"TableChartConfig": {
|
||||
"required_fields": ["chart_type", "columns"],
|
||||
"optional_fields": ["filters", "sort_by"],
|
||||
"chart_type": "table",
|
||||
"validation_rules": [
|
||||
"Must have at least one column",
|
||||
"All column labels must be unique",
|
||||
"Column names must match pattern: ^[a-zA-Z0-9_][a-zA-Z0-9_\\s\\-\\.]*$",
|
||||
],
|
||||
},
|
||||
"ColumnRef": {
|
||||
"required_fields": ["name"],
|
||||
"optional_fields": ["label", "dtype", "aggregate"],
|
||||
"aggregate_options": [
|
||||
"SUM",
|
||||
"COUNT",
|
||||
"AVG",
|
||||
"MIN",
|
||||
"MAX",
|
||||
"COUNT_DISTINCT",
|
||||
"STDDEV",
|
||||
"VAR",
|
||||
"MEDIAN",
|
||||
"PERCENTILE",
|
||||
],
|
||||
"validation_rules": [
|
||||
"Name cannot be empty and must follow pattern",
|
||||
"Labels are HTML-escaped to prevent XSS",
|
||||
"Aggregates are validated against allowed functions",
|
||||
],
|
||||
},
|
||||
"FilterConfig": {
|
||||
"required_fields": ["column", "op", "value"],
|
||||
"operator_options": ["=", ">", "<", ">=", "<=", "!="],
|
||||
"value_types": ["string", "number", "boolean"],
|
||||
"validation_rules": [
|
||||
"Column names are sanitized to prevent injection",
|
||||
"Values are checked for malicious patterns",
|
||||
"String values are HTML-escaped",
|
||||
],
|
||||
},
|
||||
"AxisConfig": {
|
||||
"optional_fields": ["title", "scale", "format"],
|
||||
"scale_options": ["linear", "log"],
|
||||
"format_examples": ["$,.2f", ".2%", ",.0f", ".1f"],
|
||||
},
|
||||
"LegendConfig": {
|
||||
"optional_fields": ["show", "position"],
|
||||
"show_default": True,
|
||||
"position_options": ["top", "bottom", "left", "right"],
|
||||
"position_default": "right",
|
||||
},
|
||||
}
|
||||
|
||||
# Best practices for each configuration type
|
||||
best_practices = {
|
||||
"xy_charts": [
|
||||
"Use descriptive labels for axes and metrics",
|
||||
"Choose appropriate aggregation functions for your data",
|
||||
"Limit the number of Y-axis metrics (3-5 maximum)",
|
||||
"Use filters to focus on relevant data",
|
||||
"Configure axis formatting for better readability",
|
||||
"Consider grouping when comparing categories",
|
||||
"Use chart kinds: line for trends, bar for comparisons, scatter plots",
|
||||
],
|
||||
"table_charts": [
|
||||
"Include essential columns only to avoid clutter",
|
||||
"Use meaningful column labels",
|
||||
"Apply sorting to highlight important data",
|
||||
"Use filters to limit result sets",
|
||||
"Mix dimensions and aggregated metrics appropriately",
|
||||
"Ensure unique labels to avoid conflicts",
|
||||
"Consider performance with large datasets",
|
||||
],
|
||||
"general": [
|
||||
"Always specify chart_type as the first field",
|
||||
"Use consistent naming conventions for columns",
|
||||
"Validate column names exist in your dataset",
|
||||
"Test configurations with actual data",
|
||||
"Consider caching for frequently accessed charts",
|
||||
"Apply security best practices - avoid user input in column names",
|
||||
],
|
||||
}
|
||||
|
||||
# Common patterns and examples
|
||||
common_patterns = {
|
||||
"time_series": {
|
||||
"description": "Standard time-based analysis",
|
||||
"x_column_types": ["date", "datetime", "timestamp"],
|
||||
"recommended_aggregations": ["SUM", "COUNT", "AVG"],
|
||||
"best_chart_types": ["line", "area", "bar"],
|
||||
},
|
||||
"categorical_analysis": {
|
||||
"description": "Comparing discrete categories",
|
||||
"x_column_types": ["string", "category", "enum"],
|
||||
"recommended_aggregations": ["SUM", "COUNT", "COUNT_DISTINCT", "AVG"],
|
||||
"best_chart_types": ["bar", "table"],
|
||||
},
|
||||
"correlation_analysis": {
|
||||
"description": "Finding relationships between variables",
|
||||
"requirements": ["Two numerical metrics"],
|
||||
"recommended_aggregations": ["AVG", "SUM", "MEDIAN"],
|
||||
"best_chart_types": ["scatter"],
|
||||
},
|
||||
}
|
||||
|
||||
resource_data = {
|
||||
"xy_chart_configs": xy_chart_configs,
|
||||
"table_chart_configs": table_chart_configs,
|
||||
"schema_reference": schema_reference,
|
||||
"best_practices": best_practices,
|
||||
"common_patterns": common_patterns,
|
||||
"metadata": {
|
||||
"version": "1.0",
|
||||
"schema_version": "ChartConfig v1.0",
|
||||
"last_updated": "2025-08-07",
|
||||
"usage_notes": [
|
||||
"All examples are valid ChartConfig objects that pass validation",
|
||||
"Copy these configurations directly into generate_chart requests",
|
||||
"Modify column names and labels to match your actual dataset",
|
||||
"Test configurations with get_dataset_info to verify columns",
|
||||
"All examples follow security best practices and input validation",
|
||||
],
|
||||
"validation_info": [
|
||||
"Column names must match: ^[a-zA-Z0-9_][a-zA-Z0-9_\\s\\-\\.]*$",
|
||||
"Labels are automatically HTML-escaped for security",
|
||||
"Filter values are sanitized to prevent injection attacks",
|
||||
"All field lengths are validated against schema limits",
|
||||
"Duplicate labels are automatically detected and rejected",
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
from superset.utils import json
|
||||
|
||||
return json.dumps(resource_data, indent=2)
|
||||
@@ -21,6 +21,8 @@ Pydantic schemas for chart-related responses
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
import re
|
||||
from datetime import datetime, timezone
|
||||
from typing import Annotated, Any, Dict, List, Literal, Protocol
|
||||
|
||||
@@ -28,12 +30,19 @@ from pydantic import (
|
||||
BaseModel,
|
||||
ConfigDict,
|
||||
Field,
|
||||
field_validator,
|
||||
model_validator,
|
||||
PositiveInt,
|
||||
)
|
||||
|
||||
from superset.daos.base import ColumnOperator, ColumnOperatorEnum
|
||||
from superset.mcp_service.common.cache_schemas import MetadataCacheControl
|
||||
from superset.mcp_service.common.cache_schemas import (
|
||||
CacheStatus,
|
||||
FormDataCacheControl,
|
||||
MetadataCacheControl,
|
||||
QueryCacheControl,
|
||||
)
|
||||
from superset.mcp_service.common.error_schemas import ChartGenerationError
|
||||
from superset.mcp_service.system.schemas import (
|
||||
PaginationInfo,
|
||||
TagInfo,
|
||||
@@ -101,6 +110,25 @@ class ChartInfo(BaseModel):
|
||||
model_config = ConfigDict(from_attributes=True, ser_json_timedelta="iso8601")
|
||||
|
||||
|
||||
class GetChartAvailableFiltersRequest(BaseModel):
|
||||
"""
|
||||
Request schema for get_chart_available_filters tool.
|
||||
|
||||
Currently has no parameters but provides consistent API for future extensibility.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(
|
||||
extra="forbid",
|
||||
str_strip_whitespace=True,
|
||||
)
|
||||
|
||||
|
||||
class ChartAvailableFiltersResponse(BaseModel):
|
||||
column_operators: Dict[str, Any] = Field(
|
||||
..., description="Available filter operators and metadata for each column"
|
||||
)
|
||||
|
||||
|
||||
class ChartError(BaseModel):
|
||||
error: str = Field(..., description="Error message")
|
||||
error_type: str = Field(..., description="Type of error")
|
||||
@@ -111,6 +139,63 @@ class ChartError(BaseModel):
|
||||
model_config = ConfigDict(ser_json_timedelta="iso8601")
|
||||
|
||||
|
||||
class ChartCapabilities(BaseModel):
|
||||
"""Describes what the chart can do for LLM understanding."""
|
||||
|
||||
supports_interaction: bool = Field(description="Chart supports user interaction")
|
||||
supports_real_time: bool = Field(description="Chart supports live data updates")
|
||||
supports_drill_down: bool = Field(
|
||||
description="Chart supports drill-down navigation"
|
||||
)
|
||||
supports_export: bool = Field(description="Chart can be exported to other formats")
|
||||
optimal_formats: List[str] = Field(description="Recommended preview formats")
|
||||
data_types: List[str] = Field(
|
||||
description="Types of data shown (time_series, categorical, etc)"
|
||||
)
|
||||
|
||||
|
||||
class ChartSemantics(BaseModel):
|
||||
"""Semantic information for LLM reasoning."""
|
||||
|
||||
primary_insight: str = Field(
|
||||
description="Main insight or pattern the chart reveals"
|
||||
)
|
||||
data_story: str = Field(description="Narrative description of what the data shows")
|
||||
recommended_actions: List[str] = Field(
|
||||
description="Suggested next steps based on data"
|
||||
)
|
||||
anomalies: List[str] = Field(description="Notable outliers or unusual patterns")
|
||||
statistical_summary: Dict[str, Any] = Field(
|
||||
description="Key statistics (mean, median, trends)"
|
||||
)
|
||||
|
||||
|
||||
class PerformanceMetadata(BaseModel):
|
||||
"""Performance information for LLM cost understanding."""
|
||||
|
||||
query_duration_ms: int = Field(description="Query execution time")
|
||||
estimated_cost: str | None = Field(None, description="Resource cost estimate")
|
||||
cache_status: str = Field(description="Cache hit/miss status")
|
||||
optimization_suggestions: List[str] = Field(
|
||||
default_factory=list, description="Performance improvement tips"
|
||||
)
|
||||
|
||||
|
||||
class AccessibilityMetadata(BaseModel):
|
||||
"""Accessibility information for inclusive visualization."""
|
||||
|
||||
color_blind_safe: bool = Field(description="Uses colorblind-safe palette")
|
||||
alt_text: str = Field(description="Screen reader description")
|
||||
high_contrast_available: bool = Field(description="High contrast version available")
|
||||
|
||||
|
||||
class VersionedResponse(BaseModel):
|
||||
"""Base class for versioned API responses."""
|
||||
|
||||
schema_version: str = Field("2.0", description="Response schema version")
|
||||
api_version: str = Field("v1", description="MCP API version")
|
||||
|
||||
|
||||
class GetChartInfoRequest(BaseModel):
|
||||
"""Request schema for get_chart_info with support for ID or UUID."""
|
||||
|
||||
@@ -124,12 +209,13 @@ def serialize_chart_object(chart: ChartLike | None) -> ChartInfo | None:
|
||||
if not chart:
|
||||
return None
|
||||
|
||||
# TODO (Phase 3): Generate MCP service screenshot URL
|
||||
# For now, use chart's native URL instead of screenshot URL
|
||||
# Screenshot functionality will be added in Phase 3 PR
|
||||
# Generate MCP service screenshot URL instead of chart's native URL
|
||||
from superset.mcp_service.utils.url_utils import get_chart_screenshot_url
|
||||
|
||||
chart_id = getattr(chart, "id", None)
|
||||
chart_url = getattr(chart, "url", None)
|
||||
screenshot_url = None
|
||||
if chart_id:
|
||||
screenshot_url = get_chart_screenshot_url(chart_id)
|
||||
|
||||
return ChartInfo(
|
||||
id=chart_id,
|
||||
@@ -137,7 +223,7 @@ def serialize_chart_object(chart: ChartLike | None) -> ChartInfo | None:
|
||||
viz_type=getattr(chart, "viz_type", None),
|
||||
datasource_name=getattr(chart, "datasource_name", None),
|
||||
datasource_type=getattr(chart, "datasource_type", None),
|
||||
url=chart_url,
|
||||
url=screenshot_url,
|
||||
description=getattr(chart, "description", None),
|
||||
cache_timeout=getattr(chart, "cache_timeout", None),
|
||||
form_data=getattr(chart, "form_data", None),
|
||||
@@ -214,6 +300,409 @@ class ChartList(BaseModel):
|
||||
model_config = ConfigDict(ser_json_timedelta="iso8601")
|
||||
|
||||
|
||||
# --- Simplified schemas for generate_chart tool ---
|
||||
|
||||
|
||||
# Common pieces
|
||||
class ColumnRef(BaseModel):
|
||||
name: str = Field(
|
||||
...,
|
||||
description="Column name",
|
||||
min_length=1,
|
||||
max_length=255,
|
||||
pattern=r"^[a-zA-Z0-9_][a-zA-Z0-9_\s\-\.]*$",
|
||||
)
|
||||
label: str | None = Field(
|
||||
None, description="Display label for the column", max_length=500
|
||||
)
|
||||
dtype: str | None = Field(None, description="Data type hint")
|
||||
aggregate: (
|
||||
Literal[
|
||||
"SUM",
|
||||
"COUNT",
|
||||
"AVG",
|
||||
"MIN",
|
||||
"MAX",
|
||||
"COUNT_DISTINCT",
|
||||
"STDDEV",
|
||||
"VAR",
|
||||
"MEDIAN",
|
||||
"PERCENTILE",
|
||||
]
|
||||
| None
|
||||
) = Field(
|
||||
None,
|
||||
description="SQL aggregation function. Only these validated functions are "
|
||||
"supported to prevent SQL errors.",
|
||||
)
|
||||
|
||||
@field_validator("name")
|
||||
@classmethod
|
||||
def sanitize_name(cls, v: str) -> str:
|
||||
"""Sanitize column name to prevent XSS and SQL injection."""
|
||||
if not v or not v.strip():
|
||||
raise ValueError("Column name cannot be empty")
|
||||
|
||||
# Length check first to prevent ReDoS attacks
|
||||
if len(v) > 255:
|
||||
raise ValueError(
|
||||
f"Column name too long ({len(v)} characters). "
|
||||
f"Maximum allowed length is 255 characters."
|
||||
)
|
||||
|
||||
# Remove HTML tags and decode entities
|
||||
sanitized = html.escape(v.strip())
|
||||
|
||||
# Check for dangerous HTML tags using substring checks (safe)
|
||||
dangerous_tags = ["<script", "</script>", "<iframe", "<object", "<embed"]
|
||||
v_lower = v.lower()
|
||||
for tag in dangerous_tags:
|
||||
if tag in v_lower:
|
||||
raise ValueError(
|
||||
"Column name contains potentially malicious script content"
|
||||
)
|
||||
|
||||
# Check URL schemes with word boundaries to match only actual URLs
|
||||
if re.search(r"\b(javascript|vbscript|data):", v, re.IGNORECASE):
|
||||
raise ValueError("Column name contains potentially malicious URL scheme")
|
||||
|
||||
# Basic SQL injection patterns (basic protection)
|
||||
# Use simple patterns without backtracking
|
||||
dangerous_patterns = [
|
||||
r"[;|&$`]", # Dangerous shell characters
|
||||
r"\b(DROP|DELETE|INSERT|UPDATE|CREATE|ALTER|EXEC|EXECUTE)\b",
|
||||
r"--", # SQL comment
|
||||
r"/\*", # SQL comment start (just check for start, not full pattern)
|
||||
]
|
||||
|
||||
for pattern in dangerous_patterns:
|
||||
if re.search(pattern, v, re.IGNORECASE):
|
||||
raise ValueError(
|
||||
"Column name contains potentially unsafe characters or SQL keywords"
|
||||
)
|
||||
|
||||
return sanitized
|
||||
|
||||
@field_validator("label")
|
||||
@classmethod
|
||||
def sanitize_label(cls, v: str | None) -> str | None:
|
||||
"""Sanitize display label to prevent XSS attacks."""
|
||||
if v is None:
|
||||
return v
|
||||
|
||||
# Strip whitespace
|
||||
v = v.strip()
|
||||
if not v:
|
||||
return None
|
||||
|
||||
# Length check first to prevent ReDoS attacks
|
||||
if len(v) > 500:
|
||||
raise ValueError(
|
||||
f"Label too long ({len(v)} characters). "
|
||||
f"Maximum allowed length is 500 characters."
|
||||
)
|
||||
|
||||
# Check for dangerous HTML tags and JavaScript protocols using substring checks
|
||||
# This avoids ReDoS vulnerabilities from regex patterns
|
||||
dangerous_tags = [
|
||||
"<script",
|
||||
"</script>",
|
||||
"<iframe",
|
||||
"</iframe>",
|
||||
"<object",
|
||||
"</object>",
|
||||
"<embed",
|
||||
"</embed>",
|
||||
"<link",
|
||||
"<meta",
|
||||
]
|
||||
|
||||
v_lower = v.lower()
|
||||
for tag in dangerous_tags:
|
||||
if tag in v_lower:
|
||||
raise ValueError(
|
||||
"Label contains potentially malicious content. "
|
||||
"HTML tags, JavaScript, and event handlers are not allowed "
|
||||
"in labels."
|
||||
)
|
||||
|
||||
# Check URL schemes and event handlers with word boundaries
|
||||
dangerous_patterns = [
|
||||
r"\b(javascript|vbscript|data):", # URL schemes
|
||||
r"on\w+\s*=", # Event handlers
|
||||
]
|
||||
for pattern in dangerous_patterns:
|
||||
if re.search(pattern, v, re.IGNORECASE):
|
||||
raise ValueError(
|
||||
"Label contains potentially malicious content. "
|
||||
"HTML tags, JavaScript, and event handlers are not allowed."
|
||||
)
|
||||
|
||||
# Filter dangerous Unicode characters
|
||||
v = re.sub(
|
||||
r"[\u200B-\u200D\uFEFF\u0000-\u0008\u000B\u000C\u000E-\u001F]", "", v
|
||||
)
|
||||
|
||||
# HTML escape the cleaned content
|
||||
sanitized = html.escape(v)
|
||||
|
||||
return sanitized if sanitized else None
|
||||
|
||||
|
||||
class AxisConfig(BaseModel):
|
||||
title: str | None = Field(None, description="Axis title", max_length=200)
|
||||
scale: Literal["linear", "log"] | None = Field(
|
||||
"linear", description="Axis scale type"
|
||||
)
|
||||
format: str | None = Field(
|
||||
None, description="Format string (e.g. '$,.2f')", max_length=50
|
||||
)
|
||||
|
||||
|
||||
class LegendConfig(BaseModel):
|
||||
show: bool = Field(True, description="Whether to show legend")
|
||||
position: Literal["top", "bottom", "left", "right"] | None = Field(
|
||||
"right", description="Legend position"
|
||||
)
|
||||
|
||||
|
||||
class FilterConfig(BaseModel):
|
||||
column: str = Field(
|
||||
..., description="Column to filter on", min_length=1, max_length=255
|
||||
)
|
||||
op: Literal["=", ">", "<", ">=", "<=", "!="] = Field(
|
||||
..., description="Filter operator"
|
||||
)
|
||||
value: str | int | float | bool = Field(..., description="Filter value")
|
||||
|
||||
@field_validator("column")
|
||||
@classmethod
|
||||
def sanitize_column(cls, v: str) -> str:
|
||||
"""Sanitize filter column name to prevent injection attacks."""
|
||||
if not v or not v.strip():
|
||||
raise ValueError("Filter column name cannot be empty")
|
||||
|
||||
# Length check first to prevent ReDoS attacks
|
||||
if len(v) > 255:
|
||||
raise ValueError(
|
||||
f"Filter column name too long ({len(v)} characters). "
|
||||
f"Maximum allowed length is 255 characters."
|
||||
)
|
||||
|
||||
# Remove HTML tags and decode entities
|
||||
sanitized = html.escape(v.strip())
|
||||
|
||||
# Check for dangerous HTML tags using substring checks (safe)
|
||||
dangerous_tags = ["<script", "</script>"]
|
||||
v_lower = v.lower()
|
||||
for tag in dangerous_tags:
|
||||
if tag in v_lower:
|
||||
raise ValueError(
|
||||
"Filter column contains potentially malicious script content"
|
||||
)
|
||||
|
||||
# Check URL schemes with word boundaries
|
||||
if re.search(r"\b(javascript|vbscript|data):", v, re.IGNORECASE):
|
||||
raise ValueError("Filter column contains potentially malicious URL scheme")
|
||||
|
||||
return sanitized
|
||||
|
||||
@staticmethod
|
||||
def _validate_string_value(v: str) -> None:
|
||||
"""Validate string filter value for security issues."""
|
||||
# Check for dangerous HTML tags and SQL procedures
|
||||
dangerous_substrings = [
|
||||
"<script",
|
||||
"</script>",
|
||||
"<iframe",
|
||||
"<object",
|
||||
"<embed",
|
||||
"xp_cmdshell",
|
||||
"sp_executesql",
|
||||
]
|
||||
v_lower = v.lower()
|
||||
for substring in dangerous_substrings:
|
||||
if substring in v_lower:
|
||||
raise ValueError(
|
||||
"Filter value contains potentially malicious content. "
|
||||
"HTML tags and JavaScript are not allowed."
|
||||
)
|
||||
|
||||
# Check URL schemes with word boundaries
|
||||
if re.search(r"\b(javascript|vbscript|data):", v, re.IGNORECASE):
|
||||
raise ValueError("Filter value contains potentially malicious URL scheme")
|
||||
|
||||
# SQL injection patterns
|
||||
sql_patterns = [
|
||||
r";\s*(DROP|DELETE|INSERT|UPDATE|CREATE|ALTER|EXEC|EXECUTE)\b",
|
||||
r"'\s*OR\s*'",
|
||||
r"'\s*AND\s*'",
|
||||
r"--\s*",
|
||||
r"/\*",
|
||||
r"UNION\s+SELECT",
|
||||
]
|
||||
for pattern in sql_patterns:
|
||||
if re.search(pattern, v, re.IGNORECASE):
|
||||
raise ValueError(
|
||||
"Filter value contains potentially malicious SQL patterns."
|
||||
)
|
||||
|
||||
# Check for other dangerous patterns
|
||||
if re.search(r"[;&|`$()]", v):
|
||||
raise ValueError(
|
||||
"Filter value contains potentially unsafe shell characters."
|
||||
)
|
||||
if re.search(r"on\w+\s*=", v, re.IGNORECASE):
|
||||
raise ValueError(
|
||||
"Filter value contains potentially malicious event handlers."
|
||||
)
|
||||
if re.search(r"\\x[0-9a-fA-F]{2}", v):
|
||||
raise ValueError("Filter value contains hex encoding which is not allowed.")
|
||||
|
||||
@field_validator("value")
|
||||
@classmethod
|
||||
def sanitize_value(cls, v: str | int | float | bool) -> str | int | float | bool:
|
||||
"""Sanitize filter value to prevent XSS and SQL injection attacks."""
|
||||
if isinstance(v, str):
|
||||
v = v.strip()
|
||||
|
||||
# Length check FIRST to prevent ReDoS attacks
|
||||
if len(v) > 1000:
|
||||
raise ValueError(
|
||||
f"Filter value too long ({len(v)} characters). "
|
||||
f"Maximum allowed length is 1000 characters."
|
||||
)
|
||||
|
||||
# Validate security
|
||||
cls._validate_string_value(v)
|
||||
|
||||
# Filter dangerous Unicode characters
|
||||
v = re.sub(
|
||||
r"[\u200B-\u200D\uFEFF\u0000-\u0008\u000B\u000C\u000E-\u001F]", "", v
|
||||
)
|
||||
|
||||
# HTML escape the cleaned content
|
||||
return html.escape(v)
|
||||
|
||||
return v # Return non-string values as-is
|
||||
|
||||
|
||||
# Actual chart types
|
||||
class TableChartConfig(BaseModel):
|
||||
chart_type: Literal["table"] = Field("table", description="Chart type")
|
||||
columns: List[ColumnRef] = Field(
|
||||
...,
|
||||
min_length=1,
|
||||
description=(
|
||||
"Columns to display. Must have at least one column. Each column must have "
|
||||
"a unique label "
|
||||
"(either explicitly set via 'label' field or auto-generated "
|
||||
"from name/aggregate)"
|
||||
),
|
||||
)
|
||||
filters: List[FilterConfig] | None = Field(None, description="Filters to apply")
|
||||
sort_by: List[str] | None = Field(None, description="Columns to sort by")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_unique_column_labels(self) -> "TableChartConfig":
|
||||
"""Ensure all column labels are unique."""
|
||||
labels_seen = set()
|
||||
duplicates = []
|
||||
|
||||
for i, col in enumerate(self.columns):
|
||||
# Generate the label that will be used (same logic as create_metric_object)
|
||||
if col.aggregate:
|
||||
label = col.label or f"{col.aggregate}({col.name})"
|
||||
else:
|
||||
label = col.label or col.name
|
||||
|
||||
if label in labels_seen:
|
||||
duplicates.append(f"columns[{i}]: '{label}'")
|
||||
else:
|
||||
labels_seen.add(label)
|
||||
|
||||
if duplicates:
|
||||
raise ValueError(
|
||||
f"Duplicate column/metric labels: {', '.join(duplicates)}. "
|
||||
f"Please make sure all columns and metrics have a unique label. "
|
||||
f"Use the 'label' field to provide custom names for columns."
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class XYChartConfig(BaseModel):
|
||||
chart_type: Literal["xy"] = Field("xy", description="Chart type")
|
||||
x: ColumnRef = Field(..., description="X-axis column")
|
||||
y: List[ColumnRef] = Field(
|
||||
...,
|
||||
min_length=1,
|
||||
description="Y-axis columns (metrics). Must have at least one Y-axis column. "
|
||||
"Each column must have a unique label "
|
||||
"that doesn't conflict with x-axis or group_by labels",
|
||||
)
|
||||
kind: Literal["line", "bar", "area", "scatter"] = Field(
|
||||
"line", description="Chart visualization type"
|
||||
)
|
||||
group_by: ColumnRef | None = Field(None, description="Column to group by")
|
||||
x_axis: AxisConfig | None = Field(None, description="X-axis configuration")
|
||||
y_axis: AxisConfig | None = Field(None, description="Y-axis configuration")
|
||||
legend: LegendConfig | None = Field(None, description="Legend configuration")
|
||||
filters: List[FilterConfig] | None = Field(None, description="Filters to apply")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_unique_column_labels(self) -> "XYChartConfig":
|
||||
"""Ensure all column labels are unique across x, y, and group_by."""
|
||||
labels_seen = {} # label -> field_name for error reporting
|
||||
duplicates = []
|
||||
|
||||
# Check X-axis label
|
||||
x_label = self.x.label or self.x.name
|
||||
labels_seen[x_label] = "x"
|
||||
|
||||
# Check Y-axis labels
|
||||
for i, col in enumerate(self.y):
|
||||
if col.aggregate:
|
||||
label = col.label or f"{col.aggregate}({col.name})"
|
||||
else:
|
||||
label = col.label or col.name
|
||||
|
||||
if label in labels_seen:
|
||||
duplicates.append(
|
||||
f"y[{i}]: '{label}' (conflicts with {labels_seen[label]})"
|
||||
)
|
||||
else:
|
||||
labels_seen[label] = f"y[{i}]"
|
||||
|
||||
# Check group_by label if present
|
||||
if self.group_by:
|
||||
group_label = self.group_by.label or self.group_by.name
|
||||
if group_label in labels_seen:
|
||||
duplicates.append(
|
||||
f"group_by: '{group_label}' "
|
||||
f"(conflicts with {labels_seen[group_label]})"
|
||||
)
|
||||
|
||||
if duplicates:
|
||||
raise ValueError(
|
||||
f"Duplicate column/metric labels: {', '.join(duplicates)}. "
|
||||
f"Please make sure all columns and metrics have a unique label. "
|
||||
f"Use the 'label' field to provide custom names for columns."
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
# Discriminated union entry point with custom error handling
|
||||
ChartConfig = Annotated[
|
||||
XYChartConfig | TableChartConfig,
|
||||
Field(
|
||||
discriminator="chart_type",
|
||||
description="Chart configuration - specify chart_type as 'xy' or 'table'",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class ListChartsRequest(MetadataCacheControl):
|
||||
"""Request schema for list_charts with clear, unambiguous types."""
|
||||
|
||||
@@ -281,3 +770,393 @@ class ListChartsRequest(MetadataCacheControl):
|
||||
"or 'filters' for precise column-based filtering, but not both."
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
# The tool input models
|
||||
class GenerateChartRequest(QueryCacheControl):
|
||||
dataset_id: int | str = Field(..., description="Dataset identifier (ID, UUID)")
|
||||
config: ChartConfig = Field(..., description="Chart configuration")
|
||||
save_chart: bool = Field(
|
||||
default=True,
|
||||
description="Whether to permanently save the chart in Superset",
|
||||
)
|
||||
generate_preview: bool = Field(
|
||||
default=True,
|
||||
description="Whether to generate a preview image",
|
||||
)
|
||||
preview_formats: List[
|
||||
Literal["url", "interactive", "ascii", "vega_lite", "table", "base64"]
|
||||
] = Field(
|
||||
default_factory=lambda: ["url"],
|
||||
description="List of preview formats to generate",
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_cache_timeout(self) -> "GenerateChartRequest":
|
||||
"""Validate cache timeout is non-negative."""
|
||||
if (
|
||||
hasattr(self, "cache_timeout")
|
||||
and self.cache_timeout is not None
|
||||
and self.cache_timeout < 0
|
||||
):
|
||||
raise ValueError(
|
||||
"cache_timeout must be non-negative (0 or positive integer)"
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class GenerateExploreLinkRequest(FormDataCacheControl):
|
||||
dataset_id: int | str = Field(..., description="Dataset identifier (ID, UUID)")
|
||||
config: ChartConfig = Field(..., description="Chart configuration")
|
||||
|
||||
|
||||
class UpdateChartRequest(QueryCacheControl):
|
||||
identifier: int | str = Field(..., description="Chart identifier (ID, UUID)")
|
||||
config: ChartConfig = Field(..., description="New chart configuration")
|
||||
chart_name: str | None = Field(
|
||||
None,
|
||||
description="New chart name (optional, will auto-generate if not provided)",
|
||||
max_length=255,
|
||||
)
|
||||
generate_preview: bool = Field(
|
||||
default=True,
|
||||
description="Whether to generate a preview after updating",
|
||||
)
|
||||
preview_formats: List[
|
||||
Literal["url", "interactive", "ascii", "vega_lite", "table", "base64"]
|
||||
] = Field(
|
||||
default_factory=lambda: ["url"],
|
||||
description="List of preview formats to generate",
|
||||
)
|
||||
|
||||
@field_validator("chart_name")
|
||||
@classmethod
|
||||
def sanitize_chart_name(cls, v: str | None) -> str | None:
|
||||
"""Sanitize chart name to prevent XSS attacks."""
|
||||
if v is None:
|
||||
return v
|
||||
|
||||
# Strip whitespace
|
||||
v = v.strip()
|
||||
if not v:
|
||||
return None
|
||||
|
||||
# Length check first to prevent ReDoS attacks
|
||||
if len(v) > 255:
|
||||
raise ValueError(
|
||||
f"Chart name too long ({len(v)} characters). "
|
||||
f"Maximum allowed length is 255 characters."
|
||||
)
|
||||
|
||||
# Check for dangerous HTML tags using substring checks (safe)
|
||||
dangerous_tags = [
|
||||
"<script",
|
||||
"</script>",
|
||||
"<iframe",
|
||||
"</iframe>",
|
||||
"<object",
|
||||
"</object>",
|
||||
"<embed",
|
||||
"</embed>",
|
||||
"<link",
|
||||
"<meta",
|
||||
]
|
||||
|
||||
v_lower = v.lower()
|
||||
for tag in dangerous_tags:
|
||||
if tag in v_lower:
|
||||
raise ValueError(
|
||||
"Chart name contains potentially malicious content. "
|
||||
"HTML tags and JavaScript are not allowed in chart names."
|
||||
)
|
||||
|
||||
# Check URL schemes with word boundaries
|
||||
if re.search(r"\b(javascript|vbscript|data):", v, re.IGNORECASE):
|
||||
raise ValueError("Chart name contains potentially malicious URL scheme")
|
||||
|
||||
# Check for event handlers with simple regex
|
||||
if re.search(r"on\w+\s*=", v, re.IGNORECASE):
|
||||
raise ValueError(
|
||||
"Chart name contains potentially malicious event handlers."
|
||||
)
|
||||
|
||||
# Filter dangerous Unicode characters
|
||||
v = re.sub(
|
||||
r"[\u200B-\u200D\uFEFF\u0000-\u0008\u000B\u000C\u000E-\u001F]", "", v
|
||||
)
|
||||
|
||||
# HTML escape the cleaned content
|
||||
sanitized = html.escape(v)
|
||||
|
||||
return sanitized if sanitized else None
|
||||
|
||||
|
||||
class UpdateChartPreviewRequest(FormDataCacheControl):
|
||||
form_data_key: str = Field(..., description="Existing form_data_key to update")
|
||||
dataset_id: int | str = Field(..., description="Dataset identifier (ID, UUID)")
|
||||
config: ChartConfig = Field(..., description="New chart configuration")
|
||||
generate_preview: bool = Field(
|
||||
default=True,
|
||||
description="Whether to generate a preview after updating",
|
||||
)
|
||||
preview_formats: List[
|
||||
Literal["url", "interactive", "ascii", "vega_lite", "table", "base64"]
|
||||
] = Field(
|
||||
default_factory=lambda: ["url"],
|
||||
description="List of preview formats to generate",
|
||||
)
|
||||
|
||||
|
||||
class GetChartDataRequest(QueryCacheControl):
|
||||
"""Request for chart data with cache control."""
|
||||
|
||||
identifier: int | str = Field(description="Chart identifier (ID, UUID)")
|
||||
limit: int | None = Field(
|
||||
default=100, description="Maximum number of data rows to return"
|
||||
)
|
||||
format: Literal["json", "csv", "excel"] = Field(
|
||||
default="json", description="Data export format"
|
||||
)
|
||||
|
||||
|
||||
class DataColumn(BaseModel):
|
||||
"""Enhanced column metadata with semantic information."""
|
||||
|
||||
name: str = Field(..., description="Column name")
|
||||
display_name: str = Field(..., description="Human-readable column name")
|
||||
data_type: str = Field(..., description="Inferred data type")
|
||||
sample_values: List[Any] = Field(description="Representative sample values")
|
||||
null_count: int = Field(description="Number of null values")
|
||||
unique_count: int = Field(description="Number of unique values")
|
||||
statistics: Dict[str, Any] | None = Field(
|
||||
None, description="Column statistics if numeric"
|
||||
)
|
||||
semantic_type: str | None = Field(
|
||||
None, description="Semantic type (currency, percentage, etc)"
|
||||
)
|
||||
|
||||
|
||||
class ChartData(BaseModel):
|
||||
"""Rich chart data response with statistical insights."""
|
||||
|
||||
# Basic information
|
||||
chart_id: int
|
||||
chart_name: str
|
||||
chart_type: str
|
||||
|
||||
# Enhanced data description
|
||||
columns: List[DataColumn] = Field(description="Rich column metadata")
|
||||
data: List[Dict[str, Any]] = Field(description="Actual data rows")
|
||||
|
||||
# Data insights
|
||||
row_count: int = Field(description="Rows returned")
|
||||
total_rows: int | None = Field(description="Total available rows")
|
||||
data_freshness: datetime | None = Field(description="When data was last updated")
|
||||
|
||||
# LLM-friendly summaries
|
||||
summary: str = Field(description="Human-readable data summary")
|
||||
insights: List[str] = Field(description="Key patterns discovered in the data")
|
||||
data_quality: Dict[str, Any] = Field(description="Data quality assessment")
|
||||
recommended_visualizations: List[str] = Field(
|
||||
description="Suggested chart types for this data"
|
||||
)
|
||||
|
||||
# Performance and metadata
|
||||
performance: PerformanceMetadata = Field(description="Query performance metrics")
|
||||
cache_status: CacheStatus | None = Field(
|
||||
None, description="Cache usage information"
|
||||
)
|
||||
|
||||
# Export format fields
|
||||
csv_data: str | None = Field(None, description="CSV content when format='csv'")
|
||||
excel_data: str | None = Field(
|
||||
None, description="Base64-encoded Excel content when format='excel'"
|
||||
)
|
||||
format: str | None = Field(
|
||||
None, description="Export format used (json, csv, excel)"
|
||||
)
|
||||
|
||||
# Inherit versioning
|
||||
schema_version: str = Field("2.0", description="Response schema version")
|
||||
api_version: str = Field("v1", description="MCP API version")
|
||||
|
||||
|
||||
class GetChartPreviewRequest(QueryCacheControl):
|
||||
"""Request for chart preview with cache control."""
|
||||
|
||||
identifier: int | str = Field(description="Chart identifier (ID, UUID)")
|
||||
format: Literal["url", "ascii", "table", "base64", "vega_lite"] = Field(
|
||||
default="url",
|
||||
description=(
|
||||
"Preview format: 'url' for image URL, 'ascii' for text art, "
|
||||
"'table' for data table, 'base64' for embedded image, "
|
||||
"'vega_lite' for interactive JSON specification"
|
||||
),
|
||||
)
|
||||
width: int | None = Field(
|
||||
default=800,
|
||||
description="Preview image width in pixels (for url/base64 formats)",
|
||||
)
|
||||
height: int | None = Field(
|
||||
default=600,
|
||||
description="Preview image height in pixels (for url/base64 formats)",
|
||||
)
|
||||
ascii_width: int | None = Field(
|
||||
default=80, description="ASCII chart width in characters (for ascii format)"
|
||||
)
|
||||
ascii_height: int | None = Field(
|
||||
default=20, description="ASCII chart height in lines (for ascii format)"
|
||||
)
|
||||
|
||||
|
||||
# Discriminated union preview formats for type safety
|
||||
class URLPreview(BaseModel):
|
||||
"""URL-based image preview format."""
|
||||
|
||||
type: Literal["url"] = "url"
|
||||
preview_url: str = Field(..., description="Direct image URL")
|
||||
width: int = Field(..., description="Image width in pixels")
|
||||
height: int = Field(..., description="Image height in pixels")
|
||||
supports_interaction: bool = Field(
|
||||
False, description="Static image, no interaction"
|
||||
)
|
||||
|
||||
|
||||
class InteractivePreview(BaseModel):
|
||||
"""Interactive HTML preview with JavaScript controls."""
|
||||
|
||||
type: Literal["interactive"] = "interactive"
|
||||
html_content: str = Field(..., description="Embeddable HTML with Plotly/D3")
|
||||
preview_url: str = Field(..., description="Iframe-compatible URL")
|
||||
width: int = Field(..., description="Viewport width")
|
||||
height: int = Field(..., description="Viewport height")
|
||||
supports_pan: bool = Field(True, description="Supports pan interaction")
|
||||
supports_zoom: bool = Field(True, description="Supports zoom interaction")
|
||||
supports_hover: bool = Field(True, description="Supports hover details")
|
||||
|
||||
|
||||
class ASCIIPreview(BaseModel):
|
||||
"""ASCII art text representation."""
|
||||
|
||||
type: Literal["ascii"] = "ascii"
|
||||
ascii_content: str = Field(..., description="Unicode art representation")
|
||||
width: int = Field(..., description="Character width")
|
||||
height: int = Field(..., description="Line height")
|
||||
supports_color: bool = Field(False, description="Uses ANSI color codes")
|
||||
|
||||
|
||||
class VegaLitePreview(BaseModel):
|
||||
"""Vega-Lite grammar of graphics specification."""
|
||||
|
||||
type: Literal["vega_lite"] = "vega_lite"
|
||||
specification: Dict[str, Any] = Field(..., description="Vega-Lite JSON spec")
|
||||
data_url: str | None = Field(None, description="External data URL")
|
||||
supports_streaming: bool = Field(False, description="Supports live data updates")
|
||||
|
||||
|
||||
class TablePreview(BaseModel):
|
||||
"""Tabular data preview format."""
|
||||
|
||||
type: Literal["table"] = "table"
|
||||
table_data: str = Field(..., description="Formatted table content")
|
||||
row_count: int = Field(..., description="Number of rows displayed")
|
||||
supports_sorting: bool = Field(False, description="Table supports sorting")
|
||||
|
||||
|
||||
# Modern discriminated union using | syntax
|
||||
ChartPreviewContent = Annotated[
|
||||
URLPreview | InteractivePreview | ASCIIPreview | VegaLitePreview | TablePreview,
|
||||
Field(discriminator="type"),
|
||||
]
|
||||
|
||||
|
||||
class GenerateChartResponse(BaseModel):
|
||||
"""Comprehensive chart creation response with rich metadata."""
|
||||
|
||||
# Core chart information
|
||||
chart: ChartInfo | None = Field(None, description="Complete chart metadata")
|
||||
|
||||
# Multiple preview formats available
|
||||
previews: Dict[str, ChartPreviewContent] = Field(
|
||||
default_factory=dict,
|
||||
description="Available preview formats keyed by format type",
|
||||
)
|
||||
|
||||
# LLM-friendly capabilities
|
||||
capabilities: ChartCapabilities | None = Field(
|
||||
None, description="Chart interaction capabilities"
|
||||
)
|
||||
semantics: ChartSemantics | None = Field(
|
||||
None, description="Semantic chart understanding"
|
||||
)
|
||||
|
||||
# Navigation and context
|
||||
explore_url: str | None = Field(None, description="Edit chart in Superset")
|
||||
embed_code: str | None = Field(None, description="HTML embed snippet")
|
||||
api_endpoints: Dict[str, str] = Field(
|
||||
default_factory=dict, description="Related API endpoints for data/updates"
|
||||
)
|
||||
|
||||
# Performance and accessibility
|
||||
performance: PerformanceMetadata | None = Field(
|
||||
None, description="Performance metrics"
|
||||
)
|
||||
accessibility: AccessibilityMetadata | None = Field(
|
||||
None, description="Accessibility info"
|
||||
)
|
||||
|
||||
# Success/error handling
|
||||
success: bool = Field(True, description="Whether chart creation succeeded")
|
||||
error: ChartGenerationError | None = Field(
|
||||
None, description="Error details if creation failed"
|
||||
)
|
||||
warnings: List[str] = Field(default_factory=list, description="Non-fatal warnings")
|
||||
|
||||
# Inherit versioning
|
||||
schema_version: str = Field("2.0", description="Response schema version")
|
||||
api_version: str = Field("v1", description="MCP API version")
|
||||
|
||||
|
||||
class ChartPreview(BaseModel):
|
||||
"""Enhanced chart preview with discriminated union content."""
|
||||
|
||||
chart_id: int
|
||||
chart_name: str
|
||||
chart_type: str = Field(description="Type of chart visualization")
|
||||
explore_url: str = Field(description="URL to open chart in Superset for editing")
|
||||
|
||||
# Type-safe preview content
|
||||
content: ChartPreviewContent = Field(
|
||||
description="Preview content in requested format"
|
||||
)
|
||||
|
||||
# Rich metadata
|
||||
chart_description: str = Field(
|
||||
description="Human-readable description of the chart"
|
||||
)
|
||||
accessibility: AccessibilityMetadata = Field(
|
||||
description="Accessibility information"
|
||||
)
|
||||
performance: PerformanceMetadata = Field(description="Performance metrics")
|
||||
|
||||
# Backward compatibility fields (populated based on content type)
|
||||
format: str | None = Field(
|
||||
None, description="Format of the preview (url, ascii, table, base64)"
|
||||
)
|
||||
preview_url: str | None = Field(None, description="Image URL for 'url' format")
|
||||
ascii_chart: str | None = Field(
|
||||
None, description="ASCII art chart for 'ascii' format"
|
||||
)
|
||||
table_data: str | None = Field(
|
||||
None, description="Formatted table for 'table' format"
|
||||
)
|
||||
width: int | None = Field(
|
||||
None, description="Width (pixels for images, characters for ASCII)"
|
||||
)
|
||||
height: int | None = Field(
|
||||
None, description="Height (pixels for images, lines for ASCII)"
|
||||
)
|
||||
|
||||
# Inherit versioning
|
||||
schema_version: str = Field("2.0", description="Response schema version")
|
||||
api_version: str = Field("v1", description="MCP API version")
|
||||
|
||||
@@ -15,10 +15,22 @@
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
from .generate_chart import generate_chart
|
||||
from .get_chart_available_filters import get_chart_available_filters
|
||||
from .get_chart_data import get_chart_data
|
||||
from .get_chart_info import get_chart_info
|
||||
from .get_chart_preview import get_chart_preview
|
||||
from .list_charts import list_charts
|
||||
from .update_chart import update_chart
|
||||
from .update_chart_preview import update_chart_preview
|
||||
|
||||
__all__ = [
|
||||
"list_charts",
|
||||
"get_chart_info",
|
||||
"get_chart_available_filters",
|
||||
"generate_chart",
|
||||
"update_chart",
|
||||
"update_chart_preview",
|
||||
"get_chart_preview",
|
||||
"get_chart_data",
|
||||
]
|
||||
|
||||
462
superset/mcp_service/chart/tool/generate_chart.py
Normal file
462
superset/mcp_service/chart/tool/generate_chart.py
Normal file
@@ -0,0 +1,462 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
"""
|
||||
MCP tool: generate_chart (simplified schema)
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import has_dataset_access, mcp_auth_hook
|
||||
from superset.mcp_service.chart.chart_utils import (
|
||||
analyze_chart_capabilities,
|
||||
analyze_chart_semantics,
|
||||
generate_chart_name,
|
||||
map_config_to_form_data,
|
||||
)
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
AccessibilityMetadata,
|
||||
GenerateChartRequest,
|
||||
GenerateChartResponse,
|
||||
PerformanceMetadata,
|
||||
URLPreview,
|
||||
)
|
||||
from superset.mcp_service.utils.url_utils import (
|
||||
get_chart_screenshot_url,
|
||||
get_superset_base_url,
|
||||
)
|
||||
from superset.utils import json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def generate_chart( # noqa: C901
|
||||
request: GenerateChartRequest, ctx: Context
|
||||
) -> GenerateChartResponse:
|
||||
"""Create and save a chart in Superset.
|
||||
|
||||
IMPORTANT BEHAVIOR:
|
||||
- Charts ARE saved by default (save_chart=True)
|
||||
- Set save_chart=False for temporary preview only
|
||||
- LLM clients MUST display returned chart URL to users
|
||||
- Embed preview_url as image: 
|
||||
|
||||
VALIDATION:
|
||||
- 5-layer pipeline: Schema, business logic, dataset, Superset compatibility, runtime
|
||||
- XSS/SQL injection prevention
|
||||
- Column existence validation with fuzzy match suggestions
|
||||
- Aggregate function type compatibility checking
|
||||
|
||||
Returns:
|
||||
- Chart ID and metadata (if saved)
|
||||
- Preview URL and explore URL
|
||||
- Detailed validation errors with suggestions
|
||||
"""
|
||||
start_time = time.time()
|
||||
await ctx.info(
|
||||
"Starting chart generation: dataset_id=%s, chart_type=%s, "
|
||||
"save_chart=%s, preview_formats=%s"
|
||||
% (
|
||||
request.dataset_id,
|
||||
request.config.chart_type,
|
||||
request.save_chart,
|
||||
request.preview_formats,
|
||||
)
|
||||
)
|
||||
await ctx.debug(
|
||||
"Chart configuration details: config=%s" % (request.config.model_dump(),)
|
||||
)
|
||||
|
||||
try:
|
||||
# Run comprehensive validation pipeline
|
||||
await ctx.report_progress(1, 5, "Running validation pipeline")
|
||||
await ctx.debug(
|
||||
"Validating chart request: dataset_id=%s" % (request.dataset_id,)
|
||||
)
|
||||
from superset.mcp_service.chart.validation import ValidationPipeline
|
||||
|
||||
is_valid, parsed_request, validation_error = (
|
||||
ValidationPipeline.validate_request(request.model_dump())
|
||||
)
|
||||
if is_valid and parsed_request is not None:
|
||||
# Use the validated request going forward
|
||||
request = parsed_request
|
||||
if not is_valid:
|
||||
execution_time = int((time.time() - start_time) * 1000)
|
||||
assert validation_error is not None # Type narrowing for mypy
|
||||
await ctx.error(
|
||||
"Chart validation failed: error=%s" % (validation_error.model_dump(),)
|
||||
)
|
||||
return GenerateChartResponse.model_validate(
|
||||
{
|
||||
"chart": None,
|
||||
"error": validation_error.model_dump(),
|
||||
"performance": {
|
||||
"query_duration_ms": execution_time,
|
||||
"cache_status": "error",
|
||||
"optimization_suggestions": [],
|
||||
},
|
||||
"success": False,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
)
|
||||
|
||||
# Map the simplified config to Superset's form_data format
|
||||
form_data = map_config_to_form_data(request.config)
|
||||
|
||||
chart = None
|
||||
chart_id = None
|
||||
explore_url = None
|
||||
form_data_key = None
|
||||
|
||||
# Save chart by default (unless save_chart=False)
|
||||
if request.save_chart:
|
||||
await ctx.report_progress(2, 5, "Creating chart in database")
|
||||
from superset.commands.chart.create import CreateChartCommand
|
||||
|
||||
# Generate a chart name
|
||||
chart_name = generate_chart_name(request.config)
|
||||
await ctx.debug("Generated chart name: chart_name=%s" % (chart_name,))
|
||||
|
||||
# Find the dataset to get its numeric ID
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
|
||||
await ctx.debug("Looking up dataset: dataset_id=%s" % (request.dataset_id,))
|
||||
dataset = None
|
||||
if isinstance(request.dataset_id, int) or (
|
||||
isinstance(request.dataset_id, str) and request.dataset_id.isdigit()
|
||||
):
|
||||
dataset_id = (
|
||||
int(request.dataset_id)
|
||||
if isinstance(request.dataset_id, str)
|
||||
else request.dataset_id
|
||||
)
|
||||
dataset = DatasetDAO.find_by_id(dataset_id)
|
||||
# SECURITY FIX: Also validate permissions for numeric ID access
|
||||
if dataset and not has_dataset_access(dataset):
|
||||
logger.warning(
|
||||
"User %s attempted to access dataset %s without permission",
|
||||
ctx.user.username if hasattr(ctx, "user") else "unknown",
|
||||
dataset_id,
|
||||
)
|
||||
dataset = None # Treat as not found
|
||||
else:
|
||||
# SECURITY FIX: Try UUID lookup with permission validation
|
||||
dataset = DatasetDAO.find_by_id(request.dataset_id, id_column="uuid")
|
||||
# Validate permissions for UUID-based access
|
||||
if dataset and not has_dataset_access(dataset):
|
||||
logger.warning(
|
||||
"User %s attempted access dataset %s via UUID",
|
||||
ctx.user.username if hasattr(ctx, "user") else "unknown",
|
||||
request.dataset_id,
|
||||
)
|
||||
dataset = None # Treat as not found
|
||||
|
||||
if not dataset:
|
||||
await ctx.error(
|
||||
"Dataset not found: dataset_id=%s" % (request.dataset_id,)
|
||||
)
|
||||
from superset.mcp_service.common.error_schemas import (
|
||||
ChartGenerationError,
|
||||
)
|
||||
|
||||
execution_time = int((time.time() - start_time) * 1000)
|
||||
error = ChartGenerationError(
|
||||
error_type="dataset_not_found",
|
||||
message=f"Dataset not found: {request.dataset_id}",
|
||||
details=(
|
||||
f"No dataset found with identifier '{request.dataset_id}'. "
|
||||
f"This could be an invalid ID/UUID or a permissions issue."
|
||||
),
|
||||
suggestions=[
|
||||
"Verify the dataset ID or UUID is correct",
|
||||
"Check that you have access to this dataset",
|
||||
"Use the list_datasets tool to find available datasets",
|
||||
"If using UUID, ensure it's the correct format",
|
||||
],
|
||||
error_code="DATASET_NOT_FOUND",
|
||||
)
|
||||
return GenerateChartResponse.model_validate(
|
||||
{
|
||||
"chart": None,
|
||||
"error": error.model_dump(),
|
||||
"performance": {
|
||||
"query_duration_ms": execution_time,
|
||||
"cache_status": "error",
|
||||
"optimization_suggestions": [],
|
||||
},
|
||||
"success": False,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
command = CreateChartCommand(
|
||||
{
|
||||
"slice_name": chart_name,
|
||||
"viz_type": form_data["viz_type"],
|
||||
"datasource_id": dataset.id,
|
||||
"datasource_type": "table",
|
||||
"params": json.dumps(form_data),
|
||||
}
|
||||
)
|
||||
|
||||
chart = command.run()
|
||||
chart_id = chart.id
|
||||
|
||||
# Ensure chart was created successfully before committing
|
||||
if not chart or not chart.id:
|
||||
raise Exception("Chart creation failed - no chart ID returned")
|
||||
|
||||
await ctx.info(
|
||||
"Chart created successfully: chart_id=%s, chart_name=%s"
|
||||
% (
|
||||
chart.id,
|
||||
chart.slice_name,
|
||||
)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Chart creation failed: %s", e)
|
||||
await ctx.error("Chart creation failed: error=%s" % (str(e),))
|
||||
raise
|
||||
# Update explore URL to use saved chart
|
||||
explore_url = f"{get_superset_base_url()}/explore/?slice_id={chart.id}"
|
||||
else:
|
||||
await ctx.report_progress(2, 5, "Generating temporary chart preview")
|
||||
# Generate explore link with cached form_data for preview-only mode
|
||||
from superset.mcp_service.chart.chart_utils import generate_explore_link
|
||||
|
||||
explore_url = generate_explore_link(request.dataset_id, form_data)
|
||||
await ctx.debug("Generated explore link: explore_url=%s" % (explore_url,))
|
||||
|
||||
# Extract form_data_key from the explore URL
|
||||
if explore_url and "form_data_key=" in explore_url:
|
||||
form_data_key = explore_url.split("form_data_key=")[1].split("&")[0]
|
||||
|
||||
# Generate semantic analysis
|
||||
capabilities = analyze_chart_capabilities(chart, request.config)
|
||||
semantics = analyze_chart_semantics(chart, request.config)
|
||||
|
||||
# Create performance metadata
|
||||
execution_time = int((time.time() - start_time) * 1000)
|
||||
performance = PerformanceMetadata(
|
||||
query_duration_ms=execution_time,
|
||||
cache_status="miss",
|
||||
optimization_suggestions=[],
|
||||
)
|
||||
|
||||
# Create accessibility metadata
|
||||
chart_name = (
|
||||
chart.slice_name
|
||||
if chart and hasattr(chart, "slice_name")
|
||||
else generate_chart_name(request.config)
|
||||
)
|
||||
accessibility = AccessibilityMetadata(
|
||||
color_blind_safe=True, # Would need actual analysis
|
||||
alt_text=f"Chart showing {chart_name}",
|
||||
high_contrast_available=False,
|
||||
)
|
||||
|
||||
# Generate previews if requested
|
||||
await ctx.report_progress(3, 5, "Generating chart previews")
|
||||
previews = {}
|
||||
if request.generate_preview:
|
||||
await ctx.debug(
|
||||
"Generating previews: formats=%s" % (str(request.preview_formats),)
|
||||
)
|
||||
try:
|
||||
for format_type in request.preview_formats:
|
||||
await ctx.debug(
|
||||
"Processing preview format: format=%s" % (format_type,)
|
||||
)
|
||||
# Skip base64 format - we never return base64
|
||||
if format_type == "base64":
|
||||
logger.info("Skipping base64 format - not supported")
|
||||
continue
|
||||
|
||||
if chart_id:
|
||||
# For saved charts, use the existing preview generation
|
||||
from superset.mcp_service.chart.tool.get_chart_preview import (
|
||||
_get_chart_preview_internal,
|
||||
GetChartPreviewRequest,
|
||||
)
|
||||
|
||||
preview_request = GetChartPreviewRequest(
|
||||
identifier=str(chart_id), format=format_type
|
||||
)
|
||||
preview_result = await _get_chart_preview_internal(
|
||||
preview_request, ctx
|
||||
)
|
||||
|
||||
if hasattr(preview_result, "content"):
|
||||
previews[format_type] = preview_result.content
|
||||
else:
|
||||
# For preview-only mode (save_chart=false)
|
||||
if format_type == "url" and form_data_key:
|
||||
# Generate screenshot URL using centralized helper
|
||||
from superset.mcp_service.utils.url_utils import (
|
||||
get_explore_screenshot_url,
|
||||
)
|
||||
|
||||
preview_url = get_explore_screenshot_url(form_data_key)
|
||||
previews[format_type] = URLPreview(
|
||||
preview_url=preview_url,
|
||||
width=800,
|
||||
height=600,
|
||||
supports_interaction=False,
|
||||
)
|
||||
elif format_type in ["ascii", "table", "vega_lite"]:
|
||||
# Generate preview from form data without saved chart
|
||||
from superset.mcp_service.chart.preview_utils import (
|
||||
generate_preview_from_form_data,
|
||||
)
|
||||
|
||||
# Convert dataset_id to int only if it's numeric
|
||||
if (
|
||||
isinstance(request.dataset_id, str)
|
||||
and request.dataset_id.isdigit()
|
||||
):
|
||||
dataset_id_for_preview = int(request.dataset_id)
|
||||
elif isinstance(request.dataset_id, int):
|
||||
dataset_id_for_preview = request.dataset_id
|
||||
else:
|
||||
# Skip preview generation for non-numeric dataset IDs
|
||||
logger.warning(
|
||||
"Cannot generate preview for non-numeric "
|
||||
)
|
||||
continue
|
||||
|
||||
preview_result = generate_preview_from_form_data(
|
||||
form_data=form_data,
|
||||
dataset_id=dataset_id_for_preview,
|
||||
preview_format=format_type,
|
||||
)
|
||||
|
||||
if not hasattr(preview_result, "error"):
|
||||
previews[format_type] = preview_result
|
||||
|
||||
except Exception as e:
|
||||
# Log warning but don't fail the entire request
|
||||
await ctx.warning("Preview generation failed: error=%s" % (str(e),))
|
||||
logger.warning("Preview generation failed: %s", e)
|
||||
|
||||
# Return enhanced data while maintaining backward compatibility
|
||||
await ctx.report_progress(4, 5, "Building response")
|
||||
result = {
|
||||
"chart": {
|
||||
"id": chart.id if chart else None,
|
||||
"slice_name": chart.slice_name
|
||||
if chart
|
||||
else generate_chart_name(request.config),
|
||||
"viz_type": chart.viz_type if chart else form_data.get("viz_type"),
|
||||
"url": explore_url,
|
||||
"uuid": str(chart.uuid) if chart and chart.uuid else None,
|
||||
"saved": request.save_chart,
|
||||
}
|
||||
if request.save_chart and chart
|
||||
else None,
|
||||
"error": None,
|
||||
# Enhanced fields for better LLM integration
|
||||
"previews": previews,
|
||||
"capabilities": capabilities.model_dump() if capabilities else None,
|
||||
"semantics": semantics.model_dump() if semantics else None,
|
||||
"explore_url": explore_url,
|
||||
"form_data_key": form_data_key,
|
||||
"api_endpoints": {
|
||||
"data": f"{get_superset_base_url()}/api/v1/chart/{chart.id}/data/"
|
||||
if chart
|
||||
else None,
|
||||
"preview": get_chart_screenshot_url(chart.id) if chart else None,
|
||||
"export": f"{get_superset_base_url()}/api/v1/chart/{chart.id}/export/"
|
||||
if chart
|
||||
else None,
|
||||
}
|
||||
if chart
|
||||
else {},
|
||||
"performance": performance.model_dump() if performance else None,
|
||||
"accessibility": accessibility.model_dump() if accessibility else None,
|
||||
"success": True,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
await ctx.report_progress(5, 5, "Chart generation completed")
|
||||
await ctx.info(
|
||||
"Chart generation completed successfully: chart_id=%s, execution_time_ms=%s"
|
||||
% (
|
||||
chart.id if chart else None,
|
||||
int((time.time() - start_time) * 1000),
|
||||
)
|
||||
)
|
||||
return GenerateChartResponse.model_validate(result)
|
||||
|
||||
except Exception as e:
|
||||
await ctx.error(
|
||||
"Chart generation failed: error=%s, execution_time_ms=%s"
|
||||
% (
|
||||
str(e),
|
||||
int((time.time() - start_time) * 1000),
|
||||
)
|
||||
)
|
||||
from superset.mcp_service.utils.error_builder import ChartErrorBuilder
|
||||
|
||||
logger.exception("Chart generation failed: %s", str(e))
|
||||
|
||||
# Extract chart_type from different sources for better error context
|
||||
chart_type = "unknown"
|
||||
try:
|
||||
if hasattr(request, "config") and hasattr(request.config, "chart_type"):
|
||||
chart_type = request.config.chart_type
|
||||
except Exception as extract_error:
|
||||
# Ignore errors when extracting chart type for error context
|
||||
logger.debug("Could not extract chart type: %s", extract_error)
|
||||
|
||||
execution_time = int((time.time() - start_time) * 1000)
|
||||
|
||||
# Build standardized error response
|
||||
error = ChartErrorBuilder.build_error(
|
||||
error_type="chart_generation_error",
|
||||
template_key="generation_failed",
|
||||
template_vars={
|
||||
"reason": str(e),
|
||||
"dataset_id": str(request.dataset_id),
|
||||
"chart_type": chart_type,
|
||||
},
|
||||
error_code="CHART_GENERATION_FAILED",
|
||||
)
|
||||
|
||||
return GenerateChartResponse.model_validate(
|
||||
{
|
||||
"chart": None,
|
||||
"error": error.model_dump(),
|
||||
"performance": {
|
||||
"query_duration_ms": execution_time,
|
||||
"cache_status": "error",
|
||||
"optimization_suggestions": [],
|
||||
},
|
||||
"success": False,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,50 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
MCP tool: get_chart_available_filters
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
ChartAvailableFiltersResponse,
|
||||
GetChartAvailableFiltersRequest,
|
||||
)
|
||||
from superset.mcp_service.mcp_core import ModelGetAvailableFiltersCore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
def get_chart_available_filters(
|
||||
request: GetChartAvailableFiltersRequest, ctx: Context
|
||||
) -> ChartAvailableFiltersResponse:
|
||||
"""Get available chart filter fields and operators."""
|
||||
from superset.daos.chart import ChartDAO
|
||||
|
||||
tool = ModelGetAvailableFiltersCore(
|
||||
dao_class=ChartDAO,
|
||||
output_schema=ChartAvailableFiltersResponse,
|
||||
logger=logger,
|
||||
)
|
||||
return tool.run_tool()
|
||||
656
superset/mcp_service/chart/tool/get_chart_data.py
Normal file
656
superset/mcp_service/chart/tool/get_chart_data.py
Normal file
@@ -0,0 +1,656 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
MCP tool: get_chart_data
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List, TYPE_CHECKING
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from superset.models.slice import Slice
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
ChartData,
|
||||
ChartError,
|
||||
DataColumn,
|
||||
GetChartDataRequest,
|
||||
PerformanceMetadata,
|
||||
)
|
||||
from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def get_chart_data( # noqa: C901
|
||||
request: GetChartDataRequest, ctx: Context
|
||||
) -> ChartData | ChartError:
|
||||
"""Get chart data by ID or UUID.
|
||||
|
||||
Returns the actual data behind a chart for LLM analysis without image rendering.
|
||||
|
||||
Supports:
|
||||
- Numeric ID or UUID lookup
|
||||
- Multiple formats: json, csv, excel
|
||||
- Cache control: use_cache, force_refresh, cache_timeout
|
||||
|
||||
Returns underlying data in requested format with cache status.
|
||||
"""
|
||||
await ctx.info(
|
||||
"Starting chart data retrieval: identifier=%s, format=%s, limit=%s"
|
||||
% (
|
||||
request.identifier,
|
||||
request.format,
|
||||
request.limit,
|
||||
)
|
||||
)
|
||||
await ctx.debug(
|
||||
"Cache settings: use_cache=%s, force_refresh=%s, cache_timeout=%s"
|
||||
% (
|
||||
request.use_cache,
|
||||
request.force_refresh,
|
||||
request.cache_timeout,
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
await ctx.report_progress(1, 4, "Looking up chart")
|
||||
from superset.daos.chart import ChartDAO
|
||||
from superset.utils import json as utils_json
|
||||
|
||||
# Find the chart
|
||||
chart = None
|
||||
if isinstance(request.identifier, int) or (
|
||||
isinstance(request.identifier, str) and request.identifier.isdigit()
|
||||
):
|
||||
chart_id = (
|
||||
int(request.identifier)
|
||||
if isinstance(request.identifier, str)
|
||||
else request.identifier
|
||||
)
|
||||
await ctx.debug(
|
||||
"Performing ID-based chart lookup: chart_id=%s" % (chart_id,)
|
||||
)
|
||||
chart = ChartDAO.find_by_id(chart_id)
|
||||
else:
|
||||
await ctx.debug(
|
||||
"Performing UUID-based chart lookup: uuid=%s" % (request.identifier,)
|
||||
)
|
||||
# Try UUID lookup using DAO flexible method
|
||||
chart = ChartDAO.find_by_id(request.identifier, id_column="uuid")
|
||||
|
||||
if not chart:
|
||||
await ctx.error("Chart not found: identifier=%s" % (request.identifier,))
|
||||
return ChartError(
|
||||
error=f"No chart found with identifier: {request.identifier}",
|
||||
error_type="NotFound",
|
||||
)
|
||||
|
||||
await ctx.info(
|
||||
"Chart found successfully: chart_id=%s, chart_name=%s, viz_type=%s"
|
||||
% (
|
||||
chart.id,
|
||||
chart.slice_name,
|
||||
chart.viz_type,
|
||||
)
|
||||
)
|
||||
logger.info("Getting data for chart %s: %s", chart.id, chart.slice_name)
|
||||
|
||||
import time
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
await ctx.report_progress(2, 4, "Preparing data query")
|
||||
# Get chart data using the existing API
|
||||
from superset.commands.chart.data.get_data_command import ChartDataCommand
|
||||
from superset.common.query_context_factory import QueryContextFactory
|
||||
|
||||
# Parse the form_data to get query context
|
||||
form_data = utils_json.loads(chart.params) if chart.params else {}
|
||||
await ctx.debug(
|
||||
"Chart form data parsed: has_filters=%s, has_groupby=%s, has_metrics=%s"
|
||||
% (
|
||||
bool(form_data.get("filters")),
|
||||
bool(form_data.get("groupby")),
|
||||
bool(form_data.get("metrics")),
|
||||
)
|
||||
)
|
||||
|
||||
# Create a proper QueryContext using the factory with cache control
|
||||
factory = QueryContextFactory()
|
||||
query_context = factory.create(
|
||||
datasource={"id": chart.datasource_id, "type": chart.datasource_type},
|
||||
queries=[
|
||||
{
|
||||
"filters": form_data.get("filters", []),
|
||||
"columns": form_data.get("groupby", []),
|
||||
"metrics": form_data.get("metrics", []),
|
||||
"row_limit": request.limit or 100,
|
||||
"order_desc": True,
|
||||
# Apply cache control from request
|
||||
"cache_timeout": request.cache_timeout,
|
||||
}
|
||||
],
|
||||
form_data=form_data,
|
||||
# Use cache unless force_refresh is True
|
||||
force=request.force_refresh,
|
||||
)
|
||||
|
||||
await ctx.report_progress(3, 4, "Executing data query")
|
||||
await ctx.debug(
|
||||
"Query execution parameters: datasource_id=%s, datasource_type=%s, "
|
||||
"row_limit=%s, force_refresh=%s"
|
||||
% (
|
||||
chart.datasource_id,
|
||||
chart.datasource_type,
|
||||
request.limit or 100,
|
||||
request.force_refresh,
|
||||
)
|
||||
)
|
||||
|
||||
# Execute the query
|
||||
command = ChartDataCommand(query_context)
|
||||
result = command.run()
|
||||
|
||||
# Handle empty query results for certain chart types
|
||||
if not result or ("queries" not in result) or len(result["queries"]) == 0:
|
||||
await ctx.warning(
|
||||
"Empty query results: chart_id=%s, chart_type=%s"
|
||||
% (chart.id, chart.viz_type)
|
||||
)
|
||||
return ChartError(
|
||||
error=f"No query results returned for chart {chart.id}. "
|
||||
f"This may occur with chart types like big_number.",
|
||||
error_type="EmptyQuery",
|
||||
)
|
||||
|
||||
# Extract data from result (we've already validated it exists above)
|
||||
query_result = result["queries"][0]
|
||||
data = query_result.get("data", [])
|
||||
raw_columns = query_result.get("colnames", [])
|
||||
|
||||
await ctx.debug(
|
||||
"Query results received: row_count=%s, column_count=%s, "
|
||||
"has_cache_key=%s"
|
||||
% (
|
||||
len(data),
|
||||
len(raw_columns),
|
||||
bool(query_result.get("cache_key")),
|
||||
)
|
||||
)
|
||||
|
||||
# Check if we have data to work with
|
||||
if not data:
|
||||
await ctx.warning("No data in query results: chart_id=%s" % (chart.id,))
|
||||
return ChartError(
|
||||
error=f"No data available for chart {chart.id}", error_type="NoData"
|
||||
)
|
||||
|
||||
# Create rich column metadata
|
||||
columns = []
|
||||
for col_name in raw_columns:
|
||||
# Sample some values for metadata
|
||||
sample_values = [
|
||||
row.get(col_name)
|
||||
for row in data[:3]
|
||||
if row.get(col_name) is not None
|
||||
]
|
||||
|
||||
# Infer data type
|
||||
data_type = "string"
|
||||
if sample_values:
|
||||
if all(isinstance(v, (int, float)) for v in sample_values):
|
||||
data_type = "numeric"
|
||||
elif all(isinstance(v, bool) for v in sample_values):
|
||||
data_type = "boolean"
|
||||
|
||||
columns.append(
|
||||
DataColumn(
|
||||
name=col_name,
|
||||
display_name=col_name.replace("_", " ").title(),
|
||||
data_type=data_type,
|
||||
sample_values=sample_values[:3],
|
||||
null_count=sum(1 for row in data if row.get(col_name) is None),
|
||||
unique_count=len({str(row.get(col_name)) for row in data}),
|
||||
)
|
||||
)
|
||||
|
||||
# Cache status information using utility function
|
||||
cache_status = get_cache_status_from_result(
|
||||
query_result, force_refresh=request.force_refresh
|
||||
)
|
||||
|
||||
# Generate insights and recommendations
|
||||
insights = []
|
||||
if len(data) > 100:
|
||||
insights.append(
|
||||
"Large dataset - consider filtering for better performance"
|
||||
)
|
||||
if len(raw_columns) > 10:
|
||||
insights.append("Many columns available - focus on key metrics")
|
||||
|
||||
# Add cache-specific insights
|
||||
if cache_status.cache_hit:
|
||||
if (
|
||||
cache_status.cache_age_seconds
|
||||
and cache_status.cache_age_seconds > 3600
|
||||
):
|
||||
hours_old = cache_status.cache_age_seconds // 3600
|
||||
insights.append(
|
||||
f"Data is from cache ({hours_old}h old) - "
|
||||
"consider refreshing for latest data"
|
||||
)
|
||||
else:
|
||||
insights.append("Data served from cache for fast response")
|
||||
else:
|
||||
insights.append("Fresh data retrieved from database")
|
||||
|
||||
recommended_visualizations = []
|
||||
if any(
|
||||
"time" in col.lower() or "date" in col.lower() for col in raw_columns
|
||||
):
|
||||
recommended_visualizations.extend(["line chart", "time series"])
|
||||
if len(raw_columns) <= 3:
|
||||
recommended_visualizations.extend(["bar chart", "scatter plot"])
|
||||
|
||||
# Performance metadata with cache awareness
|
||||
execution_time = int((time.time() - start_time) * 1000)
|
||||
performance_status = (
|
||||
"cache_hit" if cache_status.cache_hit else "fresh_query"
|
||||
)
|
||||
optimization_suggestions = []
|
||||
|
||||
if not cache_status.cache_hit and execution_time > 5000:
|
||||
optimization_suggestions.append(
|
||||
"Consider using cache for this slow query"
|
||||
)
|
||||
elif (
|
||||
cache_status.cache_hit
|
||||
and cache_status.cache_age_seconds
|
||||
and cache_status.cache_age_seconds > 86400
|
||||
):
|
||||
optimization_suggestions.append("Cache is old - consider refreshing")
|
||||
|
||||
performance = PerformanceMetadata(
|
||||
query_duration_ms=execution_time,
|
||||
cache_status=performance_status,
|
||||
optimization_suggestions=optimization_suggestions,
|
||||
)
|
||||
|
||||
# Generate comprehensive summary with cache info
|
||||
cache_info = ""
|
||||
if cache_status.cache_hit:
|
||||
age_info = (
|
||||
f" (cached {cache_status.cache_age_seconds // 60}m ago)"
|
||||
if cache_status.cache_age_seconds
|
||||
else " (cached)"
|
||||
)
|
||||
cache_info = age_info
|
||||
|
||||
summary_parts = [
|
||||
f"Chart '{chart.slice_name}' ({chart.viz_type})",
|
||||
f"Contains {len(data)} rows across {len(raw_columns)} columns"
|
||||
f"{cache_info}",
|
||||
]
|
||||
|
||||
if data and len(data) > 0:
|
||||
summary_parts.append(
|
||||
f"Sample data includes: {', '.join(raw_columns[:3])}"
|
||||
)
|
||||
|
||||
summary = ". ".join(summary_parts)
|
||||
|
||||
# Handle different export formats
|
||||
if request.format == "csv":
|
||||
return _export_data_as_csv(
|
||||
chart,
|
||||
data[: request.limit] if request.limit else data,
|
||||
raw_columns,
|
||||
cache_status,
|
||||
performance,
|
||||
)
|
||||
elif request.format == "excel":
|
||||
return _export_data_as_excel(
|
||||
chart,
|
||||
data[: request.limit] if request.limit else data,
|
||||
raw_columns,
|
||||
cache_status,
|
||||
performance,
|
||||
)
|
||||
|
||||
await ctx.report_progress(4, 4, "Building response")
|
||||
|
||||
# Calculate data quality metrics
|
||||
data_completeness = 1.0 - (
|
||||
sum(col.null_count for col in columns)
|
||||
/ max(len(data) * len(columns), 1)
|
||||
)
|
||||
|
||||
await ctx.info(
|
||||
"Chart data retrieval completed successfully: chart_id=%s, "
|
||||
"rows_returned=%s, columns_returned=%s, execution_time_ms=%s, "
|
||||
"cache_hit=%s, data_completeness=%s"
|
||||
% (
|
||||
chart.id,
|
||||
len(data),
|
||||
len(raw_columns),
|
||||
execution_time,
|
||||
cache_status.cache_hit,
|
||||
round(data_completeness, 3),
|
||||
)
|
||||
)
|
||||
|
||||
# Default JSON format
|
||||
return ChartData(
|
||||
chart_id=chart.id,
|
||||
chart_name=chart.slice_name or f"Chart {chart.id}",
|
||||
chart_type=chart.viz_type or "unknown",
|
||||
columns=columns,
|
||||
data=data[: request.limit] if request.limit else data,
|
||||
row_count=len(data),
|
||||
total_rows=query_result.get("rowcount"),
|
||||
summary=summary,
|
||||
insights=insights,
|
||||
data_quality={"completeness": data_completeness},
|
||||
recommended_visualizations=recommended_visualizations,
|
||||
data_freshness=None, # Add missing field
|
||||
performance=performance,
|
||||
cache_status=cache_status,
|
||||
)
|
||||
|
||||
except Exception as data_error:
|
||||
await ctx.error(
|
||||
"Data retrieval failed: chart_id=%s, error=%s, error_type=%s"
|
||||
% (
|
||||
chart.id,
|
||||
str(data_error),
|
||||
type(data_error).__name__,
|
||||
)
|
||||
)
|
||||
logger.error("Data retrieval error for chart %s: %s", chart.id, data_error)
|
||||
return ChartError(
|
||||
error=f"Error retrieving chart data: {str(data_error)}",
|
||||
error_type="DataError",
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
await ctx.error(
|
||||
"Chart data retrieval failed: identifier=%s, error=%s, error_type=%s"
|
||||
% (
|
||||
request.identifier,
|
||||
str(e),
|
||||
type(e).__name__,
|
||||
)
|
||||
)
|
||||
logger.error("Error in get_chart_data: %s", e)
|
||||
return ChartError(
|
||||
error=f"Failed to get chart data: {str(e)}", error_type="InternalError"
|
||||
)
|
||||
|
||||
|
||||
def _export_data_as_csv(
|
||||
chart: "Slice",
|
||||
data: List[Dict[str, Any]],
|
||||
columns: List[str],
|
||||
cache_status: Any,
|
||||
performance: Any,
|
||||
) -> "ChartData":
|
||||
"""Export chart data as CSV format."""
|
||||
import csv
|
||||
import io
|
||||
|
||||
# Create CSV content
|
||||
output = io.StringIO()
|
||||
|
||||
if data and columns:
|
||||
writer = csv.DictWriter(output, fieldnames=columns)
|
||||
writer.writeheader()
|
||||
|
||||
# Write data rows
|
||||
for row in data:
|
||||
# Ensure all values are properly formatted for CSV
|
||||
csv_row = {}
|
||||
for col in columns:
|
||||
value = row.get(col, "")
|
||||
# Handle None values and convert to string
|
||||
if value is None:
|
||||
csv_row[col] = ""
|
||||
elif isinstance(value, (list, dict)):
|
||||
csv_row[col] = str(value)
|
||||
else:
|
||||
csv_row[col] = value
|
||||
writer.writerow(csv_row)
|
||||
|
||||
csv_content = output.getvalue()
|
||||
|
||||
# Return as ChartData with CSV content in a special field
|
||||
from superset.mcp_service.chart.schemas import ChartData
|
||||
|
||||
return ChartData(
|
||||
chart_id=chart.id,
|
||||
chart_name=chart.slice_name or f"Chart {chart.id}",
|
||||
chart_type=chart.viz_type or "unknown",
|
||||
columns=[], # Not needed for CSV export
|
||||
data=[], # CSV content is in csv_data field
|
||||
row_count=len(data),
|
||||
total_rows=len(data),
|
||||
summary=f"CSV export of chart '{chart.slice_name}' with {len(data)} rows",
|
||||
insights=[f"Data exported as CSV format ({len(csv_content)} characters)"],
|
||||
data_quality={},
|
||||
recommended_visualizations=[],
|
||||
data_freshness=None,
|
||||
performance=performance,
|
||||
cache_status=cache_status,
|
||||
# Store CSV content in data field as string for the response
|
||||
csv_data=csv_content,
|
||||
format="csv",
|
||||
)
|
||||
|
||||
|
||||
def _export_data_as_excel(
|
||||
chart: "Slice",
|
||||
data: List[Dict[str, Any]],
|
||||
columns: List[str],
|
||||
cache_status: Any,
|
||||
performance: Any,
|
||||
) -> "ChartData | ChartError":
|
||||
"""Export chart data as Excel format."""
|
||||
try:
|
||||
excel_b64 = _create_excel_with_openpyxl(chart, data, columns)
|
||||
return _create_excel_chart_data(
|
||||
chart, data, excel_b64, performance, cache_status
|
||||
)
|
||||
except ImportError:
|
||||
return _try_xlsxwriter_fallback(chart, data, columns, cache_status, performance)
|
||||
|
||||
|
||||
def _create_excel_with_openpyxl(
|
||||
chart: "Slice", data: List[Dict[str, Any]], columns: List[str]
|
||||
) -> str:
|
||||
"""Create Excel file using openpyxl."""
|
||||
import base64
|
||||
import io
|
||||
|
||||
from openpyxl import Workbook
|
||||
|
||||
wb = Workbook()
|
||||
ws = wb.active
|
||||
ws.title = chart.slice_name[:31] if chart.slice_name else "Chart Data"
|
||||
|
||||
if data and columns:
|
||||
_write_excel_headers(ws, columns)
|
||||
_write_excel_data(ws, data, columns)
|
||||
|
||||
output = io.BytesIO()
|
||||
wb.save(output)
|
||||
output.seek(0)
|
||||
return base64.b64encode(output.read()).decode()
|
||||
|
||||
|
||||
def _write_excel_headers(ws: Any, columns: List[str]) -> None:
|
||||
"""Write headers to Excel worksheet."""
|
||||
for idx, col in enumerate(columns, 1):
|
||||
ws.cell(row=1, column=idx, value=col)
|
||||
|
||||
|
||||
def _write_excel_data(ws: Any, data: List[Dict[str, Any]], columns: List[str]) -> None:
|
||||
"""Write data to Excel worksheet."""
|
||||
for row_idx, row in enumerate(data, 2):
|
||||
for col_idx, col in enumerate(columns, 1):
|
||||
value = row.get(col, "")
|
||||
if value is None:
|
||||
value = ""
|
||||
elif isinstance(value, (list, dict)):
|
||||
value = str(value)
|
||||
ws.cell(row=row_idx, column=col_idx, value=value)
|
||||
|
||||
|
||||
def _try_xlsxwriter_fallback(
|
||||
chart: "Slice",
|
||||
data: List[Dict[str, Any]],
|
||||
columns: List[str],
|
||||
cache_status: Any,
|
||||
performance: Any,
|
||||
) -> "ChartData | ChartError":
|
||||
"""Try xlsxwriter as fallback for Excel export."""
|
||||
try:
|
||||
excel_b64 = _create_excel_with_xlsxwriter(chart, data, columns)
|
||||
return _create_excel_chart_data_xlsxwriter(
|
||||
chart, data, excel_b64, performance, cache_status
|
||||
)
|
||||
except ImportError:
|
||||
from superset.mcp_service.chart.schemas import ChartError
|
||||
|
||||
return ChartError(
|
||||
error="Excel export requires openpyxl or xlsxwriter package",
|
||||
error_type="ExportError",
|
||||
)
|
||||
|
||||
|
||||
def _create_excel_with_xlsxwriter(
|
||||
chart: "Slice", data: List[Dict[str, Any]], columns: List[str]
|
||||
) -> str:
|
||||
"""Create Excel file using xlsxwriter."""
|
||||
import base64
|
||||
import io
|
||||
|
||||
import xlsxwriter
|
||||
|
||||
output = io.BytesIO()
|
||||
workbook = xlsxwriter.Workbook(output, {"in_memory": True})
|
||||
sheet_name = chart.slice_name[:31] if chart.slice_name else "Chart Data"
|
||||
worksheet = workbook.add_worksheet(sheet_name)
|
||||
|
||||
if data and columns:
|
||||
_write_xlsxwriter_data(worksheet, data, columns)
|
||||
|
||||
workbook.close()
|
||||
output.seek(0)
|
||||
return base64.b64encode(output.read()).decode()
|
||||
|
||||
|
||||
def _write_xlsxwriter_data(
|
||||
worksheet: Any, data: List[Dict[str, Any]], columns: List[str]
|
||||
) -> None:
|
||||
"""Write data to xlsxwriter worksheet."""
|
||||
# Write headers
|
||||
for idx, col in enumerate(columns):
|
||||
worksheet.write(0, idx, col)
|
||||
|
||||
# Write data
|
||||
for row_idx, row in enumerate(data):
|
||||
for col_idx, col in enumerate(columns):
|
||||
value = row.get(col, "")
|
||||
if value is None:
|
||||
value = ""
|
||||
elif isinstance(value, (list, dict)):
|
||||
value = str(value)
|
||||
worksheet.write(row_idx + 1, col_idx, value)
|
||||
|
||||
|
||||
def _create_excel_chart_data(
|
||||
chart: "Slice",
|
||||
data: List[Dict[str, Any]],
|
||||
excel_b64: str,
|
||||
performance: Any,
|
||||
cache_status: Any,
|
||||
) -> "ChartData":
|
||||
"""Create ChartData response for Excel export (openpyxl)."""
|
||||
from superset.mcp_service.chart.schemas import ChartData
|
||||
|
||||
chart_name = chart.slice_name or f"Chart {chart.id}"
|
||||
summary = f"Excel export of chart '{chart.slice_name}' with {len(data)} rows"
|
||||
|
||||
return ChartData(
|
||||
chart_id=chart.id,
|
||||
chart_name=chart_name,
|
||||
chart_type=chart.viz_type or "unknown",
|
||||
columns=[],
|
||||
data=[],
|
||||
row_count=len(data),
|
||||
total_rows=len(data),
|
||||
summary=summary,
|
||||
insights=["Data exported as Excel format (base64 encoded)"],
|
||||
data_quality={},
|
||||
recommended_visualizations=[],
|
||||
data_freshness=None,
|
||||
performance=performance,
|
||||
cache_status=cache_status,
|
||||
excel_data=excel_b64,
|
||||
format="excel",
|
||||
)
|
||||
|
||||
|
||||
def _create_excel_chart_data_xlsxwriter(
|
||||
chart: "Slice",
|
||||
data: List[Dict[str, Any]],
|
||||
excel_b64: str,
|
||||
performance: Any,
|
||||
cache_status: Any,
|
||||
) -> "ChartData":
|
||||
"""Create ChartData response for Excel export (xlsxwriter)."""
|
||||
from superset.mcp_service.chart.schemas import ChartData
|
||||
|
||||
chart_name = chart.slice_name or f"Chart {chart.id}"
|
||||
summary = f"Excel export of chart '{chart.slice_name}' with {len(data)} rows"
|
||||
|
||||
return ChartData(
|
||||
chart_id=chart.id,
|
||||
chart_name=chart_name,
|
||||
chart_type=chart.viz_type or "unknown",
|
||||
columns=[],
|
||||
data=[],
|
||||
row_count=len(data),
|
||||
total_rows=len(data),
|
||||
summary=summary,
|
||||
insights=["Data exported as Excel format (base64 encoded, xlsxwriter)"],
|
||||
data_quality={},
|
||||
recommended_visualizations=[],
|
||||
data_freshness=None,
|
||||
performance=performance,
|
||||
cache_status=cache_status,
|
||||
excel_data=excel_b64,
|
||||
format="excel",
|
||||
)
|
||||
@@ -43,6 +43,14 @@ async def get_chart_info(
|
||||
) -> ChartInfo | ChartError:
|
||||
"""Get chart metadata by ID or UUID.
|
||||
|
||||
IMPORTANT FOR LLM CLIENTS:
|
||||
- ALWAYS display the chart URL when returned
|
||||
- URL field contains chart's screenshot URL for preview
|
||||
|
||||
Supports:
|
||||
- Numeric ID (e.g., 123)
|
||||
- UUID string (e.g., "a1b2c3d4-...")
|
||||
|
||||
Returns chart details including name, type, and URL.
|
||||
"""
|
||||
from superset.daos.chart import ChartDAO
|
||||
|
||||
2082
superset/mcp_service/chart/tool/get_chart_preview.py
Normal file
2082
superset/mcp_service/chart/tool/get_chart_preview.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -20,14 +20,19 @@ MCP tool: list_charts (advanced filtering with metadata cache control)
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, cast, TYPE_CHECKING
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from superset.models.slice import Slice
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
ChartFilter,
|
||||
ChartInfo,
|
||||
ChartLike,
|
||||
ChartList,
|
||||
ListChartsRequest,
|
||||
serialize_chart_object,
|
||||
@@ -89,10 +94,14 @@ async def list_charts(request: ListChartsRequest, ctx: Context) -> ChartList:
|
||||
|
||||
from superset.daos.chart import ChartDAO
|
||||
|
||||
def _serialize_chart(obj: "Slice | None", cols: Any) -> ChartInfo | None:
|
||||
"""Serialize chart object with proper type casting."""
|
||||
return serialize_chart_object(cast(ChartLike | None, obj))
|
||||
|
||||
tool = ModelListCore(
|
||||
dao_class=ChartDAO,
|
||||
output_schema=ChartInfo,
|
||||
item_serializer=lambda obj, cols: serialize_chart_object(obj) if obj else None, # type: ignore[arg-type]
|
||||
item_serializer=_serialize_chart,
|
||||
filter_type=ChartFilter,
|
||||
default_columns=DEFAULT_CHART_COLUMNS,
|
||||
search_columns=[
|
||||
|
||||
224
superset/mcp_service/chart/tool/update_chart.py
Normal file
224
superset/mcp_service/chart/tool/update_chart.py
Normal file
@@ -0,0 +1,224 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
MCP tool: update_chart
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.chart.chart_utils import (
|
||||
analyze_chart_capabilities,
|
||||
analyze_chart_semantics,
|
||||
generate_chart_name,
|
||||
map_config_to_form_data,
|
||||
)
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
AccessibilityMetadata,
|
||||
GenerateChartResponse,
|
||||
PerformanceMetadata,
|
||||
UpdateChartRequest,
|
||||
)
|
||||
from superset.mcp_service.utils.url_utils import (
|
||||
get_chart_screenshot_url,
|
||||
get_superset_base_url,
|
||||
)
|
||||
from superset.utils import json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def update_chart(
|
||||
request: UpdateChartRequest, ctx: Context
|
||||
) -> GenerateChartResponse:
|
||||
"""Update existing chart with new configuration.
|
||||
|
||||
IMPORTANT:
|
||||
- Chart must already be saved (from generate_chart with save_chart=True)
|
||||
- LLM clients MUST display updated chart URL to users
|
||||
- Embed preview_url as image: 
|
||||
|
||||
Use when:
|
||||
- Modifying existing saved chart
|
||||
- Updating title, filters, or visualization settings
|
||||
- Changing chart type or data columns
|
||||
|
||||
Returns:
|
||||
- Updated chart info and metadata
|
||||
- Preview URL and explore URL for further editing
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Find the existing chart
|
||||
from superset.daos.chart import ChartDAO
|
||||
|
||||
chart = None
|
||||
if isinstance(request.identifier, int) or (
|
||||
isinstance(request.identifier, str) and request.identifier.isdigit()
|
||||
):
|
||||
chart_id = (
|
||||
int(request.identifier)
|
||||
if isinstance(request.identifier, str)
|
||||
else request.identifier
|
||||
)
|
||||
chart = ChartDAO.find_by_id(chart_id)
|
||||
else:
|
||||
# Try UUID lookup using DAO flexible method
|
||||
chart = ChartDAO.find_by_id(request.identifier, id_column="uuid")
|
||||
|
||||
if not chart:
|
||||
return GenerateChartResponse.model_validate(
|
||||
{
|
||||
"chart": None,
|
||||
"error": f"No chart found with identifier: {request.identifier}",
|
||||
"success": False,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
)
|
||||
|
||||
# Map the new config to form_data format
|
||||
new_form_data = map_config_to_form_data(request.config)
|
||||
|
||||
# Update chart using Superset's command
|
||||
from superset.commands.chart.update import UpdateChartCommand
|
||||
|
||||
# Generate new chart name if provided, otherwise keep existing
|
||||
chart_name = (
|
||||
request.chart_name
|
||||
if request.chart_name
|
||||
else chart.slice_name or generate_chart_name(request.config)
|
||||
)
|
||||
|
||||
update_payload = {
|
||||
"slice_name": chart_name,
|
||||
"viz_type": new_form_data["viz_type"],
|
||||
"params": json.dumps(new_form_data),
|
||||
}
|
||||
|
||||
command = UpdateChartCommand(chart.id, update_payload)
|
||||
updated_chart = command.run()
|
||||
|
||||
# Generate semantic analysis
|
||||
capabilities = analyze_chart_capabilities(updated_chart, request.config)
|
||||
semantics = analyze_chart_semantics(updated_chart, request.config)
|
||||
|
||||
# Create performance metadata
|
||||
execution_time = int((time.time() - start_time) * 1000)
|
||||
performance = PerformanceMetadata(
|
||||
query_duration_ms=execution_time,
|
||||
cache_status="miss",
|
||||
optimization_suggestions=[],
|
||||
)
|
||||
|
||||
# Create accessibility metadata
|
||||
chart_name = (
|
||||
updated_chart.slice_name
|
||||
if updated_chart and hasattr(updated_chart, "slice_name")
|
||||
else generate_chart_name(request.config)
|
||||
)
|
||||
accessibility = AccessibilityMetadata(
|
||||
color_blind_safe=True, # Would need actual analysis
|
||||
alt_text=f"Updated chart showing {chart_name}",
|
||||
high_contrast_available=False,
|
||||
)
|
||||
|
||||
# Generate previews if requested
|
||||
previews = {}
|
||||
if request.generate_preview:
|
||||
try:
|
||||
from superset.mcp_service.chart.tool.get_chart_preview import (
|
||||
_get_chart_preview_internal,
|
||||
GetChartPreviewRequest,
|
||||
)
|
||||
|
||||
for format_type in request.preview_formats:
|
||||
preview_request = GetChartPreviewRequest(
|
||||
identifier=str(updated_chart.id), format=format_type
|
||||
)
|
||||
preview_result = await _get_chart_preview_internal(
|
||||
preview_request, ctx
|
||||
)
|
||||
|
||||
if hasattr(preview_result, "content"):
|
||||
previews[format_type] = preview_result.content
|
||||
|
||||
except Exception as e:
|
||||
# Log warning but don't fail the entire request
|
||||
logger.warning("Preview generation failed: %s", e)
|
||||
|
||||
# Return enhanced data
|
||||
result = {
|
||||
"chart": {
|
||||
"id": updated_chart.id,
|
||||
"slice_name": updated_chart.slice_name,
|
||||
"viz_type": updated_chart.viz_type,
|
||||
"url": (
|
||||
f"{get_superset_base_url()}/explore/?slice_id={updated_chart.id}"
|
||||
),
|
||||
"uuid": str(updated_chart.uuid) if updated_chart.uuid else None,
|
||||
"updated": True,
|
||||
},
|
||||
"error": None,
|
||||
# Enhanced fields for better LLM integration
|
||||
"previews": previews,
|
||||
"capabilities": capabilities.model_dump() if capabilities else None,
|
||||
"semantics": semantics.model_dump() if semantics else None,
|
||||
"explore_url": (
|
||||
f"{get_superset_base_url()}/explore/?slice_id={updated_chart.id}"
|
||||
),
|
||||
"api_endpoints": {
|
||||
"data": (
|
||||
f"{get_superset_base_url()}/api/v1/chart/{updated_chart.id}/data/"
|
||||
),
|
||||
"preview": get_chart_screenshot_url(updated_chart.id),
|
||||
"export": (
|
||||
f"{get_superset_base_url()}/api/v1/chart/{updated_chart.id}/export/"
|
||||
),
|
||||
},
|
||||
"performance": performance.model_dump() if performance else None,
|
||||
"accessibility": accessibility.model_dump() if accessibility else None,
|
||||
"success": True,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
return GenerateChartResponse.model_validate(result)
|
||||
|
||||
except Exception as e:
|
||||
execution_time = int((time.time() - start_time) * 1000)
|
||||
return GenerateChartResponse.model_validate(
|
||||
{
|
||||
"chart": None,
|
||||
"error": f"Chart update failed: {str(e)}",
|
||||
"performance": {
|
||||
"query_duration_ms": execution_time,
|
||||
"cache_status": "error",
|
||||
"optimization_suggestions": [],
|
||||
},
|
||||
"success": False,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
)
|
||||
168
superset/mcp_service/chart/tool/update_chart_preview.py
Normal file
168
superset/mcp_service/chart/tool/update_chart_preview.py
Normal file
@@ -0,0 +1,168 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
MCP tool: update_chart_preview
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Dict
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.chart.chart_utils import (
|
||||
analyze_chart_capabilities,
|
||||
analyze_chart_semantics,
|
||||
generate_chart_name,
|
||||
generate_explore_link,
|
||||
map_config_to_form_data,
|
||||
)
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
AccessibilityMetadata,
|
||||
PerformanceMetadata,
|
||||
UpdateChartPreviewRequest,
|
||||
URLPreview,
|
||||
)
|
||||
from superset.mcp_service.utils.url_utils import get_mcp_service_url
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
def update_chart_preview(
|
||||
request: UpdateChartPreviewRequest, ctx: Context
|
||||
) -> Dict[str, Any]:
|
||||
"""Update cached chart preview without saving.
|
||||
|
||||
IMPORTANT:
|
||||
- Modifies cached form_data from generate_chart (save_chart=False)
|
||||
- Original form_data_key is invalidated, new one returned
|
||||
- LLM clients MUST display explore_url to users
|
||||
- Embed preview_url as image: 
|
||||
|
||||
Use when:
|
||||
- Modifying preview before deciding to save
|
||||
- Iterating on chart design without creating permanent charts
|
||||
- Testing different configurations
|
||||
|
||||
Returns new form_data_key, preview images, and explore URL.
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Map the new config to form_data format
|
||||
new_form_data = map_config_to_form_data(request.config)
|
||||
|
||||
# Generate new explore link with updated form_data
|
||||
explore_url = generate_explore_link(request.dataset_id, new_form_data)
|
||||
|
||||
# Extract new form_data_key from the explore URL
|
||||
new_form_data_key = None
|
||||
if "form_data_key=" in explore_url:
|
||||
new_form_data_key = explore_url.split("form_data_key=")[1].split("&")[0]
|
||||
|
||||
# Generate semantic analysis
|
||||
capabilities = analyze_chart_capabilities(None, request.config)
|
||||
semantics = analyze_chart_semantics(None, request.config)
|
||||
|
||||
# Create performance metadata
|
||||
execution_time = int((time.time() - start_time) * 1000)
|
||||
performance = PerformanceMetadata(
|
||||
query_duration_ms=execution_time,
|
||||
cache_status="miss",
|
||||
optimization_suggestions=[],
|
||||
)
|
||||
|
||||
# Create accessibility metadata
|
||||
chart_name = generate_chart_name(request.config)
|
||||
accessibility = AccessibilityMetadata(
|
||||
color_blind_safe=True, # Would need actual analysis
|
||||
alt_text=f"Updated chart preview showing {chart_name}",
|
||||
high_contrast_available=False,
|
||||
)
|
||||
|
||||
# Generate previews if requested
|
||||
previews = {}
|
||||
if request.generate_preview and new_form_data_key:
|
||||
try:
|
||||
for format_type in request.preview_formats:
|
||||
if format_type == "url":
|
||||
# Generate screenshot URL using new form_data key
|
||||
mcp_base = get_mcp_service_url()
|
||||
preview_url = (
|
||||
f"{mcp_base}/screenshot/explore/{new_form_data_key}.png"
|
||||
)
|
||||
|
||||
previews[format_type] = URLPreview(
|
||||
preview_url=preview_url,
|
||||
width=800,
|
||||
height=600,
|
||||
supports_interaction=False,
|
||||
)
|
||||
# Other formats would need form_data execution
|
||||
# which is more complex for preview-only mode
|
||||
|
||||
except Exception as e:
|
||||
# Log warning but don't fail the entire request
|
||||
logger.warning("Preview generation failed: %s", e)
|
||||
|
||||
# Return enhanced data
|
||||
result = {
|
||||
"chart": {
|
||||
"id": None,
|
||||
"slice_name": chart_name,
|
||||
"viz_type": new_form_data.get("viz_type"),
|
||||
"url": explore_url,
|
||||
"uuid": None,
|
||||
"saved": False,
|
||||
"updated": True,
|
||||
},
|
||||
"error": None,
|
||||
# Enhanced fields for better LLM integration
|
||||
"previews": previews,
|
||||
"capabilities": capabilities.model_dump() if capabilities else None,
|
||||
"semantics": semantics.model_dump() if semantics else None,
|
||||
"explore_url": explore_url,
|
||||
"form_data_key": new_form_data_key,
|
||||
"previous_form_data_key": request.form_data_key, # For reference
|
||||
"api_endpoints": {}, # No API endpoints for unsaved charts
|
||||
"performance": performance.model_dump() if performance else None,
|
||||
"accessibility": accessibility.model_dump() if accessibility else None,
|
||||
"success": True,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
execution_time = int((time.time() - start_time) * 1000)
|
||||
return {
|
||||
"chart": None,
|
||||
"error": f"Chart preview update failed: {str(e)}",
|
||||
"performance": {
|
||||
"query_duration_ms": execution_time,
|
||||
"cache_status": "error",
|
||||
"optimization_suggestions": [],
|
||||
},
|
||||
"success": False,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
24
superset/mcp_service/chart/validation/__init__.py
Normal file
24
superset/mcp_service/chart/validation/__init__.py
Normal file
@@ -0,0 +1,24 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Chart validation module."""
|
||||
|
||||
from .dataset_validator import DatasetValidator
|
||||
from .pipeline import ValidationPipeline
|
||||
from .schema_validator import SchemaValidator
|
||||
|
||||
__all__ = ["ValidationPipeline", "SchemaValidator", "DatasetValidator"]
|
||||
329
superset/mcp_service/chart/validation/dataset_validator.py
Normal file
329
superset/mcp_service/chart/validation/dataset_validator.py
Normal file
@@ -0,0 +1,329 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Dataset-specific validation for chart configurations.
|
||||
Validates that referenced columns exist in the dataset schema.
|
||||
"""
|
||||
|
||||
import difflib
|
||||
import logging
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
ColumnRef,
|
||||
TableChartConfig,
|
||||
XYChartConfig,
|
||||
)
|
||||
from superset.mcp_service.common.error_schemas import (
|
||||
ChartGenerationError,
|
||||
ColumnSuggestion,
|
||||
DatasetContext,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DatasetValidator:
|
||||
"""Validates chart configuration against dataset schema."""
|
||||
|
||||
@staticmethod
|
||||
def validate_against_dataset(
|
||||
config: TableChartConfig | XYChartConfig, dataset_id: int | str
|
||||
) -> Tuple[bool, ChartGenerationError | None]:
|
||||
"""
|
||||
Validate chart configuration against dataset schema.
|
||||
|
||||
Args:
|
||||
config: Chart configuration to validate
|
||||
dataset_id: Dataset ID to validate against
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, error)
|
||||
"""
|
||||
# Get dataset context
|
||||
dataset_context = DatasetValidator._get_dataset_context(dataset_id)
|
||||
if not dataset_context:
|
||||
from superset.mcp_service.utils.error_builder import (
|
||||
ChartErrorBuilder,
|
||||
)
|
||||
|
||||
return False, ChartErrorBuilder.dataset_not_found_error(dataset_id)
|
||||
|
||||
# Collect all column references
|
||||
column_refs = DatasetValidator._extract_column_references(config)
|
||||
|
||||
# Validate each column exists
|
||||
invalid_columns = []
|
||||
for col_ref in column_refs:
|
||||
if not DatasetValidator._column_exists(col_ref.name, dataset_context):
|
||||
invalid_columns.append(col_ref)
|
||||
|
||||
if invalid_columns:
|
||||
# Generate suggestions for invalid columns
|
||||
suggestions_map = {}
|
||||
for col_ref in invalid_columns:
|
||||
suggestions = DatasetValidator._get_column_suggestions(
|
||||
col_ref.name, dataset_context
|
||||
)
|
||||
suggestions_map[col_ref.name] = suggestions
|
||||
|
||||
# Build error with suggestions
|
||||
return False, DatasetValidator._build_column_error(
|
||||
invalid_columns, suggestions_map, dataset_context
|
||||
)
|
||||
|
||||
# Validate aggregation compatibility
|
||||
if isinstance(config, (TableChartConfig, XYChartConfig)):
|
||||
aggregation_errors = DatasetValidator._validate_aggregations(
|
||||
column_refs, dataset_context
|
||||
)
|
||||
if aggregation_errors:
|
||||
return False, aggregation_errors[0]
|
||||
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def _get_dataset_context(dataset_id: int | str) -> DatasetContext | None:
|
||||
"""Get dataset context with column information."""
|
||||
try:
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
|
||||
# Find dataset
|
||||
if isinstance(dataset_id, int) or (
|
||||
isinstance(dataset_id, str) and dataset_id.isdigit()
|
||||
):
|
||||
dataset = DatasetDAO.find_by_id(int(dataset_id))
|
||||
else:
|
||||
# Try UUID lookup
|
||||
dataset = DatasetDAO.find_by_id(dataset_id, id_column="uuid")
|
||||
|
||||
if not dataset:
|
||||
return None
|
||||
|
||||
# Build context
|
||||
columns = []
|
||||
metrics = []
|
||||
|
||||
# Add table columns
|
||||
for col in dataset.columns:
|
||||
columns.append(
|
||||
{
|
||||
"name": col.column_name,
|
||||
"type": str(col.type) if col.type else "UNKNOWN",
|
||||
"is_temporal": col.is_temporal
|
||||
if hasattr(col, "is_temporal")
|
||||
else False,
|
||||
"is_numeric": col.is_numeric
|
||||
if hasattr(col, "is_numeric")
|
||||
else False,
|
||||
}
|
||||
)
|
||||
|
||||
# Add metrics
|
||||
for metric in dataset.metrics:
|
||||
metrics.append(
|
||||
{
|
||||
"name": metric.metric_name,
|
||||
"expression": metric.expression,
|
||||
"description": metric.description,
|
||||
}
|
||||
)
|
||||
|
||||
return DatasetContext(
|
||||
id=dataset.id,
|
||||
table_name=dataset.table_name,
|
||||
schema=dataset.schema,
|
||||
database_name=dataset.database.database_name
|
||||
if dataset.database
|
||||
else None,
|
||||
available_columns=columns,
|
||||
available_metrics=metrics,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error getting dataset context for %s: %s", dataset_id, e)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _extract_column_references(
|
||||
config: TableChartConfig | XYChartConfig,
|
||||
) -> List[ColumnRef]:
|
||||
"""Extract all column references from configuration."""
|
||||
refs = []
|
||||
|
||||
if isinstance(config, TableChartConfig):
|
||||
refs.extend(config.columns)
|
||||
elif isinstance(config, XYChartConfig):
|
||||
refs.append(config.x)
|
||||
refs.extend(config.y)
|
||||
if config.group_by:
|
||||
refs.append(config.group_by)
|
||||
|
||||
# Add filter columns
|
||||
if hasattr(config, "filters") and config.filters:
|
||||
for filter_config in config.filters:
|
||||
refs.append(ColumnRef(name=filter_config.column))
|
||||
|
||||
return refs
|
||||
|
||||
@staticmethod
|
||||
def _column_exists(column_name: str, dataset_context: DatasetContext) -> bool:
|
||||
"""Check if column exists in dataset (case-insensitive)."""
|
||||
column_lower = column_name.lower()
|
||||
|
||||
# Check regular columns
|
||||
for col in dataset_context.available_columns:
|
||||
if col["name"].lower() == column_lower:
|
||||
return True
|
||||
|
||||
# Check metrics
|
||||
for metric in dataset_context.available_metrics:
|
||||
if metric["name"].lower() == column_lower:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _get_column_suggestions(
|
||||
column_name: str, dataset_context: DatasetContext, max_suggestions: int = 3
|
||||
) -> List[ColumnSuggestion]:
|
||||
"""Get column name suggestions using fuzzy matching."""
|
||||
all_names = []
|
||||
|
||||
# Collect all column names
|
||||
for col in dataset_context.available_columns:
|
||||
all_names.append((col["name"], "column", col.get("type", "UNKNOWN")))
|
||||
|
||||
for metric in dataset_context.available_metrics:
|
||||
all_names.append((metric["name"], "metric", "METRIC"))
|
||||
|
||||
# Find close matches
|
||||
column_lower = column_name.lower()
|
||||
close_matches = difflib.get_close_matches(
|
||||
column_lower,
|
||||
[name[0].lower() for name in all_names],
|
||||
n=max_suggestions,
|
||||
cutoff=0.6,
|
||||
)
|
||||
|
||||
# Build suggestions with proper case and type info
|
||||
suggestions = []
|
||||
for match in close_matches:
|
||||
for name, col_type, data_type in all_names:
|
||||
if name.lower() == match:
|
||||
suggestions.append(
|
||||
ColumnSuggestion(name=name, type=col_type, data_type=data_type)
|
||||
)
|
||||
break
|
||||
|
||||
return suggestions
|
||||
|
||||
@staticmethod
|
||||
def _build_column_error(
|
||||
invalid_columns: List[ColumnRef],
|
||||
suggestions_map: Dict[str, List[ColumnSuggestion]],
|
||||
dataset_context: DatasetContext,
|
||||
) -> ChartGenerationError:
|
||||
"""Build error for invalid columns."""
|
||||
from superset.mcp_service.utils.error_builder import (
|
||||
ChartErrorBuilder,
|
||||
)
|
||||
|
||||
# Format error message
|
||||
if len(invalid_columns) == 1:
|
||||
col = invalid_columns[0]
|
||||
suggestions = suggestions_map.get(col.name, [])
|
||||
|
||||
if suggestions:
|
||||
return ChartErrorBuilder.column_not_found_error(
|
||||
col.name, [s.name for s in suggestions]
|
||||
)
|
||||
else:
|
||||
return ChartErrorBuilder.column_not_found_error(col.name)
|
||||
else:
|
||||
# Multiple invalid columns
|
||||
invalid_names = [col.name for col in invalid_columns]
|
||||
return ChartErrorBuilder.build_error(
|
||||
error_type="multiple_invalid_columns",
|
||||
template_key="column_not_found",
|
||||
template_vars={
|
||||
"column": ", ".join(invalid_names[:3])
|
||||
+ ("..." if len(invalid_names) > 3 else ""),
|
||||
"suggestions": "Use get_dataset_info to see all available columns",
|
||||
},
|
||||
custom_suggestions=[
|
||||
f"Invalid columns: {', '.join(invalid_names)}",
|
||||
"Check spelling and case sensitivity",
|
||||
"Use get_dataset_info to list available columns",
|
||||
],
|
||||
error_code="MULTIPLE_INVALID_COLUMNS",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _validate_aggregations(
|
||||
column_refs: List[ColumnRef], dataset_context: DatasetContext
|
||||
) -> List[ChartGenerationError]:
|
||||
"""Validate that aggregations are appropriate for column types."""
|
||||
errors = []
|
||||
|
||||
for col_ref in column_refs:
|
||||
if not col_ref.aggregate:
|
||||
continue
|
||||
|
||||
# Find column info
|
||||
col_info = None
|
||||
for col in dataset_context.available_columns:
|
||||
if col["name"].lower() == col_ref.name.lower():
|
||||
col_info = col
|
||||
break
|
||||
|
||||
if col_info:
|
||||
# Check numeric aggregates on non-numeric columns
|
||||
numeric_aggs = ["SUM", "AVG", "MIN", "MAX", "STDDEV", "VAR", "MEDIAN"]
|
||||
if (
|
||||
col_ref.aggregate in numeric_aggs
|
||||
and not col_info.get("is_numeric", False)
|
||||
and col_info.get("type", "").upper()
|
||||
not in ["INTEGER", "FLOAT", "DOUBLE", "DECIMAL", "NUMERIC"]
|
||||
):
|
||||
from superset.mcp_service.utils.error_builder import ( # noqa: E501
|
||||
ChartErrorBuilder,
|
||||
)
|
||||
|
||||
errors.append(
|
||||
ChartErrorBuilder.build_error(
|
||||
error_type="invalid_aggregation",
|
||||
template_key="incompatible_configuration",
|
||||
template_vars={
|
||||
"reason": f"Cannot apply {col_ref.aggregate} to "
|
||||
f"non-numeric column "
|
||||
f"'{col_ref.name}' (type:"
|
||||
f" {col_info.get('type', 'UNKNOWN')})",
|
||||
"primary_suggestion": "Use COUNT or COUNT_DISTINCT "
|
||||
"for text columns",
|
||||
},
|
||||
custom_suggestions=[
|
||||
"Remove the aggregate function for raw values",
|
||||
"Use COUNT to count occurrences",
|
||||
"Use COUNT_DISTINCT to count unique values",
|
||||
],
|
||||
error_code="INVALID_AGGREGATION",
|
||||
)
|
||||
)
|
||||
|
||||
return errors
|
||||
325
superset/mcp_service/chart/validation/pipeline.py
Normal file
325
superset/mcp_service/chart/validation/pipeline.py
Normal file
@@ -0,0 +1,325 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Unified validation pipeline for chart operations.
|
||||
Orchestrates schema, dataset, and runtime validations.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
ChartConfig,
|
||||
GenerateChartRequest,
|
||||
)
|
||||
from superset.mcp_service.common.error_schemas import ChartGenerationError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _redact_sql_select(error_str: str, error_str_upper: str) -> str:
|
||||
"""Redact SELECT...FROM clause content to prevent data disclosure."""
|
||||
if "SELECT" in error_str_upper and "FROM" in error_str_upper:
|
||||
select_idx = error_str_upper.find("SELECT")
|
||||
from_idx = error_str_upper.find("FROM", select_idx)
|
||||
if select_idx != -1 and from_idx != -1:
|
||||
return error_str[: select_idx + 7] + " [REDACTED] " + error_str[from_idx:]
|
||||
return error_str
|
||||
|
||||
|
||||
def _redact_sql_where(error_str: str, error_str_upper: str) -> str:
|
||||
"""Redact WHERE clause content to prevent data disclosure."""
|
||||
if "WHERE" not in error_str_upper:
|
||||
return error_str
|
||||
|
||||
where_idx = error_str_upper.find("WHERE")
|
||||
terminators = ["ORDER", "GROUP", "LIMIT", "UNION", "EXCEPT", "INTERSECT"]
|
||||
term_idx = len(error_str)
|
||||
for term in terminators:
|
||||
idx = error_str_upper.find(term, where_idx)
|
||||
if idx != -1 and idx < term_idx:
|
||||
term_idx = idx
|
||||
return error_str[: where_idx + 6] + " [REDACTED]" + error_str[term_idx:]
|
||||
|
||||
|
||||
def _get_generic_error_message(error_str: str) -> str | None:
|
||||
"""Return generic message for common error types, or None."""
|
||||
error_lower = error_str.lower()
|
||||
if "permission" in error_lower or "access" in error_lower:
|
||||
return "Validation failed due to access restrictions"
|
||||
if "database" in error_lower or "connection" in error_lower:
|
||||
return "Validation failed due to database connectivity"
|
||||
if "timeout" in error_lower:
|
||||
return "Validation timed out"
|
||||
return None
|
||||
|
||||
|
||||
def _sanitize_validation_error(error: Exception) -> str:
|
||||
"""SECURITY FIX: Sanitize validation errors to prevent disclosure."""
|
||||
error_str = str(error)
|
||||
|
||||
# SECURITY FIX: Limit length FIRST to prevent ReDoS attacks
|
||||
if len(error_str) > 200:
|
||||
error_str = error_str[:200] + "...[truncated]"
|
||||
|
||||
# Remove potentially sensitive schema information
|
||||
import re
|
||||
|
||||
sensitive_patterns = [
|
||||
(r'\btable\s+[\'"`]?(\w+)[\'"`]?', "table [REDACTED]"),
|
||||
(r'\bcolumn\s+[\'"`]?(\w+)[\'"`]?', "column [REDACTED]"),
|
||||
(r'\bdatabase\s+[\'"`]?(\w+)[\'"`]?', "database [REDACTED]"),
|
||||
(r'\bschema\s+[\'"`]?(\w+)[\'"`]?', "schema [REDACTED]"),
|
||||
]
|
||||
for pattern, replacement in sensitive_patterns:
|
||||
error_str = re.sub(pattern, replacement, error_str, flags=re.IGNORECASE)
|
||||
|
||||
# SECURITY FIX: SQL sanitization without ReDoS-vulnerable patterns
|
||||
error_str_upper = error_str.upper()
|
||||
error_str = _redact_sql_select(error_str, error_str_upper)
|
||||
error_str = _redact_sql_where(error_str, error_str_upper)
|
||||
|
||||
# Return generic message for common error types
|
||||
if generic := _get_generic_error_message(error_str):
|
||||
return generic
|
||||
|
||||
return error_str
|
||||
|
||||
|
||||
class ValidationPipeline:
|
||||
"""
|
||||
Main validation orchestrator that runs validations in sequence:
|
||||
1. Schema validation (structure and types)
|
||||
2. Dataset validation (columns exist)
|
||||
3. Runtime validation (performance, compatibility)
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def validate_request(
|
||||
request_data: Dict[str, Any],
|
||||
) -> Tuple[bool, GenerateChartRequest | None, ChartGenerationError | None]:
|
||||
"""
|
||||
Validate a chart generation request through all validation layers.
|
||||
|
||||
Args:
|
||||
request_data: Raw request data dictionary
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, parsed_request, error)
|
||||
"""
|
||||
try:
|
||||
# Layer 1: Schema validation
|
||||
from .schema_validator import SchemaValidator
|
||||
|
||||
is_valid, request, error = SchemaValidator.validate_request(request_data)
|
||||
if not is_valid:
|
||||
return False, None, error
|
||||
|
||||
# Ensure request is not None
|
||||
if request is None:
|
||||
return False, None, error
|
||||
|
||||
# Layer 2: Dataset validation
|
||||
is_valid, error = ValidationPipeline._validate_dataset(
|
||||
request.config, request.dataset_id
|
||||
)
|
||||
if not is_valid:
|
||||
return False, request, error
|
||||
|
||||
# Layer 3: Runtime validation
|
||||
is_valid, error = ValidationPipeline._validate_runtime(
|
||||
request.config, request.dataset_id
|
||||
)
|
||||
if not is_valid:
|
||||
return False, request, error
|
||||
|
||||
return True, request, None
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("Validation pipeline error")
|
||||
from superset.mcp_service.utils.error_builder import (
|
||||
ChartErrorBuilder,
|
||||
)
|
||||
|
||||
# SECURITY FIX: Sanitize validation error to prevent information disclosure
|
||||
sanitized_reason = _sanitize_validation_error(e)
|
||||
error = ChartErrorBuilder.build_error(
|
||||
error_type="validation_system_error",
|
||||
template_key="validation_error",
|
||||
template_vars={"reason": sanitized_reason},
|
||||
error_code="VALIDATION_PIPELINE_ERROR",
|
||||
)
|
||||
return False, None, error
|
||||
|
||||
@staticmethod
|
||||
def _validate_dataset(
|
||||
config: ChartConfig, dataset_id: int | str
|
||||
) -> Tuple[bool, ChartGenerationError | None]:
|
||||
"""Validate configuration against dataset schema."""
|
||||
try:
|
||||
from .dataset_validator import DatasetValidator
|
||||
|
||||
return DatasetValidator.validate_against_dataset(config, dataset_id)
|
||||
except ImportError:
|
||||
# Skip if dataset validator not available
|
||||
logger.warning(
|
||||
"Dataset validator not available, skipping dataset validation"
|
||||
)
|
||||
return True, None
|
||||
except Exception as e:
|
||||
logger.warning("Dataset validation failed: %s", e)
|
||||
# Don't fail on dataset validation errors
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def _validate_runtime(
|
||||
config: ChartConfig, dataset_id: int | str
|
||||
) -> Tuple[bool, ChartGenerationError | None]:
|
||||
"""Validate runtime issues (performance, compatibility)."""
|
||||
try:
|
||||
from .runtime import RuntimeValidator
|
||||
|
||||
return RuntimeValidator.validate_runtime_issues(config, dataset_id)
|
||||
except ImportError:
|
||||
# Skip if runtime validator not available
|
||||
logger.warning(
|
||||
"Runtime validator not available, skipping runtime validation"
|
||||
)
|
||||
return True, None
|
||||
except Exception as e:
|
||||
logger.warning("Runtime validation failed: %s", e)
|
||||
# Don't fail on runtime validation errors
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def validate_filters(
|
||||
filters: List[Any],
|
||||
) -> Tuple[bool, ChartGenerationError | None]:
|
||||
"""
|
||||
Validate filter logic for contradictions and empty results.
|
||||
|
||||
Args:
|
||||
filters: List of filter configurations
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, error)
|
||||
"""
|
||||
if not filters:
|
||||
return True, None
|
||||
|
||||
# Check for contradictory filters
|
||||
if ValidationPipeline._has_contradictory_filters(filters):
|
||||
from superset.mcp_service.utils.error_builder import (
|
||||
ChartErrorBuilder,
|
||||
)
|
||||
|
||||
return False, ChartErrorBuilder.build_error(
|
||||
error_type="contradictory_filters",
|
||||
template_key="invalid_value",
|
||||
template_vars={
|
||||
"field": "filters",
|
||||
"value": "contradictory conditions",
|
||||
"reason": "Filter conditions are logically impossible",
|
||||
"allowed_values": "non-contradictory conditions",
|
||||
"specific_suggestion": "Remove conflicting filters",
|
||||
},
|
||||
error_code="CONTRADICTORY_FILTERS",
|
||||
)
|
||||
|
||||
# Check for filters likely to return empty
|
||||
if empty_warnings := ValidationPipeline._check_empty_result_filters(filters):
|
||||
from superset.mcp_service.utils.error_builder import (
|
||||
ChartErrorBuilder,
|
||||
)
|
||||
|
||||
return False, ChartErrorBuilder.build_error(
|
||||
error_type="empty_result_warning",
|
||||
template_key="empty_result",
|
||||
template_vars={"reason": "; ".join(empty_warnings)},
|
||||
custom_suggestions=[
|
||||
"Verify filter values exist in your dataset",
|
||||
"Check for typos in filter values",
|
||||
"Use broader filter criteria",
|
||||
],
|
||||
error_code="EMPTY_RESULT_WARNING",
|
||||
)
|
||||
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def _has_contradictory_filters(filters: List[Any]) -> bool:
|
||||
"""Check if filters contain logical contradictions."""
|
||||
# Group filters by column
|
||||
column_filters: Dict[str, List[Any]] = {}
|
||||
for f in filters:
|
||||
col = f.column
|
||||
if col not in column_filters:
|
||||
column_filters[col] = []
|
||||
column_filters[col].append(f)
|
||||
|
||||
# Check for contradictions within same column
|
||||
for _col, col_filters in column_filters.items():
|
||||
# Check for > X AND < Y where X >= Y
|
||||
gt_values = [f.value for f in col_filters if f.op == ">"]
|
||||
lt_values = [f.value for f in col_filters if f.op == "<"]
|
||||
|
||||
for gt in gt_values:
|
||||
for lt in lt_values:
|
||||
try:
|
||||
if float(gt) >= float(lt):
|
||||
return True
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
|
||||
# Check for = X AND = Y where X != Y
|
||||
eq_values = [f.value for f in col_filters if f.op == "="]
|
||||
if len(eq_values) > 1 and len(set(eq_values)) > 1:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _check_empty_result_filters(filters: List[Any]) -> List[str]:
|
||||
"""Check for filter patterns that commonly result in empty results."""
|
||||
warnings = []
|
||||
|
||||
for f in filters:
|
||||
col_lower = f.column.lower()
|
||||
val_str = str(f.value).lower() if f.value is not None else ""
|
||||
|
||||
# Check for common empty result patterns
|
||||
if f.op == "=" and any(
|
||||
pattern in val_str
|
||||
for pattern in ["deleted", "archived", "inactive", "disabled"]
|
||||
):
|
||||
warnings.append(
|
||||
f"Filter '{f.column} = {f.value}' may return few or no results"
|
||||
)
|
||||
|
||||
# Check for future dates
|
||||
if "date" in col_lower and f.op in [">", ">="]:
|
||||
try:
|
||||
if "20" in val_str and int(val_str[:4]) > 2025:
|
||||
warnings.append(
|
||||
f"Filter '{f.column} {f.op} {f.value}' uses future date"
|
||||
)
|
||||
except (ValueError, IndexError):
|
||||
# Ignore invalid date formats
|
||||
pass
|
||||
|
||||
return warnings
|
||||
185
superset/mcp_service/chart/validation/runtime/__init__.py
Normal file
185
superset/mcp_service/chart/validation/runtime/__init__.py
Normal file
@@ -0,0 +1,185 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Runtime validation module for chart configurations.
|
||||
Validates performance, compatibility, and user experience issues.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List, Tuple
|
||||
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
ChartConfig,
|
||||
XYChartConfig,
|
||||
)
|
||||
from superset.mcp_service.common.error_schemas import ChartGenerationError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RuntimeValidator:
|
||||
"""Orchestrates runtime validations for chart configurations."""
|
||||
|
||||
@staticmethod
|
||||
def validate_runtime_issues(
|
||||
config: ChartConfig, dataset_id: int | str
|
||||
) -> Tuple[bool, ChartGenerationError | None]:
|
||||
"""
|
||||
Validate runtime issues that could affect chart rendering or performance.
|
||||
|
||||
Args:
|
||||
config: Chart configuration to validate
|
||||
dataset_id: Dataset identifier
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, error)
|
||||
"""
|
||||
warnings: List[str] = []
|
||||
suggestions: List[str] = []
|
||||
|
||||
# Only check XY charts for format and cardinality issues
|
||||
if isinstance(config, XYChartConfig):
|
||||
# Format-type compatibility validation
|
||||
format_warnings = RuntimeValidator._validate_format_compatibility(config)
|
||||
if format_warnings:
|
||||
warnings.extend(format_warnings)
|
||||
|
||||
# Cardinality validation
|
||||
cardinality_warnings, cardinality_suggestions = (
|
||||
RuntimeValidator._validate_cardinality(config, dataset_id)
|
||||
)
|
||||
if cardinality_warnings:
|
||||
warnings.extend(cardinality_warnings)
|
||||
suggestions.extend(cardinality_suggestions)
|
||||
|
||||
# Chart type appropriateness validation (for all chart types)
|
||||
type_warnings, type_suggestions = RuntimeValidator._validate_chart_type(
|
||||
config, dataset_id
|
||||
)
|
||||
if type_warnings:
|
||||
warnings.extend(type_warnings)
|
||||
suggestions.extend(type_suggestions)
|
||||
|
||||
# If we have warnings, return them as a validation error
|
||||
if warnings:
|
||||
from superset.mcp_service.utils.error_builder import (
|
||||
ChartErrorBuilder,
|
||||
)
|
||||
|
||||
return False, ChartErrorBuilder.build_error(
|
||||
error_type="runtime_semantic_warning",
|
||||
template_key="performance_warning",
|
||||
template_vars={
|
||||
"reason": "; ".join(warnings[:3])
|
||||
+ ("..." if len(warnings) > 3 else "")
|
||||
},
|
||||
custom_suggestions=suggestions[:5], # Limit suggestions
|
||||
error_code="RUNTIME_SEMANTIC_WARNING",
|
||||
)
|
||||
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def _validate_format_compatibility(config: XYChartConfig) -> List[str]:
|
||||
"""Validate format-type compatibility."""
|
||||
warnings: List[str] = []
|
||||
|
||||
try:
|
||||
# Import here to avoid circular imports
|
||||
from .format_validator import FormatTypeValidator
|
||||
|
||||
is_valid, format_warnings = (
|
||||
FormatTypeValidator.validate_format_compatibility(config)
|
||||
)
|
||||
if format_warnings:
|
||||
warnings.extend(format_warnings)
|
||||
except ImportError:
|
||||
logger.warning("Format validator not available")
|
||||
except Exception as e:
|
||||
logger.warning("Format validation failed: %s", e)
|
||||
|
||||
return warnings
|
||||
|
||||
@staticmethod
|
||||
def _validate_cardinality(
|
||||
config: XYChartConfig, dataset_id: int | str
|
||||
) -> Tuple[List[str], List[str]]:
|
||||
"""Validate cardinality issues."""
|
||||
warnings: List[str] = []
|
||||
suggestions: List[str] = []
|
||||
|
||||
try:
|
||||
# Import here to avoid circular imports
|
||||
from .cardinality_validator import CardinalityValidator
|
||||
|
||||
# Determine chart type for cardinality thresholds
|
||||
chart_type = config.kind if hasattr(config, "kind") else "default"
|
||||
|
||||
# Check X-axis cardinality
|
||||
is_ok, cardinality_info = CardinalityValidator.check_cardinality(
|
||||
dataset_id=dataset_id,
|
||||
x_column=config.x.name,
|
||||
chart_type=chart_type,
|
||||
group_by_column=config.group_by.name if config.group_by else None,
|
||||
)
|
||||
|
||||
if not is_ok and cardinality_info:
|
||||
warnings.extend(cardinality_info.get("warnings", []))
|
||||
suggestions.extend(cardinality_info.get("suggestions", []))
|
||||
|
||||
except ImportError:
|
||||
logger.warning("Cardinality validator not available")
|
||||
except Exception as e:
|
||||
logger.warning("Cardinality validation failed: %s", e)
|
||||
|
||||
return warnings, suggestions
|
||||
|
||||
@staticmethod
|
||||
def _validate_chart_type(
|
||||
config: ChartConfig, dataset_id: int | str
|
||||
) -> Tuple[List[str], List[str]]:
|
||||
"""Validate chart type appropriateness."""
|
||||
warnings: List[str] = []
|
||||
suggestions: List[str] = []
|
||||
|
||||
try:
|
||||
# Import here to avoid circular imports
|
||||
from .chart_type_suggester import ChartTypeSuggester
|
||||
|
||||
is_appropriate, suggestion_info = ChartTypeSuggester.analyze_and_suggest(
|
||||
config, dataset_id
|
||||
)
|
||||
|
||||
if not is_appropriate and suggestion_info:
|
||||
warnings.extend(suggestion_info.get("issues", []))
|
||||
suggestions.extend(suggestion_info.get("suggestions", []))
|
||||
|
||||
# Add recommended chart types
|
||||
recommended = suggestion_info.get("recommended_types", [])
|
||||
if recommended:
|
||||
recommendations = ", ".join(recommended)
|
||||
suggestions.append(
|
||||
f"Recommended chart types for this data: {recommendations}"
|
||||
)
|
||||
|
||||
except ImportError:
|
||||
logger.warning("Chart type suggester not available")
|
||||
except Exception as e:
|
||||
logger.warning("Chart type validation failed: %s", e)
|
||||
|
||||
return warnings, suggestions
|
||||
@@ -0,0 +1,195 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Cardinality validation to prevent unusable visualizations from high-cardinality data.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CardinalityValidator:
|
||||
"""
|
||||
Validates cardinality of dimensions to prevent charts with too many categories
|
||||
that become unreadable or cause performance issues.
|
||||
"""
|
||||
|
||||
# Thresholds for different chart types
|
||||
CARDINALITY_THRESHOLDS = {
|
||||
"bar": 50, # Bar charts become unreadable with >50 bars
|
||||
"line": 100, # Line charts can handle more points
|
||||
"scatter": 500, # Scatter plots can show many points
|
||||
"area": 30, # Area charts need fewer categories
|
||||
"table": 1000, # Tables can handle many rows with pagination
|
||||
"default": 50, # Conservative default
|
||||
}
|
||||
|
||||
# Known high-cardinality column patterns
|
||||
HIGH_CARDINALITY_PATTERNS = [
|
||||
"id",
|
||||
"uuid",
|
||||
"guid",
|
||||
"email",
|
||||
"phone",
|
||||
"address",
|
||||
"session",
|
||||
"transaction",
|
||||
"order_number",
|
||||
"invoice",
|
||||
"timestamp",
|
||||
"datetime",
|
||||
"created_at",
|
||||
"updated_at",
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def check_cardinality(
|
||||
dataset_id: int | str,
|
||||
x_column: str,
|
||||
chart_type: str = "default",
|
||||
group_by_column: str | None = None,
|
||||
) -> Tuple[bool, Dict[str, Any] | None]:
|
||||
"""
|
||||
Check cardinality of X-axis and group_by columns.
|
||||
|
||||
Returns:
|
||||
Tuple of (is_ok, warning_info)
|
||||
"""
|
||||
try:
|
||||
# Quick pattern check first (no DB query needed)
|
||||
pattern_warnings = CardinalityValidator._check_column_patterns(
|
||||
x_column, group_by_column
|
||||
)
|
||||
|
||||
if pattern_warnings:
|
||||
return False, {
|
||||
"warnings": pattern_warnings,
|
||||
"suggestions": CardinalityValidator._get_suggestions(
|
||||
x_column, chart_type, pattern_based=True
|
||||
),
|
||||
}
|
||||
|
||||
# For non-pattern columns, we could do actual cardinality check
|
||||
# but that requires DB access - for now just return OK
|
||||
# In production, you'd want to cache cardinality stats
|
||||
|
||||
return True, None
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("Cardinality check failed: %s", e)
|
||||
# Don't block on validation failures
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def _check_column_patterns(
|
||||
x_column: str, group_by_column: str | None = None
|
||||
) -> List[str]:
|
||||
"""Check for known high-cardinality column patterns."""
|
||||
warnings = []
|
||||
|
||||
x_lower = x_column.lower()
|
||||
|
||||
# Check X-axis column
|
||||
for pattern in CardinalityValidator.HIGH_CARDINALITY_PATTERNS:
|
||||
if pattern in x_lower:
|
||||
warnings.append(
|
||||
f"Column '{x_column}' appears to be a high-cardinality field "
|
||||
f"(contains '{pattern}'). This may create an unreadable chart "
|
||||
f"with too many categories on the X-axis."
|
||||
)
|
||||
break
|
||||
|
||||
# Check group_by column if present
|
||||
if group_by_column:
|
||||
group_lower = group_by_column.lower()
|
||||
for pattern in CardinalityValidator.HIGH_CARDINALITY_PATTERNS:
|
||||
if pattern in group_lower:
|
||||
warnings.append(
|
||||
f"Group by column '{group_by_column}' appears to be a "
|
||||
f"high-cardinality field (contains '{pattern}'). This may "
|
||||
f"create too many series to visualize effectively."
|
||||
)
|
||||
break
|
||||
|
||||
return warnings
|
||||
|
||||
@staticmethod
|
||||
def _get_suggestions(
|
||||
column: str, chart_type: str, pattern_based: bool = False
|
||||
) -> List[str]:
|
||||
"""Get suggestions for handling high cardinality."""
|
||||
suggestions = []
|
||||
|
||||
if pattern_based:
|
||||
# Suggestions when we detected high-cardinality patterns
|
||||
if any(p in column.lower() for p in ["id", "uuid", "guid"]):
|
||||
suggestions.extend(
|
||||
[
|
||||
"Consider using a different column for the X-axis",
|
||||
f"If you need to analyze by {column}, use filters to limit "
|
||||
f"the data",
|
||||
"A table chart might be more appropriate for ID-based data",
|
||||
]
|
||||
)
|
||||
elif any(p in column.lower() for p in ["email", "phone", "address"]):
|
||||
suggestions.extend(
|
||||
[
|
||||
"Consider grouping by a higher-level category (e.g., "
|
||||
"domain for emails)",
|
||||
f"Use filters to focus on specific {column} values",
|
||||
"Aggregate the data before visualization",
|
||||
]
|
||||
)
|
||||
elif any(
|
||||
p in column.lower() for p in ["timestamp", "datetime", "created_at"]
|
||||
):
|
||||
suggestions.extend(
|
||||
[
|
||||
"Consider truncating timestamps to date or hour level",
|
||||
"Use time-based grouping (daily, weekly, monthly)",
|
||||
"Apply date range filters to limit the data",
|
||||
]
|
||||
)
|
||||
else:
|
||||
# General high-cardinality suggestions
|
||||
threshold = CardinalityValidator.CARDINALITY_THRESHOLDS.get(chart_type, 50)
|
||||
suggestions.extend(
|
||||
[
|
||||
f"This chart type works best with fewer than {threshold} "
|
||||
f"categories",
|
||||
"Consider using filters to reduce the number of values",
|
||||
"Try grouping or categorizing the data at a higher level",
|
||||
"A table or pivot table might better display high-cardinality data",
|
||||
]
|
||||
)
|
||||
|
||||
return suggestions
|
||||
|
||||
@staticmethod
|
||||
def suggest_chart_type(cardinality: int) -> List[str]:
|
||||
"""Suggest appropriate chart types based on cardinality."""
|
||||
if cardinality <= 10:
|
||||
return ["bar", "pie", "donut", "area"]
|
||||
elif cardinality <= 30:
|
||||
return ["bar", "line", "area"]
|
||||
elif cardinality <= 100:
|
||||
return ["line", "scatter"]
|
||||
else:
|
||||
return ["table", "pivot_table", "heatmap"]
|
||||
@@ -0,0 +1,437 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Chart type suggestions based on data characteristics and user intent.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
ChartConfig,
|
||||
ColumnRef,
|
||||
TableChartConfig,
|
||||
XYChartConfig,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ChartTypeSuggester:
|
||||
"""
|
||||
Suggests appropriate chart types based on data characteristics
|
||||
and identifies potential mismatches between chart type and data.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def analyze_and_suggest(
|
||||
config: ChartConfig,
|
||||
dataset_id: int | str, # noqa: ARG002
|
||||
) -> Tuple[bool, Dict[str, Any] | None]:
|
||||
"""
|
||||
Analyze chart configuration and suggest better chart types if needed.
|
||||
|
||||
Returns:
|
||||
Tuple of (is_appropriate, suggestion_info)
|
||||
"""
|
||||
try:
|
||||
if isinstance(config, XYChartConfig):
|
||||
return ChartTypeSuggester._analyze_xy_chart(config)
|
||||
elif isinstance(config, TableChartConfig):
|
||||
return ChartTypeSuggester._analyze_table_chart(config)
|
||||
else:
|
||||
return True, None
|
||||
except Exception as e:
|
||||
logger.warning("Chart type analysis failed: %s", e)
|
||||
return True, None # Don't block on suggestion failures
|
||||
|
||||
@staticmethod
|
||||
def _analyze_xy_chart(
|
||||
config: XYChartConfig,
|
||||
) -> Tuple[bool, Dict[str, Any] | None]:
|
||||
"""Analyze XY chart appropriateness."""
|
||||
issues = []
|
||||
suggestions = []
|
||||
|
||||
x_analysis = ChartTypeSuggester._analyze_x_axis(config.x.name)
|
||||
y_analysis = ChartTypeSuggester._analyze_y_axis(config.y)
|
||||
|
||||
# Check chart type specific issues
|
||||
chart_issues, chart_suggestions = ChartTypeSuggester._check_chart_type_issues(
|
||||
config, x_analysis, y_analysis
|
||||
)
|
||||
issues.extend(chart_issues)
|
||||
suggestions.extend(chart_suggestions)
|
||||
|
||||
# Add general suggestions
|
||||
general_suggestions = ChartTypeSuggester._get_general_suggestions(
|
||||
x_analysis, y_analysis
|
||||
)
|
||||
suggestions.extend(general_suggestions)
|
||||
|
||||
if issues:
|
||||
return False, {
|
||||
"issues": issues,
|
||||
"suggestions": suggestions,
|
||||
"recommended_types": ChartTypeSuggester._get_recommended_types(
|
||||
x_analysis["is_temporal"],
|
||||
x_analysis["is_categorical"],
|
||||
y_analysis["has_count"],
|
||||
y_analysis["num_metrics"],
|
||||
),
|
||||
}
|
||||
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def _analyze_x_axis(x_name: str) -> Dict[str, Any]:
|
||||
"""Analyze X-axis characteristics."""
|
||||
x_name_lower = x_name.lower()
|
||||
return {
|
||||
"is_temporal": any(
|
||||
t in x_name_lower
|
||||
for t in [
|
||||
"date",
|
||||
"time",
|
||||
"year",
|
||||
"month",
|
||||
"day",
|
||||
"hour",
|
||||
"created",
|
||||
"updated",
|
||||
]
|
||||
),
|
||||
"is_categorical": any(
|
||||
c in x_name_lower
|
||||
for c in [
|
||||
"category",
|
||||
"type",
|
||||
"status",
|
||||
"department",
|
||||
"region",
|
||||
"country",
|
||||
"state",
|
||||
]
|
||||
),
|
||||
"is_id": any(i in x_name_lower for i in ["id", "uuid", "guid", "key"]),
|
||||
"name": x_name,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _analyze_y_axis(y_columns: List[ColumnRef]) -> Dict[str, Any]:
|
||||
"""Analyze Y-axis characteristics."""
|
||||
return {
|
||||
"has_count": any(
|
||||
col.aggregate in ["COUNT", "COUNT_DISTINCT"] for col in y_columns
|
||||
),
|
||||
"num_metrics": len(y_columns),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _check_chart_type_issues(
|
||||
config: XYChartConfig, x_analysis: Dict[str, Any], y_analysis: Dict[str, Any]
|
||||
) -> Tuple[List[str], List[str]]:
|
||||
"""Check for chart type specific issues."""
|
||||
issues = []
|
||||
suggestions = []
|
||||
|
||||
# Extract analysis values
|
||||
x_is_temporal = x_analysis["is_temporal"]
|
||||
x_is_categorical = x_analysis["is_categorical"]
|
||||
x_is_id = x_analysis["is_id"]
|
||||
num_metrics = y_analysis["num_metrics"]
|
||||
|
||||
# Check chart type specific issues by delegating to helper methods
|
||||
if config.kind == "line":
|
||||
line_issues, line_suggestions = ChartTypeSuggester._check_line_chart_issues(
|
||||
config, x_is_temporal, x_is_categorical, x_is_id
|
||||
)
|
||||
issues.extend(line_issues)
|
||||
suggestions.extend(line_suggestions)
|
||||
elif config.kind == "scatter":
|
||||
(
|
||||
scatter_issues,
|
||||
scatter_suggestions,
|
||||
) = ChartTypeSuggester._check_scatter_chart_issues(
|
||||
config, x_is_categorical, num_metrics
|
||||
)
|
||||
issues.extend(scatter_issues)
|
||||
suggestions.extend(scatter_suggestions)
|
||||
elif config.kind == "area":
|
||||
area_issues, area_suggestions = ChartTypeSuggester._check_area_chart_issues(
|
||||
config, x_is_temporal
|
||||
)
|
||||
issues.extend(area_issues)
|
||||
suggestions.extend(area_suggestions)
|
||||
elif config.kind == "bar":
|
||||
bar_issues, bar_suggestions = ChartTypeSuggester._check_bar_chart_issues(
|
||||
config, x_is_id
|
||||
)
|
||||
issues.extend(bar_issues)
|
||||
suggestions.extend(bar_suggestions)
|
||||
|
||||
return issues, suggestions
|
||||
|
||||
@staticmethod
|
||||
def _check_line_chart_issues(
|
||||
config: XYChartConfig,
|
||||
x_is_temporal: bool,
|
||||
x_is_categorical: bool,
|
||||
x_is_id: bool,
|
||||
) -> Tuple[List[str], List[str]]:
|
||||
"""Check line chart specific issues."""
|
||||
issues = []
|
||||
suggestions = []
|
||||
|
||||
if not x_is_temporal and x_is_categorical:
|
||||
issues.append(
|
||||
f"Line chart with categorical X-axis '{config.x.name}' may not "
|
||||
f"show meaningful trends"
|
||||
)
|
||||
suggestions.extend(
|
||||
[
|
||||
"Consider using a bar chart for categorical comparisons",
|
||||
"Line charts work best with temporal or continuous data",
|
||||
]
|
||||
)
|
||||
elif x_is_id:
|
||||
issues.append(
|
||||
f"Line chart with ID field '{config.x.name}' on X-axis will not "
|
||||
f"show meaningful patterns"
|
||||
)
|
||||
suggestions.extend(
|
||||
[
|
||||
"Use a table to display individual records",
|
||||
"Or aggregate the data by a meaningful dimension",
|
||||
]
|
||||
)
|
||||
|
||||
return issues, suggestions
|
||||
|
||||
@staticmethod
|
||||
def _check_scatter_chart_issues(
|
||||
config: XYChartConfig, x_is_categorical: bool, num_metrics: int
|
||||
) -> Tuple[List[str], List[str]]:
|
||||
"""Check scatter chart specific issues."""
|
||||
issues = []
|
||||
suggestions = []
|
||||
|
||||
if x_is_categorical:
|
||||
issues.append(
|
||||
f"Scatter plot with categorical X-axis '{config.x.name}' may not "
|
||||
f"effectively show correlations"
|
||||
)
|
||||
suggestions.extend(
|
||||
[
|
||||
"Scatter plots work best with two continuous variables",
|
||||
"Consider a bar chart for categorical vs numeric data",
|
||||
]
|
||||
)
|
||||
if num_metrics > 1:
|
||||
issues.append("Scatter plots with multiple Y metrics can be confusing")
|
||||
suggestions.extend(
|
||||
[
|
||||
"Consider using only one Y metric for clarity",
|
||||
"Or use a line/bar chart to compare multiple metrics",
|
||||
]
|
||||
)
|
||||
|
||||
return issues, suggestions
|
||||
|
||||
@staticmethod
|
||||
def _check_area_chart_issues(
|
||||
config: XYChartConfig, x_is_temporal: bool
|
||||
) -> Tuple[List[str], List[str]]:
|
||||
"""Check area chart specific issues."""
|
||||
issues = []
|
||||
suggestions = []
|
||||
|
||||
if not x_is_temporal:
|
||||
issues.append(
|
||||
f"Area chart with non-temporal X-axis '{config.x.name}' may be "
|
||||
f"misleading"
|
||||
)
|
||||
suggestions.extend(
|
||||
[
|
||||
"Area charts imply cumulative or part-to-whole relationships over "
|
||||
"time",
|
||||
"Consider a stacked bar chart for categorical data",
|
||||
]
|
||||
)
|
||||
|
||||
# Check for potential negative values
|
||||
for col in config.y:
|
||||
if any(term in col.name.lower() for term in ["loss", "debt", "negative"]):
|
||||
issues.append(
|
||||
f"Area chart with potentially negative values in '{col.name}' "
|
||||
f"can create visual confusion"
|
||||
)
|
||||
suggestions.extend(
|
||||
[
|
||||
"Use a line chart for data that can go negative",
|
||||
"Or ensure all values are positive before using area chart",
|
||||
]
|
||||
)
|
||||
|
||||
return issues, suggestions
|
||||
|
||||
@staticmethod
|
||||
def _check_bar_chart_issues(
|
||||
config: XYChartConfig, x_is_id: bool
|
||||
) -> Tuple[List[str], List[str]]:
|
||||
"""Check bar chart specific issues."""
|
||||
issues = []
|
||||
suggestions = []
|
||||
|
||||
if x_is_id:
|
||||
issues.append(
|
||||
f"Bar chart with ID field '{config.x.name}' may create too many bars"
|
||||
)
|
||||
suggestions.extend(
|
||||
[
|
||||
"Consider aggregating by a higher-level category",
|
||||
"Or use filters to limit the number of bars displayed",
|
||||
]
|
||||
)
|
||||
|
||||
return issues, suggestions
|
||||
|
||||
@staticmethod
|
||||
def _get_general_suggestions(
|
||||
x_analysis: Dict[str, Any], y_analysis: Dict[str, Any]
|
||||
) -> List[str]:
|
||||
"""Get general suggestions based on data patterns."""
|
||||
suggestions = []
|
||||
x_is_temporal = x_analysis["is_temporal"]
|
||||
x_is_categorical = x_analysis["is_categorical"]
|
||||
has_count = y_analysis["has_count"]
|
||||
num_metrics = y_analysis["num_metrics"]
|
||||
|
||||
if has_count and x_is_categorical:
|
||||
suggestions.append(
|
||||
"This looks like frequency analysis - bar charts work well for counts "
|
||||
"by category"
|
||||
)
|
||||
elif x_is_temporal and num_metrics == 1:
|
||||
suggestions.append(
|
||||
"Single metric over time - line charts are ideal for showing trends"
|
||||
)
|
||||
elif x_is_temporal and num_metrics > 3:
|
||||
suggestions.append(
|
||||
"Many metrics over time - consider focusing on 2-3 key metrics for "
|
||||
"clarity"
|
||||
)
|
||||
|
||||
return suggestions
|
||||
|
||||
@staticmethod
|
||||
def _analyze_table_chart(
|
||||
config: TableChartConfig,
|
||||
) -> Tuple[bool, Dict[str, Any] | None]:
|
||||
"""Analyze table chart appropriateness."""
|
||||
issues = []
|
||||
suggestions = []
|
||||
|
||||
# Count different column types
|
||||
raw_columns = sum(1 for col in config.columns if not col.aggregate)
|
||||
metric_columns = sum(1 for col in config.columns if col.aggregate)
|
||||
total_columns = len(config.columns)
|
||||
|
||||
# Check if data might be better visualized
|
||||
if metric_columns > 0 and raw_columns <= 2:
|
||||
# Mostly metrics with few dimensions - could be visualized
|
||||
issues.append(
|
||||
"Table with mostly aggregated metrics could be visualized as a chart"
|
||||
)
|
||||
suggestions.append("Consider a bar chart to compare metric values visually")
|
||||
suggestions.append("Or use a line chart if there's a time dimension")
|
||||
|
||||
# Check for ID-heavy tables
|
||||
id_columns = sum(
|
||||
1
|
||||
for col in config.columns
|
||||
if any(i in col.name.lower() for i in ["id", "uuid", "guid", "key"])
|
||||
)
|
||||
if id_columns > total_columns / 2:
|
||||
suggestions.append(
|
||||
"Table appears to be ID-heavy - ensure this is for detailed record "
|
||||
"inspection"
|
||||
)
|
||||
suggestions.append(
|
||||
"For analysis, consider aggregating by meaningful dimensions instead"
|
||||
)
|
||||
|
||||
# Very wide tables
|
||||
if total_columns > 10:
|
||||
issues.append(
|
||||
f"Table with {total_columns} columns may be difficult to read"
|
||||
)
|
||||
suggestions.append("Consider showing only the most important columns")
|
||||
suggestions.append("Or break into multiple focused views")
|
||||
|
||||
if issues:
|
||||
return False, {
|
||||
"issues": issues,
|
||||
"suggestions": suggestions,
|
||||
"recommended_types": ["table", "pivot_table"]
|
||||
if metric_columns > 0
|
||||
else ["table"],
|
||||
}
|
||||
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def _get_recommended_types(
|
||||
x_is_temporal: bool, x_is_categorical: bool, has_count: bool, num_metrics: int
|
||||
) -> List[str]:
|
||||
"""Get recommended chart types based on data characteristics."""
|
||||
recommendations = []
|
||||
|
||||
if x_is_temporal:
|
||||
recommendations.extend(["line", "area", "bar"])
|
||||
if num_metrics == 1:
|
||||
recommendations.append("scatter") # For trend analysis
|
||||
elif x_is_categorical:
|
||||
recommendations.extend(["bar", "table"])
|
||||
if has_count and num_metrics == 1:
|
||||
recommendations.append("pie") # For proportion analysis
|
||||
else:
|
||||
# Continuous or unclear X-axis
|
||||
recommendations.extend(["scatter", "line", "table"])
|
||||
|
||||
# Always include table as fallback
|
||||
if "table" not in recommendations:
|
||||
recommendations.append("table")
|
||||
|
||||
return recommendations
|
||||
|
||||
@staticmethod
|
||||
def get_chart_type_description(chart_type: str) -> str:
|
||||
"""Get a description of when to use each chart type."""
|
||||
descriptions = {
|
||||
"line": "Best for showing trends over time or continuous data",
|
||||
"bar": "Ideal for comparing values across categories",
|
||||
"area": "Shows cumulative totals and part-to-whole relationships over time",
|
||||
"scatter": "Reveals correlations between two continuous variables",
|
||||
"table": "Displays detailed data or many dimensions at once",
|
||||
"pie": "Shows proportions of a whole (use sparingly, max 5-7 slices)",
|
||||
"pivot_table": "Summarizes data across multiple dimensions",
|
||||
}
|
||||
return descriptions.get(
|
||||
chart_type, f"Visualizes data using {chart_type} format"
|
||||
)
|
||||
@@ -0,0 +1,225 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Format-type compatibility validation to prevent misleading data presentation.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
from typing import List, Tuple
|
||||
|
||||
from superset.mcp_service.chart.schemas import ColumnRef, XYChartConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FormatTypeValidator:
|
||||
"""
|
||||
Validates that format strings are appropriate for the data type and aggregation.
|
||||
Prevents issues like currency formatting on COUNT data or percentage on absolute
|
||||
values.
|
||||
"""
|
||||
|
||||
# Format patterns and their appropriate uses
|
||||
CURRENCY_PATTERNS = [
|
||||
r"\$", # Dollar sign
|
||||
r"€", # Euro
|
||||
r"£", # Pound
|
||||
r"¥", # Yen
|
||||
r"[,.]2f", # Two decimal places (common for currency)
|
||||
r"\$[,.]", # Dollar with thousands separator
|
||||
]
|
||||
|
||||
PERCENTAGE_PATTERNS = [
|
||||
r"%", # Percentage sign
|
||||
r"\.0%", # Percentage with no decimals
|
||||
r"\.1%", # Percentage with 1 decimal
|
||||
r"\.2%", # Percentage with 2 decimals
|
||||
]
|
||||
|
||||
INTEGER_PATTERNS = [
|
||||
r"\.0f", # No decimals
|
||||
r",d", # Integer with thousands separator
|
||||
r"[,.]0f", # Integer format variations
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def validate_format_compatibility(
|
||||
config: XYChartConfig,
|
||||
) -> Tuple[bool, List[str] | None]:
|
||||
"""
|
||||
Validate that axis formats are appropriate for the data types.
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, warnings_list)
|
||||
"""
|
||||
warnings = []
|
||||
|
||||
# Validate Y-axis format against metrics
|
||||
if config.y_axis and config.y_axis.format:
|
||||
y_warnings = FormatTypeValidator._validate_y_axis_format(
|
||||
config.y_axis.format, config.y
|
||||
)
|
||||
warnings.extend(y_warnings)
|
||||
|
||||
# Validate X-axis format (usually temporal or categorical)
|
||||
if config.x_axis and config.x_axis.format:
|
||||
x_warnings = FormatTypeValidator._validate_x_axis_format(
|
||||
config.x_axis.format, config.x
|
||||
)
|
||||
warnings.extend(x_warnings)
|
||||
|
||||
return len(warnings) == 0, warnings if warnings else None
|
||||
|
||||
@staticmethod
|
||||
def _validate_y_axis_format(
|
||||
format_string: str, y_columns: List[ColumnRef]
|
||||
) -> List[str]:
|
||||
"""Validate Y-axis format against the metrics."""
|
||||
warnings = []
|
||||
|
||||
warnings.extend(
|
||||
FormatTypeValidator._check_currency_format_issues(format_string, y_columns)
|
||||
)
|
||||
warnings.extend(
|
||||
FormatTypeValidator._check_percentage_format_issues(
|
||||
format_string, y_columns
|
||||
)
|
||||
)
|
||||
warnings.extend(
|
||||
FormatTypeValidator._check_decimal_format_issues(format_string, y_columns)
|
||||
)
|
||||
|
||||
return warnings
|
||||
|
||||
@staticmethod
|
||||
def _check_currency_format_issues(
|
||||
format_string: str, y_columns: List[ColumnRef]
|
||||
) -> List[str]:
|
||||
"""Check for currency format issues."""
|
||||
warnings = []
|
||||
if FormatTypeValidator._is_currency_format(format_string):
|
||||
for col in y_columns:
|
||||
if col.aggregate in ["COUNT", "COUNT_DISTINCT"]:
|
||||
warnings.append(
|
||||
f"Currency format '{format_string}' applied to {col.aggregate} "
|
||||
f"of '{col.name}'. COUNT operations return whole numbers, not "
|
||||
f"currency values. Consider using integer format like ',"
|
||||
f"d' instead."
|
||||
)
|
||||
return warnings
|
||||
|
||||
@staticmethod
|
||||
def _check_percentage_format_issues(
|
||||
format_string: str, y_columns: List[ColumnRef]
|
||||
) -> List[str]:
|
||||
"""Check for percentage format issues."""
|
||||
warnings = []
|
||||
if FormatTypeValidator._is_percentage_format(format_string):
|
||||
for col in y_columns:
|
||||
if col.aggregate in ["SUM", "COUNT", "COUNT_DISTINCT"]:
|
||||
label = col.label or f"{col.aggregate}({col.name})"
|
||||
warnings.append(
|
||||
f"Percentage format '{format_string}' applied to "
|
||||
f"{col.aggregate} of '{col.name}'. This will multiply values "
|
||||
f"by 100 and add %. "
|
||||
f"If '{label}' contains absolute values (not ratios 0-1), "
|
||||
f"consider using a numeric format instead."
|
||||
)
|
||||
return warnings
|
||||
|
||||
@staticmethod
|
||||
def _check_decimal_format_issues(
|
||||
format_string: str, y_columns: List[ColumnRef]
|
||||
) -> List[str]:
|
||||
"""Check for decimal format issues."""
|
||||
warnings = []
|
||||
if "." in format_string and any(char.isdigit() for char in format_string):
|
||||
decimal_places = FormatTypeValidator._get_decimal_places(format_string)
|
||||
if decimal_places and decimal_places > 0:
|
||||
for col in y_columns:
|
||||
if col.aggregate in ["COUNT", "COUNT_DISTINCT"]:
|
||||
warnings.append(
|
||||
f"Decimal format '{format_string}' applied to "
|
||||
f"{col.aggregate} of '{col.name}'. COUNT operations "
|
||||
f"always return "
|
||||
f"integers. Consider using integer format like ',"
|
||||
f"d' or '.0f' instead."
|
||||
)
|
||||
return warnings
|
||||
|
||||
@staticmethod
|
||||
def _validate_x_axis_format(format_string: str, x_column: ColumnRef) -> List[str]:
|
||||
"""Validate X-axis format appropriateness."""
|
||||
warnings = []
|
||||
|
||||
# Currency format on X-axis is almost always wrong
|
||||
if FormatTypeValidator._is_currency_format(format_string):
|
||||
warnings.append(
|
||||
f"Currency format '{format_string}' applied to X-axis '"
|
||||
f"{x_column.name}'. "
|
||||
f"X-axis typically shows categories, time, or dimensions, "
|
||||
f"not currency. "
|
||||
f"Consider removing the format or using a date/category format."
|
||||
)
|
||||
|
||||
# Percentage format on X-axis is unusual
|
||||
elif FormatTypeValidator._is_percentage_format(format_string):
|
||||
warnings.append(
|
||||
f"Percentage format '{format_string}' applied to X-axis '"
|
||||
f"{x_column.name}'. "
|
||||
f"This is unusual for axis labels. Consider if this is intentional."
|
||||
)
|
||||
|
||||
return warnings
|
||||
|
||||
@staticmethod
|
||||
def _is_currency_format(format_string: str) -> bool:
|
||||
"""Check if format string represents currency."""
|
||||
return any(
|
||||
re.search(pattern, format_string, re.IGNORECASE)
|
||||
for pattern in FormatTypeValidator.CURRENCY_PATTERNS
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_percentage_format(format_string: str) -> bool:
|
||||
"""Check if format string represents percentage."""
|
||||
return any(
|
||||
re.search(pattern, format_string)
|
||||
for pattern in FormatTypeValidator.PERCENTAGE_PATTERNS
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_decimal_places(format_string: str) -> int | None:
|
||||
"""Extract number of decimal places from format string."""
|
||||
if match := re.search(r"\.(\d+)f", format_string):
|
||||
return int(match.group(1))
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def suggest_format(column: ColumnRef) -> str:
|
||||
"""Suggest appropriate format based on column and aggregation."""
|
||||
if column.aggregate in ["COUNT", "COUNT_DISTINCT"]:
|
||||
return ",d" # Integer with thousands separator
|
||||
elif column.aggregate in ["AVG", "STDDEV", "VAR"]:
|
||||
return ",.2f" # Two decimals for statistical measures
|
||||
elif column.aggregate in ["SUM", "MIN", "MAX"]:
|
||||
# Could be currency or regular number, default to flexible
|
||||
return ",.2f" # Two decimals with thousands separator
|
||||
else:
|
||||
return "" # Let Superset decide
|
||||
307
superset/mcp_service/chart/validation/schema_validator.py
Normal file
307
superset/mcp_service/chart/validation/schema_validator.py
Normal file
@@ -0,0 +1,307 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Unified schema validation for chart configurations.
|
||||
Consolidates pre-validation, schema validation, and error enhancement.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, Tuple
|
||||
|
||||
from pydantic import ValidationError as PydanticValidationError
|
||||
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
GenerateChartRequest,
|
||||
)
|
||||
from superset.mcp_service.common.error_schemas import ChartGenerationError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SchemaValidator:
|
||||
"""Unified schema validator with pre-validation and enhanced error messages."""
|
||||
|
||||
@staticmethod
|
||||
def validate_request(
|
||||
request_data: Dict[str, Any],
|
||||
) -> Tuple[bool, GenerateChartRequest | None, ChartGenerationError | None]:
|
||||
"""
|
||||
Validate request data with pre-validation and enhanced error handling.
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, parsed_request, error)
|
||||
"""
|
||||
# Pre-validate to catch common issues early
|
||||
is_valid, error = SchemaValidator._pre_validate(request_data)
|
||||
if not is_valid:
|
||||
return False, None, error
|
||||
|
||||
# Try Pydantic validation
|
||||
try:
|
||||
request = GenerateChartRequest(**request_data)
|
||||
return True, request, None
|
||||
except PydanticValidationError as e:
|
||||
# Enhance the error message
|
||||
error = SchemaValidator._enhance_validation_error(e, request_data)
|
||||
return False, None, error
|
||||
|
||||
@staticmethod
|
||||
def _pre_validate(
|
||||
data: Dict[str, Any],
|
||||
) -> Tuple[bool, ChartGenerationError | None]:
|
||||
"""Pre-validate request data before Pydantic processing."""
|
||||
if not isinstance(data, dict):
|
||||
return False, ChartGenerationError(
|
||||
error_type="invalid_request_format",
|
||||
message="Request must be a JSON object",
|
||||
details="The request body must be a valid JSON object, not a string "
|
||||
"or array",
|
||||
suggestions=[
|
||||
"Ensure you're sending a JSON object with 'dataset_id' and "
|
||||
"'config' fields",
|
||||
"Check that Content-Type header is set to 'application/json'",
|
||||
],
|
||||
error_code="INVALID_REQUEST_FORMAT",
|
||||
)
|
||||
|
||||
# Check for required top-level fields
|
||||
if "dataset_id" not in data:
|
||||
return False, ChartGenerationError(
|
||||
error_type="missing_dataset_id",
|
||||
message="Missing required field: dataset_id",
|
||||
details="The 'dataset_id' field is required to identify which dataset "
|
||||
"to use",
|
||||
suggestions=[
|
||||
"Add 'dataset_id' field with the ID of your dataset",
|
||||
"Use list_datasets tool to find available dataset IDs",
|
||||
"Example: {'dataset_id': 1, 'config': {...}}",
|
||||
],
|
||||
error_code="MISSING_DATASET_ID",
|
||||
)
|
||||
|
||||
if "config" not in data:
|
||||
return False, ChartGenerationError(
|
||||
error_type="missing_config",
|
||||
message="Missing required field: config",
|
||||
details="The 'config' field is required to specify chart configuration",
|
||||
suggestions=[
|
||||
"Add 'config' field with chart type and settings",
|
||||
"Example: {'dataset_id': 1, 'config': {'chart_type': 'xy', ...}}",
|
||||
],
|
||||
error_code="MISSING_CONFIG",
|
||||
)
|
||||
|
||||
config = data.get("config", {})
|
||||
if not isinstance(config, dict):
|
||||
return False, ChartGenerationError(
|
||||
error_type="invalid_config_format",
|
||||
message="Config must be a JSON object",
|
||||
details="The 'config' field must be a valid JSON object with chart "
|
||||
"settings",
|
||||
suggestions=[
|
||||
"Ensure config is an object, not a string or array",
|
||||
"Example: 'config': {'chart_type': 'xy', 'x': {...}, 'y': [...]}",
|
||||
],
|
||||
error_code="INVALID_CONFIG_FORMAT",
|
||||
)
|
||||
|
||||
# Check chart_type early
|
||||
chart_type = config.get("chart_type")
|
||||
if not chart_type:
|
||||
return False, ChartGenerationError(
|
||||
error_type="missing_chart_type",
|
||||
message="Missing required field: chart_type",
|
||||
details="Chart configuration must specify 'chart_type' as either 'xy' "
|
||||
"or 'table'",
|
||||
suggestions=[
|
||||
"Add 'chart_type': 'xy' for line/bar/area/scatter charts",
|
||||
"Add 'chart_type': 'table' for table visualizations",
|
||||
"Example: 'config': {'chart_type': 'xy', ...}",
|
||||
],
|
||||
error_code="MISSING_CHART_TYPE",
|
||||
)
|
||||
|
||||
if chart_type not in ["xy", "table"]:
|
||||
return False, ChartGenerationError(
|
||||
error_type="invalid_chart_type",
|
||||
message=f"Invalid chart_type: '{chart_type}'",
|
||||
details=f"Chart type '{chart_type}' is not supported. Must be 'xy' or "
|
||||
f"'table'",
|
||||
suggestions=[
|
||||
"Use 'chart_type': 'xy' for line, bar, area, or scatter charts",
|
||||
"Use 'chart_type': 'table' for tabular data display",
|
||||
"Check spelling and ensure lowercase",
|
||||
],
|
||||
error_code="INVALID_CHART_TYPE",
|
||||
)
|
||||
|
||||
# Pre-validate structure based on chart type
|
||||
if chart_type == "xy":
|
||||
return SchemaValidator._pre_validate_xy_config(config)
|
||||
elif chart_type == "table":
|
||||
return SchemaValidator._pre_validate_table_config(config)
|
||||
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def _pre_validate_xy_config(
|
||||
config: Dict[str, Any],
|
||||
) -> Tuple[bool, ChartGenerationError | None]:
|
||||
"""Pre-validate XY chart configuration."""
|
||||
missing_fields = []
|
||||
|
||||
if "x" not in config:
|
||||
missing_fields.append("'x' (X-axis column)")
|
||||
if "y" not in config:
|
||||
missing_fields.append("'y' (Y-axis metrics)")
|
||||
|
||||
if missing_fields:
|
||||
return False, ChartGenerationError(
|
||||
error_type="missing_xy_fields",
|
||||
message=f"XY chart missing required "
|
||||
f"fields: {', '.join(missing_fields)}",
|
||||
details="XY charts require both X-axis (dimension) and Y-axis ("
|
||||
"metrics) specifications",
|
||||
suggestions=[
|
||||
"Add 'x' field: {'name': 'column_name'} for X-axis",
|
||||
"Add 'y' field: [{'name': 'metric_column', 'aggregate': 'SUM'}] "
|
||||
"for Y-axis",
|
||||
"Example: {'chart_type': 'xy', 'x': {'name': 'date'}, "
|
||||
"'y': [{'name': 'sales', 'aggregate': 'SUM'}]}",
|
||||
],
|
||||
error_code="MISSING_XY_FIELDS",
|
||||
)
|
||||
|
||||
# Validate Y is a list
|
||||
if not isinstance(config.get("y", []), list):
|
||||
return False, ChartGenerationError(
|
||||
error_type="invalid_y_format",
|
||||
message="Y-axis must be a list of metrics",
|
||||
details="The 'y' field must be an array of metric specifications",
|
||||
suggestions=[
|
||||
"Wrap Y-axis metric in array: 'y': [{'name': 'column', "
|
||||
"'aggregate': 'SUM'}]",
|
||||
"Multiple metrics supported: 'y': [metric1, metric2, ...]",
|
||||
],
|
||||
error_code="INVALID_Y_FORMAT",
|
||||
)
|
||||
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def _pre_validate_table_config(
|
||||
config: Dict[str, Any],
|
||||
) -> Tuple[bool, ChartGenerationError | None]:
|
||||
"""Pre-validate table chart configuration."""
|
||||
if "columns" not in config:
|
||||
return False, ChartGenerationError(
|
||||
error_type="missing_columns",
|
||||
message="Table chart missing required field: columns",
|
||||
details="Table charts require a 'columns' array to specify which "
|
||||
"columns to display",
|
||||
suggestions=[
|
||||
"Add 'columns' field with array of column specifications",
|
||||
"Example: 'columns': [{'name': 'product'}, {'name': 'sales', "
|
||||
"'aggregate': 'SUM'}]",
|
||||
"Each column can have optional 'aggregate' for metrics",
|
||||
],
|
||||
error_code="MISSING_COLUMNS",
|
||||
)
|
||||
|
||||
if not isinstance(config.get("columns", []), list):
|
||||
return False, ChartGenerationError(
|
||||
error_type="invalid_columns_format",
|
||||
message="Columns must be a list",
|
||||
details="The 'columns' field must be an array of column specifications",
|
||||
suggestions=[
|
||||
"Ensure columns is an array: 'columns': [...]",
|
||||
"Each column should be an object with 'name' field",
|
||||
],
|
||||
error_code="INVALID_COLUMNS_FORMAT",
|
||||
)
|
||||
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def _enhance_validation_error(
|
||||
error: PydanticValidationError, request_data: Dict[str, Any]
|
||||
) -> ChartGenerationError:
|
||||
"""Convert Pydantic validation errors to user-friendly messages."""
|
||||
errors = error.errors()
|
||||
|
||||
# Check for discriminated union errors (generic "'table' was expected")
|
||||
for err in errors:
|
||||
if err.get("type") == "union_tag_invalid" or "discriminator" in str(
|
||||
err.get("ctx", {})
|
||||
):
|
||||
# This is the generic union error - provide better message
|
||||
config = request_data.get("config", {})
|
||||
chart_type = config.get("chart_type", "unknown")
|
||||
|
||||
if chart_type == "xy":
|
||||
return ChartGenerationError(
|
||||
error_type="xy_validation_error",
|
||||
message="XY chart configuration validation failed",
|
||||
details="The XY chart configuration is missing required "
|
||||
"fields or has invalid structure",
|
||||
suggestions=[
|
||||
"Ensure 'x' field exists with {'name': 'column_name'}",
|
||||
"Ensure 'y' field is an array: [{'name': 'metric', "
|
||||
"'aggregate': 'SUM'}]",
|
||||
"Check that all column names are strings",
|
||||
"Verify aggregate functions are valid: SUM, COUNT, AVG, "
|
||||
"MIN, MAX",
|
||||
],
|
||||
error_code="XY_VALIDATION_ERROR",
|
||||
)
|
||||
elif chart_type == "table":
|
||||
return ChartGenerationError(
|
||||
error_type="table_validation_error",
|
||||
message="Table chart configuration validation failed",
|
||||
details="The table chart configuration is missing required "
|
||||
"fields or has invalid structure",
|
||||
suggestions=[
|
||||
"Ensure 'columns' field is an array of column "
|
||||
"specifications",
|
||||
"Each column needs {'name': 'column_name'}",
|
||||
"Optional: add 'aggregate' for metrics",
|
||||
"Example: 'columns': [{'name': 'product'}, {'name': "
|
||||
"'sales', 'aggregate': 'SUM'}]",
|
||||
],
|
||||
error_code="TABLE_VALIDATION_ERROR",
|
||||
)
|
||||
|
||||
# Default enhanced error
|
||||
error_details = []
|
||||
for err in errors[:3]: # Show first 3 errors
|
||||
loc = " -> ".join(str(location) for location in err.get("loc", []))
|
||||
msg = err.get("msg", "Validation failed")
|
||||
error_details.append(f"{loc}: {msg}")
|
||||
|
||||
return ChartGenerationError(
|
||||
error_type="validation_error",
|
||||
message="Chart configuration validation failed",
|
||||
details="; ".join(error_details),
|
||||
suggestions=[
|
||||
"Check that all required fields are present",
|
||||
"Ensure field types match the schema",
|
||||
"Use get_dataset_info to verify column names",
|
||||
"Refer to the API documentation for field requirements",
|
||||
],
|
||||
error_code="VALIDATION_ERROR",
|
||||
)
|
||||
33
superset/mcp_service/commands/create_form_data.py
Normal file
33
superset/mcp_service/commands/create_form_data.py
Normal file
@@ -0,0 +1,33 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
MCP-specific form data command that extends the base CreateFormDataCommand
|
||||
"""
|
||||
|
||||
from superset.commands.explore.form_data.create import CreateFormDataCommand
|
||||
from superset.utils.core import get_user_id
|
||||
|
||||
|
||||
class MCPCreateFormDataCommand(CreateFormDataCommand):
|
||||
"""
|
||||
MCP-specific CreateFormDataCommand that uses user_id instead of session._id
|
||||
"""
|
||||
|
||||
def _get_session_id(self) -> str:
|
||||
"""Override to use user_id instead of Flask session for MCP context."""
|
||||
return str(get_user_id())
|
||||
103
superset/mcp_service/common/error_schemas.py
Normal file
103
superset/mcp_service/common/error_schemas.py
Normal file
@@ -0,0 +1,103 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Enhanced error schemas for MCP chart generation with contextual information
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class ColumnSuggestion(BaseModel):
|
||||
"""Suggested column with context"""
|
||||
|
||||
name: str = Field(..., description="Column name")
|
||||
type: str = Field(..., description="Column data type")
|
||||
similarity_score: float = Field(..., description="Similarity score (0-1)")
|
||||
description: str | None = Field(None, description="Column description")
|
||||
|
||||
|
||||
class ValidationError(BaseModel):
|
||||
"""Individual validation error with context"""
|
||||
|
||||
field: str = Field(..., description="Field that failed validation")
|
||||
provided_value: Any = Field(..., description="Value that was provided")
|
||||
error_type: str = Field(..., description="Type of validation error")
|
||||
message: str = Field(..., description="Human-readable error message")
|
||||
suggestions: List[ColumnSuggestion] = Field(
|
||||
default_factory=list, description="Suggested alternatives"
|
||||
)
|
||||
|
||||
|
||||
class DatasetContext(BaseModel):
|
||||
"""Dataset information for error context"""
|
||||
|
||||
id: int = Field(..., description="Dataset ID")
|
||||
table_name: str = Field(..., description="Table name")
|
||||
schema: str | None = Field(None, description="Schema name")
|
||||
database_name: str = Field(..., description="Database name")
|
||||
available_columns: List[Dict[str, Any]] = Field(
|
||||
default_factory=list, description="Available columns with metadata"
|
||||
)
|
||||
available_metrics: List[Dict[str, Any]] = Field(
|
||||
default_factory=list, description="Available metrics with metadata"
|
||||
)
|
||||
|
||||
|
||||
class ChartGenerationError(BaseModel):
|
||||
"""Enhanced error response for chart generation failures"""
|
||||
|
||||
error_type: str = Field(
|
||||
..., description="Type of error (validation, execution, etc.)"
|
||||
)
|
||||
message: str = Field(..., description="High-level error message")
|
||||
details: str = Field(..., description="Detailed error explanation")
|
||||
validation_errors: List[ValidationError] = Field(
|
||||
default_factory=list, description="Specific field validation errors"
|
||||
)
|
||||
dataset_context: DatasetContext | None = Field(
|
||||
None, description="Dataset information for context"
|
||||
)
|
||||
query_info: Dict[str, Any] | None = Field(
|
||||
None, description="Query execution details"
|
||||
)
|
||||
suggestions: List[str] = Field(
|
||||
default_factory=list, description="Actionable suggestions to fix the error"
|
||||
)
|
||||
help_url: str | None = Field(
|
||||
None, description="URL to documentation for this error type"
|
||||
)
|
||||
error_code: str | None = Field(
|
||||
None, description="Unique error code for support reference"
|
||||
)
|
||||
|
||||
|
||||
class ChartGenerationResponse(BaseModel):
|
||||
"""Enhanced chart generation response with detailed error handling"""
|
||||
|
||||
success: bool = Field(..., description="Whether chart generation succeeded")
|
||||
chart: Dict[str, Any] | None = Field(
|
||||
None, description="Chart information if successful"
|
||||
)
|
||||
error: ChartGenerationError | None = Field(
|
||||
None, description="Error details if failed"
|
||||
)
|
||||
performance: Dict[str, Any] | None = Field(None, description="Performance metadata")
|
||||
schema_version: str = Field(default="2.0", description="Response schema version")
|
||||
api_version: str = Field(default="v1", description="API version")
|
||||
21
superset/mcp_service/dashboard/prompts/__init__.py
Normal file
21
superset/mcp_service/dashboard/prompts/__init__.py
Normal file
@@ -0,0 +1,21 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Dashboard prompts for Superset MCP service"""
|
||||
|
||||
# Import prompt modules here when created
|
||||
# from . import create_dashboard_guided # noqa: F401
|
||||
21
superset/mcp_service/dashboard/resources/__init__.py
Normal file
21
superset/mcp_service/dashboard/resources/__init__.py
Normal file
@@ -0,0 +1,21 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Dashboard resources for Superset MCP service"""
|
||||
|
||||
# Import resource modules here when created
|
||||
# from . import layout_templates # noqa: F401
|
||||
470
superset/mcp_service/dashboard/schemas.py
Normal file
470
superset/mcp_service/dashboard/schemas.py
Normal file
@@ -0,0 +1,470 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Pydantic schemas for dashboard-related responses
|
||||
|
||||
This module contains Pydantic models for serializing dashboard data
|
||||
in a consistent and type-safe manner.
|
||||
|
||||
Example usage:
|
||||
# For detailed dashboard info
|
||||
dashboard_info = DashboardInfo(
|
||||
id=1,
|
||||
dashboard_title="Sales Dashboard",
|
||||
published=True,
|
||||
owners=[UserInfo(id=1, username="admin")],
|
||||
charts=[ChartInfo(id=1, slice_name="Sales Chart")]
|
||||
)
|
||||
|
||||
# For dashboard list responses
|
||||
dashboard_list = DashboardList(
|
||||
dashboards=[
|
||||
DashboardInfo(
|
||||
id=1,
|
||||
dashboard_title="Sales Dashboard",
|
||||
published=True,
|
||||
tags=[TagInfo(id=1, name="sales")]
|
||||
)
|
||||
],
|
||||
count=1,
|
||||
total_count=1,
|
||||
page=0,
|
||||
page_size=10,
|
||||
total_pages=1,
|
||||
has_next=False,
|
||||
has_previous=False,
|
||||
columns_requested=["id", "dashboard_title"],
|
||||
columns_loaded=["id", "dashboard_title", "published"],
|
||||
filters_applied={"published": True},
|
||||
pagination=PaginationInfo(
|
||||
page=0,
|
||||
page_size=10,
|
||||
total_count=1,
|
||||
total_pages=1,
|
||||
has_next=False,
|
||||
has_previous=False
|
||||
),
|
||||
timestamp=datetime.now(timezone.utc)
|
||||
)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Annotated, Any, Dict, List, Literal, TYPE_CHECKING
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator, PositiveInt
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from superset.models.dashboard import Dashboard
|
||||
|
||||
from superset.daos.base import ColumnOperator, ColumnOperatorEnum
|
||||
from superset.mcp_service.chart.schemas import ChartInfo, serialize_chart_object
|
||||
from superset.mcp_service.common.cache_schemas import MetadataCacheControl
|
||||
from superset.mcp_service.system.schemas import (
|
||||
PaginationInfo,
|
||||
RoleInfo,
|
||||
TagInfo,
|
||||
UserInfo,
|
||||
)
|
||||
|
||||
|
||||
class DashboardError(BaseModel):
|
||||
"""Error response for dashboard operations"""
|
||||
|
||||
error: str = Field(..., description="Error message")
|
||||
error_type: str = Field(..., description="Type of error")
|
||||
timestamp: str | datetime | None = Field(None, description="Error timestamp")
|
||||
|
||||
model_config = ConfigDict(ser_json_timedelta="iso8601")
|
||||
|
||||
@classmethod
|
||||
def create(cls, error: str, error_type: str) -> "DashboardError":
|
||||
"""Create a standardized DashboardError with timestamp."""
|
||||
from datetime import datetime
|
||||
|
||||
return cls(error=error, error_type=error_type, timestamp=datetime.now())
|
||||
|
||||
|
||||
def serialize_user_object(user: Any) -> UserInfo | None:
|
||||
"""Serialize a user object to UserInfo"""
|
||||
if not user:
|
||||
return None
|
||||
|
||||
return UserInfo(
|
||||
id=getattr(user, "id", None),
|
||||
username=getattr(user, "username", None),
|
||||
first_name=getattr(user, "first_name", None),
|
||||
last_name=getattr(user, "last_name", None),
|
||||
email=getattr(user, "email", None),
|
||||
active=getattr(user, "active", None),
|
||||
)
|
||||
|
||||
|
||||
def serialize_tag_object(tag: Any) -> TagInfo | None:
|
||||
"""Serialize a tag object to TagInfo"""
|
||||
if not tag:
|
||||
return None
|
||||
|
||||
return TagInfo(
|
||||
id=getattr(tag, "id", None),
|
||||
name=getattr(tag, "name", None),
|
||||
type=getattr(tag, "type", None),
|
||||
description=getattr(tag, "description", None),
|
||||
)
|
||||
|
||||
|
||||
def serialize_role_object(role: Any) -> RoleInfo | None:
|
||||
"""Serialize a role object to RoleInfo"""
|
||||
if not role:
|
||||
return None
|
||||
|
||||
return RoleInfo(
|
||||
id=getattr(role, "id", None),
|
||||
name=getattr(role, "name", None),
|
||||
permissions=[perm.name for perm in getattr(role, "permissions", [])]
|
||||
if hasattr(role, "permissions")
|
||||
else None,
|
||||
)
|
||||
|
||||
|
||||
# TODO (Phase 3+): Add DashboardAvailableFilters for
|
||||
# get_dashboard_available_filters tool
|
||||
|
||||
|
||||
class DashboardFilter(ColumnOperator):
|
||||
"""
|
||||
Filter object for dashboard listing.
|
||||
col: The column to filter on. Must be one of the allowed filter fields.
|
||||
opr: The operator to use. Must be one of the supported operators.
|
||||
value: The value to filter by (type depends on col and opr).
|
||||
"""
|
||||
|
||||
col: Literal[
|
||||
"dashboard_title",
|
||||
"published",
|
||||
"favorite",
|
||||
] = Field(
|
||||
...,
|
||||
description="Column to filter on. See get_dashboard_available_filters for "
|
||||
"allowed values.",
|
||||
)
|
||||
opr: ColumnOperatorEnum = Field(
|
||||
...,
|
||||
description="Operator to use. See get_dashboard_available_filters for "
|
||||
"allowed values.",
|
||||
)
|
||||
value: str | int | float | bool | List[str | int | float | bool] = Field(
|
||||
..., description="Value to filter by (type depends on col and opr)"
|
||||
)
|
||||
|
||||
|
||||
class ListDashboardsRequest(MetadataCacheControl):
|
||||
"""Request schema for list_dashboards with clear, unambiguous types."""
|
||||
|
||||
filters: Annotated[
|
||||
List[DashboardFilter],
|
||||
Field(
|
||||
default_factory=list,
|
||||
description="List of filter objects (column, operator, value). Each "
|
||||
"filter is an object with 'col', 'opr', and 'value' properties. "
|
||||
"Cannot be used together with 'search'.",
|
||||
),
|
||||
]
|
||||
select_columns: Annotated[
|
||||
List[str],
|
||||
Field(
|
||||
default_factory=lambda: [
|
||||
"id",
|
||||
"dashboard_title",
|
||||
"slug",
|
||||
"published",
|
||||
"changed_on",
|
||||
"created_on",
|
||||
"uuid",
|
||||
],
|
||||
description="List of columns to select. Defaults to common columns "
|
||||
"if not specified.",
|
||||
),
|
||||
]
|
||||
search: Annotated[
|
||||
str | None,
|
||||
Field(
|
||||
default=None,
|
||||
description="Text search string to match against dashboard fields. "
|
||||
"Cannot be used together with 'filters'.",
|
||||
),
|
||||
]
|
||||
order_column: Annotated[
|
||||
str | None, Field(default=None, description="Column to order results by")
|
||||
]
|
||||
order_direction: Annotated[
|
||||
Literal["asc", "desc"],
|
||||
Field(
|
||||
default="asc", description="Direction to order results ('asc' or 'desc')"
|
||||
),
|
||||
]
|
||||
page: Annotated[
|
||||
PositiveInt,
|
||||
Field(default=1, description="Page number for pagination (1-based)"),
|
||||
]
|
||||
page_size: Annotated[
|
||||
PositiveInt, Field(default=10, description="Number of items per page")
|
||||
]
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_search_and_filters(self) -> "ListDashboardsRequest":
|
||||
"""Prevent using both search and filters simultaneously to avoid query
|
||||
conflicts."""
|
||||
if self.search and self.filters:
|
||||
raise ValueError(
|
||||
"Cannot use both 'search' and 'filters' parameters simultaneously. "
|
||||
"Use either 'search' for text-based searching across multiple fields, "
|
||||
"or 'filters' for precise column-based filtering, but not both."
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class GetDashboardInfoRequest(MetadataCacheControl):
|
||||
"""Request schema for get_dashboard_info with support for ID, UUID, or slug."""
|
||||
|
||||
identifier: Annotated[
|
||||
int | str,
|
||||
Field(
|
||||
description="Dashboard identifier - can be numeric ID, UUID string, or slug"
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class DashboardInfo(BaseModel):
|
||||
id: int = Field(..., description="Dashboard ID")
|
||||
dashboard_title: str = Field(..., description="Dashboard title")
|
||||
slug: str | None = Field(None, description="Dashboard slug")
|
||||
description: str | None = Field(None, description="Dashboard description")
|
||||
css: str | None = Field(None, description="Custom CSS for the dashboard")
|
||||
certified_by: str | None = Field(None, description="Who certified the dashboard")
|
||||
certification_details: str | None = Field(None, description="Certification details")
|
||||
json_metadata: str | None = Field(
|
||||
None, description="Dashboard metadata (JSON string)"
|
||||
)
|
||||
position_json: str | None = Field(None, description="Chart positions (JSON string)")
|
||||
published: bool | None = Field(
|
||||
None, description="Whether the dashboard is published"
|
||||
)
|
||||
is_managed_externally: bool | None = Field(
|
||||
None, description="Whether managed externally"
|
||||
)
|
||||
external_url: str | None = Field(None, description="External URL")
|
||||
created_on: str | datetime | None = Field(None, description="Creation timestamp")
|
||||
changed_on: str | datetime | None = Field(
|
||||
None, description="Last modification timestamp"
|
||||
)
|
||||
created_by: str | None = Field(None, description="Dashboard creator (username)")
|
||||
changed_by: str | None = Field(None, description="Last modifier (username)")
|
||||
uuid: str | None = Field(None, description="Dashboard UUID (converted to string)")
|
||||
url: str | None = Field(None, description="Dashboard URL")
|
||||
thumbnail_url: str | None = Field(None, description="Thumbnail URL")
|
||||
created_on_humanized: str | None = Field(
|
||||
None, description="Humanized creation time"
|
||||
)
|
||||
changed_on_humanized: str | None = Field(
|
||||
None, description="Humanized modification time"
|
||||
)
|
||||
chart_count: int = Field(0, description="Number of charts in the dashboard")
|
||||
owners: List[UserInfo] = Field(default_factory=list, description="Dashboard owners")
|
||||
tags: List[TagInfo] = Field(default_factory=list, description="Dashboard tags")
|
||||
roles: List[RoleInfo] = Field(default_factory=list, description="Dashboard roles")
|
||||
charts: List[ChartInfo] = Field(
|
||||
default_factory=list, description="Dashboard charts"
|
||||
)
|
||||
model_config = ConfigDict(from_attributes=True, ser_json_timedelta="iso8601")
|
||||
|
||||
|
||||
class DashboardList(BaseModel):
|
||||
dashboards: List[DashboardInfo]
|
||||
count: int
|
||||
total_count: int
|
||||
page: int
|
||||
page_size: int
|
||||
total_pages: int
|
||||
has_previous: bool
|
||||
has_next: bool
|
||||
columns_requested: List[str] | None = None
|
||||
columns_loaded: List[str] | None = None
|
||||
filters_applied: List[DashboardFilter] = Field(
|
||||
default_factory=list,
|
||||
description="List of advanced filter dicts applied to the query.",
|
||||
)
|
||||
pagination: PaginationInfo | None = None
|
||||
timestamp: datetime | None = None
|
||||
model_config = ConfigDict(ser_json_timedelta="iso8601")
|
||||
|
||||
|
||||
class AddChartToDashboardRequest(BaseModel):
|
||||
"""Request schema for adding a chart to an existing dashboard."""
|
||||
|
||||
dashboard_id: int = Field(
|
||||
..., description="ID of the dashboard to add the chart to"
|
||||
)
|
||||
chart_id: int = Field(..., description="ID of the chart to add to the dashboard")
|
||||
target_tab: str | None = Field(
|
||||
None, description="Target tab name (if dashboard has tabs)"
|
||||
)
|
||||
|
||||
|
||||
class AddChartToDashboardResponse(BaseModel):
|
||||
"""Response schema for adding chart to dashboard."""
|
||||
|
||||
dashboard: DashboardInfo | None = Field(
|
||||
None, description="The updated dashboard info, if successful"
|
||||
)
|
||||
dashboard_url: str | None = Field(
|
||||
None, description="URL to view the updated dashboard"
|
||||
)
|
||||
position: dict[str, Any] | None = Field(
|
||||
None, description="Position information for the added chart"
|
||||
)
|
||||
error: str | None = Field(None, description="Error message, if operation failed")
|
||||
|
||||
|
||||
class GenerateDashboardRequest(BaseModel):
|
||||
"""Request schema for generating a dashboard."""
|
||||
|
||||
chart_ids: List[int] = Field(
|
||||
..., description="List of chart IDs to include in the dashboard", min_length=1
|
||||
)
|
||||
dashboard_title: str = Field(..., description="Title for the new dashboard")
|
||||
description: str | None = Field(None, description="Description for the dashboard")
|
||||
published: bool = Field(
|
||||
default=True, description="Whether to publish the dashboard"
|
||||
)
|
||||
|
||||
|
||||
class GenerateDashboardResponse(BaseModel):
|
||||
"""Response schema for dashboard generation."""
|
||||
|
||||
dashboard: DashboardInfo | None = Field(
|
||||
None, description="The created dashboard info, if successful"
|
||||
)
|
||||
dashboard_url: str | None = Field(None, description="URL to view the dashboard")
|
||||
error: str | None = Field(None, description="Error message, if creation failed")
|
||||
|
||||
|
||||
# TODO (Phase 3+): Add GetDashboardAvailableFiltersRequest for
|
||||
# get_dashboard_available_filters tool
|
||||
class DashboardAvailableFilters(BaseModel):
|
||||
column_operators: Dict[str, Any] = Field(
|
||||
..., description="Available filter operators and metadata for each column"
|
||||
)
|
||||
|
||||
|
||||
class GetDashboardAvailableFiltersRequest(BaseModel):
|
||||
"""
|
||||
Request schema for get_dashboard_available_filters tool.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)
|
||||
|
||||
|
||||
def dashboard_serializer(dashboard: "Dashboard") -> DashboardInfo:
|
||||
return DashboardInfo(
|
||||
id=dashboard.id,
|
||||
dashboard_title=dashboard.dashboard_title or "Untitled",
|
||||
slug=dashboard.slug or "",
|
||||
description=dashboard.description,
|
||||
css=dashboard.css,
|
||||
certified_by=dashboard.certified_by,
|
||||
certification_details=dashboard.certification_details,
|
||||
json_metadata=dashboard.json_metadata,
|
||||
position_json=dashboard.position_json,
|
||||
published=dashboard.published,
|
||||
is_managed_externally=dashboard.is_managed_externally,
|
||||
external_url=dashboard.external_url,
|
||||
created_on=dashboard.created_on,
|
||||
changed_on=dashboard.changed_on,
|
||||
created_by=getattr(dashboard.created_by, "username", None)
|
||||
if dashboard.created_by
|
||||
else None,
|
||||
changed_by=getattr(dashboard.changed_by, "username", None)
|
||||
if dashboard.changed_by
|
||||
else None,
|
||||
uuid=str(dashboard.uuid) if dashboard.uuid else None,
|
||||
url=dashboard.url,
|
||||
thumbnail_url=dashboard.thumbnail_url,
|
||||
created_on_humanized=dashboard.created_on_humanized,
|
||||
changed_on_humanized=dashboard.changed_on_humanized,
|
||||
chart_count=len(dashboard.slices) if dashboard.slices else 0,
|
||||
owners=[
|
||||
UserInfo.model_validate(owner, from_attributes=True)
|
||||
for owner in dashboard.owners
|
||||
]
|
||||
if dashboard.owners
|
||||
else [],
|
||||
tags=[
|
||||
TagInfo.model_validate(tag, from_attributes=True) for tag in dashboard.tags
|
||||
]
|
||||
if dashboard.tags
|
||||
else [],
|
||||
roles=[
|
||||
RoleInfo.model_validate(role, from_attributes=True)
|
||||
for role in dashboard.roles
|
||||
]
|
||||
if dashboard.roles
|
||||
else [],
|
||||
charts=[serialize_chart_object(chart) for chart in dashboard.slices]
|
||||
if dashboard.slices
|
||||
else [],
|
||||
)
|
||||
|
||||
|
||||
def serialize_dashboard_object(dashboard: Any) -> DashboardInfo:
|
||||
"""Simple dashboard serializer that safely handles object attributes."""
|
||||
return DashboardInfo(
|
||||
id=getattr(dashboard, "id", None),
|
||||
dashboard_title=getattr(dashboard, "dashboard_title", None),
|
||||
slug=getattr(dashboard, "slug", None),
|
||||
url=getattr(dashboard, "url", None),
|
||||
published=getattr(dashboard, "published", None),
|
||||
changed_by_name=getattr(dashboard, "changed_by_name", None),
|
||||
changed_on=getattr(dashboard, "changed_on", None),
|
||||
changed_on_humanized=getattr(dashboard, "changed_on_humanized", None),
|
||||
created_by_name=getattr(dashboard, "created_by_name", None),
|
||||
created_on=getattr(dashboard, "created_on", None),
|
||||
created_on_humanized=getattr(dashboard, "created_on_humanized", None),
|
||||
description=getattr(dashboard, "description", None),
|
||||
css=getattr(dashboard, "css", None),
|
||||
certified_by=getattr(dashboard, "certified_by", None),
|
||||
certification_details=getattr(dashboard, "certification_details", None),
|
||||
json_metadata=getattr(dashboard, "json_metadata", None),
|
||||
position_json=getattr(dashboard, "position_json", None),
|
||||
is_managed_externally=getattr(dashboard, "is_managed_externally", None),
|
||||
external_url=getattr(dashboard, "external_url", None),
|
||||
uuid=str(getattr(dashboard, "uuid", ""))
|
||||
if getattr(dashboard, "uuid", None)
|
||||
else None,
|
||||
thumbnail_url=getattr(dashboard, "thumbnail_url", None),
|
||||
chart_count=len(getattr(dashboard, "slices", [])),
|
||||
owners=getattr(dashboard, "owners", []),
|
||||
tags=getattr(dashboard, "tags", []),
|
||||
roles=getattr(dashboard, "roles", []),
|
||||
charts=[
|
||||
serialize_chart_object(chart) for chart in getattr(dashboard, "slices", [])
|
||||
]
|
||||
if getattr(dashboard, "slices", None)
|
||||
else [],
|
||||
)
|
||||
30
superset/mcp_service/dashboard/tool/__init__.py
Normal file
30
superset/mcp_service/dashboard/tool/__init__.py
Normal file
@@ -0,0 +1,30 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
from .add_chart_to_existing_dashboard import add_chart_to_existing_dashboard
|
||||
from .generate_dashboard import generate_dashboard
|
||||
from .get_dashboard_available_filters import get_dashboard_available_filters
|
||||
from .get_dashboard_info import get_dashboard_info
|
||||
from .list_dashboards import list_dashboards
|
||||
|
||||
__all__ = [
|
||||
"list_dashboards",
|
||||
"get_dashboard_info",
|
||||
"generate_dashboard",
|
||||
"add_chart_to_existing_dashboard",
|
||||
"get_dashboard_available_filters",
|
||||
]
|
||||
@@ -0,0 +1,282 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
MCP tool: add_chart_to_existing_dashboard
|
||||
|
||||
This tool adds a chart to an existing dashboard with automatic layout positioning.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.dashboard.schemas import (
|
||||
AddChartToDashboardRequest,
|
||||
AddChartToDashboardResponse,
|
||||
DashboardInfo,
|
||||
)
|
||||
from superset.mcp_service.utils.url_utils import get_superset_base_url
|
||||
from superset.utils import json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _find_next_row_position(layout: Dict[str, Any]) -> int:
|
||||
"""
|
||||
Find the next available row position in the dashboard layout.
|
||||
|
||||
Returns:
|
||||
Row index for the new chart
|
||||
"""
|
||||
# Find existing rows
|
||||
row_indices = []
|
||||
for key in layout.keys():
|
||||
if key.startswith("ROW-") and key[4:].isdigit():
|
||||
row_indices.append(int(key[4:]))
|
||||
|
||||
# Return next available row index
|
||||
return max(row_indices) + 1 if row_indices else 0
|
||||
|
||||
|
||||
def _add_chart_to_layout(
|
||||
layout: Dict[str, Any], chart: Any, chart_id: int, row_index: int
|
||||
) -> tuple[str, str]:
|
||||
"""
|
||||
Add chart and row components to the dashboard layout.
|
||||
|
||||
Returns:
|
||||
Tuple of (chart_key, row_key)
|
||||
"""
|
||||
chart_key = f"CHART-{chart_id}"
|
||||
row_key = f"ROW-{row_index}"
|
||||
chart_width = 5 # Balanced width for good proportions
|
||||
chart_height = 50 # Good height for most chart types
|
||||
|
||||
# Add chart to layout using proper Superset structure
|
||||
layout[chart_key] = {
|
||||
"children": [],
|
||||
"id": chart_key,
|
||||
"meta": {
|
||||
"chartId": chart_id,
|
||||
"height": chart_height,
|
||||
"sliceName": chart.slice_name or f"Chart {chart_id}",
|
||||
"uuid": str(chart.uuid) if chart.uuid else f"chart-{chart_id}",
|
||||
"width": chart_width,
|
||||
},
|
||||
"parents": ["ROOT_ID", "GRID_ID", row_key],
|
||||
"type": "CHART",
|
||||
}
|
||||
|
||||
# Create row for the chart
|
||||
layout[row_key] = {
|
||||
"children": [chart_key],
|
||||
"id": row_key,
|
||||
"meta": {"background": "BACKGROUND_TRANSPARENT"},
|
||||
"parents": ["ROOT_ID", "GRID_ID"],
|
||||
"type": "ROW",
|
||||
}
|
||||
|
||||
return chart_key, row_key
|
||||
|
||||
|
||||
def _ensure_layout_structure(layout: Dict[str, Any], row_key: str) -> None:
|
||||
"""
|
||||
Ensure the dashboard layout has proper GRID and ROOT structure.
|
||||
"""
|
||||
# Ensure GRID structure exists
|
||||
if "GRID_ID" not in layout:
|
||||
layout["GRID_ID"] = {
|
||||
"children": [],
|
||||
"id": "GRID_ID",
|
||||
"parents": ["ROOT_ID"],
|
||||
"type": "GRID",
|
||||
}
|
||||
|
||||
# Add row to GRID
|
||||
if "children" not in layout["GRID_ID"]:
|
||||
layout["GRID_ID"]["children"] = []
|
||||
layout["GRID_ID"]["children"].append(row_key)
|
||||
|
||||
# Update ROOT_ID if it exists, or create it
|
||||
if "ROOT_ID" in layout:
|
||||
if "children" not in layout["ROOT_ID"]:
|
||||
layout["ROOT_ID"]["children"] = []
|
||||
if "GRID_ID" not in layout["ROOT_ID"]["children"]:
|
||||
layout["ROOT_ID"]["children"].append("GRID_ID")
|
||||
else:
|
||||
# Create ROOT_ID if it doesn't exist
|
||||
layout["ROOT_ID"] = {
|
||||
"children": ["GRID_ID"],
|
||||
"id": "ROOT_ID",
|
||||
"type": "ROOT",
|
||||
}
|
||||
|
||||
# Ensure dashboard version
|
||||
if "DASHBOARD_VERSION_KEY" not in layout:
|
||||
layout["DASHBOARD_VERSION_KEY"] = "v2"
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
def add_chart_to_existing_dashboard(
|
||||
request: AddChartToDashboardRequest, ctx: Context
|
||||
) -> AddChartToDashboardResponse:
|
||||
"""
|
||||
Add chart to existing dashboard. Auto-positions in 2-column grid.
|
||||
Returns updated dashboard info.
|
||||
"""
|
||||
try:
|
||||
from superset.commands.dashboard.update import UpdateDashboardCommand
|
||||
from superset.daos.dashboard import DashboardDAO
|
||||
|
||||
# Validate dashboard exists
|
||||
dashboard = DashboardDAO.find_by_id(request.dashboard_id)
|
||||
if not dashboard:
|
||||
return AddChartToDashboardResponse(
|
||||
dashboard=None,
|
||||
dashboard_url=None,
|
||||
position=None,
|
||||
error=f"Dashboard with ID {request.dashboard_id} not found",
|
||||
)
|
||||
|
||||
# Get chart object for SQLAlchemy relationships and validation
|
||||
from superset import db
|
||||
from superset.models.slice import Slice
|
||||
|
||||
new_chart = db.session.get(Slice, request.chart_id)
|
||||
if not new_chart:
|
||||
return AddChartToDashboardResponse(
|
||||
dashboard=None,
|
||||
dashboard_url=None,
|
||||
position=None,
|
||||
error=f"Chart with ID {request.chart_id} not found",
|
||||
)
|
||||
|
||||
# Check if chart is already in dashboard
|
||||
current_chart_ids = [slice.id for slice in dashboard.slices]
|
||||
if request.chart_id in current_chart_ids:
|
||||
return AddChartToDashboardResponse(
|
||||
dashboard=None,
|
||||
dashboard_url=None,
|
||||
position=None,
|
||||
error=(
|
||||
f"Chart {request.chart_id} is already in dashboard "
|
||||
f"{request.dashboard_id}"
|
||||
),
|
||||
)
|
||||
|
||||
# Parse current layout
|
||||
try:
|
||||
current_layout = json.loads(dashboard.position_json or "{}")
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
current_layout = {}
|
||||
|
||||
# Find position for new chart
|
||||
row_index = _find_next_row_position(current_layout)
|
||||
|
||||
# Add chart and row to layout
|
||||
chart_key, row_key = _add_chart_to_layout(
|
||||
current_layout, new_chart, request.chart_id, row_index
|
||||
)
|
||||
|
||||
# Ensure proper layout structure
|
||||
_ensure_layout_structure(current_layout, row_key)
|
||||
|
||||
# Get chart objects for SQLAlchemy relationships
|
||||
# Get existing chart objects
|
||||
existing_chart_objects = dashboard.slices
|
||||
|
||||
# Combine existing and new chart objects (new_chart was retrieved above)
|
||||
all_chart_objects = list(existing_chart_objects) + [new_chart]
|
||||
|
||||
# Prepare update data
|
||||
update_data = {
|
||||
"position_json": json.dumps(current_layout),
|
||||
"slices": all_chart_objects, # Pass ORM objects, not IDs
|
||||
}
|
||||
|
||||
# Update the dashboard
|
||||
command = UpdateDashboardCommand(request.dashboard_id, update_data)
|
||||
updated_dashboard = command.run()
|
||||
|
||||
# Convert to response format
|
||||
from superset.mcp_service.dashboard.schemas import (
|
||||
serialize_tag_object,
|
||||
serialize_user_object,
|
||||
)
|
||||
|
||||
dashboard_info = DashboardInfo(
|
||||
id=updated_dashboard.id,
|
||||
dashboard_title=updated_dashboard.dashboard_title,
|
||||
slug=updated_dashboard.slug,
|
||||
description=updated_dashboard.description,
|
||||
published=updated_dashboard.published,
|
||||
created_on=updated_dashboard.created_on,
|
||||
changed_on=updated_dashboard.changed_on,
|
||||
created_by=updated_dashboard.created_by.username
|
||||
if updated_dashboard.created_by
|
||||
else None,
|
||||
changed_by=updated_dashboard.changed_by.username
|
||||
if updated_dashboard.changed_by
|
||||
else None,
|
||||
uuid=str(updated_dashboard.uuid) if updated_dashboard.uuid else None,
|
||||
url=f"{get_superset_base_url()}/superset/dashboard/{updated_dashboard.id}/",
|
||||
chart_count=len(updated_dashboard.slices),
|
||||
owners=[
|
||||
serialize_user_object(owner)
|
||||
for owner in getattr(updated_dashboard, "owners", [])
|
||||
if serialize_user_object(owner) is not None
|
||||
],
|
||||
tags=[
|
||||
serialize_tag_object(tag)
|
||||
for tag in getattr(updated_dashboard, "tags", [])
|
||||
if serialize_tag_object(tag) is not None
|
||||
],
|
||||
roles=[],
|
||||
charts=[],
|
||||
)
|
||||
|
||||
dashboard_url = (
|
||||
f"{get_superset_base_url()}/superset/dashboard/{updated_dashboard.id}/"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Added chart %s to dashboard %s ", request.chart_id, request.dashboard_id
|
||||
)
|
||||
|
||||
# Return position info for compatibility
|
||||
position_info = {"row": row_index, "chart_key": chart_key, "row_key": row_key}
|
||||
|
||||
return AddChartToDashboardResponse(
|
||||
dashboard=dashboard_info,
|
||||
dashboard_url=dashboard_url,
|
||||
position=position_info,
|
||||
error=None,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error adding chart to dashboard: %s", e)
|
||||
return AddChartToDashboardResponse(
|
||||
dashboard=None,
|
||||
dashboard_url=None,
|
||||
position=None,
|
||||
error=f"Failed to add chart to dashboard: {str(e)}",
|
||||
)
|
||||
241
superset/mcp_service/dashboard/tool/generate_dashboard.py
Normal file
241
superset/mcp_service/dashboard/tool/generate_dashboard.py
Normal file
@@ -0,0 +1,241 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
MCP tool: generate_dashboard
|
||||
|
||||
This tool creates a new dashboard with specified charts and layout configuration.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.dashboard.schemas import (
|
||||
DashboardInfo,
|
||||
GenerateDashboardRequest,
|
||||
GenerateDashboardResponse,
|
||||
)
|
||||
from superset.mcp_service.utils.url_utils import get_superset_base_url
|
||||
from superset.utils import json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _create_dashboard_layout(chart_objects: List[Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Create a simple dashboard layout with charts arranged in a grid.
|
||||
|
||||
This creates a basic 2-column layout where charts are arranged
|
||||
vertically in alternating columns.
|
||||
|
||||
Args:
|
||||
chart_objects: List of Chart ORM objects (not IDs)
|
||||
"""
|
||||
layout: Dict[str, Any] = {}
|
||||
|
||||
# Grid configuration based on real Superset dashboard patterns
|
||||
# Use 2-chart rows with medium-sized charts (like existing dashboards)
|
||||
charts_per_row = 2
|
||||
chart_width = 5 # Balanced width for good proportions
|
||||
chart_height = 50 # Good height for most chart types
|
||||
|
||||
# Create rows with charts
|
||||
row_ids = []
|
||||
for i in range(0, len(chart_objects), charts_per_row):
|
||||
row_index = i // charts_per_row
|
||||
row_id = f"ROW-{row_index}"
|
||||
row_ids.append(row_id)
|
||||
|
||||
# Get charts for this row (up to 2 charts like real dashboards)
|
||||
row_charts = chart_objects[i : i + charts_per_row]
|
||||
chart_keys = []
|
||||
|
||||
for chart in row_charts:
|
||||
chart_key = f"CHART-{chart.id}"
|
||||
chart_keys.append(chart_key)
|
||||
|
||||
# Create chart component with standard dimensions
|
||||
layout[chart_key] = {
|
||||
"children": [],
|
||||
"id": chart_key,
|
||||
"meta": {
|
||||
"chartId": chart.id,
|
||||
"height": chart_height,
|
||||
"sliceName": chart.slice_name or f"Chart {chart.id}",
|
||||
"uuid": str(chart.uuid) if chart.uuid else f"chart-{chart.id}",
|
||||
"width": chart_width,
|
||||
},
|
||||
"parents": ["ROOT_ID", "GRID_ID", row_id],
|
||||
"type": "CHART",
|
||||
}
|
||||
|
||||
# Create row containing the charts
|
||||
layout[row_id] = {
|
||||
"children": chart_keys,
|
||||
"id": row_id,
|
||||
"meta": {"background": "BACKGROUND_TRANSPARENT"},
|
||||
"parents": ["ROOT_ID", "GRID_ID"],
|
||||
"type": "ROW",
|
||||
}
|
||||
|
||||
# Add GRID container
|
||||
layout["GRID_ID"] = {
|
||||
"children": row_ids,
|
||||
"id": "GRID_ID",
|
||||
"parents": ["ROOT_ID"],
|
||||
"type": "GRID",
|
||||
}
|
||||
|
||||
# Add root layout container
|
||||
layout["ROOT_ID"] = {
|
||||
"children": ["GRID_ID"],
|
||||
"id": "ROOT_ID",
|
||||
"type": "ROOT",
|
||||
}
|
||||
|
||||
# Add dashboard version
|
||||
layout["DASHBOARD_VERSION_KEY"] = "v2"
|
||||
|
||||
return layout
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
def generate_dashboard(
|
||||
request: GenerateDashboardRequest, ctx: Context
|
||||
) -> GenerateDashboardResponse:
|
||||
"""Create dashboard from chart IDs.
|
||||
|
||||
IMPORTANT:
|
||||
- All charts must exist and be accessible to current user
|
||||
- Charts arranged automatically in 2-column grid layout
|
||||
|
||||
Returns:
|
||||
- Dashboard ID and URL
|
||||
"""
|
||||
try:
|
||||
# Get chart objects from IDs (required for SQLAlchemy relationships)
|
||||
from superset import db
|
||||
from superset.commands.dashboard.create import CreateDashboardCommand
|
||||
from superset.models.slice import Slice
|
||||
|
||||
chart_objects = (
|
||||
db.session.query(Slice).filter(Slice.id.in_(request.chart_ids)).all()
|
||||
)
|
||||
found_chart_ids = [chart.id for chart in chart_objects]
|
||||
|
||||
# Check if all requested charts were found
|
||||
missing_chart_ids = set(request.chart_ids) - set(found_chart_ids)
|
||||
if missing_chart_ids:
|
||||
return GenerateDashboardResponse(
|
||||
dashboard=None,
|
||||
dashboard_url=None,
|
||||
error=f"Charts not found: {list(missing_chart_ids)}",
|
||||
)
|
||||
|
||||
# Create dashboard layout with chart objects
|
||||
layout = _create_dashboard_layout(chart_objects)
|
||||
|
||||
# Prepare dashboard data
|
||||
dashboard_data = {
|
||||
"dashboard_title": request.dashboard_title,
|
||||
"slug": None, # Let Superset auto-generate slug
|
||||
"css": "",
|
||||
"json_metadata": json.dumps(
|
||||
{
|
||||
"filter_scopes": {},
|
||||
"expanded_slices": {},
|
||||
"refresh_frequency": 0,
|
||||
"timed_refresh_immune_slices": [],
|
||||
"color_scheme": None,
|
||||
"label_colors": {},
|
||||
"shared_label_colors": {},
|
||||
"color_scheme_domain": [],
|
||||
"cross_filters_enabled": False,
|
||||
"native_filter_configuration": [],
|
||||
"global_chart_configuration": {
|
||||
"scope": {"rootPath": ["ROOT_ID"], "excluded": []}
|
||||
},
|
||||
"chart_configuration": {},
|
||||
}
|
||||
),
|
||||
"position_json": json.dumps(layout),
|
||||
"published": request.published,
|
||||
"slices": chart_objects, # Pass ORM objects, not IDs
|
||||
}
|
||||
|
||||
if request.description:
|
||||
dashboard_data["description"] = request.description
|
||||
|
||||
# Create the dashboard using Superset's command pattern
|
||||
command = CreateDashboardCommand(dashboard_data)
|
||||
dashboard = command.run()
|
||||
|
||||
# Convert to our response format
|
||||
from superset.mcp_service.dashboard.schemas import (
|
||||
serialize_tag_object,
|
||||
serialize_user_object,
|
||||
)
|
||||
|
||||
dashboard_info = DashboardInfo(
|
||||
id=dashboard.id,
|
||||
dashboard_title=dashboard.dashboard_title,
|
||||
slug=dashboard.slug,
|
||||
description=dashboard.description,
|
||||
published=dashboard.published,
|
||||
created_on=dashboard.created_on,
|
||||
changed_on=dashboard.changed_on,
|
||||
created_by=dashboard.created_by.username if dashboard.created_by else None,
|
||||
changed_by=dashboard.changed_by.username if dashboard.changed_by else None,
|
||||
uuid=str(dashboard.uuid) if dashboard.uuid else None,
|
||||
url=f"{get_superset_base_url()}/superset/dashboard/{dashboard.id}/",
|
||||
chart_count=len(request.chart_ids),
|
||||
owners=[
|
||||
serialize_user_object(owner)
|
||||
for owner in getattr(dashboard, "owners", [])
|
||||
if serialize_user_object(owner) is not None
|
||||
],
|
||||
tags=[
|
||||
serialize_tag_object(tag)
|
||||
for tag in getattr(dashboard, "tags", [])
|
||||
if serialize_tag_object(tag) is not None
|
||||
],
|
||||
roles=[], # Dashboard roles not typically set at creation
|
||||
charts=[], # Chart details not needed in response
|
||||
)
|
||||
|
||||
dashboard_url = f"{get_superset_base_url()}/superset/dashboard/{dashboard.id}/"
|
||||
|
||||
logger.info(
|
||||
"Created dashboard %s with %s charts", dashboard.id, len(request.chart_ids)
|
||||
)
|
||||
|
||||
return GenerateDashboardResponse(
|
||||
dashboard=dashboard_info, dashboard_url=dashboard_url, error=None
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error creating dashboard: %s", e)
|
||||
return GenerateDashboardResponse(
|
||||
dashboard=None,
|
||||
dashboard_url=None,
|
||||
error=f"Failed to create dashboard: {str(e)}",
|
||||
)
|
||||
@@ -0,0 +1,49 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
"""
|
||||
Get available filters FastMCP tool
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.dashboard.schemas import (
|
||||
DashboardAvailableFilters,
|
||||
GetDashboardAvailableFiltersRequest,
|
||||
)
|
||||
from superset.mcp_service.mcp_core import ModelGetAvailableFiltersCore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def get_dashboard_available_filters(
|
||||
request: GetDashboardAvailableFiltersRequest, ctx: Context
|
||||
) -> DashboardAvailableFilters:
|
||||
"""Get available dashboard filter fields and operators."""
|
||||
from superset.daos.dashboard import DashboardDAO
|
||||
|
||||
tool = ModelGetAvailableFiltersCore(
|
||||
dao_class=DashboardDAO,
|
||||
output_schema=DashboardAvailableFilters,
|
||||
logger=logger,
|
||||
)
|
||||
return tool.run_tool()
|
||||
101
superset/mcp_service/dashboard/tool/get_dashboard_info.py
Normal file
101
superset/mcp_service/dashboard/tool/get_dashboard_info.py
Normal file
@@ -0,0 +1,101 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Get dashboard info FastMCP tool
|
||||
|
||||
This module contains the FastMCP tool for getting detailed information
|
||||
about a specific dashboard.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.dashboard.schemas import (
|
||||
dashboard_serializer,
|
||||
DashboardError,
|
||||
DashboardInfo,
|
||||
GetDashboardInfoRequest,
|
||||
)
|
||||
from superset.mcp_service.mcp_core import ModelGetInfoCore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def get_dashboard_info(
|
||||
request: GetDashboardInfoRequest, ctx: Context
|
||||
) -> DashboardInfo | DashboardError:
|
||||
"""
|
||||
Get dashboard metadata by ID, UUID, or slug.
|
||||
|
||||
Returns title, charts, and layout details.
|
||||
"""
|
||||
await ctx.info("Retrieving dashboard information: %s" % (request.identifier,))
|
||||
await ctx.debug(
|
||||
"Metadata cache settings: use_cache=%s, refresh_metadata=%s, force_refresh=%s"
|
||||
% (request.use_cache, request.refresh_metadata, request.force_refresh)
|
||||
)
|
||||
|
||||
try:
|
||||
from superset.daos.dashboard import DashboardDAO
|
||||
|
||||
tool = ModelGetInfoCore(
|
||||
dao_class=DashboardDAO,
|
||||
output_schema=DashboardInfo,
|
||||
error_schema=DashboardError,
|
||||
serializer=dashboard_serializer,
|
||||
supports_slug=True, # Dashboards support slugs
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
result = tool.run_tool(request.identifier)
|
||||
|
||||
if isinstance(result, DashboardInfo):
|
||||
await ctx.info(
|
||||
"Dashboard information retrieved successfully: id=%s, title=%s, "
|
||||
"chart_count=%s, published=%s"
|
||||
% (
|
||||
result.id,
|
||||
result.dashboard_title,
|
||||
result.chart_count,
|
||||
result.published,
|
||||
)
|
||||
)
|
||||
else:
|
||||
await ctx.warning(
|
||||
"Dashboard retrieval failed: error_type=%s, error=%s"
|
||||
% (result.error_type, result.error)
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
await ctx.error(
|
||||
"Dashboard information retrieval failed: identifier=%s, error=%s, "
|
||||
"error_type=%s" % (request.identifier, str(e), type(e).__name__)
|
||||
)
|
||||
return DashboardError(
|
||||
error=f"Failed to get dashboard info: {str(e)}",
|
||||
error_type="InternalError",
|
||||
timestamp=datetime.now(timezone.utc),
|
||||
)
|
||||
99
superset/mcp_service/dashboard/tool/list_dashboards.py
Normal file
99
superset/mcp_service/dashboard/tool/list_dashboards.py
Normal file
@@ -0,0 +1,99 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
List dashboards FastMCP tool (Advanced with metadata cache control)
|
||||
|
||||
This module contains the FastMCP tool for listing dashboards using
|
||||
advanced filtering with clear, unambiguous request schema and metadata cache control.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.dashboard.schemas import (
|
||||
DashboardFilter,
|
||||
DashboardInfo,
|
||||
DashboardList,
|
||||
ListDashboardsRequest,
|
||||
serialize_dashboard_object,
|
||||
)
|
||||
from superset.mcp_service.mcp_core import ModelListCore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_DASHBOARD_COLUMNS = [
|
||||
"id",
|
||||
"dashboard_title",
|
||||
"slug",
|
||||
"uuid",
|
||||
"published",
|
||||
"changed_on",
|
||||
"created_on",
|
||||
]
|
||||
|
||||
SORTABLE_DASHBOARD_COLUMNS = [
|
||||
"id",
|
||||
"dashboard_title",
|
||||
"slug",
|
||||
"published",
|
||||
"changed_on",
|
||||
"created_on",
|
||||
]
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def list_dashboards(
|
||||
request: ListDashboardsRequest, ctx: Context
|
||||
) -> DashboardList:
|
||||
"""List dashboards with filtering and search. Returns dashboard metadata
|
||||
including title, slug, and charts.
|
||||
|
||||
Sortable columns for order_column: id, dashboard_title, slug, published,
|
||||
changed_on, created_on
|
||||
"""
|
||||
|
||||
from superset.daos.dashboard import DashboardDAO
|
||||
|
||||
tool = ModelListCore(
|
||||
dao_class=DashboardDAO,
|
||||
output_schema=DashboardInfo,
|
||||
item_serializer=lambda obj, cols: serialize_dashboard_object(obj),
|
||||
filter_type=DashboardFilter,
|
||||
default_columns=DEFAULT_DASHBOARD_COLUMNS,
|
||||
search_columns=[
|
||||
"dashboard_title",
|
||||
"slug",
|
||||
"uuid",
|
||||
],
|
||||
list_field_name="dashboards",
|
||||
output_list_schema=DashboardList,
|
||||
logger=logger,
|
||||
)
|
||||
return tool.run_tool(
|
||||
filters=request.filters,
|
||||
search=request.search,
|
||||
select_columns=request.select_columns,
|
||||
order_column=request.order_column,
|
||||
order_direction=request.order_direction,
|
||||
page=max(request.page - 1, 0),
|
||||
page_size=request.page_size,
|
||||
)
|
||||
21
superset/mcp_service/dataset/prompts/__init__.py
Normal file
21
superset/mcp_service/dataset/prompts/__init__.py
Normal file
@@ -0,0 +1,21 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Dataset prompts for Superset MCP service"""
|
||||
|
||||
# Import prompt modules here when created
|
||||
# from . import explore_dataset # noqa: F401
|
||||
21
superset/mcp_service/dataset/resources/__init__.py
Normal file
21
superset/mcp_service/dataset/resources/__init__.py
Normal file
@@ -0,0 +1,21 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Dataset resources for Superset MCP service"""
|
||||
|
||||
# Import resource modules here when created
|
||||
# from . import column_suggestions # noqa: F401
|
||||
349
superset/mcp_service/dataset/schemas.py
Normal file
349
superset/mcp_service/dataset/schemas.py
Normal file
@@ -0,0 +1,349 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Pydantic schemas for dataset-related responses
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Annotated, Any, Dict, List, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator, PositiveInt
|
||||
|
||||
from superset.daos.base import ColumnOperator, ColumnOperatorEnum
|
||||
from superset.mcp_service.common.cache_schemas import MetadataCacheControl
|
||||
from superset.mcp_service.system.schemas import (
|
||||
PaginationInfo,
|
||||
TagInfo,
|
||||
UserInfo,
|
||||
)
|
||||
from superset.utils import json
|
||||
|
||||
|
||||
class GetDatasetAvailableFiltersRequest(BaseModel):
|
||||
"""
|
||||
Request schema for get_dataset_available_filters tool.
|
||||
|
||||
Currently has no parameters but provides consistent API for future extensibility.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(
|
||||
extra="forbid",
|
||||
str_strip_whitespace=True,
|
||||
)
|
||||
|
||||
|
||||
class DatasetAvailableFilters(BaseModel):
|
||||
column_operators: Dict[str, List[str]] = Field(
|
||||
...,
|
||||
description="Available filter operators for each column: mapping from column "
|
||||
"name to list of supported operators",
|
||||
)
|
||||
|
||||
|
||||
class DatasetFilter(ColumnOperator):
|
||||
"""
|
||||
Filter object for dataset listing.
|
||||
col: The column to filter on. Must be one of the allowed filter fields.
|
||||
opr: The operator to use. Must be one of the supported operators.
|
||||
value: The value to filter by (type depends on col and opr).
|
||||
"""
|
||||
|
||||
col: Literal[
|
||||
"table_name",
|
||||
"schema",
|
||||
"owner",
|
||||
"favorite",
|
||||
] = Field(
|
||||
...,
|
||||
description="Column to filter on. See get_dataset_available_filters for "
|
||||
"allowed values.",
|
||||
)
|
||||
opr: ColumnOperatorEnum = Field(
|
||||
...,
|
||||
description="Operator to use. See get_dataset_available_filters for "
|
||||
"allowed values.",
|
||||
)
|
||||
value: str | int | float | bool | List[str | int | float | bool] = Field(
|
||||
..., description="Value to filter by (type depends on col and opr)"
|
||||
)
|
||||
|
||||
|
||||
class TableColumnInfo(BaseModel):
|
||||
column_name: str = Field(..., description="Column name")
|
||||
verbose_name: str | None = Field(None, description="Verbose name")
|
||||
type: str | None = Field(None, description="Column type")
|
||||
is_dttm: bool | None = Field(None, description="Is datetime column")
|
||||
groupby: bool | None = Field(None, description="Is groupable")
|
||||
filterable: bool | None = Field(None, description="Is filterable")
|
||||
description: str | None = Field(None, description="Column description")
|
||||
|
||||
|
||||
class SqlMetricInfo(BaseModel):
|
||||
metric_name: str = Field(..., description="Metric name")
|
||||
verbose_name: str | None = Field(None, description="Verbose name")
|
||||
expression: str | None = Field(None, description="SQL expression")
|
||||
description: str | None = Field(None, description="Metric description")
|
||||
d3format: str | None = Field(None, description="D3 format string")
|
||||
|
||||
|
||||
class DatasetInfo(BaseModel):
|
||||
id: int | None = Field(None, description="Dataset ID")
|
||||
table_name: str | None = Field(None, description="Table name")
|
||||
schema_name: str | None = Field(None, description="Schema name", alias="schema")
|
||||
database_name: str | None = Field(None, description="Database name")
|
||||
description: str | None = Field(None, description="Dataset description")
|
||||
changed_by: str | None = Field(None, description="Last modifier (username)")
|
||||
changed_on: str | datetime | None = Field(
|
||||
None, description="Last modification timestamp"
|
||||
)
|
||||
changed_on_humanized: str | None = Field(
|
||||
None, description="Humanized modification time"
|
||||
)
|
||||
created_by: str | None = Field(None, description="Dataset creator (username)")
|
||||
created_on: str | datetime | None = Field(None, description="Creation timestamp")
|
||||
created_on_humanized: str | None = Field(
|
||||
None, description="Humanized creation time"
|
||||
)
|
||||
tags: List[TagInfo] = Field(default_factory=list, description="Dataset tags")
|
||||
owners: List[UserInfo] = Field(
|
||||
default_factory=list, description="DatasetInfo owners"
|
||||
)
|
||||
is_virtual: bool | None = Field(
|
||||
None, description="Whether the dataset is virtual (uses SQL)"
|
||||
)
|
||||
database_id: int | None = Field(None, description="Database ID")
|
||||
uuid: str | None = Field(None, description="Dataset UUID")
|
||||
schema_perm: str | None = Field(None, description="Schema permission string")
|
||||
url: str | None = Field(None, description="Dataset URL")
|
||||
sql: str | None = Field(None, description="SQL for virtual datasets")
|
||||
main_dttm_col: str | None = Field(None, description="Main datetime column")
|
||||
offset: int | None = Field(None, description="Offset")
|
||||
cache_timeout: int | None = Field(None, description="Cache timeout")
|
||||
params: Dict[str, Any | None] | None = Field(None, description="Extra params")
|
||||
template_params: Dict[str, Any | None] | None = Field(
|
||||
None, description="Template params"
|
||||
)
|
||||
extra: Dict[str, Any | None] | None = Field(None, description="Extra metadata")
|
||||
columns: List[TableColumnInfo] = Field(
|
||||
default_factory=list, description="Columns in the dataset"
|
||||
)
|
||||
metrics: List[SqlMetricInfo] = Field(
|
||||
default_factory=list, description="Metrics in the dataset"
|
||||
)
|
||||
is_favorite: bool | None = Field(
|
||||
None, description="Whether this dataset is favorited by the current user"
|
||||
)
|
||||
model_config = ConfigDict(
|
||||
from_attributes=True,
|
||||
ser_json_timedelta="iso8601",
|
||||
populate_by_name=True, # Allow both 'schema' (alias) and 'schema_name' (field)
|
||||
)
|
||||
|
||||
|
||||
class DatasetList(BaseModel):
|
||||
datasets: List[DatasetInfo]
|
||||
count: int
|
||||
total_count: int
|
||||
page: int
|
||||
page_size: int
|
||||
total_pages: int
|
||||
has_previous: bool
|
||||
has_next: bool
|
||||
columns_requested: List[str] | None = None
|
||||
columns_loaded: List[str] | None = None
|
||||
filters_applied: List[DatasetFilter] = Field(
|
||||
default_factory=list,
|
||||
description="List of advanced filter dicts applied to the query.",
|
||||
)
|
||||
pagination: PaginationInfo | None = None
|
||||
timestamp: datetime | None = None
|
||||
model_config = ConfigDict(ser_json_timedelta="iso8601")
|
||||
|
||||
|
||||
class ListDatasetsRequest(MetadataCacheControl):
|
||||
"""Request schema for list_datasets with clear, unambiguous types."""
|
||||
|
||||
filters: Annotated[
|
||||
List[DatasetFilter],
|
||||
Field(
|
||||
default_factory=list,
|
||||
description="List of filter objects (column, operator, value). Each "
|
||||
"filter is an object with 'col', 'opr', and 'value' "
|
||||
"properties. Cannot be used together with 'search'.",
|
||||
),
|
||||
]
|
||||
select_columns: Annotated[
|
||||
List[str],
|
||||
Field(
|
||||
default_factory=lambda: [
|
||||
"id",
|
||||
"table_name",
|
||||
"schema",
|
||||
"database_name",
|
||||
"changed_by_name",
|
||||
"changed_on",
|
||||
"created_by_name",
|
||||
"created_on",
|
||||
"metrics",
|
||||
"columns",
|
||||
"uuid",
|
||||
],
|
||||
description="List of columns to select. Defaults to common columns if not "
|
||||
"specified.",
|
||||
),
|
||||
]
|
||||
search: Annotated[
|
||||
str | None,
|
||||
Field(
|
||||
default=None,
|
||||
description="Text search string to match against dataset fields. Cannot "
|
||||
"be used together with 'filters'.",
|
||||
),
|
||||
]
|
||||
order_column: Annotated[
|
||||
str | None, Field(default=None, description="Column to order results by")
|
||||
]
|
||||
order_direction: Annotated[
|
||||
Literal["asc", "desc"],
|
||||
Field(
|
||||
default="desc", description="Direction to order results ('asc' or 'desc')"
|
||||
),
|
||||
]
|
||||
page: Annotated[
|
||||
PositiveInt,
|
||||
Field(default=1, description="Page number for pagination (1-based)"),
|
||||
]
|
||||
page_size: Annotated[
|
||||
PositiveInt, Field(default=10, description="Number of items per page")
|
||||
]
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_search_and_filters(self) -> "ListDatasetsRequest":
|
||||
"""Prevent using both search and filters simultaneously to avoid query
|
||||
conflicts."""
|
||||
if self.search and self.filters:
|
||||
raise ValueError(
|
||||
"Cannot use both 'search' and 'filters' parameters simultaneously. "
|
||||
"Use either 'search' for text-based searching across multiple fields, "
|
||||
"or 'filters' for precise column-based filtering, but not both."
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class DatasetError(BaseModel):
|
||||
error: str = Field(..., description="Error message")
|
||||
error_type: str = Field(..., description="Type of error")
|
||||
timestamp: str | datetime | None = Field(None, description="Error timestamp")
|
||||
model_config = ConfigDict(ser_json_timedelta="iso8601")
|
||||
|
||||
@classmethod
|
||||
def create(cls, error: str, error_type: str) -> "DatasetError":
|
||||
"""Create a standardized DatasetError with timestamp."""
|
||||
from datetime import datetime
|
||||
|
||||
return cls(error=error, error_type=error_type, timestamp=datetime.now())
|
||||
|
||||
|
||||
class GetDatasetInfoRequest(MetadataCacheControl):
|
||||
"""Request schema for get_dataset_info with support for ID or UUID."""
|
||||
|
||||
identifier: Annotated[
|
||||
int | str,
|
||||
Field(description="Dataset identifier - can be numeric ID or UUID string"),
|
||||
]
|
||||
|
||||
|
||||
def serialize_dataset_object(dataset: Any) -> DatasetInfo | None:
|
||||
if not dataset:
|
||||
return None
|
||||
params = getattr(dataset, "params", None)
|
||||
if isinstance(params, str):
|
||||
try:
|
||||
params = json.loads(params)
|
||||
except Exception:
|
||||
params = None
|
||||
columns = [
|
||||
TableColumnInfo(
|
||||
column_name=getattr(col, "column_name", None),
|
||||
verbose_name=getattr(col, "verbose_name", None),
|
||||
type=getattr(col, "type", None),
|
||||
is_dttm=getattr(col, "is_dttm", None),
|
||||
groupby=getattr(col, "groupby", None),
|
||||
filterable=getattr(col, "filterable", None),
|
||||
description=getattr(col, "description", None),
|
||||
)
|
||||
for col in getattr(dataset, "columns", [])
|
||||
]
|
||||
metrics = [
|
||||
SqlMetricInfo(
|
||||
metric_name=getattr(metric, "metric_name", None),
|
||||
verbose_name=getattr(metric, "verbose_name", None),
|
||||
expression=getattr(metric, "expression", None),
|
||||
description=getattr(metric, "description", None),
|
||||
d3format=getattr(metric, "d3format", None),
|
||||
)
|
||||
for metric in getattr(dataset, "metrics", [])
|
||||
]
|
||||
return DatasetInfo(
|
||||
id=getattr(dataset, "id", None),
|
||||
table_name=getattr(dataset, "table_name", None),
|
||||
schema_name=getattr(dataset, "schema", None),
|
||||
database_name=getattr(dataset.database, "database_name", None)
|
||||
if getattr(dataset, "database", None)
|
||||
else None,
|
||||
description=getattr(dataset, "description", None),
|
||||
changed_by=getattr(dataset, "changed_by_name", None)
|
||||
or (str(dataset.changed_by) if getattr(dataset, "changed_by", None) else None),
|
||||
changed_on=getattr(dataset, "changed_on", None),
|
||||
changed_on_humanized=getattr(dataset, "changed_on_humanized", None),
|
||||
created_by=getattr(dataset, "created_by_name", None)
|
||||
or (str(dataset.created_by) if getattr(dataset, "created_by", None) else None),
|
||||
created_on=getattr(dataset, "created_on", None),
|
||||
created_on_humanized=getattr(dataset, "created_on_humanized", None),
|
||||
tags=[
|
||||
TagInfo.model_validate(tag, from_attributes=True)
|
||||
for tag in getattr(dataset, "tags", [])
|
||||
]
|
||||
if getattr(dataset, "tags", None)
|
||||
else [],
|
||||
owners=[
|
||||
UserInfo.model_validate(owner, from_attributes=True)
|
||||
for owner in getattr(dataset, "owners", [])
|
||||
]
|
||||
if getattr(dataset, "owners", None)
|
||||
else [],
|
||||
is_virtual=getattr(dataset, "is_virtual", None),
|
||||
database_id=getattr(dataset, "database_id", None),
|
||||
uuid=str(getattr(dataset, "uuid", ""))
|
||||
if getattr(dataset, "uuid", None)
|
||||
else None,
|
||||
schema_perm=getattr(dataset, "schema_perm", None),
|
||||
url=getattr(dataset, "url", None),
|
||||
sql=getattr(dataset, "sql", None),
|
||||
main_dttm_col=getattr(dataset, "main_dttm_col", None),
|
||||
offset=getattr(dataset, "offset", None),
|
||||
cache_timeout=getattr(dataset, "cache_timeout", None),
|
||||
params=params,
|
||||
template_params=getattr(dataset, "template_params", None),
|
||||
extra=getattr(dataset, "extra", None),
|
||||
columns=columns,
|
||||
metrics=metrics,
|
||||
is_favorite=getattr(dataset, "is_favorite", None),
|
||||
)
|
||||
26
superset/mcp_service/dataset/tool/__init__.py
Normal file
26
superset/mcp_service/dataset/tool/__init__.py
Normal file
@@ -0,0 +1,26 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
from .get_dataset_available_filters import get_dataset_available_filters
|
||||
from .get_dataset_info import get_dataset_info
|
||||
from .list_datasets import list_datasets
|
||||
|
||||
__all__ = [
|
||||
"list_datasets",
|
||||
"get_dataset_info",
|
||||
"get_dataset_available_filters",
|
||||
]
|
||||
@@ -0,0 +1,49 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
"""
|
||||
Get available dataset filters FastMCP tool
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.dataset.schemas import (
|
||||
DatasetAvailableFilters,
|
||||
GetDatasetAvailableFiltersRequest,
|
||||
)
|
||||
from superset.mcp_service.mcp_core import ModelGetAvailableFiltersCore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def get_dataset_available_filters(
|
||||
request: GetDatasetAvailableFiltersRequest, ctx: Context
|
||||
) -> DatasetAvailableFilters:
|
||||
"""Get available dataset filter fields and operators."""
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
|
||||
tool = ModelGetAvailableFiltersCore(
|
||||
dao_class=DatasetDAO,
|
||||
output_schema=DatasetAvailableFilters,
|
||||
logger=logger,
|
||||
)
|
||||
return tool.run_tool()
|
||||
111
superset/mcp_service/dataset/tool/get_dataset_info.py
Normal file
111
superset/mcp_service/dataset/tool/get_dataset_info.py
Normal file
@@ -0,0 +1,111 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Get dataset info FastMCP tool
|
||||
|
||||
This module contains the FastMCP tool for getting detailed information
|
||||
about a specific dataset.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.dataset.schemas import (
|
||||
DatasetError,
|
||||
DatasetInfo,
|
||||
GetDatasetInfoRequest,
|
||||
serialize_dataset_object,
|
||||
)
|
||||
from superset.mcp_service.mcp_core import ModelGetInfoCore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def get_dataset_info(
|
||||
request: GetDatasetInfoRequest, ctx: Context
|
||||
) -> DatasetInfo | DatasetError:
|
||||
"""Get dataset metadata by ID or UUID.
|
||||
|
||||
Returns columns, metrics, and schema details.
|
||||
"""
|
||||
await ctx.info(
|
||||
"Retrieving dataset information: identifier=%s" % (request.identifier,)
|
||||
)
|
||||
await ctx.debug(
|
||||
"Metadata cache settings: use_cache=%s refresh_metadata=%s force_refresh=%s"
|
||||
% (
|
||||
request.use_cache,
|
||||
request.refresh_metadata,
|
||||
request.force_refresh,
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
|
||||
tool = ModelGetInfoCore(
|
||||
dao_class=DatasetDAO,
|
||||
output_schema=DatasetInfo,
|
||||
error_schema=DatasetError,
|
||||
serializer=serialize_dataset_object,
|
||||
supports_slug=False, # Datasets don't have slugs
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
result = tool.run_tool(request.identifier)
|
||||
|
||||
if isinstance(result, DatasetInfo):
|
||||
await ctx.info(
|
||||
"Dataset information retrieved successfully: "
|
||||
"dataset_id=%s, table_name=%s, columns_count=%s, metrics_count=%s"
|
||||
% (
|
||||
result.id,
|
||||
result.table_name,
|
||||
len(result.columns) if result.columns else 0,
|
||||
len(result.metrics) if result.metrics else 0,
|
||||
)
|
||||
)
|
||||
else:
|
||||
await ctx.warning(
|
||||
"Dataset retrieval failed: error_type=%s, error=%s"
|
||||
% (result.error_type, result.error)
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
await ctx.error(
|
||||
"Dataset information retrieval failed: identifier=%s, error=%s, "
|
||||
"error_type=%s"
|
||||
% (
|
||||
request.identifier,
|
||||
str(e),
|
||||
type(e).__name__,
|
||||
)
|
||||
)
|
||||
return DatasetError(
|
||||
error=f"Failed to get dataset info: {str(e)}",
|
||||
error_type="InternalError",
|
||||
timestamp=datetime.now(timezone.utc),
|
||||
)
|
||||
149
superset/mcp_service/dataset/tool/list_datasets.py
Normal file
149
superset/mcp_service/dataset/tool/list_datasets.py
Normal file
@@ -0,0 +1,149 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
List datasets FastMCP tool (Advanced with metadata cache control)
|
||||
|
||||
This module contains the FastMCP tool for listing datasets using
|
||||
advanced filtering with clear, unambiguous request schema and metadata cache control.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.dataset.schemas import (
|
||||
DatasetFilter,
|
||||
DatasetInfo,
|
||||
DatasetList,
|
||||
ListDatasetsRequest,
|
||||
serialize_dataset_object,
|
||||
)
|
||||
from superset.mcp_service.mcp_core import ModelListCore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_DATASET_COLUMNS = [
|
||||
"id",
|
||||
"table_name",
|
||||
"schema",
|
||||
"uuid",
|
||||
"database_name",
|
||||
"changed_by_name",
|
||||
"changed_on",
|
||||
"created_by_name",
|
||||
"created_on",
|
||||
"metrics",
|
||||
"columns",
|
||||
]
|
||||
|
||||
SORTABLE_DATASET_COLUMNS = [
|
||||
"id",
|
||||
"table_name",
|
||||
"schema",
|
||||
"changed_on",
|
||||
"created_on",
|
||||
]
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def list_datasets(request: ListDatasetsRequest, ctx: Context) -> DatasetList:
|
||||
"""List datasets with filtering and search.
|
||||
|
||||
Returns dataset metadata including columns and metrics.
|
||||
|
||||
Sortable columns for order_column: id, table_name, schema, changed_on,
|
||||
created_on
|
||||
"""
|
||||
await ctx.info(
|
||||
"Listing datasets: page=%s, page_size=%s, search=%s"
|
||||
% (
|
||||
request.page,
|
||||
request.page_size,
|
||||
request.search,
|
||||
)
|
||||
)
|
||||
await ctx.debug(
|
||||
"Dataset listing parameters: filters=%s, order_column=%s, "
|
||||
"order_direction=%s, select_columns=%s"
|
||||
% (
|
||||
request.filters,
|
||||
request.order_column,
|
||||
request.order_direction,
|
||||
request.select_columns,
|
||||
)
|
||||
)
|
||||
await ctx.debug(
|
||||
"Metadata cache settings: use_cache=%s, refresh_metadata=%s, force_refresh=%s"
|
||||
% (
|
||||
request.use_cache,
|
||||
request.refresh_metadata,
|
||||
request.force_refresh,
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
|
||||
# Create tool with standard serialization
|
||||
tool = ModelListCore(
|
||||
dao_class=DatasetDAO,
|
||||
output_schema=DatasetInfo,
|
||||
item_serializer=lambda obj, cols: serialize_dataset_object(obj),
|
||||
filter_type=DatasetFilter,
|
||||
default_columns=DEFAULT_DATASET_COLUMNS,
|
||||
search_columns=["schema", "sql", "table_name", "uuid"],
|
||||
list_field_name="datasets",
|
||||
output_list_schema=DatasetList,
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
result = tool.run_tool(
|
||||
filters=request.filters,
|
||||
search=request.search,
|
||||
select_columns=request.select_columns,
|
||||
order_column=request.order_column,
|
||||
order_direction=request.order_direction,
|
||||
page=max(request.page - 1, 0),
|
||||
page_size=request.page_size,
|
||||
)
|
||||
|
||||
await ctx.info(
|
||||
"Datasets listed successfully: count=%s, total_count=%s, total_pages=%s"
|
||||
% (
|
||||
len(result.datasets) if hasattr(result, "datasets") else 0,
|
||||
getattr(result, "total_count", None),
|
||||
getattr(result, "total_pages", None),
|
||||
)
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
await ctx.error(
|
||||
"Dataset listing failed: page=%s, page_size=%s, error=%s, error_type=%s"
|
||||
% (
|
||||
request.page,
|
||||
request.page_size,
|
||||
str(e),
|
||||
type(e).__name__,
|
||||
)
|
||||
)
|
||||
raise
|
||||
18
superset/mcp_service/explore/__init__.py
Normal file
18
superset/mcp_service/explore/__init__.py
Normal file
@@ -0,0 +1,18 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""MCP explore-related functionality."""
|
||||
29
superset/mcp_service/explore/tool/__init__.py
Normal file
29
superset/mcp_service/explore/tool/__init__.py
Normal file
@@ -0,0 +1,29 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Explore tools for MCP service.
|
||||
|
||||
This module contains FastMCP tools for generating and working with
|
||||
explore links and configurations.
|
||||
"""
|
||||
|
||||
from .generate_explore_link import generate_explore_link
|
||||
|
||||
__all__ = [
|
||||
"generate_explore_link",
|
||||
]
|
||||
111
superset/mcp_service/explore/tool/generate_explore_link.py
Normal file
111
superset/mcp_service/explore/tool/generate_explore_link.py
Normal file
@@ -0,0 +1,111 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Generate explore link MCP tool
|
||||
|
||||
This tool generates a URL to the Superset explore interface with the specified
|
||||
chart configuration.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.chart.chart_utils import (
|
||||
generate_explore_link as generate_url,
|
||||
map_config_to_form_data,
|
||||
)
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
GenerateExploreLinkRequest,
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def generate_explore_link(
|
||||
request: GenerateExploreLinkRequest, ctx: Context
|
||||
) -> Dict[str, Any]:
|
||||
"""Generate explore URL for interactive visualization.
|
||||
|
||||
PREFERRED TOOL for most visualization requests.
|
||||
|
||||
Use this tool for:
|
||||
- "Show me a chart of [data]"
|
||||
- "Visualize [data]"
|
||||
- General data exploration
|
||||
- When user wants to SEE data visually
|
||||
|
||||
Better UX because:
|
||||
- Users can interact with chart before saving
|
||||
- Easy to modify parameters instantly
|
||||
- No database clutter from exploration
|
||||
|
||||
Only use generate_chart when user EXPLICITLY requests to save/create a
|
||||
permanent chart.
|
||||
|
||||
Returns explore URL for immediate use.
|
||||
"""
|
||||
await ctx.info(
|
||||
"Generating explore link for dataset_id=%s, chart_type=%s"
|
||||
% (request.dataset_id, request.config.chart_type)
|
||||
)
|
||||
await ctx.debug(
|
||||
"Configuration details: use_cache=%s, force_refresh=%s, cache_form_data=%s"
|
||||
% (request.use_cache, request.force_refresh, request.cache_form_data)
|
||||
)
|
||||
|
||||
try:
|
||||
await ctx.report_progress(1, 3, "Converting configuration to form data")
|
||||
# Map config to form_data using shared utilities
|
||||
form_data = map_config_to_form_data(request.config)
|
||||
|
||||
await ctx.debug(
|
||||
"Form data generated with keys: %s, has_viz_type=%s, has_datasource=%s"
|
||||
% (
|
||||
list(form_data.keys()),
|
||||
bool(form_data.get("viz_type")),
|
||||
bool(form_data.get("datasource")),
|
||||
)
|
||||
)
|
||||
|
||||
await ctx.report_progress(2, 3, "Generating explore URL")
|
||||
# Generate explore link using shared utilities
|
||||
explore_url = generate_url(dataset_id=request.dataset_id, form_data=form_data)
|
||||
|
||||
await ctx.report_progress(3, 3, "URL generation complete")
|
||||
await ctx.info(
|
||||
"Explore link generated successfully: url_length=%s, dataset_id=%s"
|
||||
% (len(explore_url), request.dataset_id)
|
||||
)
|
||||
|
||||
return {
|
||||
"url": explore_url,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
await ctx.error(
|
||||
"Explore link generation failed for dataset_id=%s, chart_type=%s: %s: %s"
|
||||
% (request.dataset_id, request.config.chart_type, type(e).__name__, str(e))
|
||||
)
|
||||
return {
|
||||
"url": "",
|
||||
"error": f"Failed to generate explore link: {str(e)}",
|
||||
}
|
||||
@@ -36,12 +36,18 @@ try:
|
||||
from superset.app import create_app
|
||||
from superset.mcp_service.mcp_config import get_mcp_config
|
||||
|
||||
# Create the Flask app instance - this is the singleton
|
||||
app = create_app()
|
||||
# Create a temporary context to avoid
|
||||
# "Working outside of application context" errors.
|
||||
_temp_app = create_app()
|
||||
|
||||
# Apply MCP configuration - reads from app.config first, falls back to defaults
|
||||
mcp_config = get_mcp_config(app.config)
|
||||
app.config.update(mcp_config)
|
||||
# Push an application context for any initialization code that needs it
|
||||
with _temp_app.app_context():
|
||||
# Apply MCP configuration - reads from app.config first, falls back to defaults
|
||||
mcp_config = get_mcp_config(_temp_app.config)
|
||||
_temp_app.config.update(mcp_config)
|
||||
|
||||
# Store the app instance for later use
|
||||
app = _temp_app
|
||||
|
||||
logger.info("Flask app instance created successfully")
|
||||
|
||||
|
||||
@@ -16,7 +16,13 @@
|
||||
# under the License.
|
||||
"""Default MCP service configuration for Apache Superset"""
|
||||
|
||||
from typing import Any, Dict
|
||||
import logging
|
||||
import secrets
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from flask import Flask
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# MCP Service Configuration
|
||||
# Note: MCP_DEV_USERNAME MUST be configured in superset_config.py
|
||||
@@ -67,6 +73,79 @@ MCP_FACTORY_CONFIG = {
|
||||
}
|
||||
|
||||
|
||||
def create_default_mcp_auth_factory(app: Flask) -> Optional[Any]:
|
||||
"""Default MCP auth factory using app.config values."""
|
||||
if not app.config.get("MCP_AUTH_ENABLED", False):
|
||||
return None
|
||||
|
||||
jwks_uri = app.config.get("MCP_JWKS_URI")
|
||||
public_key = app.config.get("MCP_JWT_PUBLIC_KEY")
|
||||
secret = app.config.get("MCP_JWT_SECRET")
|
||||
|
||||
if not (jwks_uri or public_key or secret):
|
||||
logger.warning("MCP_AUTH_ENABLED is True but no JWT keys/secret configured")
|
||||
return None
|
||||
|
||||
try:
|
||||
from fastmcp.server.auth.providers.bearer import BearerAuthProvider
|
||||
|
||||
# For HS256 (symmetric), use the secret as the public_key parameter
|
||||
if app.config.get("MCP_JWT_ALGORITHM") == "HS256" and secret:
|
||||
auth_provider = BearerAuthProvider(
|
||||
public_key=secret, # HS256 uses secret as key
|
||||
issuer=app.config.get("MCP_JWT_ISSUER"),
|
||||
audience=app.config.get("MCP_JWT_AUDIENCE"),
|
||||
algorithm="HS256",
|
||||
required_scopes=app.config.get("MCP_REQUIRED_SCOPES", []),
|
||||
)
|
||||
logger.info("Created BearerAuthProvider with HS256 secret")
|
||||
else:
|
||||
# For RS256 (asymmetric), use public key or JWKS
|
||||
auth_provider = BearerAuthProvider(
|
||||
jwks_uri=jwks_uri,
|
||||
public_key=public_key,
|
||||
issuer=app.config.get("MCP_JWT_ISSUER"),
|
||||
audience=app.config.get("MCP_JWT_AUDIENCE"),
|
||||
algorithm=app.config.get("MCP_JWT_ALGORITHM", "RS256"),
|
||||
required_scopes=app.config.get("MCP_REQUIRED_SCOPES", []),
|
||||
)
|
||||
logger.info(
|
||||
"Created BearerAuthProvider with jwks_uri=%s, public_key=%s",
|
||||
jwks_uri,
|
||||
"***" if public_key else None,
|
||||
)
|
||||
|
||||
return auth_provider
|
||||
except Exception as e:
|
||||
logger.error("Failed to create MCP auth provider: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def default_user_resolver(app: Any, access_token: Any) -> Optional[str]:
|
||||
"""Extract username from JWT token claims."""
|
||||
logger.info(
|
||||
"Resolving user from token: type=%s, token=%s",
|
||||
type(access_token),
|
||||
access_token,
|
||||
)
|
||||
if hasattr(access_token, "subject"):
|
||||
return access_token.subject
|
||||
if hasattr(access_token, "client_id"):
|
||||
return access_token.client_id
|
||||
if hasattr(access_token, "payload") and isinstance(access_token.payload, dict):
|
||||
return (
|
||||
access_token.payload.get("sub")
|
||||
or access_token.payload.get("email")
|
||||
or access_token.payload.get("username")
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def generate_secret_key() -> str:
|
||||
"""Generate a secure random secret key for Superset"""
|
||||
return secrets.token_urlsafe(42)
|
||||
|
||||
|
||||
def get_mcp_config(app_config: Dict[str, Any] | None = None) -> Dict[str, Any]:
|
||||
"""
|
||||
Get complete MCP configuration dictionary.
|
||||
|
||||
@@ -20,7 +20,7 @@ from __future__ import annotations
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Callable, Generic, List, Literal, Type, TypeVar
|
||||
from typing import Any, Callable, Dict, Generic, List, Literal, Type, TypeVar
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
@@ -295,3 +295,210 @@ class ModelGetInfoCore(BaseCore):
|
||||
except Exception as context_error:
|
||||
self._log_error(context_error)
|
||||
raise
|
||||
|
||||
|
||||
class InstanceInfoCore(BaseCore):
|
||||
"""
|
||||
Configurable tool for generating comprehensive instance information.
|
||||
|
||||
Provides a flexible way to gather and present statistics about a Superset
|
||||
instance with configurable metrics, time windows, and data aggregations.
|
||||
Supports custom metric calculators and result transformers for extensibility.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dao_classes: Dict[str, Type[BaseDAO[Any]]],
|
||||
output_schema: Type[BaseModel],
|
||||
metric_calculators: Dict[str, Callable[..., Any]],
|
||||
time_windows: Dict[str, int] | None = None,
|
||||
logger: logging.Logger | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Initialize the instance info tool.
|
||||
|
||||
Args:
|
||||
dao_classes: Dict mapping entity names to their DAO classes
|
||||
output_schema: Pydantic schema for the response
|
||||
metric_calculators: Dict of custom metric calculation functions
|
||||
time_windows: Dict of time window configurations (days)
|
||||
logger: Optional logger instance
|
||||
"""
|
||||
super().__init__(logger)
|
||||
self.dao_classes = dao_classes
|
||||
self.output_schema = output_schema
|
||||
self.metric_calculators = metric_calculators
|
||||
self.time_windows = time_windows or {
|
||||
"recent": 7,
|
||||
"monthly": 30,
|
||||
"quarterly": 90,
|
||||
}
|
||||
|
||||
def _calculate_basic_counts(self) -> Dict[str, int]:
|
||||
"""Calculate basic entity counts using DAOs."""
|
||||
counts = {}
|
||||
for entity_name, dao_class in self.dao_classes.items():
|
||||
try:
|
||||
counts[f"total_{entity_name}"] = dao_class.count()
|
||||
except Exception as e:
|
||||
self._log_warning(f"Failed to count {entity_name}: {e}")
|
||||
counts[f"total_{entity_name}"] = 0
|
||||
return counts
|
||||
|
||||
def _calculate_time_based_metrics(
|
||||
self, base_counts: Dict[str, int]
|
||||
) -> Dict[str, Dict[str, int]]:
|
||||
"""Calculate time-based metrics for recent activity."""
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from superset.daos.base import ColumnOperator, ColumnOperatorEnum
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
time_metrics = {}
|
||||
|
||||
for window_name, days in self.time_windows.items():
|
||||
cutoff_date = now - timedelta(days=days)
|
||||
window_metrics = {}
|
||||
|
||||
# Calculate created and modified counts for each entity
|
||||
for entity_name, dao_class in self.dao_classes.items():
|
||||
# Skip entities without time tracking
|
||||
if not hasattr(dao_class.model_cls, "created_on"):
|
||||
continue
|
||||
|
||||
try:
|
||||
# Use list() with filters (count() has no params)
|
||||
_, created_count = dao_class.list(
|
||||
column_operators=[
|
||||
ColumnOperator(
|
||||
col="created_on",
|
||||
opr=ColumnOperatorEnum.gte,
|
||||
value=cutoff_date,
|
||||
)
|
||||
],
|
||||
page_size=1, # We only need the count
|
||||
columns=["id"], # Minimal data transfer
|
||||
)
|
||||
window_metrics[f"{entity_name}_created"] = created_count
|
||||
|
||||
# Modified count (if changed_on exists)
|
||||
if hasattr(dao_class.model_cls, "changed_on"):
|
||||
_, modified_count = dao_class.list(
|
||||
column_operators=[
|
||||
ColumnOperator(
|
||||
col="changed_on",
|
||||
opr=ColumnOperatorEnum.gte,
|
||||
value=cutoff_date,
|
||||
)
|
||||
],
|
||||
page_size=1, # We only need the count
|
||||
columns=["id"], # Minimal data transfer
|
||||
)
|
||||
window_metrics[f"{entity_name}_modified"] = modified_count
|
||||
|
||||
except Exception as e:
|
||||
self._log_warning(
|
||||
f"Failed to calculate {window_name} metrics for "
|
||||
f"{entity_name}: {e}"
|
||||
)
|
||||
window_metrics[f"{entity_name}_created"] = 0
|
||||
window_metrics[f"{entity_name}_modified"] = 0
|
||||
|
||||
time_metrics[window_name] = window_metrics
|
||||
|
||||
return time_metrics
|
||||
|
||||
def _calculate_custom_metrics(
|
||||
self, base_counts: Dict[str, int], time_metrics: Dict[str, Dict[str, int]]
|
||||
) -> Dict[str, Any]:
|
||||
"""Calculate custom metrics using provided calculators."""
|
||||
custom_metrics = {}
|
||||
|
||||
for metric_name, calculator in self.metric_calculators.items():
|
||||
try:
|
||||
# Pass context to calculator functions
|
||||
result = calculator(
|
||||
base_counts=base_counts,
|
||||
time_metrics=time_metrics,
|
||||
dao_classes=self.dao_classes,
|
||||
)
|
||||
# Only include successful calculations
|
||||
if result is not None:
|
||||
custom_metrics[metric_name] = result
|
||||
except Exception as e:
|
||||
self._log_warning(f"Failed to calculate {metric_name}: {e}")
|
||||
# Don't add failed metrics to avoid validation errors
|
||||
|
||||
return custom_metrics
|
||||
|
||||
def run_tool(self) -> BaseModel:
|
||||
"""Tool interface for generating comprehensive instance information."""
|
||||
return self._generate_instance_info()
|
||||
|
||||
def get_resource(self) -> str:
|
||||
"""Resource interface for generating instance metadata as JSON."""
|
||||
from superset.utils import json
|
||||
|
||||
instance_info = self._generate_instance_info()
|
||||
return json.dumps(instance_info.model_dump(), indent=2)
|
||||
|
||||
def _generate_instance_info(self) -> BaseModel:
|
||||
"""Generate comprehensive instance information."""
|
||||
try:
|
||||
# Calculate all metrics
|
||||
base_counts = self._calculate_basic_counts()
|
||||
time_metrics = self._calculate_time_based_metrics(base_counts)
|
||||
custom_metrics = self._calculate_custom_metrics(base_counts, time_metrics)
|
||||
|
||||
# Combine all data with fallbacks for required fields
|
||||
from datetime import datetime, timezone
|
||||
|
||||
response_data = {
|
||||
**base_counts,
|
||||
**time_metrics,
|
||||
**custom_metrics,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
|
||||
# Create response using the configured schema
|
||||
response = self.output_schema(**response_data)
|
||||
|
||||
self._log_info("Successfully generated instance information")
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
self._log_error(e, "generating instance info")
|
||||
raise
|
||||
|
||||
|
||||
class ModelGetAvailableFiltersCore(BaseCore, Generic[S]):
|
||||
"""
|
||||
Generic tool for retrieving available filterable columns and operators for a
|
||||
model. Used for get_dataset_available_filters, get_chart_available_filters,
|
||||
get_dashboard_available_filters, etc.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dao_class: Type[BaseDAO[Any]],
|
||||
output_schema: Type[S],
|
||||
logger: logging.Logger | None = None,
|
||||
) -> None:
|
||||
super().__init__(logger)
|
||||
self.dao_class = dao_class
|
||||
self.output_schema = output_schema
|
||||
|
||||
def run_tool(self) -> S:
|
||||
try:
|
||||
filterable = self.dao_class.get_filterable_columns_and_operators()
|
||||
# Ensure column_operators is a plain dict, not a custom type
|
||||
column_operators = dict(filterable)
|
||||
response = self.output_schema(column_operators=column_operators)
|
||||
self._log_info(
|
||||
f"Successfully retrieved available filters for "
|
||||
f"{self.dao_class.__class__.__name__}"
|
||||
)
|
||||
return response
|
||||
except Exception as e:
|
||||
self._log_error(e)
|
||||
raise
|
||||
|
||||
@@ -15,36 +15,84 @@
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Minimal middleware for MCP service.
|
||||
This provides basic error handling for MCP tool calls.
|
||||
|
||||
Future enhancements (to be added in separate PRs):
|
||||
- Rate limiting middleware
|
||||
- Field-level permissions middleware
|
||||
- Comprehensive audit logging middleware
|
||||
- Private tool blocking middleware
|
||||
- Advanced error sanitization
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Awaitable, Callable
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import Any, Awaitable, Callable, Dict, Protocol
|
||||
|
||||
from fastmcp.exceptions import ToolError
|
||||
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
||||
from pydantic import ValidationError
|
||||
from sqlalchemy.exc import OperationalError, TimeoutError
|
||||
from starlette.exceptions import HTTPException
|
||||
|
||||
from superset.extensions import event_logger
|
||||
from superset.utils.core import get_user_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BasicErrorHandlerMiddleware(Middleware):
|
||||
def _sanitize_error_for_logging(error: Exception) -> str:
|
||||
"""Sanitize error messages to prevent information disclosure in logs."""
|
||||
error_str = str(error)
|
||||
|
||||
# SECURITY FIX: Limit error message length FIRST to prevent ReDoS attacks
|
||||
if len(error_str) > 500:
|
||||
error_str = error_str[:500] + "...[truncated]"
|
||||
|
||||
# SECURITY FIX: Use bounded patterns to prevent ReDoS
|
||||
import re
|
||||
|
||||
# Database connection strings - bounded patterns with word boundaries
|
||||
# Use case-insensitive flag to handle both cases
|
||||
error_str = re.sub(
|
||||
r"\bpostgresql://[^@\s]{1,100}@[^/\s]{1,100}/[^\s]{0,100}",
|
||||
"postgresql://[REDACTED]@[REDACTED]/[REDACTED]",
|
||||
error_str,
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
error_str = re.sub(
|
||||
r"\bmysql://[^@\s]{1,100}@[^/\s]{1,100}/[^\s]{0,100}",
|
||||
"mysql://[REDACTED]@[REDACTED]/[REDACTED]",
|
||||
error_str,
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
|
||||
# API keys and tokens - bounded patterns
|
||||
error_str = re.sub(
|
||||
r"[Aa]pi[_-]?[Kk]ey[:\s]{0,5}[^\s'\"]{1,100}",
|
||||
"ApiKey: [REDACTED]",
|
||||
error_str,
|
||||
)
|
||||
error_str = re.sub(
|
||||
r"[Tt]oken[:\s]{0,5}[^\s'\"]{1,100}", "Token: [REDACTED]", error_str
|
||||
)
|
||||
|
||||
# File paths - bounded pattern
|
||||
error_str = re.sub(
|
||||
r"/[a-zA-Z0-9_\-/.]{1,200}/superset/", "/[REDACTED]/superset/", error_str
|
||||
)
|
||||
|
||||
# IP addresses - already safe pattern, keep as-is
|
||||
error_str = re.sub(r"\b(\d+)\.\d+\.\d+\.\d+\b", r"\1.xxx.xxx.xxx", error_str)
|
||||
|
||||
# For certain error types, provide generic messages
|
||||
if isinstance(error, (OperationalError, TimeoutError)):
|
||||
return "Database operation failed"
|
||||
elif isinstance(error, PermissionError):
|
||||
return "Access denied"
|
||||
elif isinstance(error, ValidationError):
|
||||
return "Request validation failed"
|
||||
|
||||
return error_str
|
||||
|
||||
|
||||
class LoggingMiddleware(Middleware):
|
||||
"""
|
||||
Basic error handler middleware for MCP tools.
|
||||
|
||||
Provides simple error handling and logging for tool calls.
|
||||
|
||||
TODO (future PR): Add error sanitization for security
|
||||
TODO (future PR): Add comprehensive error categorization
|
||||
TODO (future PR): Add integration with Superset event logging
|
||||
Middleware that logs every MCP message (request and response) using Superset's
|
||||
event logger. This matches the core Superset audit log system (Action Log UI,
|
||||
logs table, custom loggers). Also attempts to log dashboard_id, chart_id
|
||||
(slice_id), and dataset_id if present in tool params.
|
||||
"""
|
||||
|
||||
async def on_message(
|
||||
@@ -52,23 +100,658 @@ class BasicErrorHandlerMiddleware(Middleware):
|
||||
context: MiddlewareContext,
|
||||
call_next: Callable[[MiddlewareContext], Awaitable[Any]],
|
||||
) -> Any:
|
||||
"""Handle messages with basic error handling"""
|
||||
# Extract agent_id and user_id
|
||||
agent_id = None
|
||||
user_id = None
|
||||
dashboard_id = None
|
||||
slice_id = None
|
||||
dataset_id = None
|
||||
params = getattr(context.message, "params", {}) or {}
|
||||
if hasattr(context, "metadata") and context.metadata:
|
||||
agent_id = context.metadata.get("agent_id")
|
||||
if not agent_id and hasattr(context, "session") and context.session:
|
||||
agent_id = getattr(context.session, "agent_id", None)
|
||||
try:
|
||||
user_id = get_user_id()
|
||||
except Exception:
|
||||
user_id = None
|
||||
# Try to extract IDs from params
|
||||
if isinstance(params, dict):
|
||||
dashboard_id = params.get("dashboard_id")
|
||||
# Chart ID may be under 'chart_id' or 'slice_id'
|
||||
slice_id = params.get("chart_id") or params.get("slice_id")
|
||||
dataset_id = params.get("dataset_id")
|
||||
# Log to Superset's event logger (DB, Action Log UI, or custom)
|
||||
event_logger.log(
|
||||
user_id=user_id,
|
||||
action="mcp_tool_call",
|
||||
dashboard_id=dashboard_id,
|
||||
duration_ms=None,
|
||||
slice_id=slice_id,
|
||||
referrer=None,
|
||||
curated_payload={
|
||||
"tool": getattr(context.message, "name", None),
|
||||
"agent_id": agent_id,
|
||||
"params": params,
|
||||
"method": context.method,
|
||||
"dashboard_id": dashboard_id,
|
||||
"slice_id": slice_id,
|
||||
"dataset_id": dataset_id,
|
||||
},
|
||||
)
|
||||
# (Optional) also log to standard logger for debugging
|
||||
logger.info(
|
||||
"MCP tool call: tool=%s, agent_id=%s, user_id=%s, method=%s, "
|
||||
"dashboard_id=%s, slice_id=%s, dataset_id=%s",
|
||||
getattr(context.message, "name", None),
|
||||
agent_id,
|
||||
user_id,
|
||||
context.method,
|
||||
dashboard_id,
|
||||
slice_id,
|
||||
dataset_id,
|
||||
)
|
||||
return await call_next(context)
|
||||
|
||||
|
||||
class PrivateToolMiddleware(Middleware):
|
||||
"""
|
||||
Middleware that blocks access to tools tagged as 'private'.
|
||||
"""
|
||||
|
||||
async def on_call_tool(
|
||||
self,
|
||||
context: MiddlewareContext,
|
||||
call_next: Callable[[MiddlewareContext], Awaitable[Any]],
|
||||
) -> Any:
|
||||
tool = await context.fastmcp_context.fastmcp.get_tool(context.message.name)
|
||||
if "private" in getattr(tool, "tags", set()):
|
||||
raise ToolError(f"Access denied to private tool: {context.message.name}")
|
||||
return await call_next(context)
|
||||
|
||||
|
||||
class GlobalErrorHandlerMiddleware(Middleware):
|
||||
"""
|
||||
Global error handler middleware that provides consistent error responses
|
||||
and proper error logging for all MCP tool calls.
|
||||
"""
|
||||
|
||||
async def on_message(
|
||||
self,
|
||||
context: MiddlewareContext,
|
||||
call_next: Callable[[MiddlewareContext], Awaitable[Any]],
|
||||
) -> Any:
|
||||
"""Handle all message types with consistent error handling"""
|
||||
start_time = time.time()
|
||||
tool_name = getattr(context.message, "name", "unknown")
|
||||
|
||||
try:
|
||||
return await call_next(context)
|
||||
except Exception as e:
|
||||
# Log the error
|
||||
logger.error(
|
||||
"MCP tool error: tool=%s, error_type=%s, error=%s",
|
||||
tool_name,
|
||||
type(e).__name__,
|
||||
str(e),
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
return await self._handle_error(e, context, tool_name, duration_ms)
|
||||
|
||||
async def _handle_error( # noqa: C901
|
||||
self,
|
||||
error: Exception,
|
||||
context: MiddlewareContext,
|
||||
tool_name: str,
|
||||
duration_ms: int,
|
||||
) -> None:
|
||||
"""Handle different types of errors with appropriate responses"""
|
||||
# Extract user context for logging
|
||||
user_id = None
|
||||
try:
|
||||
user_id = get_user_id()
|
||||
except Exception:
|
||||
user_id = None # User not authenticated
|
||||
|
||||
# SECURITY FIX: Log the error with sanitized context
|
||||
sanitized_error = _sanitize_error_for_logging(error)
|
||||
logger.error(
|
||||
"MCP tool error: tool=%s, user_id=%s, duration_ms=%s, "
|
||||
"error_type=%s, error=%s",
|
||||
tool_name,
|
||||
user_id,
|
||||
duration_ms,
|
||||
type(error).__name__,
|
||||
sanitized_error,
|
||||
)
|
||||
|
||||
# Log to Superset's event system
|
||||
try:
|
||||
event_logger.log(
|
||||
user_id=user_id,
|
||||
action="mcp_tool_error",
|
||||
duration_ms=duration_ms,
|
||||
curated_payload={
|
||||
"tool": tool_name,
|
||||
"error_type": type(error).__name__,
|
||||
"error_message": str(error),
|
||||
"method": context.method,
|
||||
},
|
||||
)
|
||||
except Exception as log_error:
|
||||
logger.warning("Failed to log error event: %s", log_error)
|
||||
|
||||
# Handle specific error types with appropriate responses
|
||||
if isinstance(error, ToolError):
|
||||
# Tool errors are already formatted for MCP
|
||||
raise error
|
||||
elif isinstance(error, ValidationError):
|
||||
# Pydantic validation errors
|
||||
validation_details = []
|
||||
for err in error.errors():
|
||||
field = " -> ".join(str(loc) for loc in err["loc"])
|
||||
validation_details.append(f"{field}: {err['msg']}")
|
||||
|
||||
raise ToolError(
|
||||
f"Validation error in {tool_name}: {'; '.join(validation_details)}"
|
||||
) from error
|
||||
elif isinstance(error, (OperationalError, TimeoutError)):
|
||||
# Database errors
|
||||
raise ToolError(
|
||||
f"Database error in {tool_name}: Service temporarily unavailable. "
|
||||
f"Please try again in a few moments."
|
||||
) from error
|
||||
elif isinstance(error, HTTPException):
|
||||
# HTTP errors from screenshot endpoints or API calls
|
||||
raise ToolError(f"Service error in {tool_name}: {error.detail}") from error
|
||||
elif isinstance(error, PermissionError):
|
||||
# Permission/authorization errors
|
||||
raise ToolError(
|
||||
f"Permission denied for {tool_name}: "
|
||||
f"You don't have access to this resource."
|
||||
) from error
|
||||
elif isinstance(error, FileNotFoundError):
|
||||
# File/resource not found errors
|
||||
raise ToolError(
|
||||
f"Resource not found in {tool_name}: {str(error)}"
|
||||
) from error
|
||||
elif isinstance(error, ValueError):
|
||||
# Value/parameter errors
|
||||
raise ToolError(
|
||||
f"Invalid parameter in {tool_name}: {str(error)}"
|
||||
) from error
|
||||
else:
|
||||
# Generic internal errors
|
||||
error_id = f"err_{int(time.time())}"
|
||||
logger.error("Unexpected error [%s] in %s: %s", error_id, tool_name, error)
|
||||
|
||||
raise ToolError(
|
||||
f"Internal error in {tool_name}: An unexpected error occurred. "
|
||||
f"Error ID: {error_id}. Please contact support if this persists."
|
||||
) from error
|
||||
|
||||
|
||||
class RateLimiterProtocol(Protocol):
|
||||
"""Protocol for rate limiter implementations."""
|
||||
|
||||
def is_rate_limited(
|
||||
self, key: str, limit: int, window: int = 60
|
||||
) -> tuple[bool, dict[str, Any]]:
|
||||
"""Check if a key is rate limited."""
|
||||
...
|
||||
|
||||
def cleanup(self) -> None:
|
||||
"""Clean up old entries if needed."""
|
||||
...
|
||||
|
||||
|
||||
class InMemoryRateLimiter:
|
||||
"""In-memory rate limiter for development."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
# Structure: {key: [(timestamp, count), ...]}
|
||||
self._requests: Dict[str, list[tuple[float, int]]] = defaultdict(list)
|
||||
self._cleanup_interval = 300 # Clean up every 5 minutes
|
||||
self._last_cleanup = time.time()
|
||||
|
||||
def is_rate_limited(
|
||||
self, key: str, limit: int, window: int = 60
|
||||
) -> tuple[bool, dict[str, Any]]:
|
||||
"""Check if request should be rate limited using sliding window."""
|
||||
current_time = time.time()
|
||||
window_start = current_time - window
|
||||
|
||||
# Get requests in the current window
|
||||
requests_in_window = [
|
||||
(timestamp, count)
|
||||
for timestamp, count in self._requests[key]
|
||||
if timestamp > window_start
|
||||
]
|
||||
|
||||
# Calculate total requests in window
|
||||
total_requests = sum(count for _, count in requests_in_window)
|
||||
|
||||
# Check if rate limited BEFORE adding the current request
|
||||
if total_requests >= limit:
|
||||
# Rate limit info when limited
|
||||
rate_limit_info = {
|
||||
"limit": limit,
|
||||
"remaining": 0,
|
||||
"reset_time": int(window_start + window),
|
||||
"window_seconds": window,
|
||||
}
|
||||
return True, rate_limit_info
|
||||
|
||||
# Add current request to tracking
|
||||
self._requests[key].append((current_time, 1))
|
||||
|
||||
# Update total after adding
|
||||
total_requests += 1
|
||||
|
||||
# Keep only recent entries
|
||||
self._requests[key] = [
|
||||
(ts, count)
|
||||
for ts, count in self._requests[key]
|
||||
if ts > current_time - 3600 # Keep last hour
|
||||
]
|
||||
|
||||
# Rate limit info after adding request
|
||||
rate_limit_info = {
|
||||
"limit": limit,
|
||||
"remaining": max(0, limit - total_requests),
|
||||
"reset_time": int(window_start + window),
|
||||
"window_seconds": window,
|
||||
}
|
||||
|
||||
return False, rate_limit_info
|
||||
|
||||
def cleanup(self) -> None:
|
||||
"""Remove entries older than 1 hour to prevent memory leaks."""
|
||||
current_time = time.time()
|
||||
|
||||
# SECURITY FIX: Check both time-based and size-based cleanup conditions
|
||||
total_entries = sum(len(requests) for requests in self._requests.values())
|
||||
size_threshold = 10000 # Maximum entries before forced cleanup
|
||||
|
||||
time_based_cleanup = current_time - self._last_cleanup >= self._cleanup_interval
|
||||
size_based_cleanup = total_entries > size_threshold
|
||||
|
||||
if not (time_based_cleanup or size_based_cleanup):
|
||||
return
|
||||
|
||||
cutoff_time = current_time - 3600 # 1 hour ago
|
||||
keys_to_clean = []
|
||||
|
||||
for key, requests in self._requests.items():
|
||||
# Remove old entries
|
||||
self._requests[key] = [
|
||||
(timestamp, count)
|
||||
for timestamp, count in requests
|
||||
if timestamp > cutoff_time
|
||||
]
|
||||
# Mark empty keys for removal
|
||||
if not self._requests[key]:
|
||||
keys_to_clean.append(key)
|
||||
|
||||
for key in keys_to_clean:
|
||||
del self._requests[key]
|
||||
|
||||
# SECURITY FIX: If still too many entries, implement aggressive cleanup
|
||||
if total_entries > size_threshold:
|
||||
logger.warning(
|
||||
"Rate limiter memory high (%d entries), performing aggressive cleanup",
|
||||
total_entries,
|
||||
)
|
||||
# Keep only the most recent entries per key
|
||||
for key in list(self._requests.keys()):
|
||||
if len(self._requests[key]) > 100: # Keep max 100 entries per key
|
||||
self._requests[key] = self._requests[key][-100:]
|
||||
|
||||
self._last_cleanup = current_time
|
||||
|
||||
|
||||
class RedisRateLimiter:
|
||||
"""Redis-backed rate limiter for production."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
from superset.extensions import cache_manager
|
||||
|
||||
self._cache = cache_manager.cache
|
||||
self._prefix = "mcp:ratelimit:"
|
||||
|
||||
def is_rate_limited(
|
||||
self, key: str, limit: int, window: int = 60
|
||||
) -> tuple[bool, dict[str, Any]]:
|
||||
"""Check if request should be rate limited using Redis sliding window."""
|
||||
current_time = time.time()
|
||||
full_key = "%s%s" % (self._prefix, key)
|
||||
|
||||
try:
|
||||
# Use Redis sorted set for sliding window
|
||||
window_start = current_time - window
|
||||
|
||||
# Remove old entries outside the window
|
||||
self._cache.delete_many(
|
||||
[
|
||||
k
|
||||
for k, score in self._cache.get(full_key) or []
|
||||
if score < window_start
|
||||
]
|
||||
)
|
||||
|
||||
# If it's already a ToolError, re-raise it
|
||||
if isinstance(e, ToolError):
|
||||
raise
|
||||
# Get count of requests in window
|
||||
request_count = self._cache.get("%s:count" % full_key) or 0
|
||||
|
||||
# Convert to ToolError for consistent error format
|
||||
raise ToolError(f"Error in {tool_name}: {str(e)}") from e
|
||||
# Rate limit info
|
||||
rate_limit_info = {
|
||||
"limit": limit,
|
||||
"remaining": max(0, limit - request_count),
|
||||
"reset_time": int(current_time + window),
|
||||
"window_seconds": window,
|
||||
}
|
||||
|
||||
if request_count >= limit:
|
||||
return True, rate_limit_info
|
||||
|
||||
# Increment counter with TTL
|
||||
new_count = (request_count or 0) + 1
|
||||
self._cache.set("%s:count" % full_key, new_count, timeout=window)
|
||||
|
||||
return False, rate_limit_info
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("Redis rate limiter error: %s, allowing request", e)
|
||||
# On Redis error, allow the request
|
||||
return False, {
|
||||
"limit": limit,
|
||||
"remaining": limit,
|
||||
"reset_time": 0,
|
||||
"window_seconds": window,
|
||||
}
|
||||
|
||||
def cleanup(self) -> None:
|
||||
"""No cleanup needed for Redis - TTL handles expiration."""
|
||||
pass
|
||||
|
||||
|
||||
def create_rate_limiter() -> RateLimiterProtocol:
|
||||
"""Factory to create appropriate rate limiter based on environment."""
|
||||
try:
|
||||
# Try to use Redis first (production)
|
||||
from superset.extensions import cache_manager
|
||||
|
||||
if cache_manager and cache_manager.cache:
|
||||
# Test Redis connectivity
|
||||
test_key = "mcp:ratelimit:test"
|
||||
cache_manager.cache.set(test_key, 1, timeout=1)
|
||||
if cache_manager.cache.get(test_key):
|
||||
cache_manager.cache.delete(test_key)
|
||||
logger.info("Using Redis for rate limiting")
|
||||
return RedisRateLimiter()
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Redis not available for rate limiting: %s, falling back to in-memory", e
|
||||
)
|
||||
|
||||
# Fallback to in-memory rate limiter (development)
|
||||
logger.info("Using in-memory rate limiter")
|
||||
return InMemoryRateLimiter()
|
||||
|
||||
|
||||
class RateLimitMiddleware(Middleware):
|
||||
"""
|
||||
Rate limiting middleware to prevent abuse of MCP tools.
|
||||
|
||||
Implements sliding window rate limiting with separate limits for:
|
||||
- Per-user limits (if authenticated)
|
||||
- Per-IP limits (for unauthenticated requests)
|
||||
- Per-tool limits (for expensive operations)
|
||||
|
||||
Configuration:
|
||||
- default_requests_per_minute: Default rate limit (60 requests/minute)
|
||||
- per_user_requests_per_minute: Rate limit per authenticated user (120/min)
|
||||
- expensive_tool_requests_per_minute: Rate limit for expensive tools (10/min)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
default_requests_per_minute: int = 60,
|
||||
per_user_requests_per_minute: int = 120,
|
||||
expensive_tool_requests_per_minute: int = 10,
|
||||
expensive_tools: list[str] | None = None,
|
||||
) -> None:
|
||||
self.default_rpm = default_requests_per_minute
|
||||
self.user_rpm = per_user_requests_per_minute
|
||||
self.expensive_rpm = expensive_tool_requests_per_minute
|
||||
self.expensive_tools = set(
|
||||
expensive_tools
|
||||
or [
|
||||
"get_chart_preview",
|
||||
"generate_chart",
|
||||
"generate_dashboard",
|
||||
"get_chart_data",
|
||||
]
|
||||
)
|
||||
|
||||
# Use hybrid rate limiter (Redis in production, in-memory in development)
|
||||
self._rate_limiter = create_rate_limiter()
|
||||
|
||||
def _get_rate_limit_key(self, context: MiddlewareContext) -> tuple[str, int]:
|
||||
"""
|
||||
Generate rate limit key and determine applicable limit.
|
||||
|
||||
Returns:
|
||||
Tuple of (key, requests_per_minute_limit)
|
||||
"""
|
||||
tool_name = getattr(context.message, "name", "unknown")
|
||||
|
||||
# Get user context
|
||||
user_id = None
|
||||
try:
|
||||
user_id = get_user_id()
|
||||
except Exception:
|
||||
user_id = None # User not authenticated
|
||||
|
||||
# Determine rate limit
|
||||
if tool_name in self.expensive_tools:
|
||||
limit = self.expensive_rpm
|
||||
key_prefix = "expensive"
|
||||
elif user_id:
|
||||
limit = self.user_rpm
|
||||
key_prefix = "user"
|
||||
else:
|
||||
limit = self.default_rpm
|
||||
key_prefix = "default"
|
||||
|
||||
# Generate key
|
||||
if user_id:
|
||||
key = f"{key_prefix}:user:{user_id}:{tool_name}"
|
||||
else:
|
||||
# Use agent_id or session info as fallback
|
||||
agent_id = None
|
||||
if hasattr(context, "metadata") and context.metadata:
|
||||
agent_id = context.metadata.get("agent_id")
|
||||
if not agent_id and hasattr(context, "session") and context.session:
|
||||
agent_id = getattr(context.session, "agent_id", None)
|
||||
|
||||
if agent_id:
|
||||
key = f"{key_prefix}:agent:{agent_id}:{tool_name}"
|
||||
else:
|
||||
key = f"{key_prefix}:anonymous:{tool_name}"
|
||||
|
||||
return key, limit
|
||||
|
||||
async def on_call_tool(
|
||||
self,
|
||||
context: MiddlewareContext,
|
||||
call_next: Callable[[MiddlewareContext], Awaitable[Any]],
|
||||
) -> Any:
|
||||
"""Check rate limits before allowing tool calls."""
|
||||
# Clean up old entries periodically (only needed for in-memory)
|
||||
self._rate_limiter.cleanup()
|
||||
|
||||
# Get rate limit key and limit
|
||||
key, limit = self._get_rate_limit_key(context)
|
||||
|
||||
# Check if rate limited
|
||||
is_limited, rate_info = self._rate_limiter.is_rate_limited(key, limit)
|
||||
|
||||
if is_limited:
|
||||
tool_name = getattr(context.message, "name", "unknown")
|
||||
|
||||
# Log rate limit event
|
||||
try:
|
||||
user_id = get_user_id() if hasattr(context, "session") else None
|
||||
event_logger.log(
|
||||
user_id=user_id,
|
||||
action="mcp_rate_limit_exceeded",
|
||||
curated_payload={
|
||||
"tool": tool_name,
|
||||
"rate_limit_key": key,
|
||||
"limit": limit,
|
||||
"window_seconds": 60,
|
||||
},
|
||||
)
|
||||
except Exception as log_error:
|
||||
logger.warning("Failed to log rate limit event: %s", log_error)
|
||||
|
||||
logger.warning(
|
||||
"Rate limit exceeded for %s: key=%s, limit=%s/min, reset_in=%ss",
|
||||
tool_name,
|
||||
key,
|
||||
limit,
|
||||
rate_info["reset_time"] - int(time.time()),
|
||||
)
|
||||
|
||||
raise ToolError(
|
||||
"Rate limit exceeded for %s. "
|
||||
"Limit: %s requests per minute. "
|
||||
"Try again in %s seconds."
|
||||
% (tool_name, limit, rate_info["reset_time"] - int(time.time()))
|
||||
)
|
||||
|
||||
# Log rate limit info for monitoring
|
||||
logger.debug(
|
||||
"Rate limit check: %s: key=%s, remaining=%s/%s",
|
||||
getattr(context.message, "name", "unknown"),
|
||||
key,
|
||||
rate_info["remaining"],
|
||||
limit,
|
||||
)
|
||||
|
||||
return await call_next(context)
|
||||
|
||||
|
||||
class FieldPermissionsMiddleware(Middleware):
|
||||
"""
|
||||
Middleware that applies field-level permissions to filter sensitive data
|
||||
from MCP tool responses based on user permissions.
|
||||
"""
|
||||
|
||||
# Map tool names to object types for permission filtering
|
||||
TOOL_OBJECT_TYPE_MAP = {
|
||||
"list_datasets": "dataset",
|
||||
"get_dataset_info": "dataset",
|
||||
"list_charts": "chart",
|
||||
"get_chart_info": "chart",
|
||||
"get_chart_data": "chart",
|
||||
"get_chart_preview": "chart",
|
||||
"update_chart": "chart",
|
||||
"generate_chart": "chart",
|
||||
"list_dashboards": "dashboard",
|
||||
"get_dashboard_info": "dashboard",
|
||||
"generate_dashboard": "dashboard",
|
||||
"add_chart_to_existing_dashboard": "dashboard",
|
||||
}
|
||||
|
||||
async def on_call_tool(
|
||||
self,
|
||||
context: MiddlewareContext,
|
||||
call_next: Callable[[MiddlewareContext], Awaitable[Any]],
|
||||
) -> Any:
|
||||
"""Apply field-level permissions to tool responses."""
|
||||
# Get the tool response first
|
||||
response = await call_next(context)
|
||||
|
||||
# Get tool name
|
||||
tool_name = getattr(context.message, "name", "unknown")
|
||||
|
||||
# Check if this tool needs field-level filtering
|
||||
object_type = self.TOOL_OBJECT_TYPE_MAP.get(tool_name)
|
||||
if not object_type:
|
||||
# No filtering needed
|
||||
return response
|
||||
|
||||
# Get current user for permissions
|
||||
try:
|
||||
user = self._get_current_user()
|
||||
except Exception as e:
|
||||
logger.warning("Could not get current user for field filtering: %s", e)
|
||||
user = None
|
||||
|
||||
# Apply field-level permissions to the response
|
||||
try:
|
||||
filtered_response = self._filter_response(response, object_type, user)
|
||||
|
||||
# Log field filtering activity for monitoring
|
||||
logger.debug(
|
||||
"Applied field-level permissions for %s (object_type=%s, user=%s)",
|
||||
tool_name,
|
||||
object_type,
|
||||
getattr(user, "username", "anonymous"),
|
||||
)
|
||||
|
||||
return filtered_response
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error applying field permissions to %s: %s", tool_name, e)
|
||||
# Return original response if filtering fails
|
||||
return response
|
||||
|
||||
def _get_current_user(self) -> Any:
|
||||
"""Get the current authenticated user."""
|
||||
try:
|
||||
from flask import g
|
||||
|
||||
return getattr(g, "user", None)
|
||||
except Exception:
|
||||
# Try to get user from core utils
|
||||
try:
|
||||
user_id = get_user_id()
|
||||
if user_id:
|
||||
from flask_appbuilder.security.sqla.models import User
|
||||
|
||||
from superset.extensions import db
|
||||
|
||||
return db.session.query(User).filter_by(id=user_id).first()
|
||||
except Exception as e:
|
||||
logger.debug("Could not get user from session: %s", e)
|
||||
return None
|
||||
|
||||
def _filter_response(self, response: Any, object_type: str, user: Any) -> Any:
|
||||
"""
|
||||
Filter response data based on object type and user permissions.
|
||||
|
||||
Args:
|
||||
response: The response object to filter
|
||||
object_type: Type of object ('dataset', 'chart', 'dashboard')
|
||||
user: User object for permission checking
|
||||
|
||||
Returns:
|
||||
Filtered response
|
||||
"""
|
||||
from superset.mcp_service.utils.permissions_utils import filter_sensitive_data
|
||||
|
||||
if not response:
|
||||
return response
|
||||
|
||||
# Handle different response types
|
||||
if hasattr(response, "model_dump"):
|
||||
# Pydantic model - convert to dict, filter, and return dict
|
||||
response_dict = response.model_dump()
|
||||
return filter_sensitive_data(response_dict, object_type, user)
|
||||
elif isinstance(response, dict):
|
||||
# Dictionary response - filter directly
|
||||
return filter_sensitive_data(response, object_type, user)
|
||||
elif isinstance(response, list):
|
||||
# List response - filter each item
|
||||
return [filter_sensitive_data(item, object_type, user) for item in response]
|
||||
else:
|
||||
# Unknown response type, return as-is
|
||||
logger.debug(
|
||||
"Unknown response type for field filtering: %s", type(response)
|
||||
)
|
||||
return response
|
||||
|
||||
35
superset/mcp_service/screenshot/__init__.py
Normal file
35
superset/mcp_service/screenshot/__init__.py
Normal file
@@ -0,0 +1,35 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Screenshot and WebDriver infrastructure for MCP service."""
|
||||
|
||||
from .pooled_screenshot import (
|
||||
PooledBaseScreenshot,
|
||||
PooledChartScreenshot,
|
||||
PooledDashboardScreenshot,
|
||||
PooledExploreScreenshot,
|
||||
)
|
||||
from .webdriver_pool import get_webdriver_pool, WebDriverPool
|
||||
|
||||
__all__ = [
|
||||
"PooledBaseScreenshot",
|
||||
"PooledChartScreenshot",
|
||||
"PooledDashboardScreenshot",
|
||||
"PooledExploreScreenshot",
|
||||
"WebDriverPool",
|
||||
"get_webdriver_pool",
|
||||
]
|
||||
483
superset/mcp_service/screenshot/pooled_screenshot.py
Normal file
483
superset/mcp_service/screenshot/pooled_screenshot.py
Normal file
@@ -0,0 +1,483 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Pooled screenshot implementations for improved performance
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from flask import current_app
|
||||
from flask_appbuilder.security.sqla.models import User
|
||||
from selenium.common.exceptions import TimeoutException, WebDriverException
|
||||
from selenium.webdriver.common.by import By
|
||||
from selenium.webdriver.support import expected_conditions
|
||||
from selenium.webdriver.support.ui import WebDriverWait
|
||||
|
||||
from superset.extensions import machine_auth_provider_factory
|
||||
from superset.mcp_service.screenshot.webdriver_pool import get_webdriver_pool
|
||||
from superset.mcp_service.utils.retry_utils import retry_screenshot_operation
|
||||
from superset.utils.screenshots import BaseScreenshot, WindowSize
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PooledBaseScreenshot(BaseScreenshot):
|
||||
"""
|
||||
Enhanced screenshot class that uses WebDriver pooling for better performance.
|
||||
|
||||
Key improvements:
|
||||
- Reuses WebDriver instances from a connection pool
|
||||
- Faster screenshot generation (no browser startup/shutdown per request)
|
||||
- Better resource management and automatic cleanup
|
||||
- Health checking and recovery of WebDriver instances
|
||||
"""
|
||||
|
||||
def get_screenshot(
|
||||
self, user: User, window_size: WindowSize | None = None
|
||||
) -> bytes | None:
|
||||
"""
|
||||
Generate screenshot using pooled WebDriver with retry logic for reliability.
|
||||
|
||||
Args:
|
||||
user: User context for authentication
|
||||
window_size: Optional window size override
|
||||
|
||||
Returns:
|
||||
Screenshot as PNG bytes or None if failed
|
||||
"""
|
||||
return retry_screenshot_operation(
|
||||
self._get_screenshot_internal, user, window_size
|
||||
)
|
||||
|
||||
def _get_screenshot_internal(
|
||||
self, user: User, window_size: WindowSize | None = None
|
||||
) -> bytes | None:
|
||||
"""
|
||||
Internal screenshot generation method with pooled WebDriver.
|
||||
This method is wrapped by retry logic in get_screenshot().
|
||||
"""
|
||||
window_size = window_size or self.window_size
|
||||
pool = get_webdriver_pool()
|
||||
|
||||
# Use pooled WebDriver
|
||||
with pool.get_driver(window_size, user.id) as driver:
|
||||
try:
|
||||
# Clear any existing cookies to ensure clean authentication
|
||||
try:
|
||||
driver.delete_all_cookies()
|
||||
logger.debug("Cleared all cookies from WebDriver")
|
||||
except Exception as e:
|
||||
logger.warning("Failed to clear cookies: %s", e)
|
||||
|
||||
# Authenticate the driver for this user
|
||||
user_name = user.username if user else "None"
|
||||
logger.debug("Authenticating WebDriver for user %s", user_name)
|
||||
machine_auth_provider_factory.instance.authenticate_webdriver(
|
||||
driver, user
|
||||
)
|
||||
|
||||
# Navigate to the URL
|
||||
logger.debug("Navigating to screenshot URL: %s", self.url)
|
||||
driver.get(self.url)
|
||||
|
||||
# Check if we were redirected to login (authentication failed)
|
||||
current_url = driver.current_url
|
||||
if "/login" in current_url:
|
||||
msg = (
|
||||
f"Authentication failed - WebDriver redirected to login page: "
|
||||
f"{current_url}"
|
||||
)
|
||||
raise Exception(msg)
|
||||
|
||||
# Take screenshot using the specific implementation
|
||||
return self._take_screenshot(driver, user)
|
||||
|
||||
except (TimeoutException, WebDriverException, OSError) as e:
|
||||
# These are retryable exceptions
|
||||
logger.warning("Retryable error taking screenshot: %s", e)
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error("Non-retryable error taking screenshot: %s", e)
|
||||
raise
|
||||
|
||||
def _take_screenshot(self, driver: Any, user: User) -> bytes | None:
|
||||
"""
|
||||
Override this method in subclasses to implement specific screenshot logic.
|
||||
|
||||
Args:
|
||||
driver: Authenticated WebDriver instance
|
||||
user: User context
|
||||
|
||||
Returns:
|
||||
Screenshot as PNG bytes
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement _take_screenshot")
|
||||
|
||||
|
||||
class PooledChartScreenshot(PooledBaseScreenshot):
|
||||
"""Pooled version of chart screenshot generation"""
|
||||
|
||||
thumbnail_type: str = "chart"
|
||||
element: str = "chart-container"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str,
|
||||
digest: str | None,
|
||||
window_size: WindowSize | None = None,
|
||||
thumb_size: WindowSize | None = None,
|
||||
):
|
||||
from superset.utils.urls import modify_url_query
|
||||
from superset.utils.webdriver import ChartStandaloneMode
|
||||
|
||||
# Chart reports are in standalone="true" mode
|
||||
url = modify_url_query(
|
||||
url,
|
||||
standalone=ChartStandaloneMode.HIDE_NAV.value,
|
||||
)
|
||||
super().__init__(url, digest)
|
||||
self.window_size = window_size or (800, 600)
|
||||
self.thumb_size = thumb_size or (400, 300)
|
||||
|
||||
def _take_screenshot(self, driver: Any, user: User) -> bytes | None:
|
||||
"""Take screenshot of chart with standard Superset chart handling"""
|
||||
try:
|
||||
# Wait for page to load
|
||||
selenium_headstart = current_app.config["SCREENSHOT_SELENIUM_HEADSTART"]
|
||||
logger.debug("Sleeping for %s seconds", selenium_headstart)
|
||||
time.sleep(selenium_headstart)
|
||||
|
||||
# Wait for chart container
|
||||
wait = WebDriverWait(driver, current_app.config["SCREENSHOT_LOCATE_WAIT"])
|
||||
element = wait.until(
|
||||
expected_conditions.presence_of_element_located(
|
||||
(By.CLASS_NAME, self.element)
|
||||
)
|
||||
)
|
||||
|
||||
# Wait for chart containers to render
|
||||
wait.until(
|
||||
expected_conditions.visibility_of_all_elements_located(
|
||||
(By.CLASS_NAME, "chart-container")
|
||||
)
|
||||
)
|
||||
|
||||
# Wait for loading to complete
|
||||
WebDriverWait(driver, current_app.config["SCREENSHOT_LOAD_WAIT"]).until_not(
|
||||
expected_conditions.presence_of_all_elements_located(
|
||||
(By.CLASS_NAME, "loading")
|
||||
)
|
||||
)
|
||||
|
||||
# Wait for animations
|
||||
animation_wait = current_app.config["SCREENSHOT_SELENIUM_ANIMATION_WAIT"]
|
||||
logger.debug("Wait %s seconds for chart animation", animation_wait)
|
||||
time.sleep(animation_wait)
|
||||
|
||||
# Handle unexpected errors if configured
|
||||
if current_app.config.get("SCREENSHOT_REPLACE_UNEXPECTED_ERRORS"):
|
||||
from superset.utils.webdriver import WebDriverSelenium
|
||||
|
||||
unexpected_errors = WebDriverSelenium.find_unexpected_errors(driver)
|
||||
if unexpected_errors:
|
||||
logger.warning(
|
||||
"%s errors found in screenshot. ", len(unexpected_errors)
|
||||
)
|
||||
|
||||
# Take screenshot
|
||||
logger.debug("Taking PNG screenshot as user %s", user.username)
|
||||
return element.screenshot_as_png
|
||||
|
||||
except TimeoutException:
|
||||
logger.exception("Timeout taking chart screenshot for URL: %s", self.url)
|
||||
raise
|
||||
except WebDriverException:
|
||||
logger.exception("WebDriver error taking screenshot for URL: %s", self.url)
|
||||
raise
|
||||
|
||||
|
||||
class PooledExploreScreenshot(PooledBaseScreenshot):
|
||||
"""
|
||||
Pooled version of explore screenshot with UI hiding functionality.
|
||||
|
||||
This class provides the same clean chart-only screenshots as the previous
|
||||
implementation but with improved performance through WebDriver pooling.
|
||||
"""
|
||||
|
||||
thumbnail_type: str = "explore"
|
||||
element: str = "chart-container"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str,
|
||||
digest: str | None,
|
||||
window_size: WindowSize | None = None,
|
||||
thumb_size: WindowSize | None = None,
|
||||
):
|
||||
super().__init__(url, digest)
|
||||
self.window_size = window_size or (1600, 1200)
|
||||
self.thumb_size = thumb_size or (800, 600)
|
||||
|
||||
def _take_screenshot(self, driver: Any, user: User) -> bytes | None:
|
||||
"""
|
||||
Take screenshot of explore page with UI elements hidden for clean chart display.
|
||||
|
||||
This implementation:
|
||||
1. Waits for chart to load
|
||||
2. Hides navigation, panels, and headers using JavaScript
|
||||
3. Takes screenshot of just the chart area
|
||||
4. Includes fallback strategies if chart container not found
|
||||
"""
|
||||
try:
|
||||
# Give explore page time to initialize
|
||||
logger.debug("Waiting for explore page to initialize")
|
||||
time.sleep(3)
|
||||
|
||||
# Wait for chart container to load and be visible
|
||||
wait = WebDriverWait(driver, 45)
|
||||
chart_element = wait.until(
|
||||
expected_conditions.visibility_of_element_located(
|
||||
(By.CLASS_NAME, "chart-container")
|
||||
)
|
||||
)
|
||||
|
||||
# Wait for any loading indicators to disappear
|
||||
WebDriverWait(driver, 30).until_not(
|
||||
expected_conditions.presence_of_all_elements_located(
|
||||
(By.CLASS_NAME, "loading")
|
||||
)
|
||||
)
|
||||
|
||||
# Hide overlapping UI elements using JavaScript
|
||||
self._hide_ui_elements(driver)
|
||||
|
||||
# Wait for UI hiding animations to complete
|
||||
time.sleep(2)
|
||||
|
||||
# Take screenshot of just the chart area
|
||||
img = chart_element.screenshot_as_png
|
||||
logger.info("Successfully captured chart-container screenshot")
|
||||
return img
|
||||
|
||||
except TimeoutException as e:
|
||||
logger.warning("Chart container not found, trying fallbacks: %s", e)
|
||||
return self._fallback_screenshot(driver)
|
||||
except Exception as e:
|
||||
logger.error("Error taking explore screenshot: %s", e)
|
||||
raise
|
||||
|
||||
def _hide_ui_elements(self, driver: Any) -> None:
|
||||
"""Hide UI elements to show only the chart"""
|
||||
hide_script = """
|
||||
// Hide the specific left panels by their exact class names
|
||||
var dataSourcePanel = document.querySelector(
|
||||
'.explore-column.data-source-selection'
|
||||
);
|
||||
if (dataSourcePanel) {
|
||||
dataSourcePanel.style.display = 'none';
|
||||
console.log('Hidden data source panel');
|
||||
}
|
||||
|
||||
var controlsPanel = document.querySelector(
|
||||
'.col-sm-3.explore-column.controls-column'
|
||||
);
|
||||
if (controlsPanel) {
|
||||
controlsPanel.style.display = 'none';
|
||||
console.log('Hidden controls panel');
|
||||
}
|
||||
|
||||
// Also try alternate selectors for the controls
|
||||
var controlsAlt = document.querySelector('.explore-column.controls-column');
|
||||
if (controlsAlt) {
|
||||
controlsAlt.style.display = 'none';
|
||||
console.log('Hidden controls panel (alt)');
|
||||
}
|
||||
|
||||
// Hide the main navigation header
|
||||
var mainHeader = document.querySelector('header.top#main-menu');
|
||||
if (mainHeader) {
|
||||
mainHeader.style.display = 'none';
|
||||
console.log('Hidden main navigation header');
|
||||
}
|
||||
|
||||
// Hide the chart header with actions (title input and save button)
|
||||
var chartHeader = document.querySelector('.header-with-actions');
|
||||
if (chartHeader) {
|
||||
chartHeader.style.display = 'none';
|
||||
console.log('Hidden chart header with actions');
|
||||
}
|
||||
|
||||
// Make sure the main chart area expands to use available space
|
||||
var mainContent = document.querySelector('.main-explore-content');
|
||||
if (mainContent) {
|
||||
mainContent.style.width = '100%';
|
||||
// Change from col-sm-7 to col-sm-12
|
||||
mainContent.className = 'main-explore-content col-sm-12';
|
||||
}
|
||||
|
||||
// Ensure chart container is visible
|
||||
var chartContainer = document.querySelector('.chart-container');
|
||||
if (chartContainer) {
|
||||
chartContainer.style.position = 'relative';
|
||||
chartContainer.style.zIndex = '999';
|
||||
chartContainer.style.width = '100%';
|
||||
}
|
||||
"""
|
||||
|
||||
try:
|
||||
driver.execute_script(hide_script)
|
||||
logger.debug("Successfully executed UI hiding script")
|
||||
except WebDriverException as e:
|
||||
logger.warning("Failed to execute UI hiding script: %s", e)
|
||||
|
||||
def _fallback_screenshot(self, driver: Any) -> bytes | None:
|
||||
"""Fallback screenshot strategies if chart container not found"""
|
||||
try:
|
||||
# Fallback 1: Try slice container
|
||||
logger.debug("Trying slice_container fallback")
|
||||
slice_element = WebDriverWait(driver, 10).until(
|
||||
expected_conditions.visibility_of_element_located(
|
||||
(By.CLASS_NAME, "slice_container")
|
||||
)
|
||||
)
|
||||
img = slice_element.screenshot_as_png
|
||||
logger.info("Successfully captured slice_container screenshot")
|
||||
return img
|
||||
|
||||
except TimeoutException:
|
||||
try:
|
||||
# Fallback 2: Try any chart-related container
|
||||
logger.debug("Trying chart-related container fallback")
|
||||
chart_containers = driver.find_elements(
|
||||
By.CSS_SELECTOR,
|
||||
".chart, .slice-container, .chart-content, "
|
||||
"[data-test='chart-container']",
|
||||
)
|
||||
if chart_containers:
|
||||
img = chart_containers[0].screenshot_as_png
|
||||
logger.info(
|
||||
"Successfully captured chart-related container screenshot"
|
||||
)
|
||||
return img
|
||||
|
||||
except Exception as e2:
|
||||
logger.warning("Chart-related container fallback failed: %s", e2)
|
||||
|
||||
# Fallback 3: Full page screenshot
|
||||
logger.warning("Using full page screenshot as final fallback")
|
||||
img = driver.get_screenshot_as_png()
|
||||
logger.info("Using full page screenshot as fallback")
|
||||
return img
|
||||
|
||||
|
||||
class PooledDashboardScreenshot(PooledBaseScreenshot):
|
||||
"""Pooled version of dashboard screenshot generation"""
|
||||
|
||||
thumbnail_type: str = "dashboard"
|
||||
element: str = "standalone"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str,
|
||||
digest: str | None,
|
||||
window_size: WindowSize | None = None,
|
||||
thumb_size: WindowSize | None = None,
|
||||
):
|
||||
from superset.utils.urls import modify_url_query
|
||||
from superset.utils.webdriver import DashboardStandaloneMode
|
||||
|
||||
# Dashboard screenshots should always capture in standalone
|
||||
url = modify_url_query(
|
||||
url,
|
||||
standalone=DashboardStandaloneMode.REPORT.value,
|
||||
)
|
||||
super().__init__(url, digest)
|
||||
self.window_size = window_size or (1600, 1200)
|
||||
self.thumb_size = thumb_size or (800, 600)
|
||||
|
||||
def _take_screenshot(self, driver: Any, user: User) -> bytes | None:
|
||||
"""Take screenshot of dashboard with standard Superset dashboard handling"""
|
||||
try:
|
||||
# Wait for page to load
|
||||
selenium_headstart = current_app.config["SCREENSHOT_SELENIUM_HEADSTART"]
|
||||
time.sleep(selenium_headstart)
|
||||
|
||||
# Wait for dashboard element
|
||||
wait = WebDriverWait(driver, current_app.config["SCREENSHOT_LOCATE_WAIT"])
|
||||
element = wait.until(
|
||||
expected_conditions.presence_of_element_located(
|
||||
(By.CLASS_NAME, self.element)
|
||||
)
|
||||
)
|
||||
|
||||
# Wait for chart containers to render
|
||||
try:
|
||||
wait.until(
|
||||
expected_conditions.visibility_of_all_elements_located(
|
||||
(By.CLASS_NAME, "chart-container")
|
||||
)
|
||||
)
|
||||
except TimeoutException:
|
||||
# Fallback for empty dashboards
|
||||
try:
|
||||
WebDriverWait(driver, 0).until(
|
||||
expected_conditions.visibility_of_all_elements_located(
|
||||
(By.CLASS_NAME, "grid-container")
|
||||
)
|
||||
)
|
||||
except TimeoutException:
|
||||
logger.exception("Dashboard failed to load at URL: %s", self.url)
|
||||
raise
|
||||
|
||||
# Wait for loading to complete
|
||||
WebDriverWait(driver, current_app.config["SCREENSHOT_LOAD_WAIT"]).until_not(
|
||||
expected_conditions.presence_of_all_elements_located(
|
||||
(By.CLASS_NAME, "loading")
|
||||
)
|
||||
)
|
||||
|
||||
# Wait for animations
|
||||
animation_wait = current_app.config["SCREENSHOT_SELENIUM_ANIMATION_WAIT"]
|
||||
time.sleep(animation_wait)
|
||||
|
||||
# Handle unexpected errors if configured
|
||||
if current_app.config.get("SCREENSHOT_REPLACE_UNEXPECTED_ERRORS"):
|
||||
from superset.utils.webdriver import WebDriverSelenium
|
||||
|
||||
unexpected_errors = WebDriverSelenium.find_unexpected_errors(driver)
|
||||
if unexpected_errors:
|
||||
logger.warning(
|
||||
"%s errors found in dashboard ", len(unexpected_errors)
|
||||
)
|
||||
|
||||
# Take screenshot
|
||||
logger.debug("Taking PNG dashboard screenshot as user %s", user.username)
|
||||
return element.screenshot_as_png
|
||||
|
||||
except TimeoutException:
|
||||
logger.exception(
|
||||
"Timeout taking dashboard screenshot for URL: %s", self.url
|
||||
)
|
||||
raise
|
||||
except WebDriverException:
|
||||
logger.exception(
|
||||
"WebDriver error taking dashboard screenshot for URL: %s", self.url
|
||||
)
|
||||
raise
|
||||
139
superset/mcp_service/screenshot/webdriver_config.py
Normal file
139
superset/mcp_service/screenshot/webdriver_config.py
Normal file
@@ -0,0 +1,139 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
WebDriver pool configuration defaults for Superset MCP service
|
||||
"""
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
# Default WebDriver pool configuration
|
||||
DEFAULT_WEBDRIVER_POOL_CONFIG = {
|
||||
# Maximum number of WebDriver instances to keep in the pool
|
||||
"MAX_POOL_SIZE": 5,
|
||||
# Maximum age of a WebDriver instance (in seconds)
|
||||
# After this time, the driver will be destroyed and recreated
|
||||
"MAX_AGE_SECONDS": 3600, # 1 hour
|
||||
# Maximum number of times a WebDriver can be reused
|
||||
# After this many uses, the driver will be destroyed and recreated
|
||||
"MAX_USAGE_COUNT": 50,
|
||||
# How long a WebDriver can sit idle before being destroyed (in seconds)
|
||||
"IDLE_TIMEOUT_SECONDS": 300, # 5 minutes
|
||||
# How often to perform health checks on WebDriver instances (in seconds)
|
||||
"HEALTH_CHECK_INTERVAL": 60, # 1 minute
|
||||
}
|
||||
|
||||
|
||||
def configure_webdriver_pool(app_config: Dict[str, Any]) -> None:
|
||||
"""
|
||||
Configure WebDriver pool settings in Superset app config.
|
||||
|
||||
This function adds WebDriver pool configuration to the Superset app config
|
||||
if it doesn't already exist, using sensible defaults.
|
||||
|
||||
Args:
|
||||
app_config: The Superset application configuration dictionary
|
||||
"""
|
||||
if "WEBDRIVER_POOL" not in app_config:
|
||||
app_config["WEBDRIVER_POOL"] = DEFAULT_WEBDRIVER_POOL_CONFIG.copy()
|
||||
else:
|
||||
# Merge with defaults for any missing keys
|
||||
for key, default_value in DEFAULT_WEBDRIVER_POOL_CONFIG.items():
|
||||
if key not in app_config["WEBDRIVER_POOL"]:
|
||||
app_config["WEBDRIVER_POOL"][key] = default_value
|
||||
|
||||
|
||||
def get_pool_stats_endpoint() -> Any:
|
||||
"""
|
||||
Create a Flask endpoint to view WebDriver pool statistics.
|
||||
|
||||
This function can be called to register a debugging endpoint
|
||||
that shows the current state of the WebDriver pool.
|
||||
|
||||
Returns:
|
||||
Flask route function for pool statistics
|
||||
"""
|
||||
|
||||
def pool_stats() -> Any:
|
||||
try:
|
||||
from flask import jsonify
|
||||
|
||||
from superset.mcp_service.screenshot.webdriver_pool import (
|
||||
get_webdriver_pool,
|
||||
)
|
||||
|
||||
pool = get_webdriver_pool()
|
||||
stats = pool.get_stats()
|
||||
|
||||
return jsonify({"webdriver_pool": stats, "status": "healthy"})
|
||||
except Exception as e:
|
||||
from flask import jsonify
|
||||
|
||||
return jsonify({"error": str(e), "status": "error"}), 500
|
||||
|
||||
return pool_stats
|
||||
|
||||
|
||||
# Performance tuning recommendations based on use case
|
||||
PERFORMANCE_CONFIGS = {
|
||||
"low_traffic": {
|
||||
"MAX_POOL_SIZE": 2,
|
||||
"MAX_AGE_SECONDS": 1800, # 30 minutes
|
||||
"MAX_USAGE_COUNT": 25,
|
||||
"IDLE_TIMEOUT_SECONDS": 180, # 3 minutes
|
||||
"HEALTH_CHECK_INTERVAL": 120, # 2 minutes
|
||||
},
|
||||
"medium_traffic": {
|
||||
"MAX_POOL_SIZE": 5,
|
||||
"MAX_AGE_SECONDS": 3600, # 1 hour
|
||||
"MAX_USAGE_COUNT": 50,
|
||||
"IDLE_TIMEOUT_SECONDS": 300, # 5 minutes
|
||||
"HEALTH_CHECK_INTERVAL": 60, # 1 minute
|
||||
},
|
||||
"high_traffic": {
|
||||
"MAX_POOL_SIZE": 10,
|
||||
"MAX_AGE_SECONDS": 7200, # 2 hours
|
||||
"MAX_USAGE_COUNT": 100,
|
||||
"IDLE_TIMEOUT_SECONDS": 600, # 10 minutes
|
||||
"HEALTH_CHECK_INTERVAL": 30, # 30 seconds
|
||||
},
|
||||
"development": {
|
||||
"MAX_POOL_SIZE": 2,
|
||||
"MAX_AGE_SECONDS": 900, # 15 minutes
|
||||
"MAX_USAGE_COUNT": 10,
|
||||
"IDLE_TIMEOUT_SECONDS": 120, # 2 minutes
|
||||
"HEALTH_CHECK_INTERVAL": 30, # 30 seconds
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def configure_for_environment(
|
||||
app_config: Dict[str, Any], environment: str = "medium_traffic"
|
||||
) -> None:
|
||||
"""
|
||||
Configure WebDriver pool for specific environment/traffic levels.
|
||||
|
||||
Args:
|
||||
app_config: The Superset application configuration dictionary
|
||||
environment: Environment type (low_traffic, medium_traffic,
|
||||
high_traffic, development)
|
||||
"""
|
||||
if environment in PERFORMANCE_CONFIGS:
|
||||
app_config["WEBDRIVER_POOL"] = PERFORMANCE_CONFIGS[environment].copy()
|
||||
else:
|
||||
# Fallback to default
|
||||
configure_webdriver_pool(app_config)
|
||||
433
superset/mcp_service/screenshot/webdriver_pool.py
Normal file
433
superset/mcp_service/screenshot/webdriver_pool.py
Normal file
@@ -0,0 +1,433 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
WebDriver connection pooling for improved screenshot performance
|
||||
"""
|
||||
|
||||
import logging
|
||||
import signal
|
||||
import threading
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from queue import Empty, Full, Queue
|
||||
from typing import Any, Dict, Generator
|
||||
|
||||
from flask import current_app
|
||||
from selenium.common.exceptions import WebDriverException
|
||||
from selenium.webdriver.remote.webdriver import WebDriver
|
||||
|
||||
from superset.utils.webdriver import WebDriverSelenium, WindowSize
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WebDriverCreationError(Exception):
|
||||
"""Exception raised when WebDriver creation times out"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def _timeout_handler(signum: int, frame: Any) -> None:
|
||||
"""Signal handler for WebDriver creation timeout"""
|
||||
raise WebDriverCreationError("WebDriver creation timed out")
|
||||
|
||||
|
||||
@dataclass
|
||||
class PooledWebDriver:
|
||||
"""Wrapper for pooled WebDriver instance with metadata"""
|
||||
|
||||
driver: WebDriver
|
||||
created_at: float
|
||||
last_used: float
|
||||
window_size: WindowSize
|
||||
user_id: int | None = None
|
||||
is_healthy: bool = True
|
||||
usage_count: int = 0
|
||||
|
||||
|
||||
class WebDriverPool:
|
||||
"""
|
||||
Connection pool for WebDriver instances to improve screenshot performance.
|
||||
|
||||
Features:
|
||||
- Reuses WebDriver instances across requests
|
||||
- Automatic health checking and recovery
|
||||
- TTL-based expiration to prevent memory leaks
|
||||
- Thread-safe operations
|
||||
- Per-user driver isolation for security
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_pool_size: int = 5,
|
||||
max_age_seconds: int = 3600, # 1 hour
|
||||
max_usage_count: int = 50, # Recreate after 50 uses
|
||||
idle_timeout_seconds: int = 300, # 5 minutes
|
||||
health_check_interval: int = 60, # 1 minute
|
||||
creation_timeout_seconds: int = 30, # SECURITY FIX: Timeout for driver creation
|
||||
):
|
||||
self.max_pool_size = max_pool_size
|
||||
self.max_age_seconds = max_age_seconds
|
||||
self.max_usage_count = max_usage_count
|
||||
self.idle_timeout_seconds = idle_timeout_seconds
|
||||
self.health_check_interval = health_check_interval
|
||||
self.creation_timeout_seconds = creation_timeout_seconds
|
||||
|
||||
# Thread-safe pool management
|
||||
self._pool: Queue[PooledWebDriver] = Queue(maxsize=max_pool_size)
|
||||
self._active_drivers: Dict[int, PooledWebDriver] = {}
|
||||
self._lock = threading.RLock()
|
||||
self._last_health_check = time.time()
|
||||
|
||||
# Pool statistics
|
||||
self._stats = {
|
||||
"created": 0,
|
||||
"destroyed": 0,
|
||||
"borrowed": 0,
|
||||
"returned": 0,
|
||||
"health_check_failures": 0,
|
||||
"evictions": 0,
|
||||
}
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get pool statistics for monitoring"""
|
||||
with self._lock:
|
||||
return {
|
||||
**self._stats,
|
||||
"pool_size": self._pool.qsize(),
|
||||
"active_count": len(self._active_drivers),
|
||||
"max_pool_size": self.max_pool_size,
|
||||
}
|
||||
|
||||
def _create_driver(
|
||||
self, window_size: WindowSize, user_id: int | None = None
|
||||
) -> PooledWebDriver:
|
||||
"""Create a new WebDriver instance with timeout protection"""
|
||||
driver = None
|
||||
old_handler = None
|
||||
|
||||
try:
|
||||
# SECURITY FIX: Set up timeout protection for driver creation
|
||||
old_handler = signal.signal(signal.SIGALRM, _timeout_handler)
|
||||
signal.alarm(self.creation_timeout_seconds)
|
||||
|
||||
driver_type = current_app.config.get("WEBDRIVER_TYPE", "firefox")
|
||||
selenium_driver = WebDriverSelenium(driver_type, window_size)
|
||||
|
||||
# Create the actual WebDriver with timeout protection
|
||||
driver = selenium_driver.create()
|
||||
driver.set_window_size(*window_size)
|
||||
|
||||
# Clear the alarm - creation successful
|
||||
signal.alarm(0)
|
||||
|
||||
pooled_driver = PooledWebDriver(
|
||||
driver=driver,
|
||||
created_at=time.time(),
|
||||
last_used=time.time(),
|
||||
window_size=window_size,
|
||||
user_id=user_id,
|
||||
is_healthy=True,
|
||||
usage_count=0,
|
||||
)
|
||||
|
||||
self._stats["created"] += 1
|
||||
logger.debug(
|
||||
"Created new WebDriver instance for window size %s", window_size
|
||||
)
|
||||
return pooled_driver
|
||||
|
||||
except WebDriverCreationError:
|
||||
logger.error(
|
||||
"WebDriver creation timed out after %s seconds",
|
||||
self.creation_timeout_seconds,
|
||||
)
|
||||
if driver:
|
||||
try:
|
||||
driver.quit()
|
||||
except Exception:
|
||||
logger.debug("Failed to cleanup driver during timeout")
|
||||
raise Exception("WebDriver creation timed out") from None
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to create WebDriver: %s", e)
|
||||
if driver:
|
||||
try:
|
||||
driver.quit()
|
||||
except Exception:
|
||||
logger.debug("Failed to cleanup driver during error")
|
||||
raise
|
||||
|
||||
finally:
|
||||
# Restore original signal handler and clear alarm
|
||||
signal.alarm(0)
|
||||
if old_handler is not None:
|
||||
signal.signal(signal.SIGALRM, old_handler)
|
||||
|
||||
def _is_driver_valid(self, pooled_driver: PooledWebDriver) -> bool:
|
||||
"""Check if a pooled driver is still valid for use"""
|
||||
now = time.time()
|
||||
|
||||
# Check age limit
|
||||
if now - pooled_driver.created_at > self.max_age_seconds:
|
||||
logger.debug("Driver expired due to age")
|
||||
return False
|
||||
|
||||
# Check usage count limit
|
||||
if pooled_driver.usage_count >= self.max_usage_count:
|
||||
logger.debug("Driver expired due to usage count")
|
||||
return False
|
||||
|
||||
# Check idle timeout
|
||||
if now - pooled_driver.last_used > self.idle_timeout_seconds:
|
||||
logger.debug("Driver expired due to idle timeout")
|
||||
return False
|
||||
|
||||
# Check if driver is healthy
|
||||
if not pooled_driver.is_healthy:
|
||||
logger.debug("Driver marked as unhealthy")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def _health_check_driver(self, pooled_driver: PooledWebDriver) -> bool:
|
||||
"""Perform health check on a WebDriver instance"""
|
||||
try:
|
||||
# Simple health check - try to get current URL
|
||||
# This will fail if the driver is dead/hung
|
||||
_ = pooled_driver.driver.current_url
|
||||
pooled_driver.is_healthy = True
|
||||
return True
|
||||
except WebDriverException:
|
||||
pooled_driver.is_healthy = False
|
||||
self._stats["health_check_failures"] += 1
|
||||
logger.warning("WebDriver failed health check")
|
||||
return False
|
||||
except Exception as e:
|
||||
pooled_driver.is_healthy = False
|
||||
self._stats["health_check_failures"] += 1
|
||||
logger.warning("WebDriver health check error: %s", e)
|
||||
return False
|
||||
|
||||
def _destroy_driver(self, pooled_driver: PooledWebDriver) -> None:
|
||||
"""Safely destroy a WebDriver instance"""
|
||||
try:
|
||||
WebDriverSelenium.destroy(pooled_driver.driver)
|
||||
self._stats["destroyed"] += 1
|
||||
logger.debug("Destroyed WebDriver instance")
|
||||
except Exception as e:
|
||||
logger.warning("Error destroying WebDriver: %s", e)
|
||||
|
||||
def _cleanup_expired_drivers(self) -> None:
|
||||
"""Remove expired drivers from the pool"""
|
||||
expired_drivers = []
|
||||
|
||||
# Check pool for expired drivers
|
||||
while not self._pool.empty():
|
||||
try:
|
||||
pooled_driver = self._pool.get_nowait()
|
||||
if self._is_driver_valid(pooled_driver):
|
||||
# Driver is still valid, put it back
|
||||
self._pool.put_nowait(pooled_driver)
|
||||
break
|
||||
else:
|
||||
# Driver is expired
|
||||
expired_drivers.append(pooled_driver)
|
||||
self._stats["evictions"] += 1
|
||||
except Empty:
|
||||
break
|
||||
except Full:
|
||||
# Pool is full, stop checking
|
||||
break
|
||||
|
||||
# Destroy expired drivers
|
||||
for pooled_driver in expired_drivers:
|
||||
self._destroy_driver(pooled_driver)
|
||||
|
||||
def _periodic_health_check(self) -> None:
|
||||
"""Perform periodic health checks if needed"""
|
||||
now = time.time()
|
||||
if now - self._last_health_check < self.health_check_interval:
|
||||
return
|
||||
|
||||
self._last_health_check = now
|
||||
logger.debug("Performing periodic WebDriver pool health check")
|
||||
|
||||
# Cleanup expired drivers
|
||||
self._cleanup_expired_drivers()
|
||||
|
||||
# Health check active drivers
|
||||
unhealthy_drivers = []
|
||||
for driver_id, pooled_driver in self._active_drivers.items():
|
||||
if not self._health_check_driver(pooled_driver):
|
||||
unhealthy_drivers.append(driver_id)
|
||||
|
||||
# Remove unhealthy active drivers
|
||||
for driver_id in unhealthy_drivers:
|
||||
pooled_driver = self._active_drivers.pop(driver_id)
|
||||
if pooled_driver:
|
||||
self._destroy_driver(pooled_driver)
|
||||
|
||||
@contextmanager
|
||||
def get_driver(
|
||||
self, window_size: WindowSize, user_id: int | None = None
|
||||
) -> Generator[WebDriver, None, None]:
|
||||
"""
|
||||
Context manager to get a WebDriver from the pool.
|
||||
|
||||
Args:
|
||||
window_size: Required window size for the driver
|
||||
user_id: Optional user ID for driver isolation
|
||||
|
||||
Yields:
|
||||
WebDriver instance ready for use
|
||||
"""
|
||||
pooled_driver = None
|
||||
driver_id = None
|
||||
|
||||
try:
|
||||
with self._lock:
|
||||
# Periodic maintenance
|
||||
self._periodic_health_check()
|
||||
|
||||
# Try to get a driver from the pool
|
||||
while not self._pool.empty():
|
||||
try:
|
||||
candidate = self._pool.get_nowait()
|
||||
|
||||
# Check if driver is valid and matches requirements
|
||||
if (
|
||||
self._is_driver_valid(candidate)
|
||||
and candidate.window_size == window_size
|
||||
):
|
||||
# Update user_id for the reused driver
|
||||
candidate.user_id = user_id
|
||||
pooled_driver = candidate
|
||||
break
|
||||
else:
|
||||
# Driver is invalid, destroy it
|
||||
self._destroy_driver(candidate)
|
||||
self._stats["evictions"] += 1
|
||||
except Empty:
|
||||
break
|
||||
|
||||
# If no suitable driver found, create a new one
|
||||
if pooled_driver is None:
|
||||
pooled_driver = self._create_driver(window_size, user_id)
|
||||
|
||||
# Mark driver as in use
|
||||
driver_id = id(pooled_driver.driver)
|
||||
pooled_driver.last_used = time.time()
|
||||
pooled_driver.usage_count += 1
|
||||
self._active_drivers[driver_id] = pooled_driver
|
||||
self._stats["borrowed"] += 1
|
||||
|
||||
# Yield the driver for use
|
||||
yield pooled_driver.driver
|
||||
|
||||
except Exception as e:
|
||||
# Mark driver as unhealthy if an error occurred
|
||||
if pooled_driver:
|
||||
pooled_driver.is_healthy = False
|
||||
logger.error("Error using pooled WebDriver: %s", e)
|
||||
raise
|
||||
|
||||
finally:
|
||||
# Return driver to pool or destroy if unhealthy
|
||||
if pooled_driver and driver_id:
|
||||
with self._lock:
|
||||
self._active_drivers.pop(driver_id, None)
|
||||
|
||||
if pooled_driver.is_healthy and self._is_driver_valid(
|
||||
pooled_driver
|
||||
):
|
||||
# Try to return to pool
|
||||
try:
|
||||
self._pool.put_nowait(pooled_driver)
|
||||
self._stats["returned"] += 1
|
||||
logger.debug("Returned WebDriver to pool")
|
||||
except Full:
|
||||
# Pool is full, destroy the driver
|
||||
self._destroy_driver(pooled_driver)
|
||||
logger.debug("Pool full, destroyed WebDriver")
|
||||
else:
|
||||
# Driver is unhealthy or expired, destroy it
|
||||
self._destroy_driver(pooled_driver)
|
||||
logger.debug("Destroyed unhealthy/expired WebDriver")
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Shutdown the pool and destroy all drivers"""
|
||||
with self._lock:
|
||||
logger.info("Shutting down WebDriver pool")
|
||||
|
||||
# Destroy all active drivers
|
||||
for pooled_driver in self._active_drivers.values():
|
||||
self._destroy_driver(pooled_driver)
|
||||
self._active_drivers.clear()
|
||||
|
||||
# Destroy all pooled drivers
|
||||
while not self._pool.empty():
|
||||
try:
|
||||
pooled_driver = self._pool.get_nowait()
|
||||
self._destroy_driver(pooled_driver)
|
||||
except Empty:
|
||||
break
|
||||
|
||||
logger.info(
|
||||
"WebDriver pool shutdown complete. Final stats: %s", self.get_stats()
|
||||
)
|
||||
|
||||
|
||||
# Global pool instance
|
||||
_global_pool: WebDriverPool | None = None
|
||||
_pool_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_webdriver_pool() -> WebDriverPool:
|
||||
"""Get or create the global WebDriver pool"""
|
||||
global _global_pool
|
||||
|
||||
if _global_pool is None:
|
||||
with _pool_lock:
|
||||
if _global_pool is None:
|
||||
# Get pool configuration from Flask config
|
||||
config = current_app.config
|
||||
pool_config = config.get("WEBDRIVER_POOL", {})
|
||||
|
||||
_global_pool = WebDriverPool(
|
||||
max_pool_size=pool_config.get("MAX_POOL_SIZE", 5),
|
||||
max_age_seconds=pool_config.get("MAX_AGE_SECONDS", 3600),
|
||||
max_usage_count=pool_config.get("MAX_USAGE_COUNT", 50),
|
||||
idle_timeout_seconds=pool_config.get("IDLE_TIMEOUT_SECONDS", 300),
|
||||
health_check_interval=pool_config.get("HEALTH_CHECK_INTERVAL", 60),
|
||||
)
|
||||
logger.info("Initialized global WebDriver pool")
|
||||
|
||||
return _global_pool
|
||||
|
||||
|
||||
def shutdown_webdriver_pool() -> None:
|
||||
"""Shutdown the global WebDriver pool"""
|
||||
global _global_pool
|
||||
|
||||
if _global_pool is not None:
|
||||
with _pool_lock:
|
||||
if _global_pool is not None:
|
||||
_global_pool.shutdown()
|
||||
_global_pool = None
|
||||
@@ -74,6 +74,8 @@ def run_server(
|
||||
|
||||
configure_logging(debug)
|
||||
|
||||
# DO NOT IMPORT TOOLS HERE!! IMPORT THEM IN app.py!!!!!
|
||||
|
||||
if use_factory_config:
|
||||
# Use factory configuration for customization
|
||||
logging.info("Creating MCP app from factory configuration...")
|
||||
|
||||
16
superset/mcp_service/sql_lab/__init__.py
Normal file
16
superset/mcp_service/sql_lab/__init__.py
Normal file
@@ -0,0 +1,16 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
221
superset/mcp_service/sql_lab/execute_sql_core.py
Normal file
221
superset/mcp_service/sql_lab/execute_sql_core.py
Normal file
@@ -0,0 +1,221 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Generic SQL execution core for MCP service.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from superset.mcp_service.mcp_core import BaseCore
|
||||
from superset.mcp_service.sql_lab.schemas import (
|
||||
ExecuteSqlRequest,
|
||||
ExecuteSqlResponse,
|
||||
)
|
||||
|
||||
|
||||
class ExecuteSqlCore(BaseCore):
|
||||
"""
|
||||
Generic tool for executing SQL queries with security validation.
|
||||
|
||||
This tool provides a high-level interface for SQL execution that can be used
|
||||
by different MCP tools or other components. It handles:
|
||||
- Database access validation
|
||||
- SQL query validation (DML permissions, disallowed functions)
|
||||
- Parameter substitution
|
||||
- Query execution with timeout
|
||||
- Result formatting
|
||||
|
||||
The tool can work in two modes:
|
||||
1. Simple mode: Direct SQL execution using sql_lab_utils (default)
|
||||
2. Command mode: Using ExecuteSqlCommand for full SQL Lab features
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
use_command_mode: bool = False,
|
||||
logger: logging.Logger | None = None,
|
||||
) -> None:
|
||||
super().__init__(logger)
|
||||
self.use_command_mode = use_command_mode
|
||||
|
||||
def run_tool(self, request: ExecuteSqlRequest) -> ExecuteSqlResponse:
|
||||
"""
|
||||
Execute SQL query and return results.
|
||||
|
||||
Args:
|
||||
request: ExecuteSqlRequest with database_id, sql, and optional parameters
|
||||
|
||||
Returns:
|
||||
ExecuteSqlResponse with success status, results, or error information
|
||||
"""
|
||||
try:
|
||||
# Import inside method to avoid initialization issues
|
||||
from superset.mcp_service.sql_lab.sql_lab_utils import check_database_access
|
||||
|
||||
# Check database access
|
||||
database = check_database_access(request.database_id)
|
||||
|
||||
if self.use_command_mode:
|
||||
# Use full SQL Lab command for complex queries
|
||||
return self._execute_with_command(request, database)
|
||||
else:
|
||||
# Use simplified execution for basic queries
|
||||
return self._execute_simple(request, database)
|
||||
|
||||
except Exception as e:
|
||||
# Handle errors and return error response with proper error types
|
||||
self._log_error(e, "executing SQL")
|
||||
return self._handle_execution_error(e)
|
||||
|
||||
def _execute_simple(
|
||||
self,
|
||||
request: ExecuteSqlRequest,
|
||||
database: Any,
|
||||
) -> ExecuteSqlResponse:
|
||||
"""Execute SQL using simplified sql_lab_utils."""
|
||||
# Import inside method to avoid initialization issues
|
||||
from superset.mcp_service.sql_lab.sql_lab_utils import execute_sql_query
|
||||
|
||||
results = execute_sql_query(
|
||||
database=database,
|
||||
sql=request.sql,
|
||||
schema=request.schema_name,
|
||||
limit=request.limit,
|
||||
timeout=request.timeout,
|
||||
parameters=request.parameters,
|
||||
)
|
||||
|
||||
return ExecuteSqlResponse(
|
||||
success=True,
|
||||
rows=results.get("rows"),
|
||||
columns=results.get("columns"),
|
||||
row_count=results.get("row_count"),
|
||||
affected_rows=results.get("affected_rows"),
|
||||
query_id=None, # Not available in simple mode
|
||||
execution_time=results.get("execution_time"),
|
||||
error=None,
|
||||
error_type=None,
|
||||
)
|
||||
|
||||
def _execute_with_command(
|
||||
self,
|
||||
request: ExecuteSqlRequest,
|
||||
database: Any,
|
||||
) -> ExecuteSqlResponse:
|
||||
"""Execute SQL using full SQL Lab command (not implemented yet)."""
|
||||
# This would use ExecuteSqlCommand for full SQL Lab features
|
||||
# Including query caching, async execution, complex parsing, etc.
|
||||
# For now, we'll fall back to simple execution
|
||||
self._log_info("Command mode not fully implemented, using simple mode")
|
||||
return self._execute_simple(request, database)
|
||||
|
||||
# Future implementation would look like:
|
||||
# context = SqlJsonExecutionContext(
|
||||
# database_id=request.database_id,
|
||||
# sql=request.sql,
|
||||
# schema=request.schema_name,
|
||||
# limit=request.limit,
|
||||
# # ... other context fields
|
||||
# )
|
||||
#
|
||||
# command = ExecuteSqlCommand(
|
||||
# execution_context=context,
|
||||
# query_dao=QueryDAO(),
|
||||
# database_dao=DatabaseDAO(),
|
||||
# # ... other dependencies
|
||||
# )
|
||||
#
|
||||
# result = command.run()
|
||||
# return self._format_command_result(result)
|
||||
|
||||
def _handle_execution_error(self, e: Exception) -> ExecuteSqlResponse:
|
||||
"""Map exceptions to error responses."""
|
||||
error_type = self._get_error_type(e)
|
||||
return ExecuteSqlResponse(
|
||||
success=False,
|
||||
error=str(e),
|
||||
error_type=error_type,
|
||||
rows=None,
|
||||
columns=None,
|
||||
row_count=None,
|
||||
affected_rows=None,
|
||||
query_id=None,
|
||||
execution_time=None,
|
||||
)
|
||||
|
||||
def _get_error_type(self, e: Exception) -> str:
|
||||
"""Determine error type from exception."""
|
||||
# Import inside method to avoid initialization issues
|
||||
from superset.exceptions import (
|
||||
SupersetDisallowedSQLFunctionException,
|
||||
SupersetDMLNotAllowedException,
|
||||
SupersetErrorException,
|
||||
SupersetSecurityException,
|
||||
SupersetTimeoutException,
|
||||
)
|
||||
|
||||
if isinstance(e, SupersetSecurityException):
|
||||
return "SECURITY_ERROR"
|
||||
elif isinstance(e, SupersetTimeoutException):
|
||||
return "TIMEOUT"
|
||||
elif isinstance(e, SupersetDMLNotAllowedException):
|
||||
return "DML_NOT_ALLOWED"
|
||||
elif isinstance(e, SupersetDisallowedSQLFunctionException):
|
||||
return "DISALLOWED_FUNCTION"
|
||||
elif isinstance(e, SupersetErrorException):
|
||||
return self._extract_superset_error_type(e)
|
||||
else:
|
||||
return "EXECUTION_ERROR"
|
||||
|
||||
def _extract_superset_error_type(self, e: Exception) -> str:
|
||||
"""Extract error type from SupersetErrorException."""
|
||||
if hasattr(e, "error") and hasattr(e.error, "error_type"):
|
||||
error_type_name = e.error.error_type.name
|
||||
# Map common error type patterns
|
||||
if "INVALID_PAYLOAD" in error_type_name:
|
||||
return "INVALID_PAYLOAD_FORMAT_ERROR"
|
||||
elif "DATABASE_NOT_FOUND" in error_type_name:
|
||||
return "DATABASE_NOT_FOUND_ERROR"
|
||||
elif "SECURITY" in error_type_name:
|
||||
return "SECURITY_ERROR"
|
||||
elif "TIMEOUT" in error_type_name:
|
||||
return "TIMEOUT"
|
||||
elif "DML_NOT_ALLOWED" in error_type_name:
|
||||
return "DML_NOT_ALLOWED"
|
||||
else:
|
||||
return error_type_name
|
||||
return "EXECUTION_ERROR"
|
||||
|
||||
def _format_command_result(
|
||||
self, command_result: dict[str, Any]
|
||||
) -> ExecuteSqlResponse:
|
||||
"""Format ExecuteSqlCommand result into ExecuteSqlResponse."""
|
||||
# This would extract relevant fields from command result
|
||||
# Placeholder implementation for future use
|
||||
return ExecuteSqlResponse(
|
||||
success=command_result.get("success", False),
|
||||
rows=command_result.get("data"),
|
||||
columns=command_result.get("columns"),
|
||||
row_count=command_result.get("row_count"),
|
||||
affected_rows=command_result.get("affected_rows"),
|
||||
query_id=command_result.get("query_id"),
|
||||
execution_time=command_result.get("execution_time"),
|
||||
error=command_result.get("error"),
|
||||
error_type=command_result.get("error_type"),
|
||||
)
|
||||
109
superset/mcp_service/sql_lab/schemas.py
Normal file
109
superset/mcp_service/sql_lab/schemas.py
Normal file
@@ -0,0 +1,109 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Schemas for SQL Lab MCP tools."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
|
||||
class ExecuteSqlRequest(BaseModel):
|
||||
"""Request schema for executing SQL queries."""
|
||||
|
||||
database_id: int = Field(
|
||||
..., description="Database connection ID to execute query against"
|
||||
)
|
||||
sql: str = Field(..., description="SQL query to execute")
|
||||
schema_name: str | None = Field(
|
||||
None, description="Schema to use for query execution", alias="schema"
|
||||
)
|
||||
limit: int = Field(
|
||||
default=1000,
|
||||
description="Maximum number of rows to return",
|
||||
ge=1,
|
||||
le=10000,
|
||||
)
|
||||
timeout: int = Field(
|
||||
default=30, description="Query timeout in seconds", ge=1, le=300
|
||||
)
|
||||
parameters: dict[str, Any] | None = Field(
|
||||
None, description="Parameters for query substitution"
|
||||
)
|
||||
|
||||
@field_validator("sql")
|
||||
@classmethod
|
||||
def sql_not_empty(cls, v: str) -> str:
|
||||
if not v or not v.strip():
|
||||
raise ValueError("SQL query cannot be empty")
|
||||
return v.strip()
|
||||
|
||||
|
||||
class ColumnInfo(BaseModel):
|
||||
"""Column metadata information."""
|
||||
|
||||
name: str = Field(..., description="Column name")
|
||||
type: str = Field(..., description="Column data type")
|
||||
is_nullable: bool | None = Field(None, description="Whether column allows NULL")
|
||||
|
||||
|
||||
class ExecuteSqlResponse(BaseModel):
|
||||
"""Response schema for SQL execution results."""
|
||||
|
||||
success: bool = Field(..., description="Whether query executed successfully")
|
||||
rows: Any | None = Field(
|
||||
None, description="Query result rows as list of dictionaries"
|
||||
)
|
||||
columns: list[ColumnInfo] | None = Field(
|
||||
None, description="Column metadata information"
|
||||
)
|
||||
row_count: int | None = Field(None, description="Number of rows returned")
|
||||
affected_rows: int | None = Field(
|
||||
None, description="Number of rows affected (for DML queries)"
|
||||
)
|
||||
query_id: str | None = Field(None, description="Query tracking ID")
|
||||
execution_time: float | None = Field(
|
||||
None, description="Query execution time in seconds"
|
||||
)
|
||||
error: str | None = Field(None, description="Error message if query failed")
|
||||
error_type: str | None = Field(None, description="Type of error if failed")
|
||||
|
||||
|
||||
class OpenSqlLabRequest(BaseModel):
|
||||
"""Request schema for opening SQL Lab with context."""
|
||||
|
||||
database_connection_id: int = Field(
|
||||
..., description="Database connection ID to use in SQL Lab"
|
||||
)
|
||||
schema_name: str | None = Field(
|
||||
None, description="Default schema to select in SQL Lab", alias="schema"
|
||||
)
|
||||
dataset_in_context: str | None = Field(
|
||||
None, description="Dataset name/table to provide as context"
|
||||
)
|
||||
sql: str | None = Field(None, description="SQL query to pre-populate in the editor")
|
||||
title: str | None = Field(None, description="Title for the SQL Lab tab/query")
|
||||
|
||||
|
||||
class SqlLabResponse(BaseModel):
|
||||
"""Response schema for SQL Lab URL generation."""
|
||||
|
||||
url: str = Field(..., description="URL to open SQL Lab with context")
|
||||
database_id: int = Field(..., description="Database ID used")
|
||||
schema_name: str | None = Field(None, description="Schema selected", alias="schema")
|
||||
title: str | None = Field(None, description="Query title")
|
||||
error: str | None = Field(None, description="Error message if failed")
|
||||
243
superset/mcp_service/sql_lab/sql_lab_utils.py
Normal file
243
superset/mcp_service/sql_lab/sql_lab_utils.py
Normal file
@@ -0,0 +1,243 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Utility functions for SQL Lab MCP tools.
|
||||
|
||||
This module contains helper functions for SQL execution, validation,
|
||||
and database access that are shared across SQL Lab tools.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def check_database_access(database_id: int) -> Any:
|
||||
"""Check if user has access to the database."""
|
||||
# Import inside function to avoid initialization issues
|
||||
from superset import db, security_manager
|
||||
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
|
||||
from superset.exceptions import SupersetErrorException, SupersetSecurityException
|
||||
from superset.models.core import Database
|
||||
|
||||
# Use session query to ensure relationships are loaded
|
||||
database = db.session.query(Database).filter_by(id=database_id).first()
|
||||
|
||||
if not database:
|
||||
raise SupersetErrorException(
|
||||
SupersetError(
|
||||
message=f"Database with ID {database_id} not found",
|
||||
error_type=SupersetErrorType.DATABASE_NOT_FOUND_ERROR,
|
||||
level=ErrorLevel.ERROR,
|
||||
)
|
||||
)
|
||||
|
||||
# Check database access permissions
|
||||
if not security_manager.can_access_database(database):
|
||||
raise SupersetSecurityException(
|
||||
SupersetError(
|
||||
message=f"Access denied to database {database.database_name}",
|
||||
error_type=SupersetErrorType.DATABASE_SECURITY_ACCESS_ERROR,
|
||||
level=ErrorLevel.ERROR,
|
||||
)
|
||||
)
|
||||
|
||||
return database
|
||||
|
||||
|
||||
def validate_sql_query(sql: str, database: Any) -> None:
|
||||
"""Validate SQL query for security and syntax."""
|
||||
# Import inside function to avoid initialization issues
|
||||
from flask import current_app as app
|
||||
|
||||
from superset.exceptions import (
|
||||
SupersetDisallowedSQLFunctionException,
|
||||
SupersetDMLNotAllowedException,
|
||||
)
|
||||
|
||||
# Simplified validation without complex parsing
|
||||
sql_upper = sql.upper().strip()
|
||||
|
||||
# Check for DML operations if not allowed
|
||||
dml_keywords = ["INSERT", "UPDATE", "DELETE", "DROP", "CREATE", "ALTER", "TRUNCATE"]
|
||||
if any(sql_upper.startswith(keyword) for keyword in dml_keywords):
|
||||
if not database.allow_dml:
|
||||
raise SupersetDMLNotAllowedException()
|
||||
|
||||
# Check for disallowed functions from config
|
||||
disallowed_functions = app.config.get("DISALLOWED_SQL_FUNCTIONS", {}).get(
|
||||
"sqlite",
|
||||
set(), # Default to sqlite for now
|
||||
)
|
||||
if disallowed_functions:
|
||||
sql_lower = sql.lower()
|
||||
for func in disallowed_functions:
|
||||
if f"{func.lower()}(" in sql_lower:
|
||||
raise SupersetDisallowedSQLFunctionException(disallowed_functions)
|
||||
|
||||
|
||||
def execute_sql_query(
|
||||
database: Any,
|
||||
sql: str,
|
||||
schema: str | None,
|
||||
limit: int,
|
||||
timeout: int,
|
||||
parameters: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Execute SQL query and return results."""
|
||||
# Import inside function to avoid initialization issues
|
||||
from superset.utils.dates import now_as_float
|
||||
|
||||
start_time = now_as_float()
|
||||
|
||||
# Apply parameters and validate
|
||||
sql = _apply_parameters(sql, parameters)
|
||||
validate_sql_query(sql, database)
|
||||
|
||||
# Apply limit for SELECT queries
|
||||
rendered_sql = _apply_limit(sql, limit)
|
||||
|
||||
# Execute and get results
|
||||
results = _execute_query(database, rendered_sql, schema, limit)
|
||||
|
||||
# Calculate execution time
|
||||
end_time = now_as_float()
|
||||
results["execution_time"] = end_time - start_time
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _apply_parameters(sql: str, parameters: dict[str, Any] | None) -> str:
|
||||
"""Apply parameters to SQL query."""
|
||||
# Import inside function to avoid initialization issues
|
||||
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
|
||||
from superset.exceptions import SupersetErrorException
|
||||
|
||||
if parameters:
|
||||
try:
|
||||
return sql.format(**parameters)
|
||||
except KeyError as e:
|
||||
raise SupersetErrorException(
|
||||
SupersetError(
|
||||
message=f"Missing parameter: {e}",
|
||||
error_type=SupersetErrorType.INVALID_PAYLOAD_FORMAT_ERROR,
|
||||
level=ErrorLevel.ERROR,
|
||||
)
|
||||
) from e
|
||||
else:
|
||||
# Check if SQL contains placeholders when no parameters provided
|
||||
import re
|
||||
|
||||
placeholders = re.findall(r"{(\w+)}", sql)
|
||||
if placeholders:
|
||||
raise SupersetErrorException(
|
||||
SupersetError(
|
||||
message=f"Missing parameter: {placeholders[0]}",
|
||||
error_type=SupersetErrorType.INVALID_PAYLOAD_FORMAT_ERROR,
|
||||
level=ErrorLevel.ERROR,
|
||||
)
|
||||
)
|
||||
return sql
|
||||
|
||||
|
||||
def _apply_limit(sql: str, limit: int) -> str:
|
||||
"""Apply limit to SELECT queries if not already present."""
|
||||
sql_lower = sql.lower().strip()
|
||||
if sql_lower.startswith("select") and "limit" not in sql_lower:
|
||||
return f"{sql.rstrip().rstrip(';')} LIMIT {limit}"
|
||||
return sql
|
||||
|
||||
|
||||
def _execute_query(
|
||||
database: Any,
|
||||
sql: str,
|
||||
schema: str | None,
|
||||
limit: int,
|
||||
) -> dict[str, Any]:
|
||||
"""Execute the query and process results."""
|
||||
# Import inside function to avoid initialization issues
|
||||
from superset.utils.core import QuerySource
|
||||
|
||||
results = {
|
||||
"rows": [],
|
||||
"columns": [],
|
||||
"row_count": 0,
|
||||
"affected_rows": None,
|
||||
"execution_time": 0.0,
|
||||
}
|
||||
|
||||
try:
|
||||
# Execute query with timeout
|
||||
with database.get_raw_connection(
|
||||
catalog=None,
|
||||
schema=schema,
|
||||
source=QuerySource.SQL_LAB,
|
||||
) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(sql)
|
||||
|
||||
# Process results based on query type
|
||||
if _is_select_query(sql):
|
||||
_process_select_results(cursor, results, limit)
|
||||
else:
|
||||
_process_dml_results(cursor, conn, results)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error executing SQL: %s", e)
|
||||
raise
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _is_select_query(sql: str) -> bool:
|
||||
"""Check if SQL is a SELECT query."""
|
||||
return sql.lower().strip().startswith("select")
|
||||
|
||||
|
||||
def _process_select_results(cursor: Any, results: dict[str, Any], limit: int) -> None:
|
||||
"""Process SELECT query results."""
|
||||
# Fetch results
|
||||
data = cursor.fetchmany(limit)
|
||||
|
||||
# Get column metadata
|
||||
column_info = []
|
||||
if cursor.description:
|
||||
for col in cursor.description:
|
||||
column_info.append(
|
||||
{
|
||||
"name": col[0],
|
||||
"type": str(col[1]) if col[1] else "unknown",
|
||||
"is_nullable": col[6] if len(col) > 6 else None,
|
||||
}
|
||||
)
|
||||
|
||||
# Set column info regardless of whether there's data
|
||||
if column_info:
|
||||
results["columns"] = column_info
|
||||
|
||||
# Convert rows to dictionaries
|
||||
column_names = [col["name"] for col in column_info]
|
||||
results["rows"] = [dict(zip(column_names, row, strict=False)) for row in data]
|
||||
results["row_count"] = len(data)
|
||||
|
||||
|
||||
def _process_dml_results(cursor: Any, conn: Any, results: dict[str, Any]) -> None:
|
||||
"""Process DML query results."""
|
||||
results["affected_rows"] = cursor.rowcount
|
||||
conn.commit() # pylint: disable=consider-using-transaction
|
||||
30
superset/mcp_service/sql_lab/tool/__init__.py
Normal file
30
superset/mcp_service/sql_lab/tool/__init__.py
Normal file
@@ -0,0 +1,30 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
SQL Lab MCP Tools
|
||||
"""
|
||||
|
||||
from superset.mcp_service.sql_lab.tool.execute_sql import execute_sql
|
||||
from superset.mcp_service.sql_lab.tool.open_sql_lab_with_context import (
|
||||
open_sql_lab_with_context,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"execute_sql",
|
||||
"open_sql_lab_with_context",
|
||||
]
|
||||
94
superset/mcp_service/sql_lab/tool/execute_sql.py
Normal file
94
superset/mcp_service/sql_lab/tool/execute_sql.py
Normal file
@@ -0,0 +1,94 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Execute SQL MCP Tool
|
||||
|
||||
Tool for executing SQL queries against databases with security validation
|
||||
and timeout protection.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.sql_lab.execute_sql_core import ExecuteSqlCore
|
||||
from superset.mcp_service.sql_lab.schemas import (
|
||||
ExecuteSqlRequest,
|
||||
ExecuteSqlResponse,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def execute_sql(request: ExecuteSqlRequest, ctx: Context) -> ExecuteSqlResponse:
|
||||
"""Execute SQL query against database.
|
||||
|
||||
Returns query results with security validation and timeout protection.
|
||||
"""
|
||||
await ctx.info(
|
||||
"Starting SQL execution: database_id=%s, timeout=%s, limit=%s, schema=%s"
|
||||
% (request.database_id, request.timeout, request.limit, request.schema_name)
|
||||
)
|
||||
|
||||
# Log SQL query details (truncated for security)
|
||||
sql_preview = request.sql[:100] + "..." if len(request.sql) > 100 else request.sql
|
||||
await ctx.debug(
|
||||
"SQL query details: sql_preview=%r, sql_length=%s, has_parameters=%s"
|
||||
% (
|
||||
sql_preview,
|
||||
len(request.sql),
|
||||
bool(request.parameters),
|
||||
)
|
||||
)
|
||||
|
||||
logger.info("Executing SQL query on database ID: %s", request.database_id)
|
||||
|
||||
try:
|
||||
# Use the ExecuteSqlCore to handle all the logic
|
||||
sql_tool = ExecuteSqlCore(use_command_mode=False, logger=logger)
|
||||
result = sql_tool.run_tool(request)
|
||||
|
||||
# Log successful execution
|
||||
if hasattr(result, "data") and result.data:
|
||||
row_count = len(result.data) if isinstance(result.data, list) else 1
|
||||
await ctx.info(
|
||||
"SQL execution completed successfully: rows_returned=%s, "
|
||||
"query_duration_ms=%s"
|
||||
% (
|
||||
row_count,
|
||||
getattr(result, "query_duration_ms", None),
|
||||
)
|
||||
)
|
||||
else:
|
||||
await ctx.info("SQL execution completed: status=no_data_returned")
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
await ctx.error(
|
||||
"SQL execution failed: error=%s, database_id=%s"
|
||||
% (
|
||||
str(e),
|
||||
request.database_id,
|
||||
)
|
||||
)
|
||||
raise
|
||||
118
superset/mcp_service/sql_lab/tool/open_sql_lab_with_context.py
Normal file
118
superset/mcp_service/sql_lab/tool/open_sql_lab_with_context.py
Normal file
@@ -0,0 +1,118 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Open SQL Lab with Context MCP Tool
|
||||
|
||||
Tool for generating SQL Lab URLs with pre-populated query and context.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.sql_lab.schemas import (
|
||||
OpenSqlLabRequest,
|
||||
SqlLabResponse,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
def open_sql_lab_with_context(
|
||||
request: OpenSqlLabRequest, ctx: Context
|
||||
) -> SqlLabResponse:
|
||||
"""Generate SQL Lab URL with pre-populated query and context.
|
||||
|
||||
Returns URL for direct navigation.
|
||||
"""
|
||||
try:
|
||||
from superset.daos.database import DatabaseDAO
|
||||
|
||||
# Validate database exists and is accessible
|
||||
database = DatabaseDAO.find_by_id(request.database_connection_id)
|
||||
if not database:
|
||||
return SqlLabResponse(
|
||||
url="",
|
||||
database_id=request.database_connection_id,
|
||||
schema_name=request.schema_name,
|
||||
title=request.title,
|
||||
error=f"Database with ID {request.database_connection_id} not found",
|
||||
)
|
||||
|
||||
# Build query parameters for SQL Lab URL
|
||||
params = {
|
||||
"dbid": str(request.database_connection_id),
|
||||
}
|
||||
|
||||
if request.schema_name:
|
||||
params["schema"] = request.schema_name
|
||||
|
||||
if request.sql:
|
||||
params["sql"] = request.sql
|
||||
|
||||
if request.title:
|
||||
params["title"] = request.title
|
||||
|
||||
if request.dataset_in_context:
|
||||
# Add dataset context as a comment in the SQL if no SQL provided
|
||||
if not request.sql:
|
||||
context_comment = (
|
||||
f"-- Context: Working with dataset '{request.dataset_in_context}'\n"
|
||||
f"-- Database: {database.database_name}\n"
|
||||
)
|
||||
if request.schema_name:
|
||||
context_comment += f"-- Schema: {request.schema_name}\n"
|
||||
table_reference = (
|
||||
f"{request.schema_name}.{request.dataset_in_context}"
|
||||
)
|
||||
else:
|
||||
table_reference = request.dataset_in_context
|
||||
|
||||
context_comment += f"\nSELECT * FROM {table_reference} LIMIT 100;"
|
||||
params["sql"] = context_comment
|
||||
|
||||
# Construct SQL Lab URL
|
||||
query_string = urlencode(params)
|
||||
url = f"/sqllab?{query_string}"
|
||||
|
||||
logger.info(
|
||||
"Generated SQL Lab URL for database %s", request.database_connection_id
|
||||
)
|
||||
|
||||
return SqlLabResponse(
|
||||
url=url,
|
||||
database_id=request.database_connection_id,
|
||||
schema_name=request.schema_name,
|
||||
title=request.title,
|
||||
error=None,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error generating SQL Lab URL: %s", e)
|
||||
return SqlLabResponse(
|
||||
url="",
|
||||
database_id=request.database_connection_id,
|
||||
schema_name=request.schema_name,
|
||||
title=request.title,
|
||||
error=f"Failed to generate SQL Lab URL: {str(e)}",
|
||||
)
|
||||
21
superset/mcp_service/system/prompts/__init__.py
Normal file
21
superset/mcp_service/system/prompts/__init__.py
Normal file
@@ -0,0 +1,21 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""System prompts for Superset MCP service"""
|
||||
|
||||
# Import to register prompts when module is imported
|
||||
from . import quickstart # noqa: F401
|
||||
94
superset/mcp_service/system/prompts/quickstart.py
Normal file
94
superset/mcp_service/system/prompts/quickstart.py
Normal file
@@ -0,0 +1,94 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
System prompts for general Superset guidance
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.prompt("superset_quickstart")
|
||||
@mcp_auth_hook
|
||||
async def superset_quickstart_prompt(
|
||||
user_type: str = "analyst", focus_area: str = "general"
|
||||
) -> str:
|
||||
"""
|
||||
Guide new users through their first Superset experience.
|
||||
|
||||
This prompt helps users:
|
||||
1. Understand what data is available
|
||||
2. Create their first visualization
|
||||
3. Build a simple dashboard
|
||||
4. Learn key Superset concepts
|
||||
|
||||
Args:
|
||||
user_type: Type of user (analyst, executive, developer)
|
||||
focus_area: Area of interest (sales, marketing, operations, general)
|
||||
"""
|
||||
# Build personalized prompt based on user type
|
||||
intro_messages = {
|
||||
"analyst": "I see you're an analyst. Let's explore the data and build some "
|
||||
"detailed visualizations.",
|
||||
"executive": "Welcome! Let's create a high-level dashboard with key business "
|
||||
"metrics.",
|
||||
"developer": "Great to have a developer here! Let's explore both the UI and "
|
||||
"API capabilities.",
|
||||
}
|
||||
|
||||
focus_examples = {
|
||||
"sales": "Since you're interested in sales, we'll focus on revenue, customer, "
|
||||
"and product metrics.",
|
||||
"marketing": "For marketing analytics, we'll look at campaigns, conversions, "
|
||||
"and customer acquisition.",
|
||||
"operations": "Let's explore operational efficiency, inventory, and process "
|
||||
"metrics.",
|
||||
"general": "We'll explore various datasets to find what's most relevant to "
|
||||
"you.",
|
||||
}
|
||||
|
||||
intro = intro_messages.get(user_type, intro_messages["analyst"])
|
||||
focus = focus_examples.get(focus_area, focus_examples["general"])
|
||||
|
||||
return f"""Welcome to Apache Superset! I'll guide you through creating your first
|
||||
dashboard.
|
||||
|
||||
{intro} {focus}
|
||||
|
||||
I'll help you through these steps:
|
||||
1. 📊 **Explore Available Data** - See what datasets you can work with
|
||||
2. 🔍 **Understand Your Data** - Examine columns, metrics, and sample data
|
||||
3. 📈 **Create Visualizations** - Build charts that tell a story
|
||||
4. 🎯 **Design a Dashboard** - Combine charts into an interactive dashboard
|
||||
5. 🚀 **Learn Advanced Features** - Discover filters, SQL Lab, and more
|
||||
|
||||
To get started, I'll use these tools:
|
||||
- `get_superset_instance_info` - Overview of your Superset instance
|
||||
- `list_datasets` - Find available datasets
|
||||
- `get_dataset_info` - Explore dataset details
|
||||
- `generate_chart` - Create visualizations
|
||||
- `generate_dashboard` - Build your dashboard
|
||||
|
||||
Let me begin by checking what's available in your Superset instance. I'll first get
|
||||
an overview, then show you the datasets filtered by your interest in {focus_area}.
|
||||
|
||||
Would you like me to start by showing you what data you can work with?"""
|
||||
21
superset/mcp_service/system/resources/__init__.py
Normal file
21
superset/mcp_service/system/resources/__init__.py
Normal file
@@ -0,0 +1,21 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""System resources for Superset MCP service"""
|
||||
|
||||
# Import to register resources when module is imported
|
||||
from . import instance_metadata # noqa: F401
|
||||
107
superset/mcp_service/system/resources/instance_metadata.py
Normal file
107
superset/mcp_service/system/resources/instance_metadata.py
Normal file
@@ -0,0 +1,107 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
System resources for providing Superset configuration and stats
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.resource("superset://instance/metadata")
|
||||
@mcp_auth_hook
|
||||
def get_instance_metadata_resource() -> str:
|
||||
"""
|
||||
Provide comprehensive metadata about the Superset instance.
|
||||
|
||||
This resource gives LLMs context about:
|
||||
- Available datasets and their popularity
|
||||
- Dashboard and chart statistics
|
||||
- Database connections
|
||||
- Popular queries and usage patterns
|
||||
- Available visualization types
|
||||
- Feature flags and configuration
|
||||
"""
|
||||
try:
|
||||
# Import the shared core and DAOs at runtime
|
||||
# Create a shared core instance for the resource
|
||||
from typing import Any, cast, Type
|
||||
|
||||
from superset.daos.base import BaseDAO
|
||||
from superset.daos.chart import ChartDAO
|
||||
from superset.daos.dashboard import DashboardDAO
|
||||
from superset.daos.database import DatabaseDAO
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
from superset.daos.tag import TagDAO
|
||||
from superset.daos.user import UserDAO
|
||||
from superset.mcp_service.mcp_core import InstanceInfoCore
|
||||
from superset.mcp_service.system.schemas import InstanceInfo
|
||||
from superset.mcp_service.system.tool.get_superset_instance_info import (
|
||||
calculate_dashboard_breakdown,
|
||||
calculate_database_breakdown,
|
||||
calculate_instance_summary,
|
||||
calculate_popular_content,
|
||||
calculate_recent_activity,
|
||||
)
|
||||
|
||||
instance_info_core = InstanceInfoCore(
|
||||
dao_classes={
|
||||
"dashboards": cast(Type[BaseDAO[Any]], DashboardDAO),
|
||||
"charts": cast(Type[BaseDAO[Any]], ChartDAO),
|
||||
"datasets": cast(Type[BaseDAO[Any]], DatasetDAO),
|
||||
"databases": cast(Type[BaseDAO[Any]], DatabaseDAO),
|
||||
"users": cast(Type[BaseDAO[Any]], UserDAO),
|
||||
"tags": cast(Type[BaseDAO[Any]], TagDAO),
|
||||
},
|
||||
output_schema=InstanceInfo,
|
||||
metric_calculators={
|
||||
"instance_summary": calculate_instance_summary,
|
||||
"recent_activity": calculate_recent_activity,
|
||||
"dashboard_breakdown": calculate_dashboard_breakdown,
|
||||
"database_breakdown": calculate_database_breakdown,
|
||||
"popular_content": calculate_popular_content,
|
||||
},
|
||||
time_windows={
|
||||
"recent": 7,
|
||||
"monthly": 30,
|
||||
"quarterly": 90,
|
||||
},
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
# Use the shared core's resource method
|
||||
return instance_info_core.get_resource()
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error generating instance metadata: %s", e)
|
||||
# Return minimal metadata on error
|
||||
from superset.utils import json
|
||||
|
||||
return json.dumps(
|
||||
{
|
||||
"error": "Unable to fetch complete metadata",
|
||||
"tips": [
|
||||
"Use list_datasets to explore available data",
|
||||
"Use get_superset_instance_info for basic stats",
|
||||
],
|
||||
}
|
||||
)
|
||||
@@ -16,10 +16,14 @@
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Pydantic schemas for shared system types (UserInfo, TagInfo, PaginationInfo)
|
||||
Pydantic schemas for system-level (instance/info) responses
|
||||
|
||||
This module contains Pydantic models for serializing Superset instance metadata and
|
||||
system-level info.
|
||||
"""
|
||||
|
||||
from typing import List
|
||||
from datetime import datetime
|
||||
from typing import Dict, List
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
@@ -36,12 +40,93 @@ class HealthCheckResponse(BaseModel):
|
||||
version: str
|
||||
python_version: str
|
||||
platform: str
|
||||
uptime_seconds: float
|
||||
uptime_seconds: float | None = None
|
||||
|
||||
|
||||
class GetSupersetInstanceInfoRequest(BaseModel):
|
||||
"""
|
||||
Request schema for get_superset_instance_info tool.
|
||||
|
||||
Currently has no parameters but provides consistent API for future extensibility.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(
|
||||
extra="forbid",
|
||||
str_strip_whitespace=True,
|
||||
)
|
||||
|
||||
|
||||
class InstanceSummary(BaseModel):
|
||||
total_dashboards: int = Field(..., description="Total number of dashboards")
|
||||
total_charts: int = Field(..., description="Total number of charts")
|
||||
total_datasets: int = Field(..., description="Total number of datasets")
|
||||
total_databases: int = Field(..., description="Total number of databases")
|
||||
total_users: int = Field(..., description="Total number of users")
|
||||
total_roles: int = Field(..., description="Total number of roles")
|
||||
total_tags: int = Field(..., description="Total number of tags")
|
||||
avg_charts_per_dashboard: float = Field(
|
||||
..., description="Average number of charts per dashboard"
|
||||
)
|
||||
|
||||
|
||||
class RecentActivity(BaseModel):
|
||||
dashboards_created_last_30_days: int = Field(
|
||||
..., description="Dashboards created in the last 30 days"
|
||||
)
|
||||
charts_created_last_30_days: int = Field(
|
||||
..., description="Charts created in the last 30 days"
|
||||
)
|
||||
datasets_created_last_30_days: int = Field(
|
||||
..., description="Datasets created in the last 30 days"
|
||||
)
|
||||
dashboards_modified_last_7_days: int = Field(
|
||||
..., description="Dashboards modified in the last 7 days"
|
||||
)
|
||||
charts_modified_last_7_days: int = Field(
|
||||
..., description="Charts modified in the last 7 days"
|
||||
)
|
||||
datasets_modified_last_7_days: int = Field(
|
||||
..., description="Datasets modified in the last 7 days"
|
||||
)
|
||||
|
||||
|
||||
class DashboardBreakdown(BaseModel):
|
||||
published: int = Field(..., description="Number of published dashboards")
|
||||
unpublished: int = Field(..., description="Number of unpublished dashboards")
|
||||
certified: int = Field(..., description="Number of certified dashboards")
|
||||
with_charts: int = Field(..., description="Number of dashboards with charts")
|
||||
without_charts: int = Field(..., description="Number of dashboards without charts")
|
||||
|
||||
|
||||
class DatabaseBreakdown(BaseModel):
|
||||
by_type: Dict[str, int] = Field(..., description="Breakdown of databases by type")
|
||||
|
||||
|
||||
class PopularContent(BaseModel):
|
||||
top_tags: List[str] = Field(..., description="Most popular tags")
|
||||
top_creators: List[str] = Field(..., description="Most active creators")
|
||||
|
||||
|
||||
class InstanceInfo(BaseModel):
|
||||
instance_summary: InstanceSummary = Field(
|
||||
..., description="Instance summary information"
|
||||
)
|
||||
recent_activity: RecentActivity = Field(
|
||||
..., description="Recent activity information"
|
||||
)
|
||||
dashboard_breakdown: DashboardBreakdown = Field(
|
||||
..., description="Dashboard breakdown information"
|
||||
)
|
||||
database_breakdown: DatabaseBreakdown = Field(
|
||||
..., description="Database breakdown by type"
|
||||
)
|
||||
popular_content: PopularContent = Field(
|
||||
..., description="Popular content information"
|
||||
)
|
||||
timestamp: datetime = Field(..., description="Response timestamp")
|
||||
|
||||
|
||||
class UserInfo(BaseModel):
|
||||
"""User information schema."""
|
||||
|
||||
id: int | None = None
|
||||
username: str | None = None
|
||||
first_name: str | None = None
|
||||
@@ -51,17 +136,19 @@ class UserInfo(BaseModel):
|
||||
|
||||
|
||||
class TagInfo(BaseModel):
|
||||
"""Tag information schema."""
|
||||
|
||||
id: int | None = None
|
||||
name: str | None = None
|
||||
type: str | None = None
|
||||
description: str | None = None
|
||||
|
||||
|
||||
class PaginationInfo(BaseModel):
|
||||
"""Pagination metadata."""
|
||||
class RoleInfo(BaseModel):
|
||||
id: int | None = None
|
||||
name: str | None = None
|
||||
permissions: List[str] | None = None
|
||||
|
||||
|
||||
class PaginationInfo(BaseModel):
|
||||
page: int = Field(..., description="Current page number")
|
||||
page_size: int = Field(..., description="Number of items per page")
|
||||
total_count: int = Field(..., description="Total number of items")
|
||||
@@ -69,11 +156,3 @@ class PaginationInfo(BaseModel):
|
||||
has_next: bool = Field(..., description="Whether there is a next page")
|
||||
has_previous: bool = Field(..., description="Whether there is a previous page")
|
||||
model_config = ConfigDict(ser_json_timedelta="iso8601")
|
||||
|
||||
|
||||
class RoleInfo(BaseModel):
|
||||
"""Role information schema (for future use)."""
|
||||
|
||||
id: int | None = None
|
||||
name: str | None = None
|
||||
permissions: List[str] | None = None
|
||||
|
||||
@@ -16,3 +16,11 @@
|
||||
# under the License.
|
||||
|
||||
"""System tools for MCP service."""
|
||||
|
||||
from .get_superset_instance_info import get_superset_instance_info
|
||||
from .health_check import health_check
|
||||
|
||||
__all__ = [
|
||||
"health_check",
|
||||
"get_superset_instance_info",
|
||||
]
|
||||
|
||||
268
superset/mcp_service/system/tool/get_superset_instance_info.py
Normal file
268
superset/mcp_service/system/tool/get_superset_instance_info.py
Normal file
@@ -0,0 +1,268 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Get Superset instance high-level information FastMCP tool using configurable
|
||||
InstanceInfoCore for flexible, extensible metrics calculation.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict
|
||||
|
||||
from fastmcp import Context
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.mcp_core import InstanceInfoCore
|
||||
from superset.mcp_service.system.schemas import (
|
||||
DashboardBreakdown,
|
||||
DatabaseBreakdown,
|
||||
GetSupersetInstanceInfoRequest,
|
||||
InstanceInfo,
|
||||
InstanceSummary,
|
||||
PopularContent,
|
||||
RecentActivity,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def calculate_dashboard_breakdown(
|
||||
base_counts: Dict[str, int],
|
||||
time_metrics: Dict[str, Dict[str, int]],
|
||||
dao_classes: Dict[str, Any],
|
||||
) -> DashboardBreakdown:
|
||||
"""Calculate detailed dashboard breakdown metrics."""
|
||||
try:
|
||||
from superset.daos.base import ColumnOperator, ColumnOperatorEnum
|
||||
from superset.extensions import db
|
||||
from superset.models.dashboard import Dashboard
|
||||
|
||||
dashboard_dao = dao_classes["dashboards"]
|
||||
|
||||
# Published vs unpublished
|
||||
published_count = dashboard_dao.count(
|
||||
column_operators=[
|
||||
ColumnOperator(col="published", opr=ColumnOperatorEnum.eq, value=True)
|
||||
]
|
||||
)
|
||||
unpublished_count = base_counts.get("total_dashboards", 0) - published_count
|
||||
|
||||
# Certified dashboards
|
||||
certified_count = dashboard_dao.count(
|
||||
column_operators=[
|
||||
ColumnOperator(
|
||||
col="certified_by", opr=ColumnOperatorEnum.is_not_null, value=None
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Dashboards with/without charts
|
||||
dashboards_with_charts = (
|
||||
db.session.query(Dashboard).join(Dashboard.slices).distinct().count()
|
||||
)
|
||||
dashboards_without_charts = (
|
||||
base_counts.get("total_dashboards", 0) - dashboards_with_charts
|
||||
)
|
||||
|
||||
return DashboardBreakdown(
|
||||
published=published_count,
|
||||
unpublished=unpublished_count,
|
||||
certified=certified_count,
|
||||
with_charts=dashboards_with_charts,
|
||||
without_charts=dashboards_without_charts,
|
||||
)
|
||||
except Exception:
|
||||
# Return empty breakdown on error
|
||||
return DashboardBreakdown(
|
||||
published=0,
|
||||
unpublished=0,
|
||||
certified=0,
|
||||
with_charts=0,
|
||||
without_charts=0,
|
||||
)
|
||||
|
||||
|
||||
def calculate_database_breakdown(
|
||||
base_counts: Dict[str, int],
|
||||
time_metrics: Dict[str, Dict[str, int]],
|
||||
dao_classes: Dict[str, Any],
|
||||
) -> DatabaseBreakdown:
|
||||
"""Calculate database type breakdown."""
|
||||
try:
|
||||
from superset.extensions import db
|
||||
from superset.models.core import Database
|
||||
|
||||
# Get database types distribution
|
||||
db_types = db.session.query(
|
||||
Database.database_name, Database.sqlalchemy_uri
|
||||
).all()
|
||||
|
||||
type_counts: Dict[str, int] = {}
|
||||
for _name, uri in db_types:
|
||||
if uri:
|
||||
# Extract database type from SQLAlchemy URI
|
||||
db_type = uri.split("://")[0] if "://" in uri else "unknown"
|
||||
type_counts[db_type] = type_counts.get(db_type, 0) + 1
|
||||
else:
|
||||
type_counts["unknown"] = type_counts.get("unknown", 0) + 1
|
||||
|
||||
return DatabaseBreakdown(by_type=type_counts)
|
||||
except Exception:
|
||||
# Return empty breakdown on error
|
||||
return DatabaseBreakdown(by_type={})
|
||||
|
||||
|
||||
def calculate_instance_summary(
|
||||
base_counts: Dict[str, int],
|
||||
time_metrics: Dict[str, Dict[str, int]],
|
||||
dao_classes: Dict[str, Any],
|
||||
) -> InstanceSummary:
|
||||
"""Calculate instance summary with computed metrics."""
|
||||
try:
|
||||
from flask_appbuilder.security.sqla.models import Role
|
||||
|
||||
from superset.extensions import db
|
||||
|
||||
# Add roles count (no DAO available)
|
||||
total_roles = db.session.query(Role).count()
|
||||
|
||||
# Calculate average charts per dashboard
|
||||
total_dashboards = base_counts.get("total_dashboards", 0)
|
||||
total_charts = base_counts.get("total_charts", 0)
|
||||
avg_charts_per_dashboard = (
|
||||
(total_charts / total_dashboards) if total_dashboards > 0 else 0
|
||||
)
|
||||
|
||||
return InstanceSummary(
|
||||
total_dashboards=total_dashboards,
|
||||
total_charts=total_charts,
|
||||
total_datasets=base_counts.get("total_datasets", 0),
|
||||
total_databases=base_counts.get("total_databases", 0),
|
||||
total_users=base_counts.get("total_users", 0),
|
||||
total_roles=total_roles,
|
||||
total_tags=base_counts.get("total_tags", 0),
|
||||
avg_charts_per_dashboard=round(avg_charts_per_dashboard, 2),
|
||||
)
|
||||
except Exception:
|
||||
# Return empty summary on error
|
||||
return InstanceSummary(
|
||||
total_dashboards=0,
|
||||
total_charts=0,
|
||||
total_datasets=0,
|
||||
total_databases=0,
|
||||
total_users=0,
|
||||
total_roles=0,
|
||||
total_tags=0,
|
||||
avg_charts_per_dashboard=0.0,
|
||||
)
|
||||
|
||||
|
||||
def calculate_recent_activity(
|
||||
base_counts: Dict[str, int],
|
||||
time_metrics: Dict[str, Dict[str, int]],
|
||||
dao_classes: Dict[str, Any],
|
||||
) -> RecentActivity:
|
||||
"""Transform time metrics into RecentActivity format."""
|
||||
monthly = time_metrics.get("monthly", {})
|
||||
recent = time_metrics.get("recent", {})
|
||||
|
||||
return RecentActivity(
|
||||
dashboards_created_last_30_days=monthly.get("dashboards_created", 0),
|
||||
charts_created_last_30_days=monthly.get("charts_created", 0),
|
||||
datasets_created_last_30_days=monthly.get("datasets_created", 0),
|
||||
dashboards_modified_last_7_days=recent.get("dashboards_modified", 0),
|
||||
charts_modified_last_7_days=recent.get("charts_modified", 0),
|
||||
datasets_modified_last_7_days=recent.get("datasets_modified", 0),
|
||||
)
|
||||
|
||||
|
||||
def calculate_popular_content(
|
||||
base_counts: Dict[str, int],
|
||||
time_metrics: Dict[str, Dict[str, int]],
|
||||
dao_classes: Dict[str, Any],
|
||||
) -> PopularContent:
|
||||
"""Calculate popular content metrics (placeholder implementation)."""
|
||||
|
||||
return PopularContent(
|
||||
top_tags=[],
|
||||
top_creators=[],
|
||||
)
|
||||
|
||||
|
||||
# Configure the instance info core
|
||||
_instance_info_core = InstanceInfoCore(
|
||||
dao_classes={
|
||||
"dashboards": None, # type: ignore[dict-item] # Will be set at runtime
|
||||
"charts": None, # type: ignore[dict-item]
|
||||
"datasets": None, # type: ignore[dict-item]
|
||||
"databases": None, # type: ignore[dict-item]
|
||||
"users": None, # type: ignore[dict-item]
|
||||
"tags": None, # type: ignore[dict-item]
|
||||
},
|
||||
output_schema=InstanceInfo,
|
||||
metric_calculators={
|
||||
"instance_summary": calculate_instance_summary,
|
||||
"recent_activity": calculate_recent_activity,
|
||||
"dashboard_breakdown": calculate_dashboard_breakdown,
|
||||
"database_breakdown": calculate_database_breakdown,
|
||||
"popular_content": calculate_popular_content,
|
||||
},
|
||||
time_windows={
|
||||
"recent": 7,
|
||||
"monthly": 30,
|
||||
"quarterly": 90,
|
||||
},
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
def get_superset_instance_info(
|
||||
request: GetSupersetInstanceInfoRequest, ctx: Context
|
||||
) -> InstanceInfo:
|
||||
"""Get Superset instance statistics.
|
||||
|
||||
Returns counts, activity metrics, and database types.
|
||||
"""
|
||||
try:
|
||||
# Import DAOs at runtime to avoid circular imports
|
||||
from superset.daos.chart import ChartDAO
|
||||
from superset.daos.dashboard import DashboardDAO
|
||||
from superset.daos.database import DatabaseDAO
|
||||
from superset.daos.dataset import DatasetDAO
|
||||
from superset.daos.tag import TagDAO
|
||||
from superset.daos.user import UserDAO
|
||||
|
||||
# Configure DAO classes at runtime
|
||||
_instance_info_core.dao_classes = {
|
||||
"dashboards": DashboardDAO,
|
||||
"charts": ChartDAO,
|
||||
"datasets": DatasetDAO,
|
||||
"databases": DatabaseDAO,
|
||||
"users": UserDAO,
|
||||
"tags": TagDAO,
|
||||
}
|
||||
|
||||
# Run the configurable core
|
||||
return _instance_info_core.run_tool()
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Unexpected error in instance info: {str(e)}"
|
||||
logger.error(error_msg, exc_info=True)
|
||||
raise
|
||||
@@ -20,15 +20,16 @@
|
||||
import datetime
|
||||
import logging
|
||||
import platform
|
||||
import time
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.auth import mcp_auth_hook
|
||||
from superset.mcp_service.system.schemas import HealthCheckResponse
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
async def health_check() -> HealthCheckResponse:
|
||||
"""
|
||||
Simple health check tool for testing the MCP service.
|
||||
@@ -47,7 +48,6 @@ async def health_check() -> HealthCheckResponse:
|
||||
version="1.0.0",
|
||||
python_version=platform.python_version(),
|
||||
platform=platform.system(),
|
||||
uptime_seconds=time.time(), # Simple uptime approximation
|
||||
)
|
||||
|
||||
logger.info("Health check completed successfully")
|
||||
@@ -63,5 +63,4 @@ async def health_check() -> HealthCheckResponse:
|
||||
version="1.0.0",
|
||||
python_version=platform.python_version(),
|
||||
platform=platform.system(),
|
||||
uptime_seconds=0.0,
|
||||
)
|
||||
|
||||
143
superset/mcp_service/utils/cache_utils.py
Normal file
143
superset/mcp_service/utils/cache_utils.py
Normal file
@@ -0,0 +1,143 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Cache utilities for MCP tools.
|
||||
|
||||
This module provides utilities for working with Superset's cache layers
|
||||
and implementing cache control in MCP tools.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict
|
||||
|
||||
from superset.mcp_service.common.cache_schemas import CacheStatus
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_cache_status_from_result(
|
||||
result: Dict[str, Any], force_refresh: bool = False
|
||||
) -> CacheStatus:
|
||||
"""
|
||||
Extract cache status information from a Superset query result.
|
||||
|
||||
Args:
|
||||
result: Query result dictionary from Superset
|
||||
force_refresh: Whether cache was force refreshed
|
||||
|
||||
Returns:
|
||||
CacheStatus object with cache usage information
|
||||
"""
|
||||
# Handle different result structures
|
||||
if "queries" in result and len(result["queries"]) > 0:
|
||||
query_result = result["queries"][0]
|
||||
else:
|
||||
query_result = result
|
||||
|
||||
cache_hit = bool(query_result.get("is_cached", False))
|
||||
|
||||
# Convert cache age to seconds if available
|
||||
cache_age_seconds = None
|
||||
if cache_age := query_result.get("cache_dttm"):
|
||||
try:
|
||||
from datetime import datetime
|
||||
|
||||
if isinstance(cache_age, str):
|
||||
cache_dt = datetime.fromisoformat(cache_age.replace("Z", "+00:00"))
|
||||
cache_age_seconds = int(
|
||||
(datetime.now(cache_dt.tzinfo) - cache_dt).total_seconds()
|
||||
)
|
||||
elif isinstance(cache_age, datetime):
|
||||
cache_age_seconds = int(
|
||||
(datetime.now(cache_age.tzinfo) - cache_age).total_seconds()
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug("Could not parse cache age: %s", e)
|
||||
|
||||
return CacheStatus(
|
||||
cache_hit=cache_hit,
|
||||
cache_type="query" if cache_hit else "none",
|
||||
cache_age_seconds=cache_age_seconds,
|
||||
refreshed=force_refresh,
|
||||
)
|
||||
|
||||
|
||||
def apply_cache_control_to_query_context(
|
||||
query_context: Dict[str, Any],
|
||||
use_cache: bool = True,
|
||||
force_refresh: bool = False,
|
||||
cache_timeout: int | None = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Apply cache control parameters to a query context.
|
||||
|
||||
Args:
|
||||
query_context: Query context dictionary
|
||||
use_cache: Whether to use cache
|
||||
force_refresh: Whether to force refresh
|
||||
cache_timeout: Cache timeout override
|
||||
|
||||
Returns:
|
||||
Modified query context with cache control applied
|
||||
"""
|
||||
if not use_cache or force_refresh:
|
||||
query_context["force"] = True
|
||||
|
||||
if cache_timeout is not None:
|
||||
# Apply to all queries in the context
|
||||
for query in query_context.get("queries", []):
|
||||
query["cache_timeout"] = cache_timeout
|
||||
|
||||
return query_context
|
||||
|
||||
|
||||
def should_use_metadata_cache(
|
||||
use_cache: bool = True,
|
||||
refresh_metadata: bool = False,
|
||||
) -> bool:
|
||||
"""
|
||||
Determine whether to use metadata cache based on cache control parameters.
|
||||
|
||||
Args:
|
||||
use_cache: Whether to use cache
|
||||
refresh_metadata: Whether to refresh metadata
|
||||
|
||||
Returns:
|
||||
True if metadata cache should be used
|
||||
"""
|
||||
return use_cache and not refresh_metadata
|
||||
|
||||
|
||||
def get_cache_key_info(cache_key: str | None) -> str | None:
|
||||
"""
|
||||
Get truncated cache key for debugging purposes.
|
||||
|
||||
Args:
|
||||
cache_key: Full cache key
|
||||
|
||||
Returns:
|
||||
Truncated cache key or None
|
||||
"""
|
||||
if not cache_key:
|
||||
return None
|
||||
|
||||
# Truncate long cache keys for readability
|
||||
if len(cache_key) > 50:
|
||||
return cache_key[:47] + "..."
|
||||
|
||||
return cache_key
|
||||
369
superset/mcp_service/utils/error_builder.py
Normal file
369
superset/mcp_service/utils/error_builder.py
Normal file
@@ -0,0 +1,369 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Unified error builder for chart operations.
|
||||
Consolidates error handling logic from multiple files.
|
||||
"""
|
||||
|
||||
import html
|
||||
import logging
|
||||
import re
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from superset.mcp_service.common.error_schemas import (
|
||||
ChartGenerationError,
|
||||
ValidationError,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _sanitize_user_input(value: Any) -> str:
|
||||
"""Sanitize user input to prevent XSS and injection attacks in error messages."""
|
||||
if value is None:
|
||||
return "None"
|
||||
|
||||
# Convert to string and limit length to prevent ReDoS attacks
|
||||
str_value = str(value)
|
||||
if len(str_value) > 200:
|
||||
str_value = str_value[:200] + "...[truncated]"
|
||||
|
||||
# HTML escape to prevent XSS
|
||||
str_value = html.escape(str_value)
|
||||
|
||||
# Check for dangerous HTML tags with simple substring checks
|
||||
html_tags = ["<script", "</script>"]
|
||||
str_lower = str_value.lower()
|
||||
for tag in html_tags:
|
||||
if tag in str_lower:
|
||||
str_value = "[FILTERED]"
|
||||
break
|
||||
|
||||
# Check for dangerous URL schemes using regex with word boundaries
|
||||
# This ensures we match actual URL schemes, not arbitrary substrings
|
||||
dangerous_url_patterns = [
|
||||
r"\b(javascript|vbscript|data):", # URL schemes
|
||||
r"on\w+\s*=", # Event handlers
|
||||
]
|
||||
for pattern in dangerous_url_patterns:
|
||||
if re.search(pattern, str_value, re.IGNORECASE):
|
||||
str_value = "[FILTERED]"
|
||||
break
|
||||
|
||||
return str_value
|
||||
|
||||
|
||||
def _sanitize_template_vars(vars_dict: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Sanitize all variables before template formatting."""
|
||||
sanitized = {}
|
||||
for key, value in vars_dict.items():
|
||||
# Only sanitize string-like values that could contain user input
|
||||
if isinstance(value, (str, int, float)) or value is None:
|
||||
sanitized[key] = _sanitize_user_input(value)
|
||||
elif isinstance(value, (list, tuple)):
|
||||
# Sanitize lists of strings
|
||||
sanitized[key] = ", ".join(
|
||||
[_sanitize_user_input(item) for item in value[:10]]
|
||||
) # Limit list size and convert to string
|
||||
else:
|
||||
# For other types, convert to string and sanitize
|
||||
sanitized[key] = _sanitize_user_input(value)
|
||||
return sanitized
|
||||
|
||||
|
||||
class ChartErrorBuilder:
|
||||
"""Unified error builder for consistent error messages across chart operations."""
|
||||
|
||||
# Error templates organized by category
|
||||
TEMPLATES = {
|
||||
# Validation errors
|
||||
"missing_field": {
|
||||
"message": "Missing required field: {field}",
|
||||
"details": "{field_description}",
|
||||
"suggestions": [
|
||||
"Add the '{field}' field to your configuration",
|
||||
"Check the API documentation for required fields",
|
||||
"{specific_suggestion}",
|
||||
],
|
||||
},
|
||||
"invalid_type": {
|
||||
"message": "Invalid type for field '{field}'",
|
||||
"details": "Expected {expected_type}, got {actual_type}",
|
||||
"suggestions": [
|
||||
"Change '{field}' to be a {expected_type}",
|
||||
"Example: {example}",
|
||||
],
|
||||
},
|
||||
"invalid_value": {
|
||||
"message": "Invalid value for '{field}'",
|
||||
"details": "Value '{value}' is not allowed. {reason}",
|
||||
"suggestions": [
|
||||
"Use one of the allowed values: {allowed_values}",
|
||||
"{specific_suggestion}",
|
||||
],
|
||||
},
|
||||
# Dataset errors
|
||||
"dataset_not_found": {
|
||||
"message": "Dataset not found: {dataset_id}",
|
||||
"details": "No dataset found with identifier '{dataset_id}'. Please "
|
||||
"verify the dataset ID or UUID is correct.",
|
||||
"suggestions": [
|
||||
"Check that the dataset ID is correct",
|
||||
"Verify you have access to this dataset",
|
||||
"Use the list_datasets tool to find available datasets",
|
||||
],
|
||||
},
|
||||
"column_not_found": {
|
||||
"message": "Column '{column}' not found in dataset",
|
||||
"details": "The column '{column}' does not exist in the dataset schema",
|
||||
"suggestions": [
|
||||
"Check column name spelling and case sensitivity",
|
||||
"Use get_dataset_info to see available columns",
|
||||
"Did you mean: {suggestions}?",
|
||||
],
|
||||
},
|
||||
# Runtime errors
|
||||
"empty_result": {
|
||||
"message": "Query would return no data",
|
||||
"details": "{reason}",
|
||||
"suggestions": [
|
||||
"Check your filter conditions",
|
||||
"Verify the data exists for your criteria",
|
||||
"Try broader filter values or remove some filters",
|
||||
],
|
||||
},
|
||||
"performance_warning": {
|
||||
"message": "Configuration may cause performance issues",
|
||||
"details": "{reason}",
|
||||
"suggestions": [
|
||||
"Consider adding filters to limit data",
|
||||
"Use aggregations to reduce data volume",
|
||||
"{specific_suggestion}",
|
||||
],
|
||||
},
|
||||
# Chart-specific errors
|
||||
"invalid_chart_type": {
|
||||
"message": "Invalid chart type: '{chart_type}'",
|
||||
"details": "Chart type must be either 'xy' or 'table'",
|
||||
"suggestions": [
|
||||
"Use 'chart_type': 'xy' for line, bar, area, or scatter charts",
|
||||
"Use 'chart_type': 'table' for tabular data display",
|
||||
],
|
||||
},
|
||||
"incompatible_configuration": {
|
||||
"message": "Chart configuration incompatible with data",
|
||||
"details": "{reason}",
|
||||
"suggestions": [
|
||||
"{primary_suggestion}",
|
||||
"Consider using a different chart type",
|
||||
"Modify your data selection or aggregation",
|
||||
],
|
||||
},
|
||||
# Chart generation errors
|
||||
"generation_failed": {
|
||||
"message": "Chart generation failed: {reason}",
|
||||
"details": "Failed to create {chart_type} chart for dataset {dataset_id}. "
|
||||
"{reason}",
|
||||
"suggestions": [
|
||||
"Check that the dataset exists and is accessible",
|
||||
"Verify chart configuration is valid for the selected chart type",
|
||||
"Ensure all referenced columns exist in the dataset",
|
||||
"Check Superset logs for detailed error information",
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def build_error(
|
||||
cls,
|
||||
error_type: str,
|
||||
template_key: str,
|
||||
template_vars: Dict[str, Any] | None = None,
|
||||
custom_suggestions: List[str] | None = None,
|
||||
error_code: str | None = None,
|
||||
validation_errors: List[ValidationError] | None = None,
|
||||
) -> ChartGenerationError:
|
||||
"""
|
||||
Build a standardized error using templates.
|
||||
|
||||
Args:
|
||||
error_type: Type of error for categorization
|
||||
template_key: Key to error template
|
||||
template_vars: Variables to format into template
|
||||
custom_suggestions: Additional suggestions to append
|
||||
error_code: Optional error code
|
||||
validation_errors: Optional list of validation errors
|
||||
|
||||
Returns:
|
||||
ChartGenerationError with formatted message
|
||||
"""
|
||||
template = cls.TEMPLATES.get(template_key, {})
|
||||
# SECURITY FIX: Sanitize template variables to prevent XSS/injection
|
||||
vars_dict = _sanitize_template_vars(template_vars or {})
|
||||
|
||||
message = cls._format_message(template, vars_dict)
|
||||
details = cls._format_details(template, vars_dict)
|
||||
suggestions = cls._format_suggestions(template, vars_dict, custom_suggestions)
|
||||
error_code = cls._generate_error_code(error_code, template_key)
|
||||
|
||||
return ChartGenerationError(
|
||||
error_type=error_type,
|
||||
message=message,
|
||||
details=details,
|
||||
suggestions=suggestions,
|
||||
error_code=error_code,
|
||||
validation_errors=validation_errors or [],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _format_message(
|
||||
cls, template: Dict[str, Any], vars_dict: Dict[str, Any]
|
||||
) -> str:
|
||||
"""Format the error message from template."""
|
||||
message_raw = template.get("message", "An error occurred")
|
||||
message: str = (
|
||||
" ".join(message_raw) if isinstance(message_raw, list) else str(message_raw)
|
||||
)
|
||||
if vars_dict:
|
||||
try:
|
||||
# SECURITY FIX: vars_dict is already sanitized by caller
|
||||
message = message.format(**vars_dict)
|
||||
except (KeyError, ValueError, TypeError) as e:
|
||||
logger.warning("Template formatting failed: %s", e)
|
||||
# Return safe fallback message
|
||||
message = "An error occurred during chart operation"
|
||||
return message
|
||||
|
||||
@classmethod
|
||||
def _format_details(
|
||||
cls, template: Dict[str, Any], vars_dict: Dict[str, Any]
|
||||
) -> str:
|
||||
"""Format the error details from template."""
|
||||
details_raw = template.get("details", "")
|
||||
details: str = (
|
||||
" ".join(details_raw) if isinstance(details_raw, list) else str(details_raw)
|
||||
)
|
||||
if vars_dict and details:
|
||||
try:
|
||||
# SECURITY FIX: vars_dict is already sanitized by caller
|
||||
details = details.format(**vars_dict)
|
||||
except (KeyError, ValueError, TypeError) as e:
|
||||
logger.warning("Template formatting failed: %s", e)
|
||||
# Return safe fallback
|
||||
details = "Additional error details unavailable"
|
||||
return details
|
||||
|
||||
@classmethod
|
||||
def _format_suggestions(
|
||||
cls,
|
||||
template: Dict[str, Any],
|
||||
vars_dict: Dict[str, Any],
|
||||
custom_suggestions: List[str] | None,
|
||||
) -> List[str]:
|
||||
"""Format suggestions from template and add custom ones."""
|
||||
suggestions = []
|
||||
for suggestion in template.get("suggestions", []):
|
||||
if vars_dict and "{" in suggestion:
|
||||
try:
|
||||
# SECURITY FIX: vars_dict is already sanitized by caller
|
||||
suggestion = suggestion.format(**vars_dict)
|
||||
if suggestion and suggestion != "None":
|
||||
suggestions.append(suggestion)
|
||||
except (KeyError, ValueError, TypeError):
|
||||
# Skip malformed suggestions rather than exposing errors
|
||||
continue
|
||||
else:
|
||||
suggestions.append(suggestion)
|
||||
|
||||
if custom_suggestions:
|
||||
# SECURITY FIX: Sanitize custom suggestions too
|
||||
sanitized_custom = [
|
||||
_sanitize_user_input(s) for s in custom_suggestions[:5]
|
||||
] # Limit count
|
||||
suggestions.extend(sanitized_custom)
|
||||
|
||||
return suggestions[:10] # Limit total suggestions to prevent response bloat
|
||||
|
||||
@classmethod
|
||||
def _generate_error_code(cls, error_code: str | None, template_key: str) -> str:
|
||||
"""Generate error code if not provided."""
|
||||
if error_code:
|
||||
return error_code
|
||||
return f"CHART_{template_key.upper()}"
|
||||
|
||||
@classmethod
|
||||
def missing_field_error(
|
||||
cls,
|
||||
field: str,
|
||||
field_description: str,
|
||||
specific_suggestion: str | None = None,
|
||||
) -> ChartGenerationError:
|
||||
"""Build a missing field error."""
|
||||
return cls.build_error(
|
||||
error_type="missing_field",
|
||||
template_key="missing_field",
|
||||
template_vars={
|
||||
"field": field,
|
||||
"field_description": field_description,
|
||||
"specific_suggestion": specific_suggestion
|
||||
or f"Add '{field}' to your request",
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def invalid_type_error(
|
||||
cls,
|
||||
field: str,
|
||||
expected_type: str,
|
||||
actual_type: str,
|
||||
example: str | None = None,
|
||||
) -> ChartGenerationError:
|
||||
"""Build an invalid type error."""
|
||||
return cls.build_error(
|
||||
error_type="invalid_type",
|
||||
template_key="invalid_type",
|
||||
template_vars={
|
||||
"field": field,
|
||||
"expected_type": expected_type,
|
||||
"actual_type": actual_type,
|
||||
"example": example or f"'{field}': <{expected_type}>",
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def column_not_found_error(
|
||||
cls, column: str, suggestions: List[str] | None = None
|
||||
) -> ChartGenerationError:
|
||||
"""Build a column not found error."""
|
||||
suggestion_text = (
|
||||
", ".join(suggestions[:3]) if suggestions else "Check available columns"
|
||||
)
|
||||
return cls.build_error(
|
||||
error_type="column_not_found",
|
||||
template_key="column_not_found",
|
||||
template_vars={"column": column, "suggestions": suggestion_text},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def dataset_not_found_error(cls, dataset_id: Any) -> ChartGenerationError:
|
||||
"""Build a dataset not found error."""
|
||||
return cls.build_error(
|
||||
error_type="dataset_not_found",
|
||||
template_key="dataset_not_found",
|
||||
template_vars={"dataset_id": str(dataset_id)},
|
||||
)
|
||||
313
superset/mcp_service/utils/permissions_utils.py
Normal file
313
superset/mcp_service/utils/permissions_utils.py
Normal file
@@ -0,0 +1,313 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Field-level permissions utilities for MCP service.
|
||||
Provides functionality to filter sensitive data based on user permissions.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, List, Optional, Set
|
||||
|
||||
from flask_appbuilder.security.sqla.models import User
|
||||
from pydantic import BaseModel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Define sensitive fields by object type
|
||||
SENSITIVE_FIELDS = {
|
||||
"dataset": {
|
||||
"sql", # Raw SQL queries may contain sensitive logic
|
||||
"extra", # May contain connection strings or credentials
|
||||
"database_id", # Internal database references
|
||||
"changed_by_fk", # Internal user references
|
||||
"created_by_fk", # Internal user references
|
||||
},
|
||||
"chart": {
|
||||
"query_context", # May contain sensitive filters or parameters
|
||||
"cache_key", # Internal cache references
|
||||
"changed_by_fk", # Internal user references
|
||||
"created_by_fk", # Internal user references
|
||||
},
|
||||
"dashboard": {
|
||||
"json_metadata", # May contain sensitive configuration
|
||||
"position_json", # Internal layout data
|
||||
"css", # May contain sensitive styling info
|
||||
"changed_by_fk", # Internal user references
|
||||
"created_by_fk", # Internal user references
|
||||
},
|
||||
"common": {
|
||||
"uuid", # Internal identifiers (keep for some use cases)
|
||||
"changed_by_fk", # Internal user references
|
||||
"created_by_fk", # Internal user references
|
||||
},
|
||||
}
|
||||
|
||||
# Permissions required to access sensitive fields
|
||||
SENSITIVE_FIELD_PERMISSIONS = {
|
||||
"sql": "can_sql_json", # SQL Lab permissions
|
||||
"extra": "can_this_form_get", # Advanced form permissions
|
||||
"database_id": "can_this_form_get", # Database access permissions
|
||||
"query_context": "can_explore_json", # Explore permissions
|
||||
"cache_key": "can_warm_up_cache", # Cache management permissions
|
||||
"json_metadata": "can_this_form_get", # Advanced dashboard permissions
|
||||
"position_json": "can_this_form_get", # Dashboard edit permissions
|
||||
"css": "can_this_form_get", # Dashboard styling permissions
|
||||
}
|
||||
|
||||
|
||||
def get_current_user() -> Optional[User]:
|
||||
"""Get the current authenticated user."""
|
||||
try:
|
||||
from flask import g
|
||||
|
||||
return getattr(g, "user", None)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def user_has_permission(
|
||||
user: Optional[User], permission: str, resource: Optional[str] = None
|
||||
) -> bool:
|
||||
"""
|
||||
Check if user has a specific permission.
|
||||
|
||||
Args:
|
||||
user: User object or None
|
||||
permission: Permission name (e.g., 'can_sql_json')
|
||||
resource: Resource name (e.g., 'Superset', 'Chart', etc.)
|
||||
|
||||
Returns:
|
||||
True if user has permission, False otherwise
|
||||
"""
|
||||
if not user:
|
||||
return False
|
||||
|
||||
try:
|
||||
# Check if user is admin (has all permissions)
|
||||
if hasattr(user, "roles"):
|
||||
for role in user.roles:
|
||||
if role.name in ("Admin", "admin"):
|
||||
return True
|
||||
|
||||
# Check specific permission
|
||||
from superset import security_manager
|
||||
|
||||
if resource:
|
||||
return security_manager.has_access(permission, resource, user)
|
||||
else:
|
||||
# Check if user has permission on any resource
|
||||
for pvm in user.get_permissions():
|
||||
if pvm.permission.name == permission:
|
||||
return True
|
||||
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Error checking permission %s for user %s: %s", permission, user, e
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def get_allowed_fields(
|
||||
object_type: str,
|
||||
user: Optional[User] = None,
|
||||
requested_fields: Optional[List[str]] = None,
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Get the set of fields that the user is allowed to access for a given object type.
|
||||
|
||||
Args:
|
||||
object_type: Type of object ('dataset', 'chart', 'dashboard')
|
||||
user: User object (if None, will try to get current user)
|
||||
requested_fields: List of fields requested (if None, all allowed fields)
|
||||
|
||||
Returns:
|
||||
Set of allowed field names
|
||||
"""
|
||||
if not user:
|
||||
user = get_current_user()
|
||||
|
||||
# Get sensitive fields for this object type
|
||||
sensitive_fields = SENSITIVE_FIELDS.get(object_type, set())
|
||||
sensitive_fields.update(SENSITIVE_FIELDS.get("common", set()))
|
||||
|
||||
# If no user, only allow non-sensitive fields
|
||||
if not user:
|
||||
if requested_fields:
|
||||
return set(requested_fields) - sensitive_fields
|
||||
else:
|
||||
# Return empty set - caller should use default safe fields
|
||||
return set()
|
||||
|
||||
# Check permissions for sensitive fields
|
||||
allowed_fields = set()
|
||||
|
||||
if requested_fields:
|
||||
for field in requested_fields:
|
||||
if field not in sensitive_fields:
|
||||
# Non-sensitive field, always allowed
|
||||
allowed_fields.add(field)
|
||||
else:
|
||||
# Check if user has permission for this sensitive field
|
||||
required_permission = SENSITIVE_FIELD_PERMISSIONS.get(field)
|
||||
if required_permission and user_has_permission(
|
||||
user, required_permission
|
||||
):
|
||||
allowed_fields.add(field)
|
||||
elif not required_permission:
|
||||
# No specific permission required, but still sensitive
|
||||
# Allow for authenticated users (basic sensitivity)
|
||||
allowed_fields.add(field)
|
||||
else:
|
||||
# No specific fields requested, return empty set
|
||||
# Caller should specify default fields
|
||||
return set()
|
||||
|
||||
return allowed_fields
|
||||
|
||||
|
||||
def filter_sensitive_data(
|
||||
data: Any,
|
||||
object_type: str,
|
||||
user: Optional[User] = None,
|
||||
allowed_fields: Optional[Set[str]] = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Filter sensitive data from an object based on user permissions.
|
||||
|
||||
Args:
|
||||
data: Data to filter (dict, Pydantic model, or list)
|
||||
object_type: Type of object ('dataset', 'chart', 'dashboard')
|
||||
user: User object (if None, will try to get current user)
|
||||
allowed_fields: Pre-computed allowed fields (optimization)
|
||||
|
||||
Returns:
|
||||
Filtered data with sensitive fields removed
|
||||
"""
|
||||
if not data:
|
||||
return data
|
||||
|
||||
if not user:
|
||||
user = get_current_user()
|
||||
|
||||
# Handle different data types
|
||||
if isinstance(data, list):
|
||||
return [
|
||||
filter_sensitive_data(item, object_type, user, allowed_fields)
|
||||
for item in data
|
||||
]
|
||||
|
||||
if isinstance(data, BaseModel):
|
||||
# Convert Pydantic model to dict for filtering
|
||||
data_dict = data.model_dump()
|
||||
filtered_dict = filter_sensitive_data(
|
||||
data_dict, object_type, user, allowed_fields
|
||||
)
|
||||
# Return as dict since we can't easily reconstruct the Pydantic model
|
||||
return filtered_dict
|
||||
|
||||
if not isinstance(data, dict):
|
||||
# Not a dict-like object, return as-is
|
||||
return data
|
||||
|
||||
# Get allowed fields if not provided
|
||||
if allowed_fields is None:
|
||||
requested_fields = list(data.keys())
|
||||
allowed_fields = get_allowed_fields(object_type, user, requested_fields)
|
||||
|
||||
# Filter the dictionary
|
||||
filtered_data = {}
|
||||
for key, value in data.items():
|
||||
if key in allowed_fields:
|
||||
filtered_data[key] = value
|
||||
else:
|
||||
# Log when we filter out sensitive data
|
||||
logger.debug(
|
||||
"Filtered sensitive field '%s' for object type '%s'", key, object_type
|
||||
)
|
||||
|
||||
return filtered_data
|
||||
|
||||
|
||||
def apply_field_permissions_to_columns(
|
||||
columns: List[str], object_type: str, user: Optional[User] = None
|
||||
) -> List[str]:
|
||||
"""
|
||||
Filter a list of column names based on field-level permissions.
|
||||
|
||||
Args:
|
||||
columns: List of column names to filter
|
||||
object_type: Type of object ('dataset', 'chart', 'dashboard')
|
||||
user: User object (if None, will try to get current user)
|
||||
|
||||
Returns:
|
||||
Filtered list of allowed column names
|
||||
"""
|
||||
allowed_fields = get_allowed_fields(object_type, user, columns)
|
||||
return [col for col in columns if col in allowed_fields]
|
||||
|
||||
|
||||
class PermissionAwareSerializer:
|
||||
"""
|
||||
Wrapper for serializers that automatically applies field-level permissions.
|
||||
"""
|
||||
|
||||
def __init__(self, object_type: str, base_serializer: Any):
|
||||
self.object_type = object_type
|
||||
self.base_serializer = base_serializer
|
||||
|
||||
def serialize(
|
||||
self, obj: Any, columns: List[str], user: Optional[User] = None
|
||||
) -> Any:
|
||||
"""
|
||||
Serialize object with field-level permissions applied.
|
||||
|
||||
Args:
|
||||
obj: Object to serialize
|
||||
columns: Requested columns
|
||||
user: User object for permission checking
|
||||
|
||||
Returns:
|
||||
Serialized object with sensitive fields filtered
|
||||
"""
|
||||
# Filter columns based on permissions
|
||||
allowed_columns = apply_field_permissions_to_columns(
|
||||
columns, self.object_type, user
|
||||
)
|
||||
|
||||
# Use base serializer with filtered columns
|
||||
serialized = self.base_serializer(obj, allowed_columns)
|
||||
|
||||
# Apply additional filtering to the serialized result
|
||||
return filter_sensitive_data(serialized, self.object_type, user)
|
||||
|
||||
|
||||
# Convenience functions for common object types
|
||||
def filter_dataset_data(data: Any, user: Optional[User] = None) -> Any:
|
||||
"""Filter sensitive data from dataset objects."""
|
||||
return filter_sensitive_data(data, "dataset", user)
|
||||
|
||||
|
||||
def filter_chart_data(data: Any, user: Optional[User] = None) -> Any:
|
||||
"""Filter sensitive data from chart objects."""
|
||||
return filter_sensitive_data(data, "chart", user)
|
||||
|
||||
|
||||
def filter_dashboard_data(data: Any, user: Optional[User] = None) -> Any:
|
||||
"""Filter sensitive data from dashboard objects."""
|
||||
return filter_sensitive_data(data, "dashboard", user)
|
||||
341
superset/mcp_service/utils/retry_utils.py
Normal file
341
superset/mcp_service/utils/retry_utils.py
Normal file
@@ -0,0 +1,341 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Retry utilities for handling transient failures in MCP service operations.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import logging
|
||||
import secrets
|
||||
import time
|
||||
from typing import Any, Callable, Type, TypeVar
|
||||
|
||||
from sqlalchemy.exc import OperationalError, TimeoutError
|
||||
from starlette.exceptions import HTTPException
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
# Default retryable exceptions
|
||||
DEFAULT_RETRYABLE_EXCEPTIONS = (
|
||||
OperationalError, # Database connection issues
|
||||
TimeoutError, # Database timeouts
|
||||
ConnectionError, # Network issues
|
||||
OSError, # File system issues (for screenshots)
|
||||
)
|
||||
|
||||
|
||||
def exponential_backoff(
|
||||
attempt: int, base_delay: float = 1.0, max_delay: float = 60.0, jitter: bool = True
|
||||
) -> float:
|
||||
"""
|
||||
Calculate exponential backoff delay with optional jitter.
|
||||
|
||||
Args:
|
||||
attempt: Current attempt number (0-based)
|
||||
base_delay: Base delay in seconds
|
||||
max_delay: Maximum delay in seconds
|
||||
jitter: Whether to add random jitter to avoid thundering herd
|
||||
|
||||
Returns:
|
||||
Delay in seconds
|
||||
"""
|
||||
delay = base_delay * (2**attempt)
|
||||
delay = min(delay, max_delay)
|
||||
|
||||
if jitter:
|
||||
# Add up to 25% jitter using cryptographically secure random
|
||||
jitter_amount = delay * 0.25
|
||||
random_gen = secrets.SystemRandom()
|
||||
delay += random_gen.uniform(-jitter_amount, jitter_amount)
|
||||
|
||||
return max(0, delay)
|
||||
|
||||
|
||||
def retry_on_exception(
|
||||
max_attempts: int = 3,
|
||||
base_delay: float = 1.0,
|
||||
max_delay: float = 60.0,
|
||||
exceptions: tuple[Type[Exception], ...] = DEFAULT_RETRYABLE_EXCEPTIONS,
|
||||
jitter: bool = True,
|
||||
) -> Callable[[Callable[..., T]], Callable[..., T]]:
|
||||
"""
|
||||
Decorator to retry function calls on specific exceptions with exponential backoff.
|
||||
|
||||
Args:
|
||||
max_attempts: Maximum number of attempts (including initial attempt)
|
||||
base_delay: Base delay in seconds for exponential backoff
|
||||
max_delay: Maximum delay in seconds between retries
|
||||
exceptions: Tuple of exception types to retry on
|
||||
jitter: Whether to add random jitter to backoff delays
|
||||
|
||||
Returns:
|
||||
Decorated function
|
||||
"""
|
||||
|
||||
def decorator(func: Callable[..., T]) -> Callable[..., T]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args: Any, **kwargs: Any) -> T:
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_attempts):
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
except exceptions as e: # noqa: PERF203
|
||||
last_exception = e
|
||||
|
||||
if attempt == max_attempts - 1:
|
||||
# Last attempt, don't wait
|
||||
break
|
||||
|
||||
delay = exponential_backoff(attempt, base_delay, max_delay, jitter)
|
||||
|
||||
logger.warning(
|
||||
"Attempt %s/%s failed for %s: %s. Retrying in %.2fs...",
|
||||
attempt + 1,
|
||||
max_attempts,
|
||||
func.__name__,
|
||||
e,
|
||||
delay,
|
||||
)
|
||||
|
||||
time.sleep(delay)
|
||||
except Exception as e:
|
||||
# Non-retryable exception, fail immediately
|
||||
logger.error("Non-retryable exception in %s: %s", func.__name__, e)
|
||||
raise
|
||||
|
||||
# All attempts failed
|
||||
if last_exception is not None:
|
||||
logger.error(
|
||||
"All %s attempts failed for %s: %s",
|
||||
max_attempts,
|
||||
func.__name__,
|
||||
last_exception,
|
||||
)
|
||||
raise last_exception
|
||||
logger.error("All %s attempts failed for %s", max_attempts, func.__name__)
|
||||
raise RuntimeError(
|
||||
f"All {max_attempts} attempts failed for {func.__name__}"
|
||||
)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def async_retry_on_exception(
|
||||
max_attempts: int = 3,
|
||||
base_delay: float = 1.0,
|
||||
max_delay: float = 60.0,
|
||||
exceptions: tuple[Type[Exception], ...] = DEFAULT_RETRYABLE_EXCEPTIONS,
|
||||
jitter: bool = True,
|
||||
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
|
||||
"""
|
||||
Async version of retry_on_exception decorator.
|
||||
"""
|
||||
|
||||
def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
|
||||
@functools.wraps(func)
|
||||
async def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_attempts):
|
||||
try:
|
||||
return await func(*args, **kwargs)
|
||||
except exceptions as e: # noqa: PERF203
|
||||
last_exception = e
|
||||
|
||||
if attempt == max_attempts - 1:
|
||||
# Last attempt, don't wait
|
||||
break
|
||||
|
||||
delay = exponential_backoff(attempt, base_delay, max_delay, jitter)
|
||||
|
||||
logger.warning(
|
||||
"Attempt %s/%s failed for %s: %s. Retrying in %.2fs...",
|
||||
attempt + 1,
|
||||
max_attempts,
|
||||
func.__name__,
|
||||
e,
|
||||
delay,
|
||||
)
|
||||
|
||||
await asyncio.sleep(delay)
|
||||
except Exception as e:
|
||||
# Non-retryable exception, fail immediately
|
||||
logger.error("Non-retryable exception in %s: %s", func.__name__, e)
|
||||
raise
|
||||
|
||||
# All attempts failed
|
||||
if last_exception is not None:
|
||||
logger.error(
|
||||
"All %s attempts failed for %s: %s",
|
||||
max_attempts,
|
||||
func.__name__,
|
||||
last_exception,
|
||||
)
|
||||
raise last_exception
|
||||
logger.error("All %s attempts failed for %s", max_attempts, func.__name__)
|
||||
raise RuntimeError(
|
||||
f"All {max_attempts} attempts failed for {func.__name__}"
|
||||
)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
class RetryableOperation:
|
||||
"""
|
||||
Context manager for retryable operations with custom logic.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operation_name: str,
|
||||
max_attempts: int = 3,
|
||||
base_delay: float = 1.0,
|
||||
max_delay: float = 60.0,
|
||||
exceptions: tuple[Type[Exception], ...] = DEFAULT_RETRYABLE_EXCEPTIONS,
|
||||
jitter: bool = True,
|
||||
) -> None:
|
||||
self.operation_name = operation_name
|
||||
self.max_attempts = max_attempts
|
||||
self.base_delay = base_delay
|
||||
self.max_delay = max_delay
|
||||
self.exceptions = exceptions
|
||||
self.jitter = jitter
|
||||
self.current_attempt = 0
|
||||
self.last_exception: Exception | None = None
|
||||
|
||||
def __enter__(self) -> "RetryableOperation":
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: Type[Exception] | None,
|
||||
exc_val: Exception | None,
|
||||
exc_tb: Any,
|
||||
) -> bool:
|
||||
if exc_type is None:
|
||||
# No exception, operation succeeded
|
||||
return False
|
||||
|
||||
if not issubclass(exc_type, self.exceptions):
|
||||
# Non-retryable exception
|
||||
logger.error(
|
||||
"Non-retryable exception in %s: %s", self.operation_name, exc_val
|
||||
)
|
||||
return False
|
||||
|
||||
self.last_exception = exc_val
|
||||
self.current_attempt += 1
|
||||
|
||||
if self.current_attempt >= self.max_attempts:
|
||||
# Max attempts reached
|
||||
logger.error(
|
||||
"All %s attempts failed for %s. ",
|
||||
self.max_attempts,
|
||||
self.operation_name,
|
||||
)
|
||||
return False
|
||||
|
||||
# Calculate delay and wait
|
||||
delay = exponential_backoff(
|
||||
self.current_attempt - 1, self.base_delay, self.max_delay, self.jitter
|
||||
)
|
||||
|
||||
logger.warning(
|
||||
"Attempt %s/%s failed for %s: %s. Retrying in %.2fs...",
|
||||
self.current_attempt,
|
||||
self.max_attempts,
|
||||
self.operation_name,
|
||||
exc_val,
|
||||
delay,
|
||||
)
|
||||
|
||||
time.sleep(delay)
|
||||
return True # Suppress the exception and continue
|
||||
|
||||
def should_retry(self) -> bool:
|
||||
"""Check if we should continue retrying"""
|
||||
return self.current_attempt < self.max_attempts
|
||||
|
||||
|
||||
# Convenience functions for common operations
|
||||
def retry_database_operation(
|
||||
func: Callable[..., T], *args: Any, max_attempts: int = 3, **kwargs: Any
|
||||
) -> T:
|
||||
"""
|
||||
Execute a database operation with retry logic.
|
||||
"""
|
||||
|
||||
@retry_on_exception(
|
||||
max_attempts=max_attempts,
|
||||
exceptions=(OperationalError, TimeoutError),
|
||||
base_delay=0.5,
|
||||
max_delay=30.0,
|
||||
)
|
||||
def _wrapped() -> T:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return _wrapped()
|
||||
|
||||
|
||||
async def async_retry_database_operation(
|
||||
func: Callable[..., Any], *args: Any, max_attempts: int = 3, **kwargs: Any
|
||||
) -> Any:
|
||||
"""
|
||||
Execute an async database operation with retry logic.
|
||||
"""
|
||||
|
||||
@async_retry_on_exception(
|
||||
max_attempts=max_attempts,
|
||||
exceptions=(OperationalError, TimeoutError),
|
||||
base_delay=0.5,
|
||||
max_delay=30.0,
|
||||
)
|
||||
async def _wrapped() -> Any:
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
return await _wrapped()
|
||||
|
||||
|
||||
def retry_screenshot_operation(
|
||||
func: Callable[..., T],
|
||||
*args: Any,
|
||||
max_attempts: int = 2, # Screenshots are expensive, fewer retries
|
||||
**kwargs: Any,
|
||||
) -> T:
|
||||
"""
|
||||
Execute a screenshot operation with retry logic.
|
||||
"""
|
||||
|
||||
@retry_on_exception(
|
||||
max_attempts=max_attempts,
|
||||
exceptions=(OSError, ConnectionError, HTTPException),
|
||||
base_delay=2.0, # Longer initial delay for screenshots
|
||||
max_delay=30.0,
|
||||
)
|
||||
def _wrapped() -> T:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return _wrapped()
|
||||
128
superset/mcp_service/utils/url_utils.py
Normal file
128
superset/mcp_service/utils/url_utils.py
Normal file
@@ -0,0 +1,128 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
URL utilities for MCP service
|
||||
"""
|
||||
|
||||
from flask import current_app
|
||||
|
||||
|
||||
def get_superset_base_url() -> str:
|
||||
"""
|
||||
Get the Superset base URL from configuration.
|
||||
|
||||
Returns:
|
||||
Base URL for Superset web server (e.g., "http://localhost:9001")
|
||||
"""
|
||||
# Default fallback to localhost:9001
|
||||
default_url = "http://localhost:9001"
|
||||
|
||||
try:
|
||||
# Try to get from configuration
|
||||
config = current_app.config
|
||||
|
||||
# Check for SUPERSET_WEBSERVER_ADDRESS first
|
||||
webserver_address = config.get("SUPERSET_WEBSERVER_ADDRESS")
|
||||
if webserver_address:
|
||||
return webserver_address
|
||||
|
||||
# Fallback to other potential config keys
|
||||
public_role_like_gamma = config.get("PUBLIC_ROLE_LIKE_GAMMA", False)
|
||||
if public_role_like_gamma:
|
||||
# If public access is enabled, might be on a different host
|
||||
webserver_protocol = config.get("ENABLE_PROXY_FIX", False)
|
||||
protocol = "https" if webserver_protocol else "http"
|
||||
host = config.get("WEBSERVER_HOST", "localhost")
|
||||
port = config.get("WEBSERVER_PORT", 9001)
|
||||
return f"{protocol}://{host}:{port}"
|
||||
|
||||
return default_url
|
||||
|
||||
except Exception:
|
||||
# If we can't access Flask config (e.g., outside app context),
|
||||
# return default
|
||||
return default_url
|
||||
|
||||
|
||||
def get_mcp_service_url() -> str:
|
||||
"""
|
||||
Get the MCP service base URL where screenshot endpoints are served.
|
||||
|
||||
The MCP service auto-detects its own host and port since it's running
|
||||
this code. Falls back to explicit configuration or default port.
|
||||
|
||||
Returns:
|
||||
Base URL for MCP service (always independent of Superset URL)
|
||||
"""
|
||||
try:
|
||||
# Try to auto-detect from Flask request context
|
||||
from flask import request
|
||||
|
||||
if request:
|
||||
# Get the host and port from the current request
|
||||
scheme = request.scheme # http or https
|
||||
host = request.host # includes port if non-standard
|
||||
return f"{scheme}://{host}"
|
||||
|
||||
except (RuntimeError, AttributeError):
|
||||
# Not in request context or Flask not available
|
||||
pass
|
||||
|
||||
try:
|
||||
# Check for explicit MCP_SERVICE_URL in config
|
||||
config = current_app.config
|
||||
mcp_service_url = config.get("MCP_SERVICE_URL")
|
||||
if mcp_service_url:
|
||||
return mcp_service_url
|
||||
|
||||
except Exception as e:
|
||||
# Log and fall back if config access fails
|
||||
import logging
|
||||
|
||||
logging.getLogger(__name__).debug("Config access failed: %s", e)
|
||||
|
||||
# Always fallback to MCP service default port (never use Superset URL)
|
||||
return "http://localhost:5008"
|
||||
|
||||
|
||||
def get_chart_screenshot_url(chart_id: int | str) -> str:
|
||||
"""
|
||||
Generate a screenshot URL for a chart using the MCP service.
|
||||
|
||||
Args:
|
||||
chart_id: Chart ID (numeric or string)
|
||||
|
||||
Returns:
|
||||
Complete URL to the chart screenshot endpoint
|
||||
"""
|
||||
mcp_base = get_mcp_service_url()
|
||||
return f"{mcp_base}/screenshot/chart/{chart_id}.png"
|
||||
|
||||
|
||||
def get_explore_screenshot_url(form_data_key: str) -> str:
|
||||
"""
|
||||
Generate a screenshot URL for an explore view using the MCP service.
|
||||
|
||||
Args:
|
||||
form_data_key: Form data key for the explore view
|
||||
|
||||
Returns:
|
||||
Complete URL to the explore screenshot endpoint
|
||||
"""
|
||||
mcp_base = get_mcp_service_url()
|
||||
return f"{mcp_base}/screenshot/explore/{form_data_key}.png"
|
||||
160
tests/unit_tests/mcp_service/chart/test_chart_schemas.py
Normal file
160
tests/unit_tests/mcp_service/chart/test_chart_schemas.py
Normal file
@@ -0,0 +1,160 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Unit tests for MCP chart schema validation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
ColumnRef,
|
||||
TableChartConfig,
|
||||
XYChartConfig,
|
||||
)
|
||||
|
||||
|
||||
class TestTableChartConfig:
|
||||
"""Test TableChartConfig validation."""
|
||||
|
||||
def test_duplicate_labels_rejected(self) -> None:
|
||||
"""Test that TableChartConfig rejects duplicate labels."""
|
||||
with pytest.raises(ValidationError, match="Duplicate column/metric labels"):
|
||||
TableChartConfig(
|
||||
columns=[
|
||||
ColumnRef(name="product_line", label="product_line"),
|
||||
ColumnRef(name="sales", aggregate="SUM", label="product_line"),
|
||||
]
|
||||
)
|
||||
|
||||
def test_unique_labels_accepted(self) -> None:
|
||||
"""Test that TableChartConfig accepts unique labels."""
|
||||
config = TableChartConfig(
|
||||
columns=[
|
||||
ColumnRef(name="product_line", label="Product Line"),
|
||||
ColumnRef(name="sales", aggregate="SUM", label="Total Sales"),
|
||||
]
|
||||
)
|
||||
assert len(config.columns) == 2
|
||||
|
||||
|
||||
class TestXYChartConfig:
|
||||
"""Test XYChartConfig validation."""
|
||||
|
||||
def test_different_labels_accepted(self) -> None:
|
||||
"""Test that different labels for x and y are accepted."""
|
||||
config = XYChartConfig(
|
||||
x=ColumnRef(name="product_line"), # Label: "product_line"
|
||||
y=[
|
||||
ColumnRef(
|
||||
name="product_line", aggregate="COUNT"
|
||||
), # Label: "COUNT(product_line)"
|
||||
],
|
||||
)
|
||||
assert config.x.name == "product_line"
|
||||
assert config.y[0].aggregate == "COUNT"
|
||||
|
||||
def test_explicit_duplicate_label_rejected(self) -> None:
|
||||
"""Test that explicit duplicate labels are rejected."""
|
||||
with pytest.raises(ValidationError, match="Duplicate column/metric labels"):
|
||||
XYChartConfig(
|
||||
x=ColumnRef(name="product_line"),
|
||||
y=[ColumnRef(name="sales", label="product_line")],
|
||||
)
|
||||
|
||||
def test_duplicate_y_axis_labels_rejected(self) -> None:
|
||||
"""Test that duplicate y-axis labels are rejected."""
|
||||
with pytest.raises(ValidationError, match="Duplicate column/metric labels"):
|
||||
XYChartConfig(
|
||||
x=ColumnRef(name="date"),
|
||||
y=[
|
||||
ColumnRef(name="sales", aggregate="SUM"),
|
||||
ColumnRef(name="revenue", aggregate="SUM", label="SUM(sales)"),
|
||||
],
|
||||
)
|
||||
|
||||
def test_unique_labels_accepted(self) -> None:
|
||||
"""Test that unique labels are accepted."""
|
||||
config = XYChartConfig(
|
||||
x=ColumnRef(name="date", label="Order Date"),
|
||||
y=[
|
||||
ColumnRef(name="sales", aggregate="SUM", label="Total Sales"),
|
||||
ColumnRef(name="profit", aggregate="AVG", label="Average Profit"),
|
||||
],
|
||||
)
|
||||
assert len(config.y) == 2
|
||||
|
||||
def test_group_by_duplicate_with_x_rejected(self) -> None:
|
||||
"""Test that group_by conflicts with x are rejected."""
|
||||
with pytest.raises(ValidationError, match="Duplicate column/metric labels"):
|
||||
XYChartConfig(
|
||||
x=ColumnRef(name="region"),
|
||||
y=[ColumnRef(name="sales", aggregate="SUM")],
|
||||
group_by=ColumnRef(name="category", label="region"),
|
||||
)
|
||||
|
||||
def test_realistic_chart_configurations(self) -> None:
|
||||
"""Test realistic chart configurations."""
|
||||
# This should work - COUNT(product_line) != product_line
|
||||
config = XYChartConfig(
|
||||
x=ColumnRef(name="product_line"),
|
||||
y=[
|
||||
ColumnRef(name="product_line", aggregate="COUNT"),
|
||||
ColumnRef(name="sales", aggregate="SUM"),
|
||||
],
|
||||
)
|
||||
assert config.x.name == "product_line"
|
||||
assert len(config.y) == 2
|
||||
|
||||
def test_time_series_chart_configuration(self) -> None:
|
||||
"""Test time series chart configuration works."""
|
||||
# This should PASS now - the chart creation logic fixes duplicates
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="order_date"),
|
||||
y=[
|
||||
ColumnRef(name="sales", aggregate="SUM", label="Total Sales"),
|
||||
ColumnRef(name="price_each", aggregate="AVG", label="Avg Price"),
|
||||
],
|
||||
kind="line",
|
||||
)
|
||||
assert config.x.name == "order_date"
|
||||
assert config.kind == "line"
|
||||
|
||||
def test_time_series_with_custom_x_axis_label(self) -> None:
|
||||
"""Test time series chart with custom x-axis label."""
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="order_date", label="Order Date"),
|
||||
y=[
|
||||
ColumnRef(name="sales", aggregate="SUM", label="Total Sales"),
|
||||
ColumnRef(name="price_each", aggregate="AVG", label="Avg Price"),
|
||||
],
|
||||
kind="line",
|
||||
)
|
||||
assert config.x.label == "Order Date"
|
||||
|
||||
def test_area_chart_configuration(self) -> None:
|
||||
"""Test area chart configuration."""
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="category"),
|
||||
y=[ColumnRef(name="sales", aggregate="SUM", label="Total Sales")],
|
||||
kind="area",
|
||||
)
|
||||
assert config.kind == "area"
|
||||
465
tests/unit_tests/mcp_service/chart/test_chart_utils.py
Normal file
465
tests/unit_tests/mcp_service/chart/test_chart_utils.py
Normal file
@@ -0,0 +1,465 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""Tests for chart utilities module"""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from superset.mcp_service.chart.chart_utils import (
|
||||
create_metric_object,
|
||||
generate_chart_name,
|
||||
generate_explore_link,
|
||||
map_config_to_form_data,
|
||||
map_filter_operator,
|
||||
map_table_config,
|
||||
map_xy_config,
|
||||
)
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
AxisConfig,
|
||||
ColumnRef,
|
||||
FilterConfig,
|
||||
LegendConfig,
|
||||
TableChartConfig,
|
||||
XYChartConfig,
|
||||
)
|
||||
|
||||
|
||||
class TestCreateMetricObject:
|
||||
"""Test create_metric_object function"""
|
||||
|
||||
def test_create_metric_object_with_aggregate(self) -> None:
|
||||
"""Test creating metric object with specified aggregate"""
|
||||
col = ColumnRef(name="revenue", aggregate="SUM", label="Total Revenue")
|
||||
result = create_metric_object(col)
|
||||
|
||||
assert result["aggregate"] == "SUM"
|
||||
assert result["column"]["column_name"] == "revenue"
|
||||
assert result["label"] == "Total Revenue"
|
||||
assert result["optionName"] == "metric_revenue"
|
||||
assert result["expressionType"] == "SIMPLE"
|
||||
|
||||
def test_create_metric_object_default_aggregate(self) -> None:
|
||||
"""Test creating metric object with default aggregate"""
|
||||
col = ColumnRef(name="orders")
|
||||
result = create_metric_object(col)
|
||||
|
||||
assert result["aggregate"] == "SUM"
|
||||
assert result["column"]["column_name"] == "orders"
|
||||
assert result["label"] == "SUM(orders)"
|
||||
assert result["optionName"] == "metric_orders"
|
||||
|
||||
|
||||
class TestMapFilterOperator:
|
||||
"""Test map_filter_operator function"""
|
||||
|
||||
def test_map_filter_operators(self) -> None:
|
||||
"""Test mapping of various filter operators"""
|
||||
assert map_filter_operator("=") == "=="
|
||||
assert map_filter_operator(">") == ">"
|
||||
assert map_filter_operator("<") == "<"
|
||||
assert map_filter_operator(">=") == ">="
|
||||
assert map_filter_operator("<=") == "<="
|
||||
assert map_filter_operator("!=") == "!="
|
||||
|
||||
def test_map_filter_operator_unknown(self) -> None:
|
||||
"""Test mapping of unknown operator returns original"""
|
||||
assert map_filter_operator("UNKNOWN") == "UNKNOWN"
|
||||
|
||||
|
||||
class TestMapTableConfig:
|
||||
"""Test map_table_config function"""
|
||||
|
||||
def test_map_table_config_basic(self) -> None:
|
||||
"""Test basic table config mapping with aggregated columns"""
|
||||
config = TableChartConfig(
|
||||
columns=[
|
||||
ColumnRef(name="product", aggregate="COUNT"),
|
||||
ColumnRef(name="revenue", aggregate="SUM"),
|
||||
]
|
||||
)
|
||||
|
||||
result = map_table_config(config)
|
||||
|
||||
assert result["viz_type"] == "table"
|
||||
assert result["query_mode"] == "aggregate"
|
||||
# Aggregated columns should be in metrics, not all_columns
|
||||
assert "all_columns" not in result
|
||||
assert len(result["metrics"]) == 2
|
||||
assert result["metrics"][0]["aggregate"] == "COUNT"
|
||||
assert result["metrics"][1]["aggregate"] == "SUM"
|
||||
|
||||
def test_map_table_config_raw_columns(self) -> None:
|
||||
"""Test table config mapping with raw columns (no aggregates)"""
|
||||
config = TableChartConfig(
|
||||
columns=[
|
||||
ColumnRef(name="product"),
|
||||
ColumnRef(name="category"),
|
||||
]
|
||||
)
|
||||
|
||||
result = map_table_config(config)
|
||||
|
||||
assert result["viz_type"] == "table"
|
||||
assert result["query_mode"] == "raw"
|
||||
# Raw columns should be in all_columns
|
||||
assert result["all_columns"] == ["product", "category"]
|
||||
assert "metrics" not in result
|
||||
|
||||
def test_map_table_config_with_filters(self) -> None:
|
||||
"""Test table config mapping with filters"""
|
||||
config = TableChartConfig(
|
||||
columns=[ColumnRef(name="product")],
|
||||
filters=[FilterConfig(column="status", op="=", value="active")],
|
||||
)
|
||||
|
||||
result = map_table_config(config)
|
||||
|
||||
assert "adhoc_filters" in result
|
||||
assert len(result["adhoc_filters"]) == 1
|
||||
filter_obj = result["adhoc_filters"][0]
|
||||
assert filter_obj["subject"] == "status"
|
||||
assert filter_obj["operator"] == "=="
|
||||
assert filter_obj["comparator"] == "active"
|
||||
assert filter_obj["expressionType"] == "SIMPLE"
|
||||
|
||||
def test_map_table_config_with_sort(self) -> None:
|
||||
"""Test table config mapping with sort"""
|
||||
config = TableChartConfig(
|
||||
columns=[ColumnRef(name="product")], sort_by=["product", "revenue"]
|
||||
)
|
||||
|
||||
result = map_table_config(config)
|
||||
assert result["order_by_cols"] == ["product", "revenue"]
|
||||
|
||||
|
||||
class TestMapXYConfig:
|
||||
"""Test map_xy_config function"""
|
||||
|
||||
def test_map_xy_config_line_chart(self) -> None:
|
||||
"""Test XY config mapping for line chart"""
|
||||
config = XYChartConfig(
|
||||
x=ColumnRef(name="date"),
|
||||
y=[ColumnRef(name="revenue", aggregate="SUM")],
|
||||
kind="line",
|
||||
)
|
||||
|
||||
result = map_xy_config(config)
|
||||
|
||||
assert result["viz_type"] == "echarts_timeseries_line"
|
||||
assert result["x_axis"] == "date"
|
||||
assert len(result["metrics"]) == 1
|
||||
assert result["metrics"][0]["aggregate"] == "SUM"
|
||||
|
||||
def test_map_xy_config_with_groupby(self) -> None:
|
||||
"""Test XY config mapping with group by"""
|
||||
config = XYChartConfig(
|
||||
x=ColumnRef(name="date"),
|
||||
y=[ColumnRef(name="revenue")],
|
||||
kind="bar",
|
||||
group_by=ColumnRef(name="region"),
|
||||
)
|
||||
|
||||
result = map_xy_config(config)
|
||||
|
||||
assert result["viz_type"] == "echarts_timeseries_bar"
|
||||
assert result["groupby"] == ["region"]
|
||||
|
||||
def test_map_xy_config_with_axes(self) -> None:
|
||||
"""Test XY config mapping with axis configurations"""
|
||||
config = XYChartConfig(
|
||||
x=ColumnRef(name="date"),
|
||||
y=[ColumnRef(name="revenue")],
|
||||
kind="area",
|
||||
x_axis=AxisConfig(title="Date", format="%Y-%m-%d"),
|
||||
y_axis=AxisConfig(title="Revenue", scale="log", format="$,.2f"),
|
||||
)
|
||||
|
||||
result = map_xy_config(config)
|
||||
|
||||
assert result["viz_type"] == "echarts_area"
|
||||
assert result["x_axis_title"] == "Date"
|
||||
assert result["x_axis_format"] == "%Y-%m-%d"
|
||||
assert result["y_axis_title"] == "Revenue"
|
||||
assert result["y_axis_format"] == "$,.2f"
|
||||
assert result["y_axis_scale"] == "log"
|
||||
|
||||
def test_map_xy_config_with_legend(self) -> None:
|
||||
"""Test XY config mapping with legend configuration"""
|
||||
config = XYChartConfig(
|
||||
x=ColumnRef(name="date"),
|
||||
y=[ColumnRef(name="revenue")],
|
||||
kind="scatter",
|
||||
legend=LegendConfig(show=False, position="top"),
|
||||
)
|
||||
|
||||
result = map_xy_config(config)
|
||||
|
||||
assert result["viz_type"] == "echarts_timeseries_scatter"
|
||||
assert result["show_legend"] is False
|
||||
assert result["legend_orientation"] == "top"
|
||||
|
||||
|
||||
class TestMapConfigToFormData:
|
||||
"""Test map_config_to_form_data function"""
|
||||
|
||||
def test_map_table_config_type(self) -> None:
|
||||
"""Test mapping table config type"""
|
||||
config = TableChartConfig(columns=[ColumnRef(name="test")])
|
||||
result = map_config_to_form_data(config)
|
||||
assert result["viz_type"] == "table"
|
||||
|
||||
def test_map_xy_config_type(self) -> None:
|
||||
"""Test mapping XY config type"""
|
||||
config = XYChartConfig(
|
||||
x=ColumnRef(name="date"), y=[ColumnRef(name="revenue")], kind="line"
|
||||
)
|
||||
result = map_config_to_form_data(config)
|
||||
assert result["viz_type"] == "echarts_timeseries_line"
|
||||
|
||||
def test_map_unsupported_config_type(self) -> None:
|
||||
"""Test mapping unsupported config type raises error"""
|
||||
with pytest.raises(ValueError, match="Unsupported config type"):
|
||||
map_config_to_form_data("invalid_config") # type: ignore
|
||||
|
||||
|
||||
class TestGenerateChartName:
|
||||
"""Test generate_chart_name function"""
|
||||
|
||||
def test_generate_table_chart_name(self) -> None:
|
||||
"""Test generating name for table chart"""
|
||||
config = TableChartConfig(
|
||||
columns=[
|
||||
ColumnRef(name="product"),
|
||||
ColumnRef(name="revenue"),
|
||||
]
|
||||
)
|
||||
|
||||
result = generate_chart_name(config)
|
||||
assert result == "Table Chart - product, revenue"
|
||||
|
||||
def test_generate_xy_chart_name(self) -> None:
|
||||
"""Test generating name for XY chart"""
|
||||
config = XYChartConfig(
|
||||
x=ColumnRef(name="date"),
|
||||
y=[ColumnRef(name="revenue"), ColumnRef(name="orders")],
|
||||
kind="line",
|
||||
)
|
||||
|
||||
result = generate_chart_name(config)
|
||||
assert result == "Line Chart - date vs revenue, orders"
|
||||
|
||||
def test_generate_chart_name_unsupported(self) -> None:
|
||||
"""Test generating name for unsupported config type"""
|
||||
result = generate_chart_name("invalid_config") # type: ignore
|
||||
assert result == "Chart"
|
||||
|
||||
|
||||
class TestGenerateExploreLink:
|
||||
"""Test generate_explore_link function"""
|
||||
|
||||
@patch("superset.mcp_service.chart.chart_utils.get_superset_base_url")
|
||||
def test_generate_explore_link_uses_base_url(self, mock_get_base_url) -> None:
|
||||
"""Test that generate_explore_link uses the configured base URL"""
|
||||
from urllib.parse import urlparse
|
||||
|
||||
mock_get_base_url.return_value = "https://superset.example.com"
|
||||
form_data = {"viz_type": "table", "metrics": ["count"]}
|
||||
|
||||
result = generate_explore_link("123", form_data)
|
||||
|
||||
# Should use the configured base URL - use urlparse to avoid CodeQL warning
|
||||
parsed_url = urlparse(result)
|
||||
expected_netloc = "superset.example.com"
|
||||
assert parsed_url.scheme == "https"
|
||||
assert parsed_url.netloc == expected_netloc
|
||||
assert "/explore/" in parsed_url.path
|
||||
assert "datasource_id=123" in result
|
||||
|
||||
@patch("superset.mcp_service.chart.chart_utils.get_superset_base_url")
|
||||
def test_generate_explore_link_fallback_url(self, mock_get_base_url) -> None:
|
||||
"""Test generate_explore_link returns fallback URL when dataset not found"""
|
||||
mock_get_base_url.return_value = "http://localhost:9001"
|
||||
form_data = {"viz_type": "table"}
|
||||
|
||||
# Mock dataset not found scenario
|
||||
with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=None):
|
||||
result = generate_explore_link("999", form_data)
|
||||
|
||||
assert (
|
||||
result
|
||||
== "http://localhost:9001/explore/?datasource_type=table&datasource_id=999"
|
||||
)
|
||||
|
||||
@patch("superset.mcp_service.chart.chart_utils.get_superset_base_url")
|
||||
@patch("superset.mcp_service.commands.create_form_data.MCPCreateFormDataCommand")
|
||||
def test_generate_explore_link_with_form_data_key(
|
||||
self, mock_command, mock_get_base_url
|
||||
) -> None:
|
||||
"""Test generate_explore_link creates form_data_key when dataset exists"""
|
||||
mock_get_base_url.return_value = "http://localhost:9001"
|
||||
mock_command.return_value.run.return_value = "test_form_data_key"
|
||||
|
||||
# Mock dataset exists
|
||||
mock_dataset = type("Dataset", (), {"id": 123})()
|
||||
with patch(
|
||||
"superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_dataset
|
||||
):
|
||||
result = generate_explore_link(123, {"viz_type": "table"})
|
||||
|
||||
assert (
|
||||
result == "http://localhost:9001/explore/?form_data_key=test_form_data_key"
|
||||
)
|
||||
mock_command.assert_called_once()
|
||||
|
||||
@patch("superset.mcp_service.chart.chart_utils.get_superset_base_url")
|
||||
def test_generate_explore_link_exception_handling(self, mock_get_base_url) -> None:
|
||||
"""Test generate_explore_link handles exceptions gracefully"""
|
||||
mock_get_base_url.return_value = "http://localhost:9001"
|
||||
|
||||
# Mock exception during form_data creation
|
||||
with patch(
|
||||
"superset.daos.dataset.DatasetDAO.find_by_id",
|
||||
side_effect=Exception("DB error"),
|
||||
):
|
||||
result = generate_explore_link("123", {"viz_type": "table"})
|
||||
|
||||
# Should fallback to basic URL
|
||||
assert (
|
||||
result
|
||||
== "http://localhost:9001/explore/?datasource_type=table&datasource_id=123"
|
||||
)
|
||||
|
||||
|
||||
class TestCriticalBugFixes:
|
||||
"""Test critical bug fixes for chart utilities."""
|
||||
|
||||
def test_time_series_aggregation_fix(self) -> None:
|
||||
"""Test that time series charts preserve temporal dimension."""
|
||||
# Create a time series chart configuration
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
kind="line",
|
||||
x=ColumnRef(name="order_date"),
|
||||
y=[ColumnRef(name="sales", aggregate="SUM", label="Total Sales")],
|
||||
)
|
||||
|
||||
form_data = map_xy_config(config)
|
||||
|
||||
# Verify the fix: x_axis should be set correctly
|
||||
assert form_data["x_axis"] == "order_date"
|
||||
|
||||
# Verify the fix: groupby should not duplicate x_axis
|
||||
# This prevents the "Duplicate column/metric labels" error
|
||||
assert "groupby" not in form_data or "order_date" not in form_data.get(
|
||||
"groupby", []
|
||||
)
|
||||
|
||||
# Verify chart type mapping
|
||||
assert form_data["viz_type"] == "echarts_timeseries_line"
|
||||
|
||||
def test_time_series_with_explicit_group_by(self) -> None:
|
||||
"""Test time series with explicit group_by different from x_axis."""
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
kind="line",
|
||||
x=ColumnRef(name="order_date"),
|
||||
y=[ColumnRef(name="sales", aggregate="SUM", label="Total Sales")],
|
||||
group_by=ColumnRef(name="category"),
|
||||
)
|
||||
|
||||
form_data = map_xy_config(config)
|
||||
|
||||
# Verify x_axis is set
|
||||
assert form_data["x_axis"] == "order_date"
|
||||
|
||||
# Verify groupby only contains the explicit group_by, not x_axis
|
||||
assert form_data.get("groupby") == ["category"]
|
||||
assert "order_date" not in form_data.get("groupby", [])
|
||||
|
||||
def test_duplicate_label_prevention(self) -> None:
|
||||
"""Test that duplicate column/metric labels are prevented."""
|
||||
# This configuration would previously cause:
|
||||
# "Duplicate column/metric labels: 'price_each'"
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="price_each", label="Price Each"), # Custom label
|
||||
y=[
|
||||
ColumnRef(name="sales", aggregate="SUM", label="Total Sales"),
|
||||
ColumnRef(name="quantity", aggregate="COUNT", label="Order Count"),
|
||||
],
|
||||
group_by=ColumnRef(name="price_each"), # Same column as x_axis
|
||||
kind="line",
|
||||
)
|
||||
|
||||
form_data = map_xy_config(config)
|
||||
|
||||
# Verify the fix: x_axis is set
|
||||
assert form_data["x_axis"] == "price_each"
|
||||
|
||||
# Verify the fix: groupby is empty because group_by == x_axis
|
||||
# This prevents the duplicate label error
|
||||
assert "groupby" not in form_data or not form_data["groupby"]
|
||||
|
||||
def test_metric_object_creation_with_labels(self) -> None:
|
||||
"""Test that metric objects are created correctly with proper labels."""
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="date"),
|
||||
y=[
|
||||
ColumnRef(name="sales", aggregate="SUM", label="Total Sales"),
|
||||
ColumnRef(name="profit", aggregate="AVG"), # No custom label
|
||||
],
|
||||
kind="bar",
|
||||
)
|
||||
|
||||
form_data = map_xy_config(config)
|
||||
|
||||
# Verify metrics are created correctly
|
||||
metrics = form_data["metrics"]
|
||||
assert len(metrics) == 2
|
||||
|
||||
# First metric with custom label
|
||||
assert metrics[0]["label"] == "Total Sales"
|
||||
assert metrics[0]["aggregate"] == "SUM"
|
||||
assert metrics[0]["column"]["column_name"] == "sales"
|
||||
|
||||
# Second metric with auto-generated label
|
||||
assert metrics[1]["label"] == "AVG(profit)"
|
||||
assert metrics[1]["aggregate"] == "AVG"
|
||||
assert metrics[1]["column"]["column_name"] == "profit"
|
||||
|
||||
def test_chart_type_mapping_comprehensive(self) -> None:
|
||||
"""Test that chart types are mapped correctly to Superset viz types."""
|
||||
test_cases = [
|
||||
("line", "echarts_timeseries_line"),
|
||||
("bar", "echarts_timeseries_bar"),
|
||||
("area", "echarts_area"),
|
||||
("scatter", "echarts_timeseries_scatter"),
|
||||
]
|
||||
|
||||
for kind, expected_viz_type in test_cases:
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="date"),
|
||||
y=[ColumnRef(name="value", aggregate="SUM")],
|
||||
kind=kind,
|
||||
)
|
||||
|
||||
form_data = map_xy_config(config)
|
||||
assert form_data["viz_type"] == expected_viz_type
|
||||
268
tests/unit_tests/mcp_service/chart/tool/test_generate_chart.py
Normal file
268
tests/unit_tests/mcp_service/chart/tool/test_generate_chart.py
Normal file
@@ -0,0 +1,268 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Unit tests for MCP generate_chart tool
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
AxisConfig,
|
||||
ColumnRef,
|
||||
FilterConfig,
|
||||
GenerateChartRequest,
|
||||
LegendConfig,
|
||||
TableChartConfig,
|
||||
XYChartConfig,
|
||||
)
|
||||
|
||||
|
||||
class TestGenerateChart:
|
||||
"""Tests for generate_chart MCP tool."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_chart_request_structure(self):
|
||||
"""Test that chart generation request structures are properly formed."""
|
||||
# Table chart request
|
||||
table_config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[
|
||||
ColumnRef(name="region", label="Region"),
|
||||
ColumnRef(name="sales", label="Sales", aggregate="SUM"),
|
||||
],
|
||||
filters=[FilterConfig(column="year", op="=", value="2024")],
|
||||
sort_by=["sales"],
|
||||
)
|
||||
table_request = GenerateChartRequest(dataset_id="1", config=table_config)
|
||||
assert table_request.dataset_id == "1"
|
||||
assert table_request.config.chart_type == "table"
|
||||
assert len(table_request.config.columns) == 2
|
||||
assert table_request.config.columns[0].name == "region"
|
||||
assert table_request.config.columns[1].aggregate == "SUM"
|
||||
|
||||
# XY chart request
|
||||
xy_config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="date"),
|
||||
y=[ColumnRef(name="value", aggregate="SUM")],
|
||||
kind="line",
|
||||
group_by=ColumnRef(name="category"),
|
||||
x_axis=AxisConfig(title="Date", format="smart_date"),
|
||||
y_axis=AxisConfig(title="Value", format="$,.2f"),
|
||||
legend=LegendConfig(show=True, position="top"),
|
||||
)
|
||||
xy_request = GenerateChartRequest(dataset_id="2", config=xy_config)
|
||||
assert xy_request.config.chart_type == "xy"
|
||||
assert xy_request.config.x.name == "date"
|
||||
assert xy_request.config.y[0].aggregate == "SUM"
|
||||
assert xy_request.config.kind == "line"
|
||||
assert xy_request.config.x_axis.title == "Date"
|
||||
assert xy_request.config.legend.show is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_chart_validation_error_handling(self):
|
||||
"""Test that validation errors are properly structured."""
|
||||
|
||||
# Create a validation error with the correct structure
|
||||
validation_error_entry = {
|
||||
"field": "x_axis",
|
||||
"provided_value": "invalid_col",
|
||||
"error_type": "column_not_found",
|
||||
"message": "Column 'invalid_col' not found",
|
||||
}
|
||||
|
||||
# Test that validation error structure is correct
|
||||
assert validation_error_entry["field"] == "x_axis"
|
||||
assert validation_error_entry["error_type"] == "column_not_found"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chart_config_variations(self):
|
||||
"""Test various chart configuration options."""
|
||||
# Test all chart types
|
||||
chart_types = ["line", "bar", "area", "scatter"]
|
||||
for chart_type in chart_types:
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="x_col"),
|
||||
y=[ColumnRef(name="y_col")],
|
||||
kind=chart_type,
|
||||
)
|
||||
assert config.kind == chart_type
|
||||
|
||||
# Test multiple Y columns
|
||||
multi_y_config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="date"),
|
||||
y=[
|
||||
ColumnRef(name="sales", aggregate="SUM"),
|
||||
ColumnRef(name="profit", aggregate="AVG"),
|
||||
ColumnRef(name="quantity", aggregate="COUNT"),
|
||||
],
|
||||
kind="line",
|
||||
)
|
||||
assert len(multi_y_config.y) == 3
|
||||
assert multi_y_config.y[1].aggregate == "AVG"
|
||||
|
||||
# Test filter operators
|
||||
operators = ["=", "!=", ">", ">=", "<", "<="]
|
||||
filters = [FilterConfig(column="col", op=op, value="val") for op in operators]
|
||||
for i, f in enumerate(filters):
|
||||
assert f.op == operators[i]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_chart_response_structure(self):
|
||||
"""Test the expected response structure for chart generation."""
|
||||
# The response should contain these fields
|
||||
_ = {
|
||||
"chart": {
|
||||
"id": int,
|
||||
"slice_name": str,
|
||||
"viz_type": str,
|
||||
"url": str,
|
||||
"uuid": str,
|
||||
"saved": bool,
|
||||
},
|
||||
"error": None,
|
||||
"success": bool,
|
||||
"schema_version": str,
|
||||
"api_version": str,
|
||||
}
|
||||
|
||||
# When chart creation succeeds, these additional fields may be present
|
||||
_ = [
|
||||
"previews",
|
||||
"capabilities",
|
||||
"semantics",
|
||||
"explore_url",
|
||||
"form_data_key",
|
||||
"api_endpoints",
|
||||
"performance",
|
||||
"accessibility",
|
||||
]
|
||||
|
||||
# This is just a structural test - actual integration tests would verify
|
||||
# the tool returns data matching this structure
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dataset_id_flexibility(self):
|
||||
"""Test that dataset_id can be string or int."""
|
||||
configs = [
|
||||
GenerateChartRequest(
|
||||
dataset_id="123",
|
||||
config=TableChartConfig(
|
||||
chart_type="table", columns=[ColumnRef(name="col1")]
|
||||
),
|
||||
),
|
||||
GenerateChartRequest(
|
||||
dataset_id="uuid-string-here",
|
||||
config=TableChartConfig(
|
||||
chart_type="table", columns=[ColumnRef(name="col1")]
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
for config in configs:
|
||||
assert isinstance(config.dataset_id, str)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_chart_flag(self):
|
||||
"""Test save_chart flag behavior."""
|
||||
# Default should be True (save chart)
|
||||
request1 = GenerateChartRequest(
|
||||
dataset_id="1",
|
||||
config=TableChartConfig(
|
||||
chart_type="table", columns=[ColumnRef(name="col1")]
|
||||
),
|
||||
)
|
||||
assert request1.save_chart is True
|
||||
|
||||
# Explicit False (preview only)
|
||||
request2 = GenerateChartRequest(
|
||||
dataset_id="1",
|
||||
config=TableChartConfig(
|
||||
chart_type="table", columns=[ColumnRef(name="col1")]
|
||||
),
|
||||
save_chart=False,
|
||||
)
|
||||
assert request2.save_chart is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preview_formats(self):
|
||||
"""Test preview format options."""
|
||||
formats = ["url", "ascii", "table"]
|
||||
request = GenerateChartRequest(
|
||||
dataset_id="1",
|
||||
config=TableChartConfig(
|
||||
chart_type="table", columns=[ColumnRef(name="col1")]
|
||||
),
|
||||
generate_preview=True,
|
||||
preview_formats=formats,
|
||||
)
|
||||
assert request.generate_preview is True
|
||||
assert set(request.preview_formats) == set(formats)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_column_ref_features(self):
|
||||
"""Test ColumnRef features like aggregation and labels."""
|
||||
# Simple column
|
||||
col1 = ColumnRef(name="region")
|
||||
assert col1.name == "region"
|
||||
assert col1.label is None
|
||||
assert col1.aggregate is None
|
||||
|
||||
# Column with aggregation
|
||||
col2 = ColumnRef(name="sales", aggregate="SUM", label="Total Sales")
|
||||
assert col2.name == "sales"
|
||||
assert col2.aggregate == "SUM"
|
||||
assert col2.label == "Total Sales"
|
||||
|
||||
# All supported aggregations
|
||||
aggs = ["SUM", "AVG", "COUNT", "MIN", "MAX", "COUNT_DISTINCT"]
|
||||
for agg in aggs:
|
||||
col = ColumnRef(name="value", aggregate=agg)
|
||||
assert col.aggregate == agg
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_axis_config_options(self):
|
||||
"""Test axis configuration options."""
|
||||
axis = AxisConfig(
|
||||
title="Sales Amount",
|
||||
format="$,.2f",
|
||||
scale="linear",
|
||||
)
|
||||
assert axis.title == "Sales Amount"
|
||||
assert axis.format == "$,.2f"
|
||||
assert axis.scale == "linear"
|
||||
|
||||
# Test different formats
|
||||
formats = ["SMART_NUMBER", "$,.2f", ",.0f", "smart_date", ".3%"]
|
||||
for fmt in formats:
|
||||
axis = AxisConfig(format=fmt)
|
||||
assert axis.format == fmt
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legend_config_options(self):
|
||||
"""Test legend configuration options."""
|
||||
positions = ["top", "bottom", "left", "right"]
|
||||
for pos in positions:
|
||||
legend = LegendConfig(show=True, position=pos)
|
||||
assert legend.position == pos
|
||||
|
||||
# Hidden legend
|
||||
legend = LegendConfig(show=False)
|
||||
assert legend.show is False
|
||||
@@ -0,0 +1,290 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Unit tests for get_chart_preview MCP tool
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
ASCIIPreview,
|
||||
GetChartPreviewRequest,
|
||||
TablePreview,
|
||||
URLPreview,
|
||||
)
|
||||
|
||||
|
||||
class TestGetChartPreview:
|
||||
"""Tests for get_chart_preview MCP tool."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_chart_preview_request_structure(self):
|
||||
"""Test that preview request structures are properly formed."""
|
||||
# Numeric ID request
|
||||
request1 = GetChartPreviewRequest(identifier=123, format="url")
|
||||
assert request1.identifier == 123
|
||||
assert request1.format == "url"
|
||||
# Default dimensions are set
|
||||
assert request1.width == 800
|
||||
assert request1.height == 600
|
||||
|
||||
# String ID request
|
||||
request2 = GetChartPreviewRequest(identifier="456", format="ascii")
|
||||
assert request2.identifier == "456"
|
||||
assert request2.format == "ascii"
|
||||
|
||||
# UUID request
|
||||
request3 = GetChartPreviewRequest(
|
||||
identifier="a1b2c3d4-e5f6-7890-abcd-ef1234567890", format="table"
|
||||
)
|
||||
assert request3.identifier == "a1b2c3d4-e5f6-7890-abcd-ef1234567890"
|
||||
assert request3.format == "table"
|
||||
|
||||
# Default format
|
||||
request4 = GetChartPreviewRequest(identifier=789)
|
||||
assert request4.format == "url" # default
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preview_format_types(self):
|
||||
"""Test different preview format types."""
|
||||
formats = ["url", "ascii", "table"]
|
||||
for fmt in formats:
|
||||
request = GetChartPreviewRequest(identifier=1, format=fmt)
|
||||
assert request.format == fmt
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_url_preview_structure(self):
|
||||
"""Test URLPreview response structure."""
|
||||
preview = URLPreview(
|
||||
preview_url="http://localhost:5008/screenshot/chart/123.png",
|
||||
width=800,
|
||||
height=600,
|
||||
supports_interaction=False,
|
||||
)
|
||||
assert preview.type == "url"
|
||||
assert preview.preview_url == "http://localhost:5008/screenshot/chart/123.png"
|
||||
assert preview.width == 800
|
||||
assert preview.height == 600
|
||||
assert preview.supports_interaction is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ascii_preview_structure(self):
|
||||
"""Test ASCIIPreview response structure."""
|
||||
ascii_art = """
|
||||
┌─────────────────────────┐
|
||||
│ Sales by Region │
|
||||
├─────────────────────────┤
|
||||
│ North ████████ 45% │
|
||||
│ South ██████ 30% │
|
||||
│ East ████ 20% │
|
||||
│ West ██ 5% │
|
||||
└─────────────────────────┘
|
||||
"""
|
||||
preview = ASCIIPreview(
|
||||
ascii_content=ascii_art.strip(),
|
||||
width=25,
|
||||
height=8,
|
||||
)
|
||||
assert preview.type == "ascii"
|
||||
assert "Sales by Region" in preview.ascii_content
|
||||
assert preview.width == 25
|
||||
assert preview.height == 8
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_table_preview_structure(self):
|
||||
"""Test TablePreview response structure."""
|
||||
table_content = """
|
||||
| Region | Sales | Profit |
|
||||
|--------|--------|--------|
|
||||
| North | 45000 | 12000 |
|
||||
| South | 30000 | 8000 |
|
||||
| East | 20000 | 5000 |
|
||||
| West | 5000 | 1000 |
|
||||
"""
|
||||
preview = TablePreview(
|
||||
table_data=table_content.strip(),
|
||||
row_count=4,
|
||||
supports_sorting=True,
|
||||
)
|
||||
assert preview.type == "table"
|
||||
assert "Region" in preview.table_data
|
||||
assert "North" in preview.table_data
|
||||
assert preview.row_count == 4
|
||||
assert preview.supports_sorting is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chart_preview_response_structure(self):
|
||||
"""Test the expected response structure for chart preview."""
|
||||
# Core fields that should always be present
|
||||
_ = [
|
||||
"chart_id",
|
||||
"chart_name",
|
||||
"chart_type",
|
||||
"explore_url",
|
||||
"content", # Union of URLPreview | ASCIIPreview | TablePreview
|
||||
"chart_description",
|
||||
"accessibility",
|
||||
"performance",
|
||||
]
|
||||
|
||||
# Additional fields that may be present for backward compatibility
|
||||
_ = [
|
||||
"format",
|
||||
"preview_url",
|
||||
"ascii_chart",
|
||||
"table_data",
|
||||
"width",
|
||||
"height",
|
||||
"schema_version",
|
||||
"api_version",
|
||||
]
|
||||
|
||||
# This is a structural test - actual integration tests would verify
|
||||
# the tool returns data matching this structure
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preview_dimensions(self):
|
||||
"""Test preview dimensions in response."""
|
||||
# Standard dimensions that may appear in preview responses
|
||||
standard_sizes = [
|
||||
(800, 600), # Default
|
||||
(1200, 800), # Large
|
||||
(400, 300), # Small
|
||||
(1920, 1080), # Full HD
|
||||
]
|
||||
|
||||
for width, height in standard_sizes:
|
||||
# URL preview with dimensions
|
||||
url_preview = URLPreview(
|
||||
preview_url="http://example.com/chart.png",
|
||||
width=width,
|
||||
height=height,
|
||||
supports_interaction=False,
|
||||
)
|
||||
assert url_preview.width == width
|
||||
assert url_preview.height == height
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_response_structures(self):
|
||||
"""Test error response structures."""
|
||||
# Error responses typically follow this structure
|
||||
error_response = {
|
||||
"error_type": "not_found",
|
||||
"message": "Chart not found",
|
||||
"details": "No chart found with ID 999",
|
||||
"chart_id": 999,
|
||||
}
|
||||
assert error_response["error_type"] == "not_found"
|
||||
assert error_response["chart_id"] == 999
|
||||
|
||||
# Preview generation error structure
|
||||
preview_error = {
|
||||
"error_type": "preview_error",
|
||||
"message": "Failed to generate preview",
|
||||
"details": "Screenshot service unavailable",
|
||||
}
|
||||
assert preview_error["error_type"] == "preview_error"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_accessibility_metadata(self):
|
||||
"""Test accessibility metadata structure."""
|
||||
from superset.mcp_service.chart.schemas import AccessibilityMetadata
|
||||
|
||||
metadata = AccessibilityMetadata(
|
||||
color_blind_safe=True,
|
||||
alt_text="Bar chart showing sales by region",
|
||||
high_contrast_available=False,
|
||||
)
|
||||
assert metadata.color_blind_safe is True
|
||||
assert "sales by region" in metadata.alt_text
|
||||
assert metadata.high_contrast_available is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_performance_metadata(self):
|
||||
"""Test performance metadata structure."""
|
||||
from superset.mcp_service.chart.schemas import PerformanceMetadata
|
||||
|
||||
metadata = PerformanceMetadata(
|
||||
query_duration_ms=150,
|
||||
cache_status="hit",
|
||||
optimization_suggestions=["Consider adding an index on date column"],
|
||||
)
|
||||
assert metadata.query_duration_ms == 150
|
||||
assert metadata.cache_status == "hit"
|
||||
assert len(metadata.optimization_suggestions) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chart_types_support(self):
|
||||
"""Test that various chart types are supported."""
|
||||
chart_types = [
|
||||
"echarts_timeseries_line",
|
||||
"echarts_timeseries_bar",
|
||||
"echarts_area",
|
||||
"echarts_timeseries_scatter",
|
||||
"table",
|
||||
"pie",
|
||||
"big_number",
|
||||
"big_number_total",
|
||||
"pivot_table_v2",
|
||||
"dist_bar",
|
||||
"box_plot",
|
||||
]
|
||||
|
||||
# All chart types should be previewable
|
||||
for _chart_type in chart_types:
|
||||
# This would be tested in integration tests
|
||||
pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ascii_art_variations(self):
|
||||
"""Test ASCII art generation for different chart types."""
|
||||
# Line chart ASCII
|
||||
_ = """
|
||||
Sales Trend
|
||||
│
|
||||
│ ╱╲
|
||||
│ ╱ ╲
|
||||
│ ╱ ╲
|
||||
│ ╱ ╲
|
||||
│ ╱ ╲
|
||||
└────────────
|
||||
Jan Feb Mar
|
||||
"""
|
||||
|
||||
# Bar chart ASCII
|
||||
_ = """
|
||||
Sales by Region
|
||||
│
|
||||
│ ████ North
|
||||
│ ███ South
|
||||
│ ██ East
|
||||
│ █ West
|
||||
└────────────
|
||||
"""
|
||||
|
||||
# Pie chart ASCII
|
||||
_ = """
|
||||
Market Share
|
||||
╭─────╮
|
||||
╱ ╲
|
||||
│ 45% │
|
||||
│ North │
|
||||
╰─────────╯
|
||||
"""
|
||||
|
||||
# These demonstrate the expected ASCII formats for different chart types
|
||||
385
tests/unit_tests/mcp_service/chart/tool/test_update_chart.py
Normal file
385
tests/unit_tests/mcp_service/chart/tool/test_update_chart.py
Normal file
@@ -0,0 +1,385 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Unit tests for update_chart MCP tool
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
AxisConfig,
|
||||
ColumnRef,
|
||||
FilterConfig,
|
||||
LegendConfig,
|
||||
TableChartConfig,
|
||||
UpdateChartRequest,
|
||||
XYChartConfig,
|
||||
)
|
||||
|
||||
|
||||
class TestUpdateChart:
|
||||
"""Tests for update_chart MCP tool."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_request_structure(self):
|
||||
"""Test that chart update request structures are properly formed."""
|
||||
# Table chart update with numeric ID
|
||||
table_config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[
|
||||
ColumnRef(name="region", label="Region"),
|
||||
ColumnRef(name="sales", label="Sales", aggregate="SUM"),
|
||||
],
|
||||
filters=[FilterConfig(column="year", op="=", value="2024")],
|
||||
sort_by=["sales"],
|
||||
)
|
||||
table_request = UpdateChartRequest(identifier=123, config=table_config)
|
||||
assert table_request.identifier == 123
|
||||
assert table_request.config.chart_type == "table"
|
||||
assert len(table_request.config.columns) == 2
|
||||
assert table_request.config.columns[0].name == "region"
|
||||
assert table_request.config.columns[1].aggregate == "SUM"
|
||||
|
||||
# XY chart update with UUID
|
||||
xy_config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="date"),
|
||||
y=[ColumnRef(name="value", aggregate="SUM")],
|
||||
kind="line",
|
||||
group_by=ColumnRef(name="category"),
|
||||
x_axis=AxisConfig(title="Date", format="smart_date"),
|
||||
y_axis=AxisConfig(title="Value", format="$,.2f"),
|
||||
legend=LegendConfig(show=True, position="top"),
|
||||
)
|
||||
xy_request = UpdateChartRequest(
|
||||
identifier="a1b2c3d4-e5f6-7890-abcd-ef1234567890", config=xy_config
|
||||
)
|
||||
assert xy_request.identifier == "a1b2c3d4-e5f6-7890-abcd-ef1234567890"
|
||||
assert xy_request.config.chart_type == "xy"
|
||||
assert xy_request.config.x.name == "date"
|
||||
assert xy_request.config.y[0].aggregate == "SUM"
|
||||
assert xy_request.config.kind == "line"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_with_chart_name(self):
|
||||
"""Test updating chart with custom chart name."""
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
)
|
||||
|
||||
# Without custom name
|
||||
request1 = UpdateChartRequest(identifier=123, config=config)
|
||||
assert request1.chart_name is None
|
||||
|
||||
# With custom name
|
||||
request2 = UpdateChartRequest(
|
||||
identifier=123, config=config, chart_name="Updated Sales Report"
|
||||
)
|
||||
assert request2.chart_name == "Updated Sales Report"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_generation(self):
|
||||
"""Test preview generation options in update request."""
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
)
|
||||
|
||||
# Default preview generation
|
||||
request1 = UpdateChartRequest(identifier=123, config=config)
|
||||
assert request1.generate_preview is True
|
||||
assert request1.preview_formats == ["url"]
|
||||
|
||||
# Custom preview formats
|
||||
request2 = UpdateChartRequest(
|
||||
identifier=123,
|
||||
config=config,
|
||||
generate_preview=True,
|
||||
preview_formats=["url", "ascii", "table"],
|
||||
)
|
||||
assert request2.generate_preview is True
|
||||
assert set(request2.preview_formats) == {"url", "ascii", "table"}
|
||||
|
||||
# Disable preview generation
|
||||
request3 = UpdateChartRequest(
|
||||
identifier=123, config=config, generate_preview=False
|
||||
)
|
||||
assert request3.generate_preview is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_identifier_types(self):
|
||||
"""Test that identifier can be int or string (UUID)."""
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
)
|
||||
|
||||
# Integer ID
|
||||
request1 = UpdateChartRequest(identifier=123, config=config)
|
||||
assert request1.identifier == 123
|
||||
assert isinstance(request1.identifier, int)
|
||||
|
||||
# String numeric ID
|
||||
request2 = UpdateChartRequest(identifier="456", config=config)
|
||||
assert request2.identifier == "456"
|
||||
assert isinstance(request2.identifier, str)
|
||||
|
||||
# UUID string
|
||||
request3 = UpdateChartRequest(
|
||||
identifier="a1b2c3d4-e5f6-7890-abcd-ef1234567890", config=config
|
||||
)
|
||||
assert request3.identifier == "a1b2c3d4-e5f6-7890-abcd-ef1234567890"
|
||||
assert isinstance(request3.identifier, str)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_config_variations(self):
|
||||
"""Test various chart configuration options in updates."""
|
||||
# Test all XY chart types
|
||||
chart_types = ["line", "bar", "area", "scatter"]
|
||||
for chart_type in chart_types:
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="x_col"),
|
||||
y=[ColumnRef(name="y_col")],
|
||||
kind=chart_type,
|
||||
)
|
||||
request = UpdateChartRequest(identifier=1, config=config)
|
||||
assert request.config.kind == chart_type
|
||||
|
||||
# Test multiple Y columns
|
||||
multi_y_config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="date"),
|
||||
y=[
|
||||
ColumnRef(name="sales", aggregate="SUM"),
|
||||
ColumnRef(name="profit", aggregate="AVG"),
|
||||
ColumnRef(name="quantity", aggregate="COUNT"),
|
||||
],
|
||||
kind="line",
|
||||
)
|
||||
request = UpdateChartRequest(identifier=1, config=multi_y_config)
|
||||
assert len(request.config.y) == 3
|
||||
assert request.config.y[1].aggregate == "AVG"
|
||||
|
||||
# Test filter operators
|
||||
operators = ["=", "!=", ">", ">=", "<", "<="]
|
||||
filters = [FilterConfig(column="col", op=op, value="val") for op in operators]
|
||||
table_config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
filters=filters,
|
||||
)
|
||||
request = UpdateChartRequest(identifier=1, config=table_config)
|
||||
assert len(request.config.filters) == 6
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_response_structure(self):
|
||||
"""Test the expected response structure for chart updates."""
|
||||
# The response should contain these fields
|
||||
expected_response = {
|
||||
"chart": {
|
||||
"id": int,
|
||||
"slice_name": str,
|
||||
"viz_type": str,
|
||||
"url": str,
|
||||
"uuid": str,
|
||||
"updated": bool,
|
||||
},
|
||||
"error": None,
|
||||
"success": bool,
|
||||
"schema_version": str,
|
||||
"api_version": str,
|
||||
}
|
||||
|
||||
# When chart update succeeds, these additional fields may be present
|
||||
optional_fields = [
|
||||
"previews",
|
||||
"capabilities",
|
||||
"semantics",
|
||||
"explore_url",
|
||||
"api_endpoints",
|
||||
"performance",
|
||||
"accessibility",
|
||||
]
|
||||
|
||||
# Validate structure expectations
|
||||
assert "chart" in expected_response
|
||||
assert "success" in expected_response
|
||||
assert len(optional_fields) > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_axis_configurations(self):
|
||||
"""Test axis configuration updates."""
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="date"),
|
||||
y=[ColumnRef(name="sales")],
|
||||
x_axis=AxisConfig(
|
||||
title="Date",
|
||||
format="smart_date",
|
||||
scale="linear",
|
||||
),
|
||||
y_axis=AxisConfig(
|
||||
title="Sales Amount",
|
||||
format="$,.2f",
|
||||
scale="log",
|
||||
),
|
||||
)
|
||||
request = UpdateChartRequest(identifier=1, config=config)
|
||||
assert request.config.x_axis.title == "Date"
|
||||
assert request.config.x_axis.format == "smart_date"
|
||||
assert request.config.x_axis.scale == "linear"
|
||||
assert request.config.y_axis.title == "Sales Amount"
|
||||
assert request.config.y_axis.format == "$,.2f"
|
||||
assert request.config.y_axis.scale == "log"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_legend_configurations(self):
|
||||
"""Test legend configuration updates."""
|
||||
positions = ["top", "bottom", "left", "right"]
|
||||
for pos in positions:
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="x"),
|
||||
y=[ColumnRef(name="y")],
|
||||
legend=LegendConfig(show=True, position=pos),
|
||||
)
|
||||
request = UpdateChartRequest(identifier=1, config=config)
|
||||
assert request.config.legend.position == pos
|
||||
assert request.config.legend.show is True
|
||||
|
||||
# Hidden legend
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="x"),
|
||||
y=[ColumnRef(name="y")],
|
||||
legend=LegendConfig(show=False),
|
||||
)
|
||||
request = UpdateChartRequest(identifier=1, config=config)
|
||||
assert request.config.legend.show is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_aggregation_functions(self):
|
||||
"""Test all supported aggregation functions in updates."""
|
||||
aggs = ["SUM", "AVG", "COUNT", "MIN", "MAX", "COUNT_DISTINCT"]
|
||||
for agg in aggs:
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="value", aggregate=agg)],
|
||||
)
|
||||
request = UpdateChartRequest(identifier=1, config=config)
|
||||
assert request.config.columns[0].aggregate == agg
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_error_responses(self):
|
||||
"""Test expected error response structures."""
|
||||
# Chart not found error
|
||||
error_response = {
|
||||
"chart": None,
|
||||
"error": "No chart found with identifier: 999",
|
||||
"success": False,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
assert error_response["success"] is False
|
||||
assert error_response["chart"] is None
|
||||
assert "chart found" in error_response["error"].lower()
|
||||
|
||||
# General update error
|
||||
update_error = {
|
||||
"chart": None,
|
||||
"error": "Chart update failed: Permission denied",
|
||||
"success": False,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
assert update_error["success"] is False
|
||||
assert "failed" in update_error["error"].lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chart_name_sanitization(self):
|
||||
"""Test that chart names are properly sanitized."""
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
)
|
||||
|
||||
# Test with potentially problematic characters
|
||||
test_names = [
|
||||
"Normal Chart Name",
|
||||
"Chart with 'quotes'",
|
||||
'Chart with "double quotes"',
|
||||
"Chart with <brackets>",
|
||||
]
|
||||
|
||||
for name in test_names:
|
||||
request = UpdateChartRequest(identifier=1, config=config, chart_name=name)
|
||||
# Chart name should be set (sanitization happens in the validator)
|
||||
assert request.chart_name is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_with_filters(self):
|
||||
"""Test updating chart with various filter configurations."""
|
||||
filters = [
|
||||
FilterConfig(column="region", op="=", value="North"),
|
||||
FilterConfig(column="sales", op=">=", value=1000),
|
||||
FilterConfig(column="date", op=">", value="2024-01-01"),
|
||||
]
|
||||
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[
|
||||
ColumnRef(name="region"),
|
||||
ColumnRef(name="sales"),
|
||||
ColumnRef(name="date"),
|
||||
],
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
request = UpdateChartRequest(identifier=1, config=config)
|
||||
assert len(request.config.filters) == 3
|
||||
assert request.config.filters[0].column == "region"
|
||||
assert request.config.filters[1].op == ">="
|
||||
assert request.config.filters[2].value == "2024-01-01"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_cache_control(self):
|
||||
"""Test cache control parameters in update request."""
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
)
|
||||
|
||||
# Default cache settings
|
||||
request1 = UpdateChartRequest(identifier=1, config=config)
|
||||
assert request1.use_cache is True
|
||||
assert request1.force_refresh is False
|
||||
assert request1.cache_timeout is None
|
||||
|
||||
# Custom cache settings
|
||||
request2 = UpdateChartRequest(
|
||||
identifier=1,
|
||||
config=config,
|
||||
use_cache=False,
|
||||
force_refresh=True,
|
||||
cache_timeout=300,
|
||||
)
|
||||
assert request2.use_cache is False
|
||||
assert request2.force_refresh is True
|
||||
assert request2.cache_timeout == 300
|
||||
@@ -0,0 +1,474 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Unit tests for update_chart_preview MCP tool
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from superset.mcp_service.chart.schemas import (
|
||||
AxisConfig,
|
||||
ColumnRef,
|
||||
FilterConfig,
|
||||
LegendConfig,
|
||||
TableChartConfig,
|
||||
UpdateChartPreviewRequest,
|
||||
XYChartConfig,
|
||||
)
|
||||
|
||||
|
||||
class TestUpdateChartPreview:
|
||||
"""Tests for update_chart_preview MCP tool."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_request_structure(self):
|
||||
"""Test that chart preview update request structures are properly formed."""
|
||||
# Table chart preview update
|
||||
table_config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[
|
||||
ColumnRef(name="region", label="Region"),
|
||||
ColumnRef(name="sales", label="Sales", aggregate="SUM"),
|
||||
],
|
||||
filters=[FilterConfig(column="year", op="=", value="2024")],
|
||||
sort_by=["sales"],
|
||||
)
|
||||
table_request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123def456", dataset_id=1, config=table_config
|
||||
)
|
||||
assert table_request.form_data_key == "abc123def456"
|
||||
assert table_request.dataset_id == 1
|
||||
assert table_request.config.chart_type == "table"
|
||||
assert len(table_request.config.columns) == 2
|
||||
assert table_request.config.columns[0].name == "region"
|
||||
|
||||
# XY chart preview update
|
||||
xy_config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="date"),
|
||||
y=[ColumnRef(name="value", aggregate="SUM")],
|
||||
kind="line",
|
||||
group_by=ColumnRef(name="category"),
|
||||
x_axis=AxisConfig(title="Date", format="smart_date"),
|
||||
y_axis=AxisConfig(title="Value", format="$,.2f"),
|
||||
legend=LegendConfig(show=True, position="top"),
|
||||
)
|
||||
xy_request = UpdateChartPreviewRequest(
|
||||
form_data_key="xyz789ghi012", dataset_id="2", config=xy_config
|
||||
)
|
||||
assert xy_request.form_data_key == "xyz789ghi012"
|
||||
assert xy_request.dataset_id == "2"
|
||||
assert xy_request.config.chart_type == "xy"
|
||||
assert xy_request.config.x.name == "date"
|
||||
assert xy_request.config.kind == "line"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_dataset_id_types(self):
|
||||
"""Test that dataset_id can be int or string (UUID)."""
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
)
|
||||
|
||||
# Integer dataset_id
|
||||
request1 = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=123, config=config
|
||||
)
|
||||
assert request1.dataset_id == 123
|
||||
assert isinstance(request1.dataset_id, int)
|
||||
|
||||
# String numeric dataset_id
|
||||
request2 = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id="456", config=config
|
||||
)
|
||||
assert request2.dataset_id == "456"
|
||||
assert isinstance(request2.dataset_id, str)
|
||||
|
||||
# UUID string dataset_id
|
||||
request3 = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123",
|
||||
dataset_id="a1b2c3d4-e5f6-7890-abcd-ef1234567890",
|
||||
config=config,
|
||||
)
|
||||
assert request3.dataset_id == "a1b2c3d4-e5f6-7890-abcd-ef1234567890"
|
||||
assert isinstance(request3.dataset_id, str)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_generation_options(self):
|
||||
"""Test preview generation options in update preview request."""
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
)
|
||||
|
||||
# Default preview generation
|
||||
request1 = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=config
|
||||
)
|
||||
assert request1.generate_preview is True
|
||||
assert request1.preview_formats == ["url"]
|
||||
|
||||
# Custom preview formats
|
||||
request2 = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123",
|
||||
dataset_id=1,
|
||||
config=config,
|
||||
generate_preview=True,
|
||||
preview_formats=["url", "ascii", "table"],
|
||||
)
|
||||
assert request2.generate_preview is True
|
||||
assert set(request2.preview_formats) == {"url", "ascii", "table"}
|
||||
|
||||
# Disable preview generation
|
||||
request3 = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123",
|
||||
dataset_id=1,
|
||||
config=config,
|
||||
generate_preview=False,
|
||||
)
|
||||
assert request3.generate_preview is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_config_variations(self):
|
||||
"""Test various chart configuration options in preview updates."""
|
||||
# Test all XY chart types
|
||||
chart_types = ["line", "bar", "area", "scatter"]
|
||||
for chart_type in chart_types:
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="x_col"),
|
||||
y=[ColumnRef(name="y_col")],
|
||||
kind=chart_type,
|
||||
)
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=config
|
||||
)
|
||||
assert request.config.kind == chart_type
|
||||
|
||||
# Test multiple Y columns
|
||||
multi_y_config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="date"),
|
||||
y=[
|
||||
ColumnRef(name="sales", aggregate="SUM"),
|
||||
ColumnRef(name="profit", aggregate="AVG"),
|
||||
ColumnRef(name="quantity", aggregate="COUNT"),
|
||||
],
|
||||
kind="line",
|
||||
)
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=multi_y_config
|
||||
)
|
||||
assert len(request.config.y) == 3
|
||||
assert request.config.y[1].aggregate == "AVG"
|
||||
|
||||
# Test filter operators
|
||||
operators = ["=", "!=", ">", ">=", "<", "<="]
|
||||
filters = [FilterConfig(column="col", op=op, value="val") for op in operators]
|
||||
table_config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
filters=filters,
|
||||
)
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=table_config
|
||||
)
|
||||
assert len(request.config.filters) == 6
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_response_structure(self):
|
||||
"""Test the expected response structure for chart preview updates."""
|
||||
# The response should contain these fields
|
||||
expected_response = {
|
||||
"chart": {
|
||||
"id": None, # No ID for unsaved previews
|
||||
"slice_name": str,
|
||||
"viz_type": str,
|
||||
"url": str,
|
||||
"uuid": None, # No UUID for unsaved previews
|
||||
"saved": bool,
|
||||
"updated": bool,
|
||||
},
|
||||
"error": None,
|
||||
"success": bool,
|
||||
"schema_version": str,
|
||||
"api_version": str,
|
||||
}
|
||||
|
||||
# When preview update succeeds, these additional fields may be present
|
||||
optional_fields = [
|
||||
"previews",
|
||||
"capabilities",
|
||||
"semantics",
|
||||
"explore_url",
|
||||
"form_data_key",
|
||||
"previous_form_data_key",
|
||||
"api_endpoints",
|
||||
"performance",
|
||||
"accessibility",
|
||||
]
|
||||
|
||||
# Validate structure expectations
|
||||
assert "chart" in expected_response
|
||||
assert "success" in expected_response
|
||||
assert len(optional_fields) > 0
|
||||
assert expected_response["chart"]["id"] is None
|
||||
assert expected_response["chart"]["uuid"] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_axis_configurations(self):
|
||||
"""Test axis configuration updates in preview."""
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="date"),
|
||||
y=[ColumnRef(name="sales")],
|
||||
x_axis=AxisConfig(
|
||||
title="Date",
|
||||
format="smart_date",
|
||||
scale="linear",
|
||||
),
|
||||
y_axis=AxisConfig(
|
||||
title="Sales Amount",
|
||||
format="$,.2f",
|
||||
scale="log",
|
||||
),
|
||||
)
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=config
|
||||
)
|
||||
assert request.config.x_axis.title == "Date"
|
||||
assert request.config.x_axis.format == "smart_date"
|
||||
assert request.config.y_axis.title == "Sales Amount"
|
||||
assert request.config.y_axis.format == "$,.2f"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_legend_configurations(self):
|
||||
"""Test legend configuration updates in preview."""
|
||||
positions = ["top", "bottom", "left", "right"]
|
||||
for pos in positions:
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="x"),
|
||||
y=[ColumnRef(name="y")],
|
||||
legend=LegendConfig(show=True, position=pos),
|
||||
)
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=config
|
||||
)
|
||||
assert request.config.legend.position == pos
|
||||
assert request.config.legend.show is True
|
||||
|
||||
# Hidden legend
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="x"),
|
||||
y=[ColumnRef(name="y")],
|
||||
legend=LegendConfig(show=False),
|
||||
)
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=config
|
||||
)
|
||||
assert request.config.legend.show is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_aggregation_functions(self):
|
||||
"""Test all supported aggregation functions in preview updates."""
|
||||
aggs = ["SUM", "AVG", "COUNT", "MIN", "MAX", "COUNT_DISTINCT"]
|
||||
for agg in aggs:
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="value", aggregate=agg)],
|
||||
)
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=config
|
||||
)
|
||||
assert request.config.columns[0].aggregate == agg
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_error_responses(self):
|
||||
"""Test expected error response structures for preview updates."""
|
||||
# General update error
|
||||
error_response = {
|
||||
"chart": None,
|
||||
"error": "Chart preview update failed: Invalid form_data_key",
|
||||
"success": False,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
assert error_response["success"] is False
|
||||
assert error_response["chart"] is None
|
||||
assert "failed" in error_response["error"].lower()
|
||||
|
||||
# Missing dataset error
|
||||
dataset_error = {
|
||||
"chart": None,
|
||||
"error": "Chart preview update failed: Dataset not found",
|
||||
"success": False,
|
||||
"schema_version": "2.0",
|
||||
"api_version": "v1",
|
||||
}
|
||||
assert dataset_error["success"] is False
|
||||
assert "dataset" in dataset_error["error"].lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_with_filters(self):
|
||||
"""Test updating preview with various filter configurations."""
|
||||
filters = [
|
||||
FilterConfig(column="region", op="=", value="North"),
|
||||
FilterConfig(column="sales", op=">=", value=1000),
|
||||
FilterConfig(column="date", op=">", value="2024-01-01"),
|
||||
]
|
||||
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[
|
||||
ColumnRef(name="region"),
|
||||
ColumnRef(name="sales"),
|
||||
ColumnRef(name="date"),
|
||||
],
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=config
|
||||
)
|
||||
assert len(request.config.filters) == 3
|
||||
assert request.config.filters[0].column == "region"
|
||||
assert request.config.filters[1].op == ">="
|
||||
assert request.config.filters[2].value == "2024-01-01"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_form_data_key_handling(self):
|
||||
"""Test form_data_key handling in preview updates."""
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
)
|
||||
|
||||
# Various form_data_key formats
|
||||
form_data_keys = [
|
||||
"abc123def456",
|
||||
"xyz-789-ghi-012",
|
||||
"key_with_underscore",
|
||||
"UPPERCASE_KEY",
|
||||
]
|
||||
|
||||
for key in form_data_keys:
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key=key, dataset_id=1, config=config
|
||||
)
|
||||
assert request.form_data_key == key
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_cache_control(self):
|
||||
"""Test cache control parameters in update preview request."""
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
)
|
||||
|
||||
# Default cache settings
|
||||
request1 = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=config
|
||||
)
|
||||
assert request1.use_cache is True
|
||||
assert request1.force_refresh is False
|
||||
assert request1.cache_form_data is True
|
||||
|
||||
# Custom cache settings
|
||||
request2 = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123",
|
||||
dataset_id=1,
|
||||
config=config,
|
||||
use_cache=False,
|
||||
force_refresh=True,
|
||||
cache_form_data=False,
|
||||
)
|
||||
assert request2.use_cache is False
|
||||
assert request2.force_refresh is True
|
||||
assert request2.cache_form_data is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_no_save_behavior(self):
|
||||
"""Test that preview updates don't create permanent charts."""
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[ColumnRef(name="col1")],
|
||||
)
|
||||
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=config
|
||||
)
|
||||
|
||||
# Preview updates should never create permanent charts
|
||||
# This is validated by checking the response structure
|
||||
expected_unsaved_fields = {
|
||||
"id": None, # No chart ID
|
||||
"uuid": None, # No UUID
|
||||
"saved": False, # Not saved
|
||||
}
|
||||
|
||||
# These expectations are validated in the response, not the request
|
||||
assert request.form_data_key == "abc123"
|
||||
assert expected_unsaved_fields["id"] is None
|
||||
assert expected_unsaved_fields["uuid"] is None
|
||||
assert expected_unsaved_fields["saved"] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_multiple_y_columns(self):
|
||||
"""Test preview updates with multiple Y-axis columns."""
|
||||
config = XYChartConfig(
|
||||
chart_type="xy",
|
||||
x=ColumnRef(name="date"),
|
||||
y=[
|
||||
ColumnRef(name="revenue", aggregate="SUM", label="Total Revenue"),
|
||||
ColumnRef(name="cost", aggregate="SUM", label="Total Cost"),
|
||||
ColumnRef(name="profit", aggregate="SUM", label="Total Profit"),
|
||||
ColumnRef(name="orders", aggregate="COUNT", label="Order Count"),
|
||||
],
|
||||
kind="line",
|
||||
)
|
||||
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=config
|
||||
)
|
||||
assert len(request.config.y) == 4
|
||||
assert request.config.y[0].name == "revenue"
|
||||
assert request.config.y[1].name == "cost"
|
||||
assert request.config.y[2].name == "profit"
|
||||
assert request.config.y[3].name == "orders"
|
||||
assert request.config.y[3].aggregate == "COUNT"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_chart_preview_table_sorting(self):
|
||||
"""Test table chart sorting in preview updates."""
|
||||
config = TableChartConfig(
|
||||
chart_type="table",
|
||||
columns=[
|
||||
ColumnRef(name="region"),
|
||||
ColumnRef(name="sales", aggregate="SUM"),
|
||||
ColumnRef(name="profit", aggregate="AVG"),
|
||||
],
|
||||
sort_by=["sales", "profit"],
|
||||
)
|
||||
|
||||
request = UpdateChartPreviewRequest(
|
||||
form_data_key="abc123", dataset_id=1, config=config
|
||||
)
|
||||
assert request.config.sort_by == ["sales", "profit"]
|
||||
assert len(request.config.columns) == 3
|
||||
23
tests/unit_tests/mcp_service/conftest.py
Normal file
23
tests/unit_tests/mcp_service/conftest.py
Normal file
@@ -0,0 +1,23 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
MCP service test configuration.
|
||||
|
||||
Tool imports are handled by app.py, not here.
|
||||
This conftest is empty to prevent test pollution.
|
||||
"""
|
||||
16
tests/unit_tests/mcp_service/dashboard/__init__.py
Normal file
16
tests/unit_tests/mcp_service/dashboard/__init__.py
Normal file
@@ -0,0 +1,16 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
16
tests/unit_tests/mcp_service/dashboard/tool/__init__.py
Normal file
16
tests/unit_tests/mcp_service/dashboard/tool/__init__.py
Normal file
@@ -0,0 +1,16 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
@@ -0,0 +1,450 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Unit tests for dashboard generation MCP tools
|
||||
"""
|
||||
|
||||
import logging
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
from fastmcp import Client
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.utils import json
|
||||
|
||||
logging.basicConfig(level=logging.DEBUG)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mcp_server():
|
||||
return mcp
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_auth():
|
||||
"""Mock authentication for all tests."""
|
||||
with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user:
|
||||
mock_user = Mock()
|
||||
mock_user.id = 1
|
||||
mock_user.username = "admin"
|
||||
mock_get_user.return_value = mock_user
|
||||
yield mock_get_user
|
||||
|
||||
|
||||
def _mock_chart(id: int = 1, slice_name: str = "Test Chart") -> Mock:
|
||||
"""Create a mock chart object."""
|
||||
chart = Mock()
|
||||
chart.id = id
|
||||
chart.slice_name = slice_name
|
||||
chart.uuid = f"chart-uuid-{id}"
|
||||
return chart
|
||||
|
||||
|
||||
def _mock_dashboard(id: int = 1, title: str = "Test Dashboard") -> Mock:
|
||||
"""Create a mock dashboard object."""
|
||||
dashboard = Mock()
|
||||
dashboard.id = id
|
||||
dashboard.dashboard_title = title
|
||||
dashboard.slug = f"test-dashboard-{id}"
|
||||
dashboard.description = "Test dashboard description"
|
||||
dashboard.published = True
|
||||
dashboard.created_on = "2024-01-01"
|
||||
dashboard.changed_on = "2024-01-01"
|
||||
dashboard.created_by = Mock()
|
||||
dashboard.created_by.username = "test_user"
|
||||
dashboard.changed_by = Mock()
|
||||
dashboard.changed_by.username = "test_user"
|
||||
dashboard.uuid = f"dashboard-uuid-{id}"
|
||||
dashboard.slices = []
|
||||
dashboard.owners = [] # Add missing owners attribute
|
||||
dashboard.tags = [] # Add missing tags attribute
|
||||
return dashboard
|
||||
|
||||
|
||||
class TestGenerateDashboard:
|
||||
"""Tests for generate_dashboard MCP tool."""
|
||||
|
||||
@patch("superset.commands.dashboard.create.CreateDashboardCommand")
|
||||
@patch("superset.db.session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_dashboard_basic(
|
||||
self, mock_db_session, mock_create_command, mcp_server
|
||||
):
|
||||
"""Test basic dashboard generation with valid charts."""
|
||||
# Mock database query for charts
|
||||
mock_query = Mock()
|
||||
mock_filter = Mock()
|
||||
mock_query.filter.return_value = mock_filter
|
||||
mock_filter.all.return_value = [
|
||||
_mock_chart(id=1, slice_name="Sales Chart"),
|
||||
_mock_chart(id=2, slice_name="Revenue Chart"),
|
||||
]
|
||||
mock_db_session.query.return_value = mock_query
|
||||
|
||||
# Mock dashboard creation
|
||||
mock_dashboard = _mock_dashboard(id=10, title="Analytics Dashboard")
|
||||
mock_create_command.return_value.run.return_value = mock_dashboard
|
||||
|
||||
request = {
|
||||
"chart_ids": [1, 2],
|
||||
"dashboard_title": "Analytics Dashboard",
|
||||
"description": "Dashboard for analytics",
|
||||
"published": True,
|
||||
}
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool("generate_dashboard", {"request": request})
|
||||
|
||||
assert result.data.error is None
|
||||
assert result.data.dashboard is not None
|
||||
assert result.data.dashboard.id == 10
|
||||
assert result.data.dashboard.dashboard_title == "Analytics Dashboard"
|
||||
assert result.data.dashboard.chart_count == 2
|
||||
assert "/superset/dashboard/10/" in result.data.dashboard_url
|
||||
|
||||
@patch("superset.db.session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_dashboard_missing_charts(self, mock_db_session, mcp_server):
|
||||
"""Test error handling when some charts don't exist."""
|
||||
# Mock database query returning only chart 1 (chart 2 missing)
|
||||
mock_query = Mock()
|
||||
mock_filter = Mock()
|
||||
mock_query.filter.return_value = mock_filter
|
||||
mock_filter.all.return_value = [
|
||||
_mock_chart(id=1),
|
||||
# Chart 2 is missing from the result
|
||||
]
|
||||
mock_db_session.query.return_value = mock_query
|
||||
|
||||
request = {"chart_ids": [1, 2], "dashboard_title": "Test Dashboard"}
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool("generate_dashboard", {"request": request})
|
||||
|
||||
assert result.data.error is not None
|
||||
assert "Charts not found: [2]" in result.data.error
|
||||
assert result.data.dashboard is None
|
||||
assert result.data.dashboard_url is None
|
||||
|
||||
@patch("superset.commands.dashboard.create.CreateDashboardCommand")
|
||||
@patch("superset.db.session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_dashboard_single_chart(
|
||||
self, mock_db_session, mock_create_command, mcp_server
|
||||
):
|
||||
"""Test dashboard generation with a single chart."""
|
||||
# Mock database query for single chart
|
||||
mock_query = Mock()
|
||||
mock_filter = Mock()
|
||||
mock_query.filter.return_value = mock_filter
|
||||
mock_filter.all.return_value = [_mock_chart(id=5, slice_name="Single Chart")]
|
||||
mock_db_session.query.return_value = mock_query
|
||||
|
||||
mock_dashboard = _mock_dashboard(id=20, title="Single Chart Dashboard")
|
||||
mock_create_command.return_value.run.return_value = mock_dashboard
|
||||
|
||||
request = {
|
||||
"chart_ids": [5],
|
||||
"dashboard_title": "Single Chart Dashboard",
|
||||
"published": False,
|
||||
}
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool("generate_dashboard", {"request": request})
|
||||
|
||||
assert result.data.error is None
|
||||
assert result.data.dashboard.chart_count == 1
|
||||
assert result.data.dashboard.published is True # From mock
|
||||
|
||||
@patch("superset.commands.dashboard.create.CreateDashboardCommand")
|
||||
@patch("superset.db.session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_dashboard_many_charts(
|
||||
self, mock_db_session, mock_create_command, mcp_server
|
||||
):
|
||||
"""Test dashboard generation with many charts (grid layout)."""
|
||||
# Mock 6 charts
|
||||
chart_ids = list(range(1, 7))
|
||||
mock_query = Mock()
|
||||
mock_filter = Mock()
|
||||
mock_query.filter.return_value = mock_filter
|
||||
mock_filter.all.return_value = [
|
||||
_mock_chart(id=i, slice_name=f"Chart {i}") for i in chart_ids
|
||||
]
|
||||
mock_db_session.query.return_value = mock_query
|
||||
|
||||
mock_dashboard = _mock_dashboard(id=30, title="Multi Chart Dashboard")
|
||||
mock_create_command.return_value.run.return_value = mock_dashboard
|
||||
|
||||
request = {"chart_ids": chart_ids, "dashboard_title": "Multi Chart Dashboard"}
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool("generate_dashboard", {"request": request})
|
||||
|
||||
assert result.data.error is None
|
||||
assert result.data.dashboard.chart_count == 6
|
||||
|
||||
# Verify CreateDashboardCommand was called with proper layout
|
||||
mock_create_command.assert_called_once()
|
||||
call_args = mock_create_command.call_args[0][0]
|
||||
|
||||
# Check position_json contains proper layout
|
||||
position_json = json.loads(call_args["position_json"])
|
||||
assert "ROOT_ID" in position_json
|
||||
assert "GRID_ID" in position_json
|
||||
assert "DASHBOARD_VERSION_KEY" in position_json
|
||||
assert position_json["DASHBOARD_VERSION_KEY"] == "v2"
|
||||
|
||||
# ROOT should only contain GRID
|
||||
assert position_json["ROOT_ID"]["children"] == ["GRID_ID"]
|
||||
|
||||
# GRID should contain rows (6 charts = 3 rows in 2-chart layout)
|
||||
grid_children = position_json["GRID_ID"]["children"]
|
||||
assert len(grid_children) == 3
|
||||
|
||||
# Check each chart has proper structure
|
||||
for i, chart_id in enumerate(chart_ids):
|
||||
chart_key = f"CHART-{chart_id}"
|
||||
row_index = i // 2 # 2 charts per row
|
||||
row_key = f"ROW-{row_index}"
|
||||
|
||||
# Chart should exist
|
||||
assert chart_key in position_json
|
||||
chart_data = position_json[chart_key]
|
||||
assert chart_data["type"] == "CHART"
|
||||
assert "meta" in chart_data
|
||||
assert chart_data["meta"]["chartId"] == chart_id
|
||||
|
||||
# Row should exist and contain charts (up to 2 per row)
|
||||
assert row_key in position_json
|
||||
row_data = position_json[row_key]
|
||||
assert row_data["type"] == "ROW"
|
||||
assert chart_key in row_data["children"]
|
||||
|
||||
@patch("superset.commands.dashboard.create.CreateDashboardCommand")
|
||||
@patch("superset.db.session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_dashboard_creation_failure(
|
||||
self, mock_db_session, mock_create_command, mcp_server
|
||||
):
|
||||
"""Test error handling when dashboard creation fails."""
|
||||
mock_query = Mock()
|
||||
mock_filter = Mock()
|
||||
mock_query.filter.return_value = mock_filter
|
||||
mock_filter.all.return_value = [_mock_chart(id=1)]
|
||||
mock_db_session.query.return_value = mock_query
|
||||
mock_create_command.return_value.run.side_effect = Exception("Creation failed")
|
||||
|
||||
request = {"chart_ids": [1], "dashboard_title": "Failed Dashboard"}
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool("generate_dashboard", {"request": request})
|
||||
|
||||
assert result.data.error is not None
|
||||
assert "Failed to create dashboard" in result.data.error
|
||||
assert result.data.dashboard is None
|
||||
|
||||
@patch("superset.commands.dashboard.create.CreateDashboardCommand")
|
||||
@patch("superset.db.session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_dashboard_minimal_request(
|
||||
self, mock_db_session, mock_create_command, mcp_server
|
||||
):
|
||||
"""Test dashboard generation with minimal required parameters."""
|
||||
# Mock database query for single chart
|
||||
mock_query = Mock()
|
||||
mock_filter = Mock()
|
||||
mock_query.filter.return_value = mock_filter
|
||||
mock_filter.all.return_value = [_mock_chart(id=3)]
|
||||
mock_db_session.query.return_value = mock_query
|
||||
|
||||
mock_dashboard = _mock_dashboard(id=40, title="Minimal Dashboard")
|
||||
mock_create_command.return_value.run.return_value = mock_dashboard
|
||||
|
||||
request = {
|
||||
"chart_ids": [3],
|
||||
"dashboard_title": "Minimal Dashboard",
|
||||
# No description, published defaults to True
|
||||
}
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool("generate_dashboard", {"request": request})
|
||||
|
||||
assert result.data.error is None
|
||||
assert result.data.dashboard.dashboard_title == "Minimal Dashboard"
|
||||
|
||||
# Check that description was not included in call
|
||||
call_args = mock_create_command.call_args[0][0]
|
||||
assert call_args["published"] is True # Default value
|
||||
assert (
|
||||
"description" not in call_args or call_args.get("description") is None
|
||||
)
|
||||
|
||||
|
||||
class TestAddChartToExistingDashboard:
|
||||
"""Tests for add_chart_to_existing_dashboard MCP tool."""
|
||||
|
||||
@patch("superset.commands.dashboard.update.UpdateDashboardCommand")
|
||||
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
|
||||
@patch("superset.db.session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_chart_to_dashboard_basic(
|
||||
self, mock_db_session, mock_find_dashboard, mock_update_command, mcp_server
|
||||
):
|
||||
"""Test adding a chart to an existing dashboard."""
|
||||
# Mock existing dashboard with some charts
|
||||
mock_dashboard = _mock_dashboard(id=1, title="Existing Dashboard")
|
||||
mock_dashboard.slices = [Mock(id=10), Mock(id=20)] # Existing charts
|
||||
mock_dashboard.position_json = json.dumps(
|
||||
{
|
||||
"ROOT_ID": {
|
||||
"children": ["CHART-10", "CHART-20"],
|
||||
"id": "ROOT_ID",
|
||||
"type": "ROOT",
|
||||
},
|
||||
"CHART-10": {"id": "CHART-10", "type": "CHART", "parents": ["ROOT_ID"]},
|
||||
"CHART-10_POSITION": {"h": 16, "w": 24, "x": 0, "y": 0},
|
||||
"CHART-20": {"id": "CHART-20", "type": "CHART", "parents": ["ROOT_ID"]},
|
||||
"CHART-20_POSITION": {"h": 16, "w": 24, "x": 24, "y": 0},
|
||||
}
|
||||
)
|
||||
mock_find_dashboard.return_value = mock_dashboard
|
||||
|
||||
# Mock chart to add
|
||||
mock_chart = _mock_chart(id=30, slice_name="New Chart")
|
||||
mock_db_session.get.return_value = mock_chart
|
||||
|
||||
# Mock updated dashboard
|
||||
updated_dashboard = _mock_dashboard(id=1, title="Existing Dashboard")
|
||||
updated_dashboard.slices = [Mock(id=10), Mock(id=20), Mock(id=30)]
|
||||
mock_update_command.return_value.run.return_value = updated_dashboard
|
||||
|
||||
request = {"dashboard_id": 1, "chart_id": 30}
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool(
|
||||
"add_chart_to_existing_dashboard", {"request": request}
|
||||
)
|
||||
|
||||
assert result.data.error is None
|
||||
assert result.data.dashboard is not None
|
||||
assert result.data.dashboard.chart_count == 3
|
||||
assert result.data.position is not None
|
||||
assert "row" in result.data.position # Should have row info
|
||||
assert "chart_key" in result.data.position
|
||||
assert "/superset/dashboard/1/" in result.data.dashboard_url
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_chart_dashboard_not_found(self, mock_find_dashboard, mcp_server):
|
||||
"""Test error when dashboard doesn't exist."""
|
||||
mock_find_dashboard.return_value = None
|
||||
|
||||
request = {"dashboard_id": 999, "chart_id": 1}
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool(
|
||||
"add_chart_to_existing_dashboard", {"request": request}
|
||||
)
|
||||
|
||||
assert result.data.error is not None
|
||||
assert "Dashboard with ID 999 not found" in result.data.error
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
|
||||
@patch("superset.db.session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_chart_chart_not_found(
|
||||
self, mock_db_session, mock_find_dashboard, mcp_server
|
||||
):
|
||||
"""Test error when chart doesn't exist."""
|
||||
mock_find_dashboard.return_value = _mock_dashboard()
|
||||
mock_db_session.get.return_value = None
|
||||
|
||||
request = {"dashboard_id": 1, "chart_id": 999}
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool(
|
||||
"add_chart_to_existing_dashboard", {"request": request}
|
||||
)
|
||||
|
||||
assert result.data.error is not None
|
||||
assert "Chart with ID 999 not found" in result.data.error
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
|
||||
@patch("superset.db.session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_chart_already_in_dashboard(
|
||||
self, mock_db_session, mock_find_dashboard, mcp_server
|
||||
):
|
||||
"""Test error when chart is already in dashboard."""
|
||||
mock_dashboard = _mock_dashboard()
|
||||
mock_dashboard.slices = [Mock(id=5)] # Chart 5 already exists
|
||||
mock_find_dashboard.return_value = mock_dashboard
|
||||
|
||||
mock_db_session.get.return_value = _mock_chart(id=5)
|
||||
|
||||
request = {"dashboard_id": 1, "chart_id": 5}
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool(
|
||||
"add_chart_to_existing_dashboard", {"request": request}
|
||||
)
|
||||
|
||||
assert result.data.error is not None
|
||||
assert "Chart 5 is already in dashboard 1" in result.data.error
|
||||
|
||||
@patch("superset.commands.dashboard.update.UpdateDashboardCommand")
|
||||
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
|
||||
@patch("superset.db.session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_chart_empty_dashboard(
|
||||
self, mock_db_session, mock_find_dashboard, mock_update_command, mcp_server
|
||||
):
|
||||
"""Test adding chart to dashboard with no existing layout."""
|
||||
mock_dashboard = _mock_dashboard(id=2)
|
||||
mock_dashboard.slices = []
|
||||
mock_dashboard.position_json = "{}" # Empty layout
|
||||
mock_find_dashboard.return_value = mock_dashboard
|
||||
|
||||
mock_chart = _mock_chart(id=15)
|
||||
mock_db_session.get.return_value = mock_chart
|
||||
|
||||
updated_dashboard = _mock_dashboard(id=2)
|
||||
updated_dashboard.slices = [Mock(id=15)]
|
||||
mock_update_command.return_value.run.return_value = updated_dashboard
|
||||
|
||||
request = {"dashboard_id": 2, "chart_id": 15}
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool(
|
||||
"add_chart_to_existing_dashboard", {"request": request}
|
||||
)
|
||||
|
||||
assert result.data.error is None
|
||||
assert "row" in result.data.position # Should have row info
|
||||
assert result.data.position.get("row") == 0 # First row
|
||||
|
||||
# Verify update was called with proper layout structure
|
||||
call_args = mock_update_command.call_args[0][1]
|
||||
layout = json.loads(call_args["position_json"])
|
||||
assert "ROOT_ID" in layout
|
||||
assert "GRID_ID" in layout
|
||||
assert "ROW-0" in layout
|
||||
assert "CHART-15" in layout
|
||||
@@ -0,0 +1,573 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Unit tests for MCP dashboard tools (list_dashboards, get_dashboard_info)
|
||||
"""
|
||||
|
||||
import logging
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
from fastmcp import Client
|
||||
from fastmcp.exceptions import ToolError
|
||||
|
||||
from superset.mcp_service.app import mcp
|
||||
from superset.mcp_service.dashboard.schemas import (
|
||||
ListDashboardsRequest,
|
||||
)
|
||||
|
||||
logging.basicConfig(level=logging.DEBUG)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mcp_server():
|
||||
return mcp
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_auth():
|
||||
"""Mock authentication for all tests."""
|
||||
with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user:
|
||||
mock_user = Mock()
|
||||
mock_user.id = 1
|
||||
mock_user.username = "admin"
|
||||
mock_get_user.return_value = mock_user
|
||||
yield mock_get_user
|
||||
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.list")
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_dashboards_basic(mock_list, mcp_server):
|
||||
dashboard = Mock()
|
||||
dashboard.id = 1
|
||||
dashboard.dashboard_title = "Test Dashboard"
|
||||
dashboard.slug = "test-dashboard"
|
||||
dashboard.url = "/dashboard/1"
|
||||
dashboard.published = True
|
||||
dashboard.changed_by_name = "admin"
|
||||
dashboard.changed_on = None
|
||||
dashboard.changed_on_humanized = None
|
||||
dashboard.created_by_name = "admin"
|
||||
dashboard.created_on = None
|
||||
dashboard.created_on_humanized = None
|
||||
dashboard.tags = []
|
||||
dashboard.owners = []
|
||||
dashboard.slices = []
|
||||
dashboard.description = None
|
||||
dashboard.css = None
|
||||
dashboard.certified_by = None
|
||||
dashboard.certification_details = None
|
||||
dashboard.json_metadata = None
|
||||
dashboard.position_json = None
|
||||
dashboard.is_managed_externally = False
|
||||
dashboard.external_url = None
|
||||
dashboard.uuid = "test-dashboard-uuid-1"
|
||||
dashboard.thumbnail_url = None
|
||||
dashboard.roles = []
|
||||
dashboard.charts = []
|
||||
dashboard._mapping = {
|
||||
"id": dashboard.id,
|
||||
"dashboard_title": dashboard.dashboard_title,
|
||||
"slug": dashboard.slug,
|
||||
"url": dashboard.url,
|
||||
"published": dashboard.published,
|
||||
"changed_by_name": dashboard.changed_by_name,
|
||||
"changed_on": dashboard.changed_on,
|
||||
"changed_on_humanized": dashboard.changed_on_humanized,
|
||||
"created_by_name": dashboard.created_by_name,
|
||||
"created_on": dashboard.created_on,
|
||||
"created_on_humanized": dashboard.created_on_humanized,
|
||||
"tags": dashboard.tags,
|
||||
"owners": dashboard.owners,
|
||||
"charts": [],
|
||||
}
|
||||
mock_list.return_value = ([dashboard], 1)
|
||||
async with Client(mcp_server) as client:
|
||||
request = ListDashboardsRequest(page=1, page_size=10)
|
||||
result = await client.call_tool(
|
||||
"list_dashboards", {"request": request.model_dump()}
|
||||
)
|
||||
dashboards = result.data.dashboards
|
||||
assert len(dashboards) == 1
|
||||
assert dashboards[0].dashboard_title == "Test Dashboard"
|
||||
assert dashboards[0].uuid == "test-dashboard-uuid-1"
|
||||
assert dashboards[0].slug == "test-dashboard"
|
||||
assert dashboards[0].published is True
|
||||
|
||||
# Verify UUID and slug are in default columns
|
||||
assert "uuid" in result.data.columns_requested
|
||||
assert "slug" in result.data.columns_requested
|
||||
assert "uuid" in result.data.columns_loaded
|
||||
assert "slug" in result.data.columns_loaded
|
||||
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.list")
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_dashboards_with_filters(mock_list, mcp_server):
|
||||
dashboard = Mock()
|
||||
dashboard.id = 1
|
||||
dashboard.dashboard_title = "Filtered Dashboard"
|
||||
dashboard.slug = "filtered-dashboard"
|
||||
dashboard.url = "/dashboard/2"
|
||||
dashboard.published = True
|
||||
dashboard.changed_by_name = "admin"
|
||||
dashboard.changed_on = None
|
||||
dashboard.changed_on_humanized = None
|
||||
dashboard.created_by_name = "admin"
|
||||
dashboard.created_on = None
|
||||
dashboard.created_on_humanized = None
|
||||
dashboard.tags = []
|
||||
dashboard.owners = []
|
||||
dashboard.slices = []
|
||||
dashboard.description = None
|
||||
dashboard.css = None
|
||||
dashboard.certified_by = None
|
||||
dashboard.certification_details = None
|
||||
dashboard.json_metadata = None
|
||||
dashboard.position_json = None
|
||||
dashboard.is_managed_externally = False
|
||||
dashboard.external_url = None
|
||||
dashboard.uuid = None
|
||||
dashboard.thumbnail_url = None
|
||||
dashboard.roles = []
|
||||
dashboard.charts = []
|
||||
dashboard._mapping = {
|
||||
"id": dashboard.id,
|
||||
"dashboard_title": dashboard.dashboard_title,
|
||||
"slug": dashboard.slug,
|
||||
"url": dashboard.url,
|
||||
"published": dashboard.published,
|
||||
"changed_by_name": dashboard.changed_by_name,
|
||||
"changed_on": dashboard.changed_on,
|
||||
"changed_on_humanized": dashboard.changed_on_humanized,
|
||||
"created_by_name": dashboard.created_by_name,
|
||||
"created_on": dashboard.created_on,
|
||||
"created_on_humanized": dashboard.created_on_humanized,
|
||||
"tags": dashboard.tags,
|
||||
"owners": dashboard.owners,
|
||||
"charts": [],
|
||||
}
|
||||
mock_list.return_value = ([dashboard], 1)
|
||||
async with Client(mcp_server) as client:
|
||||
filters = [
|
||||
{"col": "dashboard_title", "opr": "sw", "value": "Sales"},
|
||||
{"col": "published", "opr": "eq", "value": True},
|
||||
]
|
||||
request = ListDashboardsRequest(
|
||||
filters=filters,
|
||||
select_columns=["id", "dashboard_title"],
|
||||
order_column="changed_on",
|
||||
order_direction="desc",
|
||||
page=1,
|
||||
page_size=50,
|
||||
)
|
||||
result = await client.call_tool(
|
||||
"list_dashboards", {"request": request.model_dump()}
|
||||
)
|
||||
assert result.data.count == 1
|
||||
assert result.data.dashboards[0].dashboard_title == "Filtered Dashboard"
|
||||
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.list")
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_dashboards_with_string_filters(mock_list, mcp_server):
|
||||
mock_list.return_value = ([], 0)
|
||||
async with Client(mcp_server) as client: # noqa: F841
|
||||
filters = '[{"col": "dashboard_title", "opr": "sw", "value": "Sales"}]'
|
||||
|
||||
# Test that string filters cause validation error at schema level
|
||||
with pytest.raises(ValueError, match="validation error"):
|
||||
ListDashboardsRequest(filters=filters) # noqa: F841
|
||||
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.list")
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_dashboards_api_error(mock_list, mcp_server):
|
||||
mock_list.side_effect = ToolError("API request failed")
|
||||
async with Client(mcp_server) as client:
|
||||
with pytest.raises(ToolError) as excinfo: # noqa: PT012
|
||||
request = ListDashboardsRequest()
|
||||
await client.call_tool("list_dashboards", {"request": request.model_dump()})
|
||||
assert "API request failed" in str(excinfo.value)
|
||||
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.list")
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_dashboards_with_search(mock_list, mcp_server):
|
||||
dashboard = Mock()
|
||||
dashboard.id = 1
|
||||
dashboard.dashboard_title = "search_dashboard"
|
||||
dashboard.slug = "search-dashboard"
|
||||
dashboard.url = "/dashboard/1"
|
||||
dashboard.published = True
|
||||
dashboard.changed_by_name = "admin"
|
||||
dashboard.changed_on = None
|
||||
dashboard.changed_on_humanized = None
|
||||
dashboard.created_by_name = "admin"
|
||||
dashboard.created_on = None
|
||||
dashboard.created_on_humanized = None
|
||||
dashboard.tags = []
|
||||
dashboard.owners = []
|
||||
dashboard.slices = []
|
||||
dashboard.description = None
|
||||
dashboard.css = None
|
||||
dashboard.certified_by = None
|
||||
dashboard.certification_details = None
|
||||
dashboard.json_metadata = None
|
||||
dashboard.position_json = None
|
||||
dashboard.is_managed_externally = False
|
||||
dashboard.external_url = None
|
||||
dashboard.uuid = None
|
||||
dashboard.thumbnail_url = None
|
||||
dashboard.roles = []
|
||||
dashboard.charts = []
|
||||
dashboard._mapping = {
|
||||
"id": dashboard.id,
|
||||
"dashboard_title": dashboard.dashboard_title,
|
||||
"slug": dashboard.slug,
|
||||
"url": dashboard.url,
|
||||
"published": dashboard.published,
|
||||
"changed_by_name": dashboard.changed_by_name,
|
||||
"changed_on": dashboard.changed_on,
|
||||
"changed_on_humanized": dashboard.changed_on_humanized,
|
||||
"created_by_name": dashboard.created_by_name,
|
||||
"created_on": dashboard.created_on,
|
||||
"created_on_humanized": dashboard.created_on_humanized,
|
||||
"tags": dashboard.tags,
|
||||
"owners": dashboard.owners,
|
||||
"charts": [],
|
||||
}
|
||||
mock_list.return_value = ([dashboard], 1)
|
||||
async with Client(mcp_server) as client:
|
||||
request = ListDashboardsRequest(search="search_dashboard")
|
||||
result = await client.call_tool(
|
||||
"list_dashboards", {"request": request.model_dump()}
|
||||
)
|
||||
assert result.data.count == 1
|
||||
assert result.data.dashboards[0].dashboard_title == "search_dashboard"
|
||||
args, kwargs = mock_list.call_args
|
||||
assert kwargs["search"] == "search_dashboard"
|
||||
assert "dashboard_title" in kwargs["search_columns"]
|
||||
assert "slug" in kwargs["search_columns"]
|
||||
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.list")
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_dashboards_with_simple_filters(mock_list, mcp_server):
|
||||
mock_list.return_value = ([], 0)
|
||||
async with Client(mcp_server) as client:
|
||||
filters = [
|
||||
{"col": "dashboard_title", "opr": "eq", "value": "Sales"},
|
||||
{"col": "published", "opr": "eq", "value": True},
|
||||
]
|
||||
request = ListDashboardsRequest(filters=filters)
|
||||
result = await client.call_tool(
|
||||
"list_dashboards", {"request": request.model_dump()}
|
||||
)
|
||||
assert hasattr(result.data, "count")
|
||||
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_dashboard_info_success(mock_info, mcp_server):
|
||||
dashboard = Mock()
|
||||
dashboard.id = 1
|
||||
dashboard.dashboard_title = "Test Dashboard"
|
||||
dashboard.slug = "test-dashboard"
|
||||
dashboard.description = "Test description"
|
||||
dashboard.css = None
|
||||
dashboard.certified_by = None
|
||||
dashboard.certification_details = None
|
||||
dashboard.json_metadata = None
|
||||
dashboard.position_json = None
|
||||
dashboard.published = True
|
||||
dashboard.is_managed_externally = False
|
||||
dashboard.external_url = None
|
||||
dashboard.created_on = None
|
||||
dashboard.changed_on = None
|
||||
dashboard.created_by = None
|
||||
dashboard.changed_by = None
|
||||
dashboard.uuid = None
|
||||
dashboard.url = "/dashboard/1"
|
||||
dashboard.thumbnail_url = None
|
||||
dashboard.created_on_humanized = None
|
||||
dashboard.changed_on_humanized = None
|
||||
dashboard.slices = []
|
||||
dashboard.owners = []
|
||||
dashboard.tags = []
|
||||
dashboard.roles = []
|
||||
dashboard.charts = []
|
||||
dashboard._mapping = {
|
||||
"id": dashboard.id,
|
||||
"dashboard_title": dashboard.dashboard_title,
|
||||
"slug": dashboard.slug,
|
||||
"url": dashboard.url,
|
||||
"published": dashboard.published,
|
||||
"changed_by_name": dashboard.changed_by_name,
|
||||
"changed_on": dashboard.changed_on,
|
||||
"changed_on_humanized": dashboard.changed_on_humanized,
|
||||
"created_by_name": dashboard.created_by_name,
|
||||
"created_on": dashboard.created_on,
|
||||
"created_on_humanized": dashboard.created_on_humanized,
|
||||
"tags": dashboard.tags,
|
||||
"owners": dashboard.owners,
|
||||
"charts": [],
|
||||
}
|
||||
mock_info.return_value = dashboard # Only the dashboard object
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool(
|
||||
"get_dashboard_info", {"request": {"identifier": 1}}
|
||||
)
|
||||
assert result.data["dashboard_title"] == "Test Dashboard"
|
||||
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_dashboard_info_not_found(mock_info, mcp_server):
|
||||
mock_info.return_value = None # Not found returns None
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool(
|
||||
"get_dashboard_info", {"request": {"identifier": 999}}
|
||||
)
|
||||
assert result.data["error_type"] == "not_found"
|
||||
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.find_by_id")
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_dashboard_info_access_denied(mock_info, mcp_server):
|
||||
mock_info.return_value = None # Access denied returns None
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool(
|
||||
"get_dashboard_info", {"request": {"identifier": 1}}
|
||||
)
|
||||
assert result.data["error_type"] == "not_found"
|
||||
|
||||
|
||||
# TODO (Phase 3+): Add tests for get_dashboard_available_filters tool
|
||||
|
||||
|
||||
@patch("superset.mcp_service.mcp_core.ModelGetInfoCore._find_object")
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_dashboard_info_by_uuid(mock_find_object, mcp_server):
|
||||
"""Test getting dashboard info using UUID identifier."""
|
||||
dashboard = Mock()
|
||||
dashboard.id = 1
|
||||
dashboard.dashboard_title = "Test Dashboard UUID"
|
||||
dashboard.slug = "test-dashboard-uuid"
|
||||
dashboard.description = "Test description"
|
||||
dashboard.css = ""
|
||||
dashboard.certified_by = None
|
||||
dashboard.certification_details = None
|
||||
dashboard.json_metadata = "{}"
|
||||
dashboard.position_json = "{}"
|
||||
dashboard.published = True
|
||||
dashboard.is_managed_externally = False
|
||||
dashboard.external_url = None
|
||||
dashboard.created_on = None
|
||||
dashboard.changed_on = None
|
||||
dashboard.created_by = None
|
||||
dashboard.changed_by = None
|
||||
dashboard.uuid = "c3d4e5f6-g7h8-9012-cdef-gh3456789012"
|
||||
dashboard.url = "/dashboard/1"
|
||||
dashboard.thumbnail_url = None
|
||||
dashboard.created_on_humanized = "2 days ago"
|
||||
dashboard.changed_on_humanized = "1 day ago"
|
||||
dashboard.slices = []
|
||||
dashboard.owners = []
|
||||
dashboard.tags = []
|
||||
dashboard.roles = []
|
||||
|
||||
mock_find_object.return_value = dashboard
|
||||
async with Client(mcp_server) as client:
|
||||
uuid_str = "c3d4e5f6-g7h8-9012-cdef-gh3456789012"
|
||||
result = await client.call_tool(
|
||||
"get_dashboard_info", {"request": {"identifier": uuid_str}}
|
||||
)
|
||||
assert result.data["dashboard_title"] == "Test Dashboard UUID"
|
||||
|
||||
|
||||
@patch("superset.mcp_service.mcp_core.ModelGetInfoCore._find_object")
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_dashboard_info_by_slug(mock_find_object, mcp_server):
|
||||
"""Test getting dashboard info using slug identifier."""
|
||||
dashboard = Mock()
|
||||
dashboard.id = 2
|
||||
dashboard.dashboard_title = "Test Dashboard Slug"
|
||||
dashboard.slug = "test-dashboard-slug"
|
||||
dashboard.description = "Test description"
|
||||
dashboard.css = ""
|
||||
dashboard.certified_by = None
|
||||
dashboard.certification_details = None
|
||||
dashboard.json_metadata = "{}"
|
||||
dashboard.position_json = "{}"
|
||||
dashboard.published = True
|
||||
dashboard.is_managed_externally = False
|
||||
dashboard.external_url = None
|
||||
dashboard.created_on = None
|
||||
dashboard.changed_on = None
|
||||
dashboard.created_by = None
|
||||
dashboard.changed_by = None
|
||||
dashboard.uuid = "d4e5f6g7-h8i9-0123-defg-hi4567890123"
|
||||
dashboard.url = "/dashboard/2"
|
||||
dashboard.thumbnail_url = None
|
||||
dashboard.created_on_humanized = "2 days ago"
|
||||
dashboard.changed_on_humanized = "1 day ago"
|
||||
dashboard.slices = []
|
||||
dashboard.owners = []
|
||||
dashboard.tags = []
|
||||
dashboard.roles = []
|
||||
|
||||
mock_find_object.return_value = dashboard
|
||||
async with Client(mcp_server) as client:
|
||||
result = await client.call_tool(
|
||||
"get_dashboard_info", {"request": {"identifier": "test-dashboard-slug"}}
|
||||
)
|
||||
assert result.data["dashboard_title"] == "Test Dashboard Slug"
|
||||
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.list")
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_dashboards_custom_uuid_slug_columns(mock_list, mcp_server):
|
||||
"""Test that custom column selection includes UUID and slug when explicitly
|
||||
requested."""
|
||||
dashboard = Mock()
|
||||
dashboard.id = 1
|
||||
dashboard.dashboard_title = "Custom Columns Dashboard"
|
||||
dashboard.slug = "custom-dashboard"
|
||||
dashboard.uuid = "test-custom-uuid-123"
|
||||
dashboard.url = "/dashboard/1"
|
||||
dashboard.published = True
|
||||
dashboard.changed_by_name = "admin"
|
||||
dashboard.changed_on = None
|
||||
dashboard.changed_on_humanized = None
|
||||
dashboard.created_by_name = "admin"
|
||||
dashboard.created_on = None
|
||||
dashboard.created_on_humanized = None
|
||||
dashboard.tags = []
|
||||
dashboard.owners = []
|
||||
dashboard.slices = []
|
||||
dashboard.description = None
|
||||
dashboard.css = None
|
||||
dashboard.certified_by = None
|
||||
dashboard.certification_details = None
|
||||
dashboard.json_metadata = None
|
||||
dashboard.position_json = None
|
||||
dashboard.is_managed_externally = False
|
||||
dashboard.external_url = None
|
||||
dashboard.thumbnail_url = None
|
||||
dashboard.roles = []
|
||||
dashboard.charts = []
|
||||
dashboard._mapping = {
|
||||
"id": dashboard.id,
|
||||
"dashboard_title": dashboard.dashboard_title,
|
||||
"slug": dashboard.slug,
|
||||
"uuid": dashboard.uuid,
|
||||
"url": dashboard.url,
|
||||
"published": dashboard.published,
|
||||
"changed_by_name": dashboard.changed_by_name,
|
||||
"changed_on": dashboard.changed_on,
|
||||
"changed_on_humanized": dashboard.changed_on_humanized,
|
||||
"created_by_name": dashboard.created_by_name,
|
||||
"created_on": dashboard.created_on,
|
||||
"created_on_humanized": dashboard.created_on_humanized,
|
||||
"tags": dashboard.tags,
|
||||
"owners": dashboard.owners,
|
||||
"charts": [],
|
||||
}
|
||||
mock_list.return_value = ([dashboard], 1)
|
||||
async with Client(mcp_server) as client:
|
||||
request = ListDashboardsRequest(
|
||||
select_columns=["id", "dashboard_title", "uuid", "slug"],
|
||||
page=1,
|
||||
page_size=10,
|
||||
)
|
||||
result = await client.call_tool(
|
||||
"list_dashboards", {"request": request.model_dump()}
|
||||
)
|
||||
dashboards = result.data.dashboards
|
||||
assert len(dashboards) == 1
|
||||
assert dashboards[0].uuid == "test-custom-uuid-123"
|
||||
assert dashboards[0].slug == "custom-dashboard"
|
||||
|
||||
# Verify custom columns include UUID and slug
|
||||
assert "uuid" in result.data.columns_requested
|
||||
assert "slug" in result.data.columns_requested
|
||||
assert "uuid" in result.data.columns_loaded
|
||||
assert "slug" in result.data.columns_loaded
|
||||
|
||||
|
||||
class TestDashboardSortableColumns:
|
||||
"""Test sortable columns configuration for dashboard tools."""
|
||||
|
||||
def test_dashboard_sortable_columns_definition(self):
|
||||
"""Test that dashboard sortable columns are properly defined."""
|
||||
from superset.mcp_service.dashboard.tool.list_dashboards import (
|
||||
SORTABLE_DASHBOARD_COLUMNS,
|
||||
)
|
||||
|
||||
assert SORTABLE_DASHBOARD_COLUMNS == [
|
||||
"id",
|
||||
"dashboard_title",
|
||||
"slug",
|
||||
"published",
|
||||
"changed_on",
|
||||
"created_on",
|
||||
]
|
||||
# Ensure no computed properties are included
|
||||
assert "changed_on_delta_humanized" not in SORTABLE_DASHBOARD_COLUMNS
|
||||
assert "changed_by_name" not in SORTABLE_DASHBOARD_COLUMNS
|
||||
assert "uuid" not in SORTABLE_DASHBOARD_COLUMNS
|
||||
|
||||
@patch("superset.daos.dashboard.DashboardDAO.list")
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_dashboards_with_valid_order_column(self, mock_list, mcp_server):
|
||||
"""Test list_dashboards with valid order column."""
|
||||
mock_list.return_value = ([], 0)
|
||||
|
||||
async with Client(mcp_server) as client:
|
||||
# Test with valid sortable column
|
||||
request = ListDashboardsRequest(
|
||||
order_column="dashboard_title", order_direction="desc"
|
||||
)
|
||||
result = await client.call_tool(
|
||||
"list_dashboards", {"request": request.model_dump()}
|
||||
)
|
||||
|
||||
# Verify the DAO was called with the correct order column
|
||||
mock_list.assert_called_once()
|
||||
call_args = mock_list.call_args[1]
|
||||
assert call_args["order_column"] == "dashboard_title"
|
||||
assert call_args["order_direction"] == "desc"
|
||||
|
||||
# Verify the result
|
||||
assert result.data.count == 0
|
||||
assert result.data.dashboards == []
|
||||
|
||||
def test_sortable_columns_in_docstring(self):
|
||||
"""Test that sortable columns are documented in tool docstring."""
|
||||
from superset.mcp_service.dashboard.tool.list_dashboards import (
|
||||
list_dashboards,
|
||||
SORTABLE_DASHBOARD_COLUMNS,
|
||||
)
|
||||
|
||||
# Check list_dashboards docstring (stored in description after @mcp.tool)
|
||||
assert hasattr(list_dashboards, "description")
|
||||
assert "Sortable columns for order_column:" in list_dashboards.description
|
||||
for col in SORTABLE_DASHBOARD_COLUMNS:
|
||||
assert col in list_dashboards.description
|
||||
16
tests/unit_tests/mcp_service/dataset/__init__.py
Normal file
16
tests/unit_tests/mcp_service/dataset/__init__.py
Normal file
@@ -0,0 +1,16 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
16
tests/unit_tests/mcp_service/dataset/tool/__init__.py
Normal file
16
tests/unit_tests/mcp_service/dataset/tool/__init__.py
Normal file
@@ -0,0 +1,16 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
1231
tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py
Normal file
1231
tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py
Normal file
File diff suppressed because it is too large
Load Diff
16
tests/unit_tests/mcp_service/explore/__init__.py
Normal file
16
tests/unit_tests/mcp_service/explore/__init__.py
Normal file
@@ -0,0 +1,16 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user