diff --git a/bigquery/client.py b/bigquery/client.py index aa589d4..927345c 100644 --- 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 index b740414..94c7d61 100644 --- 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, [])