oxyde-postgres 0.1.0__py3-none-any.whl

oxyde/queries/mixins/mutation.py

@@ -0,0 +1,508 @@
+"""Mutation mixin for query building."""
+
+from __future__ import annotations
+
+import warnings
+from collections.abc import Iterable
+from enum import Enum
+from functools import cache
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generic,
+    Literal,
+    TypeVar,
+    get_args,
+    get_origin,
+    overload,
+)
+
+from pydantic import TypeAdapter
+
+from oxyde._msgpack import msgpack
+from oxyde.core import ir
+from oxyde.exceptions import ManagerError
+from oxyde.models.serializers import (
+    _dump_insert_data,
+    _normalize_instance,
+)
+from oxyde.models.utils import _unpack_annotated, _unwrap_optional
+from oxyde.queries.base import (
+    SupportsExecute,
+    _build_column_db_types,
+    _build_column_types,
+    _map_values_to_columns,
+    _model_key,
+    _resolve_execution_client,
+    _translate_core_error,
+)
+from oxyde.queries.expressions import F, _Expression, _serialize_value_for_ir
+from oxyde.queries.insert import InsertQuery
+
+if TYPE_CHECKING:
+    from oxyde.models.base import Model
+
+
+TModel = TypeVar("TModel", bound="Model")
+
+
+def _hydrate_models(
+    model_class: type[TModel], columns: list[str], rows: list[list[Any]]
+) -> list[TModel]:
+    """Convert raw columns + rows into validated model instances."""
+    rmap = model_class._db_meta.reverse_column_map
+    field_columns = [rmap.get(c, c) for c in columns] if rmap else columns
+    return [model_class.model_validate(dict(zip(field_columns, row))) for row in rows]
+
+
+def _decode_returning_models(
+    model_class: type[TModel], result: dict[str, Any]
+) -> list[TModel]:
+    """Decode mutation-returning result {columns, rows} into model instances."""
+    rows = result.get("rows", [])
+    if not rows:
+        return []
+    return _hydrate_models(model_class, result.get("columns", []), rows)
+
+
+def _decode_columnar_models(
+    model_class: type[TModel], result: list[Any]
+) -> list[TModel]:
+    """Decode columnar SELECT result [columns, rows] into model instances."""
+    if len(result) < 2 or not result[1]:
+        return []
+    return _hydrate_models(model_class, result[0], result[1])
+
+
+def _is_enum_annotation(annotation: Any) -> bool:
+    annotation, _ = _unpack_annotated(annotation)
+    annotation, _ = _unwrap_optional(annotation)
+    origin = get_origin(annotation)
+    if origin is list:
+        args = get_args(annotation)
+        return bool(args) and _is_enum_annotation(args[0])
+    if origin is not None:
+        return any(
+            arg is not type(None) and _is_enum_annotation(arg)
+            for arg in get_args(annotation)
+        )
+    return isinstance(annotation, type) and issubclass(annotation, Enum)
+
+
+@cache
+def _enum_type_adapter(python_type: Any) -> TypeAdapter:
+    return TypeAdapter(python_type)
+
+
+def _validate_enum_update_values(
+    model_class: type[Model],
+    values: dict[str, Any],
+) -> dict[str, Any]:
+    validated = dict(values)
+    metadata = model_class._db_meta.field_metadata
+    for field, value in values.items():
+        if isinstance(value, (F, _Expression)):
+            continue
+        meta = metadata.get(field)
+        if meta is not None and _is_enum_annotation(meta.python_type):
+            if value is None and meta.nullable:
+                validated[field] = None
+            else:
+                adapter = _enum_type_adapter(meta.python_type)
+                validated[field] = adapter.validate_python(value)
+    return validated
+
+
+class MutationMixin(Generic[TModel]):
+    """Mixin providing data mutation capabilities."""
+
+    # These attributes are defined in the base Query class
+    model_class: type[TModel]
+
+    def _build_filter_tree(self) -> ir.FilterNode | None:
+        """Must be implemented by FilteringMixin."""
+        raise NotImplementedError
+
+    async def increment(
+        self,
+        field: str,
+        by: int | float = 1,
+        *,
+        using: str | None = None,
+        client: SupportsExecute | None = None,
+    ) -> int:
+        """
+        Atomically increment a field value.
+
+        Args:
+            field: Field name to increment
+            by: Amount to increment (default 1)
+            using: Database alias
+            client: Optional database client
+
+        Returns:
+            Number of affected rows
+
+        Examples:
+            await Post.objects.filter(id=42).increment("views", by=1)
+            await User.objects.filter(is_active=True).increment("login_count")
+        """
+        exec_client = await _resolve_execution_client(using, client)
+        column_types = _build_column_types(self.model_class)
+        update_ir = ir.build_update_ir(
+            table=self.model_class.get_table_name(),
+            values={field: _serialize_value_for_ir(F(field) + by)},
+            filter_tree=self._build_filter_tree(),
+            column_types=column_types,
+            column_db_types=_build_column_db_types(self.model_class),
+            model=_model_key(self.model_class),
+        )
+        result_bytes = await exec_client.execute(update_ir)
+        result = msgpack.unpackb(result_bytes, raw=False)
+        affected: int = result.get("affected", 0)
+        return affected
+
+    @overload
+    async def update(
+        self,
+        *,
+        returning: Literal[True],
+        using: str | None = ...,
+        client: SupportsExecute | None = ...,
+        **values: Any,
+    ) -> list[TModel]: ...
+
+    @overload
+    async def update(
+        self,
+        *,
+        returning: Literal[False] = ...,
+        using: str | None = ...,
+        client: SupportsExecute | None = ...,
+        **values: Any,
+    ) -> int: ...
+
+    async def update(
+        self,
+        *,
+        returning: bool = False,
+        using: str | None = None,
+        client: SupportsExecute | None = None,
+        **values: Any,
+    ) -> int | list[TModel]:
+        """
+        Update records matching the query.
+
+        Args:
+            returning: If True, return updated model instances.
+                Default False — returns number of affected rows (Django-compatible).
+            using: Database alias
+            client: Optional database client
+            **values: Field values to update
+
+        Returns:
+            Number of affected rows (default), or list of updated model instances
+            if returning=True.
+
+        Examples:
+            count = await Post.objects.filter(id=42).update(status="published")
+            posts = await Post.objects.filter(id=42).update(
+                status="published", returning=True
+            )
+        """
+        exec_client = await _resolve_execution_client(using, client)
+        column_types = _build_column_types(self.model_class)
+        values = _validate_enum_update_values(self.model_class, values)
+        mapped_values = _map_values_to_columns(self.model_class, values)
+        serialized_values = {
+            key: _serialize_value_for_ir(value) for key, value in mapped_values.items()
+        }
+
+        update_ir = ir.build_update_ir(
+            table=self.model_class.get_table_name(),
+            values=serialized_values,
+            filter_tree=self._build_filter_tree(),
+            column_types=column_types,
+            column_db_types=_build_column_db_types(self.model_class),
+            model=_model_key(self.model_class),
+            returning=returning,
+        )
+        result_bytes = await exec_client.execute(update_ir)
+        result = msgpack.unpackb(result_bytes, raw=False)
+        if returning:
+            return _decode_returning_models(self.model_class, result)
+        affected: int = result.get("affected", 0)
+        return affected
+
+    async def delete(
+        self,
+        *,
+        using: str | None = None,
+        client: SupportsExecute | None = None,
+    ) -> int:
+        """
+        Delete records matching the query.
+
+        Args:
+            using: Database alias
+            client: Optional database client
+
+        Returns:
+            Number of affected rows
+
+        Examples:
+            await Post.objects.filter(id=42).delete()
+            await User.objects.filter(is_active=False).delete()
+        """
+        exec_client = await _resolve_execution_client(using, client)
+        delete_ir = ir.build_delete_ir(
+            table=self.model_class.get_table_name(),
+            filter_tree=self._build_filter_tree(),
+            column_types=_build_column_types(self.model_class),
+            column_db_types=_build_column_db_types(self.model_class),
+            model=_model_key(self.model_class),
+        )
+        result_bytes = await exec_client.execute(delete_ir)
+        result = msgpack.unpackb(result_bytes, raw=False)
+        affected: int = result.get("affected", 0)
+        return affected
+
+    def _primary_key_field(self) -> str | None:
+        """Get primary key field name."""
+        return self.model_class._db_meta.pk_field
+
+    async def _run_mutation(
+        self, query: Any, client: SupportsExecute
+    ) -> dict[str, Any]:
+        """Execute mutation query with error handling."""
+        try:
+            result = await query.execute(client)
+        except ManagerError:
+            raise
+        except Exception as exc:
+            translated = _translate_core_error(exc)
+            if translated is not None:
+                raise translated from exc
+            raise ManagerError(str(exc)) from exc
+        if not isinstance(result, dict):
+            raise ManagerError("Mutation response must be a dict")
+        return result
+
+    async def create(
+        self,
+        *,
+        instance: TModel | None = None,
+        using: str | None = None,
+        client: SupportsExecute | None = None,
+        _skip_hooks: bool = False,
+        **data: Any,
+    ) -> TModel:
+        """
+        Create a new record in the database.
+
+        Args:
+            instance: Model instance to create (alternative to **data)
+            using: Database alias
+            client: Optional database client
+            _skip_hooks: Skip pre_save/post_save hooks
+            **data: Field values for the new record
+
+        Returns:
+            Created model instance with populated PK
+
+        Examples:
+            user = await User.objects.create(name="Alice", email="alice@example.com")
+            # Or with instance:
+            user = User(name="Alice")
+            user = await User.objects.create(instance=user)
+        """
+        if instance is not None and data:
+            raise ManagerError(
+                "create() accepts either 'instance' or field values, not both"
+            )
+        if instance is None:
+            if not data:
+                raise ManagerError("create() requires an instance or field values")
+            instance = self.model_class(**data)
+
+        # Call pre_save hook
+        if not _skip_hooks:
+            await instance.pre_save(is_create=True, update_fields=None)
+
+        exec_client = await _resolve_execution_client(using, client)
+        payload = _dump_insert_data(instance)
+        if not payload:
+            raise ManagerError("create() requires at least one value")
+        query = InsertQuery(self.model_class).values(**payload)
+        result = await self._run_mutation(query, exec_client)
+
+        # Update instance from RETURNING * result (with Pydantic validation)
+        if "rows" in result and result["rows"]:
+            rmap = self.model_class._db_meta.reverse_column_map
+            columns = result.get("columns", [])
+            row = result["rows"][0]
+            row_dict = {rmap.get(col, col): value for col, value in zip(columns, row)}
+            instance = self.model_class.model_validate(row_dict)
+        elif "inserted_ids" in result and result["inserted_ids"]:
+            pk_field = self._primary_key_field()
+            if pk_field:
+                row_dict = instance.model_dump()
+                row_dict[pk_field] = result["inserted_ids"][0]
+                instance = self.model_class.model_validate(row_dict)
+
+        # Call post_save hook
+        if not _skip_hooks:
+            await instance.post_save(is_create=True, update_fields=None)
+
+        return instance
+
+    async def bulk_create(
+        self,
+        objects: Iterable[TModel | dict[str, Any]],
+        *,
+        using: str | None = None,
+        client: SupportsExecute | None = None,
+        batch_size: int | None = None,
+    ) -> list[TModel]:
+        """
+        Insert multiple objects efficiently.
+
+        Args:
+            objects: Iterable of model instances or dicts
+            using: Database alias
+            client: Optional database client
+            batch_size: Optional batch size. If None, inserts all in one query.
+                Use if hitting PostgreSQL parameter limits (65535).
+
+        Returns:
+            List of created model instances
+        """
+        instances = [_normalize_instance(self.model_class, obj) for obj in objects]
+        if not instances:
+            return []
+
+        exec_client = await _resolve_execution_client(using, client)
+
+        # Determine batches: all at once or user-specified batch_size
+        if batch_size is not None and batch_size <= 0:
+            raise ValueError("batch_size must be a positive integer")
+        if batch_size is None:
+            batches = [instances]
+        else:
+            batches = [
+                instances[i : i + batch_size]
+                for i in range(0, len(instances), batch_size)
+            ]
+
+        for batch in batches:
+            payloads = []
+            for instance in batch:
+                payload = _dump_insert_data(instance)
+                if not payload:
+                    raise ManagerError("bulk_create() encountered an empty payload")
+                payloads.append(payload)
+
+            query = InsertQuery(self.model_class).bulk_values(payloads)
+            result = await self._run_mutation(query, exec_client)
+
+            # Assign auto-generated PKs to instances
+            inserted_ids = result.get("inserted_ids", [])
+            pk_field = self._primary_key_field()
+
+            if inserted_ids and len(inserted_ids) != len(batch):
+                warnings.warn(
+                    f"bulk_create: received {len(inserted_ids)} IDs for {len(batch)} rows. "
+                    "Check conflict handling and the returned primary keys. "
+                    "non-sequential auto_increment, or non-integer primary keys. "
+                    "Assigned IDs may be incorrect.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+
+            if pk_field and inserted_ids:
+                for instance, pk_value in zip(batch, inserted_ids):
+                    setattr(instance, pk_field, pk_value)
+
+        return instances
+
+    async def bulk_update(
+        self,
+        objects: Iterable[TModel],
+        fields: Iterable[str],
+        *,
+        using: str | None = None,
+        client: SupportsExecute | None = None,
+    ) -> int:
+        """
+        Bulk update multiple objects with CASE WHEN.
+
+        Args:
+            objects: Instances to update
+            fields: Field names to update
+            using: Database alias
+            client: Optional database client
+
+        Returns:
+            Number of affected rows
+
+        Examples:
+            users = await User.objects.filter(is_active=True).all()
+            for user in users:
+                user.login_count += 1
+            await User.objects.bulk_update(users, ["login_count"])
+        """
+        objects_list = list(objects)
+        if not objects_list:
+            return 0
+
+        fields_list = list(fields)
+        if not fields_list:
+            raise ManagerError("bulk_update() requires at least one field")
+
+        pk_field = self._primary_key_field()
+        if not pk_field:
+            raise ManagerError("bulk_update() requires a primary key")
+
+        pk_column = self.model_class._db_meta.field_metadata[pk_field].db_column
+
+        exec_client = await _resolve_execution_client(using, client)
+
+        # Build bulk_update payload
+        bulk_entries: list[dict[str, Any]] = []
+        for obj in objects_list:
+            pk_value = getattr(obj, pk_field)
+            if pk_value is None:
+                continue
+
+            all_values = obj.model_dump(mode="python", exclude_none=False)
+
+            # Extract requested fields and map to db_column names
+            values = {
+                field_name: _serialize_value_for_ir(all_values[field_name])
+                for field_name in fields_list
+                if field_name in all_values
+            }
+            mapped_values = _map_values_to_columns(self.model_class, values)
+
+            if mapped_values:
+                bulk_entries.append(
+                    {
+                        "filters": {pk_column: pk_value},
+                        "values": mapped_values,
+                    }
+                )
+
+        if not bulk_entries:
+            return 0
+
+        column_types = _build_column_types(self.model_class)
+        update_ir = ir.build_update_ir(
+            table=self.model_class.get_table_name(),
+            bulk_update=bulk_entries,
+            column_types=column_types,
+            model=_model_key(self.model_class),
+        )
+
+        result_bytes = await exec_client.execute(update_ir)
+        result = msgpack.unpackb(result_bytes, raw=False)
+        return int(result.get("affected", 0))

oxyde/queries/mixins/pagination.py

@@ -0,0 +1,145 @@
+"""Pagination mixin for query building."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+from typing_extensions import Self
+
+if TYPE_CHECKING:
+    from typing import Literal, overload
+
+    from oxyde.models.base import Model
+    from oxyde.queries.typed import (
+        FlatValuesListQuery,
+        ValuesListQuery,
+        ValuesQuery,
+    )
+
+
+TModel = TypeVar("TModel", bound="Model")
+
+
+class PaginationMixin(Generic[TModel]):
+    """Mixin providing pagination and ordering capabilities."""
+
+    # These attributes are defined in the base Query class
+    _limit_value: int | None
+    _offset_value: int | None
+    _order_by_fields: list[tuple[str, str]]
+    _distinct: bool
+    _result_mode: str | None
+    _values_flat: bool
+    _selected_fields: list[str] | None
+
+    def _clone(self) -> Self:
+        """Must be implemented by the main Query class."""
+        raise NotImplementedError
+
+    def select(self, *fields: str) -> Self:
+        """Must be implemented by the main Query class."""
+        raise NotImplementedError
+
+    def limit(self, value: int) -> Self:
+        """Set LIMIT."""
+        if value < 0:
+            raise ValueError(f"limit() requires a non-negative value, got {value}")
+        clone = self._clone()
+        clone._limit_value = value
+        return clone
+
+    def offset(self, value: int) -> Self:
+        """Set OFFSET."""
+        if value < 0:
+            raise ValueError(f"offset() requires a non-negative value, got {value}")
+        clone = self._clone()
+        clone._offset_value = value
+        return clone
+
+    def order_by(self, *fields: str) -> Self:
+        """Set ORDER BY fields. Use "?" for random ordering (ORDER BY RANDOM())."""
+        clone = self._clone()
+        for field in fields:
+            if field == "?":
+                clone._order_by_fields.append(("?", "RANDOM"))
+            elif field.startswith("-"):
+                clone._order_by_fields.append((field[1:], "DESC"))
+            else:
+                clone._order_by_fields.append((field, "ASC"))
+        return clone
+
+    def distinct(self, distinct: bool = True) -> Self:
+        """Set DISTINCT."""
+        clone = self._clone()
+        clone._distinct = bool(distinct)
+        return clone
+
+    # values()/values_list() flip a result-mode flag on a clone at runtime;
+    # statically they are declared to return the typing facades from
+    # oxyde.queries.typed so terminal methods reflect the actual row shape
+    # (dicts / tuples / scalars) instead of model instances.
+    if TYPE_CHECKING:
+
+        def values(self, *fields: str) -> ValuesQuery[TModel]:
+            """Return results as dictionaries."""
+            ...
+
+        @overload
+        def values_list(
+            self, field: str, /, *, flat: Literal[True]
+        ) -> FlatValuesListQuery[TModel]: ...
+
+        @overload
+        def values_list(
+            self, *fields: str, flat: Literal[False] = ...
+        ) -> ValuesListQuery[TModel]: ...
+
+        @overload
+        def values_list(
+            self, *fields: str, flat: bool
+        ) -> ValuesListQuery[TModel] | FlatValuesListQuery[TModel]: ...
+
+        def values_list(self, *fields: str, flat: bool = False) -> Any:
+            """Return results as tuples (or flat list if flat=True)."""
+            ...
+
+    else:
+
+        def values(self, *fields):
+            """Return results as dictionaries."""
+            clone = self._clone()
+            if fields:
+                clone = clone.select(*fields)
+            clone._result_mode = "dict"
+            return clone
+
+        def values_list(self, *fields, flat=False):
+            """Return results as tuples (or flat list if flat=True)."""
+            clone = self._clone()
+            if fields:
+                clone = clone.select(*fields)
+            if flat and fields and len(fields) != 1:
+                raise ValueError(
+                    "flat=True is only valid when a single field is selected"
+                )
+            clone._result_mode = "list"
+            clone._values_flat = flat
+            return clone
+
+    def __getitem__(self, key: slice | int) -> Self:
+        """Support slicing: query[0:10] or query[5]."""
+        if isinstance(key, slice):
+            start = key.start or 0
+            stop = key.stop
+            if stop is None:
+                raise ValueError("Slicing a Query requires an end index")
+            if start < 0 or (stop is not None and stop < 0):
+                raise ValueError("Negative slicing is not supported")
+            clone = self.offset(start)
+            length = max(0, stop - start)
+            return clone.limit(length)
+        if isinstance(key, int):
+            if key < 0:
+                raise ValueError("Negative indexing is not supported")
+            return self.offset(key).limit(1)
+        raise TypeError("Invalid argument type for slicing Query")