mirror of
https://github.com/langgenius/dify.git
synced 2026-04-05 11:11:12 +08:00
feat: support relative mode for message clean command
This commit is contained in:
@@ -2598,15 +2598,29 @@ def migrate_oss(
|
|||||||
@click.option(
|
@click.option(
|
||||||
"--start-from",
|
"--start-from",
|
||||||
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
|
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
|
||||||
required=True,
|
required=False,
|
||||||
|
default=None,
|
||||||
help="Lower bound (inclusive) for created_at.",
|
help="Lower bound (inclusive) for created_at.",
|
||||||
)
|
)
|
||||||
@click.option(
|
@click.option(
|
||||||
"--end-before",
|
"--end-before",
|
||||||
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
|
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
|
||||||
required=True,
|
required=False,
|
||||||
|
default=None,
|
||||||
help="Upper bound (exclusive) for created_at.",
|
help="Upper bound (exclusive) for created_at.",
|
||||||
)
|
)
|
||||||
|
@click.option(
|
||||||
|
"--from-days-ago",
|
||||||
|
type=int,
|
||||||
|
default=None,
|
||||||
|
help="Relative lower bound in days ago (inclusive). Must be used with --before-days.",
|
||||||
|
)
|
||||||
|
@click.option(
|
||||||
|
"--before-days",
|
||||||
|
type=int,
|
||||||
|
default=None,
|
||||||
|
help="Relative upper bound in days ago (exclusive). Required for relative mode.",
|
||||||
|
)
|
||||||
@click.option("--batch-size", default=1000, show_default=True, help="Batch size for selecting messages.")
|
@click.option("--batch-size", default=1000, show_default=True, help="Batch size for selecting messages.")
|
||||||
@click.option(
|
@click.option(
|
||||||
"--graceful-period",
|
"--graceful-period",
|
||||||
@@ -2618,8 +2632,10 @@ def migrate_oss(
|
|||||||
def clean_expired_messages(
|
def clean_expired_messages(
|
||||||
batch_size: int,
|
batch_size: int,
|
||||||
graceful_period: int,
|
graceful_period: int,
|
||||||
start_from: datetime.datetime,
|
start_from: datetime.datetime | None,
|
||||||
end_before: datetime.datetime,
|
end_before: datetime.datetime | None,
|
||||||
|
from_days_ago: int | None,
|
||||||
|
before_days: int | None,
|
||||||
dry_run: bool,
|
dry_run: bool,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
@@ -2630,18 +2646,64 @@ def clean_expired_messages(
|
|||||||
start_at = time.perf_counter()
|
start_at = time.perf_counter()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
abs_mode = start_from is not None and end_before is not None
|
||||||
|
rel_mode = before_days is not None
|
||||||
|
|
||||||
|
if abs_mode and rel_mode:
|
||||||
|
raise click.UsageError(
|
||||||
|
"Options are mutually exclusive: use either (--start-from,--end-before) "
|
||||||
|
"or (--from-days-ago,--before-days)."
|
||||||
|
)
|
||||||
|
|
||||||
|
if from_days_ago is not None and before_days is None:
|
||||||
|
raise click.UsageError("--from-days-ago must be used together with --before-days.")
|
||||||
|
|
||||||
|
if (start_from is None) ^ (end_before is None):
|
||||||
|
raise click.UsageError("Both --start-from and --end-before are required when using absolute time range.")
|
||||||
|
|
||||||
|
if not abs_mode and not rel_mode:
|
||||||
|
raise click.UsageError(
|
||||||
|
"You must provide either (--start-from,--end-before) or (--before-days [--from-days-ago])."
|
||||||
|
)
|
||||||
|
|
||||||
|
if rel_mode:
|
||||||
|
if before_days < 0:
|
||||||
|
raise click.UsageError("--before-days must be >= 0.")
|
||||||
|
if from_days_ago is not None:
|
||||||
|
if from_days_ago < 0:
|
||||||
|
raise click.UsageError("--from-days-ago must be >= 0.")
|
||||||
|
if from_days_ago <= before_days:
|
||||||
|
raise click.UsageError("--from-days-ago must be greater than --before-days.")
|
||||||
|
|
||||||
# Create policy based on billing configuration
|
# Create policy based on billing configuration
|
||||||
# NOTE: graceful_period will be ignored when billing is disabled.
|
# NOTE: graceful_period will be ignored when billing is disabled.
|
||||||
policy = create_message_clean_policy(graceful_period_days=graceful_period)
|
policy = create_message_clean_policy(graceful_period_days=graceful_period)
|
||||||
|
|
||||||
# Create and run the cleanup service
|
# Create and run the cleanup service
|
||||||
service = MessagesCleanService.from_time_range(
|
if abs_mode:
|
||||||
policy=policy,
|
service = MessagesCleanService.from_time_range(
|
||||||
start_from=start_from,
|
policy=policy,
|
||||||
end_before=end_before,
|
start_from=start_from,
|
||||||
batch_size=batch_size,
|
end_before=end_before,
|
||||||
dry_run=dry_run,
|
batch_size=batch_size,
|
||||||
)
|
dry_run=dry_run,
|
||||||
|
)
|
||||||
|
elif from_days_ago is None:
|
||||||
|
service = MessagesCleanService.from_days(
|
||||||
|
policy=policy,
|
||||||
|
days=before_days,
|
||||||
|
batch_size=batch_size,
|
||||||
|
dry_run=dry_run,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
now = datetime.datetime.now()
|
||||||
|
service = MessagesCleanService.from_time_range(
|
||||||
|
policy=policy,
|
||||||
|
start_from=now - datetime.timedelta(days=from_days_ago),
|
||||||
|
end_before=now - datetime.timedelta(days=before_days),
|
||||||
|
batch_size=batch_size,
|
||||||
|
dry_run=dry_run,
|
||||||
|
)
|
||||||
stats = service.run()
|
stats = service.run()
|
||||||
|
|
||||||
end_at = time.perf_counter()
|
end_at = time.perf_counter()
|
||||||
|
|||||||
184
api/tests/unit_tests/commands/test_clean_expired_messages.py
Normal file
184
api/tests/unit_tests/commands/test_clean_expired_messages.py
Normal file
@@ -0,0 +1,184 @@
|
|||||||
|
import datetime
|
||||||
|
import re
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
import click
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from commands import clean_expired_messages
|
||||||
|
|
||||||
|
|
||||||
|
def _mock_service() -> MagicMock:
|
||||||
|
service = MagicMock()
|
||||||
|
service.run.return_value = {
|
||||||
|
"batches": 1,
|
||||||
|
"total_messages": 10,
|
||||||
|
"filtered_messages": 5,
|
||||||
|
"total_deleted": 5,
|
||||||
|
}
|
||||||
|
return service
|
||||||
|
|
||||||
|
|
||||||
|
def test_absolute_mode_calls_from_time_range():
|
||||||
|
policy = object()
|
||||||
|
service = _mock_service()
|
||||||
|
start_from = datetime.datetime(2024, 1, 1, 0, 0, 0)
|
||||||
|
end_before = datetime.datetime(2024, 2, 1, 0, 0, 0)
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("commands.create_message_clean_policy", return_value=policy),
|
||||||
|
patch("commands.MessagesCleanService.from_time_range", return_value=service) as mock_from_time_range,
|
||||||
|
patch("commands.MessagesCleanService.from_days") as mock_from_days,
|
||||||
|
):
|
||||||
|
clean_expired_messages.callback(
|
||||||
|
batch_size=200,
|
||||||
|
graceful_period=21,
|
||||||
|
start_from=start_from,
|
||||||
|
end_before=end_before,
|
||||||
|
from_days_ago=None,
|
||||||
|
before_days=None,
|
||||||
|
dry_run=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_from_time_range.assert_called_once_with(
|
||||||
|
policy=policy,
|
||||||
|
start_from=start_from,
|
||||||
|
end_before=end_before,
|
||||||
|
batch_size=200,
|
||||||
|
dry_run=True,
|
||||||
|
)
|
||||||
|
mock_from_days.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_relative_mode_before_days_only_calls_from_days():
|
||||||
|
policy = object()
|
||||||
|
service = _mock_service()
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("commands.create_message_clean_policy", return_value=policy),
|
||||||
|
patch("commands.MessagesCleanService.from_days", return_value=service) as mock_from_days,
|
||||||
|
patch("commands.MessagesCleanService.from_time_range") as mock_from_time_range,
|
||||||
|
):
|
||||||
|
clean_expired_messages.callback(
|
||||||
|
batch_size=500,
|
||||||
|
graceful_period=14,
|
||||||
|
start_from=None,
|
||||||
|
end_before=None,
|
||||||
|
from_days_ago=None,
|
||||||
|
before_days=30,
|
||||||
|
dry_run=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_from_days.assert_called_once_with(
|
||||||
|
policy=policy,
|
||||||
|
days=30,
|
||||||
|
batch_size=500,
|
||||||
|
dry_run=False,
|
||||||
|
)
|
||||||
|
mock_from_time_range.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_relative_mode_with_from_days_ago_calls_from_time_range():
|
||||||
|
policy = object()
|
||||||
|
service = _mock_service()
|
||||||
|
fixed_now = datetime.datetime(2024, 8, 20, 12, 0, 0)
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("commands.create_message_clean_policy", return_value=policy),
|
||||||
|
patch("commands.MessagesCleanService.from_time_range", return_value=service) as mock_from_time_range,
|
||||||
|
patch("commands.MessagesCleanService.from_days") as mock_from_days,
|
||||||
|
patch("commands.datetime", autospec=True) as mock_datetime,
|
||||||
|
):
|
||||||
|
mock_datetime.datetime.now.return_value = fixed_now
|
||||||
|
mock_datetime.timedelta = datetime.timedelta
|
||||||
|
|
||||||
|
clean_expired_messages.callback(
|
||||||
|
batch_size=1000,
|
||||||
|
graceful_period=21,
|
||||||
|
start_from=None,
|
||||||
|
end_before=None,
|
||||||
|
from_days_ago=60,
|
||||||
|
before_days=30,
|
||||||
|
dry_run=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_from_time_range.assert_called_once_with(
|
||||||
|
policy=policy,
|
||||||
|
start_from=fixed_now - datetime.timedelta(days=60),
|
||||||
|
end_before=fixed_now - datetime.timedelta(days=30),
|
||||||
|
batch_size=1000,
|
||||||
|
dry_run=False,
|
||||||
|
)
|
||||||
|
mock_from_days.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("kwargs", "message"),
|
||||||
|
[
|
||||||
|
(
|
||||||
|
{
|
||||||
|
"start_from": datetime.datetime(2024, 1, 1),
|
||||||
|
"end_before": datetime.datetime(2024, 2, 1),
|
||||||
|
"from_days_ago": None,
|
||||||
|
"before_days": 30,
|
||||||
|
},
|
||||||
|
"mutually exclusive",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
{
|
||||||
|
"start_from": datetime.datetime(2024, 1, 1),
|
||||||
|
"end_before": None,
|
||||||
|
"from_days_ago": None,
|
||||||
|
"before_days": None,
|
||||||
|
},
|
||||||
|
"Both --start-from and --end-before are required",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
{
|
||||||
|
"start_from": None,
|
||||||
|
"end_before": None,
|
||||||
|
"from_days_ago": 10,
|
||||||
|
"before_days": None,
|
||||||
|
},
|
||||||
|
"--from-days-ago must be used together with --before-days",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
{
|
||||||
|
"start_from": None,
|
||||||
|
"end_before": None,
|
||||||
|
"from_days_ago": None,
|
||||||
|
"before_days": -1,
|
||||||
|
},
|
||||||
|
"--before-days must be >= 0",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
{
|
||||||
|
"start_from": None,
|
||||||
|
"end_before": None,
|
||||||
|
"from_days_ago": 30,
|
||||||
|
"before_days": 30,
|
||||||
|
},
|
||||||
|
"--from-days-ago must be greater than --before-days",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
{
|
||||||
|
"start_from": None,
|
||||||
|
"end_before": None,
|
||||||
|
"from_days_ago": None,
|
||||||
|
"before_days": None,
|
||||||
|
},
|
||||||
|
"You must provide either (--start-from,--end-before) or (--before-days [--from-days-ago])",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_invalid_inputs_raise_usage_error(kwargs: dict, message: str):
|
||||||
|
with pytest.raises(click.UsageError, match=re.escape(message)):
|
||||||
|
clean_expired_messages.callback(
|
||||||
|
batch_size=1000,
|
||||||
|
graceful_period=21,
|
||||||
|
start_from=kwargs["start_from"],
|
||||||
|
end_before=kwargs["end_before"],
|
||||||
|
from_days_ago=kwargs["from_days_ago"],
|
||||||
|
before_days=kwargs["before_days"],
|
||||||
|
dry_run=False,
|
||||||
|
)
|
||||||
Reference in New Issue
Block a user