Merge pull request #4 from tylertreat/get_schema

Add API for retrieving table schema

Author: tylertreat-wf

bigquery/client.py

@@ -136,6 +136,31 @@ def get_query_schema(self, job_id):
 
         return query_reply['schema']['fields']
 
+    def get_table_schema(self, dataset, table):
+        """Return the table schema.
+
+        Args:
+            dataset: the dataset containing the table.
+            table: the table to get the schema for.
+
+        Returns:
+            A list of dicts that represent the table schema. If the table
+            doesn't exist, None is returned.
+        """
+
+        try:
+            result = self.bigquery.tables().get(
+                projectId=self.project_id,
+                tableId=table,
+                datasetId=dataset).execute()
+        except HttpError, e:
+            if int(e.resp['status']) == 404:
+                logger.warn('Table %s.%s does not exist', dataset, table)
+                return None
+            raise
+
+        return result['schema']['fields']
+
     def check_job(self, job_id):
         """Return the state and number of results of a query by job id.
 

bigquery/tests/test_client.py

@@ -115,8 +115,8 @@ def test_query(self):
 
         self.mock_job_collection.query.assert_called_once_with(
             projectId=self.project_id,
-            body={'query': self.query, 'timeoutMs': 10000, 'maxResults': None,
-                  'dryRun': False}
+            body={'query': self.query, 'timeoutMs': 10000, 'dryRun': False,
+                  'maxResults': None}
         )
         self.assertEquals(job_id, 'spiderman')
         self.assertEquals(results, [])
@@ -168,8 +168,8 @@ def test_query_timeout(self):
 
         self.mock_job_collection.query.assert_called_once_with(
             projectId=self.project_id,
-            body={'query': self.query, 'timeoutMs': timeout * 1000,
-                  'maxResults': None, 'dryRun': False}
+            body={'query': self.query, 'timeoutMs': timeout * 1000, 'dryRun':
+                  False, 'maxResults': None}
         )
         self.assertEquals(job_id, 'spiderman')
         self.assertEquals(results, [])
@@ -241,8 +241,8 @@ def test_query_with_results(self):
 
         self.mock_job_collection.query.assert_called_once_with(
             projectId=self.project_id,
-            body={'query': self.query, 'timeoutMs': 10000, 'maxResults': None,
-                  'dryRun': False}
+            body={'query': self.query, 'timeoutMs': 10000, 'dryRun': False,
+                  'maxResults': None}
         )
         self.assertEquals(job_id, 'spiderman')
         self.assertEquals(results, [{'foo': 10}])
@@ -564,6 +564,49 @@ def test_query_incomplete(self, get_query_mock):
         self.assertEquals(result_schema, [])
 
 
+class TestGetTableSchema(unittest.TestCase):
+
+    def setUp(self):
+        self.mock_bq_service = mock.Mock()
+        self.mock_tables = mock.Mock()
+        self.mock_bq_service.tables.return_value = self.mock_tables
+        self.table = 'table'
+        self.project = 'project'
+        self.dataset = 'dataset'
+        self.client = client.BigQueryClient(self.mock_bq_service, self.project)
+
+    def test_table_exists(self):
+        """Ensure that the table schema is returned if the table exists."""
+
+        expected = [
+            {'type': 'FLOAT', 'name': 'foo', 'mode': 'NULLABLE'},
+            {'type': 'INTEGER', 'name': 'bar', 'mode': 'NULLABLE'},
+            {'type': 'INTEGER', 'name': 'baz', 'mode': 'NULLABLE'},
+        ]
+
+        self.mock_tables.get.return_value.execute.return_value = \
+            {'schema': {'fields': expected}}
+
+        self.assertEqual(
+            expected, self.client.get_table_schema(self.dataset, self.table))
+        self.mock_tables.get.assert_called_once_with(
+            projectId=self.project, tableId=self.table, datasetId=self.dataset)
+        self.mock_tables.get.return_value.execute.assert_called_once_with()
+
+    def test_table_does_not_exist(self):
+        """Ensure that None is returned if the table doesn't exist."""
+        from apiclient.errors import HttpError
+
+        self.mock_tables.get.return_value.execute.side_effect = \
+            HttpError({'status': "404"}, '{}')
+
+        self.assertIsNone(
+            self.client.get_table_schema(self.dataset, self.table))
+        self.mock_tables.get.assert_called_once_with(
+            projectId=self.project, tableId=self.table, datasetId=self.dataset)
+        self.mock_tables.get.return_value.execute.assert_called_once_with()
+
+
 @mock.patch('bigquery.client.BigQueryClient._get_query_results')
 class TestGetQueryRows(unittest.TestCase):