diff --git a/api/commands/__init__.py b/api/commands/__init__.py index 58ca5d6ffb3..73009124062 100644 --- a/api/commands/__init__.py +++ b/api/commands/__init__.py @@ -25,6 +25,7 @@ from .plugin import ( from .rbac import ( migrate_dataset_permissions_to_rbac, migrate_member_roles_to_rbac, + migrate_only_me_resource_whitelist_scopes_to_automatic_include, migrate_resource_whitelist_scopes_to_automatic_include, ) from .retention import ( @@ -85,6 +86,7 @@ __all__ = [ "migrate_dataset_permissions_to_rbac", "migrate_knowledge_vector_database", "migrate_member_roles_to_rbac", + "migrate_only_me_resource_whitelist_scopes_to_automatic_include", "migrate_oss", "migrate_resource_whitelist_scopes_to_automatic_include", "migration_data_wizard", diff --git a/api/commands/data_migrate.py b/api/commands/data_migrate.py index 7573bc92ad8..f1a7299b5c3 100644 --- a/api/commands/data_migrate.py +++ b/api/commands/data_migrate.py @@ -7,7 +7,11 @@ from typing import cast import click -from commands.rbac import migrate_dataset_permissions_to_rbac, migrate_resource_whitelist_scopes_to_automatic_include +from commands.rbac import ( + migrate_dataset_permissions_to_rbac, + migrate_only_me_resource_whitelist_scopes_to_automatic_include, + migrate_resource_whitelist_scopes_to_automatic_include, # noqa: F401 +) from extensions.ext_database import db from graphon.model_runtime.entities.model_entities import ModelType from services.legacy_model_type_migration import ( @@ -179,4 +183,4 @@ def legacy_model_types( data_migrate.add_command(legacy_model_types) data_migrate.add_command(migrate_dataset_permissions_to_rbac) -data_migrate.add_command(migrate_resource_whitelist_scopes_to_automatic_include) +data_migrate.add_command(migrate_only_me_resource_whitelist_scopes_to_automatic_include) diff --git a/api/commands/rbac.py b/api/commands/rbac.py index 27c9f02e718..5c10c77a1dd 100644 --- a/api/commands/rbac.py +++ b/api/commands/rbac.py @@ -661,6 +661,149 @@ def migrate_resource_whitelist_scopes_to_automatic_include( ) +@click.command( + "rbac-migrate-resource-whitelist-scopes", + help=( + "Migrate RBAC app/dataset whitelist configs whose old scope is only_me to " + "automatic_include_workspace_members=true and sync workspace members into the whitelist." + ), +) +@click.option("--tenant-id", help="Only migrate resources in a single workspace.") +@click.option( + "--resource-type", + type=click.Choice(["app", "dataset", "all"]), + default="all", + show_default=True, + help="Resource type to migrate.", +) +@click.option("--resource-id", help="Only migrate a single resource. Requires --resource-type app or dataset.") +@click.option("--batch-size", default=500, show_default=True, type=click.IntRange(min=1)) +@click.option( + "--member-batch-size", + default=_RBAC_RESOURCE_ACCESS_POLICY_BATCH_SIZE, + show_default=True, + type=click.IntRange(min=1), + help="Workspace members written per default-policy call.", +) +@click.option( + "--dry-run/--apply", + default=True, + show_default=True, + help="Preview the migration without writing RBAC bindings. Use --apply to write changes.", +) +def migrate_only_me_resource_whitelist_scopes_to_automatic_include( + tenant_id: str | None, + resource_type: str, + resource_id: str | None, + batch_size: int, + member_batch_size: int, + dry_run: bool, +) -> None: + """Backfill automatic workspace-member inclusion for old RBAC only_me resource scopes.""" + if resource_id and resource_type == "all": + raise click.BadParameter("--resource-id requires --resource-type app or dataset", param_hint="--resource-id") + + click.echo(click.style("Starting RBAC only_me resource whitelist scope migration.", fg="green")) + + scanned_count = 0 + only_me_count = 0 + migrated_count = 0 + member_policy_batch_count = 0 + owner_account_ids_by_tenant_id: dict[str, str] = {} + + for ( + current_resource_type, + workspace_id, + current_resource_id, + maintainer_account_id, + ) in _iter_selected_rbac_resource_rows( + resource_type, + tenant_id=tenant_id, + resource_id=resource_id, + batch_size=batch_size, + ): + scanned_count += 1 + with session_factory.create_session() as session: + operator_account_id = maintainer_account_id or owner_account_ids_by_tenant_id.get(workspace_id) + if not operator_account_id: + operator_account_id = _owner_account_id(workspace_id, session=session) + owner_account_ids_by_tenant_id[workspace_id] = operator_account_id + + legacy_config = _resource_legacy_whitelist_config( + current_resource_type, + tenant_id=workspace_id, + operator_account_id=operator_account_id, + resource_id=current_resource_id, + ) + scope = _normalize_rbac_whitelist_scope(legacy_config.rbac_whitelist_scope) + if scope is not RBACResourceWhitelistScope.ONLY_ME: + continue + + only_me_count += 1 + _emit_resource_whitelist_scope_migration_event( + { + "event": "only_me_resource_whitelist_scope_migration_proposed_change", + "dry_run": dry_run, + "tenant_id": workspace_id, + "operator_account_id": operator_account_id, + "resource_type": current_resource_type, + "resource_id": current_resource_id, + "before": { + "rbac_whitelist_scope": scope.value, + "legacy_account_ids": sorted(set(legacy_config.account_ids)), + }, + "after": { + "automatic_include_workspace_members": True, + "default_policy_member_source": "workspace_members", + }, + } + ) + + if dry_run: + continue + + _replace_resource_whitelist( + current_resource_type, + tenant_id=workspace_id, + operator_account_id=operator_account_id, + resource_id=current_resource_id, + automatic_include_workspace_members=True, + ) + migrated_count += 1 + + for batch in _workspace_member_account_id_batches(workspace_id, member_batch_size): + _replace_resource_default_access_policies( + current_resource_type, + tenant_id=workspace_id, + operator_account_id=operator_account_id, + resource_id=current_resource_id, + account_ids=batch, + ) + member_policy_batch_count += 1 + + if scanned_count == 0: + click.echo(click.style("No RBAC resources found for migration.", fg="yellow")) + return + + if dry_run: + click.echo( + click.style( + f"Dry run completed. Scanned {scanned_count} RBAC resources, found {only_me_count} only_me resources. " + "No RBAC bindings were written.", + fg="yellow", + ) + ) + else: + click.echo( + click.style( + "RBAC only_me resource whitelist scope migration completed. " + f"Scanned {scanned_count} resources, migrated {migrated_count}, " + f"wrote {member_policy_batch_count} default-policy batches.", + fg="green", + ) + ) + + @click.command( "rbac-migrate-dataset-permissions", help=( diff --git a/api/controllers/inner_api/workspace/workspace.py b/api/controllers/inner_api/workspace/workspace.py index 7767653202a..e70a517b356 100644 --- a/api/controllers/inner_api/workspace/workspace.py +++ b/api/controllers/inner_api/workspace/workspace.py @@ -11,7 +11,12 @@ from controllers.inner_api.wraps import enterprise_inner_api_only from events.tenant_event import tenant_was_created from extensions.ext_database import db from models import Account -from services.account_service import TenantService +from models.account import TenantAccountRole +from services.account_service import ( + EnterpriseWorkspaceMemberAccountNotFoundError, + EnterpriseWorkspaceMemberWorkspaceNotFoundError, + TenantService, +) class WorkspaceCreatePayload(BaseModel): @@ -23,7 +28,16 @@ class WorkspaceOwnerlessPayload(BaseModel): name: str -register_schema_models(inner_api_ns, WorkspaceCreatePayload, WorkspaceOwnerlessPayload) +class WorkspaceMemberPayload(BaseModel): + workspace_id: str + account_id: str + email: str + role: str = TenantAccountRole.NORMAL.value + current: bool = False + operator_account_id: str | None = None + + +register_schema_models(inner_api_ns, WorkspaceCreatePayload, WorkspaceOwnerlessPayload, WorkspaceMemberPayload) @inner_api_ns.route("/enterprise/workspace") @@ -105,3 +119,51 @@ class EnterpriseWorkspaceNoOwnerEmail(Resource): "message": "enterprise workspace created.", "tenant": resp, } + + +@inner_api_ns.route("/enterprise/workspace/member") +class EnterpriseWorkspaceMember(Resource): + @setup_required + @enterprise_inner_api_only + @inner_api_ns.doc("join_enterprise_workspace_member") + @inner_api_ns.doc(description="Add an existing account to an enterprise workspace") + @inner_api_ns.expect(inner_api_ns.models[WorkspaceMemberPayload.__name__]) + @inner_api_ns.doc( + responses={ + 200: "Workspace member joined successfully", + 400: "Invalid workspace member role", + 401: "Unauthorized - invalid API key", + 404: "Workspace or account not found", + } + ) + def post(self): + args = WorkspaceMemberPayload.model_validate(inner_api_ns.payload or {}) + + try: + role = TenantAccountRole(args.role) + except ValueError: + return {"message": "invalid workspace member role."}, 400 + if role == TenantAccountRole.OWNER: + return {"message": "cannot join workspace as owner."}, 400 + + try: + membership = TenantService.join_enterprise_workspace_member( + workspace_id=args.workspace_id, + account_id=args.account_id, + email=args.email, + role=role, + operator_account_id=args.operator_account_id, + ) + except EnterpriseWorkspaceMemberWorkspaceNotFoundError: + return {"message": "workspace not found."}, 404 + except EnterpriseWorkspaceMemberAccountNotFoundError: + return {"message": "account not found."}, 404 + + return { + "message": "enterprise workspace member joined.", + "member": { + "workspace_id": membership.tenant_id, + "account_id": membership.account_id, + "role": membership.role.value, + }, + } diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index 42927654166..e29b92c8617 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -42,6 +42,7 @@ from services.account_activation_adapters import ( BillingAccountActivationEligibility, BillingWorkspaceMembershipCache, DeploymentWorkspaceInvitePolicy, + RBACWorkspaceMemberAccessSync, RegisterServiceInvitationTokenStore, ) from services.account_activation_service import AccountActivationService @@ -333,6 +334,9 @@ def build_application_services( membership_cache=BillingWorkspaceMembershipCache( enabled=deployment_edition == DeploymentEdition.CLOUD, ), + member_access_sync=RBACWorkspaceMemberAccessSync( + enabled=dify_config.RBAC_ENABLED, + ), ), app_definitions=AppDefinitionQueryService( definitions=app_definition_repository, diff --git a/api/services/account_activation_adapters.py b/api/services/account_activation_adapters.py index 62b58c1336c..708ff620966 100644 --- a/api/services/account_activation_adapters.py +++ b/api/services/account_activation_adapters.py @@ -7,6 +7,7 @@ from services.account_activation_service import ( AccountActivationEligibility, InvitationTokenStore, WorkspaceInvitePolicy, + WorkspaceMemberAccessSync, WorkspaceMembershipCache, ) from services.account_service import RegisterService @@ -68,3 +69,21 @@ class BillingWorkspaceMembershipCache(WorkspaceMembershipCache): def invalidate(self, workspace_id: str) -> None: if self._enabled: BillingService.clean_billing_info_cache(workspace_id) + + +class RBACWorkspaceMemberAccessSync(WorkspaceMemberAccessSync): + def __init__(self, *, enabled: bool) -> None: + self._enabled = enabled + + @override + def sync(self, workspace_id: str, account_id: str) -> None: + if not self._enabled: + return + + from tasks.initialize_created_app_rbac_access_task import sync_joined_workspace_member_rbac_access_task + + sync_joined_workspace_member_rbac_access_task.delay( + str(workspace_id), + str(account_id), + operator_account_id=None, + ) diff --git a/api/services/account_activation_service.py b/api/services/account_activation_service.py index d635cf24f3d..7d497471143 100644 --- a/api/services/account_activation_service.py +++ b/api/services/account_activation_service.py @@ -48,6 +48,10 @@ class WorkspaceMembershipCache(Protocol): def invalidate(self, workspace_id: str) -> None: ... +class WorkspaceMemberAccessSync(Protocol): + def sync(self, workspace_id: str, account_id: str) -> None: ... + + class InvalidInvitationError(Exception): """The invitation is invalid, stale, or missing required activation data.""" @@ -73,12 +77,14 @@ class AccountActivationService: workspace_policy: WorkspaceInvitePolicy, eligibility: AccountActivationEligibility, membership_cache: WorkspaceMembershipCache, + member_access_sync: WorkspaceMemberAccessSync, ) -> None: self._tokens = tokens self._accounts = accounts self._workspace_policy = workspace_policy self._eligibility = eligibility self._membership_cache = membership_cache + self._member_access_sync = member_access_sync def check(self, invitation: InvitationLookup) -> ActivationCheckResult: resolved = self._resolve(invitation) @@ -128,6 +134,7 @@ class AccountActivationService: raise InvalidInvitationError if result.membership_created: self._membership_cache.invalidate(invitation.workspace_id) + self._member_access_sync.sync(invitation.workspace_id, invitation.account_id) def _resolve(self, invitation: InvitationLookup) -> AccountInvitation | None: token = self._tokens.find(invitation) diff --git a/api/services/account_service.py b/api/services/account_service.py index bd0cecc7b9e..a60c2007912 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -115,6 +115,14 @@ logger = logging.getLogger(__name__) _change_email_token_adapter: TypeAdapter[ChangeEmailTokenData] = TypeAdapter(ChangeEmailTokenData) +class EnterpriseWorkspaceMemberAccountNotFoundError(Exception): + pass + + +class EnterpriseWorkspaceMemberWorkspaceNotFoundError(Exception): + pass + + class InvitationDetailDict(TypedDict): account: Account data: InvitationData @@ -1319,7 +1327,12 @@ class TenantService: @staticmethod def create_tenant_member( - tenant: Tenant, account: Account, session: Session, role: str = "normal" + tenant: Tenant, + account: Account, + session: Session, + role: str = "normal", + *, + operator_account_id: str | None = None, ) -> TenantAccountJoin: """Create tenant member""" if role == TenantAccountRole.OWNER: @@ -1334,15 +1347,67 @@ class TenantService: ) if ta: ta.role = TenantAccountRole(role) + membership_created = False else: ta = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole(role)) session.add(ta) + membership_created = True session.commit() if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: BillingService.clean_billing_info_cache(tenant.id) + if ( + membership_created + and dify_config.RBAC_ENABLED + and TenantAccountRole(role) != TenantAccountRole.OWNER + and account.status != AccountStatus.PENDING + ): + from tasks.initialize_created_app_rbac_access_task import sync_joined_workspace_member_rbac_access_task + + sync_joined_workspace_member_rbac_access_task.delay( + str(tenant.id), + str(account.id), + operator_account_id=operator_account_id, + ) return ta + @staticmethod + def join_enterprise_workspace_member( + *, + workspace_id: str, + account_id: str, + email: str, + role: TenantAccountRole, + operator_account_id: str | None, + session: Session | None = None, + ) -> TenantAccountJoin: + session = session or db.session() + tenant = session.scalar( + select(Tenant).where( + Tenant.id == workspace_id, + Tenant.status == TenantStatus.NORMAL, + ) + ) + if tenant is None: + raise EnterpriseWorkspaceMemberWorkspaceNotFoundError + + account = session.scalar( + select(Account).where( + Account.id == account_id, + Account.email == email, + ) + ) + if account is None: + raise EnterpriseWorkspaceMemberAccountNotFoundError + + return TenantService.create_tenant_member( + tenant, + account, + session=session, + role=role.value, + operator_account_id=operator_account_id, + ) + @staticmethod def get_join_tenants(account: Account, *, session: Session) -> list[Tenant]: """Get account join tenants""" @@ -2018,7 +2083,13 @@ class RegisterService: check_normalized_email=True, session=session, ) - TenantService.create_tenant_member(tenant, account, session, tenant_join_role) + TenantService.create_tenant_member( + tenant, + account, + session, + tenant_join_role, + operator_account_id=inviter.id, + ) TenantService.switch_tenant(account, tenant.id, session=session) requires_setup = True else: @@ -2031,7 +2102,13 @@ class RegisterService: requires_setup = account.status == AccountStatus.PENDING if not ta and (account.status == AccountStatus.PENDING or dify_config.RBAC_ENABLED): - TenantService.create_tenant_member(tenant, account, session, tenant_join_role) + TenantService.create_tenant_member( + tenant, + account, + session, + tenant_join_role, + operator_account_id=inviter.id, + ) # Support resend invitation email when the account is pending status if account.status != AccountStatus.PENDING: diff --git a/api/services/enterprise/rbac_service.py b/api/services/enterprise/rbac_service.py index bbfea2db812..5f206548cc5 100644 --- a/api/services/enterprise/rbac_service.py +++ b/api/services/enterprise/rbac_service.py @@ -239,6 +239,32 @@ class ResourceWhitelistConfig(_RBACModel): automatic_include_workspace_members: bool +class ResourceWhitelistConfigResource(_RBACModel): + resource_type: RBACResourceType + resource_id: str + + +class ResourceWhitelistConfigItem(_RBACModel): + resource_type: RBACResourceType + resource_id: str + automatic_include_workspace_members: bool = False + account_ids: list[str] = Field(default_factory=list) + rbac_whitelist_scope: str | None = Field( + default=None, validation_alias=AliasChoices("rbac_whitelist_scope", "scope") + ) + + @field_validator("account_ids", mode="before") + @classmethod + def _coerce_account_ids(cls, value: Any) -> list[str]: + if value is None: + return [] + return value + + +class ResourceWhitelistConfigsResponse(_RBACModel): + data: list[ResourceWhitelistConfigItem] = Field(default_factory=list) + + class _LegacyResourceWhitelistConfig(_RBACModel): """RBAC service's pre-toggle whitelist payload, used only by data migrations.""" @@ -301,6 +327,18 @@ class ReplaceUserAccessPoliciesResponse(_RBACModel): access_policies: list[AccessPolicy] = Field(default_factory=list) +class AppendAppWhitelistMembersBatchItem(_RBACModel): + app_id: str + account_ids: list[str] = Field(default_factory=list) + policy_id: str + + +class AppendDatasetWhitelistMembersBatchItem(_RBACModel): + dataset_id: str + account_ids: list[str] = Field(default_factory=list) + policy_id: str + + class MemberRolesResponse(_RBACModel): account_id: str roles: list[RBACRole] = Field(default_factory=list) @@ -1123,6 +1161,25 @@ class RBACService: ) return AccessPolicyBindingState.model_validate(data or {}) + # ------------------------------------------------------------------ + # Mixed-resource whitelist config helpers. + # ------------------------------------------------------------------ + class ResourceWhitelistConfigs: + @staticmethod + def batch_get( + tenant_id: str, + account_id: str | None, + resources: Sequence[ResourceWhitelistConfigResource], + ) -> ResourceWhitelistConfigsResponse: + data = _inner_call( + "POST", + f"{_INNER_PREFIX}/whitelist/configs", + tenant_id=tenant_id, + account_id=account_id, + json={"resources": [resource.model_dump(mode="json") for resource in resources]}, + ) + return ResourceWhitelistConfigsResponse.model_validate(data or {}) + # ------------------------------------------------------------------ # Per-app access (screenshot 1: App Access Config). # ------------------------------------------------------------------ @@ -1226,6 +1283,20 @@ class RBACService: ) return ResourceWhitelist.model_validate(data or {}) + @staticmethod + def append_whitelist_members_batch( + tenant_id: str, + account_id: str | None, + data: Sequence[AppendAppWhitelistMembersBatchItem], + ) -> None: + _inner_call( + "POST", + f"{_INNER_PREFIX}/apps/whitelist/members/batch", + tenant_id=tenant_id, + account_id=account_id, + json={"data": [item.model_dump(mode="json") for item in data]}, + ) + @staticmethod def matrix(tenant_id: str, account_id: str | None, app_id: str) -> AppAccessMatrix: data = _inner_call( @@ -1424,6 +1495,20 @@ class RBACService: ) return ResourceWhitelist.model_validate(data or {}) + @staticmethod + def append_whitelist_members_batch( + tenant_id: str, + account_id: str | None, + data: Sequence[AppendDatasetWhitelistMembersBatchItem], + ) -> None: + _inner_call( + "POST", + f"{_INNER_PREFIX}/datasets/whitelist/members/batch", + tenant_id=tenant_id, + account_id=account_id, + json={"data": [item.model_dump(mode="json") for item in data]}, + ) + @staticmethod def matrix(tenant_id: str, account_id: str | None, dataset_id: str) -> DatasetAccessMatrix: data = _inner_call( diff --git a/api/tasks/initialize_created_app_rbac_access_task.py b/api/tasks/initialize_created_app_rbac_access_task.py index cb46d34beb4..ad69f2374ac 100644 --- a/api/tasks/initialize_created_app_rbac_access_task.py +++ b/api/tasks/initialize_created_app_rbac_access_task.py @@ -1,11 +1,14 @@ """Initialize default RBAC access for existing workspace members after app creation.""" import logging +from collections.abc import Iterator from celery import shared_task +from sqlalchemy import select from configs import dify_config from extensions.ext_database import db +from models import App, Dataset, TenantAccountJoin, TenantAccountRole from services.account_service import TenantService from services.enterprise import rbac_service as enterprise_rbac_service @@ -14,6 +17,61 @@ logger = logging.getLogger(__name__) APP_RBAC_ACCOUNT_POLICY_BATCH_SIZE = 500 APP_RBAC_DEFAULT_ACCESS_POLICY_ID = "default" APP_RBAC_QUEUE = "app_rbac" +APP_RBAC_RESOURCE_CONFIG_BATCH_SIZE = 500 +APP_RBAC_MEMBER_APPEND_BATCH_SIZE = 500 + + +def _owner_account_id(tenant_id: str) -> str | None: + return db.session().scalar( + select(TenantAccountJoin.account_id) + .where(TenantAccountJoin.tenant_id == tenant_id, TenantAccountJoin.role == TenantAccountRole.OWNER) + .order_by(TenantAccountJoin.id.asc()) + .limit(1) + ) + + +def _iter_resource_config_batches( + tenant_id: str, + batch_size: int, +) -> Iterator[list[enterprise_rbac_service.ResourceWhitelistConfigResource]]: + last_app_id: str | None = None + while True: + stmt = select(App.id).where(App.tenant_id == tenant_id).order_by(App.id.asc()).limit(batch_size) + if last_app_id: + stmt = stmt.where(App.id > last_app_id) + app_ids = [str(app_id) for app_id in db.session().scalars(stmt).all()] + if not app_ids: + break + yield [ + enterprise_rbac_service.ResourceWhitelistConfigResource( + resource_type=enterprise_rbac_service.RBACResourceType.APP, + resource_id=app_id, + ) + for app_id in app_ids + ] + last_app_id = app_ids[-1] + + last_dataset_id: str | None = None + while True: + stmt = select(Dataset.id).where(Dataset.tenant_id == tenant_id).order_by(Dataset.id.asc()).limit(batch_size) + if last_dataset_id: + stmt = stmt.where(Dataset.id > last_dataset_id) + dataset_ids = [str(dataset_id) for dataset_id in db.session().scalars(stmt).all()] + if not dataset_ids: + break + yield [ + enterprise_rbac_service.ResourceWhitelistConfigResource( + resource_type=enterprise_rbac_service.RBACResourceType.DATASET, + resource_id=dataset_id, + ) + for dataset_id in dataset_ids + ] + last_dataset_id = dataset_ids[-1] + + +def _chunks[T](items: list[T], chunk_size: int) -> Iterator[list[T]]: + for index in range(0, len(items), chunk_size): + yield items[index : index + chunk_size] @shared_task(queue=APP_RBAC_QUEUE, bind=True, max_retries=3, default_retry_delay=60) @@ -65,3 +123,77 @@ def initialize_created_app_rbac_access_task( self.request.retries + 1, ) raise self.retry(exc=exc) + + +@shared_task(queue=APP_RBAC_QUEUE, bind=True, max_retries=3, default_retry_delay=60) +def sync_joined_workspace_member_rbac_access_task( + self, + tenant_id: str, + member_account_id: str, + operator_account_id: str | None = None, +) -> None: + """Grant a newly joined member default access to resources that auto-include workspace members.""" + if not dify_config.RBAC_ENABLED: + return + + try: + actor_account_id = operator_account_id or _owner_account_id(tenant_id) + if actor_account_id is None: + logger.warning( + "Skipping joined member RBAC access sync because workspace owner was not found: tenant_id=%s member=%s", + tenant_id, + member_account_id, + ) + return + + app_ids: list[str] = [] + dataset_ids: list[str] = [] + for resources in _iter_resource_config_batches(tenant_id, APP_RBAC_RESOURCE_CONFIG_BATCH_SIZE): + configs = enterprise_rbac_service.RBACService.ResourceWhitelistConfigs.batch_get( + tenant_id=tenant_id, + account_id=actor_account_id, + resources=resources, + ) + for config in configs.data: + if not config.automatic_include_workspace_members: + continue + if config.resource_type == enterprise_rbac_service.RBACResourceType.APP: + app_ids.append(config.resource_id) + elif config.resource_type == enterprise_rbac_service.RBACResourceType.DATASET: + dataset_ids.append(config.resource_id) + + for app_id_batch in _chunks(app_ids, APP_RBAC_MEMBER_APPEND_BATCH_SIZE): + enterprise_rbac_service.RBACService.AppAccess.append_whitelist_members_batch( + tenant_id=tenant_id, + account_id=actor_account_id, + data=[ + enterprise_rbac_service.AppendAppWhitelistMembersBatchItem( + app_id=app_id, + account_ids=[member_account_id], + policy_id=APP_RBAC_DEFAULT_ACCESS_POLICY_ID, + ) + for app_id in app_id_batch + ], + ) + + for dataset_id_batch in _chunks(dataset_ids, APP_RBAC_MEMBER_APPEND_BATCH_SIZE): + enterprise_rbac_service.RBACService.DatasetAccess.append_whitelist_members_batch( + tenant_id=tenant_id, + account_id=actor_account_id, + data=[ + enterprise_rbac_service.AppendDatasetWhitelistMembersBatchItem( + dataset_id=dataset_id, + account_ids=[member_account_id], + policy_id=APP_RBAC_DEFAULT_ACCESS_POLICY_ID, + ) + for dataset_id in dataset_id_batch + ], + ) + except Exception as exc: + logger.exception( + "Failed to sync joined member RBAC access; retrying: tenant_id=%s member=%s attempt=%s", + tenant_id, + member_account_id, + self.request.retries + 1, + ) + raise self.retry(exc=exc) diff --git a/api/tests/unit_tests/commands/test_legacy_model_type_migration.py b/api/tests/unit_tests/commands/test_legacy_model_type_migration.py index 39c849e69d2..71d8e0f4f1c 100644 --- a/api/tests/unit_tests/commands/test_legacy_model_type_migration.py +++ b/api/tests/unit_tests/commands/test_legacy_model_type_migration.py @@ -377,7 +377,7 @@ def test_data_migrate_group_registers_dataset_permission_rbac_migration(command_ def test_data_migrate_group_registers_resource_whitelist_scope_migration(command_module) -> None: command = command_module.data_migrate.commands["rbac-migrate-resource-whitelist-scopes"] - assert command is command_module.migrate_resource_whitelist_scopes_to_automatic_include + assert command is command_module.migrate_only_me_resource_whitelist_scopes_to_automatic_include def test_dataset_permission_rbac_migration_help_mentions_binding_clear_side_effect(command_module) -> None: @@ -596,6 +596,61 @@ def test_resource_whitelist_scope_migration_all_syncs_workspace_members( assert set(replace_policy_calls[0]["payload"].account_ids) == {"maintainer-account-1", "member-account-1"} +def test_only_me_resource_whitelist_scope_migration_syncs_workspace_members( + command_module, + rbac_session: Session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rbac_module = importlib.import_module("commands.rbac") + _persist_dataset(rbac_session, maintainer="maintainer-account-1") + rbac_session.add_all( + [ + TenantAccountJoin( + tenant_id="tenant-1", + account_id="maintainer-account-1", + role=TenantAccountRole.OWNER, + ), + TenantAccountJoin( + tenant_id="tenant-1", + account_id="member-account-1", + role=TenantAccountRole.NORMAL, + ), + ] + ) + rbac_session.commit() + replace_whitelist_calls: list[dict[str, object]] = [] + replace_policy_calls: list[dict[str, object]] = [] + + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "legacy_whitelist_config", + lambda **kwargs: SimpleNamespace(rbac_whitelist_scope="only_me", account_ids=[]), + ) + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_whitelist", + lambda **kwargs: replace_whitelist_calls.append(kwargs), + ) + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_user_access_policies", + lambda **kwargs: replace_policy_calls.append(kwargs), + ) + + command_module.migrate_only_me_resource_whitelist_scopes_to_automatic_include.callback( + tenant_id=None, + resource_type="dataset", + resource_id=None, + batch_size=500, + member_batch_size=500, + dry_run=False, + ) + + assert replace_whitelist_calls[0]["payload"].automatic_include_workspace_members is True + assert replace_policy_calls[0]["payload"].access_policy_ids == ["default"] + assert set(replace_policy_calls[0]["payload"].account_ids) == {"maintainer-account-1", "member-account-1"} + + def test_data_migrate_command_defaults_output_to_stdout_stream( command_module, monkeypatch: pytest.MonkeyPatch, diff --git a/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py b/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py index 27f0a90bd14..c92349c59ad 100644 --- a/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py +++ b/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py @@ -17,12 +17,18 @@ from sqlalchemy.orm import Session, scoped_session, sessionmaker from controllers.inner_api.workspace.workspace import ( EnterpriseWorkspace, + EnterpriseWorkspaceMember, EnterpriseWorkspaceNoOwnerEmail, WorkspaceCreatePayload, + WorkspaceMemberPayload, WorkspaceOwnerlessPayload, ) from models import Account, Tenant -from models.account import TenantStatus +from models.account import TenantAccountJoin, TenantAccountRole, TenantStatus +from services.account_service import ( + EnterpriseWorkspaceMemberAccountNotFoundError, + EnterpriseWorkspaceMemberWorkspaceNotFoundError, +) @pytest.fixture @@ -78,6 +84,32 @@ class TestWorkspaceOwnerlessPayload: assert "name" in str(exc_info.value) +class TestWorkspaceMemberPayload: + """Test WorkspaceMemberPayload Pydantic model validation""" + + def test_valid_payload(self): + data = { + "workspace_id": "workspace-id", + "account_id": "account-id", + "email": "member@example.com", + "role": "normal", + "operator_account_id": "operator-id", + } + payload = WorkspaceMemberPayload.model_validate(data) + assert payload.workspace_id == "workspace-id" + assert payload.account_id == "account-id" + assert payload.email == "member@example.com" + assert payload.role == "normal" + assert payload.current is False + assert payload.operator_account_id == "operator-id" + + def test_missing_account_id_fails_validation(self): + data = {"workspace_id": "workspace-id", "email": "member@example.com"} + with pytest.raises(ValidationError) as exc_info: + WorkspaceMemberPayload.model_validate(data) + assert "account_id" in str(exc_info.value) + + class TestEnterpriseWorkspace: """Test EnterpriseWorkspace API endpoint handler logic. @@ -196,3 +228,116 @@ class TestEnterpriseWorkspaceNoOwnerEmail: "My Workspace", is_from_dashboard=True, session=database_session() ) mock_event.send.assert_called_once_with(tenant) + + +class TestEnterpriseWorkspaceMember: + """Test EnterpriseWorkspaceMember API endpoint handler logic.""" + + @pytest.fixture + def api_instance(self): + return EnterpriseWorkspaceMember() + + def test_has_post_method(self, api_instance): + assert hasattr(api_instance, "post") + assert callable(api_instance.post) + + @patch("controllers.inner_api.workspace.workspace.TenantService") + def test_post_joins_existing_account_to_workspace(self, mock_tenant_svc, api_instance, app: Flask): + membership = TenantAccountJoin( + tenant_id="workspace-id", + account_id="account-id", + role=TenantAccountRole.NORMAL, + ) + mock_tenant_svc.join_enterprise_workspace_member.return_value = membership + + unwrapped_post = inspect.unwrap(api_instance.post) + with app.test_request_context(): + with patch("controllers.inner_api.workspace.workspace.inner_api_ns") as mock_ns: + mock_ns.payload = { + "workspace_id": "workspace-id", + "account_id": "account-id", + "email": "member@example.com", + "role": "normal", + "operator_account_id": "operator-id", + } + result = unwrapped_post(api_instance) + + assert result["message"] == "enterprise workspace member joined." + assert result["member"] == { + "workspace_id": "workspace-id", + "account_id": "account-id", + "role": "normal", + } + mock_tenant_svc.join_enterprise_workspace_member.assert_called_once_with( + workspace_id="workspace-id", + account_id="account-id", + email="member@example.com", + role=TenantAccountRole.NORMAL, + operator_account_id="operator-id", + ) + + @patch("controllers.inner_api.workspace.workspace.TenantService") + def test_post_returns_404_when_workspace_not_found(self, mock_tenant_svc, api_instance, app: Flask): + mock_tenant_svc.join_enterprise_workspace_member.side_effect = EnterpriseWorkspaceMemberWorkspaceNotFoundError + + unwrapped_post = inspect.unwrap(api_instance.post) + with app.test_request_context(): + with patch("controllers.inner_api.workspace.workspace.inner_api_ns") as mock_ns: + mock_ns.payload = { + "workspace_id": "missing-workspace", + "account_id": "account-id", + "email": "member@example.com", + "role": "normal", + } + result = unwrapped_post(api_instance) + + assert result == ({"message": "workspace not found."}, 404) + mock_tenant_svc.join_enterprise_workspace_member.assert_called_once() + + @patch("controllers.inner_api.workspace.workspace.TenantService") + def test_post_returns_404_when_account_not_found(self, mock_tenant_svc, api_instance, app: Flask): + mock_tenant_svc.join_enterprise_workspace_member.side_effect = EnterpriseWorkspaceMemberAccountNotFoundError + + unwrapped_post = inspect.unwrap(api_instance.post) + with app.test_request_context(): + with patch("controllers.inner_api.workspace.workspace.inner_api_ns") as mock_ns: + mock_ns.payload = { + "workspace_id": "workspace-id", + "account_id": "missing-account", + "email": "member@example.com", + "role": "normal", + } + result = unwrapped_post(api_instance) + + assert result == ({"message": "account not found."}, 404) + mock_tenant_svc.join_enterprise_workspace_member.assert_called_once() + + @pytest.mark.usefixtures("database_session") + def test_post_rejects_owner_role(self, api_instance, app: Flask): + unwrapped_post = inspect.unwrap(api_instance.post) + with app.test_request_context(): + with patch("controllers.inner_api.workspace.workspace.inner_api_ns") as mock_ns: + mock_ns.payload = { + "workspace_id": "workspace-id", + "account_id": "account-id", + "email": "member@example.com", + "role": "owner", + } + result = unwrapped_post(api_instance) + + assert result == ({"message": "cannot join workspace as owner."}, 400) + + @pytest.mark.usefixtures("database_session") + def test_post_rejects_invalid_role(self, api_instance, app: Flask): + unwrapped_post = inspect.unwrap(api_instance.post) + with app.test_request_context(): + with patch("controllers.inner_api.workspace.workspace.inner_api_ns") as mock_ns: + mock_ns.payload = { + "workspace_id": "workspace-id", + "account_id": "account-id", + "email": "member@example.com", + "role": "not-a-role", + } + result = unwrapped_post(api_instance) + + assert result == ({"message": "invalid workspace member role."}, 400) diff --git a/api/tests/unit_tests/extensions/test_ext_application_services.py b/api/tests/unit_tests/extensions/test_ext_application_services.py index ba2bb85b5ed..e5cf0b3b82c 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -28,6 +28,7 @@ from services.account_activation_adapters import ( BillingAccountActivationEligibility, BillingWorkspaceMembershipCache, DeploymentWorkspaceInvitePolicy, + RBACWorkspaceMemberAccessSync, RegisterServiceInvitationTokenStore, ) from services.account_avatar_file_gateway import SQLAlchemyAccountAvatarFileGateway @@ -419,6 +420,7 @@ def test_build_application_services_wires_account_activation( assert activation._eligibility._enabled is billing_enabled assert isinstance(activation._membership_cache, BillingWorkspaceMembershipCache) assert activation._membership_cache._enabled is billing_enabled + assert isinstance(activation._member_access_sync, RBACWorkspaceMemberAccessSync) def test_build_application_services_wires_data_source_api_key_auth( diff --git a/api/tests/unit_tests/services/enterprise/test_rbac_service.py b/api/tests/unit_tests/services/enterprise/test_rbac_service.py index 6bb6bea7866..93e417dbec1 100644 --- a/api/tests/unit_tests/services/enterprise/test_rbac_service.py +++ b/api/tests/unit_tests/services/enterprise/test_rbac_service.py @@ -275,6 +275,52 @@ class TestAccessPolicies: class TestResourceAccess: + def test_resource_whitelist_configs_batch_get(self, mock_send: MagicMock): + mock_send.return_value = { + "data": [ + { + "resource_type": "app", + "resource_id": "app-1", + "scope": "all", + "automatic_include_workspace_members": True, + "account_ids": ["acct-1"], + }, + { + "resource_type": "dataset", + "resource_id": "dataset-1", + "scope": "specific", + "automatic_include_workspace_members": False, + "account_ids": None, + }, + ] + } + + out = svc.RBACService.ResourceWhitelistConfigs.batch_get( + "tenant-1", + "acct-actor", + [ + svc.ResourceWhitelistConfigResource(resource_type=svc.RBACResourceType.APP, resource_id="app-1"), + svc.ResourceWhitelistConfigResource( + resource_type=svc.RBACResourceType.DATASET, + resource_id="dataset-1", + ), + ], + ) + + call = _call_args(mock_send) + assert call.method == "POST" + assert call.endpoint == "/rbac/whitelist/configs" + assert call.json == { + "resources": [ + {"resource_type": "app", "resource_id": "app-1"}, + {"resource_type": "dataset", "resource_id": "dataset-1"}, + ] + } + assert [item.resource_id for item in out.data] == ["app-1", "dataset-1"] + assert out.data[0].automatic_include_workspace_members is True + assert out.data[0].rbac_whitelist_scope == "all" + assert out.data[1].account_ids == [] + def test_app_whitelist_resources(self, mock_send: MagicMock): mock_send.return_value = {"unrestricted": True, "resource_ids": ["app-1", "app-2"]} @@ -390,6 +436,50 @@ class TestResourceAccess: assert call.json == {"access_policy_ids": ["policy-1"]} assert out.access_policies[0].id == "policy-1" + def test_app_append_whitelist_members_batch(self, mock_send: MagicMock): + mock_send.return_value = None + + svc.RBACService.AppAccess.append_whitelist_members_batch( + "tenant-1", + "acct-actor", + [ + svc.AppendAppWhitelistMembersBatchItem( + app_id="app-1", + account_ids=["acct-1", "acct-2"], + policy_id="policy-1", + ) + ], + ) + + call = _call_args(mock_send) + assert call.method == "POST" + assert call.endpoint == "/rbac/apps/whitelist/members/batch" + assert call.json == { + "data": [{"app_id": "app-1", "account_ids": ["acct-1", "acct-2"], "policy_id": "policy-1"}] + } + + def test_dataset_append_whitelist_members_batch(self, mock_send: MagicMock): + mock_send.return_value = None + + svc.RBACService.DatasetAccess.append_whitelist_members_batch( + "tenant-1", + "acct-actor", + [ + svc.AppendDatasetWhitelistMembersBatchItem( + dataset_id="dataset-1", + account_ids=["acct-1", "acct-2"], + policy_id="policy-1", + ) + ], + ) + + call = _call_args(mock_send) + assert call.method == "POST" + assert call.endpoint == "/rbac/datasets/whitelist/members/batch" + assert call.json == { + "data": [{"dataset_id": "dataset-1", "account_ids": ["acct-1", "acct-2"], "policy_id": "policy-1"}] + } + def test_dataset_whitelist(self, mock_send: MagicMock): mock_send.return_value = {"account_ids": ["acct-2"], "automatic_include_workspace_members": False} diff --git a/api/tests/unit_tests/services/test_account_activation_adapters.py b/api/tests/unit_tests/services/test_account_activation_adapters.py index 402ee708d33..7bfdd87ddf4 100644 --- a/api/tests/unit_tests/services/test_account_activation_adapters.py +++ b/api/tests/unit_tests/services/test_account_activation_adapters.py @@ -4,6 +4,7 @@ from services.account_activation_adapters import ( BillingAccountActivationEligibility, BillingWorkspaceMembershipCache, DeploymentWorkspaceInvitePolicy, + RBACWorkspaceMemberAccessSync, RegisterServiceInvitationTokenStore, ) from services.entities.account_activation_entities import InvitationLookup, InvitationToken @@ -72,3 +73,21 @@ def test_workspace_policy_delegates_to_existing_policy_owner() -> None: DeploymentWorkspaceInvitePolicy().ensure_allowed("workspace-1") ensure_allowed.assert_called_once_with("workspace-1") + + +def test_rbac_member_access_sync_skips_gateway_when_disabled() -> None: + with patch( + "tasks.initialize_created_app_rbac_access_task.sync_joined_workspace_member_rbac_access_task.delay" + ) as delay: + RBACWorkspaceMemberAccessSync(enabled=False).sync("workspace-1", "account-1") + + delay.assert_not_called() + + +def test_rbac_member_access_sync_enqueues_joined_member_sync_when_enabled() -> None: + with patch( + "tasks.initialize_created_app_rbac_access_task.sync_joined_workspace_member_rbac_access_task.delay" + ) as delay: + RBACWorkspaceMemberAccessSync(enabled=True).sync("workspace-1", "account-1") + + delay.assert_called_once_with("workspace-1", "account-1", operator_account_id=None) diff --git a/api/tests/unit_tests/services/test_account_activation_service.py b/api/tests/unit_tests/services/test_account_activation_service.py index cd43170fb9e..43a67f59651 100644 --- a/api/tests/unit_tests/services/test_account_activation_service.py +++ b/api/tests/unit_tests/services/test_account_activation_service.py @@ -12,6 +12,7 @@ from services.account_activation_service import ( InvitationAccountMismatchError, InvitationTokenStore, WorkspaceInvitePolicy, + WorkspaceMemberAccessSync, WorkspaceMembershipCache, ) from services.entities.account_activation_entities import ( @@ -55,12 +56,13 @@ def _invitation( ) -def _service() -> tuple[AccountActivationService, Mock, Mock, Mock, Mock, Mock]: +def _service() -> tuple[AccountActivationService, Mock, Mock, Mock, Mock, Mock, Mock]: tokens = Mock(spec=InvitationTokenStore) accounts = Mock(spec=AccountActivationRepository) policy = Mock(spec=WorkspaceInvitePolicy) eligibility = Mock(spec=AccountActivationEligibility) membership_cache = Mock(spec=WorkspaceMembershipCache) + member_access_sync = Mock(spec=WorkspaceMemberAccessSync) eligibility.get_freeze_type.return_value = None service = AccountActivationService( tokens=tokens, @@ -68,13 +70,14 @@ def _service() -> tuple[AccountActivationService, Mock, Mock, Mock, Mock, Mock]: workspace_policy=policy, eligibility=eligibility, membership_cache=membership_cache, + member_access_sync=member_access_sync, ) - return service, tokens, accounts, policy, eligibility, membership_cache + return service, tokens, accounts, policy, eligibility, membership_cache, member_access_sync class TestCheckInvitation: def test_returns_invalid_without_touching_database_when_token_is_missing(self) -> None: - service, tokens, accounts, policy, _, _ = _service() + service, tokens, accounts, policy, _, _, _ = _service() tokens.find.return_value = None result = service.check(_lookup()) @@ -85,7 +88,7 @@ class TestCheckInvitation: policy.ensure_allowed.assert_not_called() def test_does_not_repeat_database_lookup_for_normalized_email(self) -> None: - service, tokens, accounts, policy, _, _ = _service() + service, tokens, accounts, policy, _, _, _ = _service() token = _token() tokens.find.return_value = token accounts.resolve.return_value = None @@ -97,7 +100,7 @@ class TestCheckInvitation: policy.ensure_allowed.assert_not_called() def test_falls_back_to_normalized_email_and_applies_workspace_policy(self) -> None: - service, tokens, accounts, policy, _, _ = _service() + service, tokens, accounts, policy, _, _, _ = _service() upper_case_token = InvitationToken( account_id="account-1", email="Invitee@Example.com", @@ -124,7 +127,7 @@ class TestCheckInvitation: class TestActivateInvitation: def test_rejects_authenticated_account_mismatch_before_side_effects(self) -> None: - service, tokens, accounts, _, eligibility, _ = _service() + service, tokens, accounts, _, eligibility, _, member_access_sync = _service() tokens.find.return_value = _token() accounts.resolve.return_value = _invitation() @@ -137,9 +140,10 @@ class TestActivateInvitation: eligibility.get_freeze_type.assert_not_called() tokens.revoke.assert_not_called() accounts.activate.assert_not_called() + member_access_sync.sync.assert_not_called() def test_rejects_frozen_account_without_consuming_token(self) -> None: - service, tokens, accounts, _, eligibility, _ = _service() + service, tokens, accounts, _, eligibility, _, member_access_sync = _service() tokens.find.return_value = _token() accounts.resolve.return_value = _invitation() eligibility.get_freeze_type.return_value = "freeze" @@ -150,9 +154,10 @@ class TestActivateInvitation: eligibility.get_freeze_type.assert_called_once_with("invitee@example.com") tokens.revoke.assert_not_called() accounts.activate.assert_not_called() + member_access_sync.sync.assert_not_called() def test_requires_all_setup_fields_before_consuming_token(self) -> None: - service, tokens, accounts, _, _, _ = _service() + service, tokens, accounts, _, _, _, member_access_sync = _service() tokens.find.return_value = _token() accounts.resolve.return_value = _invitation() @@ -164,9 +169,10 @@ class TestActivateInvitation: tokens.revoke.assert_not_called() accounts.activate.assert_not_called() + member_access_sync.sync.assert_not_called() def test_rejects_suspended_email_domain_without_consuming_token(self) -> None: - service, tokens, accounts, _, eligibility, _ = _service() + service, tokens, accounts, _, eligibility, _, member_access_sync = _service() tokens.find.return_value = _token() accounts.resolve.return_value = _invitation() eligibility.get_freeze_type.return_value = "email_domain_suspended" @@ -177,9 +183,10 @@ class TestActivateInvitation: eligibility.get_freeze_type.assert_called_once_with("invitee@example.com") tokens.revoke.assert_not_called() accounts.activate.assert_not_called() + member_access_sync.sync.assert_not_called() def test_activates_anonymous_invitation_and_invalidates_new_membership_cache(self) -> None: - service, tokens, accounts, _, eligibility, membership_cache = _service() + service, tokens, accounts, _, eligibility, membership_cache, member_access_sync = _service() tokens.find.return_value = _token() invitation = _invitation(role="owner") accounts.resolve.return_value = invitation @@ -201,9 +208,10 @@ class TestActivateInvitation: setup=AccountSetup(name="John Doe", interface_language="en-US", timezone="UTC"), ) membership_cache.invalidate.assert_called_once_with("workspace-1") + member_access_sync.sync.assert_called_once_with("workspace-1", "account-1") def test_preserves_existing_membership_cache_and_ignores_setup_fields(self) -> None: - service, tokens, accounts, _, _, membership_cache = _service() + service, tokens, accounts, _, _, membership_cache, member_access_sync = _service() tokens.find.return_value = _token() invitation = _invitation( account_status="active", @@ -225,3 +233,4 @@ class TestActivateInvitation: accounts.activate.assert_called_once_with(invitation, role="editor", setup=None) membership_cache.invalidate.assert_not_called() + member_access_sync.sync.assert_called_once_with("workspace-1", "account-1") diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index fb0dbdf4952..43f07765fef 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -21,7 +21,13 @@ from models.account import ( TenantStatus, ) from models.model import DifySetup -from services.account_service import AccountService, RegisterService, TenantService +from services.account_service import ( + AccountService, + EnterpriseWorkspaceMemberAccountNotFoundError, + EnterpriseWorkspaceMemberWorkspaceNotFoundError, + RegisterService, + TenantService, +) from services.enterprise.rbac_service import MembersInRole, Paginated from services.errors.account import ( AccountAlreadyInTenantError, @@ -945,6 +951,150 @@ class TestTenantService: assert persisted_tenant_account_join.account_id == account_id assert persisted_tenant_account_join.role == TenantAccountRole.NORMAL + def test_create_tenant_member_queues_joined_member_rbac_sync( + self, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + """New regular members are synced into auto-included RBAC resource whitelists.""" + import tasks.initialize_created_app_rbac_access_task as rbac_task_module + + delay = MagicMock() + with sqlite_session_factory() as service_session: + tenant = Tenant(name="Test Workspace") + account = Account(name="Test User", email="test@example.com") + service_session.add_all([tenant, account]) + service_session.flush() + tenant_id = tenant.id + account_id = account.id + service_session.commit() + + with ( + patch("services.account_service.dify_config.RBAC_ENABLED", True), + patch.object(rbac_task_module.sync_joined_workspace_member_rbac_access_task, "delay", delay), + ): + TenantService.create_tenant_member( + tenant, + account, + service_session, + "normal", + operator_account_id="operator-1", + ) + + delay.assert_called_once_with( + tenant_id, + account_id, + operator_account_id="operator-1", + ) + + def test_create_tenant_member_does_not_queue_pending_member_rbac_sync( + self, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + """Pending invited members are synced after activation, not at invitation time.""" + import tasks.initialize_created_app_rbac_access_task as rbac_task_module + + delay = MagicMock() + with sqlite_session_factory() as service_session: + tenant = Tenant(name="Test Workspace") + account = Account(name="Test User", email="test@example.com", status=AccountStatus.PENDING) + service_session.add_all([tenant, account]) + service_session.flush() + service_session.commit() + + with ( + patch("services.account_service.dify_config.RBAC_ENABLED", True), + patch.object(rbac_task_module.sync_joined_workspace_member_rbac_access_task, "delay", delay), + ): + TenantService.create_tenant_member( + tenant, + account, + service_session, + "normal", + operator_account_id="operator-1", + ) + + delay.assert_not_called() + + def test_join_enterprise_workspace_member_success(self, sqlite_session_factory: sessionmaker[Session]) -> None: + with sqlite_session_factory() as service_session: + tenant = Tenant(name="Test Workspace", status=TenantStatus.NORMAL) + account = Account(name="Test User", email="test@example.com") + service_session.add_all([tenant, account]) + service_session.flush() + tenant_id = tenant.id + account_id = account.id + service_session.commit() + + with patch("services.account_service.TenantService.create_tenant_member") as create_tenant_member: + membership = TenantAccountJoin( + tenant_id=tenant_id, + account_id=account_id, + role=TenantAccountRole.NORMAL, + ) + create_tenant_member.return_value = membership + + result = TenantService.join_enterprise_workspace_member( + workspace_id=tenant_id, + account_id=account_id, + email="test@example.com", + role=TenantAccountRole.NORMAL, + operator_account_id="operator-1", + session=service_session, + ) + + assert result is membership + create_tenant_member.assert_called_once() + call_args = create_tenant_member.call_args + assert call_args.args[0].id == tenant_id + assert call_args.args[1].id == account_id + assert call_args.kwargs == { + "session": service_session, + "role": "normal", + "operator_account_id": "operator-1", + } + + def test_join_enterprise_workspace_member_raises_when_workspace_missing( + self, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + with sqlite_session_factory() as service_session: + account = Account(name="Test User", email="test@example.com") + service_session.add(account) + service_session.flush() + account_id = account.id + service_session.commit() + + with pytest.raises(EnterpriseWorkspaceMemberWorkspaceNotFoundError): + TenantService.join_enterprise_workspace_member( + workspace_id="missing-workspace", + account_id=account_id, + email="test@example.com", + role=TenantAccountRole.NORMAL, + operator_account_id=None, + session=service_session, + ) + + def test_join_enterprise_workspace_member_raises_when_account_missing( + self, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + with sqlite_session_factory() as service_session: + tenant = Tenant(name="Test Workspace", status=TenantStatus.NORMAL) + service_session.add(tenant) + service_session.flush() + tenant_id = tenant.id + service_session.commit() + + with pytest.raises(EnterpriseWorkspaceMemberAccountNotFoundError): + TenantService.join_enterprise_workspace_member( + workspace_id=tenant_id, + account_id="missing-account", + email="test@example.com", + role=TenantAccountRole.NORMAL, + operator_account_id=None, + session=service_session, + ) + # ==================== Member Removal Tests ==================== def test_remove_pending_member_deletes_orphaned_account( @@ -2169,7 +2319,13 @@ class TestRegisterService: "add", session=sqlite_session, ) - mock_create_member.assert_called_once_with(mock_tenant, mock_new_account, sqlite_session, "normal") + mock_create_member.assert_called_once_with( + mock_tenant, + mock_new_account, + sqlite_session, + "normal", + operator_account_id=mock_inviter.id, + ) mock_switch_tenant.assert_called_once_with(mock_new_account, mock_tenant.id, session=sqlite_session) mock_generate_token.assert_called_once_with( mock_tenant, mock_new_account, "normal", requires_setup=True @@ -2214,7 +2370,13 @@ class TestRegisterService: # Verify results assert result == "invite-token-123" - mock_create_member.assert_called_once_with(mock_tenant, mock_existing_account, sqlite_session, "normal") + mock_create_member.assert_called_once_with( + mock_tenant, + mock_existing_account, + sqlite_session, + "normal", + operator_account_id=mock_inviter.id, + ) mock_generate_token.assert_called_once_with( mock_tenant, mock_existing_account, "normal", requires_setup=True ) @@ -2360,7 +2522,11 @@ class TestRegisterService: assert result == "rbac-token" mock_create_member.assert_called_once_with( - mock_tenant, mock_new_account, sqlite_session, TenantAccountRole.NORMAL.value + mock_tenant, + mock_new_account, + sqlite_session, + TenantAccountRole.NORMAL.value, + operator_account_id=mock_inviter.id, ) mock_rbac_service.MemberRoles.replace.assert_called_once_with( tenant_id=mock_tenant.id, @@ -2409,6 +2575,7 @@ class TestRegisterService: mock_existing_account, sqlite_session, TenantAccountRole.NORMAL.value, + operator_account_id=mock_inviter.id, ) mock_rbac_service.MemberRoles.replace.assert_called_once_with( tenant_id=mock_tenant.id, @@ -2458,6 +2625,7 @@ class TestRegisterService: mock_existing_account, sqlite_session, TenantAccountRole.NORMAL.value, + operator_account_id=mock_inviter.id, ) mock_rbac_service.MemberRoles.replace.assert_called_once_with( tenant_id=mock_tenant.id, @@ -2506,7 +2674,13 @@ class TestRegisterService: ) assert result == "legacy-token" - mock_create_member.assert_called_once_with(mock_tenant, mock_new_account, sqlite_session, "editor") + mock_create_member.assert_called_once_with( + mock_tenant, + mock_new_account, + sqlite_session, + "editor", + operator_account_id=mock_inviter.id, + ) mock_rbac_service.MemberRoles.replace.assert_not_called() # ==================== Token Management Tests ==================== diff --git a/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py b/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py index bc04cd5a800..54eaf46f37b 100644 --- a/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py +++ b/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py @@ -11,6 +11,12 @@ def test_initialize_created_app_rbac_access_task_uses_rbac_queue(): assert initialize_created_app_rbac_access_task.queue == APP_RBAC_QUEUE +def test_sync_joined_workspace_member_rbac_access_task_uses_rbac_queue(): + from tasks.initialize_created_app_rbac_access_task import sync_joined_workspace_member_rbac_access_task + + assert sync_joined_workspace_member_rbac_access_task.queue == APP_RBAC_QUEUE + + def test_initialize_created_app_rbac_access_task_batches_workspace_members(monkeypatch: pytest.MonkeyPatch): import tasks.initialize_created_app_rbac_access_task as task_module from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task @@ -67,3 +73,70 @@ def test_initialize_created_app_rbac_access_task_retries_on_failure(monkeypatch: retry.assert_called_once() assert isinstance(retry.call_args.kwargs["exc"], ConnectionError) + + +def test_sync_joined_workspace_member_rbac_access_task_appends_auto_included_resources( + monkeypatch: pytest.MonkeyPatch, +): + import tasks.initialize_created_app_rbac_access_task as task_module + from tasks.initialize_created_app_rbac_access_task import sync_joined_workspace_member_rbac_access_task + + rbac = task_module.enterprise_rbac_service + resources = [ + rbac.ResourceWhitelistConfigResource(resource_type=rbac.RBACResourceType.APP, resource_id="app-1"), + rbac.ResourceWhitelistConfigResource(resource_type=rbac.RBACResourceType.DATASET, resource_id="dataset-1"), + rbac.ResourceWhitelistConfigResource(resource_type=rbac.RBACResourceType.APP, resource_id="app-2"), + ] + configs = rbac.ResourceWhitelistConfigsResponse( + data=[ + rbac.ResourceWhitelistConfigItem( + resource_type=rbac.RBACResourceType.APP, + resource_id="app-1", + automatic_include_workspace_members=True, + ), + rbac.ResourceWhitelistConfigItem( + resource_type=rbac.RBACResourceType.DATASET, + resource_id="dataset-1", + automatic_include_workspace_members=True, + ), + rbac.ResourceWhitelistConfigItem( + resource_type=rbac.RBACResourceType.APP, + resource_id="app-2", + automatic_include_workspace_members=False, + ), + ] + ) + batch_get = MagicMock(return_value=configs) + app_append = MagicMock() + dataset_append = MagicMock() + + monkeypatch.setattr(task_module.dify_config, "RBAC_ENABLED", True) + monkeypatch.setattr(task_module, "_iter_resource_config_batches", lambda tenant_id, batch_size: iter([resources])) + monkeypatch.setattr(rbac.RBACService.ResourceWhitelistConfigs, "batch_get", batch_get) + monkeypatch.setattr(rbac.RBACService.AppAccess, "append_whitelist_members_batch", app_append) + monkeypatch.setattr(rbac.RBACService.DatasetAccess, "append_whitelist_members_batch", dataset_append) + + sync_joined_workspace_member_rbac_access_task.run("tenant-1", "member-1", "actor-1") + + batch_get.assert_called_once_with( + tenant_id="tenant-1", + account_id="actor-1", + resources=resources, + ) + app_append.assert_called_once() + app_call = app_append.call_args.kwargs + assert app_call["tenant_id"] == "tenant-1" + assert app_call["account_id"] == "actor-1" + assert len(app_call["data"]) == 1 + assert app_call["data"][0].app_id == "app-1" + assert app_call["data"][0].account_ids == ["member-1"] + assert app_call["data"][0].policy_id == task_module.APP_RBAC_DEFAULT_ACCESS_POLICY_ID + + dataset_append.assert_called_once() + dataset_call = dataset_append.call_args.kwargs + assert dataset_call["tenant_id"] == "tenant-1" + assert dataset_call["account_id"] == "actor-1" + assert len(dataset_call["data"]) == 1 + assert dataset_call["data"][0].dataset_id == "dataset-1" + assert dataset_call["data"][0].account_ids == ["member-1"] + assert dataset_call["data"][0].policy_id == task_module.APP_RBAC_DEFAULT_ACCESS_POLICY_ID