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

mongo_aggro/stages/join.py

@@ -0,0 +1,221 @@
+"""Join-related MongoDB aggregation pipeline stages.
+
+This module contains stages for joining and combining data from multiple
+collections: Lookup, UnionWith, and GraphLookup.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from mongo_aggro.base import Pipeline
+
+
+class Lookup(BaseModel):
+    """
+    $lookup stage - performs a left outer join.
+
+    Example:
+        >>> # Simple lookup
+        >>> Lookup(
+        ...     from_collection="products",
+        ...     local_field="product_id",
+        ...     foreign_field="_id",
+        ...     as_field="product"
+        ... ).model_dump()
+        {"$lookup": {
+            "from": "products",
+            "localField": "product_id",
+            "foreignField": "_id",
+            "as": "product"
+        }}
+
+        >>> # With pipeline
+        >>> Lookup(
+        ...     from_collection="orders",
+        ...     let={"customerId": "$_id"},
+        ...     pipeline=Pipeline([Match(query={"status": "active"})]),
+        ...     as_field="orders"
+        ... ).model_dump()
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    from_collection: str = Field(
+        ...,
+        validation_alias="from",
+        serialization_alias="from",
+        description="Foreign collection name",
+    )
+    local_field: str | None = Field(
+        default=None,
+        validation_alias="localField",
+        serialization_alias="localField",
+        description="Local field for join",
+    )
+    foreign_field: str | None = Field(
+        default=None,
+        validation_alias="foreignField",
+        serialization_alias="foreignField",
+        description="Foreign field for join",
+    )
+    let: dict[str, Any] | None = Field(
+        default=None, description="Variables for pipeline"
+    )
+    pipeline: Pipeline | list[dict[str, Any]] | None = Field(
+        default=None, description="Sub-pipeline for complex joins"
+    )
+    as_field: str = Field(
+        ...,
+        validation_alias="as",
+        serialization_alias="as",
+        description="Output array field name",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {
+            "from": self.from_collection,
+            "as": self.as_field,
+        }
+
+        if self.local_field is not None:
+            result["localField"] = self.local_field
+        if self.foreign_field is not None:
+            result["foreignField"] = self.foreign_field
+        if self.let is not None:
+            result["let"] = self.let
+        if self.pipeline is not None:
+            if isinstance(self.pipeline, Pipeline):
+                result["pipeline"] = self.pipeline.to_list()
+            else:
+                result["pipeline"] = self.pipeline
+
+        return {"$lookup": result}
+
+
+class UnionWith(BaseModel):
+    """
+    $unionWith stage - combines pipeline results with another collection.
+
+    Example:
+        >>> UnionWith(collection="archive").model_dump()
+        {"$unionWith": "archive"}
+
+        >>> UnionWith(
+        ...     collection="archive",
+        ...     pipeline=Pipeline([Match(query={"year": 2023})])
+        ... ).model_dump()
+        {"$unionWith": {"coll": "archive", "pipeline": [...]}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    collection: str = Field(
+        ...,
+        validation_alias="coll",
+        serialization_alias="coll",
+        description="Collection to union",
+    )
+    pipeline: Pipeline | list[dict[str, Any]] | None = Field(
+        default=None, description="Pipeline for the other collection"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        if self.pipeline is None:
+            return {"$unionWith": self.collection}
+
+        pl = (
+            self.pipeline.to_list()
+            if isinstance(self.pipeline, Pipeline)
+            else self.pipeline
+        )
+        return {"$unionWith": {"coll": self.collection, "pipeline": pl}}
+
+
+class GraphLookup(BaseModel):
+    """
+    $graphLookup stage - performs recursive search.
+
+    Example:
+        >>> GraphLookup(
+        ...     from_collection="employees",
+        ...     start_with="$reportsTo",
+        ...     connect_from_field="reportsTo",
+        ...     connect_to_field="name",
+        ...     as_field="reportingHierarchy"
+        ... ).model_dump()
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    from_collection: str = Field(
+        ...,
+        validation_alias="from",
+        serialization_alias="from",
+        description="Collection to search",
+    )
+    start_with: Any = Field(
+        ...,
+        validation_alias="startWith",
+        serialization_alias="startWith",
+        description="Expression for starting point",
+    )
+    connect_from_field: str = Field(
+        ...,
+        validation_alias="connectFromField",
+        serialization_alias="connectFromField",
+        description="Field to recurse from",
+    )
+    connect_to_field: str = Field(
+        ...,
+        validation_alias="connectToField",
+        serialization_alias="connectToField",
+        description="Field to match",
+    )
+    as_field: str = Field(
+        ...,
+        validation_alias="as",
+        serialization_alias="as",
+        description="Output array field",
+    )
+    max_depth: int | None = Field(
+        default=None,
+        validation_alias="maxDepth",
+        serialization_alias="maxDepth",
+        description="Maximum recursion depth",
+    )
+    depth_field: str | None = Field(
+        default=None,
+        validation_alias="depthField",
+        serialization_alias="depthField",
+        description="Field for recursion depth",
+    )
+    restrict_search_with_match: dict[str, Any] | None = Field(
+        default=None,
+        validation_alias="restrictSearchWithMatch",
+        serialization_alias="restrictSearchWithMatch",
+        description="Additional match conditions",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {
+            "from": self.from_collection,
+            "startWith": self.start_with,
+            "connectFromField": self.connect_from_field,
+            "connectToField": self.connect_to_field,
+            "as": self.as_field,
+        }
+        if self.max_depth is not None:
+            result["maxDepth"] = self.max_depth
+        if self.depth_field is not None:
+            result["depthField"] = self.depth_field
+        if self.restrict_search_with_match is not None:
+            result["restrictSearchWithMatch"] = self.restrict_search_with_match
+        return {"$graphLookup": result}
+
+
+__all__ = [
+    "Lookup",
+    "UnionWith",
+    "GraphLookup",
+]

mongo_aggro/stages/misc.py

@@ -0,0 +1,45 @@
+"""Miscellaneous MongoDB aggregation pipeline stages.
+
+This module contains less commonly used stages:
+ListClusterCatalog and QuerySettings.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+
+
+class ListClusterCatalog(BaseModel):
+    """
+    $listClusterCatalog stage - lists collections in a cluster.
+
+    Example:
+        >>> ListClusterCatalog().model_dump()
+        {"$listClusterCatalog": {}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$listClusterCatalog": {}}
+
+
+class QuerySettings(BaseModel):
+    """
+    $querySettings stage - returns query settings (MongoDB 8.0+).
+
+    Example:
+        >>> QuerySettings().model_dump()
+        {"$querySettings": {}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$querySettings": {}}
+
+
+__all__ = [
+    "ListClusterCatalog",
+    "QuerySettings",
+]

mongo_aggro/stages/output.py

@@ -0,0 +1,136 @@
+"""Output and sampling MongoDB aggregation pipeline stages.
+
+This module contains stages for writing output and sampling:
+Out, Merge, Sample, and Documents.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class Sample(BaseModel):
+    """
+    $sample stage - randomly selects documents.
+
+    Example:
+        >>> Sample(size=10).model_dump()
+        {"$sample": {"size": 10}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    size: int = Field(..., gt=0, description="Number of documents to sample")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$sample": {"size": self.size}}
+
+
+class Out(BaseModel):
+    """
+    $out stage - writes results to a collection.
+
+    Example:
+        >>> Out(collection="results").model_dump()
+        {"$out": "results"}
+
+        >>> Out(collection="results", db="analytics").model_dump()
+        {"$out": {"db": "analytics", "coll": "results"}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    collection: str = Field(..., description="Output collection name")
+    db: str | None = Field(default=None, description="Output database name")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        if self.db is not None:
+            return {"$out": {"db": self.db, "coll": self.collection}}
+        return {"$out": self.collection}
+
+
+class Merge(BaseModel):
+    """
+    $merge stage - writes results to a collection with merge behavior.
+
+    Example:
+        >>> Merge(
+        ...     into="reports",
+        ...     on="_id",
+        ...     when_matched="merge",
+        ...     when_not_matched="insert"
+        ... ).model_dump()
+        {"$merge": {
+            "into": "reports",
+            "on": "_id",
+            "whenMatched": "merge",
+            "whenNotMatched": "insert"
+        }}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    into: str | dict[str, str] = Field(
+        ..., description="Target collection (or {db, coll})"
+    )
+    on: str | list[str] | None = Field(
+        default=None, description="Field(s) to match on"
+    )
+    let: dict[str, Any] | None = Field(
+        default=None, description="Variables for pipeline"
+    )
+    when_matched: str | list[dict[str, Any]] | None = Field(
+        default=None,
+        validation_alias="whenMatched",
+        serialization_alias="whenMatched",
+        description="Action when matched (replace, keepExisting, merge, fail, "
+        "or pipeline)",
+    )
+    when_not_matched: str | None = Field(
+        default=None,
+        validation_alias="whenNotMatched",
+        serialization_alias="whenNotMatched",
+        description="Action when not matched (insert, discard, fail)",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {"into": self.into}
+        if self.on is not None:
+            result["on"] = self.on
+        if self.let is not None:
+            result["let"] = self.let
+        if self.when_matched is not None:
+            result["whenMatched"] = self.when_matched
+        if self.when_not_matched is not None:
+            result["whenNotMatched"] = self.when_not_matched
+        return {"$merge": result}
+
+
+class Documents(BaseModel):
+    """
+    $documents stage - returns literal documents.
+
+    Example:
+        >>> Documents(documents=[
+        ...     {"x": 1, "y": 2},
+        ...     {"x": 3, "y": 4}
+        ... ]).model_dump()
+        {"$documents": [{"x": 1, "y": 2}, {"x": 3, "y": 4}]}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    documents: list[dict[str, Any]] = Field(
+        ..., description="Documents to return"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$documents": self.documents}
+
+
+__all__ = [
+    "Out",
+    "Merge",
+    "Sample",
+    "Documents",
+]