diff --git a/README.md b/README.md index 0756ada..3834210 100644 --- a/README.md +++ b/README.md @@ -316,7 +316,7 @@ enriched_activities = enricher.enrich_activities(activities) #### Change how models are retrieved -The enrich class that comes with the packages tries to minimise the amount of database queries. The models are grouped by their class and then retrieved with a pk__in query. You can implement a different approach to retrieve the instances of a model subclassing the ```stream_django.enrich.Enrich``` class. +The enrich class that comes with the packages tries to minimise the amount of database queries. The models are grouped by their class and then retrieved with a bulk query. You can implement a different approach to retrieve the instances of a model subclassing the ```stream_django.enrich.Enrich``` class. To change the retrieval for every model you should override the ```fetch_model_instances``` method; in alternative you can change how certain models' are retrieved by implementing the hook function ```fetch__instances``` @@ -326,10 +326,10 @@ class MyEnrich(Enrich): Overwrites how model instances are fetched from the database ''' - def fetch_model_instances(self, modelClass, pks): + def fetch_model_instances(self, modelClass, lookup_values): ''' - returns a dict {id:modelInstance} with instances of model modelClass - and pk in pks + returns a dict {lookup_value:modelInstance} with instances of model modelClass + and lookup_value in lookup_values ''' ... @@ -338,7 +338,7 @@ class AnotherEnrich(Enrich): Overwrites how Likes instances are fetched from the database ''' - def fetch_like_instances(self, pks): + def fetch_like_instances(self, lookup_values): return {l.id: l for l in Like.objects.cached_likes(ids)} ``` diff --git a/stream_django/activity.py b/stream_django/activity.py index 9499ed8..3b5da15 100644 --- a/stream_django/activity.py +++ b/stream_django/activity.py @@ -20,17 +20,30 @@ def create_model_reference(model_instance): creates a reference to a model instance that can be stored in activities >>> from core.models import Like - >>> like = Like.object.get(id=1) + >>> like = Like.object.get(=1) where lookup_field defaults to 'pk' >>> create_reference(like) core.Like:1 ''' + try: + field_name = model_instance.activity_lookup_field() + content_id = getattr(model_instance, field_name) + except AttributeError: + content_id = model_instance.pk + content_type = model_content_type(model_instance.__class__) - content_id = model_instance.pk return '%s:%s' % (content_type, content_id) class Activity(object): + + @classmethod + def activity_lookup_field(cls): + + ''' + The name of the field to use for model lookups + ''' + return 'pk' @property def activity_author_feed(self): @@ -78,7 +91,12 @@ def activity_object_attr(self): @property def activity_actor_id(self): - return self.activity_actor_attr.pk + try: + field_name = self.activity_actor_attr.activity_lookup_field() + actor_id = getattr(self.activity_actor_attr, field_name) + except AttributeError: + actor_id = self.activity_actor_attr.pk + return actor_id @property def activity_actor(self): diff --git a/stream_django/enrich.py b/stream_django/enrich.py index 1af2045..1c8c8e0 100644 --- a/stream_django/enrich.py +++ b/stream_django/enrich.py @@ -3,8 +3,10 @@ import operator import itertools + try: from django.apps import apps + from django.core.exceptions import FieldDoesNotExist get_model = apps.get_model except ImportError: from django.db.models.loading import get_model @@ -91,18 +93,23 @@ def _collect_references(self, activities, fields): model_references[f_ct].append(f_id) return model_references - def fetch_model_instances(self, modelClass, pks): + def fetch_model_instances(self, modelClass, lookup_values): ''' - returns a dict {id:modelInstance} with instances of model modelClass - and pk in pks + returns a dict {:modelInstance} with instances of model modelClass + and lookup_value in lookup_values (i.e., {1:} or {:}) ''' + try: + lookup_field_name = modelClass.activity_lookup_field() + except AttributeError: + lookup_field_name = 'pk' + hook_function_name = 'fetch_%s_instances' % (modelClass._meta.object_name.lower(), ) if hasattr(self, hook_function_name): - return getattr(self, hook_function_name)(pks) + return getattr(self, hook_function_name)(lookup_values) qs = modelClass.objects if hasattr(modelClass, 'activity_related_models') and modelClass.activity_related_models() is not None: qs = qs.select_related(*modelClass.activity_related_models()) - return qs.in_bulk(pks) + return qs.in_bulk(lookup_values, field_name=lookup_field_name) def _fetch_objects(self, references): objects = defaultdict(list) @@ -119,8 +126,17 @@ def _inject_objects(self, activities, objects, fields): continue f_ct, f_id = activity[field].split(':') model = get_model(*f_ct.split('.')) - f_id = model._meta.pk.to_python(f_id) - + + try: + lookup_field_name = model.activity_lookup_field() + if lookup_field_name == 'pk': + f_id = model._meta.pk.to_python(f_id) + else: + lookup_field = model._meta.get_field(lookup_field_name) + f_id = lookup_field.to_python(f_id) + except (AttributeError, FieldDoesNotExist): + f_id = model._meta.pk.to_python(f_id) + instance = objects[f_ct].get(f_id) if instance is None: activity.track_not_enriched_field(field, activity[field])