Add "or_" filtering (PrefectHQ/orion#2139)

Add or_ filtering

src/prefect/orion/schemas/filters.py

@@ -14,6 +14,7 @@
 
 import prefect.orion.schemas as schemas
 from prefect.orion.utilities.schemas import PrefectBaseModel
+from prefect.utilities.collections import AutoEnum
 
 if TYPE_CHECKING:
     from prefect.orion.database.interface import OrionDBInterface
@@ -23,6 +24,13 @@
 #       present in the schemas module
 
 
+class Operator(AutoEnum):
+    """Operators for combining filter criteria."""
+
+    and_ = AutoEnum.auto()
+    or_ = AutoEnum.auto()
+
+
 class PrefectFilterBaseModel(PrefectBaseModel):
     """Base model for Prefect filters"""
 
@@ -32,13 +40,30 @@ class Config:
     def as_sql_filter(self, db: "OrionDBInterface") -> BooleanClauseList:
         """Generate SQL filter from provided filter parameters. If no filters parameters are available, return a TRUE filter."""
         filters = self._get_filter_list(db)
-        return sa.and_(*filters) if filters else sa.and_(True)
+        if not filters:
+            return True
+        return sa.and_(*filters)
 
     def _get_filter_list(self, db: "OrionDBInterface") -> List:
         """Return a list of boolean filter statements based on filter parameters"""
         raise NotImplementedError("_get_filter_list must be implemented")
 
 
+class PrefectOperatorFilterBaseModel(PrefectFilterBaseModel):
+    """Base model for Prefect filters that combines criteria with a user-provided operator"""
+
+    operator: Operator = Field(
+        default=Operator.and_,
+        description="Operator for combining filter criteria. Defaults to 'and_'.",
+    )
+
+    def as_sql_filter(self, db: "OrionDBInterface") -> BooleanClauseList:
+        filters = self._get_filter_list(db)
+        if not filters:
+            return True
+        return sa.and_(*filters) if self.operator == Operator.and_ else sa.or_(*filters)
+
+
 class FlowFilterId(PrefectFilterBaseModel):
     """Filter by `Flow.id`."""
 
@@ -79,7 +104,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class FlowFilterTags(PrefectFilterBaseModel):
+class FlowFilterTags(PrefectOperatorFilterBaseModel):
     """Filter by `Flow.tags`."""
 
     all_: List[str] = Field(
@@ -100,7 +125,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class FlowFilter(PrefectFilterBaseModel):
+class FlowFilter(PrefectOperatorFilterBaseModel):
     """Filter for flows. Only flows matching all criteria will be returned."""
 
     id: Optional[FlowFilterId] = Field(
@@ -169,7 +194,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class FlowRunFilterTags(PrefectFilterBaseModel):
+class FlowRunFilterTags(PrefectOperatorFilterBaseModel):
     """Filter by `FlowRun.tags`."""
 
     all_: List[str] = Field(
@@ -194,7 +219,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class FlowRunFilterDeploymentId(PrefectFilterBaseModel):
+class FlowRunFilterDeploymentId(PrefectOperatorFilterBaseModel):
     """Filter by `FlowRun.deployment_id`."""
 
     any_: List[UUID] = Field(
@@ -340,7 +365,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class FlowRunFilterParentTaskRunId(PrefectFilterBaseModel):
+class FlowRunFilterParentTaskRunId(PrefectOperatorFilterBaseModel):
     """Filter by `FlowRun.parent_task_run_id`."""
 
     any_: List[UUID] = Field(
@@ -363,7 +388,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class FlowRunFilter(PrefectFilterBaseModel):
+class FlowRunFilter(PrefectOperatorFilterBaseModel):
     """Filter flow runs. Only flow runs matching all criteria will be returned"""
 
     id: Optional[FlowRunFilterId] = Field(
@@ -464,7 +489,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class TaskRunFilterTags(PrefectFilterBaseModel):
+class TaskRunFilterTags(PrefectOperatorFilterBaseModel):
     """Filter by `TaskRun.tags`."""
 
     all_: List[str] = Field(
@@ -573,7 +598,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class TaskRunFilter(PrefectFilterBaseModel):
+class TaskRunFilter(PrefectOperatorFilterBaseModel):
     """Filter task runs. Only task runs matching all criteria will be returned"""
 
     id: Optional[TaskRunFilterId] = Field(
@@ -669,7 +694,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class DeploymentFilterTags(PrefectFilterBaseModel):
+class DeploymentFilterTags(PrefectOperatorFilterBaseModel):
     """Filter by `Deployment.tags`."""
 
     all_: List[str] = Field(
@@ -694,7 +719,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class DeploymentFilter(PrefectFilterBaseModel):
+class DeploymentFilter(PrefectOperatorFilterBaseModel):
     """Filter for deployments. Only deployments matching all criteria will be returned."""
 
     id: Optional[DeploymentFilterId] = Field(
@@ -808,7 +833,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class LogFilter(PrefectFilterBaseModel):
+class LogFilter(PrefectOperatorFilterBaseModel):
     """Filter logs. Only logs matching all criteria will be returned"""
 
     level: Optional[LogFilterLevel] = Field(
@@ -936,7 +961,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class BlockSchemaFilter(PrefectFilterBaseModel):
+class BlockSchemaFilter(PrefectOperatorFilterBaseModel):
     """Filter BlockSchemas"""
 
     block_type_id: Optional[BlockSchemaFilterBlockTypeId] = Field(
@@ -989,7 +1014,7 @@ def _get_filter_list(self, db: "OrionDBInterface") -> List:
         return filters
 
 
-class BlockDocumentFilter(PrefectFilterBaseModel):
+class BlockDocumentFilter(PrefectOperatorFilterBaseModel):
     """Filter BlockDocuments. Only BlockDocuments matching all criteria will be returned"""
 
     is_anonymous: Optional[BlockDocumentFilterIsAnonymous] = Field(

tests/orion/models/test_filters.py

@@ -352,6 +352,16 @@ class TestCountFlowsModels:
             ),
             1,
         ],
+        # empty filter
+        [dict(flow_filter=filters.FlowFilter()), 4],
+        # multiple empty filters
+        [
+            dict(
+                flow_filter=filters.FlowFilter(),
+                flow_run_filter=filters.FlowRunFilter(),
+            ),
+            4,
+        ],
     ]
 
     @pytest.mark.parametrize("kwargs,expected", params)
@@ -549,6 +559,16 @@ class TestCountFlowRunModels:
             ),
             1,
         ],
+        # empty filter
+        [dict(flow_filter=filters.FlowFilter()), 12],
+        # multiple empty filters
+        [
+            dict(
+                flow_filter=filters.FlowFilter(),
+                flow_run_filter=filters.FlowRunFilter(),
+            ),
+            12,
+        ],
     ]
 
     @pytest.mark.parametrize("kwargs,expected", params)
@@ -724,6 +744,16 @@ class TestCountTaskRunsModels:
             ),
             0,
         ],
+        # empty filter
+        [dict(flow_filter=filters.FlowFilter()), 10],
+        # multiple empty filters
+        [
+            dict(
+                flow_filter=filters.FlowFilter(),
+                flow_run_filter=filters.FlowRunFilter(),
+            ),
+            10,
+        ],
     ]
 
     @pytest.mark.parametrize("kwargs,expected", params)
@@ -875,6 +905,16 @@ class TestCountDeploymentModels:
             ),
             0,
         ],
+        # empty filter
+        [dict(flow_filter=filters.FlowFilter()), 3],
+        # multiple empty filters
+        [
+            dict(
+                flow_filter=filters.FlowFilter(),
+                flow_run_filter=filters.FlowRunFilter(),
+            ),
+            3,
+        ],
     ]
 
     @pytest.mark.parametrize("kwargs,expected", params)

tests/orion/models/test_flow_runs.py

@@ -905,6 +905,17 @@ async def test_read_flow_runs_filters_by_multiple_criteria(self, flow, session):
         )
         assert len(result) == 0
 
+        # filter using OR
+        result = await models.flow_runs.read_flow_runs(
+            session=session,
+            flow_run_filter=schemas.filters.FlowRunFilter(
+                operator="or_",
+                id=schemas.filters.FlowRunFilterId(any_=[flow_run_2.id]),
+                tags=schemas.filters.FlowRunFilterTags(all_=["blue"]),
+            ),
+        )
+        assert {res.id for res in result} == {flow_run_1.id, flow_run_2.id}
+
     async def test_read_flow_runs_filters_by_flow_criteria(self, flow, session):
         flow_run_1 = await models.flow_runs.create_flow_run(
             session=session,