mindtrace-database 0.10.0__tar.gz → 0.11.0__tar.gz

{mindtrace_database-0.10.0 → mindtrace_database-0.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mindtrace-database
-Version: 0.10.0
+Version: 0.11.0
 Summary: Database functionality for Mindtrace
 Author: Mindtrace Team
 License-Expression: Apache-2.0
@@ -11,13 +11,13 @@ Classifier: Programming Language :: Python :: 3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: beanie<2,>=1.29.0
-Requires-Dist: mindtrace-core>=0.10.0
+Requires-Dist: mindtrace-core>=0.11.0
 Requires-Dist: pydantic>=2.11.1
 Requires-Dist: redis>=4.0.0
 Requires-Dist: redis-om>=0.3.5
 Requires-Dist: motor>=3.3.0
 Requires-Dist: pymongo>=4.14.0
-Requires-Dist: mindtrace-registry>=0.10.0
+Requires-Dist: mindtrace-registry>=0.11.0
 Dynamic: license-file
 
 [![PyPI version](https://img.shields.io/pypi/v/mindtrace-database)](https://pypi.org/project/mindtrace-database/)

{mindtrace_database-0.10.0 → mindtrace_database-0.11.0}/mindtrace/database/backends/mongo_odm.py

@@ -1,9 +1,14 @@
 import asyncio
-from typing import Dict, List, Optional, Type, TypeVar
+from collections.abc import AsyncIterator, Iterator
+from typing import Any, Dict, List, Optional, Type, TypeVar
 
 from beanie import Document, PydanticObjectId, init_beanie
+from beanie.odm.utils.dump import get_dict
+from bson import ObjectId
+from bson.errors import InvalidId
 from motor.motor_asyncio import AsyncIOMotorClient
 from pydantic import BaseModel
+from pymongo import MongoClient
 from pymongo.errors import DuplicateKeyError
 
 from mindtrace.database.backends.mindtrace_odm import InitMode, MindtraceODM
@@ -44,6 +49,12 @@ class MindtraceDocument(Document):
 
 ModelType = TypeVar("ModelType", bound=MindtraceDocument)
 
+# Beanie binds each Document class to one database *process-wide*. A second ``init_beanie`` for the same
+# ``model_cls`` (whether another ``db_name``, another URI, or another ``MongoMindtraceODM`` instance)
+# would re-bind that class and break earlier lakes. Only the first ODM per ``model_cls`` runs
+# ``init_beanie``; subsequent backends set ``_motor_routing`` and use Motor on ``self.client[self.db_name]``.
+_BEANIE_INITIALIZED_MODEL_CLASSES: set[Type[Any]] = set()
+
 
 class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
     """
@@ -87,6 +98,8 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
         allow_index_dropping: bool = False,
         auto_init: bool = False,
         init_mode: InitMode | None = None,
+        *,
+        client: AsyncIOMotorClient | None = None,
     ):
         """
         Initialize the MongoDB ODM backend.
@@ -127,8 +140,17 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
         if not db_uri or not db_name:
             raise ValueError("db_uri and db_name are required")
 
-        self.client = AsyncIOMotorClient(db_uri)
+        self.db_uri = db_uri
+        if client is not None:
+            self.client = client
+            self._owns_client = False
+        else:
+            self.client = AsyncIOMotorClient(db_uri)
+            self._owns_client = True
+        self._sync_client: MongoClient | None = None
+        self._closed = False
         self.db_name = db_name
+        self._motor_routing = False
         self._allow_index_dropping = allow_index_dropping
         self._is_initialized = False
         self._model_odms: Dict[str, "MongoMindtraceODM"] = {}
@@ -152,9 +174,8 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
                     allow_index_dropping=allow_index_dropping,
                     auto_init=False,  # We'll initialize all together
                     init_mode=init_mode,
+                    client=self.client,
                 )
-                # Share the same client instance
-                odm.client = self.client
                 # Store parent reference for initialization delegation
                 odm._parent_odm = self
                 self._model_odms[name] = odm
@@ -192,6 +213,56 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
             # Defer initialization - operations will auto-init on first use
             self._needs_init = True
 
+    def _collection_name(self) -> str:
+        if self.model_cls is None:
+            raise ValueError("model_cls is required for collection name resolution")
+        return getattr(self.model_cls.Settings, "name")
+
+    def _motor_collection(self):
+        return self.client[self.db_name][self._collection_name()]
+
+    def _sync_collection(self):
+        if self._sync_client is None:
+            self._sync_client = MongoClient(self.db_uri)
+        return self._sync_client[self.db_name][self._collection_name()]
+
+    def _mongo_doc_to_model(self, doc: dict | None) -> T | None:
+        if doc is None:
+            return None
+        payload = dict(doc)
+        if "_id" in payload and "id" not in payload:
+            payload["id"] = payload["_id"]
+        return self.model_cls.model_validate(payload)
+
+    def _motor_model_dump(self, doc: BaseModel) -> dict[str, Any]:
+        """Serialize a model for raw Motor insert/replace.
+
+        Pydantic v2 ``model_dump`` on Beanie ``Document`` can fail (class-level
+        descriptors such as ``ExpressionField``). Beanie's BSON ``get_dict`` matches
+        ``Document.insert`` and avoids that serializer path.
+        """
+        if isinstance(doc, Document):
+            return get_dict(doc, to_db=True, keep_nulls=doc.get_settings().keep_nulls)
+        return doc.model_dump(mode="python", by_alias=True, serialize_as_any=True)
+
+    def _motor_patch_fields(self, obj: BaseModel) -> dict[str, Any]:
+        """Shallow field map from ``obj`` for merge-into-document updates (Motor mode)."""
+        if isinstance(obj, Document):
+            return {name: getattr(obj, name) for name in obj.model_fields if name != "id"}
+        return obj.model_dump(exclude={"id"}, serialize_as_any=True)
+
+    @staticmethod
+    def _to_object_id(id: str | PydanticObjectId | ObjectId) -> ObjectId:
+        # ``PydanticObjectId`` subclasses ``ObjectId``; check it first so we normalize via ``str``.
+        if isinstance(id, PydanticObjectId):
+            return ObjectId(str(id))
+        if isinstance(id, ObjectId):
+            return id
+        try:
+            return ObjectId(str(id))
+        except InvalidId as exc:
+            raise DocumentNotFoundError(f"Invalid document id {id!r}") from exc
+
     async def _do_initialize(self):
         """Internal method to perform the actual initialization."""
         if not self._is_initialized:
@@ -199,7 +270,12 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
                 # Multi-model mode: initialize all models together
                 document_models = list(self._models.values())
             else:
-                # Single model mode
+                # Single model mode: at most one init_beanie per Document subclass in this process.
+                if self.model_cls in _BEANIE_INITIALIZED_MODEL_CLASSES:
+                    self._motor_routing = True
+                    self._is_initialized = True
+                    return
+                _BEANIE_INITIALIZED_MODEL_CLASSES.add(self.model_cls)
                 document_models = [self.model_cls]
 
             await init_beanie(
@@ -253,6 +329,31 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
             self._allow_index_dropping = allow_index_dropping
         await self._do_initialize()
 
+    def close(self) -> None:
+        """Release resources held by this ODM.
+
+        Closes the lazily-created synchronous client if present, and closes
+        the owning motor client if this ODM created it. ODMs constructed with
+        an injected ``client=`` hold a non-owning reference and do not close
+        it on shutdown — responsibility stays with whoever created the client.
+
+        Idempotent: safe to call multiple times.
+        """
+        if self._closed:
+            return
+        if self._sync_client is not None:
+            try:
+                self._sync_client.close()
+            except Exception:
+                pass
+            self._sync_client = None
+        if self._owns_client and self.client is not None:
+            try:
+                self.client.close()
+            except Exception:
+                pass
+        self._closed = True
+
     def __getattr__(self, name: str):
         """Support attribute-based access to model-specific ODMs in multi-model mode.
 
@@ -340,11 +441,22 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
         doc.id = None
 
         try:
+            if self._motor_routing:
+                raw = self._motor_model_dump(doc)
+                raw.pop("_id", None)
+                if raw.get("id") is None:
+                    raw.pop("id", None)
+                result = await self._motor_collection().insert_one(raw)
+                inserted = await self._motor_collection().find_one({"_id": result.inserted_id})
+                if inserted is None:
+                    raise DuplicateInsertError("insert_one did not persist a readable document")
+                out = self._mongo_doc_to_model(inserted)
+                if out is None:
+                    raise DuplicateInsertError("Could not deserialize inserted document")
+                return out
             return await doc.insert()
         except DuplicateKeyError as e:
             raise DuplicateInsertError(f"Duplicate key error: {str(e)}")
-        except Exception as e:
-            raise DuplicateInsertError(str(e))
 
     async def get(self, id: str | PydanticObjectId, fetch_links: bool = False) -> T:
         """
@@ -381,6 +493,18 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
         if not self._is_initialized:
             await self.initialize()
 
+        if self._motor_routing:
+            if fetch_links:
+                raise NotImplementedError("fetch_links is not supported for this database routing mode")
+            oid = self._to_object_id(id)
+            doc = await self._motor_collection().find_one({"_id": oid})
+            if not doc:
+                raise DocumentNotFoundError(f"Object with id {id} not found")
+            out = self._mongo_doc_to_model(doc)
+            if out is None:
+                raise DocumentNotFoundError(f"Object with id {id} not found")
+            return out
+
         doc = await self.model_cls.get(id, fetch_links=fetch_links)
         if not doc:
             raise DocumentNotFoundError(f"Object with id {id} not found")
@@ -421,6 +545,29 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
         if not self._is_initialized:
             await self.initialize()
 
+        if self._motor_routing:
+            if isinstance(obj, self.model_cls):
+                if not obj.id:
+                    raise DocumentNotFoundError("Document must have an id to be updated")
+                oid = self._to_object_id(obj.id)
+                raw = self._motor_model_dump(obj)
+                await self._motor_collection().replace_one({"_id": oid}, raw)
+                return obj
+            if not hasattr(obj, "id") or not obj.id:
+                raise DocumentNotFoundError("Document must have an id to be updated")
+            oid = self._to_object_id(obj.id)
+            existing = await self._motor_collection().find_one({"_id": oid})
+            if not existing:
+                raise DocumentNotFoundError(f"Object with id {obj.id} not found")
+            doc = self._mongo_doc_to_model(existing)
+            if doc is None:
+                raise DocumentNotFoundError(f"Object with id {obj.id} not found")
+            for key, value in self._motor_patch_fields(obj).items():
+                setattr(doc, key, value)
+            raw = self._motor_model_dump(doc)
+            await self._motor_collection().replace_one({"_id": oid}, raw)
+            return doc
+
         # Check if obj is already a document instance
         if isinstance(obj, self.model_cls):
             # If it's already a document instance, just save it
@@ -471,6 +618,13 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
         if not self._is_initialized:
             await self.initialize()
 
+        if self._motor_routing:
+            oid = self._to_object_id(id)
+            result = await self._motor_collection().delete_one({"_id": oid})
+            if result.deleted_count == 0:
+                raise DocumentNotFoundError(f"Object with id {id} not found")
+            return
+
         doc = await self.model_cls.get(id)
         if doc:
             await doc.delete()
@@ -502,6 +656,10 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
         if not self._is_initialized:
             await self.initialize()
 
+        if self._motor_routing:
+            raw = await self._motor_collection().find({}).to_list(length=None)
+            return [m for m in (self._mongo_doc_to_model(d) for d in raw) if m is not None]
+
         return await self.model_cls.find_all().to_list()
 
     async def find(self, *args, fetch_links: bool = False, **kwargs) -> List[T]:
@@ -541,10 +699,136 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
         # Remove fetch_links from kwargs if present (it's a parameter, not a query field)
         kwargs_without_fetch_links = {k: v for k, v in kwargs.items() if k != "fetch_links"}
 
+        if self._motor_routing:
+            if fetch_links:
+                raise NotImplementedError("fetch_links is not supported for this database routing mode")
+            if len(args) != 1 or not isinstance(args[0], dict) or kwargs_without_fetch_links:
+                raise NotImplementedError(
+                    "Only a single dict filter is supported for find() in this database routing mode"
+                )
+            raw = await self._motor_collection().find(args[0]).to_list(length=None)
+            return [m for m in (self._mongo_doc_to_model(d) for d in raw) if m is not None]
+
         # In Beanie, fetch_links is passed as a parameter to find(), not called as a method
         query = self.model_cls.find(*args, fetch_links=fetch_links, **kwargs_without_fetch_links)
         return await query.to_list()
 
+    async def find_iter(
+        self,
+        query: dict[str, Any] | None = None,
+        *,
+        sort: list[tuple[str, int]] | None = None,
+        batch_size: int | None = None,
+        limit: int | None = None,
+    ) -> AsyncIterator[T]:
+        """
+        Lazily iterate over matching documents using a Mongo cursor.
+
+        This is intended for large scans where materializing the full result set
+        would be wasteful.
+        """
+        if self._models is not None:
+            raise ValueError("Cannot use find_iter() in multi-model mode. Use db.model_name.find_iter() instead.")
+
+        if not self._is_initialized:
+            await self.initialize()
+
+        cursor = self._motor_collection().find(query or {})
+        if sort:
+            cursor = cursor.sort(sort)
+        if batch_size is not None:
+            cursor = cursor.batch_size(batch_size)
+        if limit is not None:
+            cursor = cursor.limit(limit)
+
+        try:
+            async for raw in cursor:
+                model = self._mongo_doc_to_model(raw)
+                if model is not None:
+                    yield model
+        finally:
+            await cursor.close()
+
+    async def find_window(
+        self,
+        query: dict[str, Any] | None = None,
+        *,
+        sort: list[tuple[str, int]] | None = None,
+        limit: int | None = None,
+    ) -> List[T]:
+        """
+        Return a bounded, sorted window of matching documents.
+
+        This is useful for page-oriented callers that need ``limit + 1`` rows to
+        determine whether another page exists without loading the full collection.
+        """
+        if self._models is not None:
+            raise ValueError("Cannot use find_window() in multi-model mode. Use db.model_name.find_window() instead.")
+
+        if not self._is_initialized:
+            await self.initialize()
+
+        cursor = self._motor_collection().find(query or {})
+        if sort:
+            cursor = cursor.sort(sort)
+        raw = await cursor.to_list(length=limit)
+        return [m for m in (self._mongo_doc_to_model(d) for d in raw) if m is not None]
+
+    def find_iter_sync(
+        self,
+        query: dict[str, Any] | None = None,
+        *,
+        sort: list[tuple[str, int]] | None = None,
+        batch_size: int | None = None,
+        limit: int | None = None,
+    ) -> Iterator[T]:
+        """Lazily iterate over matching documents using a synchronous Mongo cursor."""
+        if self._models is not None:
+            raise ValueError(
+                "Cannot use find_iter_sync() in multi-model mode. Use db.model_name.find_iter_sync() instead."
+            )
+
+        try:
+            asyncio.get_running_loop()
+            raise RuntimeError("find_iter_sync() called from async context. Use async for item in find_iter() instead.")
+        except RuntimeError as exc:
+            if "no running event loop" not in str(exc).lower():
+                raise
+
+        if not self._is_initialized:
+            self.initialize_sync()
+
+        cursor = self._sync_collection().find(query or {})
+        if sort:
+            cursor = cursor.sort(sort)
+        if batch_size is not None:
+            cursor = cursor.batch_size(batch_size)
+        if limit is not None:
+            cursor = cursor.limit(limit)
+
+        def _generator() -> Iterator[T]:
+            try:
+                for raw in cursor:
+                    model = self._mongo_doc_to_model(raw)
+                    if model is not None:
+                        yield model
+            finally:
+                cursor.close()
+
+        return _generator()
+
+    async def count_documents(self, query: dict[str, Any] | None = None) -> int:
+        """Return the number of documents matching ``query`` without loading them."""
+        if self._models is not None:
+            raise ValueError(
+                "Cannot use count_documents() in multi-model mode. Use db.model_name.count_documents() instead."
+            )
+
+        if not self._is_initialized:
+            await self.initialize()
+
+        return await self._motor_collection().count_documents(query or {})
+
     async def aggregate(self, pipeline: list) -> List[T]:
         """
         Execute a MongoDB aggregation pipeline.
@@ -568,6 +852,8 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
         # Auto-initialize if needed (backward compatible)
         if not self._is_initialized:
             await self.initialize()
+        if self._motor_routing:
+            return await self._motor_collection().aggregate(pipeline).to_list(None)
         return await self.model_cls.get_motor_collection().aggregate(pipeline).to_list(None)
 
     def get_raw_model(self) -> Type[T]:

{mindtrace_database-0.10.0 → mindtrace_database-0.11.0}/mindtrace/database/backends/unified_odm.py

@@ -532,6 +532,7 @@ class UnifiedMindtraceODM(MindtraceODM):
         self.unified_model_cls = unified_model_cls
         self._unified_models = unified_models
         self._model_odms: Dict[str, "UnifiedMindtraceODM"] = {}
+        self._closed = False
 
         # Support multi-model mode with unified models
         if unified_models is not None:
@@ -634,6 +635,27 @@ class UnifiedMindtraceODM(MindtraceODM):
             return self._model_odms[name]
         raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
 
+    def close(self) -> None:
+        """Release resources held by the underlying backends.
+
+        Closes any child UnifiedMindtraceODM instances (multi-model mode)
+        and each backend that exposes a ``close`` method. Idempotent.
+        """
+        if self._closed:
+            return
+        for odm in self._model_odms.values():
+            try:
+                odm.close()
+            except Exception:
+                pass
+        for backend in (self.mongo_backend, self.redis_backend):
+            if backend is not None and hasattr(backend, "close"):
+                try:
+                    backend.close()
+                except Exception:
+                    pass
+        self._closed = True
+
     def _get_active_backend(self):
         """
         Get the currently active backend based on preference and availability.

{mindtrace_database-0.10.0 → mindtrace_database-0.11.0}/mindtrace_database.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mindtrace-database
-Version: 0.10.0
+Version: 0.11.0
 Summary: Database functionality for Mindtrace
 Author: Mindtrace Team
 License-Expression: Apache-2.0
@@ -11,13 +11,13 @@ Classifier: Programming Language :: Python :: 3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: beanie<2,>=1.29.0
-Requires-Dist: mindtrace-core>=0.10.0
+Requires-Dist: mindtrace-core>=0.11.0
 Requires-Dist: pydantic>=2.11.1
 Requires-Dist: redis>=4.0.0
 Requires-Dist: redis-om>=0.3.5
 Requires-Dist: motor>=3.3.0
 Requires-Dist: pymongo>=4.14.0
-Requires-Dist: mindtrace-registry>=0.10.0
+Requires-Dist: mindtrace-registry>=0.11.0
 Dynamic: license-file
 
 [![PyPI version](https://img.shields.io/pypi/v/mindtrace-database)](https://pypi.org/project/mindtrace-database/)

{mindtrace_database-0.10.0 → mindtrace_database-0.11.0}/mindtrace_database.egg-info/requires.txt

@@ -1,8 +1,8 @@
 beanie<2,>=1.29.0
-mindtrace-core>=0.10.0
+mindtrace-core>=0.11.0
 pydantic>=2.11.1
 redis>=4.0.0
 redis-om>=0.3.5
 motor>=3.3.0
 pymongo>=4.14.0
-mindtrace-registry>=0.10.0
+mindtrace-registry>=0.11.0

{mindtrace_database-0.10.0 → mindtrace_database-0.11.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mindtrace-database"
-version = "0.10.0"
+version = "0.11.0"
 description = "Database functionality for Mindtrace"
 readme = "README.md"
 license = "Apache-2.0"
@@ -13,13 +13,13 @@ classifiers = [
 ]
 dependencies = [
     "beanie>=1.29.0, <2",
-    "mindtrace-core>=0.10.0",
+    "mindtrace-core>=0.11.0",
     "pydantic>=2.11.1",
     "redis>=4.0.0",
     "redis-om>=0.3.5",
     "motor>=3.3.0",
     "pymongo>=4.14.0",
-    "mindtrace-registry>=0.10.0",
+    "mindtrace-registry>=0.11.0",
 ]
 
 [project.urls]