Repository: langgenius/dify
File: api/tasks/remove_app_and_related_data_task.py (Python)

Code before:
    return delete_draft_variables_batch(app_id, batch_size=1000)


def delete_draft_variables_batch(app_id: str, batch_size: int = 1000) -> int:
    """
    Delete draft variables for an app in batches.

    This function now handles cleanup of associated Offload data including:
    - WorkflowDraftVariableFile records
    - UploadFile records
    - Object storage files

    Args:
        app_id: The ID of the app whose draft variables should be deleted
        batch_size: Number of records to delete per batch

    Returns:
        Total number of records deleted
    """
    if batch_size <= 0:
        raise ValueError("batch_size must be positive")

    total_deleted = 0
    total_files_deleted = 0

    with session_factory.create_session() as session:
        while True:
            # Get a batch of draft variable IDs along with their file_ids
            query_sql = """
                SELECT id, file_id FROM workflow_draft_variables
                WHERE app_id = :app_id
                LIMIT :batch_size
            """
            result = session.execute(sa.text(query_sql), {"app_id": app_id, "batch_size": batch_size})

            rows = list(result)
            if not rows:
                break

            draft_var_ids = [row[0] for row in rows]
            file_ids = [row[1] for row in rows if row[1] is not None]

            # Clean up associated Offload data first
            if file_ids:
                files_deleted = _delete_draft_variable_offload_data(session, file_ids)
                total_files_deleted += files_deleted

            # Delete the draft variables
            delete_sql = """
                DELETE FROM workflow_draft_variables
                WHERE id IN :ids

Review comment (suggestion):
I suggest use context manager to ensure the context management stay close to the session creation, avoiding forgetting calls to `session.commit`. Like the following:

```py
with session_factory.create_session() as session, session.begin():
```

Code after:
    return delete_draft_variables_batch(app_id, batch_size=1000)


def delete_draft_variables_batch(app_id: str, batch_size: int = 1000) -> int:
    """
    Delete draft variables for an app in batches.

    This function now handles cleanup of associated Offload data including:
    - WorkflowDraftVariableFile records
    - UploadFile records
    - Object storage files

    Args:
        app_id: The ID of the app whose draft variables should be deleted
        batch_size: Number of records to delete per batch

    Returns:
        Total number of records deleted
    """
    if batch_size <= 0:
        raise ValueError("batch_size must be positive")

    total_deleted = 0
    total_files_deleted = 0

    while True:
        with session_factory.create_session() as session, session.begin():
            # Get a batch of draft variable IDs along with their file_ids
            query_sql = """
                SELECT id, file_id FROM workflow_draft_variables
                WHERE app_id = :app_id
                LIMIT :batch_size
            """
            result = session.execute(sa.text(query_sql), {"app_id": app_id, "batch_size": batch_size})

            rows = list(result)
            if not rows:
                break

            draft_var_ids = [row[0] for row in rows]
            file_ids = [row[1] for row in rows if row[1] is not None]

            # Clean up associated Offload data first
            if file_ids:
                files_deleted = _delete_draft_variable_offload_data(session, file_ids)
                total_files_deleted += files_deleted

            # Delete the draft variables
            delete_sql = """
                DELETE FROM workflow_draft_variables
                WHERE id IN :ids