b47669403b
This commit: 1. Formats all of our existing code using `black`. 2. Adds a note about using `black` to `CONTRIBUTING.rst`. 3. Runs `black --check` as part of CI (failing builds that aren't properly formatted).
1079 lines
36 KiB
Python
1079 lines
36 KiB
Python
# -*- coding: utf-8 -*-
|
|
import unittest
|
|
from datetime import datetime
|
|
|
|
from nose.plugins.skip import SkipTest
|
|
from pymongo.errors import OperationFailure
|
|
import pymongo
|
|
from six import iteritems
|
|
|
|
from mongoengine import *
|
|
from mongoengine.connection import get_db
|
|
|
|
__all__ = ("IndexesTest",)
|
|
|
|
|
|
class IndexesTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.connection = connect(db="mongoenginetest")
|
|
self.db = get_db()
|
|
|
|
class Person(Document):
|
|
name = StringField()
|
|
age = IntField()
|
|
|
|
non_field = True
|
|
|
|
meta = {"allow_inheritance": True}
|
|
|
|
self.Person = Person
|
|
|
|
def tearDown(self):
|
|
self.connection.drop_database(self.db)
|
|
|
|
def test_indexes_document(self):
|
|
"""Ensure that indexes are used when meta[indexes] is specified for
|
|
Documents
|
|
"""
|
|
self._index_test(Document)
|
|
|
|
def test_indexes_dynamic_document(self):
|
|
"""Ensure that indexes are used when meta[indexes] is specified for
|
|
Dynamic Documents
|
|
"""
|
|
self._index_test(DynamicDocument)
|
|
|
|
def _index_test(self, InheritFrom):
|
|
class BlogPost(InheritFrom):
|
|
date = DateTimeField(db_field="addDate", default=datetime.now)
|
|
category = StringField()
|
|
tags = ListField(StringField())
|
|
meta = {"indexes": ["-date", "tags", ("category", "-date")]}
|
|
|
|
expected_specs = [
|
|
{"fields": [("addDate", -1)]},
|
|
{"fields": [("tags", 1)]},
|
|
{"fields": [("category", 1), ("addDate", -1)]},
|
|
]
|
|
self.assertEqual(expected_specs, BlogPost._meta["index_specs"])
|
|
|
|
BlogPost.ensure_indexes()
|
|
info = BlogPost.objects._collection.index_information()
|
|
# _id, '-date', 'tags', ('cat', 'date')
|
|
self.assertEqual(len(info), 4)
|
|
info = [value["key"] for key, value in iteritems(info)]
|
|
for expected in expected_specs:
|
|
self.assertIn(expected["fields"], info)
|
|
|
|
def _index_test_inheritance(self, InheritFrom):
|
|
class BlogPost(InheritFrom):
|
|
date = DateTimeField(db_field="addDate", default=datetime.now)
|
|
category = StringField()
|
|
tags = ListField(StringField())
|
|
meta = {
|
|
"indexes": ["-date", "tags", ("category", "-date")],
|
|
"allow_inheritance": True,
|
|
}
|
|
|
|
expected_specs = [
|
|
{"fields": [("_cls", 1), ("addDate", -1)]},
|
|
{"fields": [("_cls", 1), ("tags", 1)]},
|
|
{"fields": [("_cls", 1), ("category", 1), ("addDate", -1)]},
|
|
]
|
|
self.assertEqual(expected_specs, BlogPost._meta["index_specs"])
|
|
|
|
BlogPost.ensure_indexes()
|
|
info = BlogPost.objects._collection.index_information()
|
|
# _id, '-date', 'tags', ('cat', 'date')
|
|
# NB: there is no index on _cls by itself, since
|
|
# the indices on -date and tags will both contain
|
|
# _cls as first element in the key
|
|
self.assertEqual(len(info), 4)
|
|
info = [value["key"] for key, value in iteritems(info)]
|
|
for expected in expected_specs:
|
|
self.assertIn(expected["fields"], info)
|
|
|
|
class ExtendedBlogPost(BlogPost):
|
|
title = StringField()
|
|
meta = {"indexes": ["title"]}
|
|
|
|
expected_specs.append({"fields": [("_cls", 1), ("title", 1)]})
|
|
self.assertEqual(expected_specs, ExtendedBlogPost._meta["index_specs"])
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
ExtendedBlogPost.ensure_indexes()
|
|
info = ExtendedBlogPost.objects._collection.index_information()
|
|
info = [value["key"] for key, value in iteritems(info)]
|
|
for expected in expected_specs:
|
|
self.assertIn(expected["fields"], info)
|
|
|
|
def test_indexes_document_inheritance(self):
|
|
"""Ensure that indexes are used when meta[indexes] is specified for
|
|
Documents
|
|
"""
|
|
self._index_test_inheritance(Document)
|
|
|
|
def test_indexes_dynamic_document_inheritance(self):
|
|
"""Ensure that indexes are used when meta[indexes] is specified for
|
|
Dynamic Documents
|
|
"""
|
|
self._index_test_inheritance(DynamicDocument)
|
|
|
|
def test_inherited_index(self):
|
|
"""Ensure index specs are inhertited correctly"""
|
|
|
|
class A(Document):
|
|
title = StringField()
|
|
meta = {"indexes": [{"fields": ("title",)}], "allow_inheritance": True}
|
|
|
|
class B(A):
|
|
description = StringField()
|
|
|
|
self.assertEqual(A._meta["index_specs"], B._meta["index_specs"])
|
|
self.assertEqual(
|
|
[{"fields": [("_cls", 1), ("title", 1)]}], A._meta["index_specs"]
|
|
)
|
|
|
|
def test_index_no_cls(self):
|
|
"""Ensure index specs are inhertited correctly"""
|
|
|
|
class A(Document):
|
|
title = StringField()
|
|
meta = {
|
|
"indexes": [{"fields": ("title",), "cls": False}],
|
|
"allow_inheritance": True,
|
|
"index_cls": False,
|
|
}
|
|
|
|
self.assertEqual([("title", 1)], A._meta["index_specs"][0]["fields"])
|
|
A._get_collection().drop_indexes()
|
|
A.ensure_indexes()
|
|
info = A._get_collection().index_information()
|
|
self.assertEqual(len(info.keys()), 2)
|
|
|
|
class B(A):
|
|
c = StringField()
|
|
d = StringField()
|
|
meta = {
|
|
"indexes": [{"fields": ["c"]}, {"fields": ["d"], "cls": True}],
|
|
"allow_inheritance": True,
|
|
}
|
|
|
|
self.assertEqual([("c", 1)], B._meta["index_specs"][1]["fields"])
|
|
self.assertEqual([("_cls", 1), ("d", 1)], B._meta["index_specs"][2]["fields"])
|
|
|
|
def test_build_index_spec_is_not_destructive(self):
|
|
class MyDoc(Document):
|
|
keywords = StringField()
|
|
|
|
meta = {"indexes": ["keywords"], "allow_inheritance": False}
|
|
|
|
self.assertEqual(MyDoc._meta["index_specs"], [{"fields": [("keywords", 1)]}])
|
|
|
|
# Force index creation
|
|
MyDoc.ensure_indexes()
|
|
|
|
self.assertEqual(MyDoc._meta["index_specs"], [{"fields": [("keywords", 1)]}])
|
|
|
|
def test_embedded_document_index_meta(self):
|
|
"""Ensure that embedded document indexes are created explicitly
|
|
"""
|
|
|
|
class Rank(EmbeddedDocument):
|
|
title = StringField(required=True)
|
|
|
|
class Person(Document):
|
|
name = StringField(required=True)
|
|
rank = EmbeddedDocumentField(Rank, required=False)
|
|
|
|
meta = {"indexes": ["rank.title"], "allow_inheritance": False}
|
|
|
|
self.assertEqual([{"fields": [("rank.title", 1)]}], Person._meta["index_specs"])
|
|
|
|
Person.drop_collection()
|
|
|
|
# Indexes are lazy so use list() to perform query
|
|
list(Person.objects)
|
|
info = Person.objects._collection.index_information()
|
|
info = [value["key"] for key, value in iteritems(info)]
|
|
self.assertIn([("rank.title", 1)], info)
|
|
|
|
def test_explicit_geo2d_index(self):
|
|
"""Ensure that geo2d indexes work when created via meta[indexes]
|
|
"""
|
|
|
|
class Place(Document):
|
|
location = DictField()
|
|
meta = {"allow_inheritance": True, "indexes": ["*location.point"]}
|
|
|
|
self.assertEqual(
|
|
[{"fields": [("location.point", "2d")]}], Place._meta["index_specs"]
|
|
)
|
|
|
|
Place.ensure_indexes()
|
|
info = Place._get_collection().index_information()
|
|
info = [value["key"] for key, value in iteritems(info)]
|
|
self.assertIn([("location.point", "2d")], info)
|
|
|
|
def test_explicit_geo2d_index_embedded(self):
|
|
"""Ensure that geo2d indexes work when created via meta[indexes]
|
|
"""
|
|
|
|
class EmbeddedLocation(EmbeddedDocument):
|
|
location = DictField()
|
|
|
|
class Place(Document):
|
|
current = DictField(field=EmbeddedDocumentField("EmbeddedLocation"))
|
|
meta = {"allow_inheritance": True, "indexes": ["*current.location.point"]}
|
|
|
|
self.assertEqual(
|
|
[{"fields": [("current.location.point", "2d")]}], Place._meta["index_specs"]
|
|
)
|
|
|
|
Place.ensure_indexes()
|
|
info = Place._get_collection().index_information()
|
|
info = [value["key"] for key, value in iteritems(info)]
|
|
self.assertIn([("current.location.point", "2d")], info)
|
|
|
|
def test_explicit_geosphere_index(self):
|
|
"""Ensure that geosphere indexes work when created via meta[indexes]
|
|
"""
|
|
|
|
class Place(Document):
|
|
location = DictField()
|
|
meta = {"allow_inheritance": True, "indexes": ["(location.point"]}
|
|
|
|
self.assertEqual(
|
|
[{"fields": [("location.point", "2dsphere")]}], Place._meta["index_specs"]
|
|
)
|
|
|
|
Place.ensure_indexes()
|
|
info = Place._get_collection().index_information()
|
|
info = [value["key"] for key, value in iteritems(info)]
|
|
self.assertIn([("location.point", "2dsphere")], info)
|
|
|
|
def test_explicit_geohaystack_index(self):
|
|
"""Ensure that geohaystack indexes work when created via meta[indexes]
|
|
"""
|
|
raise SkipTest(
|
|
"GeoHaystack index creation is not supported for now"
|
|
"from meta, as it requires a bucketSize parameter."
|
|
)
|
|
|
|
class Place(Document):
|
|
location = DictField()
|
|
name = StringField()
|
|
meta = {"indexes": [(")location.point", "name")]}
|
|
|
|
self.assertEqual(
|
|
[{"fields": [("location.point", "geoHaystack"), ("name", 1)]}],
|
|
Place._meta["index_specs"],
|
|
)
|
|
|
|
Place.ensure_indexes()
|
|
info = Place._get_collection().index_information()
|
|
info = [value["key"] for key, value in iteritems(info)]
|
|
self.assertIn([("location.point", "geoHaystack")], info)
|
|
|
|
def test_create_geohaystack_index(self):
|
|
"""Ensure that geohaystack indexes can be created
|
|
"""
|
|
|
|
class Place(Document):
|
|
location = DictField()
|
|
name = StringField()
|
|
|
|
Place.create_index({"fields": (")location.point", "name")}, bucketSize=10)
|
|
info = Place._get_collection().index_information()
|
|
info = [value["key"] for key, value in iteritems(info)]
|
|
self.assertIn([("location.point", "geoHaystack"), ("name", 1)], info)
|
|
|
|
def test_dictionary_indexes(self):
|
|
"""Ensure that indexes are used when meta[indexes] contains
|
|
dictionaries instead of lists.
|
|
"""
|
|
|
|
class BlogPost(Document):
|
|
date = DateTimeField(db_field="addDate", default=datetime.now)
|
|
category = StringField()
|
|
tags = ListField(StringField())
|
|
meta = {"indexes": [{"fields": ["-date"], "unique": True, "sparse": True}]}
|
|
|
|
self.assertEqual(
|
|
[{"fields": [("addDate", -1)], "unique": True, "sparse": True}],
|
|
BlogPost._meta["index_specs"],
|
|
)
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
info = BlogPost.objects._collection.index_information()
|
|
# _id, '-date'
|
|
self.assertEqual(len(info), 2)
|
|
|
|
# Indexes are lazy so use list() to perform query
|
|
list(BlogPost.objects)
|
|
info = BlogPost.objects._collection.index_information()
|
|
info = [
|
|
(value["key"], value.get("unique", False), value.get("sparse", False))
|
|
for key, value in iteritems(info)
|
|
]
|
|
self.assertIn(([("addDate", -1)], True, True), info)
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
def test_abstract_index_inheritance(self):
|
|
class UserBase(Document):
|
|
user_guid = StringField(required=True)
|
|
meta = {
|
|
"abstract": True,
|
|
"indexes": ["user_guid"],
|
|
"allow_inheritance": True,
|
|
}
|
|
|
|
class Person(UserBase):
|
|
name = StringField()
|
|
|
|
meta = {"indexes": ["name"]}
|
|
|
|
Person.drop_collection()
|
|
|
|
Person(name="test", user_guid="123").save()
|
|
|
|
self.assertEqual(1, Person.objects.count())
|
|
info = Person.objects._collection.index_information()
|
|
self.assertEqual(
|
|
sorted(info.keys()), ["_cls_1_name_1", "_cls_1_user_guid_1", "_id_"]
|
|
)
|
|
|
|
def test_disable_index_creation(self):
|
|
"""Tests setting auto_create_index to False on the connection will
|
|
disable any index generation.
|
|
"""
|
|
|
|
class User(Document):
|
|
meta = {
|
|
"allow_inheritance": True,
|
|
"indexes": ["user_guid"],
|
|
"auto_create_index": False,
|
|
}
|
|
user_guid = StringField(required=True)
|
|
|
|
class MongoUser(User):
|
|
pass
|
|
|
|
User.drop_collection()
|
|
|
|
User(user_guid="123").save()
|
|
MongoUser(user_guid="123").save()
|
|
|
|
self.assertEqual(2, User.objects.count())
|
|
info = User.objects._collection.index_information()
|
|
self.assertEqual(list(info.keys()), ["_id_"])
|
|
|
|
User.ensure_indexes()
|
|
info = User.objects._collection.index_information()
|
|
self.assertEqual(sorted(info.keys()), ["_cls_1_user_guid_1", "_id_"])
|
|
|
|
def test_embedded_document_index(self):
|
|
"""Tests settings an index on an embedded document
|
|
"""
|
|
|
|
class Date(EmbeddedDocument):
|
|
year = IntField(db_field="yr")
|
|
|
|
class BlogPost(Document):
|
|
title = StringField()
|
|
date = EmbeddedDocumentField(Date)
|
|
|
|
meta = {"indexes": ["-date.year"]}
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
info = BlogPost.objects._collection.index_information()
|
|
self.assertEqual(sorted(info.keys()), ["_id_", "date.yr_-1"])
|
|
|
|
def test_list_embedded_document_index(self):
|
|
"""Ensure list embedded documents can be indexed
|
|
"""
|
|
|
|
class Tag(EmbeddedDocument):
|
|
name = StringField(db_field="tag")
|
|
|
|
class BlogPost(Document):
|
|
title = StringField()
|
|
tags = ListField(EmbeddedDocumentField(Tag))
|
|
|
|
meta = {"indexes": ["tags.name"]}
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
info = BlogPost.objects._collection.index_information()
|
|
# we don't use _cls in with list fields by default
|
|
self.assertEqual(sorted(info.keys()), ["_id_", "tags.tag_1"])
|
|
|
|
post1 = BlogPost(
|
|
title="Embedded Indexes tests in place",
|
|
tags=[Tag(name="about"), Tag(name="time")],
|
|
)
|
|
post1.save()
|
|
|
|
def test_recursive_embedded_objects_dont_break_indexes(self):
|
|
class RecursiveObject(EmbeddedDocument):
|
|
obj = EmbeddedDocumentField("self")
|
|
|
|
class RecursiveDocument(Document):
|
|
recursive_obj = EmbeddedDocumentField(RecursiveObject)
|
|
meta = {"allow_inheritance": True}
|
|
|
|
RecursiveDocument.ensure_indexes()
|
|
info = RecursiveDocument._get_collection().index_information()
|
|
self.assertEqual(sorted(info.keys()), ["_cls_1", "_id_"])
|
|
|
|
def test_covered_index(self):
|
|
"""Ensure that covered indexes can be used
|
|
"""
|
|
|
|
class Test(Document):
|
|
a = IntField()
|
|
b = IntField()
|
|
|
|
meta = {"indexes": ["a"], "allow_inheritance": False}
|
|
|
|
Test.drop_collection()
|
|
|
|
obj = Test(a=1)
|
|
obj.save()
|
|
|
|
# Need to be explicit about covered indexes as mongoDB doesn't know if
|
|
# the documents returned might have more keys in that here.
|
|
query_plan = Test.objects(id=obj.id).exclude("a").explain()
|
|
self.assertEqual(
|
|
query_plan.get("queryPlanner")
|
|
.get("winningPlan")
|
|
.get("inputStage")
|
|
.get("stage"),
|
|
"IDHACK",
|
|
)
|
|
|
|
query_plan = Test.objects(id=obj.id).only("id").explain()
|
|
self.assertEqual(
|
|
query_plan.get("queryPlanner")
|
|
.get("winningPlan")
|
|
.get("inputStage")
|
|
.get("stage"),
|
|
"IDHACK",
|
|
)
|
|
|
|
query_plan = Test.objects(a=1).only("a").exclude("id").explain()
|
|
self.assertEqual(
|
|
query_plan.get("queryPlanner")
|
|
.get("winningPlan")
|
|
.get("inputStage")
|
|
.get("stage"),
|
|
"IXSCAN",
|
|
)
|
|
self.assertEqual(
|
|
query_plan.get("queryPlanner").get("winningPlan").get("stage"), "PROJECTION"
|
|
)
|
|
|
|
query_plan = Test.objects(a=1).explain()
|
|
self.assertEqual(
|
|
query_plan.get("queryPlanner")
|
|
.get("winningPlan")
|
|
.get("inputStage")
|
|
.get("stage"),
|
|
"IXSCAN",
|
|
)
|
|
self.assertEqual(
|
|
query_plan.get("queryPlanner").get("winningPlan").get("stage"), "FETCH"
|
|
)
|
|
|
|
def test_index_on_id(self):
|
|
class BlogPost(Document):
|
|
meta = {"indexes": [["categories", "id"]]}
|
|
|
|
title = StringField(required=True)
|
|
description = StringField(required=True)
|
|
categories = ListField()
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
indexes = BlogPost.objects._collection.index_information()
|
|
self.assertEqual(
|
|
indexes["categories_1__id_1"]["key"], [("categories", 1), ("_id", 1)]
|
|
)
|
|
|
|
def test_hint(self):
|
|
TAGS_INDEX_NAME = "tags_1"
|
|
|
|
class BlogPost(Document):
|
|
tags = ListField(StringField())
|
|
meta = {"indexes": [{"fields": ["tags"], "name": TAGS_INDEX_NAME}]}
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
for i in range(10):
|
|
tags = [("tag %i" % n) for n in range(i % 2)]
|
|
BlogPost(tags=tags).save()
|
|
|
|
# Hinting by shape should work.
|
|
self.assertEqual(BlogPost.objects.hint([("tags", 1)]).count(), 10)
|
|
|
|
# Hinting by index name should work.
|
|
self.assertEqual(BlogPost.objects.hint(TAGS_INDEX_NAME).count(), 10)
|
|
|
|
# Clearing the hint should work fine.
|
|
self.assertEqual(BlogPost.objects.hint().count(), 10)
|
|
self.assertEqual(BlogPost.objects.hint([("ZZ", 1)]).hint().count(), 10)
|
|
|
|
# Hinting on a non-existent index shape should fail.
|
|
with self.assertRaises(OperationFailure):
|
|
BlogPost.objects.hint([("ZZ", 1)]).count()
|
|
|
|
# Hinting on a non-existent index name should fail.
|
|
with self.assertRaises(OperationFailure):
|
|
BlogPost.objects.hint("Bad Name").count()
|
|
|
|
# Invalid shape argument (missing list brackets) should fail.
|
|
with self.assertRaises(ValueError):
|
|
BlogPost.objects.hint(("tags", 1)).count()
|
|
|
|
def test_unique(self):
|
|
"""Ensure that uniqueness constraints are applied to fields.
|
|
"""
|
|
|
|
class BlogPost(Document):
|
|
title = StringField()
|
|
slug = StringField(unique=True)
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
post1 = BlogPost(title="test1", slug="test")
|
|
post1.save()
|
|
|
|
# Two posts with the same slug is not allowed
|
|
post2 = BlogPost(title="test2", slug="test")
|
|
self.assertRaises(NotUniqueError, post2.save)
|
|
self.assertRaises(NotUniqueError, BlogPost.objects.insert, post2)
|
|
|
|
# Ensure backwards compatibility for errors
|
|
self.assertRaises(OperationError, post2.save)
|
|
|
|
def test_primary_key_unique_not_working(self):
|
|
"""Relates to #1445"""
|
|
|
|
class Blog(Document):
|
|
id = StringField(primary_key=True, unique=True)
|
|
|
|
Blog.drop_collection()
|
|
|
|
with self.assertRaises(OperationFailure) as ctx_err:
|
|
Blog(id="garbage").save()
|
|
|
|
# One of the errors below should happen. Which one depends on the
|
|
# PyMongo version and dict order.
|
|
err_msg = str(ctx_err.exception)
|
|
self.assertTrue(
|
|
any(
|
|
[
|
|
"The field 'unique' is not valid for an _id index specification"
|
|
in err_msg,
|
|
"The field 'background' is not valid for an _id index specification"
|
|
in err_msg,
|
|
"The field 'sparse' is not valid for an _id index specification"
|
|
in err_msg,
|
|
]
|
|
)
|
|
)
|
|
|
|
def test_unique_with(self):
|
|
"""Ensure that unique_with constraints are applied to fields.
|
|
"""
|
|
|
|
class Date(EmbeddedDocument):
|
|
year = IntField(db_field="yr")
|
|
|
|
class BlogPost(Document):
|
|
title = StringField()
|
|
date = EmbeddedDocumentField(Date)
|
|
slug = StringField(unique_with="date.year")
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
post1 = BlogPost(title="test1", date=Date(year=2009), slug="test")
|
|
post1.save()
|
|
|
|
# day is different so won't raise exception
|
|
post2 = BlogPost(title="test2", date=Date(year=2010), slug="test")
|
|
post2.save()
|
|
|
|
# Now there will be two docs with the same slug and the same day: fail
|
|
post3 = BlogPost(title="test3", date=Date(year=2010), slug="test")
|
|
self.assertRaises(OperationError, post3.save)
|
|
|
|
def test_unique_embedded_document(self):
|
|
"""Ensure that uniqueness constraints are applied to fields on embedded documents.
|
|
"""
|
|
|
|
class SubDocument(EmbeddedDocument):
|
|
year = IntField(db_field="yr")
|
|
slug = StringField(unique=True)
|
|
|
|
class BlogPost(Document):
|
|
title = StringField()
|
|
sub = EmbeddedDocumentField(SubDocument)
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
post1 = BlogPost(title="test1", sub=SubDocument(year=2009, slug="test"))
|
|
post1.save()
|
|
|
|
# sub.slug is different so won't raise exception
|
|
post2 = BlogPost(title="test2", sub=SubDocument(year=2010, slug="another-slug"))
|
|
post2.save()
|
|
|
|
# Now there will be two docs with the same sub.slug
|
|
post3 = BlogPost(title="test3", sub=SubDocument(year=2010, slug="test"))
|
|
self.assertRaises(NotUniqueError, post3.save)
|
|
|
|
def test_unique_embedded_document_in_list(self):
|
|
"""
|
|
Ensure that the uniqueness constraints are applied to fields in
|
|
embedded documents, even when the embedded documents in in a
|
|
list field.
|
|
"""
|
|
|
|
class SubDocument(EmbeddedDocument):
|
|
year = IntField(db_field="yr")
|
|
slug = StringField(unique=True)
|
|
|
|
class BlogPost(Document):
|
|
title = StringField()
|
|
subs = ListField(EmbeddedDocumentField(SubDocument))
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
post1 = BlogPost(
|
|
title="test1",
|
|
subs=[
|
|
SubDocument(year=2009, slug="conflict"),
|
|
SubDocument(year=2009, slug="conflict"),
|
|
],
|
|
)
|
|
post1.save()
|
|
|
|
post2 = BlogPost(title="test2", subs=[SubDocument(year=2014, slug="conflict")])
|
|
|
|
self.assertRaises(NotUniqueError, post2.save)
|
|
|
|
def test_unique_embedded_document_in_sorted_list(self):
|
|
"""
|
|
Ensure that the uniqueness constraints are applied to fields in
|
|
embedded documents, even when the embedded documents in a sorted list
|
|
field.
|
|
"""
|
|
|
|
class SubDocument(EmbeddedDocument):
|
|
year = IntField()
|
|
slug = StringField(unique=True)
|
|
|
|
class BlogPost(Document):
|
|
title = StringField()
|
|
subs = SortedListField(EmbeddedDocumentField(SubDocument), ordering="year")
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
post1 = BlogPost(
|
|
title="test1",
|
|
subs=[
|
|
SubDocument(year=2009, slug="conflict"),
|
|
SubDocument(year=2009, slug="conflict"),
|
|
],
|
|
)
|
|
post1.save()
|
|
|
|
# confirm that the unique index is created
|
|
indexes = BlogPost._get_collection().index_information()
|
|
self.assertIn("subs.slug_1", indexes)
|
|
self.assertTrue(indexes["subs.slug_1"]["unique"])
|
|
|
|
post2 = BlogPost(title="test2", subs=[SubDocument(year=2014, slug="conflict")])
|
|
|
|
self.assertRaises(NotUniqueError, post2.save)
|
|
|
|
def test_unique_embedded_document_in_embedded_document_list(self):
|
|
"""
|
|
Ensure that the uniqueness constraints are applied to fields in
|
|
embedded documents, even when the embedded documents in an embedded
|
|
list field.
|
|
"""
|
|
|
|
class SubDocument(EmbeddedDocument):
|
|
year = IntField()
|
|
slug = StringField(unique=True)
|
|
|
|
class BlogPost(Document):
|
|
title = StringField()
|
|
subs = EmbeddedDocumentListField(SubDocument)
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
post1 = BlogPost(
|
|
title="test1",
|
|
subs=[
|
|
SubDocument(year=2009, slug="conflict"),
|
|
SubDocument(year=2009, slug="conflict"),
|
|
],
|
|
)
|
|
post1.save()
|
|
|
|
# confirm that the unique index is created
|
|
indexes = BlogPost._get_collection().index_information()
|
|
self.assertIn("subs.slug_1", indexes)
|
|
self.assertTrue(indexes["subs.slug_1"]["unique"])
|
|
|
|
post2 = BlogPost(title="test2", subs=[SubDocument(year=2014, slug="conflict")])
|
|
|
|
self.assertRaises(NotUniqueError, post2.save)
|
|
|
|
def test_unique_with_embedded_document_and_embedded_unique(self):
|
|
"""Ensure that uniqueness constraints are applied to fields on
|
|
embedded documents. And work with unique_with as well.
|
|
"""
|
|
|
|
class SubDocument(EmbeddedDocument):
|
|
year = IntField(db_field="yr")
|
|
slug = StringField(unique=True)
|
|
|
|
class BlogPost(Document):
|
|
title = StringField(unique_with="sub.year")
|
|
sub = EmbeddedDocumentField(SubDocument)
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
post1 = BlogPost(title="test1", sub=SubDocument(year=2009, slug="test"))
|
|
post1.save()
|
|
|
|
# sub.slug is different so won't raise exception
|
|
post2 = BlogPost(title="test2", sub=SubDocument(year=2010, slug="another-slug"))
|
|
post2.save()
|
|
|
|
# Now there will be two docs with the same sub.slug
|
|
post3 = BlogPost(title="test3", sub=SubDocument(year=2010, slug="test"))
|
|
self.assertRaises(NotUniqueError, post3.save)
|
|
|
|
# Now there will be two docs with the same title and year
|
|
post3 = BlogPost(title="test1", sub=SubDocument(year=2009, slug="test-1"))
|
|
self.assertRaises(NotUniqueError, post3.save)
|
|
|
|
def test_ttl_indexes(self):
|
|
class Log(Document):
|
|
created = DateTimeField(default=datetime.now)
|
|
meta = {"indexes": [{"fields": ["created"], "expireAfterSeconds": 3600}]}
|
|
|
|
Log.drop_collection()
|
|
|
|
# Indexes are lazy so use list() to perform query
|
|
list(Log.objects)
|
|
info = Log.objects._collection.index_information()
|
|
self.assertEqual(3600, info["created_1"]["expireAfterSeconds"])
|
|
|
|
def test_index_drop_dups_silently_ignored(self):
|
|
class Customer(Document):
|
|
cust_id = IntField(unique=True, required=True)
|
|
meta = {
|
|
"indexes": ["cust_id"],
|
|
"index_drop_dups": True,
|
|
"allow_inheritance": False,
|
|
}
|
|
|
|
Customer.drop_collection()
|
|
Customer.objects.first()
|
|
|
|
def test_unique_and_indexes(self):
|
|
"""Ensure that 'unique' constraints aren't overridden by
|
|
meta.indexes.
|
|
"""
|
|
|
|
class Customer(Document):
|
|
cust_id = IntField(unique=True, required=True)
|
|
meta = {"indexes": ["cust_id"], "allow_inheritance": False}
|
|
|
|
Customer.drop_collection()
|
|
cust = Customer(cust_id=1)
|
|
cust.save()
|
|
|
|
cust_dupe = Customer(cust_id=1)
|
|
with self.assertRaises(NotUniqueError):
|
|
cust_dupe.save()
|
|
|
|
cust = Customer(cust_id=2)
|
|
cust.save()
|
|
|
|
# duplicate key on update
|
|
with self.assertRaises(NotUniqueError):
|
|
cust.cust_id = 1
|
|
cust.save()
|
|
|
|
def test_primary_save_duplicate_update_existing_object(self):
|
|
"""If you set a field as primary, then unexpected behaviour can occur.
|
|
You won't create a duplicate but you will update an existing document.
|
|
"""
|
|
|
|
class User(Document):
|
|
name = StringField(primary_key=True)
|
|
password = StringField()
|
|
|
|
User.drop_collection()
|
|
|
|
user = User(name="huangz", password="secret")
|
|
user.save()
|
|
|
|
user = User(name="huangz", password="secret2")
|
|
user.save()
|
|
|
|
self.assertEqual(User.objects.count(), 1)
|
|
self.assertEqual(User.objects.get().password, "secret2")
|
|
|
|
def test_unique_and_primary_create(self):
|
|
"""Create a new record with a duplicate primary key
|
|
throws an exception
|
|
"""
|
|
|
|
class User(Document):
|
|
name = StringField(primary_key=True)
|
|
password = StringField()
|
|
|
|
User.drop_collection()
|
|
|
|
User.objects.create(name="huangz", password="secret")
|
|
with self.assertRaises(NotUniqueError):
|
|
User.objects.create(name="huangz", password="secret2")
|
|
|
|
self.assertEqual(User.objects.count(), 1)
|
|
self.assertEqual(User.objects.get().password, "secret")
|
|
|
|
def test_index_with_pk(self):
|
|
"""Ensure you can use `pk` as part of a query"""
|
|
|
|
class Comment(EmbeddedDocument):
|
|
comment_id = IntField(required=True)
|
|
|
|
try:
|
|
|
|
class BlogPost(Document):
|
|
comments = EmbeddedDocumentField(Comment)
|
|
meta = {
|
|
"indexes": [
|
|
{"fields": ["pk", "comments.comment_id"], "unique": True}
|
|
]
|
|
}
|
|
|
|
except UnboundLocalError:
|
|
self.fail("Unbound local error at index + pk definition")
|
|
|
|
info = BlogPost.objects._collection.index_information()
|
|
info = [value["key"] for key, value in iteritems(info)]
|
|
index_item = [("_id", 1), ("comments.comment_id", 1)]
|
|
self.assertIn(index_item, info)
|
|
|
|
def test_compound_key_embedded(self):
|
|
class CompoundKey(EmbeddedDocument):
|
|
name = StringField(required=True)
|
|
term = StringField(required=True)
|
|
|
|
class ReportEmbedded(Document):
|
|
key = EmbeddedDocumentField(CompoundKey, primary_key=True)
|
|
text = StringField()
|
|
|
|
my_key = CompoundKey(name="n", term="ok")
|
|
report = ReportEmbedded(text="OK", key=my_key).save()
|
|
|
|
self.assertEqual(
|
|
{"text": "OK", "_id": {"term": "ok", "name": "n"}}, report.to_mongo()
|
|
)
|
|
self.assertEqual(report, ReportEmbedded.objects.get(pk=my_key))
|
|
|
|
def test_compound_key_dictfield(self):
|
|
class ReportDictField(Document):
|
|
key = DictField(primary_key=True)
|
|
text = StringField()
|
|
|
|
my_key = {"name": "n", "term": "ok"}
|
|
report = ReportDictField(text="OK", key=my_key).save()
|
|
|
|
self.assertEqual(
|
|
{"text": "OK", "_id": {"term": "ok", "name": "n"}}, report.to_mongo()
|
|
)
|
|
|
|
# We can't directly call ReportDictField.objects.get(pk=my_key),
|
|
# because dicts are unordered, and if the order in MongoDB is
|
|
# different than the one in `my_key`, this test will fail.
|
|
self.assertEqual(report, ReportDictField.objects.get(pk__name=my_key["name"]))
|
|
self.assertEqual(report, ReportDictField.objects.get(pk__term=my_key["term"]))
|
|
|
|
def test_string_indexes(self):
|
|
class MyDoc(Document):
|
|
provider_ids = DictField()
|
|
meta = {"indexes": ["provider_ids.foo", "provider_ids.bar"]}
|
|
|
|
info = MyDoc.objects._collection.index_information()
|
|
info = [value["key"] for key, value in iteritems(info)]
|
|
self.assertIn([("provider_ids.foo", 1)], info)
|
|
self.assertIn([("provider_ids.bar", 1)], info)
|
|
|
|
def test_sparse_compound_indexes(self):
|
|
class MyDoc(Document):
|
|
provider_ids = DictField()
|
|
meta = {
|
|
"indexes": [
|
|
{"fields": ("provider_ids.foo", "provider_ids.bar"), "sparse": True}
|
|
]
|
|
}
|
|
|
|
info = MyDoc.objects._collection.index_information()
|
|
self.assertEqual(
|
|
[("provider_ids.foo", 1), ("provider_ids.bar", 1)],
|
|
info["provider_ids.foo_1_provider_ids.bar_1"]["key"],
|
|
)
|
|
self.assertTrue(info["provider_ids.foo_1_provider_ids.bar_1"]["sparse"])
|
|
|
|
def test_text_indexes(self):
|
|
class Book(Document):
|
|
title = DictField()
|
|
meta = {"indexes": ["$title"]}
|
|
|
|
indexes = Book.objects._collection.index_information()
|
|
self.assertIn("title_text", indexes)
|
|
key = indexes["title_text"]["key"]
|
|
self.assertIn(("_fts", "text"), key)
|
|
|
|
def test_hashed_indexes(self):
|
|
class Book(Document):
|
|
ref_id = StringField()
|
|
meta = {"indexes": ["#ref_id"]}
|
|
|
|
indexes = Book.objects._collection.index_information()
|
|
self.assertIn("ref_id_hashed", indexes)
|
|
self.assertIn(("ref_id", "hashed"), indexes["ref_id_hashed"]["key"])
|
|
|
|
def test_indexes_after_database_drop(self):
|
|
"""
|
|
Test to ensure that indexes are re-created on a collection even
|
|
after the database has been dropped.
|
|
|
|
Issue #812
|
|
"""
|
|
# Use a new connection and database since dropping the database could
|
|
# cause concurrent tests to fail.
|
|
connection = connect(
|
|
db="tempdatabase", alias="test_indexes_after_database_drop"
|
|
)
|
|
|
|
class BlogPost(Document):
|
|
title = StringField()
|
|
slug = StringField(unique=True)
|
|
|
|
meta = {"db_alias": "test_indexes_after_database_drop"}
|
|
|
|
try:
|
|
BlogPost.drop_collection()
|
|
|
|
# Create Post #1
|
|
post1 = BlogPost(title="test1", slug="test")
|
|
post1.save()
|
|
|
|
# Drop the Database
|
|
connection.drop_database("tempdatabase")
|
|
|
|
# Re-create Post #1
|
|
post1 = BlogPost(title="test1", slug="test")
|
|
post1.save()
|
|
|
|
# Create Post #2
|
|
post2 = BlogPost(title="test2", slug="test")
|
|
self.assertRaises(NotUniqueError, post2.save)
|
|
finally:
|
|
# Drop the temporary database at the end
|
|
connection.drop_database("tempdatabase")
|
|
|
|
def test_index_dont_send_cls_option(self):
|
|
"""
|
|
Ensure that 'cls' option is not sent through ensureIndex. We shouldn't
|
|
send internal MongoEngine arguments that are not a part of the index
|
|
spec.
|
|
|
|
This is directly related to the fact that MongoDB doesn't validate the
|
|
options that are passed to ensureIndex. For more details, see:
|
|
https://jira.mongodb.org/browse/SERVER-769
|
|
"""
|
|
|
|
class TestDoc(Document):
|
|
txt = StringField()
|
|
|
|
meta = {
|
|
"allow_inheritance": True,
|
|
"indexes": [{"fields": ("txt",), "cls": False}],
|
|
}
|
|
|
|
class TestChildDoc(TestDoc):
|
|
txt2 = StringField()
|
|
|
|
meta = {"indexes": [{"fields": ("txt2",), "cls": False}]}
|
|
|
|
TestDoc.drop_collection()
|
|
TestDoc.ensure_indexes()
|
|
TestChildDoc.ensure_indexes()
|
|
|
|
index_info = TestDoc._get_collection().index_information()
|
|
for key in index_info:
|
|
del index_info[key][
|
|
"v"
|
|
] # drop the index version - we don't care about that here
|
|
if "ns" in index_info[key]:
|
|
del index_info[key][
|
|
"ns"
|
|
] # drop the index namespace - we don't care about that here, MongoDB 3+
|
|
if "dropDups" in index_info[key]:
|
|
del index_info[key][
|
|
"dropDups"
|
|
] # drop the index dropDups - it is deprecated in MongoDB 3+
|
|
|
|
self.assertEqual(
|
|
index_info,
|
|
{
|
|
"txt_1": {"key": [("txt", 1)], "background": False},
|
|
"_id_": {"key": [("_id", 1)]},
|
|
"txt2_1": {"key": [("txt2", 1)], "background": False},
|
|
"_cls_1": {"key": [("_cls", 1)], "background": False},
|
|
},
|
|
)
|
|
|
|
def test_compound_index_underscore_cls_not_overwritten(self):
|
|
"""
|
|
Test that the compound index doesn't get another _cls when it is specified
|
|
"""
|
|
|
|
class TestDoc(Document):
|
|
shard_1 = StringField()
|
|
txt_1 = StringField()
|
|
|
|
meta = {
|
|
"collection": "test",
|
|
"allow_inheritance": True,
|
|
"sparse": True,
|
|
"shard_key": "shard_1",
|
|
"indexes": [("shard_1", "_cls", "txt_1")],
|
|
}
|
|
|
|
TestDoc.drop_collection()
|
|
TestDoc.ensure_indexes()
|
|
|
|
index_info = TestDoc._get_collection().index_information()
|
|
self.assertIn("shard_1_1__cls_1_txt_1_1", index_info)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|