Merge pull request tylertreat#107 from cbhimavarapu-datalicious/master

tylertreat · web-flow · commit 6a3276eacb2f · 2016-08-03T09:07:01.000-05:00
Allowing Write query result to table to accept maximumBillingTire argument
diff --git a/bigquery/client.py b/bigquery/client.py
@@ -1048,7 +1048,8 @@ def write_to_table(
             priority=None,
             create_disposition=None,
             write_disposition=None,
-            use_legacy_sql=None
+            use_legacy_sql=None,
+            maximum_billing_tier=None
     ):
         """
         Write query result to table. If dataset or table is not provided,
@@ -1077,9 +1078,15 @@ def write_to_table(
             One of the JOB_CREATE_* constants
         write_disposition : str, optional
             One of the JOB_WRITE_* constants
-        use_legacy_sql:
-            If False, the query will use BigQuery's standard SQL (https://cloud.google.com/bigquery/sql-reference/)
-
+        use_legacy_sql: bool, optional
+            If False, the query will use BigQuery's standard SQL
+            (https://cloud.google.com/bigquery/sql-reference/)
+        maximum_billing_tier : integer, optional
+            Limits the billing tier for this job. Queries that have resource
+            usage beyond this tier will fail (without incurring a charge). If
+            unspecified, this will be set to your project default. For more
+            information,
+            see https://cloud.google.com/bigquery/pricing#high-compute
 
         Returns
         -------
@@ -1106,6 +1113,9 @@ def write_to_table(
         if allow_large_results is not None:
             configuration['allowLargeResults'] = allow_large_results
 
+        if maximum_billing_tier is not None:
+            configuration['maximumBillingTier'] = maximum_billing_tier
+
         if use_query_cache is not None:
             configuration['useQueryCache'] = use_query_cache
 
diff --git a/bigquery/tests/test_client.py b/bigquery/tests/test_client.py
@@ -1117,6 +1117,7 @@ def setUp(self):
         self.project_id = 'project'
         self.dataset_id = 'dataset'
         self.table_id = 'table'
+        self.maximum_billing_tier = 1000
         self.external_udf_uris = ['gs://bucket/external_udf.js']
         self.use_query_cache = False
         self.priority = "INTERACTIVE"
@@ -1162,6 +1163,44 @@ def test_write(self):
 
         self.assertEqual(result, expected_result)
 
+    def test_write_maxbilltier(self):
+        """ Ensure that write is working when maximumBillingTier is set"""
+        expected_result = {
+            'status': {'state': u'RUNNING'},
+        }
+
+        body = {
+            "configuration": {
+                "query": {
+                    "destinationTable": {
+                        "projectId": self.project_id,
+                        "datasetId": self.dataset_id,
+                        "tableId": self.table_id
+                    },
+                    "query": self.query,
+                    "userDefinedFunctionResources": [{
+                        "resourceUri": self.external_udf_uris[0]
+                    }],
+                    "useQueryCache": self.use_query_cache,
+                    "priority": self.priority,
+                    "maximumBillingTier": self.maximum_billing_tier
+                }
+            }
+        }
+
+        self.mock_api.jobs().insert().execute.return_value = expected_result
+        result = self.client.write_to_table(
+            self.query, self.dataset_id, self.table_id, priority=self.priority,
+            external_udf_uris=self.external_udf_uris, use_query_cache=False,
+            maximum_billing_tier=self.maximum_billing_tier)
+
+        self.mock_api.jobs().insert.assert_called_with(
+            projectId=self.project_id,
+            body=body
+        )
+
+        self.assertEqual(result, expected_result)
+
     def test_write_http_error(self):
         """ Test write with http error"""
         expected_result = {
