diff --git a/api/core/app/apps/advanced_chat/app_generator.py b/api/core/app/apps/advanced_chat/app_generator.py index f88034e2b43..42233b22477 100644 --- a/api/core/app/apps/advanced_chat/app_generator.py +++ b/api/core/app/apps/advanced_chat/app_generator.py @@ -242,6 +242,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): workflow_triggered_from = WorkflowRunTriggeredFrom.APP_RUN workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=workflow_triggered_from, @@ -249,6 +250,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): # Create workflow node execution repository workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -375,6 +377,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): # Create workflow execution(aka workflow run) repository workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, @@ -382,6 +385,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): # Create workflow node execution repository workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, @@ -466,6 +470,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): # Create workflow execution(aka workflow run) repository workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, @@ -473,6 +478,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): # Create workflow node execution repository workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, diff --git a/api/core/app/apps/pipeline/pipeline_generator.py b/api/core/app/apps/pipeline/pipeline_generator.py index 7cf7949a4d5..dca2eb9055f 100644 --- a/api/core/app/apps/pipeline/pipeline_generator.py +++ b/api/core/app/apps/pipeline/pipeline_generator.py @@ -216,6 +216,7 @@ class PipelineGenerator(BaseAppGenerator): session_factory = sessionmaker(bind=db.engine, expire_on_commit=False) workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=pipeline.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=workflow_triggered_from, @@ -223,6 +224,7 @@ class PipelineGenerator(BaseAppGenerator): workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=pipeline.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.RAG_PIPELINE_RUN, @@ -425,6 +427,7 @@ class PipelineGenerator(BaseAppGenerator): workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=pipeline.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowRunTriggeredFrom.RAG_PIPELINE_DEBUGGING, @@ -432,6 +435,7 @@ class PipelineGenerator(BaseAppGenerator): workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=pipeline.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, @@ -524,6 +528,7 @@ class PipelineGenerator(BaseAppGenerator): workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=pipeline.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowRunTriggeredFrom.RAG_PIPELINE_DEBUGGING, @@ -531,6 +536,7 @@ class PipelineGenerator(BaseAppGenerator): workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=pipeline.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, diff --git a/api/core/app/apps/workflow/app_generator.py b/api/core/app/apps/workflow/app_generator.py index e8eca44cf88..3dc2cf99e04 100644 --- a/api/core/app/apps/workflow/app_generator.py +++ b/api/core/app/apps/workflow/app_generator.py @@ -243,6 +243,7 @@ class WorkflowAppGenerator(BaseAppGenerator): workflow_triggered_from = WorkflowRunTriggeredFrom.APP_RUN workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=workflow_triggered_from, @@ -250,6 +251,7 @@ class WorkflowAppGenerator(BaseAppGenerator): # Create workflow node execution repository workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -470,6 +472,7 @@ class WorkflowAppGenerator(BaseAppGenerator): # Create workflow execution(aka workflow run) repository workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, @@ -477,6 +480,7 @@ class WorkflowAppGenerator(BaseAppGenerator): # Create workflow node execution repository workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, @@ -560,6 +564,7 @@ class WorkflowAppGenerator(BaseAppGenerator): # Create workflow execution(aka workflow run) repository workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, @@ -567,6 +572,7 @@ class WorkflowAppGenerator(BaseAppGenerator): # Create workflow node execution repository workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=application_generate_entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, diff --git a/api/core/repositories/celery_workflow_execution_repository.py b/api/core/repositories/celery_workflow_execution_repository.py index bf7c8d48f8e..bc72742beb7 100644 --- a/api/core/repositories/celery_workflow_execution_repository.py +++ b/api/core/repositories/celery_workflow_execution_repository.py @@ -13,7 +13,6 @@ from sqlalchemy.orm import sessionmaker from core.repositories.factory import WorkflowExecutionRepository from graphon.entities import WorkflowExecution -from libs.helper import extract_tenant_id from models import Account, CreatorUserRole, EndUser from models.enums import WorkflowRunTriggeredFrom from tasks.workflow_execution_tasks import ( @@ -47,6 +46,7 @@ class CeleryWorkflowExecutionRepository(WorkflowExecutionRepository): def __init__( self, session_factory: sessionmaker | Engine, + tenant_id: str, user: Account | EndUser, app_id: str | None, triggered_from: WorkflowRunTriggeredFrom | None, @@ -56,7 +56,8 @@ class CeleryWorkflowExecutionRepository(WorkflowExecutionRepository): Args: session_factory: SQLAlchemy sessionmaker or engine for fallback operations - user: Account or EndUser object containing tenant_id, user ID, and role information + tenant_id: Tenant that owns the workflow execution + user: Account or EndUser used for creator attribution app_id: App ID for filtering by application (can be None) triggered_from: Source of the execution trigger (DEBUGGING or APP_RUN) """ @@ -71,10 +72,8 @@ class CeleryWorkflowExecutionRepository(WorkflowExecutionRepository): f"Invalid session_factory type {type(session_factory).__name__}; expected sessionmaker or Engine" ) - # Extract tenant_id from user - tenant_id = extract_tenant_id(user) if not tenant_id: - raise ValueError("User must have a tenant_id or current_tenant_id") + raise ValueError("tenant_id is required") self._tenant_id = tenant_id # Store app context diff --git a/api/core/repositories/celery_workflow_node_execution_repository.py b/api/core/repositories/celery_workflow_node_execution_repository.py index f48d92f8797..3b152a0e232 100644 --- a/api/core/repositories/celery_workflow_node_execution_repository.py +++ b/api/core/repositories/celery_workflow_node_execution_repository.py @@ -17,7 +17,6 @@ from core.repositories.factory import ( WorkflowNodeExecutionRepository, ) from graphon.entities import WorkflowNodeExecution -from libs.helper import extract_tenant_id from models import Account, CreatorUserRole, EndUser from models.workflow import WorkflowNodeExecutionTriggeredFrom from tasks.workflow_node_execution_tasks import ( @@ -54,6 +53,7 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): def __init__( self, session_factory: sessionmaker | Engine, + tenant_id: str, user: Account | EndUser, app_id: str | None, triggered_from: WorkflowNodeExecutionTriggeredFrom | None, @@ -63,7 +63,8 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): Args: session_factory: SQLAlchemy sessionmaker or engine for fallback operations - user: Account or EndUser object containing tenant_id, user ID, and role information + tenant_id: Tenant that owns the workflow node execution + user: Account or EndUser used for creator attribution app_id: App ID for filtering by application (can be None) triggered_from: Source of the execution trigger (SINGLE_STEP or WORKFLOW_RUN) """ @@ -78,10 +79,8 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): f"Invalid session_factory type {type(session_factory).__name__}; expected sessionmaker or Engine" ) - # Extract tenant_id from user - tenant_id = extract_tenant_id(user) if not tenant_id: - raise ValueError("User must have a tenant_id or current_tenant_id") + raise ValueError("tenant_id is required") self._tenant_id = tenant_id # Store app context diff --git a/api/core/repositories/factory.py b/api/core/repositories/factory.py index 4e83e707997..8a37ba933fc 100644 --- a/api/core/repositories/factory.py +++ b/api/core/repositories/factory.py @@ -62,6 +62,7 @@ class DifyCoreRepositoryFactory: def create_workflow_execution_repository( cls, session_factory: sessionmaker | Engine, + tenant_id: str, user: Account | EndUser, app_id: str, triggered_from: WorkflowRunTriggeredFrom, @@ -71,7 +72,8 @@ class DifyCoreRepositoryFactory: Args: session_factory: SQLAlchemy sessionmaker or engine - user: Account or EndUser object + tenant_id: Tenant that owns the workflow execution + user: Account or EndUser used for creator attribution app_id: Application ID triggered_from: Source of the execution trigger @@ -87,6 +89,7 @@ class DifyCoreRepositoryFactory: repository_class = import_string(class_path) return repository_class( session_factory=session_factory, + tenant_id=tenant_id, user=user, app_id=app_id, triggered_from=triggered_from, @@ -98,6 +101,7 @@ class DifyCoreRepositoryFactory: def create_workflow_node_execution_repository( cls, session_factory: sessionmaker | Engine, + tenant_id: str, user: Account | EndUser, app_id: str, triggered_from: WorkflowNodeExecutionTriggeredFrom, @@ -107,7 +111,8 @@ class DifyCoreRepositoryFactory: Args: session_factory: SQLAlchemy sessionmaker or engine - user: Account or EndUser object + tenant_id: Tenant that owns the workflow node execution + user: Account or EndUser used for creator attribution app_id: Application ID triggered_from: Source of the execution trigger @@ -123,6 +128,7 @@ class DifyCoreRepositoryFactory: repository_class = import_string(class_path) return repository_class( session_factory=session_factory, + tenant_id=tenant_id, user=user, app_id=app_id, triggered_from=triggered_from, diff --git a/api/core/repositories/sqlalchemy_workflow_execution_repository.py b/api/core/repositories/sqlalchemy_workflow_execution_repository.py index 0e9f842731d..fa68dfa84a0 100644 --- a/api/core/repositories/sqlalchemy_workflow_execution_repository.py +++ b/api/core/repositories/sqlalchemy_workflow_execution_repository.py @@ -13,7 +13,6 @@ from core.repositories.factory import WorkflowExecutionRepository from graphon.entities import WorkflowExecution from graphon.enums import WorkflowExecutionStatus, WorkflowType from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter -from libs.helper import extract_tenant_id from models import ( Account, CreatorUserRole, @@ -40,6 +39,7 @@ class SQLAlchemyWorkflowExecutionRepository(WorkflowExecutionRepository): def __init__( self, session_factory: sessionmaker | Engine, + tenant_id: str, user: Account | EndUser, app_id: str | None, triggered_from: WorkflowRunTriggeredFrom | None, @@ -49,7 +49,8 @@ class SQLAlchemyWorkflowExecutionRepository(WorkflowExecutionRepository): Args: session_factory: SQLAlchemy sessionmaker or engine for creating sessions - user: Account or EndUser object containing tenant_id, user ID, and role information + tenant_id: Tenant that owns the workflow execution + user: Account or EndUser used for creator attribution app_id: App ID for filtering by application (can be None) triggered_from: Source of the execution trigger (DEBUGGING or APP_RUN) """ @@ -64,10 +65,8 @@ class SQLAlchemyWorkflowExecutionRepository(WorkflowExecutionRepository): f"Invalid session_factory type {type(session_factory).__name__}; expected sessionmaker or Engine" ) - # Extract tenant_id from user - tenant_id = extract_tenant_id(user) if not tenant_id: - raise ValueError("User must have a tenant_id or current_tenant_id") + raise ValueError("tenant_id is required") self._tenant_id = tenant_id # Store app context diff --git a/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py b/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py index 65324028f82..15dca1ff8cf 100644 --- a/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py +++ b/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py @@ -23,7 +23,6 @@ from graphon.entities import WorkflowNodeExecution from graphon.enums import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus from graphon.model_runtime.utils.encoders import jsonable_encoder from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter -from libs.helper import extract_tenant_id from libs.uuid_utils import uuidv7 from models import ( Account, @@ -63,6 +62,7 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository) def __init__( self, session_factory: sessionmaker | Engine, + tenant_id: str, user: Account | EndUser, app_id: str | None, triggered_from: WorkflowNodeExecutionTriggeredFrom | None, @@ -72,7 +72,8 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository) Args: session_factory: SQLAlchemy sessionmaker or engine for creating sessions - user: Account or EndUser object containing tenant_id, user ID, and role information + tenant_id: Tenant that owns the workflow node execution + user: Account or EndUser used for creator attribution app_id: App ID for filtering by application (can be None) triggered_from: Source of the execution trigger (SINGLE_STEP or WORKFLOW_RUN) """ @@ -87,10 +88,8 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository) f"Invalid session_factory type {type(session_factory).__name__}; expected sessionmaker or Engine" ) - # Extract tenant_id from user - tenant_id = extract_tenant_id(user) if not tenant_id: - raise ValueError("User must have a tenant_id or current_tenant_id") + raise ValueError("tenant_id is required") self._tenant_id = tenant_id # Store app context @@ -299,6 +298,7 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository) content=value_json.encode("utf-8"), mimetype="application/json", user=self._user, + tenant_id=self._tenant_id, ) offload = WorkflowNodeExecutionOffload( id=uuidv7(), diff --git a/api/extensions/logstore/repositories/logstore_workflow_execution_repository.py b/api/extensions/logstore/repositories/logstore_workflow_execution_repository.py index d0653bec3f7..640b9c4ff66 100644 --- a/api/extensions/logstore/repositories/logstore_workflow_execution_repository.py +++ b/api/extensions/logstore/repositories/logstore_workflow_execution_repository.py @@ -12,7 +12,6 @@ from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchem from extensions.logstore.aliyun_logstore import AliyunLogStore from graphon.entities import WorkflowExecution from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter -from libs.helper import extract_tenant_id from models import ( Account, CreatorUserRole, @@ -27,6 +26,7 @@ class LogstoreWorkflowExecutionRepository(WorkflowExecutionRepository): def __init__( self, session_factory: sessionmaker | Engine, + tenant_id: str, user: Account | EndUser, app_id: str | None, triggered_from: WorkflowRunTriggeredFrom | None, @@ -36,7 +36,8 @@ class LogstoreWorkflowExecutionRepository(WorkflowExecutionRepository): Args: session_factory: SQLAlchemy sessionmaker or engine for creating sessions - user: Account or EndUser object containing tenant_id, user ID, and role information + tenant_id: Tenant that owns the workflow execution + user: Account or EndUser used for creator attribution app_id: App ID for filtering by application (can be None) triggered_from: Source of the execution trigger (DEBUGGING or APP_RUN) """ @@ -47,10 +48,8 @@ class LogstoreWorkflowExecutionRepository(WorkflowExecutionRepository): # Note: Project/logstore/index initialization is done at app startup via ext_logstore self.logstore_client = AliyunLogStore() - # Extract tenant_id from user - tenant_id = extract_tenant_id(user) if not tenant_id: - raise ValueError("User must have a tenant_id or current_tenant_id") + raise ValueError("tenant_id is required") self._tenant_id = tenant_id # Store app context @@ -64,7 +63,13 @@ class LogstoreWorkflowExecutionRepository(WorkflowExecutionRepository): self._creator_user_role = CreatorUserRole.ACCOUNT if isinstance(user, Account) else CreatorUserRole.END_USER # Initialize SQL repository for dual-write support - self.sql_repository = SQLAlchemyWorkflowExecutionRepository(session_factory, user, app_id, triggered_from) + self.sql_repository = SQLAlchemyWorkflowExecutionRepository( + session_factory=session_factory, + tenant_id=tenant_id, + user=user, + app_id=app_id, + triggered_from=triggered_from, + ) # Control flag for dual-write (write to both LogStore and SQL database) # Set to True to enable dual-write for safe migration, False to use LogStore only diff --git a/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py b/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py index 6daac8a0a51..1c62df3746d 100644 --- a/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py +++ b/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py @@ -26,7 +26,6 @@ from graphon.entities import WorkflowNodeExecution from graphon.enums import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus from graphon.model_runtime.utils.encoders import jsonable_encoder from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter -from libs.helper import extract_tenant_id from models import ( Account, CreatorUserRole, @@ -109,6 +108,7 @@ class LogstoreWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): def __init__( self, session_factory: sessionmaker | Engine, + tenant_id: str, user: Account | EndUser, app_id: str | None, triggered_from: WorkflowNodeExecutionTriggeredFrom | None, @@ -118,7 +118,8 @@ class LogstoreWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): Args: session_factory: SQLAlchemy sessionmaker or engine for creating sessions - user: Account or EndUser object containing tenant_id, user ID, and role information + tenant_id: Tenant that owns the workflow node execution + user: Account or EndUser used for creator attribution app_id: App ID for filtering by application (can be None) triggered_from: Source of the execution trigger (SINGLE_STEP or WORKFLOW_RUN) """ @@ -128,10 +129,8 @@ class LogstoreWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): # Initialize LogStore client self.logstore_client = AliyunLogStore() - # Extract tenant_id from user - tenant_id = extract_tenant_id(user) if not tenant_id: - raise ValueError("User must have a tenant_id or current_tenant_id") + raise ValueError("tenant_id is required") self._tenant_id = tenant_id # Store app context @@ -145,7 +144,13 @@ class LogstoreWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): self._creator_user_role = CreatorUserRole.ACCOUNT if isinstance(user, Account) else CreatorUserRole.END_USER # Initialize SQL repository for dual-write support - self.sql_repository = SQLAlchemyWorkflowNodeExecutionRepository(session_factory, user, app_id, triggered_from) + self.sql_repository = SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=session_factory, + tenant_id=tenant_id, + user=user, + app_id=app_id, + triggered_from=triggered_from, + ) # Control flag for dual-write (write to both LogStore and SQL database) # Set to True to enable dual-write for safe migration, False to use LogStore only diff --git a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/aliyun_trace.py b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/aliyun_trace.py index 8a5eef69b9e..b2de91860e0 100644 --- a/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/aliyun_trace.py +++ b/api/providers/trace/trace-aliyun/src/dify_trace_aliyun/aliyun_trace.py @@ -296,6 +296,7 @@ class AliyunDataTrace(BaseTraceInstance): session_factory = sessionmaker(bind=db.engine) workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=trace_info.tenant_id, user=service_account, app_id=app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_phoenix_trace.py b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_phoenix_trace.py index 9f779fc43de..4199ec131a0 100644 --- a/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_phoenix_trace.py +++ b/api/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/arize_phoenix_trace.py @@ -845,6 +845,7 @@ class ArizePhoenixDataTrace(BaseTraceInstance): workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=trace_info.tenant_id, user=service_account, app_id=app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/langfuse_trace.py b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/langfuse_trace.py index caa9b67c662..c1b0b0e9132 100644 --- a/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/langfuse_trace.py +++ b/api/providers/trace/trace-langfuse/src/dify_trace_langfuse/langfuse_trace.py @@ -186,6 +186,7 @@ class LangFuseDataTrace(BaseTraceInstance): workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=trace_info.tenant_id, user=service_account, app_id=app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/langsmith_trace.py b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/langsmith_trace.py index 4260e5d6ba6..8ce59dfcf13 100644 --- a/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/langsmith_trace.py +++ b/api/providers/trace/trace-langsmith/src/dify_trace_langsmith/langsmith_trace.py @@ -154,6 +154,7 @@ class LangSmithDataTrace(BaseTraceInstance): workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=trace_info.tenant_id, user=service_account, app_id=app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/providers/trace/trace-opik/src/dify_trace_opik/opik_trace.py b/api/providers/trace/trace-opik/src/dify_trace_opik/opik_trace.py index c9356ab3687..2e9ab7ccee8 100644 --- a/api/providers/trace/trace-opik/src/dify_trace_opik/opik_trace.py +++ b/api/providers/trace/trace-opik/src/dify_trace_opik/opik_trace.py @@ -174,6 +174,7 @@ class OpikDataTrace(BaseTraceInstance): workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=trace_info.tenant_id, user=service_account, app_id=app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/providers/trace/trace-tencent/src/dify_trace_tencent/tencent_trace.py b/api/providers/trace/trace-tencent/src/dify_trace_tencent/tencent_trace.py index 9221ee9d7e9..2d677e2ae84 100644 --- a/api/providers/trace/trace-tencent/src/dify_trace_tencent/tencent_trace.py +++ b/api/providers/trace/trace-tencent/src/dify_trace_tencent/tencent_trace.py @@ -252,18 +252,9 @@ class TencentDataTrace(BaseTraceInstance): if not service_account: raise ValueError(f"Creator account not found for app {app_id}") - current_tenant = session.scalar( - select(TenantAccountJoin) - .where(TenantAccountJoin.account_id == service_account.id, TenantAccountJoin.current.is_(True)) - .limit(1) - ) - if not current_tenant: - raise ValueError(f"Current tenant not found for account {service_account.id}") - - service_account.set_tenant_id_with_session(current_tenant.tenant_id, session=session) - repository = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=session_maker, + tenant_id=app.tenant_id, user=service_account, app_id=trace_info.metadata.get("app_id"), triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_tencent_trace.py b/api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_tencent_trace.py index 4aae27f6b4c..71bcabd1416 100644 --- a/api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_tencent_trace.py +++ b/api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_tencent_trace.py @@ -18,7 +18,7 @@ from core.ops.entities.trace_entity import ( ) from graphon.entities import WorkflowNodeExecution from graphon.enums import BuiltinNodeTypes -from models import Account, App, TenantAccountJoin +from models import Account, App logger = logging.getLogger(__name__) @@ -420,20 +420,18 @@ class TestTencentDataTrace: app = MagicMock(spec=App) app.id = "app-1" app.created_by = "user-1" + app.tenant_id = "tenant-1" account = MagicMock(spec=Account) account.id = "user-1" - tenant_join = MagicMock(spec=TenantAccountJoin) - tenant_join.tenant_id = "tenant-1" - mock_executions = [MagicMock()] with patch("dify_trace_tencent.tencent_trace.db") as mock_db: mock_db.engine = "engine" with patch("dify_trace_tencent.tencent_trace.Session") as mock_session_ctx: session = mock_session_ctx.return_value.__enter__.return_value - session.scalar.side_effect = [app, account, tenant_join] + session.scalar.side_effect = [app, account] with patch("dify_trace_tencent.tencent_trace.SQLAlchemyWorkflowNodeExecutionRepository") as mock_repo: mock_repo.return_value.get_by_workflow_execution.return_value = mock_executions @@ -441,7 +439,7 @@ class TestTencentDataTrace: results = tencent_data_trace._get_workflow_node_executions(trace_info) assert results == mock_executions - account.set_tenant_id_with_session.assert_called_once_with("tenant-1", session=session) + assert mock_repo.call_args.kwargs["tenant_id"] == "tenant-1" def test_get_workflow_node_executions_no_app_id(self, tencent_data_trace, caplog: pytest.LogCaptureFixture): trace_info = MagicMock(spec=WorkflowTraceInfo) diff --git a/api/providers/trace/trace-weave/src/dify_trace_weave/weave_trace.py b/api/providers/trace/trace-weave/src/dify_trace_weave/weave_trace.py index 03a13a636c2..1cd4e3b3b5b 100644 --- a/api/providers/trace/trace-weave/src/dify_trace_weave/weave_trace.py +++ b/api/providers/trace/trace-weave/src/dify_trace_weave/weave_trace.py @@ -159,6 +159,7 @@ class WeaveDataTrace(BaseTraceInstance): workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=trace_info.tenant_id, user=service_account, app_id=app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/services/file_service.py b/api/services/file_service.py index 92dbe0dc826..fba7a760dfa 100644 --- a/api/services/file_service.py +++ b/api/services/file_service.py @@ -53,6 +53,7 @@ class FileService: content: bytes, mimetype: str, user: Account | EndUser, + tenant_id: str | None = None, source: Literal["datasets"] | None = None, source_url: str = "", ) -> UploadFile: @@ -84,16 +85,16 @@ class FileService: # generate file key file_uuid = str(uuid.uuid4()) - current_tenant_id = extract_tenant_id(user) + resource_tenant_id = tenant_id if tenant_id is not None else extract_tenant_id(user) - file_key = "upload_files/" + (current_tenant_id or "") + "/" + file_uuid + "." + extension + file_key = "upload_files/" + (resource_tenant_id or "") + "/" + file_uuid + "." + extension # save file to storage storage.save(file_key, content) # save file to db upload_file = UploadFile( - tenant_id=current_tenant_id or "", + tenant_id=resource_tenant_id or "", storage_type=StorageType(dify_config.STORAGE_TYPE), key=file_key, name=filename, diff --git a/api/services/rag_pipeline/rag_pipeline.py b/api/services/rag_pipeline/rag_pipeline.py index f700e0088e3..6d7f3a01af4 100644 --- a/api/services/rag_pipeline/rag_pipeline.py +++ b/api/services/rag_pipeline/rag_pipeline.py @@ -586,6 +586,7 @@ class RagPipelineService: repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=db.engine, + tenant_id=pipeline.tenant_id, user=account, app_id=pipeline.id, triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, @@ -1377,6 +1378,7 @@ class RagPipelineService: # Create repository and save the node execution repository = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=db.engine, + tenant_id=pipeline.tenant_id, user=current_user, app_id=pipeline.id, triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index c8d77cb7f29..4f9d85719b9 100644 --- a/api/services/workflow_service.py +++ b/api/services/workflow_service.py @@ -1039,6 +1039,7 @@ class WorkflowService: # Create repository and save the node execution repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=db.engine, + tenant_id=app_model.tenant_id, user=account, app_id=app_model.id, triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, diff --git a/api/tasks/app_generate/workflow_execute_task.py b/api/tasks/app_generate/workflow_execute_task.py index 213f37a8850..d76066a8aa7 100644 --- a/api/tasks/app_generate/workflow_execute_task.py +++ b/api/tasks/app_generate/workflow_execute_task.py @@ -248,7 +248,6 @@ class _AppRunner: case _Account(): with self._session() as session: user: Account = session.get(Account, user_params.user_id) - user.set_tenant_id_with_session(self._exec_params.tenant_id, session=session) return user case _: raise AssertionError(f"user should only be _Account or _EndUser, got {type(user_params)}") @@ -257,10 +256,7 @@ class _AppRunner: def _resolve_user_for_run(session: Session, workflow_run: WorkflowRun) -> Account | EndUser | None: role = CreatorUserRole(workflow_run.created_by_role) if role == CreatorUserRole.ACCOUNT: - user = session.get(Account, workflow_run.created_by) - if user: - user.set_tenant_id_with_session(workflow_run.tenant_id, session=session) - return user + return session.get(Account, workflow_run.created_by) return session.get(EndUser, workflow_run.created_by) @@ -617,12 +613,14 @@ def _resume_advanced_chat( workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=app_model.id, triggered_from=triggered_from, ) workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=app_model.id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -684,12 +682,14 @@ def _resume_workflow( workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=app_model.id, triggered_from=triggered_from, ) workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=app_model.id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/tasks/async_workflow_tasks.py b/api/tasks/async_workflow_tasks.py index 319c5379f57..a5f879be301 100644 --- a/api/tasks/async_workflow_tasks.py +++ b/api/tasks/async_workflow_tasks.py @@ -247,12 +247,14 @@ def resume_workflow_execution(task_data_dict: dict[str, Any]) -> None: workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=generate_entity.app_config.app_id, triggered_from=WorkflowRunTriggeredFrom(workflow_run.triggered_from), ) workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=app_model.tenant_id, user=user, app_id=generate_entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py b/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py index d78c30ebfee..c0e447f5afe 100644 --- a/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py +++ b/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py @@ -146,6 +146,7 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any], session_factory = sessionmaker(bind=db.engine, expire_on_commit=False) workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=pipeline.tenant_id, user=account, app_id=entity.app_config.app_id, triggered_from=WorkflowRunTriggeredFrom.RAG_PIPELINE_RUN, @@ -154,6 +155,7 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any], workflow_node_execution_repository = ( DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=pipeline.tenant_id, user=account, app_id=entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.RAG_PIPELINE_RUN, diff --git a/api/tasks/rag_pipeline/rag_pipeline_run_task.py b/api/tasks/rag_pipeline/rag_pipeline_run_task.py index 6b51a03c1ad..1e8ab2b79f0 100644 --- a/api/tasks/rag_pipeline/rag_pipeline_run_task.py +++ b/api/tasks/rag_pipeline/rag_pipeline_run_task.py @@ -160,6 +160,7 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any], session_factory = sessionmaker(bind=db.engine, expire_on_commit=False) workflow_execution_repository = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=session_factory, + tenant_id=pipeline.tenant_id, user=account, app_id=entity.app_config.app_id, triggered_from=WorkflowRunTriggeredFrom.RAG_PIPELINE_RUN, @@ -168,6 +169,7 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any], workflow_node_execution_repository = ( DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, + tenant_id=pipeline.tenant_id, user=account, app_id=entity.app_config.app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.RAG_PIPELINE_RUN, diff --git a/api/tests/test_containers_integration_tests/core/workflow/test_human_input_resume_node_execution.py b/api/tests/test_containers_integration_tests/core/workflow/test_human_input_resume_node_execution.py index 60f3996e2d0..1b74f4a5851 100644 --- a/api/tests/test_containers_integration_tests/core/workflow/test_human_input_resume_node_execution.py +++ b/api/tests/test_containers_integration_tests/core/workflow/test_human_input_resume_node_execution.py @@ -267,12 +267,14 @@ class TestHumanInputResumeNodeExecutionIntegration: ) execution_repo = SQLAlchemyWorkflowExecutionRepository( session_factory=self.session.get_bind(), + tenant_id=self.tenant.id, user=self.account, app_id=self.app.id, triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, ) node_execution_repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=self.session.get_bind(), + tenant_id=self.tenant.id, user=self.account, app_id=self.app.id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_workflow_node_execution_repository.py b/api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_workflow_node_execution_repository.py index fa78f1c28b1..f5c624d3154 100644 --- a/api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_workflow_node_execution_repository.py +++ b/api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_workflow_node_execution_repository.py @@ -40,8 +40,11 @@ def _create_account_with_tenant(session: Session) -> Account: def _make_repo(session: Session, account: Account, app_id: str) -> SQLAlchemyWorkflowNodeExecutionRepository: engine = session.get_bind() assert isinstance(engine, Engine) + tenant_id = account.current_tenant_id + assert tenant_id is not None return SQLAlchemyWorkflowNodeExecutionRepository( session_factory=sessionmaker(bind=engine, expire_on_commit=False), + tenant_id=tenant_id, user=account, app_id=app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py index a857dfdca0f..50c5a6811db 100644 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py @@ -81,13 +81,15 @@ def test_generate_includes_parent_trace_context_in_extras(monkeypatch): "core.app.apps.workflow.app_generator.file_factory.build_from_mappings", lambda *args, **kwargs: [] ) monkeypatch.setattr("core.app.apps.workflow.app_generator.TraceQueueManager", MagicMock()) + workflow_execution_factory = MagicMock(return_value=MagicMock()) + workflow_node_execution_factory = MagicMock(return_value=MagicMock()) monkeypatch.setattr( "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - MagicMock(return_value=MagicMock()), + workflow_execution_factory, ) monkeypatch.setattr( "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - MagicMock(return_value=MagicMock()), + workflow_node_execution_factory, ) monkeypatch.setattr("core.app.apps.workflow.app_generator.db", SimpleNamespace(engine=MagicMock())) monkeypatch.setattr(generator, "_prepare_user_inputs", lambda *, user_inputs, **kwargs: user_inputs) @@ -134,6 +136,8 @@ def test_generate_includes_parent_trace_context_in_extras(monkeypatch): "parent_node_execution_id": "outer-node-execution-1", } assert extras["trace_session_id"] == "session-1" + assert workflow_execution_factory.call_args.kwargs["tenant_id"] == "tenant-1" + assert workflow_node_execution_factory.call_args.kwargs["tenant_id"] == "tenant-1" def test_resume_delegates_to_generate(mocker: MockerFixture): diff --git a/api/tests/unit_tests/core/repositories/test_celery_workflow_execution_repository.py b/api/tests/unit_tests/core/repositories/test_celery_workflow_execution_repository.py index 3d3322094eb..b158061dbe6 100644 --- a/api/tests/unit_tests/core/repositories/test_celery_workflow_execution_repository.py +++ b/api/tests/unit_tests/core/repositories/test_celery_workflow_execution_repository.py @@ -17,6 +17,8 @@ from libs.datetime_utils import naive_utc_now from models import Account, EndUser from models.enums import WorkflowRunTriggeredFrom +RESOURCE_TENANT_ID = "resource-tenant-id" + @pytest.fixture def mock_session_factory(): @@ -71,12 +73,13 @@ class TestCeleryWorkflowExecutionRepository: repo = CeleryWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id=app_id, triggered_from=triggered_from, ) - assert repo._tenant_id == mock_account.current_tenant_id + assert repo._tenant_id == RESOURCE_TENANT_ID assert repo._app_id == app_id assert repo._triggered_from == triggered_from assert repo._creator_user_id == mock_account.id @@ -86,13 +89,14 @@ class TestCeleryWorkflowExecutionRepository: """Test repository initialization basic functionality.""" repo = CeleryWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, ) # Verify basic initialization - assert repo._tenant_id == mock_account.current_tenant_id + assert repo._tenant_id == RESOURCE_TENANT_ID assert repo._app_id == "test-app" assert repo._triggered_from == WorkflowRunTriggeredFrom.DEBUGGING @@ -100,12 +104,13 @@ class TestCeleryWorkflowExecutionRepository: """Test repository initialization with EndUser.""" repo = CeleryWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_end_user, app_id="test-app", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, ) - assert repo._tenant_id == mock_end_user.tenant_id + assert repo._tenant_id == RESOURCE_TENANT_ID def test_init_without_tenant_id_raises_error(self, mock_session_factory): """Test that initialization fails without tenant_id.""" @@ -114,19 +119,37 @@ class TestCeleryWorkflowExecutionRepository: user.current_tenant_id = None user.id = str(uuid4()) - with pytest.raises(ValueError, match="User must have a tenant_id"): + with pytest.raises(ValueError, match="tenant_id is required"): CeleryWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id="", user=user, app_id="test-app", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, ) + def test_init_uses_resource_tenant_when_account_has_no_current_tenant(self, mock_session_factory): + user = Mock(spec=Account) + user.current_tenant_id = None + user.id = str(uuid4()) + + repo = CeleryWorkflowExecutionRepository( + session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=user, + app_id="test-app", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + + assert repo._tenant_id == RESOURCE_TENANT_ID + assert repo._creator_user_id == user.id + @patch("core.repositories.celery_workflow_execution_repository.save_workflow_execution_task") def test_save_queues_celery_task(self, mock_task, mock_session_factory, mock_account, sample_workflow_execution): """Test that save operation queues a Celery task without tracking.""" repo = CeleryWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, @@ -139,7 +162,7 @@ class TestCeleryWorkflowExecutionRepository: call_args = mock_task.delay.call_args[1] assert call_args["execution_data"] == sample_workflow_execution.model_dump() - assert call_args["tenant_id"] == mock_account.current_tenant_id + assert call_args["tenant_id"] == RESOURCE_TENANT_ID assert call_args["app_id"] == "test-app" assert call_args["triggered_from"] == WorkflowRunTriggeredFrom.APP_RUN assert call_args["creator_user_id"] == mock_account.id @@ -156,6 +179,7 @@ class TestCeleryWorkflowExecutionRepository: repo = CeleryWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, @@ -171,6 +195,7 @@ class TestCeleryWorkflowExecutionRepository: """Test that save operation works in fire-and-forget mode.""" repo = CeleryWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, @@ -187,6 +212,7 @@ class TestCeleryWorkflowExecutionRepository: """Test multiple save operations work correctly.""" repo = CeleryWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, @@ -224,6 +250,7 @@ class TestCeleryWorkflowExecutionRepository: """Test save operation with different user types.""" repo = CeleryWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=mock_end_user.tenant_id, user=mock_end_user, app_id="test-app", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, diff --git a/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py b/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py index 05b4f3a053c..018626c450c 100644 --- a/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py @@ -21,6 +21,8 @@ from libs.datetime_utils import naive_utc_now from models import Account, EndUser from models.workflow import WorkflowNodeExecutionTriggeredFrom +RESOURCE_TENANT_ID = "resource-tenant-id" + @pytest.fixture def mock_session_factory(): @@ -79,12 +81,13 @@ class TestCeleryWorkflowNodeExecutionRepository: repo = CeleryWorkflowNodeExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id=app_id, triggered_from=triggered_from, ) - assert repo._tenant_id == mock_account.current_tenant_id + assert repo._tenant_id == RESOURCE_TENANT_ID assert repo._app_id == app_id assert repo._triggered_from == triggered_from assert repo._creator_user_id == mock_account.id @@ -94,6 +97,7 @@ class TestCeleryWorkflowNodeExecutionRepository: """Test repository initialization with cache properly initialized.""" repo = CeleryWorkflowNodeExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, @@ -106,12 +110,13 @@ class TestCeleryWorkflowNodeExecutionRepository: """Test repository initialization with EndUser.""" repo = CeleryWorkflowNodeExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_end_user, app_id="test-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) - assert repo._tenant_id == mock_end_user.tenant_id + assert repo._tenant_id == RESOURCE_TENANT_ID def test_init_without_tenant_id_raises_error(self, mock_session_factory): """Test that initialization fails without tenant_id.""" @@ -120,14 +125,31 @@ class TestCeleryWorkflowNodeExecutionRepository: user.current_tenant_id = None user.id = str(uuid4()) - with pytest.raises(ValueError, match="User must have a tenant_id"): + with pytest.raises(ValueError, match="tenant_id is required"): CeleryWorkflowNodeExecutionRepository( session_factory=mock_session_factory, + tenant_id="", user=user, app_id="test-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) + def test_init_uses_resource_tenant_when_account_has_no_current_tenant(self, mock_session_factory): + user = Mock(spec=Account) + user.current_tenant_id = None + user.id = str(uuid4()) + + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=user, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + assert repo._tenant_id == RESOURCE_TENANT_ID + assert repo._creator_user_id == user.id + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") def test_save_caches_and_queues_celery_task( self, mock_task, mock_session_factory, mock_account, sample_workflow_node_execution @@ -135,6 +157,7 @@ class TestCeleryWorkflowNodeExecutionRepository: """Test that save operation caches execution and queues a Celery task.""" repo = CeleryWorkflowNodeExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -147,7 +170,7 @@ class TestCeleryWorkflowNodeExecutionRepository: call_args = mock_task.delay.call_args[1] assert call_args["execution_data"] == sample_workflow_node_execution.model_dump() - assert call_args["tenant_id"] == mock_account.current_tenant_id + assert call_args["tenant_id"] == RESOURCE_TENANT_ID assert call_args["app_id"] == "test-app" assert call_args["triggered_from"] == WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN assert call_args["creator_user_id"] == mock_account.id @@ -172,6 +195,7 @@ class TestCeleryWorkflowNodeExecutionRepository: repo = CeleryWorkflowNodeExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -187,6 +211,7 @@ class TestCeleryWorkflowNodeExecutionRepository: """Test that get_by_workflow_execution retrieves executions from cache.""" repo = CeleryWorkflowNodeExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -209,6 +234,7 @@ class TestCeleryWorkflowNodeExecutionRepository: """Test get_by_workflow_execution without order configuration.""" repo = CeleryWorkflowNodeExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -224,6 +250,7 @@ class TestCeleryWorkflowNodeExecutionRepository: """Test cache operations work correctly.""" repo = CeleryWorkflowNodeExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -245,6 +272,7 @@ class TestCeleryWorkflowNodeExecutionRepository: """Test multiple executions for the same workflow.""" repo = CeleryWorkflowNodeExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -296,6 +324,7 @@ class TestCeleryWorkflowNodeExecutionRepository: """Test ordering functionality works correctly.""" repo = CeleryWorkflowNodeExecutionRepository( session_factory=mock_session_factory, + tenant_id=mock_account.current_tenant_id, user=mock_account, app_id="test-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/tests/unit_tests/core/repositories/test_factory.py b/api/tests/unit_tests/core/repositories/test_factory.py index f686ab13a84..1b47cb64a70 100644 --- a/api/tests/unit_tests/core/repositories/test_factory.py +++ b/api/tests/unit_tests/core/repositories/test_factory.py @@ -22,6 +22,8 @@ from models import Account, EndUser from models.enums import WorkflowRunTriggeredFrom from models.workflow import WorkflowNodeExecutionTriggeredFrom +RESOURCE_TENANT_ID = "resource-tenant-id" + class TestRepositoryFactory: """Test cases for RepositoryFactory.""" @@ -72,6 +74,7 @@ class TestRepositoryFactory: with patch("core.repositories.factory.import_string", return_value=mock_repository_class, autospec=True): result = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_user, app_id=app_id, triggered_from=triggered_from, @@ -80,6 +83,7 @@ class TestRepositoryFactory: # Verify the repository was created with correct parameters mock_repository_class.assert_called_once_with( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_user, app_id=app_id, triggered_from=triggered_from, @@ -98,6 +102,7 @@ class TestRepositoryFactory: with pytest.raises(RepositoryImportError) as exc_info: DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_user, app_id="test-app-id", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, @@ -122,6 +127,7 @@ class TestRepositoryFactory: with pytest.raises(RepositoryImportError) as exc_info: DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_user, app_id="test-app-id", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, @@ -149,6 +155,7 @@ class TestRepositoryFactory: with patch("core.repositories.factory.import_string", return_value=mock_repository_class, autospec=True): result = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_user, app_id=app_id, triggered_from=triggered_from, @@ -157,6 +164,7 @@ class TestRepositoryFactory: # Verify the repository was created with correct parameters mock_repository_class.assert_called_once_with( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_user, app_id=app_id, triggered_from=triggered_from, @@ -175,6 +183,7 @@ class TestRepositoryFactory: with pytest.raises(RepositoryImportError) as exc_info: DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_user, app_id="test-app-id", triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, @@ -199,6 +208,7 @@ class TestRepositoryFactory: with pytest.raises(RepositoryImportError) as exc_info: DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_user, app_id="test-app-id", triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, @@ -232,6 +242,7 @@ class TestRepositoryFactory: with patch("core.repositories.factory.import_string", return_value=mock_repository_class, autospec=True): result = DifyCoreRepositoryFactory.create_workflow_execution_repository( session_factory=mock_engine, # Using Engine instead of sessionmaker + tenant_id=RESOURCE_TENANT_ID, user=mock_user, app_id=app_id, triggered_from=triggered_from, @@ -240,6 +251,7 @@ class TestRepositoryFactory: # Verify the repository was created with correct parameters mock_repository_class.assert_called_once_with( session_factory=mock_engine, + tenant_id=RESOURCE_TENANT_ID, user=mock_user, app_id=app_id, triggered_from=triggered_from, diff --git a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_execution_repository.py b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_execution_repository.py index a08c5729cb1..f247525aedd 100644 --- a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_execution_repository.py +++ b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_execution_repository.py @@ -12,6 +12,8 @@ from graphon.enums import WorkflowExecutionStatus, WorkflowType from models import Account, CreatorUserRole, EndUser, WorkflowRun from models.enums import WorkflowRunTriggeredFrom +RESOURCE_TENANT_ID = "resource-tenant-id" + @pytest.fixture def mock_session_factory(): @@ -74,11 +76,15 @@ class TestSQLAlchemyWorkflowExecutionRepository: triggered_from = WorkflowRunTriggeredFrom.APP_RUN repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=mock_session_factory, user=mock_account, app_id=app_id, triggered_from=triggered_from + session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=mock_account, + app_id=app_id, + triggered_from=triggered_from, ) assert repo._session_factory == mock_session_factory - assert repo._tenant_id == mock_account.current_tenant_id + assert repo._tenant_id == RESOURCE_TENANT_ID assert repo._app_id == app_id assert repo._triggered_from == triggered_from assert repo._creator_user_id == mock_account.id @@ -87,6 +93,7 @@ class TestSQLAlchemyWorkflowExecutionRepository: def test_init_with_engine(self, mock_engine, mock_account): repo = SQLAlchemyWorkflowExecutionRepository( session_factory=mock_engine, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test_app_id", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, @@ -98,28 +105,60 @@ class TestSQLAlchemyWorkflowExecutionRepository: def test_init_invalid_session_factory(self, mock_account): with pytest.raises(ValueError, match="Invalid session_factory type"): SQLAlchemyWorkflowExecutionRepository( - session_factory="invalid", user=mock_account, app_id=None, triggered_from=None + session_factory="invalid", + tenant_id=RESOURCE_TENANT_ID, + user=mock_account, + app_id=None, + triggered_from=None, ) def test_init_no_tenant_id(self, mock_session_factory): user = MagicMock(spec=Account) user.current_tenant_id = None - with pytest.raises(ValueError, match="User must have a tenant_id"): + with pytest.raises(ValueError, match="tenant_id is required"): SQLAlchemyWorkflowExecutionRepository( - session_factory=mock_session_factory, user=user, app_id=None, triggered_from=None + session_factory=mock_session_factory, + tenant_id="", + user=user, + app_id=None, + triggered_from=None, ) + def test_init_uses_resource_tenant_when_account_has_no_current_tenant(self, mock_session_factory): + user = MagicMock(spec=Account) + user.current_tenant_id = None + user.id = str(uuid4()) + + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=mock_session_factory, + tenant_id="resource-tenant-id", + user=user, + app_id="test-app", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + + assert repo._tenant_id == "resource-tenant-id" + assert repo._creator_user_id == user.id + def test_init_with_end_user(self, mock_session_factory, mock_end_user): repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=mock_session_factory, user=mock_end_user, app_id=None, triggered_from=None + session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=mock_end_user, + app_id=None, + triggered_from=None, ) - assert repo._tenant_id == mock_end_user.tenant_id + assert repo._tenant_id == RESOURCE_TENANT_ID assert repo._creator_user_role == CreatorUserRole.END_USER def test_to_domain_model(self, mock_session_factory, mock_account): repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=mock_session_factory, user=mock_account, app_id=None, triggered_from=None + session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=mock_account, + app_id=None, + triggered_from=None, ) db_model = MagicMock(spec=WorkflowRun) @@ -149,6 +188,7 @@ class TestSQLAlchemyWorkflowExecutionRepository: def test_to_db_model(self, mock_session_factory, mock_account, sample_workflow_execution): repo = SQLAlchemyWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test_app", triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, @@ -171,6 +211,7 @@ class TestSQLAlchemyWorkflowExecutionRepository: def test_to_db_model_edge_cases(self, mock_session_factory, mock_account, sample_workflow_execution): repo = SQLAlchemyWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test_app", triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, @@ -193,6 +234,7 @@ class TestSQLAlchemyWorkflowExecutionRepository: def test_to_db_model_app_id_none(self, mock_session_factory, mock_account, sample_workflow_execution): repo = SQLAlchemyWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id=None, triggered_from=WorkflowRunTriggeredFrom.APP_RUN, @@ -204,7 +246,11 @@ class TestSQLAlchemyWorkflowExecutionRepository: def test_to_db_model_missing_context(self, mock_session_factory, mock_account, sample_workflow_execution): repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=mock_session_factory, user=mock_account, app_id=None, triggered_from=None + session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=mock_account, + app_id=None, + triggered_from=None, ) # Test triggered_from missing @@ -224,6 +270,7 @@ class TestSQLAlchemyWorkflowExecutionRepository: def test_save(self, mock_session_factory, mock_account, sample_workflow_execution): repo = SQLAlchemyWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test_app", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, @@ -245,6 +292,7 @@ class TestSQLAlchemyWorkflowExecutionRepository: ): repo = SQLAlchemyWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test_app", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, @@ -267,6 +315,7 @@ class TestSQLAlchemyWorkflowExecutionRepository: ): repo = SQLAlchemyWorkflowExecutionRepository( session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=mock_account, app_id="test_app", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, diff --git a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py index 6af7b02d4cc..3289e2b50e8 100644 --- a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py @@ -33,6 +33,8 @@ from models import Account, EndUser from models.enums import ExecutionOffLoadType from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionOffload, WorkflowNodeExecutionTriggeredFrom +RESOURCE_TENANT_ID = "tenant" + def _mock_account(*, tenant_id: str = "tenant", user_id: str = "user") -> Account: user = Mock(spec=Account) @@ -107,6 +109,7 @@ def test_init_accepts_engine_and_sessionmaker_and_sets_role(monkeypatch: pytest. engine: Engine = create_engine("sqlite:///:memory:") repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=engine, + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -116,6 +119,7 @@ def test_init_accepts_engine_and_sessionmaker_and_sets_role(monkeypatch: pytest. sm = Mock(spec=sessionmaker) repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=sm, + tenant_id=RESOURCE_TENANT_ID, user=_mock_end_user(), app_id="app", triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, @@ -131,6 +135,7 @@ def test_init_rejects_invalid_session_factory_type(monkeypatch: pytest.MonkeyPat with pytest.raises(ValueError, match="Invalid session_factory type"): SQLAlchemyWorkflowNodeExecutionRepository( # type: ignore[arg-type] session_factory=object(), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -144,15 +149,36 @@ def test_init_requires_tenant_id(monkeypatch: pytest.MonkeyPatch) -> None: ) user = _mock_account() user.current_tenant_id = None - with pytest.raises(ValueError, match="User must have a tenant_id"): + with pytest.raises(ValueError, match="tenant_id is required"): SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id="", user=user, app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) +def test_init_uses_resource_tenant_when_account_has_no_current_tenant(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "core.repositories.sqlalchemy_workflow_node_execution_repository.FileService", + lambda *_: SimpleNamespace(upload_file=Mock()), + ) + user = _mock_account() + user.current_tenant_id = None + + repo = SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=Mock(spec=sessionmaker), + tenant_id="resource-tenant-id", + user=user, + app_id="app-id", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + assert repo._tenant_id == "resource-tenant-id" + assert repo._creator_user_id == user.id + + def test_create_truncator_uses_config(monkeypatch: pytest.MonkeyPatch) -> None: created: dict[str, Any] = {} @@ -177,6 +203,7 @@ def test_create_truncator_uses_config(monkeypatch: pytest.MonkeyPatch) -> None: repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -206,6 +233,7 @@ def test_to_db_model_requires_constructor_context(monkeypatch: pytest.MonkeyPatc ) repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -214,6 +242,9 @@ def test_to_db_model_requires_constructor_context(monkeypatch: pytest.MonkeyPatc # Happy path: deterministic json dump should be sorted db_model = repo._to_db_model(execution) + assert db_model.tenant_id == RESOURCE_TENANT_ID + assert db_model.created_by == "user" + assert db_model.created_by_role.value == "account" assert json.loads(db_model.inputs or "{}") == {"a": 2, "b": 1} assert json.loads(db_model.execution_metadata or "{}")["total_tokens"] == 1 @@ -229,6 +260,7 @@ def test_to_db_model_requires_creator_user_id_and_role(monkeypatch: pytest.Monke ) repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id="app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -256,6 +288,7 @@ def test_is_duplicate_key_error_and_regenerate_id( ) repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -285,6 +318,7 @@ def test_persist_to_database_updates_existing_and_inserts_new(monkeypatch: pytes session = MagicMock() repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=_session_factory(session), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -318,6 +352,7 @@ def test_truncate_and_upload_returns_none_when_no_values_or_not_truncated(monkey ) repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id="app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -337,8 +372,10 @@ def test_truncate_and_upload_uploads_and_builds_offload(monkeypatch: pytest.Monk uploaded: dict[str, Any] = {} class FakeFileService: - def upload_file(self, *, filename: str, content: bytes, mimetype: str, user: Any): # type: ignore[no-untyped-def] - uploaded.update({"filename": filename, "content": content, "mimetype": mimetype, "user": user}) + def upload_file(self, *, filename: str, content: bytes, mimetype: str, user: Any, tenant_id: str): # type: ignore[no-untyped-def] + uploaded.update( + {"filename": filename, "content": content, "mimetype": mimetype, "user": user, "tenant_id": tenant_id} + ) return SimpleNamespace(id="file-id", key="file-key") monkeypatch.setattr( @@ -348,6 +385,7 @@ def test_truncate_and_upload_uploads_and_builds_offload(monkeypatch: pytest.Monk repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id="app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -363,6 +401,7 @@ def test_truncate_and_upload_uploads_and_builds_offload(monkeypatch: pytest.Monk assert result is not None assert result.truncated_value == {"truncated": True} assert uploaded["filename"].startswith("node_execution_exec_inputs.json") + assert uploaded["tenant_id"] == RESOURCE_TENANT_ID assert result.offload.file_id == "file-id" assert result.offload.type_ == ExecutionOffLoadType.INPUTS @@ -374,6 +413,7 @@ def test_to_domain_model_loads_offloaded_files(monkeypatch: pytest.MonkeyPatch) ) repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -428,6 +468,7 @@ def test_to_domain_model_returns_early_when_no_offload_data(monkeypatch: pytest. ) repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -491,6 +532,7 @@ def test_save_execution_data_handles_existing_db_model_and_truncation(monkeypatc repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=_session_factory(session), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id="app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -538,6 +580,7 @@ def test_save_execution_data_truncates_outputs_and_process_data(monkeypatch: pyt repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=_session_factory(session), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id="app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -583,6 +626,7 @@ def test_save_execution_data_handles_missing_db_model(monkeypatch: pytest.Monkey repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=_session_factory(session), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -608,6 +652,7 @@ def test_save_retries_duplicate_and_logs_non_duplicate( ) repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -647,6 +692,7 @@ def test_save_logs_and_reraises_on_unexpected_error( ) repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -691,6 +737,7 @@ def test_get_db_models_by_workflow_run_orders_and_caches(monkeypatch: pytest.Mon repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=_session_factory(session), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id="app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -728,6 +775,7 @@ def test_get_db_models_by_workflow_run_uses_asc_order(monkeypatch: pytest.Monkey session.scalars.return_value.all.return_value = [] repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=_session_factory(session), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, @@ -743,6 +791,7 @@ def test_get_by_workflow_run_maps_to_domain(monkeypatch: pytest.MonkeyPatch) -> repo = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=Mock(spec=sessionmaker), + tenant_id=RESOURCE_TENANT_ID, user=_mock_account(), app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py b/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py index abdbc72085f..e70745f70b3 100644 --- a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py +++ b/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py @@ -35,6 +35,7 @@ class TestWorkflowNodeExecutionConflictHandling: # Create repository instance self.repository = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=self.mock_session_factory, + tenant_id="test-tenant-id", user=self.mock_user, app_id="test-app-id", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_truncation.py b/api/tests/unit_tests/core/repositories/test_workflow_node_execution_truncation.py index 5af1376a0a9..93ea4271093 100644 --- a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_truncation.py +++ b/api/tests/unit_tests/core/repositories/test_workflow_node_execution_truncation.py @@ -128,6 +128,7 @@ class TestSQLAlchemyWorkflowNodeExecutionRepositoryTruncation: """Create a repository instance for testing.""" return SQLAlchemyWorkflowNodeExecutionRepository( session_factory=MagicMock(spec=Engine), + tenant_id="test-tenant-id", user=mock_user(), app_id="test-app-id", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, diff --git a/api/tests/unit_tests/services/test_file_service.py b/api/tests/unit_tests/services/test_file_service.py index 41b86fda0cb..1954b4d8f0e 100644 --- a/api/tests/unit_tests/services/test_file_service.py +++ b/api/tests/unit_tests/services/test_file_service.py @@ -84,6 +84,27 @@ class TestFileService: mock_db_session.add.assert_called_once_with(result) mock_db_session.commit.assert_called_once() + def test_upload_file_uses_explicit_resource_tenant(self, file_service: FileService): + user = MagicMock(spec=Account) + user.id = "user-id" + + with ( + patch("services.file_service.storage") as mock_storage, + patch("services.file_service.extract_tenant_id") as mock_extract_tenant_id, + patch("services.file_service.file_helpers.get_signed_file_url"), + ): + result = file_service.upload_file( + filename="test.txt", + content=b"test", + mimetype="text/plain", + user=user, + tenant_id="resource-tenant-id", + ) + + assert result.tenant_id == "resource-tenant-id" + assert mock_storage.save.call_args.args[0].startswith("upload_files/resource-tenant-id/") + mock_extract_tenant_id.assert_not_called() + def test_upload_file_invalid_characters(self, file_service): with pytest.raises(ValueError, match="Filename contains invalid characters"): file_service.upload_file(filename="invalid/file.txt", content=b"", mimetype="text/plain", user=MagicMock()) diff --git a/api/tests/unit_tests/tasks/test_workflow_execute_task.py b/api/tests/unit_tests/tasks/test_workflow_execute_task.py index d5a32e082ee..544aaa73234 100644 --- a/api/tests/unit_tests/tasks/test_workflow_execute_task.py +++ b/api/tests/unit_tests/tasks/test_workflow_execute_task.py @@ -451,6 +451,46 @@ def test_app_runner_streaming_failure_publishes_started_then_failed_workflow_fin assert finished_payload["data"]["files"] == [] +def test_app_runner_resolves_account_without_switching_tenant(monkeypatch: pytest.MonkeyPatch): + exec_params = AppExecutionParams( + app_id="app-id", + workflow_id="workflow-id", + tenant_id="resource-tenant-id", + app_mode=AppMode.WORKFLOW, + user={"TYPE": "account", "user_id": "user-id"}, + args={"inputs": {}}, + invoke_from=InvokeFrom.EXPLORE, + streaming=True, + workflow_run_id="workflow-run-id", + ) + runner = _AppRunner(session_factory=MagicMock(), exec_params=exec_params) + account = MagicMock() + session = MagicMock() + session.get.return_value = account + monkeypatch.setattr(runner, "_session", lambda: nullcontext(session)) + + resolved_user = runner._resolve_user() + + assert resolved_user is account + account.set_tenant_id_with_session.assert_not_called() + + +def test_resolve_account_for_run_without_switching_tenant(): + account = MagicMock() + session = MagicMock() + session.get.return_value = account + workflow_run = MagicMock( + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-id", + tenant_id="resource-tenant-id", + ) + + resolved_user = workflow_execute_task_module._resolve_user_for_run(session, workflow_run) + + assert resolved_user is account + account.set_tenant_id_with_session.assert_not_called() + + def test_app_runner_streaming_failure_keeps_existing_pre_runtime_helper_behavior( mock_topic: MagicMock, monkeypatch: pytest.MonkeyPatch, @@ -717,7 +757,7 @@ def test_resume_advanced_chat_publishes_events_for_originally_blocking_runs(monk session = MagicMock() _resume_advanced_chat( - app_model=SimpleNamespace(id="app-id"), + app_model=SimpleNamespace(id="app-id", tenant_id="resource-tenant-id"), workflow=workflow, user=MagicMock(), conversation=SimpleNamespace(id="conversation-id"), @@ -773,7 +813,7 @@ def test_resume_workflow_publishes_events_for_originally_blocking_runs(monkeypat pause_entity = MagicMock() _resume_workflow( - app_model=SimpleNamespace(id="app-id"), + app_model=SimpleNamespace(id="app-id", tenant_id="resource-tenant-id"), workflow=workflow, user=MagicMock(), generate_entity=generate_entity, @@ -829,7 +869,7 @@ def test_resume_workflow_ignores_missing_old_pause_after_repause(monkeypatch: py pause_entity = MagicMock() _resume_workflow( - app_model=SimpleNamespace(id="app-id"), + app_model=SimpleNamespace(id="app-id", tenant_id="resource-tenant-id"), workflow=workflow, user=MagicMock(), generate_entity=generate_entity,