feat: add configurable timeout to querybook (#1207)

* feat: add configurable timeout to querybook * remove print * remove unless comment
pinterest · Mar 31, 2023 · 4e234d6 · 4e234d6
1 parent 354b08d
commit 4e234d6
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 5 deletions.
diff --git 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
@@ -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
@@ -51,6 +51,13 @@
 <p>See [here](https://prestodb.github.io/docs/current/installation/jdbc.html) for more details.</p>""",
         ),
     ),
+    (
+        "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
@@ -44,10 +44,6 @@ function SimpleFormField<T>({
     value: T;
     onChange: onChangeFunc<T>;
 }) {
-    const onFieldChange = React.useCallback(
-        (newVal: any) => onChange('', newVal),
-        [onChange]
-    );
     const {
         description,
         hidden,
@@ -56,6 +52,18 @@ function SimpleFormField<T>({
         field_type: fieldType,
         options,
     } = formField;
+
+    const onFieldChange = React.useCallback(
+        (newVal: any) => {
+            if (fieldType === 'number') {
+                newVal = Number(newVal);
+            }
+
+            return onChange('', newVal);
+        },
+        [onChange, fieldType]
+    );
+
     let controlDOM: React.ReactChild;
     if (fieldType === 'string' || fieldType === 'number') {
         const inputProps = {