-
Notifications
You must be signed in to change notification settings - Fork 261
/
documents.py
185 lines (153 loc) · 5.77 KB
/
documents.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
from __future__ import unicode_literals
from django.db import models
from django.utils.six import add_metaclass, iteritems
from elasticsearch.helpers import bulk
from elasticsearch_dsl import DocType as DSLDocType
from elasticsearch_dsl.document import DocTypeMeta as DSLDocTypeMeta
from elasticsearch_dsl.field import Field
from .apps import DEDConfig
from .exceptions import ModelFieldNotMappedError, RedeclaredFieldError
from .fields import (
BooleanField,
DateField,
DEDField,
DoubleField,
FileField,
IntegerField,
LongField,
ShortField,
StringField,
)
from .indices import Index
from .registries import registry
model_field_class_to_field_class = {
models.AutoField: IntegerField,
models.BigIntegerField: LongField,
models.BooleanField: BooleanField,
models.CharField: StringField,
models.DateField: DateField,
models.DateTimeField: DateField,
models.EmailField: StringField,
models.FileField: FileField,
models.FilePathField: StringField,
models.FloatField: DoubleField,
models.ImageField: FileField,
models.IntegerField: IntegerField,
models.NullBooleanField: BooleanField,
models.PositiveIntegerField: IntegerField,
models.PositiveSmallIntegerField: ShortField,
models.SlugField: StringField,
models.SmallIntegerField: ShortField,
models.TextField: StringField,
models.TimeField: LongField,
models.URLField: StringField,
}
class DocTypeMeta(DSLDocTypeMeta):
def __new__(cls, name, bases, attrs):
"""
Subclass default DocTypeMeta to generate ES fields from django
models fields
"""
super_new = super(DocTypeMeta, cls).__new__
parents = [b for b in bases if isinstance(b, DocTypeMeta)]
if not parents:
return super_new(cls, name, bases, attrs)
model = attrs['Meta'].model
ignore_signals = getattr(attrs['Meta'], "ignore_signals", False)
auto_refresh = getattr(
attrs['Meta'], 'auto_refresh', DEDConfig.auto_refresh_enabled()
)
model_field_names = getattr(attrs['Meta'], "fields", [])
related_models = getattr(attrs['Meta'], "related_models", [])
class_fields = set(
name for name, field in iteritems(attrs)
if isinstance(field, Field)
)
cls = super_new(cls, name, bases, attrs)
cls._doc_type.model = model
cls._doc_type.ignore_signals = ignore_signals
cls._doc_type.auto_refresh = auto_refresh
cls._doc_type.related_models = related_models
doc = cls()
fields = model._meta.get_fields()
fields_lookup = dict((field.name, field) for field in fields)
for field_name in model_field_names:
if field_name in class_fields:
raise RedeclaredFieldError(
"You cannot redeclare the field named '{}' on {}"
.format(field_name, cls.__name__)
)
field_instance = doc.to_field(field_name,
fields_lookup[field_name])
cls._doc_type.mapping.field(field_name, field_instance)
cls._doc_type._fields = (
lambda: cls._doc_type.mapping.properties.properties.to_dict())
if getattr(cls._doc_type, 'index'):
index = Index(cls._doc_type.index)
index.doc_type(cls)
registry.register(index, doc)
return cls
@add_metaclass(DocTypeMeta)
class DocType(DSLDocType):
def __eq__(self, other):
return id(self) == id(other)
def __hash__(self):
return id(self)
def get_queryset(self):
"""
Return the queryset that should be indexed by this doc type.
"""
qs = self._doc_type.model._default_manager
return qs
def prepare(self, instance):
"""
Take a model instance, and turn it into a dict that can be serialized
based on the fields defined on this DocType subclass
"""
data = {}
for name, field in iteritems(self._doc_type._fields()):
if not isinstance(field, DEDField):
continue
if field._path == []:
field._path = [name]
prep_func = getattr(
self,
"prepare_" + name,
field.get_value_from_instance
)
data[name] = prep_func(instance)
return data
def to_field(self, field_name, model_field):
"""
Returns the elasticsearch field instance appropriate for the model
field class. This is a good place to hook into if you have more complex
model field to ES field logic
"""
try:
return model_field_class_to_field_class[
model_field.__class__](attr=field_name)
except KeyError:
raise ModelFieldNotMappedError(
"Cannot convert model field {} "
"to an Elasticsearch field!".format(field_name)
)
def bulk(self, actions, **kwargs):
return bulk(client=self.connection, actions=actions, **kwargs)
def update(self, thing, refresh=None, action='index', **kwargs):
"""
Update each document in ES for a model, iterable of models or queryset
"""
if refresh is True or (
refresh is None and self._doc_type.auto_refresh
):
kwargs['refresh'] = True
if isinstance(thing, models.Model):
thing = [thing]
actions = ({
'_op_type': action,
'_index': str(self._doc_type.index),
'_type': self._doc_type.mapping.doc_type,
'_id': model.pk,
'_source': self.prepare(model) if action != 'delete' else None,
} for model in thing)
return self.bulk(actions, **kwargs)