feat: plugin auto upgrade strategy (#19758)
Co-authored-by: Joel <iamjoel007@gmail.com> Co-authored-by: crazywoola <427733928@qq.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Novice <novice12185727@gmail.com>
This commit is contained in:

committed by
GitHub

parent
e6913744ae
commit
eaae79a581
@@ -471,6 +471,16 @@ APP_MAX_ACTIVE_REQUESTS=0
|
||||
# Celery beat configuration
|
||||
CELERY_BEAT_SCHEDULER_TIME=1
|
||||
|
||||
# Celery schedule tasks configuration
|
||||
ENABLE_CLEAN_EMBEDDING_CACHE_TASK=false
|
||||
ENABLE_CLEAN_UNUSED_DATASETS_TASK=false
|
||||
ENABLE_CREATE_TIDB_SERVERLESS_TASK=false
|
||||
ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK=false
|
||||
ENABLE_CLEAN_MESSAGES=false
|
||||
ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK=false
|
||||
ENABLE_DATASETS_QUEUE_MONITOR=false
|
||||
ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK=true
|
||||
|
||||
# Position configuration
|
||||
POSITION_TOOL_PINS=
|
||||
POSITION_TOOL_INCLUDES=
|
||||
|
@@ -74,7 +74,12 @@
|
||||
10. If you need to handle and debug the async tasks (e.g. dataset importing and documents indexing), please start the worker service.
|
||||
|
||||
```bash
|
||||
uv run celery -A app.celery worker -P gevent -c 1 --loglevel INFO -Q dataset,generation,mail,ops_trace,app_deletion
|
||||
uv run celery -A app.celery worker -P gevent -c 1 --loglevel INFO -Q dataset,generation,mail,ops_trace,app_deletion,plugin
|
||||
```
|
||||
|
||||
Addition, if you want to debug the celery scheduled tasks, you can use the following command in another terminal:
|
||||
```bash
|
||||
uv run celery -A app.celery beat
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
@@ -832,6 +832,41 @@ class CeleryBeatConfig(BaseSettings):
|
||||
)
|
||||
|
||||
|
||||
class CeleryScheduleTasksConfig(BaseSettings):
|
||||
ENABLE_CLEAN_EMBEDDING_CACHE_TASK: bool = Field(
|
||||
description="Enable clean embedding cache task",
|
||||
default=False,
|
||||
)
|
||||
ENABLE_CLEAN_UNUSED_DATASETS_TASK: bool = Field(
|
||||
description="Enable clean unused datasets task",
|
||||
default=False,
|
||||
)
|
||||
ENABLE_CREATE_TIDB_SERVERLESS_TASK: bool = Field(
|
||||
description="Enable create tidb service job task",
|
||||
default=False,
|
||||
)
|
||||
ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK: bool = Field(
|
||||
description="Enable update tidb service job status task",
|
||||
default=False,
|
||||
)
|
||||
ENABLE_CLEAN_MESSAGES: bool = Field(
|
||||
description="Enable clean messages task",
|
||||
default=False,
|
||||
)
|
||||
ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK: bool = Field(
|
||||
description="Enable mail clean document notify task",
|
||||
default=False,
|
||||
)
|
||||
ENABLE_DATASETS_QUEUE_MONITOR: bool = Field(
|
||||
description="Enable queue monitor task",
|
||||
default=False,
|
||||
)
|
||||
ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK: bool = Field(
|
||||
description="Enable check upgradable plugin task",
|
||||
default=True,
|
||||
)
|
||||
|
||||
|
||||
class PositionConfig(BaseSettings):
|
||||
POSITION_PROVIDER_PINS: str = Field(
|
||||
description="Comma-separated list of pinned model providers",
|
||||
@@ -961,5 +996,6 @@ class FeatureConfig(
|
||||
# hosted services config
|
||||
HostedServiceConfig,
|
||||
CeleryBeatConfig,
|
||||
CeleryScheduleTasksConfig,
|
||||
):
|
||||
pass
|
||||
|
@@ -12,7 +12,8 @@ from controllers.console.wraps import account_initialization_required, setup_req
|
||||
from core.model_runtime.utils.encoders import jsonable_encoder
|
||||
from core.plugin.impl.exc import PluginDaemonClientSideError
|
||||
from libs.login import login_required
|
||||
from models.account import TenantPluginPermission
|
||||
from models.account import TenantPluginAutoUpgradeStrategy, TenantPluginPermission
|
||||
from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
|
||||
from services.plugin.plugin_parameter_service import PluginParameterService
|
||||
from services.plugin.plugin_permission_service import PluginPermissionService
|
||||
from services.plugin.plugin_service import PluginService
|
||||
@@ -534,6 +535,114 @@ class PluginFetchDynamicSelectOptionsApi(Resource):
|
||||
return jsonable_encoder({"options": options})
|
||||
|
||||
|
||||
class PluginChangePreferencesApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
def post(self):
|
||||
user = current_user
|
||||
if not user.is_admin_or_owner:
|
||||
raise Forbidden()
|
||||
|
||||
req = reqparse.RequestParser()
|
||||
req.add_argument("permission", type=dict, required=True, location="json")
|
||||
req.add_argument("auto_upgrade", type=dict, required=True, location="json")
|
||||
args = req.parse_args()
|
||||
|
||||
tenant_id = user.current_tenant_id
|
||||
|
||||
permission = args["permission"]
|
||||
|
||||
install_permission = TenantPluginPermission.InstallPermission(permission.get("install_permission", "everyone"))
|
||||
debug_permission = TenantPluginPermission.DebugPermission(permission.get("debug_permission", "everyone"))
|
||||
|
||||
auto_upgrade = args["auto_upgrade"]
|
||||
|
||||
strategy_setting = TenantPluginAutoUpgradeStrategy.StrategySetting(
|
||||
auto_upgrade.get("strategy_setting", "fix_only")
|
||||
)
|
||||
upgrade_time_of_day = auto_upgrade.get("upgrade_time_of_day", 0)
|
||||
upgrade_mode = TenantPluginAutoUpgradeStrategy.UpgradeMode(auto_upgrade.get("upgrade_mode", "exclude"))
|
||||
exclude_plugins = auto_upgrade.get("exclude_plugins", [])
|
||||
include_plugins = auto_upgrade.get("include_plugins", [])
|
||||
|
||||
# set permission
|
||||
set_permission_result = PluginPermissionService.change_permission(
|
||||
tenant_id,
|
||||
install_permission,
|
||||
debug_permission,
|
||||
)
|
||||
if not set_permission_result:
|
||||
return jsonable_encoder({"success": False, "message": "Failed to set permission"})
|
||||
|
||||
# set auto upgrade strategy
|
||||
set_auto_upgrade_strategy_result = PluginAutoUpgradeService.change_strategy(
|
||||
tenant_id,
|
||||
strategy_setting,
|
||||
upgrade_time_of_day,
|
||||
upgrade_mode,
|
||||
exclude_plugins,
|
||||
include_plugins,
|
||||
)
|
||||
if not set_auto_upgrade_strategy_result:
|
||||
return jsonable_encoder({"success": False, "message": "Failed to set auto upgrade strategy"})
|
||||
|
||||
return jsonable_encoder({"success": True})
|
||||
|
||||
|
||||
class PluginFetchPreferencesApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
def get(self):
|
||||
tenant_id = current_user.current_tenant_id
|
||||
|
||||
permission = PluginPermissionService.get_permission(tenant_id)
|
||||
permission_dict = {
|
||||
"install_permission": TenantPluginPermission.InstallPermission.EVERYONE,
|
||||
"debug_permission": TenantPluginPermission.DebugPermission.EVERYONE,
|
||||
}
|
||||
|
||||
if permission:
|
||||
permission_dict["install_permission"] = permission.install_permission
|
||||
permission_dict["debug_permission"] = permission.debug_permission
|
||||
|
||||
auto_upgrade = PluginAutoUpgradeService.get_strategy(tenant_id)
|
||||
auto_upgrade_dict = {
|
||||
"strategy_setting": TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED,
|
||||
"upgrade_time_of_day": 0,
|
||||
"upgrade_mode": TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
|
||||
"exclude_plugins": [],
|
||||
"include_plugins": [],
|
||||
}
|
||||
|
||||
if auto_upgrade:
|
||||
auto_upgrade_dict = {
|
||||
"strategy_setting": auto_upgrade.strategy_setting,
|
||||
"upgrade_time_of_day": auto_upgrade.upgrade_time_of_day,
|
||||
"upgrade_mode": auto_upgrade.upgrade_mode,
|
||||
"exclude_plugins": auto_upgrade.exclude_plugins,
|
||||
"include_plugins": auto_upgrade.include_plugins,
|
||||
}
|
||||
|
||||
return jsonable_encoder({"permission": permission_dict, "auto_upgrade": auto_upgrade_dict})
|
||||
|
||||
|
||||
class PluginAutoUpgradeExcludePluginApi(Resource):
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
def post(self):
|
||||
# exclude one single plugin
|
||||
tenant_id = current_user.current_tenant_id
|
||||
|
||||
req = reqparse.RequestParser()
|
||||
req.add_argument("plugin_id", type=str, required=True, location="json")
|
||||
args = req.parse_args()
|
||||
|
||||
return jsonable_encoder({"success": PluginAutoUpgradeService.exclude_plugin(tenant_id, args["plugin_id"])})
|
||||
|
||||
|
||||
api.add_resource(PluginDebuggingKeyApi, "/workspaces/current/plugin/debugging-key")
|
||||
api.add_resource(PluginListApi, "/workspaces/current/plugin/list")
|
||||
api.add_resource(PluginListLatestVersionsApi, "/workspaces/current/plugin/list/latest-versions")
|
||||
@@ -560,3 +669,7 @@ api.add_resource(PluginChangePermissionApi, "/workspaces/current/plugin/permissi
|
||||
api.add_resource(PluginFetchPermissionApi, "/workspaces/current/plugin/permission/fetch")
|
||||
|
||||
api.add_resource(PluginFetchDynamicSelectOptionsApi, "/workspaces/current/plugin/parameters/dynamic-options")
|
||||
|
||||
api.add_resource(PluginFetchPreferencesApi, "/workspaces/current/plugin/preferences/fetch")
|
||||
api.add_resource(PluginChangePreferencesApi, "/workspaces/current/plugin/preferences/change")
|
||||
api.add_resource(PluginAutoUpgradeExcludePluginApi, "/workspaces/current/plugin/preferences/autoupgrade/exclude")
|
||||
|
@@ -25,9 +25,29 @@ def batch_fetch_plugin_manifests(plugin_ids: list[str]) -> Sequence[MarketplaceP
|
||||
url = str(marketplace_api_url / "api/v1/plugins/batch")
|
||||
response = requests.post(url, json={"plugin_ids": plugin_ids})
|
||||
response.raise_for_status()
|
||||
|
||||
return [MarketplacePluginDeclaration(**plugin) for plugin in response.json()["data"]["plugins"]]
|
||||
|
||||
|
||||
def batch_fetch_plugin_manifests_ignore_deserialization_error(
|
||||
plugin_ids: list[str],
|
||||
) -> Sequence[MarketplacePluginDeclaration]:
|
||||
if len(plugin_ids) == 0:
|
||||
return []
|
||||
|
||||
url = str(marketplace_api_url / "api/v1/plugins/batch")
|
||||
response = requests.post(url, json={"plugin_ids": plugin_ids})
|
||||
response.raise_for_status()
|
||||
result: list[MarketplacePluginDeclaration] = []
|
||||
for plugin in response.json()["data"]["plugins"]:
|
||||
try:
|
||||
result.append(MarketplacePluginDeclaration(**plugin))
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def record_install_plugin_event(plugin_unique_identifier: str):
|
||||
url = str(marketplace_api_url / "api/v1/stats/plugins/install_count")
|
||||
response = requests.post(url, json={"unique_identifier": plugin_unique_identifier})
|
||||
|
@@ -22,7 +22,7 @@ if [[ "${MODE}" == "worker" ]]; then
|
||||
|
||||
exec celery -A app.celery worker -P ${CELERY_WORKER_CLASS:-gevent} $CONCURRENCY_OPTION \
|
||||
--max-tasks-per-child ${MAX_TASK_PRE_CHILD:-50} --loglevel ${LOG_LEVEL:-INFO} \
|
||||
-Q ${CELERY_QUEUES:-dataset,mail,ops_trace,app_deletion}
|
||||
-Q ${CELERY_QUEUES:-dataset,mail,ops_trace,app_deletion,plugin}
|
||||
|
||||
elif [[ "${MODE}" == "beat" ]]; then
|
||||
exec celery -A app.celery beat --loglevel ${LOG_LEVEL:-INFO}
|
||||
|
@@ -64,49 +64,62 @@ def init_app(app: DifyApp) -> Celery:
|
||||
celery_app.set_default()
|
||||
app.extensions["celery"] = celery_app
|
||||
|
||||
imports = [
|
||||
"schedule.clean_embedding_cache_task",
|
||||
"schedule.clean_unused_datasets_task",
|
||||
"schedule.create_tidb_serverless_task",
|
||||
"schedule.update_tidb_serverless_status_task",
|
||||
"schedule.clean_messages",
|
||||
"schedule.mail_clean_document_notify_task",
|
||||
"schedule.queue_monitor_task",
|
||||
]
|
||||
imports = []
|
||||
day = dify_config.CELERY_BEAT_SCHEDULER_TIME
|
||||
beat_schedule = {
|
||||
"clean_embedding_cache_task": {
|
||||
|
||||
# if you add a new task, please add the switch to CeleryScheduleTasksConfig
|
||||
beat_schedule = {}
|
||||
if dify_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK:
|
||||
imports.append("schedule.clean_embedding_cache_task")
|
||||
beat_schedule["clean_embedding_cache_task"] = {
|
||||
"task": "schedule.clean_embedding_cache_task.clean_embedding_cache_task",
|
||||
"schedule": timedelta(days=day),
|
||||
},
|
||||
"clean_unused_datasets_task": {
|
||||
}
|
||||
if dify_config.ENABLE_CLEAN_UNUSED_DATASETS_TASK:
|
||||
imports.append("schedule.clean_unused_datasets_task")
|
||||
beat_schedule["clean_unused_datasets_task"] = {
|
||||
"task": "schedule.clean_unused_datasets_task.clean_unused_datasets_task",
|
||||
"schedule": timedelta(days=day),
|
||||
},
|
||||
"create_tidb_serverless_task": {
|
||||
}
|
||||
if dify_config.ENABLE_CREATE_TIDB_SERVERLESS_TASK:
|
||||
imports.append("schedule.create_tidb_serverless_task")
|
||||
beat_schedule["create_tidb_serverless_task"] = {
|
||||
"task": "schedule.create_tidb_serverless_task.create_tidb_serverless_task",
|
||||
"schedule": crontab(minute="0", hour="*"),
|
||||
},
|
||||
"update_tidb_serverless_status_task": {
|
||||
}
|
||||
if dify_config.ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK:
|
||||
imports.append("schedule.update_tidb_serverless_status_task")
|
||||
beat_schedule["update_tidb_serverless_status_task"] = {
|
||||
"task": "schedule.update_tidb_serverless_status_task.update_tidb_serverless_status_task",
|
||||
"schedule": timedelta(minutes=10),
|
||||
},
|
||||
"clean_messages": {
|
||||
}
|
||||
if dify_config.ENABLE_CLEAN_MESSAGES:
|
||||
imports.append("schedule.clean_messages")
|
||||
beat_schedule["clean_messages"] = {
|
||||
"task": "schedule.clean_messages.clean_messages",
|
||||
"schedule": timedelta(days=day),
|
||||
},
|
||||
# every Monday
|
||||
"mail_clean_document_notify_task": {
|
||||
}
|
||||
if dify_config.ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK:
|
||||
imports.append("schedule.mail_clean_document_notify_task")
|
||||
beat_schedule["mail_clean_document_notify_task"] = {
|
||||
"task": "schedule.mail_clean_document_notify_task.mail_clean_document_notify_task",
|
||||
"schedule": crontab(minute="0", hour="10", day_of_week="1"),
|
||||
},
|
||||
"datasets-queue-monitor": {
|
||||
}
|
||||
if dify_config.ENABLE_DATASETS_QUEUE_MONITOR:
|
||||
imports.append("schedule.queue_monitor_task")
|
||||
beat_schedule["datasets-queue-monitor"] = {
|
||||
"task": "schedule.queue_monitor_task.queue_monitor_task",
|
||||
"schedule": timedelta(
|
||||
minutes=dify_config.QUEUE_MONITOR_INTERVAL if dify_config.QUEUE_MONITOR_INTERVAL else 30
|
||||
),
|
||||
},
|
||||
}
|
||||
}
|
||||
if dify_config.ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK:
|
||||
imports.append("schedule.check_upgradable_plugin_task")
|
||||
beat_schedule["check_upgradable_plugin_task"] = {
|
||||
"task": "schedule.check_upgradable_plugin_task.check_upgradable_plugin_task",
|
||||
"schedule": crontab(minute="*/15"),
|
||||
}
|
||||
|
||||
celery_app.conf.update(beat_schedule=beat_schedule, imports=imports)
|
||||
|
||||
return celery_app
|
||||
|
@@ -0,0 +1,42 @@
|
||||
"""add_tenant_plugin_autoupgrade_table
|
||||
|
||||
Revision ID: 8bcc02c9bd07
|
||||
Revises: 375fe79ead14
|
||||
Create Date: 2025-07-23 15:08:50.161441
|
||||
|
||||
"""
|
||||
from alembic import op
|
||||
import models as models
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = '8bcc02c9bd07'
|
||||
down_revision = '375fe79ead14'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table('tenant_plugin_auto_upgrade_strategies',
|
||||
sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
|
||||
sa.Column('tenant_id', models.types.StringUUID(), nullable=False),
|
||||
sa.Column('strategy_setting', sa.String(length=16), server_default='fix_only', nullable=False),
|
||||
sa.Column('upgrade_time_of_day', sa.Integer(), nullable=False),
|
||||
sa.Column('upgrade_mode', sa.String(length=16), server_default='exclude', nullable=False),
|
||||
sa.Column('exclude_plugins', sa.ARRAY(sa.String(length=255)), nullable=False),
|
||||
sa.Column('include_plugins', sa.ARRAY(sa.String(length=255)), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False),
|
||||
sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False),
|
||||
sa.PrimaryKeyConstraint('id', name='tenant_plugin_auto_upgrade_strategy_pkey'),
|
||||
sa.UniqueConstraint('tenant_id', name='unique_tenant_plugin_auto_upgrade_strategy')
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade():
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
|
||||
op.drop_table('tenant_plugin_auto_upgrade_strategies')
|
||||
# ### end Alembic commands ###
|
@@ -297,6 +297,40 @@ class TenantPluginPermission(Base):
|
||||
)
|
||||
|
||||
id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
|
||||
tenant_id: Mapped[str] = mapped_column(StringUUID)
|
||||
install_permission: Mapped[InstallPermission] = mapped_column(db.String(16), server_default="everyone")
|
||||
debug_permission: Mapped[DebugPermission] = mapped_column(db.String(16), server_default="noone")
|
||||
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
install_permission: Mapped[InstallPermission] = mapped_column(
|
||||
db.String(16), nullable=False, server_default="everyone"
|
||||
)
|
||||
debug_permission: Mapped[DebugPermission] = mapped_column(db.String(16), nullable=False, server_default="noone")
|
||||
|
||||
|
||||
class TenantPluginAutoUpgradeStrategy(Base):
|
||||
class StrategySetting(enum.StrEnum):
|
||||
DISABLED = "disabled"
|
||||
FIX_ONLY = "fix_only"
|
||||
LATEST = "latest"
|
||||
|
||||
class UpgradeMode(enum.StrEnum):
|
||||
ALL = "all"
|
||||
PARTIAL = "partial"
|
||||
EXCLUDE = "exclude"
|
||||
|
||||
__tablename__ = "tenant_plugin_auto_upgrade_strategies"
|
||||
__table_args__ = (
|
||||
db.PrimaryKeyConstraint("id", name="tenant_plugin_auto_upgrade_strategy_pkey"),
|
||||
db.UniqueConstraint("tenant_id", name="unique_tenant_plugin_auto_upgrade_strategy"),
|
||||
)
|
||||
|
||||
id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
|
||||
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
|
||||
strategy_setting: Mapped[StrategySetting] = mapped_column(db.String(16), nullable=False, server_default="fix_only")
|
||||
upgrade_time_of_day: Mapped[int] = mapped_column(db.Integer, nullable=False, default=0) # seconds of the day
|
||||
upgrade_mode: Mapped[UpgradeMode] = mapped_column(db.String(16), nullable=False, server_default="exclude")
|
||||
exclude_plugins: Mapped[list[str]] = mapped_column(
|
||||
db.ARRAY(db.String(255)), nullable=False
|
||||
) # plugin_id (author/name)
|
||||
include_plugins: Mapped[list[str]] = mapped_column(
|
||||
db.ARRAY(db.String(255)), nullable=False
|
||||
) # plugin_id (author/name)
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
|
||||
|
49
api/schedule/check_upgradable_plugin_task.py
Normal file
49
api/schedule/check_upgradable_plugin_task.py
Normal file
@@ -0,0 +1,49 @@
|
||||
import time
|
||||
|
||||
import click
|
||||
|
||||
import app
|
||||
from extensions.ext_database import db
|
||||
from models.account import TenantPluginAutoUpgradeStrategy
|
||||
from tasks.process_tenant_plugin_autoupgrade_check_task import process_tenant_plugin_autoupgrade_check_task
|
||||
|
||||
AUTO_UPGRADE_MINIMAL_CHECKING_INTERVAL = 15 * 60 # 15 minutes
|
||||
|
||||
|
||||
@app.celery.task(queue="plugin")
|
||||
def check_upgradable_plugin_task():
|
||||
click.echo(click.style("Start check upgradable plugin.", fg="green"))
|
||||
start_at = time.perf_counter()
|
||||
|
||||
now_seconds_of_day = time.time() % 86400 - 30 # we assume the tz is UTC
|
||||
click.echo(click.style("Now seconds of day: {}".format(now_seconds_of_day), fg="green"))
|
||||
|
||||
strategies = (
|
||||
db.session.query(TenantPluginAutoUpgradeStrategy)
|
||||
.filter(
|
||||
TenantPluginAutoUpgradeStrategy.upgrade_time_of_day >= now_seconds_of_day,
|
||||
TenantPluginAutoUpgradeStrategy.upgrade_time_of_day
|
||||
< now_seconds_of_day + AUTO_UPGRADE_MINIMAL_CHECKING_INTERVAL,
|
||||
TenantPluginAutoUpgradeStrategy.strategy_setting
|
||||
!= TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
for strategy in strategies:
|
||||
process_tenant_plugin_autoupgrade_check_task.delay(
|
||||
strategy.tenant_id,
|
||||
strategy.strategy_setting,
|
||||
strategy.upgrade_time_of_day,
|
||||
strategy.upgrade_mode,
|
||||
strategy.exclude_plugins,
|
||||
strategy.include_plugins,
|
||||
)
|
||||
|
||||
end_at = time.perf_counter()
|
||||
click.echo(
|
||||
click.style(
|
||||
"Checked upgradable plugin success latency: {}".format(end_at - start_at),
|
||||
fg="green",
|
||||
)
|
||||
)
|
@@ -29,6 +29,7 @@ from models.account import (
|
||||
Tenant,
|
||||
TenantAccountJoin,
|
||||
TenantAccountRole,
|
||||
TenantPluginAutoUpgradeStrategy,
|
||||
TenantStatus,
|
||||
)
|
||||
from models.model import DifySetup
|
||||
@@ -828,6 +829,17 @@ class TenantService:
|
||||
db.session.add(tenant)
|
||||
db.session.commit()
|
||||
|
||||
plugin_upgrade_strategy = TenantPluginAutoUpgradeStrategy(
|
||||
tenant_id=tenant.id,
|
||||
strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
|
||||
upgrade_time_of_day=0,
|
||||
upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
|
||||
exclude_plugins=[],
|
||||
include_plugins=[],
|
||||
)
|
||||
db.session.add(plugin_upgrade_strategy)
|
||||
db.session.commit()
|
||||
|
||||
tenant.encrypt_public_key = generate_key_pair(tenant.id)
|
||||
db.session.commit()
|
||||
return tenant
|
||||
|
87
api/services/plugin/plugin_auto_upgrade_service.py
Normal file
87
api/services/plugin/plugin_auto_upgrade_service.py
Normal file
@@ -0,0 +1,87 @@
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from extensions.ext_database import db
|
||||
from models.account import TenantPluginAutoUpgradeStrategy
|
||||
|
||||
|
||||
class PluginAutoUpgradeService:
|
||||
@staticmethod
|
||||
def get_strategy(tenant_id: str) -> TenantPluginAutoUpgradeStrategy | None:
|
||||
with Session(db.engine) as session:
|
||||
return (
|
||||
session.query(TenantPluginAutoUpgradeStrategy)
|
||||
.filter(TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id)
|
||||
.first()
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def change_strategy(
|
||||
tenant_id: str,
|
||||
strategy_setting: TenantPluginAutoUpgradeStrategy.StrategySetting,
|
||||
upgrade_time_of_day: int,
|
||||
upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode,
|
||||
exclude_plugins: list[str],
|
||||
include_plugins: list[str],
|
||||
) -> bool:
|
||||
with Session(db.engine) as session:
|
||||
exist_strategy = (
|
||||
session.query(TenantPluginAutoUpgradeStrategy)
|
||||
.filter(TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id)
|
||||
.first()
|
||||
)
|
||||
if not exist_strategy:
|
||||
strategy = TenantPluginAutoUpgradeStrategy(
|
||||
tenant_id=tenant_id,
|
||||
strategy_setting=strategy_setting,
|
||||
upgrade_time_of_day=upgrade_time_of_day,
|
||||
upgrade_mode=upgrade_mode,
|
||||
exclude_plugins=exclude_plugins,
|
||||
include_plugins=include_plugins,
|
||||
)
|
||||
session.add(strategy)
|
||||
else:
|
||||
exist_strategy.strategy_setting = strategy_setting
|
||||
exist_strategy.upgrade_time_of_day = upgrade_time_of_day
|
||||
exist_strategy.upgrade_mode = upgrade_mode
|
||||
exist_strategy.exclude_plugins = exclude_plugins
|
||||
exist_strategy.include_plugins = include_plugins
|
||||
|
||||
session.commit()
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def exclude_plugin(tenant_id: str, plugin_id: str) -> bool:
|
||||
with Session(db.engine) as session:
|
||||
exist_strategy = (
|
||||
session.query(TenantPluginAutoUpgradeStrategy)
|
||||
.filter(TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id)
|
||||
.first()
|
||||
)
|
||||
if not exist_strategy:
|
||||
# create for this tenant
|
||||
PluginAutoUpgradeService.change_strategy(
|
||||
tenant_id,
|
||||
TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
|
||||
0,
|
||||
TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
|
||||
[plugin_id],
|
||||
[],
|
||||
)
|
||||
return True
|
||||
else:
|
||||
if exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE:
|
||||
if plugin_id not in exist_strategy.exclude_plugins:
|
||||
new_exclude_plugins = exist_strategy.exclude_plugins.copy()
|
||||
new_exclude_plugins.append(plugin_id)
|
||||
exist_strategy.exclude_plugins = new_exclude_plugins
|
||||
elif exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL:
|
||||
if plugin_id in exist_strategy.include_plugins:
|
||||
new_include_plugins = exist_strategy.include_plugins.copy()
|
||||
new_include_plugins.remove(plugin_id)
|
||||
exist_strategy.include_plugins = new_include_plugins
|
||||
elif exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL:
|
||||
exist_strategy.upgrade_mode = TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE
|
||||
exist_strategy.exclude_plugins = [plugin_id]
|
||||
|
||||
session.commit()
|
||||
return True
|
166
api/tasks/process_tenant_plugin_autoupgrade_check_task.py
Normal file
166
api/tasks/process_tenant_plugin_autoupgrade_check_task.py
Normal file
@@ -0,0 +1,166 @@
|
||||
import traceback
|
||||
import typing
|
||||
|
||||
import click
|
||||
from celery import shared_task # type: ignore
|
||||
|
||||
from core.helper import marketplace
|
||||
from core.helper.marketplace import MarketplacePluginDeclaration
|
||||
from core.plugin.entities.plugin import PluginInstallationSource
|
||||
from core.plugin.impl.plugin import PluginInstaller
|
||||
from models.account import TenantPluginAutoUpgradeStrategy
|
||||
|
||||
RETRY_TIMES_OF_ONE_PLUGIN_IN_ONE_TENANT = 3
|
||||
|
||||
|
||||
cached_plugin_manifests: dict[str, typing.Union[MarketplacePluginDeclaration, None]] = {}
|
||||
|
||||
|
||||
def marketplace_batch_fetch_plugin_manifests(
|
||||
plugin_ids_plain_list: list[str],
|
||||
) -> list[MarketplacePluginDeclaration]:
|
||||
global cached_plugin_manifests
|
||||
# return marketplace.batch_fetch_plugin_manifests(plugin_ids_plain_list)
|
||||
not_included_plugin_ids = [
|
||||
plugin_id for plugin_id in plugin_ids_plain_list if plugin_id not in cached_plugin_manifests
|
||||
]
|
||||
if not_included_plugin_ids:
|
||||
manifests = marketplace.batch_fetch_plugin_manifests_ignore_deserialization_error(not_included_plugin_ids)
|
||||
for manifest in manifests:
|
||||
cached_plugin_manifests[manifest.plugin_id] = manifest
|
||||
|
||||
if (
|
||||
len(manifests) == 0
|
||||
): # this indicates that the plugin not found in marketplace, should set None in cache to prevent future check
|
||||
for plugin_id in not_included_plugin_ids:
|
||||
cached_plugin_manifests[plugin_id] = None
|
||||
|
||||
result: list[MarketplacePluginDeclaration] = []
|
||||
for plugin_id in plugin_ids_plain_list:
|
||||
final_manifest = cached_plugin_manifests.get(plugin_id)
|
||||
if final_manifest is not None:
|
||||
result.append(final_manifest)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@shared_task(queue="plugin")
|
||||
def process_tenant_plugin_autoupgrade_check_task(
|
||||
tenant_id: str,
|
||||
strategy_setting: TenantPluginAutoUpgradeStrategy.StrategySetting,
|
||||
upgrade_time_of_day: int,
|
||||
upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode,
|
||||
exclude_plugins: list[str],
|
||||
include_plugins: list[str],
|
||||
):
|
||||
try:
|
||||
manager = PluginInstaller()
|
||||
|
||||
click.echo(
|
||||
click.style(
|
||||
"Checking upgradable plugin for tenant: {}".format(tenant_id),
|
||||
fg="green",
|
||||
)
|
||||
)
|
||||
|
||||
if strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED:
|
||||
return
|
||||
|
||||
# get plugin_ids to check
|
||||
plugin_ids: list[tuple[str, str, str]] = [] # plugin_id, version, unique_identifier
|
||||
click.echo(click.style("Upgrade mode: {}".format(upgrade_mode), fg="green"))
|
||||
|
||||
if upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL and include_plugins:
|
||||
all_plugins = manager.list_plugins(tenant_id)
|
||||
|
||||
for plugin in all_plugins:
|
||||
if plugin.source == PluginInstallationSource.Marketplace and plugin.plugin_id in include_plugins:
|
||||
plugin_ids.append(
|
||||
(
|
||||
plugin.plugin_id,
|
||||
plugin.version,
|
||||
plugin.plugin_unique_identifier,
|
||||
)
|
||||
)
|
||||
|
||||
elif upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE:
|
||||
# get all plugins and remove excluded plugins
|
||||
all_plugins = manager.list_plugins(tenant_id)
|
||||
plugin_ids = [
|
||||
(plugin.plugin_id, plugin.version, plugin.plugin_unique_identifier)
|
||||
for plugin in all_plugins
|
||||
if plugin.source == PluginInstallationSource.Marketplace and plugin.plugin_id not in exclude_plugins
|
||||
]
|
||||
elif upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL:
|
||||
all_plugins = manager.list_plugins(tenant_id)
|
||||
plugin_ids = [
|
||||
(plugin.plugin_id, plugin.version, plugin.plugin_unique_identifier)
|
||||
for plugin in all_plugins
|
||||
if plugin.source == PluginInstallationSource.Marketplace
|
||||
]
|
||||
|
||||
if not plugin_ids:
|
||||
return
|
||||
|
||||
plugin_ids_plain_list = [plugin_id for plugin_id, _, _ in plugin_ids]
|
||||
|
||||
manifests = marketplace_batch_fetch_plugin_manifests(plugin_ids_plain_list)
|
||||
|
||||
if not manifests:
|
||||
return
|
||||
|
||||
for manifest in manifests:
|
||||
for plugin_id, version, original_unique_identifier in plugin_ids:
|
||||
if manifest.plugin_id != plugin_id:
|
||||
continue
|
||||
|
||||
try:
|
||||
current_version = version
|
||||
latest_version = manifest.latest_version
|
||||
|
||||
def fix_only_checker(latest_version, current_version):
|
||||
latest_version_tuple = tuple(int(val) for val in latest_version.split("."))
|
||||
current_version_tuple = tuple(int(val) for val in current_version.split("."))
|
||||
|
||||
if (
|
||||
latest_version_tuple[0] == current_version_tuple[0]
|
||||
and latest_version_tuple[1] == current_version_tuple[1]
|
||||
):
|
||||
return latest_version_tuple[2] != current_version_tuple[2]
|
||||
return False
|
||||
|
||||
version_checker = {
|
||||
TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST: lambda latest_version,
|
||||
current_version: latest_version != current_version,
|
||||
TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY: fix_only_checker,
|
||||
}
|
||||
|
||||
if version_checker[strategy_setting](latest_version, current_version):
|
||||
# execute upgrade
|
||||
new_unique_identifier = manifest.latest_package_identifier
|
||||
|
||||
marketplace.record_install_plugin_event(new_unique_identifier)
|
||||
click.echo(
|
||||
click.style(
|
||||
"Upgrade plugin: {} -> {}".format(original_unique_identifier, new_unique_identifier),
|
||||
fg="green",
|
||||
)
|
||||
)
|
||||
task_start_resp = manager.upgrade_plugin(
|
||||
tenant_id,
|
||||
original_unique_identifier,
|
||||
new_unique_identifier,
|
||||
PluginInstallationSource.Marketplace,
|
||||
{
|
||||
"plugin_unique_identifier": new_unique_identifier,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
click.echo(click.style("Error when upgrading plugin: {}".format(e), fg="red"))
|
||||
traceback.print_exc()
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
click.echo(click.style("Error when checking upgradable plugin: {}".format(e), fg="red"))
|
||||
traceback.print_exc()
|
||||
return
|
Reference in New Issue
Block a user