Merge pull request tylertreat#98 from ruxandraburtica/97-all-tables-for-dataset

tylertreat · tylertreat · commit 0f91d2f0b56b · 2016-05-16T11:59:46.000-05:00
Added get_all_tables method
diff --git a/bigquery/client.py b/bigquery/client.py
@@ -1249,8 +1249,31 @@ def push_rows(self, dataset, table, rows, insert_id_key=None,
                     }]
                 }
 
+    def get_all_tables(self, dataset_id):
+        """Retrieve a list of tables for the dataset.
+
+        Parameters
+        ----------
+        dataset_id : str
+            The dataset to retrieve table data for.
+
+        Returns
+        -------
+        A ``list`` with all table names
+        """
+        tables_data = self._get_all_tables_for_dataset(dataset_id)
+
+        tables = []
+        for table in tables_data['tables']:
+            table_name = table.get('tableReference', {}).get('tableId')
+            if table_name:
+                tables.append(table_name)
+        return tables
+
     def _get_all_tables(self, dataset_id, cache=False):
-        """Retrieve a list of all tables for the dataset.
+        """Retrieve the list of tables for dataset, that respect the formats:
+            * appid_YYYY_MM
+            * YYYY_MM_appid
 
         Parameters
         ----------
@@ -1272,23 +1295,39 @@ def _get_all_tables(self, dataset_id, cache=False):
                 do_fetch = False
 
         if do_fetch:
-            result = self.bigquery.tables().list(
-                projectId=self.project_id,
-                datasetId=dataset_id).execute()
-
-            page_token = result.get('nextPageToken')
-            while page_token:
-                res = self.bigquery.tables().list(
-                    projectId=self.project_id,
-                    datasetId=dataset_id,
-                    pageToken=page_token
-                ).execute()
-                page_token = res.get('nextPageToken')
-                result['tables'] += res.get('tables', [])
+            result = self._get_all_tables_for_dataset(dataset_id)
             self.cache[dataset_id] = (datetime.now(), result)
 
         return self._parse_table_list_response(result)
 
+    def _get_all_tables_for_dataset(self, dataset_id):
+        """Retrieve a list of all tables for the dataset.
+
+        Parameters
+        ----------
+        dataset_id : str
+            The dataset to retrieve table names for
+
+        Returns
+        -------
+        dict
+            A ``dict`` containing tables key with all tables
+        """
+        result = self.bigquery.tables().list(
+            projectId=self.project_id,
+            datasetId=dataset_id).execute()
+
+        page_token = result.get('nextPageToken')
+        while page_token:
+            res = self.bigquery.tables().list(
+                projectId=self.project_id,
+                datasetId=dataset_id,
+                pageToken=page_token
+            ).execute()
+            page_token = res.get('nextPageToken')
+            result['tables'] += res.get('tables', [])
+        return result
+
     def _parse_table_list_response(self, list_response):
         """Parse the response received from calling list on tables.
 
diff --git a/bigquery/tests/test_client.py b/bigquery/tests/test_client.py
@@ -1290,12 +1290,21 @@ def test_not_inside_range(self):
                 "tableId": "appspot_6_2013_06"
             }
         },
+        {
+            "kind": "bigquery#table",
+            "id": "project:dataset.table_not_matching_naming",
+            "tableReference": {
+                "projectId": "project",
+                "datasetId": "dataset",
+                "tableId": "table_not_matching_naming"
+            }
+        },
         {
             "kind": "bigquery#table",
             "id": "bad table data"
-        }
+        },
     ],
-    "totalItems": 8
+    "totalItems": 9
 }
 
 
@@ -2191,7 +2200,7 @@ def test_request_data_with_options(self):
 
 class TestGetAllTables(unittest.TestCase):
 
-    def test_get_tables(self):
+    def test_get_all_tables(self):
         """Ensure get_all_tables fetches table names from BigQuery."""
 
         mock_execute = mock.Mock()
@@ -2205,6 +2214,29 @@ def test_get_tables(self):
 
         bq = client.BigQueryClient(mock_bq_service, 'project')
 
+        expected_result = [
+            '2013_05_appspot', '2013_06_appspot_1', '2013_06_appspot_2',
+            '2013_06_appspot_3', '2013_06_appspot_4', '2013_06_appspot_5',
+            'appspot_6_2013_06', 'table_not_matching_naming'
+        ]
+
+        tables = bq.get_all_tables('dataset')
+        self.assertEquals(expected_result, tables)
+
+    def test_get_tables(self):
+        """Ensure _get_all_tables fetches table names from BigQuery."""
+
+        mock_execute = mock.Mock()
+        mock_execute.execute.return_value = FULL_TABLE_LIST_RESPONSE
+
+        mock_tables = mock.Mock()
+        mock_tables.list.return_value = mock_execute
+
+        mock_bq_service = mock.Mock()
+        mock_bq_service.tables.return_value = mock_tables
+
+        bq = client.BigQueryClient(mock_bq_service, 'project')
+
         expected_result = {
             'appspot-3': {'2013_06_appspot_3': 1370044800},
             'appspot-2': {'2013_06_appspot_2': 1370044800},
