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,1335 @@
1
+ """
2
+ This file contains the AssertionInput class and related classes, which are used to
3
+ validate and represent the input for creating an Assertion in DataHub.
4
+ """
5
+
6
+ import random
7
+ import string
8
+ from abc import ABC, abstractmethod
9
+ from datetime import datetime
10
+ from enum import Enum
11
+ from typing import Literal, Optional, TypeAlias, Union
12
+
13
+ import pydantic
14
+ import pytz
15
+ import tzlocal
16
+ from avrogen.dict_wrapper import DictWrapper
17
+ from croniter import croniter
18
+ from pydantic import BaseModel, Extra, ValidationError
19
+
20
+ from acryl_datahub_cloud.sdk.entities.assertion import (
21
+ Assertion,
22
+ AssertionActionsInputType,
23
+ AssertionInfoInputType,
24
+ TagsInputType,
25
+ )
26
+ from acryl_datahub_cloud.sdk.entities.monitor import Monitor
27
+ from acryl_datahub_cloud.sdk.errors import (
28
+ SDKNotYetSupportedError,
29
+ SDKUsageError,
30
+ SDKUsageErrorWithExamples,
31
+ )
32
+ from datahub.emitter.enum_helpers import get_enum_options
33
+ from datahub.emitter.mce_builder import make_ts_millis, parse_ts_millis
34
+ from datahub.metadata import schema_classes as models
35
+ from datahub.metadata.urns import AssertionUrn, CorpUserUrn, DatasetUrn
36
+ from datahub.sdk import Dataset
37
+ from datahub.sdk.entity_client import EntityClient
38
+
39
+ # TODO: Import ASSERTION_MONITOR_DEFAULT_TRAINING_LOOKBACK_WINDOW_DAYS from datahub_executor.config
40
+ ASSERTION_MONITOR_DEFAULT_TRAINING_LOOKBACK_WINDOW_DAYS = 60
41
+
42
+ DEFAULT_NAME_PREFIX = "New Assertion"
43
+ DEFAULT_NAME_SUFFIX_LENGTH = 8
44
+
45
+ DEFAULT_SCHEDULE = models.CronScheduleClass(
46
+ cron="0 * * * *", # Every hour, matches the UI default
47
+ timezone=str(
48
+ tzlocal.get_localzone()
49
+ ), # User local timezone, matches the UI default
50
+ )
51
+
52
+
53
+ class AbstractDetectionMechanism(BaseModel, ABC):
54
+ type: str
55
+
56
+ class Config:
57
+ extra = Extra.forbid
58
+
59
+
60
+ class _InformationSchema(AbstractDetectionMechanism):
61
+ type: Literal["information_schema"] = "information_schema"
62
+
63
+
64
+ class _AuditLog(AbstractDetectionMechanism):
65
+ type: Literal["audit_log"] = "audit_log"
66
+
67
+
68
+ # Keep this in sync with the allowed field types in the UI, currently in
69
+ # datahub-web-react/src/app/entity/shared/tabs/Dataset/Validations/assertion/builder/constants.ts: LAST_MODIFIED_FIELD_TYPES
70
+ LAST_MODIFIED_ALLOWED_FIELD_TYPES = [models.DateTypeClass(), models.TimeTypeClass()]
71
+
72
+
73
+ class _LastModifiedColumn(AbstractDetectionMechanism):
74
+ type: Literal["last_modified_column"] = "last_modified_column"
75
+ column_name: str
76
+ additional_filter: Optional[str] = None
77
+
78
+
79
+ # Keep this in sync with the allowed field types in the UI, currently in
80
+ # datahub-web-react/src/app/entity/shared/tabs/Dataset/Validations/assertion/builder/constants.ts: HIGH_WATERMARK_FIELD_TYPES
81
+ HIGH_WATERMARK_ALLOWED_FIELD_TYPES = [
82
+ models.NumberTypeClass(),
83
+ models.DateTypeClass(),
84
+ models.TimeTypeClass(),
85
+ ]
86
+
87
+
88
+ class _HighWatermarkColumn(AbstractDetectionMechanism):
89
+ type: Literal["high_watermark_column"] = "high_watermark_column"
90
+ column_name: str
91
+ additional_filter: Optional[str] = None
92
+
93
+
94
+ class _DataHubOperation(AbstractDetectionMechanism):
95
+ type: Literal["datahub_operation"] = "datahub_operation"
96
+
97
+
98
+ class _Query(AbstractDetectionMechanism):
99
+ # COUNT(*) query
100
+ type: Literal["query"] = "query"
101
+ additional_filter: Optional[str] = None
102
+
103
+
104
+ class _DatasetProfile(AbstractDetectionMechanism):
105
+ type: Literal["dataset_profile"] = "dataset_profile"
106
+
107
+
108
+ # Keep these two lists in sync:
109
+ _DETECTION_MECHANISM_CONCRETE_TYPES = (
110
+ _InformationSchema,
111
+ _AuditLog,
112
+ _LastModifiedColumn,
113
+ _HighWatermarkColumn,
114
+ _DataHubOperation,
115
+ _Query,
116
+ _DatasetProfile,
117
+ )
118
+ _DetectionMechanismTypes = Union[
119
+ _InformationSchema,
120
+ _AuditLog,
121
+ _LastModifiedColumn,
122
+ _HighWatermarkColumn,
123
+ _DataHubOperation,
124
+ _Query,
125
+ _DatasetProfile,
126
+ ]
127
+
128
+ _DETECTION_MECHANISM_TYPES_WITH_ADDITIONAL_FILTER = (
129
+ _LastModifiedColumn,
130
+ _HighWatermarkColumn,
131
+ _Query,
132
+ )
133
+
134
+
135
+ class DetectionMechanism:
136
+ # To have a more enum-like user experience even with sub parameters, we define the detection mechanisms as class attributes.
137
+ # The options with sub parameters are the classes themselves so that parameters can be applied, and the rest are already instantiated instances of the classes.
138
+ INFORMATION_SCHEMA = _InformationSchema()
139
+ AUDIT_LOG = _AuditLog()
140
+ LAST_MODIFIED_COLUMN = _LastModifiedColumn
141
+ HIGH_WATERMARK_COLUMN = _HighWatermarkColumn
142
+ DATAHUB_OPERATION = _DataHubOperation()
143
+ QUERY = _Query
144
+ DATASET_PROFILE = _DatasetProfile()
145
+
146
+ _DETECTION_MECHANISM_EXAMPLES = {
147
+ "Information Schema from string": "information_schema",
148
+ "Information Schema from DetectionMechanism": "DetectionMechanism.INFORMATION_SCHEMA",
149
+ "Audit Log from string": "audit_log",
150
+ "Audit Log from DetectionMechanism": "DetectionMechanism.AUDIT_LOG",
151
+ "Last Modified Column from dict": {
152
+ "type": "last_modified_column",
153
+ "column_name": "last_modified",
154
+ "additional_filter": "last_modified > '2021-01-01'",
155
+ },
156
+ "Last Modified Column from DetectionMechanism": "DetectionMechanism.LAST_MODIFIED_COLUMN(column_name='last_modified', additional_filter='last_modified > 2021-01-01')",
157
+ "High Watermark Column from dict": {
158
+ "type": "high_watermark_column",
159
+ "column_name": "id",
160
+ "additional_filter": "id > 1000",
161
+ },
162
+ "High Watermark Column from DetectionMechanism": "DetectionMechanism.HIGH_WATERMARK_COLUMN(column_name='id', additional_filter='id > 1000')",
163
+ "DataHub Operation from string": "datahub_operation",
164
+ "DataHub Operation from DetectionMechanism": "DetectionMechanism.DATAHUB_OPERATION",
165
+ "Query from string": "query",
166
+ "Query from dict": {
167
+ "type": "query",
168
+ "additional_filter": "id > 1000",
169
+ },
170
+ "Query from DetectionMechanism (with optional additional filter)": "DetectionMechanism.QUERY(additional_filter='id > 1000')",
171
+ "Dataset Profile from string": "dataset_profile",
172
+ "Dataset Profile from DetectionMechanism": "DetectionMechanism.DATASET_PROFILE",
173
+ }
174
+
175
+ @staticmethod
176
+ def parse(
177
+ detection_mechanism_config: Optional[
178
+ Union[str, dict[str, str], _DetectionMechanismTypes]
179
+ ] = None,
180
+ ) -> _DetectionMechanismTypes:
181
+ if detection_mechanism_config is None:
182
+ return DEFAULT_DETECTION_MECHANISM
183
+ if isinstance(detection_mechanism_config, _DETECTION_MECHANISM_CONCRETE_TYPES):
184
+ return detection_mechanism_config
185
+ elif isinstance(detection_mechanism_config, str):
186
+ return DetectionMechanism._try_parse_from_string(detection_mechanism_config)
187
+ elif isinstance(detection_mechanism_config, dict):
188
+ return DetectionMechanism._try_parse_from_dict(detection_mechanism_config)
189
+ else:
190
+ raise SDKUsageErrorWithExamples(
191
+ msg=f"Invalid detection mechanism: {detection_mechanism_config}",
192
+ examples=DetectionMechanism._DETECTION_MECHANISM_EXAMPLES,
193
+ )
194
+
195
+ @staticmethod
196
+ def _try_parse_from_string(
197
+ detection_mechanism_config: str,
198
+ ) -> _DetectionMechanismTypes:
199
+ try:
200
+ return_value = getattr(
201
+ DetectionMechanism, detection_mechanism_config.upper()
202
+ )
203
+ if isinstance(return_value, pydantic.main.ModelMetaclass):
204
+ try:
205
+ # We try to instantiate here to let pydantic raise a helpful error
206
+ # about which parameters are missing
207
+ return_value = return_value()
208
+ except ValidationError as e:
209
+ raise SDKUsageErrorWithExamples(
210
+ msg=f"Detection mechanism type '{detection_mechanism_config}' requires additional parameters: {e}",
211
+ examples=DetectionMechanism._DETECTION_MECHANISM_EXAMPLES,
212
+ ) from e
213
+ return return_value
214
+ except AttributeError as e:
215
+ raise SDKUsageErrorWithExamples(
216
+ msg=f"Invalid detection mechanism type: {detection_mechanism_config}",
217
+ examples=DetectionMechanism._DETECTION_MECHANISM_EXAMPLES,
218
+ ) from e
219
+
220
+ @staticmethod
221
+ def _try_parse_from_dict(
222
+ detection_mechanism_config: dict[str, str],
223
+ ) -> _DetectionMechanismTypes:
224
+ try:
225
+ detection_mechanism_type = detection_mechanism_config.pop("type")
226
+ except KeyError as e:
227
+ raise SDKUsageErrorWithExamples(
228
+ msg="Detection mechanism type is required if using a dict to create a DetectionMechanism",
229
+ examples=DetectionMechanism._DETECTION_MECHANISM_EXAMPLES,
230
+ ) from e
231
+ try:
232
+ detection_mechanism_obj = getattr(
233
+ DetectionMechanism, detection_mechanism_type.upper()
234
+ )
235
+ except AttributeError as e:
236
+ raise SDKUsageErrorWithExamples(
237
+ msg=f"Invalid detection mechanism type: {detection_mechanism_type}",
238
+ examples=DetectionMechanism._DETECTION_MECHANISM_EXAMPLES,
239
+ ) from e
240
+
241
+ try:
242
+ return detection_mechanism_obj(**detection_mechanism_config)
243
+ except TypeError as e:
244
+ if "object is not callable" not in e.args[0]:
245
+ raise e
246
+ if detection_mechanism_config:
247
+ # If we are here in the TypeError case, the detection mechanism is an instance of a class,
248
+ # not a class itself, so we can't instantiate it with the config dict.
249
+ # In this case, the config dict should be empty after the type is popped.
250
+ # If it is not empty, we raise an error.
251
+ raise SDKUsageErrorWithExamples(
252
+ msg=f"Invalid additional fields specified for detection mechanism '{detection_mechanism_type}': {detection_mechanism_config}",
253
+ examples=DetectionMechanism._DETECTION_MECHANISM_EXAMPLES,
254
+ ) from e
255
+ return detection_mechanism_obj
256
+ except ValidationError as e:
257
+ raise SDKUsageErrorWithExamples(
258
+ msg=f"Invalid detection mechanism type '{detection_mechanism_type}': {detection_mechanism_config} {e}",
259
+ examples=DetectionMechanism._DETECTION_MECHANISM_EXAMPLES,
260
+ ) from e
261
+
262
+
263
+ DEFAULT_DETECTION_MECHANISM = DetectionMechanism.INFORMATION_SCHEMA
264
+
265
+ DetectionMechanismInputTypes: TypeAlias = Union[
266
+ str, dict[str, str], _DetectionMechanismTypes, None
267
+ ]
268
+
269
+
270
+ class InferenceSensitivity(Enum):
271
+ HIGH = "high"
272
+ MEDIUM = "medium"
273
+ LOW = "low"
274
+
275
+ @staticmethod
276
+ def parse(
277
+ sensitivity: Optional[
278
+ Union[
279
+ str,
280
+ int,
281
+ "InferenceSensitivity",
282
+ models.AssertionMonitorSensitivityClass,
283
+ ]
284
+ ],
285
+ ) -> "InferenceSensitivity":
286
+ if sensitivity is None:
287
+ return DEFAULT_SENSITIVITY
288
+ EXAMPLES = {
289
+ "High sensitivity from string": "high",
290
+ "High sensitivity from enum": "InferenceSensitivity.HIGH",
291
+ "Medium sensitivity from string": "medium",
292
+ "Medium sensitivity from enum": "InferenceSensitivity.MEDIUM",
293
+ "Low sensitivity from string": "low",
294
+ "Low sensitivity from enum": "InferenceSensitivity.LOW",
295
+ "Sensitivity from int (1-3: low, 4-6: medium, 7-10: high)": "10",
296
+ }
297
+
298
+ if isinstance(sensitivity, InferenceSensitivity):
299
+ return sensitivity
300
+ if isinstance(sensitivity, models.AssertionMonitorSensitivityClass):
301
+ sensitivity = sensitivity.level
302
+ if isinstance(sensitivity, int):
303
+ if (sensitivity < 1) or (sensitivity > 10):
304
+ raise SDKUsageErrorWithExamples(
305
+ msg=f"Invalid inference sensitivity: {sensitivity}",
306
+ examples=EXAMPLES,
307
+ )
308
+ elif sensitivity < 4:
309
+ return InferenceSensitivity.LOW
310
+ elif sensitivity < 7:
311
+ return InferenceSensitivity.MEDIUM
312
+ else:
313
+ return InferenceSensitivity.HIGH
314
+ try:
315
+ return InferenceSensitivity(sensitivity)
316
+ except ValueError as e:
317
+ raise SDKUsageErrorWithExamples(
318
+ msg=f"Invalid inference sensitivity: {sensitivity}",
319
+ examples=EXAMPLES,
320
+ ) from e
321
+
322
+ @staticmethod
323
+ def to_int(sensitivity: "InferenceSensitivity") -> int:
324
+ return {
325
+ InferenceSensitivity.HIGH: 10,
326
+ InferenceSensitivity.MEDIUM: 5,
327
+ InferenceSensitivity.LOW: 1,
328
+ }[sensitivity]
329
+
330
+
331
+ DEFAULT_SENSITIVITY = InferenceSensitivity.MEDIUM
332
+
333
+
334
+ class FixedRangeExclusionWindow(BaseModel):
335
+ type: Literal["fixed_range_exclusion_window"] = "fixed_range_exclusion_window"
336
+ start: datetime
337
+ end: datetime
338
+
339
+
340
+ ExclusionWindowTypes: TypeAlias = Union[
341
+ FixedRangeExclusionWindow,
342
+ # Add other exclusion window types here as they are added to the SDK.
343
+ ]
344
+
345
+ FIXED_RANGE_EXCLUSION_WINDOW_EXAMPLES = {
346
+ "Exclusion Window from datetimes": {
347
+ "start": "datetime(2025, 1, 1, 0, 0, 0)",
348
+ "end": "datetime(2025, 1, 2, 0, 0, 0)",
349
+ },
350
+ "Exclusion Window from strings": {
351
+ "start": "2025-01-01T00:00:00",
352
+ "end": "2025-01-02T00:00:00",
353
+ },
354
+ "Exclusion Window from object": "ExclusionWindow(start=datetime(2025, 1, 1, 0, 0, 0), end=datetime(2025, 1, 2, 0, 0, 0))",
355
+ }
356
+ FixedRangeExclusionWindowInputTypes: TypeAlias = Union[
357
+ dict[str, datetime],
358
+ dict[str, str],
359
+ list[dict[str, datetime]],
360
+ list[dict[str, str]],
361
+ FixedRangeExclusionWindow,
362
+ list[FixedRangeExclusionWindow],
363
+ ]
364
+
365
+ ExclusionWindowInputTypes: TypeAlias = Union[
366
+ models.AssertionExclusionWindowClass,
367
+ list[models.AssertionExclusionWindowClass],
368
+ FixedRangeExclusionWindowInputTypes,
369
+ # Add other exclusion window types here as they are added to the SDK.
370
+ ]
371
+
372
+ IterableExclusionWindowInputTypes: TypeAlias = Union[
373
+ list[dict[str, datetime]],
374
+ list[dict[str, str]],
375
+ list[FixedRangeExclusionWindow],
376
+ list[models.AssertionExclusionWindowClass],
377
+ ]
378
+
379
+
380
+ def _try_parse_exclusion_window(
381
+ config: Optional[ExclusionWindowInputTypes],
382
+ ) -> Union[FixedRangeExclusionWindow, list[FixedRangeExclusionWindow], None]:
383
+ if config is None:
384
+ return []
385
+ if isinstance(config, dict):
386
+ return [FixedRangeExclusionWindow(**config)]
387
+ if isinstance(config, FixedRangeExclusionWindow):
388
+ return [config]
389
+ elif isinstance(config, models.AssertionExclusionWindowClass):
390
+ assert config.fixedRange is not None
391
+ return [
392
+ FixedRangeExclusionWindow(
393
+ start=parse_ts_millis(config.fixedRange.startTimeMillis),
394
+ end=parse_ts_millis(config.fixedRange.endTimeMillis),
395
+ )
396
+ ]
397
+ elif isinstance(config, list):
398
+ return _try_parse_list_of_exclusion_windows(config)
399
+ else:
400
+ raise SDKUsageErrorWithExamples(
401
+ msg=f"Invalid exclusion window: {config}",
402
+ examples=FIXED_RANGE_EXCLUSION_WINDOW_EXAMPLES,
403
+ )
404
+
405
+
406
+ def _try_parse_list_of_exclusion_windows(
407
+ config: IterableExclusionWindowInputTypes,
408
+ ) -> Union[list[FixedRangeExclusionWindow], None]:
409
+ if all(isinstance(item, models.AssertionExclusionWindowClass) for item in config):
410
+ exclusion_windows = []
411
+ for item in config:
412
+ assert isinstance(item, models.AssertionExclusionWindowClass)
413
+ assert item.fixedRange is not None
414
+ exclusion_windows.append(
415
+ FixedRangeExclusionWindow(
416
+ start=parse_ts_millis(item.fixedRange.startTimeMillis),
417
+ end=parse_ts_millis(item.fixedRange.endTimeMillis),
418
+ )
419
+ )
420
+ return exclusion_windows
421
+ else:
422
+ exclusion_windows = []
423
+ for item in config:
424
+ if isinstance(item, dict):
425
+ try:
426
+ exclusion_windows.append(FixedRangeExclusionWindow(**item))
427
+ except ValidationError as e:
428
+ raise SDKUsageErrorWithExamples(
429
+ msg=f"Invalid exclusion window: {item}",
430
+ examples=FIXED_RANGE_EXCLUSION_WINDOW_EXAMPLES,
431
+ ) from e
432
+ elif isinstance(item, FixedRangeExclusionWindow):
433
+ exclusion_windows.append(item)
434
+ elif item is None:
435
+ pass
436
+ else:
437
+ raise SDKUsageErrorWithExamples(
438
+ msg=f"Invalid exclusion window: {item}",
439
+ examples=FIXED_RANGE_EXCLUSION_WINDOW_EXAMPLES,
440
+ )
441
+ return exclusion_windows
442
+
443
+
444
+ class AssertionIncidentBehavior(Enum):
445
+ RAISE_ON_FAIL = "raise_on_fail"
446
+ RESOLVE_ON_PASS = "resolve_on_pass"
447
+
448
+
449
+ ASSERTION_INCIDENT_BEHAVIOR_EXAMPLES = {
450
+ "Raise on fail from string": "raise_on_fail",
451
+ "Raise on fail from enum": "AssertionIncidentBehavior.RAISE_ON_FAIL",
452
+ "Resolve on pass from string": "resolve_on_pass",
453
+ "Resolve on pass from enum": "AssertionIncidentBehavior.RESOLVE_ON_PASS",
454
+ }
455
+
456
+ AssertionIncidentBehaviorInputTypes: TypeAlias = Union[
457
+ str,
458
+ list[str],
459
+ AssertionIncidentBehavior,
460
+ list[AssertionIncidentBehavior],
461
+ None,
462
+ ]
463
+
464
+
465
+ def _try_parse_incident_behavior(
466
+ config: AssertionIncidentBehaviorInputTypes,
467
+ ) -> Union[AssertionIncidentBehavior, list[AssertionIncidentBehavior], None]:
468
+ if config is None:
469
+ return []
470
+ if isinstance(config, str):
471
+ try:
472
+ return [AssertionIncidentBehavior(config)]
473
+ except ValueError as e:
474
+ raise SDKUsageErrorWithExamples(
475
+ msg=f"Invalid incident behavior: {config}",
476
+ examples=ASSERTION_INCIDENT_BEHAVIOR_EXAMPLES,
477
+ ) from e
478
+ if isinstance(config, AssertionIncidentBehavior):
479
+ return [config]
480
+ elif isinstance(config, list):
481
+ incident_behaviors = []
482
+ for item in config:
483
+ if isinstance(item, str):
484
+ try:
485
+ incident_behaviors.append(AssertionIncidentBehavior(item))
486
+ except ValueError as e:
487
+ raise SDKUsageErrorWithExamples(
488
+ msg=f"Invalid incident behavior: {item}",
489
+ examples=ASSERTION_INCIDENT_BEHAVIOR_EXAMPLES,
490
+ ) from e
491
+ elif isinstance(item, AssertionIncidentBehavior):
492
+ incident_behaviors.append(item)
493
+ else:
494
+ raise SDKUsageErrorWithExamples(
495
+ msg=f"Invalid incident behavior: {item}",
496
+ examples=ASSERTION_INCIDENT_BEHAVIOR_EXAMPLES,
497
+ )
498
+ return incident_behaviors
499
+ else:
500
+ raise SDKUsageErrorWithExamples(
501
+ msg=f"Invalid incident behavior: {config}",
502
+ examples=ASSERTION_INCIDENT_BEHAVIOR_EXAMPLES,
503
+ )
504
+
505
+
506
+ def _generate_default_name(prefix: str, suffix_length: int) -> str:
507
+ return f"{prefix}-{''.join(random.choices(string.ascii_letters + string.digits, k=suffix_length))}"
508
+
509
+
510
+ TRAINING_DATA_LOOKBACK_DAYS_EXAMPLES = {
511
+ "Training data lookback days from int": ASSERTION_MONITOR_DEFAULT_TRAINING_LOOKBACK_WINDOW_DAYS,
512
+ f"Training data lookback days from None (uses default of {ASSERTION_MONITOR_DEFAULT_TRAINING_LOOKBACK_WINDOW_DAYS} days)": None,
513
+ }
514
+
515
+
516
+ def _try_parse_training_data_lookback_days(
517
+ training_data_lookback_days: Optional[int],
518
+ ) -> int:
519
+ if training_data_lookback_days is None:
520
+ return ASSERTION_MONITOR_DEFAULT_TRAINING_LOOKBACK_WINDOW_DAYS
521
+ if isinstance(training_data_lookback_days, str):
522
+ try:
523
+ training_data_lookback_days = int(training_data_lookback_days)
524
+ except ValueError as e:
525
+ raise SDKUsageErrorWithExamples(
526
+ msg=f"Invalid training data lookback days: {training_data_lookback_days}",
527
+ examples=TRAINING_DATA_LOOKBACK_DAYS_EXAMPLES,
528
+ ) from e
529
+ if not isinstance(training_data_lookback_days, int):
530
+ raise SDKUsageErrorWithExamples(
531
+ msg=f"Invalid training data lookback days: {training_data_lookback_days}",
532
+ examples=TRAINING_DATA_LOOKBACK_DAYS_EXAMPLES,
533
+ )
534
+ if training_data_lookback_days < 0:
535
+ raise SDKUsageError("Training data lookback days must be non-negative")
536
+ return training_data_lookback_days
537
+
538
+
539
+ def _validate_cron_schedule(schedule: str, timezone: str) -> None:
540
+ """We are using the POSIX.1-2017 standard for cron expressions.
541
+
542
+ Note: We are using the croniter library for cron parsing which is different from executor, which uses apscheduler, so there is a risk of mismatch here.
543
+ """
544
+ try:
545
+ # Validate timezone - pytz.timezone() raises UnknownTimeZoneError for invalid timezones
546
+ # Skip timezone validation when empty
547
+ if timezone:
548
+ pytz.timezone(timezone)
549
+
550
+ # Validate 5-field cron expression only (POSIX.1-2017 standard)
551
+ fields = schedule.strip().split()
552
+ if len(fields) != 5:
553
+ raise ValueError("POSIX.1-2017 requires exactly 5 fields")
554
+
555
+ # POSIX.1-2017 specific validation: Sunday must be 0, not 7
556
+ # However croniter accepts 7 as Sunday, so custom check is needed here.
557
+ # Check the day-of-week field (5th field, index 4)
558
+ dow_field = fields[4]
559
+ if "7" in dow_field:
560
+ # Check if 7 appears as a standalone value or in ranges
561
+ import re
562
+
563
+ # Match 7 as standalone, in lists, or in ranges
564
+ if re.search(r"\b7\b|7-|,7,|^7,|,7$|-7\b", dow_field):
565
+ raise ValueError(
566
+ "POSIX.1-2017 standard: Sunday must be represented as 0, not 7"
567
+ )
568
+
569
+ # Validate cron expression - croniter constructor validates the expression
570
+ croniter(schedule)
571
+
572
+ except Exception as e:
573
+ raise SDKUsageError(
574
+ f"Invalid cron expression or timezone: {schedule} {timezone}, please use a POSIX.1-2017 compatible cron expression and timezone."
575
+ ) from e
576
+
577
+
578
+ def _try_parse_schedule(
579
+ schedule: Optional[Union[str, models.CronScheduleClass]],
580
+ ) -> Optional[models.CronScheduleClass]:
581
+ if schedule is None:
582
+ return None
583
+ if isinstance(schedule, str):
584
+ _validate_cron_schedule(schedule, "UTC")
585
+ return models.CronScheduleClass(
586
+ cron=schedule,
587
+ timezone="UTC",
588
+ )
589
+ if isinstance(schedule, models.CronScheduleClass):
590
+ _validate_cron_schedule(schedule.cron, schedule.timezone)
591
+ return schedule
592
+
593
+
594
+ FieldSpecType = Union[models.FreshnessFieldSpecClass, models.SchemaFieldSpecClass]
595
+
596
+
597
+ class _AssertionInput(ABC):
598
+ def __init__(
599
+ self,
600
+ *,
601
+ # Required fields
602
+ dataset_urn: Union[str, DatasetUrn],
603
+ entity_client: EntityClient, # Needed to get the schema field spec for the detection mechanism if needed
604
+ # Optional fields
605
+ urn: Optional[
606
+ Union[str, AssertionUrn]
607
+ ] = None, # Can be None if the assertion is not yet created
608
+ display_name: Optional[str] = None,
609
+ enabled: bool = True,
610
+ schedule: Optional[Union[str, models.CronScheduleClass]] = None,
611
+ detection_mechanism: DetectionMechanismInputTypes = None,
612
+ sensitivity: Optional[Union[str, InferenceSensitivity]] = None,
613
+ exclusion_windows: Optional[ExclusionWindowInputTypes] = None,
614
+ training_data_lookback_days: Optional[int] = None,
615
+ incident_behavior: Optional[
616
+ Union[AssertionIncidentBehavior, list[AssertionIncidentBehavior]]
617
+ ] = None,
618
+ tags: Optional[TagsInputType] = None,
619
+ source_type: str = models.AssertionSourceTypeClass.NATIVE, # Verified on init to be a valid enum value
620
+ created_by: Union[str, CorpUserUrn],
621
+ created_at: datetime,
622
+ updated_by: Union[str, CorpUserUrn],
623
+ updated_at: datetime,
624
+ ):
625
+ """
626
+ Create an AssertionInput object.
627
+
628
+ Args:
629
+ dataset_urn: The urn of the dataset to be monitored.
630
+ entity_client: The entity client to be used for creating the assertion.
631
+ urn: The urn of the assertion. If not provided, a random urn will be generated.
632
+ display_name: The display name of the assertion. If not provided, a random display name will be generated.
633
+ enabled: Whether the assertion is enabled. Defaults to True.
634
+ detection_mechanism: The detection mechanism to be used for the assertion.
635
+ sensitivity: The sensitivity to be applied to the assertion.
636
+ exclusion_windows: The exclusion windows to be applied to the assertion. If not provided, no exclusion windows will be applied.
637
+ training_data_lookback_days: The training data lookback days to be applied to the assertion.
638
+ incident_behavior: The incident behavior to be applied to the assertion.
639
+ tags: The tags to be applied to the assertion.
640
+ source_type: The source type of the assertion. Defaults to models.AssertionSourceTypeClass.NATIVE.
641
+ created_by: The actor that created the assertion.
642
+ created_at: The timestamp of the assertion creation.
643
+ updated_by: The actor that last updated the assertion.
644
+ updated_at: The timestamp of the assertion last update.
645
+ """
646
+ self.dataset_urn = DatasetUrn.from_string(dataset_urn)
647
+ self.entity_client = entity_client
648
+ self.urn = AssertionUrn(urn) if urn else None
649
+ self.display_name = (
650
+ display_name
651
+ if display_name is not None
652
+ else _generate_default_name(DEFAULT_NAME_PREFIX, DEFAULT_NAME_SUFFIX_LENGTH)
653
+ )
654
+ self.enabled = enabled
655
+ self.schedule = _try_parse_schedule(schedule)
656
+ self.detection_mechanism = DetectionMechanism.parse(detection_mechanism)
657
+ self.sensitivity = InferenceSensitivity.parse(sensitivity)
658
+ self.exclusion_windows = _try_parse_exclusion_window(exclusion_windows)
659
+ self.training_data_lookback_days = _try_parse_training_data_lookback_days(
660
+ training_data_lookback_days
661
+ )
662
+ self.incident_behavior = _try_parse_incident_behavior(incident_behavior)
663
+ self.tags = tags
664
+ if source_type not in get_enum_options(models.AssertionSourceTypeClass):
665
+ raise SDKUsageError(
666
+ msg=f"Invalid source type: {source_type}, valid options are {get_enum_options(models.AssertionSourceTypeClass)}",
667
+ )
668
+ self.source_type = source_type
669
+ self.created_by = created_by
670
+ self.created_at = created_at
671
+ self.updated_by = updated_by
672
+ self.updated_at = updated_at
673
+
674
+ self.cached_dataset: Optional[Dataset] = None
675
+
676
+ def to_assertion_and_monitor_entities(self) -> tuple[Assertion, Monitor]:
677
+ """
678
+ Convert the assertion input to an assertion and monitor entity.
679
+
680
+ Returns:
681
+ A tuple of (assertion, monitor) entities.
682
+ """
683
+ assertion = self.to_assertion_entity()
684
+ monitor = self.to_monitor_entity(assertion.urn)
685
+ return assertion, monitor
686
+
687
+ def to_assertion_entity(self) -> Assertion:
688
+ """
689
+ Convert the assertion input to an assertion entity.
690
+
691
+ Returns:
692
+ The created assertion entity.
693
+ """
694
+ on_success, on_failure = self._convert_incident_behavior()
695
+ filter = self._create_filter_from_detection_mechanism()
696
+
697
+ return Assertion(
698
+ id=self.urn,
699
+ info=self._create_assertion_info(filter),
700
+ description=self.display_name,
701
+ on_success=on_success,
702
+ on_failure=on_failure,
703
+ tags=self._convert_tags(),
704
+ source=self._convert_source(),
705
+ last_updated=self._convert_last_updated(),
706
+ )
707
+
708
+ def _convert_incident_behavior(
709
+ self,
710
+ ) -> tuple[
711
+ Optional[AssertionActionsInputType],
712
+ Optional[AssertionActionsInputType],
713
+ ]:
714
+ """
715
+ Convert incident behavior to on_success and on_failure actions.
716
+
717
+ Returns:
718
+ A tuple of (on_success, on_failure) actions.
719
+ """
720
+ if not self.incident_behavior:
721
+ return None, None
722
+
723
+ behaviors = (
724
+ [self.incident_behavior]
725
+ if isinstance(self.incident_behavior, AssertionIncidentBehavior)
726
+ else self.incident_behavior
727
+ )
728
+
729
+ on_success: Optional[AssertionActionsInputType] = [
730
+ models.AssertionActionClass(
731
+ type=models.AssertionActionTypeClass.RESOLVE_INCIDENT
732
+ )
733
+ for behavior in behaviors
734
+ if behavior == AssertionIncidentBehavior.RESOLVE_ON_PASS
735
+ ] or None
736
+
737
+ on_failure: Optional[AssertionActionsInputType] = [
738
+ models.AssertionActionClass(
739
+ type=models.AssertionActionTypeClass.RAISE_INCIDENT
740
+ )
741
+ for behavior in behaviors
742
+ if behavior == AssertionIncidentBehavior.RAISE_ON_FAIL
743
+ ] or None
744
+
745
+ return on_success, on_failure
746
+
747
+ def _create_filter_from_detection_mechanism(
748
+ self,
749
+ ) -> Optional[models.DatasetFilterClass]:
750
+ """
751
+ Create a filter from the detection mechanism if it has an additional filter.
752
+
753
+ Returns:
754
+ A DatasetFilterClass if the detection mechanism has an additional filter, None otherwise.
755
+ """
756
+ if not isinstance(
757
+ self.detection_mechanism,
758
+ _DETECTION_MECHANISM_TYPES_WITH_ADDITIONAL_FILTER,
759
+ ):
760
+ return None
761
+
762
+ additional_filter = self.detection_mechanism.additional_filter
763
+ if not additional_filter:
764
+ return None
765
+
766
+ return models.DatasetFilterClass(
767
+ type=models.DatasetFilterTypeClass.SQL,
768
+ sql=additional_filter,
769
+ )
770
+
771
+ def _convert_tags(self) -> Optional[TagsInputType]:
772
+ """
773
+ Convert the tags input into a standardized format.
774
+
775
+ Returns:
776
+ A list of tags or None if no tags are provided.
777
+
778
+ Raises:
779
+ SDKUsageErrorWithExamples: If the tags input is invalid.
780
+ """
781
+ if not self.tags:
782
+ return None
783
+
784
+ if isinstance(self.tags, str):
785
+ return [self.tags]
786
+ elif isinstance(self.tags, list):
787
+ return self.tags
788
+ else:
789
+ raise SDKUsageErrorWithExamples(
790
+ msg=f"Invalid tags: {self.tags}",
791
+ examples={
792
+ "Tags from string": "urn:li:tag:my_tag_1",
793
+ "Tags from list": [
794
+ "urn:li:tag:my_tag_1",
795
+ "urn:li:tag:my_tag_2",
796
+ ],
797
+ },
798
+ )
799
+
800
+ def _convert_source(self) -> models.AssertionSourceClass:
801
+ """
802
+ Convert the source input into a models.AssertionSourceClass.
803
+ """
804
+ return models.AssertionSourceClass(
805
+ type=self.source_type,
806
+ created=models.AuditStampClass(
807
+ time=make_ts_millis(self.created_at),
808
+ actor=str(self.created_by),
809
+ ),
810
+ )
811
+
812
+ def _convert_last_updated(self) -> tuple[datetime, str]:
813
+ """
814
+ Convert the last updated input into a tuple of (datetime, str).
815
+
816
+ Validation is handled in the Assertion entity constructor.
817
+ """
818
+ return (self.updated_at, str(self.updated_by))
819
+
820
+ def to_monitor_entity(self, assertion_urn: AssertionUrn) -> Monitor:
821
+ """
822
+ Convert the assertion input to a monitor entity.
823
+
824
+ Args:
825
+ assertion_urn: The URN of the assertion to monitor.
826
+
827
+ Returns:
828
+ A Monitor entity configured with the assertion input parameters.
829
+ """
830
+ source_type, field = self._convert_assertion_source_type_and_field()
831
+ return Monitor(
832
+ id=(self.dataset_urn, assertion_urn),
833
+ info=self._create_monitor_info(
834
+ assertion_urn=assertion_urn,
835
+ status=self._convert_monitor_status(),
836
+ schedule=self._convert_schedule(),
837
+ source_type=source_type,
838
+ field=field,
839
+ sensitivity=self._convert_sensitivity(),
840
+ exclusion_windows=self._convert_exclusion_windows(),
841
+ ),
842
+ )
843
+
844
+ def _convert_monitor_status(self) -> models.MonitorStatusClass:
845
+ """
846
+ Convert the enabled flag into a MonitorStatusClass.
847
+
848
+ Returns:
849
+ A MonitorStatusClass with ACTIVE or INACTIVE mode based on the enabled flag.
850
+ """
851
+ return models.MonitorStatusClass(
852
+ mode=models.MonitorModeClass.ACTIVE
853
+ if self.enabled
854
+ else models.MonitorModeClass.INACTIVE,
855
+ )
856
+
857
+ def _convert_exclusion_windows(
858
+ self,
859
+ ) -> list[models.AssertionExclusionWindowClass]:
860
+ """
861
+ Convert exclusion windows into AssertionExclusionWindowClass objects including generating display names for them.
862
+
863
+ Returns:
864
+ A list of AssertionExclusionWindowClass objects.
865
+
866
+ Raises:
867
+ SDKUsageErrorWithExamples: If an exclusion window is of an invalid type.
868
+ """
869
+ exclusion_windows: list[models.AssertionExclusionWindowClass] = []
870
+ if self.exclusion_windows:
871
+ for window in self.exclusion_windows:
872
+ if not isinstance(window, FixedRangeExclusionWindow):
873
+ raise SDKUsageErrorWithExamples(
874
+ msg=f"Invalid exclusion window type: {window}",
875
+ examples=FIXED_RANGE_EXCLUSION_WINDOW_EXAMPLES,
876
+ )
877
+ # To match the UI, we generate a display name for the exclusion window.
878
+ # See here for the UI code: https://github.com/acryldata/datahub-fork/blob/acryl-main/datahub-web-react/src/app/entityV2/shared/tabs/Dataset/Validations/assertion/builder/steps/inferred/common/ExclusionWindowAdjuster.tsx#L31
879
+ # Copied here for reference: displayName: `${dayjs(startTime).format('MMM D, h:mm A')} - ${dayjs(endTime).format('MMM D, h:mm A')}`,
880
+ generated_display_name = f"{window.start.strftime('%b %-d, %-I:%M %p')} - {window.end.strftime('%b %-d, %-I:%M %p')}"
881
+ exclusion_windows.append(
882
+ models.AssertionExclusionWindowClass(
883
+ type=models.AssertionExclusionWindowTypeClass.FIXED_RANGE, # Currently only fixed range is supported
884
+ displayName=generated_display_name,
885
+ fixedRange=models.AbsoluteTimeWindowClass(
886
+ startTimeMillis=make_ts_millis(window.start),
887
+ endTimeMillis=make_ts_millis(window.end),
888
+ ),
889
+ )
890
+ )
891
+ return exclusion_windows
892
+
893
+ def _convert_sensitivity(self) -> models.AssertionMonitorSensitivityClass:
894
+ """
895
+ Convert sensitivity into an AssertionMonitorSensitivityClass.
896
+
897
+ Returns:
898
+ An AssertionMonitorSensitivityClass with the appropriate sensitivity.
899
+ """
900
+ return models.AssertionMonitorSensitivityClass(
901
+ level=InferenceSensitivity.to_int(self.sensitivity),
902
+ )
903
+
904
+ def _get_schema_field_spec(self, column_name: str) -> models.SchemaFieldSpecClass:
905
+ """
906
+ Get the schema field spec for the detection mechanism if needed.
907
+ """
908
+ # Only fetch the dataset if it's not already cached.
909
+ # Also we only fetch the dataset if it's needed for the detection mechanism.
910
+ if self.cached_dataset is None:
911
+ self.cached_dataset = self.entity_client.get(self.dataset_urn)
912
+
913
+ # Handle case where dataset doesn't exist
914
+ if self.cached_dataset is None:
915
+ raise SDKUsageError(
916
+ f"Dataset {self.dataset_urn} not found. Cannot validate column {column_name}."
917
+ )
918
+
919
+ # TODO: Make a public accessor for _schema_dict in the SDK
920
+ schema_fields = self.cached_dataset._schema_dict()
921
+ field = schema_fields.get(column_name)
922
+ if field:
923
+ return models.SchemaFieldSpecClass(
924
+ path=field.fieldPath,
925
+ type=field.type.type.__class__.__name__,
926
+ nativeType=field.nativeDataType,
927
+ )
928
+ else:
929
+ raise SDKUsageError(
930
+ msg=f"Column {column_name} not found in dataset {self.dataset_urn}",
931
+ )
932
+
933
+ def _validate_field_type(
934
+ self,
935
+ field_spec: models.SchemaFieldSpecClass,
936
+ column_name: str,
937
+ allowed_types: list[DictWrapper],
938
+ field_type_name: str,
939
+ ) -> None:
940
+ """
941
+ Validate that a field has an allowed type.
942
+
943
+ Args:
944
+ field_spec: The field specification to validate
945
+ column_name: The name of the column for error messages
946
+ allowed_types: List of allowed field types
947
+ field_type_name: Human-readable name of the field type for error messages
948
+
949
+ Raises:
950
+ SDKUsageError: If the field has an invalid type
951
+ """
952
+ allowed_type_names = [t.__class__.__name__ for t in allowed_types]
953
+ if field_spec.type not in allowed_type_names:
954
+ raise SDKUsageError(
955
+ msg=f"Column {column_name} with type {field_spec.type} does not have an allowed type for a {field_type_name} in dataset {self.dataset_urn}. "
956
+ f"Allowed types are {allowed_type_names}.",
957
+ )
958
+
959
+ def _create_monitor_info(
960
+ self,
961
+ assertion_urn: AssertionUrn,
962
+ status: models.MonitorStatusClass,
963
+ schedule: models.CronScheduleClass,
964
+ source_type: Union[str, models.DatasetFreshnessSourceTypeClass],
965
+ field: Optional[FieldSpecType],
966
+ sensitivity: models.AssertionMonitorSensitivityClass,
967
+ exclusion_windows: list[models.AssertionExclusionWindowClass],
968
+ ) -> models.MonitorInfoClass:
969
+ """
970
+ Create a MonitorInfoClass with all the necessary components.
971
+
972
+ Args:
973
+ status: The monitor status.
974
+ schedule: The monitor schedule.
975
+ source_type: The source type.
976
+ field: Optional field specification.
977
+ sensitivity: The monitor sensitivity.
978
+ exclusion_windows: List of exclusion windows.
979
+
980
+ Returns:
981
+ A MonitorInfoClass configured with all the provided components.
982
+ """
983
+ return models.MonitorInfoClass(
984
+ type=models.MonitorTypeClass.ASSERTION,
985
+ status=status,
986
+ assertionMonitor=models.AssertionMonitorClass(
987
+ assertions=[
988
+ models.AssertionEvaluationSpecClass(
989
+ assertion=str(assertion_urn),
990
+ schedule=schedule,
991
+ parameters=self._get_assertion_evaluation_parameters(
992
+ str(source_type), field
993
+ ),
994
+ )
995
+ ],
996
+ settings=models.AssertionMonitorSettingsClass(
997
+ adjustmentSettings=models.AssertionAdjustmentSettingsClass(
998
+ sensitivity=sensitivity,
999
+ exclusionWindows=exclusion_windows,
1000
+ trainingDataLookbackWindowDays=self.training_data_lookback_days,
1001
+ ),
1002
+ ),
1003
+ ),
1004
+ )
1005
+
1006
+ @abstractmethod
1007
+ def _create_assertion_info(
1008
+ self, filter: Optional[models.DatasetFilterClass]
1009
+ ) -> AssertionInfoInputType:
1010
+ """Create assertion info specific to the assertion type."""
1011
+ pass
1012
+
1013
+ @abstractmethod
1014
+ def _convert_schedule(self) -> models.CronScheduleClass:
1015
+ """Convert schedule to appropriate format for the assertion type."""
1016
+ pass
1017
+
1018
+ @abstractmethod
1019
+ def _get_assertion_evaluation_parameters(
1020
+ self, source_type: str, field: Optional[FieldSpecType]
1021
+ ) -> models.AssertionEvaluationParametersClass:
1022
+ """Get evaluation parameters specific to the assertion type."""
1023
+ pass
1024
+
1025
+ @abstractmethod
1026
+ def _convert_assertion_source_type_and_field(
1027
+ self,
1028
+ ) -> tuple[str, Optional[FieldSpecType]]:
1029
+ """Convert detection mechanism to source type and field spec."""
1030
+ pass
1031
+
1032
+
1033
+ class _SmartFreshnessAssertionInput(_AssertionInput):
1034
+ def __init__(
1035
+ self,
1036
+ *,
1037
+ # Required fields
1038
+ dataset_urn: Union[str, DatasetUrn],
1039
+ entity_client: EntityClient, # Needed to get the schema field spec for the detection mechanism if needed
1040
+ # Optional fields
1041
+ urn: Optional[Union[str, AssertionUrn]] = None,
1042
+ display_name: Optional[str] = None,
1043
+ enabled: bool = True,
1044
+ schedule: Optional[Union[str, models.CronScheduleClass]] = None,
1045
+ detection_mechanism: DetectionMechanismInputTypes = None,
1046
+ sensitivity: Optional[Union[str, InferenceSensitivity]] = None,
1047
+ exclusion_windows: Optional[ExclusionWindowInputTypes] = None,
1048
+ training_data_lookback_days: Optional[int] = None,
1049
+ incident_behavior: Optional[
1050
+ Union[AssertionIncidentBehavior, list[AssertionIncidentBehavior]]
1051
+ ] = None,
1052
+ tags: Optional[TagsInputType] = None,
1053
+ created_by: Union[str, CorpUserUrn],
1054
+ created_at: datetime,
1055
+ updated_by: Union[str, CorpUserUrn],
1056
+ updated_at: datetime,
1057
+ ):
1058
+ super().__init__(
1059
+ dataset_urn=dataset_urn,
1060
+ entity_client=entity_client,
1061
+ urn=urn,
1062
+ display_name=display_name,
1063
+ enabled=enabled,
1064
+ schedule=schedule
1065
+ if schedule is not None
1066
+ else DEFAULT_SCHEDULE, # Use provided schedule or default for create case
1067
+ detection_mechanism=detection_mechanism,
1068
+ sensitivity=sensitivity,
1069
+ exclusion_windows=exclusion_windows,
1070
+ training_data_lookback_days=training_data_lookback_days,
1071
+ incident_behavior=incident_behavior,
1072
+ tags=tags,
1073
+ source_type=models.AssertionSourceTypeClass.INFERRED, # Smart assertions are of type inferred, not native
1074
+ created_by=created_by,
1075
+ created_at=created_at,
1076
+ updated_by=updated_by,
1077
+ updated_at=updated_at,
1078
+ )
1079
+
1080
+ def _create_assertion_info(
1081
+ self, filter: Optional[models.DatasetFilterClass]
1082
+ ) -> AssertionInfoInputType:
1083
+ """
1084
+ Create a FreshnessAssertionInfoClass for a smart freshness assertion.
1085
+
1086
+ Args:
1087
+ filter: Optional filter to apply to the assertion.
1088
+
1089
+ Returns:
1090
+ A FreshnessAssertionInfoClass configured for smart freshness.
1091
+ """
1092
+ return models.FreshnessAssertionInfoClass(
1093
+ type=models.FreshnessAssertionTypeClass.DATASET_CHANGE, # Currently only dataset change is supported
1094
+ entity=str(self.dataset_urn),
1095
+ # schedule (optional, must be left empty for smart freshness assertions - managed by the AI inference engine)
1096
+ filter=filter,
1097
+ )
1098
+
1099
+ def _convert_schedule(self) -> models.CronScheduleClass:
1100
+ """Create a schedule for a smart freshness assertion.
1101
+
1102
+ For create case, uses DEFAULT_SCHEDULE. For update case, preserves existing schedule.
1103
+
1104
+ Returns:
1105
+ A CronScheduleClass with appropriate schedule settings.
1106
+ """
1107
+ assert self.schedule is not None, (
1108
+ "Schedule should never be None due to constructor logic"
1109
+ )
1110
+ return self.schedule
1111
+
1112
+ def _get_assertion_evaluation_parameters(
1113
+ self, source_type: str, field: Optional[FieldSpecType]
1114
+ ) -> models.AssertionEvaluationParametersClass:
1115
+ # Ensure field is either None or FreshnessFieldSpecClass
1116
+ freshness_field = None
1117
+ if field is not None:
1118
+ if not isinstance(field, models.FreshnessFieldSpecClass):
1119
+ raise SDKUsageError(
1120
+ f"Expected FreshnessFieldSpecClass for freshness assertion, got {type(field).__name__}"
1121
+ )
1122
+ freshness_field = field
1123
+
1124
+ return models.AssertionEvaluationParametersClass(
1125
+ type=models.AssertionEvaluationParametersTypeClass.DATASET_FRESHNESS,
1126
+ datasetFreshnessParameters=models.DatasetFreshnessAssertionParametersClass(
1127
+ sourceType=source_type, field=freshness_field
1128
+ ),
1129
+ )
1130
+
1131
+ def _convert_assertion_source_type_and_field(
1132
+ self,
1133
+ ) -> tuple[str, Optional[FieldSpecType]]:
1134
+ """
1135
+ Convert detection mechanism into source type and field specification for freshness assertions.
1136
+
1137
+ Returns:
1138
+ A tuple of (source_type, field) where field may be None.
1139
+ Note that the source_type is a string, not a models.DatasetFreshnessSourceTypeClass (or other assertion source type) since
1140
+ the source type is not a enum in the code generated from the DatasetFreshnessSourceType enum in the PDL.
1141
+
1142
+ Raises:
1143
+ SDKNotYetSupportedError: If the detection mechanism is not supported.
1144
+ SDKUsageError: If the field (column) is not found in the dataset,
1145
+ and the detection mechanism requires a field. Also if the field
1146
+ is not an allowed type for the detection mechanism.
1147
+ """
1148
+ source_type = models.DatasetFreshnessSourceTypeClass.INFORMATION_SCHEMA
1149
+ field = None
1150
+
1151
+ if isinstance(self.detection_mechanism, _LastModifiedColumn):
1152
+ source_type = models.DatasetFreshnessSourceTypeClass.FIELD_VALUE
1153
+ field = self._create_field_spec(
1154
+ self.detection_mechanism.column_name,
1155
+ LAST_MODIFIED_ALLOWED_FIELD_TYPES,
1156
+ "last modified column",
1157
+ models.FreshnessFieldKindClass.LAST_MODIFIED,
1158
+ )
1159
+ elif isinstance(self.detection_mechanism, _InformationSchema):
1160
+ source_type = models.DatasetFreshnessSourceTypeClass.INFORMATION_SCHEMA
1161
+ elif isinstance(self.detection_mechanism, _DataHubOperation):
1162
+ source_type = models.DatasetFreshnessSourceTypeClass.DATAHUB_OPERATION
1163
+ elif isinstance(self.detection_mechanism, _AuditLog):
1164
+ source_type = models.DatasetFreshnessSourceTypeClass.AUDIT_LOG
1165
+ else:
1166
+ raise SDKNotYetSupportedError(
1167
+ f"Detection mechanism {self.detection_mechanism} not yet supported for smart freshness assertions"
1168
+ )
1169
+
1170
+ return source_type, field
1171
+
1172
+ def _create_field_spec(
1173
+ self,
1174
+ column_name: str,
1175
+ allowed_types: list[DictWrapper], # TODO: Use the type from the PDL
1176
+ field_type_name: str,
1177
+ kind: str,
1178
+ ) -> models.FreshnessFieldSpecClass:
1179
+ """
1180
+ Create a field specification for a column, validating its type.
1181
+
1182
+ Args:
1183
+ column_name: The name of the column to create a spec for
1184
+ allowed_types: List of allowed field types
1185
+ field_type_name: Human-readable name of the field type for error messages
1186
+ kind: The kind of field to create
1187
+
1188
+ Returns:
1189
+ A FreshnessFieldSpecClass for the column
1190
+
1191
+ Raises:
1192
+ SDKUsageError: If the column is not found or has an invalid type
1193
+ """
1194
+ SUPPORTED_KINDS = [
1195
+ models.FreshnessFieldKindClass.LAST_MODIFIED,
1196
+ models.FreshnessFieldKindClass.HIGH_WATERMARK,
1197
+ ]
1198
+ if kind not in SUPPORTED_KINDS:
1199
+ raise SDKUsageError(
1200
+ msg=f"Invalid kind: {kind}. Must be one of {SUPPORTED_KINDS}",
1201
+ )
1202
+
1203
+ field_spec = self._get_schema_field_spec(column_name)
1204
+ self._validate_field_type(
1205
+ field_spec, column_name, allowed_types, field_type_name
1206
+ )
1207
+ return models.FreshnessFieldSpecClass(
1208
+ path=field_spec.path,
1209
+ type=field_spec.type,
1210
+ nativeType=field_spec.nativeType,
1211
+ kind=kind,
1212
+ )
1213
+
1214
+
1215
+ class _SmartVolumeAssertionInput(_AssertionInput):
1216
+ def __init__(
1217
+ self,
1218
+ *,
1219
+ # Required fields
1220
+ dataset_urn: Union[str, DatasetUrn],
1221
+ entity_client: EntityClient, # Needed to get the schema field spec for the detection mechanism if needed
1222
+ # Optional fields
1223
+ urn: Optional[Union[str, AssertionUrn]] = None,
1224
+ display_name: Optional[str] = None,
1225
+ enabled: bool = True,
1226
+ schedule: Optional[Union[str, models.CronScheduleClass]] = None,
1227
+ detection_mechanism: DetectionMechanismInputTypes = None,
1228
+ sensitivity: Optional[Union[str, InferenceSensitivity]] = None,
1229
+ exclusion_windows: Optional[ExclusionWindowInputTypes] = None,
1230
+ training_data_lookback_days: Optional[int] = None,
1231
+ incident_behavior: Optional[
1232
+ Union[AssertionIncidentBehavior, list[AssertionIncidentBehavior]]
1233
+ ] = None,
1234
+ tags: Optional[TagsInputType] = None,
1235
+ created_by: Union[str, CorpUserUrn],
1236
+ created_at: datetime,
1237
+ updated_by: Union[str, CorpUserUrn],
1238
+ updated_at: datetime,
1239
+ ):
1240
+ super().__init__(
1241
+ dataset_urn=dataset_urn,
1242
+ entity_client=entity_client,
1243
+ urn=urn,
1244
+ display_name=display_name,
1245
+ enabled=enabled,
1246
+ schedule=schedule,
1247
+ detection_mechanism=detection_mechanism,
1248
+ sensitivity=sensitivity,
1249
+ exclusion_windows=exclusion_windows,
1250
+ training_data_lookback_days=training_data_lookback_days,
1251
+ incident_behavior=incident_behavior,
1252
+ tags=tags,
1253
+ source_type=models.AssertionSourceTypeClass.INFERRED, # Smart assertions are of type inferred, not native
1254
+ created_by=created_by,
1255
+ created_at=created_at,
1256
+ updated_by=updated_by,
1257
+ updated_at=updated_at,
1258
+ )
1259
+
1260
+ def _create_assertion_info(
1261
+ self, filter: Optional[models.DatasetFilterClass]
1262
+ ) -> AssertionInfoInputType:
1263
+ """
1264
+ Create a VolumeAssertionInfoClass for a smart volume assertion.
1265
+
1266
+ Args:
1267
+ filter: Optional filter to apply to the assertion.
1268
+
1269
+ Returns:
1270
+ A VolumeAssertionInfoClass configured for smart volume.
1271
+ """
1272
+ return models.VolumeAssertionInfoClass(
1273
+ type=models.VolumeAssertionTypeClass.ROW_COUNT_TOTAL, # Currently only ROW_COUNT_TOTAL is supported for smart volume
1274
+ entity=str(self.dataset_urn),
1275
+ filter=filter,
1276
+ )
1277
+
1278
+ def _convert_schedule(self) -> models.CronScheduleClass:
1279
+ """Create a schedule for a smart freshness assertion.
1280
+
1281
+ Since the schedule is not used for smart freshness assertions, we return a default schedule.
1282
+
1283
+ Returns:
1284
+ A CronScheduleClass with appropriate schedule settings.
1285
+ """
1286
+ if self.schedule is None:
1287
+ return DEFAULT_SCHEDULE
1288
+
1289
+ return models.CronScheduleClass(
1290
+ cron=self.schedule.cron,
1291
+ timezone=self.schedule.timezone,
1292
+ )
1293
+
1294
+ def _get_assertion_evaluation_parameters(
1295
+ self, source_type: str, field: Optional[FieldSpecType]
1296
+ ) -> models.AssertionEvaluationParametersClass:
1297
+ return models.AssertionEvaluationParametersClass(
1298
+ type=models.AssertionEvaluationParametersTypeClass.DATASET_VOLUME,
1299
+ datasetVolumeParameters=models.DatasetVolumeAssertionParametersClass(
1300
+ sourceType=source_type,
1301
+ ),
1302
+ )
1303
+
1304
+ def _convert_assertion_source_type_and_field(
1305
+ self,
1306
+ ) -> tuple[str, Optional[FieldSpecType]]:
1307
+ """
1308
+ Convert detection mechanism into source type and field specification for volume assertions.
1309
+
1310
+ Returns:
1311
+ A tuple of (source_type, field) where field may be None.
1312
+ Note that the source_type is a string, not a models.DatasetFreshnessSourceTypeClass (or other assertion source type) since
1313
+ the source type is not a enum in the code generated from the DatasetFreshnessSourceType enum in the PDL.
1314
+
1315
+ Raises:
1316
+ SDKNotYetSupportedError: If the detection mechanism is not supported.
1317
+ SDKUsageError: If the field (column) is not found in the dataset,
1318
+ and the detection mechanism requires a field. Also if the field
1319
+ is not an allowed type for the detection mechanism.
1320
+ """
1321
+ source_type = models.DatasetVolumeSourceTypeClass.INFORMATION_SCHEMA
1322
+ field = None
1323
+
1324
+ if isinstance(self.detection_mechanism, _Query):
1325
+ source_type = models.DatasetVolumeSourceTypeClass.QUERY
1326
+ elif isinstance(self.detection_mechanism, _InformationSchema):
1327
+ source_type = models.DatasetVolumeSourceTypeClass.INFORMATION_SCHEMA
1328
+ elif isinstance(self.detection_mechanism, _DatasetProfile):
1329
+ source_type = models.DatasetVolumeSourceTypeClass.DATAHUB_DATASET_PROFILE
1330
+ else:
1331
+ raise SDKNotYetSupportedError(
1332
+ f"Detection mechanism {self.detection_mechanism} not yet supported for smart volume assertions"
1333
+ )
1334
+
1335
+ return source_type, field