ethyca-fides 2.67.0rc1__py2.py3-none-any.whl → 2.67.1b0__py2.py3-none-any.whl

fides/api/service/connectors/query_configs/bigquery_query_config.py

@@ -93,7 +93,7 @@ class BigQueryQueryConfig(QueryStringWithoutTuplesOverrideQueryConfig):
 
         return where_clauses
 
-    def _generate_table_name(self) -> str:
+    def generate_table_name(self) -> str:
         """
         Prepends the dataset ID and project ID to the base table name
         if the BigQuery namespace meta is provided.
@@ -116,7 +116,7 @@ class BigQueryQueryConfig(QueryStringWithoutTuplesOverrideQueryConfig):
         Returns a query string with backtick formatting for tables that have the same names as
         BigQuery reserved words.
         """
-        return f'SELECT {field_list} FROM `{self._generate_table_name()}` WHERE ({" OR ".join(clauses)})'
+        return f'SELECT {field_list} FROM `{self.generate_table_name()}` WHERE ({" OR ".join(clauses)})'
 
     def generate_masking_stmt(
         self,
@@ -197,7 +197,7 @@ class BigQueryQueryConfig(QueryStringWithoutTuplesOverrideQueryConfig):
             )
             return []
 
-        table = Table(self._generate_table_name(), MetaData(bind=client), autoload=True)
+        table = Table(self.generate_table_name(), MetaData(bind=client), autoload=True)
         where_clauses: List[ColumnElement] = [
             table.c[k] == v for k, v in non_empty_reference_field_keys.items()
         ]
@@ -256,7 +256,7 @@ class BigQueryQueryConfig(QueryStringWithoutTuplesOverrideQueryConfig):
             )
             return []
 
-        table = Table(self._generate_table_name(), MetaData(bind=client), autoload=True)
+        table = Table(self.generate_table_name(), MetaData(bind=client), autoload=True)
 
         # Build individual reference clauses
         where_clauses: List[ColumnElement] = []

fides/api/service/connectors/query_configs/snowflake_query_config.py

@@ -30,7 +30,7 @@ class SnowflakeQueryConfig(SQLQueryConfig):
         """Returns field names in clauses surrounded by quotation marks as required by Snowflake syntax."""
         return f'"{string_path}" {operator} (:{operand})'
 
-    def _generate_table_name(self) -> str:
+    def generate_table_name(self) -> str:
         """
         Prepends the dataset name and schema to the base table name
         if the Snowflake namespace meta is provided.
@@ -57,7 +57,7 @@ class SnowflakeQueryConfig(SQLQueryConfig):
         clauses: List[str],
     ) -> str:
         """Returns a query string with double quotation mark formatting as required by Snowflake syntax."""
-        return f'SELECT {field_list} FROM {self._generate_table_name()} WHERE ({" OR ".join(clauses)})'
+        return f'SELECT {field_list} FROM {self.generate_table_name()} WHERE ({" OR ".join(clauses)})'
 
     def format_key_map_for_update_stmt(self, param_map: Dict[str, Any]) -> List[str]:
         """Adds the appropriate formatting for update statements in this datastore."""
@@ -69,4 +69,4 @@ class SnowflakeQueryConfig(SQLQueryConfig):
         where_clauses: List[str],
     ) -> str:
         """Returns a parameterized update statement in Snowflake dialect."""
-        return f'UPDATE {self._generate_table_name()} SET {", ".join(update_clauses)} WHERE {" AND ".join(where_clauses)}'
+        return f'UPDATE {self.generate_table_name()} SET {", ".join(update_clauses)} WHERE {" AND ".join(where_clauses)}'

fides/api/service/connectors/snowflake_connector.py

@@ -3,11 +3,11 @@ from typing import Any, Dict, Union
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from snowflake.sqlalchemy import URL as Snowflake_URL
+from sqlalchemy import text
 from sqlalchemy.orm import Session
 
 from fides.api.graph.execution import ExecutionNode
 from fides.api.schemas.connection_configuration import SnowflakeSchema
-from fides.api.service.connectors.query_configs.query_config import SQLQueryConfig
 from fides.api.service.connectors.query_configs.snowflake_query_config import (
     SnowflakeQueryConfig,
 )
@@ -69,10 +69,63 @@ class SnowflakeConnector(SQLConnector):
                 connect_args["private_key"] = private_key
         return connect_args
 
-    def query_config(self, node: ExecutionNode) -> SQLQueryConfig:
+    def query_config(self, node: ExecutionNode) -> SnowflakeQueryConfig:
         """Query wrapper corresponding to the input execution_node."""
 
         db: Session = Session.object_session(self.configuration)
         return SnowflakeQueryConfig(
             node, SQLConnector.get_namespace_meta(db, node.address.dataset)
         )
+
+    def get_qualified_table_name(self, node: ExecutionNode) -> str:
+        """Get fully qualified Snowflake table name using existing query config logic"""
+        query_config = self.query_config(node)
+        return query_config.generate_table_name()
+
+    def table_exists(self, qualified_table_name: str) -> bool:
+        """
+        Check if table exists in Snowflake using the proper three-part naming convention.
+
+        Snowflake supports database.schema.table naming, and the generic SQLConnector
+        table_exists method doesn't handle quoted identifiers properly.
+        """
+        try:
+            client = self.create_client()
+            with client.connect() as connection:
+                # Remove quotes and split the parts
+                clean_name = qualified_table_name.replace('"', "")
+                parts = clean_name.split(".")
+
+                if len(parts) == 1:
+                    # Simple table name - use current schema
+                    table_name = parts[0]
+                    result = connection.execute(text(f'DESC TABLE "{table_name}"'))
+                elif len(parts) == 2:
+                    # schema.table format
+                    schema_name, table_name = parts
+                    result = connection.execute(
+                        text(f'DESC TABLE "{schema_name}"."{table_name}"')
+                    )
+                elif len(parts) >= 3:
+                    # database.schema.table format
+                    database_name, schema_name, table_name = (
+                        parts[-3],
+                        parts[-2],
+                        parts[-1],
+                    )
+                    # Use the database.schema.table format
+                    result = connection.execute(
+                        text(
+                            f'DESC TABLE "{database_name}"."{schema_name}"."{table_name}"'
+                        )
+                    )
+                else:
+                    return False
+
+                # If we get here without an exception, the table exists
+                result.close()
+                return True
+
+        except Exception:
+            # Table doesn't exist or other error
+            return False

fides/api/service/connectors/sql_connector.py

@@ -6,7 +6,7 @@ import paramiko
 import sshtunnel  # type: ignore
 from aiohttp.client_exceptions import ClientResponseError
 from loguru import logger
-from sqlalchemy import Column, select
+from sqlalchemy import Column, inspect, select
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.engine import (  # type: ignore
     Connection,
@@ -22,6 +22,7 @@ from sqlalchemy.sql.elements import TextClause
 from fides.api.common_exceptions import (
     ConnectionException,
     SSHTunnelConfigNotFoundException,
+    TableNotFound,
 )
 from fides.api.graph.execution import ExecutionNode
 from fides.api.models.connectionconfig import ConnectionConfig, ConnectionTestStatus
@@ -189,14 +190,28 @@ class SQLConnector(BaseConnector[Engine]):
 
         logger.info("Starting data retrieval for {}", node.address)
         with client.connect() as connection:
-            self.set_schema(connection)
-            if (
-                query_config.partitioning
-            ):  # only BigQuery supports partitioning, for now
-                return self.partitioned_retrieval(query_config, connection, stmt)
-
-            results = connection.execute(stmt)
-            return self.cursor_result_to_rows(results)
+            try:
+                self.set_schema(connection)
+                if (
+                    query_config.partitioning
+                ):  # only BigQuery supports partitioning, for now
+                    return self.partitioned_retrieval(query_config, connection, stmt)
+
+                results = connection.execute(stmt)
+                return self.cursor_result_to_rows(results)
+            except Exception as exc:
+                # Check if table exists using qualified table name
+                qualified_table_name = self.get_qualified_table_name(node)
+                if not self.table_exists(qualified_table_name):
+                    # Central decision point - will raise TableNotFound or ConnectionException
+                    self.handle_table_not_found(
+                        node=node,
+                        table_name=qualified_table_name,
+                        operation_context="data retrieval",
+                        original_exception=exc,
+                    )
+                # Table exists or can't check - re-raise original exception
+                raise
 
     def mask_data(
         self,
@@ -290,3 +305,86 @@ class SQLConnector(BaseConnector[Engine]):
         raise NotImplementedError(
             "Partitioned retrieval is only supported for BigQuery currently!"
         )
+
+    def get_qualified_table_name(self, node: ExecutionNode) -> str:
+        """
+        Get the fully qualified table name for this database.
+
+        Default: Returns the simple collection name
+        Override: Database-specific connectors can implement namespace resolution
+        """
+        return node.collection.name
+
+    def table_exists(self, qualified_table_name: str) -> bool:
+        """
+        Check if table exists using SQLAlchemy introspection.
+
+        This is a generic implementation that should work for most SQL databases.
+        Override: Connectors can implement database-specific table existence checking
+        """
+        try:
+            client = self.create_client()
+            with client.connect() as connection:
+                inspector = inspect(connection)
+
+                # For simple table names
+                if "." not in qualified_table_name:
+                    return inspector.has_table(qualified_table_name)
+
+                # For qualified names like schema.table or database.schema.table
+                parts = qualified_table_name.split(".")
+
+                if len(parts) == 2:
+                    # schema.table format
+                    schema_name, table_name = parts
+                    return inspector.has_table(table_name, schema=schema_name)
+
+                if len(parts) >= 3:
+                    # database.schema.table format (use schema.table)
+                    schema_name, table_name = parts[-2], parts[-1]
+                    return inspector.has_table(table_name, schema=schema_name)
+
+                # Fallback for unexpected format
+                return inspector.has_table(qualified_table_name)
+
+        except Exception as exc:
+            # Graceful fallback - if we can't check, assume table exists
+            # to preserve existing behavior for connectors that don't implement this
+            logger.error("Unable to check if table exists, assuming it does: {}", exc)
+            return True
+
+    def handle_table_not_found(
+        self,
+        node: ExecutionNode,
+        table_name: str,
+        operation_context: str,
+        original_exception: Optional[Exception] = None,
+    ) -> None:
+        """
+        Central decision point for table-not-found scenarios.
+
+        Raises TableNotFound (for collection skipping) or ConnectionException (for hard errors).
+        The raised exception will be caught by the @retry decorator in graph_task.py.
+
+        Args:
+            node: The ExecutionNode being processed
+            table_name: Name of the missing table
+            operation_context: Context like "data retrieval" or "data masking"
+            original_exception: The original exception that triggered this check
+        """
+        if node.has_outgoing_dependencies():
+            # Collection has dependencies - cannot skip safely
+            error_msg = (
+                f"Table '{table_name}' did not exist during {operation_context}. "
+                f"Cannot skip collection '{node.address}' because other collections depend on it."
+            )
+            if original_exception:
+                raise ConnectionException(error_msg) from original_exception
+            raise ConnectionException(error_msg)
+
+        # Safe to skip - raise TableNotFound for @retry decorator to catch
+        skip_msg = f"Table '{table_name}' did not exist during {operation_context}."
+        if original_exception:
+            raise TableNotFound(skip_msg) from original_exception
+
+        raise TableNotFound(skip_msg)

fides/api/service/privacy_request/request_runner_service.py

@@ -80,6 +80,7 @@ from fides.api.util.cache import get_all_masking_secret_keys
 from fides.api.util.collection_util import Row
 from fides.api.util.logger import Pii, _log_exception, _log_warning
 from fides.api.util.logger_context_utils import LoggerContextKeys, log_context
+from fides.api.util.memory_watchdog import memory_limiter
 from fides.common.api.v1.urn_registry import (
     PRIVACY_REQUEST_TRANSFER_TO_PARENT,
     V1_URL_PREFIX,
@@ -358,8 +359,8 @@ def upload_and_save_access_results(  # pylint: disable=R0912
 
 
 @celery_app.task(base=DatabaseTask, bind=True)
-# TODO: Add log_context back in, this is just for some temporary testing
-# @log_context(capture_args={"privacy_request_id": LoggerContextKeys.privacy_request_id})
+@memory_limiter
+@log_context(capture_args={"privacy_request_id": LoggerContextKeys.privacy_request_id})
 def run_privacy_request(
     self: DatabaseTask,
     privacy_request_id: str,

fides/api/service/privacy_request/request_service.py

@@ -3,11 +3,12 @@ from __future__ import annotations
 import json
 from asyncio import sleep
 from datetime import datetime, timedelta
-from typing import Any, Dict, Optional, Set
+from typing import Any, Dict, List, Optional, Set
 
 from httpx import AsyncClient
 from loguru import logger
 from sqlalchemy import text
+from sqlalchemy.orm import Session
 from sqlalchemy.sql.elements import TextClause
 
 from fides.api.common_exceptions import PrivacyRequestNotFound
@@ -31,6 +32,9 @@ from fides.api.util.cache import (
     celery_tasks_in_flight,
     get_async_task_tracking_cache_key,
     get_cache,
+    get_privacy_request_retry_count,
+    increment_privacy_request_retry_count,
+    reset_privacy_request_retry_count,
 )
 from fides.api.util.lock import redis_lock
 from fides.common.api.v1.urn_registry import PRIVACY_REQUESTS, V1_URL_PREFIX
@@ -350,10 +354,17 @@ def initiate_interrupted_task_requeue_poll() -> None:
 
 
 def get_cached_task_id(entity_id: str) -> Optional[str]:
-    """Gets the cached task ID for a privacy request or request task by ID."""
+    """Gets the cached task ID for a privacy request or request task by ID.
+
+    Raises Exception if cache operations fail, allowing callers to handle cache failures appropriately.
+    """
     cache: FidesopsRedis = get_cache()
-    task_id = cache.get(get_async_task_tracking_cache_key(entity_id))
-    return task_id
+    try:
+        task_id = cache.get(get_async_task_tracking_cache_key(entity_id))
+        return task_id
+    except Exception as exc:
+        logger.error(f"Failed to get cached task ID for entity {entity_id}: {exc}")
+        raise
 
 
 REQUEUE_INTERRUPTED_TASKS_LOCK = "requeue_interrupted_tasks_lock"
@@ -393,6 +404,115 @@ def _get_task_ids_from_dsr_queue(
     return queued_tasks_ids
 
 
+def _cancel_interrupted_tasks_and_error_privacy_request(
+    db: Session, privacy_request: PrivacyRequest, error_message: Optional[str] = None
+) -> None:
+    """
+    Cancel all tasks associated with an interrupted privacy request and set the privacy request to error state.
+
+    This function:
+    1. Logs the error message (either provided or default)
+    2. Revokes the main privacy request task and all associated request tasks
+    3. Sets the privacy request status to error
+    4. Creates an error log entry
+
+    Args:
+        db: Database session
+        privacy_request: The privacy request to cancel and error
+        error_message: Optional error message to log. If not provided, uses default message.
+    """
+    if error_message:
+        logger.error(error_message)
+    else:
+        logger.error(
+            f"Canceling interrupted tasks and marking privacy request {privacy_request.id} as error"
+        )
+
+    # Cancel all associated Celery tasks
+    privacy_request.cancel_celery_tasks()
+
+    # Set privacy request to error state using the existing method
+    try:
+        privacy_request.error_processing(db)
+        logger.info(
+            f"Privacy request {privacy_request.id} marked as error due to task interruption"
+        )
+    except Exception as exc:
+        logger.error(
+            f"Failed to mark privacy request {privacy_request.id} as error: {exc}"
+        )
+
+
+def _handle_privacy_request_requeue(
+    db: Session, privacy_request: PrivacyRequest
+) -> None:
+    """Handle retry logic for a privacy request - either requeue or cancel based on retry count."""
+    try:
+        # Check retry count and either requeue or cancel based on limit
+        current_retry_count = get_privacy_request_retry_count(privacy_request.id)
+        max_retries = CONFIG.execution.privacy_request_requeue_retry_count
+
+        if current_retry_count < max_retries:
+            # Increment retry count and attempt requeue
+            new_retry_count = increment_privacy_request_retry_count(privacy_request.id)
+            logger.info(
+                f"Requeuing privacy request {privacy_request.id} "
+                f"(attempt {new_retry_count}/{max_retries})"
+            )
+
+            from fides.service.privacy_request.privacy_request_service import (  # pylint: disable=cyclic-import
+                PrivacyRequestError,
+                _requeue_privacy_request,
+            )
+
+            try:
+                _requeue_privacy_request(db, privacy_request)
+            except PrivacyRequestError as exc:
+                # If requeue fails, cancel tasks and set to error state
+                _cancel_interrupted_tasks_and_error_privacy_request(
+                    db, privacy_request, exc.message
+                )
+        else:
+            # Exceeded retry limit, cancel tasks and set to error state
+            _cancel_interrupted_tasks_and_error_privacy_request(
+                db,
+                privacy_request,
+                f"Privacy request {privacy_request.id} exceeded max retry attempts "
+                f"({max_retries}), canceling tasks and setting to error state",
+            )
+            # Reset retry count since we're giving up
+            reset_privacy_request_retry_count(privacy_request.id)
+
+    except Exception as cache_exc:
+        # If cache operations fail (Redis down, network issues, etc.), fail safe by canceling
+        _cancel_interrupted_tasks_and_error_privacy_request(
+            db,
+            privacy_request,
+            f"Cache operation failed for privacy request {privacy_request.id}, "
+            f"failing safe by canceling tasks: {cache_exc}",
+        )
+
+
+def _get_request_task_ids_in_progress(
+    db: Session, privacy_request_id: str
+) -> List[str]:
+    """Get the IDs of request tasks that are currently in progress for a privacy request."""
+    request_tasks_in_progress = (
+        db.query(RequestTask.id)
+        .filter(RequestTask.privacy_request_id == privacy_request_id)
+        .filter(
+            RequestTask.status.in_(
+                [
+                    ExecutionLogStatus.in_processing,
+                    ExecutionLogStatus.pending,
+                ]
+            )
+        )
+        .all()
+    )
+    return [task[0] for task in request_tasks_in_progress]
+
+
 # pylint: disable=too-many-branches
 @celery_app.task(base=DatabaseTask, bind=True)
 def requeue_interrupted_tasks(self: DatabaseTask) -> None:
@@ -442,17 +562,40 @@ def requeue_interrupted_tasks(self: DatabaseTask) -> None:
             )
 
             # Get task IDs from the queue in a memory-efficient way
-            queued_tasks_ids = _get_task_ids_from_dsr_queue(redis_conn)
+            try:
+                queued_tasks_ids = _get_task_ids_from_dsr_queue(redis_conn)
+            except Exception as queue_exc:
+                logger.warning(
+                    f"Failed to get task IDs from queue, skipping queue state checks: {queue_exc}"
+                )
+                return
 
             # Check each privacy request
             for privacy_request in in_progress_requests:
                 should_requeue = False
                 logger.debug(f"Checking tasks for privacy request {privacy_request.id}")
 
-                task_id = get_cached_task_id(privacy_request.id)
+                try:
+                    task_id = get_cached_task_id(privacy_request.id)
+                except Exception as cache_exc:
+                    # If we can't get the task ID due to cache failure, fail safe by canceling
+                    _cancel_interrupted_tasks_and_error_privacy_request(
+                        db,
+                        privacy_request,
+                        f"Cache failure when getting task ID for privacy request {privacy_request.id}, "
+                        f"failing safe by canceling tasks: {cache_exc}",
+                    )
+                    continue
 
                 # If the task ID is not cached, we can't check if it's running
+                # This means the request is stuck - cancel it
                 if not task_id:
+                    _cancel_interrupted_tasks_and_error_privacy_request(
+                        db,
+                        privacy_request,
+                        f"No task ID found for privacy request {privacy_request.id}, "
+                        f"request is stuck without a running task - canceling",
+                    )
                     continue
 
                 # Check if the main privacy request task is active
@@ -470,30 +613,36 @@ def requeue_interrupted_tasks(self: DatabaseTask) -> None:
                         )
                         should_requeue = True
 
-                    request_tasks_in_progress = (
-                        db.query(RequestTask.id)
-                        .filter(RequestTask.privacy_request_id == privacy_request.id)
-                        .filter(
-                            RequestTask.status.in_(
-                                [
-                                    ExecutionLogStatus.in_processing,
-                                    ExecutionLogStatus.pending,
-                                ]
-                            )
-                        )
-                        .all()
+                    request_task_ids_in_progress = _get_request_task_ids_in_progress(
+                        db, privacy_request.id
                     )
-                    request_task_ids_in_progress = [
-                        task[0] for task in request_tasks_in_progress
-                    ]
 
                     # Check each individual request task
                     for request_task_id in request_task_ids_in_progress:
-                        subtask_id = get_cached_task_id(request_task_id)
+                        try:
+                            subtask_id = get_cached_task_id(request_task_id)
+                        except Exception as cache_exc:
+                            # If we can't get the subtask ID due to cache failure, fail safe by canceling
+                            _cancel_interrupted_tasks_and_error_privacy_request(
+                                db,
+                                privacy_request,
+                                f"Cache failure when getting subtask ID for request task {request_task_id} "
+                                f"(privacy request {privacy_request.id}), failing safe by canceling tasks: {cache_exc}",
+                            )
+                            should_requeue = False
+                            break
 
                         # If the task ID is not cached, we can't check if it's running
+                        # This means the subtask is stuck - cancel the entire privacy request
                         if not subtask_id:
-                            continue
+                            _cancel_interrupted_tasks_and_error_privacy_request(
+                                db,
+                                privacy_request,
+                                f"No task ID found for request task {request_task_id} "
+                                f"(privacy request {privacy_request.id}), subtask is stuck - canceling privacy request",
+                            )
+                            should_requeue = False
+                            break
 
                         if (
                             subtask_id not in queued_tasks_ids
@@ -507,12 +656,4 @@ def requeue_interrupted_tasks(self: DatabaseTask) -> None:
 
                 # Requeue the privacy request if needed
                 if should_requeue:
-                    from fides.service.privacy_request.privacy_request_service import (  # pylint: disable=cyclic-import
-                        PrivacyRequestError,
-                        _requeue_privacy_request,
-                    )
-
-                    try:
-                        _requeue_privacy_request(db, privacy_request)
-                    except PrivacyRequestError as exc:
-                        logger.error(exc.message)
+                    _handle_privacy_request_requeue(db, privacy_request)

fides/api/task/execute_request_tasks.py

@@ -36,6 +36,7 @@ from fides.api.tasks import DSR_QUEUE_NAME, DatabaseTask, celery_app
 from fides.api.util.cache import cache_task_tracking_key
 from fides.api.util.collection_util import Row
 from fides.api.util.logger_context_utils import LoggerContextKeys, log_context
+from fides.api.util.memory_watchdog import memory_limiter
 
 # DSR 3.0 task functions
 
@@ -255,6 +256,7 @@ def queue_downstream_tasks(
 
 
 @celery_app.task(base=DatabaseTask, bind=True)
+@memory_limiter
 @log_context(
     capture_args={
         "privacy_request_id": LoggerContextKeys.privacy_request_id,
@@ -319,6 +321,7 @@ def run_access_node(
 
 
 @celery_app.task(base=DatabaseTask, bind=True)
+@memory_limiter
 @log_context(
     capture_args={
         "privacy_request_id": LoggerContextKeys.privacy_request_id,
@@ -391,6 +394,7 @@ def run_erasure_node(
 
 
 @celery_app.task(base=DatabaseTask, bind=True)
+@memory_limiter
 @log_context(
     capture_args={
         "privacy_request_id": LoggerContextKeys.privacy_request_id,