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

squirrels/_parameter_configs.py

@@ -9,13 +9,13 @@ from fastapi import Query
 from pydantic.fields import Field
 import polars as pl, re
 
-from . import parameter_options as _po, parameters as p, data_sources as d, _utils as u, _constants as c
+from . import _data_sources as d, _parameter_options as po, _parameters as p, _utils as u, _constants as c
 from ._exceptions import InvalidInputError
-from ._auth import BaseUser
+from ._schemas.auth_models import AbstractUser
 from ._connection_set import ConnectionSet
 from ._seeds import Seeds
 
-ParamOptionType = TypeVar("ParamOptionType", bound=_po.ParameterOption)
+ParamOptionType = TypeVar("ParamOptionType", bound=po.ParameterOption)
 
 
 @dataclass
@@ -56,12 +56,20 @@ class ParameterConfigBase(metaclass=ABCMeta):
     user_attribute: str | None = field(default=None, kw_only=True)
     parent_name: str | None = field(default=None, kw_only=True)
 
-    def _get_user_group(self, user: BaseUser | None) -> Any:
-        if self.user_attribute is not None:
-            if user is None:
-                raise u.ConfigurationError(f"Public datasets (accessible without authentication) cannot use parameter " +
-                                           f"'{self.name}' because 'user_attribute' is defined on this parameter.")
-            return getattr(user, self.user_attribute)
+    def _get_user_group(self, user: AbstractUser) -> Any:
+        if self.user_attribute is None:
+            return None
+        
+        final_object = user
+        attribute = self.user_attribute
+        try:
+            if "." in attribute:
+                parts = attribute.split(".", 1)
+                final_object = getattr(final_object, parts[0])
+                attribute = parts[1]
+            return getattr(final_object, attribute)
+        except AttributeError:
+            raise u.ConfigurationError(f"User attribute '{self.user_attribute}' is not valid")
         
     def copy(self):
         """
@@ -108,16 +116,16 @@ class ParameterConfig(Generic[ParamOptionType], ParameterConfigBase):
         pass
     
     def _invalid_input_error(self, selection: str, more_details: str = '') -> InvalidInputError:
-        return InvalidInputError(200, f'Selected value "{selection}" is not valid for parameter "{self.name}". ' + more_details)
+        return InvalidInputError(400, "invalid_parameter_selection", f'Selected value "{selection}" is not valid for parameter "{self.name}". ' + more_details)
     
     @abstractmethod
     def with_selection(
-        self, selection: str | None, user: BaseUser | None, parent_param: p._SelectionParameter | None
+        self, selection: str | None, user: AbstractUser, parent_param: p._SelectionParameter | None
     ) -> p.Parameter:
         pass
     
     def _get_options_iterator(
-        self, all_options: Sequence[ParamOptionType], user: BaseUser | None, parent_param: p._SelectionParameter | None
+        self, all_options: Sequence[ParamOptionType], user: AbstractUser, parent_param: p._SelectionParameter | None
     ) -> Iterator[ParamOptionType]:
         user_group = self._get_user_group(user)
         selected_parent_option_ids = frozenset(parent_param._get_selected_ids_as_list()) if parent_param else None
@@ -129,7 +137,7 @@ class ParameterConfig(Generic[ParamOptionType], ParameterConfigBase):
 
 
 @dataclass
-class SelectionParameterConfig(ParameterConfig[_po.SelectParameterOption]):
+class SelectionParameterConfig(ParameterConfig[po.SelectParameterOption]):
     """
     Abstract class for select parameter classes (single-select, multi-select, etc)
     """
@@ -138,7 +146,7 @@ class SelectionParameterConfig(ParameterConfig[_po.SelectParameterOption]):
 
     @abstractmethod
     def __init__(
-        self, name: str, label: str, all_options: Sequence[_po.SelectParameterOption | dict], *,
+        self, name: str, label: str, all_options: Sequence[po.SelectParameterOption | dict], *,
         description: str = "", user_attribute: str | None = None, parent_name: str | None = None
     ) -> None:
         super().__init__(name, label, all_options, description=description, user_attribute=user_attribute, parent_name=parent_name)
@@ -147,16 +155,16 @@ class SelectionParameterConfig(ParameterConfig[_po.SelectParameterOption]):
 
     @staticmethod
     def ParameterOption(*args, **kwargs):
-        return _po.SelectParameterOption(*args, **kwargs)
+        return po.SelectParameterOption(*args, **kwargs)
     
     def _add_child_mutate(self, child: ParameterConfigBase):
         self.children[child.name] = child
         self.trigger_refresh = True
     
-    def _get_options(self, user: BaseUser | None, parent_param: p._SelectionParameter | None) -> Sequence[_po.SelectParameterOption]:
+    def _get_options(self, user: AbstractUser, parent_param: p._SelectionParameter | None) -> Sequence[po.SelectParameterOption]:
         return tuple(self._get_options_iterator(self.all_options, user, parent_param))
     
-    def _get_default_ids_iterator(self, options: Sequence[_po.SelectParameterOption]) -> Iterator[str]:
+    def _get_default_ids_iterator(self, options: Sequence[po.SelectParameterOption]) -> Iterator[str]:
         return (x._identifier for x in options if x._is_default)
     
     def copy(self) -> Self:
@@ -175,7 +183,7 @@ class SingleSelectParameterConfig(SelectionParameterConfig):
     """
     
     def __init__(
-        self, name: str, label: str, all_options: Sequence[_po.SelectParameterOption | dict], *, description: str = "", 
+        self, name: str, label: str, all_options: Sequence[po.SelectParameterOption | dict], *, description: str = "", 
         user_attribute: str | None = None, parent_name: str | None = None
     ) -> None:
         super().__init__(name, label, all_options, description=description, user_attribute=user_attribute, parent_name=parent_name)
@@ -189,7 +197,7 @@ class SingleSelectParameterConfig(SelectionParameterConfig):
         return d.SelectDataSource(*args, **kwargs)
     
     def with_selection(
-        self, selection: str | None, user: BaseUser | None, parent_param: p._SelectionParameter | None
+        self, selection: str | None, user: AbstractUser, parent_param: p._SelectionParameter | None
     ) -> p.SingleSelectParameter:
         options = self._get_options(user, parent_param)
         if selection is None:
@@ -217,7 +225,7 @@ class MultiSelectParameterConfig(SelectionParameterConfig):
     none_is_all: bool = field(default=True, kw_only=True)
 
     def __init__(
-        self, name: str, label: str, all_options: Sequence[_po.SelectParameterOption | dict], *, description: str = "", 
+        self, name: str, label: str, all_options: Sequence[po.SelectParameterOption | dict], *, description: str = "", 
         show_select_all: bool = True, order_matters: bool = False, none_is_all: bool = True,
         user_attribute: str | None = None, parent_name: str | None = None
     ) -> None:
@@ -237,7 +245,7 @@ class MultiSelectParameterConfig(SelectionParameterConfig):
         return d.SelectDataSource(*args, **kwargs)
     
     def with_selection(
-        self, selection: str | None, user: BaseUser | None, parent_param: p._SelectionParameter | None
+        self, selection: str | None, user: AbstractUser, parent_param: p._SelectionParameter | None
     ) -> p.MultiSelectParameter:
         options = self._get_options(user, parent_param)
         if selection is None:
@@ -268,14 +276,14 @@ class _DateTypeParameterConfig(ParameterConfig[ParamOptionType]):
 
 
 @dataclass
-class DateParameterConfig(_DateTypeParameterConfig[_po.DateParameterOption]):
+class DateParameterConfig(_DateTypeParameterConfig[po.DateParameterOption]):
     """
     Class to define configurations for date parameter widgets.
     """
-    _all_options: Sequence[_po.DateParameterOption] = field(repr=False)
+    _all_options: Sequence[po.DateParameterOption] = field(repr=False)
     
     def __init__(
-        self, name: str, label: str, all_options: Sequence[_po.DateParameterOption | dict], *, 
+        self, name: str, label: str, all_options: Sequence[po.DateParameterOption | dict], *, 
         description: str = "", user_attribute: str | None = None, parent_name: str | None = None
     ) -> None:
         super().__init__(name, label, all_options, description=description, user_attribute=user_attribute, parent_name=parent_name)
@@ -286,16 +294,16 @@ class DateParameterConfig(_DateTypeParameterConfig[_po.DateParameterOption]):
 
     @staticmethod
     def ParameterOption(*args, **kwargs):
-        return _po.DateParameterOption(*args, **kwargs)
+        return po.DateParameterOption(*args, **kwargs)
     
     @staticmethod
     def DataSource(*args, **kwargs):
         return d.DateDataSource(*args, **kwargs)
     
     def with_selection(
-        self, selection: str | None, user: BaseUser | None, parent_param: p._SelectionParameter | None
+        self, selection: str | None, user: AbstractUser, parent_param: p._SelectionParameter | None
     ) -> p.DateParameter:
-        curr_option: _po.DateParameterOption | None = next(self._get_options_iterator(self.all_options, user, parent_param), None)
+        curr_option: po.DateParameterOption | None = next(self._get_options_iterator(self.all_options, user, parent_param), None)
         selected_date = curr_option._default_date if selection is None and curr_option is not None else selection
         return p.DateParameter(self, curr_option, selected_date)
     
@@ -307,14 +315,14 @@ class DateParameterConfig(_DateTypeParameterConfig[_po.DateParameterOption]):
 
 
 @dataclass
-class DateRangeParameterConfig(_DateTypeParameterConfig[_po.DateRangeParameterOption]):
+class DateRangeParameterConfig(_DateTypeParameterConfig[po.DateRangeParameterOption]):
     """
     Class to define configurations for date range parameter widgets.
     """
-    _all_options: Sequence[_po.DateRangeParameterOption] = field(repr=False)
+    _all_options: Sequence[po.DateRangeParameterOption] = field(repr=False)
     
     def __init__(
-        self, name: str, label: str, all_options: Sequence[_po.DateRangeParameterOption | dict], *,
+        self, name: str, label: str, all_options: Sequence[po.DateRangeParameterOption | dict], *,
         description: str = "", user_attribute: str | None = None, parent_name: str | None = None
     ) -> None:
         super().__init__(name, label, all_options, description=description, user_attribute=user_attribute, parent_name=parent_name)
@@ -325,16 +333,16 @@ class DateRangeParameterConfig(_DateTypeParameterConfig[_po.DateRangeParameterOp
 
     @staticmethod
     def ParameterOption(*args, **kwargs):
-        return _po.DateRangeParameterOption(*args, **kwargs)
+        return po.DateRangeParameterOption(*args, **kwargs)
     
     @staticmethod
     def DataSource(*args, **kwargs):
         return d.DateRangeDataSource(*args, **kwargs)
     
     def with_selection(
-        self, selection: str | None, user: BaseUser | None, parent_param: p._SelectionParameter | None
+        self, selection: str | None, user: AbstractUser, parent_param: p._SelectionParameter | None
     ) -> p.DateRangeParameter:
-        curr_option: _po.DateRangeParameterOption | None = next(self._get_options_iterator(self.all_options, user, parent_param), None)
+        curr_option: po.DateRangeParameterOption | None = next(self._get_options_iterator(self.all_options, user, parent_param), None)
         if selection is None:
             if curr_option is not None:
                 selected_start_date = curr_option._default_start_date
@@ -370,14 +378,14 @@ class _NumericParameterConfig(ParameterConfig[ParamOptionType]):
 
 
 @dataclass
-class NumberParameterConfig(_NumericParameterConfig[_po.NumberParameterOption]):
+class NumberParameterConfig(_NumericParameterConfig[po.NumberParameterOption]):
     """
     Class to define configurations for number parameter widgets.
     """
-    _all_options: Sequence[_po.NumberParameterOption] = field(repr=False)
+    _all_options: Sequence[po.NumberParameterOption] = field(repr=False)
     
     def __init__(
-        self, name: str, label: str, all_options: Sequence[_po.NumberParameterOption | dict], *,
+        self, name: str, label: str, all_options: Sequence[po.NumberParameterOption | dict], *,
         description: str = "", user_attribute: str | None = None, parent_name: str | None = None
     ) -> None:
         super().__init__(name, label, all_options, description=description, user_attribute=user_attribute, parent_name=parent_name)
@@ -388,16 +396,16 @@ class NumberParameterConfig(_NumericParameterConfig[_po.NumberParameterOption]):
     
     @staticmethod
     def ParameterOption(*args, **kwargs):
-        return _po.NumberParameterOption(*args, **kwargs)
+        return po.NumberParameterOption(*args, **kwargs)
     
     @staticmethod
     def DataSource(*args, **kwargs):
         return d.NumberDataSource(*args, **kwargs)
     
     def with_selection(
-        self, selection: str | None, user: BaseUser | None, parent_param: p._SelectionParameter | None
+        self, selection: str | None, user: AbstractUser, parent_param: p._SelectionParameter | None
     ) -> p.NumberParameter:
-        curr_option: _po.NumberParameterOption | None = next(self._get_options_iterator(self.all_options, user, parent_param), None)
+        curr_option: po.NumberParameterOption | None = next(self._get_options_iterator(self.all_options, user, parent_param), None)
         selected_value = curr_option._default_value if selection is None and curr_option is not None else selection
         return p.NumberParameter(self, curr_option, selected_value)
     
@@ -409,14 +417,14 @@ class NumberParameterConfig(_NumericParameterConfig[_po.NumberParameterOption]):
 
 
 @dataclass
-class NumberRangeParameterConfig(_NumericParameterConfig[_po.NumberRangeParameterOption]):
+class NumberRangeParameterConfig(_NumericParameterConfig[po.NumberRangeParameterOption]):
     """
     Class to define configurations for number range parameter widgets.
     """
-    _all_options: Sequence[_po.NumberRangeParameterOption] = field(repr=False)
+    _all_options: Sequence[po.NumberRangeParameterOption] = field(repr=False)
     
     def __init__(
-        self, name: str, label: str, all_options: Sequence[_po.NumberRangeParameterOption | dict], *, 
+        self, name: str, label: str, all_options: Sequence[po.NumberRangeParameterOption | dict], *, 
         description: str = "", user_attribute: str | None = None, parent_name: str | None = None
     ) -> None:
         super().__init__(name, label, all_options, description=description, user_attribute=user_attribute, parent_name=parent_name)
@@ -427,16 +435,16 @@ class NumberRangeParameterConfig(_NumericParameterConfig[_po.NumberRangeParamete
 
     @staticmethod
     def ParameterOption(*args, **kwargs):
-        return _po.NumberRangeParameterOption(*args, **kwargs)
+        return po.NumberRangeParameterOption(*args, **kwargs)
     
     @staticmethod
     def DataSource(*args, **kwargs):
         return d.NumberRangeDataSource(*args, **kwargs)
     
     def with_selection(
-        self, selection: str | None, user: BaseUser | None, parent_param: p._SelectionParameter | None
+        self, selection: str | None, user: AbstractUser, parent_param: p._SelectionParameter | None
     ) -> p.NumberRangeParameter:
-        curr_option: _po.NumberRangeParameterOption | None = next(self._get_options_iterator(self.all_options, user, parent_param), None)
+        curr_option: po.NumberRangeParameterOption | None = next(self._get_options_iterator(self.all_options, user, parent_param), None)
         if selection is None:
             if curr_option is not None:
                 selected_lower_value = curr_option._default_lower_value
@@ -458,15 +466,15 @@ class NumberRangeParameterConfig(_NumericParameterConfig[_po.NumberRangeParamete
 
 
 @dataclass
-class TextParameterConfig(ParameterConfig[_po.TextParameterOption]):
+class TextParameterConfig(ParameterConfig[po.TextParameterOption]):
     """
     Class to define configurations for text parameter widgets.
     """
-    _all_options: Sequence[_po.TextParameterOption] = field(repr=False)
+    _all_options: Sequence[po.TextParameterOption] = field(repr=False)
     input_type: str
     
     def __init__(
-        self, name: str, label: str, all_options: Sequence[_po.TextParameterOption | dict], *, description: str = "", 
+        self, name: str, label: str, all_options: Sequence[po.TextParameterOption | dict], *, description: str = "", 
         input_type: str = "text", user_attribute: str | None = None, parent_name: str | None = None
     ) -> None:
         super().__init__(name, label, all_options, description=description, user_attribute=user_attribute, parent_name=parent_name)
@@ -517,16 +525,16 @@ class TextParameterConfig(ParameterConfig[_po.TextParameterOption]):
 
     @staticmethod
     def ParameterOption(*args, **kwargs):
-        return _po.TextParameterOption(*args, **kwargs)
+        return po.TextParameterOption(*args, **kwargs)
     
     @staticmethod
     def DataSource(*args, **kwargs):
         return d.TextDataSource(*args, **kwargs)
     
     def with_selection(
-        self, selection: str | None, user: BaseUser | None, parent_param: p._SelectionParameter | None
+        self, selection: str | None, user: AbstractUser, parent_param: p._SelectionParameter | None
     ) -> p.TextParameter:
-        curr_option: _po.TextParameterOption | None = next(self._get_options_iterator(self.all_options, user, parent_param), None)
+        curr_option: po.TextParameterOption | None = next(self._get_options_iterator(self.all_options, user, parent_param), None)
         entered_text = curr_option._default_text if selection is None and curr_option is not None else selection
         return p.TextParameter(self, curr_option, entered_text)
     
@@ -550,6 +558,8 @@ class DataSourceParameterConfig(Generic[ParamConfigType], ParameterConfigBase):
         super().__init__(name, label, description=description, user_attribute=user_attribute, parent_name=parent_name)
         self.parameter_type = parameter_type
         if isinstance(data_source, dict):
+            if "source" in data_source:
+                data_source["source"] = d.SourceEnum(data_source["source"])
             data_source = parameter_type.DataSource(**data_source)
         self.data_source = data_source
         self.extra_args = extra_args
@@ -557,15 +567,26 @@ class DataSourceParameterConfig(Generic[ParamConfigType], ParameterConfigBase):
     def convert(self, df: pl.DataFrame) -> ParamConfigType:
         return self.data_source._convert(self, df)
     
-    def get_dataframe(self, default_conn_name: str, conn_set: ConnectionSet, seeds: Seeds) -> pl.DataFrame:
+    def get_dataframe(self, default_conn_name: str, conn_set: ConnectionSet, seeds: Seeds, datalake_db_path: str = "") -> pl.DataFrame:
         datasource = self.data_source
         query = datasource._get_query()
-        if datasource._is_from_seeds:
+        if datasource._source == d.SourceEnum.SEEDS:
             df = seeds.run_query(query)
-        else:
+        elif datasource._source == d.SourceEnum.VDL:
+            vdl_conn = u.create_duckdb_connection(datalake_db_path)
+            try:
+                # Query the VDL database
+                df = vdl_conn.sql(query).pl()
+            except Exception as e:
+                raise u.ConfigurationError(f'Error executing query for datasource parameter "{self.name}" from VDL') from e
+            finally:
+                vdl_conn.close()
+        else:  # source == "connection"
+            conn_name = None
             try:
                 conn_name = datasource._get_connection_name(default_conn_name)
                 df = conn_set.run_sql_query_from_conn_name(query, conn_name)
             except RuntimeError as e:
-                raise u.ConfigurationError(f'Error executing query for datasource parameter "{self.name}"') from e
+                ending = f' "{conn_name}"' if conn_name is not None else ""
+                raise u.ConfigurationError(f'Error executing query for datasource parameter "{self.name}" from connection{ending}') from e
         return df