Merge pull request tylertreat#113 from yu-sa/master

tylertreat · web-flow · commit e4df25336372 · 2016-10-04T09:23:38.000-05:00
Add the time partitioning
diff --git a/bigquery/client.py b/bigquery/client.py
@@ -520,7 +520,8 @@ def get_table(self, dataset, table):
 
         return table
 
-    def create_table(self, dataset, table, schema, expiration_time=None):
+    def create_table(self, dataset, table, schema,
+                     expiration_time=None, time_partitioning=False):
         """Create a new table in the dataset.
 
         Parameters
@@ -533,6 +534,8 @@ def create_table(self, dataset, table, schema, expiration_time=None):
             The table schema
         expiration_time : float, optional
             The expiry time in milliseconds since the epoch.
+        time_partitioning : bool, optional
+            Create a time partitioning.
 
         Returns
         -------
@@ -553,6 +556,9 @@ def create_table(self, dataset, table, schema, expiration_time=None):
         if expiration_time is not None:
             body['expirationTime'] = expiration_time
 
+        if time_partitioning:
+            body['timePartitioning'] = "DAY"
+
         try:
             table = self.bigquery.tables().insert(
                 projectId=self.project_id,
diff --git a/bigquery/tests/test_client.py b/bigquery/tests/test_client.py
@@ -1675,6 +1675,7 @@ def setUp(self):
                 'datasetId': self.dataset}
         }
         self.expiration_time = 1437513693000
+        self.time_partitioning = True
 
     def test_table_create_failed(self):
         """Ensure that if creating the table fails, False is returned,
@@ -1748,6 +1749,27 @@ def test_table_create_body_with_expiration_time(self):
 
         self.mock_tables.insert.return_value.execute.assert_called_with()
 
+    def test_table_create_body_with_time_partitioning(self):
+        """Ensure that if time_partitioning has specified,
+        it passed to the body."""
+
+        self.mock_tables.insert.return_value.execute.side_effect = [{
+            'status': 'foo'}, {'status': 'bar'}]
+
+        self.client.create_table(self.dataset, self.table,
+                                 self.schema,
+                                 time_partitioning=self.time_partitioning)
+
+        body = self.body.copy()
+        body.update({
+            'timePartitioning': "DAY"
+        })
+
+        self.mock_tables.insert.assert_called_with(
+            projectId=self.project, datasetId=self.dataset, body=body)
+
+        self.mock_tables.insert.return_value.execute.assert_called_with()
+
 
 class TestUpdateTable(unittest.TestCase):
 
