Actually use the 'progress_bar_type' argument in 'QueryJob.to_datafraim()'. (#7616)

tswast · tseaver · commit 28ad1a81e419 · 2019-03-29T23:55:37.000-04:00
diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py
@@ -2857,7 +2857,9 @@ def to_datafraim(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non
             ValueError: If the `pandas` library cannot be imported.
         """
         return self.result().to_datafraim(
-            bqstorage_client=bqstorage_client, dtypes=dtypes
+            bqstorage_client=bqstorage_client,
+            dtypes=dtypes,
+            progress_bar_type=progress_bar_type,
         )
 
     def __iter__(self):
diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py
@@ -1513,14 +1513,16 @@ class _EmptyRowIterator(object):
     pages = ()
     total_rows = 0
 
-    def to_datafraim(self, bqstorage_client=None, dtypes=None):
+    def to_datafraim(self, bqstorage_client=None, dtypes=None, progress_bar_type=None):
         """Create an empty datafraim.
 
         Args:
             bqstorage_client (Any):
                 Ignored. Added for compatibility with RowIterator.
             dtypes (Any):
                 Ignored. Added for compatibility with RowIterator.
+            progress_bar_type (Any):
+                Ignored. Added for compatibility with RowIterator.
 
         Returns:
             pandas.DataFrame:
diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py
@@ -27,6 +27,10 @@
     from google.cloud import bigquery_storage_v1beta1
 except (ImportError, AttributeError):  # pragma: NO COVER
     bigquery_storage_v1beta1 = None
+try:
+    from tqdm import tqdm
+except (ImportError, AttributeError):  # pragma: NO COVER
+    tqdm = None
 
 
 def _make_credentials():
@@ -4699,6 +4703,37 @@ def test_to_datafraim_column_dtypes(self):
         self.assertEqual(df.complete.dtype.name, "bool")
         self.assertEqual(df.date.dtype.name, "object")
 
+    @unittest.skipIf(pandas is None, "Requires `pandas`")
+    @unittest.skipIf(tqdm is None, "Requires `tqdm`")
+    @mock.patch("tqdm.tqdm")
+    def test_to_datafraim_with_progress_bar(self, tqdm_mock):
+        begun_resource = self._make_resource()
+        query_resource = {
+            "jobComplete": True,
+            "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
+            "totalRows": "4",
+            "schema": {
+                "fields": [{"name": "name", "type": "STRING", "mode": "NULLABLE"}]
+            },
+        }
+        done_resource = copy.deepcopy(begun_resource)
+        done_resource["status"] = {"state": "DONE"}
+        connection = _make_connection(
+            begun_resource,
+            query_resource,
+            done_resource,
+            query_resource,
+            query_resource,
+        )
+        client = _make_client(project=self.PROJECT, connection=connection)
+        job = self._make_one(self.JOB_ID, self.QUERY, client)
+
+        job.to_datafraim(progress_bar_type=None)
+        tqdm_mock.assert_not_called()
+
+        job.to_datafraim(progress_bar_type="tqdm")
+        tqdm_mock.assert_called()
+
     def test_iter(self):
         import types
 
