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

squirrels/_model_builder.py

@@ -0,0 +1,111 @@
+from dataclasses import dataclass, field
+import asyncio, shutil, duckdb, time
+
+from . import _utils as u, _connection_set as cs, _models as m
+from ._exceptions import InvalidInputError
+
+
+@dataclass
+class ModelBuilder:
+    _duckdb_venv_path: str
+    _conn_set: cs.ConnectionSet
+    _static_models: dict[str, m.StaticModel]
+    _conn_args: cs.ConnectionsArgs = field(default_factory=lambda: cs.ConnectionsArgs(".", {}, {}))
+    _logger: u.Logger = field(default_factory=lambda: u.Logger(""))
+    
+    def _attach_connections(self, duckdb_conn: duckdb.DuckDBPyConnection) -> dict[str, str]:
+        dialect_by_conn_name: dict[str, str] = {}
+        for conn_name, conn_props in self._conn_set.get_connections_as_dict().items():
+            if not isinstance(conn_props, m.ConnectionProperties):
+                continue
+            dialect = conn_props.dialect
+            attach_uri = conn_props.attach_uri_for_duckdb
+            if attach_uri is None:
+                continue # skip unsupported dialects
+            attach_stmt = f"ATTACH IF NOT EXISTS '{attach_uri}' AS db_{conn_name} (TYPE {dialect}, READ_ONLY)"
+            u.run_duckdb_stmt(self._logger, duckdb_conn, attach_stmt, redacted_values=[attach_uri])
+            dialect_by_conn_name[conn_name] = dialect
+        return dialect_by_conn_name
+
+    async def _build_models(self, duckdb_conn: duckdb.DuckDBPyConnection, select: str | None, full_refresh: bool) -> None:
+        """
+        Compile and construct the build models as DuckDB tables.
+        """
+        # Compile the build models
+        models_list = self._static_models.values() if select is None else [self._static_models[select]]
+        for model in models_list:
+            model.compile_for_build(self._conn_args, self._static_models)
+
+        # Find all terminal nodes
+        terminal_nodes = set()
+        if select is None:
+            for model in models_list:
+                terminal_nodes.update(model.get_terminal_nodes_for_build(set()))
+            for model in models_list:
+                model.confirmed_no_cycles = False
+        else:
+            terminal_nodes.add(select)
+
+        # Run the build models
+        coroutines = []
+        for model_name in terminal_nodes:
+            model = self._static_models[model_name]
+            coro = model.build_model(duckdb_conn, full_refresh)
+            coroutines.append(coro)
+        await u.asyncio_gather(coroutines)
+
+    async def build(self, full_refresh: bool, select: str | None, stage_file: bool) -> None:
+        start = time.time()
+
+        # Create target folder if it doesn't exist
+        duckdb_path = u.Path(self._duckdb_venv_path)
+        duckdb_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Delete any existing DuckDB file if full refresh is requested
+        duckdb_dev_path = u.Path(self._duckdb_venv_path + ".dev")
+        duckdb_stg_path = u.Path(self._duckdb_venv_path + ".stg")
+
+        # If the development copy is already in use, a concurrent build is not allowed
+        duckdb_dev_lock_path = u.Path(self._duckdb_venv_path + ".dev.lock")
+        if duckdb_dev_lock_path.exists():
+            raise InvalidInputError(60, "An existing build process is already running and a concurrent build is not allowed")
+        duckdb_dev_lock_path.touch(exist_ok=False)
+        
+        # Ensure the lock file is deleted even if an exception is raised
+        try:
+            # If not full refresh, create a development copy of the existing virtual data environment
+            if not full_refresh:
+                if duckdb_stg_path.exists():
+                    duckdb_stg_path.replace(duckdb_dev_path)
+                elif duckdb_path.exists():
+                    shutil.copy(duckdb_path, duckdb_dev_path)
+                
+            self._logger.log_activity_time("creating development copy of virtual data environment", start)
+        
+            # Connect to DuckDB file
+            duckdb_conn = u.create_duckdb_connection(duckdb_dev_path)
+            
+        except Exception:
+            duckdb_dev_lock_path.unlink()
+            raise
+        
+        # Sometimes code after conn.close() doesn't run (as if the python process is killed but no error is raised)
+        # Using a new try block to ensure the lock file is removed before closing the connection
+        try:
+            # Attach connections
+            self._attach_connections(duckdb_conn)
+
+            # Construct build models
+            await self._build_models(duckdb_conn, select, full_refresh)
+        
+        finally:
+            duckdb_dev_lock_path.unlink()
+            duckdb_conn.close()
+
+        # Rename duckdb_dev_path to duckdb_path (or duckdb_stg_path if stage_file is True)
+        if stage_file:
+            duckdb_dev_path.replace(duckdb_stg_path)
+        else:
+            duckdb_dev_path.replace(duckdb_path)
+
+        self._logger.log_activity_time("TOTAL TIME to build virtual data environment", start)

squirrels/_model_configs.py

@@ -0,0 +1,74 @@
+from enum import Enum
+from pydantic import BaseModel, Field
+
+from . import _constants as c
+
+
+class ColumnCategory(Enum):
+    DIMENSION = "dimension"
+    MEASURE = "measure"
+    MISC = "misc"
+
+
+class ColumnConfig(BaseModel):
+    name: str = Field(description="The name of the column")
+    type: str = Field(default="", description="The type of the column such as 'string', 'integer', 'float', 'boolean', 'datetime', etc.")
+    condition: str = Field(default="", description="The condition of when the column is included")
+    description: str = Field(default="", description="The description of the column")
+    category: ColumnCategory = Field(default=ColumnCategory.MISC, description="The category of the column, either 'dimension', 'measure', or 'misc'")
+    depends_on: set[str] = Field(default_factory=set, description="List of dependent columns")
+    pass_through: bool = Field(default=False, description="Whether the column should be passed through to the federate")
+
+
+class ModelConfig(BaseModel):
+    description: str = Field(default="", description="The description of the model")
+    columns: list[ColumnConfig] = Field(default_factory=list, description="The columns of the model")
+
+
+class SeedConfig(ModelConfig):
+    cast_column_types: bool = Field(default=False, description="Whether the column types should be cast to the appropriate type")
+
+
+class ConnectionInterface(BaseModel):
+    connection: str | None = Field(default=None, description="The connection name of the source model / database view")
+
+    def finalize_connection(self, env_vars: dict[str, str]):
+        if self.connection is None:
+            self.connection = env_vars.get(c.SQRL_CONNECTIONS_DEFAULT_NAME_USED, "default")
+        return self
+
+    def get_connection(self) -> str:
+        assert self.connection is not None, "Connection must be set"
+        return self.connection
+    
+
+class QueryModelConfig(ModelConfig):
+    depends_on: set[str] = Field(default_factory=set, description="The dependencies of the model")
+
+
+class BuildModelConfig(QueryModelConfig):
+    materialization: str = Field(default="TABLE", description="The materialization of the model (ignored if Python model which is always a table)")
+
+    def get_sql_for_build(self, model_name: str, select_query: str) -> str:
+        if self.materialization.upper() == "TABLE":
+            materialization = "TABLE"
+        elif self.materialization.upper() == "VIEW":
+            materialization = "VIEW"
+        else:
+            raise ValueError(f"Invalid materialization: {self.materialization}")
+        
+        create_prefix = f"CREATE OR REPLACE {materialization} {model_name} AS\n"
+        return create_prefix + select_query
+
+
+class DbviewModelConfig(ConnectionInterface, QueryModelConfig):
+    translate_to_duckdb: bool = Field(default=False, description="Whether to translate the query to DuckDB and use DuckDB tables at runtime")
+
+
+class FederateModelConfig(QueryModelConfig):
+    eager: bool = Field(default=False, description="Whether the model should be materialized for SQL models")
+
+    def get_sql_for_create(self, model_name: str, select_query: str) -> str:
+        materialization = "TABLE" if self.eager else "VIEW"
+        create_prefix = f"CREATE {materialization} {model_name} AS\n"
+        return create_prefix + select_query

squirrels/_model_queries.py

@@ -0,0 +1,52 @@
+from abc import ABCMeta
+from dataclasses import dataclass, field
+from typing import Callable, Generic, TypeVar, Any
+import polars as pl, pandas as pd
+
+from .arguments.run_time_args import BuildModelArgs
+from ._model_configs import ModelConfig
+
+
+# Input query file classes
+
+@dataclass(frozen=True)
+class QueryFile(metaclass=ABCMeta):
+    filepath: str
+    raw_query: Any
+
+@dataclass(frozen=True)
+class SqlQueryFile(QueryFile):
+    raw_query: str
+
+@dataclass(frozen=True)
+class PyQueryFile(QueryFile):
+    raw_query: Callable[[BuildModelArgs], pl.LazyFrame | pd.DataFrame]
+
+
+Q = TypeVar('Q', bound=QueryFile)
+M = TypeVar('M', bound=ModelConfig)
+
+@dataclass(frozen=True)
+class QueryFileWithConfig(Generic[Q, M]):
+    query_file: Q
+    config: M
+
+
+# Compiled query classes
+
+@dataclass
+class Query(metaclass=ABCMeta):
+    query: Any
+
+@dataclass
+class WorkInProgress(Query):
+    query: None = field(default=None, init=False)
+
+@dataclass
+class SqlModelQuery(Query):
+    query: str
+    is_duckdb: bool
+
+@dataclass
+class PyModelQuery(Query):
+    query: Callable[[], pl.LazyFrame | pd.DataFrame]