stackraise 0.1.0__py3-none-any.whl

stackraise/db/document.py

@@ -0,0 +1,282 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Annotated, Any, Awaitable, Callable, ClassVar, Mapping, Optional, Self
+
+from pydantic import Field
+from pydantic_core import core_schema
+
+from .adapter import Adapter, DocumentProtocol
+from .collection import Collection
+from .id import Id
+from .exceptions import NotFoundError
+from .protocols import QueryLike, ensure_mongo_query
+
+import stackraise.model as model
+
+
+class DocumentMeta(type(model.Dto)):
+    pass
+
+
+class Document(model.Dto, metaclass=DocumentMeta):
+    """
+    Base class for all documents in the database.
+    This class provides a common interface for all documents, including
+    methods for storing, updating, and deleting documents.
+    """
+
+    class Reference[T: DocumentProtocol]:
+        __slots__ = ("_id",)
+        collection: ClassVar[Collection[T]]
+        _id: Id
+        _sync_item: ClassVar[T|None] = None
+
+        @classmethod
+        def __get_pydantic_core_schema__(cls, *_):
+
+            def js_parse(value: str):
+                assert isinstance(value, str), f"Bad value type {type(value)}"
+                return cls(Id(value))
+
+            def js_serial(ref: Document.Reference):
+                return str(ref.id)
+
+            def py_parse(value: Document.Reference | Id | str):
+                if isinstance(value, Document.Reference):
+                    assert type(value) is cls, f"Bad type {type(value)} for {cls}"
+                    return value
+                if not isinstance(value, Id):
+                    value = Id(value)
+                return cls(value)
+
+            def py_serial(ref: Document.Reference):
+                return ref.id
+
+            return core_schema.json_or_python_schema(
+                # JSON
+                json_schema=core_schema.no_info_plain_validator_function(
+                    js_parse,
+                    serialization=core_schema.plain_serializer_function_ser_schema(
+                        js_serial
+                    ),
+                ),
+                # PYTHON
+                python_schema=core_schema.no_info_plain_validator_function(
+                    py_parse,
+                    serialization=core_schema.plain_serializer_function_ser_schema(
+                        py_serial
+                    ),
+                ),
+            )
+
+        @classmethod
+        def __get_pydantic_json_schema__(cls, _, handler):
+            return handler(core_schema.str_schema())
+
+        def __repr__(self):
+            return f"{type(self).__qualname__}({self._id!s})"
+
+        def __init__(self, id: Id):
+            assert isinstance(id, Id), f"Bad id type {type(id)}"
+            self._id = id
+
+        def __eq__(self, other) -> bool:
+            if not isinstance(other, type(self)):
+                return False
+            return self._id == other._id
+
+        def __hash__(self) -> int:
+            return hash(self._id)
+
+        @property
+        def id(self):
+            return self._id
+
+        @property
+        def created_at(self) -> datetime:
+            return self.id.generation_time
+        
+        def to_mongo_query(self) -> Mapping[str, Any]:
+            """Return a MongoDB query for this reference."""
+            return {"_id": self._id}
+
+        def sync_fetch(self) -> T:
+            """Fetch the document referenced by this reference."""
+            return self._sync_item
+
+        async def fetch(self, not_found_error: bool=True) -> T:
+            """Fetch the document referenced by this reference."""
+            if (sync_item := self._sync_item) is not None:
+                return sync_item
+            return await self.collection.fetch_by_id(self._id, not_found_error=not_found_error)
+
+        async def delete(self):
+            """Delete the document referenced by this reference."""
+            return await self.collection.delete_by_id(self._id)
+
+        async def exists(self) -> bool:
+            """ Check if the document exists in the database. """
+            count = await self.collection.count(self)
+            return count == 1
+
+        async def complies(self, conditions: QueryLike) -> bool:
+            """
+            Check if the document referenced by this reference complies with the given conditions.
+            Args:
+                conditions (QueryLike): The conditions to check against the document.
+            Returns:
+                bool: True if the document complies with the conditions, False otherwise.
+            """
+            query = ensure_mongo_query(conditions)
+            count = await self.collection.count(query | {"_id": self._id})
+            return count == 1
+
+        async def assign(self, **values):
+            """
+            Assign values to the document referenced by this reference.
+            Args:
+                **values: The values to assign to the document.
+            """
+
+            result = await self.collection._update_one(
+                {"_id": self._id},
+                {"$set": values},
+                upsert=True,
+            )
+
+            if result.upserted_id != self.id:
+                raise NotFoundError(self)
+
+        async def update(self, **values):
+            result = await self.collection._update_one(
+                {"_id": self._id}, {"$set": values}, upsert=False
+            )
+
+            if result.matched_count != 1:
+                raise NotFoundError(self)
+
+        ## mutate() as doc creara un contexto de mutacion
+        ## doc puede ser valuado leido y escrito de diferentes maneras
+        ## despues del contexto de mutacion la operacion será realizada
+
+    def __init_subclass__(
+        cls,
+        abstract=False,
+        collection: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init_subclass__(**kwargs)
+        ##
+        ## el subclassing de Document tiene propieadades especiales:
+        ##  - Se crea una clase de referencia para el nuevo documento
+        ##  - Se instancia un repositorio para el nuevo documento
+        ##
+
+        # Make a repository class for the new persistent class
+        # Instance the repository
+
+        collection = Collection(Adapter(cls), collection)
+
+        # Make a reference class for the new persistent class
+
+        reference_cls = type(
+            "Reference",
+            tuple(
+                vars(base)["Reference"]
+                for base in cls.__mro__
+                if "Reference" in vars(base)
+                and issubclass(vars(base)["Reference"], Document.Reference)
+            ),
+            {
+                "__module__": cls.__module__,
+                "__qualname__": f"{cls.__qualname__}.Reference",
+                "document_class": cls,
+                "collection": collection,
+            },
+        )
+
+        setattr(cls, "Reference", reference_cls)
+        setattr(cls, "Ref", reference_cls)
+        setattr(cls, "collection", collection)
+
+    type Ref = Reference[Self]
+    collection: ClassVar[Collection[Self]]
+
+    id: Annotated[
+        Optional[Id],
+        Field(
+            None,
+            title="Unique object Id",
+            alias="_id",
+        ),
+    ]
+
+    # @computed_field
+    # def kind(self) -> str:
+    #     """Return the kind of the document."""
+    #     return self.__class__.__name__
+
+    @property
+    def ref(self) -> Ref | None:
+        """Return a reference to the document."""
+        if self.id is None:
+            return None
+        return self.Reference(self.id)
+
+    async def __prepare_for_storage__(self):
+        """
+        Hook to be called before persisting the object.
+        This method can be overridden in subclasses to perform custom validation
+        or processing before the object is stored in the database.
+        For example, you can check for uniqueness constraints or perform
+        additional validation on the object's attributes.
+        This method is called automatically by the `insert` and `update` methods.
+
+        TODO: esto es parte del document protocol
+        """
+
+    @classmethod
+    async def __handle_post_deletion__(cls, ref: Ref):
+        """
+        Hook to be called after the document is deleted.
+        This method can be overridden in subclasses to perform custom cleanup
+        or processing after the object is deleted from the database.
+        For example, you can delete related documents or perform additional
+        cleanup tasks.
+        This method is called automatically by the `delete` method.
+        """
+
+    async def store(self) -> Self:
+        """Persist the object in the database.shipment is
+
+        This method will use the insert procedure if the object does not have
+        an identifier defined (id is None) or the update(upsert) procedure if
+        the object already has an identifier.
+
+        Returns:
+            RefBase[Self]: A reference to the persisted object.
+        """
+        if self.id:
+            return await self.update()
+        else:
+            return await self.insert()
+
+    async def insert(self, *, with_id: Optional[Id] = None) -> Self:
+        return await self.collection.insert_item(self, with_id=with_id)
+
+    async def update(self):
+        return await self.collection.update_item(self)
+
+    async def delete(self):
+        """
+        Deletes the document from the database.
+        """
+        doc_id = self.id
+        if doc_id is None:
+            raise NotFoundError(self)
+            
+        self.id = None
+        await self.collection.delete_by_id(doc_id)
+
+

stackraise/db/exceptions.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+import stackraise.db as db
+
+class NotFoundError(Exception):
+    """Exception raised when a document is not found in the database."""
+
+    def __init__(self, ref: db.Document.Ref):
+        super().__init__(f"Document with reference {ref} not found.")
+        self.ref = ref

stackraise/db/id.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+from typing_extensions import deprecated
+
+from bson import ObjectId
+from bson.errors import InvalidId
+from pydantic_core import core_schema, SchemaSerializer
+from stackraise.model.validation import validation_error
+
+__all__ = ["Id"]
+
+class Id(ObjectId):
+    _GENERATE_NEW = object()  # Sentinel para generar nuevo ID
+    
+    def __init__(self, value=None):
+        # Permitir generación explícita de nuevo ID con sentinel
+        if value is Id._GENERATE_NEW:
+            super().__init__()
+            return
+        # Evitar que None o vacío generen un ObjectId aleatorio accidentalmente
+        if value is None:
+            raise ValueError("Valor no seleccionado")
+        try:
+            super().__init__(value)
+        except InvalidId as e:
+            raise ValueError(*e.args) from InvalidId
+
+    @classmethod
+    def new(cls):
+        """Genera un nuevo ObjectId único."""
+        return cls(cls._GENERATE_NEW)
+
+    @classmethod
+    @deprecated("you must not use this method")
+    def from_str(cls, val: str):
+        assert isinstance(val, str), f"from_str receive {val}"
+        return cls(val)
+
+    @classmethod
+    @deprecated("you must not use this method")
+    def from_oid(cls, val: Id | ObjectId | str):
+        # assert isinstance(val, ObjectId), f"Id from python is {val}"
+        return val if isinstance(val, cls) else cls(val)
+
+    @property
+    def created_at(self):
+        return self.generation_time
+
+    def to_mongo_query(self): # query protocol
+        return {"_id": self}
+    
+    @property
+    def value(self) -> str:
+        """Return the string representation of the Id"""
+        return str(self)
+
+    SCHEMA = core_schema.json_or_python_schema(
+        # JSON
+        json_schema=core_schema.no_info_plain_validator_function(
+            lambda s: Id.from_str(s),
+        ),
+        # PYTHON
+        python_schema=core_schema.no_info_plain_validator_function(
+            lambda v: Id.from_oid(v),
+        ),
+        serialization=core_schema.plain_serializer_function_ser_schema(
+            str, when_used="json"  # as str
+        ),
+    )
+
+    __pydantic_serializer__ = SchemaSerializer(SCHEMA)
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, *_):
+        return cls.SCHEMA
+
+    @classmethod
+    def __get_pydantic_json_schema__(cls, _, handler):
+        # return {"type": "string"}
+        return handler(core_schema.str_schema())

stackraise/db/index.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from logging import getLogger as get_logger
+
+#import stackraise.db as db
+import stackraise.inflection as inflection
+from pymongo.asynchronous.client_session import \
+    AsyncClientSession as MongoSession
+from pymongo.asynchronous.collection import AsyncCollection as MongoCollection
+
+from .protocols import DocumentProtocol
+
+log = get_logger(__name__)
+
+_indices: dict[type[DocumentProtocol], list[Index]] = {}
+
+
+class Index:
+    def __init__(self, args: list[str], kwargs: dict[str, int] = {}, **options):
+        fields = {field: 1 for field in args}
+        fields.update(kwargs)
+        self.fields = {inflection.to_camelcase(field): value for field, value in fields.items()}
+        self.options = options
+
+    def __call__[T: type[DocumentProtocol]](self, cls: T) -> T:
+        _indices.setdefault(cls, []).append(self)
+        return cls
+
+def index(*args, **kwargs: int) -> Index:
+    return Index(args, kwargs, unique=False)
+
+def unique_index(*args, **kwargs: int) -> Index:
+    return Index(args, unique=True)
+
+def text_index(*fields, **options) -> Index:
+    return Index([], {field: 'text' for field in fields}, **options)
+
+
+async def _update_indices(
+    document_class: type[DocumentProtocol],
+    collection: MongoCollection,
+    session: MongoSession,
+):
+    """
+    Applies the indices to the collection.
+    """
+
+    desired_indices = _indices.get(document_class, [])
+
+    existing_indices = await collection.index_information()
+    # Map existing indices by their key tuple (excluding _id_)
+    existing_keys = {
+        tuple(idx['key']): name
+        for name, idx in existing_indices.items()
+        if name != '_id_'
+    }
+
+    # Map desired indices by their key tuple
+    desired_keys = {}
+    for idx in desired_indices:
+        # Convert fields dict to tuple of (field, direction) preserving order
+        key_tuple = tuple(idx.fields.items())
+        desired_keys[key_tuple] = idx
+
+    # Indices to drop: present in existing but not in desired
+    to_drop = [existing_keys[k] for k in existing_keys if k not in desired_keys]
+    # Indices to create: present in desired but not in existing
+    to_create = [desired_keys[k] for k in desired_keys if k not in existing_keys]
+
+    # Drop obsolete indices
+    for name in to_drop:
+        log.debug(f"Dropping index {name} from collection {collection.name}")
+        await collection.drop_index(name, session=session)
+
+    # Create new indices
+    for idx in to_create:
+        log.debug(f"Creating index {idx.fields} on collection {collection.name}")
+        await collection.create_index(
+            list(idx.fields.items()),
+            session=session,
+            **idx.options
+        )
+

stackraise/db/persistence.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+from anyio import create_task_group
+from contextlib import asynccontextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field as dc_field
+from functools import cached_property, wraps
+from typing import Annotated, Awaitable, Callable, Optional, Tuple, TypedDict
+
+from gridfs import AsyncGridFS
+
+from pydantic import BaseModel, Field, MongoDsn
+from pymongo import AsyncMongoClient
+from pymongo.asynchronous.database import AsyncDatabase
+
+from pymongo.asynchronous.client_session import AsyncClientSession
+from pymongo.errors import OperationFailure
+from pymongo.read_concern import ReadConcern
+from pymongo.read_preferences import ReadPreference
+from pymongo.write_concern import WriteConcern
+from pymongo import AsyncMongoClient
+
+import stackraise.event as ev
+import stackraise.model as model
+import stackraise.db as db
+
+from .protocols import get_collection_instances
+
+_current_context: ContextVar[Optional[Tuple[Persistence, AsyncClientSession]]] = (
+    ContextVar(
+        f"{__name__}._current_context",
+        default=None,
+    )
+)
+
+_startup_tasks: Callable[[], Awaitable] = []
+
+
+class ChangeEvent(TypedDict):
+    op: str
+    collection: str
+    refs: list[str]
+
+
+change_event_emitter: ev.EventEmitter[ChangeEvent] = ev.EventEmitter(
+    "persistence.change_event"
+)
+
+
+@dataclass(frozen=True)
+class Persistence:
+    """
+    A class representing the persistence layer for the ctrl.
+    """
+
+    class Settings(BaseModel):
+        """Persistence layer settings"""
+
+        mongo_dsn: Annotated[
+            MongoDsn,
+            Field(
+                MongoDsn("mongodb://localhost/test"),
+                title="Mongo database DSN",
+                description="The DSN for the MongoDB database. ",
+            ),
+        ]
+
+        direct_connection: Annotated[
+            bool,
+            Field(
+                True,
+                title="Direct connection",
+                description="Whether to connect directly to the MongoDB instance.",
+            ),
+        ]
+
+        causal_consistency: Annotated[
+            bool,
+            Field(
+                True,
+                title="Causal consistency",
+                description="Whether to enable causal consistency.",
+            ),
+        ]
+
+        # TODO: Default transaction options
+
+    settings: Settings = dc_field(default_factory=Settings)
+
+    @cached_property
+    def client(self):
+        # print(str(self.settings.mongo_dsn))
+        return AsyncMongoClient(
+            str(self.settings.mongo_dsn),
+            directConnection=self.settings.direct_connection,
+        )
+        # return AsyncMongoClient()
+
+    @cached_property
+    def database(self) -> AsyncDatabase:
+        # TODO: read / write preference and concern from settings
+        return self.client.get_default_database()
+
+    @cached_property
+    def fs(self):
+        return AsyncGridFS(self.database)
+
+    @asynccontextmanager
+    async def lifespan(self):
+        async with self.client:
+            async with self.session():
+
+                async with create_task_group() as tg:
+                    tg.start_soon(self._watch_task)
+
+                    for collection in get_collection_instances():
+                        tg.start_soon(collection._startup_task, self)
+
+                    yield
+
+                    tg.cancel_scope.cancel()
+
+    async def _watch_task(self):
+        """
+        Background task to watch changes in the database.
+        Only works with MongoDB replica sets. Gracefully disables on standalone instances.
+        """
+        try:
+            async with await self.database.watch() as change_stream:
+                async for change in change_stream:
+
+                    # Emit change events
+                    change_event = ChangeEvent(
+                        op=change["operationType"],
+                        collection=change["ns"]["coll"],
+                        refs=[str(change["documentKey"]["_id"])],
+                    )
+
+                    await change_event_emitter.emit(change_event)
+        except OperationFailure as e:
+            # Change Streams require replica set (code 40573)
+            if e.code == 40573:
+                print(
+                    "⚠️  MongoDB Change Streams disabled: running in standalone mode (replica set required)"
+                )
+                return  # Exit gracefully
+            else:
+                raise
+
+    @asynccontextmanager
+    async def session(self):
+        """Enter the persistence session context."""
+        async with self.client.start_session(
+            causal_consistency=self.settings.causal_consistency,
+            default_transaction_options=None,  # TODO: From settings
+        ) as session:
+            session_token = _current_context.set((self, session))
+            try:
+                yield session
+            finally:
+                _current_context.reset(session_token)
+
+
+def current_context():
+    """
+    Get the current persistence context.
+    """
+    context = _current_context.get()
+    if context is None:
+        raise RuntimeError("No persistence context is currently set.")
+    return context
+
+
+def current_database() -> AsyncDatabase:
+    """
+    Get the current database from the context.
+    """
+    persistence, _ = current_context()
+    return persistence.database
+
+
+def current_fs() -> AsyncGridFS:
+    """
+    Get the current GridFS instance from the context.
+    """
+    persistence, _ = current_context()
+    return persistence.fs
+
+
+def current_session() -> AsyncClientSession:
+    """
+    Get the current session from the context.
+    """
+    _, session = _current_context.get()
+    return session
+
+
+def in_transaction() -> bool:
+    """
+    Check if the current session is in a transaction.
+    """
+    session = current_session()
+    return session is not None and session.in_transaction
+
+
+def transaction(
+    read_concern: ReadConcern | None = None,
+    write_concern: WriteConcern | None = None,
+    read_preference: ReadPreference | None = None,
+    max_commit_time_ms: int | None = None,
+):
+    """
+    Decorator to run a function within a transaction.
+    """
+
+    def decorator(fn):
+
+        @wraps(fn)
+        async def wrapper(*args, **kwargs):
+            session = current_session()
+
+            if session.in_transaction:
+                return await coro(*args, **kwargs)
+
+            async def coro(_session):
+                return await fn(*args, **kwargs)
+
+            return await session.with_transaction(
+                coro,
+                read_concern=read_concern,
+                write_concern=write_concern,
+                read_preference=read_preference,
+                max_commit_time_ms=max_commit_time_ms,
+            )
+
+        return wrapper
+
+    return decorator