factor querying and job insertion into shared method

cpdean · cpdean · commit 408f6ec9120c · 2015-05-02T14:54:47.000-04:00
diff --git a/bigquery/client.py b/bigquery/client.py
@@ -120,17 +120,22 @@ def __init__(self, bq_service, project_id, swallow_results=True):
         self.swallow_results = swallow_results
         self.cache = {}
 
-    def query(self, query, max_results=None, timeout=0, dry_run=False):
-        """Submit a query to BigQuery.
+    def _submit_query_job(self, query_data):
+
+        """ Submit a query job to BigQuery.
+
+            This is similar to BigQueryClient.query, but gives the user
+            direct access to the query method on the offical BigQuery
+            python client.
+
+            For fine-grained control over a query job, see:
+            https://google-api-client-libraries.appspot.com/documentation/bigquery/v2/python/latest/bigquery_v2.jobs.html#query
+
+
 
         Args:
-            query: BigQuery query string.
-            max_results: maximum number of rows to return per page of results.
-            timeout: how long to wait for the query to complete, in seconds,
-                     before the request times out and returns.
-            dry_run: 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 message it would if it wasn't a dry run.
+            query_data: query object as per "configuration.query" in
+                        https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.query
 
         Returns:
             job id and query results if query completed. If dry_run is True,
@@ -141,21 +146,15 @@ def query(self, query, max_results=None, timeout=0, dry_run=False):
             BigQueryTimeoutException on timeout
         """
 
-        logging.debug('Executing query: %s' % query)
+        logging.debug('Submitting query job: %s' % query_data)
 
         job_collection = self.bigquery.jobs()
-        query_data = {
-            'query': query,
-            'timeoutMs': timeout * 1000,
-            'dryRun': dry_run,
-            'maxResults': max_results,
-        }
 
         try:
             query_reply = job_collection.query(
                 projectId=self.project_id, body=query_data).execute()
         except HttpError as e:
-            if dry_run:
+            if query_data.get("dryRun", False):
                 return None, json.loads(e.content)
             raise
 
@@ -166,12 +165,75 @@ def query(self, query, max_results=None, timeout=0, dry_run=False):
 
         # raise exceptions if it's not an async query
         # and job is not completed after timeout
-        if not job_complete and timeout:
+        if not job_complete and query_data.get("timeoutMs", False):
             logging.error('BigQuery job %s timeout' % job_id)
             raise BigQueryTimeoutException()
 
         return job_id, [self._transform_row(row, schema) for row in rows]
 
+    def _insert_job(self, body_object):
+
+        """ Submit a job to BigQuery
+
+            Direct proxy to the insert() method of the offical BigQuery
+            python client.
+
+            Able to submit load, link, query, copy, or extract jobs.
+
+            For more details, see:
+            https://google-api-client-libraries.appspot.com/documentation/bigquery/v2/python/latest/bigquery_v2.jobs.html#insert
+
+
+        Args:
+            body_object: body object passed to bigquery.jobs().insert()
+
+        Returns:
+            response of the bigquery.jobs().insert().execute() call
+
+        Raises:
+            BigQueryTimeoutException on timeout
+        """
+
+        logging.debug('Submitting job: %s' % body_object)
+
+        job_collection = self.bigquery.jobs()
+
+        return job_collection.insert(
+            projectId=self.project_id,
+            body=body_object
+        ).execute()
+
+    def query(self, query, max_results=None, timeout=0, dry_run=False):
+        """Submit a query to BigQuery.
+
+        Args:
+            query: BigQuery query string.
+            max_results: maximum number of rows to return per page of results.
+            timeout: how long to wait for the query to complete, in seconds,
+                     before the request times out and returns.
+            dry_run: 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 message it would if it wasn't a dry run.
+
+        Returns:
+            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
+            or a dict containing the response if invalid.
+
+        Raises:
+            BigQueryTimeoutException on timeout
+        """
+
+        logging.debug('Executing query: %s' % query)
+
+        query_data = {
+            'query': query,
+            'timeoutMs': timeout * 1000,
+            'dryRun': dry_run,
+            'maxResults': max_results,
+        }
+        return self._submit_query_job(query_data)
+
     def get_query_schema(self, job_id):
         """Retrieve the schema of a query by job id.
 
@@ -534,9 +596,7 @@ def import_data_from_uris(
         }
 
         logging.debug("Creating load job %s" % body)
-        job_resource = self.bigquery.jobs() \
-            .insert(projectId=self.project_id, body=body) \
-            .execute()
+        job_resource = self._insert_job(body)
         self._raise_insert_exception_if_error(job_resource)
         return job_resource
 
@@ -620,9 +680,7 @@ def export_data_to_uris(
         }
 
         logging.info("Creating export job %s" % body)
-        job_resource = self.bigquery.jobs() \
-            .insert(projectId=self.project_id, body=body) \
-            .execute()
+        job_resource = self._insert_job(body)
         self._raise_insert_exception_if_error(job_resource)
         return job_resource
 
@@ -696,9 +754,7 @@ def write_to_table(
         }
 
         logging.info("Creating write to table job %s" % body)
-        job_resource = self.bigquery.jobs() \
-            .insert(projectId=self.project_id, body=body) \
-            .execute()
+        job_resource = self._insert_job(body)
         self._raise_insert_exception_if_error(job_resource)
         return job_resource
 
