squirrels 0.5.0b2__py3-none-any.whl → 0.5.0b4__py3-none-any.whl

squirrels/_api_routes/datasets.py

@@ -0,0 +1,242 @@
+"""
+Dataset routes for parameters and results
+"""
+from typing import Callable, Any
+from pydantic import Field, BaseModel
+from fastapi import FastAPI, Depends, Request
+from fastapi.responses import JSONResponse, HTMLResponse
+from fastapi.security import HTTPBearer
+
+from mcp.server.fastmcp import FastMCP, Context
+from dataclasses import asdict
+from cachetools import TTLCache
+from textwrap import dedent
+
+import time
+
+from .. import _constants as c, _utils as u
+from .._schemas import response_models as rm
+from .._exceptions import ConfigurationError, InvalidInputError
+from .._dataset_types import DatasetResult
+from .._schemas.query_param_models import get_query_models_for_parameters, get_query_models_for_dataset
+from .._auth import BaseUser
+from .base import RouteBase
+
+
+class DatasetRoutes(RouteBase):
+    """Dataset parameter and result routes"""
+    
+    def __init__(self, get_bearer_token: HTTPBearer, project, no_cache: bool = False):
+        super().__init__(get_bearer_token, project, no_cache)
+        
+        # Setup caches
+        dataset_results_cache_size = int(self.env_vars.get(c.SQRL_DATASETS_CACHE_SIZE, 128))
+        dataset_results_cache_ttl = int(self.env_vars.get(c.SQRL_DATASETS_CACHE_TTL_MINUTES, 60))
+        self.dataset_results_cache = TTLCache(maxsize=dataset_results_cache_size, ttl=dataset_results_cache_ttl*60)
+        
+    async def _get_dataset_results_helper(
+        self, dataset: str, user: BaseUser | None, selections: tuple[tuple[str, Any], ...]
+    ) -> DatasetResult:
+        """Helper to get dataset results"""
+        return await self.project.dataset(dataset, selections=dict(selections), user=user)
+
+    async def _get_dataset_results_cachable(
+        self, dataset: str, user: BaseUser | None, selections: tuple[tuple[str, Any], ...]
+    ) -> DatasetResult:
+        """Cachable version of dataset results helper"""
+        return await self.do_cachable_action(self.dataset_results_cache, self._get_dataset_results_helper, dataset, user, selections)
+    
+    async def _get_dataset_results_definition(
+        self, dataset_name: str, user: BaseUser | None, all_request_params: dict, params: dict
+    ) -> rm.DatasetResultModel:
+        """Get dataset results definition"""
+        self._validate_request_params(all_request_params, params)
+
+        get_dataset_function = self._get_dataset_results_helper if self.no_cache else self._get_dataset_results_cachable
+        uncached_keys = {"x_verify_params", "x_orientation", "x_select", "x_limit", "x_offset"}
+        selections = self.get_selections_as_immutable(params, uncached_keys)
+        result = await get_dataset_function(dataset_name, user, selections)
+        
+        orientation = params.get("x_orientation", "records")
+        raw_select: list[str] | None = params.get("x_select")
+        select = tuple(raw_select) if raw_select is not None else tuple()
+        limit = params.get("x_limit", 1000)
+        offset = params.get("x_offset", 0)
+        return rm.DatasetResultModel(**result.to_json(orientation, select, limit, offset)) 
+    
+    def setup_routes(
+        self, app: FastAPI, mcp: FastMCP, project_metadata_path: str, project_name: str, project_version: str, 
+        param_fields: dict, get_parameters_definition: Callable
+    ) -> None:
+        """Setup dataset routes"""
+        
+        dataset_results_path = project_metadata_path + '/dataset/{dataset}'
+        dataset_parameters_path = dataset_results_path + '/parameters'
+        
+        def validate_parameters_list(parameters: list[str] | None, entity_type: str, dataset_name: str) -> None:
+            if parameters is None:
+                return
+            for param in parameters:
+                if param not in param_fields:
+                    all_params = list(param_fields.keys())
+                    raise ConfigurationError(
+                        f"{entity_type} '{dataset_name}' use parameter '{param}' which doesn't exist. Available parameters are:"
+                        f"\n  {all_params}"
+                    )
+        
+        async def get_dataset_parameters_updates(dataset_name: str, user: BaseUser | None, all_request_params: dict, params: dict):
+            parameters_list = self.manifest_cfg.datasets[dataset_name].parameters
+            scope = self.manifest_cfg.datasets[dataset_name].scope
+            result = await get_parameters_definition(
+                parameters_list, "dataset", dataset_name, scope, user, all_request_params, params
+            )
+            return result
+
+        # Dataset parameters and results APIs
+        for dataset_name, dataset_config in self.manifest_cfg.datasets.items():
+            dataset_normalized = u.normalize_name_for_api(dataset_name)
+            curr_parameters_path = dataset_parameters_path.format(dataset=dataset_normalized)
+            curr_results_path = dataset_results_path.format(dataset=dataset_normalized)
+
+            validate_parameters_list(dataset_config.parameters, "Dataset", dataset_name)
+
+            QueryModelForGetParams, QueryModelForPostParams = get_query_models_for_parameters(dataset_config.parameters, param_fields)
+            QueryModelForGetDataset, QueryModelForPostDataset = get_query_models_for_dataset(dataset_config.parameters, param_fields)
+
+            @app.get(curr_parameters_path, tags=[f"Dataset '{dataset_name}'"], description=self._parameters_description, response_class=JSONResponse)
+            async def get_dataset_parameters(
+                request: Request, params: QueryModelForGetParams, user=Depends(self.get_current_user) # type: ignore
+            ) -> rm.ParametersModel:
+                start = time.time()
+                curr_dataset_name = self.get_name_from_path_section(request, -2)
+                result = await get_dataset_parameters_updates(curr_dataset_name, user, dict(request.query_params), asdict(params))
+                self.log_activity_time("GET REQUEST for PARAMETERS", start, request)
+                return result
+
+            @app.post(curr_parameters_path, tags=[f"Dataset '{dataset_name}'"], description=self._parameters_description, response_class=JSONResponse)
+            async def get_dataset_parameters_with_post(
+                request: Request, params: QueryModelForPostParams, user=Depends(self.get_current_user) # type: ignore
+            ) -> rm.ParametersModel:
+                start = time.time()
+                curr_dataset_name = self.get_name_from_path_section(request, -2)
+                payload: dict = await request.json()
+                result = await get_dataset_parameters_updates(curr_dataset_name, user, payload, params.model_dump())
+                self.log_activity_time("POST REQUEST for PARAMETERS", start, request)
+                return result
+            
+            @app.get(curr_results_path, tags=[f"Dataset '{dataset_name}'"], description=dataset_config.description, response_class=JSONResponse)
+            async def get_dataset_results(
+                request: Request, params: QueryModelForGetDataset, user=Depends(self.get_current_user) # type: ignore
+            ) -> rm.DatasetResultModel:
+                start = time.time()
+                curr_dataset_name = self.get_name_from_path_section(request, -1)
+                result = await self._get_dataset_results_definition(curr_dataset_name, user, dict(request.query_params), asdict(params))
+                self.log_activity_time("GET REQUEST for DATASET RESULTS", start, request)
+                return result
+            
+            @app.post(curr_results_path, tags=[f"Dataset '{dataset_name}'"], description=dataset_config.description, response_class=JSONResponse)
+            async def get_dataset_results_with_post(
+                request: Request, params: QueryModelForPostDataset, user=Depends(self.get_current_user) # type: ignore
+            ) -> rm.DatasetResultModel:
+                start = time.time()
+                curr_dataset_name = self.get_name_from_path_section(request, -1)
+                payload: dict = await request.json()
+                result = await self._get_dataset_results_definition(curr_dataset_name, user, payload, params.model_dump())
+                self.log_activity_time("POST REQUEST for DATASET RESULTS", start, request)
+                return result
+    
+        # Setup MCP tools
+
+        @mcp.tool(
+            name=f"get_dataset_parameters_for_{project_name}_{project_version}",
+            description=dedent(f"""
+            Use this tool to get updates for dataset parameters in the Squirrels project "{project_name}" when a selection is to be made on a parameter with "trigger_refresh" as true.
+
+            For example, suppose there are two parameters, "country" and "city", and the user selects "United States" for "country". If "country" has the "trigger_refresh" field as true, then this tool will be called to get the updates for other parameters such as "city".
+
+            Do not use this tool on parameters whose "trigger_refresh" field is false!
+            """).strip()
+        )
+        async def get_dataset_parameters_tool(
+            ctx: Context,
+            dataset: str = Field(description="The name of the dataset whose parameters the trigger parameter will update"),
+            parameter_name: str = Field(description="The name of the parameter triggering the refresh"),
+            selected_ids: list[str] = Field(description="The ID(s) of the selected option(s) for the parameter"),
+        ):
+            user = self.get_user_from_tool_ctx(ctx)
+            dataset_name = u.normalize_name(dataset)
+            payload = {
+                "x_parent_param": parameter_name,
+                parameter_name: selected_ids
+            }
+            return await get_dataset_parameters_updates(dataset_name, user, payload, payload)
+        
+        @mcp.tool(
+            name=f"get_dataset_results_for_{project_name}_{project_version}",
+            description=dedent(f"""
+            Use this tool to get the dataset results as a JSON object for a dataset in the Squirrels project "{project_name}".
+            - Use the "offset" and "limit" arguments to limit the number of rows you require
+            - The "limit" argument controls the number of rows returned. The maximum allowed value is 100. If the 'total_num_rows' field in the response is greater than 100, let the user know that only 100 rows are shown and clarify if they would like to see more.
+            """).strip()
+        )
+        async def get_dataset_results_tool(
+            ctx: Context,
+            dataset: str = Field(description="The name of the dataset to get results for"),
+            parameters: dict[str, Any] = Field(description=dedent("""
+            Key-value pairs for parameter name and selected value. The selected value to provide depends on the parameter widget type:
+            - For single select, use a string for the ID of the selected value
+            - For multi select, use an array of strings for the IDs of the selected values
+            - For date, use a string like "YYYY-MM-DD"
+            - For date ranges, use array of strings like ["YYYY-MM-DD","YYYY-MM-DD"]
+            - For number, use a number like 1
+            - For number ranges, use array of numbers like [1,100]
+            - For text, use a string for the text value
+            - Complex objects are NOT supported""").strip()),
+            offset: int = Field(0, description="The number of rows to skip from first row. Default is 0."),
+            limit: int = Field(100, description="The maximum number of rows to return. Default is 100. Maximum allowed value is 100."),
+        ):
+            if limit > 100:
+                raise ValueError("The maximum number of rows to return is 100.")
+
+            user = self.get_user_from_tool_ctx(ctx)
+            dataset_name = u.normalize_name(dataset)
+            params = {
+                **parameters,
+                "x_orientation": "rows",
+                "x_offset": offset,
+                "x_limit": limit
+            }
+            result = await self._get_dataset_results_definition(dataset_name, user, params, params)
+            return result
+        
+        # Setup UI for tool results
+        mcp_tool_results_ui_path = project_metadata_path + "/mcp/tool-results-ui"
+
+        @app.get(mcp_tool_results_ui_path + "/list-tools", tags=["MCP Supplements"])
+        async def list_tools():
+            return ["get_dataset_results"]
+
+        class ToolResultBody(BaseModel):
+            """Flexible model for tool results - accepts any additional fields"""
+            
+            class Config:
+                extra = "allow"  # Allow additional fields not defined in the model
+
+        @app.post(mcp_tool_results_ui_path + "/tool/{tool_name}", tags=["MCP Supplements"])
+        async def tool_results_ui(tool_name: str, tool_result: ToolResultBody):
+            if tool_name == "get_dataset_results":
+                # Convert Pydantic model to dict to access any extra fields
+                tool_result_dict = tool_result.model_dump()
+                
+                # Prepare template context
+                context = {
+                    "schema": tool_result_dict.get("schema", {}),
+                    "data": tool_result_dict.get("data", []),
+                }
+                
+                # Render HTML template
+                html_content = self.templates.get_template("dataset_results.html").render(context)
+                return HTMLResponse(content=html_content, status_code=200)
+            else:
+                raise InvalidInputError(400, "Invalid tool name", f"Tool name '{tool_name}' not supported for UI")
+        

squirrels/_api_routes/oauth2.py

@@ -0,0 +1,300 @@
+from fastapi import FastAPI, Depends, Request, Query, Response, APIRouter, Form
+from fastapi.responses import RedirectResponse, HTMLResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from typing import Annotated, cast
+
+from .base import RouteBase
+from .._schemas.auth_models import (
+    ClientRegistrationRequest, ClientUpdateRequest, ClientRegistrationResponse, ClientDetailsResponse, ClientUpdateResponse, 
+    TokenResponse, OAuthServerMetadata
+)
+from .._exceptions import InvalidInputError
+
+
+class OAuth2Routes(RouteBase):
+    """OAuth2 routes"""
+    
+    def __init__(self, get_bearer_token: HTTPBearer, project, no_cache: bool = False):
+        super().__init__(get_bearer_token, project, no_cache)
+    
+    def serve_login_page(self, auth_path: str, request: Request, client_id: str) -> HTMLResponse:
+        """Helper function to serve the login page with optional error message"""
+        # Get client information for display
+        client_details = self.authenticator.get_oauth_client_details(client_id)
+        client_name = client_details.client_name if client_details else None
+        project_name = self.manifest_cfg.project_variables.label
+        
+        # Get available login providers
+        providers = []
+        for provider in self.authenticator.auth_providers:
+            provider_login_url = f"{auth_path}/providers/{provider.name}/login"
+            providers.append({
+                "name": provider.name,
+                "label": provider.label,
+                "icon": provider.icon,
+                "login_url": provider_login_url
+            })
+        
+        # Template context
+        context = {
+            "request": request,
+            "project_name": project_name,
+            "client_name": client_name,
+            "providers": providers,
+            "login_url": f"{auth_path}/login",
+            "return_url": str(request.url),
+        }
+        
+        return HTMLResponse(
+            content=self.templates.get_template("oauth_login.html").render(context),
+            status_code=200
+        )
+    
+    def setup_routes(self, app: FastAPI) -> None:
+        """Setup all OAuth2 routes"""
+
+        router_path = "/api/auth/oauth2"
+        router = APIRouter(prefix=router_path)
+        
+        # Create user models
+        class UserInfoModel(self.UserInfoModel):
+            username: str
+
+        # Authorization dependency for client management
+        get_client_token = HTTPBearer(auto_error=False)
+        
+        async def validate_client_registration_token(
+            client_id: str, auth: HTTPAuthorizationCredentials = Depends(get_client_token),
+        ) -> None:
+            """Validate Bearer token for client management operations"""
+
+            if not auth or not auth.scheme == "Bearer":
+                raise InvalidInputError(401, "invalid_client", 
+                    "Missing or invalid authorization header. Use 'Authorization: Bearer <registration_access_token>'"
+                )
+            
+            token = auth.credentials
+            is_valid = self.authenticator.validate_registration_access_token(client_id, token)
+            if not is_valid:
+                raise InvalidInputError(401, "invalid_token", "Invalid registration access token for this client")
+        
+        def validate_oauth_client_credentials(client_id: str | None, client_secret: str | None) -> str:
+            """
+            Validate OAuth client credentials from form data or Authorization header.
+            Returns the validated client_id.
+            """
+            
+            # Validate client credentials
+            if not client_id or not client_secret or not self.authenticator.validate_client_credentials(client_id, client_secret):
+                raise InvalidInputError(400, "invalid_client", "Invalid client credentials")
+            
+            return cast(str, client_id)
+        
+        # Client Registration Endpoint
+        client_management_path = '/client/{client_id}'
+        
+        @router.post("/register", description="Register a new OAuth client", tags=["OAuth2"])
+        async def register_oauth_client(request: ClientRegistrationRequest) -> ClientRegistrationResponse:
+            """Register a new OAuth client and return client credentials"""
+            
+            # Register the client using the authenticator
+            client_registration_response = self.authenticator.register_oauth_client(
+                request, client_management_path_format=router_path+client_management_path
+            )
+
+            return client_registration_response
+            
+        # Client Management Endpoints
+        @router.get(client_management_path, description="Get OAuth client registration details", tags=["OAuth2"])
+        async def get_oauth_client(
+            client_id: str, _: Annotated[None, Depends(validate_client_registration_token)]
+        ) -> ClientDetailsResponse:
+            """Get OAuth client registration details"""
+
+            client_details = self.authenticator.get_oauth_client_details(client_id)
+            
+            return client_details
+        
+        @router.put(client_management_path, description="Update OAuth client registration", tags=["OAuth2"])
+        async def update_oauth_client(
+            client_id: str, request: ClientUpdateRequest, _: Annotated[None, Depends(validate_client_registration_token)]
+        ) -> ClientUpdateResponse:
+            """Update OAuth client registration and rotate access token"""
+
+            # Update the client and get new registration access token
+            client_details = self.authenticator.update_oauth_client_with_token_rotation(client_id, request)
+            
+            return client_details
+        
+        @router.delete(client_management_path, description="Revoke OAuth client registration", tags=["OAuth2"], responses={
+            204: { "description": "OAuth client registration revoked successfully" }
+        })
+        async def revoke_oauth_client(
+            client_id: str, _: Annotated[None, Depends(validate_client_registration_token)]
+        ) -> Response:
+            """Revoke (deactivate) OAuth client registration"""
+
+            self.authenticator.revoke_oauth_client(client_id)
+            return Response(status_code=204)
+            
+        # Authorization Endpoint
+        @router.get("/authorize", description="OAuth 2.1 Authorization Endpoint", tags=["OAuth2"], response_model=None)
+        async def authorize_endpoint(
+            request: Request,
+            response_type: str = Query(default="code", description="OAuth response type"),
+            client_id: str = Query(..., description="OAuth client identifier"),
+            redirect_uri: str = Query(..., description="URI to redirect after authorization"),
+            scope: str = Query(default="read", description="Requested scope"),
+            state: str | None = Query(default=None, description="State parameter for CSRF protection"),
+            code_challenge: str = Query(..., description="PKCE code challenge (required)"),
+            code_challenge_method: str = Query(default="S256", description="PKCE code challenge method"),
+            user: UserInfoModel | None = Depends(self.get_current_user)
+        ):
+            """OAuth 2.1 authorization endpoint for initiating authorization code flow"""
+            
+            try:
+                # Validate response_type
+                if response_type != "code":
+                    raise InvalidInputError(400, "unsupported_response_type", "Only 'code' response type is supported")
+                
+                # Check if user is authenticated
+                if user is None:
+                    # User is not authenticated - serve login page
+                    return self.serve_login_page("/api/auth", request, client_id)
+                
+                # TODO: Serve a page with an "authorize" button even if user is already authenticated
+                # Ex. if not request.session.get("authorization_approved"), redirect to a page with button that submits to "/approve-authorization"
+                
+                # User is authenticated - generate authorization code
+                authorization_code = self.authenticator.create_authorization_code(
+                    client_id=client_id,
+                    username=user.username,
+                    redirect_uri=redirect_uri,
+                    scope=scope,
+                    code_challenge=code_challenge,
+                    code_challenge_method=code_challenge_method
+                )
+                
+                # Redirect back to client with authorization code
+                success_params = f"?code={authorization_code}"
+                if state:
+                    success_params += f"&state={state}"
+                
+                return RedirectResponse(url=f"{redirect_uri}{success_params}")
+                
+            except InvalidInputError as e:
+                if e.error == "invalid_request":
+                    error_params = f"?error={e.error}&error_description={e.error_description.replace(' ', '+')}"
+                    if state:
+                        error_params += f"&state={state}"
+                    return RedirectResponse(url=f"{redirect_uri}{error_params}")
+                else:
+                    raise e
+
+        # Token Endpoint
+        @router.post("/token", description="OAuth 2.1 Token Endpoint", tags=["OAuth2"])
+        async def token_endpoint(
+            grant_type: str = Form(...),
+            code: str | None = Form(default=None),
+            redirect_uri: str | None = Form(default=None),
+            code_verifier: str | None = Form(default=None),
+            refresh_token: str | None = Form(default=None),
+            client_id: str | None = Form(default=None),
+            client_secret: str | None = Form(default=None)
+        ) -> TokenResponse:
+            """OAuth 2.1 token endpoint for exchanging authorization code or refresh token for access token"""
+            
+            # Validate client credentials
+            auth_client_id = validate_oauth_client_credentials(client_id, client_secret)
+            
+            # Get token expiry configuration
+            expiry_mins = self._get_access_token_expiry_minutes()
+            
+            if grant_type == "authorization_code":
+                # Validate required parameters for authorization code flow
+                if not all([code, redirect_uri, code_verifier, auth_client_id]):
+                    raise InvalidInputError(400, "invalid_request", "Missing required parameters for authorization_code grant")
+                
+                # Type casts since we validated above
+                code = cast(str, code)
+                redirect_uri = cast(str, redirect_uri)
+                code_verifier = cast(str, code_verifier)
+                auth_client_id = cast(str, auth_client_id)
+                
+                # Exchange authorization code for tokens
+                token_response = self.authenticator.exchange_authorization_code(
+                    code=code,
+                    client_id=auth_client_id,
+                    redirect_uri=redirect_uri,
+                    code_verifier=code_verifier,
+                    access_token_expiry_minutes=expiry_mins
+                )
+                
+                return token_response
+                
+            elif grant_type == "refresh_token":
+                # Validate required parameters for refresh token flow
+                if not all([refresh_token, auth_client_id]):
+                    raise InvalidInputError(400, "invalid_request", "Missing required parameters for refresh_token grant")
+                
+                # Type casts since we validated above
+                refresh_token = cast(str, refresh_token)
+                auth_client_id = cast(str, auth_client_id)
+                
+                # Refresh access token
+                token_response = self.authenticator.refresh_oauth_access_token(
+                    refresh_token=refresh_token,
+                    client_id=auth_client_id,
+                    access_token_expiry_minutes=expiry_mins
+                )
+                
+                return token_response
+            
+            else:
+                raise InvalidInputError(400, "unsupported_grant_type", f"Grant type '{grant_type}' is not supported")
+
+        # Token Revocation Endpoint
+        @router.post("/token/revoke", description="OAuth 2.1 Token Revocation Endpoint", tags=["OAuth2"])
+        async def revoke_endpoint(
+            token: str = Form(..., description="The token to be revoked"),
+            token_type_hint: str | None = Form(default=None, description="Hint about the type of token being revoked"),
+            client_id: str | None = Form(default=None),
+            client_secret: str | None = Form(default=None)
+        ) -> Response:
+            """OAuth 2.1 token revocation endpoint for revoking refresh tokens"""
+            
+            # Validate client credentials
+            auth_client_id = validate_oauth_client_credentials(client_id, client_secret)
+            
+            # Revoke the token (per RFC 7009, always return 200 regardless of token validity)
+            try:
+                self.authenticator.revoke_oauth_token(auth_client_id, token, token_type_hint)
+            except InvalidInputError:
+                # Per OAuth spec, revocation endpoint should return 200 even for invalid tokens
+                pass
+            
+            return Response(status_code=200)
+
+        # Authorization Server Metadata Endpoint (well-known endpoint)
+        @app.get("/.well-known/oauth-authorization-server", tags=["OAuth2"], description="OAuth 2.1 Authorization Server Metadata")
+        async def authorization_server_metadata(request: Request) -> OAuthServerMetadata:
+            """OAuth 2.1 Authorization Server Metadata endpoint (RFC 8414)"""
+            
+            # Get the base URL from the request
+            scheme = "http" if request.url.hostname in ("localhost", "127.0.0.1") else "https"
+            base_url = scheme + "://" + request.url.netloc
+            
+            return OAuthServerMetadata(
+                issuer=base_url,
+                authorization_endpoint=f"{base_url}{router_path}/authorize",
+                token_endpoint=f"{base_url}{router_path}/token",
+                revocation_endpoint=f"{base_url}{router_path}/token/revoke",
+                registration_endpoint=f"{base_url}{router_path}/register",
+                scopes_supported=["read"],
+                response_types_supported=["code"],
+                grant_types_supported=["authorization_code", "refresh_token"],
+                token_endpoint_auth_methods_supported=["client_secret_post"],
+                code_challenge_methods_supported=["S256"]
+            )
+
+        app.include_router(router)