Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

More Cleanup #159

Merged
merged 2 commits into from
Jul 1, 2023
Merged

More Cleanup #159

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
573a1a5
Command scope is now an enum
lorinkoz Jul 1, 2023
b1aa22f
Tweaked enforcing of allow_wildcards in tenant commands
lorinkoz Jul 1, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 64 additions & 40 deletions django_pgschemas/management/commands/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,18 +1,39 @@
import enum

from django.conf import settings
from django.core.management.base import BaseCommand, CommandError
from django.db.models import CharField, Q, Value as V
from django.db.models import CharField, Q
from django.db.models import Value as V
from django.db.models.functions import Concat
from django.db.utils import ProgrammingError

from ...schema import get_current_schema
from ...utils import create_schema, dynamic_models_exist, get_clone_reference, get_tenant_model
from ._executors import parallel, sequential

EXECUTORS = {"sequential": sequential, "parallel": parallel}

class CommandScope(enum.Enum):
ALL = "all"
DYNAMIC = "dynamic"
STATIC = "static"

@classmethod
def allow_static(cls):
return [cls.ALL, cls.STATIC]

@classmethod
def allow_dynamic(cls):
return [cls.ALL, cls.DYNAMIC]


EXECUTORS = {
"sequential": sequential,
"parallel": parallel,
}


class WrappedSchemaOption:
scope = "all"
scope = CommandScope.ALL
specific_schemas = None

allow_interactive = True
Expand All @@ -30,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",
Expand Down Expand Up @@ -102,7 +126,7 @@ def get_executor_from_options(self, **options):
return EXECUTORS["parallel"] if options.get("parallel") else EXECUTORS["sequential"]

def get_scope_display(self):
return "|".join(self.specific_schemas or []) or self.scope
return "|".join(self.specific_schemas or []) or self.scope.value

def _get_schemas_from_options(self, **options):
schemas = options.get("schemas") or []
Expand All @@ -112,8 +136,8 @@ def _get_schemas_from_options(self, **options):
include_dynamic_schemas = options.get("dynamic_schemas") or False
include_tenant_schemas = options.get("tenant_schemas") or False
dynamic_ready = dynamic_models_exist()
allow_static = self.scope in ["all", "static"]
allow_dynamic = self.scope in ["all", "dynamic"]
allow_static = self.scope in CommandScope.allow_static()
allow_dynamic = self.scope in CommandScope.allow_dynamic()
clone_reference = get_clone_reference()

if (
Expand Down Expand Up @@ -148,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:
Expand Down Expand Up @@ -236,8 +260,8 @@ def handle_tenant(self, tenant, *args, **options):


class StaticTenantCommand(TenantCommand):
scope = "static"
scope = CommandScope.STATIC


class DynamicTenantCommand(TenantCommand):
scope = "dynamic"
scope = CommandScope.DYNAMIC
24 changes: 8 additions & 16 deletions dpgs_sandbox/tests/test_tenant_commands.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from django.core.management.base import CommandError
from django.test import TestCase

from django_pgschemas.management.commands import CommandScope
from django_pgschemas.management.commands.whowill import Command as WhoWillCommand
from django_pgschemas.utils import get_domain_model, get_tenant_model

Expand Down Expand Up @@ -39,52 +40,43 @@ def test_no_schema_provided(self):
self.assertEqual(str(ctx.exception), "No schema provided")

def test_no_all_schemas_allowed(self):
command = WhoWillCommand()
command.scope = "unknown-scope"
with self.assertRaises(CommandError) as ctx:
management.call_command(command, all_schemas=True, verbosity=0)
self.assertEqual(str(ctx.exception), "Including all schemas is NOT allowed")
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()
command.scope = "dynamic"
command.scope = CommandScope.DYNAMIC
with self.assertRaises(CommandError) as ctx:
management.call_command(command, static_schemas=True, verbosity=0)
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()
command.scope = "static"
command.scope = CommandScope.STATIC
with self.assertRaises(CommandError) as ctx:
management.call_command(command, dynamic_schemas=True, verbosity=0)
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()
command.scope = "static"
command.scope = CommandScope.STATIC
with self.assertRaises(CommandError) as ctx:
management.call_command(command, tenant_schemas=True, verbosity=0)
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:
Expand Down