[SHARE- ][Fix] Use a slightly better count estimate function

Author: chrisseto

requirements.txt

@@ -14,8 +14,6 @@ django-oauth-toolkit==0.10.0
 django-revproxy==0.9.12
 django-mptt==0.8.5
 # Django 1.9 compatibility
-# django-postgres-fuzzycount==0.1.6
-git+https://github.com/stephenmcd/django-postgres-fuzzycount.git@686053fe130c742b9bb0cce338d2aa7eb24fb91c
 # Fix the Choices for TypedModel.type when using Django 1.9
 # django-typed-models==0.5.0
 git+https://github.com/craigds/django-typed-models.git@7c1eaec06e4a5db188b58cee5d86fae6897208b1

share/migrations/0018_fuzzycount.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.9.7 on 2017-02-03 16:22
+from __future__ import unicode_literals
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('share', '0017_merge'),
+    ]
+
+    operations = [
+        migrations.RunSQL(
+            sql='''
+                CREATE OR REPLACE FUNCTION count_estimate(query text) RETURNS INTEGER AS
+                $func$
+                DECLARE
+                    rec   record;
+                    ROWS  INTEGER;
+                BEGIN
+                    FOR rec IN EXECUTE 'EXPLAIN ' || query LOOP
+                        ROWS := SUBSTRING(rec."QUERY PLAN" FROM ' rows=([[:digit:]]+)');
+                        EXIT WHEN ROWS IS NOT NULL;
+                    END LOOP;
+
+                    RETURN ROWS - 1;
+                END
+                $func$ LANGUAGE plpgsql;
+            ''',
+            reverse_sql='DROP FUNCTION count_estimate();',
+        ),
+    ]

share/models/base.py

@@ -11,18 +11,19 @@
 from django.db.models.fields import AutoField
 from django.utils.translation import ugettext_lazy as _
 
-from fuzzycount import FuzzyCountManager
+from typedmodels import models as typedmodels
 
 from db.deletion import DATABASE_CASCADE
 
 from share.models import fields
 from share.models.change import Change
+from share.models.fuzzycount import FuzzyCountManager
 from share.models.sql import ShareObjectManager
-from typedmodels import models as typedmodels
 
 
 class ShareObjectVersion(models.Model):
     action = models.TextField(max_length=10)
+    objects = FuzzyCountManager()
 
     class Meta:
         abstract = True

share/models/celery.py

@@ -1,11 +1,12 @@
-from fuzzycount import FuzzyCountManager
 from model_utils import Choices
 
 from django.conf import settings
 from django.db import models
 from django.utils.translation import ugettext as _
 from typedmodels.models import TypedModel
 
+from share.models.fuzzycount import FuzzyCountManager
+
 
 class CeleryTask(TypedModel):
     STATUS = Choices(

share/models/change.py

@@ -2,8 +2,6 @@
 
 from model_utils import Choices
 
-from fuzzycount import FuzzyCountManager
-
 from django.apps import apps
 from django.contrib.contenttypes.fields import GenericForeignKey
 from django.contrib.contenttypes.models import ContentType
@@ -14,6 +12,7 @@
 from django.utils import timezone
 from django.utils.translation import ugettext as _
 
+from share.models.fuzzycount import FuzzyCountManager
 from share.models import NormalizedData
 from share.util import IDObfuscator
 
@@ -114,7 +113,7 @@ def _resolve_ref(self, ref):
             raise Exception('Could not resolve reference {}'.format(ref)) from ex
 
     def __repr__(self):
-        return '<{}({}, {}, {} changes)>'.format(self.__class__.__name__, self.STATUS[self.status].upper(), self.normalized_data.source, self.changes.count())
+        return '<{}({}, {}, {} changes)>'.format(self.__class__.__name__, self.STATUS[self.status].upper(), self.normalized_data.source, self.changes.exact_count())
 
 
 class Change(models.Model):

share/models/core.py

@@ -13,11 +13,12 @@
 from django.dispatch import receiver
 from django.utils import timezone
 from django.utils.translation import ugettext_lazy as _
-from fuzzycount import FuzzyCountManager
 from oauth2_provider.models import AccessToken, Application
 
 from osf_oauth2_adapter.apps import OsfOauth2AdapterConfig
+
 from share.models.fields import DateTimeAwareJSONField, ShareURLField
+from share.models.fuzzycount import FuzzyCountManager
 from share.models.validators import JSONLDValidator
 
 logger = logging.getLogger(__name__)

share/models/fuzzycount.py

@@ -0,0 +1,14 @@
+from django.db import connections
+from django.db.models import QuerySet, Manager
+
+
+class FuzzyCountQuerySet(QuerySet):
+    def count(self):
+        cursor = connections[self.db].cursor()
+        cursor.execute('SELECT count_estimate(%s);', (cursor.mogrify(*self.query.sql_with_params()).decode(), ))
+        return int(cursor.fetchone()[0])
+
+    def exact_count(self):
+        return super().count()
+
+FuzzyCountManager = Manager.from_queryset(FuzzyCountQuerySet)

share/models/meta.py

@@ -6,6 +6,7 @@
 
 from share.models.base import ShareObject
 from share.models.fields import ShareForeignKey
+from share.models.fuzzycount import FuzzyCountManager
 from share.util import strip_whitespace
 
 
@@ -49,7 +50,7 @@ class Disambiguation:
         all = ('name',)
 
 
-class SubjectManager(models.Manager):
+class SubjectManager(FuzzyCountManager):
     def get_by_natural_key(self, subject):
         return self.get(name=subject)
 

share/models/sql.py

@@ -1,9 +1,9 @@
 from django.db import connections
-from django.db.models.manager import Manager
 from django.db.models.sql import InsertQuery
 from django.db.models.sql.compiler import SQLInsertCompiler
 
-from fuzzycount import FuzzyCountQuerySet
+from share.models.fuzzycount import FuzzyCountManager
+from share.models.fuzzycount import FuzzyCountQuerySet
 
 
 class SQLInsertReturnVersionCompiler(SQLInsertCompiler):
@@ -84,7 +84,7 @@ def _insert(self, objs, fields, return_id=False, raw=False, using=None):
     _insert.queryset_only = False
 
 
-class ShareObjectManager(Manager):
+class ShareObjectManager(FuzzyCountManager):
     use_for_related_fields = True
 
     def get_queryset(self):