acryl-datahub-cloud 0.3.12rc7__py3-none-any.whl → 0.3.12rc8__py3-none-any.whl

acryl_datahub_cloud/sdk/assertions_client.py

@@ -39,13 +39,20 @@ from acryl_datahub_cloud.sdk.assertion_input.smart_column_metric_assertion_input
     _SmartColumnMetricAssertionInput,
 )
 from acryl_datahub_cloud.sdk.assertion_input.sql_assertion_input import (
+    SqlAssertionChangeType,
     SqlAssertionCriteria,
+    SqlAssertionOperator,
+    SqlAssertionType,
     _SqlAssertionInput,
 )
 from acryl_datahub_cloud.sdk.assertion_input.volume_assertion_input import (
+    RowCountChange,
     RowCountTotal,
     VolumeAssertionDefinition,
+    VolumeAssertionDefinitionChangeKind,
     VolumeAssertionDefinitionInputTypes,
+    VolumeAssertionDefinitionParameters,
+    VolumeAssertionDefinitionType,
     VolumeAssertionOperator,
     _VolumeAssertionDefinitionTypes,
     _VolumeAssertionInput,
@@ -518,7 +525,7 @@ class AssertionsClient:
         updated_by: Optional[Union[str, CorpUserUrn]],
         now_utc: datetime,
         schedule: Optional[Union[str, models.CronScheduleClass]],
-        definition: Optional[VolumeAssertionDefinitionInputTypes],
+        definition: VolumeAssertionDefinitionInputTypes,
         use_backend_definition: bool = False,
     ) -> Union[VolumeAssertion, _VolumeAssertionInput]:
         # 1. Retrieve any existing assertion and monitor entities:
@@ -549,6 +556,9 @@ class AssertionsClient:
             logger.info(
                 f"No existing assertion entity found for assertion urn {urn}, creating a new assertion with a generated urn"
             )
+            # Extract criteria from definition to call the new signature
+            parsed_definition = VolumeAssertionDefinition.parse(definition)
+            assert isinstance(parsed_definition, (RowCountTotal, RowCountChange))
             return self._create_volume_assertion(
                 dataset_urn=dataset_urn,
                 display_name=display_name,
@@ -557,7 +567,12 @@ class AssertionsClient:
                 tags=tags,
                 created_by=updated_by,
                 schedule=schedule,
-                definition=definition,
+                criteria_type=parsed_definition.type,
+                criteria_change_type=parsed_definition.kind
+                if isinstance(parsed_definition, RowCountChange)
+                else None,
+                criteria_operator=parsed_definition.operator,
+                criteria_parameters=parsed_definition.parameters,
             )
 
         # 3. Check for any issues e.g. different dataset urns
@@ -654,7 +669,10 @@ class AssertionsClient:
             return self._create_sql_assertion(
                 dataset_urn=dataset_urn,
                 display_name=display_name,
-                criteria=criteria,
+                criteria_type=criteria.type,
+                criteria_change_type=criteria.change_type,
+                criteria_operator=criteria.operator,
+                criteria_parameters=criteria.parameters,
                 statement=statement,
                 incident_behavior=incident_behavior,
                 tags=tags,
@@ -1165,7 +1183,7 @@ class AssertionsClient:
                 "criteria",
                 assertion_input,
                 existing_assertion,
-                existing_assertion.criteria if existing_assertion else None,
+                existing_assertion._criteria if existing_assertion else None,
             ),
             statement=_merge_field(
                 statement,
@@ -1683,7 +1701,12 @@ class AssertionsClient:
         tags: Optional[TagsInputType] = None,
         created_by: Optional[Union[str, CorpUserUrn]] = None,
         schedule: Optional[Union[str, models.CronScheduleClass]] = None,
-        definition: Optional[VolumeAssertionDefinitionInputTypes] = None,
+        criteria_type: Union[str, VolumeAssertionDefinitionType],
+        criteria_change_type: Optional[
+            Union[str, VolumeAssertionDefinitionChangeKind]
+        ] = None,
+        criteria_operator: Union[str, VolumeAssertionOperator],
+        criteria_parameters: VolumeAssertionDefinitionParameters,
     ) -> VolumeAssertion:
         """Create a volume assertion.
 
@@ -1723,25 +1746,23 @@ class AssertionsClient:
                 The format is a cron expression, e.g. "0 * * * *" for every hour using UTC timezone.
                 Alternatively, a models.CronScheduleClass object can be provided with string parameters
                 cron and timezone. Use `from datahub.metadata import schema_classes as models` to import the class.
-            definition: The volume assertion definition. Must be provided and include type, operator,
-                and parameters. Can be provided as:
-                - A typed volume assertion object (RowCountTotal or RowCountChange)
-                - A dictionary with keys: type, operator, parameters (and kind for row_count_change)
-
-                Example dictionary for row count total:
-                    {
-                        "type": "row_count_total",
-                        "operator": "GREATER_THAN_OR_EQUAL_TO",
-                        "parameters": 100
-                    }
-
-                Example dictionary for row count change:
-                    {
-                        "type": "row_count_change",
-                        "kind": "percent",
-                        "operator": "BETWEEN",
-                        "parameters": (10, 50)
-                    }
+            criteria_type: The type of volume assertion. Must be either VolumeAssertionDefinitionType.ROW_COUNT_TOTAL or VolumeAssertionDefinitionType.ROW_COUNT_CHANGE.
+                Raw string values are also accepted: "ROW_COUNT_TOTAL" or "ROW_COUNT_CHANGE".
+            criteria_change_type: Required when criteria_type is VolumeAssertionDefinitionType.ROW_COUNT_CHANGE. Must be either VolumeAssertionDefinitionChangeKind.ABSOLUTE
+                or VolumeAssertionDefinitionChangeKind.PERCENT. Optional (ignored) when criteria_type is VolumeAssertionDefinitionType.ROW_COUNT_TOTAL.
+                Raw string values are also accepted: "ABSOLUTE" or "PERCENTAGE".
+            criteria_operator: The comparison operator for the assertion. Must be a VolumeAssertionOperator value:
+                - VolumeAssertionOperator.GREATER_THAN_OR_EQUAL_TO
+                - VolumeAssertionOperator.LESS_THAN_OR_EQUAL_TO
+                - VolumeAssertionOperator.BETWEEN
+                Raw string values are also accepted: "GREATER_THAN_OR_EQUAL_TO", "LESS_THAN_OR_EQUAL_TO", "BETWEEN".
+            criteria_parameters: The parameters for the assertion. For single-value operators
+                (GREATER_THAN_OR_EQUAL_TO, LESS_THAN_OR_EQUAL_TO), provide a single number.
+                For BETWEEN operator, provide a tuple of two numbers (min_value, max_value).
+
+                Examples:
+                - For single value: 100 or 50.5
+                - For BETWEEN: (10, 100) or (5.0, 15.5)
 
         Returns:
             VolumeAssertion: The created assertion.
@@ -1753,6 +1774,17 @@ class AssertionsClient:
                 f"Created by is not set, using {DEFAULT_CREATED_BY} as a placeholder"
             )
             created_by = DEFAULT_CREATED_BY
+
+        # Create definition from individual criteria parameters
+        # The dictionary object will be fully validated down in the _VolumeAssertionInput class
+        definition: dict[str, Any] = {
+            "type": criteria_type,
+            "operator": criteria_operator,
+            "parameters": criteria_parameters,
+        }
+        if criteria_type == VolumeAssertionDefinitionType.ROW_COUNT_CHANGE:
+            definition["kind"] = criteria_change_type
+
         assertion_input = _VolumeAssertionInput(
             urn=None,
             entity_client=self.client.entities,
@@ -1789,7 +1821,12 @@ class AssertionsClient:
         dataset_urn: Union[str, DatasetUrn],
         display_name: Optional[str] = None,
         enabled: bool = True,
-        criteria: SqlAssertionCriteria,
+        criteria_type: Union[SqlAssertionType, str],
+        criteria_change_type: Optional[Union[SqlAssertionChangeType, str]] = None,
+        criteria_operator: Union[SqlAssertionOperator, str],
+        criteria_parameters: Union[
+            Union[float, int], tuple[Union[float, int], Union[float, int]]
+        ],
         statement: str,
         incident_behavior: Optional[
             Union[AssertionIncidentBehavior, list[AssertionIncidentBehavior]]
@@ -1805,25 +1842,23 @@ class AssertionsClient:
             display_name: The display name of the assertion. If not provided, a random display
                 name will be generated.
             enabled: Whether the assertion is enabled. Defaults to True.
-            criteria: The criteria to be used for the assertion. This is of type SqlAssertionCriteria. It has the following fields:
-                - type: The type of sql assertion. Valid values are:
-                    - "METRIC" -> Looks at the current value of the metric.
-                    - "METRIC_CHANGE" -> Looks at the change in the metric between the current and previous run.
-                - change_type: The change type of the assertion, if the type is "METRIC_CHANGE". Valid values are:
-                    - "ABSOLUTE" -> Looks at the absolute change in the metric.
-                    - "PERCENTAGE" -> Looks at the percentage change in the metric.
-                - operator: The operator to be used for the assertion. Valid values are:
-                    - "GREATER_THAN" -> The metric value is greater than the threshold.
-                    - "LESS_THAN" -> The metric value is less than the threshold.
-                    - "GREATER_THAN_OR_EQUAL_TO" -> The metric value is greater than or equal to the threshold.
-                    - "LESS_THAN_OR_EQUAL_TO" -> The metric value is less than or equal to the threshold.
-                    - "EQUAL_TO" -> The metric value is equal to the threshold.
-                    - "NOT_EQUAL_TO" -> The metric value is not equal to the threshold.
-                    - "BETWEEN" -> The metric value is between the two thresholds.
-                - parameters: The parameters to be used for the assertion. This is of type SqlAssertionParameters. It has the following fields:
-                    - value: The value of the metric. This can be a single value or a tuple range.
-                        - If the operator is "BETWEEN", the value is a tuple of two values, with format min, max.
-                        - If the operator is not "BETWEEN", the value is a single value.
+            criteria_type: The type of sql assertion. Valid values are:
+                - "METRIC" -> Looks at the current value of the metric.
+                - "METRIC_CHANGE" -> Looks at the change in the metric between the current and previous run.
+            criteria_change_type: The change type of the assertion, if the type is "METRIC_CHANGE". Valid values are:
+                - "ABSOLUTE" -> Looks at the absolute change in the metric.
+                - "PERCENTAGE" -> Looks at the percentage change in the metric.
+            criteria_operator: The operator to be used for the assertion. Valid values are:
+                - "GREATER_THAN" -> The metric value is greater than the threshold.
+                - "LESS_THAN" -> The metric value is less than the threshold.
+                - "GREATER_THAN_OR_EQUAL_TO" -> The metric value is greater than or equal to the threshold.
+                - "LESS_THAN_OR_EQUAL_TO" -> The metric value is less than or equal to the threshold.
+                - "EQUAL_TO" -> The metric value is equal to the threshold.
+                - "NOT_EQUAL_TO" -> The metric value is not equal to the threshold.
+                - "BETWEEN" -> The metric value is between the two thresholds.
+            criteria_parameters: The parameters to be used for the assertion. This can be a single value or a tuple range.
+                - If the operator is "BETWEEN", the value is a tuple of two values, with format min, max.
+                - If the operator is not "BETWEEN", the value is a single value.
             statement: The statement to be used for the assertion.
             incident_behavior: The incident behavior to be applied to the assertion. Valid values are:
                 - "raise_on_fail" or AssertionIncidentBehavior.RAISE_ON_FAIL
@@ -1850,6 +1885,12 @@ class AssertionsClient:
                 f"Created by is not set, using {DEFAULT_CREATED_BY} as a placeholder"
             )
             created_by = DEFAULT_CREATED_BY
+        criteria = SqlAssertionCriteria(
+            type=criteria_type,
+            change_type=criteria_change_type,
+            operator=criteria_operator,
+            parameters=criteria_parameters,
+        )
         assertion_input = _SqlAssertionInput(
             urn=None,
             entity_client=self.client.entities,
@@ -2959,7 +3000,12 @@ class AssertionsClient:
         tags: Optional[TagsInputType] = None,
         updated_by: Optional[Union[str, CorpUserUrn]] = None,
         schedule: Optional[Union[str, models.CronScheduleClass]] = None,
-        definition: Optional[VolumeAssertionDefinitionInputTypes] = None,
+        criteria_type: Optional[Union[str, VolumeAssertionDefinitionType]] = None,
+        criteria_change_type: Optional[
+            Union[str, VolumeAssertionDefinitionChangeKind]
+        ] = None,
+        criteria_operator: Optional[Union[str, VolumeAssertionOperator]] = None,
+        criteria_parameters: Optional[VolumeAssertionDefinitionParameters] = None,
     ) -> VolumeAssertion:
         """Upsert and merge a volume assertion.
 
@@ -3010,25 +3056,28 @@ class AssertionsClient:
                 The format is a cron expression, e.g. "0 * * * *" for every hour using UTC timezone.
                 Alternatively, a models.CronScheduleClass object can be provided with string parameters
                 cron and timezone. Use `from datahub.metadata import schema_classes as models` to import the class.
-            definition: The volume assertion definition. Can be provided as:
-                - A typed volume assertion object (RowCountTotal or RowCountChange)
-                - A dictionary with keys: type, operator, parameters (and kind for row_count_change)
-                - None to preserve the existing definition from the backend (for update operations)
-
-                Example dictionary for row count total:
-                    {
-                        "type": "row_count_total",
-                        "operator": "GREATER_THAN_OR_EQUAL_TO",
-                        "parameters": 100
-                    }
-
-                Example dictionary for row count change:
-                    {
-                        "type": "row_count_change",
-                        "kind": "absolute",
-                        "operator": "LESS_THAN_OR_EQUAL_TO",
-                        "parameters": 50
-                    }
+            criteria_type: Optional type of volume assertion. Must be either VolumeAssertionDefinitionType.ROW_COUNT_TOTAL or VolumeAssertionDefinitionType.ROW_COUNT_CHANGE.
+                Raw string values are also accepted: "ROW_COUNT_TOTAL" or "ROW_COUNT_CHANGE".
+                If not provided, the existing definition from the backend will be preserved (for update operations).
+                Required when creating a new assertion (when urn is None).
+            criteria_change_type: Optional change type for row count change assertions. Must be either VolumeAssertionDefinitionChangeKind.ABSOLUTE
+                or VolumeAssertionDefinitionChangeKind.PERCENT. Required when criteria_type is VolumeAssertionDefinitionType.ROW_COUNT_CHANGE. Ignored when criteria_type
+                is VolumeAssertionDefinitionType.ROW_COUNT_TOTAL. If not provided, existing value is preserved for updates.
+                Raw string values are also accepted: "ABSOLUTE" or "PERCENTAGE".
+            criteria_operator: Optional comparison operator for the assertion. Must be a VolumeAssertionOperator value:
+                - VolumeAssertionOperator.GREATER_THAN_OR_EQUAL_TO
+                - VolumeAssertionOperator.LESS_THAN_OR_EQUAL_TO
+                - VolumeAssertionOperator.BETWEEN
+                Raw string values are also accepted: "GREATER_THAN_OR_EQUAL_TO", "LESS_THAN_OR_EQUAL_TO", "BETWEEN".
+                If not provided, existing value is preserved for updates. Required when creating a new assertion.
+            criteria_parameters: Optional parameters for the assertion. For single-value operators
+                (GREATER_THAN_OR_EQUAL_TO, LESS_THAN_OR_EQUAL_TO), provide a single number.
+                For BETWEEN operator, provide a tuple of two numbers (min_value, max_value).
+                If not provided, existing value is preserved for updates. Required when creating a new assertion.
+
+                Examples:
+                - For single value: 100 or 50.5
+                - For BETWEEN: (10, 100) or (5.0, 15.5)
 
         Returns:
             VolumeAssertion: The created or updated assertion.
@@ -3042,9 +3091,46 @@ class AssertionsClient:
             )
             updated_by = DEFAULT_CREATED_BY
 
-        # 1. If urn is not set, create a new assertion
+        # 1. Validate criteria parameters if any are provided
+        if (
+            criteria_type is not None
+            or criteria_operator is not None
+            or criteria_parameters is not None
+        ) and (
+            criteria_type is None
+            or criteria_operator is None
+            or criteria_parameters is None
+            or (
+                criteria_type == VolumeAssertionDefinitionType.ROW_COUNT_CHANGE
+                and criteria_change_type is None
+            )
+        ):
+            raise SDKUsageError(
+                "When providing volume assertion criteria, all required parameters must be provided "
+                "(criteria_type, criteria_operator, criteria_parameters must be provided, "
+                "and criteria_change_type is required when criteria_type is 'row_count_change')"
+            )
+
+        # Assert the invariant: if criteria_type is provided, all required parameters are provided
+        assert criteria_type is None or (
+            criteria_operator is not None
+            and criteria_parameters is not None
+            and (
+                criteria_type != VolumeAssertionDefinitionType.ROW_COUNT_CHANGE
+                or criteria_change_type is not None
+            )
+        ), "criteria fields already validated"
+
+        # 2. If urn is not set, create a new assertion
         if urn is None:
+            if criteria_type is None:
+                raise SDKUsageError(
+                    "Volume assertion criteria are required when creating a new assertion"
+                )
             logger.info("URN is not set, creating a new assertion")
+            # Type narrowing: we know these are not None because of validation above
+            assert criteria_operator is not None
+            assert criteria_parameters is not None
             return self._create_volume_assertion(
                 dataset_urn=dataset_urn,
                 display_name=display_name,
@@ -3054,25 +3140,40 @@ class AssertionsClient:
                 tags=tags,
                 created_by=updated_by,
                 schedule=schedule,
-                definition=definition,
+                criteria_type=criteria_type,
+                criteria_change_type=criteria_change_type,
+                criteria_operator=criteria_operator,
+                criteria_parameters=criteria_parameters,
             )
 
         # 2. If urn is set, prepare definition for validation
-        # We use temporary default definition if None is provided, just to pass the _VolumeAssertionInput validation.
+        # If criteria parameters are provided, create definition from them
+        # Otherwise, we use temporary default definition if None is provided, just to pass the _VolumeAssertionInput validation.
         # However, we keep memory of this in use_backend_definition flag, so we can later
         # fail if there is no definition in backend (basically, there is no assertion). That would mean that
         # this is a creation case and the user missed the definition parameter, which is required.
         # Likely this pattern never happened before because there is no a publicly documented default definition
         # that we can use as fallback.
-        use_backend_definition = definition is None
-        temp_definition = (
-            definition
-            if definition is not None
-            else RowCountTotal(
-                operator=VolumeAssertionOperator.GREATER_THAN_OR_EQUAL_TO,
-                parameters=0,  # Temporary placeholder
-            )
-        )
+        if criteria_type is not None:
+            # Create definition from individual criteria parameters
+            temp_definition: dict[str, Any] = {
+                "type": criteria_type,
+                "operator": criteria_operator,
+                "parameters": criteria_parameters,
+            }
+
+            if criteria_type == VolumeAssertionDefinitionType.ROW_COUNT_CHANGE:
+                temp_definition["kind"] = criteria_change_type
+
+            use_backend_definition = False
+        else:
+            # No criteria provided, use backend definition
+            use_backend_definition = True
+            temp_definition = {
+                "type": VolumeAssertionDefinitionType.ROW_COUNT_TOTAL,
+                "operator": VolumeAssertionOperator.GREATER_THAN_OR_EQUAL_TO,
+                "parameters": 0,  # Temporary placeholder
+            }
 
         # 3. Create assertion input with effective definition
         assertion_input = _VolumeAssertionInput(
@@ -3100,7 +3201,7 @@ class AssertionsClient:
                 display_name=display_name,
                 enabled=enabled,
                 detection_mechanism=detection_mechanism,
-                definition=definition,
+                definition=temp_definition,
                 use_backend_definition=use_backend_definition,
                 incident_behavior=incident_behavior,
                 tags=tags,
@@ -3141,7 +3242,12 @@ class AssertionsClient:
         display_name: Optional[str] = None,
         enabled: Optional[bool] = None,
         statement: str,
-        criteria: SqlAssertionCriteria,
+        criteria_type: Union[SqlAssertionType, str],
+        criteria_change_type: Optional[Union[SqlAssertionChangeType, str]] = None,
+        criteria_operator: Union[SqlAssertionOperator, str],
+        criteria_parameters: Union[
+            Union[float, int], tuple[Union[float, int], Union[float, int]]
+        ],
         incident_behavior: Optional[
             Union[AssertionIncidentBehavior, list[AssertionIncidentBehavior]]
         ] = None,
@@ -3171,25 +3277,23 @@ class AssertionsClient:
                 will be generated.
             enabled: Whether the assertion is enabled. If not provided, the existing value
                 will be preserved.
-            criteria: The criteria to be used for the assertion. This is of type SqlAssertionCriteria. It has the following fields:
-                - type: The type of sql assertion. Valid values are:
-                    - "METRIC" -> Looks at the current value of the metric.
-                    - "METRIC_CHANGE" -> Looks at the change in the metric between the current and previous run.
-                - change_type: The change type of the assertion, if the type is "METRIC_CHANGE". Valid values are:
-                    - "ABSOLUTE" -> Looks at the absolute change in the metric.
-                    - "PERCENTAGE" -> Looks at the percentage change in the metric.
-                - operator: The operator to be used for the assertion. Valid values are:
-                    - "GREATER_THAN" -> The metric value is greater than the threshold.
-                    - "LESS_THAN" -> The metric value is less than the threshold.
-                    - "GREATER_THAN_OR_EQUAL_TO" -> The metric value is greater than or equal to the threshold.
-                    - "LESS_THAN_OR_EQUAL_TO" -> The metric value is less than or equal to the threshold.
-                    - "EQUAL_TO" -> The metric value is equal to the threshold.
-                    - "NOT_EQUAL_TO" -> The metric value is not equal to the threshold.
-                    - "BETWEEN" -> The metric value is between the two thresholds.
-                - parameters: The parameters to be used for the assertion. This is of type SqlAssertionParameters. It has the following fields:
-                    - value: The value of the metric. This can be a single value or a tuple range.
-                        - If the operator is "BETWEEN", the value is a tuple of two values, with format min, max.
-                        - If the operator is not "BETWEEN", the value is a single value.
+            criteria_type: The type of sql assertion. Valid values are:
+                - "METRIC" -> Looks at the current value of the metric.
+                - "METRIC_CHANGE" -> Looks at the change in the metric between the current and previous run.
+            criteria_change_type: The change type of the assertion, if the type is "METRIC_CHANGE". Valid values are:
+                - "ABSOLUTE" -> Looks at the absolute change in the metric.
+                - "PERCENTAGE" -> Looks at the percentage change in the metric.
+            criteria_operator: The operator to be used for the assertion. Valid values are:
+                - "GREATER_THAN" -> The metric value is greater than the threshold.
+                - "LESS_THAN" -> The metric value is less than the threshold.
+                - "GREATER_THAN_OR_EQUAL_TO" -> The metric value is greater than or equal to the threshold.
+                - "LESS_THAN_OR_EQUAL_TO" -> The metric value is less than or equal to the threshold.
+                - "EQUAL_TO" -> The metric value is equal to the threshold.
+                - "NOT_EQUAL_TO" -> The metric value is not equal to the threshold.
+                - "BETWEEN" -> The metric value is between the two thresholds.
+            criteria_parameters: The parameters to be used for the assertion. This can be a single value or a tuple range.
+                - If the operator is "BETWEEN", the value is a tuple of two values, with format min, max.
+                - If the operator is not "BETWEEN", the value is a single value.
             statement: The SQL statement to be used for the assertion.
                 - "SELECT COUNT(*) FROM table WHERE column > 100"
             incident_behavior: The incident behavior to be applied to the assertion. Valid values are:
@@ -3228,7 +3332,10 @@ class AssertionsClient:
                 dataset_urn=dataset_urn,
                 display_name=display_name,
                 enabled=enabled if enabled is not None else True,
-                criteria=criteria,
+                criteria_type=criteria_type,
+                criteria_change_type=criteria_change_type,
+                criteria_operator=criteria_operator,
+                criteria_parameters=criteria_parameters,
                 statement=statement,
                 incident_behavior=incident_behavior,
                 tags=tags,
@@ -3237,6 +3344,12 @@ class AssertionsClient:
             )
 
         # 2. If urn is set, first validate the input:
+        criteria = SqlAssertionCriteria(
+            type=criteria_type,
+            change_type=criteria_change_type,
+            operator=criteria_operator,
+            parameters=criteria_parameters,
+        )
         assertion_input = _SqlAssertionInput(
             urn=urn,
             entity_client=self.client.entities,

acryl_datahub_cloud/sdk/subscription_client.py

@@ -16,7 +16,7 @@ from datahub.sdk.main_client import DataHubClient
 
 logger = logging.getLogger(__name__)
 
-SubscriberInputType: TypeAlias = Union[CorpUserUrn, CorpGroupUrn]
+SubscriberInputType: TypeAlias = Union[CorpUserUrn, CorpGroupUrn, str]
 
 
 ASSERTION_RELATED_ENTITY_CHANGE_TYPES = {
@@ -73,6 +73,7 @@ class SubscriptionClient:
                  For datasets: subscription applies to all assertions on the dataset.
                  For assertions: subscription applies only to that specific assertion.
             subscriber_urn: The URN of the user or group that will receive notifications.
+                           Can be a string (valid corpuser or corpGroup URN) or URN object.
             entity_change_types: Specific change types to subscribe to. If None, defaults are:
                                 - Dataset: all existing change types
                                 - Assertion: assertion-related types (ASSERTION_PASSED,
@@ -91,6 +92,9 @@ class SubscriptionClient:
         # Parse URN string if needed
         parsed_urn = self._maybe_parse_urn(urn)
 
+        # Parse subscriber URN string if needed
+        parsed_subscriber_urn = self._maybe_parse_subscriber_urn(subscriber_urn)
+
         dataset_urn: DatasetUrn
         assertion_urn: Optional[AssertionUrn]
         dataset_urn, assertion_urn = (
@@ -100,7 +104,7 @@ class SubscriptionClient:
         )
 
         logger.info(
-            f"Subscribing to dataset={dataset_urn} assertion={assertion_urn} for subscriber={subscriber_urn} with change types: {entity_change_types}"
+            f"Subscribing to dataset={dataset_urn} assertion={assertion_urn} for subscriber={parsed_subscriber_urn} with change types: {entity_change_types}"
         )
 
         # Get entity change types (use all if none provided)
@@ -111,16 +115,16 @@ class SubscriptionClient:
 
         existing_subscriptions = self.client.resolve.subscription(  # type: ignore[attr-defined]
             entity_urn=dataset_urn.urn(),
-            actor_urn=subscriber_urn.urn(),
+            actor_urn=parsed_subscriber_urn.urn(),
         )
         if not existing_subscriptions:
             # new subscription
             subscription = Subscription(
                 info=models.SubscriptionInfoClass(
                     entityUrn=dataset_urn.urn(),
-                    actorUrn=subscriber_urn.urn(),
+                    actorUrn=parsed_subscriber_urn.urn(),
                     actorType=CorpUserUrn.ENTITY_TYPE
-                    if isinstance(subscriber_urn, CorpUserUrn)
+                    if isinstance(parsed_subscriber_urn, CorpUserUrn)
                     else CorpGroupUrn.ENTITY_TYPE,
                     types=[
                         models.SubscriptionTypeClass.ENTITY_CHANGE,
@@ -158,7 +162,7 @@ class SubscriptionClient:
             return
         else:
             raise SdkUsageError(
-                f"We have a mesh here - {len(existing_subscriptions)} subscriptions found for dataset={dataset_urn} assertion={assertion_urn} and subscriber={subscriber_urn}!"
+                f"We have a mesh here - {len(existing_subscriptions)} subscriptions found for dataset={dataset_urn} assertion={assertion_urn} and subscriber={parsed_subscriber_urn}!"
             )
 
     def list_subscriptions(
@@ -178,7 +182,8 @@ class SubscriptionClient:
             entity_change_types: Optional filter to return only subscriptions for specific
                                 change types. If None, returns subscriptions for all change types.
             subscriber_urn: Optional filter to return only subscriptions for a specific user
-                           or group. If None, returns subscriptions for all subscribers.
+                           or group. Can be a string (valid corpuser or corpGroup URN) or URN object.
+                           If None, returns subscriptions for all subscribers.
 
         Returns:
             List[Subscription]: List of matching subscription objects.
@@ -227,6 +232,7 @@ class SubscriptionClient:
         Args:
             urn: URN (string or URN object) of the dataset or assertion to unsubscribe from.
             subscriber_urn: User or group URN to unsubscribe.
+                           Can be a string (valid corpuser or corpGroup URN) or URN object.
             entity_change_types: Specific change types to remove. If None, defaults are:
                                 - Dataset: all existing change types in the subscription
                                 - Assertion: assertion-related types (ASSERTION_PASSED,
@@ -248,6 +254,9 @@ class SubscriptionClient:
         # Parse URN string if needed
         parsed_urn = self._maybe_parse_urn(urn)
 
+        # Parse subscriber URN string if needed
+        parsed_subscriber_urn = self._maybe_parse_subscriber_urn(subscriber_urn)
+
         dataset_urn: DatasetUrn
         assertion_urn: Optional[AssertionUrn]
         dataset_urn, assertion_urn = (
@@ -257,23 +266,23 @@ class SubscriptionClient:
         )
 
         logger.info(
-            f"Unsubscribing from dataset={dataset_urn}{f' assertion={assertion_urn}' if assertion_urn else ''} for subscriber={subscriber_urn} with change types: {entity_change_types}"
+            f"Unsubscribing from dataset={dataset_urn}{f' assertion={assertion_urn}' if assertion_urn else ''} for subscriber={parsed_subscriber_urn} with change types: {entity_change_types}"
         )
 
         # Find existing subscription
         existing_subscription_urns = self.client.resolve.subscription(  # type: ignore[attr-defined]
             entity_urn=dataset_urn.urn(),
-            actor_urn=subscriber_urn.urn(),
+            actor_urn=parsed_subscriber_urn.urn(),
         )
 
         if not existing_subscription_urns:
             logger.info(
-                f"No subscription found for dataset={dataset_urn} and subscriber={subscriber_urn}"
+                f"No subscription found for dataset={dataset_urn} and subscriber={parsed_subscriber_urn}"
             )
             return
         elif len(existing_subscription_urns) > 1:
             raise SdkUsageError(
-                f"Multiple subscriptions found for dataset={dataset_urn} and subscriber={subscriber_urn}. "
+                f"Multiple subscriptions found for dataset={dataset_urn} and subscriber={parsed_subscriber_urn}. "
                 f"Expected at most 1, got {len(existing_subscription_urns)}"
             )
 
@@ -659,6 +668,33 @@ class SubscriptionClient:
                 f"Unsupported URN type. Only dataset and assertion URNs are supported, got: {urn}"
             )
 
+    def _maybe_parse_subscriber_urn(
+        self, subscriber_urn: SubscriberInputType
+    ) -> Union[CorpUserUrn, CorpGroupUrn]:
+        """Parse subscriber URN string into appropriate URN object if needed.
+
+        Args:
+            subscriber_urn: String URN or URN object (CorpUserUrn or CorpGroupUrn)
+
+        Returns:
+            Parsed URN object (CorpUserUrn or CorpGroupUrn)
+
+        Raises:
+            SdkUsageError: If the URN string format is invalid or unsupported
+        """
+        if isinstance(subscriber_urn, (CorpUserUrn, CorpGroupUrn)):
+            return subscriber_urn
+
+        # Try to determine URN type from string format
+        if ":corpuser:" in subscriber_urn:
+            return CorpUserUrn.from_string(subscriber_urn)
+        elif ":corpGroup:" in subscriber_urn:
+            return CorpGroupUrn.from_string(subscriber_urn)
+        else:
+            raise SdkUsageError(
+                f"Unsupported subscriber URN type. Only corpuser and corpGroup URNs are supported, got: {subscriber_urn}"
+            )
+
     def _fetch_dataset_from_assertion(
         self, assertion_urn: AssertionUrn
     ) -> Tuple[DatasetUrn, AssertionUrn]: