Skip to content

Commit c8ba881

Browse files
dwmclarydwmclary
committed
refork and fix of #2015 
1 parent 565c992 commit c8ba881

File tree

5 files changed

+251
-6
lines changed

5 files changed

+251
-6
lines changed

‎gcloud/bigquery/job.py‎

Lines changed: 70 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,42 @@
2727
from gcloud.bigquery._helpers import _TypedProperty
2828

2929

30+
class UDFResource(object):
31+
"""Describe a single user-defined function (UDF) resource.
32+
:type udf_type: str
33+
:param udf_type: the type of the resource ('inlineCode' or 'resourceUri')
34+
35+
:type value: str
36+
:param value: the inline code or resource URI
37+
38+
See
39+
https://cloud.google.com/bigquery/user-defined-functions#api
40+
"""
41+
def __init__(self, udf_type, value):
42+
self.udf_type = udf_type
43+
self.value = value
44+
45+
def __eq__(self, other):
46+
return(
47+
self.udf_type == other.udf_type and
48+
self.value == other.value)
49+
50+
51+
def _build_udf_resources(resources):
52+
"""
53+
:type resources: sequence of :class:`UDFResource`
54+
:param resources: fields to be appended
55+
56+
:rtype: mapping
57+
:returns: a mapping describing userDefinedFunctionResources for the query.
58+
"""
59+
udfs = []
60+
for resource in resources:
61+
udf = {resource.udf_type: resource.value}
62+
udfs.append(udf)
63+
return udfs
64+
65+
3066
class Compression(_EnumProperty):
3167
"""Pseudo-enum for ``compression`` properties."""
3268
GZIP = 'GZIP'
@@ -888,14 +924,44 @@ class QueryJob(_AsyncJob):
888924
:type client: :class:`gcloud.bigquery.client.Client`
889925
:param client: A client which holds credentials and project configuration
890926
for the dataset (which requires a project).
927+
928+
:type udf_resources: tuple
929+
:param udf_resources: An iterable of
930+
:class:`gcloud.bigquery.job.UDFResource`
931+
(empty by default)
891932
"""
892933
_JOB_TYPE = 'query'
934+
_UDF_KEY = 'userDefinedFunctionResources'
935+
_udf_resources = None
893936

894-
def __init__(self, name, query, client):
937+
def __init__(self, name, query, client, udf_resources=()):
895938
super(QueryJob, self).__init__(name, client)
896939
self.query = query
940+
self.udf_resources = udf_resources
897941
self._configuration = _AsyncQueryConfiguration()
898942

943+
@property
944+
def udf_resources(self):
945+
"""Property for list of UDF resources attached to a query
946+
See
947+
https://cloud.google.com/bigquery/user-defined-functions#api
948+
"""
949+
return list(self._udf_resources)
950+
951+
@udf_resources.setter
952+
def udf_resources(self, value):
953+
"""Update queries UDF resources
954+
955+
:type value: list of :class:`UDFResource`
956+
:param value: an object which defines the type and value of a resource
957+
958+
:raises: TypeError if 'value' is not a sequence, or ValueError if
959+
any item in the sequence is not a UDFResource
960+
"""
961+
if not all(isinstance(u, UDFResource) for u in value):
962+
raise ValueError("udf items must be UDFResource")
963+
self._udf_resources = tuple(value)
964+
899965
allow_large_results = _TypedProperty('allow_large_results', bool)
900966
"""See:
901967
https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.query.allowLargeResults
@@ -979,6 +1045,9 @@ def _populate_config_resource(self, configuration):
9791045
configuration['useLegacySql'] = self.use_legacy_sql
9801046
if self.write_disposition is not None:
9811047
configuration['writeDisposition'] = self.write_disposition
1048+
if len(self._udf_resources) > 0:
1049+
configuration[self._UDF_KEY] = _build_udf_resources(
1050+
self._udf_resources)
9821051

9831052
def _build_resource(self):
9841053
"""Generate a resource for :meth:`begin`."""

‎gcloud/bigquery/query.py‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
from gcloud.bigquery._helpers import _rows_from_json
2121
from gcloud.bigquery.dataset import Dataset
2222
from gcloud.bigquery.job import QueryJob
23+
from gcloud.bigquery.job import UDFResource
24+
from gcloud.bigquery.job import _build_udf_resources
2325
from gcloud.bigquery.table import _parse_schema_resource
2426

2527

@@ -46,12 +48,22 @@ class QueryResults(object):
4648
:type client: :class:`gcloud.bigquery.client.Client`
4749
:param client: A client which holds credentials and project configuration
4850
for the dataset (which requires a project).
51+
52+
:type udf_resources: tuple
53+
:param udf_resources: An iterable of
54+
:class:`gcloud.bigquery.job.UDFResource`
55+
(empty by default)
4956
"""
50-
def __init__(self, query, client):
57+
58+
_UDF_KEY = 'userDefinedFunctionResources'
59+
_udf_resources = None
60+
61+
def __init__(self, query, client, udf_resources=()):
5162
self._client = client
5263
self._properties = {}
5364
self.query = query
5465
self._configuration = _SyncQueryConfiguration()
66+
self.udf_resources = udf_resources
5567
self._job = None
5668

5769
@property
@@ -204,6 +216,28 @@ def schema(self):
204216
"""
205217
return _parse_schema_resource(self._properties.get('schema', {}))
206218

219+
@property
220+
def udf_resources(self):
221+
"""Property for list of UDF resources attached to a query
222+
See
223+
https://cloud.google.com/bigquery/user-defined-functions#api
224+
"""
225+
return list(self._udf_resources)
226+
227+
@udf_resources.setter
228+
def udf_resources(self, value):
229+
"""Update queries UDF resources
230+
231+
:type value: list of :class:`UDFResource`
232+
:param value: an object which defines the type and value of a resource
233+
234+
:raises: TypeError if 'value' is not a sequence, or ValueError if
235+
any item in the sequence is not a UDFResource
236+
"""
237+
if not all(isinstance(udf, UDFResource) for udf in value):
238+
raise ValueError("udf items must be UDFResource")
239+
self._udf_resources = tuple(value)
240+
207241
default_dataset = _TypedProperty('default_dataset', Dataset)
208242
"""See:
209243
https://cloud.google.com/bigquery/docs/reference/v2/jobs/query#defaultDataset
@@ -277,6 +311,9 @@ def _build_resource(self):
277311
if self.dry_run is not None:
278312
resource['dryRun'] = self.dry_run
279313

314+
if len(self._udf_resources) > 0:
315+
resource[self._UDF_KEY] = _build_udf_resources(self._udf_resources)
316+
280317
return resource
281318

282319
def run(self, client=None):

‎gcloud/bigquery/test_job.py‎

Lines changed: 58 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1384,7 +1384,7 @@ def test_begin_w_bound_client(self):
13841384
job = self._makeOne(self.JOB_NAME, self.QUERY, client)
13851385

13861386
job.begin()
1387-
1387+
self.assertEqual(job.udf_resources, [])
13881388
self.assertEqual(len(conn._requested), 1)
13891389
req = conn._requested[0]
13901390
self.assertEqual(req['method'], 'POST')
@@ -1396,7 +1396,7 @@ def test_begin_w_bound_client(self):
13961396
},
13971397
'configuration': {
13981398
'query': {
1399-
'query': self.QUERY,
1399+
'query': self.QUERY
14001400
},
14011401
},
14021402
}
@@ -1468,6 +1468,62 @@ def test_begin_w_alternate_client(self):
14681468
self.assertEqual(req['data'], SENT)
14691469
self._verifyResourceProperties(job, RESOURCE)
14701470

1471+
def test_begin_w_bound_client_and_udf(self):
1472+
from gcloud.bigquery.job import UDFResource
1473+
RESOURCE_URI = 'gs://some-bucket/js/lib.js'
1474+
PATH = 'projects/%s/jobs' % self.PROJECT
1475+
RESOURCE = self._makeResource()
1476+
# Ensure None for missing server-set props
1477+
del RESOURCE['statistics']['creationTime']
1478+
del RESOURCE['etag']
1479+
del RESOURCE['selfLink']
1480+
del RESOURCE['user_email']
1481+
conn = _Connection(RESOURCE)
1482+
client = _Client(project=self.PROJECT, connection=conn)
1483+
job = self._makeOne(self.JOB_NAME, self.QUERY, client,
1484+
udf_resources=[
1485+
UDFResource("resourceUri", RESOURCE_URI)
1486+
])
1487+
1488+
job.begin()
1489+
1490+
self.assertEqual(len(conn._requested), 1)
1491+
req = conn._requested[0]
1492+
self.assertEqual(req['method'], 'POST')
1493+
self.assertEqual(req['path'], '/%s' % PATH)
1494+
self.assertEqual(job.udf_resources,
1495+
[UDFResource("resourceUri", RESOURCE_URI)])
1496+
SENT = {
1497+
'jobReference': {
1498+
'projectId': self.PROJECT,
1499+
'jobId': self.JOB_NAME,
1500+
},
1501+
'configuration': {
1502+
'query': {
1503+
'query': self.QUERY,
1504+
'userDefinedFunctionResources':
1505+
[{'resourceUri': RESOURCE_URI}]
1506+
},
1507+
},
1508+
}
1509+
self.assertEqual(req['data'], SENT)
1510+
self._verifyResourceProperties(job, RESOURCE)
1511+
1512+
def test_begin_w_bad_udf(self):
1513+
RESOURCE = self._makeResource()
1514+
# Ensure None for missing server-set props
1515+
del RESOURCE['statistics']['creationTime']
1516+
del RESOURCE['etag']
1517+
del RESOURCE['selfLink']
1518+
del RESOURCE['user_email']
1519+
conn = _Connection(RESOURCE)
1520+
client = _Client(project=self.PROJECT, connection=conn)
1521+
job = self._makeOne(self.JOB_NAME, self.QUERY, client)
1522+
1523+
with self.assertRaises(ValueError):
1524+
job.udf_resources = ["foo"]
1525+
self.assertEqual(job.udf_resources, [])
1526+
14711527
def test_exists_miss_w_bound_client(self):
14721528
PATH = 'projects/%s/jobs/%s' % (self.PROJECT, self.JOB_NAME)
14731529
conn = _Connection()

‎gcloud/bigquery/test_query.py‎

Lines changed: 83 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,7 @@ def test_run_w_bound_client(self):
181181
conn = _Connection(RESOURCE)
182182
client = _Client(project=self.PROJECT, connection=conn)
183183
query = self._makeOne(self.QUERY, client)
184-
184+
self.assertEqual(query.udf_resources, [])
185185
query.run()
186186

187187
self.assertEqual(len(conn._requested), 1)
@@ -233,6 +233,88 @@ def test_run_w_alternate_client(self):
233233
self.assertEqual(req['data'], SENT)
234234
self._verifyResourceProperties(query, RESOURCE)
235235

236+
def test_run_w_inline_udf(self):
237+
from gcloud.bigquery.query import UDFResource
238+
INLINE_UDF_CODE = 'var someCode = "here";'
239+
PATH = 'projects/%s/queries' % self.PROJECT
240+
RESOURCE = self._makeResource(complete=False)
241+
conn = _Connection(RESOURCE)
242+
client = _Client(project=self.PROJECT, connection=conn)
243+
query = self._makeOne(self.QUERY, client)
244+
query.udf_resources = [UDFResource("inlineCode", INLINE_UDF_CODE)]
245+
246+
query.run()
247+
248+
self.assertEqual(len(conn._requested), 1)
249+
req = conn._requested[0]
250+
self.assertEqual(req['method'], 'POST')
251+
self.assertEqual(req['path'], '/%s' % PATH)
252+
SENT = {'query': self.QUERY,
253+
'userDefinedFunctionResources':
254+
[{'inlineCode': INLINE_UDF_CODE}]}
255+
self.assertEqual(req['data'], SENT)
256+
self._verifyResourceProperties(query, RESOURCE)
257+
258+
def test_run_w_udf_resource_uri(self):
259+
from gcloud.bigquery.job import UDFResource
260+
RESOURCE_URI = 'gs://some-bucket/js/lib.js'
261+
PATH = 'projects/%s/queries' % self.PROJECT
262+
RESOURCE = self._makeResource(complete=False)
263+
conn = _Connection(RESOURCE)
264+
client = _Client(project=self.PROJECT, connection=conn)
265+
query = self._makeOne(self.QUERY, client)
266+
query.udf_resources = [UDFResource("resourceUri", RESOURCE_URI)]
267+
268+
query.run()
269+
270+
self.assertEqual(len(conn._requested), 1)
271+
req = conn._requested[0]
272+
self.assertEqual(req['method'], 'POST')
273+
self.assertEqual(req['path'], '/%s' % PATH)
274+
SENT = {'query': self.QUERY,
275+
'userDefinedFunctionResources':
276+
[{'resourceUri': RESOURCE_URI}]}
277+
self.assertEqual(req['data'], SENT)
278+
self._verifyResourceProperties(query, RESOURCE)
279+
280+
def test_run_w_mixed_udfs(self):
281+
from gcloud.bigquery.job import UDFResource
282+
RESOURCE_URI = 'gs://some-bucket/js/lib.js'
283+
INLINE_UDF_CODE = 'var someCode = "here";'
284+
PATH = 'projects/%s/queries' % self.PROJECT
285+
RESOURCE = self._makeResource(complete=False)
286+
conn = _Connection(RESOURCE)
287+
client = _Client(project=self.PROJECT, connection=conn)
288+
query = self._makeOne(self.QUERY, client)
289+
query.udf_resources = [UDFResource("resourceUri", RESOURCE_URI),
290+
UDFResource("inlineCode", INLINE_UDF_CODE)]
291+
292+
query.run()
293+
294+
self.assertEqual(len(conn._requested), 1)
295+
req = conn._requested[0]
296+
self.assertEqual(req['method'], 'POST')
297+
self.assertEqual(req['path'], '/%s' % PATH)
298+
self.assertEqual(query.udf_resources,
299+
[UDFResource("resourceUri", RESOURCE_URI),
300+
UDFResource("inlineCode", INLINE_UDF_CODE)])
301+
SENT = {'query': self.QUERY,
302+
'userDefinedFunctionResources': [
303+
{'resourceUri': RESOURCE_URI},
304+
{"inlineCode": INLINE_UDF_CODE}]}
305+
self.assertEqual(req['data'], SENT)
306+
self._verifyResourceProperties(query, RESOURCE)
307+
308+
def test_run_w_bad_udfs(self):
309+
RESOURCE = self._makeResource(complete=False)
310+
conn = _Connection(RESOURCE)
311+
client = _Client(project=self.PROJECT, connection=conn)
312+
query = self._makeOne(self.QUERY, client)
313+
314+
with self.assertRaises(ValueError):
315+
query.udf_resources = ["foo"]
316+
self.assertEqual(query.udf_resources, [])
317+
236318
def test_fetch_data_query_not_yet_run(self):
237319
conn = _Connection()
238320
client = _Client(project=self.PROJECT, connection=conn)

‎tox.ini‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[tox]
22
envlist =
3-
py27,py34,py35,cover,docs,lint
3+
py26,py27,py34,py35,cover,docs,lint
44

55
[testing]
66
deps =
@@ -43,6 +43,7 @@ deps =
4343
basepython =
4444
python2.6
4545
deps =
46+
ordereddict
4647
{[testing]deps}
4748
setenv =
4849
PYTHONPATH = {toxinidir}/_testing

0 commit comments

Comments
 (0)