Support both v1 stable and beta1 BQ Storage client

google/cloud/bigquery/_pandas_helpers.py

@@ -577,7 +577,19 @@ 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)
+    # Passing a BQ Storage client in implies that the BigQuery Storage library
+    # is available and can be imported.
+    from google.cloud import bigquery_storage_v1beta1
+
+    # We want to preserve comaptibility with the v1beta1 BQ Storage clients,
+    # thus adjust constructing the rowstream if needed.
+    # The assumption is that the caller provides a BQ Storage `session` that is
+    # compatible with the version of the BQ Storage client passed in.
+    if isinstance(bqstorage_client, bigquery_storage_v1beta1.BigQueryStorageClient):
+        position = bigquery_storage_v1beta1.types.StreamPosition(stream=stream)
+        rowstream = bqstorage_client.read_rows(position).rows(session)
+    else:
+        rowstream = bqstorage_client.read_rows(stream.name).rows(session)
 
     for page in rowstream.pages:
         if download_state.done:
@@ -609,29 +621,57 @@ def _download_table_bqstorage(
     page_to_item=None,
 ):
     """Use (faster, but billable) BQ Storage API to construct DataFrame."""
+
+    # Passing a BQ Storage client in implies that the BigQuery Storage library
+    # is available and can be imported.
+    from google.cloud import bigquery_storage_v1
+    from google.cloud import bigquery_storage_v1beta1
+
     if "$" in table.table_id:
         raise ValueError(
             "Reading from a specific partition is not currently supported."
         )
     if "@" in table.table_id:
         raise ValueError("Reading from a specific snapshot is not currently supported.")
 
-    requested_session = bigquery_storage_v1.types.ReadSession(
-        table=table.to_bqstorage(),
-        data_format=bigquery_storage_v1.enums.DataFormat.ARROW,
-    )
-
-    if selected_fields is not None:
-        for field in selected_fields:
-            requested_session.read_options.selected_fields.append(field.name)
-
     requested_streams = 1 if preserve_order else 0
 
-    session = bqstorage_client.create_read_session(
-        parent="projects/{}".format(project_id),
-        read_session=requested_session,
-        max_stream_count=requested_streams,
-    )
+    # We want to preserve comaptibility with the v1beta1 BQ Storage clients,
+    # thus adjust the session creation if needed.
+    if isinstance(bqstorage_client, bigquery_storage_v1beta1.BigQueryStorageClient):
+        warnings.warn(
+            "Support for BigQuery Storage v1beta1 clients is deprecated, please "
+            "consider upgrading the client to BigQuery Storage v1 stable version.",
+            category=DeprecationWarning,
+        )
+        read_options = bigquery_storage_v1beta1.types.TableReadOptions()
+
+        if selected_fields is not None:
+            for field in selected_fields:
+                read_options.selected_fields.append(field.name)
+
+        session = bqstorage_client.create_read_session(
+            table.to_bqstorage(v1beta1=True),
+            "projects/{}".format(project_id),
+            format_=bigquery_storage_v1beta1.enums.DataFormat.ARROW,
+            read_options=read_options,
+            requested_streams=requested_streams,
+        )
+    else:
+        requested_session = bigquery_storage_v1.types.ReadSession(
+            table=table.to_bqstorage(),
+            data_format=bigquery_storage_v1.enums.DataFormat.ARROW,
+        )
+        if selected_fields is not None:
+            for field in selected_fields:
+                requested_session.read_options.selected_fields.append(field.name)
+
+        session = bqstorage_client.create_read_session(
+            parent="projects/{}".format(project_id),
+            read_session=requested_session,
+            max_stream_count=requested_streams,
+        )
+
     _LOGGER.debug(
         "Started reading table '{}.{}.{}' with BQ Storage API session '{}'.".format(
             table.project, table.dataset_id, table.table_id, session.name

google/cloud/bigquery/dbapi/cursor.py

@@ -15,6 +15,7 @@
 """Cursor for the Google BigQuery DB-API."""
 
 import collections
+import warnings
 
 try:
     from collections import abc as collections_abc
@@ -268,28 +269,54 @@ def _bqstorage_fetch(self, bqstorage_client):
                 A sequence of rows, represented as dictionaries.
         """
         # Hitting this code path with a BQ Storage client instance implies that
-        # bigquery_storage_v1 can indeed be imported here without errors.
+        # bigquery_storage_v1* can indeed be imported here without errors.
         from google.cloud import bigquery_storage_v1
+        from google.cloud import bigquery_storage_v1beta1
 
         table_reference = self._query_job.destination
 
-        requested_session = bigquery_storage_v1.types.ReadSession(
-            table=table_reference.to_bqstorage(),
-            data_format=bigquery_storage_v1.enums.DataFormat.AVRO,
+        is_v1beta1_client = isinstance(
+            bqstorage_client, bigquery_storage_v1beta1.BigQueryStorageClient
         )
 
-        read_session = bqstorage_client.create_read_session(
-            parent="projects/{}".format(table_reference.project),
-            read_session=requested_session,
-            # a single stream only, as DB API is not well-suited for multithreading
-            max_stream_count=1,
-        )
+        # We want to preserve comaptibility with the v1beta1 BQ Storage clients,
+        # thus adjust the session creation if needed.
+        if is_v1beta1_client:
+            warnings.warn(
+                "Support for BigQuery Storage v1beta1 clients is deprecated, please "
+                "consider upgrading the client to BigQuery Storage v1 stable version.",
+                category=DeprecationWarning,
+            )
+            read_session = bqstorage_client.create_read_session(
+                table_reference.to_bqstorage(v1beta1=True),
+                "projects/{}".format(table_reference.project),
+                # a single stream only, as DB API is not well-suited for multithreading
+                requested_streams=1,
+            )
+        else:
+            requested_session = bigquery_storage_v1.types.ReadSession(
+                table=table_reference.to_bqstorage(),
+                data_format=bigquery_storage_v1.enums.DataFormat.AVRO,
+            )
+            read_session = bqstorage_client.create_read_session(
+                parent="projects/{}".format(table_reference.project),
+                read_session=requested_session,
+                # a single stream only, as DB API is not well-suited for multithreading
+                max_stream_count=1,
+            )
 
         if not read_session.streams:
             return iter([])  # empty table, nothing to read
 
-        stream_name = read_session.streams[0].name
-        read_rows_stream = bqstorage_client.read_rows(stream_name)
+        if is_v1beta1_client:
+            read_position = bigquery_storage_v1beta1.types.StreamPosition(
+                stream=read_session.streams[0],
+            )
+            read_rows_stream = bqstorage_client.read_rows(read_position)
+        else:
+            stream_name = read_session.streams[0].name
+            read_rows_stream = bqstorage_client.read_rows(stream_name)
+
         rows_iterable = read_rows_stream.rows(read_session)
         return rows_iterable
 

google/cloud/bigquery/table.py

@@ -25,6 +25,12 @@
 
 import six
 
+try:
+    # Needed for the to_bqstorage() method.
+    from google.cloud import bigquery_storage_v1beta1
+except ImportError:  # pragma: NO COVER
+    bigquery_storage_v1beta1 = None
+
 try:
     import pandas
 except ImportError:  # pragma: NO COVER
@@ -221,7 +227,7 @@ def to_api_repr(self):
             "tableId": self._table_id,
         }
 
-    def to_bqstorage(self):
+    def to_bqstorage(self, v1beta1=False):
         """Construct a BigQuery Storage API representation of this table.
 
         Install the ``google-cloud-bigquery-storage`` package to use this
@@ -235,15 +241,37 @@ def to_bqstorage(self):
         :class:`google.cloud.bigquery_storage_v1.types.ReadSession.TableModifiers`
         to select a specific snapshot to read from.
 
+        Args:
+            v1beta1 (Optiona[bool]):
+                If :data:`True`, return representation compatible with BigQuery
+                Storage ``v1beta1`` version. Defaults to :data:`False`.
+
         Returns:
-            str: A reference to this table in the BigQuery Storage API.
+            Union[str, google.cloud.bigquery_storage_v1beta1.types.TableReference:]:
+                A reference to this table in the BigQuery Storage API.
+
+        Raises:
+            ValueError:
+                If ``v1beta1`` compatibility is requested, but the
+                :mod:`google.cloud.bigquery_storage_v1beta1` module	cannot be imported.
         """
+        if v1beta1 and bigquery_storage_v1beta1 is None:
+            raise ValueError(_NO_BQSTORAGE_ERROR)
+
         table_id, _, _ = self._table_id.partition("@")
         table_id, _, _ = table_id.partition("$")
 
-        table_ref = "projects/{}/datasets/{}/tables/{}".format(
-            self._project, self._dataset_id, table_id,
-        )
+        if v1beta1:
+            table_ref = bigquery_storage_v1beta1.types.TableReference(
+                project_id=self._project,
+                dataset_id=self._dataset_id,
+                table_id=table_id,
+            )
+        else:
+            table_ref = "projects/{}/datasets/{}/tables/{}".format(
+                self._project, self._dataset_id, table_id,
+            )
+
         return table_ref
 
     def _key(self):
@@ -849,13 +877,19 @@ def to_api_repr(self):
         """
         return copy.deepcopy(self._properties)
 
-    def to_bqstorage(self):
+    def to_bqstorage(self, v1beta1=False):
         """Construct a BigQuery Storage API representation of this table.
 
+        Args:
+            v1beta1 (Optiona[bool]):
+                If :data:`True`, return representation compatible with BigQuery
+                Storage ``v1beta1`` version. Defaults to :data:`False`.
+
         Returns:
-            str: A reference to this table in the BigQuery Storage API.
+            Union[str, google.cloud.bigquery_storage_v1beta1.types.TableReference:]:
+                A reference to this table in the BigQuery Storage API.
         """
-        return self.reference.to_bqstorage()
+        return self.reference.to_bqstorage(v1beta1=v1beta1)
 
     def _build_resource(self, filter_fields):
         """Generate a resource for ``update``."""
@@ -1063,13 +1097,19 @@ def from_string(cls, full_table_id):
             {"tableReference": TableReference.from_string(full_table_id).to_api_repr()}
         )
 
-    def to_bqstorage(self):
+    def to_bqstorage(self, v1beta1=False):
         """Construct a BigQuery Storage API representation of this table.
 
+        Args:
+            v1beta1 (Optiona[bool]):
+                If :data:`True`, return representation compatible with BigQuery
+                Storage ``v1beta1`` version. Defaults to :data:`False`.
+
         Returns:
-            str: A reference to this table in the BigQuery Storage API.
+            Union[str, google.cloud.bigquery_storage_v1beta1.types.TableReference:]:
+                A reference to this table in the BigQuery Storage API.
         """
-        return self.reference.to_bqstorage()
+        return self.reference.to_bqstorage(v1beta1=v1beta1)
 
 
 def _row_from_mapping(mapping, schema):