Clear DB page doesn't work with PostgreSQL and SQLite

silk/management/commands/silk_clear_request_log.py

@@ -1,43 +1,16 @@
-from django.conf import settings
 from django.core.management.base import BaseCommand
-from django.db import connection
 
 import silk.models
+from silk.utils.data_deletion import delete_model
 
 
 class Command(BaseCommand):
     help = "Clears silk's log of requests."
 
-    @staticmethod
-    def delete_model(model):
-        engine = settings.DATABASES['default']['ENGINE']
-        table = model._meta.db_table
-        if 'mysql' in engine or 'postgresql' in engine:
-            # Use "TRUNCATE" on the table
-            with connection.cursor() as cursor:
-                if 'mysql' in engine:
-                    cursor.execute("SET FOREIGN_KEY_CHECKS=0;")
-                    cursor.execute("TRUNCATE TABLE {0}".format(table))
-                    cursor.execute("SET FOREIGN_KEY_CHECKS=1;")
-                elif 'postgres' in engine:
-                    cursor.execute("ALTER TABLE {0} DISABLE TRIGGER USER;".format(table))
-                    cursor.execute("TRUNCATE TABLE {0} CASCADE".format(table))
-                    cursor.execute("ALTER TABLE {0} ENABLE TRIGGER USER;".format(table))
-            return
-
-        # Manually delete rows because sqlite does not support TRUNCATE and
-        # oracle doesn't provide good support for disabling foreign key checks
-        while True:
-            items_to_delete = list(
-                model.objects.values_list('pk', flat=True).all()[:1000])
-            if not items_to_delete:
-                break
-            model.objects.filter(pk__in=items_to_delete).delete()
-
     def handle(self, *args, **options):
         # Django takes a long time to traverse foreign key relations,
         # so delete in the order that makes it easy.
-        Command.delete_model(silk.models.Profile)
-        Command.delete_model(silk.models.SQLQuery)
-        Command.delete_model(silk.models.Response)
-        Command.delete_model(silk.models.Request)
+        delete_model(silk.models.Profile)
+        delete_model(silk.models.SQLQuery)
+        delete_model(silk.models.Response)
+        delete_model(silk.models.Request)

silk/utils/data_deletion.py

@@ -0,0 +1,28 @@
+from django.conf import settings
+from django.db import connection
+
+
+def delete_model(model):
+    engine = settings.DATABASES['default']['ENGINE']
+    table = model._meta.db_table
+    if 'mysql' in engine or 'postgresql' in engine:
+        # Use "TRUNCATE" on the table
+        with connection.cursor() as cursor:
+            if 'mysql' in engine:
+                cursor.execute("SET FOREIGN_KEY_CHECKS=0;")
+                cursor.execute("TRUNCATE TABLE {0}".format(table))
+                cursor.execute("SET FOREIGN_KEY_CHECKS=1;")
+            elif 'postgres' in engine:
+                cursor.execute("ALTER TABLE {0} DISABLE TRIGGER USER;".format(table))
+                cursor.execute("TRUNCATE TABLE {0} CASCADE".format(table))
+                cursor.execute("ALTER TABLE {0} ENABLE TRIGGER USER;".format(table))
+        return
+
+    # Manually delete rows because sqlite does not support TRUNCATE and
+    # oracle doesn't provide good support for disabling foreign key checks
+    while True:
+        items_to_delete = list(
+            model.objects.values_list('pk', flat=True).all()[:1000])
+        if not items_to_delete:
+            break
+        model.objects.filter(pk__in=items_to_delete).delete()

silk/views/clear_db.py

@@ -1,22 +1,14 @@
-from django.db import connection
 from django.shortcuts import render
 from django.utils.decorators import method_decorator
 from django.views.generic import View
 
 from silk.auth import login_possibly_required, permissions_possibly_required
 from silk.models import Request, Response, SQLQuery, Profile
+from silk.utils.data_deletion import delete_model
 
 
 class ClearDBView(View):
 
-    def _truncate_tables(self, models):
-        raw_query = 'TRUNCATE TABLE {0};'
-        truncate_query = [raw_query.format(m._meta.db_table) for m in models]
-        truncate_query = ' '.join(truncate_query)
-        query = 'SET FOREIGN_KEY_CHECKS = 0; {0} SET FOREIGN_KEY_CHECKS = 1;'.format(truncate_query)
-        cursor = connection.cursor()
-        cursor.execute(query)
-
     @method_decorator(login_possibly_required)
     @method_decorator(permissions_possibly_required)
     def get(self, request, *_, **kwargs):
@@ -27,7 +19,10 @@ def get(self, request, *_, **kwargs):
     def post(self, request, *_, **kwargs):
         context = {}
         if 'clear_all' in request.POST:
-            self._truncate_tables([Response, SQLQuery, Profile, Request])
+            delete_model(Profile)
+            delete_model(SQLQuery)
+            delete_model(Response)
+            delete_model(Request)
             tables = ['Response', 'SQLQuery', 'Profile', 'Request']
             context['msg'] = 'Cleared data for following silk tables: {0}'.format(', '.join(tables))
         return render(request, 'silk/clear_db.html', context=context)