ethyca-fides 2.68.1b2__py2.py3-none-any.whl → 2.68.1b4__py2.py3-none-any.whl

fides/api/oauth/utils.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import json
-from datetime import datetime
+from datetime import datetime, timedelta
 from functools import update_wrapper
 from types import FunctionType
 from typing import Any, Callable, Dict, List, Optional, Tuple
@@ -32,7 +32,11 @@ from fides.api.models.pre_approval_webhook import PreApprovalWebhook
 from fides.api.models.privacy_request import RequestTask
 from fides.api.oauth.roles import get_scopes_from_roles
 from fides.api.request_context import set_user_id
-from fides.api.schemas.external_https import RequestTaskJWE, WebhookJWE
+from fides.api.schemas.external_https import (
+    DownloadTokenJWE,
+    RequestTaskJWE,
+    WebhookJWE,
+)
 from fides.api.schemas.oauth import OAuth2ClientCredentialsBearer
 from fides.common.api.v1.urn_registry import TOKEN, V1_URL_PREFIX
 from fides.config import CONFIG, FidesConfig
@@ -49,63 +53,28 @@ oauth2_scheme = OAuth2ClientCredentialsBearer(
 
 def extract_payload(jwe_string: str, encryption_key: str) -> str:
     """Given a jwe, extracts the payload and returns it in string form."""
-    return jwe.decrypt(jwe_string, encryption_key)
+    try:
+        decrypted_payload = jwe.decrypt(jwe_string, encryption_key)
+        return decrypted_payload.decode("utf-8")
+    except exceptions.JWEError as e:
+        logger.debug("Failed to decrypt JWE: {}", e)
+        raise e
 
 
-def is_token_expired(issued_at: datetime | None, token_duration_min: int) -> bool:
-    """Returns True if the datetime is earlier than token_duration_min ago."""
-    if not issued_at:
+def is_token_expired(
+    issued_at: Optional[datetime], token_duration_minutes: int
+) -> bool:
+    """Check if a token has expired based on its issued_at timestamp and duration."""
+    if issued_at is None:
         return True
+    expiration_time = issued_at + timedelta(minutes=token_duration_minutes)
+    return datetime.now() > expiration_time
 
-    return (datetime.now() - issued_at).total_seconds() / 60.0 > token_duration_min
-
-
-def copy_func(source_function: Callable) -> Callable:
-    """Based on http://stackoverflow.com/a/6528148/190597 (Glenn Maynard)"""
-    target_function = FunctionType(
-        source_function.__code__,
-        source_function.__globals__,
-        name=source_function.__name__,
-        argdefs=source_function.__defaults__,
-        closure=source_function.__closure__,
-    )
-    updated_target_function: Callable = update_wrapper(target_function, source_function)
-    updated_target_function.__kwdefaults__ = source_function.__kwdefaults__
-    return updated_target_function
-
-
-async def get_current_user(
-    security_scopes: SecurityScopes,
-    authorization: str = Security(oauth2_scheme),
-    db: Session = Depends(get_db),
-) -> FidesUser:
-    """A wrapper around verify_oauth_client that returns that client's user if one exists."""
-    client = await verify_oauth_client(
-        security_scopes=security_scopes,
-        authorization=authorization,
-        db=db,
-    )
-
-    if client.id == CONFIG.security.oauth_root_client_id:
-        return FidesUser(
-            id=CONFIG.security.oauth_root_client_id,
-            username=CONFIG.security.root_username,
-            created_at=datetime.utcnow(),
-        )
-
-    return client.user  # type: ignore[attr-defined]
-
-
-def is_callback_token_expired(issued_at: datetime | None) -> bool:
-    """Returns True if the token is older than the expiration of the redis cache.  We
-    can't resume executing the privacy request if the identity data is gone.
-    """
-    if not issued_at:
-        return True
 
-    return (
-        datetime.now() - issued_at
-    ).total_seconds() / 60.0 > CONFIG.execution.privacy_request_delay_timeout
+def is_callback_token_expired(issued_at: datetime) -> bool:
+    """Check if a callback token has expired (24 hours)."""
+    expiration_time = issued_at + timedelta(hours=24)
+    return datetime.now() > expiration_time
 
 
 def _get_webhook_jwe_or_error(
@@ -114,9 +83,13 @@ def _get_webhook_jwe_or_error(
     if authorization is None:
         raise AuthenticationError(detail="Authentication Failure")
 
-    token_data = json.loads(
-        extract_payload(authorization, CONFIG.security.app_encryption_key)
-    )
+    try:
+        token_data = json.loads(
+            extract_payload(authorization, CONFIG.security.app_encryption_key)
+        )
+    except exceptions.JWEError:
+        raise AuthorizationError(detail="Not Authorized for this action")
+
     try:
         token = WebhookJWE(**token_data)
     except ValidationError:
@@ -160,6 +133,98 @@ def _get_request_task_jwe_or_error(
     return token
 
 
+def validate_download_token(token: str, privacy_request_id: str) -> DownloadTokenJWE:
+    """
+    Validate a download token for accessing privacy request packages.
+
+    Args:
+        token: The JWE token to validate
+        privacy_request_id: The privacy request ID the token should grant access to
+
+    Returns:
+        The validated DownloadTokenJWE object
+
+    Raises:
+        AuthenticationError: If token is invalid or expired
+        AuthorizationError: If token doesn't grant access to the requested privacy request
+    """
+    if not token:
+        raise AuthenticationError(detail="Download token is required")
+
+    # Check if token looks like a JWE (should have 5 parts separated by dots)
+    if token.count(".") != 4:
+        raise AuthenticationError(detail="Invalid download token format")
+
+    try:
+        token_data = json.loads(
+            extract_payload(token, CONFIG.security.app_encryption_key)
+        )
+    except exceptions.JWEError:
+        raise AuthenticationError(detail="Invalid download token format")
+
+    try:
+        download_token = DownloadTokenJWE(**token_data)
+    except ValidationError:
+        raise AuthenticationError(detail="Invalid download token structure")
+
+    # Verify the token grants access to the requested privacy request
+    if download_token.privacy_request_id != privacy_request_id:
+        raise AuthorizationError(
+            detail="Download token does not grant access to this privacy request"
+        )
+
+    # Verify the token has the required scope
+    required_scope = "privacy-request-access-results:read"
+    if required_scope not in download_token.scopes:
+        raise AuthorizationError(detail="Download token lacks required permissions")
+
+    # Check if the token has expired
+    try:
+        expiration_time = datetime.fromisoformat(download_token.exp)
+        if datetime.now() > expiration_time:
+            raise AuthenticationError(detail="Download token has expired")
+    except (ValueError, TypeError):
+        raise AuthenticationError(detail="Invalid token expiration format")
+
+    return download_token
+
+
+def copy_func(source_function: Callable) -> Callable:
+    """Based on http://stackoverflow.com/a/6528148/190597 (Glenn Maynard)"""
+    target_function = FunctionType(
+        source_function.__code__,
+        source_function.__globals__,
+        name=source_function.__name__,
+        argdefs=source_function.__defaults__,
+        closure=source_function.__closure__,
+    )
+    updated_target_function: Callable = update_wrapper(target_function, source_function)
+    updated_target_function.__kwdefaults__ = source_function.__kwdefaults__
+    return updated_target_function
+
+
+async def get_current_user(
+    security_scopes: SecurityScopes,
+    authorization: str = Security(oauth2_scheme),
+    db: Session = Depends(get_db),
+) -> FidesUser:
+    """A wrapper around verify_oauth_client that returns that client's user if one exists."""
+    client = await verify_oauth_client(
+        security_scopes=security_scopes,
+        authorization=authorization,
+        db=db,
+    )
+
+    if client.id == CONFIG.security.oauth_root_client_id:
+        return FidesUser(
+            id=CONFIG.security.oauth_root_client_id,
+            username=CONFIG.security.root_username,
+            created_at=datetime.utcnow(),
+        )
+
+    return client.user  # type: ignore[attr-defined]
+
+
 def verify_callback_oauth_policy_pre_webhook(
     security_scopes: SecurityScopes,
     authorization: str = Security(oauth2_scheme),

fides/api/schemas/application_config.py

@@ -87,6 +87,12 @@ class AdminUIConfig(FidesSchema):
     model_config = ConfigDict(extra="forbid")
 
 
+class PrivacyCenterConfig(FidesSchema):
+    url: SerializeAsAny[Optional[AnyHttpUrlStringRemovesSlash]] = None
+
+    model_config = ConfigDict(extra="forbid")
+
+
 class ConsentConfig(FidesSchema):
     override_vendor_purposes: Optional[bool]
     model_config = ConfigDict(extra="forbid")
@@ -119,6 +125,7 @@ class ApplicationConfig(FidesSchema):
     security: Optional[SecurityApplicationConfig] = None
     consent: Optional[ConsentConfig] = None
     admin_ui: Optional[AdminUIConfig] = None
+    privacy_center: Optional[PrivacyCenterConfig] = None
 
     @model_validator(mode="before")
     @classmethod

fides/api/schemas/connection_configuration/connection_type_system_map.py

@@ -4,6 +4,8 @@ from pydantic import BaseModel, ConfigDict
 
 from fides.api.models.connectionconfig import ConnectionType
 from fides.api.schemas.connection_configuration.enums.system_type import SystemType
+from fides.api.schemas.enums.connection_category import ConnectionCategory
+from fides.api.schemas.enums.integration_feature import IntegrationFeature
 from fides.api.schemas.policy import ActionType
 
 
@@ -19,4 +21,8 @@ class ConnectionSystemTypeMap(BaseModel):
     authorization_required: Optional[bool] = False
     user_guide: Optional[str] = None
     supported_actions: List[ActionType]
+    # New fields for enhanced display information
+    category: Optional[ConnectionCategory] = None
+    tags: Optional[List[str]] = None
+    enabled_features: Optional[List[IntegrationFeature]] = None
     model_config = ConfigDict(use_enum_values=True, from_attributes=True)

fides/api/schemas/enums/connection_category.py

@@ -0,0 +1,20 @@
+from enum import Enum
+
+
+class ConnectionCategory(str, Enum):
+    """
+    Categories for connection types, matching frontend ConnectionCategory enum
+    """
+
+    DATA_CATALOG = "DATA_CATALOG"
+    DATA_WAREHOUSE = "DATA_WAREHOUSE"
+    DATABASE = "DATABASE"
+    IDENTITY_PROVIDER = "IDENTITY_PROVIDER"
+    WEBSITE = "WEBSITE"
+    CRM = "CRM"
+    MANUAL = "MANUAL"
+    MARKETING = "MARKETING"
+    ANALYTICS = "ANALYTICS"
+    ECOMMERCE = "ECOMMERCE"
+    COMMUNICATION = "COMMUNICATION"
+    CUSTOM = "CUSTOM"  # Fallback for uncategorized/custom uploaded integrations

fides/api/schemas/enums/integration_feature.py

@@ -0,0 +1,23 @@
+from enum import Enum
+
+
+class IntegrationFeature(str, Enum):
+    """
+    Features that can be enabled for different integration types.
+    These control which tabs and functionality are available in the integration detail view.
+    """
+
+    # Enables data discovery and monitoring functionality - shows "Data discovery" tab
+    DATA_DISCOVERY = "DATA_DISCOVERY"
+
+    # Enables data synchronization to external systems - shows "Data sync" tab
+    DATA_SYNC = "DATA_SYNC"
+
+    # Enables task/workflow management for manual processes - shows "Tasks" tab
+    TASKS = "TASKS"
+
+    # Indicates integration doesn't require connection testing - shows "Details" tab instead of "Connection" tab
+    WITHOUT_CONNECTION = "WITHOUT_CONNECTION"
+
+    # Enables Data Subject Request automation for SAAS integrations
+    DSR_AUTOMATION = "DSR_AUTOMATION"

fides/api/schemas/external_https.py

@@ -39,3 +39,12 @@ class RequestTaskJWE(BaseModel):
     request_task_id: str
     scopes: List[str]
     iat: str
+
+
+class DownloadTokenJWE(BaseModel):
+    """Describes JWE that is given to users to access their privacy request download package"""
+
+    privacy_request_id: str
+    scopes: List[str]
+    iat: str
+    exp: str

fides/api/schemas/privacy_center_config.py

@@ -1,3 +1,4 @@
+from abc import ABC
 from typing import Any, Dict, List, Literal, Optional, Union
 
 from pydantic import ConfigDict, Field, field_validator, model_validator
@@ -12,22 +13,10 @@ class CustomIdentity(FidesSchema):
     label: str
 
 
-class LocationIdentityField(FidesSchema):
-    """Location field configuration that extends the useful parts of CustomPrivacyRequestField"""
-
-    label: str
-    required: Optional[bool] = True
-    default_value: Optional[str] = None
-    query_param_key: Optional[str] = None
-    ip_geolocation_hint: Optional[bool] = False
-    # Note: We intentionally omit 'hidden' field as it doesn't make sense for location identity input
-
-
 class IdentityInputs(FidesSchema):
     name: Optional[RequiredType] = None
     email: Optional[RequiredType] = None
     phone: Optional[RequiredType] = None
-    location: Optional[Union[RequiredType, LocationIdentityField]] = None
     model_config = ConfigDict(extra="allow")
 
     def __init__(self, **data: Any):
@@ -42,13 +31,12 @@ class IdentityInputs(FidesSchema):
                         f'Custom identity "{field}" must be an instance of CustomIdentity '
                         '(e.g. {"label": "Field label"})'
                     )
-            elif field == "location" and isinstance(value, dict):
-                # Handle location field as LocationIdentityField
-                data[field] = LocationIdentityField(**value)
         super().__init__(**data)
 
 
-class CustomPrivacyRequestField(FidesSchema):
+class BaseCustomPrivacyRequestField(FidesSchema, ABC):
+    """Abstract base class for all custom privacy request fields"""
+
     label: str
     required: Optional[bool] = True
     default_value: Optional[str] = None
@@ -69,6 +57,41 @@ class CustomPrivacyRequestField(FidesSchema):
         return values
 
 
+class CustomPrivacyRequestField(BaseCustomPrivacyRequestField):
+    """Regular custom privacy request field supporting text, select, and multiselect types"""
+
+    field_type: Optional[Literal["text", "select", "multiselect"]] = None
+    options: Optional[List[str]] = None
+
+
+class LocationCustomPrivacyRequestField(BaseCustomPrivacyRequestField):
+    """Location field that doesn't support options and includes IP geolocation hint"""
+
+    field_type: Literal["location"] = "location"
+    ip_geolocation_hint: Optional[bool] = False
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_location_field(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        # Ensure options is not provided for location fields
+        if "options" in values:
+            raise ValueError(
+                "LocationCustomPrivacyRequestField does not support options"
+            )
+
+        # This field cannot be hidden
+        if values.get("hidden"):
+            raise ValueError("Custom location fields cannot be hidden")
+
+        return values
+
+
+# Create a simple union type - Pydantic will use the field_type to determine which model to use
+CustomPrivacyRequestFieldUnion = Union[
+    LocationCustomPrivacyRequestField, CustomPrivacyRequestField
+]
+
+
 class PrivacyRequestOption(FidesSchema):
     locations: Optional[Union[List[PrivacyNoticeRegion], Literal["fallback"]]] = None
     policy_key: Optional[str] = None
@@ -79,7 +102,9 @@ class PrivacyRequestOption(FidesSchema):
     confirm_button_text: Optional[str] = Field(alias="confirmButtonText", default=None)
     cancel_button_text: Optional[str] = Field(alias="cancelButtonText", default=None)
     identity_inputs: Optional[IdentityInputs] = None
-    custom_privacy_request_fields: Optional[Dict[str, CustomPrivacyRequestField]] = None
+    custom_privacy_request_fields: Optional[
+        Dict[str, CustomPrivacyRequestFieldUnion]
+    ] = None
 
 
 class ConsentConfigButton(FidesSchema):
@@ -89,7 +114,9 @@ class ConsentConfigButton(FidesSchema):
     cancel_button_text: Optional[str] = Field(alias="cancelButtonText", default=None)
     icon_path: str
     identity_inputs: IdentityInputs
-    custom_privacy_request_fields: Optional[Dict[str, CustomPrivacyRequestField]] = None
+    custom_privacy_request_fields: Optional[
+        Dict[str, CustomPrivacyRequestFieldUnion]
+    ] = None
     title: str
     modal_title: Optional[str] = Field(alias="modalTitle", default=None)
 
@@ -174,7 +201,9 @@ class PartialPrivacyRequestOption(FidesSchema):
     policy_key: str
     title: str
     identity_inputs: Optional[IdentityInputs] = None
-    custom_privacy_request_fields: Optional[Dict[str, CustomPrivacyRequestField]] = None
+    custom_privacy_request_fields: Optional[
+        Dict[str, CustomPrivacyRequestFieldUnion]
+    ] = None
 
 
 class PartialPrivacyCenterConfig(FidesSchema):

fides/api/schemas/privacy_request.py

@@ -18,6 +18,7 @@ from fides.api.schemas.user import PrivacyRequestReviewer
 from fides.api.util.collection_util import Row
 from fides.api.util.encryption.aes_gcm_encryption_scheme import verify_encryption_key
 from fides.api.util.enums import ColumnSort
+from fides.api.util.text import normalize_location_code
 from fides.config import CONFIG
 
 
@@ -103,6 +104,7 @@ class PrivacyRequestCreate(FidesSchema):
     property_id: Optional[str] = None
     consent_preferences: Optional[List[Consent]] = None  # TODO Slated for deprecation
     source: Optional[PrivacyRequestSource] = None
+    location: Optional[str] = None
 
     @field_validator("encryption_key")
     @classmethod
@@ -114,6 +116,18 @@ class PrivacyRequestCreate(FidesSchema):
             verify_encryption_key(value.encode(CONFIG.security.encoding))
         return value
 
+    @field_validator("location")
+    @classmethod
+    def validate_location(
+        cls: Type["PrivacyRequestCreate"], value: Optional[str] = None
+    ) -> Optional[str]:
+        """Validate and normalize location to ISO 3166 format"""
+        if value is None:
+            return None
+
+        # nuance here is the validator will coalesce values from the less strict format to ISO 3166 (i.e. "us_ca" -> "US-CA")
+        return normalize_location_code(value)
+
 
 class PrivacyRequestResubmit(PrivacyRequestCreate):
     """Schema used to copy a privacy request for resubmission"""
@@ -321,6 +335,7 @@ class PrivacyRequestResponse(FidesSchema):
     custom_privacy_request_fields_approved_by: Optional[str] = None
     custom_privacy_request_fields_approved_at: Optional[datetime] = None
     source: Optional[PrivacyRequestSource] = None
+    location: Optional[str] = None
     deleted_at: Optional[datetime] = None
     deleted_by: Optional[str] = None
     finalized_at: Optional[datetime] = None
@@ -437,6 +452,7 @@ class PrivacyRequestFilter(FidesSchema):
     errored_lt: Optional[datetime] = None
     errored_gt: Optional[datetime] = None
     external_id: Optional[str] = None
+    location: Optional[str] = None
     action_type: Optional[ActionType] = None
     verbose: Optional[bool] = False
     include_identities: Optional[bool] = False

fides/api/schemas/saas/display_info.py

@@ -0,0 +1,19 @@
+from typing import List, Optional
+
+from pydantic import BaseModel, ConfigDict
+
+from fides.api.schemas.enums.connection_category import ConnectionCategory
+from fides.api.schemas.enums.integration_feature import IntegrationFeature
+
+
+class SaaSDisplayInfo(BaseModel):
+    """
+    Optional display information for SAAS integrations to enhance frontend presentation.
+    When not provided, smart defaults will be inferred based on the integration type.
+    """
+
+    category: Optional[ConnectionCategory] = None
+    tags: Optional[List[str]] = None
+    enabled_features: Optional[List[IntegrationFeature]] = None
+
+    model_config = ConfigDict(use_enum_values=True)

fides/api/schemas/saas/saas_config.py

@@ -17,6 +17,7 @@ from fides.api.graph.config import (
 from fides.api.schemas.base_class import FidesSchema
 from fides.api.schemas.limiter.rate_limit_config import RateLimitConfig
 from fides.api.schemas.policy import ActionType
+from fides.api.schemas.saas.display_info import SaaSDisplayInfo
 from fides.api.schemas.saas.shared_schemas import HTTPMethod
 from fides.api.service.saas_request.saas_request_override_factory import (
     SaaSRequestOverrideFactory,
@@ -428,6 +429,7 @@ class SaaSConfig(SaaSConfigBase):
     rate_limit_config: Optional[RateLimitConfig] = None
     consent_requests: Optional[ConsentRequestMap] = None
     user_guide: Optional[str] = None
+    display_info: Optional[SaaSDisplayInfo] = None
 
     @property
     def top_level_endpoint_dict(self) -> Dict[str, Endpoint]:

fides/api/schemas/storage/storage.py

@@ -37,6 +37,8 @@ class StorageDetails(Enum):
     NAMING = "naming"
     MAX_RETRIES = "max_retries"
     AUTH_METHOD = "auth_method"
+    ENABLE_STREAMING = "enable_streaming"
+    ENABLE_ACCESS_PACKAGE_REDIRECT = "enable_access_package_redirect"
     model_config = ConfigDict(extra="forbid")
 
 
@@ -58,6 +60,8 @@ class StorageDetailsS3(FileBasedStorageDetails):
     auth_method: AWSAuthMethod
     bucket: str
     max_retries: Optional[int] = 0
+    enable_streaming: Optional[bool] = False
+    enable_access_package_redirect: Optional[bool] = False
     model_config = ConfigDict(use_enum_values=True)
 
 

fides/api/service/async_dsr/async_dsr_service.py

@@ -0,0 +1,75 @@
+from loguru import logger
+from sqlalchemy.orm import Session
+
+from fides.api.common_exceptions import PrivacyRequestError
+from fides.api.models.connectionconfig import ConnectionConfig
+from fides.api.models.datasetconfig import DatasetConfig
+from fides.api.models.privacy_request import PrivacyRequest, RequestTask
+from fides.api.schemas.privacy_request import PrivacyRequestStatus
+from fides.api.task.task_resources import TaskResources
+
+
+def requeue_polling_request(
+    db: Session,
+    async_task: RequestTask,
+) -> None:
+    """Re-queue a Privacy request that polls async tasks for a given privacy request"""
+    # Check that the privacy request is approved or in processing
+    privacy_request: PrivacyRequest = async_task.privacy_request
+
+    if privacy_request.status not in [
+        PrivacyRequestStatus.approved,
+        PrivacyRequestStatus.in_processing,
+    ]:
+        raise PrivacyRequestError(
+            f"Cannot re-queue privacy request {privacy_request.id} with status {privacy_request.status.value}"
+        )
+
+    logger.info(
+        "Polling starting for {} task {} {}",
+        async_task.action_type,
+        async_task.collection_address,
+        async_task.id,
+    )
+
+    connection_config = get_connection_config_from_task(db, async_task)
+
+    with TaskResources(
+        privacy_request,
+        privacy_request.policy,
+        [connection_config],
+        async_task,
+        db,
+    ) as resources:
+        # graph_task: GraphTask = create_graph_task(db, async_task, resources)
+        # Currently, upstream tasks and "input keys" (which are built by data dependencies)
+        # are the same, but they may not be the same in the future.
+        # upstream_tasks = async_task.upstream_tasks_objects(db)
+        # upstream_access_data: List[List[Row]] = _build_upstream_access_data(
+        #    graph_task.execution_node.input_keys, upstream_tasks
+        # )
+        # TODO: Implement the polling strategy
+        logger.info(f"found resources: {resources}")
+        return None
+
+
+def get_connection_config_from_task(
+    db: Session, request_task: RequestTask
+) -> ConnectionConfig:
+    dataset_config = DatasetConfig.filter(
+        db=db,
+        conditions=(DatasetConfig.fides_key == request_task.dataset_name),
+    ).first()
+    if not dataset_config:
+        raise PrivacyRequestError(
+            f"DatasetConfig with fides_key {request_task.dataset_name} not found."
+        )
+    connection_config = ConnectionConfig.get(
+        db=db, object_id=dataset_config.connection_config_id
+    )
+    if not connection_config:
+        raise PrivacyRequestError(
+            f"ConnectionConfig with id {dataset_config.connection_config_id} not found."
+        )
+
+    return connection_config

fides/api/service/connectors/saas_connector.py

@@ -11,7 +11,7 @@ from sqlalchemy.orm import Session
 from starlette.status import HTTP_204_NO_CONTENT
 
 from fides.api.common_exceptions import (
-    AwaitingAsyncTaskCallback,
+    AwaitingAsyncTask,
     FidesopsException,
     PostProcessingException,
     SkippingConsentPropagation,
@@ -229,7 +229,6 @@ class SaaSConnector(BaseConnector[AuthenticatedClient], Contextualizable):
         query_config: SaaSQueryConfig = self.query_config(node)
 
         # generate initial set of requests if read request is defined, otherwise raise an exception
-
         # An endpoint can be defined with multiple 'read' requests if the data for a single
         # collection can be accessed in multiple ways for example:
         #
@@ -336,8 +335,8 @@ class SaaSConnector(BaseConnector[AuthenticatedClient], Contextualizable):
         if awaiting_async_callback:
             # If a read request was marked to expect async results, original response data here is ignored.
             # We'll instead use the data received in the callback URL later.
-            # Raising an AwaitingAsyncTaskCallback to put this task in an awaiting_processing state
-            raise AwaitingAsyncTaskCallback()
+            # Raising an AwaitingAsyncTask to put this task in an awaiting_processing state
+            raise AwaitingAsyncTask()
 
         return rows
 
@@ -641,8 +640,8 @@ class SaaSConnector(BaseConnector[AuthenticatedClient], Contextualizable):
         if awaiting_async_callback:
             # Asynchronous masking request detected in saas config.
             # If the masking request was marked to expect async results, original responses are ignored
-            # and we raise an AwaitingAsyncTaskCallback to put this task in an awaiting_processing state.
-            raise AwaitingAsyncTaskCallback()
+            # and we raise an AwaitingAsyncTask to put this task in an awaiting_processing state.
+            raise AwaitingAsyncTask()
         return rows_updated
 
     @staticmethod