squirrels 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

squirrels/_api_routes/oauth2.py

@@ -0,0 +1,298 @@
+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, AbstractUser
+)
+from .._exceptions import InvalidInputError
+from .. import _utils as u
+
+
+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, squirrels_version_path: str) -> None:
+        """Setup all OAuth2 routes"""
+
+        auth_path = squirrels_version_path + "/auth"
+        router_path = "/oauth2"
+        router = APIRouter(prefix=router_path)
+        
+        # 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("/client", 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: AbstractUser = 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.access_level == "guest":
+                    # User is not authenticated - serve login page
+                    return self.serve_login_page(auth_path, 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 = u.get_scheme(request.url.hostname)
+            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}/client",
+                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)

squirrels/_api_routes/project.py

@@ -0,0 +1,252 @@
+"""
+Project metadata routes
+"""
+from typing import Any
+from fastapi import FastAPI, Depends, Request
+from fastapi.responses import JSONResponse
+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 _utils as u, _constants as c
+from .._schemas import response_models as rm
+from .._parameter_sets import ParameterSet
+from .._exceptions import ConfigurationError, InvalidInputError
+from .._manifest import PermissionScope, AuthenticationEnforcement
+from .._version import __version__
+from .._schemas.query_param_models import get_query_models_for_parameters
+from .._schemas.auth_models import AbstractUser
+from .base import RouteBase
+
+
+class ProjectRoutes(RouteBase):
+    """Project metadata and data catalog routes"""
+    
+    def __init__(self, get_bearer_token: HTTPBearer, project, no_cache: bool = False):
+        super().__init__(get_bearer_token, project, no_cache)
+        
+        # Setup caches
+        parameters_cache_size = int(self.env_vars.get(c.SQRL_PARAMETERS_CACHE_SIZE, 1024))
+        parameters_cache_ttl = int(self.env_vars.get(c.SQRL_PARAMETERS_CACHE_TTL_MINUTES, 60))
+        self.parameters_cache = TTLCache(maxsize=parameters_cache_size, ttl=parameters_cache_ttl*60)
+
+    async def _get_parameters_helper(
+        self, parameters_tuple: tuple[str, ...] | None, entity_type: str, entity_name: str, entity_scope: PermissionScope,
+        user: AbstractUser, selections: tuple[tuple[str, Any], ...]
+    ) -> ParameterSet:
+        """Helper for getting parameters"""
+        selections_dict = dict(selections)
+        if "x_parent_param" not in selections_dict:
+            if len(selections_dict) > 1:
+                raise InvalidInputError(400, "invalid_input_for_cascading_parameters", f"The parameters endpoint takes at most 1 widget parameter selection (unless x_parent_param is provided). Got {selections_dict}")
+            elif len(selections_dict) == 1:
+                parent_param = next(iter(selections_dict))
+                selections_dict["x_parent_param"] = parent_param
+        
+        parent_param = selections_dict.get("x_parent_param")
+        if parent_param is not None and parent_param not in selections_dict:
+                # this condition is possible for multi-select parameters with empty selection
+            selections_dict[parent_param] = list()
+        
+        if not self.authenticator.can_user_access_scope(user, entity_scope):
+            raise self.project._permission_error(user, entity_type, entity_name, entity_scope.name)
+        
+        param_set = self.param_cfg_set.apply_selections(parameters_tuple, selections_dict, user, parent_param=parent_param)
+        return param_set
+
+    async def _get_parameters_cachable(
+        self, parameters_tuple: tuple[str, ...] | None, entity_type: str, entity_name: str, entity_scope: PermissionScope,
+        user: AbstractUser, selections: tuple[tuple[str, Any], ...]
+    ) -> ParameterSet:
+        """Cachable version of parameters helper"""
+        return await self.do_cachable_action(
+            self.parameters_cache, self._get_parameters_helper, parameters_tuple, entity_type, entity_name, entity_scope, user, selections
+        )
+        
+    def setup_routes(
+        self, app: FastAPI, mcp: FastMCP, project_metadata_path: str, project_name: str, project_version: str, project_label: str, param_fields: dict
+    ):
+        """Setup project metadata routes"""
+
+        elevated_access_level = self.project._elevated_access_level
+        if elevated_access_level != "admin":
+            self.logger.warning(f"{c.SQRL_PERMISSIONS_ELEVATED_ACCESS_LEVEL} has been set to a non-admin access level. For security reasons, DO NOT expose the APIs for this app publicly!")
+        
+        # Project metadata endpoint
+        @app.get(project_metadata_path, tags=["Project Metadata"], response_class=JSONResponse)
+        async def get_project_metadata(request: Request) -> rm.ProjectModel:
+            return rm.ProjectModel(
+                name=project_name,
+                version=project_version,
+                label=self.manifest_cfg.project_variables.label,
+                description=self.manifest_cfg.project_variables.description,
+                elevated_access_level=elevated_access_level,
+                redoc_path=project_metadata_path + "/redoc",
+                swagger_path=project_metadata_path + "/docs",
+                mcp_server_path=project_metadata_path + "/mcp",
+                squirrels_version=__version__
+            )
+        
+        # Data catalog endpoint
+        data_catalog_path = project_metadata_path + '/data-catalog'
+        
+        async def get_data_catalog0(user: AbstractUser) -> rm.CatalogModel:
+            parameters = self.param_cfg_set.apply_selections(None, {}, user)
+            parameters_model = parameters.to_api_response_model0()
+            full_parameters_list = [p.name for p in parameters_model.parameters]
+            user_has_elevated_privileges = u.user_has_elevated_privileges(user.access_level, elevated_access_level)
+
+            dataset_items: list[rm.DatasetItemModel] = []
+            for name, config in self.manifest_cfg.datasets.items():
+                if self.authenticator.can_user_access_scope(user, config.scope):
+                    name_for_api = u.normalize_name_for_api(name)
+                    metadata = self.project.dataset_metadata(name).to_json()
+                    parameters = config.parameters if config.parameters is not None else full_parameters_list
+                    
+                    # Build dataset-specific configurables list
+                    if user_has_elevated_privileges:
+                        dataset_configurables_defaults = self.manifest_cfg.get_default_configurables(name)
+                        dataset_configurables_list = [
+                            rm.ConfigurableDefaultModel(name=name, default=default)
+                            for name, default in dataset_configurables_defaults.items()
+                        ]
+                    else:
+                        dataset_configurables_list = []
+                    
+                    dataset_items.append(rm.DatasetItemModel(
+                        name=name, label=config.label, 
+                        description=config.description,
+                        schema=metadata["schema"], # type: ignore
+                        configurables=dataset_configurables_list,
+                        parameters=parameters,
+                        parameters_path=f"{project_metadata_path}/dataset/{name_for_api}/parameters",
+                        result_path=f"{project_metadata_path}/dataset/{name_for_api}"
+                    ))
+            
+            dashboard_items: list[rm.DashboardItemModel] = []
+            for name, dashboard in self.project._dashboards.items():
+                config = dashboard.config
+                if self.authenticator.can_user_access_scope(user, config.scope):
+                    name_for_api = u.normalize_name_for_api(name)
+
+                    try:
+                        dashboard_format = self.project._dashboards[name].get_dashboard_format()
+                    except KeyError:
+                        raise ConfigurationError(f"No dashboard file found for: {name}")
+                    
+                    parameters = config.parameters if config.parameters is not None else full_parameters_list
+                    dashboard_items.append(rm.DashboardItemModel(
+                        name=name, label=config.label, 
+                        description=config.description, 
+                        result_format=dashboard_format,
+                        parameters=parameters,
+                        parameters_path=f"{project_metadata_path}/dashboard/{name_for_api}/parameters",
+                        result_path=f"{project_metadata_path}/dashboard/{name_for_api}"
+                    ))
+            
+            if user_has_elevated_privileges:
+                compiled_dag = await self.project._get_compiled_dag(user)
+                connections_items = self.project._get_all_connections()
+                data_models = self.project._get_all_data_models(compiled_dag)
+                lineage_items = self.project._get_all_data_lineage(compiled_dag)
+                configurables_list = [
+                    rm.ConfigurableItemModel(name=name, label=cfg.label, default=cfg.default, description=cfg.description)
+                    for name, cfg in self.manifest_cfg.configurables.items()
+                ]
+            else:
+                connections_items = []
+                data_models = []
+                lineage_items = []
+                configurables_list = []
+
+            return rm.CatalogModel(
+                parameters=parameters_model.parameters, 
+                datasets=dataset_items, 
+                dashboards=dashboard_items,
+                connections=connections_items,
+                models=data_models,
+                lineage=lineage_items,
+                configurables=configurables_list,
+            )
+        
+        @app.get(data_catalog_path, tags=["Project Metadata"], summary="Get catalog of datasets and dashboards available for user")
+        async def get_data_catalog(request: Request, user: AbstractUser = Depends(self.get_current_user)) -> rm.CatalogModel:
+            """
+            Get catalog of datasets and dashboards available for the authenticated user.
+            
+            For admin users, this endpoint will also return detailed information about all models and their lineage in the project.
+            """
+            start = time.time()
+
+            # If authentication is required, require user to be authenticated to access catalog
+            if self.manifest_cfg.authentication.enforcement == AuthenticationEnforcement.REQUIRED and user.access_level == "guest":
+                raise InvalidInputError(401, "user_required", "Authentication is required to access the data catalog")
+            data_catalog = await get_data_catalog0(user)
+            
+            self.log_activity_time("GET REQUEST for DATA CATALOG", start, request)
+            return data_catalog
+        
+        @mcp.tool(
+            name=f"get_data_catalog_from_{project_name}", 
+            title=f"Get Data Catalog (Project: {project_label})",
+            description=dedent(f"""
+            Use this tool to get the details of all datasets and parameters you can access in the Squirrels project '{project_name}'.
+            
+            Unless the data catalog for this project has already been provided, use this tool at the start of each conversation.
+            """).strip()
+        )
+        async def get_data_catalog_tool(ctx: Context) -> rm.CatalogModelForTool:
+            headers = self.get_headers_from_tool_ctx(ctx)
+            user = self.get_user_from_tool_headers(headers)
+            data_catalog = await get_data_catalog0(user)
+            return rm.CatalogModelForTool(parameters=data_catalog.parameters, datasets=data_catalog.datasets)
+        
+        # Project-level parameters endpoints
+        project_level_parameters_path = project_metadata_path + '/parameters'
+        parameters_description = "Selections of one parameter may cascade the available options in another parameter. " \
+                "For example, if the dataset has parameters for 'country' and 'city', available options for 'city' would " \
+                "depend on the selected option 'country'. If a parameter has 'trigger_refresh' as true, provide the parameter " \
+                "selection to this endpoint whenever it changes to refresh the parameter options of children parameters."
+
+        QueryModelForGetProjectParams, QueryModelForPostProjectParams = get_query_models_for_parameters(None, param_fields)
+
+        async def get_parameters_definition(
+            parameters_list: list[str] | None, entity_type: str, entity_name: str, entity_scope: PermissionScope,
+            user: AbstractUser, all_request_params: dict, params: dict, *, headers: dict[str, str]
+        ) -> rm.ParametersModel:
+            self._validate_request_params(all_request_params, params, headers)
+
+            get_parameters_function = self._get_parameters_helper if self.no_cache else self._get_parameters_cachable
+            selections = self.get_selections_as_immutable(params, uncached_keys={"x_verify_params"})
+            parameters_tuple = tuple(parameters_list) if parameters_list is not None else None
+            result = await get_parameters_function(parameters_tuple, entity_type, entity_name, entity_scope, user, selections)
+            return result.to_api_response_model0()
+
+        @app.get(project_level_parameters_path, tags=["Project Metadata"], description=parameters_description)
+        async def get_project_parameters(
+            request: Request, params: QueryModelForGetProjectParams, user=Depends(self.get_current_user) # type: ignore
+        ) -> rm.ParametersModel:
+            start = time.time()
+            result = await get_parameters_definition(
+                None, "project", "", PermissionScope.PUBLIC, user, dict(request.query_params), asdict(params), headers=dict(request.headers)
+            )
+            self.log_activity_time("GET REQUEST for PROJECT PARAMETERS", start, request)
+            return result
+
+        @app.post(project_level_parameters_path, tags=["Project Metadata"], description=parameters_description)
+        async def get_project_parameters_with_post(
+            request: Request, params: QueryModelForPostProjectParams, user=Depends(self.get_current_user) # type: ignore
+        ) -> rm.ParametersModel:
+            start = time.time()
+            payload: dict = await request.json()
+            result = await get_parameters_definition(
+                None, "project", "", PermissionScope.PUBLIC, user, payload, params.model_dump(), headers=dict(request.headers)
+            )
+            self.log_activity_time("POST REQUEST for PROJECT PARAMETERS", start, request)
+            return result
+        
+        return get_parameters_definition
+