squirrels 0.1.1.post1__py3-none-any.whl → 0.2.0.dev0__py3-none-any.whl

squirrels/__init__.py

@@ -1,18 +1,12 @@
-from .parameter_options import SelectParameterOption, DateParameterOption, NumberParameterOption, NumRangeParameterOption
-from .parameters import Parameter, SingleSelectParameter, MultiSelectParameter, DateParameter, NumberParameter, NumRangeParameter, DataSourceParameter
-from .data_sources import SelectionDataSource, DateDataSource, NumberDataSource, NumRangeDataSource
-from .connection_set import ConnectionSet
+__version__ = '0.2.0'
 
+from typing import Union
+from sqlalchemy import Engine, Pool
+from pandas import DataFrame
 
-def get_credential(key: str):
-    """
-    Gets the username and password that was set through "$squirrels set-credential [key]"
-
-    Parameters:
-        key (str): The credential key
-    
-    Returns:
-        Credential: Object with attributes "username" and "password"
-    """
-    from ._credentials_manager import squirrels_config_io
-    return squirrels_config_io.get_credential(key)
+from .arguments.init_time_args import ConnectionsArgs, ParametersArgs
+from .arguments.run_time_args import ContextArgs, ModelDepsArgs, ModelArgs
+from .parameter_options import SelectParameterOption, DateParameterOption, DateRangeParameterOption, NumberParameterOption, NumberRangeParameterOption
+from .parameters import Parameter, SingleSelectParameter, MultiSelectParameter, DateParameter, DateRangeParameter, NumberParameter, NumberRangeParameter
+from .data_sources import SingleSelectDataSource, MultiSelectDataSource, DateDataSource, DateRangeDataSource, NumberDataSource, NumberRangeDataSource
+from .user_base import User, WrongPassword

squirrels/_api_server.py

@@ -1,134 +1,288 @@
-from typing import Dict, List, Tuple, Set
-from fastapi import FastAPI, Request, HTTPException
-from fastapi.datastructures import QueryParams
+from typing import Iterable, Optional, Mapping, Callable, Coroutine, TypeVar, Any
+from fastapi import Depends, FastAPI, Request, HTTPException, status
 from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.templating import Jinja2Templates
 from fastapi.staticfiles import StaticFiles
-from cachetools.func import ttl_cache
-import os, traceback
+from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
+from fastapi.middleware.cors import CORSMiddleware
+from cachetools import TTLCache
+import os, traceback, pandas as pd
 
-from squirrels import _constants as c, _utils
-from squirrels._version import major_version
-from squirrels._manifest import Manifest
-from squirrels.connection_set import ConnectionSet
-from squirrels._renderer import RendererIOWrapper, Renderer
+from . import _constants as c, _utils as u
+from ._version import sq_major_version
+from ._manifest import ManifestIO
+from ._authenticator import User, Authenticator
+from ._timer import timer, time
+from ._parameter_sets import ParameterSet
+from ._models import ModelsIO
 
 
 class ApiServer:
-    def __init__(self, manifest: Manifest, conn_set: ConnectionSet, no_cache: bool, debug: bool) -> None:
+    def __init__(self, no_cache: bool, debug: bool) -> None:
         """
         Constructor for ApiServer
 
         Parameters:
-            manifest (Manifest): Manifest object produced from squirrels.yaml
-            conn_set (ConnectionSet): Set of all connection pools defined in connections.py
             no_cache (bool): Whether to disable caching
             debug (bool): Set to True to show "hidden" parameters in the /parameters endpoint response
         """
-        self.manifest = manifest
-        self.conn_set = conn_set
         self.no_cache = no_cache
         self.debug = debug
+        self.dataset_configs = ManifestIO.obj.datasets
         
-        self.datasets = manifest.get_all_dataset_names()
-        self.renderers: Dict[str, Renderer] = {}
-        for dataset in self.datasets:
-            rendererIO = RendererIOWrapper(dataset, manifest, conn_set)
-            self.renderers[dataset] = rendererIO.renderer
-        
-    def _get_parameters_helper(self, dataset: str, query_params: Set[Tuple[str, str]]) -> Dict:
-        if len(query_params) > 1:
-            raise _utils.InvalidInputError("The /parameters endpoint takes at most 1 query parameter")
-        renderer = self.renderers[dataset]
-        parameters = renderer.apply_selections(dict(query_params), updates_only = True)
-        return parameters.to_json_dict(self.debug)
+        token_expiry_minutes = ManifestIO.obj.settings.get(c.AUTH_TOKEN_EXPIRE_SETTING, 30)
+        self.authenticator = Authenticator(token_expiry_minutes)
     
-    def _get_results_helper(self, dataset: str, query_params: Set[Tuple[str, str]]) -> Dict:
-        renderer = self.renderers[dataset]
-        _, _, _, _, df = renderer.load_results(dict(query_params))
-        return _utils.df_to_json(df)
-
-    def _apply_api_function(self, api_function):
-        try:
-            return api_function()
-        except _utils.InvalidInputError as e:
-            traceback.print_exc()
-            raise HTTPException(status_code=400, detail="Invalid User Input: "+str(e)) from e
-        except _utils.ConfigurationError as e:
-            traceback.print_exc()
-            raise HTTPException(status_code=500, detail="Squirrels Configuration Error: "+str(e)) from e
-        except Exception as e:
-            traceback.print_exc()
-            raise HTTPException(status_code=500, detail="Squirrels Framework Error: "+str(e)) from e
-    
-    def _apply_dataset_api_function(self, api_function, dataset: str, raw_query_params: QueryParams):
-        dataset = _utils.normalize_name(dataset)
-        query_params = set()
-        for key, val in raw_query_params.items():
-            query_params.add((_utils.normalize_name(key), val))
-        query_params = frozenset(query_params)
-        return self._apply_api_function(lambda: api_function(dataset, query_params))
-    
-    def run(self, uvicorn_args: List[str]) -> None:
+    def run(self, uvicorn_args: list[str]) -> None:
         """
         Runs the API server with uvicorn for CLI "squirrels run"
 
         Parameters:
-            uvicorn_args (List[str]): List of arguments to pass to uvicorn.run. Currently only supports "host" and "port"
+            uvicorn_args: List of arguments to pass to uvicorn.run. Currently only supports "host" and "port"
         """
+        start = time.time()
         app = FastAPI()
 
-        squirrels_version_path = f'/squirrels{major_version}'
-        config_base_path = _utils.normalize_name_for_api(self.manifest.get_base_path())
-        base_path = squirrels_version_path + config_base_path
+        app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
+
+        squirrels_version_path = f'/squirrels-v{sq_major_version}'
+        partial_base_path = f'/{ManifestIO.obj.project_variables.get_name()}/v{ManifestIO.obj.project_variables.get_major_version()}'
+        base_path = squirrels_version_path + u.normalize_name_for_api(partial_base_path)
 
-        static_dir = _utils.join_paths(os.path.dirname(__file__), 'package_data', 'static')
+        static_dir = u.join_paths(os.path.dirname(__file__), c.PACKAGE_DATA_FOLDER, c.STATIC_FOLDER)
         app.mount('/static', StaticFiles(directory=static_dir), name='static')
 
-        templates_dir = _utils.join_paths(os.path.dirname(__file__), 'package_data', 'templates')
+        templates_dir = u.join_paths(os.path.dirname(__file__), c.PACKAGE_DATA_FOLDER, c.TEMPLATES_FOLDER)
         templates = Jinja2Templates(directory=templates_dir)
 
+        # Exception handlers
+        @app.exception_handler(u.InvalidInputError)
+        async def invalid_input_error_handler(request: Request, exc: u.InvalidInputError):
+            traceback.print_exc()
+            return JSONResponse(status_code=status.HTTP_400_BAD_REQUEST, 
+                                content={"message": f"Invalid user input: {str(exc)}"})
+
+        @app.exception_handler(u.ConfigurationError)
+        async def configuration_error_handler(request: Request, exc: u.InvalidInputError):
+            traceback.print_exc()
+            return JSONResponse(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 
+                                content={"message": f"Squirrels configuration error: {str(exc)}"})
+        
+        @app.exception_handler(NotImplementedError)
+        async def not_implemented_error_handler(request: Request, exc: u.InvalidInputError):
+            traceback.print_exc()
+            return JSONResponse(status_code=status.HTTP_501_NOT_IMPLEMENTED, 
+                                content={"message": f"Not implemented error: {str(exc)}"})
+        
+        # Helpers
+        T = TypeVar('T')
+        
+        def get_versioning_request_header(headers: Mapping, header_key: str):
+            header_value = headers.get(header_key)
+            if header_value is None:
+                return None
+            
+            try:
+                result = int(header_value)
+            except ValueError:
+                raise u.InvalidInputError(f"Request header '{header_key}' must be an integer. Got '{header_value}'")
+            
+            if result < 0 or result > int(sq_major_version):
+                raise u.InvalidInputError(f"Request header '{header_key}' not in valid range. Got '{result}'")
+            
+            return result
+
+        REQUEST_VERSION_REQUEST_HEADER = "squirrels-request-version"
+        def get_request_version_header(headers: Mapping):
+            return get_versioning_request_header(headers, REQUEST_VERSION_REQUEST_HEADER)
+        
+        RESPONSE_VERSION_REQUEST_HEADER = "squirrels-response-version"
+        def process_based_on_response_version_header(headers: Mapping, processes: dict[str, Callable[[], T]]) -> T:
+            response_version = get_versioning_request_header(headers, RESPONSE_VERSION_REQUEST_HEADER)
+            if response_version is None or response_version >= 0:
+                return processes[0]()
+            else:
+                raise u.InvalidInputError(f'Invalid value for "{RESPONSE_VERSION_REQUEST_HEADER}" header: {response_version}')
+        
+        def can_user_access_dataset(user: Optional[User], dataset: str):
+            dataset_scope = self.dataset_configs[dataset].scope
+            return self.authenticator.can_user_access_scope(user, dataset_scope)
+
+        async def apply_dataset_api_function(
+            api_function: Callable[..., Coroutine[Any, Any, T]], user: Optional[User], dataset: str, headers: Mapping, params: Mapping
+        ) -> T:
+            dataset_normalized = u.normalize_name(dataset)
+            if not can_user_access_dataset(user, dataset_normalized):
+                raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED,
+                                    detail="Could not validate credentials",
+                                    headers={"WWW-Authenticate": "Bearer"})
+            
+            request_version = get_request_version_header(headers)
+            
+            # Changing selections into a cachable "frozenset" that will later be converted to dictionary
+            selections = set()
+            for key, val in params.items():
+                if not isinstance(val, str):
+                    val = tuple(val)
+                selections.add((u.normalize_name(key), val))
+            selections = frozenset(selections)
+            
+            return await api_function(user, dataset_normalized, selections, request_version)
+
+        # Login
+        token_path = base_path + '/token'
+
+        oauth2_scheme = OAuth2PasswordBearer(tokenUrl=token_path, auto_error=False)
+
+        @app.post(token_path)
+        async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
+            user: Optional[User] = self.authenticator.authenticate_user(form_data.username, form_data.password)
+            if not user:
+                raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, 
+                                    detail="Incorrect username or password",
+                                    headers={"WWW-Authenticate": "Bearer"})
+            access_token, expiry = self.authenticator.create_access_token(user)
+            return {
+                "access_token": access_token, 
+                "token_type": "bearer", 
+                "username": user.username,
+                "expiry_time": expiry
+            }
+        
+        async def get_current_user(token: str = Depends(oauth2_scheme)) -> Optional[User]:
+            user = self.authenticator.get_user_from_token(token)
+            return user
+
+        async def do_cachable_action(cache: TTLCache, action: Callable[..., Coroutine[Any, Any, T]], *args) -> T:
+            cache_key = tuple(args)
+            result = cache.get(cache_key)
+            if result is None:
+                result = await action(*args)
+                cache[cache_key] = result
+            return result
+
         # Parameters API
         parameters_path = base_path + '/{dataset}/parameters'
         
-        parameters_cache_size = self.manifest.get_setting(c.PARAMETERS_CACHE_SIZE_SETTING, 1024)
-        parameters_cache_ttl = self.manifest.get_setting(c.PARAMETERS_CACHE_TTL_SETTING, 24*60*60)
+        parameters_cache_size = ManifestIO.obj.settings.get(c.PARAMETERS_CACHE_SIZE_SETTING, 1024)
+        parameters_cache_ttl = ManifestIO.obj.settings.get(c.PARAMETERS_CACHE_TTL_SETTING, 0)
+    
+        async def get_parameters_helper(
+            user: Optional[User], dataset: str, selections: Iterable[tuple[str, str]], request_version: Optional[int]
+        ) -> ParameterSet:
+            if len(selections) > 1:
+                raise u.InvalidInputError(f"The /parameters endpoint takes at most 1 query parameter. Got {dict(selections)}")
+            dag = ModelsIO.GenerateDAG(dataset)
+            dag.apply_selections(user, dict(selections), request_version=request_version)
+            return dag.parameter_set
 
-        @ttl_cache(maxsize=parameters_cache_size, ttl=parameters_cache_ttl)
-        def get_parameters_cachable(*args):
-            return self._get_parameters_helper(*args)
+        params_cache = TTLCache(maxsize=parameters_cache_size, ttl=parameters_cache_ttl*60)
+
+        async def get_parameters_cachable(*args) -> T:
+            return await do_cachable_action(params_cache, get_parameters_helper, *args)
+        
+        async def get_parameters_definition(dataset: str, user: Optional[User], headers: Mapping, params: Mapping):
+            api_function = get_parameters_helper if self.no_cache else get_parameters_cachable
+            result = await apply_dataset_api_function(api_function, user, dataset, headers, params)
+            return process_based_on_response_version_header(headers, {
+                0: result.to_json_dict0
+            })
         
         @app.get(parameters_path, response_class=JSONResponse)
-        async def get_parameters(dataset: str, request: Request):
-            api_function = self._get_parameters_helper if self.no_cache else get_parameters_cachable
-            return self._apply_dataset_api_function(api_function, dataset, request.query_params)
+        async def get_parameters(dataset: str, request: Request, user: Optional[User] = Depends(get_current_user)):
+            start = time.time()
+            result = await get_parameters_definition(dataset, user, request.headers, request.query_params)
+            timer.add_activity_time("GET REQUEST total time for PARAMETERS", start)
+            return result
+
+        @app.post(parameters_path, response_class=JSONResponse)
+        async def get_parameters_with_post(dataset: str, request: Request, user: Optional[User] = Depends(get_current_user)):
+            start = time.time()
+            request_body = await request.json()
+            result = await get_parameters_definition(dataset, user, request.headers, request_body)
+            timer.add_activity_time("POST REQUEST total time for PARAMETERS", start)
+            return result
 
         # Results API
         results_path = base_path + '/{dataset}'
 
-        results_cache_size = self.manifest.get_setting(c.RESULTS_CACHE_SIZE_SETTING, 128)
-        results_cache_ttl = self.manifest.get_setting(c.RESULTS_CACHE_TTL_SETTING, 60*60)
+        results_cache_size = ManifestIO.obj.settings.get(c.RESULTS_CACHE_SIZE_SETTING, 128)
+        results_cache_ttl = ManifestIO.obj.settings.get(c.RESULTS_CACHE_TTL_SETTING, 0)
+    
+        async def get_results_helper(
+            user: Optional[User], dataset: str, selections: Iterable[tuple[str, str]], request_version: Optional[int]
+        ) -> pd.DataFrame:
+            dag = ModelsIO.GenerateDAG(dataset)
+            await dag.execute(ModelsIO.context_func, user, dict(selections), request_version=request_version)
+            return dag.target_model.result
+
+        results_cache = TTLCache(maxsize=results_cache_size, ttl=results_cache_ttl*60)
 
-        @ttl_cache(maxsize=results_cache_size, ttl=results_cache_ttl)
-        def get_results_cachable(*args):
-            return self._get_results_helper(*args)
+        async def get_results_cachable(*args):
+            return await do_cachable_action(results_cache, get_results_helper, *args)
+        
+        async def get_results_definition(dataset: str, user: Optional[User], headers: Mapping, params: Mapping):
+            api_function = get_results_helper if self.no_cache else get_results_cachable
+            result = await apply_dataset_api_function(api_function, user, dataset, headers, params)
+            return process_based_on_response_version_header(headers, {
+                0: lambda: u.df_to_json0(result)
+            })
         
         @app.get(results_path, response_class=JSONResponse)
-        async def get_results(dataset: str, request: Request):
-            api_function = self._get_results_helper if self.no_cache else get_results_cachable
-            return self._apply_dataset_api_function(api_function, dataset, request.query_params)
+        async def get_results(dataset: str, request: Request, user: Optional[User] = Depends(get_current_user)):
+            start = time.time()
+            result = await get_results_definition(dataset, user, request.headers, request.query_params)
+            timer.add_activity_time("GET REQUEST total time for DATASET", start)
+            return result
+        
+        @app.post(results_path, response_class=JSONResponse)
+        async def get_results_with_post(dataset: str, request: Request, user: Optional[User] = Depends(get_current_user)):
+            start = time.time()
+            request_body = await request.json()
+            result = await get_results_definition(dataset, user, request.headers, request_body)
+            timer.add_activity_time("POST REQUEST total time for DATASET", start)
+            return result
         
         # Catalog API
+        def get_catalog0(user: Optional[User]):
+            datasets_info = []
+            for dataset_name, dataset_config in self.dataset_configs.items():
+                if can_user_access_dataset(user, dataset_name):
+                    dataset_normalized = u.normalize_name_for_api(dataset_name)
+                    datasets_info.append({
+                        'name': dataset_name,
+                        'label': dataset_config.label,
+                        'parameters_path': parameters_path.format(dataset=dataset_normalized),
+                        'result_path': results_path.format(dataset=dataset_normalized),
+                        'first_minor_version': 0
+                    })
+            
+            return {
+                'products': [{
+                    'name': ManifestIO.obj.project_variables.get_name(),
+                    'label': ManifestIO.obj.project_variables.get_label(),
+                    'versions': [{
+                        'major_version': ManifestIO.obj.project_variables.get_major_version(),
+                        'latest_minor_version': ManifestIO.obj.project_variables.get_minor_version(),
+                        'datasets': datasets_info
+                    }]
+                }]
+            }
+        
         @app.get(squirrels_version_path, response_class=JSONResponse)
-        async def get_catalog():
-            api_function = lambda: self.manifest.get_catalog(parameters_path, results_path)
-            return self._apply_api_function(api_function)
+        async def get_catalog(request: Request, user: Optional[User] = Depends(get_current_user)):
+            return process_based_on_response_version_header(request.headers, {
+                0: lambda: get_catalog0(user)
+            })
         
         # Squirrels UI
         @app.get('/', response_class=HTMLResponse)
         async def get_ui(request: Request):
-            return templates.TemplateResponse('index.html', {'request': request, 'catalog_path': squirrels_version_path})
+            return templates.TemplateResponse('index.html', {
+                'request': request, 'catalog_path': squirrels_version_path, 'token_path': token_path
+            })
         
         # Run API server
         import uvicorn
+        timer.add_activity_time("creating app for api server", start)
         uvicorn.run(app, host=uvicorn_args.host, port=uvicorn_args.port)

squirrels/_authenticator.py

@@ -0,0 +1,84 @@
+from typing import Optional
+from datetime import datetime, timedelta, timezone
+from jose import JWTError, jwt
+import secrets
+
+from . import _utils as u, _constants as c
+from ._py_module import PyModule
+from .user_base import User, WrongPassword
+from ._environcfg import EnvironConfigIO
+from ._manifest import DatasetScope
+
+
+class Authenticator:
+    
+    @classmethod
+    def get_auth_helper(cls, default_auth_helper = None):
+        auth_module_path = u.join_paths(c.PYCONFIG_FOLDER, c.AUTH_FILE)
+        return PyModule(auth_module_path, default_class=default_auth_helper)
+
+    def __init__(self, token_expiry_minutes: int, auth_helper = None) -> None:
+        self.token_expiry_minutes = token_expiry_minutes
+        self.auth_helper = self.get_auth_helper(auth_helper)
+        self.secret_key = self._get_secret_key()
+        self.algorithm = "HS256"
+
+    def _get_secret_key(self):
+        secret_key = EnvironConfigIO.obj.get_secret(c.JWT_SECRET_KEY, default_factory=lambda: secrets.token_hex(32))
+        return secret_key
+
+    def authenticate_user(self, username: str, password: str) -> Optional[User]:
+        if self.auth_helper:
+            user_cls = self.auth_helper.get_func_or_class("User", default_attr=User)
+            get_user = self.auth_helper.get_func_or_class(c.GET_USER_FUNC)
+            try:
+                real_user = get_user(username, password)
+            except Exception as e:
+                raise u.FileExecutionError(f'Failed to run "{c.GET_USER_FUNC}" in {c.AUTH_FILE}', e)
+        else:
+            user_cls = User
+            real_user = None
+        
+        if isinstance(real_user, User):
+            return real_user
+        
+        if not isinstance(real_user, WrongPassword):
+            fake_users = EnvironConfigIO.obj.get_users()
+            if username in fake_users and secrets.compare_digest(fake_users[username][c.USER_PWD_KEY], password):
+                is_internal = fake_users[username].get("is_internal", False)
+                user = user_cls(username, is_internal=is_internal)
+                try:
+                    return user.with_attributes(fake_users[username])
+                except Exception as e:
+                    raise u.FileExecutionError(f'Failed to create user from User model in {c.AUTH_FILE}', e)
+        
+        return None
+    
+    def create_access_token(self, user: User) -> str:
+        expire = datetime.now(timezone.utc) + timedelta(minutes=self.token_expiry_minutes)
+        to_encode = {**vars(user), "exp": expire}
+        encoded_jwt = jwt.encode(to_encode, self.secret_key, algorithm=self.algorithm)
+        return encoded_jwt, expire
+    
+    def get_user_from_token(self, token: Optional[str]) -> Optional[User]:
+        if token is not None:
+            try:
+                payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm])
+                payload.pop("exp")
+                if self.auth_helper is not None:
+                    user_cls: User = self.auth_helper.get_func_or_class("User", default_attr=User)
+                    return user_cls._FromDict(payload)
+                else:
+                    return User._FromDict(payload)
+            except JWTError:
+                return None
+
+    def can_user_access_scope(self, user: Optional[User], scope: DatasetScope) -> bool:
+        if user is None:
+            user_level = DatasetScope.PUBLIC
+        elif not user.is_internal:
+            user_level = DatasetScope.PROTECTED
+        else:
+            user_level = DatasetScope.PRIVATE
+        
+        return user_level.value >= scope.value