mindtrace-database 0.9.4__tar.gz → 0.10.1__tar.gz

{mindtrace_database-0.9.4 → mindtrace_database-0.10.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mindtrace-database
-Version: 0.9.4
+Version: 0.10.1
 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.9.4
+Requires-Dist: mindtrace-core>=0.10.1
 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.9.4
+Requires-Dist: mindtrace-registry>=0.10.1
 Dynamic: license-file
 
 [![PyPI version](https://img.shields.io/pypi/v/mindtrace-database)](https://pypi.org/project/mindtrace-database/)

{mindtrace_database-0.9.4 → mindtrace_database-0.10.1}/mindtrace/database/__init__.py

@@ -4,7 +4,6 @@ from beanie import Link
 from mindtrace.database.backends.mindtrace_odm import InitMode, MindtraceODM
 from mindtrace.database.backends.mongo_odm import MindtraceDocument, MongoMindtraceODM
 from mindtrace.database.backends.redis_odm import MindtraceRedisDocument, RedisMindtraceODM
-from mindtrace.database.backends.registry_odm import RegistryMindtraceODM
 from mindtrace.database.backends.unified_odm import (
     BackendType,
     UnifiedMindtraceDocument,
@@ -27,3 +26,12 @@ __all__ = [
     "UnifiedMindtraceDocument",
     "UnifiedMindtraceODM",
 ]
+
+
+def __getattr__(name):
+    if name == "RegistryMindtraceODM":
+        from mindtrace.database.backends.registry_odm import RegistryMindtraceODM
+
+        globals()["RegistryMindtraceODM"] = RegistryMindtraceODM
+        return RegistryMindtraceODM
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

{mindtrace_database-0.9.4 → mindtrace_database-0.10.1}/mindtrace/database/backends/mongo_odm.py

@@ -1,7 +1,10 @@
 import asyncio
-from typing import Dict, List, Optional, Type, TypeVar
+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.errors import DuplicateKeyError
@@ -44,6 +47,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):
     """
@@ -129,6 +138,7 @@ class MongoMindtraceODM[T: MindtraceDocument](MindtraceODM):
 
         self.client = AsyncIOMotorClient(db_uri)
         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"] = {}
@@ -192,6 +202,51 @@ 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 _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 +254,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(
@@ -340,6 +400,19 @@ 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)}")
@@ -381,6 +454,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 +506,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 +579,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 +617,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,6 +660,16 @@ 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()
@@ -568,6 +697,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.9.4 → mindtrace_database-0.10.1}/mindtrace/database/backends/redis_odm.py

@@ -777,25 +777,13 @@ class RedisMindtraceODM(MindtraceODM):
 
         # After saving, ensure the index is working - if it has 0 docs, recreate it
         # This handles the case where index was created before documents were inserted
+        # NOTE: Do NOT run Migrator().run() here. It can detect a schema hash mismatch
+        # (e.g. if test cleanup deleted the hash key) and DROP+CREATE the index after
+        # the doc was saved, causing a race where find()/all() returns empty before
+        # re-indexing completes. _do_initialize() (via self.initialize()) and
+        # _ensure_index_has_documents() are sufficient.
         try:
             self._ensure_index_has_documents(self.model_cls)
-            # After ensuring index, try running Migrator to make redis-om aware of it
-            # This ensures redis-om's find() method can use the index
-            try:
-                import os
-
-                original_redis_url = os.environ.get("REDIS_OM_URL", None)
-                if self.redis_url:
-                    os.environ["REDIS_OM_URL"] = self.redis_url
-                try:
-                    Migrator().run()
-                finally:
-                    if original_redis_url:
-                        os.environ["REDIS_OM_URL"] = original_redis_url
-                    elif "REDIS_OM_URL" in os.environ:
-                        del os.environ["REDIS_OM_URL"]
-            except Exception:
-                pass  # Don't fail if Migrator fails
         except Exception:
             pass  # Don't fail insert if index check fails
 

{mindtrace_database-0.9.4 → mindtrace_database-0.10.1}/mindtrace_database.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mindtrace-database
-Version: 0.9.4
+Version: 0.10.1
 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.9.4
+Requires-Dist: mindtrace-core>=0.10.1
 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.9.4
+Requires-Dist: mindtrace-registry>=0.10.1
 Dynamic: license-file
 
 [![PyPI version](https://img.shields.io/pypi/v/mindtrace-database)](https://pypi.org/project/mindtrace-database/)

{mindtrace_database-0.9.4 → mindtrace_database-0.10.1}/mindtrace_database.egg-info/requires.txt

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

{mindtrace_database-0.9.4 → mindtrace_database-0.10.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mindtrace-database"
-version = "0.9.4"
+version = "0.10.1"
 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.9.4",
+    "mindtrace-core>=0.10.1",
     "pydantic>=2.11.1",
     "redis>=4.0.0",
     "redis-om>=0.3.5",
     "motor>=3.3.0",
     "pymongo>=4.14.0",
-    "mindtrace-registry>=0.9.4",
+    "mindtrace-registry>=0.10.1",
 ]
 
 [project.urls]