mirror of
https://github.com/apache/superset.git
synced 2026-04-08 02:45:22 +00:00
504 lines
15 KiB
Python
504 lines
15 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file to you under
|
|
# the Apache License, Version 2.0 (the "License"); you may not
|
|
# use this file except in compliance with the License. You may obtain
|
|
# a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
from collections.abc import Iterator
|
|
from uuid import UUID
|
|
|
|
import pytest
|
|
from sqlalchemy.orm.session import Session
|
|
from superset_core.api.tasks import TaskProperties, TaskScope, TaskStatus
|
|
|
|
from superset.commands.tasks.exceptions import TaskNotAbortableError
|
|
from superset.models.tasks import Task
|
|
from superset.tasks.utils import get_active_dedup_key, get_finished_dedup_key
|
|
|
|
# Test constants
|
|
TASK_UUID = UUID("e7765491-40c1-4f35-a4f5-06308e79310e")
|
|
TASK_ID = 42
|
|
TEST_TASK_TYPE = "test_type"
|
|
TEST_TASK_KEY = "test-key"
|
|
TEST_USER_ID = 1
|
|
|
|
|
|
def create_task(
|
|
session: Session,
|
|
*,
|
|
task_id: int | None = None,
|
|
task_uuid: UUID | None = None,
|
|
task_key: str = TEST_TASK_KEY,
|
|
task_type: str = TEST_TASK_TYPE,
|
|
scope: TaskScope = TaskScope.PRIVATE,
|
|
status: TaskStatus = TaskStatus.PENDING,
|
|
user_id: int | None = TEST_USER_ID,
|
|
properties: TaskProperties | None = None,
|
|
use_finished_dedup_key: bool = False,
|
|
) -> Task:
|
|
"""Helper to create a task with sensible defaults for testing."""
|
|
if use_finished_dedup_key:
|
|
dedup_key = get_finished_dedup_key(task_uuid or TASK_UUID)
|
|
else:
|
|
dedup_key = get_active_dedup_key(
|
|
scope=scope,
|
|
task_type=task_type,
|
|
task_key=task_key,
|
|
user_id=user_id,
|
|
)
|
|
|
|
task = Task(
|
|
task_type=task_type,
|
|
task_key=task_key,
|
|
scope=scope.value,
|
|
status=status.value,
|
|
dedup_key=dedup_key,
|
|
user_id=user_id,
|
|
)
|
|
if task_id is not None:
|
|
task.id = task_id
|
|
if task_uuid:
|
|
task.uuid = task_uuid
|
|
if properties:
|
|
task.update_properties(properties)
|
|
|
|
session.add(task)
|
|
session.flush()
|
|
return task
|
|
|
|
|
|
@pytest.fixture
|
|
def session_with_task(session: Session) -> Iterator[Session]:
|
|
"""Create a session with Task and TaskSubscriber tables."""
|
|
from superset.models.task_subscribers import TaskSubscriber
|
|
|
|
engine = session.get_bind()
|
|
Task.metadata.create_all(engine)
|
|
TaskSubscriber.metadata.create_all(engine)
|
|
|
|
yield session
|
|
session.rollback()
|
|
|
|
|
|
def test_find_by_task_key_active(session_with_task: Session) -> None:
|
|
"""Test finding active task by task_key"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
create_task(session_with_task)
|
|
|
|
result = TaskDAO.find_by_task_key(
|
|
task_type=TEST_TASK_TYPE,
|
|
task_key=TEST_TASK_KEY,
|
|
scope=TaskScope.PRIVATE,
|
|
user_id=TEST_USER_ID,
|
|
)
|
|
|
|
assert result is not None
|
|
assert result.task_key == TEST_TASK_KEY
|
|
assert result.task_type == TEST_TASK_TYPE
|
|
assert result.status == TaskStatus.PENDING.value
|
|
|
|
|
|
def test_find_by_task_key_not_found(session_with_task: Session) -> None:
|
|
"""Test finding task by task_key returns None when not found"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
result = TaskDAO.find_by_task_key(
|
|
task_type=TEST_TASK_TYPE,
|
|
task_key="nonexistent-key",
|
|
scope=TaskScope.PRIVATE,
|
|
user_id=TEST_USER_ID,
|
|
)
|
|
|
|
assert result is None
|
|
|
|
|
|
def test_find_by_task_key_finished_not_found(session_with_task: Session) -> None:
|
|
"""Test that find_by_task_key returns None for finished tasks.
|
|
|
|
Finished tasks have a different dedup_key format (UUID-based),
|
|
so they won't be found by the active task lookup.
|
|
"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
create_task(
|
|
session_with_task,
|
|
task_key="finished-key",
|
|
status=TaskStatus.SUCCESS,
|
|
use_finished_dedup_key=True,
|
|
task_uuid=TASK_UUID,
|
|
)
|
|
|
|
# Should not find SUCCESS task via active lookup
|
|
result = TaskDAO.find_by_task_key(
|
|
task_type=TEST_TASK_TYPE,
|
|
task_key="finished-key",
|
|
scope=TaskScope.PRIVATE,
|
|
user_id=TEST_USER_ID,
|
|
)
|
|
assert result is None
|
|
|
|
|
|
def test_create_task_success(session_with_task: Session) -> None:
|
|
"""Test successful task creation."""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
result = TaskDAO.create_task(
|
|
task_type=TEST_TASK_TYPE,
|
|
task_key=TEST_TASK_KEY,
|
|
scope=TaskScope.PRIVATE,
|
|
user_id=TEST_USER_ID,
|
|
)
|
|
|
|
assert result is not None
|
|
assert result.task_key == TEST_TASK_KEY
|
|
assert result.task_type == TEST_TASK_TYPE
|
|
assert result.status == TaskStatus.PENDING.value
|
|
assert isinstance(result, Task)
|
|
|
|
|
|
def test_create_task_with_user_id(session_with_task: Session) -> None:
|
|
"""Test task creation with explicit user_id."""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
result = TaskDAO.create_task(
|
|
task_type=TEST_TASK_TYPE,
|
|
task_key="user-task",
|
|
scope=TaskScope.PRIVATE,
|
|
user_id=42,
|
|
)
|
|
|
|
assert result is not None
|
|
assert result.user_id == 42
|
|
# Creator should be auto-subscribed
|
|
assert len(result.subscribers) == 1
|
|
assert result.subscribers[0].user_id == 42
|
|
|
|
|
|
def test_create_task_with_properties(session_with_task: Session) -> None:
|
|
"""Test task creation with properties."""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
result = TaskDAO.create_task(
|
|
task_type=TEST_TASK_TYPE,
|
|
task_key="props-task",
|
|
scope=TaskScope.PRIVATE,
|
|
user_id=TEST_USER_ID,
|
|
properties={"timeout": 300},
|
|
)
|
|
|
|
assert result is not None
|
|
assert result.properties_dict.get("timeout") == 300
|
|
|
|
|
|
def test_abort_task_pending_success(session_with_task: Session) -> None:
|
|
"""Test successful abort of pending task - goes directly to ABORTED"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_key="pending-task",
|
|
status=TaskStatus.PENDING,
|
|
)
|
|
|
|
result = TaskDAO.abort_task(task.uuid, skip_base_filter=True)
|
|
|
|
assert result is not None
|
|
assert result.status == TaskStatus.ABORTED.value
|
|
|
|
|
|
def test_abort_task_in_progress_abortable(session_with_task: Session) -> None:
|
|
"""Test abort of in-progress task with abort handler.
|
|
|
|
Should transition to ABORTING status.
|
|
"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_key="abortable-task",
|
|
status=TaskStatus.IN_PROGRESS,
|
|
properties={"is_abortable": True},
|
|
)
|
|
|
|
result = TaskDAO.abort_task(task.uuid, skip_base_filter=True)
|
|
|
|
assert result is not None
|
|
# Should set status to ABORTING, not ABORTED
|
|
assert result.status == TaskStatus.ABORTING.value
|
|
|
|
|
|
def test_abort_task_in_progress_not_abortable(session_with_task: Session) -> None:
|
|
"""Test abort of in-progress task without abort handler - raises error"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_key="non-abortable-task",
|
|
status=TaskStatus.IN_PROGRESS,
|
|
properties={"is_abortable": False},
|
|
)
|
|
|
|
with pytest.raises(TaskNotAbortableError):
|
|
TaskDAO.abort_task(task.uuid, skip_base_filter=True)
|
|
|
|
|
|
def test_abort_task_in_progress_is_abortable_none(session_with_task: Session) -> None:
|
|
"""Test abort of in-progress task with is_abortable not set - raises error"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_key="no-abortable-prop-task",
|
|
status=TaskStatus.IN_PROGRESS,
|
|
# Empty properties - no is_abortable key
|
|
)
|
|
|
|
with pytest.raises(TaskNotAbortableError):
|
|
TaskDAO.abort_task(task.uuid, skip_base_filter=True)
|
|
|
|
|
|
def test_abort_task_already_aborting(session_with_task: Session) -> None:
|
|
"""Test abort of already aborting task - idempotent success"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_key="aborting-task",
|
|
status=TaskStatus.ABORTING,
|
|
)
|
|
|
|
result = TaskDAO.abort_task(task.uuid, skip_base_filter=True)
|
|
|
|
# Idempotent - returns task without error
|
|
assert result is not None
|
|
assert result.status == TaskStatus.ABORTING.value
|
|
|
|
|
|
def test_abort_task_not_found(session_with_task: Session) -> None:
|
|
"""Test abort fails when task not found"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
result = TaskDAO.abort_task(UUID("00000000-0000-0000-0000-000000000000"))
|
|
|
|
assert result is None
|
|
|
|
|
|
def test_abort_task_already_finished(session_with_task: Session) -> None:
|
|
"""Test abort fails when task already finished"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_key="finished-task",
|
|
status=TaskStatus.SUCCESS,
|
|
use_finished_dedup_key=True,
|
|
task_uuid=TASK_UUID,
|
|
)
|
|
|
|
result = TaskDAO.abort_task(task.uuid, skip_base_filter=True)
|
|
|
|
assert result is None
|
|
|
|
|
|
def test_add_subscriber(session_with_task: Session) -> None:
|
|
"""Test adding a subscriber to a task"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_key="shared-task",
|
|
scope=TaskScope.SHARED,
|
|
user_id=None,
|
|
)
|
|
|
|
# Add subscriber
|
|
result = TaskDAO.add_subscriber(task.id, user_id=TEST_USER_ID)
|
|
assert result is True
|
|
|
|
# Verify subscriber was added
|
|
session_with_task.refresh(task)
|
|
assert len(task.subscribers) == 1
|
|
assert task.subscribers[0].user_id == TEST_USER_ID
|
|
|
|
|
|
def test_add_subscriber_idempotent(session_with_task: Session) -> None:
|
|
"""Test adding same subscriber twice is idempotent"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_key="shared-task-2",
|
|
scope=TaskScope.SHARED,
|
|
user_id=None,
|
|
)
|
|
|
|
# Add subscriber twice
|
|
result1 = TaskDAO.add_subscriber(task.id, user_id=TEST_USER_ID)
|
|
result2 = TaskDAO.add_subscriber(task.id, user_id=TEST_USER_ID)
|
|
|
|
assert result1 is True
|
|
assert result2 is False # Already subscribed
|
|
|
|
# Verify only one subscriber
|
|
session_with_task.refresh(task)
|
|
assert len(task.subscribers) == 1
|
|
|
|
|
|
def test_remove_subscriber(session_with_task: Session) -> None:
|
|
"""Test removing a subscriber from a task"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_key="shared-task-3",
|
|
scope=TaskScope.SHARED,
|
|
user_id=None,
|
|
)
|
|
|
|
TaskDAO.add_subscriber(task.id, user_id=TEST_USER_ID)
|
|
session_with_task.refresh(task)
|
|
assert len(task.subscribers) == 1
|
|
|
|
# Remove subscriber
|
|
result = TaskDAO.remove_subscriber(task.id, user_id=TEST_USER_ID)
|
|
|
|
assert result is not None
|
|
assert len(result.subscribers) == 0
|
|
|
|
|
|
def test_remove_subscriber_not_subscribed(session_with_task: Session) -> None:
|
|
"""Test removing non-existent subscriber returns None"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_key="shared-task-4",
|
|
scope=TaskScope.SHARED,
|
|
user_id=None,
|
|
)
|
|
|
|
# Try to remove non-existent subscriber
|
|
result = TaskDAO.remove_subscriber(task.id, user_id=999)
|
|
|
|
assert result is None
|
|
|
|
|
|
def test_get_status(session_with_task: Session) -> None:
|
|
"""Test get_status returns status string when task found by UUID"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_uuid=TASK_UUID,
|
|
task_key="status-task",
|
|
status=TaskStatus.IN_PROGRESS,
|
|
)
|
|
|
|
result = TaskDAO.get_status(task.uuid)
|
|
|
|
assert result == TaskStatus.IN_PROGRESS.value
|
|
|
|
|
|
def test_get_status_not_found(session_with_task: Session) -> None:
|
|
"""Test get_status returns None when task not found"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
result = TaskDAO.get_status(UUID("00000000-0000-0000-0000-000000000000"))
|
|
|
|
assert result is None
|
|
|
|
|
|
def test_conditional_status_update_non_terminal_state_keeps_dedup_key(
|
|
session_with_task: Session,
|
|
) -> None:
|
|
"""Test that conditional_status_update preserves dedup_key for
|
|
non-terminal transitions"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
# Create task in PENDING state
|
|
task = create_task(
|
|
session_with_task,
|
|
task_uuid=TASK_UUID,
|
|
task_key="non-terminal-test-task",
|
|
status=TaskStatus.PENDING,
|
|
)
|
|
|
|
# Store original active dedup_key
|
|
original_dedup_key = task.dedup_key
|
|
|
|
# Transition to non-terminal state (IN_PROGRESS)
|
|
result = TaskDAO.conditional_status_update(
|
|
task_uuid=TASK_UUID,
|
|
new_status=TaskStatus.IN_PROGRESS,
|
|
expected_status=TaskStatus.PENDING,
|
|
set_started_at=True,
|
|
)
|
|
|
|
# Should succeed
|
|
assert result is True
|
|
|
|
# Refresh task and verify dedup_key was NOT changed
|
|
session_with_task.refresh(task)
|
|
assert task.status == TaskStatus.IN_PROGRESS.value
|
|
assert task.dedup_key == original_dedup_key # Should remain the same
|
|
assert task.started_at is not None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"terminal_state",
|
|
[
|
|
TaskStatus.SUCCESS,
|
|
TaskStatus.FAILURE,
|
|
TaskStatus.ABORTED,
|
|
TaskStatus.TIMED_OUT,
|
|
],
|
|
)
|
|
def test_conditional_status_update_terminal_state_updates_dedup_key(
|
|
session_with_task: Session, terminal_state: TaskStatus
|
|
) -> None:
|
|
"""Test that terminal states (SUCCESS, FAILURE, ABORTED, TIMED_OUT)
|
|
update dedup_key"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
task = create_task(
|
|
session_with_task,
|
|
task_uuid=TASK_UUID,
|
|
task_key=f"terminal-test-{terminal_state.value}",
|
|
status=TaskStatus.IN_PROGRESS,
|
|
)
|
|
|
|
original_dedup_key = task.dedup_key
|
|
expected_finished_key = get_finished_dedup_key(TASK_UUID)
|
|
|
|
# Transition to terminal state
|
|
result = TaskDAO.conditional_status_update(
|
|
task_uuid=TASK_UUID,
|
|
new_status=terminal_state,
|
|
expected_status=TaskStatus.IN_PROGRESS,
|
|
set_ended_at=True,
|
|
)
|
|
|
|
assert result is True, f"Failed to update to {terminal_state.value}"
|
|
|
|
# Verify dedup_key was updated
|
|
session_with_task.refresh(task)
|
|
assert task.status == terminal_state.value
|
|
assert task.dedup_key == expected_finished_key, (
|
|
f"dedup_key not updated for {terminal_state.value}"
|
|
)
|
|
assert task.dedup_key != original_dedup_key, (
|
|
f"dedup_key should have changed for {terminal_state.value}"
|
|
)
|