Merge pull request #20 from sumedhsakdeo/dry_run

Google BigQuery API Client

Author: mxmzdlv

packages/sqlalchemy-bigquery/README.rst

@@ -1,18 +1,26 @@
-SQLAlchemy dialect for BigQuery.
+SQLAlchemy dialect and API client for BigQuery.
 
 
 Usage
 =====
 
 .. code-block:: python
 
+SQLAchemy
+_________
+
     from sqlalchemy import *
     from sqlalchemy.engine import create_engine
     from sqlalchemy.schema import *
     engine = create_engine('bigquery://project')
     table = Table('dataset.table', MetaData(bind=engine), autoload=True)
     print(select([func.count('*')], from_obj=table).scalar())
 
+API Client
+_________
+    from pybigquery.api import ApiClient
+    api_client = ApiClient()
+    print(api_client.dry_run_query(query=sqlstr).total_bytes_processed)
 
 Project
 _______

packages/sqlalchemy-bigquery/pybigquery/api.py

@@ -0,0 +1,22 @@
+"""Integration with BigQuery API."""
+
+from __future__ import absolute_import
+from __future__ import unicode_literals
+
+from google.cloud.bigquery import Client, QueryJobConfig
+
+
+class ApiClient(object):
+    def __init__(self, credentials_path=None):
+        self.credentials_path = credentials_path
+        if self.credentials_path:
+            self.client = Client.from_service_account_json(
+                self.credentials_path)
+        else:
+            self.client = Client()
+
+    def dry_run_query(self, query):
+        job_config = QueryJobConfig()
+        job_config.dry_run = True
+        job_config.use_query_cache = False
+        return self.client.query(query=(query), job_config=job_config)

packages/sqlalchemy-bigquery/test/test_sqlalchemy_bigquery.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
 from __future__ import unicode_literals
 
+from google.api_core.exceptions import BadRequest
+from pybigquery.api import ApiClient
 from sqlalchemy.engine import create_engine
 from sqlalchemy.schema import Table, MetaData, Column
 from sqlalchemy import types, func, case, inspect
@@ -131,6 +133,19 @@ def query(table):
     )
     return query
 
+@pytest.fixture(scope='session')
+def api_client():
+    return ApiClient()
+
+def test_dry_run(engine, api_client):
+    sql = 'SELECT * FROM test_pybigquery.sample_one_row'
+    assert api_client.dry_run_query(sql).total_bytes_processed == 112
+
+    sql = 'SELECT * FROM sample_one_row'
+    with pytest.raises(BadRequest) as excinfo:
+        api_client.dry_run_query(sql)
+
+    assert 'Table name "sample_one_row" cannot be resolved: dataset name is missing.' in str(excinfo.value.message)
 
 def test_reflect_select(engine, table):
     assert len(table.c) == 14