diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index b9e83f30..26b6340b 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -4,6 +4,8 @@ Changelog 0.2.1 / 2017-??-?? ------------------ +- :func:`read_gbq` now handles query configuration `query.timeoutMs` and stop waiting. (:issue:`76`) + 0.2.0 / 2017-07-24 ------------------ diff --git a/pandas_gbq/gbq.py b/pandas_gbq/gbq.py index 452de5ef..9d1e5cac 100644 --- a/pandas_gbq/gbq.py +++ b/pandas_gbq/gbq.py @@ -172,6 +172,13 @@ class NotFoundException(ValueError): pass +class QueryTimeout(ValueError): + """ + Raised when the query job timeout + """ + pass + + class StreamingInsertError(ValueError): """ Raised when BigQuery reports a streaming insert error. @@ -536,6 +543,11 @@ def run_query(self, query, **kwargs): while not query_reply.get('jobComplete', False): self.print_elapsed_seconds(' Elapsed', 's. Waiting...') + + timeout_ms = job_config['query'].get('timeoutMs') + if timeout_ms and timeout_ms < self.get_elapsed_seconds() * 1000: + raise QueryTimeout('Query timeout: {} ms'.format(timeout_ms)) + try: query_reply = job_collection.getQueryResults( projectId=job_reference['projectId'], diff --git a/pandas_gbq/tests/test_gbq.py b/pandas_gbq/tests/test_gbq.py index 930a9296..c9ac31dd 100644 --- a/pandas_gbq/tests/test_gbq.py +++ b/pandas_gbq/tests/test_gbq.py @@ -884,6 +884,19 @@ def test_configuration_raises_value_error_with_multiple_config(self): private_key=_get_private_key_path(), configuration=config) + def test_timeout_configuration(self): + sql_statement = 'SELECT 1' + config = { + 'query': { + "timeoutMs": 1 + } + } + # Test that QueryTimeout error raises + with pytest.raises(gbq.QueryTimeout): + gbq.read_gbq(sql_statement, project_id=_get_project_id(), + private_key=_get_private_key_path(), + configuration=config) + def test_query_response_bytes(self): assert self.gbq_connector.sizeof_fmt(999) == "999.0 B" assert self.gbq_connector.sizeof_fmt(1024) == "1.0 KB"