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

squirrels/_seeds.py

@@ -0,0 +1,97 @@
+from dataclasses import dataclass
+import os
+import re
+import time
+import glob
+import json
+
+import polars as pl
+
+from ._exceptions import ConfigurationError
+from . import _utils as u, _constants as c, _model_configs as mc
+from ._env_vars import SquirrelsEnvVars
+
+
+@dataclass
+class Seed:
+    config: mc.SeedConfig
+    df: pl.LazyFrame
+
+    def __post_init__(self):
+        if self.config.cast_column_types:
+            exprs = []
+            for col_config in self.config.columns:
+                col_type = col_config.type.lower()
+                if col_type.startswith("decimal"):
+                    polars_dtype = self._parse_decimal_type(col_type)
+                else:
+                    try:
+                        polars_dtype = u.sqrl_dtypes_to_polars_dtypes[col_type]
+                    except KeyError as e:
+                        raise ConfigurationError(f"Unknown column type: '{col_type}'") from e
+                
+                exprs.append(pl.col(col_config.name).cast(polars_dtype))
+
+            self.df = self.df.with_columns(*exprs)
+
+    @staticmethod
+    def _parse_decimal_type(col_type: str) -> pl.Decimal:
+        """Parse a decimal type string and return the appropriate polars Decimal type.
+        
+        Supports formats: "decimal" or "decimal(precision, scale)"
+        """
+
+        # Match decimal(precision, scale) pattern
+        match = re.match(r"decimal\s*\(\s*(\d+)\s*,\s*(\d+)\s*\)", col_type)
+        if match:
+            precision = int(match.group(1))
+            scale = int(match.group(2))
+            return pl.Decimal(precision=precision, scale=scale)
+        
+        if col_type == "decimal":
+            return pl.Decimal(precision=18, scale=2)
+        
+        raise ConfigurationError(f"Unknown column type: '{col_type}'")
+
+
+@dataclass
+class Seeds:
+    _data: dict[str, Seed]
+
+    def run_query(self, sql_query: str) -> pl.DataFrame:
+        dataframes = {key: seed.df for key, seed in self._data.items()}
+        return u.run_sql_on_dataframes(sql_query, dataframes)
+
+    def get_dataframes(self) -> dict[str, Seed]:
+        return self._data.copy()
+
+
+class SeedsIO:
+
+    @classmethod
+    def load_files(cls, logger: u.Logger, env_vars: SquirrelsEnvVars) -> Seeds:
+        start = time.time()
+        project_path = env_vars.project_path
+        infer_schema_setting: bool = env_vars.seeds_infer_schema
+        na_values_setting: list[str] = env_vars.seeds_na_values
+
+        seeds_dict = {}
+        csv_files = glob.glob(os.path.join(project_path, c.SEEDS_FOLDER, '**/*.csv'), recursive=True)
+        for csv_file in csv_files:
+            config_file = os.path.splitext(csv_file)[0] + '.yml'
+            config_dict = u.load_yaml_config(config_file) if os.path.exists(config_file) else {}
+            config = mc.SeedConfig(**config_dict)
+
+            file_stem = os.path.splitext(os.path.basename(csv_file))[0]
+            infer_schema = not config.cast_column_types and infer_schema_setting
+            df = pl.read_csv(
+                csv_file, try_parse_dates=True,
+                infer_schema=infer_schema,
+                null_values=na_values_setting
+            ).lazy()
+
+            seeds_dict[file_stem] = Seed(config, df)
+
+        seeds = Seeds(seeds_dict)
+        logger.log_activity_time("loading seed files", start)
+        return seeds

squirrels/_sources.py

@@ -0,0 +1,112 @@
+from typing import Any
+from pydantic import BaseModel, Field, model_validator
+import time, yaml
+
+from . import _utils as u, _constants as c, _model_configs as mc
+from ._env_vars import SquirrelsEnvVars
+
+
+class UpdateHints(BaseModel):
+    increasing_column: str | None = Field(default=None)
+    strictly_increasing: bool = Field(default=True, description="Delete the max value of the increasing column, ignored if selective_overwrite_value is set")
+    selective_overwrite_value: Any = Field(default=None, description="Delete all values of the increasing column greater than or equal to this value")
+
+
+class Source(mc.ConnectionInterface, mc.ModelConfig):
+    table: str | None = Field(default=None)
+    load_to_vdl: bool = Field(default=False, description="Whether to load the data to the 'virtual data lake' (VDL)")
+    primary_key: list[str] = Field(default_factory=list)
+    update_hints: UpdateHints = Field(default_factory=UpdateHints)
+
+    def finalize_table(self, source_name: str):
+        if self.table is None:
+            self.table = source_name
+        return self
+    
+    def get_table(self) -> str:
+        assert self.table is not None, "Table must be set"
+        return self.table
+    
+    def get_cols_for_create_table_stmt(self) -> str:
+        cols_clause = ", ".join([f"{col.name} {col.type}" for col in self.columns])
+        return cols_clause
+    
+    def get_max_incr_col_query(self, source_name: str) -> str:
+        return f"SELECT max({self.update_hints.increasing_column}) FROM {source_name}"
+    
+    def get_query_for_upsert(self, dialect: str, conn_name: str, table_name: str, max_value_of_increasing_col: Any | None, *, full_refresh: bool = True) -> str:
+        select_cols = ", ".join([col.name for col in self.columns])
+        if full_refresh or max_value_of_increasing_col is None:
+            return f"SELECT {select_cols} FROM db_{conn_name}.{table_name}"
+        
+        increasing_col = self.update_hints.increasing_column
+        increasing_col_type = next(col.type for col in self.columns if col.name == increasing_col)
+        where_cond = f"{increasing_col}::{increasing_col_type} > '{max_value_of_increasing_col}'::{increasing_col_type}"
+        
+        # TODO: figure out if using pushdown query is worth it
+        # if dialect in ['postgres', 'mysql']:
+        #     pushdown_query = f"SELECT {select_cols} FROM {table_name} WHERE {where_cond}"
+        #     transpiled_query = sqlglot.transpile(pushdown_query, read='duckdb', write=dialect)[0].replace("'", "''")
+        #     return f"FROM {dialect}_query('db_{conn_name}', '{transpiled_query}')"
+        
+        return f"SELECT {select_cols} FROM db_{conn_name}.{table_name} WHERE {where_cond}"
+    
+
+class Sources(BaseModel):
+    sources: dict[str, Source] = Field(default_factory=dict)
+    
+    @model_validator(mode="before")
+    @classmethod
+    def convert_sources_list_to_dict(cls, data: dict[str, Any]) -> dict[str, Any]:
+        if "sources" in data and isinstance(data["sources"], list):
+            # Convert list of sources to dictionary
+            sources_dict = {}
+            for source in data["sources"]:
+                if isinstance(source, dict) and "name" in source:
+                    name = source.pop("name")  # Remove name from source config
+                    if name in sources_dict:
+                        raise u.ConfigurationError(f"Duplicate source name found: {name}")
+                    sources_dict[name] = source
+                else:
+                    raise u.ConfigurationError(f"All sources must have a name field in sources file")
+            data["sources"] = sources_dict
+        return data
+    
+    @model_validator(mode="after")
+    def validate_column_types(self):
+        for source_name, source in self.sources.items():
+            for col in source.columns:
+                if not col.type:
+                    raise u.ConfigurationError(f"Column '{col.name}' in source '{source_name}' must have a type specified")
+        return self
+    
+    def finalize_null_fields(self, env_vars: SquirrelsEnvVars):
+        default_conn_name = env_vars.connections_default_name_used
+        for source_name, source in self.sources.items():
+            source.finalize_connection(default_conn_name=default_conn_name)
+            source.finalize_table(source_name)
+        return self
+
+
+class SourcesIO:
+    @classmethod
+    def load_file(cls, logger: u.Logger, env_vars: SquirrelsEnvVars, env_vars_unformatted: dict[str, str]) -> Sources:
+        start = time.time()
+        
+        sources_path = u.Path(env_vars.project_path, c.MODELS_FOLDER, c.SOURCES_FILE)
+        if sources_path.exists():
+            raw_content = u.read_file(sources_path)
+            rendered = u.render_string(raw_content, project_path=env_vars.project_path, env_vars=env_vars_unformatted)
+            sources_data = yaml.safe_load(rendered) or {}
+        else:
+            sources_data = {}
+        
+        if not isinstance(sources_data, dict):
+            raise u.ConfigurationError(
+                f"Parsed content from YAML file must be a dictionary. Got: {sources_data}"
+            )
+        
+        sources = Sources(**sources_data).finalize_null_fields(env_vars)
+        
+        logger.log_activity_time("loading sources", start)
+        return sources