SQL Operations

def select(*entities: Any) -> Union[Select, SelectOfScalar]:
    """
    Create a SELECT statement with proper typing based on selected entities.
    
    Parameters:
        *entities: Model classes, columns, or expressions to select
        
    Returns:
        Select object for multiple columns/entities
        SelectOfScalar object for single scalar values
    """

class Select(SelectBase[_T]):
    """
    Type-safe SELECT statement for multiple columns or entities.
    
    Provides methods for building complex queries with proper type inference.
    """

class SelectOfScalar(SelectBase[_T]):
    """
    SELECT statement for scalar (single value) results.
    
    Used when selecting single columns or aggregate functions.
    """

# Select entire model
statement = select(Hero)  # Returns Select[Hero]
heroes = session.exec(statement).all()  # List[Hero]

# Select specific columns
statement = select(Hero.name, Hero.age)  # Returns Select[Tuple[str, Optional[int]]]
results = session.exec(statement).all()  # List[Tuple[str, Optional[int]]]

# Select scalar values
statement = select(Hero.name)  # Returns SelectOfScalar[str]
names = session.exec(statement).all()  # List[str]

# Aggregate functions
statement = select(func.count(Hero.id))  # Returns SelectOfScalar[int]
count = session.exec(statement).one()  # int

# Complex queries with joins
statement = select(Hero, Team).join(Team)
results = session.exec(statement).all()  # List[Tuple[Hero, Team]]

def and_(*clauses) -> BooleanClauseList:
    """
    Create an AND expression from multiple conditions.
    
    Parameters:
        *clauses: Boolean expressions to combine with AND
        
    Returns:
        BooleanClauseList representing the AND operation
    """

def or_(*clauses) -> BooleanClauseList:
    """
    Create an OR expression from multiple conditions.
    
    Parameters:
        *clauses: Boolean expressions to combine with OR
        
    Returns:
        BooleanClauseList representing the OR operation
    """

def not_(clause) -> UnaryExpression:
    """
    Create a NOT expression to negate a condition.
    
    Parameters:
        clause: Boolean expression to negate
        
    Returns:
        UnaryExpression representing the NOT operation
    """

# AND conditions
statement = select(Hero).where(
    and_(Hero.age > 18, Hero.name.contains("Spider"))
)

# OR conditions  
statement = select(Hero).where(
    or_(Hero.name == "Spider-Boy", Hero.name == "Spider-Man")
)

# NOT conditions
statement = select(Hero).where(
    not_(Hero.secret_name.contains("secret"))
)

# Complex combinations
statement = select(Hero).where(
    and_(
        Hero.age > 18,
        or_(Hero.name.contains("Spider"), Hero.name.contains("Super"))
    )
)

def between(expr, lower_bound, upper_bound) -> BinaryExpression:
    """
    Create a BETWEEN expression for range queries.
    
    Parameters:
        expr: Column or expression to test
        lower_bound: Lower bound of the range
        upper_bound: Upper bound of the range
        
    Returns:
        BinaryExpression for the BETWEEN operation
    """

def all_(expr) -> CollectionAggregate[bool]:
    """
    Create an ALL expression for subquery comparisons.
    
    Parameters:
        expr: Expression or subquery to compare against
        
    Returns:
        CollectionAggregate for ALL comparison
    """

def any_(expr) -> CollectionAggregate[bool]:
    """
    Create an ANY expression for subquery comparisons.
    
    Parameters:
        expr: Expression or subquery to compare against
        
    Returns:
        CollectionAggregate for ANY comparison
    """

# BETWEEN queries
statement = select(Hero).where(between(Hero.age, 18, 65))

# Subquery with ANY
subquery = select(Team.id).where(Team.name.contains("Avengers"))
statement = select(Hero).where(Hero.team_id.in_(subquery))

# Using ANY with arrays (PostgreSQL)
statement = select(Hero).where(Hero.id == any_([1, 2, 3, 4]))

def asc(column) -> UnaryExpression:
    """
    Create an ascending order expression.
    
    Parameters:
        column: Column or expression to sort by
        
    Returns:
        UnaryExpression for ascending order
    """

def desc(column) -> UnaryExpression:
    """
    Create a descending order expression.
    
    Parameters:
        column: Column or expression to sort by
        
    Returns:
        UnaryExpression for descending order
    """

def nulls_first(column) -> UnaryExpression:
    """
    Specify that NULL values should appear first in ordering.
    
    Parameters:
        column: Column expression to modify
        
    Returns:
        UnaryExpression with NULLS FIRST ordering
    """

def nulls_last(column) -> UnaryExpression:
    """
    Specify that NULL values should appear last in ordering.
    
    Parameters:
        column: Column expression to modify
        
    Returns:
        UnaryExpression with NULLS LAST ordering
    """

# Basic ordering
statement = select(Hero).order_by(asc(Hero.name))
statement = select(Hero).order_by(desc(Hero.age))

# Multiple columns
statement = select(Hero).order_by(asc(Hero.team_id), desc(Hero.age))

# NULL handling
statement = select(Hero).order_by(nulls_last(asc(Hero.age)))

def cast(expression, type_) -> Cast:
    """
    Create a CAST expression to convert types.
    
    Parameters:
        expression: Expression to cast
        type_: Target SQLAlchemy type
        
    Returns:
        Cast expression
    """

def type_coerce(expression, type_) -> TypeCoerce:
    """
    Coerce expression to a specific type for Python typing.
    
    Parameters:
        expression: Expression to coerce
        type_: Target SQLAlchemy type
        
    Returns:
        TypeCoerce expression
    """

def distinct(expr) -> UnaryExpression:
    """
    Create a DISTINCT expression.
    
    Parameters:
        expr: Expression to make distinct
        
    Returns:
        UnaryExpression for DISTINCT operation
    """

# Type casting
statement = select(cast(Hero.age, String))

# Distinct values
statement = select(distinct(Hero.team_id))

# Type coercion for better typing
from sqlalchemy import JSON
json_data = type_coerce(Hero.metadata, JSON)
statement = select(json_data["key"])

def case(*whens, value=None, else_=None) -> Case:
    """
    Create a CASE expression for conditional logic.
    
    Parameters:
        *whens: Tuples of (condition, result) pairs
        value: Column to test against (for simple CASE)
        else_: Default value if no conditions match
        
    Returns:
        Case expression
    """

# Conditional selection
hero_status = case(
    (Hero.age < 18, "Young"),
    (Hero.age < 65, "Adult"),
    else_="Senior"
)
statement = select(Hero.name, hero_status.label("status"))

# Simple CASE with value
team_type = case(
    ("Avengers", "Heroes"),
    ("X-Men", "Mutants"),
    value=Team.name,
    else_="Other"
)

def over(element, partition_by=None, order_by=None) -> Over:
    """
    Create an OVER clause for window functions.
    
    Parameters:
        element: Function or expression to apply window to
        partition_by: Columns to partition by
        order_by: Columns to order by within partitions
        
    Returns:
        Over expression for window functions
    """

def within_group(element, *order_by) -> WithinGroup:
    """
    Create a WITHIN GROUP clause for ordered-set aggregates.
    
    Parameters:
        element: Aggregate function
        *order_by: Ordering expressions
        
    Returns:
        WithinGroup expression
    """

# Window functions
from sqlalchemy import func

# Row numbers within partitions
row_num = func.row_number().over(
    partition_by=Hero.team_id,
    order_by=desc(Hero.age)
)
statement = select(Hero.name, row_num.label("rank"))

# Running totals
running_count = func.count().over(
    order_by=Hero.id,
    rows=(None, 0)  # Unbounded preceding to current row
)

def col(column_expression) -> Mapped:
    """
    Create a column reference with proper typing.
    
    Parameters:
        column_expression: Column to reference
        
    Returns:
        Mapped column reference
    """

def collate(expression, collation) -> BinaryExpression:
    """
    Apply a collation to a string expression.
    
    Parameters:
        expression: String expression
        collation: Collation name
        
    Returns:
        BinaryExpression with collation
    """

def extract(field, expr) -> Extract:
    """
    Extract a date/time component from a datetime expression.
    
    Parameters:
        field: Component to extract ("year", "month", "day", etc.)
        expr: Datetime expression
        
    Returns:
        Extract expression
    """

# Column references
hero_name = col(Hero.name)
statement = select(hero_name).where(hero_name.like("Spider%"))

# Date/time extraction
year_born = extract("year", Hero.birth_date)
statement = select(Hero.name, year_born.label("birth_year"))

# String collation
statement = select(Hero).where(
    collate(Hero.name, "NOCASE") == "spider-man"
)