From b80f7279cd7611a4c87294d0035c071502b5d869 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Fri, 9 Jul 2021 16:14:03 -0500
Subject: [PATCH 01/10] refactor: omit `read_session` with latest
 google-cloud-bigquery-storage

`read_session` is unnecessary as of `google-cloud-bigquery-storage>=2.6.0`.
This will allow us to more loudly deprecate the use of `rows(read_session)`.

Rather than require 2.6.0, version switches will allow us to keep our
requirements range wider. Will want to give this version some time to bake
before making it required.
---
 google/cloud/bigquery/_helpers.py        |  7 +++
 google/cloud/bigquery/_pandas_helpers.py | 10 +++-
 tests/unit/test__pandas_helpers.py       | 65 ++++++++++++++++++++++++
 3 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
index 77054542a..343ef67d9 100644
--- a/google/cloud/bigquery/_helpers.py
+++ b/google/cloud/bigquery/_helpers.py
@@ -42,6 +42,8 @@
 )
 
 _MIN_BQ_STORAGE_VERSION = pkg_resources.parse_version("2.0.0")
+_BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION = pkg_resources.parse_version("2.6.0")
+_IS_BQ_STORAGE_READ_SESSION_OPTIONAL = False
 
 
 def _verify_bq_storage_version():
@@ -54,6 +56,8 @@ def _verify_bq_storage_version():
     in setup.py, the the calling code can use this helper to verify the version
     compatibility at runtime.
     """
+    global _IS_BQ_STORAGE_READ_SESSION_OPTIONAL
+
     from google.cloud import bigquery_storage
 
     installed_version = pkg_resources.parse_version(
@@ -67,6 +71,9 @@ def _verify_bq_storage_version():
         )
         raise LegacyBigQueryStorageError(msg)
 
+    if installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION:
+        _IS_BQ_STORAGE_READ_SESSION_OPTIONAL = True
+
 
 def _not_null(value, field):
     """Check whether 'value' should be coerced to 'field' type."""
diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py
index e93a99eba..cfd28b1b3 100644
--- a/google/cloud/bigquery/_pandas_helpers.py
+++ b/google/cloud/bigquery/_pandas_helpers.py
@@ -41,6 +41,7 @@
     # Having BQ Storage available implies that pyarrow >=1.0.0 is available, too.
     _ARROW_COMPRESSION_SUPPORT = True
 
+from google.cloud.bigquery import _helpers
 from google.cloud.bigquery import schema
 
 
@@ -590,7 +591,14 @@ def _bqstorage_page_to_dataframe(column_names, dtypes, page):
 def _download_table_bqstorage_stream(
     download_state, bqstorage_client, session, stream, worker_queue, page_to_item
 ):
-    rowstream = bqstorage_client.read_rows(stream.name).rows(session)
+    reader = bqstorage_client.read_rows(stream.name)
+
+    # Avoid deprecation warnings for passing in unnecessary read session.
+    # https://github.com/googleapis/python-bigquery-storage/issues/229
+    if _helpers._IS_BQ_STORAGE_READ_SESSION_OPTIONAL:
+        rowstream = reader.rows()
+    else:
+        rowstream = reader.rows(session)
 
     for page in rowstream.pages:
         if download_state.done:
diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py
index 39a3d845b..027cee18e 100644
--- a/tests/unit/test__pandas_helpers.py
+++ b/tests/unit/test__pandas_helpers.py
@@ -39,6 +39,7 @@
 import pytz
 
 from google import api_core
+from google.cloud.bigquery import _helpers
 from google.cloud.bigquery import schema
 from google.cloud.bigquery._pandas_helpers import _BIGNUMERIC_SUPPORT
 
@@ -1271,6 +1272,70 @@ def test_dataframe_to_parquet_dict_sequence_schema(module_under_test):
     assert schema_arg == expected_schema_arg
 
 
+@pytest.mark.skipif(
+    bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`"
+)
+def test__download_table_bqstorage_stream_includes_read_session(
+    monkeypatch, module_under_test
+):
+    import google.cloud.bigquery_storage_v1.reader
+    import google.cloud.bigquery_storage_v1.types
+
+    monkeypatch.setattr(bigquery_storage, "__version__", "2.5.0")
+    bqstorage_client = mock.create_autospec(
+        bigquery_storage.BigQueryReadClient, instance=True
+    )
+    reader = mock.create_autospec(
+        google.cloud.bigquery_storage_v1.reader.ReadRowsStream, instance=True
+    )
+    bqstorage_client.read_rows.return_value = reader
+    session = google.cloud.bigquery_storage_v1.types.ReadSession()
+    _helpers._verify_bq_storage_version()
+
+    module_under_test._download_table_bqstorage_stream(
+        module_under_test._DownloadState(),
+        bqstorage_client,
+        session,
+        google.cloud.bigquery_storage_v1.types.ReadStream(name="test"),
+        queue.Queue(),
+        lambda item: item,
+    )
+
+    reader.rows.assert_called_once_with(session)
+
+
+@pytest.mark.skipif(
+    bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`"
+)
+def test__download_table_bqstorage_stream_omits_read_session(
+    monkeypatch, module_under_test
+):
+    import google.cloud.bigquery_storage_v1.reader
+    import google.cloud.bigquery_storage_v1.types
+
+    monkeypatch.setattr(bigquery_storage, "__version__", "2.6.0")
+    bqstorage_client = mock.create_autospec(
+        bigquery_storage.BigQueryReadClient, instance=True
+    )
+    reader = mock.create_autospec(
+        google.cloud.bigquery_storage_v1.reader.ReadRowsStream, instance=True
+    )
+    bqstorage_client.read_rows.return_value = reader
+    session = google.cloud.bigquery_storage_v1.types.ReadSession()
+    _helpers._verify_bq_storage_version()
+
+    module_under_test._download_table_bqstorage_stream(
+        module_under_test._DownloadState(),
+        bqstorage_client,
+        session,
+        google.cloud.bigquery_storage_v1.types.ReadStream(name="test"),
+        queue.Queue(),
+        lambda item: item,
+    )
+
+    reader.rows.assert_called_once_with()
+
+
 @pytest.mark.parametrize(
     "stream_count,maxsize_kwarg,expected_call_count,expected_maxsize",
     [

From c153afb6389b000d5bc6f0f0028e8199422de097 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Tue, 13 Jul 2021 14:32:12 -0500
Subject: [PATCH 02/10] optimize _verify_bq_storage_version

---
 google/cloud/bigquery/_helpers.py        | 8 +++++++-
 google/cloud/bigquery/_pandas_helpers.py | 1 +
 tests/unit/test__pandas_helpers.py       | 4 ++--
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
index 343ef67d9..28c51f0a0 100644
--- a/google/cloud/bigquery/_helpers.py
+++ b/google/cloud/bigquery/_helpers.py
@@ -44,6 +44,7 @@
 _MIN_BQ_STORAGE_VERSION = pkg_resources.parse_version("2.0.0")
 _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION = pkg_resources.parse_version("2.6.0")
 _IS_BQ_STORAGE_READ_SESSION_OPTIONAL = False
+_HAS_VERIFIED_BQ_STORAGE = False
 
 
 def _verify_bq_storage_version():
@@ -56,7 +57,10 @@ def _verify_bq_storage_version():
     in setup.py, the the calling code can use this helper to verify the version
     compatibility at runtime.
     """
-    global _IS_BQ_STORAGE_READ_SESSION_OPTIONAL
+    global _IS_BQ_STORAGE_READ_SESSION_OPTIONAL, _HAS_VERIFIED_BQ_STORAGE
+
+    if _HAS_VERIFIED_BQ_STORAGE:
+        return
 
     from google.cloud import bigquery_storage
 
@@ -74,6 +78,8 @@ def _verify_bq_storage_version():
     if installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION:
         _IS_BQ_STORAGE_READ_SESSION_OPTIONAL = True
 
+    _HAS_VERIFIED_BQ_STORAGE = True
+
 
 def _not_null(value, field):
     """Check whether 'value' should be coerced to 'field' type."""
diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py
index cb3109333..56883f0a6 100644
--- a/google/cloud/bigquery/_pandas_helpers.py
+++ b/google/cloud/bigquery/_pandas_helpers.py
@@ -595,6 +595,7 @@ def _download_table_bqstorage_stream(
 
     # Avoid deprecation warnings for passing in unnecessary read session.
     # https://github.com/googleapis/python-bigquery-storage/issues/229
+    _helpers._verify_bq_storage_version()
     if _helpers._IS_BQ_STORAGE_READ_SESSION_OPTIONAL:
         rowstream = reader.rows()
     else:
diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py
index 847821817..b31ae8b98 100644
--- a/tests/unit/test__pandas_helpers.py
+++ b/tests/unit/test__pandas_helpers.py
@@ -1321,6 +1321,7 @@ def test__download_table_bqstorage_stream_includes_read_session(
     import google.cloud.bigquery_storage_v1.reader
     import google.cloud.bigquery_storage_v1.types
 
+    monkeypatch.setattr(_helpers, "_HAS_VERIFIED_BQ_STORAGE", False)
     monkeypatch.setattr(bigquery_storage, "__version__", "2.5.0")
     bqstorage_client = mock.create_autospec(
         bigquery_storage.BigQueryReadClient, instance=True
@@ -1330,7 +1331,6 @@ def test__download_table_bqstorage_stream_includes_read_session(
     )
     bqstorage_client.read_rows.return_value = reader
     session = google.cloud.bigquery_storage_v1.types.ReadSession()
-    _helpers._verify_bq_storage_version()
 
     module_under_test._download_table_bqstorage_stream(
         module_under_test._DownloadState(),
@@ -1353,6 +1353,7 @@ def test__download_table_bqstorage_stream_omits_read_session(
     import google.cloud.bigquery_storage_v1.reader
     import google.cloud.bigquery_storage_v1.types
 
+    monkeypatch.setattr(_helpers, "_HAS_VERIFIED_BQ_STORAGE", False)
     monkeypatch.setattr(bigquery_storage, "__version__", "2.6.0")
     bqstorage_client = mock.create_autospec(
         bigquery_storage.BigQueryReadClient, instance=True
@@ -1362,7 +1363,6 @@ def test__download_table_bqstorage_stream_omits_read_session(
     )
     bqstorage_client.read_rows.return_value = reader
     session = google.cloud.bigquery_storage_v1.types.ReadSession()
-    _helpers._verify_bq_storage_version()
 
     module_under_test._download_table_bqstorage_stream(
         module_under_test._DownloadState(),

From bcfa3e05dfa786988b1c9daf9f62a904c3f469c4 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Tue, 13 Jul 2021 14:40:57 -0500
Subject: [PATCH 03/10] fix failing tests due to optimization

---
 tests/unit/test__helpers.py        | 5 +++--
 tests/unit/test__pandas_helpers.py | 5 +++++
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py
index c62947d37..3a97c96fb 100644
--- a/tests/unit/test__helpers.py
+++ b/tests/unit/test__helpers.py
@@ -28,9 +28,10 @@
 @unittest.skipIf(bigquery_storage is None, "Requires `google-cloud-bigquery-storage`")
 class Test_verify_bq_storage_version(unittest.TestCase):
     def _call_fut(self):
-        from google.cloud.bigquery._helpers import _verify_bq_storage_version
+        from google.cloud.bigquery import _helpers
 
-        return _verify_bq_storage_version()
+        _helpers._HAS_VERIFIED_BQ_STORAGE = False
+        return _helpers._verify_bq_storage_version()
 
     def test_raises_no_error_w_recent_bqstorage(self):
         from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py
index b31ae8b98..34a29adbe 100644
--- a/tests/unit/test__pandas_helpers.py
+++ b/tests/unit/test__pandas_helpers.py
@@ -46,6 +46,8 @@
 
 try:
     from google.cloud import bigquery_storage
+
+    _helpers._verify_bq_storage_version()
 except ImportError:  # pragma: NO COVER
     bigquery_storage = None
 
@@ -1347,6 +1349,9 @@ def test__download_table_bqstorage_stream_includes_read_session(
 @pytest.mark.skipif(
     bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`"
 )
+@pytest.mark.skipif(
+    not _helpers._IS_BQ_STORAGE_READ_SESSION_OPTIONAL, reason="read_session is required"
+)
 def test__download_table_bqstorage_stream_omits_read_session(
     monkeypatch, module_under_test
 ):

From d4211b2869e365c415e0ef6ad840e40bb3783f76 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Tue, 13 Jul 2021 14:54:01 -0500
Subject: [PATCH 04/10] fix unit tests

---
 google/cloud/bigquery/_helpers.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
index 28c51f0a0..754902c03 100644
--- a/google/cloud/bigquery/_helpers.py
+++ b/google/cloud/bigquery/_helpers.py
@@ -77,6 +77,10 @@ def _verify_bq_storage_version():
 
     if installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION:
         _IS_BQ_STORAGE_READ_SESSION_OPTIONAL = True
+    else:
+        # Defaults to False, but set it just in case we're mocking out the
+        # version in unit tests.
+        _IS_BQ_STORAGE_READ_SESSION_OPTIONAL = False
 
     _HAS_VERIFIED_BQ_STORAGE = True
 

From 6c18969808ab63523d975fe467e0f06dc83fec0c Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Wed, 14 Jul 2021 15:11:33 -0500
Subject: [PATCH 05/10] create BQStorageVersions class for version comparisons

---
 google/cloud/bigquery/_helpers.py        | 60 ++++++++++++------------
 google/cloud/bigquery/_pandas_helpers.py |  3 +-
 google/cloud/bigquery/client.py          |  4 +-
 google/cloud/bigquery/table.py           |  2 +-
 tests/unit/test__helpers.py              |  6 +--
 tests/unit/test__pandas_helpers.py       |  9 ++--
 tests/unit/test_client.py                |  4 +-
 tests/unit/test_magics.py                |  2 +-
 tests/unit/test_table.py                 |  2 +-
 9 files changed, 46 insertions(+), 46 deletions(-)

diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
index 754902c03..d60b61429 100644
--- a/google/cloud/bigquery/_helpers.py
+++ b/google/cloud/bigquery/_helpers.py
@@ -43,46 +43,46 @@
 
 _MIN_BQ_STORAGE_VERSION = pkg_resources.parse_version("2.0.0")
 _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION = pkg_resources.parse_version("2.6.0")
-_IS_BQ_STORAGE_READ_SESSION_OPTIONAL = False
-_HAS_VERIFIED_BQ_STORAGE = False
 
 
-def _verify_bq_storage_version():
-    """Verify that a recent enough version of BigQuery Storage extra is installed.
+class BQStorageVersions:
+    def __init__(self):
+        self._installed_version = None
 
-    The function assumes that google-cloud-bigquery-storage extra is installed, and
-    should thus be used in places where this assumption holds.
+    @property
+    def installed_version(self):
+        if self._installed_version is None:
+            from google.cloud import bigquery_storage
 
-    Because `pip` can install an outdated version of this extra despite the constraints
-    in setup.py, the the calling code can use this helper to verify the version
-    compatibility at runtime.
-    """
-    global _IS_BQ_STORAGE_READ_SESSION_OPTIONAL, _HAS_VERIFIED_BQ_STORAGE
+            self._installed_version = pkg_resources.parse_version(
+                getattr(bigquery_storage, "__version__", "legacy")
+            )
 
-    if _HAS_VERIFIED_BQ_STORAGE:
-        return
+        return self._installed_version
 
-    from google.cloud import bigquery_storage
+    @property
+    def is_read_session_optional(self):
+        return self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION
 
-    installed_version = pkg_resources.parse_version(
-        getattr(bigquery_storage, "__version__", "legacy")
-    )
+    def verify_version(self):
+        """Verify that a recent enough version of BigQuery Storage extra is installed.
 
-    if installed_version < _MIN_BQ_STORAGE_VERSION:
-        msg = (
-            "Dependency google-cloud-bigquery-storage is outdated, please upgrade "
-            f"it to version >= 2.0.0 (version found: {installed_version})."
-        )
-        raise LegacyBigQueryStorageError(msg)
+        The function assumes that google-cloud-bigquery-storage extra is installed, and
+        should thus be used in places where this assumption holds.
+
+        Because `pip` can install an outdated version of this extra despite the constraints
+        in setup.py, the the calling code can use this helper to verify the version
+        compatibility at runtime.
+        """
+        if self.installed_version < _MIN_BQ_STORAGE_VERSION:
+            msg = (
+                "Dependency google-cloud-bigquery-storage is outdated, please upgrade "
+                f"it to version >= 2.0.0 (version found: {self.installed_version})."
+            )
+            raise LegacyBigQueryStorageError(msg)
 
-    if installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION:
-        _IS_BQ_STORAGE_READ_SESSION_OPTIONAL = True
-    else:
-        # Defaults to False, but set it just in case we're mocking out the
-        # version in unit tests.
-        _IS_BQ_STORAGE_READ_SESSION_OPTIONAL = False
 
-    _HAS_VERIFIED_BQ_STORAGE = True
+BQ_STORAGE_VERSIONS = BQStorageVersions()
 
 
 def _not_null(value, field):
diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py
index 56883f0a6..2ff96da4d 100644
--- a/google/cloud/bigquery/_pandas_helpers.py
+++ b/google/cloud/bigquery/_pandas_helpers.py
@@ -595,8 +595,7 @@ def _download_table_bqstorage_stream(
 
     # Avoid deprecation warnings for passing in unnecessary read session.
     # https://github.com/googleapis/python-bigquery-storage/issues/229
-    _helpers._verify_bq_storage_version()
-    if _helpers._IS_BQ_STORAGE_READ_SESSION_OPTIONAL:
+    if _helpers.BQ_STORAGE_VERSIONS.is_read_session_optional:
         rowstream = reader.rows()
     else:
         rowstream = reader.rows(session)
diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py
index de259abce..8572ba911 100644
--- a/google/cloud/bigquery/client.py
+++ b/google/cloud/bigquery/client.py
@@ -61,7 +61,7 @@
 from google.cloud.bigquery._helpers import _get_sub_prop
 from google.cloud.bigquery._helpers import _record_field_to_json
 from google.cloud.bigquery._helpers import _str_or_none
-from google.cloud.bigquery._helpers import _verify_bq_storage_version
+from google.cloud.bigquery._helpers import BQ_STORAGE_VERSIONS
 from google.cloud.bigquery._helpers import _verify_job_config_type
 from google.cloud.bigquery._http import Connection
 from google.cloud.bigquery import _pandas_helpers
@@ -508,7 +508,7 @@ def _ensure_bqstorage_client(
             return None
 
         try:
-            _verify_bq_storage_version()
+            BQ_STORAGE_VERSIONS.verify_version()
         except LegacyBigQueryStorageError as exc:
             warnings.warn(str(exc))
             return None
diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py
index 765110ae6..2d9c15f50 100644
--- a/google/cloud/bigquery/table.py
+++ b/google/cloud/bigquery/table.py
@@ -1565,7 +1565,7 @@ def _validate_bqstorage(self, bqstorage_client, create_bqstorage_client):
             return False
 
         try:
-            _helpers._verify_bq_storage_version()
+            _helpers.BQ_STORAGE_VERSIONS.verify_version()
         except LegacyBigQueryStorageError as exc:
             warnings.warn(str(exc))
             return False
diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py
index 3a97c96fb..2ce443771 100644
--- a/tests/unit/test__helpers.py
+++ b/tests/unit/test__helpers.py
@@ -26,12 +26,12 @@
 
 
 @unittest.skipIf(bigquery_storage is None, "Requires `google-cloud-bigquery-storage`")
-class Test_verify_bq_storage_version(unittest.TestCase):
+class TestBQStorageVersions(unittest.TestCase):
     def _call_fut(self):
         from google.cloud.bigquery import _helpers
 
-        _helpers._HAS_VERIFIED_BQ_STORAGE = False
-        return _helpers._verify_bq_storage_version()
+        _helpers.BQ_STORAGE_VERSIONS._installed_version = None
+        return _helpers.BQ_STORAGE_VERSIONS.verify_version()
 
     def test_raises_no_error_w_recent_bqstorage(self):
         from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py
index 34a29adbe..1186cf5a8 100644
--- a/tests/unit/test__pandas_helpers.py
+++ b/tests/unit/test__pandas_helpers.py
@@ -47,7 +47,7 @@
 try:
     from google.cloud import bigquery_storage
 
-    _helpers._verify_bq_storage_version()
+    _helpers.BQ_STORAGE_VERSIONS.verify_version()
 except ImportError:  # pragma: NO COVER
     bigquery_storage = None
 
@@ -1323,7 +1323,7 @@ def test__download_table_bqstorage_stream_includes_read_session(
     import google.cloud.bigquery_storage_v1.reader
     import google.cloud.bigquery_storage_v1.types
 
-    monkeypatch.setattr(_helpers, "_HAS_VERIFIED_BQ_STORAGE", False)
+    monkeypatch.setattr(_helpers.BQ_STORAGE_VERSIONS, "_installed_version", None)
     monkeypatch.setattr(bigquery_storage, "__version__", "2.5.0")
     bqstorage_client = mock.create_autospec(
         bigquery_storage.BigQueryReadClient, instance=True
@@ -1350,7 +1350,8 @@ def test__download_table_bqstorage_stream_includes_read_session(
     bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`"
 )
 @pytest.mark.skipif(
-    not _helpers._IS_BQ_STORAGE_READ_SESSION_OPTIONAL, reason="read_session is required"
+    not _helpers.BQ_STORAGE_VERSIONS.is_read_session_optional,
+    reason="read_session is required",
 )
 def test__download_table_bqstorage_stream_omits_read_session(
     monkeypatch, module_under_test
@@ -1358,7 +1359,7 @@ def test__download_table_bqstorage_stream_omits_read_session(
     import google.cloud.bigquery_storage_v1.reader
     import google.cloud.bigquery_storage_v1.types
 
-    monkeypatch.setattr(_helpers, "_HAS_VERIFIED_BQ_STORAGE", False)
+    monkeypatch.setattr(_helpers.BQ_STORAGE_VERSIONS, "_installed_version", None)
     monkeypatch.setattr(bigquery_storage, "__version__", "2.6.0")
     bqstorage_client = mock.create_autospec(
         bigquery_storage.BigQueryReadClient, instance=True
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index 2be8daab6..6b62eb85b 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -663,7 +663,7 @@ def test_ensure_bqstorage_client_obsolete_dependency(self):
         client = self._make_one(project=self.PROJECT, credentials=creds)
 
         patcher = mock.patch(
-            "google.cloud.bigquery.client._verify_bq_storage_version",
+            "google.cloud.bigquery.client.BQ_STORAGE_VERSIONS.verify_version",
             side_effect=LegacyBigQueryStorageError("BQ Storage too old"),
         )
         with patcher, warnings.catch_warnings(record=True) as warned:
@@ -700,7 +700,7 @@ def test_ensure_bqstorage_client_existing_client_check_fails(self):
         mock_storage_client = mock.sentinel.mock_storage_client
 
         patcher = mock.patch(
-            "google.cloud.bigquery.client._verify_bq_storage_version",
+            "google.cloud.bigquery.client.BQ_STORAGE_VERSIONS.verify_version",
             side_effect=LegacyBigQueryStorageError("BQ Storage too old"),
         )
         with patcher, warnings.catch_warnings(record=True) as warned:
diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py
index 5e9bf28a9..d030482cc 100644
--- a/tests/unit/test_magics.py
+++ b/tests/unit/test_magics.py
@@ -368,7 +368,7 @@ def test__make_bqstorage_client_true_obsolete_dependency():
     )
 
     patcher = mock.patch(
-        "google.cloud.bigquery.client._verify_bq_storage_version",
+        "google.cloud.bigquery.client.BQ_STORAGE_VERSIONS.verify_version",
         side_effect=LegacyBigQueryStorageError("BQ Storage too old"),
     )
     with patcher, warnings.catch_warnings(record=True) as warned:
diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py
index b30f16fe0..37650cd27 100644
--- a/tests/unit/test_table.py
+++ b/tests/unit/test_table.py
@@ -1889,7 +1889,7 @@ def test__validate_bqstorage_returns_false_w_warning_if_obsolete_version(self):
         iterator = self._make_one(first_page_response=None)  # not cached
 
         patcher = mock.patch(
-            "google.cloud.bigquery.table._helpers._verify_bq_storage_version",
+            "google.cloud.bigquery.table._helpers.BQ_STORAGE_VERSIONS.verify_version",
             side_effect=LegacyBigQueryStorageError("BQ Storage too old"),
         )
         with patcher, warnings.catch_warnings(record=True) as warned:

From 237cdd16f9162ef338daeefb9901eeb8a6feed77 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Wed, 14 Jul 2021 16:41:50 -0500
Subject: [PATCH 06/10] add type annotations

Also, use packaging directly, since that's all pkg_resources does
https://github.com/pypa/setuptools/blob/a4dbe3457d89cf67ee3aa571fdb149e6eb544e88/pkg_resources/__init__.py\#L112
---
 google/cloud/bigquery/_helpers.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
index d60b61429..9dbc2fe8f 100644
--- a/google/cloud/bigquery/_helpers.py
+++ b/google/cloud/bigquery/_helpers.py
@@ -26,7 +26,7 @@
 from google.cloud._helpers import _RFC3339_MICROS
 from google.cloud._helpers import _RFC3339_NO_FRACTION
 from google.cloud._helpers import _to_bytes
-import pkg_resources
+import packaging.version
 
 from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
 
@@ -41,8 +41,8 @@
     re.VERBOSE,
 )
 
-_MIN_BQ_STORAGE_VERSION = pkg_resources.parse_version("2.0.0")
-_BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION = pkg_resources.parse_version("2.6.0")
+_MIN_BQ_STORAGE_VERSION = packaging.version.Version("2.0.0")
+_BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION = packaging.version.Version("2.6.0")
 
 
 class BQStorageVersions:
@@ -50,18 +50,18 @@ def __init__(self):
         self._installed_version = None
 
     @property
-    def installed_version(self):
+    def installed_version(self) -> packaging.version.Version:
         if self._installed_version is None:
             from google.cloud import bigquery_storage
 
-            self._installed_version = pkg_resources.parse_version(
+            self._installed_version = packaging.version.Version(
                 getattr(bigquery_storage, "__version__", "legacy")
             )
 
         return self._installed_version
 
     @property
-    def is_read_session_optional(self):
+    def is_read_session_optional(self) -> bool:
         return self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION
 
     def verify_version(self):
@@ -73,6 +73,9 @@ def verify_version(self):
         Because `pip` can install an outdated version of this extra despite the constraints
         in setup.py, the the calling code can use this helper to verify the version
         compatibility at runtime.
+
+        Raises:
+            LegacyBigQueryStorageError: If google-cloud-bigquery-storage is outdated.
         """
         if self.installed_version < _MIN_BQ_STORAGE_VERSION:
             msg = (

From 0d4938764b6b2a1103f5dcf44fbcc87fdb26a186 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Wed, 14 Jul 2021 16:55:53 -0500
Subject: [PATCH 07/10] allow legacy versions

---
 google/cloud/bigquery/_helpers.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
index 9dbc2fe8f..b9d3774a7 100644
--- a/google/cloud/bigquery/_helpers.py
+++ b/google/cloud/bigquery/_helpers.py
@@ -19,6 +19,7 @@
 import decimal
 import math
 import re
+from typing import Union
 
 from google.cloud._helpers import UTC
 from google.cloud._helpers import _date_from_iso8601_date
@@ -50,11 +51,13 @@ def __init__(self):
         self._installed_version = None
 
     @property
-    def installed_version(self) -> packaging.version.Version:
+    def installed_version(
+        self,
+    ) -> Union[packaging.version.LegacyVersion, packaging.version.Version]:
         if self._installed_version is None:
             from google.cloud import bigquery_storage
 
-            self._installed_version = packaging.version.Version(
+            self._installed_version = packaging.version.parse(
                 getattr(bigquery_storage, "__version__", "legacy")
             )
 

From c50b08c400407272b5de2f29ec3f366440ae0ff5 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Thu, 15 Jul 2021 09:49:30 -0500
Subject: [PATCH 08/10] fix coverage

---
 google/cloud/bigquery/_helpers.py | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
index b9d3774a7..56e677849 100644
--- a/google/cloud/bigquery/_helpers.py
+++ b/google/cloud/bigquery/_helpers.py
@@ -65,20 +65,26 @@ def installed_version(
 
     @property
     def is_read_session_optional(self) -> bool:
-        return self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION
+        return (
+            self.installed_version is not None
+            and self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION
+        )
 
     def verify_version(self):
-        """Verify that a recent enough version of BigQuery Storage extra is installed.
+        """Verify that a recent enough version of BigQuery Storage extra is
+        installed.
 
-        The function assumes that google-cloud-bigquery-storage extra is installed, and
-        should thus be used in places where this assumption holds.
+        The function assumes that google-cloud-bigquery-storage extra is
+        installed, and should thus be used in places where this assumption
+        holds.
 
-        Because `pip` can install an outdated version of this extra despite the constraints
-        in setup.py, the the calling code can use this helper to verify the version
-        compatibility at runtime.
+        Because `pip` can install an outdated version of this extra despite the
+        constraints in `setup.py`, the calling code can use this helper to
+        verify the version compatibility at runtime.
 
         Raises:
-            LegacyBigQueryStorageError: If google-cloud-bigquery-storage is outdated.
+            LegacyBigQueryStorageError:
+                If the google-cloud-bigquery-storage package is outdated.
         """
         if self.installed_version < _MIN_BQ_STORAGE_VERSION:
             msg = (

From 44ef80b9db7b6adfdac7b090cbd5a292d182d235 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Thu, 15 Jul 2021 09:51:25 -0500
Subject: [PATCH 09/10] fix coverage

---
 google/cloud/bigquery/_helpers.py  | 5 +----
 tests/unit/test__pandas_helpers.py | 8 +++-----
 2 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
index 56e677849..1836a01a1 100644
--- a/google/cloud/bigquery/_helpers.py
+++ b/google/cloud/bigquery/_helpers.py
@@ -65,10 +65,7 @@ def installed_version(
 
     @property
     def is_read_session_optional(self) -> bool:
-        return (
-            self.installed_version is not None
-            and self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION
-        )
+        return self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION
 
     def verify_version(self):
         """Verify that a recent enough version of BigQuery Storage extra is
diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py
index 1186cf5a8..ed90f9a24 100644
--- a/tests/unit/test__pandas_helpers.py
+++ b/tests/unit/test__pandas_helpers.py
@@ -1347,11 +1347,9 @@ def test__download_table_bqstorage_stream_includes_read_session(
 
 
 @pytest.mark.skipif(
-    bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`"
-)
-@pytest.mark.skipif(
-    not _helpers.BQ_STORAGE_VERSIONS.is_read_session_optional,
-    reason="read_session is required",
+    bigquery_storage is None
+    or not _helpers.BQ_STORAGE_VERSIONS.is_read_session_optional,
+    reason="Requires `google-cloud-bigquery-storage` >= 2.6.0",
 )
 def test__download_table_bqstorage_stream_omits_read_session(
     monkeypatch, module_under_test

From 9d161c5500c5320a6bf516ac14136b7974dd7c31 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Thu, 15 Jul 2021 14:29:49 -0500
Subject: [PATCH 10/10] add tests for version helpers

---
 google/cloud/bigquery/_helpers.py  | 18 ++++++++++++-----
 tests/unit/test__helpers.py        | 32 +++++++++++++++++++++++++++++-
 tests/unit/test__pandas_helpers.py |  4 ++--
 3 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
index 1836a01a1..bf0f80e22 100644
--- a/google/cloud/bigquery/_helpers.py
+++ b/google/cloud/bigquery/_helpers.py
@@ -19,7 +19,6 @@
 import decimal
 import math
 import re
-from typing import Union
 
 from google.cloud._helpers import UTC
 from google.cloud._helpers import _date_from_iso8601_date
@@ -47,24 +46,33 @@
 
 
 class BQStorageVersions:
+    """Version comparisons for google-cloud-bigqueyr-storage package."""
+
     def __init__(self):
         self._installed_version = None
 
     @property
-    def installed_version(
-        self,
-    ) -> Union[packaging.version.LegacyVersion, packaging.version.Version]:
+    def installed_version(self) -> packaging.version.Version:
+        """Return the parsed version of google-cloud-bigquery-storage."""
         if self._installed_version is None:
             from google.cloud import bigquery_storage
 
             self._installed_version = packaging.version.parse(
-                getattr(bigquery_storage, "__version__", "legacy")
+                # Use 0.0.0, since it is earlier than any released version.
+                # Legacy versions also have the same property, but
+                # creating a LegacyVersion has been deprecated.
+                # https://github.com/pypa/packaging/issues/321
+                getattr(bigquery_storage, "__version__", "0.0.0")
             )
 
         return self._installed_version
 
     @property
     def is_read_session_optional(self) -> bool:
+        """True if read_session is optional to rows().
+
+        See: https://github.com/googleapis/python-bigquery-storage/pull/228
+        """
         return self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION
 
     def verify_version(self):
diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py
index 2ce443771..af026ccbe 100644
--- a/tests/unit/test__helpers.py
+++ b/tests/unit/test__helpers.py
@@ -27,6 +27,11 @@
 
 @unittest.skipIf(bigquery_storage is None, "Requires `google-cloud-bigquery-storage`")
 class TestBQStorageVersions(unittest.TestCase):
+    def _object_under_test(self):
+        from google.cloud.bigquery import _helpers
+
+        return _helpers.BQStorageVersions()
+
     def _call_fut(self):
         from google.cloud.bigquery import _helpers
 
@@ -54,10 +59,35 @@ def test_raises_error_w_unknown_bqstorage_version(self):
 
         with mock.patch("google.cloud.bigquery_storage", autospec=True) as fake_module:
             del fake_module.__version__
-            error_pattern = r"version found: legacy"
+            error_pattern = r"version found: 0.0.0"
             with self.assertRaisesRegex(LegacyBigQueryStorageError, error_pattern):
                 self._call_fut()
 
+    def test_installed_version_returns_cached(self):
+        versions = self._object_under_test()
+        versions._installed_version = object()
+        assert versions.installed_version is versions._installed_version
+
+    def test_installed_version_returns_parsed_version(self):
+        versions = self._object_under_test()
+
+        with mock.patch("google.cloud.bigquery_storage.__version__", new="1.2.3"):
+            version = versions.installed_version
+
+        assert version.major == 1
+        assert version.minor == 2
+        assert version.micro == 3
+
+    def test_is_read_session_optional_true(self):
+        versions = self._object_under_test()
+        with mock.patch("google.cloud.bigquery_storage.__version__", new="2.6.0"):
+            assert versions.is_read_session_optional
+
+    def test_is_read_session_optional_false(self):
+        versions = self._object_under_test()
+        with mock.patch("google.cloud.bigquery_storage.__version__", new="2.5.0"):
+            assert not versions.is_read_session_optional
+
 
 class Test_not_null(unittest.TestCase):
     def _call_fut(self, value, field):
diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py
index ed90f9a24..0ba671cd9 100644
--- a/tests/unit/test__pandas_helpers.py
+++ b/tests/unit/test__pandas_helpers.py
@@ -1340,7 +1340,7 @@ def test__download_table_bqstorage_stream_includes_read_session(
         session,
         google.cloud.bigquery_storage_v1.types.ReadStream(name="test"),
         queue.Queue(),
-        lambda item: item,
+        mock.Mock(),
     )
 
     reader.rows.assert_called_once_with(session)
@@ -1374,7 +1374,7 @@ def test__download_table_bqstorage_stream_omits_read_session(
         session,
         google.cloud.bigquery_storage_v1.types.ReadStream(name="test"),
         queue.Queue(),
-        lambda item: item,
+        mock.Mock(),
     )
 
     reader.rows.assert_called_once_with()