Session Management

class Session(SQLAlchemySession):
    """
    Enhanced SQLAlchemy Session with improved typing for SQLModel.
    
    Provides the same functionality as SQLAlchemy Session but with better
    type inference when executing SQLModel select statements.
    """
    
    def exec(
        self,
        statement: Select[_TSelectParam],
        *,
        params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
        execution_options: Mapping[str, Any] = {},
        bind_arguments: Optional[Dict[str, Any]] = None,
        _parent_execute_state: Optional[Any] = None,
        _add_event: Optional[Any] = None,
    ) -> TupleResult[_TSelectParam]:
        """
        Execute a select statement and return results with proper typing.
        
        Parameters:
            statement: SQLModel Select statement to execute
            params: Parameters for the statement
            execution_options: Execution-specific options
            bind_arguments: Bind-specific arguments
            
        Returns:
            TupleResult for regular selects, ScalarResult for scalar selects
        """
    
    def exec(
        self,
        statement: SelectOfScalar[_TSelectParam],
        *,
        params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
        execution_options: Mapping[str, Any] = {},
        bind_arguments: Optional[Dict[str, Any]] = None,
        _parent_execute_state: Optional[Any] = None,
        _add_event: Optional[Any] = None,
    ) -> ScalarResult[_TSelectParam]:
        """Execute a scalar select statement."""
    
    def exec(
        self,
        statement: Executable,
        *,
        params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
        execution_options: Mapping[str, Any] = {},
        bind_arguments: Optional[Dict[str, Any]] = None,
        _parent_execute_state: Optional[Any] = None,
        _add_event: Optional[Any] = None,
    ) -> Result[Any]:
        """Execute any other executable statement."""
    
    def execute(
        self,
        statement: Executable,
        params: Optional[Any] = None,
        *,
        execution_options: Any = {},
        bind_arguments: Optional[Dict[str, Any]] = None,
        _parent_execute_state: Optional[Any] = None,
        _add_event: Optional[Any] = None,
    ) -> Result[Any]:
        """
        🚨 You probably want to use `session.exec()` instead of `session.execute()`.

        This is the original SQLAlchemy `session.execute()` method that returns objects
        of type `Row`, and that you have to call `scalars()` to get the model objects.
        """
    
    def query(self, *entities: Any, **kwargs: Any) -> Any:
        """
        🚨 You probably want to use `session.exec()` instead of `session.query()`.

        `session.exec()` is SQLModel's own short version with increased type
        annotations.
        """

from sqlmodel import Session, create_engine, select

# Create engine and session
engine = create_engine("sqlite:///database.db")

# Basic query execution
with Session(engine) as session:
    # Type-safe query execution
    statement = select(Hero).where(Hero.name == "Spider-Boy")
    result = session.exec(statement)
    hero = result.first()  # Returns Hero | None with proper typing
    
    # Scalar queries
    count_statement = select(func.count(Hero.id))
    count_result = session.exec(count_statement)
    count = count_result.one()  # Returns int with proper typing

# Transaction management
with Session(engine) as session:
    # Create new record
    hero = Hero(name="New Hero", secret_name="Secret")
    session.add(hero)
    
    # Update existing record
    statement = select(Hero).where(Hero.id == 1)
    existing_hero = session.exec(statement).first()
    if existing_hero:
        existing_hero.age = 25
        session.add(existing_hero)
    
    # Commit all changes
    session.commit()
    
    # Session automatically closes and rolls back on exception

# Repository pattern with type safety
class HeroRepository:
    def __init__(self, session: Session):
        self.session = session
    
    def get_by_id(self, hero_id: int) -> Optional[Hero]:
        statement = select(Hero).where(Hero.id == hero_id)
        result = self.session.exec(statement)
        return result.first()  # Typed as Optional[Hero]
    
    def get_all(self) -> List[Hero]:
        statement = select(Hero)
        result = self.session.exec(statement)
        return result.all()  # Typed as List[Hero]
    
    def create(self, hero: Hero) -> Hero:
        self.session.add(hero)
        self.session.commit()
        self.session.refresh(hero)
        return hero

# Dependency injection (FastAPI example)
def get_session():
    with Session(engine) as session:
        yield session

@app.get("/heroes/{hero_id}")
def get_hero(hero_id: int, session: Session = Depends(get_session)) -> Hero:
    statement = select(Hero).where(Hero.id == hero_id)
    result = session.exec(statement)
    hero = result.first()
    if not hero:
        raise HTTPException(status_code=404, detail="Hero not found")
    return hero  # Automatically serialized as JSON by FastAPI