oxyde-postgres 0.1.0__py3-none-any.whl

oxyde/queries/aggregates.py

@@ -0,0 +1,145 @@
+"""SQL aggregate and scalar functions for SELECT queries.
+
+This module provides Python wrappers for SQL aggregate functions
+used with annotate() and direct aggregate queries.
+
+Aggregate Functions:
+    Count(field="*", distinct=False) → COUNT(*) or COUNT(DISTINCT field)
+    Sum(field)                       → SUM(field)
+    Avg(field)                       → AVG(field)
+    Max(field)                       → MAX(field)
+    Min(field)                       → MIN(field)
+
+Scalar Functions:
+    Concat(*fields, separator="")    → Concatenate strings
+    Coalesce(*fields)                → First non-NULL value
+
+Raw SQL:
+    RawSQL(sql)                      → Inject raw SQL expression
+
+Usage with annotate():
+    # Add computed columns to query results
+    query = User.objects.annotate(
+        post_count=Count("id"),
+        total_score=Sum("score"),
+    )
+
+Usage with group_by():
+    query = Order.objects.annotate(
+        total=Sum("amount"),
+    ).group_by("customer_id")
+
+Direct aggregates (execute immediately):
+    total = await User.objects.filter(active=True).count()
+    avg_age = await User.objects.avg("age")
+    max_score = await User.objects.max("score")
+
+IR Format:
+    Count("id").to_ir() returns:
+    {"func": "count", "field": "id"}
+
+    Count("id", distinct=True).to_ir() returns:
+    {"func": "count", "field": "id", "distinct": True}
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class Aggregate:
+    """Base class for aggregate functions."""
+
+    def __init__(self, field: str, *, distinct: bool = False):
+        self.field = field
+        self.distinct = distinct
+        self.func_name = self.__class__.__name__.lower()
+
+    def to_ir(self) -> dict[str, Any]:
+        """Convert to IR representation."""
+        ir: dict[str, Any] = {
+            "func": self.func_name,
+            "field": self.field,
+        }
+        if self.distinct:
+            ir["distinct"] = True
+        return ir
+
+
+class Count(Aggregate):
+    """COUNT aggregate function."""
+
+    def __init__(self, field: str = "*", *, distinct: bool = False):
+        super().__init__(field, distinct=distinct)
+
+
+class Sum(Aggregate):
+    """SUM aggregate function."""
+
+
+class Avg(Aggregate):
+    """AVG aggregate function."""
+
+
+class Max(Aggregate):
+    """MAX aggregate function."""
+
+
+class Min(Aggregate):
+    """MIN aggregate function."""
+
+
+class Concat(Aggregate):
+    """String concatenation."""
+
+    def __init__(self, *fields: str, separator: str = ""):
+        self.fields = fields
+        self.separator = separator
+        self.func_name = "concat"
+
+    def to_ir(self) -> dict[str, Any]:
+        return {
+            "func": "concat",
+            "fields": list(self.fields),
+            "separator": self.separator,
+        }
+
+
+class Coalesce(Aggregate):
+    """COALESCE function - returns first non-NULL value."""
+
+    def __init__(self, *fields: str):
+        self.fields = fields
+        self.func_name = "coalesce"
+
+    def to_ir(self) -> dict[str, Any]:
+        return {
+            "func": "coalesce",
+            "fields": list(self.fields),
+        }
+
+
+class RawSQL:
+    """Raw SQL expression."""
+
+    def __init__(self, sql: str):
+        self.sql = sql
+
+    def to_ir(self) -> dict[str, Any]:
+        return {
+            "type": "raw",
+            "sql": self.sql,
+        }
+
+
+__all__ = [
+    "Aggregate",
+    "Count",
+    "Sum",
+    "Avg",
+    "Max",
+    "Min",
+    "Concat",
+    "Coalesce",
+    "RawSQL",
+]

oxyde/queries/base.py

@@ -0,0 +1,249 @@
+"""Shared utilities, protocols, and helpers for query builders.
+
+This module provides common functionality used across all query types
+(SELECT, INSERT, UPDATE, DELETE).
+
+Protocol:
+    SupportsExecute: Interface for objects that can execute queries.
+        - AsyncDatabase implements this
+        - AsyncTransaction implements this
+        Any object with async execute(ir: dict) -> bytes works.
+
+Type Variables:
+    TQuery: TypeVar bound to Query for method chaining with correct return types.
+
+Helper Functions:
+    _model_key(model_class) -> str:
+        Returns "{module}.{qualname}" for model identification.
+
+    _resolve_registered_model(key) -> type[Model]:
+        Look up model by key or class name in registry.
+
+    _primary_key_meta(model_class) -> ColumnMeta:
+        Get metadata for primary key field.
+
+    _collect_model_columns(model_class) -> list[(field, column)]:
+        Get all (field_name, db_column) pairs.
+
+    _map_values_to_columns(model_class, values) -> dict:
+        Convert field names to database column names in a dict.
+
+    _resolve_execution_client(using, client) -> SupportsExecute:
+        Resolve connection from using="alias" or client object.
+        Checks active transaction first.
+
+    _resolve_pool_name(using, client) -> str:
+        Get pool name for schema/explain operations.
+
+Caching:
+    _TYPE_ADAPTER_CACHE: Thread-safe cache for Pydantic TypeAdapter instances.
+    Used for efficient list[Model] validation in fetch_models().
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Protocol,
+    TypeVar,
+    runtime_checkable,
+)
+
+from pydantic import TypeAdapter
+
+from oxyde.core import wrapper as _core
+from oxyde.db.pool import AsyncDatabase
+from oxyde.db.registry import get_connection
+from oxyde.db.transaction import AsyncTransaction, get_active_transaction
+from oxyde.exceptions import (
+    CheckViolationError,
+    FieldLookupError,
+    ForeignKeyViolationError,
+    IntegrityError,
+    ManagerError,
+    NotNullViolationError,
+    UniqueViolationError,
+)
+from oxyde.models.metadata import ColumnMeta
+from oxyde.models.registry import registered_tables
+from oxyde.models.serializers import _get_virtual_fields
+
+if TYPE_CHECKING:
+    from oxyde.models.base import Model
+    from oxyde.queries.select import Query
+
+# Global cache for TypeAdapter instances (thread-safe)
+_TYPE_ADAPTER_CACHE: dict[type, TypeAdapter] = {}
+_TYPE_ADAPTER_LOCK = threading.Lock()
+
+# Core exception → oxyde exception, most specific first (core classes form
+# the same hierarchy, so IntegrityError must be checked last).
+_CORE_ERROR_MAP: tuple[tuple[type[BaseException], type[IntegrityError]], ...] = (
+    (_core.UniqueViolationError, UniqueViolationError),
+    (_core.ForeignKeyViolationError, ForeignKeyViolationError),
+    (_core.NotNullViolationError, NotNullViolationError),
+    (_core.CheckViolationError, CheckViolationError),
+    (_core.IntegrityError, IntegrityError),
+)
+
+
+def _translate_core_error(exc: BaseException) -> ManagerError | None:
+    """Map a Rust-core exception onto the oxyde hierarchy, if it is one.
+
+    Classification is done by the driver (sqlx ErrorKind) — no message-text
+    matching happens anywhere on the Python side.
+    """
+    for core_cls, oxyde_cls in _CORE_ERROR_MAP:
+        if isinstance(exc, core_cls):
+            return oxyde_cls(str(exc))
+    return None
+
+
+def _model_key(model_class: type[Model]) -> str:
+    """Return fully qualified identifier for a model."""
+    return f"{model_class.__module__}.{model_class.__qualname__}"
+
+
+@runtime_checkable
+class SupportsExecute(Protocol):
+    """Protocol for objects that can execute queries."""
+
+    async def execute(self, ir: dict[str, Any]) -> bytes: ...
+
+
+# TypeVar for query builders (covariant for subclassing)
+TQuery = TypeVar("TQuery", bound="Query")
+
+
+def _resolve_registered_model(model_key: str) -> type[Model]:
+    """Resolve a model by its fully qualified key or simple class name."""
+    tables = registered_tables()
+    # Try exact match first
+    model = tables.get(model_key)
+    if model is not None:
+        return model
+    # Fallback: search by simple class name (for forward refs and test classes)
+    for key, table_model in tables.items():
+        if key.endswith(f".{model_key}") or table_model.__name__ == model_key:
+            return table_model
+    raise FieldLookupError(f"Related model '{model_key}' is not registered")
+
+
+def _primary_key_meta(model_class: type[Model]) -> ColumnMeta:
+    """Get primary key metadata from model (uses cached pk_field)."""
+    pk = model_class._db_meta.pk_field
+    if pk is not None:
+        return model_class._db_meta.field_metadata[pk]
+    raise FieldLookupError(f"{model_class.__name__} has no primary key field")
+
+
+def _build_column_types(model_class: type[Model]) -> dict[str, dict] | None:
+    """Get cached column_types mapping (ColumnTypeSpec dicts) from model metadata."""
+    return model_class._db_meta.column_types
+
+
+def _build_column_db_types(model_class: type[Model]) -> dict[str, str] | None:
+    """Get cached db_type overrides by column from model metadata."""
+    return model_class._db_meta.column_db_types
+
+
+def _collect_model_columns(model_class: type[Model]) -> list[tuple[str, str]]:
+    """Collect all (field_name, db_column) pairs from model.
+
+    Excludes virtual fields (reverse FK, M2M, FK model fields) using
+    _get_virtual_fields() as the single source of truth.
+    """
+    virtual = _get_virtual_fields(model_class)
+    return [
+        (meta.name, meta.db_column)
+        for meta in model_class._db_meta.field_metadata.values()
+        if meta.name not in virtual
+    ]
+
+
+def _map_values_to_columns(
+    model_class: type[Model],
+    values: dict[str, Any],
+) -> dict[str, Any]:
+    """Map field names to database column names."""
+    metadata = model_class._db_meta.field_metadata
+    mapped: dict[str, Any] = {}
+    for key, value in values.items():
+        meta = metadata.get(key)
+        column = meta.db_column if meta else key
+        mapped[column] = value
+    return mapped
+
+
+async def _resolve_execution_client(
+    using: str | None,
+    client: SupportsExecute | None,
+) -> SupportsExecute:
+    """Resolve the execution client from using/client parameters."""
+    if client is not None and using is not None:
+        raise ManagerError("Provide either 'client' or 'using', not both")
+    if client is not None:
+        return client
+
+    alias = using or "default"
+    active_tx = get_active_transaction(alias)
+    if active_tx is not None:
+        return active_tx
+    return await get_connection(alias)
+
+
+def _resolve_pool_name(
+    using: str | None,
+    client: SupportsExecute | None,
+) -> str:
+    """
+    Resolve pool name from using or client parameters.
+
+    Args:
+        using: Database alias
+        client: Database client (AsyncDatabase or AsyncTransaction)
+
+    Returns:
+        Pool name string
+
+    Raises:
+        ManagerError: If both using and client are provided
+    """
+    if client is not None and using is not None:
+        raise ManagerError("Provide either 'client' or 'using', not both")
+
+    if client is not None:
+        # Try to get pool name from client
+        if isinstance(client, AsyncDatabase):
+            return client.name
+        elif isinstance(client, AsyncTransaction):
+            return client._database.name
+        else:
+            # Unknown client type, try to get name attribute
+            if hasattr(client, "name"):
+                return str(client.name)
+            elif hasattr(client, "_database") and hasattr(client._database, "name"):
+                return str(client._database.name)
+            else:
+                ctype = type(client).__name__
+                raise ManagerError(f"Cannot determine pool name from client: {ctype}")
+
+    # Use alias
+    return using or "default"
+
+
+__all__ = [
+    "SupportsExecute",
+    "TQuery",
+    "_model_key",
+    "_resolve_registered_model",
+    "_primary_key_meta",
+    "_collect_model_columns",
+    "_map_values_to_columns",
+    "_resolve_execution_client",
+    "_resolve_pool_name",
+    "_TYPE_ADAPTER_CACHE",
+    "_TYPE_ADAPTER_LOCK",
+]

oxyde/queries/conditions.py

@@ -0,0 +1,92 @@
+"""Condition class representing a single WHERE clause comparison.
+
+Condition is the leaf node of the filter tree. It represents a single
+comparison like "age >= 18" or "name LIKE '%alice%'".
+
+Attributes:
+    field: Python field name (e.g., "age")
+    operator: SQL operator ("=", ">", "LIKE", "IN", "IS NULL", etc.)
+    value: Comparison value (can be None for IS NULL)
+    column: Database column name (may differ from field if db_column set)
+    escape: Optional LIKE escape character for pattern filters
+
+Operators:
+    "="        → Equality
+    ">", ">="  → Greater than
+    "<", "<="  → Less than
+    "LIKE"     → Pattern match (case-sensitive)
+    "ILIKE"    → Pattern match (case-insensitive, PostgreSQL)
+    "IN"       → Value in list
+    "BETWEEN"  → Range (value is [low, high])
+    "IS NULL"  → NULL check
+    "IS NOT NULL" → NOT NULL check
+    "ARRAY_CONTAINS"     → Array holds every given element
+    "ARRAY_CONTAINED_BY" → Every stored element is among the given ones
+    "ARRAY_OVERLAP"      → At least one element in common
+
+Creation:
+    Conditions are typically created via Field comparison operators:
+        User.name == "Alice"  → Condition("name", "=", "Alice")
+        User.age >= 18        → Condition("age", ">=", 18)
+
+    Or via lookups module:
+        _build_lookup_conditions(User, "name", "contains", "alice", meta)
+        → Condition("name", "LIKE", "%alice%")
+
+IR Format:
+    to_ir() returns dict for Rust codec. For a LIKE condition with an
+    escape character:
+    {
+        "type": "condition",
+        "field": "name",
+        "op": "LIKE",
+        "value": "A\\_%",
+        "column": "name",
+        "escape": "\\"
+    }
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from oxyde.core import ir
+from oxyde.core.ir import FilterNode
+from oxyde.core.types import serialize_value
+
+
+class Condition:
+    """Represents a filter condition."""
+
+    def __init__(
+        self,
+        field: str,
+        operator: str,
+        value: Any,
+        *,
+        column: str | None = None,
+        escape: str | None = None,
+    ):
+        self.field = field
+        self.operator = operator
+        self.value = value
+        self.column = column or field
+        self.escape = escape
+
+    def to_ir(self) -> FilterNode:
+        """Convert to IR format expected by the Rust codec.
+
+        Values are serialized via TYPE_REGISTRY so that date, UUID, Decimal
+        etc. are converted to msgpack-safe representations before hitting
+        the IR dict.
+        """
+        return ir.filter_condition(
+            field=self.field,
+            operator=self.operator,
+            value=serialize_value(self.value),
+            column=self.column,
+            escape=self.escape,
+        )
+
+
+__all__ = ["Condition"]

oxyde/queries/expressions.py

@@ -0,0 +1,171 @@
+"""F-expressions for database-side field references and arithmetic.
+
+F() allows referencing database columns in UPDATE statements for atomic
+operations that execute entirely in the database, avoiding race conditions.
+
+Basic Usage:
+    F("price")           → Reference column "price"
+    F("price") * 1.1     → price * 1.1
+    F("count") + 1       → count + 1
+    F("a") + F("b")      → a + b
+
+Supported Operations:
+    F("x") + y   → Addition
+    F("x") - y   → Subtraction
+    F("x") * y   → Multiplication
+    F("x") / y   → Division
+    -F("x")      → Negation
+
+    Operands can be: F(), numbers, or other _Expression objects.
+
+Use Cases:
+    # Atomic increment (no race condition)
+    await User.objects.filter(id=1).update(login_count=F("login_count") + 1)
+
+    # Percentage increase
+    await Product.objects.filter(category="sale").update(price=F("price") * 0.9)
+
+    # Combine columns
+    await Order.objects.update(total=F("price") * F("quantity"))
+
+Internal Structure:
+    F("name") creates:
+        _Expression({"type": "column", "name": "name"})
+
+    F("x") + 5 creates:
+        _Expression({
+            "type": "op",
+            "op": "add",
+            "lhs": {"type": "column", "name": "x"},
+            "rhs": {"type": "value", "value": 5}
+        })
+
+    _serialize_value_for_ir() wraps expressions in {"__expr__": ...}
+    for the Rust IR parser to recognize them. Literal nodes carry an
+    optional value_type ColumnTypeSpec dict ({"kind": "decimal"}).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from oxyde.core.column_types import spec_for_literal
+from oxyde.core.types import serialize_value as _serialize_type_value
+
+
+class _Expression:
+    """Represents an arithmetic expression usable in mutation statements."""
+
+    __slots__ = ("_expr",)
+
+    def __init__(self, expr: dict[str, Any]):
+        self._expr = expr
+
+    def _binary(self, op: str, other: Any) -> _Expression:
+        rhs = _coerce_expression(other)
+        return _Expression(
+            {
+                "type": "op",
+                "op": op,
+                "lhs": self._expr,
+                "rhs": rhs._expr,
+            }
+        )
+
+    def __add__(self, other: Any) -> _Expression:
+        return self._binary("add", other)
+
+    def __radd__(self, other: Any) -> _Expression:
+        return _coerce_expression(other)._binary("add", self)
+
+    def __sub__(self, other: Any) -> _Expression:
+        return self._binary("sub", other)
+
+    def __rsub__(self, other: Any) -> _Expression:
+        return _coerce_expression(other)._binary("sub", self)
+
+    def __mul__(self, other: Any) -> _Expression:
+        return self._binary("mul", other)
+
+    def __rmul__(self, other: Any) -> _Expression:
+        return _coerce_expression(other)._binary("mul", self)
+
+    def __truediv__(self, other: Any) -> _Expression:
+        return self._binary("div", other)
+
+    def __rtruediv__(self, other: Any) -> _Expression:
+        return _coerce_expression(other)._binary("div", self)
+
+    def __neg__(self) -> _Expression:
+        return _Expression({"type": "neg", "expr": self._expr})
+
+
+def _coerce_expression(value: Any) -> _Expression:
+    """Coerce a value to an _Expression.
+
+    For known Python types, attach a `value_type` ColumnTypeSpec dict.
+    The Rust side uses it to recover typed bindings (Decimal, UUID,
+    datetime, ...) that msgpack-encode as strings via `serialize_value`.
+    """
+    if isinstance(value, _Expression):
+        return value
+    if isinstance(value, F):
+        return value._expression
+    node: dict[str, Any] = {"type": "value", "value": value}
+    spec = spec_for_literal(type(value))
+    if spec is not None:
+        node["value_type"] = spec
+    return _Expression(node)
+
+
+class F:
+    """Field expression helper for arithmetic updates."""
+
+    __slots__ = ("name", "_expression")
+
+    def __init__(self, name: str):
+        self.name = name
+        self._expression = _Expression({"type": "column", "name": name})
+
+    def __add__(self, other: Any) -> _Expression:
+        return self._expression.__add__(other)
+
+    def __radd__(self, other: Any) -> _Expression:
+        return self._expression.__radd__(other)
+
+    def __sub__(self, other: Any) -> _Expression:
+        return self._expression.__sub__(other)
+
+    def __rsub__(self, other: Any) -> _Expression:
+        return self._expression.__rsub__(other)
+
+    def __mul__(self, other: Any) -> _Expression:
+        return self._expression.__mul__(other)
+
+    def __rmul__(self, other: Any) -> _Expression:
+        return self._expression.__rmul__(other)
+
+    def __truediv__(self, other: Any) -> _Expression:
+        return self._expression.__truediv__(other)
+
+    def __rtruediv__(self, other: Any) -> _Expression:
+        return self._expression.__rtruediv__(other)
+
+    def __neg__(self) -> _Expression:
+        return -self._expression
+
+
+def _serialize_value_for_ir(value: Any) -> Any:
+    """Serialize values for IR, handling expressions and non-msgpack types."""
+    if isinstance(value, _Expression):
+        return {"__expr__": value._expr}
+    if isinstance(value, F):
+        return {"__expr__": value._expression._expr}
+    if isinstance(value, list):
+        return [_serialize_value_for_ir(item) for item in value]
+    if isinstance(value, dict):
+        return {key: _serialize_value_for_ir(val) for key, val in value.items()}
+    return _serialize_type_value(value)
+
+
+__all__ = ["F", "_Expression", "_coerce_expression", "_serialize_value_for_ir"]