Skip to content

Commit

Permalink
fix: timeout context manager on Windows (apache#13041)
Browse files Browse the repository at this point in the history
* fix: timeout decorator in Windows

* Fix lint
  • Loading branch information
betodealmeida authored and amitmiran137 committed Feb 14, 2021
1 parent a2d5196 commit 90aff84
Showing 1 changed file with 31 additions and 1 deletion.
32 changes: 31 additions & 1 deletion superset/utils/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import json
import logging
import os
import platform
import re
import signal
import smtplib
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.
"""
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 90aff84

Please sign in to comment.