Mirror of https://github.com/roostorg/osprey
github.com/roostorg/osprey
Add error_on_empty parameter to HasLabel UDF (#128)
+141
-1
+40
osprey_worker/src/osprey/engine/stdlib/udfs/labels.py
+40
osprey_worker/src/osprey/engine/stdlib/udfs/labels.py
···
107
107
REMOVED = 'removed'
108
108
109
109
110
+
class EmptyEntityError(Exception):
111
+
"""Raised when an entity has no labels and error_on_empty is True."""
112
+
113
+
def __init__(self, entity: EntityT[Any], label: str):
114
+
self.entity = entity
115
+
self.label = label
116
+
super().__init__(
117
+
f"Entity '{entity.type}/{entity.id}' has no labels. "
118
+
f'This may indicate the labels service failed to fetch data. '
119
+
f"(Checked for label: '{label}')"
120
+
)
121
+
122
+
110
123
class HasLabelArguments(ArgumentsBase):
111
124
entity: EntityT[Any]
112
125
"""An Entity to check for a label on."""
···
118
131
"""Optional: A specific status to check for. Default is 'added'."""
119
132
min_label_age: Optional[TimeDeltaT] = None
120
133
"""Optional: Checks to see if the label was added after a period of time"""
134
+
error_on_empty: bool = False
135
+
"""Optional: If True, raise EmptyEntityError when the entity has no labels at all.
136
+
137
+
WARNING: Only use this for safety-critical rules where a false negative (due to labels
138
+
service returning empty data on failure) could cause dangerous rule evaluations, such as
139
+
incorrectly allowing a known-bad entity through. Do not use this for general label checks.
140
+
141
+
This parameter should only be used when the entity type is guaranteed to have at least one
142
+
label in the labels service. If the entity type is not guaranteed to have labels, the rule
143
+
should add a dummy/sentinel label to the entity before calling HasLabel with error_on_empty=True.
144
+
"""
121
145
122
146
123
147
@dataclass
···
128
152
status: str
129
153
min_label_age: Optional[TimeDeltaT]
130
154
desired_status: Optional[_SimpleStatus]
155
+
error_on_empty: bool
131
156
132
157
133
158
class HasLabel(
···
165
190
166
191
validation_context.add_error(message='unknown label', span=arguments.label.argument_span, hint=hint)
167
192
193
+
def _check_error_on_empty(
194
+
self, entity: EntityT[Any], label: str, entity_labels: EntityLabels, error_on_empty: bool
195
+
) -> None:
196
+
"""Fail-closed check for labels service data integrity.
197
+
198
+
When error_on_empty is True, raises EmptyEntityError if the entity has zero labels.
199
+
This catches cases where the labels service may have failed to fetch data and returned
200
+
an empty default response. Without this check, `not HasLabel(...)` would incorrectly
201
+
evaluate to True, potentially allowing dangerous entities through safety rules.
202
+
"""
203
+
if error_on_empty and len(entity_labels.labels) == 0:
204
+
raise EmptyEntityError(entity, label)
205
+
168
206
def _execute(
169
207
self, execution_context: ExecutionContext, arguments: BatchableHasLabelArguments, entity_labels: EntityLabels
170
208
) -> bool:
209
+
self._check_error_on_empty(arguments.entity, arguments.label, entity_labels, arguments.error_on_empty)
171
210
desired_manual = _ManualType.get(arguments.manual)
172
211
desired_delay = TimeDeltaT.inner_from_optional(arguments.min_label_age)
173
212
label_state = entity_labels.labels.get(arguments.label)
···
237
276
status=arguments.status.value,
238
277
min_label_age=arguments.min_label_age,
239
278
desired_status=self.desired_status,
279
+
error_on_empty=arguments.error_on_empty,
240
280
)
241
281
242
282
def get_batch_routing_key(self, arguments: BatchableHasLabelArguments) -> str:
+101
-1
osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_labels.py
+101
-1
osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_labels.py
···
23
23
from osprey.engine.language_types.labels import LabelStatus
24
24
from osprey.engine.stdlib import get_config_registry
25
25
from osprey.engine.stdlib.udfs.entity import Entity
26
-
from osprey.engine.stdlib.udfs.labels import HasLabel, LabelAdd, LabelRemove
26
+
from osprey.engine.stdlib.udfs.labels import EmptyEntityError, HasLabel, LabelAdd, LabelRemove
27
27
from osprey.engine.stdlib.udfs.rules import Rule, WhenRules
28
28
from osprey.engine.stdlib.udfs.time_delta import TimeDelta
29
29
from osprey.engine.udf.registry import UDFRegistry
···
573
573
}
574
574
)
575
575
assert set(result.extracted_features['__entity_label_mutations']) == set(entity_label_mutation)
576
+
577
+
578
+
def test_error_on_empty_raises_when_entity_has_no_labels(execute_with_result: ExecuteWithResultFunction) -> None:
579
+
"""error_on_empty=True should raise EmptyEntityError when entity has no labels at all."""
580
+
empty_labels = EntityLabels(labels={})
581
+
label_provider = StaticLabelProvider({EntityT('MyEntity', 'my_id'): empty_labels})
582
+
583
+
result = execute_with_result(
584
+
source_with_labels_config(
585
+
"""
586
+
L = HasLabel(
587
+
entity=Entity(type='MyEntity', id='my_id'),
588
+
label='my_label',
589
+
error_on_empty=True,
590
+
)
591
+
""",
592
+
labels={'my_label'},
593
+
),
594
+
udf_helpers=UDFHelpers().set_udf_helper(HasLabel, label_provider),
595
+
)
596
+
597
+
assert len(result.error_infos) > 0
598
+
error = result.error_infos[0].error
599
+
assert isinstance(error, EmptyEntityError)
600
+
assert 'MyEntity/my_id' in str(error)
601
+
assert 'my_label' in str(error)
602
+
603
+
604
+
def test_error_on_empty_returns_true_when_label_exists(execute: ExecuteFunction) -> None:
605
+
"""error_on_empty=True should return True when the entity has labels and the requested label exists."""
606
+
labels = EntityLabels(
607
+
labels={'my_label': LabelState(status=LabelStatus.ADDED, reasons=LabelReasons({'TestReason': LabelReason()}))}
608
+
)
609
+
label_provider = StaticLabelProvider({EntityT('MyEntity', 'my_id'): labels})
610
+
611
+
data = execute(
612
+
source_with_labels_config(
613
+
"""
614
+
L = HasLabel(
615
+
entity=Entity(type='MyEntity', id='my_id'),
616
+
label='my_label',
617
+
error_on_empty=True,
618
+
)
619
+
""",
620
+
labels={'my_label'},
621
+
),
622
+
udf_helpers=UDFHelpers().set_udf_helper(HasLabel, label_provider),
623
+
)
624
+
625
+
assert data == {'L': True}
626
+
627
+
628
+
def test_error_on_empty_returns_false_when_label_missing_but_entity_has_other_labels(
629
+
execute: ExecuteFunction,
630
+
) -> None:
631
+
"""error_on_empty=True should return False when entity has labels but not the requested one."""
632
+
labels = EntityLabels(
633
+
labels={
634
+
'other_label': LabelState(status=LabelStatus.ADDED, reasons=LabelReasons({'TestReason': LabelReason()}))
635
+
}
636
+
)
637
+
label_provider = StaticLabelProvider({EntityT('MyEntity', 'my_id'): labels})
638
+
639
+
data = execute(
640
+
source_with_labels_config(
641
+
"""
642
+
L = HasLabel(
643
+
entity=Entity(type='MyEntity', id='my_id'),
644
+
label='my_label',
645
+
error_on_empty=True,
646
+
)
647
+
""",
648
+
labels={'my_label', 'other_label'},
649
+
),
650
+
udf_helpers=UDFHelpers().set_udf_helper(HasLabel, label_provider),
651
+
)
652
+
653
+
assert data == {'L': False}
654
+
655
+
656
+
def test_error_on_empty_false_returns_false_when_entity_has_no_labels(execute: ExecuteFunction) -> None:
657
+
"""error_on_empty=False (default) should return False when entity has no labels."""
658
+
empty_labels = EntityLabels(labels={})
659
+
label_provider = StaticLabelProvider({EntityT('MyEntity', 'my_id'): empty_labels})
660
+
661
+
data = execute(
662
+
source_with_labels_config(
663
+
"""
664
+
L = HasLabel(
665
+
entity=Entity(type='MyEntity', id='my_id'),
666
+
label='my_label',
667
+
error_on_empty=False,
668
+
)
669
+
""",
670
+
labels={'my_label'},
671
+
),
672
+
udf_helpers=UDFHelpers().set_udf_helper(HasLabel, label_provider),
673
+
)
674
+
675
+
assert data == {'L': False}