diff --git a/api/controllers/console/workspace/plugin.py b/api/controllers/console/workspace/plugin.py index c0a4734828..4f3b56b68b 100644 --- a/api/controllers/console/workspace/plugin.py +++ b/api/controllers/console/workspace/plugin.py @@ -12,8 +12,9 @@ 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 services.plugin.plugin_parameter_service import PluginParameterService +from models.account import TenantPluginAutoUpgradeStrategy, TenantPluginPermission +from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService from services.plugin.plugin_permission_service import PluginPermissionService from services.plugin.plugin_service import PluginService @@ -532,6 +533,116 @@ class PluginFetchDynamicSelectOptionsApi(Resource): raise ValueError(e) 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) + + if not permission: + permission = { + "install_permission": TenantPluginPermission.InstallPermission.EVERYONE, + "debug_permission": TenantPluginPermission.DebugPermission.EVERYONE, + } + else: + permission = { + "install_permission": permission.install_permission, + "debug_permission": permission.debug_permission, + } + + auto_upgrade = PluginAutoUpgradeService.get_strategy(tenant_id) + if not auto_upgrade: + auto_upgrade = { + "strategy_setting": TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + "upgrade_time_of_day": 0, + "upgrade_mode": TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + "exclude_plugins": [], + "include_plugins": [], + } + else: + auto_upgrade = { + "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, "auto_upgrade": auto_upgrade}) + + +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") @@ -560,3 +671,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") diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index 6279b1ad36..be89705858 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -72,6 +72,7 @@ def init_app(app: DifyApp) -> Celery: "schedule.clean_messages", "schedule.mail_clean_document_notify_task", "schedule.queue_monitor_task", + "schedule.check_upgradable_plugin_task", ] day = dify_config.CELERY_BEAT_SCHEDULER_TIME beat_schedule = { @@ -106,6 +107,11 @@ def init_app(app: DifyApp) -> Celery: minutes=dify_config.QUEUE_MONITOR_INTERVAL if dify_config.QUEUE_MONITOR_INTERVAL else 30 ), }, + # every 15 minutes + "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) diff --git a/api/migrations/versions/2025_05_15_1635-16081485540c_.py b/api/migrations/versions/2025_05_15_1635-16081485540c_.py new file mode 100644 index 0000000000..3db4754e2d --- /dev/null +++ b/api/migrations/versions/2025_05_15_1635-16081485540c_.py @@ -0,0 +1,41 @@ +"""empty message + +Revision ID: 16081485540c +Revises: d28f2004b072 +Create Date: 2025-05-15 16:35:39.113777 + +""" +from alembic import op +import models as models +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '16081485540c' +down_revision = 'd28f2004b072' +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 ### diff --git a/api/models/account.py b/api/models/account.py index 7ffeefa980..1fa27f71dd 100644 --- a/api/models/account.py +++ b/api/models/account.py @@ -299,3 +299,35 @@ class TenantPluginPermission(Base): 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()) diff --git a/api/schedule/check_upgradable_plugin_task.py b/api/schedule/check_upgradable_plugin_task.py new file mode 100644 index 0000000000..b709689d27 --- /dev/null +++ b/api/schedule/check_upgradable_plugin_task.py @@ -0,0 +1,153 @@ +import time +import traceback + +import click + +import app +from core.helper import marketplace +from core.plugin.entities.plugin import PluginInstallationSource +from core.plugin.impl.plugin import PluginInstaller +from extensions.ext_database import db +from models.account import TenantPluginAutoUpgradeStrategy + +AUTO_UPGRADE_MINIMAL_CHECKING_INTERVAL = 15 * 60 # 15 minutes + +RETRY_TIMES_OF_ONE_PLUGIN_IN_ONE_TENANT = 3 + + +@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 # we assume the tz is UTC + click.echo(click.style("Now seconds of day: {}".format(now_seconds_of_day), fg="green")) + + # get strategies that set to be performed in the next AUTO_UPGRADE_MINIMAL_CHECKING_INTERVAL + 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() + ) + + manager = PluginInstaller() + + for strategy in strategies: + try: + tenant_id = strategy.tenant_id + strategy_setting = strategy.strategy_setting + upgrade_mode = strategy.upgrade_mode + exclude_plugins = strategy.exclude_plugins + include_plugins = strategy.include_plugins + + if strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED: + continue + + # get plugins that need to be checked + plugin_ids: list[tuple[str, str, str]] = [] # plugin_id, version, unique_identifier + + 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 the exclude 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: + continue + + plugin_ids_plain_list = [plugin_id for plugin_id, _, _ in plugin_ids] + + click.echo(click.style("Fetching manifests for plugins: {}".format(plugin_ids_plain_list), fg="green")) + + # fetch latest versions from marketplace + manifests = marketplace.batch_fetch_plugin_manifests(plugin_ids_plain_list) + + 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 + + # @yeuoly review here + 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() + continue + + end_at = time.perf_counter() + click.echo( + click.style( + "Checked upgradable plugin success latency: {}".format(end_at - start_at), + fg="green", + ) + ) diff --git a/api/services/plugin/plugin_auto_upgrade_service.py b/api/services/plugin/plugin_auto_upgrade_service.py new file mode 100644 index 0000000000..d05292a4bc --- /dev/null +++ b/api/services/plugin/plugin_auto_upgrade_service.py @@ -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], + ) -> None: + 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