Set a timeout when waiting for query job to finish. (GoogleCloudPlatform#2064)

lbristol88 · tswast · commit 18cbdf30f34c · 2019-03-22T14:15:23.000-07:00
diff --git a/appengine/standard_python37/bigquery/.gcloudignore b/appengine/standard_python37/bigquery/.gcloudignore
@@ -0,0 +1,19 @@
+# This file specifies files that are *not* uploaded to Google Cloud Platform
+# using gcloud. It follows the same syntax as .gitignore, with the addition of
+# "#!include" directives (which insert the entries of the given .gitignore-style
+# file at that point).
+#
+# For more information, run:
+#   $ gcloud topic gcloudignore
+#
+.gcloudignore
+# If you would like to upload your .git directory, .gitignore file or files
+# from your .gitignore file, remove the corresponding line
+# below:
+.git
+.gitignore
+
+# Python pycache:
+__pycache__/
+# Ignored by the build system
+/setup.cfg
diff --git a/appengine/standard_python37/bigquery/main.py b/appengine/standard_python37/bigquery/main.py
@@ -13,9 +13,12 @@
 # limitations under the License.
 
 # [START gae_python37_bigquery]
+import concurrent.futures
+
 from flask import Flask, render_template
 from google.cloud import bigquery
 
+
 app = Flask(__name__)
 bigquery_client = bigquery.Client()
 
@@ -34,7 +37,12 @@ def main():
         LIMIT 10
     """)
 
-    results = query_job.result()
+    try:
+        # Set a timeout because queries could take longer than one minute.
+        results = query_job.result(timeout=30)
+    except concurrent.futures.TimeoutError:
+        return render_template('timeout.html', job_id=query_job.job_id)
+
     return render_template('query_result.html', results=results)
 
 
diff --git a/appengine/standard_python37/bigquery/main_test.py b/appengine/standard_python37/bigquery/main_test.py
@@ -12,13 +12,38 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import concurrent.futures
+from unittest import mock
 
-def test_main():
+from google.cloud import bigquery
+import pytest
+
+
+@pytest.fixture
+def flask_client():
     import main
 
     main.app.testing = True
-    client = main.app.test_client()
+    return main.app.test_client()
+
 
-    r = client.get('/')
+def test_main(flask_client):
+    r = flask_client.get('/')
     assert r.status_code == 200
     assert 'Query Result' in r.data.decode('utf-8')
+
+
+def test_main_timeout(flask_client, monkeypatch):
+    import main
+
+    fake_job = mock.create_autospec(bigquery.QueryJob)
+    fake_job.result.side_effect = concurrent.futures.TimeoutError()
+
+    def fake_query(query):
+        return fake_job
+
+    monkeypatch.setattr(main.bigquery_client, 'query', fake_query)
+
+    r = flask_client.get('/')
+    assert r.status_code == 200
+    assert 'Query Timeout' in r.data.decode('utf-8')
diff --git a/appengine/standard_python37/bigquery/templates/query_result.html b/appengine/standard_python37/bigquery/templates/query_result.html
@@ -1,21 +1,31 @@
-<!doctype <!DOCTYPE html>
-<html>
-<head>
-    <meta charset="utf-8" />
-    <title>Query Result</title>
-</head>
-<body>
-    <table>
-        <tr>
-            <th>URL</th>
-            <th>View Count</th>
-        </tr>
-        {% for result in results %}
-        <tr>
-            <td>{{ result[0] }}</td>
-            <td>{{ result[1] }}</td>
-        </tr>
-        {% endfor %}
-    </table>
-</body>
-</html>
+<!DOCTYPE html>
+{#
+Copyright 2019 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+#}
+<meta charset="utf-8" />
+<title>Query Result</title>
+
+<table>
+    <tr>
+        <th>URL</th>
+        <th>View Count</th>
+    </tr>
+    {% for result in results %}
+    <tr>
+        <td>{{ result[0] }}</td>
+        <td>{{ result[1] }}</td>
+    </tr>
+    {% endfor %}
+</table>
diff --git a/appengine/standard_python37/bigquery/templates/timeout.html b/appengine/standard_python37/bigquery/templates/timeout.html
@@ -0,0 +1,20 @@
+<!DOCTYPE html>
+{#
+Copyright 2019 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+#}
+<meta charset="utf-8" />
+<title>Query Timeout</title>
+
+<p>Query job {{ job_id }} timed out. 
