Merge pull request tylertreat#62 from hagino3000/add_table_expiration_time

tylertreat · tylertreat · commit 0f5a1df4d2d3 · 2015-07-21T22:00:36.000-05:00
Add expiration_time option for create_table
diff --git a/bigquery/client.py b/bigquery/client.py
@@ -384,13 +384,14 @@ def get_table(self, dataset, table):
 
         return table
 
-    def create_table(self, dataset, table, schema):
+    def create_table(self, dataset, table, schema, expiration_time=None):
         """Create a new table in the dataset.
 
         Args:
             dataset: the dataset to create the table in.
             table: the name of table to create.
             schema: table schema dict.
+            expiration_time: the expiry time in milliseconds since the epoch.
 
         Returns:
             bool indicating if the table was successfully created or not,
@@ -406,6 +407,9 @@ def create_table(self, dataset, table, schema):
             }
         }
 
+        if expiration_time is not None:
+            body['expirationTime'] = expiration_time
+
         try:
             table = self.bigquery.tables().insert(
                 projectId=self.project_id,
diff --git a/bigquery/tests/test_client.py b/bigquery/tests/test_client.py
@@ -65,7 +65,7 @@ def test_initialize_readonly(self, mock_build, mock_return_cred):
         mock_return_cred.assert_called_once_with()
         mock_cred.assert_called_once_with(service_account, key,
                                           scope=BIGQUERY_SCOPE_READ_ONLY)
-        mock_cred.authorize.assert_called_once()
+        self.assertTrue(mock_cred.return_value.authorize.called)
         mock_build.assert_called_once_with('bigquery', 'v2', http=mock_http)
         self.assertEquals(mock_bq, bq_client.bigquery)
         self.assertEquals(project_id, bq_client.project_id)
@@ -95,7 +95,7 @@ def test_initialize_read_write(self, mock_build, mock_return_cred):
         mock_return_cred.assert_called_once_with()
         mock_cred.assert_called_once_with(service_account, key,
                                           scope=BIGQUERY_SCOPE)
-        mock_cred.authorize.assert_called_once()
+        self.assertTrue(mock_cred.return_value.authorize.called)
         mock_build.assert_called_once_with('bigquery', 'v2', http=mock_http)
         self.assertEquals(mock_bq, bq_client.bigquery)
         self.assertEquals(project_id, bq_client.project_id)
@@ -130,7 +130,7 @@ def test_initialize_key_file(self, mock_open, mock_build,
         mock_return_cred.assert_called_once_with()
         mock_cred.assert_called_once_with(service_account, key,
                                           scope=BIGQUERY_SCOPE)
-        mock_cred.authorize.assert_called_once()
+        self.assertTrue(mock_cred.return_value.authorize.called)
         mock_build.assert_called_once_with('bigquery', 'v2', http=mock_http)
         self.assertEquals(mock_bq, bq_client.bigquery)
         self.assertEquals(project_id, bq_client.project_id)
@@ -378,7 +378,7 @@ def test_get_response(self):
             projectId=self.project_id, jobId=job_id, startIndex=offset,
             maxResults=limit, pageToken=page_token, timeoutMs=1000)
 
-        mock_query_job.execute.assert_called_once()
+        mock_query_job.execute.assert_called_once_with()
         self.assertEquals(actual, mock_query_reply)
 
 
@@ -1482,6 +1482,7 @@ def setUp(self):
                 'tableId': self.table, 'projectId': self.project,
                 'datasetId': self.dataset}
         }
+        self.expiration_time = 1437513693000
 
     def test_table_create_failed(self):
         """Ensure that if creating the table fails, False is returned,
@@ -1535,6 +1536,26 @@ def test_table_create_success(self):
 
         self.mock_tables.insert.return_value.execute.assert_called_with()
 
+    def test_table_create_body_with_expiration_time(self):
+        """Ensure that if expiration_time 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, self.expiration_time)
+
+        body = self.body.copy()
+        body.update({
+            'expirationTime': self.expiration_time
+        })
+
+        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 TestCreateView(unittest.TestCase):
 
