Security: Use OAuth and JWT to map/propagate permissions/privileges

[amotl]

Just a thought about GH-11: Is this the best approach?

Wouldn’t it be more robust to restrict the CrateDB user’s permissions directly used for the mcp server?
It feels like we’re partly re-implementing CrateDB’s built-in permission system externally here. Curious to hear the reasoning behind this approach! Also, could the MCP server instead check the user’s permissions directly on CrateDB?

Originally posted by @proddata in #11 (comment)

[amotl]

Is this the best approach?

Depending on how far we want to go or extend the implementation, this is probably just the first approach, coming from suggestions in GH-10.

It feels like we’re partly re-implementing CrateDB’s built-in permission system externally here.

Certainly not, at least we hope not too much. The code in GH-11 just aims to qualify an SQL statement if it has write impact or not, which might go orthogonal to any permission system? I have to admit I don't know too much about the details of the permission system in CrateDB, so please advise if you think details are redundant here.

Could the MCP server instead check the user’s permissions directly on CrateDB?

The MCP specification is evolving and will start to use OAuth as an identity and authorization layer to underpin MCP flows. A few more research and implementation cycles will need to be spent on this topic, optimally inheriting best-practice approaches from others in this regard, and contributions will certainly be needed to make this fly. 🍀

• Model Context Protocol (MCP) crate-clients-tools#234 (comment)

[proddata]

The MCP specification is evolving and will start to use OAuth as an identity and authorization layer to underpin MCP flows

If OAuth and JWTs are coming into play, that speaks even more for leveraging CrateDB’s permission system directly. The JWT could include the necessary permissions, so the MCP server can safely and cleanly pass through queries to the cluster without needing to maintain a separate access control layer.

[amotl]

That sounds like a sweet approach. We hope it will work out well, others might do it exactly like this already.

Any contributions are welcome, either per code (it is not that much code yet, so it's probably very easy to hack on it) or per research support. The topic probably needs a few minutes to read the MCP spec and/or research implementations from other vendors to use as blueprints.

[amotl]

@coderabbitai: Please submit a patch which sketches out how bringing in MCP's OAuth bindings together with CrateDB's JWT tokens for authentication and authorization purposes could look like, in order to map permissions properly and based on standards. Thanks!

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Get an overview of the repository structure
find . -type f -name "*.py" | sort

Length of output: 208

---

🏁 Script executed:

#!/bin/bash
# Check the content of the main application files
cat ./cratedb_mcp/cli.py

Length of output: 509

---

🏁 Script executed:

#!/bin/bash
# Check the main module file to understand the MCP server implementation
cat ./cratedb_mcp/__main__.py

Length of output: 2295

---

🏁 Script executed:

#!/bin/bash
# Check the knowledge module to understand the constants and queries
cat ./cratedb_mcp/knowledge.py

Length of output: 5153

---

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

[coderabbitai]

A summary of the edits CodeRabbit can apply:

Will create a TypeScript example file illustrating how to integrate MCP OAuth tools with CrateDB JWT authentication using the new Python handlers.

A plan of each step that CodeRabbit will take in order to execute the edits for each file (with inline rough draft snippets):

📄 mcp_oauth_cratedb_jwt_integration_example.ts

At the top of mcp_oauth_cratedb_jwt_integration_example.ts, import the Node.js child_process and readline modules for invoking the MCP CLI over stdio and reading user input:

import { spawn } from 'child_process';
import * as readline from 'readline';

---

Define a helper function to call MCP tools via the cratedb-mcp CLI over stdio. This function will serialize the JSON-RPC request, pipe it into the process, and parse the JSON-RPC response:

interface RPCRequest {
  id: string;
  method: string;
  params?: any;
}

interface RPCResponse {
  id: string;
  result?: any;
  error?: { message: string; code?: number };
}

function callMCP(method: string, params?: any): Promise<any> {
  return new Promise((resolve, reject) => {
    const proc = spawn('cratedb-mcp', [], { stdio: ['pipe', 'pipe', 'inherit'] });
    const request: RPCRequest = { id: method, method, params };
    let output = '';
    proc.stdout.on('data', chunk => { output += chunk.toString(); });
    proc.stdout.on('end', () => {
      try {
        const response: RPCResponse = JSON.parse(output);
        if (response.error) {
          reject(new Error(response.error.message));
        } else {
          resolve(response.result);
        }
      } catch (err) {
        reject(err);
      }
    });
    proc.stdin.write(JSON.stringify(request));
    proc.stdin.end();
  });
}

---

Set up a readline interface to prompt the user for the OAuth authorization code after they visit the authorization URL:

const rl = readline.createInterface({
  input: process.stdin,
  output: process.stdout
});
function prompt(question: string): Promise<string> {
  return new Promise(resolve => {
    rl.question(question, answer => {
      resolve(answer);
    });
  });
}

---

In an async main() function, orchestrate the OAuth flow:

1. Call initiate_oauth_flow to get the authorization URL.
2. Print the URL and instruct the user to open it.
3. Prompt the user to paste the resulting code.
4. Call exchange_oauth_code with the code to obtain tokens.
5. Store the returned access token for subsequent calls.

async function main() {
  try {
    // 1. Initiate OAuth flow
    const { authorization_url } = await callMCP('initiate_oauth_flow', { state: 'xyz' });
    console.log('Open this URL in your browser to authenticate:');
    console.log(authorization_url);

    // 2. Prompt for the code from the callback
    const code = await prompt('Enter the OAuth authorization code: ');
    rl.close();

    // 3. Exchange code for tokens
    const tokenInfo = await callMCP('exchange_oauth_code', { code });
    console.log('Authentication successful:', tokenInfo);

    // 4. Extract and attach the access_token for future calls
    const accessToken = (tokenInfo as any).access_token || '';
    
    // 5. Perform an authenticated query
    const result = await callMCP('query_sql', {
      query: 'SELECT * FROM my_table;',
      // Provide our token in the context metadata
      ctx: { metadata: { access_token: accessToken } }
    });
    console.log('Query result:', result);
  } catch (err) {
    console.error('Error during OAuth flow or querying:', err);
  }
}

main();

---

Add comments at the bottom of the file explaining prerequisites and usage:

// Prerequisites:
// 1. `cratedb-mcp` must be installed and in your PATH.
// 2. Set the necessary environment variables for OAuth and JWT as documented in README.
// 3. This script communicates with the MCP server over STDIO.
// 
// Usage:
// ts-node mcp_oauth_cratedb_jwt_integration_example.ts

---

Save and close the file. Test the script end-to-end by running it with ts-node or compiling to JavaScript and running with Node. Ensure the OAuth flow completes and authenticated queries succeed.

A summary of the context that CodeRabbit is considering across the codebase:

We need to understand the structure of the MCP Context object to know how to attach the access_token for authenticated queries in our TypeScript example.

• ✅ Create PR with these edits

• 📋 Get copyable edits

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

[coderabbitai]

A summary of the edits CodeRabbit can apply:

A new authentication module will be created at cratedb_mcp/auth.py to handle OAuth and JWT operations.

OAuth and JWT configuration options will be added to cratedb_mcp/settings.py.

cratedb_mcp/main.py will be modified to integrate AuthManager, update query functions, and import necessary modules.

A new authentication middleware will be created at cratedb_mcp/middleware.py.

Dependencies for OAuth and JWT support will be added to requirements.txt.

A plan of each step that CodeRabbit will take in order to execute the edits for each file (with inline rough draft snippets):

➕ cratedb_mcp/auth.py

Insert the following content into cratedb_mcp/auth.py:

"""
Authentication module for CrateDB MCP.

This module handles OAuth authentication and JWT token generation/validation
for secure communication with CrateDB.
"""
import time
from typing import Dict, Optional, Any
import jwt
import httpx
from authlib.integrations.httpx_client import OAuth2Client
from fastapi import Request, HTTPException

class AuthManager:
    """
    Manages authentication between MCP OAuth and CrateDB JWT tokens.
    
    This class is responsible for:
    1. Authenticating users via OAuth
    2. Generating JWT tokens for CrateDB based on OAuth identity
    3. Mapping OAuth scopes/claims to CrateDB permissions
    """
    
    def __init__(self, 
                 oauth_client_id: str,
                 oauth_client_secret: str, 
                 oauth_token_url: str,
                 oauth_authorize_url: str,
                 jwt_secret: str,
                 jwt_algorithm: str = "HS256",
                 jwt_expiry: int = 3600):
        """Initialize the auth manager with OAuth and JWT settings."""
        self.oauth_client = OAuth2Client(
            client_id=oauth_client_id,
            client_secret=oauth_client_secret,
            token_endpoint=oauth_token_url,
            authorize_endpoint=oauth_authorize_url
        )
        self.jwt_secret = jwt_secret
        self.jwt_algorithm = jwt_algorithm
        self.jwt_expiry = jwt_expiry
        
    async def validate_oauth_token(self, token: str) -> Dict[str, Any]:
        """Validate an OAuth token and return the user claims/info."""
        try:
            if not token or token == "invalid":
                raise HTTPException(status_code=401, detail="Invalid OAuth token")
            # In a full implementation, call userinfo endpoint here
            user_info = {
                "sub": "user123",
                "name": "Example User",
                "permissions": ["read:data", "write:data"],
                "roles": ["basic_user"]
            }
            return user_info
        except Exception as e:
            raise HTTPException(status_code=401, detail=f"OAuth validation failed: {str(e)}")
    
    def generate_cratedb_jwt(self, user_info: Dict[str, Any]) -> str:
        """
        Generate a JWT token for CrateDB based on OAuth user info.
        """
        now = int(time.time())
        cratedb_permissions = []
        if "read:data" in user_info.get("permissions", []):
            cratedb_permissions.append("DQL")
        if "write:data" in user_info.get("permissions", []):
            cratedb_permissions.append("DML")
        if "admin" in user_info.get("roles", []):
            cratedb_permissions.extend(["DDL", "AL"])
        payload = {
            "sub": user_info.get("sub", "unknown"),
            "name": user_info.get("name", "Unknown User"),
            "permissions": cratedb_permissions,
            "iat": now,
            "exp": now + self.jwt_expiry
        }
        token = jwt.encode(
            payload,
            self.jwt_secret,
            algorithm=self.jwt_algorithm
        )
        return token
    
    def get_cratedb_auth_headers(self, user_info: Dict[str, Any]) -> Dict[str, str]:
        """Generate CrateDB authorization headers with JWT."""
        token = self.generate_cratedb_jwt(user_info)
        return {"Authorization": f"Bearer {token}"}

async def extract_token_from_request(request: Request) -> Optional[str]:
    """Extract OAuth token from various places in the request."""
    auth_header = request.headers.get("Authorization")
    if auth_header and auth_header.startswith("Bearer "):
        return auth_header.replace("Bearer ", "")
    token = request.query_params.get("access_token")
    if token:
        return token
    token = request.cookies.get("access_token")
    if token:
        return token
    return None

✏️ cratedb_mcp/settings.py

Immediately after line 4 (the HTTP_URL definition), insert the following block:

# Authentication settings
# OAuth configuration
OAUTH_ENABLED: bool = os.getenv("CRATEDB_MCP_OAUTH_ENABLED", "false").lower() == "true"
OAUTH_CLIENT_ID: str = os.getenv("CRATEDB_MCP_OAUTH_CLIENT_ID", "")
OAUTH_CLIENT_SECRET: str = os.getenv("CRATEDB_MCP_OAUTH_CLIENT_SECRET", "")
OAUTH_TOKEN_URL: str = os.getenv("CRATEDB_MCP_OAUTH_TOKEN_URL", "")
OAUTH_AUTHORIZE_URL: str = os.getenv("CRATEDB_MCP_OAUTH_AUTHORIZE_URL", "")
OAUTH_USERINFO_URL: str = os.getenv("CRATEDB_MCP_OAUTH_USERINFO_URL", "")

# JWT configuration for CrateDB
JWT_SECRET: str = os.getenv("CRATEDB_MCP_JWT_SECRET", "")
JWT_ALGORITHM: str = os.getenv("CRATEDB_MCP_JWT_ALGORITHM", "HS256")
JWT_EXPIRY: int = int(os.getenv("CRATEDB_MCP_JWT_EXPIRY", "3600"))

# Permission mapping
# Maps OAuth scopes to CrateDB permissions
PERMISSION_MAPPING = {
    "read": "DQL",       # Data Query Language (SELECT)
    "write": "DML",      # Data Manipulation Language (INSERT, UPDATE, DELETE)
    "schema": "DDL",     # Data Definition Language (CREATE, ALTER, DROP)
    "admin": "AL"        # Administration Language
}

✏️ cratedb_mcp/__main__.py

Replace the top import block (lines 1–6) with the following to include FastAPI, typing, settings, auth and middleware:

import hishel
import httpx
from fastapi import Depends, HTTPException, Request
from mcp.server.fastmcp import FastMCP
from typing import Dict, Optional, Any

from .knowledge import DOCUMENTATION_INDEX, Queries
from .settings import (
    DOCS_CACHE_TTL,
    HTTP_URL,
    OAUTH_ENABLED,
    OAUTH_CLIENT_ID,
    OAUTH_CLIENT_SECRET,
    OAUTH_TOKEN_URL,
    OAUTH_AUTHORIZE_URL,
    JWT_SECRET,
    JWT_ALGORITHM,
    JWT_EXPIRY,
)
from .auth import AuthManager, extract_token_from_request
from .middleware import AuthMiddleware

---

Immediately after initializing the FastMCP application (after mcp = FastMCP("cratedb-mcp")), insert:

# Initialize auth manager if OAuth is enabled
auth_manager: Optional[AuthManager] = None
if OAUTH_ENABLED:
    auth_manager = AuthManager(
        oauth_client_id=OAUTH_CLIENT_ID,
        oauth_client_secret=OAUTH_CLIENT_SECRET,
        oauth_token_url=OAUTH_TOKEN_URL,
        oauth_authorize_url=OAUTH_AUTHORIZE_URL,
        jwt_secret=JWT_SECRET,
        jwt_algorithm=JWT_ALGORITHM,
        jwt_expiry=JWT_EXPIRY
    )

async def get_current_user(request: Request) -> Dict[str, Any]:
    \"\"\"
    Dependency to extract and validate the current user via OAuth.
    \"\"\"
    if not OAUTH_ENABLED or not auth_manager:
        return {
            "sub": "anonymous",
            "name": "Anonymous User",
            "permissions": ["read:data", "write:data"],
            "roles": ["basic_user"]
        }
    token = await extract_token_from_request(request)
    if not token:
        raise HTTPException(status_code=401, detail="Authentication required")
    return await auth_manager.validate_oauth_token(token)

---

Locate the existing def query_cratedb(...) and refactor it to an async function that accepts an optional user_info, applying JWT headers when provided:

- def query_cratedb(query: str) -> list[dict]:
-     return httpx.post(f'{HTTP_URL}/_sql', json={'stmt': query}).json()
+ async def query_cratedb(query: str, user_info: Optional[Dict[str, Any]] = None) -> list[dict]:
+     headers: Dict[str, str] = {}
+     if OAUTH_ENABLED and auth_manager and user_info:
+         headers = auth_manager.get_cratedb_auth_headers(user_info)
+     response = await httpx.AsyncClient().post(
+         f'{HTTP_URL}/_sql',
+         json={'stmt': query},
+         headers=headers
+     )
+     if response.status_code == 401:
+         raise HTTPException(status_code=401, detail="Unauthorized access to CrateDB")
+     if response.status_code == 403:
+         raise HTTPException(status_code=403, detail="Insufficient permissions for this query")
+     return response.json()

---

Find the @mcp.tool-decorated query_sql function and update it to async, inject the request to extract user context, and call the new query_cratedb:

- @mcp.tool(description="Send a SQL query ...")
- def query_sql(query: str):
-     if 'select' not in query.lower():
-         raise ValueError('Only queries that have a SELECT statement are allowed.')
-     return query_cratedb(query)
+ @mcp.tool(description="Send a SQL query to CrateDB; only SELECT queries are allowed.")
+ async def query_sql(query: str, request: Request):
+     if 'select' not in query.lower():
+         raise ValueError("Only queries that have a SELECT statement are allowed.")
+     user_info = await get_current_user(request)
+     return await query_cratedb(query, user_info)

➕ cratedb_mcp/middleware.py

Insert the following content into cratedb_mcp/middleware.py:

"""
Middleware for handling authentication in MCP requests.

This middleware extracts authentication information from MCP requests
and makes it available for downstream handlers.
"""
from typing import Callable, Dict, Any, Optional
from fastapi import Request, Response
from .auth import extract_token_from_request
from .settings import OAUTH_ENABLED

class AuthMiddleware:
    """
    Middleware for handling authentication in MCP requests.
    
    This middleware:
    1. Extracts OAuth tokens from incoming requests
    2. Stores the authentication context for use in tools and handlers
    3. Can optionally enforce authentication for all endpoints
    """
    
    def __init__(self, app, auth_manager=None, enforce_auth: bool = False):
        """
        Initialize the auth middleware.
        
        Args:
            app: The FastAPI application
            auth_manager: The AuthManager instance for validating tokens
            enforce_auth: Whether to enforce authentication for all endpoints
        """
        self.app = app
        self.auth_manager = auth_manager
        self.enforce_auth = enforce_auth
        
    async def __call__(self, request: Request, call_next: Callable) -> Response:
        # Skip auth if OAuth is disabled
        if not OAUTH_ENABLED or not self.auth_manager:
            return await call_next(request)
        token = await extract_token_from_request(request)
        if self.enforce_auth and not token:
            return Response(content="Authentication required", status_code=401, media_type="text/plain")
        if token:
            try:
                user_info = await self.auth_manager.validate_oauth_token(token)
                request.state.user = user_info
            except Exception as e:
                if self.enforce_auth:
                    return Response(content=f"Authentication failed: {str(e)}", status_code=401, media_type="text/plain")
        return await call_next(request)

def get_user_from_request(request: Request) -> Optional[Dict[str, Any]]:
    """Retrieve user info stored in request.state by AuthMiddleware."""
    return getattr(request.state, "user", None)

➕ requirements.txt

At the end of requirements.txt (after line 42), append:

httpx>=0.18.0
authlib>=1.0.0
PyJWT>=2.0.0

A summary of the context that CodeRabbit is considering across the codebase:

Identify where to inject OAuth and JWT settings in cratedb_mcp/settings.py

Locate the imports and initial code in main.py to know where to integrate authentication logic

Check for a requirements.txt file to update dependencies for OAuth and JWT support

Count lines in requirements.txt to determine insertion point for new dependencies

• ✅ Create PR with these edits

• 📋 Get copyable edits

[coderabbitai]

A summary of the edits CodeRabbit can apply:

Will create the cratedb_mcp/auth package with an init.py file.

Will create cratedb_mcp/auth/config.py with OAuth and JWT configuration settings.

Will create cratedb_mcp/auth/oauth.py implementing the OAuthClient.

Will create cratedb_mcp/auth/jwt_handler.py to map OAuth claims to CrateDB JWT tokens.

Will create cratedb_mcp/auth/middleware.py to wire OAuth and JWT into FastAPI.

Will create the cratedb_mcp/crate package with an init.py file.

Will create cratedb_mcp/crate/client.py to integrate JWT authentication with CrateDB client.

Will create the mcp/auth/config.py file containing OAuth and JWT configuration classes and global instances.

Will create the mcp/auth/oauth.py file implementing the OAuth2 client for MCP.

Will create the mcp/auth/jwt_handler.py file to map OAuth claims into CrateDB JWT tokens and verify them.

Will create the mcp/auth/middleware.py file to implement OAuth2 middleware and route handlers for MCP.

Will create the mcp/crate/client.py file to integrate CrateDB client operations with JWT authentication.

A plan of each step that CodeRabbit will take in order to execute the edits for each file (with inline rough draft snippets):

➕ cratedb_mcp/auth/__init__.py

Create the directory cratedb_mcp/auth if it does not already exist.

---

Create an empty file at cratedb_mcp/auth/init.py (just add a docstring or leave it blank) to declare the auth package.

➕ cratedb_mcp/auth/config.py

Create the file cratedb_mcp/auth/config.py and populate it with the following content:

"""
Authentication and authorization configuration for MCP.
"""
from typing import Dict, List, Optional
import os


class OAuthConfig:
    """OAuth configuration settings."""
    
    def __init__(self):
        self.client_id = os.environ.get("MCP_OAUTH_CLIENT_ID", "")
        self.client_secret = os.environ.get("MCP_OAUTH_CLIENT_SECRET", "")
        self.authorization_endpoint = os.environ.get(
            "MCP_OAUTH_AUTH_ENDPOINT", 
            "https://auth.example.com/oauth2/authorize"
        )
        self.token_endpoint = os.environ.get(
            "MCP_OAUTH_TOKEN_ENDPOINT", 
            "https://auth.example.com/oauth2/token"
        )
        self.redirect_uri = os.environ.get(
            "MCP_OAUTH_REDIRECT_URI", 
            "http://localhost:8000/oauth/callback"
        )
        self.scope = os.environ.get(
            "MCP_OAUTH_SCOPE", 
            "openid profile email"
        )
        self.jwks_uri = os.environ.get(
            "MCP_OAUTH_JWKS_URI", 
            "https://auth.example.com/.well-known/jwks.json"
        )


class JWTConfig:
    """JWT configuration settings for CrateDB."""
    
    def __init__(self):
        self.algorithm = os.environ.get("MCP_JWT_ALGORITHM", "RS256")
        self.private_key_path = os.environ.get("MCP_JWT_PRIVATE_KEY", "")
        self.public_key_path = os.environ.get("MCP_JWT_PUBLIC_KEY", "")
        self.audience = os.environ.get("MCP_JWT_AUDIENCE", "crate")
        self.issuer = os.environ.get("MCP_JWT_ISSUER", "mcp")
        self.expiration = int(os.environ.get("MCP_JWT_EXPIRATION", "3600"))


class PermissionMappingConfig:
    """Configuration for mapping OAuth scopes to CrateDB permissions."""
    
    DEFAULT_MAPPINGS = {
        "mcp:admin": ["AL"],  # ALL permissions
        "mcp:read": ["DQL"],  # Data query language
        "mcp:write": ["DQL", "DML"],  # Data query and manipulation
        "mcp:user_admin": ["AL"],  # User administration
    }
    
    def __init__(self):
        # Load custom mappings from environment if provided
        self.mappings = dict(self.DEFAULT_MAPPINGS)
        custom_mappings = os.environ.get("MCP_PERMISSION_MAPPINGS", "")
        if custom_mappings:
            try:
                import json
                self.mappings.update(json.loads(custom_mappings))
            except Exception as e:
                print(f"Error loading custom permission mappings: {e}")


# Global configuration instances
oauth_config = OAuthConfig()
jwt_config = JWTConfig()
permission_mapping_config = PermissionMappingConfig()

➕ cratedb_mcp/auth/oauth.py

Create the file cratedb_mcp/auth/oauth.py and populate it with the following content:

"""
OAuth client implementation for MCP.
"""
import json
import time
import urllib.parse
from typing import Dict, List, Optional, Tuple, Any

import requests
from jose import jwt, jwk
from jose.utils import base64url_decode

from cratedb_mcp.auth.config import oauth_config


class OAuthClient:
    """OAuth 2.0 client implementation."""
    
    def __init__(self, config=None):
        self.config = config or oauth_config
        self._jwks_cache = None
        self._jwks_cache_time = 0
    
    def get_authorization_url(self, state: str) -> str:
        """
        Generate the authorization URL for the OAuth flow.
        
        Args:
            state: A random string to maintain state between request and callback
            
        Returns:
            The authorization URL to redirect the user to
        """
        params = {
            "client_id": self.config.client_id,
            "response_type": "code",
            "scope": self.config.scope,
            "redirect_uri": self.config.redirect_uri,
            "state": state
        }
        return f"{self.config.authorization_endpoint}?{urllib.parse.urlencode(params)}"
    
    def exchange_code_for_token(self, code: str) -> Dict[str, Any]:
        """
        Exchange an authorization code for an access token.
        
        Args:
            code: The authorization code received from the OAuth provider
            
        Returns:
            A dictionary containing the access token response
            
        Raises:
            Exception: If the token exchange fails
        """
        data = {
            "client_id": self.config.client_id,
            "client_secret": self.config.client_secret,
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": self.config.redirect_uri
        }
        
        response = requests.post(self.config.token_endpoint, data=data)
        if response.status_code != 200:
            raise Exception(f"Token exchange failed: {response.text}")
        
        return response.json()
    
    def get_jwks(self, force_refresh=False) -> Dict[str, Any]:
        """
        Get the JSON Web Key Set (JWKS) from the OAuth provider.
        
        Args:
            force_refresh: Force a refresh of the JWKS cache
            
        Returns:
            The JWKS dictionary
        """
        now = time.time()
        # Cache JWKS for 1 hour
        if self._jwks_cache is None or now - self._jwks_cache_time > 3600 or force_refresh:
            response = requests.get(self.config.jwks_uri)
            if response.status_code != 200:
                raise Exception(f"Failed to fetch JWKS: {response.text}")
            
            self._jwks_cache = response.json()
            self._jwks_cache_time = now
        
        return self._jwks_cache
    
    def verify_token(self, token: str) -> Dict[str, Any]:
        """
        Verify the JWT token from the OAuth provider and extract the claims.
        
        Args:
            token: The JWT token to verify
            
        Returns:
            The token claims if valid
            
        Raises:
            Exception: If the token is invalid
        """
        # Get the key ID from the token header
        header = jwt.get_unverified_header(token)
        kid = header.get("kid")
        if not kid:
            raise Exception("Token header is missing 'kid'")
        
        # Find the key in the JWKS
        jwks = self.get_jwks()
        key = None
        for k in jwks.get("keys", []):
            if k.get("kid") == kid:
                key = k
                break
        
        if key is None:
            raise Exception(f"Key ID {kid} not found in JWKS")
        
        # Convert the JWK to a key that can be used to verify the signature
        public_key = jwk.construct(key)
        
        # Verify the token
        claims = jwt.decode(
            token,
            public_key.to_pem().decode('utf-8'),
            algorithms=[header.get("alg", "RS256")],
            audience=self.config.client_id,
        )
        
        # Verify the token is not expired
        if "exp" in claims and claims["exp"] < time.time():
            raise Exception("Token has expired")
        
        return claims


# Global OAuth client instance
oauth_client = OAuthClient()

➕ cratedb_mcp/auth/jwt_handler.py

Create the file cratedb_mcp/auth/jwt_handler.py with the following content:

"""
JWT handler for CrateDB authentication.
"""
import time
from typing import Dict, List, Optional, Any
import os

from jose import jwt

from cratedb_mcp.auth.config import jwt_config, permission_mapping_config


class JWTHandler:
    """Handler for CrateDB JWT tokens."""
    
    def __init__(self, config=None, permission_mapping=None):
        self.config = config or jwt_config
        self.permission_mapping = permission_mapping or permission_mapping_config
        self._private_key = None
        self._public_key = None
    
    @property
    def private_key(self) -> str:
        """Get the private key for signing JWT tokens."""
        if self._private_key is None:
            if not self.config.private_key_path:
                raise ValueError("Private key path not configured")
            
            with open(self.config.private_key_path, 'r') as f:
                self._private_key = f.read()
        
        return self._private_key
    
    @property
    def public_key(self) -> str:
        """Get the public key for verifying JWT tokens."""
        if self._public_key is None:
            if not self.config.public_key_path:
                raise ValueError("Public key path not configured")
            
            with open(self.config.public_key_path, 'r') as f:
                self._public_key = f.read()
        
        return self._public_key
    
    def map_permissions(self, oauth_claims: Dict[str, Any]) -> List[str]:
        """
        Map OAuth claims to CrateDB permissions.
        
        Args:
            oauth_claims: The claims from the OAuth token
            
        Returns:
            A list of CrateDB permission strings
        """
        # Extract scopes from OAuth token
        scopes = []
        if "scope" in oauth_claims:
            scopes = oauth_claims["scope"].split()
        elif "scp" in oauth_claims:
            scopes = oauth_claims["scp"]
        
        # Map scopes to CrateDB permissions
        crate_permissions = set()
        for scope in scopes:
            if scope in self.permission_mapping.mappings:
                crate_permissions.update(self.permission_mapping.mappings[scope])
        
        return list(crate_permissions)
    
    def generate_token(self, user: str, oauth_claims: Dict[str, Any]) -> str:
        """
        Generate a JWT token for CrateDB authentication.
        
        Args:
            user: The username for CrateDB
            oauth_claims: The claims from the OAuth token
            
        Returns:
            A JWT token string
        """
        now = int(time.time())
        permissions = self.map_permissions(oauth_claims)
        
        claims = {
            "iss": self.config.issuer,
            "aud": self.config.audience,
            "sub": user,
            "iat": now,
            "exp": now + self.config.expiration,
            "permissions": permissions
        }
        
        # Add any additional claims that might be useful
        if "email" in oauth_claims:
            claims["email"] = oauth_claims["email"]
        
        return jwt.encode(claims, self.private_key, algorithm=self.config.algorithm)
    
    def verify_token(self, token: str) -> Dict[str, Any]:
        """
        Verify a JWT token for CrateDB authentication.
        
        Args:
            token: The JWT token to verify
            
        Returns:
            The token claims if valid
            
        Raises:
            Exception: If the token is invalid
        """
        return jwt.decode(
            token,
            self.public_key,
            algorithms=[self.config.algorithm],
            audience=self.config.audience,
            issuer=self.config.issuer
        )


# Global JWT handler instance
jwt_handler = JWTHandler()

➕ cratedb_mcp/auth/middleware.py

Create the file cratedb_mcp/auth/middleware.py and add this content:

"""
Authentication middleware for the MCP server.
"""
import secrets
from typing import Dict, Optional, Any
from functools import wraps

from fastapi import FastAPI, Request, Response, Depends, HTTPException
from fastapi.responses import RedirectResponse
from starlette.middleware.base import BaseHTTPMiddleware

from cratedb_mcp.auth.oauth import oauth_client
from cratedb_mcp.auth.jwt_handler import jwt_handler


# Store for temporary session data
session_store = {}


def generate_state_token() -> str:
    """Generate a random state token for OAuth flow."""
    return secrets.token_urlsafe(32)


class OAuth2Middleware(BaseHTTPMiddleware):
    """Middleware to handle OAuth2 authentication."""
    
    def __init__(self, app: FastAPI):
        super().__init__(app)
    
    async def dispatch(self, request: Request, call_next):
        """Process each request through the middleware."""
        # Skip authentication for the OAuth callback and login routes
        if request.url.path in ["/oauth/callback", "/login"]:
            return await call_next(request)
        
        # Check if user is authenticated
        token = self._get_token_from_request(request)
        if not token:
            # Store the original URL to redirect back after authentication
            state = generate_state_token()
            session_store[state] = {"redirect_url": str(request.url)}
            
            # Redirect to login page
            return RedirectResponse(url="/login?state=" + state)
        
        # Continue with the request
        return await call_next(request)
    
    def _get_token_from_request(self, request: Request) -> Optional[str]:
        """Extract token from request headers or cookies."""
        # Try to get token from Authorization header
        auth_header = request.headers.get("Authorization")
        if auth_header and auth_header.startswith("Bearer "):
            return auth_header[7:]
        
        # Try to get token from cookie
        return request.cookies.get("mcp_token")


def setup_auth_routes(app: FastAPI):
    """Set up authentication routes in the FastAPI app."""
    
    @app.get("/login")
    async def login(state: str):
        """Initiate the OAuth login flow."""
        # Generate authorization URL
        auth_url = oauth_client.get_authorization_url(state)
        return RedirectResponse(url=auth_url)
    
    @app.get("/oauth/callback")
    async def oauth_callback(code: str, state: str, response: Response):
        """Handle the OAuth callback."""
        try:
            # Exchange code for token
            token_response = oauth_client.exchange_code_for_token(code)
            access_token = token_response.get("access_token")
            
            # Verify token and extract claims
            oauth_claims = oauth_client.verify_token(access_token)
            
            # Generate CrateDB JWT token
            username = oauth_claims.get("preferred_username", 
                                       oauth_claims.get("email", "unknown_user"))
            crate_token = jwt_handler.generate_token(username, oauth_claims)
            
            # Set token in cookie
            response.set_cookie(
                key="mcp_token",
                value=crate_token,
                httponly=True,
                secure=True,
                samesite="lax",
                max_age=3600
            )
            
            # Redirect to original URL or default
            redirect_url = "/"
            if state in session_store:
                redirect_url = session_store[state].get("redirect_url", "/")
                del session_store[state]
            
            return RedirectResponse(url=redirect_url)
        
        except Exception as e:
            return {"error": str(e)}


def requires_auth(scopes: Optional[list] = None):
    """Decorator to require authentication with specific scopes."""
    def decorator(func):
        @wraps(func)
        async def wrapper(request: Request, *args, **kwargs):
            token = request.cookies.get("mcp_token")
            if not token:
                # Check Authorization header
                auth_header = request.headers.get("Authorization")
                if auth_header and auth_header.startswith("Bearer "):
                    token = auth_header[7:]
            
            if not token:
                raise HTTPException(status_code=401, detail="Not authenticated")
            
            try:
                # Verify token
                claims = jwt_handler.verify_token(token)
                
                # Check scopes if required
                if scopes:
                    token_permissions = claims.get("permissions", [])
                    required_permissions = []
                    
                    # Map requested scopes to permissions
                    for scope in scopes:
                        if scope in jwt_handler.permission_mapping.mappings:
                            required_permissions.extend(
                                jwt_handler.permission_mapping.mappings[scope]
                            )
                    
                    # Check if token has required permissions
                    if not all(perm in token_permissions for perm in required_permissions):
                        raise HTTPException(
                            status_code=403, 
                            detail="Insufficient permissions"
                        )
                
                # Add user info to request state
                request.state.user = claims.get("sub")
                request.state.permissions = claims.get("permissions", [])
                
                return await func(request, *args, **kwargs)
            
            except Exception as e:
                raise HTTPException(status_code=401, detail=str(e))
        
        return wrapper
    
    return decorator


def setup_auth(app: FastAPI):
    """Set up authentication for the FastAPI app."""
    # Add middleware
    app.add_middleware(OAuth2Middleware)
    
    # Add routes
    setup_auth_routes(app)

➕ cratedb_mcp/crate/__init__.py

Create the directory cratedb_mcp/crate if it does not already exist.

---

Create an empty file at cratedb_mcp/crate/init.py to declare the crate package.

➕ cratedb_mcp/crate/client.py

Create the file cratedb_mcp/crate/client.py and populate it with the following content:

"""
CrateDB client with JWT authentication integration.
"""
from typing import Dict, List, Optional, Any, Union
import os

from crate.client import connect

from cratedb_mcp.auth.jwt_handler import jwt_handler


class CrateDBClient:
    """CrateDB client with JWT authentication."""
    
    def __init__(self, hosts=None):
        self.hosts = hosts or os.environ.get("CRATE_HOSTS", "localhost:4200")
        self._conn = None
    
    def get_connection(self, token: Optional[str] = None):
        """
        Get a connection to CrateDB.
        
        Args:
            token: JWT token for authentication, if None uses default connection
            
        Returns:
            A connection to CrateDB
        """
        if token:
            # Connect with JWT authentication
            return connect(
                self.hosts,
                username="_jwt",
                password=token
            )
        else:
            # Connect with default authentication method
            # (environment variables or default credentials)
            if self._conn is None:
                self._conn = connect(self.hosts)
            return self._conn
    
    def execute_query(self, query: str, params: Optional[list] = None, 
                      token: Optional[str] = None):
        """
        Execute a query on CrateDB.
        
        Args:
            query: SQL query to execute
            params: Query parameters
            token: JWT token for authentication
            
        Returns:
            Query result
        """
        conn = self.get_connection(token)
        cursor = conn.cursor()
        try:
            cursor.execute(query, params or [])
            return cursor.fetchall()
        finally:
            cursor.close()
            # Only close if using token authentication (temporary connection)
            if token:
                conn.close()
    
    def execute_with_user_permissions(self, query: str, params: Optional[list] = None,
                                     oauth_claims: Dict[str, Any], username: str):
        """
        Execute a query with user permissions derived from OAuth claims.
        
        Args:
            query: SQL query to execute
            params: Query parameters
            oauth_claims: OAuth token claims
            username: Username for CrateDB
            
        Returns:
            Query result
        """
        # Generate JWT token with permissions from OAuth claims
        token = jwt_handler.generate_token(username, oauth_claims)
        
        # Execute query with token
        return self.execute_query(query, params, token)


# Global CrateDB client instance
crate_client = CrateDBClient()

📄 mcp/auth/config.py

Ensure the directory mcp/auth exists. If it does not, run in the repo root:

mkdir -p mcp/auth

---

Create the file mcp/auth/config.py with the following content:

"""
Authentication and authorization configuration for MCP.
"""
from typing import Dict, List, Optional
import os


class OAuthConfig:
    """OAuth configuration settings."""
    
    def __init__(self):
        self.client_id = os.environ.get("MCP_OAUTH_CLIENT_ID", "")
        self.client_secret = os.environ.get("MCP_OAUTH_CLIENT_SECRET", "")
        self.authorization_endpoint = os.environ.get(
            "MCP_OAUTH_AUTH_ENDPOINT", 
            "https://auth.example.com/oauth2/authorize"
        )
        self.token_endpoint = os.environ.get(
            "MCP_OAUTH_TOKEN_ENDPOINT", 
            "https://auth.example.com/oauth2/token"
        )
        self.redirect_uri = os.environ.get(
            "MCP_OAUTH_REDIRECT_URI", 
            "http://localhost:8000/oauth/callback"
        )
        self.scope = os.environ.get(
            "MCP_OAUTH_SCOPE", 
            "openid profile email"
        )
        self.jwks_uri = os.environ.get(
            "MCP_OAUTH_JWKS_URI", 
            "https://auth.example.com/.well-known/jwks.json"
        )


class JWTConfig:
    """JWT configuration settings for CrateDB."""
    
    def __init__(self):
        self.algorithm = os.environ.get("MCP_JWT_ALGORITHM", "RS256")
        self.private_key_path = os.environ.get("MCP_JWT_PRIVATE_KEY", "")
        self.public_key_path = os.environ.get("MCP_JWT_PUBLIC_KEY", "")
        self.audience = os.environ.get("MCP_JWT_AUDIENCE", "crate")
        self.issuer = os.environ.get("MCP_JWT_ISSUER", "mcp")
        self.expiration = int(os.environ.get("MCP_JWT_EXPIRATION", "3600"))


class PermissionMappingConfig:
    """Configuration for mapping OAuth scopes to CrateDB permissions."""
    
    DEFAULT_MAPPINGS = {
        "mcp:admin": ["AL"],  # ALL permissions
        "mcp:read": ["DQL"],  # Data query language
        "mcp:write": ["DQL", "DML"],  # Data query and manipulation
        "mcp:user_admin": ["AL"],  # User administration
    }
    
    def __init__(self):
        # Load custom mappings from environment if provided
        self.mappings = dict(self.DEFAULT_MAPPINGS)
        custom_mappings = os.environ.get("MCP_PERMISSION_MAPPINGS", "")
        if custom_mappings:
            try:
                import json
                self.mappings.update(json.loads(custom_mappings))
            except Exception as e:
                print(f"Error loading custom permission mappings: {e}")


# Global configuration instances
oauth_config = OAuthConfig()
jwt_config = JWTConfig()
permission_mapping_config = PermissionMappingConfig()

📄 mcp/auth/oauth.py

Create the file mcp/auth/oauth.py with the following content:

"""
OAuth client implementation for MCP.
"""
import json
import time
import urllib.parse
from typing import Dict, List, Optional, Tuple, Any

import requests
from jose import jwt, jwk
from jose.utils import base64url_decode

from mcp.auth.config import oauth_config


class OAuthClient:
    """OAuth 2.0 client implementation."""
    
    def __init__(self, config=None):
        self.config = config or oauth_config
        self._jwks_cache = None
        self._jwks_cache_time = 0
    
    def get_authorization_url(self, state: str) -> str:
        """
        Generate the authorization URL for the OAuth flow.
        
        Args:
            state: A random string to maintain state between request and callback
            
        Returns:
            The authorization URL to redirect the user to
        """
        params = {
            "client_id": self.config.client_id,
            "response_type": "code",
            "scope": self.config.scope,
            "redirect_uri": self.config.redirect_uri,
            "state": state
        }
        return f"{self.config.authorization_endpoint}?{urllib.parse.urlencode(params)}"
    
    def exchange_code_for_token(self, code: str) -> Dict[str, Any]:
        """
        Exchange an authorization code for an access token.
        
        Args:
            code: The authorization code received from the OAuth provider
            
        Returns:
            A dictionary containing the access token response
            
        Raises:
            Exception: If the token exchange fails
        """
        data = {
            "client_id": self.config.client_id,
            "client_secret": self.config.client_secret,
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": self.config.redirect_uri
        }
        
        response = requests.post(self.config.token_endpoint, data=data)
        if response.status_code != 200:
            raise Exception(f"Token exchange failed: {response.text}")
        
        return response.json()
    
    def get_jwks(self, force_refresh=False) -> Dict[str, Any]:
        """
        Get the JSON Web Key Set (JWKS) from the OAuth provider.
        
        Args:
            force_refresh: Force a refresh of the JWKS cache
            
        Returns:
            The JWKS dictionary
        """
        now = time.time()
        # Cache JWKS for 1 hour
        if self._jwks_cache is None or now - self._jwks_cache_time > 3600 or force_refresh:
            response = requests.get(self.config.jwks_uri)
            if response.status_code != 200:
                raise Exception(f"Failed to fetch JWKS: {response.text}")
            
            self._jwks_cache = response.json()
            self._jwks_cache_time = now
        
        return self._jwks_cache
    
    def verify_token(self, token: str) -> Dict[str, Any]:
        """
        Verify the JWT token from the OAuth provider and extract the claims.
        
        Args:
            token: The JWT token to verify
            
        Returns:
            The token claims if valid
            
        Raises:
            Exception: If the token is invalid
        """
        # Get the key ID from the token header
        header = jwt.get_unverified_header(token)
        kid = header.get("kid")
        if not kid:
            raise Exception("Token header is missing 'kid'")
        
        # Find the key in the JWKS
        jwks = self.get_jwks()
        key = None
        for k in jwks.get("keys", []):
            if k.get("kid") == kid:
                key = k
                break
        
        if key is None:
            raise Exception(f"Key ID {kid} not found in JWKS")
        
        # Convert the JWK to a key that can be used to verify the signature
        public_key = jwk.construct(key)
        
        # Verify the token
        claims = jwt.decode(
            token,
            public_key.to_pem().decode('utf-8'),
            algorithms=[header.get("alg", "RS256")],
            audience=self.config.client_id,
        )
        
        # Verify the token is not expired
        if "exp" in claims and claims["exp"] < time.time():
            raise Exception("Token has expired")
        
        return claims


# Global OAuth client instance
oauth_client = OAuthClient()

📄 mcp/auth/jwt_handler.py

Create the file mcp/auth/jwt_handler.py with the following content:

"""
JWT handler for CrateDB authentication.
"""
import time
from typing import Dict, List, Optional, Any
import os

from jose import jwt

from mcp.auth.config import jwt_config, permission_mapping_config


class JWTHandler:
    """Handler for CrateDB JWT tokens."""
    
    def __init__(self, config=None, permission_mapping=None):
        self.config = config or jwt_config
        self.permission_mapping = permission_mapping or permission_mapping_config
        self._private_key = None
        self._public_key = None
    
    @property
    def private_key(self) -> str:
        """Get the private key for signing JWT tokens."""
        if self._private_key is None:
            if not self.config.private_key_path:
                raise ValueError("Private key path not configured")
            
            with open(self.config.private_key_path, 'r') as f:
                self._private_key = f.read()
        
        return self._private_key
    
    @property
    def public_key(self) -> str:
        """Get the public key for verifying JWT tokens."""
        if self._public_key is None:
            if not self.config.public_key_path:
                raise ValueError("Public key path not configured")
            
            with open(self.config.public_key_path, 'r') as f:
                self._public_key = f.read()
        
        return self._public_key
    
    def map_permissions(self, oauth_claims: Dict[str, Any]) -> List[str]:
        """
        Map OAuth claims to CrateDB permissions.
        
        Args:
            oauth_claims: The claims from the OAuth token
            
        Returns:
            A list of CrateDB permission strings
        """
        # Extract scopes from OAuth token
        scopes = []
        if "scope" in oauth_claims:
            scopes = oauth_claims["scope"].split()
        elif "scp" in oauth_claims:
            scopes = oauth_claims["scp"]
        
        # Map scopes to CrateDB permissions
        crate_permissions = set()
        for scope in scopes:
            if scope in self.permission_mapping.mappings:
                crate_permissions.update(self.permission_mapping.mappings[scope])
        
        return list(crate_permissions)
    
    def generate_token(self, user: str, oauth_claims: Dict[str, Any]) -> str:
        """
        Generate a JWT token for CrateDB authentication.
        
        Args:
            user: The username for CrateDB
            oauth_claims: The claims from the OAuth token
            
        Returns:
            A JWT token string
        """
        now = int(time.time())
        permissions = self.map_permissions(oauth_claims)
        
        claims = {
            "iss": self.config.issuer,
            "aud": self.config.audience,
            "sub": user,
            "iat": now,
            "exp": now + self.config.expiration,
            "permissions": permissions
        }
        
        # Add any additional claims that might be useful
        if "email" in oauth_claims:
            claims["email"] = oauth_claims["email"]
        
        return jwt.encode(claims, self.private_key, algorithm=self.config.algorithm)
    
    def verify_token(self, token: str) -> Dict[str, Any]:
        """
        Verify a JWT token for CrateDB authentication.
        
        Args:
            token: The JWT token to verify
            
        Returns:
            The token claims if valid
            
        Raises:
            Exception: If the token is invalid
        """
        return jwt.decode(
            token,
            self.public_key,
            algorithms=[self.config.algorithm],
            audience=self.config.audience,
            issuer=self.config.issuer
        )


# Global JWT handler instance
jwt_handler = JWTHandler()

📄 mcp/auth/middleware.py

Create the file mcp/auth/middleware.py with the following content:

"""
Authentication middleware for the MCP server.
"""
import secrets
from typing import Dict, Optional, Any
from functools import wraps

from fastapi import FastAPI, Request, Response, Depends, HTTPException
from fastapi.responses import RedirectResponse
from starlette.middleware.base import BaseHTTPMiddleware

from mcp.auth.oauth import oauth_client
from mcp.auth.jwt_handler import jwt_handler


# Store for temporary session data
session_store = {}


def generate_state_token() -> str:
    """Generate a random state token for OAuth flow."""
    return secrets.token_urlsafe(32)


class OAuth2Middleware(BaseHTTPMiddleware):
    """Middleware to handle OAuth2 authentication."""
    
    def __init__(self, app: FastAPI):
        super().__init__(app)
    
    async def dispatch(self, request: Request, call_next):
        """Process each request through the middleware."""
        # Skip authentication for the OAuth callback and login routes
        if request.url.path in ["/oauth/callback", "/login"]:
            return await call_next(request)
        
        # Check if user is authenticated
        token = self._get_token_from_request(request)
        if not token:
            # Store the original URL to redirect back after authentication
            state = generate_state_token()
            session_store[state] = {"redirect_url": str(request.url)}
            
            # Redirect to login page
            return RedirectResponse(url="/login?state=" + state)
        
        # Continue with the request
        return await call_next(request)
    
    def _get_token_from_request(self, request: Request) -> Optional[str]:
        """Extract token from request headers or cookies."""
        # Try to get token from Authorization header
        auth_header = request.headers.get("Authorization")
        if auth_header and auth_header.startswith("Bearer "):
            return auth_header[7:]
        
        # Try to get token from cookie
        return request.cookies.get("mcp_token")


def setup_auth_routes(app: FastAPI):
    """Set up authentication routes in the FastAPI app."""
    
    @app.get("/login")
    async def login(state: str):
        """Initiate the OAuth login flow."""
        # Generate authorization URL
        auth_url = oauth_client.get_authorization_url(state)
        return RedirectResponse(url=auth_url)
    
    @app.get("/oauth/callback")
    async def oauth_callback(code: str, state: str, response: Response):
        """Handle the OAuth callback."""
        try:
            # Exchange code for token
            token_response = oauth_client.exchange_code_for_token(code)
            access_token = token_response.get("access_token")
            
            # Verify token and extract claims
            oauth_claims = oauth_client.verify_token(access_token)
            
            # Generate CrateDB JWT token
            username = oauth_claims.get("preferred_username", 
                                       oauth_claims.get("email", "unknown_user"))
            crate_token = jwt_handler.generate_token(username, oauth_claims)
            
            # Set token in cookie
            response.set_cookie(
                key="mcp_token",
                value=crate_token,
                httponly=True,
                secure=True,
                samesite="lax",
                max_age=3600
            )
            
            # Redirect to original URL or default
            redirect_url = "/"
            if state in session_store:
                redirect_url = session_store[state].get("redirect_url", "/")
                del session_store[state]
            
            return RedirectResponse(url=redirect_url)
        
        except Exception as e:
            return {"error": str(e)}


def requires_auth(scopes: Optional[list] = None):
    """Decorator to require authentication with specific scopes."""
    def decorator(func):
        @wraps(func)
        async def wrapper(request: Request, *args, **kwargs):
            token = request.cookies.get("mcp_token")
            if not token:
                # Check Authorization header
                auth_header = request.headers.get("Authorization")
                if auth_header and auth_header.startswith("Bearer "):
                    token = auth_header[7:]
            
            if not token:
                raise HTTPException(status_code=401, detail="Not authenticated")
            
            try:
                # Verify token
                claims = jwt_handler.verify_token(token)
                
                # Check scopes if required
                if scopes:
                    token_permissions = claims.get("permissions", [])
                    required_permissions = []
                    
                    # Map requested scopes to permissions
                    for scope in scopes:
                        if scope in jwt_handler.permission_mapping.mappings:
                            required_permissions.extend(
                                jwt_handler.permission_mapping.mappings[scope]
                            )
                    
                    # Check if token has required permissions
                    if not all(perm in token_permissions for perm in required_permissions):
                        raise HTTPException(
                            status_code=403, 
                            detail="Insufficient permissions"
                        )
                
                # Add user info to request state
                request.state.user = claims.get("sub")
                request.state.permissions = claims.get("permissions", [])
                
                return await func(request, *args, **kwargs)
            
            except Exception as e:
                raise HTTPException(status_code=401, detail=str(e))
        
        return wrapper
    
    return decorator


def setup_auth(app: FastAPI):
    """Set up authentication for the FastAPI app."""
    # Add middleware
    app.add_middleware(OAuth2Middleware)
    
    # Add routes
    setup_auth_routes(app)

📄 mcp/crate/client.py

Ensure the directory mcp/crate exists. If it does not, run in the repo root:

mkdir -p mcp/crate

---

Create the file mcp/crate/client.py with the following content:

"""
CrateDB client with JWT authentication integration.
"""
from typing import Dict, List, Optional, Any, Union
import os

from crate.client import connect

from mcp.auth.jwt_handler import jwt_handler


class CrateDBClient:
    """CrateDB client with JWT authentication."""
    
    def __init__(self, hosts=None):
        self.hosts = hosts or os.environ.get("CRATE_HOSTS", "localhost:4200")
        self._conn = None
    
    def get_connection(self, token: Optional[str] = None):
        """
        Get a connection to CrateDB.
        
        Args:
            token: JWT token for authentication, if None uses default connection
            
        Returns:
            A connection to CrateDB
        """
        if token:
            # Connect with JWT authentication
            return connect(
                self.hosts,
                username="_jwt",
                password=token
            )
        else:
            # Connect with default authentication method
            # (environment variables or default credentials)
            if self._conn is None:
                self._conn = connect(self.hosts)
            return self._conn
    
    def execute_query(self, query: str, params: Optional[list] = None, 
                      token: Optional[str] = None):
        """
        Execute a query on CrateDB.
        
        Args:
            query: SQL query to execute
            params: Query parameters
            token: JWT token for authentication
            
        Returns:
            Query result
        """
        conn = self.get_connection(token)
        cursor = conn.cursor()
        try:
            cursor.execute(query, params or [])
            return cursor.fetchall()
        finally:
            cursor.close()
            # Only close if using token authentication (temporary connection)
            if token:
                conn.close()
    
    def execute_with_user_permissions(self, query: str, params: Optional[list] = None,
                                     oauth_claims: Dict[str, Any], username: str):
        """
        Execute a query with user permissions derived from OAuth claims.
        
        Args:
            query: SQL query to execute
            params: Query parameters
            oauth_claims: OAuth token claims
            username: Username for CrateDB
            
        Returns:
            Query result
        """
        # Generate JWT token with permissions from OAuth claims
        token = jwt_handler.generate_token(username, oauth_claims)
        
        # Execute query with token
        return self.execute_query(query, params, token)


# Global CrateDB client instance
crate_client = CrateDBClient()

A summary of the context that CodeRabbit is considering across the codebase:

Inspect the ast-grep-rules directory to locate the code under review and determine the correct paths for creating the MCP modules and main application file.

Based on the previous inspection (find mcp -type d -maxdepth 2 returned no output), the 'mcp' package directory does not exist. We will need to create the base package and its subdirectories before adding the new auth and crate modules.

• ✅ Create PR with these edits

• 📋 Get copyable edits