Configure Elasticsearch _id dynamically from document (#272)

* added Document method 'document_id' to customize documents' ids if needed

* renamed new method to 'generate_id' and added integration tests

* fixed compat with python 2

* Update tests/test_integration.py

Co-authored-by: Safwan Rahman <safwan.rahman15@gmail.com>

* Update tests/test_integration.py

Co-authored-by: Safwan Rahman <safwan.rahman15@gmail.com>

* generate_id is now a class method

* improved integration tests for generate_id method of Document class

* added commented broken unit tests for review / fix

* Fixing pull requests tests and documentations

Co-authored-by: Safwan Rahman <safwan.rahman15@gmail.com>

Author: Jean-François Schaff

django_elasticsearch_dsl/documents.py

@@ -156,11 +156,20 @@ def parallel_bulk(self, actions, **kwargs):
         # the result is currently not used upstream anyway.
         return (1, [])
 
+    @classmethod
+    def generate_id(cls, object_instance):
+        """
+        The default behavior is to use the Django object's pk (id) as the 
+        elasticseach index id (_id). If needed, this method can be overloaded 
+        to change this default behavior.
+        """
+        return object_instance.pk
+
     def _prepare_action(self, object_instance, action):
         return {
             '_op_type': action,
             '_index': self._index._name,
-            '_id': object_instance.pk,
+            '_id': self.generate_id(object_instance),
             '_source': (
                 self.prepare(object_instance) if action != 'delete' else None
             ),

docs/source/fields.rst

@@ -233,3 +233,46 @@ Available Fields
 ``properties`` is a dict where the key is a field name, and the value is a field
 instance.
 
+
+Document id
+===========
+
+The elasticsearch document id (``_id``) is not strictly speaking a field, as it is not 
+part of the document itself. The default behavior of ``django_elasticsearch_dsl``
+is to use the primary key of the model as the document's id (``pk`` or ``id``).
+Nevertheless, it can sometimes be useful to change this default behavior. For this, one
+can redefine the ``generate_id(cls, instance)`` class method of the ``Document`` class.
+
+For example, to use an article's slug as the elasticsearch ``_id`` instead of the 
+article's integer id, one could use:
+
+.. code-block:: python
+
+    # models.py
+
+    from django.db import models
+
+    class Article(models.Model):
+        # ... #
+
+        slug = models.SlugField(
+            max_length=255,
+            unique=True,
+        )
+
+        # ... #
+
+
+    # documents.py
+
+    from .models import Article
+
+    class ArticleDocument(Document):
+        class Django:
+            model = Article
+
+        # ... #
+
+        @classmethod
+        def generate_id(cls, article):
+            return article.slug

tests/documents.py

@@ -2,7 +2,7 @@
 from django_elasticsearch_dsl import Document, Index, fields
 from django_elasticsearch_dsl.registries import registry
 
-from .models import Ad, Category, Car, Manufacturer
+from .models import Ad, Category, Car, Manufacturer, Article
 
 index_settings = {
     'number_of_shards': 1,
@@ -148,5 +148,34 @@ class Index:
         settings = index_settings
 
 
+@registry.register_document
+class ArticleDocument(Document):
+    class Django:
+        model = Article
+        fields = [
+            'slug',
+        ]
+
+    class Index:
+        name = 'test_articles'
+        settings = index_settings
+
+@registry.register_document
+class ArticleWithSlugAsIdDocument(Document):
+    class Django:
+        model = Article
+        fields = [
+            'slug',
+        ]
+
+    class Index:
+        name = 'test_articles_with_slugs_as_doc_ids'
+        settings = index_settings
+
+    @classmethod
+    def generate_id(cls, article):
+        return article.slug
+
+
 ad_index = AdDocument._index
 car_index = CarDocument._index

tests/models.py

@@ -87,3 +87,17 @@ class Meta:
 
     def __str__(self):
         return self.title
+
+
+@python_2_unicode_compatible
+class Article(models.Model):
+    slug = models.CharField(
+        max_length=255,
+        unique=True,
+    )
+
+    class Meta:
+        app_label = 'tests'
+
+    def __str__(self):
+        return self.slug

tests/test_documents.py

@@ -1,3 +1,4 @@
+import json
 from unittest import TestCase
 
 from django.db import models
@@ -6,12 +7,15 @@
 from mock import patch, Mock, PropertyMock
 
 from django_elasticsearch_dsl import fields
-from django_elasticsearch_dsl.documents import DocType
+from django_elasticsearch_dsl.documents import DocType, Document
 from django_elasticsearch_dsl.exceptions import (ModelFieldNotMappedError,
                                                  RedeclaredFieldError)
 from django_elasticsearch_dsl.registries import registry
 from tests import ES_MAJOR_VERSION
 
+from .models import Article
+from .documents import ArticleDocument, ArticleWithSlugAsIdDocument
+
 
 class Car(models.Model):
     name = models.CharField(max_length=255)
@@ -346,3 +350,64 @@ def test_init_prepare_results(self):
         self.assertEqual(sorted([tuple(x) for x in m.method_calls], key=lambda _: _[0]),
                          [('name', (), {}),  ('price', (), {}), ('type', (), {})]
         )
+
+    # Mock the elasticsearch connection because we need to execute the bulk so that the generator
+    # got iterated and generate_id called.
+    # If we mock the bulk in django_elasticsearch_dsl.document
+    # the actual bulk will be never called and the test will fail
+    @patch('elasticsearch_dsl.connections.Elasticsearch.bulk')
+    def test_default_generate_id_is_called(self, _):
+        article = Article(
+            id=124594,
+            slug='some-article',
+        )
+        @registry.register_document
+        class ArticleDocument(DocType):
+            class Django:
+                model = Article
+                fields = [
+                    'slug',
+                ]
+
+            class Index:
+                name = 'test_articles'
+                settings = {
+                    'number_of_shards': 1,
+                    'number_of_replicas': 0,
+                }
+
+        with patch.object(ArticleDocument, 'generate_id',
+                          return_value=article.id) as patched_method:
+            d = ArticleDocument()
+            d.update(article)
+            patched_method.assert_called()
+
+    @patch('elasticsearch_dsl.connections.Elasticsearch.bulk')
+    def test_custom_generate_id_is_called(self, mock_bulk):
+        article = Article(
+            id=54218,
+            slug='some-article-2',
+        )
+
+        @registry.register_document
+        class ArticleDocument(DocType):
+            class Django:
+                model = Article
+                fields = [
+                    'slug',
+                ]
+
+            class Index:
+                name = 'test_articles'
+
+            @classmethod
+            def generate_id(cls, article):
+                return article.slug
+
+        d = ArticleDocument()
+        d.update(article)
+
+        # Get the data from the elasticsearch low level API because
+        # The generator get executed there.
+        data = json.loads(mock_bulk.call_args[0][0].split("\n")[0])
+        assert data["index"]["_id"] == article.slug

tests/test_integration.py

@@ -7,6 +7,7 @@
 from django.utils.translation import ugettext_lazy as _
 from six import StringIO
 
+from elasticsearch.exceptions import NotFoundError
 from elasticsearch_dsl import Index as DSLIndex
 from django_elasticsearch_dsl.test import ESTestCase
 from tests import ES_MAJOR_VERSION
@@ -18,9 +19,11 @@
     CarDocument,
     CarWithPrepareDocument,
     ManufacturerDocument,
+    ArticleDocument,
+    ArticleWithSlugAsIdDocument,
     index_settings
 )
-from .models import Car, Manufacturer, Ad, Category, COUNTRIES
+from .models import Car, Manufacturer, Ad, Category, Article, COUNTRIES
 
 
 @unittest.skipUnless(
@@ -352,3 +355,43 @@ def test_queryset_iterator_queries(self):
                 set([ad3.pk, self.ad1.pk, self.ad2.pk])
             )
 
+    def test_default_document_id(self):
+        obj_id = 12458
+        article_slug = "some-article"
+        article = Article(
+            id=obj_id,
+            slug=article_slug,
+        )
+
+        # saving should create two documents (in the two indices): one with the 
+        # Django object's id as the ES doc _id, and the other with the slug 
+        # as the ES _id
+        article.save()
+
+        # assert that the document's id is the id of the Django object
+        try:
+            es_obj = ArticleDocument.get(id=obj_id)
+        except NotFoundError:
+            self.fail("document with _id {} not found").format(obj_id)
+        self.assertEqual(es_obj.slug, article.slug)
+
+    def test_custom_document_id(self):
+        article_slug = "my-very-first-article"
+        article = Article(
+            slug=article_slug,
+        )
+
+        # saving should create two documents (in the two indices): one with the 
+        # Django object's id as the ES doc _id, and the other with the slug 
+        # as the ES _id
+        article.save()
+
+        # assert that the document's id is its the slug
+        try:
+            es_obj = ArticleWithSlugAsIdDocument.get(id=article_slug)
+        except NotFoundError:
+            self.fail(
+                "document with _id '{}' not found: "
+                "using a custom id is broken".format(article_slug)
+            )
+        self.assertEqual(es_obj.slug, article.slug)