Created new route for results and added redirect for main. (GoogleCloudPlatform#2069)

lbristol88 · tswast · commit 8fba904626b9 · 2019-04-02T16:15:53.000-07:00
* Set a timeout when waiting for query job to finish.

* Transferred results from main function into its own route. Added redirect from main.
diff --git a/appengine/standard_python37/bigquery/main.py b/appengine/standard_python37/bigquery/main.py
@@ -15,17 +15,18 @@
 # [START gae_python37_bigquery]
 import concurrent.futures
 
-from flask import Flask, render_template
+import flask
 from google.cloud import bigquery
 
 
-app = Flask(__name__)
+app = flask.Flask(__name__)
 bigquery_client = bigquery.Client()
 
 
-@app.route('/')
+@app.route("/")
 def main():
-    query_job = bigquery_client.query("""
+    query_job = bigquery_client.query(
+        """
         SELECT
         CONCAT(
             'https://stackoverflow.com/questions/',
@@ -35,20 +36,43 @@ def main():
         WHERE tags like '%google-bigquery%'
         ORDER BY view_count DESC
         LIMIT 10
-    """)
+    """
+    )
+
+    return flask.redirect(
+        flask.url_for(
+            "results",
+            project_id=query_job.project,
+            job_id=query_job.job_id,
+            location=query_job.location,
+        )
+    )
+
+
+@app.route("/results")
+def results():
+    project_id = flask.request.args.get("project_id")
+    job_id = flask.request.args.get("job_id")
+    location = flask.request.args.get("location")
+
+    query_job = bigquery_client.get_job(
+        job_id,
+        project=project_id,
+        location=location,
+    )
 
     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 flask.render_template("timeout.html", job_id=query_job.job_id)
 
-    return render_template('query_result.html', results=results)
+    return flask.render_template("query_result.html", results=results)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     # This is used when running locally only. When deploying to Google App
     # Engine, a webserver process such as Gunicorn will serve the app. This
     # can be configured by adding an `entrypoint` to app.yaml.
-    app.run(host='127.0.0.1', port=8080, debug=True)
+    app.run(host="127.0.0.1", port=8080, debug=True)
 # [END gae_python37_bigquery]
diff --git a/appengine/standard_python37/bigquery/main_test.py b/appengine/standard_python37/bigquery/main_test.py
@@ -28,22 +28,46 @@ def flask_client():
 
 
 def test_main(flask_client):
-    r = flask_client.get('/')
+    r = flask_client.get("/")
+    assert r.status_code == 302
+    assert "/results" in r.headers.get("location", "")
+
+
+def test_results(flask_client, monkeypatch):
+    import main
+
+    fake_job = mock.create_autospec(bigquery.QueryJob)
+    fake_rows = [("example1.com", "42"), ("example2.com", "38")]
+    fake_job.result.return_value = fake_rows
+
+    def fake_get_job(self, job_id, **kwargs):
+        return fake_job
+
+    monkeypatch.setattr(main.bigquery.Client, "get_job", fake_get_job)
+
+    r = flask_client.get(
+        "/results?project_id=123&job_id=456&location=my_location"
+    )
+    response_body = r.data.decode("utf-8")
+
     assert r.status_code == 200
-    assert 'Query Result' in r.data.decode('utf-8')
+    assert "Query Result" in response_body  # verifies header
+    assert "example2.com" in response_body
+    assert "42" in response_body
 
 
-def test_main_timeout(flask_client, monkeypatch):
+def test_results_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):
+    def fake_get_job(self, job_id, **kwargs):
         return fake_job
 
-    monkeypatch.setattr(main.bigquery_client, 'query', fake_query)
+    monkeypatch.setattr(main.bigquery.Client, "get_job", fake_get_job)
+
+    r = flask_client.get("/results", follow_redirects=True)
 
-    r = flask_client.get('/')
     assert r.status_code == 200
-    assert 'Query Timeout' in r.data.decode('utf-8')
+    assert "Query Timeout" in r.data.decode("utf-8")
