Improve error handling in the pgautopartition command

SectorLabs · Nov 3, 2019 · 8c5a781 · 8c5a781
1 parent 37e5135
commit 8c5a781
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/psqlextra/management/commands/pgautopartition.py b/psqlextra/management/commands/pgautopartition.py
@@ -1,15 +1,11 @@
 from django.apps import apps
 from django.core.management.base import BaseCommand
+from django.db.utils import NotSupportedError, OperationalError
 
 from psqlextra.auto_partition import (
     PostgresAutoPartitioningIntervalUnit,
     postgres_auto_partition,
 )
-from psqlextra.models import PostgresPartitionedModel
-
-
-class PostgresAutoPartitioningError(RuntimeError):
-    pass
 
 
 class Command(BaseCommand):
@@ -66,8 +62,11 @@ def handle(self, *app_labels, **options):
         interval = options.get("interval")
 
         model = apps.get_model(app_label, model_name)
-        if not issubclass(model, PostgresPartitionedModel):
-            raise PostgresAutoPartitioningError(
+        if not model:
+            raise OperationalError(f"Cannot find a model named '{model_name}'")
+
+        if not issubclass(model, PostgresMaterializedViewModel):
+            raise NotSupportedError(
                 f"Model {model.__name__} is not a `PostgresPartitionedModel`"
             )
 

diff --git a/psqlextra/management/commands/pgrefreshmv.py b/psqlextra/management/commands/pgrefreshmv.py
@@ -1,4 +1,3 @@
-
 from django.apps import apps
 from django.core.management.base import BaseCommand
 from django.db.utils import NotSupportedError, OperationalError