From 90aff84dbbf949e92cca511d89180e798a6727c2 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Wed, 10 Feb 2021 10:40:28 -0800
Subject: [PATCH] fix: timeout context manager on Windows (#13041)

* fix: timeout decorator in Windows

* Fix lint
---
 superset/utils/core.py | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/superset/utils/core.py b/superset/utils/core.py
index ceec948a393f9..c37a13622ac1b 100644
--- a/superset/utils/core.py
+++ b/superset/utils/core.py
@@ -23,6 +23,7 @@
 import json
 import logging
 import os
+import platform
 import re
 import signal
 import smtplib
@@ -82,6 +83,7 @@
 from sqlalchemy.sql.type_api import Variant
 from sqlalchemy.types import TEXT, TypeDecorator
 
+import _thread  # pylint: disable=C0411
 from superset.errors import ErrorLevel, SupersetErrorType
 from superset.exceptions import (
     CertificateException,
@@ -710,7 +712,7 @@ def validate_json(obj: Union[bytes, bytearray, str]) -> None:
             raise SupersetException("JSON is not valid")
 
 
-class timeout:  # pylint: disable=invalid-name
+class SigalrmTimeout:
     """
     To be used in a ``with`` block and timeout its content.
     """
@@ -749,6 +751,34 @@ def __exit__(  # pylint: disable=redefined-outer-name,unused-variable,redefined-
             logger.exception(ex)
 
 
+class TimerTimeout:
+    def __init__(self, seconds: int = 1, error_message: str = "Timeout") -> None:
+        self.seconds = seconds
+        self.error_message = error_message
+        self.timer = threading.Timer(seconds, _thread.interrupt_main)
+
+    def __enter__(self) -> None:
+        self.timer.start()
+
+    def __exit__(  # pylint: disable=redefined-outer-name,unused-variable,redefined-builtin
+        self, type: Any, value: Any, traceback: TracebackType
+    ) -> None:
+        self.timer.cancel()
+        if type is KeyboardInterrupt:  # raised by _thread.interrupt_main
+            raise SupersetTimeoutException(
+                error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR,
+                message=self.error_message,
+                level=ErrorLevel.ERROR,
+                extra={"timeout": self.seconds},
+            )
+
+
+# Windows has no support for SIGALRM, so we use the timer based timeout
+timeout: Union[Type[TimerTimeout], Type[SigalrmTimeout]] = (
+    TimerTimeout if platform.system() == "Windows" else SigalrmTimeout
+)
+
+
 def pessimistic_connection_handling(some_engine: Engine) -> None:
     @event.listens_for(some_engine, "engine_connect")
     def ping_connection(  # pylint: disable=unused-variable