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

squirrels/_data_sources.py

@@ -1,10 +1,16 @@
 from __future__ import annotations
 from dataclasses import dataclass
+from enum import Enum
 import polars as pl, typing as t, abc
 
 from . import _parameter_configs as pc, _parameter_options as po
 from ._exceptions import ConfigurationError
 
+class SourceEnum(Enum):
+    CONNECTION = "connection"
+    SEEDS = "seeds"
+    VDL = "vdl"
+
 
 @dataclass
 class DataSource(metaclass=abc.ABCMeta):
@@ -13,19 +19,19 @@ class DataSource(metaclass=abc.ABCMeta):
     """
     _table_or_query: str
     _id_col: str | None
-    _is_from_seeds: bool
+    _source: SourceEnum
     _user_group_col: str | None
     _parent_id_col: str | None
     _connection: str | None
 
     @abc.abstractmethod
     def __init__(
-        self, table_or_query: str, *, id_col: str | None = None, from_seeds: bool = False, user_group_col: str | None = None, 
-        parent_id_col: str | None = None, connection: str | None = None, **kwargs
+        self, table_or_query: str, *, id_col: str | None = None, source: SourceEnum = SourceEnum.CONNECTION, 
+        user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None, **kwargs
     ) -> None:
         self._table_or_query = table_or_query
         self._id_col = id_col
-        self._is_from_seeds = from_seeds
+        self._source = source
         self._user_group_col = user_group_col
         self._parent_id_col = parent_id_col
         self._connection = connection
@@ -100,12 +106,12 @@ class _SelectionDataSource(DataSource):
     @abc.abstractmethod
     def __init__(
         self, table_or_query: str, id_col: str, options_col: str, *, order_by_col: str | None = None, 
-        is_default_col: str | None = None, custom_cols: dict[str, str] = {}, from_seeds: bool = False, 
+        is_default_col: str | None = None, custom_cols: dict[str, str] = {}, source: SourceEnum = SourceEnum.CONNECTION, 
         user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None, 
         **kwargs
     ) -> None:
         super().__init__(
-            table_or_query, id_col=id_col, from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, 
+            table_or_query, id_col=id_col, source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, 
             connection=connection
         )
         self._options_col = options_col
@@ -151,7 +157,7 @@ class SelectDataSource(_SelectionDataSource):
 
     def __init__(
             self, table_or_query: str, id_col: str, options_col: str, *, order_by_col: str | None = None, 
-            is_default_col: str | None = None, custom_cols: dict[str, str] = {}, from_seeds: bool = False, 
+            is_default_col: str | None = None, custom_cols: dict[str, str] = {}, source: SourceEnum = SourceEnum.CONNECTION, 
             user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None, 
             **kwargs
         ) -> None:
@@ -165,19 +171,19 @@ class SelectDataSource(_SelectionDataSource):
             order_by_col: The column name to order the options by. Orders by the id_col instead if this is None
             is_default_col: The column name that indicates which options are the default
             custom_cols: Dictionary of attribute to column name for custom fields for the SelectParameterOption
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that must be selected for this option to be valid
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
             table_or_query, id_col, options_col, order_by_col=order_by_col, is_default_col=is_default_col, custom_cols=custom_cols,
-            from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, connection=connection
+            source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, connection=connection
         )
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.SelectionParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a SingleSelectParameterConfig or MultiSelectParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a SingleSelectParameterConfig or MultiSelectParameterConfig
 
         Arguments:
             ds_param: The parameter to convert
@@ -212,7 +218,7 @@ class DateDataSource(DataSource):
     def __init__(
         self, table_or_query: str, default_date_col: str, *, min_date_col: str | None = None, 
         max_date_col: str | None = None, date_format: str = '%Y-%m-%d', id_col: str | None = None, 
-        from_seeds: bool = False, user_group_col: str | None = None, parent_id_col: str | None = None, 
+        source: SourceEnum = SourceEnum.CONNECTION, user_group_col: str | None = None, parent_id_col: str | None = None, 
         connection: str | None = None, **kwargs
     ) -> None:
         """
@@ -223,13 +229,13 @@ class DateDataSource(DataSource):
             default_date_col: The column name of the default date
             date_format: The format of the default date(s). Defaults to '%Y-%m-%d'
             id_col: The column name of the id
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that the default date belongs to
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
-            table_or_query, id_col=id_col, from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, 
+            table_or_query, id_col=id_col, source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, 
             connection=connection
         )
         self._default_date_col = default_date_col
@@ -239,7 +245,7 @@ class DateDataSource(DataSource):
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.DateParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a DateParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a DateParameterConfig
 
         Arguments:
             ds_param: The parameter to convert
@@ -281,7 +287,7 @@ class DateRangeDataSource(DataSource):
 
     def __init__(
         self, table_or_query: str, default_start_date_col: str, default_end_date_col: str, *, date_format: str = '%Y-%m-%d',
-        min_date_col: str | None = None, max_date_col: str | None = None, id_col: str | None = None, from_seeds: bool = False, 
+        min_date_col: str | None = None, max_date_col: str | None = None, id_col: str | None = None, source: SourceEnum = SourceEnum.CONNECTION, 
         user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None, **kwargs
     ) -> None:
         """
@@ -293,13 +299,13 @@ class DateRangeDataSource(DataSource):
             default_end_date_col: The column name of the default end date
             date_format: The format of the default date(s). Defaults to '%Y-%m-%d'
             id_col: The column name of the id
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that the default date belongs to
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
-            table_or_query, id_col=id_col, from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, 
+            table_or_query, id_col=id_col, source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, 
             connection=connection
         )
         self._default_start_date_col = default_start_date_col
@@ -310,7 +316,7 @@ class DateRangeDataSource(DataSource):
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.DateRangeParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a DateRangeParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a DateRangeParameterConfig
 
         Arguments:
             ds_param: The parameter to convert
@@ -354,11 +360,11 @@ class _NumericDataSource(DataSource):
     @abc.abstractmethod
     def __init__(
         self, table_or_query: str, min_value_col: str, max_value_col: str, *, increment_col: str | None = None, 
-        id_col: str | None = None, from_seeds: bool = False, user_group_col: str | None = None, 
+        id_col: str | None = None, source: SourceEnum = SourceEnum.CONNECTION, user_group_col: str | None = None, 
         parent_id_col: str | None = None, connection: str | None = None, **kwargs
     ) -> None:
         super().__init__(
-            table_or_query, id_col=id_col, from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, 
+            table_or_query, id_col=id_col, source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, 
             connection=connection
         )
         self._min_value_col = min_value_col
@@ -375,7 +381,7 @@ class NumberDataSource(_NumericDataSource):
 
     def __init__(
         self, table_or_query: str, min_value_col: str, max_value_col: str, *, increment_col: str | None = None,
-        default_value_col: str | None = None, id_col: str | None = None, from_seeds: bool = False, 
+        default_value_col: str | None = None, id_col: str | None = None, source: SourceEnum = SourceEnum.CONNECTION, 
         user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None, **kwargs
     ) -> None:
         """
@@ -388,20 +394,20 @@ class NumberDataSource(_NumericDataSource):
             increment_col: The column name of the increment value. Defaults to column of 1's if None
             default_value_col: The column name of the default value. Defaults to min_value_col if None
             id_col: The column name of the id
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that the default value belongs to
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
-            table_or_query, min_value_col, max_value_col, increment_col=increment_col, id_col=id_col, from_seeds=from_seeds,
+            table_or_query, min_value_col, max_value_col, increment_col=increment_col, id_col=id_col, source=source,
             user_group_col=user_group_col, parent_id_col=parent_id_col, connection=connection
         )
         self._default_value_col = default_value_col
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.NumberParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a NumberParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a NumberParameterConfig
 
         Arguments:
             ds_param: The parameter to convert
@@ -443,7 +449,7 @@ class NumberRangeDataSource(_NumericDataSource):
     def __init__(
         self, table_or_query: str, min_value_col: str, max_value_col: str, *, increment_col: str | None = None,
         default_lower_value_col: str | None = None, default_upper_value_col: str | None = None, id_col: str | None = None, 
-        from_seeds: bool = False, user_group_col: str | None = None, parent_id_col: str | None = None, 
+        source: SourceEnum = SourceEnum.CONNECTION, user_group_col: str | None = None, parent_id_col: str | None = None, 
         connection: str | None = None, **kwargs
     ) -> None:
         """
@@ -457,13 +463,13 @@ class NumberRangeDataSource(_NumericDataSource):
             default_lower_value_col: The column name of the default lower value. Defaults to min_value_col if None
             default_upper_value_col: The column name of the default upper value. Defaults to max_value_col if None
             id_col: The column name of the id
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that the default value belongs to
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
-            table_or_query, min_value_col, max_value_col, increment_col=increment_col, id_col=id_col, from_seeds=from_seeds, 
+            table_or_query, min_value_col, max_value_col, increment_col=increment_col, id_col=id_col, source=source, 
             user_group_col=user_group_col, parent_id_col=parent_id_col, connection=connection
         )
         self._default_lower_value_col = default_lower_value_col
@@ -471,7 +477,7 @@ class NumberRangeDataSource(_NumericDataSource):
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.NumberRangeParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a NumberRangeParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a NumberRangeParameterConfig
 
         Arguments:
             ds_param: The parameter to convert
@@ -511,7 +517,7 @@ class TextDataSource(DataSource):
     _default_text_col: str
 
     def __init__(
-        self, table_or_query: str, default_text_col: str, *, id_col: str | None = None, from_seeds: bool = False, 
+        self, table_or_query: str, default_text_col: str, *, id_col: str | None = None, source: SourceEnum = SourceEnum.CONNECTION, 
         user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None,
         **kwargs
     ) -> None:
@@ -522,20 +528,20 @@ class TextDataSource(DataSource):
             table_or_query: Either the name of the table to use, or a query to run
             default_text_col: The column name of the default text
             id_col: The column name of the id
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that the default date belongs to
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
-            table_or_query, id_col=id_col, from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, 
+            table_or_query, id_col=id_col, source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, 
             connection=connection
         )
         self._default_text_col = default_text_col
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.TextParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a TextParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a TextParameterConfig
 
         Arguments:
             ds_param: The parameter to convert

squirrels/_dataset_types.py

@@ -35,19 +35,17 @@ class DatasetMetadata:
 @dataclass
 class DatasetResult(DatasetMetadata):
     df: pl.DataFrame
-    to_json: Callable[[str, tuple[str, ...], int, int], dict] = field(init=False)
+    to_json: Callable[[str, int, int], dict] = field(init=False)
 
     def __post_init__(self):
         self.to_json = lru_cache()(self._to_json)
     
-    def _to_json(self, orientation: Literal["records", "rows", "columns"], select: tuple[str, ...], limit: int, offset: int) -> dict:
+    def _to_json(self, orientation: Literal["records", "rows", "columns"], limit: int, offset: int) -> dict:
         df = self.df.lazy()
         if offset > 0:
             df = df.filter(pl.col("_row_num") > offset)
         if limit > 0:
             df = df.limit(limit)
-        if select:
-            df = df.select(select)
         df = df.collect()
         
         if orientation == "columns":

squirrels/_initializer.py

@@ -266,7 +266,7 @@ class Initializer:
         
         self._copy_database_file(c.EXPENSES_DB)
         
-        print(f"\nSuccessfully created new Squirrels project in current directory!\n")
+        print(f"\nSuccessfully created new Squirrels project!\n")
     
     def get_file(self, args):
         if args.file_name == c.DOTENV_FILE:

squirrels/_logging.py

@@ -0,0 +1,117 @@
+from pathlib import Path
+from logging.handlers import RotatingFileHandler
+from uuid import uuid4
+import logging as l, json
+
+from . import _constants as c, _utils as u
+from ._request_context import get_request_id
+
+
+class _BaseFormatter(l.Formatter):
+    def _format_helper(self, level_for_print: str, record: l.LogRecord) -> str:
+        # Save original levelname
+        original_levelname = record.levelname
+
+        # Add padding to the levelname for printing
+        visible_length = len(record.levelname) + 1
+        padding_needed = max(1, 9 - visible_length)
+        padded_level = f"{level_for_print}:{' ' * padding_needed}"
+        record.levelname = padded_level
+        
+        # Format the message
+        formatted = super().format(record)
+        
+        # Append request ID if available
+        request_id = get_request_id()
+        request_id_str = f" [req_id: {request_id}]" if request_id else ""
+        formatted = formatted.format(request_id=request_id_str)
+        
+        # Restore original levelname
+        record.levelname = original_levelname
+        
+        return formatted
+
+
+class _ColoredFormatter(_BaseFormatter):
+    """Custom formatter that adds colors to log levels for terminal output"""
+    
+    # ANSI color codes
+    COLORS = {
+        'DEBUG': '\033[36m',      # Cyan
+        'INFO': '\033[32m',       # Green
+        'WARNING': '\033[33m',    # Yellow
+        'ERROR': '\033[31m',      # Red
+        'CRITICAL': '\033[35m',   # Magenta
+    }
+    RESET = '\033[0m'
+    BOLD = '\033[1m'
+    
+    def format(self, record: l.LogRecord) -> str:
+        # Add color to levelname with colon and padding
+        color = self.COLORS.get(record.levelname, '')
+        colored_level = f"{color}{record.levelname}{self.RESET}"
+        return self._format_helper(colored_level, record)
+
+
+class _PlainFormatter(_BaseFormatter):
+    """Custom formatter that adds colon to log levels for file output"""
+    
+    def format(self, record: l.LogRecord) -> str:
+        return self._format_helper(record.levelname, record)
+
+
+class _CustomJsonFormatter(l.Formatter):
+    def format(self, record: l.LogRecord) -> str:
+        super().format(record)
+        request_id = get_request_id()
+        info = {
+            "timestamp": self.formatTime(record),
+            "level": record.levelname,
+            "message": record.getMessage(),
+            "request_id": request_id,
+        }
+        output = {
+            "data": record.__dict__.get("data", {}),
+            "info": info
+        }
+        return json.dumps(output)
+
+
+def get_logger(
+    base_path: str, log_to_file: bool, log_level: str, log_format: str, log_file_size_mb: int, log_file_backup_count: int
+) -> u.Logger:
+    logger = u.Logger(name=uuid4().hex, level=log_level.upper())
+
+    # Determine the formatter based on log_format
+    if log_format.lower() == "json":
+        stdout_formatter = _CustomJsonFormatter()
+        file_formatter = _CustomJsonFormatter()
+    elif log_format.lower() == "text":
+        # Use colored formatter for stdout, plain formatter with colon for file
+        format_string = "%(levelname)s [%(asctime)s]{request_id} %(message)s"
+        stdout_formatter = _ColoredFormatter(format_string, datefmt="%Y-%m-%d %H:%M:%S")
+        file_formatter = _PlainFormatter(format_string, datefmt="%Y-%m-%d %H:%M:%S")
+    else:
+        raise ValueError("log_format must be either 'text' or 'json'")
+    
+    if log_to_file:
+        log_file_path = Path(base_path, c.LOGS_FOLDER, c.LOGS_FILE)
+        log_file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Rotating file handler
+        file_handler = RotatingFileHandler(
+            log_file_path,
+            maxBytes=log_file_size_mb * 1024 * 1024,
+            backupCount=log_file_backup_count
+        )
+        file_handler.setLevel(log_level.upper())
+        file_handler.setFormatter(file_formatter)
+        logger.addHandler(file_handler)
+    
+    else:
+        stdout_handler = l.StreamHandler()
+        stdout_handler.setLevel(log_level.upper())
+        stdout_handler.setFormatter(stdout_formatter)
+        logger.addHandler(stdout_handler)
+    
+    return logger