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

mongo_aggro/stages/search.py

@@ -0,0 +1,315 @@
+"""Atlas Search MongoDB aggregation pipeline stages.
+
+This module contains stages for Atlas Search and Vector Search:
+Search, SearchMeta, VectorSearch, ListSearchIndexes, and RankFusion.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class ListSearchIndexes(BaseModel):
+    """
+    $listSearchIndexes stage - lists Atlas Search indexes.
+
+    Example:
+        >>> ListSearchIndexes().model_dump()
+        {"$listSearchIndexes": {}}
+
+        >>> ListSearchIndexes(id="index_id").model_dump()
+        {"$listSearchIndexes": {"id": "index_id"}}
+
+        >>> ListSearchIndexes(name="index_name").model_dump()
+        {"$listSearchIndexes": {"name": "index_name"}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    id: str | None = Field(
+        default=None,
+        description="Search index ID to filter",
+    )
+    name: str | None = Field(
+        default=None,
+        description="Search index name to filter",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        if self.id is not None:
+            result["id"] = self.id
+        if self.name is not None:
+            result["name"] = self.name
+        return {"$listSearchIndexes": result}
+
+
+class Search(BaseModel):
+    """
+    $search stage - Atlas full-text search.
+
+    Example:
+        >>> Search(index="default", text={"query": "coffee", "path": "title"})
+        {"$search": {"index": "default", "text": {"query": "coffee", ...}}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    index: str | None = Field(
+        default=None,
+        description="Name of the Atlas Search index",
+    )
+    text: dict[str, Any] | None = Field(
+        default=None,
+        description="Text search operator",
+    )
+    compound: dict[str, Any] | None = Field(
+        default=None,
+        description="Compound search operator",
+    )
+    autocomplete: dict[str, Any] | None = Field(
+        default=None,
+        description="Autocomplete search operator",
+    )
+    phrase: dict[str, Any] | None = Field(
+        default=None,
+        description="Phrase search operator",
+    )
+    wildcard: dict[str, Any] | None = Field(
+        default=None,
+        description="Wildcard search operator",
+    )
+    regex: dict[str, Any] | None = Field(
+        default=None,
+        description="Regex search operator",
+    )
+    near: dict[str, Any] | None = Field(
+        default=None,
+        description="Near search operator",
+    )
+    range: dict[str, Any] | None = Field(
+        default=None,
+        description="Range search operator",
+    )
+    exists: dict[str, Any] | None = Field(
+        default=None,
+        description="Exists search operator",
+    )
+    equals: dict[str, Any] | None = Field(
+        default=None,
+        description="Equals search operator",
+    )
+    more_like_this: dict[str, Any] | None = Field(
+        default=None,
+        serialization_alias="moreLikeThis",
+        description="More like this search operator",
+    )
+    query_string: dict[str, Any] | None = Field(
+        default=None,
+        serialization_alias="queryString",
+        description="Query string search operator",
+    )
+    highlight: dict[str, Any] | None = Field(
+        default=None,
+        description="Highlight options",
+    )
+    count: dict[str, Any] | None = Field(
+        default=None,
+        description="Count options",
+    )
+    return_stored_source: bool | None = Field(
+        default=None,
+        serialization_alias="returnStoredSource",
+        description="Return stored source",
+    )
+
+    def _add_operators(self, result: dict[str, Any]) -> None:
+        """Add search operators to result dict."""
+        if self.text is not None:
+            result["text"] = self.text
+        if self.compound is not None:
+            result["compound"] = self.compound
+        if self.autocomplete is not None:
+            result["autocomplete"] = self.autocomplete
+        if self.phrase is not None:
+            result["phrase"] = self.phrase
+        if self.wildcard is not None:
+            result["wildcard"] = self.wildcard
+        if self.regex is not None:
+            result["regex"] = self.regex
+        if self.near is not None:
+            result["near"] = self.near
+        if self.range is not None:
+            result["range"] = self.range
+
+    def _add_advanced(self, result: dict[str, Any]) -> None:
+        """Add advanced search options to result dict."""
+        if self.exists is not None:
+            result["exists"] = self.exists
+        if self.equals is not None:
+            result["equals"] = self.equals
+        if self.more_like_this is not None:
+            result["moreLikeThis"] = self.more_like_this
+        if self.query_string is not None:
+            result["queryString"] = self.query_string
+        if self.highlight is not None:
+            result["highlight"] = self.highlight
+        if self.count is not None:
+            result["count"] = self.count
+        if self.return_stored_source is not None:
+            result["returnStoredSource"] = self.return_stored_source
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        if self.index is not None:
+            result["index"] = self.index
+        self._add_operators(result)
+        self._add_advanced(result)
+        return {"$search": result}
+
+
+class SearchMeta(BaseModel):
+    """
+    $searchMeta stage - returns Atlas Search metadata.
+
+    Example:
+        >>> SearchMeta(index="default", count={"type": "total"}).model_dump()
+        {"$searchMeta": {"index": "default", "count": {"type": "total"}}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    index: str | None = Field(
+        default=None,
+        description="Name of the Atlas Search index",
+    )
+    count: dict[str, Any] | None = Field(
+        default=None,
+        description="Count options",
+    )
+    facet: dict[str, Any] | None = Field(
+        default=None,
+        description="Facet options",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        if self.index is not None:
+            result["index"] = self.index
+        if self.count is not None:
+            result["count"] = self.count
+        if self.facet is not None:
+            result["facet"] = self.facet
+        return {"$searchMeta": result}
+
+
+class VectorSearch(BaseModel):
+    """
+    $vectorSearch stage - Atlas vector search (MongoDB 7.0.2+).
+
+    Example:
+        >>> VectorSearch(
+        ...     index="vector_index",
+        ...     path="embedding",
+        ...     query_vector=[0.1, 0.2, 0.3],
+        ...     num_candidates=100,
+        ...     limit=10
+        ... ).model_dump()
+        {"$vectorSearch": {
+            "index": "vector_index",
+            "path": "embedding",
+            "queryVector": [0.1, 0.2, 0.3],
+            "numCandidates": 100,
+            "limit": 10
+        }}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    index: str = Field(
+        ...,
+        description="Name of the Atlas Vector Search index",
+    )
+    path: str = Field(
+        ...,
+        description="Field path containing the vector",
+    )
+    query_vector: list[float] = Field(
+        ...,
+        serialization_alias="queryVector",
+        description="Query vector for similarity search",
+    )
+    num_candidates: int = Field(
+        ...,
+        serialization_alias="numCandidates",
+        description="Number of candidates to consider",
+    )
+    limit: int = Field(
+        ...,
+        description="Maximum number of results to return",
+    )
+    filter: dict[str, Any] | None = Field(
+        default=None,
+        description="Pre-filter for vector search",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {
+            "index": self.index,
+            "path": self.path,
+            "queryVector": self.query_vector,
+            "numCandidates": self.num_candidates,
+            "limit": self.limit,
+        }
+        if self.filter is not None:
+            result["filter"] = self.filter
+        return {"$vectorSearch": result}
+
+
+class RankFusion(BaseModel):
+    """
+    $rankFusion stage - combines ranked results from multiple pipelines.
+
+    Example:
+        >>> RankFusion(
+        ...     input={"search": [...], "vector": [...]},
+        ...     combination={"weights": {"search": 0.7, "vector": 0.3}}
+        ... ).model_dump()
+        {"$rankFusion": {
+            "input": {"search": [...], "vector": [...]},
+            "combination": {"weights": {"search": 0.7, "vector": 0.3}}
+        }}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    input: dict[str, list[dict[str, Any]]] = Field(
+        ...,
+        description="Named input pipelines",
+    )
+    combination: dict[str, Any] | None = Field(
+        default=None,
+        description="Combination options",
+    )
+    score_details: bool | None = Field(
+        default=None,
+        serialization_alias="scoreDetails",
+        description="Include score details",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {"input": self.input}
+        if self.combination is not None:
+            result["combination"] = self.combination
+        if self.score_details is not None:
+            result["scoreDetails"] = self.score_details
+        return {"$rankFusion": result}
+
+
+__all__ = [
+    "Search",
+    "SearchMeta",
+    "VectorSearch",
+    "ListSearchIndexes",
+    "RankFusion",
+]

mongo_aggro/stages/session.py

@@ -0,0 +1,111 @@
+"""Session-related MongoDB aggregation pipeline stages.
+
+This module contains stages for listing sessions and sampled queries:
+ListSessions, ListLocalSessions, and ListSampledQueries.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class ListSessions(BaseModel):
+    """
+    $listSessions stage - lists all sessions in system.sessions.
+
+    Example:
+        >>> ListSessions().model_dump()
+        {"$listSessions": {}}
+
+        >>> ListSessions(users=[{"user": "admin", "db": "admin"}]).model_dump()
+        {"$listSessions": {"users": [{"user": "admin", "db": "admin"}]}}
+
+        >>> ListSessions(all_users=True).model_dump()
+        {"$listSessions": {"allUsers": True}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    users: list[dict[str, str]] | None = Field(
+        default=None,
+        description="List of users to filter sessions",
+    )
+    all_users: bool | None = Field(
+        default=None,
+        serialization_alias="allUsers",
+        description="Return sessions for all users",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        if self.users is not None:
+            result["users"] = self.users
+        if self.all_users is not None:
+            result["allUsers"] = self.all_users
+        return {"$listSessions": result}
+
+
+class ListLocalSessions(BaseModel):
+    """
+    $listLocalSessions stage - lists local sessions (db.aggregate only).
+
+    Example:
+        >>> ListLocalSessions().model_dump()
+        {"$listLocalSessions": {}}
+
+        >>> ListLocalSessions(all_users=True).model_dump()
+        {"$listLocalSessions": {"allUsers": True}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    users: list[dict[str, str]] | None = Field(
+        default=None,
+        description="List of users to filter sessions",
+    )
+    all_users: bool | None = Field(
+        default=None,
+        serialization_alias="allUsers",
+        description="Return sessions for all users",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        if self.users is not None:
+            result["users"] = self.users
+        if self.all_users is not None:
+            result["allUsers"] = self.all_users
+        return {"$listLocalSessions": result}
+
+
+class ListSampledQueries(BaseModel):
+    """
+    $listSampledQueries stage - lists sampled queries.
+
+    Example:
+        >>> ListSampledQueries().model_dump()
+        {"$listSampledQueries": {}}
+
+        >>> ListSampledQueries(namespace="db.collection").model_dump()
+        {"$listSampledQueries": {"namespace": "db.collection"}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    namespace: str | None = Field(
+        default=None,
+        description="Namespace to filter sampled queries",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        if self.namespace is not None:
+            result["namespace"] = self.namespace
+        return {"$listSampledQueries": result}
+
+
+__all__ = [
+    "ListSessions",
+    "ListLocalSessions",
+    "ListSampledQueries",
+]

mongo_aggro/stages/stats.py

@@ -0,0 +1,152 @@
+"""Statistics and diagnostics MongoDB aggregation pipeline stages.
+
+This module contains stages for collection/index statistics and operations:
+CollStats, IndexStats, PlanCacheStats, and CurrentOp.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class CollStats(BaseModel):
+    """
+    $collStats stage - returns collection statistics.
+
+    Example:
+        >>> CollStats(lat_stats={"histograms": True}).model_dump()
+        {"$collStats": {"latencyStats": {"histograms": True}}}
+
+        >>> CollStats(storage_stats={}).model_dump()
+        {"$collStats": {"storageStats": {}}}
+
+        >>> CollStats(count={}).model_dump()
+        {"$collStats": {"count": {}}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    lat_stats: dict[str, Any] | None = Field(
+        default=None,
+        serialization_alias="latencyStats",
+        description="Latency statistics options",
+    )
+    storage_stats: dict[str, Any] | None = Field(
+        default=None,
+        serialization_alias="storageStats",
+        description="Storage statistics options",
+    )
+    count: dict[str, Any] | None = Field(
+        default=None,
+        description="Document count options",
+    )
+    query_exec_stats: dict[str, Any] | None = Field(
+        default=None,
+        serialization_alias="queryExecStats",
+        description="Query execution statistics options",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        if self.lat_stats is not None:
+            result["latencyStats"] = self.lat_stats
+        if self.storage_stats is not None:
+            result["storageStats"] = self.storage_stats
+        if self.count is not None:
+            result["count"] = self.count
+        if self.query_exec_stats is not None:
+            result["queryExecStats"] = self.query_exec_stats
+        return {"$collStats": result}
+
+
+class IndexStats(BaseModel):
+    """
+    $indexStats stage - returns index usage statistics.
+
+    Example:
+        >>> IndexStats().model_dump()
+        {"$indexStats": {}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$indexStats": {}}
+
+
+class PlanCacheStats(BaseModel):
+    """
+    $planCacheStats stage - returns plan cache information.
+
+    Example:
+        >>> PlanCacheStats().model_dump()
+        {"$planCacheStats": {}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$planCacheStats": {}}
+
+
+class CurrentOp(BaseModel):
+    """
+    $currentOp stage - returns current operations (db.aggregate only).
+
+    Example:
+        >>> CurrentOp().model_dump()
+        {"$currentOp": {}}
+
+        >>> CurrentOp(all_users=True, idle_connections=True).model_dump()
+        {"$currentOp": {"allUsers": True, "idleConnections": True}}
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="forbid")
+
+    all_users: bool | None = Field(
+        default=None,
+        serialization_alias="allUsers",
+        description="Return operations for all users",
+    )
+    idle_connections: bool | None = Field(
+        default=None,
+        serialization_alias="idleConnections",
+        description="Include idle connections",
+    )
+    idle_cursors: bool | None = Field(
+        default=None,
+        serialization_alias="idleCursors",
+        description="Include idle cursors",
+    )
+    idle_sessions: bool | None = Field(
+        default=None,
+        serialization_alias="idleSessions",
+        description="Include idle sessions",
+    )
+    local_ops: bool | None = Field(
+        default=None,
+        serialization_alias="localOps",
+        description="Return local operations only",
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        if self.all_users is not None:
+            result["allUsers"] = self.all_users
+        if self.idle_connections is not None:
+            result["idleConnections"] = self.idle_connections
+        if self.idle_cursors is not None:
+            result["idleCursors"] = self.idle_cursors
+        if self.idle_sessions is not None:
+            result["idleSessions"] = self.idle_sessions
+        if self.local_ops is not None:
+            result["localOps"] = self.local_ops
+        return {"$currentOp": result}
+
+
+__all__ = [
+    "CollStats",
+    "IndexStats",
+    "PlanCacheStats",
+    "CurrentOp",
+]