Tweaked enforcing of allow_wildcards in tenant commands

lorinkoz · Jul 1, 2023 · b1aa22f · b1aa22f
1 parent 573a1a5
commit b1aa22f
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 40 deletions.
diff --git a/django_pgschemas/management/commands/__init__.py b/django_pgschemas/management/commands/__init__.py
@@ -51,41 +51,44 @@ def add_arguments(self, parser):
         parser.add_argument(
             "-s", "--schema", nargs="+", dest="schemas", help="Schema(s) to execute the current command"
         )
-        parser.add_argument(
-            "-as",
-            "--include-all-schemas",
-            action="store_true",
-            dest="all_schemas",
-            help="Include all schemas when executing the current command",
-        )
-        parser.add_argument(
-            "-ss",
-            "--include-static-schemas",
-            action="store_true",
-            dest="static_schemas",
-            help="Include all static schemas when executing the current command",
-        )
-        parser.add_argument(
-            "-ds",
-            "--include-dynamic-schemas",
-            action="store_true",
-            dest="dynamic_schemas",
-            help="Include all dynamic schemas when executing the current command",
-        )
-        parser.add_argument(
-            "-ts",
-            "--include-tenant-schemas",
-            action="store_true",
-            dest="tenant_schemas",
-            help="Include all tenant-like schemas when executing the current command",
-        )
         parser.add_argument(
             "-x",
             "--exclude-schema",
             nargs="+",
             dest="excluded_schemas",
             help="Schema(s) to exclude when executing the current command",
         )
+
+        if self.allow_wildcards:
+            parser.add_argument(
+                "-as",
+                "--include-all-schemas",
+                action="store_true",
+                dest="all_schemas",
+                help="Include all schemas when executing the current command",
+            )
+            parser.add_argument(
+                "-ss",
+                "--include-static-schemas",
+                action="store_true",
+                dest="static_schemas",
+                help="Include all static schemas when executing the current command",
+            )
+            parser.add_argument(
+                "-ds",
+                "--include-dynamic-schemas",
+                action="store_true",
+                dest="dynamic_schemas",
+                help="Include all dynamic schemas when executing the current command",
+            )
+            parser.add_argument(
+                "-ts",
+                "--include-tenant-schemas",
+                action="store_true",
+                dest="tenant_schemas",
+                help="Include all tenant-like schemas when executing the current command",
+            )
+
         parser.add_argument(
             "--parallel",
             dest="parallel",
@@ -169,19 +172,19 @@ def _get_schemas_from_options(self, **options):
         schemas_to_return = set()
 
         if include_all_schemas:
-            if not self.allow_wildcards or (not allow_static and not allow_dynamic):
+            if not allow_static and not allow_dynamic:
                 raise CommandError("Including all schemas is NOT allowed")
             schemas_to_return = schemas_to_return.union(static_schemas + list(dynamic_schemas))
         if include_static_schemas:
-            if not self.allow_wildcards or not allow_static:
+            if not allow_static:
                 raise CommandError("Including static schemas is NOT allowed")
             schemas_to_return = schemas_to_return.union(static_schemas)
         if include_dynamic_schemas:
-            if not self.allow_wildcards or not allow_dynamic:
+            if not allow_dynamic:
                 raise CommandError("Including dynamic schemas is NOT allowed")
             schemas_to_return = schemas_to_return.union(dynamic_schemas)
         if include_tenant_schemas:
-            if not self.allow_wildcards or not allow_dynamic:
+            if not allow_dynamic:
                 raise CommandError("Including tenant-like schemas is NOT allowed")
             schemas_to_return = schemas_to_return.union(dynamic_schemas)
             if clone_reference:

diff --git a/dpgs_sandbox/tests/test_tenant_commands.py b/dpgs_sandbox/tests/test_tenant_commands.py
@@ -42,9 +42,8 @@ def test_no_schema_provided(self):
     def test_no_all_schemas_allowed(self):
         command = WhoWillCommand()
         command.allow_wildcards = False
-        with self.assertRaises(CommandError) as ctx:
+        with self.assertRaises(TypeError):
             management.call_command(command, all_schemas=True, verbosity=0)
-        self.assertEqual(str(ctx.exception), "Including all schemas is NOT allowed")
 
     def test_no_static_schemas_allowed(self):
         command = WhoWillCommand()
@@ -54,9 +53,8 @@ def test_no_static_schemas_allowed(self):
         self.assertEqual(str(ctx.exception), "Including static schemas is NOT allowed")
         command = WhoWillCommand()
         command.allow_wildcards = False
-        with self.assertRaises(CommandError) as ctx:
+        with self.assertRaises(TypeError):
             management.call_command(command, static_schemas=True, verbosity=0)
-        self.assertEqual(str(ctx.exception), "Including static schemas is NOT allowed")
 
     def test_no_dynamic_schemas_allowed(self):
         command = WhoWillCommand()
@@ -66,9 +64,8 @@ def test_no_dynamic_schemas_allowed(self):
         self.assertEqual(str(ctx.exception), "Including dynamic schemas is NOT allowed")
         command = WhoWillCommand()
         command.allow_wildcards = False
-        with self.assertRaises(CommandError) as ctx:
+        with self.assertRaises(TypeError):
             management.call_command(command, dynamic_schemas=True, verbosity=0)
-        self.assertEqual(str(ctx.exception), "Including dynamic schemas is NOT allowed")
 
     def test_no_tenant_like_schemas_allowed(self):
         command = WhoWillCommand()
@@ -78,9 +75,8 @@ def test_no_tenant_like_schemas_allowed(self):
         self.assertEqual(str(ctx.exception), "Including tenant-like schemas is NOT allowed")
         command = WhoWillCommand()
         command.allow_wildcards = False
-        with self.assertRaises(CommandError) as ctx:
+        with self.assertRaises(TypeError):
             management.call_command(command, tenant_schemas=True, verbosity=0)
-        self.assertEqual(str(ctx.exception), "Including tenant-like schemas is NOT allowed")
 
     def test_nonexisting_schema(self):
         with self.assertRaises(CommandError) as ctx: