fix: timeout context manager on Windows (#13041)

* fix: timeout decorator in Windows

* Fix lint
This commit is contained in:
Beto Dealmeida 2021-02-10 10:40:28 -08:00 committed by GitHub
parent 249b51bfba
commit 7b8a0e184b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 31 additions and 1 deletions

View File

@ -23,6 +23,7 @@ import hashlib
import json
import logging
import os
import platform
import re
import signal
import smtplib
@ -82,6 +83,7 @@ from sqlalchemy.engine.reflection import Inspector
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 @@ class timeout: # pylint: disable=invalid-name
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