add external_udf_uris support to query(), with docstring and

Robin Thomas · Robin Thomas · commit a9e37c839bc2 · 2016-10-26T12:21:22.000-04:00
test coverage. Change write_to_table code to use the same
list comprehension to construct userDefinedFunctionResources.
diff --git a/bigquery/client.py b/bigquery/client.py
@@ -276,7 +276,7 @@ def _insert_job(self, body_object):
             body=body_object
         ).execute()
 
-    def query(self, query, max_results=None, timeout=0, dry_run=False, use_legacy_sql=None):
+    def query(self, query, max_results=None, timeout=0, dry_run=False, use_legacy_sql=None, external_udf_uris=None):
         """Submit a query to BigQuery.
 
         Parameters
@@ -294,6 +294,9 @@ def query(self, query, max_results=None, timeout=0, dry_run=False, use_legacy_sq
             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/)
+        external_udf_uris : list, optional
+            Contains external UDF URIs. If given, URIs must be Google Cloud
+            Storage and have .js extensions.
 
 
         Returns
@@ -321,6 +324,10 @@ def query(self, query, max_results=None, timeout=0, dry_run=False, use_legacy_sq
         if use_legacy_sql is not None:
             query_data['useLegacySql'] = use_legacy_sql
 
+        if external_udf_uris:
+            query_data['userDefinedFunctionResources'] = \
+                [ {'resourceUri': u} for u in external_udf_uris ]
+
         return self._submit_query_job(query_data)
 
     def get_query_schema(self, job_id):
@@ -1048,7 +1055,7 @@ def write_to_table(
             query,
             dataset=None,
             table=None,
-            external_udf_uris=[],
+            external_udf_uris=None,
             allow_large_results=None,
             use_query_cache=None,
             priority=None,
@@ -1073,7 +1080,7 @@ def write_to_table(
         table : str, optional
             String id of the table
         external_udf_uris : list, optional
-            Contains extternal UDF URIs. If given, URIs must be Google Cloud
+            Contains external UDF URIs. If given, URIs must be Google Cloud
             Storage and have .js extensions.
         allow_large_results : bool, optional
             Whether or not to allow large results
@@ -1144,13 +1151,9 @@ def write_to_table(
         if write_disposition:
             configuration['writeDisposition'] = write_disposition
 
-        configuration['userDefinedFunctionResources'] = []
-        for external_udf_uri in external_udf_uris:
-            configuration['userDefinedFunctionResources'].append(
-                {
-                    "resourceUri": external_udf_uri
-                }
-            )
+        if external_udf_uris:
+            configuration['userDefinedFunctionResources'] = \
+                [ {'resourceUri': u} for u in external_udf_uris ]
 
         body = {
             "configuration": {
diff --git a/bigquery/tests/test_client.py b/bigquery/tests/test_client.py
@@ -259,6 +259,7 @@ def setUp(self):
 
         self.query = 'foo'
         self.project_id = 'project'
+        self.external_udf_uris = ['gs://bucket/external_udf.js']
         self.client = client.BigQueryClient(self.mock_bq_service,
                                             self.project_id)
 
@@ -276,12 +277,17 @@ def test_query(self):
 
         self.mock_job_collection.query.return_value = mock_query_job
 
-        job_id, results = self.client.query(self.query)
+        job_id, results = self.client.query(self.query, external_udf_uris=self.external_udf_uris)
 
         self.mock_job_collection.query.assert_called_once_with(
             projectId=self.project_id,
-            body={'query': self.query, 'timeoutMs': 0, 'dryRun': False,
-                  'maxResults': None}
+            body={
+                'query': self.query,
+                'userDefinedFunctionResources': [ {'resourceUri': u} for u in self.external_udf_uris ],
+                'timeoutMs': 0,
+                'dryRun': False,
+                'maxResults': None
+            }
         )
         self.assertEquals(job_id, 'spiderman')
         self.assertEquals(results, [])
