squirrels 0.1.0__py3-none-any.whl → 0.6.0.post0__py3-none-any.whl

squirrels/_model_configs.py

@@ -0,0 +1,74 @@
+from enum import Enum
+from pydantic import BaseModel, Field
+
+from ._env_vars import SquirrelsEnvVars
+
+
+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: list[str] = Field(default_factory=list, description="The condition(s) of when the column is included. Only for documentation purposes.")
+    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, *, default_conn_name: str = "default"):
+        if self.connection is None:
+            self.connection = default_conn_name
+        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="VIEW", 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\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 always be materialized in memory 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\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]