Add support for specifying location

Author: mxmzdlv

dev_requirements.txt

@@ -1,5 +1,5 @@
 sqlalchemy>=1.1.9
-google-cloud-bigquery>=0.30.0
+google-cloud-bigquery>=1.5.0
 future==0.16.0
 
 pytest==3.2.2

pybigquery/api.py

@@ -7,13 +7,14 @@
 
 
 class ApiClient(object):
-    def __init__(self, credentials_path=None):
+    def __init__(self, credentials_path=None, location=None):
         self.credentials_path = credentials_path
+        self.location = location
         if self.credentials_path:
             self.client = Client.from_service_account_json(
-                self.credentials_path)
+                self.credentials_path, location=self.location)
         else:
-            self.client = Client()
+            self.client = Client(location=self.location)
 
     def dry_run_query(self, query):
         job_config = QueryJobConfig()

pybigquery/sqlalchemy_bigquery.py

@@ -193,20 +193,22 @@ class BigQueryDialect(DefaultDialect):
     supports_simple_order_by_label = True
     postfetch_lastrowid = False
 
-    def __init__(self, arraysize=5000, credentials_path=None, *args, **kwargs):
+    def __init__(self, arraysize=5000, credentials_path=None, location=None, *args, **kwargs):
         super(BigQueryDialect, self).__init__(*args, **kwargs)
         self.arraysize = arraysize
         self.credentials_path = credentials_path
+        self.location = location
 
     @classmethod
     def dbapi(cls):
         return dbapi
 
     def create_connect_args(self, url):
         if self.credentials_path:
-            client = bigquery.Client.from_service_account_json(self.credentials_path)
+            client = bigquery.Client.from_service_account_json(
+                self.credentials_path, location=self.location)
         else:
-            client = bigquery.Client(url.host)
+            client = bigquery.Client(url.host, location=self.location)
         return ([client], {})
 
     def _json_deserializer(self, row):

scripts/load_test_data.sh

@@ -1,7 +1,12 @@
 bq mk test_pybigquery
+bq mk --data_location=asia-northeast1 test_pybigquery_location
+
 bq rm -f -t test_pybigquery.sample
 bq rm -f -t test_pybigquery.sample_one_row
 bq rm -f -t test_pybigquery.sample_dml
+bq rm -f -t test_pybigquery_location.sample_one_row
+
 bq load --source_format=NEWLINE_DELIMITED_JSON --schema=$(dirname $0)/schema.json test_pybigquery.sample $(dirname $0)/sample.json
 bq load --source_format=NEWLINE_DELIMITED_JSON --schema=$(dirname $0)/schema.json test_pybigquery.sample_one_row $(dirname $0)/sample_one_row.json
+bq --location=asia-northeast1 load --source_format=NEWLINE_DELIMITED_JSON --schema=$(dirname $0)/schema.json test_pybigquery_location.sample_one_row $(dirname $0)/sample_one_row.json
 bq mk --schema=$(dirname $0)/schema.json -t test_pybigquery.sample_dml

setup.py

@@ -8,23 +8,23 @@ def readme():
 
 setup(
     name="pybigquery",
-    version='0.4.1',
+    version='0.4.2',
     description="SQLAlchemy dialect for BigQuery",
     long_description=readme(),
     long_description_content_type="text/x-rst",
     author="Maxim Zudilov",
     author_email="maxim.zudilov@gmail.com",
     packages=['pybigquery'],
     url="https://github.com/mxmzdlv/pybigquery",
-    download_url='https://github.com/mxmzdlv/pybigquery/archive/v0.4.1.tar.gz',
+    download_url='https://github.com/mxmzdlv/pybigquery/archive/v0.4.2.tar.gz',
     keywords=['bigquery', 'sqlalchemy'],
     classifiers=[
         "Intended Audience :: Developers",
         "Topic :: Database :: Front-Ends"
     ],
     install_requires=[
         'sqlalchemy>=1.1.9',
-        'google-cloud-bigquery>=0.30.0',
+        'google-cloud-bigquery>=1.5.0',
         'future',
     ],
     tests_require=[

test/test_sqlalchemy_bigquery.py

@@ -91,6 +91,12 @@ def engine():
     return engine
 
 
+@pytest.fixture(scope='session')
+def engine_with_location():
+    engine = create_engine('bigquery://', echo=True, location="asia-northeast1")
+    return engine
+
+
 @pytest.fixture(scope='session')
 def table(engine):
     return Table('test_pybigquery.sample', MetaData(bind=engine), autoload=True)
@@ -133,20 +139,28 @@ 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
+    assert api_client.dry_run_query(sql).total_bytes_processed == 132
 
     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_dataset_location(engine_with_location):
+    rows = engine_with_location.execute('SELECT * FROM test_pybigquery_location.sample_one_row').fetchall()
+    assert list(rows[0]) == ONE_ROW_CONTENTS
+
+
 def test_reflect_select(engine, table):
     assert len(table.c) == 14