feat: crud for auto upgrade strategy

9 months ago · 71b3d6ad9c
parent fccb00c281
commit 71b3d6ad9c
4 changed files with 218 additions and 7 deletions
--- a/api/controllers/console/workspace/plugin.py
+++ b/api/controllers/console/workspace/plugin.py
@ -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
@ -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")
--- a/api/migrations/versions/2025_05_15_1635-16081485540c_.py
+++ 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 ###
--- a/api/models/account.py
+++ b/api/models/account.py
@ -300,8 +300,9 @@ class TenantPluginPermission(Base):
    )
    debug_permission: Mapped[DebugPermission] = mapped_column(db.String(16), nullable=False, server_default="noone")

+
 class TenantPluginAutoUpgradeStrategy(Base):
-    class StrategyOption(enum.StrEnum):
+    class StrategySetting(enum.StrEnum):
        DISABLED = "disabled"
        FIX_ONLY = "fix_only"
        LATEST = "latest"
@ -318,10 +319,14 @@ class TenantPluginAutoUpgradeStrategy(Base):

    id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
-    strategy_option: Mapped[StrategyOption] = mapped_column(db.String(16), nullable=False)
-    upgrade_time_of_day: Mapped[int] = mapped_column(db.Integer, nullable=False)
-    upgrade_mode: Mapped[UpgradeMode] = mapped_column(db.String(16), nullable=False)
-    exclude_plugins: Mapped[list[str]] = mapped_column(db.ARRAY(db.String(255)), nullable=False)
-    include_plugins: Mapped[list[str]] = mapped_column(db.ARRAY(db.String(255)), 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())
--- a/api/services/plugin/plugin_auto_upgrade_service.py
+++ b/api/services/plugin/plugin_auto_upgrade_service.py
@ -0,0 +1,50 @@
+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