use function instead of class

Author: hsheth2

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

@@ -46,8 +46,8 @@
     parse_statement,
 )
 from datahub.utilities.cooperative_timeout import (
-    CooperativeTimeout,
     CooperativeTimeoutError,
+    cooperative_timeout,
 )
 
 logger = logging.getLogger(__name__)
@@ -873,7 +873,7 @@ def _sqlglot_lineage_inner(
     column_lineage: Optional[List[_ColumnLineageInfo]] = None
     try:
         if select_statement is not None:
-            with CooperativeTimeout(
+            with cooperative_timeout(
                 timeout=SQL_LINEAGE_TIMEOUT_SECONDS
                 if SQL_LINEAGE_TIMEOUT_ENABLED
                 else None

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

@@ -1,8 +1,7 @@
+import contextlib
 import threading
 import time
-from contextlib import AbstractContextManager
-from types import TracebackType
-from typing import Optional, Type
+from typing import Iterator, Optional
 
 _cooperation = threading.local()
 
@@ -19,51 +18,46 @@ def cooperate() -> None:
         raise CooperativeTimeoutError("CooperativeTimeout deadline exceeded")
 
 
-class CooperativeTimeout(AbstractContextManager):
+@contextlib.contextmanager
+def cooperative_timeout(timeout: Optional[float] = None) -> Iterator[None]:
     """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.
+    Useful 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.
 
     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.
+
+    Raises:
+        RuntimeError: If a cooperative timeout is already active.
+        CooperativeTimeoutError: If the cooperative timeout is exceeded.
     """
 
-    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
+    # Getting code to time out in Python is actually rather tricky, and so this felt
+    # like the most straightforward approach. Other 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.
+
+    if hasattr(_cooperation, "deadline"):
+        raise RuntimeError("cooperative timeout already active")
+
+    if timeout is not None:
+        _cooperation.deadline = time.perf_counter_ns() + timeout * 1_000_000_000
+        yield
+        del _cooperation.deadline
+
+    else:
+        # No-op.
+        yield