From 584f5a7049a587f1da398ad421f9afd40a13fde8 Mon Sep 17 00:00:00 2001
From: Karl Hobley <karl@torchbox.com>
Date: Fri, 20 Jun 2014 16:22:35 +0100
Subject: [PATCH] Index fields with correct type in ElasticSearch

Previously, everything was converted to a string before indexing in ElasticSearch.
This caused issues where certian filters may not work as expected (such as a greater than filter on an integer field)

This commit changes this by adding type conversion into the ElasticSearch backend.
---
 .../wagtailsearch/backends/elasticsearch.py   | 38 ++++++++++++++-----
 wagtail/wagtailsearch/indexed.py              | 20 ++++++++++
 2 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/wagtail/wagtailsearch/backends/elasticsearch.py b/wagtail/wagtailsearch/backends/elasticsearch.py
index 8fb706bbae..7ffebd0454 100644
--- a/wagtail/wagtailsearch/backends/elasticsearch.py
+++ b/wagtail/wagtailsearch/backends/elasticsearch.py
@@ -13,6 +13,32 @@ from wagtail.wagtailsearch.utils import normalise_query_string
 
 
 class ElasticSearchMapping(object):
+    TYPE_MAP = {
+        'AutoField': 'integer',
+        'BinaryField': 'binary',
+        'BooleanField': 'boolean',
+        'CharField': 'string',
+        'CommaSeparatedIntegerField': 'string',
+        'DateField': 'date',
+        'DateTimeField': 'date',
+        'DecimalField': 'double',
+        'FileField': 'string',
+        'FilePathField': 'string',
+        'FloatField': 'double',
+        'IntegerField': 'integer',
+        'BigIntegerField': 'long',
+        'IPAddressField': 'string',
+        'GenericIPAddressField': 'string',
+        'NullBooleanField': 'boolean',
+        'OneToOneField': 'integer',
+        'PositiveIntegerField': 'integer',
+        'PositiveSmallIntegerField': 'integer',
+        'SlugField': 'string',
+        'SmallIntegerField': 'integer',
+        'TextField': 'string',
+        'TimeField': 'date',
+    }
+
     def __init__(self, model):
         self.model = model
 
@@ -20,7 +46,7 @@ class ElasticSearchMapping(object):
         return self.model.indexed_get_content_type()
 
     def get_field_mapping(self, field):
-        mapping = {'type': 'string'}
+        mapping = {'type': self.TYPE_MAP.get(field.get_type(self.model), 'string')}
 
         if isinstance(field, SearchField):
             if field.boost:
@@ -60,14 +86,8 @@ class ElasticSearchMapping(object):
     def get_document(self, obj):
         # Build document
         doc = dict(pk=str(obj.pk), content_type=self.model.indexed_get_content_type())
-        for field in [field.get_attname(self.model) for field in self.model.get_search_fields()]:
-            if hasattr(obj, field):
-                doc[field] = getattr(obj, field)
-
-                # Check if this field is callable
-                if hasattr(doc[field], '__call__'):
-                    # Call it
-                    doc[field] = doc[field]()
+        for field in self.model.get_search_fields():
+            doc[field.get_index_name(self.model)] = field.get_value(obj)
 
         return doc
 
diff --git a/wagtail/wagtailsearch/indexed.py b/wagtail/wagtailsearch/indexed.py
index 4b636497bc..b290f440b4 100644
--- a/wagtail/wagtailsearch/indexed.py
+++ b/wagtail/wagtailsearch/indexed.py
@@ -138,6 +138,26 @@ class BaseField(object):
     def get_index_name(self, cls):
         return self.get_attname(cls) + self.suffix
 
+    def get_type(self, cls):
+        if 'type' in self.kwargs:
+            return self.kwargs['type']
+
+        try:
+            field = self.get_field(cls)
+            return field.get_internal_type()
+        except models.fields.FieldDoesNotExist:
+            return 'CharField'
+
+    def get_value(self, obj):
+        try:
+            field = self.get_field(obj.__class__)
+            return field._get_val_from_obj(obj)
+        except models.fields.FieldDoesNotExist:
+            value = getattr(obj, self.field_name, None)
+            if hasattr(value, '__call__'):
+                value = value()
+            return value
+
     def __repr__(self):
         return '<%s: %s>' % (self.__class__.__name__, self.field_name)