jira/apps/api/plane/bgtasks/dummy_data_task.py

# Copyright (c) 2023-present Plane Software, Inc. and contributors
# SPDX-License-Identifier: AGPL-3.0-only
# See the LICENSE file for details.

# Python imports
import uuid
import random
from datetime import datetime, timedelta

# Django imports
from django.db.models import Max

# Third party imports
from celery import shared_task
from faker import Faker

# Module imports
from plane.db.models import (
    Workspace,
    User,
    Project,
    ProjectMember,
    State,
    StateGroup,
    Label,
    Cycle,
    Module,
    Issue,
    IssueSequence,
    IssueAssignee,
    IssueLabel,
    IssueActivity,
    CycleIssue,
    ModuleIssue,
    Page,
    ProjectPage,
    PageLabel,
    Intake,
    IntakeIssue,
)
from plane.db.models.intake import SourceType


def create_project(workspace, user_id):
    fake = Faker()
    name = fake.name()
    unique_id = str(uuid.uuid4())[:5]

    project = Project.objects.create(
        workspace=workspace,
        name=f"{name}_{unique_id}",
        identifier=name[: random.randint(2, 12 if len(name) - 1 >= 12 else len(name) - 1)].upper(),
        created_by_id=user_id,
        intake_view=True,
    )

    # Add current member as project member
    _ = ProjectMember.objects.create(project=project, member_id=user_id, role=20)

    return project


def create_project_members(workspace, project, members):
    members = User.objects.filter(email__in=members)

    _ = ProjectMember.objects.bulk_create(
        [
            ProjectMember(
                project=project,
                workspace=workspace,
                member=member,
                role=20,
                sort_order=random.randint(0, 65535),
            )
            for member in members
        ],
        ignore_conflicts=True,
    )
    return


def create_states(workspace, project, user_id):
    states = [
        {
            "name": "Backlog",
            "color": "#A3A3A3",
            "sequence": 15000,
            "group": "backlog",
            "default": True,
        },
        {"name": "Todo", "color": "#3A3A3A", "sequence": 25000, "group": "unstarted"},
        {
            "name": "In Progress",
            "color": "#F59E0B",
            "sequence": 35000,
            "group": "started",
        },
        {"name": "Done", "color": "#16A34A", "sequence": 45000, "group": "completed"},
        {
            "name": "Cancelled",
            "color": "#EF4444",
            "sequence": 55000,
            "group": "cancelled",
        },
    ]

    states = State.objects.bulk_create(
        [
            State(
                name=state["name"],
                color=state["color"],
                project=project,
                sequence=state["sequence"],
                workspace=workspace,
                group=state["group"],
                default=state.get("default", False),
                created_by_id=user_id,
            )
            for state in states
        ]
    )

    return states


def create_labels(workspace, project, user_id):
    fake = Faker()
    Faker.seed(0)

    return Label.objects.bulk_create(
        [
            Label(
                name=fake.color_name(),
                color=fake.hex_color(),
                project=project,
                workspace=workspace,
                created_by_id=user_id,
                sort_order=random.randint(0, 65535),
            )
            for _ in range(0, 50)
        ],
        ignore_conflicts=True,
    )


def create_cycles(workspace, project, user_id, cycle_count):
    fake = Faker()
    Faker.seed(0)

    cycles = []
    used_date_ranges = set()  # Track used date ranges

    while len(cycles) <= cycle_count:
        # Generate a start date, allowing for None
        start_date_option = [None, fake.date_this_year()]
        start_date = start_date_option[random.randint(0, 1)]

        # Initialize end_date based on start_date
        end_date = (
            None
            if start_date is None
            else fake.date_between_dates(
                date_start=start_date,
                date_end=datetime.now().date().replace(month=12, day=31),
            )
        )

        # Ensure end_date is strictly after start_date if start_date is not None
        while start_date is not None and (end_date <= start_date or (start_date, end_date) in used_date_ranges):
            end_date = fake.date_this_year()

        # Add the unique date range to the set
        (used_date_ranges.add((start_date, end_date)) if (end_date is not None and start_date is not None) else None)

        # Append the cycle with unique date range
        cycles.append(
            Cycle(
                name=fake.name(),
                owned_by_id=user_id,
                sort_order=random.randint(0, 65535),
                start_date=start_date,
                end_date=end_date,
                project=project,
                workspace=workspace,
            )
        )

    return Cycle.objects.bulk_create(cycles, ignore_conflicts=True)


def create_modules(workspace, project, user_id, module_count):
    fake = Faker()
    Faker.seed(0)

    modules = []
    for _ in range(0, module_count):
        start_date = [None, fake.date_this_year()][random.randint(0, 1)]
        end_date = (
            None
            if start_date is None
            else fake.date_between_dates(
                date_start=start_date,
                date_end=datetime.now().date().replace(month=12, day=31),
            )
        )

        modules.append(
            Module(
                name=fake.name(),
                sort_order=random.randint(0, 65535),
                start_date=start_date,
                target_date=end_date,
                project=project,
                workspace=workspace,
            )
        )

    return Module.objects.bulk_create(modules, ignore_conflicts=True)


def create_pages(workspace, project, user_id, pages_count):
    fake = Faker()
    Faker.seed(0)

    pages = []
    for _ in range(0, pages_count):
        text = fake.text(max_nb_chars=60000)
        pages.append(
            Page(
                name=fake.name(),
                workspace=workspace,
                owned_by_id=user_id,
                access=random.randint(0, 1),
                color=fake.hex_color(),
                description_html=f"<p>{text}</p>",
                archived_at=None,
                is_locked=False,
            )
        )
    # Bulk create pages
    pages = Page.objects.bulk_create(pages, ignore_conflicts=True)
    # Add Page to project
    ProjectPage.objects.bulk_create(
        [ProjectPage(page=page, project=project, workspace=workspace) for page in pages],
        batch_size=1000,
    )


def create_page_labels(workspace, project, user_id, pages_count):
    # labels
    labels = Label.objects.filter(project=project).values_list("id", flat=True)
    pages = random.sample(
        list(Page.objects.filter(projects__id=project.id).values_list("id", flat=True)),
        int(pages_count / 2),
    )

    # Bulk page labels
    bulk_page_labels = []
    for page in pages:
        for label in random.sample(list(labels), random.randint(0, len(labels) - 1)):
            bulk_page_labels.append(PageLabel(page_id=page, label_id=label, workspace=workspace))

    # Page labels
    PageLabel.objects.bulk_create(bulk_page_labels, batch_size=1000, ignore_conflicts=True)


def create_issues(workspace, project, user_id, issue_count):
    fake = Faker()
    Faker.seed(0)

    states = (
        State.objects.filter(workspace=workspace, project=project)
        .exclude(group=StateGroup.TRIAGE.value)
        .values_list("id", flat=True)
    )
    creators = ProjectMember.objects.filter(workspace=workspace, project=project).values_list("member_id", flat=True)

    issues = []

    # Get the maximum sequence_id
    last_id = IssueSequence.objects.filter(project=project).aggregate(largest=Max("sequence"))["largest"]

    last_id = 1 if last_id is None else last_id + 1

    # Get the maximum sort order
    largest_sort_order = Issue.objects.filter(
        project=project, state_id=states[random.randint(0, len(states) - 1)]
    ).aggregate(largest=Max("sort_order"))["largest"]

    largest_sort_order = 65535 if largest_sort_order is None else largest_sort_order + 10000

    for _ in range(0, issue_count):
        start_date = [None, fake.date_this_year()][random.randint(0, 1)]
        end_date = (
            None
            if start_date is None
            else fake.date_between_dates(
                date_start=start_date,
                date_end=datetime.now().date().replace(month=12, day=31),
            )
        )

        text = fake.text(max_nb_chars=3000)
        issues.append(
            Issue(
                state_id=states[random.randint(0, len(states) - 1)],
                project=project,
                workspace=workspace,
                name=text[:254],
                description_html=f"<p>{text}</p>",
                description_stripped=text,
                sequence_id=last_id,
                sort_order=largest_sort_order,
                start_date=start_date,
                target_date=end_date,
                priority=["urgent", "high", "medium", "low", "none"][random.randint(0, 4)],
                created_by_id=creators[random.randint(0, len(creators) - 1)],
            )
        )

        largest_sort_order = largest_sort_order + random.randint(0, 1000)
        last_id = last_id + 1

    issues = Issue.objects.bulk_create(issues, ignore_conflicts=True, batch_size=1000)
    # Sequences
    _ = IssueSequence.objects.bulk_create(
        [
            IssueSequence(
                issue=issue,
                sequence=issue.sequence_id,
                project=project,
                workspace=workspace,
            )
            for issue in issues
        ],
        batch_size=100,
    )

    # Track the issue activities
    IssueActivity.objects.bulk_create(
        [
            IssueActivity(
                issue=issue,
                actor_id=user_id,
                project=project,
                workspace=workspace,
                comment="created the issue",
                verb="created",
                created_by_id=user_id,
            )
            for issue in issues
        ],
        batch_size=100,
    )
    return issues


def create_intake_issues(workspace, project, user_id, intake_issue_count):
    issues = create_issues(workspace, project, user_id, intake_issue_count)
    intake, create = Intake.objects.get_or_create(name="Intake", project=project, is_default=True)
    IntakeIssue.objects.bulk_create(
        [
            IntakeIssue(
                issue=issue,
                intake=intake,
                status=(status := [-2, -1, 0, 1, 2][random.randint(0, 4)]),
                snoozed_till=(datetime.now() + timedelta(days=random.randint(1, 30)) if status == 0 else None),
                source=SourceType.IN_APP,
                workspace=workspace,
                project=project,
            )
            for issue in issues
        ],
        batch_size=100,
    )


def create_issue_parent(workspace, project, user_id, issue_count):
    parent_count = issue_count / 4

    parent_issues = Issue.objects.filter(project=project).values_list("id", flat=True)[: int(parent_count)]
    sub_issues = Issue.objects.filter(project=project).exclude(pk__in=parent_issues)[: int(issue_count / 2)]

    bulk_sub_issues = []
    for sub_issue in sub_issues:
        sub_issue.parent_id = parent_issues[random.randint(0, int(parent_count - 1))]

    Issue.objects.bulk_update(bulk_sub_issues, ["parent"], batch_size=1000)


def create_issue_assignees(workspace, project, user_id, issue_count):
    # assignees
    assignees = ProjectMember.objects.filter(project=project).values_list("member_id", flat=True)
    issues = random.sample(
        list(Issue.objects.filter(project=project).values_list("id", flat=True)),
        int(issue_count / 2),
    )

    # Bulk issue
    bulk_issue_assignees = []
    for issue in issues:
        for assignee in random.sample(list(assignees), random.randint(0, len(assignees) - 1)):
            bulk_issue_assignees.append(
                IssueAssignee(
                    issue_id=issue,
                    assignee_id=assignee,
                    project=project,
                    workspace=workspace,
                )
            )

    # Issue assignees
    IssueAssignee.objects.bulk_create(bulk_issue_assignees, batch_size=1000, ignore_conflicts=True)


def create_issue_labels(workspace, project, user_id, issue_count):
    # labels
    labels = Label.objects.filter(project=project).values_list("id", flat=True)
    # issues = random.sample(
    #     list(
    #         Issue.objects.filter(project=project).values_list("id", flat=True)
    #     ),
    #     int(issue_count / 2),
    # )
    issues = list(Issue.objects.filter(project=project).values_list("id", flat=True))
    shuffled_labels = list(labels)

    # Bulk issue
    bulk_issue_labels = []
    for issue in issues:
        random.shuffle(shuffled_labels)
        for label in random.sample(shuffled_labels, random.randint(0, 5)):
            bulk_issue_labels.append(IssueLabel(issue_id=issue, label_id=label, project=project, workspace=workspace))

    # Issue labels
    IssueLabel.objects.bulk_create(bulk_issue_labels, batch_size=1000, ignore_conflicts=True)


def create_cycle_issues(workspace, project, user_id, issue_count):
    # assignees
    cycles = Cycle.objects.filter(project=project).values_list("id", flat=True)
    issues = random.sample(
        list(Issue.objects.filter(project=project).values_list("id", flat=True)),
        int(issue_count / 2),
    )

    # Bulk issue
    bulk_cycle_issues = []
    for issue in issues:
        cycle = cycles[random.randint(0, len(cycles) - 1)]
        bulk_cycle_issues.append(CycleIssue(cycle_id=cycle, issue_id=issue, project=project, workspace=workspace))

    # Issue assignees
    CycleIssue.objects.bulk_create(bulk_cycle_issues, batch_size=1000, ignore_conflicts=True)


def create_module_issues(workspace, project, user_id, issue_count):
    # assignees
    modules = Module.objects.filter(project=project).values_list("id", flat=True)
    # issues = random.sample(
    #     list(
    #         Issue.objects.filter(project=project).values_list("id", flat=True)
    #     ),
    #     int(issue_count / 2),
    # )
    issues = list(Issue.objects.filter(project=project).values_list("id", flat=True))

    shuffled_modules = list(modules)

    # Bulk issue
    bulk_module_issues = []
    for issue in issues:
        random.shuffle(shuffled_modules)
        for module in random.sample(shuffled_modules, random.randint(0, 5)):
            bulk_module_issues.append(
                ModuleIssue(
                    module_id=module,
                    issue_id=issue,
                    project=project,
                    workspace=workspace,
                )
            )
    # Issue assignees
    ModuleIssue.objects.bulk_create(bulk_module_issues, batch_size=1000, ignore_conflicts=True)


@shared_task
def create_dummy_data(
    slug,
    email,
    members,
    issue_count,
    cycle_count,
    module_count,
    pages_count,
    intake_issue_count,
):
    workspace = Workspace.objects.get(slug=slug)

    user = User.objects.get(email=email)
    user_id = user.id

    # Create a project
    project = create_project(workspace=workspace, user_id=user_id)

    # create project members
    create_project_members(workspace=workspace, project=project, members=members)

    # Create states
    create_states(workspace=workspace, project=project, user_id=user_id)

    # Create labels
    create_labels(workspace=workspace, project=project, user_id=user_id)

    # create cycles
    create_cycles(workspace=workspace, project=project, user_id=user_id, cycle_count=cycle_count)

    # create modules
    create_modules(workspace=workspace, project=project, user_id=user_id, module_count=module_count)

    # create pages
    create_pages(workspace=workspace, project=project, user_id=user_id, pages_count=pages_count)

    # create page labels
    create_page_labels(workspace=workspace, project=project, user_id=user_id, pages_count=pages_count)

    # create issues
    create_issues(workspace=workspace, project=project, user_id=user_id, issue_count=issue_count)

    # create intake issues
    create_intake_issues(
        workspace=workspace,
        project=project,
        user_id=user_id,
        intake_issue_count=intake_issue_count,
    )

    # create issue parent
    create_issue_parent(workspace=workspace, project=project, user_id=user_id, issue_count=issue_count)

    # create issue assignees
    create_issue_assignees(workspace=workspace, project=project, user_id=user_id, issue_count=issue_count)

    # create issue labels
    create_issue_labels(workspace=workspace, project=project, user_id=user_id, issue_count=issue_count)

    # create cycle issues
    create_cycle_issues(workspace=workspace, project=project, user_id=user_id, issue_count=issue_count)

    # create module issues
    create_module_issues(workspace=workspace, project=project, user_id=user_id, issue_count=issue_count)

    return