autonomous-app 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

autonomous/db/context_managers.py

@@ -0,0 +1,313 @@
+import contextlib
+import threading
+from contextlib import contextmanager
+
+from pymongo.read_concern import ReadConcern
+from pymongo.write_concern import WriteConcern
+
+from autonomous.db.base.fields import _no_dereference_for_fields
+from autonomous.db.common import _import_class
+from autonomous.db.connection import DEFAULT_CONNECTION_NAME, get_db
+from autonomous.db.pymongo_support import count_documents
+
+__all__ = (
+    "switch_db",
+    "switch_collection",
+    "no_dereference",
+    "no_sub_classes",
+    "query_counter",
+    "set_write_concern",
+    "set_read_write_concern",
+    "no_dereferencing_active_for_class",
+)
+
+
+class MyThreadLocals(threading.local):
+    def __init__(self):
+        # {DocCls: count} keeping track of classes with an active no_dereference context
+        self.no_dereferencing_class = {}
+
+
+thread_locals = MyThreadLocals()
+
+
+def no_dereferencing_active_for_class(cls):
+    return cls in thread_locals.no_dereferencing_class
+
+
+def _register_no_dereferencing_for_class(cls):
+    thread_locals.no_dereferencing_class.setdefault(cls, 0)
+    thread_locals.no_dereferencing_class[cls] += 1
+
+
+def _unregister_no_dereferencing_for_class(cls):
+    thread_locals.no_dereferencing_class[cls] -= 1
+    if thread_locals.no_dereferencing_class[cls] == 0:
+        thread_locals.no_dereferencing_class.pop(cls)
+
+
+class switch_db:
+    """switch_db alias context manager.
+
+    Example ::
+
+        # Register connections
+        register_connection('default', 'mongoenginetest')
+        register_connection('testdb-1', 'mongoenginetest2')
+
+        class Group(Document):
+            name = StringField()
+
+        Group(name='test').save()  # Saves in the default db
+
+        with switch_db(Group, 'testdb-1') as Group:
+            Group(name='hello testdb!').save()  # Saves in testdb-1
+    """
+
+    def __init__(self, cls, db_alias):
+        """Construct the switch_db context manager
+
+        :param cls: the class to change the registered db
+        :param db_alias: the name of the specific database to use
+        """
+        self.cls = cls
+        self.collection = cls._get_collection()
+        self.db_alias = db_alias
+        self.ori_db_alias = cls._meta.get("db_alias", DEFAULT_CONNECTION_NAME)
+
+    def __enter__(self):
+        """Change the db_alias and clear the cached collection."""
+        self.cls._meta["db_alias"] = self.db_alias
+        self.cls._collection = None
+        return self.cls
+
+    def __exit__(self, t, value, traceback):
+        """Reset the db_alias and collection."""
+        self.cls._meta["db_alias"] = self.ori_db_alias
+        self.cls._collection = self.collection
+
+
+class switch_collection:
+    """switch_collection alias context manager.
+
+    Example ::
+
+        class Group(Document):
+            name = StringField()
+
+        Group(name='test').save()  # Saves in the default db
+
+        with switch_collection(Group, 'group1') as Group:
+            Group(name='hello testdb!').save()  # Saves in group1 collection
+    """
+
+    def __init__(self, cls, collection_name):
+        """Construct the switch_collection context manager.
+
+        :param cls: the class to change the registered db
+        :param collection_name: the name of the collection to use
+        """
+        self.cls = cls
+        self.ori_collection = cls._get_collection()
+        self.ori_get_collection_name = cls._get_collection_name
+        self.collection_name = collection_name
+
+    def __enter__(self):
+        """Change the _get_collection_name and clear the cached collection."""
+
+        @classmethod
+        def _get_collection_name(cls):
+            return self.collection_name
+
+        self.cls._get_collection_name = _get_collection_name
+        self.cls._collection = None
+        return self.cls
+
+    def __exit__(self, t, value, traceback):
+        """Reset the collection."""
+        self.cls._collection = self.ori_collection
+        self.cls._get_collection_name = self.ori_get_collection_name
+
+
+@contextlib.contextmanager
+def no_dereference(cls):
+    """no_dereference context manager.
+
+    Turns off all dereferencing in Documents for the duration of the context
+    manager::
+
+        with no_dereference(Group):
+            Group.objects()
+    """
+    try:
+        cls = cls
+
+        ReferenceField = _import_class("ReferenceField")
+        GenericReferenceField = _import_class("GenericReferenceField")
+        ComplexBaseField = _import_class("ComplexBaseField")
+
+        deref_fields = [
+            field
+            for name, field in cls._fields.items()
+            if isinstance(
+                field, (ReferenceField, GenericReferenceField, ComplexBaseField)
+            )
+        ]
+
+        _register_no_dereferencing_for_class(cls)
+
+        with _no_dereference_for_fields(*deref_fields):
+            yield None
+    finally:
+        _unregister_no_dereferencing_for_class(cls)
+
+
+class no_sub_classes:
+    """no_sub_classes context manager.
+
+    Only returns instances of this class and no sub (inherited) classes::
+
+        with no_sub_classes(Group) as Group:
+            Group.objects.find()
+    """
+
+    def __init__(self, cls):
+        """Construct the no_sub_classes context manager.
+
+        :param cls: the class to turn querying subclasses on
+        """
+        self.cls = cls
+        self.cls_initial_subclasses = None
+
+    def __enter__(self):
+        """Change the objects default and _auto_dereference values."""
+        self.cls_initial_subclasses = self.cls._subclasses
+        self.cls._subclasses = (self.cls._class_name,)
+        return self.cls
+
+    def __exit__(self, t, value, traceback):
+        """Reset the default and _auto_dereference values."""
+        self.cls._subclasses = self.cls_initial_subclasses
+
+
+class query_counter:
+    """Query_counter context manager to get the number of queries.
+    This works by updating the `profiling_level` of the database so that all queries get logged,
+    resetting the db.system.profile collection at the beginning of the context and counting the new entries.
+
+    This was designed for debugging purpose. In fact it is a global counter so queries issued by other threads/processes
+    can interfere with it
+
+    Usage:
+
+    .. code-block:: python
+
+        class User(Document):
+            name = StringField()
+
+        with query_counter() as q:
+            user = User(name='Bob')
+            assert q == 0       # no query fired yet
+            user.save()
+            assert q == 1       # 1 query was fired, an 'insert'
+            user_bis = User.objects().first()
+            assert q == 2       # a 2nd query was fired, a 'find_one'
+
+    Be aware that:
+
+    - Iterating over large amount of documents (>101) makes pymongo issue `getmore` queries to fetch the next batch of documents (https://www.mongodb.com/docs/manual/tutorial/iterate-a-cursor/#cursor-batches)
+    - Some queries are ignored by default by the counter (killcursors, db.system.indexes)
+    """
+
+    def __init__(self, alias=DEFAULT_CONNECTION_NAME):
+        self.db = get_db(alias=alias)
+        self.initial_profiling_level = None
+        self._ctx_query_counter = 0  # number of queries issued by the context
+
+        self._ignored_query = {
+            "ns": {"$ne": "%s.system.indexes" % self.db.name},
+            "op": {"$ne": "killcursors"},  # MONGODB < 3.2
+            "command.killCursors": {"$exists": False},  # MONGODB >= 3.2
+        }
+
+    def _turn_on_profiling(self):
+        profile_update_res = self.db.command({"profile": 0})
+        self.initial_profiling_level = profile_update_res["was"]
+
+        self.db.system.profile.drop()
+        self.db.command({"profile": 2})
+
+    def _resets_profiling(self):
+        self.db.command({"profile": self.initial_profiling_level})
+
+    def __enter__(self):
+        self._turn_on_profiling()
+        return self
+
+    def __exit__(self, t, value, traceback):
+        self._resets_profiling()
+
+    def __eq__(self, value):
+        counter = self._get_count()
+        return value == counter
+
+    def __ne__(self, value):
+        return not self.__eq__(value)
+
+    def __lt__(self, value):
+        return self._get_count() < value
+
+    def __le__(self, value):
+        return self._get_count() <= value
+
+    def __gt__(self, value):
+        return self._get_count() > value
+
+    def __ge__(self, value):
+        return self._get_count() >= value
+
+    def __int__(self):
+        return self._get_count()
+
+    def __repr__(self):
+        """repr query_counter as the number of queries."""
+        return "%s" % self._get_count()
+
+    def _get_count(self):
+        """Get the number of queries by counting the current number of entries in db.system.profile
+        and substracting the queries issued by this context. In fact everytime this is called, 1 query is
+        issued so we need to balance that
+        """
+        count = (
+            count_documents(self.db.system.profile, self._ignored_query)
+            - self._ctx_query_counter
+        )
+        self._ctx_query_counter += (
+            1  # Account for the query we just issued to gather the information
+        )
+        return count
+
+
+@contextmanager
+def set_write_concern(collection, write_concerns):
+    combined_concerns = dict(collection.write_concern.document.items())
+    combined_concerns.update(write_concerns)
+    yield collection.with_options(write_concern=WriteConcern(**combined_concerns))
+
+
+@contextmanager
+def set_read_write_concern(collection, write_concerns, read_concerns):
+    combined_write_concerns = dict(collection.write_concern.document.items())
+
+    if write_concerns is not None:
+        combined_write_concerns.update(write_concerns)
+
+    combined_read_concerns = dict(collection.read_concern.document.items())
+
+    if read_concerns is not None:
+        combined_read_concerns.update(read_concerns)
+
+    yield collection.with_options(
+        write_concern=WriteConcern(**combined_write_concerns),
+        read_concern=ReadConcern(**combined_read_concerns),
+    )

autonomous/db/dereference.py

@@ -0,0 +1,291 @@
+from bson import SON, DBRef
+
+from autonomous import log
+from autonomous.db.base import (
+    BaseDict,
+    BaseList,
+    EmbeddedDocumentList,
+    TopLevelDocumentMetaclass,
+    get_document,
+)
+from autonomous.db.base.datastructures import LazyReference
+from autonomous.db.connection import get_db
+from autonomous.db.document import Document, EmbeddedDocument
+from autonomous.db.fields import (
+    DictField,
+    GenericReferenceField,
+    ListField,
+    MapField,
+    ReferenceField,
+)
+from autonomous.db.queryset import QuerySet
+
+
+class DeReference:
+    def __call__(self, items, max_depth=1, instance=None, name=None):
+        """
+        Cheaply dereferences the items to a set depth.
+        Also handles the conversion of complex data types.
+
+        :param items: The iterable (dict, list, queryset) to be dereferenced.
+        :param max_depth: The maximum depth to recurse to
+        :param instance: The owning instance used for tracking changes by
+            :class:`~autonomous.db.base.ComplexBaseField`
+        :param name: The name of the field, used for tracking changes by
+            :class:`~autonomous.db.base.ComplexBaseField`
+        :param get: A boolean determining if being called by __get__
+        """
+        if items is None or isinstance(items, str):
+            return items
+
+        # cheapest way to convert a queryset to a list
+        # list(queryset) uses a count() query to determine length
+        if isinstance(items, QuerySet):
+            items = [i for i in items]
+
+        self.max_depth = max_depth
+        doc_type = None
+
+        if instance and isinstance(
+            instance, (Document, EmbeddedDocument, TopLevelDocumentMetaclass)
+        ):
+            doc_type = instance._fields.get(name)
+            while hasattr(doc_type, "field"):
+                doc_type = doc_type.field
+            if isinstance(doc_type, (ReferenceField)):
+                field = doc_type
+                doc_type = doc_type.document_type
+                is_list = not hasattr(items, "items")
+
+                if is_list and all(i.__class__ == doc_type for i in items):
+                    return items
+                elif not is_list and all(
+                    i.__class__ == doc_type for i in items.values()
+                ):
+                    return items
+                elif not field.dbref:
+                    # We must turn the ObjectIds into DBRefs
+
+                    # Recursively dig into the sub items of a list/dict
+                    # to turn the ObjectIds into DBRefs
+                    def _get_items_from_list(items):
+                        new_items = []
+                        for v in items:
+                            value = v
+                            if isinstance(v, dict):
+                                value = _get_items_from_dict(v)
+                            elif isinstance(v, list):
+                                value = _get_items_from_list(v)
+                            elif not isinstance(v, (DBRef, Document)):
+                                value = field.to_python(v)
+                            new_items.append(value)
+                        return new_items
+
+                    def _get_items_from_dict(items):
+                        new_items = {}
+                        for k, v in items.items():
+                            value = v
+                            if isinstance(v, list):
+                                value = _get_items_from_list(v)
+                            elif isinstance(v, dict):
+                                value = _get_items_from_dict(v)
+                            elif not isinstance(v, (DBRef, Document)):
+                                value = field.to_python(v)
+                            new_items[k] = value
+                        return new_items
+
+                    if not hasattr(items, "items"):
+                        items = _get_items_from_list(items)
+                    else:
+                        items = _get_items_from_dict(items)
+        self.reference_map = self._find_references(items)
+        self.object_map = self._fetch_objects(doc_type=doc_type)
+        return self._attach_objects(items, 0, instance, name)
+
+    def _find_references(self, items, depth=0):
+        """
+        Recursively finds all db references to be dereferenced
+
+        :param items: The iterable (dict, list, queryset)
+        :param depth: The current depth of recursion
+        """
+        reference_map = {}
+        if not items or depth >= self.max_depth:
+            return reference_map
+
+        # Determine the iterator to use
+        if isinstance(items, dict):
+            iterator = items.values()
+        else:
+            iterator = items
+
+        # Recursively find dbreferences
+        depth += 1
+        for item in iterator:
+            if isinstance(item, (Document, EmbeddedDocument)):
+                for field_name, field in item._fields.items():
+                    v = item._data.get(field_name, None)
+                    if isinstance(v, LazyReference):
+                        # LazyReference inherits DBRef but should not be dereferenced here !
+                        continue
+                    elif isinstance(v, DBRef):
+                        reference_map.setdefault(field.document_type, set()).add(v.id)
+                    elif isinstance(v, (dict, SON)) and "_ref" in v:
+                        reference_map.setdefault(get_document(v["_cls"]), set()).add(
+                            v["_ref"].id
+                        )
+                    elif isinstance(v, (dict, list, tuple)) and depth <= self.max_depth:
+                        field_cls = getattr(
+                            getattr(field, "field", None), "document_type", None
+                        )
+                        references = self._find_references(v, depth)
+                        for key, refs in references.items():
+                            if isinstance(
+                                field_cls, (Document, TopLevelDocumentMetaclass)
+                            ):
+                                key = field_cls
+                            reference_map.setdefault(key, set()).update(refs)
+            elif isinstance(item, LazyReference):
+                # LazyReference inherits DBRef but should not be dereferenced here !
+                continue
+            elif isinstance(item, DBRef):
+                reference_map.setdefault(item.collection, set()).add(item.id)
+            elif isinstance(item, (dict, SON)) and "_ref" in item:
+                reference_map.setdefault(get_document(item["_cls"]), set()).add(
+                    item["_ref"].id
+                )
+            elif isinstance(item, (dict, list, tuple)) and depth - 1 <= self.max_depth:
+                references = self._find_references(item, depth - 1)
+                for key, refs in references.items():
+                    reference_map.setdefault(key, set()).update(refs)
+
+        return reference_map
+
+    def _fetch_objects(self, doc_type=None):
+        """Fetch all references and convert to their document objects"""
+        object_map = {}
+        for collection, dbrefs in self.reference_map.items():
+            if hasattr(collection, "objects"):
+                col_name = collection._get_collection_name()
+                refs = [
+                    dbref for dbref in dbrefs if (col_name, dbref) not in object_map
+                ]
+                references = collection.objects.in_bulk(refs)
+                for key, doc in references.items():
+                    object_map[(col_name, key)] = doc
+            else:  # Generic reference: use the refs data to convert to document
+                if isinstance(doc_type, (ListField, DictField, MapField)):
+                    continue
+
+                refs = [
+                    dbref for dbref in dbrefs if (collection, dbref) not in object_map
+                ]
+                # log(doc_type)
+                if isinstance(doc_type, Document):
+                    references = doc_type._get_db()[collection].find(
+                        {"_id": {"$in": refs}}
+                    )
+                    for ref in references:
+                        doc = doc_type._from_son(ref)
+                        object_map[(collection, doc.id)] = doc
+                else:
+                    references = get_db()[collection].find({"_id": {"$in": refs}})
+                    for ref in references:
+                        if "_cls" in ref:
+                            doc = get_document(ref["_cls"])._from_son(ref)
+                        elif doc_type is None:
+                            doc = get_document(
+                                "".join(x.capitalize() for x in collection.split("_"))
+                            )._from_son(ref)
+                        else:
+                            doc = doc_type._from_son(ref)
+                        object_map[(collection, doc.id)] = doc
+        return object_map
+
+    def _attach_objects(self, items, depth=0, instance=None, name=None):
+        """
+        Recursively finds all db references to be dereferenced
+
+        :param items: The iterable (dict, list, queryset)
+        :param depth: The current depth of recursion
+        :param instance: The owning instance used for tracking changes by
+            :class:`~autonomous.db.base.ComplexBaseField`
+        :param name: The name of the field, used for tracking changes by
+            :class:`~autonomous.db.base.ComplexBaseField`
+        """
+        if not items:
+            if isinstance(items, (BaseDict, BaseList)):
+                return items
+
+            if instance:
+                if isinstance(items, dict):
+                    return BaseDict(items, instance, name)
+                else:
+                    return BaseList(items, instance, name)
+
+        if isinstance(items, (dict, SON)):
+            if "_ref" in items:
+                return self.object_map.get(
+                    (items["_ref"].collection, items["_ref"].id), items
+                )
+            elif "_cls" in items:
+                doc = get_document(items["_cls"])._from_son(items)
+                _cls = doc._data.pop("_cls", None)
+                del items["_cls"]
+                doc._data = self._attach_objects(doc._data, depth, doc, None)
+                if _cls is not None:
+                    doc._data["_cls"] = _cls
+                return doc
+
+        if not hasattr(items, "items"):
+            is_list = True
+            list_type = BaseList
+            if isinstance(items, EmbeddedDocumentList):
+                list_type = EmbeddedDocumentList
+            as_tuple = isinstance(items, tuple)
+            iterator = enumerate(items)
+            data = []
+        else:
+            is_list = False
+            iterator = items.items()
+            data = {}
+
+        depth += 1
+        for k, v in iterator:
+            if is_list:
+                data.append(v)
+            else:
+                data[k] = v
+
+            if k in self.object_map and not is_list:
+                data[k] = self.object_map[k]
+            elif isinstance(v, (Document, EmbeddedDocument)):
+                for field_name in v._fields:
+                    v = data[k]._data.get(field_name, None)
+                    if isinstance(v, DBRef):
+                        data[k]._data[field_name] = self.object_map.get(
+                            (v.collection, v.id), v
+                        )
+                    elif isinstance(v, (dict, SON)) and "_ref" in v:
+                        data[k]._data[field_name] = self.object_map.get(
+                            (v["_ref"].collection, v["_ref"].id), v
+                        )
+                    elif isinstance(v, (dict, list, tuple)) and depth <= self.max_depth:
+                        item_name = f"{name}.{k}.{field_name}"
+                        data[k]._data[field_name] = self._attach_objects(
+                            v, depth, instance=instance, name=item_name
+                        )
+            elif isinstance(v, (dict, list, tuple)) and depth <= self.max_depth:
+                item_name = f"{name}.{k}" if name else name
+                data[k] = self._attach_objects(
+                    v, depth - 1, instance=instance, name=item_name
+                )
+            elif isinstance(v, DBRef) and hasattr(v, "id"):
+                data[k] = self.object_map.get((v.collection, v.id), v)
+
+        if instance and name:
+            if is_list:
+                return tuple(data) if as_tuple else list_type(data, instance, name)
+            return BaseDict(data, instance, name)
+        depth += 1
+        return data