policyengine 3.0.0__py3-none-any.whl → 3.1.1__py3-none-any.whl

policyengine/database/database.py

@@ -1,251 +0,0 @@
-from typing import Any
-
-from sqlmodel import Session, SQLModel
-
-from .aggregate import aggregate_table_link
-from .baseline_parameter_value_table import baseline_parameter_value_table_link
-from .baseline_variable_table import baseline_variable_table_link
-from .dataset_table import dataset_table_link
-from .dynamic_table import dynamic_table_link
-from .link import TableLink
-
-# Import all table links
-from .model_table import model_table_link
-from .model_version_table import model_version_table_link
-from .parameter_table import parameter_table_link
-from .parameter_value_table import parameter_value_table_link
-from .policy_table import policy_table_link
-from .report_element_table import report_element_table_link
-from .report_table import report_table_link
-from .simulation_table import simulation_table_link
-from .user_table import user_table_link
-from .versioned_dataset_table import versioned_dataset_table_link
-
-
-class Database:
-    url: str
-
-    _model_table_links: list[TableLink] = []
-
-    def __init__(self, url: str):
-        self.url = url
-        self.engine = self._create_engine()
-        self.session = Session(self.engine)
-
-        for link in [
-            model_table_link,
-            model_version_table_link,
-            dataset_table_link,
-            versioned_dataset_table_link,
-            policy_table_link,
-            dynamic_table_link,
-            parameter_table_link,
-            parameter_value_table_link,
-            baseline_parameter_value_table_link,
-            baseline_variable_table_link,
-            simulation_table_link,
-            aggregate_table_link,
-            user_table_link,
-            report_table_link,
-            report_element_table_link,
-        ]:
-            self.register_table(link)
-
-    def _create_engine(self):
-        from sqlmodel import create_engine
-
-        return create_engine(self.url, echo=False)
-
-    def create_tables(self):
-        """Create all database tables."""
-        SQLModel.metadata.create_all(self.engine)
-
-    def drop_tables(self):
-        """Drop all database tables."""
-        SQLModel.metadata.drop_all(self.engine)
-
-    def reset(self):
-        """Drop and recreate all tables."""
-        self.drop_tables()
-        self.create_tables()
-
-    def __enter__(self):
-        """Context manager entry - creates a session."""
-        self.session = Session(self.engine)
-        return self.session
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        """Context manager exit - closes the session."""
-        if exc_type:
-            self.session.rollback()
-        else:
-            self.session.commit()
-        self.session.close()
-
-    def register_table(self, link: TableLink):
-        self._model_table_links.append(link)
-        # Create the table if not exists
-        link.table_cls.metadata.create_all(self.engine)
-
-    def get(self, model_cls: type, **kwargs):
-        table_link = next(
-            (
-                link
-                for link in self._model_table_links
-                if link.model_cls == model_cls
-            ),
-            None,
-        )
-        if table_link is not None:
-            return table_link.get(self, **kwargs)
-
-    def set(self, object: Any, commit: bool = True):
-        table_link = next(
-            (
-                link
-                for link in self._model_table_links
-                if link.model_cls is type(object)
-            ),
-            None,
-        )
-        if table_link is not None:
-            table_link.set(self, object, commit=commit)
-
-    def register_model_version(self, model_version):
-        """Register a model version with its model and seed objects.
-        This replaces all existing parameters, baseline parameter values,
-        and baseline variables for this model version."""
-        # Add or update the model directly to avoid conflicts
-        from policyengine.utils.compress import compress_data
-
-        from .baseline_parameter_value_table import BaselineParameterValueTable
-        from .baseline_variable_table import BaselineVariableTable
-        from .model_table import ModelTable
-        from .model_version_table import ModelVersionTable
-        from .parameter_table import ParameterTable
-
-        existing_model = (
-            self.session.query(ModelTable)
-            .filter(ModelTable.id == model_version.model.id)
-            .first()
-        )
-        if not existing_model:
-            model_table = ModelTable(
-                id=model_version.model.id,
-                name=model_version.model.name,
-                description=model_version.model.description,
-                simulation_function=(
-                    lambda m: compress_data(m.simulation_function)
-                )(model_version.model),
-            )
-            self.session.add(model_table)
-            self.session.flush()
-
-        # Add or update the model version
-        existing_version = (
-            self.session.query(ModelVersionTable)
-            .filter(ModelVersionTable.id == model_version.id)
-            .first()
-        )
-        if not existing_version:
-            version_table = ModelVersionTable(
-                id=model_version.id,
-                model_id=model_version.model.id,
-                version=model_version.version,
-                description=model_version.description,
-                created_at=model_version.created_at,
-            )
-            self.session.add(version_table)
-            self.session.flush()
-
-        # Get seed objects from the model
-        seed_objects = model_version.model.create_seed_objects(model_version)
-
-        # Delete ALL existing seed data for this model (not just this version)
-        # This ensures we start fresh with the new version's data
-        # Order matters due to foreign key constraints
-
-        # First delete baseline parameter values (they reference parameters)
-        self.session.query(BaselineParameterValueTable).filter(
-            BaselineParameterValueTable.model_id == model_version.model.id
-        ).delete()
-
-        # Then delete baseline variables for this model
-        self.session.query(BaselineVariableTable).filter(
-            BaselineVariableTable.model_id == model_version.model.id
-        ).delete()
-
-        # Finally delete all parameters for this model
-        self.session.query(ParameterTable).filter(
-            ParameterTable.model_id == model_version.model.id
-        ).delete()
-
-        self.session.commit()
-
-        # Add all parameters first
-        for parameter in seed_objects.parameters:
-            # We need to add directly to session to avoid the autoflush issue
-            from .parameter_table import ParameterTable
-
-            param_table = ParameterTable(
-                id=parameter.id,
-                model_id=parameter.model.id,  # Now required as part of composite key
-                description=parameter.description,
-                data_type=parameter.data_type.__name__
-                if parameter.data_type
-                else None,
-            )
-            self.session.add(param_table)
-
-        # Flush parameters to database so they exist for foreign key constraints
-        self.session.flush()
-
-        # Add all baseline parameter values
-        for baseline_param_value in seed_objects.baseline_parameter_values:
-            import math
-            from uuid import uuid4
-
-            from .baseline_parameter_value_table import (
-                BaselineParameterValueTable,
-            )
-
-            # Handle special float values that JSON doesn't support
-            value = baseline_param_value.value
-            if isinstance(value, float):
-                if math.isinf(value):
-                    value = "Infinity" if value > 0 else "-Infinity"
-                elif math.isnan(value):
-                    value = "NaN"
-
-            bpv_table = BaselineParameterValueTable(
-                id=str(uuid4()),
-                parameter_id=baseline_param_value.parameter.id,
-                model_id=baseline_param_value.parameter.model.id,  # Add model_id
-                model_version_id=baseline_param_value.model_version.id,
-                value=value,
-                start_date=baseline_param_value.start_date,
-                end_date=baseline_param_value.end_date,
-            )
-            self.session.add(bpv_table)
-
-        # Add all baseline variables
-        for baseline_variable in seed_objects.baseline_variables:
-            from .baseline_variable_table import BaselineVariableTable
-
-            bv_table = BaselineVariableTable(
-                id=baseline_variable.id,
-                model_id=baseline_variable.model_version.model.id,  # Add model_id
-                model_version_id=baseline_variable.model_version.id,
-                entity=baseline_variable.entity,
-                label=baseline_variable.label,
-                description=baseline_variable.description,
-                data_type=(lambda bv: compress_data(bv.data_type))(
-                    baseline_variable
-                )
-                if baseline_variable.data_type
-                else None,
-            )
-            self.session.add(bv_table)
-
-        # Commit everything at once
-        self.session.commit()

policyengine/database/dataset_table.py

@@ -1,41 +0,0 @@
-from uuid import uuid4
-
-from sqlmodel import Field, SQLModel
-
-from policyengine.models import Dataset
-from policyengine.utils.compress import compress_data, decompress_data
-
-from .link import TableLink
-
-
-class DatasetTable(SQLModel, table=True):
-    __tablename__ = "datasets"
-
-    id: str = Field(default_factory=lambda: str(uuid4()), primary_key=True)
-    name: str = Field(nullable=False)
-    description: str | None = Field(default=None)
-    version: str | None = Field(default=None)
-    versioned_dataset_id: str | None = Field(
-        default=None, foreign_key="versioned_datasets.id", ondelete="SET NULL"
-    )
-    year: int | None = Field(default=None)
-    data: bytes | None = Field(default=None)
-    model_id: str | None = Field(
-        default=None, foreign_key="models.id", ondelete="SET NULL"
-    )
-
-
-dataset_table_link = TableLink(
-    model_cls=Dataset,
-    table_cls=DatasetTable,
-    model_to_table_custom_transforms=dict(
-        versioned_dataset_id=lambda d: d.versioned_dataset.id
-        if d.versioned_dataset
-        else None,
-        model_id=lambda d: d.model.id if d.model else None,
-        data=lambda d: compress_data(d.data) if d.data else None,
-    ),
-    table_to_model_custom_transforms=dict(
-        data=lambda b: decompress_data(b) if b else None,
-    ),
-)

policyengine/database/dynamic_table.py

@@ -1,34 +0,0 @@
-from datetime import datetime
-from uuid import uuid4
-
-from sqlmodel import Field, SQLModel
-
-from policyengine.models import Dynamic
-from policyengine.utils.compress import compress_data, decompress_data
-
-from .link import TableLink
-
-
-class DynamicTable(SQLModel, table=True):
-    __tablename__ = "dynamics"
-
-    id: str = Field(default_factory=lambda: str(uuid4()), primary_key=True)
-    name: str = Field(nullable=False)
-    description: str | None = Field(default=None)
-    simulation_modifier: bytes | None = Field(default=None)
-    created_at: datetime = Field(default_factory=datetime.now)
-    updated_at: datetime = Field(default_factory=datetime.now)
-
-
-dynamic_table_link = TableLink(
-    model_cls=Dynamic,
-    table_cls=DynamicTable,
-    model_to_table_custom_transforms=dict(
-        simulation_modifier=lambda d: compress_data(d.simulation_modifier)
-        if d.simulation_modifier
-        else None,
-    ),
-    table_to_model_custom_transforms=dict(
-        simulation_modifier=lambda b: decompress_data(b) if b else None,
-    ),
-)

policyengine/database/link.py

@@ -1,82 +0,0 @@
-from collections.abc import Callable
-from typing import TYPE_CHECKING
-
-from pydantic import BaseModel
-from sqlmodel import SQLModel, select
-
-if TYPE_CHECKING:
-    from .database import Database
-
-
-class TableLink(BaseModel):
-    model_cls: type[BaseModel]
-    table_cls: type[SQLModel]
-    model_to_table_custom_transforms: dict[str, Callable] | None = None
-    table_to_model_custom_transforms: dict[str, Callable] | None = None
-    primary_key: str | tuple[str, ...] = (
-        "id"  # Allow multiple strings in tuple
-    )
-
-    def get(self, database: "Database", **kwargs):
-        statement = select(self.table_cls).filter_by(**kwargs)
-        result = database.session.exec(statement).first()
-        if result is None:
-            return None
-        model_data = result.model_dump()
-        if self.table_to_model_custom_transforms:
-            for (
-                field,
-                transform,
-            ) in self.table_to_model_custom_transforms.items():
-                model_data[field] = transform(getattr(result, field))
-
-        # Only include fields that exist in the model class
-        valid_fields = {
-            field_name for field_name in self.model_cls.__annotations__.keys()
-        }
-        filtered_model_data = {
-            k: v for k, v in model_data.items() if k in valid_fields
-        }
-        return self.model_cls(**filtered_model_data)
-
-    def set(self, database: "Database", obj: BaseModel, commit: bool = True):
-        model_data = obj.model_dump()
-        if self.model_to_table_custom_transforms:
-            for (
-                field,
-                transform,
-            ) in self.model_to_table_custom_transforms.items():
-                model_data[field] = transform(obj)
-        # Only include fields that exist in the table class
-        valid_fields = {
-            field_name for field_name in self.table_cls.__annotations__.keys()
-        }
-        filtered_model_data = {
-            k: v for k, v in model_data.items() if k in valid_fields
-        }
-        table_obj = self.table_cls(**filtered_model_data)
-
-        # Check if already exists using primary key
-        query = select(self.table_cls)
-        if isinstance(self.primary_key, tuple):
-            for key in self.primary_key:
-                query = query.where(
-                    getattr(self.table_cls, key) == getattr(table_obj, key)
-                )
-        else:
-            query = query.where(
-                getattr(self.table_cls, self.primary_key)
-                == getattr(table_obj, self.primary_key)
-            )
-
-        existing = database.session.exec(query).first()
-        if existing:
-            # Update existing record
-            for key, value in filtered_model_data.items():
-                setattr(existing, key, value)
-            database.session.add(existing)
-        else:
-            database.session.add(table_obj)
-
-        if commit:
-            database.session.commit()

policyengine/database/model_table.py

@@ -1,27 +0,0 @@
-from sqlmodel import Field, SQLModel
-
-from policyengine.models import Model
-from policyengine.utils.compress import compress_data, decompress_data
-
-from .link import TableLink
-
-
-class ModelTable(SQLModel, table=True, extend_existing=True):
-    __tablename__ = "models"
-
-    id: str = Field(primary_key=True)
-    name: str = Field(nullable=False)
-    description: str | None = Field(default=None)
-    simulation_function: bytes
-
-
-model_table_link = TableLink(
-    model_cls=Model,
-    table_cls=ModelTable,
-    model_to_table_custom_transforms=dict(
-        simulation_function=lambda m: compress_data(m.simulation_function),
-    ),
-    table_to_model_custom_transforms=dict(
-        simulation_function=lambda b: decompress_data(b),
-    ),
-)

policyengine/database/model_version_table.py

@@ -1,28 +0,0 @@
-from datetime import datetime
-from uuid import uuid4
-
-from sqlmodel import Field, SQLModel
-
-from policyengine.models import ModelVersion
-
-from .link import TableLink
-
-
-class ModelVersionTable(SQLModel, table=True):
-    __tablename__ = "model_versions"
-
-    id: str = Field(default_factory=lambda: str(uuid4()), primary_key=True)
-    model_id: str = Field(foreign_key="models.id", ondelete="CASCADE")
-    version: str = Field(nullable=False)
-    description: str | None = Field(default=None)
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-model_version_table_link = TableLink(
-    model_cls=ModelVersion,
-    table_cls=ModelVersionTable,
-    model_to_table_custom_transforms=dict(
-        model_id=lambda model_version: model_version.model.id,
-    ),
-    table_to_model_custom_transforms={},
-)

policyengine/database/parameter_table.py

@@ -1,31 +0,0 @@
-from sqlmodel import Field, SQLModel
-
-from policyengine.models import Parameter
-
-from .link import TableLink
-
-
-class ParameterTable(SQLModel, table=True):
-    __tablename__ = "parameters"
-    __table_args__ = ({"extend_existing": True},)
-
-    id: str = Field(primary_key=True)  # Parameter name
-    model_id: str = Field(
-        primary_key=True, foreign_key="models.id"
-    )  # Part of composite key
-    description: str | None = Field(default=None)
-    data_type: str | None = Field(nullable=True)  # Data type name
-
-
-parameter_table_link = TableLink(
-    model_cls=Parameter,
-    table_cls=ParameterTable,
-    primary_key=("id", "model_id"),  # Composite primary key
-    model_to_table_custom_transforms=dict(
-        data_type=lambda p: p.data_type.__name__ if p.data_type else None,
-        model_id=lambda p: p.model.id if p.model else None,
-    ),
-    table_to_model_custom_transforms=dict(
-        data_type=lambda t: eval(t.data_type) if t.data_type else None
-    ),
-)

policyengine/database/parameter_value_table.py

@@ -1,62 +0,0 @@
-from datetime import datetime
-from typing import Any
-from uuid import uuid4
-
-from sqlmodel import JSON, Column, Field, SQLModel
-
-from policyengine.models import ParameterValue
-
-from .link import TableLink
-
-
-class ParameterValueTable(SQLModel, table=True):
-    __tablename__ = "parameter_values"
-    __table_args__ = ({"extend_existing": True},)
-
-    id: str = Field(default_factory=lambda: str(uuid4()), primary_key=True)
-    parameter_id: str = Field(nullable=False)  # Part of composite foreign key
-    model_id: str = Field(nullable=False)  # Part of composite foreign key
-    value: Any | None = Field(
-        default=None, sa_column=Column(JSON)
-    )  # JSON field for any type
-    start_date: datetime = Field(nullable=False)
-    end_date: datetime | None = Field(default=None)
-
-
-def transform_value_to_table(pv):
-    """Transform value for storage, handling special float values."""
-    import math
-
-    value = pv.value
-    if isinstance(value, float):
-        if math.isinf(value):
-            return "Infinity" if value > 0 else "-Infinity"
-        elif math.isnan(value):
-            return "NaN"
-    return value
-
-
-def transform_value_from_table(table_row):
-    """Transform value from storage, converting special strings back to floats."""
-    value = table_row.value
-    if value == "Infinity":
-        return float("inf")
-    elif value == "-Infinity":
-        return float("-inf")
-    elif value == "NaN":
-        return float("nan")
-    return value
-
-
-parameter_value_table_link = TableLink(
-    model_cls=ParameterValue,
-    table_cls=ParameterValueTable,
-    model_to_table_custom_transforms=dict(
-        parameter_id=lambda pv: pv.parameter.id,
-        model_id=lambda pv: pv.parameter.model.id,  # Add model_id from parameter
-        value=transform_value_to_table,
-    ),
-    table_to_model_custom_transforms=dict(
-        value=transform_value_from_table,
-    ),
-)

policyengine/database/policy_table.py

@@ -1,34 +0,0 @@
-from datetime import datetime
-from uuid import uuid4
-
-from sqlmodel import Field, SQLModel
-
-from policyengine.models import Policy
-from policyengine.utils.compress import compress_data, decompress_data
-
-from .link import TableLink
-
-
-class PolicyTable(SQLModel, table=True):
-    __tablename__ = "policies"
-
-    id: str = Field(default_factory=lambda: str(uuid4()), primary_key=True)
-    name: str = Field(nullable=False)
-    description: str | None = Field(default=None)
-    simulation_modifier: bytes | None = Field(default=None)
-    created_at: datetime = Field(default_factory=datetime.now)
-    updated_at: datetime = Field(default_factory=datetime.now)
-
-
-policy_table_link = TableLink(
-    model_cls=Policy,
-    table_cls=PolicyTable,
-    model_to_table_custom_transforms=dict(
-        simulation_modifier=lambda p: compress_data(p.simulation_modifier)
-        if p.simulation_modifier
-        else None,
-    ),
-    table_to_model_custom_transforms=dict(
-        simulation_modifier=lambda b: decompress_data(b) if b else None,
-    ),
-)

policyengine/database/report_element_table.py

@@ -1,48 +0,0 @@
-import uuid
-from datetime import datetime
-
-from sqlmodel import Field, SQLModel
-
-from policyengine.models.report_element import ReportElement
-
-from .link import TableLink
-
-
-class ReportElementTable(SQLModel, table=True, extend_existing=True):
-    __tablename__ = "report_elements"
-
-    id: str = Field(
-        primary_key=True, default_factory=lambda: str(uuid.uuid4())
-    )
-    label: str = Field(nullable=False)
-    type: str = Field(nullable=False)  # "chart" or "markdown"
-
-    # Data source
-    data_table: str | None = Field(default=None)  # "aggregates"
-
-    # Chart configuration
-    chart_type: str | None = Field(
-        default=None
-    )  # "bar", "line", "scatter", "area", "pie"
-    x_axis_variable: str | None = Field(default=None)
-    y_axis_variable: str | None = Field(default=None)
-    group_by: str | None = Field(default=None)
-    color_by: str | None = Field(default=None)
-    size_by: str | None = Field(default=None)
-
-    # Markdown specific
-    markdown_content: str | None = Field(default=None)
-
-    # Metadata
-    report_id: str | None = Field(default=None, foreign_key="reports.id")
-    user_id: str | None = Field(default=None, foreign_key="users.id")
-    position: int | None = Field(default=None)
-    visible: bool | None = Field(default=True)
-    created_at: datetime = Field(default_factory=datetime.utcnow)
-    updated_at: datetime = Field(default_factory=datetime.utcnow)
-
-
-report_element_table_link = TableLink(
-    model_cls=ReportElement,
-    table_cls=ReportElementTable,
-)

policyengine/database/report_table.py

@@ -1,24 +0,0 @@
-import uuid
-from datetime import datetime
-
-from sqlmodel import Field, SQLModel
-
-from policyengine.models.report import Report
-
-from .link import TableLink
-
-
-class ReportTable(SQLModel, table=True, extend_existing=True):
-    __tablename__ = "reports"
-
-    id: str = Field(
-        primary_key=True, default_factory=lambda: str(uuid.uuid4())
-    )
-    label: str = Field(nullable=False)
-    created_at: datetime = Field(default_factory=datetime.utcnow)
-
-
-report_table_link = TableLink(
-    model_cls=Report,
-    table_cls=ReportTable,
-)