squirrels 0.5.0b2__py3-none-any.whl → 0.5.0b3__py3-none-any.whl

squirrels/_project.py

@@ -1,5 +1,6 @@
 from dotenv import dotenv_values
 from uuid import uuid4
+from pathlib import Path
 import asyncio, typing as t, functools as ft, shutil, json, os
 import logging as l, matplotlib.pyplot as plt, networkx as nx, polars as pl
 import sqlglot, sqlglot.expressions
@@ -9,7 +10,7 @@ from ._model_builder import ModelBuilder
 from ._exceptions import InvalidInputError, ConfigurationError
 from . import _utils as u, _constants as c, _manifest as mf, _connection_set as cs, _api_response_models as arm
 from . import _seeds as s, _models as m, _model_configs as mc, _model_queries as mq, _sources as so
-from . import _parameter_sets as ps, _dashboards_io as d, dashboards as dash, dataset_result as dr
+from . import _parameter_sets as ps, _dashboards_io as d, _dashboard_types as dash, _dataset_types as dr
 
 T = t.TypeVar("T", bound=dash.Dashboard)
 M = t.TypeVar("M", bound=m.DataModel)
@@ -70,7 +71,7 @@ class SquirrelsProject:
             raise ValueError("log_format must be either 'text' or 'json'")
             
         if log_file:
-            path = u.Path(base_path, c.LOGS_FOLDER, log_file)
+            path = Path(base_path, c.LOGS_FOLDER, log_file)
             path.parent.mkdir(parents=True, exist_ok=True)
 
             handler = l.FileHandler(path)
@@ -131,6 +132,10 @@ class SquirrelsProject:
     def _auth(self) -> Authenticator:
         return Authenticator(self._logger, self._filepath, self._env_vars)
     
+    @ft.cached_property
+    def User(self) -> t.Type[BaseUser]:
+        return self._auth.User
+    
     @ft.cached_property
     def _param_args(self) -> ps.ParametersArgs:
         return ps.ParameterConfigsSetIO.get_param_args(self._conn_args)
@@ -143,12 +148,32 @@ class SquirrelsProject:
     
     @ft.cached_property
     def _j2_env(self) -> u.EnvironmentWithMacros:
-        return u.EnvironmentWithMacros(self._logger, loader=u.j2.FileSystemLoader(self._filepath))
+        env = u.EnvironmentWithMacros(self._logger, loader=u.j2.FileSystemLoader(self._filepath))
+
+        def value_to_str(value: t.Any, attribute: str | None = None) -> str:
+            if attribute is None:
+                return str(value)
+            else:
+                return str(getattr(value, attribute))
+
+        def join(value: list[t.Any], d: str = ", ", attribute: str | None = None) -> str:
+            return d.join(map(lambda x: value_to_str(x, attribute), value))
+
+        def quote(value: t.Any, q: str = "'", attribute: str | None = None) -> str:
+            return q + value_to_str(value, attribute) + q
+        
+        def quote_and_join(value: list[t.Any], q: str = "'", d: str = ", ", attribute: str | None = None) -> str:
+            return d.join(map(lambda x: quote(x, q, attribute), value))
+        
+        env.filters["join"] = join
+        env.filters["quote"] = quote
+        env.filters["quote_and_join"] = quote_and_join
+        return env
 
     @ft.cached_property
     def _duckdb_venv_path(self) -> str:
         duckdb_filepath_setting_val = self._env_vars.get(c.SQRL_DUCKDB_VENV_DB_FILE_PATH, f"{c.TARGET_FOLDER}/{c.DUCKDB_VENV_FILE}")
-        return str(u.Path(self._filepath, duckdb_filepath_setting_val))
+        return str(Path(self._filepath, duckdb_filepath_setting_val))
     
     def close(self) -> None:
         """
@@ -255,7 +280,7 @@ class SquirrelsProject:
         dag = m.DAG(None, fake_target_model, models_dict, self._duckdb_venv_path, self._logger)
         return dag
     
-    def _draw_dag(self, dag: m.DAG, output_folder: u.Path) -> None:
+    def _draw_dag(self, dag: m.DAG, output_folder: Path) -> None:
         color_map = {
             m.ModelType.SEED: "green", m.ModelType.DBVIEW: "red", m.ModelType.FEDERATE: "skyblue",
             m.ModelType.BUILD: "purple", m.ModelType.SOURCE: "orange"
@@ -275,7 +300,7 @@ class SquirrelsProject:
         
         fig.tight_layout()
         plt.margins(x=0.1, y=0.1)
-        fig.savefig(u.Path(output_folder, "dag.png"))
+        fig.savefig(Path(output_folder, "dag.png"))
         plt.close(fig)
     
     async def _get_compiled_dag(self, *, sql_query: str | None = None, selections: dict[str, t.Any] = {}, user: BaseUser | None = None) -> m.DAG:
@@ -371,28 +396,28 @@ class SquirrelsProject:
             runquery=runquery, recurse=recurse, default_traits=self._manifest_cfg.get_default_traits()
         )
         
-        output_folder = u.Path(self._filepath, c.TARGET_FOLDER, c.COMPILE_FOLDER, dataset, test_set)
+        output_folder = Path(self._filepath, c.TARGET_FOLDER, c.COMPILE_FOLDER, dataset, test_set)
         if output_folder.exists():
             shutil.rmtree(output_folder)
         output_folder.mkdir(parents=True, exist_ok=True)
         
         def write_placeholders() -> None:
-            output_filepath = u.Path(output_folder, "placeholders.json")
+            output_filepath = Path(output_folder, "placeholders.json")
             with open(output_filepath, 'w') as f:
                 json.dump(dag.placeholders, f, indent=4)
         
         def write_model_outputs(model: m.DataModel) -> None:
             assert isinstance(model, m.QueryModel)
             subfolder = c.DBVIEWS_FOLDER if model.model_type == m.ModelType.DBVIEW else c.FEDERATES_FOLDER
-            subpath = u.Path(output_folder, subfolder)
+            subpath = Path(output_folder, subfolder)
             subpath.mkdir(parents=True, exist_ok=True)
             if isinstance(model.compiled_query, mq.SqlModelQuery):
-                output_filepath = u.Path(subpath, model.name+'.sql')
+                output_filepath = Path(subpath, model.name+'.sql')
                 query = model.compiled_query.query
                 with open(output_filepath, 'w') as f:
                     f.write(query)
             if runquery and isinstance(model.result, pl.LazyFrame):
-                output_filepath = u.Path(subpath, model.name+'.csv')
+                output_filepath = Path(subpath, model.name+'.csv')
                 model.result.collect().write_csv(output_filepath)
 
         write_placeholders()
@@ -520,7 +545,7 @@ class SquirrelsProject:
         )
     
     async def dashboard(
-        self, name: str, *, selections: dict[str, t.Any] = {}, user: BaseUser | None = None, dashboard_type: t.Type[T] = dash.Dashboard
+        self, name: str, *, selections: dict[str, t.Any] = {}, user: BaseUser | None = None, dashboard_type: t.Type[T] = dash.PngDashboard
     ) -> T:
         """
         Async method to retrieve a dashboard given parameter selections.

squirrels/_utils.py

@@ -290,7 +290,8 @@ def load_yaml_config(filepath: FilePath) -> dict:
 
 
 def run_duckdb_stmt(
-    logger: Logger, duckdb_conn: duckdb.DuckDBPyConnection, stmt: str, *, params: dict[str, Any] | None = None, redacted_values: list[str] = []
+    logger: Logger, duckdb_conn: duckdb.DuckDBPyConnection, stmt: str, *, params: dict[str, Any] | None = None, 
+    model_name: str | None = None, redacted_values: list[str] = []
 ) -> duckdb.DuckDBPyConnection:
     """
     Runs a statement on a DuckDB connection
@@ -306,7 +307,8 @@ def run_duckdb_stmt(
     for value in redacted_values:
         redacted_stmt = redacted_stmt.replace(value, "[REDACTED]")
     
-    logger.info(f"Running statement: {redacted_stmt}", extra={"data": {"params": params}})
+    for_model_name = f" for model '{model_name}'" if model_name is not None else ""
+    logger.info(f"Running SQL statement{for_model_name}:\n{redacted_stmt}", extra={"data": {"params": params}})
     try:
         return duckdb_conn.execute(stmt, params)
     except duckdb.ParserException as e:

squirrels/arguments.py

@@ -0,0 +1,2 @@
+from ._arguments._init_time_args import ConnectionsArgs, ParametersArgs, BuildModelArgs
+from ._arguments._run_time_args import ContextArgs, ModelArgs, DashboardArgs

squirrels/connections.py

@@ -0,0 +1 @@
+from ._manifest import ConnectionProperties, ConnectionTypeEnum

squirrels/dashboards.py

@@ -1,82 +1 @@
-import matplotlib.figure as _figure, io as _io, abc as _abc, typing as _t
-
-from . import _constants as c
-
-
-class Dashboard(metaclass=_abc.ABCMeta):
-    """
-    Abstract parent class for all Dashboard classes.
-    """
-    
-    @property
-    @_abc.abstractmethod
-    def _content(self) -> bytes | str:
-        pass
-    
-    @property
-    @_abc.abstractmethod
-    def _format(self) -> str:
-        pass
-
-
-class PngDashboard(Dashboard):
-    """
-    Instantiate a Dashboard in PNG format from a matplotlib figure or bytes
-    """
-    
-    def __init__(self, content: _figure.Figure | _io.BytesIO | bytes) -> None:
-        """
-        Constructor for PngDashboard
-
-        Arguments:
-            content: The content of the dashboard as a matplotlib.figure.Figure or bytes
-        """
-        if isinstance(content, _figure.Figure):
-            buffer = _io.BytesIO()
-            content.savefig(buffer, format=c.PNG)
-            content = buffer.getvalue()
-        
-        if isinstance(content, _io.BytesIO):
-            content = content.getvalue()
-        
-        self.__content = content
-
-    @property
-    def _content(self) -> bytes:
-        return self.__content
-    
-    @property
-    def _format(self) -> _t.Literal['png']:
-        return c.PNG
-    
-    def _repr_png_(self):
-        return self._content
-    
-
-class HtmlDashboard(Dashboard):
-    """
-    Instantiate a Dashboard from an HTML string
-    """
-
-    def __init__(self, content: _io.StringIO | str) -> None:
-        """
-        Constructor for HtmlDashboard
-
-        Arguments:
-            content: The content of the dashboard as HTML string
-        """
-        if isinstance(content, _io.StringIO):
-            content = content.getvalue()
-        
-        self.__content = content
-
-    @property
-    def _content(self) -> str:
-        return self.__content
-    
-    @property
-    def _format(self) -> _t.Literal['html']:
-        return c.HTML
-    
-    def _repr_html_(self):
-        return self._content
+from ._dashboard_types import PngDashboard, HtmlDashboard