Add bq storage client

Author: alonme

dev_requirements.txt

@@ -1,5 +1,6 @@
 sqlalchemy>=1.1.9
-google-cloud-bigquery>=1.6.0
+google-cloud-bigquery[bqstorage]>=1.25.0
+google-cloud-bigquery-storage[fastavro]>=1.0.0
 future==0.16.0
 
 pytest==3.2.2

pybigquery/parse_url.py

@@ -34,6 +34,7 @@ def parse_url(url):
     dataset_id = url.database or None
     arraysize = None
     credentials_path = None
+    use_bqstorage_api = None
 
     # location
     if 'location' in query:
@@ -51,14 +52,22 @@ def parse_url(url):
         except ValueError:
             raise ValueError("invalid int in url query arraysize: " + str_arraysize)
 
+    # use_bqstorage_api
+    if 'use_bqstorage_api' in query:
+        str_use_bqstorage_api = query.pop('use_bqstorage_api')
+        try:
+            use_bqstorage_api = parse_boolean(str_use_bqstorage_api)
+        except ValueError:
+            raise ValueError("invalid boolean in url query for use_bqstorage_api: " + str_use_bqstorage_api)
+
     # if only these "non-config" values were present, the dict will now be empty
     if not query:
         # if a dataset_id exists, we need to return a job_config that isn't None
         # so it can be updated with a dataset reference from the client
         if dataset_id:
-            return project_id, location, dataset_id, arraysize, credentials_path, QueryJobConfig()
+            return project_id, location, dataset_id, arraysize, credentials_path, use_bqstorage_api, QueryJobConfig()
         else:
-            return project_id, location, dataset_id, arraysize, credentials_path, None
+            return project_id, location, dataset_id, arraysize, credentials_path, use_bqstorage_api, None
 
     job_config = QueryJobConfig()
 
@@ -170,4 +179,4 @@ def parse_url(url):
         except AttributeError:
             raise ValueError("invalid write_disposition in url query: " + query['write_disposition'])
 
-    return project_id, location, dataset_id, arraysize, credentials_path, job_config
+    return project_id, location, dataset_id, arraysize, credentials_path, use_bqstorage_api, job_config

pybigquery/sqlalchemy_bigquery.py

@@ -4,7 +4,7 @@
 from __future__ import unicode_literals
 
 from google import auth
-from google.cloud import bigquery
+from google.cloud import bigquery, bigquery_storage_v1beta1
 from google.cloud.bigquery import dbapi, QueryJobConfig
 from google.cloud.bigquery.schema import SchemaField
 from google.cloud.bigquery.table import EncryptionConfiguration
@@ -320,7 +320,7 @@ def create_credentials(credentials_path, credentials_info):
         return credentials
 
     def create_connect_args(self, url):
-        project_id, location, dataset_id, arraysize, credentials_path, default_query_job_config = parse_url(url)
+        project_id, location, dataset_id, arraysize, credentials_path, use_bqstorage_api, default_query_job_config = parse_url(url)
 
         self.arraysize = self.arraysize or arraysize
         self.location = location or self.location
@@ -341,8 +341,12 @@ def create_connect_args(self, url):
             location=self.location,
             default_query_job_config=default_query_job_config
         )
+        
+        storage_client = None 
+        if use_bqstorage_api:
+            storage_client = bigquery_storage_v1beta1.BigQueryStorageClient(credentials=credentials)
 
-        return ([client], {})
+        return ([client, storage_client], {})
 
     def _json_deserializer(self, row):
         """JSON deserializer for RECORD types.

setup.py

@@ -25,7 +25,8 @@ def readme():
     ],
     install_requires=[
         'sqlalchemy>=1.1.9',
-        'google-cloud-bigquery>=1.6.0',
+        'google-cloud-bigquery[bqstorage]>=1.25.0',
+        'google-cloud-bigquery-storage[fastavro]>=1.0.0',
         'future',
     ],
     tests_require=[

test/test_parse_url.py

@@ -31,6 +31,7 @@ def url_with_everything():
         '?credentials_path=/some/path/to.json'
         '&location=some-location'
         '&arraysize=1000'
+        '&use_bqstorage_api=True'
         '&clustering_fields=a,b,c'
         '&create_disposition=CREATE_IF_NEEDED'
         '&destination=different-project.different-dataset.table'
@@ -46,12 +47,13 @@ def url_with_everything():
 
 
 def test_basic(url_with_everything):
-    project_id, location, dataset_id, arraysize, credentials_path, job_config = parse_url(url_with_everything)
+    project_id, location, dataset_id, arraysize, credentials_path, use_bqstorage_api, job_config = parse_url(url_with_everything)
 
     assert project_id == 'some-project'
     assert location == 'some-location'
     assert dataset_id == 'some-dataset'
     assert arraysize == 1000
+    assert use_bqstorage_api is True
     assert credentials_path == '/some/path/to.json'
     assert isinstance(job_config, QueryJobConfig)
 
@@ -69,7 +71,7 @@ def test_basic(url_with_everything):
     ('write_disposition', 'WRITE_APPEND'),
 ])
 def test_all_values(url_with_everything, param, value):
-    job_config = parse_url(url_with_everything)[5]
+    job_config = parse_url(url_with_everything)[6]
 
     config_value = getattr(job_config, param)
     if callable(value):
@@ -85,6 +87,7 @@ def test_all_values(url_with_everything, param, value):
 
 @pytest.mark.parametrize("param, value", [
     ('arraysize', 'not-int'),
+    ('use_bqstorage_api', 'not-bool'),
     ('create_disposition', 'not-attribute'),
     ('destination', 'not.fully-qualified'),
     ('dry_run', 'not-bool'),
@@ -108,24 +111,26 @@ def test_empty_url():
         assert value is None
 
 def test_empty_with_non_config():
-    url = parse_url(make_url('bigquery:///?location=some-location&arraysize=1000&credentials_path=/some/path/to.json'))
-    project_id, location, dataset_id, arraysize, credentials_path, job_config = url
+    url = parse_url(make_url('bigquery:///?location=some-location&arraysize=1000&use_bqstorage_api=False&credentials_path=/some/path/to.json'))
+    project_id, location, dataset_id, arraysize, credentials_path, use_bqstorage_api, job_config = url
 
     assert project_id is None
     assert location == 'some-location'
     assert dataset_id is None
     assert arraysize == 1000
+    assert use_bqstorage_api is False
     assert credentials_path == '/some/path/to.json'
     assert job_config is None
 
 def test_only_dataset():
     url = parse_url(make_url('bigquery:///some-dataset'))
-    project_id, location, dataset_id, arraysize, credentials_path, job_config = url
+    project_id, location, dataset_id, arraysize, credentials_path, use_bqstorage_api, job_config = url
 
     assert project_id is None
     assert location is None
     assert dataset_id == 'some-dataset'
     assert arraysize is None
+    assert use_bqstorage_api is None
     assert credentials_path is None
     assert isinstance(job_config, QueryJobConfig)
     # we can't actually test that the dataset is on the job_config,