opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-08-28 19:23:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: new user add automatic add to app / dataset whitelist (#41363)

Browse Source
This commit is contained in:
wangxiaolei
2026-08-27 08:51:05 +00:00
committed by GitHub
parent a38ddd9d02
commit d84088ef74
18 changed files with 1127 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/commands/__init__.py
+2
View File
@@ -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",
api/commands/data_migrate.py
+6 -2
View File
@@ -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)
api/commands/rbac.py
+143
View File
@@ -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=(
api/controllers/inner_api/workspace/workspace.py
+64 -2
View File
@@ -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,
},
}
api/extensions/ext_application_services.py
+4
View File
@@ -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,
api/services/account_activation_adapters.py
+19
View File
@@ -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,
)
api/services/account_activation_service.py
+7
View File
@@ -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)
api/services/account_service.py
+80 -3
View File
@@ -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:
api/services/enterprise/rbac_service.py
+85
View File
@@ -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(
api/tasks/initialize_created_app_rbac_access_task.py
+132
View File
@@ -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)
api/tests/unit_tests/commands/test_legacy_model_type_migration.py
+56 -1
View File
@@ -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,
api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py
+146 -1
View File
@@ -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)
api/tests/unit_tests/extensions/test_ext_application_services.py
+2
View File
@@ -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(
api/tests/unit_tests/services/enterprise/test_rbac_service.py
+90
View File
@@ -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}
api/tests/unit_tests/services/test_account_activation_adapters.py
+19
View File
@@ -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)
api/tests/unit_tests/services/test_account_activation_service.py
+20 -11
View File
@@ -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")
api/tests/unit_tests/services/test_account_service.py
+179 -5
View File
@@ -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 ====================
api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py
+73
View File
@@ -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