acryl-datahub-cloud 0.3.12rc4__py3-none-any.whl → 0.3.12rc5__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.

@@ -4,13 +4,16 @@ import logging
4
4
  from datetime import datetime, timezone
5
5
  from typing import TYPE_CHECKING, Any, Optional, Union
6
6
 
7
- from acryl_datahub_cloud.sdk.assertion import (
7
+ from acryl_datahub_cloud.sdk.assertion.assertion_base import (
8
8
  AssertionMode,
9
9
  FreshnessAssertion,
10
10
  SmartFreshnessAssertion,
11
11
  SmartVolumeAssertion,
12
12
  _AssertionPublic,
13
13
  )
14
+ from acryl_datahub_cloud.sdk.assertion.smart_column_metric_assertion import (
15
+ SmartColumnMetricAssertion,
16
+ )
14
17
  from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
15
18
  AssertionIncidentBehavior,
16
19
  DetectionMechanismInputTypes,
@@ -24,6 +27,15 @@ from acryl_datahub_cloud.sdk.assertion_input.assertion_input import (
24
27
  from acryl_datahub_cloud.sdk.assertion_input.freshness_assertion_input import (
25
28
  _FreshnessAssertionInput,
26
29
  )
30
+ from acryl_datahub_cloud.sdk.assertion_input.smart_column_metric_assertion_input import (
31
+ MetricInputType,
32
+ OperatorInputType,
33
+ RangeInputType,
34
+ RangeTypeInputType,
35
+ ValueInputType,
36
+ ValueTypeInputType,
37
+ _SmartColumnMetricAssertionInput,
38
+ )
27
39
  from acryl_datahub_cloud.sdk.entities.assertion import Assertion, TagsInputType
28
40
  from acryl_datahub_cloud.sdk.entities.monitor import Monitor
29
41
  from acryl_datahub_cloud.sdk.errors import SDKUsageError
@@ -1427,6 +1439,719 @@ class AssertionsClient:
1427
1439
 
1428
1440
  return SmartVolumeAssertion._from_entities(assertion_entity, monitor_entity)
1429
1441
 
1442
+ def sync_smart_column_metric_assertion(
1443
+ self,
1444
+ *,
1445
+ dataset_urn: Union[str, DatasetUrn],
1446
+ column_name: str,
1447
+ metric_type: MetricInputType,
1448
+ operator: OperatorInputType,
1449
+ value: Optional[ValueInputType] = None,
1450
+ value_type: Optional[ValueTypeInputType] = None,
1451
+ range: Optional[RangeInputType] = None,
1452
+ range_type: Optional[RangeTypeInputType] = None,
1453
+ urn: Optional[Union[str, AssertionUrn]] = None,
1454
+ display_name: Optional[str] = None,
1455
+ enabled: Optional[bool] = None,
1456
+ detection_mechanism: DetectionMechanismInputTypes = None,
1457
+ sensitivity: Optional[Union[str, InferenceSensitivity]] = None,
1458
+ exclusion_windows: Optional[ExclusionWindowInputTypes] = None,
1459
+ training_data_lookback_days: Optional[int] = None,
1460
+ incident_behavior: Optional[
1461
+ Union[AssertionIncidentBehavior, list[AssertionIncidentBehavior]]
1462
+ ] = None,
1463
+ tags: Optional[TagsInputType] = None,
1464
+ updated_by: Optional[Union[str, CorpUserUrn]] = None,
1465
+ schedule: Optional[Union[str, models.CronScheduleClass]] = None,
1466
+ ) -> SmartColumnMetricAssertion:
1467
+ """Upsert and merge a smart column metric assertion.
1468
+
1469
+ Note: keyword arguments are required.
1470
+
1471
+ Upsert and merge is a combination of create and update. If the assertion does not exist,
1472
+ it will be created. If it does exist, it will be updated.
1473
+
1474
+ Existing assertion fields will be updated if the input value is not None. If the input value is None, the existing value
1475
+ will be preserved. If the input value can be un-set e.g. by passing an empty list or
1476
+ empty string.
1477
+
1478
+ Args:
1479
+ dataset_urn: The urn of the dataset to be monitored. (Required)
1480
+ column_name: The name of the column to be monitored. (Required)
1481
+ metric_type: The type of the metric to be monitored. (Required)
1482
+ operator: The operator to be used for the assertion. (Required)
1483
+ value: The value to be used for the assertion. (Required if operator requires a value)
1484
+ value_type: The type of the value to be used for the assertion. (Required if operator requires a value)
1485
+ range: The range to be used for the assertion. (Required if operator requires a range)
1486
+ range_type: The type of the range to be used for the assertion. (Required if operator requires a range)
1487
+ urn: The urn of the assertion. If not provided, a urn will be generated and the assertion
1488
+ will be _created_ in the DataHub instance.
1489
+ display_name: The display name of the assertion. If not provided, a random display name
1490
+ will be generated.
1491
+ enabled: Whether the assertion is enabled. If not provided, the existing value
1492
+ will be preserved.
1493
+ detection_mechanism: The detection mechanism to be used for the assertion. Valid values are:
1494
+ - All rows query datahub dataset profile:
1495
+ - "all_rows_query_datahub_dataset_profile" or DetectionMechanism.ALL_ROWS_QUERY_DATAHUB_DATASET_PROFILE
1496
+
1497
+ - All rows query:
1498
+ - "all_rows_query" or DetectionMechanism.ALL_ROWS_QUERY
1499
+ - with optional additional filter: DetectionMechanism.ALL_ROWS_QUERY(additional_filter='last_modified > 2021-01-01')
1500
+ - Or as a dict: {
1501
+ "type": "all_rows_query",
1502
+ "additional_filter": "last_modified > '2021-01-01'", # optional
1503
+ }
1504
+
1505
+ - Changed rows query:
1506
+ - For changed rows query, you need to pass a supported column type (Number, Date or Time)
1507
+ - DetectionMechanism.CHANGED_ROWS_QUERY(column_name='last_modified')
1508
+ - With optional additional filter: DetectionMechanism.CHANGED_ROWS_QUERY(column_name='last_modified', additional_filter='last_modified > 2021-01-01')
1509
+ - Or as a dict: {
1510
+ "type": "changed_rows_query",
1511
+ "column_name": "last_modified",
1512
+ "additional_filter": "last_modified > '2021-01-01'", # optional
1513
+ }
1514
+
1515
+ sensitivity: The sensitivity to be applied to the assertion. Valid values are:
1516
+ - "low" or InferenceSensitivity.LOW
1517
+ - "medium" or InferenceSensitivity.MEDIUM
1518
+ - "high" or InferenceSensitivity.HIGH
1519
+ exclusion_windows: The exclusion windows to be applied to the assertion, currently only
1520
+ fixed range exclusion windows are supported. Valid values are:
1521
+ - from datetime.datetime objects: {
1522
+ "start": "datetime(2025, 1, 1, 0, 0, 0)",
1523
+ "end": "datetime(2025, 1, 2, 0, 0, 0)",
1524
+ }
1525
+ - from string datetimes: {
1526
+ "start": "2025-01-01T00:00:00",
1527
+ "end": "2025-01-02T00:00:00",
1528
+ }
1529
+ - from FixedRangeExclusionWindow objects: FixedRangeExclusionWindow(
1530
+ start=datetime(2025, 1, 1, 0, 0, 0),
1531
+ end=datetime(2025, 1, 2, 0, 0, 0)
1532
+ )
1533
+ training_data_lookback_days: The training data lookback days to be applied to the
1534
+ assertion as an integer.
1535
+ incident_behavior: The incident behavior to be applied to the assertion. Valid values are:
1536
+ - "raise_on_fail" or AssertionIncidentBehavior.RAISE_ON_FAIL
1537
+ - "resolve_on_pass" or AssertionIncidentBehavior.RESOLVE_ON_PASS
1538
+ tags: The tags to be applied to the assertion. Valid values are:
1539
+ - a list of strings (strings will be converted to TagUrn objects)
1540
+ - a list of TagUrn objects
1541
+ - a list of TagAssociationClass objects
1542
+ updated_by: Optional urn of the user who updated the assertion. The format is
1543
+ "urn:li:corpuser:<username>", which you can find on the Users & Groups page.
1544
+ The default is the datahub system user.
1545
+ TODO: Retrieve the SDK user as the default instead of the datahub system user.
1546
+ schedule: Optional cron formatted schedule for the assertion. If not provided, a default
1547
+ schedule of every 6 hours will be used. The schedule determines when the assertion will be evaluated.
1548
+ The format is a cron expression, e.g. "0 * * * *" for every hour using UTC timezone.
1549
+ Alternatively, a models.CronScheduleClass object can be provided with string parameters
1550
+ cron and timezone. Use `from datahub.metadata import schema_classes as models` to import the class.
1551
+
1552
+ Returns:
1553
+ SmartColumnMetricAssertion: The created or updated assertion.
1554
+ """
1555
+ _print_experimental_warning()
1556
+ now_utc = datetime.now(timezone.utc)
1557
+
1558
+ if updated_by is None:
1559
+ logger.warning(
1560
+ f"updated_by is not set, using {DEFAULT_CREATED_BY} as a placeholder"
1561
+ )
1562
+ updated_by = DEFAULT_CREATED_BY
1563
+
1564
+ # 1. If urn is not set, create a new assertion
1565
+ if urn is None:
1566
+ logger.info("URN is not set, creating a new assertion")
1567
+ return self._create_smart_column_metric_assertion(
1568
+ dataset_urn=dataset_urn,
1569
+ column_name=column_name,
1570
+ metric_type=metric_type,
1571
+ operator=operator,
1572
+ value=value,
1573
+ value_type=value_type,
1574
+ range=range,
1575
+ range_type=range_type,
1576
+ display_name=display_name,
1577
+ enabled=enabled if enabled is not None else True,
1578
+ detection_mechanism=detection_mechanism,
1579
+ sensitivity=sensitivity,
1580
+ exclusion_windows=exclusion_windows,
1581
+ training_data_lookback_days=training_data_lookback_days,
1582
+ incident_behavior=incident_behavior,
1583
+ tags=tags,
1584
+ created_by=updated_by,
1585
+ schedule=schedule,
1586
+ )
1587
+
1588
+ # 2. If urn is set, first validate the input:
1589
+ assertion_input = _SmartColumnMetricAssertionInput(
1590
+ urn=urn,
1591
+ entity_client=self.client.entities,
1592
+ dataset_urn=dataset_urn,
1593
+ column_name=column_name,
1594
+ metric_type=metric_type,
1595
+ operator=operator,
1596
+ value=value,
1597
+ value_type=value_type,
1598
+ range=range,
1599
+ range_type=range_type,
1600
+ display_name=display_name,
1601
+ detection_mechanism=detection_mechanism,
1602
+ sensitivity=sensitivity,
1603
+ exclusion_windows=exclusion_windows,
1604
+ training_data_lookback_days=training_data_lookback_days,
1605
+ incident_behavior=incident_behavior,
1606
+ tags=tags,
1607
+ created_by=updated_by, # This will be overridden by the actual created_by
1608
+ created_at=now_utc, # This will be overridden by the actual created_at
1609
+ updated_by=updated_by,
1610
+ updated_at=now_utc,
1611
+ schedule=schedule,
1612
+ )
1613
+
1614
+ # 3. Merge the assertion input with the existing assertion and monitor entities or create a new assertion
1615
+ # if the assertion does not exist:
1616
+ merged_assertion_input_or_created_assertion = (
1617
+ self._retrieve_and_merge_smart_column_metric_assertion_and_monitor(
1618
+ assertion_input=assertion_input,
1619
+ dataset_urn=dataset_urn,
1620
+ column_name=column_name,
1621
+ metric_type=metric_type,
1622
+ operator=operator,
1623
+ value=value,
1624
+ value_type=value_type,
1625
+ range=range,
1626
+ range_type=range_type,
1627
+ urn=urn,
1628
+ display_name=display_name,
1629
+ enabled=enabled,
1630
+ detection_mechanism=detection_mechanism,
1631
+ sensitivity=sensitivity,
1632
+ exclusion_windows=exclusion_windows,
1633
+ training_data_lookback_days=training_data_lookback_days,
1634
+ incident_behavior=incident_behavior,
1635
+ tags=tags,
1636
+ updated_by=updated_by,
1637
+ now_utc=now_utc,
1638
+ schedule=schedule,
1639
+ )
1640
+ )
1641
+
1642
+ # Return early if we created a new assertion in the merge:
1643
+ if isinstance(merged_assertion_input_or_created_assertion, _AssertionPublic):
1644
+ # We know this is the correct type because we passed the assertion_class parameter
1645
+ assert isinstance(
1646
+ merged_assertion_input_or_created_assertion, SmartColumnMetricAssertion
1647
+ )
1648
+ return merged_assertion_input_or_created_assertion
1649
+
1650
+ # 4. Upsert the assertion and monitor entities:
1651
+ assertion_entity, monitor_entity = (
1652
+ merged_assertion_input_or_created_assertion.to_assertion_and_monitor_entities()
1653
+ )
1654
+ # If assertion upsert fails, we won't try to upsert the monitor
1655
+ self.client.entities.upsert(assertion_entity)
1656
+ # TODO: Wrap monitor upsert in a try-except and delete the assertion if monitor upsert fails (once delete is implemented https://linear.app/acryl-data/issue/OBS-1350/add-delete-method-to-entity-clientpy)
1657
+ # try:
1658
+ self.client.entities.upsert(monitor_entity)
1659
+ # except Exception as e:
1660
+ # logger.error(f"Error upserting monitor: {e}")
1661
+ # self.client.entities.delete(assertion_entity)
1662
+ # raise e
1663
+
1664
+ return SmartColumnMetricAssertion._from_entities(
1665
+ assertion_entity, monitor_entity
1666
+ )
1667
+
1668
+ def _create_smart_column_metric_assertion(
1669
+ self,
1670
+ *,
1671
+ dataset_urn: Union[str, DatasetUrn],
1672
+ column_name: str,
1673
+ metric_type: MetricInputType,
1674
+ operator: OperatorInputType,
1675
+ value: Optional[ValueInputType] = None,
1676
+ value_type: Optional[ValueTypeInputType] = None,
1677
+ range: Optional[RangeInputType] = None,
1678
+ range_type: Optional[RangeTypeInputType] = None,
1679
+ display_name: Optional[str] = None,
1680
+ enabled: bool = True,
1681
+ detection_mechanism: DetectionMechanismInputTypes = None,
1682
+ sensitivity: Optional[Union[str, InferenceSensitivity]] = None,
1683
+ exclusion_windows: Optional[ExclusionWindowInputTypes] = None,
1684
+ training_data_lookback_days: Optional[int] = None,
1685
+ incident_behavior: Optional[
1686
+ Union[AssertionIncidentBehavior, list[AssertionIncidentBehavior]]
1687
+ ] = None,
1688
+ tags: Optional[TagsInputType] = None,
1689
+ created_by: Optional[Union[str, CorpUserUrn]] = None,
1690
+ schedule: Optional[Union[str, models.CronScheduleClass]] = None,
1691
+ ) -> SmartColumnMetricAssertion:
1692
+ """Create a smart column metric assertion.
1693
+
1694
+ Note: keyword arguments are required.
1695
+
1696
+ Args:
1697
+ dataset_urn: The urn of the dataset to be monitored. (Required)
1698
+ column_name: The name of the column to be monitored. (Required)
1699
+ metric_type: The type of the metric to be monitored. (Required)
1700
+ operator: The operator to be used for the assertion. (Required)
1701
+ value: The value to be used for the assertion. (Required if operator requires a value)
1702
+ value_type: The type of the value to be used for the assertion. (Required if operator requires a value)
1703
+ range: The range to be used for the assertion. (Required if operator requires a range)
1704
+ range_type: The type of the range to be used for the assertion. (Required if operator requires a range)
1705
+ display_name: The display name of the assertion. If not provided, a random display
1706
+ name will be generated.
1707
+ enabled: Whether the assertion is enabled. Defaults to True.
1708
+ detection_mechanism: The detection mechanism to be used for the assertion. Information
1709
+ schema is recommended. Valid values are:
1710
+ - "information_schema" or DetectionMechanism.INFORMATION_SCHEMA
1711
+ - "audit_log" or DetectionMechanism.AUDIT_LOG
1712
+ - {
1713
+ "type": "last_modified_column",
1714
+ "column_name": "last_modified",
1715
+ "additional_filter": "last_modified > '2021-01-01'",
1716
+ } or DetectionMechanism.LAST_MODIFIED_COLUMN(column_name='last_modified',
1717
+ additional_filter='last_modified > 2021-01-01')
1718
+ - {
1719
+ "type": "high_watermark_column",
1720
+ "column_name": "id",
1721
+ "additional_filter": "id > 1000",
1722
+ } or DetectionMechanism.HIGH_WATERMARK_COLUMN(column_name='id',
1723
+ additional_filter='id > 1000')
1724
+ - "datahub_operation" or DetectionMechanism.DATAHUB_OPERATION
1725
+ sensitivity: The sensitivity to be applied to the assertion. Valid values are:
1726
+ - "low" or InferenceSensitivity.LOW
1727
+ - "medium" or InferenceSensitivity.MEDIUM
1728
+ - "high" or InferenceSensitivity.HIGH
1729
+ exclusion_windows: The exclusion windows to be applied to the assertion, currently only
1730
+ fixed range exclusion windows are supported. Valid values are:
1731
+ - from datetime.datetime objects: {
1732
+ "start": "datetime(2025, 1, 1, 0, 0, 0)",
1733
+ "end": "datetime(2025, 1, 2, 0, 0, 0)",
1734
+ }
1735
+ - from string datetimes: {
1736
+ "start": "2025-01-01T00:00:00",
1737
+ "end": "2025-01-02T00:00:00",
1738
+ }
1739
+ - from FixedRangeExclusionWindow objects: FixedRangeExclusionWindow(
1740
+ start=datetime(2025, 1, 1, 0, 0, 0),
1741
+ end=datetime(2025, 1, 2, 0, 0, 0)
1742
+ )
1743
+ training_data_lookback_days: The training data lookback days to be applied to the
1744
+ assertion as an integer.
1745
+ incident_behavior: The incident behavior to be applied to the assertion. Valid values are:
1746
+ - "raise_on_fail" or AssertionIncidentBehavior.RAISE_ON_FAIL
1747
+ - "resolve_on_pass" or AssertionIncidentBehavior.RESOLVE_ON_PASS
1748
+ tags: The tags to be applied to the assertion. Valid values are:
1749
+ - a list of strings (strings will be converted to TagUrn objects)
1750
+ - a list of TagUrn objects
1751
+ - a list of TagAssociationClass objects
1752
+ created_by: Optional urn of the user who created the assertion. The format is
1753
+ "urn:li:corpuser:<username>", which you can find on the Users & Groups page.
1754
+ The default is the datahub system user.
1755
+ TODO: Retrieve the SDK user as the default instead of the datahub system user.
1756
+ schedule: Optional cron formatted schedule for the assertion. If not provided, a default
1757
+ schedule will be used. The schedule determines when the assertion will be evaluated.
1758
+ The format is a cron expression, e.g. "0 * * * *" for every hour using UTC timezone.
1759
+ Alternatively, a models.CronScheduleClass object can be provided with string parameters
1760
+ cron and timezone. Use `from datahub.metadata import schema_classes as models` to import the class.
1761
+
1762
+ Returns:
1763
+ SmartVolumeAssertion: The created assertion.
1764
+ """
1765
+ _print_experimental_warning()
1766
+ now_utc = datetime.now(timezone.utc)
1767
+ if created_by is None:
1768
+ logger.warning(
1769
+ f"Created by is not set, using {DEFAULT_CREATED_BY} as a placeholder"
1770
+ )
1771
+ created_by = DEFAULT_CREATED_BY
1772
+ assertion_input = _SmartColumnMetricAssertionInput(
1773
+ urn=None,
1774
+ entity_client=self.client.entities,
1775
+ dataset_urn=dataset_urn,
1776
+ column_name=column_name,
1777
+ metric_type=metric_type,
1778
+ operator=operator,
1779
+ value=value,
1780
+ value_type=value_type,
1781
+ range=range,
1782
+ range_type=range_type,
1783
+ display_name=display_name,
1784
+ enabled=enabled,
1785
+ detection_mechanism=detection_mechanism,
1786
+ sensitivity=sensitivity,
1787
+ exclusion_windows=exclusion_windows,
1788
+ training_data_lookback_days=training_data_lookback_days,
1789
+ incident_behavior=incident_behavior,
1790
+ tags=tags,
1791
+ created_by=created_by,
1792
+ created_at=now_utc,
1793
+ updated_by=created_by,
1794
+ updated_at=now_utc,
1795
+ schedule=schedule,
1796
+ )
1797
+ assertion_entity, monitor_entity = (
1798
+ assertion_input.to_assertion_and_monitor_entities()
1799
+ )
1800
+ # If assertion creation fails, we won't try to create the monitor
1801
+ self.client.entities.create(assertion_entity)
1802
+ # TODO: Wrap monitor creation in a try-except and delete the assertion if monitor creation fails (once delete is implemented https://linear.app/acryl-data/issue/OBS-1350/add-delete-method-to-entity-clientpy)
1803
+ # try:
1804
+ self.client.entities.create(monitor_entity)
1805
+ # except Exception as e:
1806
+ # logger.error(f"Error creating monitor: {e}")
1807
+ # self.client.entities.delete(assertion_entity)
1808
+ # raise e
1809
+ return SmartColumnMetricAssertion._from_entities(
1810
+ assertion_entity, monitor_entity
1811
+ )
1812
+
1813
+ def _retrieve_and_merge_smart_column_metric_assertion_and_monitor(
1814
+ self,
1815
+ assertion_input: _SmartColumnMetricAssertionInput,
1816
+ dataset_urn: Union[str, DatasetUrn],
1817
+ column_name: str,
1818
+ metric_type: MetricInputType,
1819
+ operator: OperatorInputType,
1820
+ value: Optional[ValueInputType],
1821
+ value_type: Optional[ValueTypeInputType],
1822
+ range: Optional[RangeInputType],
1823
+ range_type: Optional[RangeTypeInputType],
1824
+ urn: Union[str, AssertionUrn],
1825
+ display_name: Optional[str],
1826
+ enabled: Optional[bool],
1827
+ detection_mechanism: DetectionMechanismInputTypes,
1828
+ sensitivity: Optional[Union[str, InferenceSensitivity]],
1829
+ exclusion_windows: Optional[ExclusionWindowInputTypes],
1830
+ training_data_lookback_days: Optional[int],
1831
+ incident_behavior: Optional[
1832
+ Union[AssertionIncidentBehavior, list[AssertionIncidentBehavior]]
1833
+ ],
1834
+ tags: Optional[TagsInputType],
1835
+ updated_by: Optional[Union[str, CorpUserUrn]],
1836
+ now_utc: datetime,
1837
+ schedule: Optional[Union[str, models.CronScheduleClass]],
1838
+ ) -> Union[SmartColumnMetricAssertion, _SmartColumnMetricAssertionInput]:
1839
+ # 1. Retrieve any existing assertion and monitor entities:
1840
+ maybe_assertion_entity, monitor_urn, maybe_monitor_entity = (
1841
+ self._retrieve_assertion_and_monitor(assertion_input)
1842
+ )
1843
+
1844
+ # 2.1 If the assertion and monitor entities exist, create an assertion object from them:
1845
+ if maybe_assertion_entity and maybe_monitor_entity:
1846
+ existing_assertion = SmartColumnMetricAssertion._from_entities(
1847
+ maybe_assertion_entity, maybe_monitor_entity
1848
+ )
1849
+ # 2.2 If the assertion exists but the monitor does not, create a placeholder monitor entity to be able to create the assertion:
1850
+ elif maybe_assertion_entity and not maybe_monitor_entity:
1851
+ monitor_mode = (
1852
+ "ACTIVE" if enabled else "INACTIVE" if enabled is not None else "ACTIVE"
1853
+ )
1854
+ existing_assertion = SmartColumnMetricAssertion._from_entities(
1855
+ maybe_assertion_entity,
1856
+ Monitor(id=monitor_urn, info=("ASSERTION", monitor_mode)),
1857
+ )
1858
+ # 2.3 If the assertion does not exist, create a new assertion with a generated urn and return the assertion input:
1859
+ elif not maybe_assertion_entity:
1860
+ logger.info(
1861
+ f"No existing assertion entity found for assertion urn {urn}, creating a new assertion with a generated urn"
1862
+ )
1863
+ return self._create_smart_column_metric_assertion(
1864
+ dataset_urn=dataset_urn,
1865
+ column_name=column_name,
1866
+ metric_type=metric_type,
1867
+ operator=operator,
1868
+ value=value,
1869
+ value_type=value_type,
1870
+ range=range,
1871
+ range_type=range_type,
1872
+ schedule=schedule,
1873
+ display_name=display_name,
1874
+ detection_mechanism=detection_mechanism,
1875
+ sensitivity=sensitivity,
1876
+ exclusion_windows=exclusion_windows,
1877
+ training_data_lookback_days=training_data_lookback_days,
1878
+ incident_behavior=incident_behavior,
1879
+ tags=tags,
1880
+ created_by=updated_by,
1881
+ )
1882
+
1883
+ # 3. Check for any issues e.g. different dataset urns
1884
+ if (
1885
+ existing_assertion
1886
+ and hasattr(existing_assertion, "dataset_urn")
1887
+ and existing_assertion.dataset_urn != assertion_input.dataset_urn
1888
+ ):
1889
+ raise SDKUsageError(
1890
+ f"Dataset URN mismatch, existing assertion: {existing_assertion.dataset_urn} != new assertion: {dataset_urn}"
1891
+ )
1892
+
1893
+ # 4. Merge the existing assertion with the validated input:
1894
+ merged_assertion_input = self._merge_smart_column_metric_input(
1895
+ dataset_urn=dataset_urn,
1896
+ column_name=column_name,
1897
+ metric_type=metric_type,
1898
+ operator=operator,
1899
+ value=value,
1900
+ value_type=value_type,
1901
+ range=range,
1902
+ range_type=range_type,
1903
+ urn=urn,
1904
+ display_name=display_name,
1905
+ enabled=enabled,
1906
+ schedule=schedule,
1907
+ detection_mechanism=detection_mechanism,
1908
+ sensitivity=sensitivity,
1909
+ exclusion_windows=exclusion_windows,
1910
+ training_data_lookback_days=training_data_lookback_days,
1911
+ incident_behavior=incident_behavior,
1912
+ tags=tags,
1913
+ now_utc=now_utc,
1914
+ assertion_input=assertion_input,
1915
+ maybe_assertion_entity=maybe_assertion_entity,
1916
+ maybe_monitor_entity=maybe_monitor_entity,
1917
+ existing_assertion=existing_assertion,
1918
+ )
1919
+
1920
+ return merged_assertion_input
1921
+
1922
+ def _merge_smart_column_metric_input(
1923
+ self,
1924
+ dataset_urn: Union[str, DatasetUrn],
1925
+ column_name: str,
1926
+ metric_type: MetricInputType,
1927
+ operator: OperatorInputType,
1928
+ value: Optional[ValueInputType],
1929
+ value_type: Optional[ValueTypeInputType],
1930
+ range: Optional[RangeInputType],
1931
+ range_type: Optional[RangeTypeInputType],
1932
+ urn: Union[str, AssertionUrn],
1933
+ display_name: Optional[str],
1934
+ enabled: Optional[bool],
1935
+ detection_mechanism: DetectionMechanismInputTypes,
1936
+ sensitivity: Optional[Union[str, InferenceSensitivity]],
1937
+ exclusion_windows: Optional[ExclusionWindowInputTypes],
1938
+ training_data_lookback_days: Optional[int],
1939
+ incident_behavior: Optional[
1940
+ Union[AssertionIncidentBehavior, list[AssertionIncidentBehavior]]
1941
+ ],
1942
+ tags: Optional[TagsInputType],
1943
+ schedule: Optional[Union[str, models.CronScheduleClass]],
1944
+ now_utc: datetime,
1945
+ assertion_input: _SmartColumnMetricAssertionInput,
1946
+ maybe_assertion_entity: Optional[Assertion],
1947
+ maybe_monitor_entity: Optional[Monitor],
1948
+ existing_assertion: SmartColumnMetricAssertion,
1949
+ ) -> _SmartColumnMetricAssertionInput:
1950
+ """Merge the input with the existing assertion and monitor entities.
1951
+
1952
+ Args:
1953
+ dataset_urn: The urn of the dataset to be monitored.
1954
+ column_name: The name of the column to be monitored.
1955
+ metric_type: The type of the metric to be monitored.
1956
+ operator: The operator to be used for the assertion.
1957
+ value: The value to be used for the assertion.
1958
+ value_type: The type of the value to be used for the assertion.
1959
+ range: The range to be used for the assertion.
1960
+ range_type: The type of the range to be used for the assertion.
1961
+ urn: The urn of the assertion.
1962
+ display_name: The display name of the assertion.
1963
+ enabled: Whether the assertion is enabled.
1964
+ detection_mechanism: The detection mechanism to be used for the assertion.
1965
+ sensitivity: The sensitivity to be applied to the assertion.
1966
+ exclusion_windows: The exclusion windows to be applied to the assertion.
1967
+ training_data_lookback_days: The training data lookback days to be applied to the assertion.
1968
+ incident_behavior: The incident behavior to be applied to the assertion.
1969
+ tags: The tags to be applied to the assertion.
1970
+ now_utc: The current UTC time from when the function is called.
1971
+ assertion_input: The validated input to the function.
1972
+ maybe_assertion_entity: The existing assertion entity from the DataHub instance.
1973
+ maybe_monitor_entity: The existing monitor entity from the DataHub instance.
1974
+ existing_assertion: The existing assertion from the DataHub instance.
1975
+
1976
+ Returns:
1977
+ The merged assertion input.
1978
+ """
1979
+ merged_assertion_input = _SmartColumnMetricAssertionInput(
1980
+ urn=urn,
1981
+ entity_client=self.client.entities,
1982
+ dataset_urn=dataset_urn,
1983
+ column_name=_merge_field(
1984
+ input_field_value=column_name,
1985
+ input_field_name="column_name",
1986
+ validated_assertion_input=assertion_input,
1987
+ validated_existing_assertion=existing_assertion,
1988
+ existing_entity_value=SmartColumnMetricAssertion._get_column_name(
1989
+ maybe_assertion_entity
1990
+ )
1991
+ if maybe_assertion_entity
1992
+ else None,
1993
+ ),
1994
+ metric_type=_merge_field(
1995
+ input_field_value=metric_type,
1996
+ input_field_name="metric_type",
1997
+ validated_assertion_input=assertion_input,
1998
+ validated_existing_assertion=existing_assertion,
1999
+ existing_entity_value=SmartColumnMetricAssertion._get_metric_type(
2000
+ maybe_assertion_entity
2001
+ )
2002
+ if maybe_assertion_entity
2003
+ else None,
2004
+ ),
2005
+ operator=_merge_field(
2006
+ input_field_value=operator,
2007
+ input_field_name="operator",
2008
+ validated_assertion_input=assertion_input,
2009
+ validated_existing_assertion=existing_assertion,
2010
+ existing_entity_value=SmartColumnMetricAssertion._get_operator(
2011
+ maybe_assertion_entity
2012
+ )
2013
+ if maybe_assertion_entity
2014
+ else None,
2015
+ ),
2016
+ value=_merge_field(
2017
+ input_field_value=value,
2018
+ input_field_name="value",
2019
+ validated_assertion_input=assertion_input,
2020
+ validated_existing_assertion=existing_assertion,
2021
+ existing_entity_value=SmartColumnMetricAssertion._get_value(
2022
+ maybe_assertion_entity
2023
+ )
2024
+ if maybe_assertion_entity
2025
+ else None,
2026
+ ),
2027
+ value_type=_merge_field(
2028
+ input_field_value=value_type,
2029
+ input_field_name="value_type",
2030
+ validated_assertion_input=assertion_input,
2031
+ validated_existing_assertion=existing_assertion,
2032
+ existing_entity_value=SmartColumnMetricAssertion._get_value_type(
2033
+ maybe_assertion_entity
2034
+ )
2035
+ if maybe_assertion_entity
2036
+ else None,
2037
+ ),
2038
+ range=_merge_field(
2039
+ input_field_value=range,
2040
+ input_field_name="range",
2041
+ validated_assertion_input=assertion_input,
2042
+ validated_existing_assertion=existing_assertion,
2043
+ existing_entity_value=SmartColumnMetricAssertion._get_range(
2044
+ maybe_assertion_entity
2045
+ )
2046
+ if maybe_assertion_entity
2047
+ else None,
2048
+ ),
2049
+ range_type=_merge_field(
2050
+ input_field_value=range_type,
2051
+ input_field_name="range_type",
2052
+ validated_assertion_input=assertion_input,
2053
+ validated_existing_assertion=existing_assertion,
2054
+ existing_entity_value=SmartColumnMetricAssertion._get_range_type(
2055
+ maybe_assertion_entity
2056
+ )
2057
+ if maybe_assertion_entity
2058
+ else None,
2059
+ ),
2060
+ display_name=_merge_field(
2061
+ input_field_value=display_name,
2062
+ input_field_name="display_name",
2063
+ validated_assertion_input=assertion_input,
2064
+ validated_existing_assertion=existing_assertion,
2065
+ existing_entity_value=maybe_assertion_entity.description
2066
+ if maybe_assertion_entity
2067
+ else None,
2068
+ ),
2069
+ enabled=_merge_field(
2070
+ input_field_value=enabled,
2071
+ input_field_name="enabled",
2072
+ validated_assertion_input=assertion_input,
2073
+ validated_existing_assertion=existing_assertion,
2074
+ existing_entity_value=existing_assertion.mode == AssertionMode.ACTIVE
2075
+ if existing_assertion
2076
+ else None,
2077
+ ),
2078
+ schedule=_merge_field(
2079
+ input_field_value=schedule,
2080
+ input_field_name="schedule",
2081
+ validated_assertion_input=assertion_input,
2082
+ validated_existing_assertion=existing_assertion,
2083
+ existing_entity_value=existing_assertion.schedule
2084
+ if existing_assertion
2085
+ else None,
2086
+ ),
2087
+ detection_mechanism=_merge_field(
2088
+ input_field_value=detection_mechanism,
2089
+ input_field_name="detection_mechanism",
2090
+ validated_assertion_input=assertion_input,
2091
+ validated_existing_assertion=existing_assertion,
2092
+ existing_entity_value=SmartColumnMetricAssertion._get_detection_mechanism(
2093
+ maybe_assertion_entity, maybe_monitor_entity, default=None
2094
+ )
2095
+ if maybe_assertion_entity and maybe_monitor_entity
2096
+ else None,
2097
+ ),
2098
+ sensitivity=_merge_field(
2099
+ input_field_value=sensitivity,
2100
+ input_field_name="sensitivity",
2101
+ validated_assertion_input=assertion_input,
2102
+ validated_existing_assertion=existing_assertion,
2103
+ existing_entity_value=maybe_monitor_entity.sensitivity
2104
+ if maybe_monitor_entity
2105
+ else None,
2106
+ ),
2107
+ exclusion_windows=_merge_field(
2108
+ input_field_value=exclusion_windows,
2109
+ input_field_name="exclusion_windows",
2110
+ validated_assertion_input=assertion_input,
2111
+ validated_existing_assertion=existing_assertion,
2112
+ existing_entity_value=maybe_monitor_entity.exclusion_windows
2113
+ if maybe_monitor_entity
2114
+ else None,
2115
+ ),
2116
+ training_data_lookback_days=_merge_field(
2117
+ input_field_value=training_data_lookback_days,
2118
+ input_field_name="training_data_lookback_days",
2119
+ validated_assertion_input=assertion_input,
2120
+ validated_existing_assertion=existing_assertion,
2121
+ existing_entity_value=maybe_monitor_entity.training_data_lookback_days
2122
+ if maybe_monitor_entity
2123
+ else None,
2124
+ ),
2125
+ incident_behavior=_merge_field(
2126
+ input_field_value=incident_behavior,
2127
+ input_field_name="incident_behavior",
2128
+ validated_assertion_input=assertion_input,
2129
+ validated_existing_assertion=existing_assertion,
2130
+ existing_entity_value=SmartColumnMetricAssertion._get_incident_behavior(
2131
+ maybe_assertion_entity
2132
+ )
2133
+ if maybe_assertion_entity
2134
+ else None,
2135
+ ),
2136
+ tags=_merge_field(
2137
+ input_field_value=tags,
2138
+ input_field_name="tags",
2139
+ validated_assertion_input=assertion_input,
2140
+ validated_existing_assertion=existing_assertion,
2141
+ existing_entity_value=maybe_assertion_entity.tags
2142
+ if maybe_assertion_entity
2143
+ else None,
2144
+ ),
2145
+ created_by=existing_assertion.created_by
2146
+ or DEFAULT_CREATED_BY, # Override with the existing assertion's created_by or the default created_by if not set
2147
+ created_at=existing_assertion.created_at
2148
+ or now_utc, # Override with the existing assertion's created_at or now if not set
2149
+ updated_by=assertion_input.updated_by, # Override with the input's updated_by
2150
+ updated_at=assertion_input.updated_at, # Override with the input's updated_at (now)
2151
+ )
2152
+
2153
+ return merged_assertion_input
2154
+
1430
2155
  def sync_freshness_assertion(
1431
2156
  self,
1432
2157
  *,