diff --git a/appengine/standard_python37/bigquery/main.py b/appengine/standard_python37/bigquery/main.py index e8d35ad414f..fc2ec8d4ad0 100644 --- 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 index 83806860cf1..b00f7cd8aac 100644 --- 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")