Merge pull request #2276 from telligent-data/bigquery-add-billing-tie…

…r-bytes-billed Bigquery: Adding configuration properties for maximumBillingTier and maximumBytesBilled for QueryJobs
googleapis · Sep 8, 2016 · dc6e6d5 · dc6e6d5
2 parents a8338fd + a57a2b3
commit dc6e6d5
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 0 deletions.
diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py
@@ -933,6 +933,8 @@ class _AsyncQueryConfiguration(object):
     _use_query_cache = None
     _use_legacy_sql = None
     _write_disposition = None
+    _maximum_billing_tier = None
+    _maximum_bytes_billed = None
 
 
 class QueryJob(_AsyncJob):
@@ -1010,6 +1012,16 @@ def __init__(self, name, query, client, udf_resources=()):
     https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.query.writeDisposition
     """
 
+    maximum_billing_tier = _TypedProperty('maximum_billing_tier', int)
+    """See:
+    https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.query.maximumBillingTier
+    """
+
+    maximum_bytes_billed = _TypedProperty('maximum_bytes_billed', int)
+    """See:
+    https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.query.maximumBytesBilled
+    """
+
     def _destination_table_resource(self):
         """Create a JSON resource for the destination table.
 
@@ -1047,6 +1059,10 @@ 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 self.maximum_billing_tier is not None:
+            configuration['maximumBillingTier'] = self.maximum_billing_tier
+        if self.maximum_bytes_billed is not None:
+            configuration['maximumBytesBilled'] = self.maximum_bytes_billed
         if len(self._udf_resources) > 0:
             configuration[self._UDF_KEY] = _build_udf_resources(
                 self._udf_resources)

diff --git a/unit_tests/bigquery/test_job.py b/unit_tests/bigquery/test_job.py
@@ -1324,11 +1324,24 @@ def _verifyBooleanResourceProperties(self, job, config):
         else:
             self.assertTrue(job.use_legacy_sql is None)
 
+    def _verifyIntegerResourceProperties(self, job, config):
+        if 'maximumBillingTier' in config:
+            self.assertEqual(job.maximum_billing_tier,
+                             config['maximumBillingTier'])
+        else:
+            self.assertTrue(job.maximum_billing_tier is None)
+        if 'maximumBytesBilled' in config:
+            self.assertEqual(job.maximum_bytes_billed,
+                             config['maximumBytesBilled'])
+        else:
+            self.assertTrue(job.maximum_bytes_billed is None)
+
     def _verifyResourceProperties(self, job, resource):
         self._verifyReadonlyResourceProperties(job, resource)
 
         config = resource.get('configuration', {}).get('query')
         self._verifyBooleanResourceProperties(job, config)
+        self._verifyIntegerResourceProperties(job, config)
 
         if 'createDisposition' in config:
             self.assertEqual(job.create_disposition,
@@ -1387,6 +1400,8 @@ def test_ctor(self):
         self.assertTrue(job.use_query_cache is None)
         self.assertTrue(job.use_legacy_sql is None)
         self.assertTrue(job.write_disposition is None)
+        self.assertTrue(job.maximum_billing_tier is None)
+        self.assertTrue(job.maximum_bytes_billed is None)
 
     def test_from_api_repr_missing_identity(self):
         self._setUpConstants()
@@ -1498,6 +1513,8 @@ def test_begin_w_alternate_client(self):
             'useQueryCache': True,
             'useLegacySql': True,
             'writeDisposition': 'WRITE_TRUNCATE',
+            'maximumBillingTier': 4,
+            'maximumBytesBilled': 123456
         }
         RESOURCE['configuration']['query'] = QUERY_CONFIGURATION
         conn1 = _Connection()
@@ -1518,6 +1535,8 @@ def test_begin_w_alternate_client(self):
         job.use_query_cache = True
         job.use_legacy_sql = True
         job.write_disposition = 'WRITE_TRUNCATE'
+        job.maximum_billing_tier = 4
+        job.maximum_bytes_billed = 123456
 
         job.begin(client=client2)