Add part 2

Author: ParentJA

README.md

@@ -1 +1,15 @@
 # django-postgres-elasticsearch
+
+If you're starting with **Part 2** or later, don't forget to run the following commands first:
+
+```bash
+$ docker-compose exec server python manage.py createsuperuser
+```
+
+```bash
+$ docker-compose exec server \
+    python manage.py loaddata \
+        catalog/fixtures/wines.json \
+        --app catalog \
+        --format json
+```

part2/server/catalog/admin.py

@@ -1,12 +1,22 @@
 from django.contrib import admin
 
-from .models import Wine
+from .models import Wine, WineSearchWord
 
 
 @admin.register(Wine)
 class WineAdmin(admin.ModelAdmin):
-    fields = ('id', 'country', 'description', 'points', 'price', 'variety', 'winery',)
+    fields = (
+        'id', 'country', 'description', 'points', 'price', 'variety', 'winery',
+        'search_vector',
+    )
     list_display = ('id', 'country', 'points', 'price', 'variety', 'winery',)
     list_filter = ('country', 'variety', 'winery',)
     ordering = ('variety',)
     readonly_fields = ('id',)
+
+
+@admin.register(WineSearchWord)
+class WineSearchWordAdmin(admin.ModelAdmin):
+    fields = ('word',)
+    list_display = ('word',)
+    ordering = ('word',)

part2/server/catalog/apps.py

@@ -3,3 +3,6 @@
 
 class CatalogConfig(AppConfig):
     name = 'catalog'
+
+    def ready(self):
+        from . import signals

part2/server/catalog/filters.py

@@ -1,21 +1,25 @@
-from django.db.models import Q
-
 from django_filters.rest_framework import CharFilter, FilterSet
 
-from .models import Wine
+from .models import Wine, WineSearchWord
 
 
 class WineFilterSet(FilterSet):
     query = CharFilter(method='filter_query')
 
     def filter_query(self, queryset, name, value):
-        search_query = Q(
-            Q(variety__contains=value) |
-            Q(winery__contains=value) |
-            Q(description__contains=value)
-        )
-        return queryset.filter(search_query)
+        return queryset.search(value)
 
     class Meta:
         model = Wine
         fields = ('query', 'country', 'points',)
+
+
+class WineSearchWordFilterSet(FilterSet):
+    query = CharFilter(method='filter_query')
+
+    def filter_query(self, queryset, name, value):
+        return queryset.search(value)
+
+    class Meta:
+        model = WineSearchWord
+        fields = ('query',)

part2/server/catalog/fixtures/test_wines.json

@@ -23,6 +23,18 @@
       "winery": "Barnard Griffin"
     }
   },
+  {
+    "model": "catalog.Wine",
+    "fields": {
+      "country": "France",
+      "description": "A creamy wine with full Chardonnay flavors.",
+      "id": "000bbdff-30fc-4897-81c1-7947e11e6d1a",
+      "points": 89,
+      "price": 33.00,
+      "variety": "Champagne Blend",
+      "winery": "Phillipe Gonet"
+    }
+  },
   {
     "model": "catalog.Wine",
     "fields": {

part2/server/catalog/migrations/0002_search_vector.py

@@ -0,0 +1,24 @@
+# Generated by Django 3.1.7 on 2021-03-23 01:48
+
+import django.contrib.postgres.indexes
+import django.contrib.postgres.search
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('catalog', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='wine',
+            name='search_vector',
+            field=django.contrib.postgres.search.SearchVectorField(blank=True, null=True),
+        ),
+        migrations.AddIndex(
+            model_name='wine',
+            index=django.contrib.postgres.indexes.GinIndex(fields=['search_vector'], name='search_vector_index'),
+        ),
+    ]

part2/server/catalog/migrations/0003_update_search_vector.py

@@ -0,0 +1,22 @@
+from django.contrib.postgres.search import SearchVector
+from django.db import migrations
+
+
+def update_search_vector(apps, schema_editor):
+    Wine = apps.get_model('catalog', 'Wine')
+    Wine.objects.all().update(search_vector=(
+        SearchVector('variety', weight='A') +
+        SearchVector('winery', weight='A') +
+        SearchVector('description', weight='B')
+    ))
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('catalog', '0002_search_vector'),
+    ]
+
+    operations = [
+        migrations.RunPython(update_search_vector, elidable=True),
+    ]

part2/server/catalog/migrations/0004_winesearchword.py

@@ -0,0 +1,20 @@
+# Generated by Django 3.1.7 on 2021-03-23 02:03
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('catalog', '0003_update_search_vector'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='WineSearchWord',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('word', models.CharField(max_length=255, unique=True)),
+            ],
+        ),
+    ]

part2/server/catalog/migrations/0005_update_wine_search_word.py

@@ -0,0 +1,32 @@
+from django.contrib.postgres.operations import TrigramExtension
+from django.db import connection, migrations
+
+
+def update_wine_search_word(apps, schema_editor):
+    sql = """
+        INSERT INTO catalog_winesearchword (word)
+        SELECT word FROM ts_stat('
+          SELECT to_tsvector(''simple'', winery) ||
+                 to_tsvector(''simple'', coalesce(description, ''''))
+            FROM catalog_wine
+        ');
+    """
+    with connection.cursor() as cursor:
+        cursor.execute(sql)
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('catalog', '0004_winesearchword'),
+    ]
+
+    operations = [
+        TrigramExtension(),
+        migrations.RunSQL(sql="""
+            CREATE INDEX IF NOT EXISTS wine_search_word_trigram_index
+                ON catalog_winesearchword
+             USING gin (word gin_trgm_ops);
+        """, elidable=True),
+        migrations.RunPython(update_wine_search_word, elidable=True),
+    ]

part2/server/catalog/models.py

@@ -1,20 +1,64 @@
 import uuid
 
+from django.contrib.postgres.indexes import GinIndex
+from django.contrib.postgres.search import (
+    SearchQuery, SearchRank, SearchVectorField, TrigramSimilarity
+)
 from django.db import models
+from django.db.models import F, Q
+
+
+class WineQuerySet(models.query.QuerySet):
+    def search(self, query):
+        search_query = Q(
+            Q(search_vector=SearchQuery(query))
+        )
+        return self.annotate(
+            variety_headline=SearchHeadline(F('variety'), SearchQuery(query)),
+            winery_headline=SearchHeadline(F('winery'), SearchQuery(query)),
+            description_headline=SearchHeadline(F('description'), SearchQuery(query)),
+            search_rank=SearchRank(F('search_vector'), SearchQuery(query))
+        ).filter(search_query).order_by('-search_rank', 'id')
 
 
 class Wine(models.Model):
-    id = models.UUIDField(
-        primary_key=True, default=uuid.uuid4, editable=False
-    )
+    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
     country = models.CharField(max_length=255)
     description = models.TextField(null=True, blank=True)
     points = models.IntegerField()
-    price = models.DecimalField(
-        decimal_places=2, max_digits=10, null=True, blank=True
-    )
+    price = models.DecimalField(decimal_places=2, max_digits=10, null=True, blank=True)
     variety = models.CharField(max_length=255)
     winery = models.CharField(max_length=255)
+    search_vector = SearchVectorField(null=True, blank=True)
+
+    objects = WineQuerySet.as_manager()
+
+    class Meta:
+        indexes = [
+            GinIndex(fields=['search_vector'], name='search_vector_index')
+        ]
 
     def __str__(self):
         return f'{self.id}'
+
+
+class SearchHeadline(models.Func):
+    function = 'ts_headline'
+    output_field = models.TextField()
+    template = '%(function)s(%(expressions)s, \'StartSel = <mark>, StopSel = </mark>, HighlightAll=TRUE\')'
+
+
+class WineSearchWordQuerySet(models.query.QuerySet):
+    def search(self, query):
+        return self.annotate(
+            similarity=TrigramSimilarity('word', query)
+        ).filter(similarity__gte=0.3).order_by('-similarity')
+
+
+class WineSearchWord(models.Model):
+    word = models.CharField(max_length=255, unique=True)
+
+    objects = WineSearchWordQuerySet.as_manager()
+
+    def __str__(self):
+        return self.word

part2/server/catalog/serializers.py

@@ -1,9 +1,34 @@
 from rest_framework import serializers
 
-from .models import Wine
+from .models import Wine, WineSearchWord
 
 
 class WineSerializer(serializers.ModelSerializer):
+    variety = serializers.SerializerMethodField()
+    winery = serializers.SerializerMethodField()
+    description = serializers.SerializerMethodField()
+
+    def get_variety(self, obj):
+        if hasattr(obj, 'variety_headline'):
+            return getattr(obj, 'variety_headline')
+        return getattr(obj, 'variety')
+
+    def get_winery(self, obj):
+        if hasattr(obj, 'winery_headline'):
+            return getattr(obj, 'winery_headline')
+        return getattr(obj, 'winery')
+
+    def get_description(self, obj):
+        if hasattr(obj, 'description_headline'):
+            return getattr(obj, 'description_headline')
+        return getattr(obj, 'description')
+
     class Meta:
         model = Wine
         fields = ('id', 'country', 'description', 'points', 'price', 'variety', 'winery',)
+
+
+class WineSearchWordSerializer(serializers.ModelSerializer):
+    class Meta:
+        model = WineSearchWord
+        fields = ('word',)

part2/server/catalog/signals.py

@@ -0,0 +1,27 @@
+from django.contrib.postgres.search import SearchVector
+from django.db import connection
+from django.db.models.signals import post_save
+from django.dispatch import receiver
+
+from .models import Wine
+
+
+@receiver(post_save, sender=Wine, dispatch_uid='on_wine_save')
+def on_wine_save(sender, instance, *args, **kwargs):
+    sender.objects.filter(pk=instance.id).update(search_vector=(
+        SearchVector('variety', weight='A') +
+        SearchVector('winery', weight='A') +
+        SearchVector('description', weight='B')
+    ))
+
+    with connection.cursor() as cursor:
+        cursor.execute("""
+            INSERT INTO catalog_winesearchword (word)
+            SELECT word FROM ts_stat('
+              SELECT to_tsvector(''simple'', winery) ||
+                     to_tsvector(''simple'', coalesce(description, ''''))
+                FROM catalog_wine
+               WHERE id = '%s'
+            ')
+            ON CONFLICT (word) DO NOTHING;
+        """, [str(instance.id),])