mongo-aggro 0.1.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

mongo_aggro/stages/__init__.py

@@ -0,0 +1,110 @@
+"""MongoDB aggregation pipeline stages.
+
+This package provides typed stage classes for building MongoDB aggregation
+pipelines. All stages are re-exported here for convenient access.
+"""
+
+from mongo_aggro.stages.array import Unwind
+from mongo_aggro.stages.change import ChangeStream, ChangeStreamSplitLargeEvent
+from mongo_aggro.stages.core import (
+    Count,
+    Group,
+    Limit,
+    Match,
+    Project,
+    Skip,
+    Sort,
+)
+from mongo_aggro.stages.geo import GeoNear
+from mongo_aggro.stages.group import Bucket, BucketAuto, Facet, SortByCount
+from mongo_aggro.stages.join import GraphLookup, Lookup, UnionWith
+from mongo_aggro.stages.misc import ListClusterCatalog, QuerySettings
+from mongo_aggro.stages.output import Documents, Merge, Out, Sample
+from mongo_aggro.stages.search import (
+    ListSearchIndexes,
+    RankFusion,
+    Search,
+    SearchMeta,
+    VectorSearch,
+)
+from mongo_aggro.stages.session import (
+    ListLocalSessions,
+    ListSampledQueries,
+    ListSessions,
+)
+from mongo_aggro.stages.stats import (
+    CollStats,
+    CurrentOp,
+    IndexStats,
+    PlanCacheStats,
+)
+from mongo_aggro.stages.transform import (
+    AddFields,
+    Redact,
+    ReplaceRoot,
+    ReplaceWith,
+    Set,
+    Unset,
+)
+from mongo_aggro.stages.window import Densify, Fill, SetWindowFields
+
+__all__ = [
+    # core
+    "Match",
+    "Project",
+    "Group",
+    "Sort",
+    "Limit",
+    "Skip",
+    "Count",
+    # array
+    "Unwind",
+    # join
+    "Lookup",
+    "UnionWith",
+    "GraphLookup",
+    # transform
+    "AddFields",
+    "Set",
+    "Unset",
+    "ReplaceRoot",
+    "ReplaceWith",
+    "Redact",
+    # group
+    "Facet",
+    "Bucket",
+    "BucketAuto",
+    "SortByCount",
+    # output
+    "Out",
+    "Merge",
+    "Sample",
+    "Documents",
+    # window
+    "SetWindowFields",
+    "Densify",
+    "Fill",
+    # geo
+    "GeoNear",
+    # stats
+    "CollStats",
+    "IndexStats",
+    "PlanCacheStats",
+    "CurrentOp",
+    # session
+    "ListSessions",
+    "ListLocalSessions",
+    "ListSampledQueries",
+    # change
+    "ChangeStream",
+    "ChangeStreamSplitLargeEvent",
+    # search
+    "Search",
+    "SearchMeta",
+    "VectorSearch",
+    "ListSearchIndexes",
+    "RankFusion",
+    # misc
+    "ListClusterCatalog",
+    "QuerySettings",
+]

mongo_aggro/stages/array.py

@@ -0,0 +1,69 @@
+"""Array-related MongoDB aggregation pipeline stages.
+
+This module contains stages for working with array fields, primarily Unwind.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class Unwind(BaseModel):
+    """
+    $unwind stage - deconstructs an array field.
+
+    Example:
+        >>> Unwind(path="cars").model_dump()
+        {"$unwind": "$cars"}
+
+        >>> # With options
+        >>> Unwind(
+        ...     path="items",
+        ...     include_array_index="itemIndex",
+        ...     preserve_null_and_empty=True
+        ... ).model_dump()
+        {"$unwind": {
+            "path": "$items",
+            "includeArrayIndex": "itemIndex",
+            "preserveNullAndEmptyArrays": true
+        }}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    path: str = Field(..., description="Array field path (without $)")
+    include_array_index: str | None = Field(
+        default=None,
+        validation_alias="includeArrayIndex",
+        serialization_alias="includeArrayIndex",
+        description="Name of index field",
+    )
+    preserve_null_and_empty: bool | None = Field(
+        default=None,
+        validation_alias="preserveNullAndEmptyArrays",
+        serialization_alias="preserveNullAndEmptyArrays",
+        description="Output doc if array is null/empty/missing",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        field_path = (
+            f"${self.path}" if not self.path.startswith("$") else self.path
+        )
+
+        if (
+            self.include_array_index is None
+            and self.preserve_null_and_empty is None
+        ):
+            return {"$unwind": field_path}
+
+        result: dict[str, Any] = {"path": field_path}
+        if self.include_array_index is not None:
+            result["includeArrayIndex"] = self.include_array_index
+        if self.preserve_null_and_empty is not None:
+            result["preserveNullAndEmptyArrays"] = self.preserve_null_and_empty
+        return {"$unwind": result}
+
+
+__all__ = [
+    "Unwind",
+]

mongo_aggro/stages/change.py

@@ -0,0 +1,109 @@
+"""Change stream MongoDB aggregation pipeline stages.
+
+This module contains stages for change stream operations:
+ChangeStream and ChangeStreamSplitLargeEvent.
+"""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class ChangeStream(BaseModel):
+    """
+    $changeStream stage - returns a change stream cursor.
+
+    Must be the first stage in the pipeline.
+
+    Example:
+        >>> ChangeStream().model_dump()
+        {"$changeStream": {}}
+
+        >>> ChangeStream(full_document="updateLookup").model_dump()
+        {"$changeStream": {"fullDocument": "updateLookup"}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    full_document: (
+        Literal["default", "updateLookup", "whenAvailable", "required"] | None
+    ) = Field(
+        default=None,
+        serialization_alias="fullDocument",
+        description="Full document option for update events",
+    )
+    full_document_before_change: (
+        Literal["off", "whenAvailable", "required"] | None
+    ) = Field(
+        default=None,
+        serialization_alias="fullDocumentBeforeChange",
+        description="Include pre-image of modified document",
+    )
+    resume_after: dict[str, Any] | None = Field(
+        default=None,
+        serialization_alias="resumeAfter",
+        description="Resume token to resume change stream",
+    )
+    start_after: dict[str, Any] | None = Field(
+        default=None,
+        serialization_alias="startAfter",
+        description="Resume token to start after",
+    )
+    start_at_operation_time: Any | None = Field(
+        default=None,
+        serialization_alias="startAtOperationTime",
+        description="Timestamp to start watching changes",
+    )
+    all_changes_for_cluster: bool | None = Field(
+        default=None,
+        serialization_alias="allChangesForCluster",
+        description="Watch all changes for the cluster",
+    )
+    show_expanded_events: bool | None = Field(
+        default=None,
+        serialization_alias="showExpandedEvents",
+        description="Show expanded change events",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        if self.full_document is not None:
+            result["fullDocument"] = self.full_document
+        if self.full_document_before_change is not None:
+            result["fullDocumentBeforeChange"] = (
+                self.full_document_before_change
+            )
+        if self.resume_after is not None:
+            result["resumeAfter"] = self.resume_after
+        if self.start_after is not None:
+            result["startAfter"] = self.start_after
+        if self.start_at_operation_time is not None:
+            result["startAtOperationTime"] = self.start_at_operation_time
+        if self.all_changes_for_cluster is not None:
+            result["allChangesForCluster"] = self.all_changes_for_cluster
+        if self.show_expanded_events is not None:
+            result["showExpandedEvents"] = self.show_expanded_events
+        return {"$changeStream": result}
+
+
+class ChangeStreamSplitLargeEvent(BaseModel):
+    """
+    $changeStreamSplitLargeEvent stage - splits large change events.
+
+    Must be the last stage in a $changeStream pipeline.
+
+    Example:
+        >>> ChangeStreamSplitLargeEvent().model_dump()
+        {"$changeStreamSplitLargeEvent": {}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$changeStreamSplitLargeEvent": {}}
+
+
+__all__ = [
+    "ChangeStream",
+    "ChangeStreamSplitLargeEvent",
+]

mongo_aggro/stages/core.py

@@ -0,0 +1,170 @@
+"""Core MongoDB aggregation pipeline stages.
+
+This module contains the fundamental stages used in most aggregation pipelines:
+Match, Project, Group, Sort, Limit, Skip, and Count.
+"""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class Match(BaseModel):
+    """
+    $match stage - filters documents by specified criteria.
+
+    Example:
+        >>> Match(query={"status": "active"}).model_dump()
+        {"$match": {"status": "active"}}
+
+        >>> # With logical operators
+        >>> Match(query={"$and": [{"status": "active"}, {"age": {"$gt": 18}}]})
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    query: dict[str, Any] = Field(..., description="Query filter conditions")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$match": self.query}
+
+
+class Project(BaseModel):
+    """
+    $project stage - shapes documents by including/excluding fields.
+
+    Example:
+        >>> Project(fields={"name": 1, "year": 1, "_id": 0}).model_dump()
+        {"$project": {"name": 1, "year": 1, "_id": 0}}
+
+        >>> # With expressions
+        >>> Project(fields={"fullName": {"$concat": ["$first", " ", "$last"]}})
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    fields: dict[str, Any] = Field(
+        ..., description="Field projections (1=include, 0=exclude, or expr)"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$project": self.fields}
+
+
+class Group(BaseModel):
+    """
+    $group stage - groups documents by specified expression.
+
+    Example:
+        >>> Group(
+        ...     id="$category",
+        ...     total={"$sum": "$quantity"},
+        ...     count={"$sum": 1}
+        ... ).model_dump()
+        {
+            "$group": {
+                "_id": "$category",
+                "total": {"$sum": "$quantity"},
+                "count": {"$sum": 1}
+            }
+        }
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    id: Any = Field(
+        ...,
+        validation_alias="_id",
+        serialization_alias="_id",
+        description="Grouping expression",
+    )
+    accumulators: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Accumulator expressions (e.g., $sum, $avg)",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result = {"_id": self.id}
+        result.update(self.accumulators)
+        return {"$group": result}
+
+
+class Sort(BaseModel):
+    """
+    $sort stage - sorts documents.
+
+    Example:
+        >>> Sort(fields={"age": -1, "name": 1}).model_dump()
+        {"$sort": {"age": -1, "name": 1}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    fields: dict[str, Literal[-1, 1]] = Field(
+        ..., description="Sort specification (1=asc, -1=desc)"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$sort": self.fields}
+
+
+class Limit(BaseModel):
+    """
+    $limit stage - limits the number of documents.
+
+    Example:
+        >>> Limit(count=10).model_dump()
+        {"$limit": 10}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    count: int = Field(..., gt=0, description="Maximum number of documents")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$limit": self.count}
+
+
+class Skip(BaseModel):
+    """
+    $skip stage - skips a number of documents.
+
+    Example:
+        >>> Skip(count=5).model_dump()
+        {"$skip": 5}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    count: int = Field(..., ge=0, description="Number of documents to skip")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$skip": self.count}
+
+
+class Count(BaseModel):
+    """
+    $count stage - counts documents.
+
+    Example:
+        >>> Count(field="total").model_dump()
+        {"$count": "total"}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    field: str = Field(..., description="Output field name for count")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$count": self.field}
+
+
+__all__ = [
+    "Match",
+    "Project",
+    "Group",
+    "Sort",
+    "Limit",
+    "Skip",
+    "Count",
+]

mongo_aggro/stages/geo.py

@@ -0,0 +1,93 @@
+"""Geospatial MongoDB aggregation pipeline stages.
+
+This module contains stages for geospatial queries: GeoNear.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class GeoNear(BaseModel):
+    """
+    $geoNear stage - returns documents near a geographic point.
+
+    Example:
+        >>> GeoNear(
+        ...     near={"type": "Point", "coordinates": [-73.99, 40.73]},
+        ...     distance_field="dist.calculated",
+        ...     spherical=True,
+        ...     max_distance=5000
+        ... ).model_dump()
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    near: dict[str, Any] | list[float] = Field(
+        ..., description="GeoJSON point or legacy coordinates"
+    )
+    distance_field: str = Field(
+        ...,
+        validation_alias="distanceField",
+        serialization_alias="distanceField",
+        description="Field for calculated distance",
+    )
+    spherical: bool | None = Field(
+        default=None, description="Use spherical geometry"
+    )
+    max_distance: float | None = Field(
+        default=None,
+        validation_alias="maxDistance",
+        serialization_alias="maxDistance",
+        description="Max distance in meters",
+    )
+    min_distance: float | None = Field(
+        default=None,
+        validation_alias="minDistance",
+        serialization_alias="minDistance",
+        description="Min distance in meters",
+    )
+    query: dict[str, Any] | None = Field(
+        default=None, description="Additional query filter"
+    )
+    distance_multiplier: float | None = Field(
+        default=None,
+        validation_alias="distanceMultiplier",
+        serialization_alias="distanceMultiplier",
+        description="Multiplier for distances",
+    )
+    include_locs: str | None = Field(
+        default=None,
+        validation_alias="includeLocs",
+        serialization_alias="includeLocs",
+        description="Field for matched location",
+    )
+    key: str | None = Field(
+        default=None, description="Geospatial index to use"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {
+            "near": self.near,
+            "distanceField": self.distance_field,
+        }
+        if self.spherical is not None:
+            result["spherical"] = self.spherical
+        if self.max_distance is not None:
+            result["maxDistance"] = self.max_distance
+        if self.min_distance is not None:
+            result["minDistance"] = self.min_distance
+        if self.query is not None:
+            result["query"] = self.query
+        if self.distance_multiplier is not None:
+            result["distanceMultiplier"] = self.distance_multiplier
+        if self.include_locs is not None:
+            result["includeLocs"] = self.include_locs
+        if self.key is not None:
+            result["key"] = self.key
+        return {"$geoNear": result}
+
+
+__all__ = [
+    "GeoNear",
+]

mongo_aggro/stages/group.py

@@ -0,0 +1,154 @@
+"""Grouping and bucketing MongoDB aggregation pipeline stages.
+
+This module contains stages for advanced grouping operations:
+Facet, Bucket, BucketAuto, and SortByCount.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from mongo_aggro.base import Pipeline
+
+
+class SortByCount(BaseModel):
+    """
+    $sortByCount stage - groups and counts by field, sorted by count.
+
+    Example:
+        >>> SortByCount(field="category").model_dump()
+        {"$sortByCount": "$category"}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    field: str = Field(..., description="Field to group and count by")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        field_path = (
+            f"${self.field}" if not self.field.startswith("$") else self.field
+        )
+        return {"$sortByCount": field_path}
+
+
+class Facet(BaseModel):
+    """
+    $facet stage - processes multiple pipelines within a single stage.
+
+    Example:
+        >>> Facet(pipelines={
+        ...     "byCategory": Pipeline([Group(id="$category")]),
+        ...     "byYear": Pipeline([Group(id="$year")])
+        ... }).model_dump()
+        {"$facet": {
+            "byCategory": [{"$group": {"_id": "$category"}}],
+            "byYear": [{"$group": {"_id": "$year"}}]
+        }}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    pipelines: dict[str, Pipeline | list[dict[str, Any]]] = Field(
+        ..., description="Named pipelines"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, list[dict[str, Any]]] = {}
+        for name, pipeline in self.pipelines.items():
+            if isinstance(pipeline, Pipeline):
+                result[name] = pipeline.to_list()
+            else:
+                result[name] = pipeline
+        return {"$facet": result}
+
+
+class Bucket(BaseModel):
+    """
+    $bucket stage - categorizes documents into buckets.
+
+    Example:
+        >>> Bucket(
+        ...     group_by="$price",
+        ...     boundaries=[0, 100, 500, 1000],
+        ...     default="Other",
+        ...     output={"count": {"$sum": 1}}
+        ... ).model_dump()
+        {"$bucket": {
+            "groupBy": "$price",
+            "boundaries": [0, 100, 500, 1000],
+            "default": "Other",
+            "output": {"count": {"$sum": 1}}
+        }}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    group_by: str | dict[str, Any] = Field(
+        ...,
+        validation_alias="groupBy",
+        serialization_alias="groupBy",
+        description="Expression to group by",
+    )
+    boundaries: list[Any] = Field(..., description="Bucket boundaries")
+    default: Any | None = Field(
+        default=None, description="Default bucket for non-matching docs"
+    )
+    output: dict[str, Any] | None = Field(
+        default=None, description="Output document specification"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {
+            "groupBy": self.group_by,
+            "boundaries": self.boundaries,
+        }
+        if self.default is not None:
+            result["default"] = self.default
+        if self.output is not None:
+            result["output"] = self.output
+        return {"$bucket": result}
+
+
+class BucketAuto(BaseModel):
+    """
+    $bucketAuto stage - automatically categorizes into specified buckets.
+
+    Example:
+        >>> BucketAuto(group_by="$age", buckets=5).model_dump()
+        {"$bucketAuto": {"groupBy": "$age", "buckets": 5}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    group_by: str | dict[str, Any] = Field(
+        ...,
+        validation_alias="groupBy",
+        serialization_alias="groupBy",
+        description="Expression to group by",
+    )
+    buckets: int = Field(..., gt=0, description="Number of buckets")
+    output: dict[str, Any] | None = Field(
+        default=None, description="Output document specification"
+    )
+    granularity: str | None = Field(
+        default=None, description="Preferred number series"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {
+            "groupBy": self.group_by,
+            "buckets": self.buckets,
+        }
+        if self.output is not None:
+            result["output"] = self.output
+        if self.granularity is not None:
+            result["granularity"] = self.granularity
+        return {"$bucketAuto": result}
+
+
+__all__ = [
+    "Facet",
+    "Bucket",
+    "BucketAuto",
+    "SortByCount",
+]