Add quickstart demonstrating most BQ Storage API read features (#7223)

* Add quickstart demonstrating most BQ Storage API read features This snippet will be included at the bottom of the client library index page for the BigQuery Storage API. It will also be included in the docs at cloud.google.com. Adds a new session for nox to run the tests for samples in the docs directory. Acts as a system test of the rows() method. * s/handles reconnecting/reconnects/ * Add note about optional dependencies to readme.

Author: tswast

bigquery_storage/README.rst

@@ -8,7 +8,7 @@ Python Client for BigQuery Storage API (`Alpha`_)
 
 .. _Alpha: https://github.com/GoogleCloudPlatform/google-cloud-python/blob/master/README.rst
 .. _BigQuery Storage API: https://cloud.google.com/bigquery
-.. _Client Library Documentation: https://googlecloudplatform.github.io/google-cloud-python/stable/bigquery_storage/index.html
+.. _Client Library Documentation: https://googlecloudplatform.github.io/google-cloud-python/latest/bigquery_storage/index.html
 .. _Product Documentation: https://cloud.google.com/bigquery
 
 Quick Start
@@ -70,6 +70,22 @@ Windows
  <your-env>\Scripts\activate
  <your-env>\Scripts\pip.exe install google-cloud-bigquery-storage
 
+Optional Dependencies
+^^^^^^^^^^^^^^^^^^^^^
+
+Several features of ``google-cloud-bigquery-storage`` require additional
+dependencies.
+
+* Parse Avro blocks in a ``read_rows()`` stream using `fastavro
+ <https://fastavro.readthedocs.io/en/latest/>`_.
+
+ ``pip install google-cloud-bigquery-storage[fastavro]``
+
+* Write rows to a `pandas <http://pandas.pydata.org/pandas-docs/stable/>`_
+ dataframe.
+
+ ``pip install google-cloud-bigquery-storage[pandas,fastavro]``
+
 Next Steps
 ~~~~~~~~~~
 

bigquery_storage/docs/index.rst

@@ -8,3 +8,12 @@ API Reference
  gapic/v1beta1/api
  gapic/v1beta1/reader
  gapic/v1beta1/types
+
+Example Usage
+-------------
+
+.. literalinclude:: quickstart.py
+ :language: python
+ :dedent: 4
+ :start-after: [START bigquerystorage_quickstart]
+ :end-before: [END bigquerystorage_quickstart]

bigquery_storage/docs/quickstart.py

@@ -0,0 +1,96 @@
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+
+
+
+def main(project_id='your-project-id', snapshot_millis=0):
+ # [START bigquerystorage_quickstart]
+ from google.cloud import bigquery_storage_v1beta1
+
+
+ # TODO(developer): Set the project_id variable.
+ # project_id = 'your-project-id'
+ #
+ # The read session is created in this project. This project can be
+ # different from that which contains the table.
+
+ client = bigquery_storage_v1beta1.BigQueryStorageClient()
+
+ # This example reads baby name data from the public datasets.
+ table_ref = bigquery_storage_v1beta1.types.TableReference()
+ table_ref.project_id = "bigquery-public-data"
+ table_ref.dataset_id = "usa_names"
+ table_ref.table_id = "usa_1910_current"
+
+ # We limit the output columns to a subset of those allowed in the table,
+ # and set a simple filter to only report names from the state of
+ # Washington (WA).
+ read_options = bigquery_storage_v1beta1.types.TableReadOptions()
+ read_options.selected_fields.append("name")
+ read_options.selected_fields.append("number")
+ read_options.selected_fields.append("state")
+ read_options.row_restriction = 'state = "WA"'
+
+ # Set a snapshot time if it's been specified.
+ modifiers = None
+ if snapshot_millis > 0:
+ modifiers = bigquery_storage_v1beta1.types.TableModifiers()
+ modifiers.snapshot_time.FromMilliseconds(snapshot_millis)
+
+ parent = "projects/{}".format(project_id)
+ session = client.create_read_session(
+ table_ref,
+ parent,
+ table_modifiers=modifiers,
+ read_options=read_options) # API request.
+
+ # We'll use only a single stream for reading data from the table. Because
+ # of dynamic sharding, this will yield all the rows in the table. However,
+ # if you wanted to fan out multiple readers you could do so by having a
+ # reader process each individual stream.
+ reader = client.read_rows(
+ bigquery_storage_v1beta1.types.StreamPosition(
+ stream=session.streams[0],
+ )
+ )
+
+ # The read stream contains blocks of Avro-encoded bytes. The rows() method
+ # uses the fastavro library to parse these blocks as an interable of Python
+ # dictionaries. Install fastavro with the following command:
+ #
+ # pip install google-cloud-bigquery-storage[fastavro]
+ rows = reader.rows(session)
+
+ # Do any local processing by iterating over the rows. The
+ # google-cloud-bigquery-storage client reconnects to the API after any
+ # transient network errors or timeouts.
+ names = set()
+ states = set()
+
+ for row in rows:
+ names.add(row["name"])
+ states.add(row["state"])
+
+ print("Got {} unique names in states: {}".format(len(names), states))
+ # [END bigquerystorage_quickstart]
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument('project_id')
+ parser.add_argument('--snapshot_millis', default=0, type=int)
+ args = parser.parse_args()
+ main(project_id=args.project_id)

bigquery_storage/docs/quickstart_test.py

@@ -0,0 +1,43 @@
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+import os
+
+import pytest
+
+import quickstart
+
+
+def now_millis():
+ return int(
+ (datetime.datetime.utcnow() - datetime.datetime(1970, 1, 1)).total_seconds()
+ * 1000)
+
+
+@pytest.fixture()
+def project_id():
+ return os.environ["PROJECT_ID"]
+
+
+def test_quickstart_wo_snapshot(capsys, project_id):
+ quickstart.main(project_id)
+ out, _ = capsys.readouterr()
+ assert 'WA' in out
+
+
+def test_quickstart_with_snapshot(capsys, project_id):
+ quickstart.main(project_id, now_millis() - 5000)
+ out, _ = capsys.readouterr()
+ assert 'WA' in out

bigquery_storage/noxfile.py

@@ -127,6 +127,27 @@ def system(session):
  session.run('py.test', '--quiet', 'tests/system/')
 
 
+@nox.session(python=['2.7', '3.6'])
+def snippets(session):
+ """Run the snippets test suite."""
+
+ # Sanity check: Only run snippets tests if the environment variable is set.
+ if not os.environ.get('GOOGLE_APPLICATION_CREDENTIALS', ''):
+ session.skip('Credentials must be set via environment variable.')
+
+ # Install all test dependencies, then install this package into the
+ # virtualenv's dist-packages.
+ session.install('pytest')
+ session.install('-e', os.path.join('..', 'test_utils'))
+ for local_dep in LOCAL_DEPS:
+ session.install('-e', local_dep)
+ session.install('-e', '.[pandas,fastavro]')
+
+ # Run py.test against the snippets tests.
+ session.run(
+ 'py.test', 'docs', *session.posargs)
+
+
 @nox.session(python='3.6')
 def docs(session):
  """Build the docs."""