squirrels 0.5.0b4__py3-none-any.whl → 0.5.1__py3-none-any.whl

squirrels/_api_routes/data_management.py

@@ -1,7 +1,7 @@
 """
 Data management routes for build and query models
 """
-from typing import Callable, Any
+from typing import Any
 from fastapi import FastAPI, Depends, Request, Response, status
 from fastapi.responses import JSONResponse
 from fastapi.security import HTTPBearer
@@ -9,13 +9,12 @@ from dataclasses import asdict
 from cachetools import TTLCache
 import time
 
-from .. import _constants as c
+from .. import _constants as c, _utils as u
 from .._schemas import response_models as rm
 from .._exceptions import InvalidInputError
-from .._auth import BaseUser
-from .._manifest import PermissionScope
+from .._schemas.auth_models import AbstractUser
 from .._dataset_types import DatasetResult
-from .._schemas.query_param_models import get_query_models_for_querying_models
+from .._schemas.query_param_models import get_query_models_for_querying_models, get_query_models_for_compiled_models
 from .base import RouteBase
 
 
@@ -31,55 +30,74 @@ class DataManagementRoutes(RouteBase):
         self.query_models_cache = TTLCache(maxsize=dataset_results_cache_size, ttl=dataset_results_cache_ttl*60)
         
     async def _query_models_helper(
-        self, sql_query: str, user: BaseUser | None, selections: tuple[tuple[str, Any], ...]
+        self, sql_query: str, user: AbstractUser, selections: tuple[tuple[str, Any], ...], configurables: tuple[tuple[str, str], ...]
     ) -> DatasetResult:
         """Helper to query models"""
-        return await self.project.query_models(sql_query, selections=dict(selections), user=user)
+        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: BaseUser | None, selections: tuple[tuple[str, Any], ...]
+        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)
+        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: BaseUser | None, all_request_params: dict, params: dict
+        self, user: AbstractUser, all_request_params: dict, params: dict, *, headers: dict[str, str]
     ) -> rm.DatasetResultModel:
         """Query models definition"""
-        self._validate_request_params(all_request_params, params)
+        self._validate_request_params(all_request_params, params, headers)
 
-        if not self.authenticator.can_user_access_scope(user, PermissionScope.PRIVATE):
-            raise InvalidInputError(403, "Forbidden to query data models", f"User '{user}' does not have permission to query data models")
+        if not u.user_has_elevated_privileges(user.access_level, self.project._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 must be provided", "SQL query must be provided")
+            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_verify_params", "x_sql_query", "x_orientation", "x_limit", "x_offset"}
         selections = self.get_selections_as_immutable(params, uncached_keys)
-        result = await query_models_function(sql_query, user, selections)
+        configurables = self.get_configurables_from_headers(headers)
+        result = await query_models_function(sql_query, user, selections, configurables)
         
-        orientation = params.get("x_orientation", "records")
+        orientation_header = headers.get("x-orientation")
+        orientation = str(orientation_header).lower() if orientation_header is not None else params.get("x_orientation", "records")
         limit = params.get("x_limit", 1000)
         offset = params.get("x_offset", 0)
-        return rm.DatasetResultModel(**result.to_json(orientation, tuple(), limit, offset)) 
+        return rm.DatasetResultModel(**result.to_json(orientation, limit, offset)) 
     
+    async def _get_compiled_model_definition(
+        self, model_name: str, user: AbstractUser, all_request_params: dict, 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.project._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={"x_verify_params"})
+        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, project_metadata_path: str, param_fields: dict) -> None:
         """Setup data management routes"""
         
         # Build project endpoint
         build_path = project_metadata_path + '/build'
         
-        @app.post(build_path, tags=["Data Management"], summary="Build or update the virtual data environment for the project")
+        @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)): # type: ignore
-            if not self.authenticator.can_user_access_scope(user, PermissionScope.PRIVATE):
-                raise InvalidInputError(403, "Forbidden to build", f"User '{user}' does not have permission to build the virtual data environment")
-            await self.project.build(stage_file=True)
+            if not u.user_has_elevated_privileges(user.access_level, self.project._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 models endpoints
-        query_models_path = project_metadata_path + '/query-models'
+        # Query result endpoints
+        query_models_path = project_metadata_path + '/query-result'
         QueryModelForQueryModels, QueryModelForPostQueryModels = get_query_models_for_querying_models(param_fields)
 
         @app.get(query_models_path, tags=["Data Management"], response_class=JSONResponse)
@@ -87,8 +105,8 @@ class DataManagementRoutes(RouteBase):
             request: Request, params: QueryModelForQueryModels, user=Depends(self.get_current_user)  # type: ignore
         ) -> rm.DatasetResultModel:
             start = time.time()
-            result = await self._query_models_definition(user, dict(request.query_params), asdict(params))
-            self.log_activity_time("GET REQUEST for QUERY MODELS", start, request)
+            result = await self._query_models_definition(user, dict(request.query_params), 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)
@@ -97,7 +115,34 @@ class DataManagementRoutes(RouteBase):
         ) -> rm.DatasetResultModel:
             start = time.time()
             payload: dict = await request.json()
-            result = await self._query_models_definition(user, payload, params.model_dump())
-            self.log_activity_time("POST REQUEST for QUERY MODELS", start, request)
+            result = await self._query_models_definition(user, payload, params.model_dump(), headers=dict(request.headers))
+            self.logger.log_activity_time("POST REQUEST for QUERY MODELS", start)
+            return result
+
+        # Compiled models endpoints - TODO: remove duplication
+        compiled_models_path = project_metadata_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: str, params: QueryModelForGetCompiled, user=Depends(self.get_current_user)
+        ) -> rm.CompiledQueryModel:
+            start = time.time()
+            result = await self._get_compiled_model_definition(model_name, user, dict(request.query_params), 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: str, params: QueryModelForPostCompiled, user=Depends(self.get_current_user)
+        ) -> rm.CompiledQueryModel:
+            start = time.time()
+            payload: dict = await request.json()
+            result = await self._get_compiled_model_definition(model_name, user, payload, 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

@@ -1,10 +1,10 @@
 """
 Dataset routes for parameters and results
 """
-from typing import Callable, Any
-from pydantic import Field, BaseModel
+from typing import Callable, Coroutine, Any
+from pydantic import Field
 from fastapi import FastAPI, Depends, Request
-from fastapi.responses import JSONResponse, HTMLResponse
+from fastapi.responses import JSONResponse
 from fastapi.security import HTTPBearer
 
 from mcp.server.fastmcp import FastMCP, Context
@@ -12,14 +12,14 @@ from dataclasses import asdict
 from cachetools import TTLCache
 from textwrap import dedent
 
-import time
+import time, json
 
 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 .._schemas.auth_models import AbstractUser
 from .base import RouteBase
 
 
@@ -34,39 +34,57 @@ class DatasetRoutes(RouteBase):
         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)
         
+        # Setup max rows for AI
+        self.max_rows_for_ai = int(self.env_vars.get(c.SQRL_DATASETS_MAX_ROWS_FOR_AI, 100))
+        
     async def _get_dataset_results_helper(
-        self, dataset: str, user: BaseUser | None, selections: tuple[tuple[str, Any], ...]
+        self, dataset: str, user: AbstractUser, selections: tuple[tuple[str, Any], ...], configurables: tuple[tuple[str, str], ...]
     ) -> DatasetResult:
         """Helper to get dataset results"""
-        return await self.project.dataset(dataset, selections=dict(selections), user=user)
+        # 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(dataset, user=user, selections=dict(selections), configurables=cfg_filtered)
 
     async def _get_dataset_results_cachable(
-        self, dataset: str, user: BaseUser | None, selections: tuple[tuple[str, Any], ...]
+        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)
+        return await self.do_cachable_action(self.dataset_results_cache, self._get_dataset_results_helper, dataset, user, selections, configurables)
     
     async def _get_dataset_results_definition(
-        self, dataset_name: str, user: BaseUser | None, all_request_params: dict, params: dict
+        self, dataset_name: str, user: AbstractUser, all_request_params: dict, params: dict, headers: dict[str, str]
     ) -> rm.DatasetResultModel:
         """Get dataset results definition"""
-        self._validate_request_params(all_request_params, params)
+        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_verify_params", "x_orientation", "x_select", "x_limit", "x_offset"}
+        uncached_keys = {"x_verify_params", "x_orientation", "x_sql_query", "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()
+        user_has_elevated_privileges = u.user_has_elevated_privileges(user.access_level, self.project._elevated_access_level)
+        configurables = self.get_configurables_from_headers(headers) if user_has_elevated_privileges else tuple()
+        result = await get_dataset_function(dataset_name, user, selections, configurables)
+        
+        # Apply optional final SQL transformation before select/limit/offset
+        sql_query = params.get("x_sql_query")
+        if sql_query:
+            try:
+                transformed = u.run_sql_on_dataframes(sql_query, {"result": result.df.lazy()})
+            except Exception as e:
+                raise InvalidInputError(400, "invalid_sql_query", "Failed to run provided SQL on the dataset result") from e
+            
+            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)
+        
+        orientation_header = headers.get("x-orientation")
+        orientation = str(orientation_header).lower() if orientation_header is not None else params.get("x_orientation", "records")
         limit = params.get("x_limit", 1000)
         offset = params.get("x_offset", 0)
-        return rm.DatasetResultModel(**result.to_json(orientation, select, limit, offset)) 
+        return rm.DatasetResultModel(**result.to_json(orientation, 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
+        self, app: FastAPI, mcp: FastMCP, project_metadata_path: str, project_name: str, project_label: str, 
+        param_fields: dict, get_parameters_definition: Callable[..., Coroutine[Any, Any, rm.ParametersModel]]
     ) -> None:
         """Setup dataset routes"""
         
@@ -84,19 +102,19 @@ class DatasetRoutes(RouteBase):
                         f"\n  {all_params}"
                     )
         
-        async def get_dataset_parameters_updates(dataset_name: str, user: BaseUser | None, all_request_params: dict, params: dict):
+        async def get_dataset_parameters_updates(dataset_name: str, user: AbstractUser, all_request_params: dict, params: dict, headers: dict[str, str]):
             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
+                parameters_list, "dataset", dataset_name, scope, user, all_request_params, params, headers=headers
             )
             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)
+            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)
 
@@ -109,8 +127,10 @@ class DatasetRoutes(RouteBase):
             ) -> 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)
+                result = await get_dataset_parameters_updates(curr_dataset_name, user, dict(request.query_params), asdict(params), dict(request.headers))
+                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)
@@ -120,8 +140,10 @@ class DatasetRoutes(RouteBase):
                 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)
+                result = await get_dataset_parameters_updates(curr_dataset_name, user, payload, params.model_dump(), dict(request.headers))
+                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)
@@ -130,8 +152,12 @@ class DatasetRoutes(RouteBase):
             ) -> 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)
+                result = await self._get_dataset_results_definition(
+                    curr_dataset_name, user, dict(request.query_params), 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)
@@ -141,18 +167,23 @@ class DatasetRoutes(RouteBase):
                 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)
+                result = await self._get_dataset_results_definition(
+                    curr_dataset_name, user, payload, 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
     
         # Setup MCP tools
 
         @mcp.tool(
-            name=f"get_dataset_parameters_for_{project_name}_{project_version}",
+            name=f"get_dataset_parameters_from_{project_name}",
+            title=f"Get Dataset Parameters Updates (Project: {project_label})",
             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".
+            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 should 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()
@@ -162,28 +193,30 @@ class DatasetRoutes(RouteBase):
             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)
+        ) -> rm.ParametersModel:
+            headers = self.get_headers_from_tool_ctx(ctx)
+            user = self.get_user_from_tool_headers(headers)
             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)
+            return await get_dataset_parameters_updates(dataset_name, user, payload, payload, headers)
         
         @mcp.tool(
-            name=f"get_dataset_results_for_{project_name}_{project_version}",
+            name=f"get_dataset_results_from_{project_name}",
+            title=f"Get Dataset Results (Project: {project_label})",
             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.
+            - The "limit" argument controls the number of rows returned. The maximum allowed value is {self.max_rows_for_ai}. If the 'total_num_rows' field in the response is greater than {self.max_rows_for_ai}, let the user know that only {self.max_rows_for_ai} 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:
+            parameters: str = Field(description=dedent("""
+            A JSON object (as string) containing 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"
@@ -191,52 +224,42 @@ class DatasetRoutes(RouteBase):
             - 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)
+            - Complex objects are NOT supported
+            """).strip()),
+            sql_query: str | None = Field(None, description=dedent("""
+            A custom DuckDB SQL query to execute on the final dataset result. 
+            - Use table name 'result' to reference the dataset result.
+            - Use this to apply transformations to the dataset result if needed (such as filtering, sorting, or selecting columns).
+            - If not provided, the dataset result is returned as is.
+            """).strip()),
+            offset: int = Field(0, description="The number of rows to skip from first row. Applied after final SQL. Default is 0."),
+            limit: int = Field(self.max_rows_for_ai, description=f"The maximum number of rows to return. Applied after final SQL. Default is {self.max_rows_for_ai}. Maximum allowed value is {self.max_rows_for_ai}."),
+        ) -> rm.DatasetResultModel:
+            if limit > self.max_rows_for_ai:
+                raise ValueError(f"The maximum number of rows to return is {self.max_rows_for_ai}.")
+
+            headers = self.get_headers_from_tool_ctx(ctx)
+            user = self.get_user_from_tool_headers(headers)
             dataset_name = u.normalize_name(dataset)
-            params = {
-                **parameters,
-                "x_orientation": "rows",
+            
+            try:
+                params = json.loads(parameters)
+            except json.JSONDecodeError:
+                params = None # error handled below
+            
+            if not isinstance(params, dict):
+                raise InvalidInputError(400, "invalid_parameters", f"The 'parameters' argument must be a JSON object.")
+            
+            params.update({
+                "x_sql_query": sql_query,
                 "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"""
+            # Set default orientation as rows if not provided
+            if "x-orientation" not in headers:
+                headers["x-orientation"] = "rows"
             
-            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")
+            result = await self._get_dataset_results_definition(dataset_name, user, params, params, headers)
+            return result
         

squirrels/_api_routes/oauth2.py

@@ -6,9 +6,10 @@ from typing import Annotated, cast
 from .base import RouteBase
 from .._schemas.auth_models import (
     ClientRegistrationRequest, ClientUpdateRequest, ClientRegistrationResponse, ClientDetailsResponse, ClientUpdateResponse, 
-    TokenResponse, OAuthServerMetadata
+    TokenResponse, OAuthServerMetadata, AbstractUser
 )
 from .._exceptions import InvalidInputError
+from .. import _utils as u
 
 
 class OAuth2Routes(RouteBase):
@@ -50,16 +51,13 @@ class OAuth2Routes(RouteBase):
             status_code=200
         )
     
-    def setup_routes(self, app: FastAPI) -> None:
+    def setup_routes(self, app: FastAPI, squirrels_version_path: str) -> None:
         """Setup all OAuth2 routes"""
 
-        router_path = "/api/auth/oauth2"
+        auth_path = squirrels_version_path + "/auth"
+        router_path = "/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)
         
@@ -93,7 +91,7 @@ class OAuth2Routes(RouteBase):
         # Client Registration Endpoint
         client_management_path = '/client/{client_id}'
         
-        @router.post("/register", description="Register a new OAuth client", tags=["OAuth2"])
+        @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"""
             
@@ -148,7 +146,7 @@ class OAuth2Routes(RouteBase):
             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)
+            user: AbstractUser = Depends(self.get_current_user)
         ):
             """OAuth 2.1 authorization endpoint for initiating authorization code flow"""
             
@@ -158,9 +156,9 @@ class OAuth2Routes(RouteBase):
                     raise InvalidInputError(400, "unsupported_response_type", "Only 'code' response type is supported")
                 
                 # Check if user is authenticated
-                if user is None:
+                if user.access_level == "guest":
                     # User is not authenticated - serve login page
-                    return self.serve_login_page("/api/auth", request, client_id)
+                    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"
@@ -281,7 +279,7 @@ class OAuth2Routes(RouteBase):
             """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"
+            scheme = u.get_scheme(request.url.hostname)
             base_url = scheme + "://" + request.url.netloc
             
             return OAuthServerMetadata(
@@ -289,7 +287,7 @@ class OAuth2Routes(RouteBase):
                 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",
+                registration_endpoint=f"{base_url}{router_path}/client",
                 scopes_supported=["read"],
                 response_types_supported=["code"],
                 grant_types_supported=["authorization_code", "refresh_token"],