Merge pull request #153 from egeucak/master

Dry runs return bytes processed, and cache hit now

Author: tylertreat

bigquery/client.py

@@ -243,7 +243,7 @@ def _submit_query_job(self, query_data):
         -------
         tuple
             job id and query results if query completed. If dry_run is True,
-            job id will be None and results will be empty if the query is valid
+            job id will be None and results will be [cacheHit and totalBytesProcessed] if the query is valid
             or a dict containing the response if invalid.
 
         Raises
@@ -269,13 +269,17 @@ def _submit_query_job(self, query_data):
         schema = query_reply.get('schema', {'fields': None})['fields']
         rows = query_reply.get('rows', [])
         job_complete = query_reply.get('jobComplete', False)
+        cache_hit = query_reply['cacheHit']
+        total_bytes_processed = query_reply['totalBytesProcessed']
 
         # raise exceptions if it's not an async query
         # and job is not completed after timeout
         if not job_complete and query_data.get("timeoutMs", False):
             logger.error('BigQuery job %s timeout' % job_id)
             raise BigQueryTimeoutException()
-
+        
+        if query_data.get("dryRun", True):
+            return job_id, [cache_hit, total_bytes_processed]
         return job_id, [self._transform_row(row, schema) for row in rows]
 
     def _get_job_reference(self, job_id):
@@ -345,8 +349,8 @@ def query(self, query, max_results=None, timeout=0, dry_run=False, use_legacy_sq
             How long to wait for the query to complete, in seconds before
             the request times out and returns.
         dry_run : bool, optional
-            If True, the query isn't actually run. A valid query will return an
-            empty response, while an invalid one will return the same error
+            If True, the query isn't actually run. A valid query will return 
+            cache hit, and total bytes processed, while an invalid one will return the same error
             message it would if it wasn't a dry run.
         use_legacy_sql : bool, optional. Default True.
             If False, the query will use BigQuery's standard SQL (https://cloud.google.com/bigquery/sql-reference/)
@@ -359,7 +363,7 @@ def query(self, query, max_results=None, timeout=0, dry_run=False, use_legacy_sq
         -------
         tuple
             (job id, query results) if the query completed. If dry_run is True,
-            job id will be None and results will be empty if the query is valid
+            job id will be None and results will be [cacheHit and totalBytesProcessed] if the query is valid
             or a ``dict`` containing the response if invalid.
 
         Raises

bigquery/tests/test_client.py

@@ -297,7 +297,9 @@ def test_query(self):
 
         mock_query_job.execute.return_value = {
             'jobReference': expected_job_ref,
-            'jobComplete': True
+            'jobComplete': True,
+            'cacheHit': False,
+            'totalBytesProcessed': 0
         }
 
         self.mock_job_collection.query.return_value = mock_query_job
@@ -329,6 +331,8 @@ def test_query_max_results_set(self):
         mock_query_job.execute.return_value = {
             'jobReference': expected_job_ref,
             'jobComplete': True,
+            'cacheHit': False,
+            'totalBytesProcessed': 0
         }
 
         self.mock_job_collection.query.return_value = mock_query_job
@@ -357,6 +361,8 @@ def test_query_timeout_set(self):
         mock_query_job.execute.return_value = {
             'jobReference': expected_job_ref,
             'jobComplete': True,
+            'cacheHit': False,
+            'totalBytesProcessed': 0
         }
 
         self.mock_job_collection.query.return_value = mock_query_job
@@ -382,6 +388,8 @@ def test_sync_query_timeout(self):
         mock_query_job.execute.return_value = {
             'jobReference': expected_job_ref,
             'jobComplete': False,
+            'cacheHit': False,
+            'totalBytesProcessed': 0
         }
 
         self.mock_job_collection.query.return_value = mock_query_job
@@ -400,6 +408,8 @@ def test_async_query_timeout(self):
         mock_query_job.execute.return_value = {
             'jobReference': expected_job_ref,
             'jobComplete': False,
+            'cacheHit': False,
+            'totalBytesProcessed': 0
         }
 
         self.mock_job_collection.query.return_value = mock_query_job
@@ -409,14 +419,18 @@ def test_async_query_timeout(self):
         self.assertEquals(results, [])
 
     def test_query_dry_run_valid(self):
-        """Ensure that None and an empty list is returned from the query when
+        """Ensure that None and [cacheHit, totalBytesProcessed] is returned from the query when
         dry_run is True and the query is valid.
         """
 
         mock_query_job = mock.Mock()
 
-        mock_query_job.execute.return_value = {'jobReference': {},
-                                               'jobComplete': True}
+        mock_query_job.execute.return_value = {
+            'jobReference': {},
+            'jobComplete': True, 
+            'cacheHit': False,
+            'totalBytesProcessed': 0
+        }
 
         self.mock_job_collection.query.return_value = mock_query_job
 
@@ -428,7 +442,7 @@ def test_query_dry_run_valid(self):
                   'dryRun': True}
         )
         self.assertIsNone(job_id)
-        self.assertEqual([], results)
+        self.assertEqual([False, 0], results)
 
     def test_query_dry_run_invalid(self):
         """Ensure that None and a dict is returned from the query when dry_run
@@ -468,6 +482,8 @@ def test_query_with_results(self):
             'schema': {'fields': [{'name': 'foo', 'type': 'INTEGER'}]},
             'rows': [{'f': [{'v': 10}]}],
             'jobComplete': True,
+            'cacheHit': False,
+            'totalBytesProcessed': 0
         }
 
         self.mock_job_collection.query.return_value = mock_query_job
@@ -491,7 +507,9 @@ def test_query_with_using_legacy_sql(self):
 
         mock_query_job.execute.return_value = {
             'jobReference': expected_job_ref,
-            'jobComplete': True
+            'jobComplete': True,
+            'cacheHit': False,
+            'totalBytesProcessed': 0
         }
 
         self.mock_job_collection.query.return_value = mock_query_job
@@ -873,7 +891,7 @@ def test_json_job_body_constructed_correctly(self):
         body = {
             "jobReference": {
                 "projectId": self.project_id,
-                "jobId": "job"
+                "jobId": "job",
             },
             "configuration": {
                 "load": {