diff --git a/querybook/server/datasources/query_execution.py b/querybook/server/datasources/query_execution.py index 147e1be46..83ddeb180 100644 --- a/querybook/server/datasources/query_execution.py +++ b/querybook/server/datasources/query_execution.py @@ -181,7 +181,7 @@ def cancel_query_and_notify(): "RETRY", # Very unlikely case, because query normally do not retry ): - task.revoke() # last attempt to cancel it + task.revoke(terminate=True) # last attempt to cancel it cancel_query_and_notify() elif task.state == "ABORTED": # In this case, the task is already aborted, but the status is running diff --git a/querybook/server/lib/query_executor/clients/presto.py b/querybook/server/lib/query_executor/clients/presto.py index 2b31f27a3..8346cb2ca 100644 --- a/querybook/server/lib/query_executor/clients/presto.py +++ b/querybook/server/lib/query_executor/clients/presto.py @@ -51,6 +51,7 @@ def __init__( password: Optional[str] = None, proxy_user: Optional[str] = None, impersonate: bool = False, + connection_timeout: int = None, *args: Any, **kwargs: Any ) -> None: @@ -59,6 +60,10 @@ def __init__( host = presto_conf.host port = 8080 if not presto_conf.port else presto_conf.port + requests_kwargs = {} + if connection_timeout and connection_timeout > 0: + requests_kwargs["timeout"] = connection_timeout + connection = presto.connect( host, port=port, @@ -69,6 +74,7 @@ def __init__( schema=presto_conf.schema, source="querybook", protocol=presto_conf.protocol, + requests_kwargs=requests_kwargs, ) self._connection = connection super(PrestoClient, self).__init__() diff --git a/querybook/server/lib/query_executor/executor_template/templates.py b/querybook/server/lib/query_executor/executor_template/templates.py index 30815f20e..d57c6ab7f 100644 --- a/querybook/server/lib/query_executor/executor_template/templates.py +++ b/querybook/server/lib/query_executor/executor_template/templates.py @@ -51,6 +51,13 @@
See [here](https://prestodb.github.io/docs/current/installation/jdbc.html) for more details.
""", ), ), + ( + "connection_timeout", + FormField( + field_type=FormFieldType.Number, + helper="Max number of seconds before timing out the request", + ), + ), ) trino_executor_template = StructFormField( diff --git a/querybook/webapp/ui/SmartForm/SmartForm.tsx b/querybook/webapp/ui/SmartForm/SmartForm.tsx index 2123484a7..fb486febe 100644 --- a/querybook/webapp/ui/SmartForm/SmartForm.tsx +++ b/querybook/webapp/ui/SmartForm/SmartForm.tsx @@ -44,10 +44,6 @@ function SimpleFormField