added udf support to client.write_table and tests fixed

sshrestha-datalicious · sshrestha-datalicious · commit 8f365c6dcc6b · 2016-01-21T12:11:15.000+11:00
diff --git a/bigquery/client.py b/bigquery/client.py
@@ -861,6 +861,7 @@ def write_to_table(
             query,
             dataset=None,
             table=None,
+            external_udf_uris=[],
             allow_large_results=None,
             use_query_cache=None,
             priority=None,
@@ -919,6 +920,14 @@ 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
+                }
+            )
+
         body = {
             "configuration": {
                 'query': configuration
@@ -1230,7 +1239,7 @@ def _transform_row(self, row, schema):
 
             elif col_dict['type'] == 'BOOLEAN':
                 row_value = row_value in ('True', 'true', 'TRUE')
-            
+
             elif col_dict['type'] == 'TIMESTAMP':
                 row_value = float(row_value)
 
diff --git a/bigquery/tests/test_client.py b/bigquery/tests/test_client.py
@@ -1012,6 +1012,7 @@ def setUp(self):
         self.project_id = 'project'
         self.dataset_id = 'dataset'
         self.table_id = 'table'
+        self.external_udf_uris = ['gs://bucket/external_udf.js']
         self.use_query_cache = False
         self.priority = "INTERACTIVE"
         self.client = client.BigQueryClient(self.mock_api,
@@ -1032,6 +1033,9 @@ def test_write(self):
                         "tableId": self.table_id
                     },
                     "query": self.query,
+                    "userDefinedFunctionResources": [{
+                        "resourceUri": self.external_udf_uris[0]
+                    }],
                     "useQueryCache": self.use_query_cache,
                     "priority": self.priority,
                 }
@@ -1042,6 +1046,7 @@ def test_write(self):
         result = self.client.write_to_table(self.query,
                                             self.dataset_id,
                                             self.table_id,
+                                            external_udf_uris=self.external_udf_uris,
                                             use_query_cache=False,
                                             priority=self.priority)
 
