acryl-datahub-cloud 0.3.11.1rc8__py3-none-any.whl → 0.3.12rc3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of acryl-datahub-cloud might be problematic. Click here for more details.

Files changed (82) hide show
  1. acryl_datahub_cloud/_codegen_config.json +1 -1
  2. acryl_datahub_cloud/action_request/action_request_owner_source.py +36 -6
  3. acryl_datahub_cloud/datahub_forms_notifications/__init__.py +0 -0
  4. acryl_datahub_cloud/datahub_forms_notifications/forms_notifications_source.py +524 -0
  5. acryl_datahub_cloud/datahub_forms_notifications/get_search_results_total.gql +14 -0
  6. acryl_datahub_cloud/datahub_forms_notifications/query.py +17 -0
  7. acryl_datahub_cloud/datahub_forms_notifications/scroll_forms_for_notification.gql +29 -0
  8. acryl_datahub_cloud/datahub_forms_notifications/send_form_notification_request.gql +5 -0
  9. acryl_datahub_cloud/datahub_usage_reporting/query_builder.py +48 -8
  10. acryl_datahub_cloud/datahub_usage_reporting/usage_feature_reporter.py +49 -40
  11. acryl_datahub_cloud/metadata/_urns/urn_defs.py +2014 -1958
  12. acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/application/__init__.py +19 -0
  13. acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/assertion/__init__.py +2 -2
  14. acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/form/__init__.py +8 -0
  15. acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/notification/__init__.py +19 -0
  16. acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/settings/global/__init__.py +2 -0
  17. acryl_datahub_cloud/metadata/schema.avsc +26713 -26274
  18. acryl_datahub_cloud/metadata/schema_classes.py +1302 -777
  19. acryl_datahub_cloud/metadata/schemas/ApplicationKey.avsc +31 -0
  20. acryl_datahub_cloud/metadata/schemas/ApplicationProperties.avsc +72 -0
  21. acryl_datahub_cloud/metadata/schemas/Applications.avsc +38 -0
  22. acryl_datahub_cloud/metadata/schemas/AssertionAnalyticsRunEvent.avsc +220 -208
  23. acryl_datahub_cloud/metadata/schemas/AssertionInfo.avsc +36 -7
  24. acryl_datahub_cloud/metadata/schemas/AssertionKey.avsc +1 -1
  25. acryl_datahub_cloud/metadata/schemas/AssertionRunEvent.avsc +40 -8
  26. acryl_datahub_cloud/metadata/schemas/{AssertionSummary.avsc → AssertionRunSummary.avsc} +2 -2
  27. acryl_datahub_cloud/metadata/schemas/AssertionsSummary.avsc +14 -0
  28. acryl_datahub_cloud/metadata/schemas/ChartKey.avsc +1 -0
  29. acryl_datahub_cloud/metadata/schemas/ConstraintInfo.avsc +12 -1
  30. acryl_datahub_cloud/metadata/schemas/ContainerKey.avsc +1 -0
  31. acryl_datahub_cloud/metadata/schemas/CorpGroupKey.avsc +2 -1
  32. acryl_datahub_cloud/metadata/schemas/CorpUserKey.avsc +2 -1
  33. acryl_datahub_cloud/metadata/schemas/DashboardKey.avsc +1 -0
  34. acryl_datahub_cloud/metadata/schemas/DataFlowKey.avsc +1 -0
  35. acryl_datahub_cloud/metadata/schemas/DataHubIngestionSourceKey.avsc +2 -1
  36. acryl_datahub_cloud/metadata/schemas/DataHubPolicyInfo.avsc +12 -1
  37. acryl_datahub_cloud/metadata/schemas/DataJobKey.avsc +1 -0
  38. acryl_datahub_cloud/metadata/schemas/DataProductKey.avsc +1 -0
  39. acryl_datahub_cloud/metadata/schemas/DataProductProperties.avsc +1 -1
  40. acryl_datahub_cloud/metadata/schemas/DatasetKey.avsc +1 -0
  41. acryl_datahub_cloud/metadata/schemas/FormAssignmentStatus.avsc +36 -0
  42. acryl_datahub_cloud/metadata/schemas/FormInfo.avsc +6 -0
  43. acryl_datahub_cloud/metadata/schemas/FormKey.avsc +3 -1
  44. acryl_datahub_cloud/metadata/schemas/FormNotifications.avsc +69 -0
  45. acryl_datahub_cloud/metadata/schemas/FormSettings.avsc +30 -0
  46. acryl_datahub_cloud/metadata/schemas/GlobalSettingsInfo.avsc +22 -0
  47. acryl_datahub_cloud/metadata/schemas/GlossaryTermKey.avsc +1 -0
  48. acryl_datahub_cloud/metadata/schemas/MLFeatureKey.avsc +1 -0
  49. acryl_datahub_cloud/metadata/schemas/MLFeatureTableKey.avsc +1 -0
  50. acryl_datahub_cloud/metadata/schemas/MLModelGroupKey.avsc +1 -0
  51. acryl_datahub_cloud/metadata/schemas/MLModelKey.avsc +1 -0
  52. acryl_datahub_cloud/metadata/schemas/MLPrimaryKeyKey.avsc +1 -0
  53. acryl_datahub_cloud/metadata/schemas/MetadataChangeEvent.avsc +12 -1
  54. acryl_datahub_cloud/metadata/schemas/MonitorAnomalyEvent.avsc +21 -9
  55. acryl_datahub_cloud/metadata/schemas/MonitorInfo.avsc +39 -10
  56. acryl_datahub_cloud/metadata/schemas/MonitorSuiteInfo.avsc +1 -1
  57. acryl_datahub_cloud/metadata/schemas/NotebookKey.avsc +1 -0
  58. acryl_datahub_cloud/metadata/schemas/NotificationRequest.avsc +1 -0
  59. acryl_datahub_cloud/metadata/schemas/Operation.avsc +17 -0
  60. acryl_datahub_cloud/metadata/schemas/SubscriptionInfo.avsc +3 -3
  61. acryl_datahub_cloud/metadata/schemas/__init__.py +3 -3
  62. acryl_datahub_cloud/notifications/__init__.py +0 -0
  63. acryl_datahub_cloud/notifications/notification_recipient_builder.py +399 -0
  64. acryl_datahub_cloud/sdk/__init__.py +25 -0
  65. acryl_datahub_cloud/sdk/assertion.py +767 -0
  66. acryl_datahub_cloud/sdk/assertion_input.py +1335 -0
  67. acryl_datahub_cloud/sdk/assertions_client.py +1153 -0
  68. acryl_datahub_cloud/sdk/entities/__init__.py +0 -0
  69. acryl_datahub_cloud/sdk/entities/assertion.py +425 -0
  70. acryl_datahub_cloud/sdk/entities/monitor.py +291 -0
  71. acryl_datahub_cloud/sdk/entities/subscription.py +84 -0
  72. acryl_datahub_cloud/sdk/errors.py +34 -0
  73. acryl_datahub_cloud/sdk/resolver_client.py +39 -0
  74. acryl_datahub_cloud/sdk/subscription_client.py +678 -0
  75. {acryl_datahub_cloud-0.3.11.1rc8.dist-info → acryl_datahub_cloud-0.3.12rc3.dist-info}/METADATA +44 -39
  76. {acryl_datahub_cloud-0.3.11.1rc8.dist-info → acryl_datahub_cloud-0.3.12rc3.dist-info}/RECORD +79 -55
  77. {acryl_datahub_cloud-0.3.11.1rc8.dist-info → acryl_datahub_cloud-0.3.12rc3.dist-info}/WHEEL +1 -1
  78. {acryl_datahub_cloud-0.3.11.1rc8.dist-info → acryl_datahub_cloud-0.3.12rc3.dist-info}/entry_points.txt +1 -0
  79. acryl_datahub_cloud/_sdk_extras/__init__.py +0 -4
  80. acryl_datahub_cloud/_sdk_extras/assertion.py +0 -15
  81. acryl_datahub_cloud/_sdk_extras/assertions_client.py +0 -23
  82. {acryl_datahub_cloud-0.3.11.1rc8.dist-info → acryl_datahub_cloud-0.3.12rc3.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,767 @@
1
+ """
2
+ This module contains the classes that represent assertions. These
3
+ classes are used to provide a user-friendly interface for creating and
4
+ managing assertions.
5
+
6
+ The actual Assertion Entity classes are defined in `metadata-ingestion/src/datahub/sdk`.
7
+ """
8
+
9
+ import logging
10
+ from abc import ABC, abstractmethod
11
+ from datetime import datetime
12
+ from enum import Enum
13
+ from typing import Optional, Union
14
+
15
+ from typing_extensions import Self
16
+
17
+ from acryl_datahub_cloud.sdk.assertion_input import (
18
+ ASSERTION_MONITOR_DEFAULT_TRAINING_LOOKBACK_WINDOW_DAYS,
19
+ DEFAULT_DETECTION_MECHANISM,
20
+ DEFAULT_SCHEDULE,
21
+ DEFAULT_SENSITIVITY,
22
+ AssertionIncidentBehavior,
23
+ DetectionMechanism,
24
+ ExclusionWindowTypes,
25
+ FixedRangeExclusionWindow,
26
+ InferenceSensitivity,
27
+ _DetectionMechanismTypes,
28
+ )
29
+ from acryl_datahub_cloud.sdk.entities.assertion import Assertion
30
+ from acryl_datahub_cloud.sdk.entities.monitor import (
31
+ Monitor,
32
+ _get_nested_field_for_entity_with_default,
33
+ )
34
+ from acryl_datahub_cloud.sdk.errors import SDKNotYetSupportedError
35
+ from datahub.emitter.mce_builder import parse_ts_millis
36
+ from datahub.metadata import schema_classes as models
37
+ from datahub.metadata.urns import AssertionUrn, CorpUserUrn, DatasetUrn, TagUrn
38
+
39
+ logger = logging.getLogger(__name__)
40
+
41
+
42
+ class AssertionMode(Enum):
43
+ """
44
+ The mode of an assertion, e.g. whether it is active or inactive.
45
+ """
46
+
47
+ # Note: Modeled here after MonitorStatus but called AssertionMode in this user facing interface
48
+ # to keep all naming related to assertions.
49
+ ACTIVE = "ACTIVE"
50
+ INACTIVE = "INACTIVE"
51
+ # PASSIVE = "PASSIVE" # Not supported in the user facing interface.
52
+
53
+
54
+ class _HasSchedule:
55
+ """
56
+ Mixin class that provides schedule functionality for assertions.
57
+ """
58
+
59
+ def __init__(self, schedule: models.CronScheduleClass) -> None:
60
+ self._schedule = schedule
61
+
62
+ @property
63
+ def schedule(self) -> models.CronScheduleClass:
64
+ return self._schedule
65
+
66
+ @staticmethod
67
+ def _get_schedule(monitor: Monitor) -> models.CronScheduleClass:
68
+ """Get the schedule from the monitor."""
69
+ assertion_evaluation_specs = _get_nested_field_for_entity_with_default(
70
+ monitor,
71
+ "info.assertionMonitor.assertions",
72
+ [],
73
+ )
74
+ if len(assertion_evaluation_specs) == 0:
75
+ return DEFAULT_SCHEDULE
76
+ assertion_evaluation_spec = assertion_evaluation_specs[0]
77
+ schedule = assertion_evaluation_spec.schedule
78
+ if schedule is None:
79
+ return DEFAULT_SCHEDULE
80
+ return schedule
81
+
82
+
83
+ class _HasSmartFunctionality:
84
+ """
85
+ Mixin class that provides smart functionality for assertions.
86
+ """
87
+
88
+ _SUPPORTED_WITH_FILTER_ASSERTION_TYPES = (
89
+ models.FreshnessAssertionInfoClass,
90
+ models.VolumeAssertionInfoClass,
91
+ )
92
+
93
+ def __init__(
94
+ self,
95
+ *,
96
+ sensitivity: InferenceSensitivity = DEFAULT_SENSITIVITY,
97
+ exclusion_windows: list[ExclusionWindowTypes],
98
+ training_data_lookback_days: int = ASSERTION_MONITOR_DEFAULT_TRAINING_LOOKBACK_WINDOW_DAYS,
99
+ incident_behavior: list[AssertionIncidentBehavior],
100
+ detection_mechanism: Optional[
101
+ _DetectionMechanismTypes
102
+ ] = DEFAULT_DETECTION_MECHANISM,
103
+ ) -> None:
104
+ """
105
+ Initialize the smart functionality mixin.
106
+
107
+ Args:
108
+ sensitivity: The sensitivity of the assertion (low, medium, high).
109
+ exclusion_windows: The exclusion windows of the assertion.
110
+ training_data_lookback_days: The max number of days of data to use for training the assertion.
111
+ incident_behavior: Whether to raise or resolve an incident when the assertion fails / passes.
112
+ detection_mechanism: The detection mechanism of the assertion.
113
+ **kwargs: Additional arguments to pass to the parent class (_Assertion).
114
+ """
115
+ self._sensitivity = sensitivity
116
+ self._exclusion_windows = exclusion_windows
117
+ self._training_data_lookback_days = training_data_lookback_days
118
+ self._incident_behavior = incident_behavior
119
+ self._detection_mechanism = detection_mechanism
120
+
121
+ @property
122
+ def sensitivity(self) -> InferenceSensitivity:
123
+ return self._sensitivity
124
+
125
+ @property
126
+ def exclusion_windows(self) -> list[ExclusionWindowTypes]:
127
+ return self._exclusion_windows
128
+
129
+ @property
130
+ def training_data_lookback_days(self) -> int:
131
+ return self._training_data_lookback_days
132
+
133
+ @property
134
+ def incident_behavior(self) -> list[AssertionIncidentBehavior]:
135
+ return self._incident_behavior
136
+
137
+ @property
138
+ def detection_mechanism(self) -> Optional[_DetectionMechanismTypes]:
139
+ return self._detection_mechanism
140
+
141
+ @staticmethod
142
+ def _get_sensitivity(monitor: Monitor) -> InferenceSensitivity:
143
+ # 1. Check if the monitor has a sensitivity field
144
+ raw_sensitivity = _get_nested_field_for_entity_with_default(
145
+ monitor,
146
+ "info.assertionMonitor.settings.adjustmentSettings.sensitivity.level",
147
+ DEFAULT_SENSITIVITY,
148
+ )
149
+
150
+ # 2. Convert the raw sensitivity to the SDK sensitivity enum (1-3: LOW, 4-6: MEDIUM, 7-10: HIGH)
151
+ return InferenceSensitivity.parse(raw_sensitivity)
152
+
153
+ @staticmethod
154
+ def _get_exclusion_windows(monitor: Monitor) -> list[ExclusionWindowTypes]:
155
+ # 1. Check if the monitor has an exclusion windows field
156
+ raw_windows = monitor.exclusion_windows or []
157
+
158
+ # 2. Convert the raw exclusion windows to the SDK exclusion windows
159
+ exclusion_windows = []
160
+ for raw_window in raw_windows:
161
+ if raw_window.type == models.AssertionExclusionWindowTypeClass.FIXED_RANGE:
162
+ if raw_window.fixedRange is None:
163
+ logger.warning(
164
+ f"Monitor {monitor.urn} has a fixed range exclusion window with no fixed range, skipping"
165
+ )
166
+ continue
167
+ exclusion_windows.append(
168
+ FixedRangeExclusionWindow(
169
+ start=parse_ts_millis(raw_window.fixedRange.startTimeMillis),
170
+ end=parse_ts_millis(raw_window.fixedRange.endTimeMillis),
171
+ )
172
+ )
173
+ else:
174
+ raise SDKNotYetSupportedError(
175
+ f"AssertionExclusionWindowType {raw_window.type}"
176
+ )
177
+ return exclusion_windows
178
+
179
+ @staticmethod
180
+ def _get_training_data_lookback_days(monitor: Monitor) -> int:
181
+ retrieved = monitor.training_data_lookback_days
182
+ if (
183
+ retrieved is None
184
+ ): # Explicitly check for None since retrieved can be 0 which is falsy
185
+ return ASSERTION_MONITOR_DEFAULT_TRAINING_LOOKBACK_WINDOW_DAYS
186
+ assert isinstance(retrieved, int)
187
+ return retrieved
188
+
189
+ @staticmethod
190
+ def _get_detection_mechanism(
191
+ assertion: Assertion,
192
+ monitor: Monitor,
193
+ default: Optional[_DetectionMechanismTypes] = DEFAULT_DETECTION_MECHANISM,
194
+ ) -> Optional[_DetectionMechanismTypes]:
195
+ """Get the detection mechanism from the monitor and assertion."""
196
+ if not _HasSmartFunctionality._has_valid_monitor_info(monitor):
197
+ return default
198
+
199
+ # 1. Check if the assertion has a parameters field
200
+ def _warn_and_return_default_detection_mechanism(
201
+ field_name: str,
202
+ default: Optional[_DetectionMechanismTypes] = DEFAULT_DETECTION_MECHANISM,
203
+ ) -> Optional[_DetectionMechanismTypes]:
204
+ logger.warning(
205
+ f"Monitor {monitor.urn} does not have an `{field_name}` field, defaulting detection mechanism to {default}"
206
+ )
207
+ return default
208
+
209
+ parameters = _HasSmartFunctionality._get_assertion_parameters(monitor, default)
210
+ if parameters is None:
211
+ return _warn_and_return_default_detection_mechanism("parameters", default)
212
+
213
+ # 2. Convert the raw detection mechanism to the SDK detection mechanism
214
+ if parameters.type in [
215
+ models.AssertionEvaluationParametersTypeClass.DATASET_FRESHNESS,
216
+ models.AssertionEvaluationParametersTypeClass.DATASET_VOLUME,
217
+ ]:
218
+ if assertion.info is None:
219
+ return _warn_and_return_default_detection_mechanism("info", default)
220
+ if isinstance(assertion.info, models.VolumeAssertionInfoClass):
221
+ return _HasSmartFunctionality._get_volume_detection_mechanism(
222
+ assertion, parameters, default
223
+ )
224
+ elif isinstance(assertion.info, models.FreshnessAssertionInfoClass):
225
+ return _HasSmartFunctionality._get_freshness_detection_mechanism(
226
+ assertion, parameters, default
227
+ )
228
+ # TODO: Consider moving the detection mechanism logic to the assertion classes themselves e.g. _get_assertion_specific_detection_mechanism as an abstract method
229
+ # TODO: Add support here for other detection mechanisms when other assertion types are supported
230
+ else:
231
+ raise SDKNotYetSupportedError(
232
+ f"AssertionType {type(assertion.info).__name__}"
233
+ )
234
+ else:
235
+ raise SDKNotYetSupportedError(
236
+ f"AssertionEvaluationParametersType {parameters.type} not supported"
237
+ )
238
+
239
+ @staticmethod
240
+ def _has_valid_monitor_info(monitor: Monitor) -> bool:
241
+ """Check if monitor has valid info and assertion monitor."""
242
+
243
+ def _warn_and_return_false(field_name: str) -> bool:
244
+ logger.warning(
245
+ f"Monitor {monitor.urn} does not have an `{field_name}` field, defaulting detection mechanism to {DEFAULT_DETECTION_MECHANISM}"
246
+ )
247
+ return False
248
+
249
+ if monitor.info is None:
250
+ return _warn_and_return_false("info")
251
+ if monitor.info.assertionMonitor is None:
252
+ return _warn_and_return_false("assertionMonitor")
253
+ if (
254
+ monitor.info.assertionMonitor.assertions is None
255
+ or len(monitor.info.assertionMonitor.assertions) == 0
256
+ ):
257
+ return _warn_and_return_false("assertionMonitor.assertions")
258
+
259
+ return True
260
+
261
+ @staticmethod
262
+ def _get_assertion_parameters(
263
+ monitor: Monitor,
264
+ default: Optional[_DetectionMechanismTypes] = DEFAULT_DETECTION_MECHANISM,
265
+ ) -> Optional[models.AssertionEvaluationParametersClass]:
266
+ """Get the assertion parameters from the monitor."""
267
+ # We know these are not None from _has_valid_monitor_info check
268
+ assert (
269
+ monitor is not None
270
+ and monitor.info is not None
271
+ and monitor.info.assertionMonitor is not None
272
+ )
273
+ assertion_monitor = monitor.info.assertionMonitor
274
+ assert (
275
+ assertion_monitor is not None and assertion_monitor.assertions is not None
276
+ )
277
+ assertions = assertion_monitor.assertions
278
+
279
+ if assertions[0].parameters is None:
280
+ logger.warning(
281
+ f"Monitor {monitor.urn} does not have a assertionMonitor.assertions[0].parameters, defaulting detection mechanism to {default}"
282
+ )
283
+ return None
284
+ return assertions[0].parameters
285
+
286
+ @staticmethod
287
+ def _get_freshness_detection_mechanism(
288
+ assertion: Assertion,
289
+ parameters: models.AssertionEvaluationParametersClass,
290
+ default: Optional[_DetectionMechanismTypes] = DEFAULT_DETECTION_MECHANISM,
291
+ ) -> Optional[_DetectionMechanismTypes]:
292
+ """Get the detection mechanism for freshness assertions."""
293
+ if parameters.datasetFreshnessParameters is None:
294
+ logger.warning(
295
+ f"Monitor does not have datasetFreshnessParameters, defaulting detection mechanism to {DEFAULT_DETECTION_MECHANISM}"
296
+ )
297
+ return default
298
+
299
+ source_type = parameters.datasetFreshnessParameters.sourceType
300
+ if source_type == models.DatasetFreshnessSourceTypeClass.INFORMATION_SCHEMA:
301
+ return DetectionMechanism.INFORMATION_SCHEMA
302
+ elif source_type == models.DatasetFreshnessSourceTypeClass.AUDIT_LOG:
303
+ return DetectionMechanism.AUDIT_LOG
304
+ elif source_type == models.DatasetFreshnessSourceTypeClass.FIELD_VALUE:
305
+ return _HasSmartFunctionality._get_field_value_detection_mechanism(
306
+ assertion, parameters
307
+ )
308
+ elif source_type == models.DatasetFreshnessSourceTypeClass.DATAHUB_OPERATION:
309
+ return DetectionMechanism.DATAHUB_OPERATION
310
+ elif source_type == models.DatasetFreshnessSourceTypeClass.FILE_METADATA:
311
+ raise SDKNotYetSupportedError("FILE_METADATA DatasetFreshnessSourceType")
312
+ else:
313
+ raise SDKNotYetSupportedError(f"DatasetFreshnessSourceType {source_type}")
314
+
315
+ @staticmethod
316
+ def _get_volume_detection_mechanism(
317
+ assertion: Assertion,
318
+ parameters: models.AssertionEvaluationParametersClass,
319
+ default: Optional[_DetectionMechanismTypes] = DEFAULT_DETECTION_MECHANISM,
320
+ ) -> _DetectionMechanismTypes:
321
+ """Get the detection mechanism for volume assertions."""
322
+ if parameters.datasetVolumeParameters is None:
323
+ logger.warning(
324
+ f"Monitor does not have datasetVolumeParameters, defaulting detection mechanism to {DEFAULT_DETECTION_MECHANISM}"
325
+ )
326
+ if default is None:
327
+ return DEFAULT_DETECTION_MECHANISM
328
+ else:
329
+ return default
330
+
331
+ source_type = parameters.datasetVolumeParameters.sourceType
332
+ if source_type == models.DatasetVolumeSourceTypeClass.INFORMATION_SCHEMA:
333
+ return DetectionMechanism.INFORMATION_SCHEMA
334
+ elif source_type == models.DatasetVolumeSourceTypeClass.QUERY:
335
+ additional_filter = _HasSmartFunctionality._get_additional_filter(assertion)
336
+ return DetectionMechanism.QUERY(additional_filter=additional_filter)
337
+ elif source_type == models.DatasetVolumeSourceTypeClass.DATAHUB_DATASET_PROFILE:
338
+ return DetectionMechanism.DATASET_PROFILE
339
+ else:
340
+ raise SDKNotYetSupportedError(f"DatasetVolumeSourceType {source_type}")
341
+
342
+ @staticmethod
343
+ def _get_field_value_detection_mechanism(
344
+ assertion: Assertion,
345
+ parameters: models.AssertionEvaluationParametersClass,
346
+ ) -> _DetectionMechanismTypes:
347
+ """Get the detection mechanism for field value based freshness."""
348
+ # We know datasetFreshnessParameters is not None from _get_freshness_detection_mechanism check
349
+ assert parameters.datasetFreshnessParameters is not None
350
+ field = parameters.datasetFreshnessParameters.field
351
+
352
+ if field is None or field.kind is None:
353
+ logger.warning(
354
+ f"Monitor does not have valid field info, defaulting detection mechanism to {DEFAULT_DETECTION_MECHANISM}"
355
+ )
356
+ return DEFAULT_DETECTION_MECHANISM
357
+
358
+ column_name = field.path
359
+ additional_filter = _HasSmartFunctionality._get_additional_filter(assertion)
360
+
361
+ if field.kind == models.FreshnessFieldKindClass.LAST_MODIFIED:
362
+ return DetectionMechanism.LAST_MODIFIED_COLUMN(
363
+ column_name=column_name, additional_filter=additional_filter
364
+ )
365
+ elif field.kind == models.FreshnessFieldKindClass.HIGH_WATERMARK:
366
+ return DetectionMechanism.HIGH_WATERMARK_COLUMN(
367
+ column_name=column_name, additional_filter=additional_filter
368
+ )
369
+ else:
370
+ raise SDKNotYetSupportedError(f"FreshnessFieldKind {field.kind}")
371
+
372
+ @staticmethod
373
+ def _get_additional_filter(assertion: Assertion) -> Optional[str]:
374
+ """Get the additional filter SQL from the assertion."""
375
+ if assertion.info is None:
376
+ logger.warning(
377
+ f"Assertion {assertion.urn} does not have an info, defaulting additional filter to None"
378
+ )
379
+ return None
380
+ if (
381
+ not isinstance(
382
+ assertion.info,
383
+ _HasSmartFunctionality._SUPPORTED_WITH_FILTER_ASSERTION_TYPES,
384
+ )
385
+ or assertion.info.filter is None
386
+ ):
387
+ logger.warning(
388
+ f"Assertion {assertion.urn} does not have a filter, defaulting additional filter to None"
389
+ )
390
+ return None
391
+ if assertion.info.filter.type != models.DatasetFilterTypeClass.SQL:
392
+ raise SDKNotYetSupportedError(
393
+ f"DatasetFilterType {assertion.info.filter.type}"
394
+ )
395
+ return assertion.info.filter.sql
396
+
397
+
398
+ class _AssertionPublic(ABC):
399
+ """
400
+ Abstract base class that represents a public facing assertion and contains the common properties of all assertions.
401
+ """
402
+
403
+ def __init__(
404
+ self,
405
+ *,
406
+ urn: AssertionUrn,
407
+ dataset_urn: DatasetUrn,
408
+ display_name: str,
409
+ mode: AssertionMode,
410
+ tags: list[TagUrn],
411
+ created_by: Optional[CorpUserUrn] = None,
412
+ created_at: Union[datetime, None] = None,
413
+ updated_by: Optional[CorpUserUrn] = None,
414
+ updated_at: Optional[datetime] = None,
415
+ ):
416
+ """
417
+ Initialize the public facing assertion class.
418
+
419
+ Args:
420
+ urn: The urn of the assertion.
421
+ dataset_urn: The urn of the dataset that the assertion is for.
422
+ display_name: The display name of the assertion.
423
+ mode: The mode of the assertion (active, inactive).
424
+ tags: The tags of the assertion.
425
+ created_by: The urn of the user that created the assertion.
426
+ created_at: The timestamp of when the assertion was created.
427
+ updated_by: The urn of the user that updated the assertion.
428
+ updated_at: The timestamp of when the assertion was updated.
429
+ """
430
+ self._urn = urn
431
+ self._dataset_urn = dataset_urn
432
+ self._display_name = display_name
433
+ self._mode = mode
434
+ self._created_by = created_by
435
+ self._created_at = created_at
436
+ self._updated_by = updated_by
437
+ self._updated_at = updated_at
438
+ self._tags = tags
439
+
440
+ @property
441
+ def urn(self) -> AssertionUrn:
442
+ return self._urn
443
+
444
+ @property
445
+ def dataset_urn(self) -> DatasetUrn:
446
+ return self._dataset_urn
447
+
448
+ @property
449
+ def display_name(self) -> str:
450
+ return self._display_name
451
+
452
+ @property
453
+ def mode(self) -> AssertionMode:
454
+ return self._mode
455
+
456
+ @property
457
+ def created_by(self) -> Optional[CorpUserUrn]:
458
+ return self._created_by
459
+
460
+ @property
461
+ def created_at(self) -> Union[datetime, None]:
462
+ return self._created_at
463
+
464
+ @property
465
+ def updated_by(self) -> Optional[CorpUserUrn]:
466
+ return self._updated_by
467
+
468
+ @property
469
+ def updated_at(self) -> Union[datetime, None]:
470
+ return self._updated_at
471
+
472
+ @property
473
+ def tags(self) -> list[TagUrn]:
474
+ return self._tags
475
+
476
+ @staticmethod
477
+ def _get_incident_behavior(assertion: Assertion) -> list[AssertionIncidentBehavior]:
478
+ incident_behaviors = []
479
+ for action in assertion.on_failure + assertion.on_success:
480
+ if action.type == models.AssertionActionTypeClass.RAISE_INCIDENT:
481
+ incident_behaviors.append(AssertionIncidentBehavior.RAISE_ON_FAIL)
482
+ elif action.type == models.AssertionActionTypeClass.RESOLVE_INCIDENT:
483
+ incident_behaviors.append(AssertionIncidentBehavior.RESOLVE_ON_PASS)
484
+
485
+ return incident_behaviors
486
+
487
+ @staticmethod
488
+ def _get_created_by(assertion: Assertion) -> Optional[CorpUserUrn]:
489
+ if assertion.source is None:
490
+ logger.warning(f"Assertion {assertion.urn} does not have a source")
491
+ return None
492
+ if isinstance(assertion.source, models.AssertionSourceClass):
493
+ if assertion.source.created is None:
494
+ logger.warning(
495
+ f"Assertion {assertion.urn} does not have a created by in the source"
496
+ )
497
+ return None
498
+ return CorpUserUrn.from_string(assertion.source.created.actor)
499
+ elif isinstance(assertion.source, models.AssertionSourceTypeClass):
500
+ logger.warning(
501
+ f"Assertion {assertion.urn} has a source type with no created by"
502
+ )
503
+ return None
504
+ return None
505
+
506
+ @staticmethod
507
+ def _get_created_at(assertion: Assertion) -> Union[datetime, None]:
508
+ if assertion.source is None:
509
+ logger.warning(f"Assertion {assertion.urn} does not have a source")
510
+ return None
511
+ if isinstance(assertion.source, models.AssertionSourceClass):
512
+ if assertion.source.created is None:
513
+ logger.warning(
514
+ f"Assertion {assertion.urn} does not have a created by in the source"
515
+ )
516
+ return None
517
+ return parse_ts_millis(assertion.source.created.time)
518
+ elif isinstance(assertion.source, models.AssertionSourceTypeClass):
519
+ logger.warning(
520
+ f"Assertion {assertion.urn} has a source type with no created by"
521
+ )
522
+ return None
523
+ return None
524
+
525
+ @staticmethod
526
+ def _get_updated_by(assertion: Assertion) -> Optional[CorpUserUrn]:
527
+ if assertion.last_updated is None:
528
+ logger.warning(f"Assertion {assertion.urn} does not have a last updated")
529
+ return None
530
+ return CorpUserUrn.from_string(assertion.last_updated.actor)
531
+
532
+ @staticmethod
533
+ def _get_updated_at(assertion: Assertion) -> Union[datetime, None]:
534
+ if assertion.last_updated is None:
535
+ logger.warning(f"Assertion {assertion.urn} does not have a last updated")
536
+ return None
537
+ return parse_ts_millis(assertion.last_updated.time)
538
+
539
+ @staticmethod
540
+ def _get_tags(assertion: Assertion) -> list[TagUrn]:
541
+ return [TagUrn.from_string(t.tag) for t in assertion.tags or []]
542
+
543
+ @staticmethod
544
+ def _get_mode(monitor: Monitor) -> AssertionMode:
545
+ if monitor.info is None:
546
+ logger.warning(
547
+ f"Monitor {monitor.urn} does not have a info, defaulting status to INACTIVE"
548
+ )
549
+ return AssertionMode.INACTIVE
550
+ return AssertionMode(monitor.info.status.mode)
551
+
552
+ @classmethod
553
+ @abstractmethod
554
+ def _from_entities(
555
+ cls,
556
+ assertion: Assertion,
557
+ monitor: Monitor,
558
+ ) -> Self:
559
+ """
560
+ Create an assertion from the assertion and monitor entities.
561
+
562
+ Note: This is a private method since it is intended to be called internally by the client.
563
+ """
564
+ pass
565
+
566
+
567
+ class SmartFreshnessAssertion(_HasSchedule, _HasSmartFunctionality, _AssertionPublic):
568
+ """
569
+ A class that represents a smart freshness assertion.
570
+ """
571
+
572
+ def __init__(
573
+ self,
574
+ *,
575
+ urn: AssertionUrn,
576
+ dataset_urn: DatasetUrn,
577
+ display_name: str,
578
+ mode: AssertionMode,
579
+ schedule: models.CronScheduleClass = DEFAULT_SCHEDULE,
580
+ sensitivity: InferenceSensitivity = DEFAULT_SENSITIVITY,
581
+ exclusion_windows: list[ExclusionWindowTypes],
582
+ training_data_lookback_days: int = ASSERTION_MONITOR_DEFAULT_TRAINING_LOOKBACK_WINDOW_DAYS,
583
+ incident_behavior: list[AssertionIncidentBehavior],
584
+ detection_mechanism: Optional[
585
+ _DetectionMechanismTypes
586
+ ] = DEFAULT_DETECTION_MECHANISM,
587
+ tags: list[TagUrn],
588
+ created_by: Optional[CorpUserUrn] = None,
589
+ created_at: Union[datetime, None] = None,
590
+ updated_by: Optional[CorpUserUrn] = None,
591
+ updated_at: Optional[datetime] = None,
592
+ ):
593
+ """
594
+ Initialize a smart freshness assertion.
595
+
596
+ Note: Values can be accessed, but not set on the assertion object.
597
+ To update an assertion, use the `upsert_*` method.
598
+ Args:
599
+ urn: The urn of the assertion.
600
+ dataset_urn: The urn of the dataset that the assertion is for.
601
+ display_name: The display name of the assertion.
602
+ mode: The mode of the assertion (active, inactive).
603
+ schedule: The schedule of the assertion.
604
+ sensitivity: The sensitivity of the assertion (low, medium, high).
605
+ exclusion_windows: The exclusion windows of the assertion.
606
+ training_data_lookback_days: The max number of days of data to use for training the assertion.
607
+ incident_behavior: Whether to raise or resolve an incident when the assertion fails / passes.
608
+ detection_mechanism: The detection mechanism of the assertion.
609
+ tags: The tags applied to the assertion.
610
+ created_by: The urn of the user that created the assertion.
611
+ created_at: The timestamp of when the assertion was created.
612
+ updated_by: The urn of the user that updated the assertion.
613
+ updated_at: The timestamp of when the assertion was updated.
614
+ """
615
+ # Initialize the mixins first
616
+ _HasSchedule.__init__(self, schedule=schedule)
617
+ _HasSmartFunctionality.__init__(
618
+ self,
619
+ sensitivity=sensitivity,
620
+ exclusion_windows=exclusion_windows,
621
+ training_data_lookback_days=training_data_lookback_days,
622
+ incident_behavior=incident_behavior,
623
+ detection_mechanism=detection_mechanism,
624
+ )
625
+ # Then initialize the parent class
626
+ _AssertionPublic.__init__(
627
+ self,
628
+ urn=urn,
629
+ dataset_urn=dataset_urn,
630
+ display_name=display_name,
631
+ mode=mode,
632
+ created_by=created_by,
633
+ created_at=created_at,
634
+ updated_by=updated_by,
635
+ updated_at=updated_at,
636
+ tags=tags,
637
+ )
638
+
639
+ @classmethod
640
+ def _from_entities(cls, assertion: Assertion, monitor: Monitor) -> Self:
641
+ """
642
+ Create a smart freshness assertion from the assertion and monitor entities.
643
+
644
+ Note: This is a private method since it is intended to be called internally by the client.
645
+ """
646
+ return cls(
647
+ urn=assertion.urn,
648
+ dataset_urn=assertion.dataset,
649
+ display_name=assertion.description or "",
650
+ mode=cls._get_mode(monitor),
651
+ schedule=cls._get_schedule(monitor),
652
+ sensitivity=cls._get_sensitivity(monitor),
653
+ exclusion_windows=cls._get_exclusion_windows(monitor),
654
+ training_data_lookback_days=cls._get_training_data_lookback_days(monitor),
655
+ incident_behavior=cls._get_incident_behavior(assertion),
656
+ detection_mechanism=cls._get_detection_mechanism(assertion, monitor),
657
+ created_by=cls._get_created_by(assertion),
658
+ created_at=cls._get_created_at(assertion),
659
+ updated_by=cls._get_updated_by(assertion),
660
+ updated_at=cls._get_updated_at(assertion),
661
+ tags=cls._get_tags(assertion),
662
+ )
663
+
664
+
665
+ class SmartVolumeAssertion(_HasSchedule, _HasSmartFunctionality, _AssertionPublic):
666
+ """
667
+ A class that represents a smart volume assertion.
668
+ """
669
+
670
+ def __init__(
671
+ self,
672
+ *,
673
+ urn: AssertionUrn,
674
+ dataset_urn: DatasetUrn,
675
+ display_name: str,
676
+ mode: AssertionMode,
677
+ schedule: models.CronScheduleClass,
678
+ sensitivity: InferenceSensitivity = DEFAULT_SENSITIVITY,
679
+ exclusion_windows: list[ExclusionWindowTypes],
680
+ training_data_lookback_days: int = ASSERTION_MONITOR_DEFAULT_TRAINING_LOOKBACK_WINDOW_DAYS,
681
+ incident_behavior: list[AssertionIncidentBehavior],
682
+ detection_mechanism: Optional[
683
+ _DetectionMechanismTypes
684
+ ] = DEFAULT_DETECTION_MECHANISM,
685
+ tags: list[TagUrn],
686
+ created_by: Optional[CorpUserUrn] = None,
687
+ created_at: Union[datetime, None] = None,
688
+ updated_by: Optional[CorpUserUrn] = None,
689
+ updated_at: Optional[datetime] = None,
690
+ ):
691
+ """
692
+ Initialize a smart volume assertion.
693
+
694
+ Note: Values can be accessed, but not set on the assertion object.
695
+ To update an assertion, use the `upsert_*` method.
696
+ Args:
697
+ urn: The urn of the assertion.
698
+ dataset_urn: The urn of the dataset that the assertion is for.
699
+ display_name: The display name of the assertion.
700
+ mode: The mode of the assertion (active, inactive).
701
+ schedule: The schedule of the assertion.
702
+ sensitivity: The sensitivity of the assertion (low, medium, high).
703
+ exclusion_windows: The exclusion windows of the assertion.
704
+ training_data_lookback_days: The max number of days of data to use for training the assertion.
705
+ incident_behavior: Whether to raise or resolve an incident when the assertion fails / passes.
706
+ detection_mechanism: The detection mechanism of the assertion.
707
+ tags: The tags applied to the assertion.
708
+ created_by: The urn of the user that created the assertion.
709
+ created_at: The timestamp of when the assertion was created.
710
+ updated_by: The urn of the user that updated the assertion.
711
+ updated_at: The timestamp of when the assertion was updated.
712
+ """
713
+ # Initialize the mixins first
714
+ _HasSchedule.__init__(self, schedule=schedule)
715
+ _HasSmartFunctionality.__init__(
716
+ self,
717
+ sensitivity=sensitivity,
718
+ exclusion_windows=exclusion_windows,
719
+ training_data_lookback_days=training_data_lookback_days,
720
+ incident_behavior=incident_behavior,
721
+ detection_mechanism=detection_mechanism,
722
+ )
723
+ # Then initialize the parent class
724
+ _AssertionPublic.__init__(
725
+ self,
726
+ urn=urn,
727
+ dataset_urn=dataset_urn,
728
+ display_name=display_name,
729
+ mode=mode,
730
+ created_by=created_by,
731
+ created_at=created_at,
732
+ updated_by=updated_by,
733
+ updated_at=updated_at,
734
+ tags=tags,
735
+ )
736
+
737
+ @classmethod
738
+ def _from_entities(cls, assertion: Assertion, monitor: Monitor) -> Self:
739
+ """
740
+ Create a smart freshness assertion from the assertion and monitor entities.
741
+
742
+ Note: This is a private method since it is intended to be called internally by the client.
743
+ """
744
+ return cls(
745
+ urn=assertion.urn,
746
+ dataset_urn=assertion.dataset,
747
+ display_name=assertion.description or "",
748
+ mode=cls._get_mode(monitor),
749
+ schedule=cls._get_schedule(monitor),
750
+ sensitivity=cls._get_sensitivity(monitor),
751
+ exclusion_windows=cls._get_exclusion_windows(monitor),
752
+ training_data_lookback_days=cls._get_training_data_lookback_days(monitor),
753
+ incident_behavior=cls._get_incident_behavior(assertion),
754
+ detection_mechanism=cls._get_detection_mechanism(assertion, monitor),
755
+ created_by=cls._get_created_by(assertion),
756
+ created_at=cls._get_created_at(assertion),
757
+ updated_by=cls._get_updated_by(assertion),
758
+ updated_at=cls._get_updated_at(assertion),
759
+ tags=cls._get_tags(assertion),
760
+ )
761
+
762
+
763
+ AssertionTypes = Union[
764
+ SmartFreshnessAssertion,
765
+ SmartVolumeAssertion,
766
+ # TODO: Add other assertion types here as we add them.
767
+ ]