setup cooperative timeout

Author: hsheth2

metadata-ingestion/setup.py

@@ -94,8 +94,7 @@
 sqlglot_lib = {
     # Using an Acryl fork of sqlglot.
     # https://github.com/tobymao/sqlglot/compare/main...hsheth2:sqlglot:hsheth?expand=1
-    "acryl-sqlglot==21.1.2.dev9",
-    "wrapt_timeout_decorator",
+    "acryl-sqlglot==21.1.2.dev10",
 }
 
 sql_common = (

metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py

@@ -12,7 +12,6 @@
 import sqlglot.optimizer.annotate_types
 import sqlglot.optimizer.optimizer
 import sqlglot.optimizer.qualify
-import wrapt_timeout_decorator
 
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.ingestion.graph.client import DataHubGraph
@@ -46,6 +45,10 @@
     is_dialect_instance,
     parse_statement,
 )
+from datahub.utilities.cooperative_timeout import (
+    CooperativeTimeout,
+    CooperativeTimeoutError,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -311,9 +314,6 @@ class SqlUnderstandingError(Exception):
 
 
 # TODO: Break this up into smaller functions.
-@wrapt_timeout_decorator.timeout(
-    SQL_LINEAGE_TIMEOUT_SECONDS if SQL_LINEAGE_TIMEOUT_ENABLED else None,
-)
 def _column_level_lineage(  # noqa: C901
     statement: sqlglot.exp.Expression,
     dialect: sqlglot.Dialect,
@@ -873,19 +873,27 @@ def _sqlglot_lineage_inner(
     column_lineage: Optional[List[_ColumnLineageInfo]] = None
     try:
         if select_statement is not None:
-            column_lineage = _column_level_lineage(
-                select_statement,
-                dialect=dialect,
-                table_schemas=table_name_schema_mapping,
-                output_table=downstream_table,
-                default_db=default_db,
-                default_schema=default_schema,
-            )
+            with CooperativeTimeout(
+                timeout=SQL_LINEAGE_TIMEOUT_SECONDS
+                if SQL_LINEAGE_TIMEOUT_ENABLED
+                else None
+            ):
+                column_lineage = _column_level_lineage(
+                    select_statement,
+                    dialect=dialect,
+                    table_schemas=table_name_schema_mapping,
+                    output_table=downstream_table,
+                    default_db=default_db,
+                    default_schema=default_schema,
+                )
     except UnsupportedStatementTypeError as e:
         # Inject details about the outer statement type too.
         e.args = (f"{e.args[0]} (outer statement type: {type(statement)})",)
         debug_info.column_error = e
         logger.debug(debug_info.column_error)
+    except CooperativeTimeoutError as e:
+        logger.debug(f"Timed out while generating column-level lineage: {e}")
+        debug_info.column_error = e
     except SqlUnderstandingError as e:
         logger.debug(f"Failed to generate column-level lineage: {e}", exc_info=True)
         debug_info.column_error = e

metadata-ingestion/src/datahub/utilities/cooperative_timeout.py

@@ -0,0 +1,69 @@
+import threading
+import time
+from contextlib import AbstractContextManager
+from types import TracebackType
+from typing import Optional, Type
+
+_cooperation = threading.local()
+
+
+class CooperativeTimeoutError(TimeoutError):
+    """An exception raised when a cooperative timeout is exceeded."""
+
+
+def cooperate() -> None:
+    """Method to be called periodically to cooperate with the timeout mechanism."""
+
+    deadline = getattr(_cooperation, "deadline", None)
+    if deadline is not None and deadline < time.perf_counter_ns():
+        raise CooperativeTimeoutError("CooperativeTimeout deadline exceeded")
+
+
+class CooperativeTimeout(AbstractContextManager):
+    """A cooperative timeout mechanism.
+
+    Getting code to time out in Python is actually rather tricky. Common approaches include:
+
+      - Using the signal module to set a signal handler that raises an exception
+        after a certain time. Unfortunately, this approach only works on the main
+        thread, and is not available on Windows.
+      - Creating a separate process to run the code and then killing it if it hasn't
+        finished by the deadline. This usually requires that all arguments/return
+        types are pickleable so that they can be passed between processes. Overall,
+        this approach is heavy-handed and can be tricky to implement correctly.
+      - Using `threading` is not an option, since Python threads are not interruptible
+        (unless you're willing to use some hacks https://stackoverflow.com/a/61528202).
+        Attempting to forcibly terminate a thread can deadlock on the GIL.
+
+    In cases where (1) we have control over the code that we want to time out and
+    (2) we can modify it to regularly and reliably call a specific function, we can
+    use a cooperative timeout mechanism instead.
+
+    This is not reentrant and cannot be used in nested contexts. It can be used
+    in multi-threaded contexts, so long as the cooperative function is called
+    from the same thread that created the timeout.
+
+    Args:
+        timeout: The timeout in seconds. If None, the timeout is disabled.
+    """
+
+    def __init__(self, timeout: Optional[None] = None):
+        self.timeout = timeout
+
+    def __enter__(self) -> "CooperativeTimeout":
+        if hasattr(_cooperation, "deadline"):
+            raise RuntimeError("CooperativeTimeout already active")
+        if self.timeout is not None:
+            _cooperation.deadline = (
+                time.perf_counter_ns() + self.timeout * 1_000_000_000
+            )
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        if self.timeout is not None:
+            del _cooperation.deadline