Export the logic to limit database query time

MANIFEST

@@ -4,6 +4,7 @@ src/auditlog/__init__.py
 src/auditlog/admin.py
 src/auditlog/apps.py
 src/auditlog/context.py
+src/auditlog/count.py
 src/auditlog/diff.py
 src/auditlog/filters.py
 src/auditlog/middleware.py

src/auditlog/admin.py

@@ -1,30 +1,20 @@
 from django.conf import settings
 from django.contrib import admin
 from django.core.paginator import Paginator
-from django.db import connection, transaction, OperationalError
 from django.utils.functional import cached_property
 
+from .count import limit_query_time
 from .models import LogEntry
 from .mixins import LogEntryAdminMixin
 from .filters import ResourceTypeFilter, FieldFilter, get_timestamp_filter
 
 
 class TimeLimitedPaginator(Paginator):
-    """A PostgreSQL-specific paginator with a hard time limit for total count of pages.
-
-    Courtesy of https://medium.com/@hakibenita/optimizing-django-admin-paginator-53c4eb6bfca3
-    """
-    DEFAULT_PAGE_COUNT = 10000
-
+    """A PostgreSQL-specific paginator with a hard time limit for total count of pages."""
     @cached_property
+    @limit_query_time(getattr(settings, 'AUDITLOG_PAGINATOR_TIMEOUT', 500), default=100000)
     def count(self):
-        timeout = getattr(settings, 'AUDITLOG_PAGINATOR_TIMEOUT', 500)  # ms
-        with transaction.atomic(), connection.cursor() as cursor:
-            cursor.execute('SET LOCAL statement_timeout TO %s;', (timeout,))
-            try:
-                return super().count
-            except OperationalError:
-                return self.per_page * self.DEFAULT_PAGE_COUNT
+        return super().count
 
 
 class LogEntryAdmin(admin.ModelAdmin, LogEntryAdminMixin):

src/auditlog/count.py

@@ -0,0 +1,23 @@
+from django.db import connection, transaction, OperationalError
+
+
+def limit_query_time(timeout, default=None):
+    """A PostgreSQL-specific decorator with a hard time limit and a default return value.
+
+    Timeout in milliseconds.
+
+    Courtesy of https://medium.com/@hakibenita/optimizing-django-admin-paginator-53c4eb6bfca3
+    """
+
+    def decorator(function):
+        def _limit_query_time(*args, **kwargs):
+            with transaction.atomic(), connection.cursor() as cursor:
+                cursor.execute('SET LOCAL statement_timeout TO %s;', (timeout,))
+                try:
+                    return function(*args, **kwargs)
+                except OperationalError:
+                    return default
+
+        return _limit_query_time
+
+    return decorator