ENH: Add timeout support (#76)

* Add timeout support

* Fix a pep8 issue

* Check before fetch results

* Add test

* timeoutMs -> timeout_ms

* Add changelog

Author: Takashi Nishibayashi

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
 ------------------
 

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'],

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"