Add support for signals

Author: Photonios

docs/features.md

@@ -101,3 +101,81 @@ Note that a single call to `upsert` results in a single `INSERT INTO ... ON CONF
     )
 
     assert obj1.id == obj2.id
+
+# Signals
+Django has long supported "signals", a feature that can be really useful. It allows you to get notified in a callback when a certain event occurs. One of the most common use cases is connecting to model signals. These signals get triggered when a model gets saved, or deleted.
+
+Django's built-in signals have one major flaw:
+
+* The `QuerySet.update(..)` method does **not** emit **any** signals.
+
+Because of this limitation, Django's signals cannot be reliably used to be signaled on model changes. `django-postges-extra` adds three new signals which are much more primitive, but work reliably across the board.
+
+The signals defined by this library are completely valid, standard Django signals. Therefore, their documentation also applies: [https://docs.djangoproject.com/en/1.10/topics/signals/](https://docs.djangoproject.com/en/1.10/topics/signals/).
+
+Each of the signals send upon model modification send one parameter containing the value of the primary key of the row that was affected. Therefore the signal's signature looks like this:
+
+    def my_receiver(sender, pk: int):
+        # pk is the primary key, a keyword argument
+
+* `psqlextra.signals.create`
+    * Send **after** a new model instance was created.
+        
+            from django.db import models
+            from psqlextra.models import PostgresModel
+            from psqlextra import signals
+    
+            class MyModel(PostgresModel):
+                myfield = models.CharField(max_length=255, unique=True)
+                
+            def on_create(sender, **kwargs):
+                 print('model created with pk %d' % kwargs['pk'])
+                
+            signals.create.connect(MyModel, on_create, weak=False)
+
+            # this will trigger the signal
+            instance = MyModel(myfield='cookies')
+            instance.save()
+            
+            # but so will this
+            MyModel.objects.create(myfield='cheers')
+
+* `psqlextra.signals.update`
+    * Send **after** a new model instance was updated.
+        
+            from django.db import models
+            from psqlextra.models import PostgresModel
+            from psqlextra import signals
+    
+            class MyModel(PostgresModel):
+                myfield = models.CharField(max_length=255, unique=True)
+                
+            def on_update(sender, **kwargs):
+                 print('model updated with pk %d' % kwargs['pk'])
+                
+            signals.update.connect(MyModel, on_update, weak=False)
+
+            # for every row that is affected, the signal will be send
+            MyModel.objects.filter(myfield='cookies').update(myfield='cheers')
+            
+* `psqlextra.signals.delete`
+    * Send **before** a new model instance is deleted.
+        
+            from django.db import models
+            from psqlextra.models import PostgresModel
+            from psqlextra import signals
+    
+            class MyModel(PostgresModel):
+                myfield = models.CharField(max_length=255, unique=True)
+                
+            def on_delete(sender, **kwargs):
+                 print('model deleted with pk %d' % kwargs['pk'])
+                
+            signals.delete.connect(MyModel, on_update, weak=False)
+
+            # for every row that is affected, the signal will be send
+            MyModel.objects.filter(myfield='cookies').delete()
+            
+            # in this case, a single row is deleted, the signal will be send
+            # for this particular row
+            MyModel.objects.get(id=1).delete()

psqlextra/compiler.py

@@ -1,5 +1,26 @@
 from django.core.exceptions import SuspiciousOperation
-from django.db.models.sql.compiler import SQLInsertCompiler
+from django.db.models.sql.compiler import SQLInsertCompiler, SQLUpdateCompiler
+
+
+class PostgresSQLReturningUpdateCompiler(SQLUpdateCompiler):
+    """Compiler for SQL UPDATE statements that return
+    the primary keys of the affected rows."""
+
+    def execute_sql(self, _result_type):
+        sql, params = self.as_sql()
+        sql += self._form_returning()
+
+        with self.connection.cursor() as cursor:
+            cursor.execute(sql, params)
+            primary_keys = cursor.fetchall()
+
+        return primary_keys
+
+    def _form_returning(self):
+        """Builds the RETURNING part of the query."""
+
+        qn = self.connection.ops.quote_name
+        return 'RETURNING %s' % qn(self.query.model._meta.pk.name)
 
 
 class PostgresSQLUpsertCompiler(SQLInsertCompiler):

psqlextra/manager.py

@@ -3,12 +3,46 @@
 import django
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
-from django.db import models
+from django.db import models, transaction
+from django.db.models.sql import UpdateQuery
+from django.db.models.sql.constants import CURSOR
 
-from .compiler import PostgresSQLUpsertCompiler
+from . import signals
+from .compiler import (PostgresSQLReturningUpdateCompiler,
+                       PostgresSQLUpsertCompiler)
 from .query import PostgresUpsertQuery
 
 
+class PostgresQuerySet(models.QuerySet):
+    """Adds support for PostgreSQL specifics."""
+
+    def update(self, **fields):
+        """Updates all rows that match the filter."""
+
+        # build up the query to execute
+        self._for_write = True
+        query = self.query.clone(UpdateQuery)
+        query._annotations = None
+        query.add_update_values(fields)
+
+        # build the compiler for form the query
+        connection = django.db.connections[self.db]
+        compiler = PostgresSQLReturningUpdateCompiler(query, connection, self.db)
+
+        # execute the query
+        with transaction.atomic(using=self.db, savepoint=False):
+            rows = compiler.execute_sql(CURSOR)
+        self._result_cache = None
+
+        # send out a signal for each row
+        for row in rows:
+            signals.update.send(self.model, pk=row[0])
+
+        # the original update(..) returns the amount of rows
+        # affected, let's do the same
+        return len(rows)
+
+
 class PostgresManager(models.Manager):
     """Adds support for PostgreSQL specifics."""
 
@@ -27,6 +61,35 @@ def __init__(self, *args, **kwargs):
                 'the \'psqlextra.backend\'. Set DATABASES.ENGINE.'
             ) % db_backend)
 
+        # hook into django signals to then trigger our own
+
+        def on_model_save(sender, **kwargs):
+            """When a model gets created or updated."""
+
+            created, instance = kwargs['created'], kwargs['instance']
+
+            if created:
+                signals.create.send(sender, pk=instance.pk)
+            else:
+                signals.update.send(sender, pk=instance.pk)
+
+        django.db.models.signals.post_save.connect(
+            on_model_save, sender=self.model, weak=False)
+
+        def on_model_delete(sender, **kwargs):
+            """When a model gets deleted."""
+
+            instance = kwargs['instance']
+            signals.delete.send(sender, pk=instance.pk)
+
+        django.db.models.signals.pre_delete.connect(
+            on_model_delete, sender=self.model, weak=False)
+
+    def get_queryset(self):
+        """Gets the query set to be used on this manager."""
+
+        return PostgresQuerySet(self.model, using=self._db)
+
     def upsert(self, conflict_target: List, fields: Dict) -> int:
         """Creates a new record or updates the existing one
         with the specified data.

psqlextra/signals.py

@@ -0,0 +1,5 @@
+from django.dispatch import Signal
+
+create = Signal(providing_args=['pk'])
+update = Signal(providing_args=['pk'])
+delete = Signal(providing_args=['pk'])

tests/test_signals.py

@@ -0,0 +1,110 @@
+from unittest.mock import Mock
+
+import pytest
+from django.db import models
+from django.test import TestCase
+
+from psqlextra import signals
+
+from .fake_model import get_fake_model
+
+
+@pytest.mark.django_db
+class SignalsTestCase(TestCase):
+    """Tests whether all signals work correctly."""
+
+    def mock_signal_handler(self, signal):
+        """Creates a new model and attaches the specified
+        signal to a mocked signal handler.
+
+        Returns:
+            The created model and the mocked signal handler.
+        """
+
+        model = get_fake_model({
+            'title': models.CharField(max_length=255)
+        })
+
+        signal_handler = Mock()
+        signal.connect(signal_handler, sender=model, weak=False)
+
+        return model, signal_handler
+
+    def test_create(self):
+        """Tests whether the create signal is properly emitted
+        when using QuerySet.create."""
+
+        model, signal_handler = self.mock_signal_handler(signals.create)
+        instance = model.objects.create(title='beer')
+
+        assert signal_handler.call_count == 1
+        assert signal_handler.call_args[1]['pk'] == instance.pk
+
+    def test_model_save_create(self):
+        """Tests whether the create signal is properly
+        emitted when using Model.save()."""
+
+        model, signal_handler = self.mock_signal_handler(signals.create)
+
+        instance = model(title='beer')
+        instance.save()
+
+        assert signal_handler.call_count == 1
+        assert signal_handler.call_args[1]['pk'] == instance.pk
+
+    def test_model_save_update(self):
+        """Tests whether the update signal properly
+        emitted when using Model.save()."""
+
+        model, signal_handler = self.mock_signal_handler(signals.update)
+
+        instance = model(title='beer')
+        instance.save()  # create
+        instance.save()  # update
+
+        assert signal_handler.call_count == 1
+        assert signal_handler.call_args[1]['pk'] == instance.pk
+
+    def test_model_delete(self):
+        """Tests whether the delete signal properly
+        emitted when using Model.delete()."""
+
+        model, signal_handler = self.mock_signal_handler(signals.delete)
+        instance = model.objects.create(title='beer')
+        instance_pk = instance.pk
+        instance.delete()
+
+        assert signal_handler.call_count == 1
+        assert signal_handler.call_args[1]['pk'] == instance_pk
+
+    def test_query_set_delete(self):
+        """Tests whether the delete signal is emitted
+        for each row that is deleted."""
+
+        model, signal_handler = self.mock_signal_handler(signals.delete)
+
+        instance_1 = model.objects.create(title='beer')
+        instance_1_pk = instance_1.pk
+        instance_2 = model.objects.create(title='more boar')
+        instance_2_pk = instance_2.pk
+
+        model.objects.all().delete()
+
+        assert signal_handler.call_count == 2
+        assert signal_handler.call_args_list[0][1]['pk'] == instance_1_pk
+        assert signal_handler.call_args_list[1][1]['pk'] == instance_2_pk
+
+    def test_query_set_update(self):
+        """Tests whether the update signal is emitted
+        for each row that has been updated."""
+
+        model, signal_handler = self.mock_signal_handler(signals.update)
+
+        instance_1 = model.objects.create(title='beer')
+        instance_2 = model.objects.create(title='more boar')
+
+        model.objects.all().update(title='cookies')
+
+        assert signal_handler.call_count == 2
+        assert signal_handler.call_args_list[0][1]['pk'] == instance_1.pk
+        assert signal_handler.call_args_list[1][1]['pk'] == instance_2.pk