base_commands.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. from django.core.management.base import BaseCommand
  2. from model_bakery import baker
  3. from model_bakery.random_gen import gen_json, gen_slug
  4. from organizations_ext.models import Organization
  5. from projects.models import Project
  6. baker.generators.add("organizations.fields.SlugField", gen_slug)
  7. baker.generators.add("organizations_ext.fields.OrganizationSlugField", gen_slug)
  8. baker.generators.add("django.db.models.JSONField", gen_json)
  9. class MakeSampleCommand(BaseCommand):
  10. organization = None
  11. project = None
  12. batch_size = 10000
  13. def add_org_project_arguments(self, parser):
  14. parser.add_argument("--org", type=str, help="Organization slug")
  15. parser.add_argument("--project", type=str, help="Project slug")
  16. def add_arguments(self, parser):
  17. parser.add_argument("--quantity", type=int, default=1000)
  18. self.add_org_project_arguments(parser)
  19. def handle(self, *args, **options):
  20. self.organization = self.get_organization(options.get("org"))
  21. self.project = self.get_project(options.get("project"))
  22. def get_organization(self, org: str):
  23. if org:
  24. return Organization.objects.get(slug=org)
  25. organization = Organization.objects.first()
  26. if not organization:
  27. organization = baker.make("organizations_ext.Organization")
  28. return organization
  29. def get_project(self, project: str):
  30. if project:
  31. return Project.objects.get(slug=project, organization=self.organization)
  32. project = Project.objects.filter(organization=self.organization).first()
  33. if not project:
  34. project = baker.make("projects.Project", organization=self.organization)
  35. return project
  36. def progress_tick(self):
  37. self.stdout.write(self.style.NOTICE("."), ending="")
  38. def success_message(self, message: str):
  39. self.stdout.write(self.style.SUCCESS(message))