├── tests ├── __init__.py ├── es_configs │ ├── index_1_v1.json │ └── index_2_v1.json ├── celery_app.py ├── settings.py ├── models.py ├── base.py ├── test_commands.py └── test_mixins.py ├── rubber ├── management │ ├── __init__.py │ ├── commands │ │ ├── __init__.py │ │ ├── es_create_index.py │ │ └── es_create_documents.py │ └── base.py ├── version.py ├── models.py ├── __init__.py ├── tasks.py ├── mixins.py └── apps.py ├── setup.cfg ├── .gitignore ├── versions.cfg ├── LICENSE.txt ├── setup.py ├── README.md └── bootstrap.py /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /rubber/management/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/es_configs/index_1_v1.json: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/es_configs/index_2_v1.json: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /rubber/management/commands/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /rubber/version.py: -------------------------------------------------------------------------------- 1 | __version__ = '1.0' 2 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | description-file = README.md 3 | -------------------------------------------------------------------------------- /rubber/models.py: -------------------------------------------------------------------------------- 1 | """ 2 | Empty models file for older django versions. 3 | """ 4 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .installed.cfg 2 | *.pyc 3 | *~ 4 | bin/ 5 | cover/ 6 | build/ 7 | test/ 8 | develop-eggs/ 9 | dist/ 10 | django_rubber.egg-info/ 11 | eggs/ 12 | logs/ 13 | parts/ 14 | ve/ 15 | .* 16 | -------------------------------------------------------------------------------- /rubber/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Init for rubber. 3 | """ 4 | from .version import __version__ # noqa 5 | 6 | from rubber.apps import get_rubber_config # noqa 7 | 8 | default_app_config = 'rubber.apps.RubberConfig' 9 | -------------------------------------------------------------------------------- /versions.cfg: -------------------------------------------------------------------------------- 1 | [versions] 2 | celery = 4.4.3 3 | django = 3.0.6 4 | elasticsearch = 7.7.1 5 | elasticsearch-dsl = 7.3.0 6 | nose = 1.3.7 7 | nose-sfd = 0.4 8 | -------------------------------------------------------------------------------- /tests/celery_app.py: -------------------------------------------------------------------------------- 1 | """ 2 | Celery app for rubber tests. 3 | """ 4 | from django.conf import settings 5 | 6 | from celery import Celery 7 | 8 | 9 | app = Celery('rubber') 10 | 11 | app.config_from_object('django.conf:settings') 12 | app.autodiscover_tasks(lambda: settings.INSTALLED_APPS) 13 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2015 Laurent Guilbert 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 6 | 7 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 8 | 9 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 10 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | """ 2 | Setup for rubber. 3 | """ 4 | import sys 5 | 6 | from setuptools import find_packages 7 | from setuptools import setup 8 | 9 | exec(open('rubber/version.py').read()) 10 | 11 | install_requires = [ 12 | 'elasticsearch', 13 | 'elasticsearch-dsl', 14 | 'celery', 15 | 'six', 16 | 'tqdm', 17 | ] 18 | if sys.version_info.major == 2: 19 | install_requires.append('futures') 20 | 21 | setup( 22 | name='django-rubber', 23 | version=__version__, # noqa 24 | keywords='django, elasticsearch', 25 | author='Laurent Guilbert', 26 | author_email='laurent@guilbert.me', 27 | url='https://github.com/liberation/django-rubber', 28 | description="No-frills Elasticsearch's wrapper for your Django project.", 29 | license='MIT License', 30 | classifiers=( 31 | 'Framework :: Django', 32 | 'Environment :: Web Environment', 33 | 'Programming Language :: Python', 34 | 'Programming Language :: Python :: 3', 35 | 'Intended Audience :: Developers', 36 | 'Operating System :: OS Independent', 37 | 'Topic :: Software Development :: Libraries :: Python Modules', 38 | ), 39 | zip_safe=False, 40 | include_package_data=True, 41 | packages=find_packages(), 42 | install_requires=install_requires, 43 | ) 44 | -------------------------------------------------------------------------------- /tests/settings.py: -------------------------------------------------------------------------------- 1 | """ 2 | Test settings for rubber. 3 | """ 4 | import os 5 | 6 | SITE_ROOT = os.path.realpath(os.path.dirname(__file__)) 7 | 8 | DEBUG = True 9 | 10 | DATABASES = { 11 | 'default': { 12 | 'NAME': 'rubber.db', 13 | 'ENGINE': 'django.db.backends.sqlite3', 14 | 'TEST': { 15 | 'NAME': 'rubber.db', 16 | }, 17 | 'TEST_NAME': 'rubber.db', 18 | } 19 | } 20 | 21 | INSTALLED_APPS = ( 22 | 'django.contrib.contenttypes', 23 | 'rubber', 24 | 'tests', 25 | ) 26 | 27 | LOGGING = { 28 | 'version': 1, 29 | 'handlers': { 30 | 'console': { 31 | 'class': 'logging.StreamHandler', 32 | }, 33 | }, 34 | 'loggers': { 35 | 'rubber.tasks': { 36 | 'handlers': ['console'], 37 | 'level': 'DEBUG', 38 | }, 39 | } 40 | } 41 | 42 | SECRET_KEY = 'secret-key' 43 | 44 | ################################################## 45 | # Rubber # 46 | ################################################## 47 | 48 | RUBBER = { 49 | 'MODELS': [ 50 | 'tests.models.Token', 51 | ], 52 | 'CONFIG_ROOT': os.path.join(SITE_ROOT, 'es_configs'), 53 | 'OPTIONS': { 54 | 'disabled': False, 55 | 'fail_silently': True, 56 | }, 57 | } 58 | 59 | ################################################## 60 | # Celery # 61 | ################################################## 62 | 63 | CELERY_ALWAYS_EAGER = True 64 | CELERY_EAGER_PROPAGATES_EXCEPTIONS = True 65 | 66 | from . import celery_app # noqa 67 | -------------------------------------------------------------------------------- /rubber/tasks.py: -------------------------------------------------------------------------------- 1 | """ 2 | Celery tasks for rubber. 3 | """ 4 | import logging 5 | 6 | from django.contrib.contenttypes.models import ContentType 7 | 8 | from celery import shared_task 9 | 10 | from rubber import get_rubber_config 11 | 12 | 13 | logger = logging.getLogger(__name__) 14 | rubber_config = get_rubber_config() 15 | 16 | 17 | @shared_task 18 | def es_bulk(body, fail_silently=None): 19 | if fail_silently is None: 20 | fail_silently = rubber_config.should_fail_silently 21 | try: 22 | rubber_config.es.bulk(body=body) 23 | except: 24 | if fail_silently: 25 | logger.error( 26 | "Exception occured in es_bulk.", 27 | exc_info=True, 28 | extra={'body': body} 29 | ) 30 | else: 31 | raise 32 | 33 | 34 | @shared_task 35 | def es_index_object(content_type_id, object_id, fail_silently=None): 36 | if fail_silently is None: 37 | fail_silently = rubber_config.should_fail_silently 38 | try: 39 | content_type = ContentType.objects.get_for_id(content_type_id) 40 | obj = content_type.model_class()._default_manager.get(pk=object_id) 41 | if not obj.is_indexable(): 42 | return 43 | rubber_config.es.bulk(body=obj.get_es_index_body()) 44 | except: 45 | if fail_silently: 46 | logger.error( 47 | "Exception occured while indexing object.", 48 | exc_info=True, 49 | extra={ 50 | 'content_type_id': content_type_id, 51 | 'object_id': object_id, 52 | } 53 | ) 54 | else: 55 | raise 56 | -------------------------------------------------------------------------------- /rubber/management/commands/es_create_index.py: -------------------------------------------------------------------------------- 1 | """ 2 | Management command for rubber. 3 | """ 4 | import os 5 | import sys 6 | 7 | from rubber.management.base import ESBaseCommand 8 | 9 | 10 | class Command(ESBaseCommand): 11 | def add_arguments(self, parser): 12 | parser.add_argument('indexes', 13 | action='store', 14 | type=str, 15 | nargs='*', 16 | help=( 17 | "List of indexes names to create" 18 | ) 19 | ) 20 | parser.add_argument('--dry-run', 21 | action='store_true', 22 | dest='dry_run', 23 | default=False, 24 | help=( 25 | "Run the command in dry run mode without actually changing " 26 | "anything." 27 | ) 28 | ) 29 | 30 | def run(self, *args, **options): 31 | if len(options['indexes']) == 0: 32 | self.print_error("Please provide at least one index.") 33 | sys.exit(1) 34 | for index in options['indexes']: 35 | config_path = os.path.join( 36 | self.rubber_config.config_root, 37 | '{0}.json'.format(index) 38 | ) 39 | self.print_info(u"Using config file : {0}".format(config_path)) 40 | body = None 41 | try: 42 | with open(config_path, 'r') as f: 43 | body = f.read() 44 | except IOError: 45 | self.print_error("Config file does not exist.") 46 | continue 47 | self.rubber_config.es.indices.create(index=index, body=body) 48 | self.print_success(u"Index {0} created.".format(index)) 49 | -------------------------------------------------------------------------------- /tests/models.py: -------------------------------------------------------------------------------- 1 | """ 2 | Models for rubber tests. 3 | """ 4 | from django.db import models 5 | 6 | from elasticsearch_dsl import Document 7 | from elasticsearch_dsl import Text 8 | 9 | from rubber.mixins import ESIndexableMixin 10 | 11 | 12 | class TokenDocType(Document): 13 | name = Text() 14 | number = Text() 15 | multi = Text(multi=True) 16 | 17 | class Meta: 18 | doc_type = 'token' 19 | 20 | class Index: 21 | name = 'index_2' 22 | 23 | 24 | class TokenSerializer(object): 25 | 26 | def __init__(self, token, *args, **kwargs): 27 | if token.name == 'raise_exception': 28 | raise RuntimeError 29 | self.token = token 30 | 31 | @property 32 | def data(self): 33 | return { 34 | 'name': self.token.name, 35 | 'number': self.token.number, 36 | 'multi': ['item_1', 'item_2'] 37 | } 38 | 39 | 40 | class Token(ESIndexableMixin, models.Model): 41 | modified_at = models.DateTimeField(auto_now=True) 42 | name = models.CharField(default='token', max_length=200) 43 | number = models.IntegerField(default=42) 44 | 45 | def __unicode__(self): 46 | return self.name 47 | 48 | def get_es_indexers(self): 49 | return { 50 | 'INDEX_1': { 51 | 'version': 1, 52 | 'index': 'index_1', 53 | 'serializer': TokenSerializer, 54 | 'doc_type': 'token' 55 | }, 56 | 'INDEX_2': { 57 | 'version': 1, 58 | 'dsl_doc_type': TokenDocType, 59 | 'dsl_doc_type_mapping': self.dsl_doc_type_mapping 60 | }, 61 | } 62 | 63 | def dsl_doc_type_mapping(self): 64 | if self.name == 'raise_exception': 65 | raise RuntimeError 66 | doc = TokenDocType() 67 | doc.name = self.name 68 | doc.number = self.number 69 | doc.multi.append('item_1') 70 | doc.multi.append('item_2') 71 | return doc 72 | 73 | def is_indexable(self): 74 | if self.name == 'not_indexable': 75 | return False 76 | return True 77 | -------------------------------------------------------------------------------- /tests/base.py: -------------------------------------------------------------------------------- 1 | """ 2 | Base test case for rubber. 3 | """ 4 | from django.test import TransactionTestCase 5 | 6 | from rubber import get_rubber_config 7 | 8 | rubber_config = get_rubber_config() 9 | 10 | 11 | class BaseTestCase(TransactionTestCase): 12 | 13 | def setUp(self): 14 | super(BaseTestCase, self).setUp() 15 | self.rubber_config = rubber_config 16 | 17 | def refresh(self): 18 | rubber_config.es.indices.refresh('_all') 19 | 20 | def docExists(self, obj, indexer_index, doc_id=None): 21 | doc = obj.get_es_doc(indexer_index) 22 | if doc is not None: 23 | return True 24 | else: 25 | return False 26 | 27 | def aliasExists(self, index, name): 28 | return rubber_config.es.indices.exists_alias( 29 | index=index, name=name) 30 | 31 | def indexExists(self, index): 32 | return rubber_config.es.indices.exists(index=index) 33 | 34 | def typeExists(self, index, doc_type): 35 | return rubber_config.es.indices.exists_type( 36 | index=index, 37 | doc_type=doc_type 38 | ) 39 | 40 | def assertAliasExists(self, index, name): 41 | self.assertTrue(self.aliasExists(index, name)) 42 | 43 | def assertAliasDoesntExist(self, index, name): 44 | self.assertFalse(self.aliasExists(index, name)) 45 | 46 | def assertIndexExists(self, index): 47 | self.assertTrue(self.indexExists(index)) 48 | 49 | def assertIndexDoesntExist(self, index): 50 | self.assertFalse(self.indexExists(index)) 51 | 52 | def assertTypeExists(self, index, doc_type): 53 | self.assertTrue(self.typeExists(index, doc_type)) 54 | 55 | def assertTypeDoesntExist(self, index, doc_type): 56 | self.assertFalse(self.typeExists(index, doc_type)) 57 | 58 | def assertDocExists(self, obj, indexer_index): 59 | self.assertTrue(self.docExists(obj, indexer_index)) 60 | 61 | def assertDocDoesntExist(self, obj, indexer_index): 62 | self.assertFalse(self.docExists(obj, indexer_index)) 63 | 64 | def createIndex(self, index): 65 | rubber_config.es.indices.create(index=index) 66 | self.refresh() 67 | 68 | def deleteIndex(self, index): 69 | rubber_config.es.indices.delete(index=index, ignore=404) 70 | self.refresh() 71 | -------------------------------------------------------------------------------- /rubber/management/base.py: -------------------------------------------------------------------------------- 1 | """ 2 | Base management command for rubber. 3 | """ 4 | from __future__ import print_function 5 | from optparse import make_option 6 | import sys 7 | import traceback 8 | 9 | import six 10 | 11 | from django.core.management.base import BaseCommand 12 | 13 | from rubber import get_rubber_config 14 | 15 | 16 | class ESBaseCommand(BaseCommand): 17 | required_options = [] 18 | 19 | def add_arguments(self, parser): 20 | parser.add_argument('--dry-run', 21 | action='store_true', 22 | dest='dry_run', 23 | default=False, 24 | help=( 25 | "Run the command in dry run mode without actually changing " 26 | "anything." 27 | ) 28 | ) 29 | 30 | def handle(self, *args, **options): 31 | self.rubber_config = get_rubber_config() 32 | try: 33 | self.parse_options(**options) 34 | self.run(*args, **options) 35 | except Exception as exc: 36 | if options.get('traceback', False): 37 | traceback.print_exc() 38 | self.print_error(repr(exc)) 39 | sys.exit(1) 40 | 41 | def parse_options(self, **options): 42 | for required_option in self.required_options: 43 | if options.get(required_option) is None: 44 | self.print_error("{0} is required (use -h for help).".format( 45 | required_option)) 46 | sys.exit(1) 47 | 48 | for key, value in six.iteritems(options): 49 | setattr(self, key, value) 50 | 51 | try: 52 | self.verbosity = int(self.verbosity) 53 | except ValueError: 54 | self.verbosity = 1 55 | 56 | ################################################## 57 | # Print # 58 | ################################################## 59 | 60 | GREEN = '\033[92m' 61 | BLUE = '\033[94m' 62 | YELLOW = '\033[93m' 63 | RED = '\033[91m' 64 | DIM = '\033[2m' 65 | BOLD = '\033[1m' 66 | PURPLE = '\033[35m' 67 | UNDERLINE = '\033[4m' 68 | RESET = '\033[0m' 69 | 70 | def confirm(self, message): # pragma: no cover 71 | if self.yes: 72 | return True 73 | message += u" [Y/n]" 74 | self.print_warning(message) 75 | try: 76 | choice = raw_input() 77 | except NameError: 78 | choice = input() 79 | if choice != 'Y': 80 | self.print_error("Operation canceled.") 81 | sys.exit(1) 82 | 83 | def print_normal(self, message, verbosity=1): 84 | if self.verbosity >= verbosity: 85 | print(message) 86 | 87 | def print_info(self, message, verbosity=1): 88 | if self.verbosity >= verbosity: 89 | print(u"{0}{1}{2}".format(self.BLUE, message, self.RESET)) 90 | 91 | def print_success(self, message, verbosity=1): 92 | if self.verbosity >= verbosity: 93 | print(u"{0}{1}{2}".format(self.GREEN, message, self.RESET)) 94 | 95 | def print_error(self, message): 96 | print( 97 | u"{0}{1}{2}".format(self.RED, message, self.RESET), 98 | file=sys.stderr 99 | ) 100 | 101 | def print_warning(self, message): # pragma: no cover 102 | print(u"{0}{1}{2}".format(self.YELLOW, message, self.RESET)) 103 | -------------------------------------------------------------------------------- /tests/test_commands.py: -------------------------------------------------------------------------------- 1 | """ 2 | Test management commands for rubber. 3 | """ 4 | from django.conf import settings 5 | from django.core.management import call_command 6 | 7 | from tests.base import BaseTestCase 8 | from tests.models import Token 9 | 10 | 11 | class TestCreateIndex(BaseTestCase): 12 | 13 | def tearDown(self): 14 | super(TestCreateIndex, self).tearDown() 15 | # Delete remnants of previous tests. 16 | self.deleteIndex('index_1_v1') 17 | 18 | def test_create_index(self): 19 | # Index name is required. 20 | with self.assertRaises(SystemExit): 21 | call_command('es_create_index') 22 | 23 | # Fail silently if file doesn't exist. 24 | call_command('es_create_index', 'index_1_v404') 25 | 26 | call_command('es_create_index', 'index_1_v1') 27 | 28 | 29 | class TestCreateDocuments(BaseTestCase): 30 | 31 | def setUp(self): 32 | super(TestCreateDocuments, self).setUp() 33 | self.createIndex('index_1_v1') 34 | self.createIndex('index_2_v1') 35 | self.refresh() 36 | 37 | def tearDown(self): 38 | super(TestCreateDocuments, self).tearDown() 39 | # Delete remnants of previous tests. 40 | self.deleteIndex('index_1_v1') 41 | self.deleteIndex('index_2_v1') 42 | 43 | def test_es_create_documents_models(self): 44 | settings.RUBBER['OPTIONS']['disabled'] = True 45 | token = Token.objects.create() 46 | settings.RUBBER['OPTIONS']['disabled'] = False 47 | 48 | call_command('es_create_documents', models='YOLO') 49 | self.assertDocDoesntExist(token, 'INDEX_1') 50 | self.assertDocDoesntExist(token, 'INDEX_2') 51 | 52 | call_command('es_create_documents', models='tests.models.Token') 53 | self.assertDocExists(token, 'INDEX_2') 54 | self.assertDocExists(token, 'INDEX_2') 55 | 56 | def test_es_create_documents_from(self): 57 | settings.RUBBER['OPTIONS']['disabled'] = True 58 | token = Token.objects.create() 59 | settings.RUBBER['OPTIONS']['disabled'] = False 60 | self.assertDocDoesntExist(token, 'INDEX_1') 61 | self.assertDocDoesntExist(token, 'INDEX_2') 62 | 63 | with self.assertRaises(SystemExit): 64 | call_command('es_create_documents', from_date='foobar') 65 | 66 | call_command('es_create_documents', from_date='2008-09-03T20:56:35') 67 | 68 | def test_es_create_documents(self): 69 | settings.RUBBER['OPTIONS']['disabled'] = True 70 | token = Token.objects.create() 71 | settings.RUBBER['OPTIONS']['disabled'] = False 72 | self.assertDocDoesntExist(token, 'INDEX_1') 73 | self.assertDocDoesntExist(token, 'INDEX_2') 74 | 75 | # Dry run. 76 | call_command('es_create_documents', dry_run=True) 77 | self.assertDocDoesntExist(token, 'INDEX_1') 78 | self.assertDocDoesntExist(token, 'INDEX_2') 79 | 80 | call_command('es_create_documents') 81 | self.assertDocExists(token, 'INDEX_1') 82 | self.assertDocExists(token, 'INDEX_2') 83 | 84 | settings.RUBBER['OPTIONS']['disabled'] = True 85 | token = Token.objects.create(name='raise_exception') 86 | settings.RUBBER['OPTIONS']['disabled'] = False 87 | 88 | # def test_es_create_indices(self): 89 | # # Dry run. 90 | # call_command('es_create_indices', dry_run=True) 91 | # self.assertIndexDoesntExist(Token.get_es_index()) 92 | # 93 | # call_command('es_create_indices') 94 | # self.assertIndexExists(Token.get_es_index()) 95 | # 96 | # # Skip already created indices silently. 97 | # call_command('es_create_indices') 98 | -------------------------------------------------------------------------------- /rubber/mixins.py: -------------------------------------------------------------------------------- 1 | """ 2 | Mixins for rubber. 3 | """ 4 | import logging 5 | import json 6 | from uuid import UUID 7 | 8 | 9 | from django.contrib.contenttypes.models import ContentType 10 | 11 | from elasticsearch_dsl.serializer import serializer as dsl_serializer 12 | 13 | from rubber import get_rubber_config 14 | from rubber.tasks import es_bulk 15 | from rubber.tasks import es_index_object 16 | 17 | logger = logging.getLogger(__name__) 18 | rubber_config = get_rubber_config() 19 | 20 | 21 | class EnhancedJsonEncoder(json.JSONEncoder): 22 | def default(self, obj): 23 | if isinstance(obj, UUID): 24 | return str(obj) 25 | # Let the base class default method raise the TypeError 26 | return json.JSONEncoder.default(self, obj) 27 | 28 | 29 | class ESIndexableMixin(object): 30 | """ 31 | Provide the required methods and attributes to index django models. 32 | """ 33 | es_indexers = {} 34 | es_reference_date = 'modified_at' 35 | 36 | @classmethod 37 | def get_indexable_queryset(cls): 38 | return cls._default_manager.all() 39 | 40 | def get_es_indexers(self): 41 | return self.es_indexers 42 | 43 | def is_indexable(self): 44 | return True 45 | 46 | def get_es_indexer_meta(self, indexer): 47 | version = indexer.get('version') 48 | if 'dsl_doc_type' in indexer: 49 | index = indexer['dsl_doc_type']._index._name 50 | doc_type = indexer['dsl_doc_type']._doc_type.name 51 | else: 52 | index = indexer['index'] 53 | doc_type = indexer['doc_type'] 54 | if version is not None: 55 | index = '{0}_v{1}'.format(index, version) 56 | return (index, doc_type, version) 57 | 58 | def get_es_doc(self, indexer_key): 59 | if not self.pk: 60 | return None 61 | indexer = self.get_es_indexers()[indexer_key] 62 | index, doc_type, version = self.get_es_indexer_meta(indexer) 63 | result = rubber_config.es.get( 64 | index=index, 65 | id=self.pk, 66 | ignore=404 67 | ) 68 | if 'found' not in result or result['found'] is False: 69 | return None 70 | return result 71 | 72 | def get_es_index_body(self): 73 | requests = [] 74 | for _, indexer in iter(self.get_es_indexers().items()): 75 | index, doc_type, version = self.get_es_indexer_meta(indexer) 76 | requests.append({ 77 | 'index': { 78 | '_index': index, 79 | '_id': self.pk 80 | } 81 | }) 82 | if 'dsl_doc_type' in indexer: 83 | doc = indexer['dsl_doc_type_mapping']() 84 | requests.append(doc) 85 | else: 86 | body = indexer['serializer'](self, context={'request': None}).data 87 | requests.append(body) 88 | return u"\n".join([ 89 | dsl_serializer.dumps(request) for request in requests 90 | ]) 91 | 92 | def get_es_delete_body(self): 93 | requests = [] 94 | for _, indexer in iter(self.get_es_indexers().items()): 95 | index, doc_type, version = self.get_es_indexer_meta(indexer) 96 | requests.append({ 97 | 'delete': { 98 | '_index': index, 99 | '_id': self.pk 100 | } 101 | }) 102 | return u"\n".join([json.dumps(request, cls=EnhancedJsonEncoder) for request in requests]) 103 | 104 | def es_index(self, is_async=True, countdown=0): 105 | if rubber_config.is_disabled or not self.is_indexable(): 106 | return 107 | content_type = ContentType.objects.get_for_model(self) 108 | if is_async: 109 | es_index_object.apply_async( 110 | args=(content_type.pk, self.pk,), 111 | countdown=countdown, 112 | queue=rubber_config.celery_queue 113 | ) 114 | else: 115 | if rubber_config.should_fail_silently: 116 | es_index_object.apply(args=(content_type.pk, self.pk,)) 117 | else: 118 | es_index_object.run(content_type.pk, self.pk) 119 | 120 | def es_delete(self, is_async=True): 121 | if rubber_config.is_disabled: 122 | return 123 | body = self.get_es_delete_body() 124 | if is_async: 125 | es_bulk.apply_async( 126 | args=(body,), 127 | queue=rubber_config.celery_queue 128 | ) 129 | else: 130 | if rubber_config.should_fail_silently: 131 | es_bulk.apply((body,)) 132 | else: 133 | es_bulk.run(body) 134 | -------------------------------------------------------------------------------- /rubber/management/commands/es_create_documents.py: -------------------------------------------------------------------------------- 1 | """ 2 | Management command for rubber. 3 | """ 4 | from datetime import datetime 5 | from optparse import make_option 6 | import concurrent.futures as futures 7 | import sys 8 | 9 | from django.core.paginator import Paginator 10 | 11 | from tqdm import tqdm 12 | 13 | from rubber.management.base import ESBaseCommand 14 | 15 | 16 | class Command(ESBaseCommand): 17 | def add_arguments(self, parser): 18 | parser.add_argument('--dry-run', 19 | action='store_true', 20 | dest='dry_run', 21 | default=False, 22 | help=( 23 | "Run the command in dry run mode without actually changing " 24 | "anything." 25 | ) 26 | ) 27 | parser.add_argument('--show-tqdm', 28 | action='store_true', 29 | dest='show_tqdm', 30 | default=False, 31 | help="Show tqdm progress bar." 32 | ) 33 | parser.add_argument('--from', 34 | action='store', 35 | type=str, 36 | dest='from_date', 37 | help=( 38 | "Filter queryset by date. " 39 | "Must be formatted as YYYY-MM-DDTHH:MM:SS." 40 | ) 41 | ) 42 | parser.add_argument('--models', 43 | action='store', 44 | type=str, 45 | dest='models', 46 | help=( 47 | "Comma separated list of models to be indexed. It must match " 48 | "at least one of the models defined in RUBBER settings." 49 | ) 50 | ) 51 | 52 | def get_models_paths(self): 53 | if not self.models: 54 | return self.rubber_config.models_paths 55 | models_paths = [ 56 | model_path for model_path in self.rubber_config.models_paths 57 | if model_path in self.models.split(',') 58 | ] 59 | return models_paths 60 | 61 | def get_from_date(self): 62 | if not self.from_date: 63 | return None 64 | try: 65 | from_date = datetime.strptime( 66 | self.from_date, 67 | '%Y-%m-%dT%H:%M:%S' 68 | ) 69 | except Exception as exc: 70 | self.print_error(exc) 71 | sys.exit(1) 72 | else: 73 | return from_date 74 | return None 75 | 76 | def run(self, *args, **options): 77 | from_date = self.get_from_date() 78 | if from_date is not None: 79 | self.print_info(u"Reference date : {0}".format(from_date)) 80 | 81 | models_paths = self.get_models_paths() 82 | indexable_models = self.rubber_config.get_models_from_paths( 83 | models_paths) 84 | self.print_info(u"Models : {0}".format(indexable_models)) 85 | 86 | for model in indexable_models: 87 | self.print_success( 88 | u"Indexing model: '{0}'.".format(model.__name__)) 89 | queryset = model.get_indexable_queryset() 90 | 91 | if from_date is not None and model.es_reference_date is not None: 92 | filter_dict = {} 93 | filter_name = '{0}__gt'.format(model.es_reference_date) 94 | filter_dict[filter_name] = from_date 95 | queryset = queryset.filter(**filter_dict) 96 | 97 | max_bulk_size = 200 98 | paginator = Paginator(queryset, max_bulk_size) 99 | if self.show_tqdm: 100 | pbar = tqdm(total=paginator.count) 101 | else: 102 | total = 0 103 | executor = futures.ThreadPoolExecutor(max_workers=8) 104 | for page_number in paginator.page_range: 105 | page = paginator.page(page_number) 106 | if len(page.object_list) == 0: 107 | continue 108 | tasks = [ 109 | executor.submit(lambda obj: obj.get_es_index_body(), obj) 110 | for obj in page.object_list 111 | ] 112 | requests = [] 113 | for task in futures.as_completed(tasks): 114 | requests.append(task.result()) 115 | if self.show_tqdm: 116 | pbar.update(1) 117 | try: 118 | body = u"\n".join(requests) 119 | if not self.dry_run: 120 | self.rubber_config.es.bulk(body=body) 121 | except Exception as exc: 122 | self.print_error(exc) 123 | if not self.show_tqdm: 124 | total += len(page.object_list) 125 | self.print_info(total) 126 | executor.shutdown() 127 | if self.show_tqdm: 128 | pbar.close() 129 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |
2 |
3 |