squirrels 0.5.0b3__py3-none-any.whl → 0.6.0.post0__py3-none-any.whl

squirrels/_api_routes/dashboards.py

@@ -0,0 +1,156 @@
+"""
+Dashboard routes for parameters and results
+"""
+from typing import Callable, Coroutine, Any
+from fastapi import FastAPI, Depends, Request
+from fastapi.responses import Response, HTMLResponse
+from fastapi.security import HTTPBearer
+from dataclasses import asdict
+from cachetools import TTLCache
+import time
+
+from .. import _constants as c, _utils as u
+from .._schemas import response_models as rm
+from .._exceptions import ConfigurationError
+from .._dashboards import Dashboard
+from .._schemas.query_param_models import get_query_models_for_parameters, get_query_models_for_dashboard
+from .._schemas.auth_models import AbstractUser
+from .base import RouteBase
+
+
+class DashboardRoutes(RouteBase):
+    """Dashboard 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
+        self.dashboard_results_cache = TTLCache(
+            maxsize=self.env_vars.dashboards_cache_size, 
+            ttl=self.env_vars.dashboards_cache_ttl_minutes*60
+        )
+        
+    async def _get_dashboard_results_helper(
+        self, dashboard: str, user: AbstractUser, selections: tuple[tuple[str, Any], ...], configurables: tuple[tuple[str, str], ...]
+    ) -> Dashboard:
+        """Helper to get dashboard results"""
+        cfg_filtered = {k: v for k, v in dict(configurables).items() if k in self.manifest_cfg.configurables}
+        return await self.project.dashboard(dashboard, user=user, selections=dict(selections), configurables=cfg_filtered)
+    
+    async def _get_dashboard_results_cachable(
+        self, dashboard: str, user: AbstractUser, selections: tuple[tuple[str, Any], ...], configurables: tuple[tuple[str, str], ...]
+    ) -> Dashboard:
+        """Cachable version of dashboard results helper"""
+        return await self.do_cachable_action(self.dashboard_results_cache, self._get_dashboard_results_helper, dashboard, user, selections, configurables)
+    
+    async def _get_dashboard_results_definition(
+        self, dashboard_name: str, user: AbstractUser, params: dict, headers: dict[str, str]
+    ) -> Response:
+        """Get dashboard results definition"""
+        get_dashboard_function = self._get_dashboard_results_helper if self.no_cache else self._get_dashboard_results_cachable
+        selections = self.get_selections_as_immutable(params, uncached_keys=set())
+        
+        user_has_elevated_privileges = u.user_has_elevated_privileges(user.access_level, self.env_vars.elevated_access_level)
+        configurables = self.get_configurables_from_headers(headers) if user_has_elevated_privileges else tuple()
+        dashboard_obj = await get_dashboard_function(dashboard_name, user, selections, configurables)
+        
+        if dashboard_obj._format == c.PNG:
+            assert isinstance(dashboard_obj._content, bytes)
+            result = Response(dashboard_obj._content, media_type="image/png")
+        elif dashboard_obj._format == c.HTML:
+            result = HTMLResponse(dashboard_obj._content)
+        else:
+            raise NotImplementedError()
+        return result 
+    
+    def setup_routes(
+        self, app: FastAPI, param_fields: dict, 
+        get_parameters_definition: Callable[..., Coroutine[Any, Any, rm.ParametersModel]]
+    ) -> None:
+        """Setup dashboard routes"""
+        
+        dashboard_results_path = '/dashboards/{dashboard}'
+        dashboard_parameters_path = dashboard_results_path + '/parameters'
+        
+        def validate_parameters_list(parameters: list[str] | None, entity_type: str, dashboard_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} '{dashboard_name}' use parameter '{param}' which doesn't exist. Available parameters are:"
+                        f"\n  {all_params}"
+                    )
+        
+        # Dashboard parameters and results APIs
+        for dashboard_name, dashboard in self.project._dashboards.items():
+            dashboard_name_for_api = u.normalize_name_for_api(dashboard_name)
+            curr_parameters_path = dashboard_parameters_path.format(dashboard=dashboard_name_for_api)
+            curr_results_path = dashboard_results_path.format(dashboard=dashboard_name_for_api)
+
+            validate_parameters_list(dashboard.config.parameters, "Dashboard", dashboard_name)
+            
+            QueryModelForGetParams, QueryModelForPostParams = get_query_models_for_parameters(param_fields, dashboard.config.parameters)
+            QueryModelForGetDash, QueryModelForPostDash = get_query_models_for_dashboard(param_fields, dashboard.config.parameters)
+
+            @app.get(curr_parameters_path, tags=[f"Dashboard '{dashboard_name}'"], description=self._parameters_description)
+            async def get_dashboard_parameters(
+                request: Request, params: QueryModelForGetParams, user=Depends(self.get_current_user)
+            ) -> rm.ParametersModel:
+                start = time.time()
+                curr_dashboard_name = self.get_name_from_path_section(request, -2)
+                parameters_list = self.project._dashboards[curr_dashboard_name].config.parameters    
+                scope = self.project._dashboards[curr_dashboard_name].config.scope
+                result = await get_parameters_definition(
+                    parameters_list, "dashboard", curr_dashboard_name, scope, user, asdict(params)
+                )
+                self.logger.log_activity_time(
+                    "GET REQUEST for PARAMETERS", start, additional_data={"dashboard_name": curr_dashboard_name}
+                )
+                return result
+
+            @app.post(curr_parameters_path, tags=[f"Dashboard '{dashboard_name}'"], description=self._parameters_description)
+            async def get_dashboard_parameters_with_post(
+                request: Request, params: QueryModelForPostParams, user=Depends(self.get_current_user)
+            ) -> rm.ParametersModel:
+                start = time.time()
+                curr_dashboard_name = self.get_name_from_path_section(request, -2)
+                parameters_list = self.project._dashboards[curr_dashboard_name].config.parameters
+                scope = self.project._dashboards[curr_dashboard_name].config.scope
+                result = await get_parameters_definition(
+                    parameters_list, "dashboard", curr_dashboard_name, scope, user, params.model_dump()
+                )
+                self.logger.log_activity_time(
+                    "POST REQUEST for PARAMETERS", start, additional_data={"dashboard_name": curr_dashboard_name}
+                )
+                return result
+            
+            @app.get(curr_results_path, tags=[f"Dashboard '{dashboard_name}'"], description=dashboard.config.description)
+            async def get_dashboard_results(
+                request: Request, params: QueryModelForGetDash, user=Depends(self.get_current_user)
+            ) -> Response:
+                start = time.time()
+                curr_dashboard_name = self.get_name_from_path_section(request, -1)
+                result = await self._get_dashboard_results_definition(
+                    curr_dashboard_name, user, asdict(params), headers=dict(request.headers)
+                )
+                self.logger.log_activity_time(
+                    "GET REQUEST for DASHBOARD RESULTS", start, additional_data={"dashboard_name": curr_dashboard_name}
+                )
+                return result
+
+            @app.post(curr_results_path, tags=[f"Dashboard '{dashboard_name}'"], description=dashboard.config.description, response_class=Response)
+            async def get_dashboard_results_with_post(
+                request: Request, params: QueryModelForPostDash, user=Depends(self.get_current_user)
+            ) -> Response:
+                start = time.time()
+                curr_dashboard_name = self.get_name_from_path_section(request, -1)
+                result = await self._get_dashboard_results_definition(
+                    curr_dashboard_name, user, params.model_dump(), headers=dict(request.headers)
+                )
+                self.logger.log_activity_time(
+                    "POST REQUEST for DASHBOARD RESULTS", start, additional_data={"dashboard_name": curr_dashboard_name}
+                )
+                return result
+    

squirrels/_api_routes/data_management.py

@@ -0,0 +1,148 @@
+"""
+Data management routes for build and query models
+"""
+from typing import Any, Annotated
+from fastapi import FastAPI, Depends, Request, Response, status, Path
+from fastapi.responses import JSONResponse
+from fastapi.security import HTTPBearer
+from dataclasses import asdict
+from cachetools import TTLCache
+import time
+
+from .. import _constants as c, _utils as u
+from .._schemas import response_models as rm
+from .._exceptions import InvalidInputError
+from .._schemas.auth_models import AbstractUser
+from .._dataset_types import DatasetResult
+from .._schemas.query_param_models import get_query_models_for_querying_models, get_query_models_for_compiled_models
+from .base import RouteBase
+
+
+class DataManagementRoutes(RouteBase):
+    """Data management routes for build and query operations"""
+    
+    def __init__(self, get_bearer_token: HTTPBearer, project, no_cache: bool = False):
+        super().__init__(get_bearer_token, project, no_cache)
+        
+        # Setup cache (same settings as dataset results cache)
+        self.query_models_cache = TTLCache(
+            maxsize=self.env_vars.datasets_cache_size, 
+            ttl=self.env_vars.datasets_cache_ttl_minutes*60
+        )
+        
+    async def _query_models_helper(
+        self, sql_query: str, user: AbstractUser, selections: tuple[tuple[str, Any], ...], configurables: tuple[tuple[str, str], ...]
+    ) -> DatasetResult:
+        """Helper to query models"""
+        cfg_filtered = {k: v for k, v in dict(configurables).items() if k in self.manifest_cfg.configurables}
+        return await self.project.query_models(sql_query, user=user, selections=dict(selections), configurables=cfg_filtered)
+
+    async def _query_models_cachable(
+        self, sql_query: str, user: AbstractUser, selections: tuple[tuple[str, Any], ...], configurables: tuple[tuple[str, str], ...]
+    ) -> DatasetResult:
+        """Cachable version of query models helper"""
+        return await self.do_cachable_action(self.query_models_cache, self._query_models_helper, sql_query, user, selections, configurables)
+
+    async def _query_models_definition(
+        self, user: AbstractUser, params: dict, *, headers: dict[str, str]
+    ) -> rm.DatasetResultModel:
+        """Query models definition"""
+        if not u.user_has_elevated_privileges(user.access_level, self.env_vars.elevated_access_level):
+            raise InvalidInputError(403, "unauthorized_access_to_query_models", f"User '{user}' does not have permission to query data models")
+        
+        sql_query = params.get("x_sql_query")
+        if sql_query is None:
+            raise InvalidInputError(400, "sql_query_required", "SQL query must be provided")
+        
+        query_models_function = self._query_models_helper if self.no_cache else self._query_models_cachable
+        uncached_keys = {"x_sql_query", "x_orientation", "x_offset", "x_limit"}
+        selections = self.get_selections_as_immutable(params, uncached_keys)
+        configurables = self.get_configurables_from_headers(headers)
+        result = await query_models_function(sql_query, user, selections, configurables)
+        
+        result_format = self.extract_orientation_offset_and_limit(params)
+        return rm.DatasetResultModel(**result.to_json(result_format)) 
+    
+    async def _get_compiled_model_definition(
+        self, model_name: str, user: AbstractUser, params: dict, *, headers: dict[str, str]
+    ) -> rm.CompiledQueryModel:
+        """Get compiled model definition"""
+        normalized_model_name = u.normalize_name(model_name)
+        # self._validate_request_params(all_request_params, params, headers)
+
+        # Internal users only
+        if not u.user_has_elevated_privileges(user.access_level, self.env_vars.elevated_access_level):
+            raise InvalidInputError(403, "unauthorized_access_to_compile_model", f"User '{user}' does not have permission to fetch compiled SQL")
+        
+        selections = self.get_selections_as_immutable(params, uncached_keys=set())
+        configurables = self.get_configurables_from_headers(headers)
+        cfg_filtered = {k: v for k, v in dict(configurables).items() if k in self.manifest_cfg.configurables}
+        return await self.project.get_compiled_model_query(normalized_model_name, user=user, selections=dict(selections), configurables=cfg_filtered)
+        
+    def setup_routes(self, app: FastAPI, param_fields: dict) -> None:
+        """Setup data management routes"""
+        
+        # Build project endpoint
+        build_path = '/build'
+        
+        @app.post(build_path, tags=["Data Management"], summary="Build or update the Virtual Data Lake (VDL) for the project")
+        async def build(user=Depends(self.get_current_user)) -> Response:
+            if not u.user_has_elevated_privileges(user.access_level, self.env_vars.elevated_access_level):
+                raise InvalidInputError(403, "unauthorized_access_to_build_model", f"User '{user}' does not have permission to build the virtual data lake (VDL)")
+            await self.project.build()
+            return Response(status_code=status.HTTP_200_OK)
+        
+        # Query result endpoints
+        query_models_path = '/query-result'
+        QueryModelForQueryModels, QueryModelForPostQueryModels = get_query_models_for_querying_models(param_fields)
+
+        @app.get(query_models_path, tags=["Data Management"], response_class=JSONResponse)
+        async def query_models(
+            request: Request, params: QueryModelForQueryModels, user=Depends(self.get_current_user)
+        ) -> rm.DatasetResultModel:
+            start = time.time()
+            result = await self._query_models_definition(user, asdict(params), headers=dict(request.headers))
+            self.logger.log_activity_time("GET REQUEST for QUERY MODELS", start)
+            return result
+        
+        @app.post(query_models_path, tags=["Data Management"], response_class=JSONResponse)
+        async def query_models_with_post(
+            request: Request, params: QueryModelForPostQueryModels, user=Depends(self.get_current_user)
+        ) -> rm.DatasetResultModel:
+            start = time.time()
+            result = await self._query_models_definition(user, params.model_dump(), headers=dict(request.headers))
+            self.logger.log_activity_time("POST REQUEST for QUERY MODELS", start)
+            return result
+
+        # Compiled models endpoints
+        compiled_models_path = '/compiled-models/{model_name}'
+        QueryModelForGetCompiled, QueryModelForPostCompiled = get_query_models_for_compiled_models(param_fields)
+
+        @app.get(compiled_models_path, tags=["Data Management"], response_class=JSONResponse, summary="Get compiled definition for a model")
+        async def get_compiled_model(
+            request: Request, 
+            model_name: Annotated[str, Path(description="The name of the model. Both snake case (with underscores) and kebab case (with dashes) are supported")], 
+            params: QueryModelForGetCompiled, 
+            user=Depends(self.get_current_user)
+        ) -> rm.CompiledQueryModel:
+            start = time.time()
+            result = await self._get_compiled_model_definition(model_name, user, asdict(params), headers=dict(request.headers))
+            self.logger.log_activity_time(
+                "GET REQUEST for GET COMPILED MODEL", start, additional_data={"model_name": model_name}
+            )
+            return result
+
+        @app.post(compiled_models_path, tags=["Data Management"], response_class=JSONResponse, summary="Get compiled definition for a model")
+        async def get_compiled_model_with_post(
+            request: Request, 
+            model_name: Annotated[str, Path(description="The name of the model. Both snake case (with underscores) and kebab case (with dashes) are supported")], 
+            params: QueryModelForPostCompiled, 
+            user=Depends(self.get_current_user)
+        ) -> rm.CompiledQueryModel:
+            start = time.time()
+            result = await self._get_compiled_model_definition(model_name, user, params.model_dump(), headers=dict(request.headers))
+            self.logger.log_activity_time(
+                "POST REQUEST for GET COMPILED MODEL", start, additional_data={"model_name": model_name}
+            )
+            return result
+    

squirrels/_api_routes/datasets.py

@@ -0,0 +1,220 @@
+"""
+Dataset routes for parameters and results
+"""
+from typing import Callable, Coroutine, Any
+from fastapi import FastAPI, Depends, Request
+from fastapi.responses import JSONResponse
+from fastapi.security import HTTPBearer
+from dataclasses import asdict
+from cachetools import TTLCache
+
+import time
+import polars as pl
+
+from .. import _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 .._schemas.auth_models import AbstractUser
+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
+        self.dataset_results_cache = TTLCache(
+            maxsize=self.env_vars.datasets_cache_size, 
+            ttl=self.env_vars.datasets_cache_ttl_minutes*60
+        )
+        
+        # Setup max rows
+        self.max_result_rows = self.env_vars.datasets_max_rows_output
+        
+        # Setup SQL query timeout
+        self.sql_timeout_seconds = self.env_vars.datasets_sql_timeout_seconds
+        
+    async def _get_dataset_results_helper(
+        self, dataset: str, user: AbstractUser, selections: tuple[tuple[str, Any], ...], configurables: tuple[tuple[str, str], ...]
+    ) -> DatasetResult:
+        """Helper to get dataset results"""
+        # Only pass configurables that are defined in manifest
+        cfg_filtered = {k: v for k, v in dict(configurables).items() if k in self.manifest_cfg.configurables}
+        return await self.project._dataset_result(dataset, user=user, selections=dict(selections), configurables=cfg_filtered)
+
+    async def _get_dataset_results_cachable(
+        self, dataset: str, user: AbstractUser, selections: tuple[tuple[str, Any], ...], configurables: tuple[tuple[str, str], ...]
+    ) -> 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, configurables
+        )
+    
+    async def _get_dataset_result_object(
+        self, dataset_name: str, user: AbstractUser, params: dict, configurables: tuple[tuple[str, str], ...]
+    ) -> DatasetResult:
+        """Get dataset result object"""
+        # self._validate_request_params(all_request_params, params, headers)
+
+        get_dataset_function = self._get_dataset_results_helper if self.no_cache else self._get_dataset_results_cachable
+        uncached_keys = {"x_sql_query", "x_orientation", "x_offset", "x_limit"}
+        selections = self.get_selections_as_immutable(params, uncached_keys)
+        
+        user_has_elevated_privileges = u.user_has_elevated_privileges(user.access_level, self.env_vars.elevated_access_level)
+        configurables_if_elevated: tuple[tuple[str, str], ...] = configurables if user_has_elevated_privileges else tuple()
+        result = await get_dataset_function(dataset_name, user, selections, configurables_if_elevated)
+        
+        # Apply optional final SQL transformation before select/limit/offset
+        sql_query = params.get("x_sql_query")
+        if sql_query:
+            try:
+                transformed = await u.run_polars_sql_on_dataframes(
+                    sql_query, {"result": result.df.lazy()}, timeout_seconds=self.sql_timeout_seconds, max_rows=self.max_result_rows+1
+                )
+            except Exception as e:
+                raise InvalidInputError(400, "invalid_sql_query", "Failed to run provided Polars SQL on the dataset result") from e
+            
+            # Enforce max result rows on transformed result
+            row_count = transformed.select(pl.len()).item()
+            if row_count > self.max_result_rows:
+                raise InvalidInputError(
+                    413,
+                    "dataset_result_too_large",
+                    f"The transformed dataset result exceeds the maximum allowed of {self.max_result_rows} rows."
+                )
+            
+            transformed = transformed.drop("_row_num", strict=False).with_row_index("_row_num", offset=1)
+            result = DatasetResult(target_model_config=result.target_model_config, df=transformed)
+        
+        return result
+
+    async def _get_dataset_results_definition(
+        self, dataset_name: str, user: AbstractUser, params: dict, headers: dict[str, str]
+    ) -> rm.DatasetResultModel:
+        """Get dataset results definition"""        
+        configurables = self.get_configurables_from_headers(headers)
+        result = await self._get_dataset_result_object(dataset_name, user, params, configurables)
+
+        result_format = self.extract_orientation_offset_and_limit(params)
+        return rm.DatasetResultModel(**result.to_json(result_format)) 
+    
+    def setup_routes(
+        self, app: FastAPI, param_fields: dict, 
+        get_parameters_definition: Callable[..., Coroutine[Any, Any, rm.ParametersModel]]
+    ) -> None:
+        """Setup dataset routes"""
+        
+        dataset_results_path = '/datasets/{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: AbstractUser, 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, params
+            )
+            return result
+
+        # Dataset parameters and results APIs
+        for dataset_name, dataset_config in self.manifest_cfg.datasets.items():
+            dataset_name_for_api = u.normalize_name_for_api(dataset_name)
+            curr_parameters_path = dataset_parameters_path.format(dataset=dataset_name_for_api)
+            curr_results_path = dataset_results_path.format(dataset=dataset_name_for_api)
+
+            validate_parameters_list(dataset_config.parameters, "Dataset", dataset_name)
+
+            QueryModelForGetParams, QueryModelForPostParams = get_query_models_for_parameters(param_fields, dataset_config.parameters)
+            QueryModelForGetDataset, QueryModelForPostDataset = get_query_models_for_dataset(param_fields, dataset_config.parameters)
+
+            @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)
+            ) -> 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, asdict(params))
+                self.logger.log_activity_time(
+                    "GET REQUEST for PARAMETERS", start, additional_data={"dataset_name": curr_dataset_name}
+                )
+                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)
+            ) -> 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, params.model_dump())
+                self.logger.log_activity_time(
+                    "POST REQUEST for PARAMETERS", start, additional_data={"dataset_name": curr_dataset_name}
+                )
+                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)
+            ) -> 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, asdict(params), headers=dict(request.headers)
+                )
+                self.logger.log_activity_time(
+                    "GET REQUEST for DATASET RESULTS", start, additional_data={"dataset_name": curr_dataset_name}
+                )
+                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)
+            ) -> 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, params.model_dump(), headers=dict(request.headers)
+                )
+                self.logger.log_activity_time(
+                    "POST REQUEST for DATASET RESULTS", start, additional_data={"dataset_name": curr_dataset_name}
+                )
+                return result
+    
+        # MCP-callable methods (exposed as instance attributes for McpServerBuilder)
+        
+        async def get_dataset_parameters_for_mcp(
+            dataset: str, parameter_name: str, selected_ids: str | list[str], user: AbstractUser
+        ) -> rm.ParametersModel:
+            """Get dataset parameter updates for MCP tools. Takes user and headers."""
+            dataset_name = u.normalize_name(dataset)
+            parameters = {
+                "x_parent_param": parameter_name,
+                parameter_name: selected_ids
+            }
+            return await get_dataset_parameters_updates(dataset_name, user, parameters)
+        
+        async def get_dataset_results_for_mcp(
+            dataset: str, parameters: dict[str, Any], sql_query: str | None, user: AbstractUser, configurables: tuple[tuple[str, str], ...]
+        ) -> DatasetResult:
+            """Get dataset results for MCP tools. Takes user and configurables."""
+            dataset_name = u.normalize_name(dataset)
+            parameters.update({ "x_sql_query": sql_query })
+            return await self._get_dataset_result_object(dataset_name, user, parameters, configurables)
+        
+        # Store the MCP functions as instance attributes for access by McpServerBuilder
+        self._get_dataset_parameters_for_mcp = get_dataset_parameters_for_mcp
+        self._get_dataset_results_for_mcp = get_dataset_results_for_mcp
+