Merge pull request #2029 from dwmclary/add_udf_support.

Add UDFResource support for QueryJob / Query.

Author: tseaver

gcloud/bigquery/job.py

@@ -27,6 +27,58 @@
 from gcloud.bigquery._helpers import _TypedProperty
 
 
+class UDFResource(object):
+    """Describe a single user-defined function (UDF) resource.
+    :type udf_type: str
+    :param udf_type: the type of the resource ('inlineCode' or 'resourceUri')
+
+    :type value: str
+    :param value: the inline code or resource URI
+
+    See
+    https://cloud.google.com/bigquery/user-defined-functions#api
+    """
+    def __init__(self, udf_type, value):
+        self.udf_type = udf_type
+        self.value = value
+
+    def __eq__(self, other):
+        return(
+            self.udf_type == other.udf_type and
+            self.value == other.value)
+
+
+def _build_udf_resources(resources):
+    """
+    :type resources: sequence of :class:`UDFResource`
+    :param resources: fields to be appended
+
+    :rtype: mapping
+    :returns: a mapping describing userDefinedFunctionResources for the query.
+    """
+    udfs = []
+    for resource in resources:
+        udf = {resource.udf_type: resource.value}
+        udfs.append(udf)
+    return udfs
+
+
+class UDFResourcesProperty(object):
+    """Custom property type for :class:`QueryJob` / :class:`.query.Query`."""
+
+    def __get__(self, instance, owner):
+        """Descriptor protocal:  accesstor"""
+        if instance is None:
+            return self
+        return list(instance._udf_resources)
+
+    def __set__(self, instance, value):
+        """Descriptor protocal:  mutator"""
+        if not all(isinstance(u, UDFResource) for u in value):
+            raise ValueError("udf items must be UDFResource")
+        instance._udf_resources = tuple(value)
+
+
 class Compression(_EnumProperty):
     """Pseudo-enum for ``compression`` properties."""
     GZIP = 'GZIP'
@@ -888,12 +940,19 @@ class QueryJob(_AsyncJob):
     :type client: :class:`gcloud.bigquery.client.Client`
     :param client: A client which holds credentials and project configuration
                    for the dataset (which requires a project).
+
+    :type udf_resources: tuple
+    :param udf_resources: An iterable of
+                        :class:`gcloud.bigquery.job.UDFResource`
+                        (empty by default)
     """
     _JOB_TYPE = 'query'
+    _UDF_KEY = 'userDefinedFunctionResources'
 
-    def __init__(self, name, query, client):
+    def __init__(self, name, query, client, udf_resources=()):
         super(QueryJob, self).__init__(name, client)
         self.query = query
+        self.udf_resources = udf_resources
         self._configuration = _AsyncQueryConfiguration()
 
     allow_large_results = _TypedProperty('allow_large_results', bool)
@@ -926,6 +985,8 @@ def __init__(self, name, query, client):
     https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.query.priority
     """
 
+    udf_resources = UDFResourcesProperty()
+
     use_query_cache = _TypedProperty('use_query_cache', bool)
     """See:
     https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.query.useQueryCache
@@ -979,6 +1040,9 @@ def _populate_config_resource(self, configuration):
             configuration['useLegacySql'] = self.use_legacy_sql
         if self.write_disposition is not None:
             configuration['writeDisposition'] = self.write_disposition
+        if len(self._udf_resources) > 0:
+            configuration[self._UDF_KEY] = _build_udf_resources(
+                self._udf_resources)
 
     def _build_resource(self):
         """Generate a resource for :meth:`begin`."""

gcloud/bigquery/query.py

@@ -20,6 +20,8 @@
 from gcloud.bigquery._helpers import _rows_from_json
 from gcloud.bigquery.dataset import Dataset
 from gcloud.bigquery.job import QueryJob
+from gcloud.bigquery.job import UDFResourcesProperty
+from gcloud.bigquery.job import _build_udf_resources
 from gcloud.bigquery.table import _parse_schema_resource
 
 
@@ -46,12 +48,21 @@ class QueryResults(object):
     :type client: :class:`gcloud.bigquery.client.Client`
     :param client: A client which holds credentials and project configuration
                    for the dataset (which requires a project).
+
+    :type udf_resources: tuple
+    :param udf_resources: An iterable of
+                        :class:`gcloud.bigquery.job.UDFResource`
+                        (empty by default)
     """
-    def __init__(self, query, client):
+
+    _UDF_KEY = 'userDefinedFunctionResources'
+
+    def __init__(self, query, client, udf_resources=()):
         self._client = client
         self._properties = {}
         self.query = query
         self._configuration = _SyncQueryConfiguration()
+        self.udf_resources = udf_resources
         self._job = None
 
     @property
@@ -229,6 +240,8 @@ def schema(self):
     https://cloud.google.com/bigquery/docs/reference/v2/jobs/query#timeoutMs
     """
 
+    udf_resources = UDFResourcesProperty()
+
     use_query_cache = _TypedProperty('use_query_cache', bool)
     """See:
     https://cloud.google.com/bigquery/docs/reference/v2/jobs/query#useQueryCache
@@ -277,6 +290,9 @@ def _build_resource(self):
         if self.dry_run is not None:
             resource['dryRun'] = self.dry_run
 
+        if len(self._udf_resources) > 0:
+            resource[self._UDF_KEY] = _build_udf_resources(self._udf_resources)
+
         return resource
 
     def run(self, client=None):

gcloud/bigquery/test_job.py

@@ -15,6 +15,76 @@
 import unittest2
 
 
+class Test_UDFResourcesProperty(unittest2.TestCase):
+
+    def _getTargetClass(self):
+        from gcloud.bigquery.job import UDFResourcesProperty
+        return UDFResourcesProperty
+
+    def _makeOne(self, *args, **kw):
+        return self._getTargetClass()(*args, **kw)
+
+    def _descriptor_and_klass(self):
+        descriptor = self._makeOne()
+
+        class _Test(object):
+            _udf_resources = ()
+            udf_resources = descriptor
+
+        return descriptor, _Test
+
+    def test_class_getter(self):
+        descriptor, klass = self._descriptor_and_klass()
+        self.assertTrue(klass.udf_resources is descriptor)
+
+    def test_instance_getter_empty(self):
+        _, klass = self._descriptor_and_klass()
+        instance = klass()
+        self.assertEqual(instance.udf_resources, [])
+
+    def test_instance_getter_w_non_empty_list(self):
+        from gcloud.bigquery.job import UDFResource
+        RESOURCE_URI = 'gs://some-bucket/js/lib.js'
+        udf_resources = [UDFResource("resourceUri", RESOURCE_URI)]
+        _, klass = self._descriptor_and_klass()
+        instance = klass()
+        instance._udf_resources = tuple(udf_resources)
+
+        self.assertEqual(instance.udf_resources, udf_resources)
+
+    def test_instance_setter_w_empty_list(self):
+        from gcloud.bigquery.job import UDFResource
+        RESOURCE_URI = 'gs://some-bucket/js/lib.js'
+        udf_resources = [UDFResource("resourceUri", RESOURCE_URI)]
+        _, klass = self._descriptor_and_klass()
+        instance = klass()
+        instance._udf_resources = udf_resources
+
+        instance.udf_resources = []
+
+        self.assertEqual(instance.udf_resources, [])
+
+    def test_instance_setter_w_valid_udf(self):
+        from gcloud.bigquery.job import UDFResource
+        RESOURCE_URI = 'gs://some-bucket/js/lib.js'
+        udf_resources = [UDFResource("resourceUri", RESOURCE_URI)]
+        _, klass = self._descriptor_and_klass()
+        instance = klass()
+
+        instance.udf_resources = udf_resources
+
+        self.assertEqual(instance.udf_resources, udf_resources)
+
+    def test_instance_setter_w_bad_udfs(self):
+        _, klass = self._descriptor_and_klass()
+        instance = klass()
+
+        with self.assertRaises(ValueError):
+            instance.udf_resources = ["foo"]
+
+        self.assertEqual(instance.udf_resources, [])
+
+
 class _Base(object):
     PROJECT = 'project'
     SOURCE1 = 'http://example.com/source1.csv'
@@ -1384,7 +1454,7 @@ def test_begin_w_bound_client(self):
         job = self._makeOne(self.JOB_NAME, self.QUERY, client)
 
         job.begin()
-
+        self.assertEqual(job.udf_resources, [])
         self.assertEqual(len(conn._requested), 1)
         req = conn._requested[0]
         self.assertEqual(req['method'], 'POST')
@@ -1396,7 +1466,7 @@ def test_begin_w_bound_client(self):
             },
             'configuration': {
                 'query': {
-                    'query': self.QUERY,
+                    'query': self.QUERY
                 },
             },
         }
@@ -1468,6 +1538,47 @@ def test_begin_w_alternate_client(self):
         self.assertEqual(req['data'], SENT)
         self._verifyResourceProperties(job, RESOURCE)
 
+    def test_begin_w_bound_client_and_udf(self):
+        from gcloud.bigquery.job import UDFResource
+        RESOURCE_URI = 'gs://some-bucket/js/lib.js'
+        PATH = 'projects/%s/jobs' % self.PROJECT
+        RESOURCE = self._makeResource()
+        # Ensure None for missing server-set props
+        del RESOURCE['statistics']['creationTime']
+        del RESOURCE['etag']
+        del RESOURCE['selfLink']
+        del RESOURCE['user_email']
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        job = self._makeOne(self.JOB_NAME, self.QUERY, client,
+                            udf_resources=[
+                                UDFResource("resourceUri", RESOURCE_URI)
+                            ])
+
+        job.begin()
+
+        self.assertEqual(len(conn._requested), 1)
+        req = conn._requested[0]
+        self.assertEqual(req['method'], 'POST')
+        self.assertEqual(req['path'], '/%s' % PATH)
+        self.assertEqual(job.udf_resources,
+                         [UDFResource("resourceUri", RESOURCE_URI)])
+        SENT = {
+            'jobReference': {
+                'projectId': self.PROJECT,
+                'jobId': self.JOB_NAME,
+            },
+            'configuration': {
+                'query': {
+                    'query': self.QUERY,
+                    'userDefinedFunctionResources':
+                        [{'resourceUri': RESOURCE_URI}]
+                },
+            },
+        }
+        self.assertEqual(req['data'], SENT)
+        self._verifyResourceProperties(job, RESOURCE)
+
     def test_exists_miss_w_bound_client(self):
         PATH = 'projects/%s/jobs/%s' % (self.PROJECT, self.JOB_NAME)
         conn = _Connection()

gcloud/bigquery/test_query.py

@@ -181,7 +181,7 @@ def test_run_w_bound_client(self):
         conn = _Connection(RESOURCE)
         client = _Client(project=self.PROJECT, connection=conn)
         query = self._makeOne(self.QUERY, client)
-
+        self.assertEqual(query.udf_resources, [])
         query.run()
 
         self.assertEqual(len(conn._requested), 1)
@@ -233,6 +233,78 @@ def test_run_w_alternate_client(self):
         self.assertEqual(req['data'], SENT)
         self._verifyResourceProperties(query, RESOURCE)
 
+    def test_run_w_inline_udf(self):
+        from gcloud.bigquery.query import UDFResource
+        INLINE_UDF_CODE = 'var someCode = "here";'
+        PATH = 'projects/%s/queries' % self.PROJECT
+        RESOURCE = self._makeResource(complete=False)
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        query = self._makeOne(self.QUERY, client)
+        query.udf_resources = [UDFResource("inlineCode", INLINE_UDF_CODE)]
+
+        query.run()
+
+        self.assertEqual(len(conn._requested), 1)
+        req = conn._requested[0]
+        self.assertEqual(req['method'], 'POST')
+        self.assertEqual(req['path'], '/%s' % PATH)
+        SENT = {'query': self.QUERY,
+                'userDefinedFunctionResources':
+                [{'inlineCode': INLINE_UDF_CODE}]}
+        self.assertEqual(req['data'], SENT)
+        self._verifyResourceProperties(query, RESOURCE)
+
+    def test_run_w_udf_resource_uri(self):
+        from gcloud.bigquery.job import UDFResource
+        RESOURCE_URI = 'gs://some-bucket/js/lib.js'
+        PATH = 'projects/%s/queries' % self.PROJECT
+        RESOURCE = self._makeResource(complete=False)
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        query = self._makeOne(self.QUERY, client)
+        query.udf_resources = [UDFResource("resourceUri", RESOURCE_URI)]
+
+        query.run()
+
+        self.assertEqual(len(conn._requested), 1)
+        req = conn._requested[0]
+        self.assertEqual(req['method'], 'POST')
+        self.assertEqual(req['path'], '/%s' % PATH)
+        SENT = {'query': self.QUERY,
+                'userDefinedFunctionResources':
+                [{'resourceUri': RESOURCE_URI}]}
+        self.assertEqual(req['data'], SENT)
+        self._verifyResourceProperties(query, RESOURCE)
+
+    def test_run_w_mixed_udfs(self):
+        from gcloud.bigquery.job import UDFResource
+        RESOURCE_URI = 'gs://some-bucket/js/lib.js'
+        INLINE_UDF_CODE = 'var someCode = "here";'
+        PATH = 'projects/%s/queries' % self.PROJECT
+        RESOURCE = self._makeResource(complete=False)
+        conn = _Connection(RESOURCE)
+        client = _Client(project=self.PROJECT, connection=conn)
+        query = self._makeOne(self.QUERY, client)
+        query.udf_resources = [UDFResource("resourceUri", RESOURCE_URI),
+                               UDFResource("inlineCode", INLINE_UDF_CODE)]
+
+        query.run()
+
+        self.assertEqual(len(conn._requested), 1)
+        req = conn._requested[0]
+        self.assertEqual(req['method'], 'POST')
+        self.assertEqual(req['path'], '/%s' % PATH)
+        self.assertEqual(query.udf_resources,
+                         [UDFResource("resourceUri", RESOURCE_URI),
+                          UDFResource("inlineCode", INLINE_UDF_CODE)])
+        SENT = {'query': self.QUERY,
+                'userDefinedFunctionResources': [
+                    {'resourceUri': RESOURCE_URI},
+                    {"inlineCode": INLINE_UDF_CODE}]}
+        self.assertEqual(req['data'], SENT)
+        self._verifyResourceProperties(query, RESOURCE)
+
     def test_fetch_data_query_not_yet_run(self):
         conn = _Connection()
         client = _Client(project=self.PROJECT, connection=conn)

tox.ini

@@ -44,6 +44,8 @@ basepython =
     python2.6
 deps =
     {[testing]deps}
+#   ordereddict needed for google.protobuf, which doesn't declare it.
+    ordereddict
 setenv =
     PYTHONPATH = {toxinidir}/_testing