refactor: filter_for_user pattern to support ServicePrincipal (HEXA-1619) (#1835)

* fix: add is_service_principal attribute to authentication

* fix: inherit from ServicePrincipal in PipelineRunUser class

* fix: update filter_for_user methods to use UserInterface and include archived workspaces

* test: add unit tests for workspace and connection filtering logic

* feat: add ServicePrincipal mixin and update filter_for_user methods to support UserInterface

* test: add unit tests for organization filtering logic based on user roles

* fix: update filter_for_user methods to use UserInterface and support ServicePrincipal

* fix: update filter_for_user methods to use UserInterface for improved user handling

* fix: update filter_for_user method to use UserInterface for improved user handling

* fix: add workspace_id to pipeline_run for accurate user filtering

* fix: refactor user filtering to utilize UserInterface and improve organization access

* test: add test for pipeline user visibility of organization shared datasets

* fix: add real_user property to ServicePrincipal for accurate user attribution

* fix: add slug property to workspaces for improved identification

* fix: test

* feat: add workspace and workspace_id properties to enhance ServicePrincipal

* fix: remove include_archived parameter from filter_for_user calls for workspaces

* fix: simplify permission checks for ServicePrincipal by removing unnecessary workspace queries

* fix: restrict organization invitations visibility for non-direct members

* Update backend/hexa/user_management/models.py

Co-authored-by: Bram Jans <ik@bramjans.com>

* fix: restrict organization dataset visibility for external collaborators

---------

Co-authored-by: Bram Jans <ik@bramjans.com>

Author: yolanfery

backend/hexa/core/tests/test_pipeline_user_filtering.py

@@ -28,8 +28,17 @@ def setUp(self):
         self.pipeline_run = MagicMock(PipelineRun)
         self.pipeline_run.pipeline = MagicMock(Pipeline)
         self.pipeline_run.pipeline.workspace = self.WORKSPACE1
+        self.pipeline_run.pipeline.workspace_id = self.WORKSPACE1.id
+        self.pipeline_run.user = self.USER_ROOT
         self.pipeline_user = PipelineRunUser(self.pipeline_run)
 
+    def test_real_user_returns_triggering_user(self):
+        self.assertEqual(self.pipeline_user.real_user, self.USER_ROOT)
+
+    def test_real_user_is_none_for_scheduled_run(self):
+        self.pipeline_run.user = None
+        self.assertIsNone(self.pipeline_user.real_user)
+
     def test_pipeline_user_filtering(self):
         Pipeline.objects.create(
             name="Test Pipeline1",

backend/hexa/datasets/models.py

@@ -18,7 +18,12 @@
 from hexa.datasets.api import get_blob
 from hexa.files import storage
 from hexa.metadata.models import MetadataMixin
-from hexa.user_management.models import OrganizationMembershipRole, User
+from hexa.user_management.models import (
+    Organization,
+    ServicePrincipal,
+    User,
+    UserInterface,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -37,90 +42,63 @@ def create_dataset_slug(name: str, workspace):
 
 
 class DatasetQuerySet(BaseQuerySet):
-    @staticmethod
-    def _workspace_query(user):
-        return Q(links__workspace__members=user)
-
-    @staticmethod
-    def _org_shared_query(user):
-        return Q(
-            shared_with_organization=True,
-            workspace__organization__organizationmembership__user=user,
-        )
-
-    @staticmethod
-    def _org_admin_or_owner_query(user):
-        return Q(
-            workspace__organization__organizationmembership__user=user,
-            workspace__organization__organizationmembership__role__in=[
-                OrganizationMembershipRole.ADMIN,
-                OrganizationMembershipRole.OWNER,
-            ],
-        )
-
     @staticmethod
     def optimize_query(qs: models.QuerySet) -> models.QuerySet:
         return qs.select_related(
             "workspace", "workspace__organization", "created_by"
         ).prefetch_related("versions", "links", "links__workspace")
 
-    def filter_for_user(self, user: AnonymousUser | User):
-        from hexa.pipelines.authentication import PipelineRunUser
+    def filter_for_user(self, user: AnonymousUser | UserInterface) -> models.QuerySet:
+        from hexa.workspaces.models import Workspace
 
-        if isinstance(user, PipelineRunUser):
-            return self._filter_for_user_and_query_object(
-                user, models.Q(links__workspace=user.pipeline_run.pipeline.workspace)
-            )
-        else:
-            return self.optimize_query(
-                self._filter_for_user_and_query_object(
-                    user,
-                    self._workspace_query(user) | self._org_shared_query(user),
-                    return_all_if_superuser=True,
-                    return_all_if_organization_admin_or_owner=True,
+        accessible_workspaces = Workspace.objects.filter_for_user(user)
+        accessible_organizations = Organization.objects.filter_for_user(
+            user, direct_membership_only=True
+        )
+        return self.optimize_query(
+            self.filter(
+                Q(links__workspace__in=accessible_workspaces)
+                | Q(
+                    shared_with_organization=True,
+                    workspace__organization__in=accessible_organizations,
                 )
-            )
+            ).distinct()
+        )
 
     def filter_for_workspace_slugs(
-        self, user: AnonymousUser | User, workspace_slugs: list[str]
+        self, user: AnonymousUser | UserInterface, workspace_slugs: list[str]
     ):
-        return self.optimize_query(
-            self._filter_for_user_and_query_object(
-                user,
-                self._workspace_query(user) & Q(workspace__slug__in=workspace_slugs)
-                | (
-                    self._org_shared_query(user)
-                    & Q(links__workspace__slug__in=workspace_slugs)
-                )
-                | (
-                    self._org_admin_or_owner_query(user)
-                    & Q(workspace__slug__in=workspace_slugs)
-                ),
-                return_all_if_superuser=False,
-            )
+        # Datasets auto-link to their primary workspace, so `links__workspace`
+        # uniformly covers both primary and shared scopes — no separate
+        # primary-workspace clause needed.
+        return (
+            self.filter_for_user(user)
+            .filter(links__workspace__slug__in=workspace_slugs)
+            .distinct()
         )
 
 
 class DatasetManager(models.Manager):
     def create_if_has_perm(
         self,
-        principal: User,
+        principal: UserInterface,
         workspace: any,
         *,
         name: str,
         description: str,
         files: list[dict] | None = None,
     ):
-        from hexa.pipelines.authentication import PipelineRunUser
-
-        # FIXME: Use a generic permission system instead of differencing between User and PipelineRunUser
-        if isinstance(principal, PipelineRunUser):
-            if principal.pipeline_run.pipeline.workspace != workspace:
+        if isinstance(principal, ServicePrincipal):
+            if principal.workspace_id != workspace.pk:
                 raise PermissionDenied
         elif not principal.has_perm("datasets.create_dataset", workspace):
             raise PermissionDenied
 
-        created_by = principal if not isinstance(principal, PipelineRunUser) else None
+        created_by = (
+            principal.real_user
+            if isinstance(principal, ServicePrincipal)
+            else principal
+        )
 
         with transaction.atomic():
             dataset = self.create(
@@ -224,7 +202,7 @@ def link(self, principal: User, workspace: any):
 
 
 class DatasetVersionQuerySet(BaseQuerySet):
-    def filter_for_user(self, user: AnonymousUser | User):
+    def filter_for_user(self, user: AnonymousUser | UserInterface):
         # TODO: It should also check workspace where it's added
         return self._filter_for_user_and_query_object(
             user,
@@ -236,25 +214,24 @@ def filter_for_user(self, user: AnonymousUser | User):
 class DatasetVersionManager(models.Manager):
     def create_if_has_perm(
         self,
-        principal: User,
+        principal: UserInterface,
         dataset: Dataset,
         *,
         name: str,
         changelog: str,
         files: list[dict] | None = None,
     ):
-        # FIXME: Use a generic permission system instead of differencing between User and PipelineRunUser
-        from hexa.pipelines.authentication import PipelineRunUser
-
-        if isinstance(principal, PipelineRunUser):
-            if principal.pipeline_run.pipeline.workspace != dataset.workspace:
+        if isinstance(principal, ServicePrincipal):
+            if principal.workspace_id != dataset.workspace_id:
                 raise PermissionDenied
         elif not principal.has_perm("datasets.create_dataset_version", dataset):
             raise PermissionDenied
-        created_by = principal if not isinstance(principal, PipelineRunUser) else None
-        pipeline_run = (
-            principal.pipeline_run if isinstance(principal, PipelineRunUser) else None
+        created_by = (
+            principal.real_user
+            if isinstance(principal, ServicePrincipal)
+            else principal
         )
+        pipeline_run = getattr(principal, "pipeline_run", None)
 
         uploaded_uris = []
         with transaction.atomic():
@@ -357,7 +334,7 @@ def get_file_by_name(self, name: str):
 
 
 class DatasetVersionFileQuerySet(BaseQuerySet):
-    def filter_for_user(self, user: AnonymousUser | User):
+    def filter_for_user(self, user: AnonymousUser | UserInterface):
         return self._filter_for_user_and_query_object(
             user,
             models.Q(
@@ -373,27 +350,25 @@ def filter_by_filename(self, filename: str):
 class DatasetVersionFileManager(models.Manager):
     def create_if_has_perm(
         self,
-        principal: User,
+        principal: UserInterface,
         dataset_version: DatasetVersion,
         *,
         uri: str,
         content_type: str,
     ):
-        # FIXME: Use a generic permission system instead of differencing between User and PipelineRunUser
-        from hexa.pipelines.authentication import PipelineRunUser
-
-        if isinstance(principal, PipelineRunUser):
-            if (
-                principal.pipeline_run.pipeline.workspace
-                != dataset_version.dataset.workspace
-            ):
+        if isinstance(principal, ServicePrincipal):
+            if principal.workspace_id != dataset_version.dataset.workspace_id:
                 raise PermissionDenied
         elif not principal.has_perm(
             "datasets.create_dataset_version_file", dataset_version
         ):
             raise PermissionDenied
 
-        created_by = principal if not isinstance(principal, PipelineRunUser) else None
+        created_by = (
+            principal.real_user
+            if isinstance(principal, ServicePrincipal)
+            else principal
+        )
 
         return self.create(
             dataset_version=dataset_version,
@@ -591,35 +566,22 @@ def filter_for_workspaces(self, workspaces, pinned=None, query=None):
         )
         return qs
 
-    def filter_for_user(self, user: AnonymousUser | User):
-        # FIXME: Use a generic permission system instead of differencing between User and PipelineRunUser
-        from hexa.pipelines.authentication import PipelineRunUser
-
-        if isinstance(user, PipelineRunUser):
-            workspace = user.pipeline_run.pipeline.workspace
-            return self.optimize_query(
-                self._filter_for_user_and_query_object(
-                    user,
-                    models.Q(workspace=workspace)
-                    | models.Q(
-                        dataset__shared_with_organization=True,
-                        workspace__organization=workspace.organization,
-                    ),
-                )
-            )
-        else:
-            return self.optimize_query(
-                self._filter_for_user_and_query_object(
-                    user,
-                    models.Q(workspace__members=user)
-                    | models.Q(
-                        dataset__shared_with_organization=True,
-                        workspace__organization__organizationmembership__user=user,
-                    ),
-                    return_all_if_superuser=True,
-                    return_all_if_organization_admin_or_owner=True,
+    def filter_for_user(self, user: AnonymousUser | UserInterface) -> models.QuerySet:
+        from hexa.workspaces.models import Workspace
+
+        accessible_workspaces = Workspace.objects.filter_for_user(user)
+        accessible_organizations = Organization.objects.filter_for_user(
+            user, direct_membership_only=True
+        )
+        return self.optimize_query(
+            self.filter(
+                Q(workspace__in=accessible_workspaces)
+                | Q(
+                    dataset__shared_with_organization=True,
+                    workspace__organization__in=accessible_organizations,
                 )
-            )
+            ).distinct()
+        )
 
 
 class DatasetLink(Base):

backend/hexa/datasets/tests/test_models.py

@@ -449,6 +449,45 @@ def test_dataset_filter_for_user_excludes_non_organization_members(self):
         datasets = Dataset.objects.filter_for_user(self.USER_SERENA)
         self.assertNotIn(org_dataset, datasets)
 
+    def test_dataset_filter_for_user_excludes_external_collaborator_in_same_org(self):
+        """An external collaborator — a member of one workspace in the
+        organization but not a member of the organization itself — must not see
+        organization-shared datasets belonging to other workspaces in that org.
+        """
+        shared_workspace = Workspace.objects.create_if_has_perm(
+            self.USER_ADMIN,
+            name="Shared Dataset Workspace",
+            description="Holds the organization-shared dataset",
+            organization=self.ORGANIZATION,
+        )
+        org_dataset = Dataset.objects.create_if_has_perm(
+            self.USER_ADMIN,
+            shared_workspace,
+            name="Org Dataset",
+            description="Organization shared dataset",
+        )
+        org_dataset.shared_with_organization = True
+        org_dataset.save()
+
+        collaborator = User.objects.create_user("collaborator@example.com", "password")
+        collaborator_workspace = Workspace.objects.create_if_has_perm(
+            self.USER_ADMIN,
+            name="Collaborator Workspace",
+            description="The only workspace the collaborator belongs to",
+            organization=self.ORGANIZATION,
+        )
+        WorkspaceMembership.objects.create(
+            workspace=collaborator_workspace,
+            user=collaborator,
+            role=WorkspaceMembershipRole.EDITOR,
+        )
+
+        self.assertNotIn(org_dataset, Dataset.objects.filter_for_user(collaborator))
+        self.assertNotIn(
+            org_dataset.links.first(),
+            DatasetLink.objects.filter_for_user(collaborator),
+        )
+
     def test_dataset_filter_for_user_excludes_non_shared_datasets(self):
         other_workspace = Workspace.objects.create_if_has_perm(
             self.USER_ADMIN,
@@ -898,3 +937,41 @@ def test_pipeline_user_cannot_access_non_shared_dataset_links_from_same_organiza
         links = DatasetLink.objects.filter_for_user(self.pipeline_user)
         link_datasets = [link.dataset for link in links]
         self.assertNotIn(self.DATASET_IN_WORKSPACE_3, link_datasets)
+
+    def test_pipeline_user_sees_org_shared_datasets_via_dataset_queryset(self):
+        """A pipeline run sees datasets shared with its workspace's organization."""
+        self.DATASET_IN_OTHER_WORKSPACE_SAME_ORG.shared_with_organization = True
+        self.DATASET_IN_OTHER_WORKSPACE_SAME_ORG.save()
+        self.DATASET_IN_DIFFERENT_ORG.shared_with_organization = True
+        self.DATASET_IN_DIFFERENT_ORG.save()
+
+        datasets = Dataset.objects.filter_for_user(self.pipeline_user)
+
+        self.assertIn(self.DATASET_IN_PIPELINE_WORKSPACE, datasets)
+        self.assertIn(self.DATASET_IN_OTHER_WORKSPACE_SAME_ORG, datasets)
+        # Same-org but not shared → still hidden.
+        self.assertNotIn(self.DATASET_IN_WORKSPACE_3, datasets)
+        # Different-org dataset → hidden even when shared.
+        self.assertNotIn(self.DATASET_IN_DIFFERENT_ORG, datasets)
+
+    def test_pipeline_user_attributes_dataset_to_triggering_user(self):
+        """`created_by` is set to the user who triggered the run."""
+        self.pipeline_run.user = self.USER_ADMIN
+        dataset = Dataset.objects.create_if_has_perm(
+            self.pipeline_user,
+            self.WORKSPACE,
+            name="Pipeline-created Dataset",
+            description="",
+        )
+        self.assertEqual(dataset.created_by, self.USER_ADMIN)
+
+    def test_pipeline_user_leaves_created_by_null_for_scheduled_run(self):
+        """A scheduled run (no triggering user) leaves created_by null."""
+        self.pipeline_run.user = None
+        dataset = Dataset.objects.create_if_has_perm(
+            self.pipeline_user,
+            self.WORKSPACE,
+            name="Scheduled-created Dataset",
+            description="",
+        )
+        self.assertIsNone(dataset.created_by)