diff --git a/README.rst b/README.rst
index 69ea3020..c626bcc7 100644
--- a/README.rst
+++ b/README.rst
@@ -70,6 +70,64 @@ By default, ``arraysize`` is set to ``5000``. ``arraysize`` is used to set the b
engine = create_engine('bigquery://project', arraysize=1000)
+Adding a Default Dataset
+________________________
+
+If you want to have the ``Client`` use a default dataset, specify it as the "database" portion of the connection string.
+
+.. code-block:: python
+
+ engine = create_engine('bigquery://project/dataset')
+
+When using a default dataset, don't include the dataset name in the table name, e.g.:
+
+.. code-block:: python
+
+ table = Table('table_name')
+
+Note that specyfing a default dataset doesn't restrict execution of queries to that particular dataset when using raw queries, e.g.:
+
+.. code-block:: python
+
+ # Set default dataset to dataset_a
+ engine = create_engine('bigquery://project/dataset_a')
+
+ # This will still execute and return rows from dataset_b
+ engine.execute('SELECT * FROM dataset_b.table').fetchall()
+
+
+Connection String Parameters
+____________________________
+
+There are many situations where you can't call ``create_engine`` directly, such as when using tools like `Flask SQLAlchemy `_. For situations like these, or for situations where you want the ``Client`` to have a `default_query_job_config `_, you can pass many arguments in the query of the connection string.
+
+The ``credentials_path``, ``location``, and ``arraysize`` parameters are used by this library, and the rest are used to create a `QueryJobConfig `_
+
+Note that if you want to use query strings, it will be more reliable if you use three slashes, so ``'bigquery:///?a=b'`` will work reliably, but ``'bigquery://?a=b'`` might be interpreted as having a "database" of ``?a=b``, depending on the system being used to parse the connection string.
+
+Here are examples of all the supported arguments. Any not present are either for legacy sql (which isn't supported by this library), or are too complex and are not implemented.
+
+.. code-block:: python
+
+ engine = create_engine(
+ 'bigquery://some-project/some-dataset' '?'
+ 'credentials_path=/some/path/to.json' '&'
+ 'location=some-location' '&'
+ 'arraysize=1000' '&'
+ 'clustering_fields=a,b,c' '&'
+ 'create_disposition=CREATE_IF_NEEDED' '&'
+ 'destination=different-project.different-dataset.table' '&'
+ 'destination_encryption_configuration=some-configuration' '&'
+ 'dry_run=true' '&'
+ 'labels=a:b,c:d' '&'
+ 'maximum_bytes_billed=1000' '&'
+ 'priority=INTERACTIVE' '&'
+ 'schema_update_options=ALLOW_FIELD_ADDITION,ALLOW_FIELD_RELAXATION' '&'
+ 'use_query_cache=true' '&'
+ 'write_disposition=WRITE_APPEND'
+ )
+
+
Requirements
============
diff --git a/dev_requirements.txt b/dev_requirements.txt
index c444d8cf..b32e352b 100644
--- a/dev_requirements.txt
+++ b/dev_requirements.txt
@@ -1,5 +1,5 @@
sqlalchemy>=1.1.9
-google-cloud-bigquery>=1.5.0
+google-cloud-bigquery>=1.6.0
future==0.16.0
pytest==3.2.2
diff --git a/pybigquery/parse_url.py b/pybigquery/parse_url.py
new file mode 100644
index 00000000..b9e6c5f6
--- /dev/null
+++ b/pybigquery/parse_url.py
@@ -0,0 +1,172 @@
+import re
+GROUP_DELIMITER = re.compile(r'\s*\,\s*')
+KEY_VALUE_DELIMITER = re.compile(r'\s*\:\s*')
+
+from google.cloud.bigquery import QueryJobConfig
+from google.cloud.bigquery.job import CreateDisposition, WriteDisposition, QueryPriority, SchemaUpdateOption
+
+from google.cloud.bigquery.table import EncryptionConfiguration, TableReference
+from google.cloud.bigquery.dataset import DatasetReference
+
+def parse_boolean(bool_string):
+ bool_string = bool_string.lower()
+ if bool_string == 'true':
+ return True
+ elif bool_string == 'false':
+ return False
+ else:
+ raise ValueError()
+
+def parse_url(url):
+ query = url.query
+
+ # use_legacy_sql (legacy)
+ if 'use_legacy_sql' in query: raise ValueError("legacy sql is not supported by this dialect")
+ # allow_large_results (legacy)
+ if 'allow_large_results' in query: raise ValueError("allow_large_results is only allowed for legacy sql, which is not supported by this dialect")
+ # flatten_results (legacy)
+ if 'flatten_results' in query: raise ValueError("flatten_results is only allowed for legacy sql, which is not supported by this dialect")
+ # maximum_billing_tier (deprecated)
+ if 'maximum_billing_tier' in query: raise ValueError("maximum_billing_tier is a deprecated argument")
+
+ location = None
+ dataset_id = url.database or None
+ arraysize = None
+ credentials_path = None
+
+ # location
+ if 'location' in query:
+ location = query.pop('location')
+
+ # credentials_path
+ if 'credentials_path' in query:
+ credentials_path = query.pop('credentials_path')
+
+ # arraysize
+ if 'arraysize' in query:
+ str_arraysize = query.pop('arraysize')
+ try:
+ arraysize = int(str_arraysize)
+ except ValueError:
+ raise ValueError("invalid int in url query arraysize: " + str_arraysize)
+
+ # 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 location, dataset_id, arraysize, credentials_path, QueryJobConfig()
+ else:
+ return location, dataset_id, arraysize, credentials_path, None
+
+ job_config = QueryJobConfig()
+
+ # clustering_fields list(str)
+ if 'clustering_fields' in query:
+ clustering_fields = GROUP_DELIMITER.split(query['clustering_fields'])
+ job_config.clustering_fields = list(clustering_fields)
+
+ # create_disposition
+ if 'create_disposition' in query:
+ create_disposition = query['create_disposition']
+ try:
+ job_config.create_disposition = getattr(CreateDisposition, create_disposition)
+ except AttributeError:
+ raise ValueError("invalid create_disposition in url query: " + create_disposition)
+
+ # default_dataset
+ if 'default_dataset' in query or 'dataset_id' in query or 'project_id' in query:
+ raise ValueError("don't pass default_dataset, dataset_id, project_id in url query, instead use the url host and database")
+
+ # destination
+ if 'destination' in query:
+ dest_project = None
+ dest_dataset = None
+ dest_table = None
+
+ try:
+ dest_project, dest_dataset, dest_table = query['destination'].split('.')
+ except ValueError:
+ raise ValueError("url query destination parameter should be fully qualified with project, dataset, and table")
+
+ job_config.destination = TableReference(DatasetReference(dest_project, dest_dataset), dest_table)
+
+ # destination_encryption_configuration
+ if 'destination_encryption_configuration' in query:
+ job_config.destination_encryption_configuration = EncryptionConfiguration(query['destination_encryption_configuration'])
+
+ # dry_run
+ if 'dry_run' in query:
+ try:
+ job_config.dry_run = parse_boolean(query['dry_run'])
+ except ValueError:
+ raise ValueError("invalid boolean in url query for dry_run: " + query['dry_run'])
+
+ # labels
+ if 'labels' in query:
+ label_groups = GROUP_DELIMITER.split(query['labels'])
+ labels = {}
+ for label_group in label_groups:
+ try:
+ key, value = KEY_VALUE_DELIMITER.split(label_group)
+ except ValueError:
+ raise ValueError("malformed url query in labels: " + label_group)
+ labels[key] = value
+
+ job_config.labels = labels
+
+ # maximum_bytes_billed
+ if 'maximum_bytes_billed' in query:
+ try:
+ job_config.maximum_bytes_billed = int(query['maximum_bytes_billed'])
+ except ValueError:
+ raise ValueError("invalid int in url query maximum_bytes_billed: " + query['maximum_bytes_billed'])
+
+ # priority
+ if 'priority' in query:
+ try:
+ job_config.priority = getattr(QueryPriority, query['priority'])
+ except AttributeError:
+ raise ValueError("invalid priority in url query: " + query['priority'])
+
+ # query_parameters
+ if 'query_parameters' in query:
+ raise NotImplementedError("url query query_parameters not implemented")
+
+ # schema_update_options
+ if 'schema_update_options' in query:
+ schema_update_options = GROUP_DELIMITER.split(query['schema_update_options'])
+ try:
+ job_config.schema_update_options = [
+ getattr(SchemaUpdateOption, schema_update_option) for schema_update_option in schema_update_options
+ ]
+ except AttributeError:
+ raise ValueError("invalid schema_update_options in url query: " + query['schema_update_options'])
+
+ # table_definitions
+ if 'table_definitions' in query:
+ raise NotImplementedError("url query table_definitions not implemented")
+
+ # time_partitioning
+ if 'time_partitioning' in query:
+ raise NotImplementedError("url query time_partitioning not implemented")
+
+ # udf_resources
+ if 'udf_resources' in query:
+ raise NotImplementedError("url query udf_resources not implemented")
+
+ # use_query_cache
+ if 'use_query_cache' in query:
+ try:
+ job_config.use_query_cache = parse_boolean(query['use_query_cache'])
+ except ValueError:
+ raise ValueError("invalid boolean in url query for use_query_cache: " + query['use_query_cache'])
+
+ # write_disposition
+ if 'write_disposition' in query:
+ try:
+ job_config.write_disposition = getattr(WriteDisposition, query['write_disposition'])
+ except AttributeError:
+ raise ValueError("invalid write_disposition in url query: " + query['write_disposition'])
+
+ return location, dataset_id, arraysize, credentials_path, job_config
diff --git a/pybigquery/sqlalchemy_bigquery.py b/pybigquery/sqlalchemy_bigquery.py
index 71052219..a2cbd07b 100644
--- a/pybigquery/sqlalchemy_bigquery.py
+++ b/pybigquery/sqlalchemy_bigquery.py
@@ -3,9 +3,11 @@
from __future__ import absolute_import
from __future__ import unicode_literals
-from google.cloud.bigquery import dbapi
-from google.cloud.bigquery.schema import SchemaField
from google.cloud import bigquery
+from google.cloud.bigquery import dbapi, QueryJobConfig
+from google.cloud.bigquery.schema import SchemaField
+from google.cloud.bigquery.table import EncryptionConfiguration
+from google.cloud.bigquery.dataset import DatasetReference
from google.oauth2 import service_account
from google.api_core.exceptions import NotFound
from sqlalchemy.exc import NoSuchTableError
@@ -17,8 +19,9 @@
from sqlalchemy.sql import elements
import re
+from .parse_url import parse_url
-FIELD_ILLEGAL_CHARACTERS = re.compile('[^\w]+')
+FIELD_ILLEGAL_CHARACTERS = re.compile(r'[^\w]+')
class UniversalSet(object):
@@ -87,7 +90,6 @@ def format_label(self, label, name=None):
result = self.quote(name)
return result
-
_type_map = {
'STRING': types.String,
'BOOLEAN': types.Boolean,
@@ -207,25 +209,47 @@ def __init__(
self.credentials_path = credentials_path
self.credentials_info = credentials_info
self.location = location
+ self.dataset_id = None
@classmethod
def dbapi(cls):
return dbapi
def create_connect_args(self, url):
+ location, dataset_id, arraysize, credentials_path, default_query_job_config = parse_url(url)
+
+ self.arraysize = self.arraysize or arraysize
+ self.location = location or self.location
+ self.credentials_path = credentials_path or self.credentials_path
+ self.dataset_id = dataset_id
+
if self.credentials_path:
client = bigquery.Client.from_service_account_json(
- self.credentials_path, location=self.location)
+ self.credentials_path,
+ location=self.location,
+ default_query_job_config=default_query_job_config
+ )
elif self.credentials_info:
credentials = service_account.Credentials.from_service_account_info(
- self.credentials_info)
+ self.credentials_info
+ )
client = bigquery.Client(
+ project=self.credentials_info.get('project_id'),
credentials=credentials,
location=self.location,
- project=self.credentials_info.get('project_id'),
+ default_query_job_config=default_query_job_config,
)
else:
- client = bigquery.Client(url.host, location=self.location)
+ client = bigquery.Client(
+ project=url.host,
+ location=self.location,
+ default_query_job_config=default_query_job_config
+ )
+
+ # if dataset_id is set, then we know the job_config isn't None
+ if dataset_id:
+ default_query_job_config.default_dataset = client.dataset(dataset_id)
+
return ([client], {})
def _json_deserializer(self, row):
@@ -254,10 +278,17 @@ def _get_table(self, connection, table_name, schema=None):
if isinstance(connection, Engine):
connection = connection.connect()
- project, dataset, table_name_prepared = self._split_table_name(table_name)
- if dataset is None and schema is not None:
- dataset = schema
+ table_name_prepared = dataset = project = None
+ if self.dataset_id:
table_name_prepared = table_name
+ dataset = self.dataset_id
+ project = None
+
+ else:
+ project, dataset, table_name_prepared = self._split_table_name(table_name)
+ if dataset is None and schema is not None:
+ table_name_prepared = table_name
+ dataset = schema
table = connection.connection._client.dataset(dataset, project=project).table(table_name_prepared)
try:
@@ -328,7 +359,10 @@ def get_schema_names(self, connection, **kw):
connection = connection.connect()
datasets = connection.connection._client.list_datasets()
- return [d.dataset_id for d in datasets]
+ if self.dataset_id is not None:
+ return [d.dataset_id for d in datasets if d.dataset_id == self.dataset_id]
+ else:
+ return [d.dataset_id for d in datasets]
def get_table_names(self, connection, schema=None, **kw):
if isinstance(connection, Engine):
@@ -339,9 +373,17 @@ def get_table_names(self, connection, schema=None, **kw):
for d in datasets:
if schema is not None and d.dataset_id != schema:
continue
+
+ if self.dataset_id is not None and d.dataset_id != self.dataset_id:
+ continue
+
tables = connection.connection._client.list_tables(d.reference)
for t in tables:
- result.append(d.dataset_id + '.' + t.table_id)
+ if self.dataset_id is None:
+ table_name = d.dataset_id + '.' + t.table_id
+ else:
+ table_name = t.table_id
+ result.append(table_name)
return result
def do_rollback(self, dbapi_connection):
diff --git a/setup.py b/setup.py
index 89b166fa..159b9ef4 100644
--- a/setup.py
+++ b/setup.py
@@ -25,7 +25,7 @@ def readme():
],
install_requires=[
'sqlalchemy>=1.1.9',
- 'google-cloud-bigquery>=1.5.0',
+ 'google-cloud-bigquery>=1.6.0',
'future',
],
tests_require=[
diff --git a/test/test_parse_url.py b/test/test_parse_url.py
new file mode 100644
index 00000000..74cfa4c2
--- /dev/null
+++ b/test/test_parse_url.py
@@ -0,0 +1,154 @@
+import pytest
+from sqlalchemy.engine.url import make_url
+from google.cloud.bigquery import QueryJobConfig
+from google.cloud.bigquery.table import EncryptionConfiguration, TableReference
+from google.cloud.bigquery.dataset import DatasetReference
+
+from pybigquery.parse_url import parse_url
+
+
+@pytest.fixture(scope='session')
+def url_with_everything():
+ # create_disposition
+ # CREATE_IF_NEEDED
+ # CREATE_NEVER
+
+ # write_disposition
+ # WRITE_APPEND
+ # WRITE_TRUNCATE
+ # WRITE_EMPTY
+
+ # priority
+ # INTERACTIVE
+ # BATCH
+
+ # schema_update_option
+ # ALLOW_FIELD_ADDITION
+ # ALLOW_FIELD_RELAXATION
+
+ return make_url(
+ 'bigquery://some-project/some-dataset'
+ '?credentials_path=/some/path/to.json'
+ '&location=some-location'
+ '&arraysize=1000'
+ '&clustering_fields=a,b,c'
+ '&create_disposition=CREATE_IF_NEEDED'
+ '&destination=different-project.different-dataset.table'
+ '&destination_encryption_configuration=some-configuration'
+ '&dry_run=true'
+ '&labels=a:b,c:d'
+ '&maximum_bytes_billed=1000'
+ '&priority=INTERACTIVE'
+ '&schema_update_options=ALLOW_FIELD_ADDITION,ALLOW_FIELD_RELAXATION'
+ '&use_query_cache=true'
+ '&write_disposition=WRITE_APPEND'
+ )
+
+
+def test_basic(url_with_everything):
+ location, dataset_id, arraysize, credentials_path, job_config = parse_url(url_with_everything)
+
+ assert location == 'some-location'
+ assert dataset_id == 'some-dataset'
+ assert arraysize == 1000
+ assert credentials_path == '/some/path/to.json'
+ assert isinstance(job_config, QueryJobConfig)
+
+@pytest.mark.parametrize('param, value', [
+ ('clustering_fields', ['a', 'b', 'c']),
+ ('create_disposition', 'CREATE_IF_NEEDED'),
+ ('destination', TableReference(DatasetReference('different-project', 'different-dataset'), 'table')),
+ ('destination_encryption_configuration', lambda enc: enc.kms_key_name == EncryptionConfiguration('some-configuration').kms_key_name),
+ ('dry_run', True),
+ ('labels', { 'a': 'b', 'c': 'd' }),
+ ('maximum_bytes_billed', 1000),
+ ('priority', 'INTERACTIVE'),
+ ('schema_update_options', ['ALLOW_FIELD_ADDITION', 'ALLOW_FIELD_RELAXATION']),
+ ('use_query_cache', True),
+ ('write_disposition', 'WRITE_APPEND'),
+])
+def test_all_values(url_with_everything, param, value):
+ job_config = parse_url(url_with_everything)[4]
+
+ config_value = getattr(job_config, param)
+ if callable(value):
+ assert value(config_value)
+ else:
+ assert config_value == value
+
+# def test_malformed():
+# location, dataset_id, arraysize, credentials_path, job_config = parse_url(make_url('bigquery:///?credentials_path=a'))
+
+# print(credentials_path)
+# assert False
+
+@pytest.mark.parametrize("param, value", [
+ ('arraysize', 'not-int'),
+ ('create_disposition', 'not-attribute'),
+ ('destination', 'not.fully-qualified'),
+ ('dry_run', 'not-bool'),
+ ('labels', 'not-key-value'),
+ ('maximum_bytes_billed', 'not-int'),
+ ('priority', 'not-attribute'),
+ ('schema_update_options', 'not-attribute'),
+ ('use_query_cache', 'not-bool'),
+ ('write_disposition', 'not-attribute'),
+])
+def test_bad_values(param, value):
+ url = make_url('bigquery:///?' + param + '=' + value)
+ with pytest.raises(ValueError):
+ parse_url(url)
+
+def test_empty_url():
+ for value in parse_url(make_url('bigquery://')):
+ assert value is None
+
+ for value in parse_url(make_url('bigquery:///')):
+ 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'))
+ location, dataset_id, arraysize, credentials_path, job_config = url
+
+ assert location == 'some-location'
+ assert dataset_id is None
+ assert arraysize == 1000
+ assert credentials_path == '/some/path/to.json'
+ assert job_config is None
+
+def test_only_dataset():
+ url = parse_url(make_url('bigquery:///some-dataset'))
+ location, dataset_id, arraysize, credentials_path, job_config = url
+
+ assert location is None
+ assert dataset_id == 'some-dataset'
+ assert arraysize 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,
+ # since we take care of that afterwards, when we have a client to fill in the project
+
+@pytest.mark.parametrize('disallowed_arg', [
+ 'use_legacy_sql',
+ 'allow_large_results',
+ 'flatten_results',
+ 'maximum_billing_tier',
+ 'default_dataset',
+ 'dataset_id',
+ 'project_id',
+])
+def test_disallowed(disallowed_arg):
+ url = make_url('bigquery://some-project/some-dataset/?' + disallowed_arg + '=' + 'whatever')
+ with pytest.raises(ValueError):
+ parse_url(url)
+
+@pytest.mark.parametrize('not_implemented_arg', [
+ 'query_parameters',
+ 'table_definitions',
+ 'time_partitioning',
+ 'udf_resources',
+])
+def test_not_implemented(not_implemented_arg):
+ url = make_url('bigquery://some-project/some-dataset/?' + not_implemented_arg + '=' + 'whatever')
+ with pytest.raises(NotImplementedError):
+ parse_url(url)
diff --git a/test/test_sqlalchemy_bigquery.py b/test/test_sqlalchemy_bigquery.py
index 162936a8..0f38e1fc 100644
--- a/test/test_sqlalchemy_bigquery.py
+++ b/test/test_sqlalchemy_bigquery.py
@@ -91,6 +91,12 @@ def engine():
return engine
+@pytest.fixture(scope='session')
+def engine_using_test_dataset():
+ engine = create_engine('bigquery:///test_pybigquery', echo=True)
+ return engine
+
+
@pytest.fixture(scope='session')
def engine_with_location():
engine = create_engine('bigquery://', echo=True, location="asia-northeast1")
@@ -102,6 +108,11 @@ def table(engine):
return Table('test_pybigquery.sample', MetaData(bind=engine), autoload=True)
+@pytest.fixture(scope='session')
+def table_using_test_dataset(engine_using_test_dataset):
+ return Table('sample', MetaData(bind=engine_using_test_dataset), autoload=True)
+
+
@pytest.fixture(scope='session')
def table_one_row(engine):
return Table('test_pybigquery.sample_one_row', MetaData(bind=engine), autoload=True)
@@ -119,24 +130,37 @@ def session(engine):
return session
+@pytest.fixture(scope='session')
+def session_using_test_dataset(engine_using_test_dataset):
+ Session = sessionmaker(bind=engine_using_test_dataset)
+ session = Session()
+ return session
+
+
@pytest.fixture(scope='session')
def inspector(engine):
return inspect(engine)
+@pytest.fixture(scope='session')
+def inspector_using_test_dataset(engine_using_test_dataset):
+ return inspect(engine_using_test_dataset)
+
@pytest.fixture(scope='session')
-def query(table):
- col1 = literal_column("TIMESTAMP_TRUNC(timestamp, DAY)").label("timestamp_label")
- col2 = func.sum(table.c.integer)
- query = (
- select([
- col1,
- col2,
- ])
- .where(col1 < '2017-01-01 00:00:00')
- .group_by(col1)
- .order_by(col2)
- )
+def query():
+ def query(table):
+ col1 = literal_column("TIMESTAMP_TRUNC(timestamp, DAY)").label("timestamp_label")
+ col2 = func.sum(table.c.integer)
+ query = (
+ select([
+ col1,
+ col2,
+ ])
+ .where(col1 < '2017-01-01 00:00:00')
+ .group_by(col1)
+ .order_by(col2)
+ )
+ return query
return query
@@ -156,32 +180,46 @@ def test_dry_run(engine, api_client):
assert 'Table name "sample_one_row" cannot be resolved: dataset name is missing.' in str(excinfo.value.message)
+def test_engine_with_dataset(engine_using_test_dataset):
+ rows = engine_using_test_dataset.execute('SELECT * FROM sample_one_row').fetchall()
+ assert list(rows[0]) == ONE_ROW_CONTENTS
+
+ table_one_row = Table('sample_one_row', MetaData(bind=engine_using_test_dataset), autoload=True)
+ rows = table_one_row.select().execute().fetchall()
+ assert list(rows[0]) == ONE_ROW_CONTENTS_EXPANDED
+
+ # Table name shouldn't include dataset
+ with pytest.raises(Exception):
+ table_one_row = Table('test_pybigquery.sample_one_row', MetaData(bind=engine_using_test_dataset), autoload=True)
+
+
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
+def test_reflect_select(table, table_using_test_dataset):
+ for table in [table, table_using_test_dataset]:
+ assert len(table.c) == 14
- assert isinstance(table.c.integer, Column)
- assert isinstance(table.c.integer.type, types.Integer)
- assert isinstance(table.c.timestamp.type, types.TIMESTAMP)
- assert isinstance(table.c.string.type, types.String)
- assert isinstance(table.c.float.type, types.Float)
- assert isinstance(table.c.boolean.type, types.Boolean)
- assert isinstance(table.c.date.type, types.DATE)
- assert isinstance(table.c.datetime.type, types.DATETIME)
- assert isinstance(table.c.time.type, types.TIME)
- assert isinstance(table.c.bytes.type, types.BINARY)
- assert isinstance(table.c['record.age'].type, types.Integer)
- assert isinstance(table.c['record.name'].type, types.String)
- assert isinstance(table.c['nested_record.record.age'].type, types.Integer)
- assert isinstance(table.c['nested_record.record.name'].type, types.String)
- assert isinstance(table.c.array.type, types.ARRAY)
+ assert isinstance(table.c.integer, Column)
+ assert isinstance(table.c.integer.type, types.Integer)
+ assert isinstance(table.c.timestamp.type, types.TIMESTAMP)
+ assert isinstance(table.c.string.type, types.String)
+ assert isinstance(table.c.float.type, types.Float)
+ assert isinstance(table.c.boolean.type, types.Boolean)
+ assert isinstance(table.c.date.type, types.DATE)
+ assert isinstance(table.c.datetime.type, types.DATETIME)
+ assert isinstance(table.c.time.type, types.TIME)
+ assert isinstance(table.c.bytes.type, types.BINARY)
+ assert isinstance(table.c['record.age'].type, types.Integer)
+ assert isinstance(table.c['record.name'].type, types.String)
+ assert isinstance(table.c['nested_record.record.age'].type, types.Integer)
+ assert isinstance(table.c['nested_record.record.name'].type, types.String)
+ assert isinstance(table.c.array.type, types.ARRAY)
- rows = table.select().execute().fetchall()
- assert len(rows) == 1000
+ rows = table.select().execute().fetchall()
+ assert len(rows) == 1000
def test_content_from_raw_queries(engine):
@@ -226,65 +264,85 @@ def test_reflect_dataset_does_not_exist(engine):
Table('dataset_does_not_exist.table_does_not_exist', MetaData(bind=engine), autoload=True)
-def test_tables_list(engine):
+def test_tables_list(engine, engine_using_test_dataset):
tables = engine.table_names()
assert 'test_pybigquery.sample' in tables
assert 'test_pybigquery.sample_one_row' in tables
assert 'test_pybigquery.sample_dml' in tables
+ tables = engine_using_test_dataset.table_names()
+ assert 'sample' in tables
+ assert 'sample_one_row' in tables
+ assert 'sample_dml' in tables
-def test_group_by(session, table):
+
+def test_group_by(session, table, session_using_test_dataset, table_using_test_dataset):
"""labels in SELECT clause should be correclty formatted (dots are replaced with underscores)"""
- result = session.query(table.c.string, func.count(table.c.integer)).group_by(table.c.string).all()
+ for session, table in [(session, table), (session_using_test_dataset, table_using_test_dataset)]:
+ result = session.query(table.c.string, func.count(table.c.integer)).group_by(table.c.string).all()
assert len(result) > 0
-def test_session_query(session, table):
- col_concat = func.concat(table.c.string).label('concat')
- result = (
- session
- .query(
- table.c.string,
- col_concat,
- func.avg(table.c.integer),
- func.sum(case([(table.c.boolean == True, 1)], else_=0))
+def test_session_query(session, table, session_using_test_dataset, table_using_test_dataset):
+ for session, table in [(session, table), (session_using_test_dataset, table_using_test_dataset)]:
+ col_concat = func.concat(table.c.string).label('concat')
+ result = (
+ session
+ .query(
+ table.c.string,
+ col_concat,
+ func.avg(table.c.integer),
+ func.sum(case([(table.c.boolean == True, 1)], else_=0))
+ )
+ .group_by(table.c.string, col_concat)
+ .having(func.avg(table.c.integer) > 10)
+
+ ).all()
+ assert len(result) > 0
+
+
+def test_labels(session, table, session_using_test_dataset, table_using_test_dataset):
+ for session, table in [(session, table), (session_using_test_dataset, table_using_test_dataset)]:
+ result = (
+ session
+ .query(
+ # Valid
+ table.c.string.label('abc'),
+ # Invalid, needs to start with underscore
+ table.c.string.label('123'),
+ # Valid
+ table.c.string.label('_123abc'),
+ # Invalid, contains illegal characters
+ table.c.string.label('!@#$%^&*()~`'),
+ )
)
- .group_by(table.c.string, col_concat)
- .having(func.avg(table.c.integer) > 10)
-
- ).all()
- assert len(result) > 0
+ result = result.all()
+ assert len(result) > 0
-def test_labels(session, table):
- result = (
- session
- .query(
- # Valid
- table.c.string.label('abc'),
- # Invalid, needs to start with underscore
- table.c.string.label('123'),
- # Valid
- table.c.string.label('_123abc'),
- # Invalid, contains illegal characters
- table.c.string.label('!@#$%^&*()~`'),
- )
- )
- result = result.all()
+def test_custom_expression(engine, engine_using_test_dataset, table, table_using_test_dataset, query):
+ """GROUP BY clause should use labels instead of expressions"""
+ q = query(table)
+ result = engine.execute(q).fetchall()
assert len(result) > 0
+ q = query(table_using_test_dataset)
+ result = engine_using_test_dataset.execute(q).fetchall()
-def test_custom_expression(engine, query):
- """GROUP BY clause should use labels instead of expressions"""
- result = engine.execute(query).fetchall()
assert len(result) > 0
-def test_compiled_query_literal_binds(engine, query):
- compiled = query.compile(engine, compile_kwargs={"literal_binds": True})
+def test_compiled_query_literal_binds(engine, engine_using_test_dataset, table, table_using_test_dataset, query):
+ q = query(table)
+ compiled = q.compile(engine, compile_kwargs={"literal_binds": True})
result = engine.execute(compiled).fetchall()
assert len(result) > 0
+ q = query(table_using_test_dataset)
+ compiled = q.compile(engine_using_test_dataset, compile_kwargs={"literal_binds": True})
+ result = engine_using_test_dataset.execute(compiled).fetchall()
+ assert len(result) > 0
+
def test_joins(session, table, table_one_row):
result = (session.query(table.c.string, func.count(table_one_row.c.integer))
@@ -319,20 +377,29 @@ def test_dml(engine, session, table_dml):
assert len(result) == 0
-def test_schemas_names(inspector):
+def test_schemas_names(inspector, inspector_using_test_dataset):
datasets = inspector.get_schema_names()
assert 'test_pybigquery' in datasets
+ datasets = inspector_using_test_dataset.get_schema_names()
+ assert 'test_pybigquery' in datasets
-def test_table_names_in_schema(inspector):
+
+def test_table_names_in_schema(inspector, inspector_using_test_dataset):
tables = inspector.get_table_names('test_pybigquery')
assert 'test_pybigquery.sample' in tables
assert 'test_pybigquery.sample_one_row' in tables
assert 'test_pybigquery.sample_dml' in tables
assert len(tables) == 3
+ tables = inspector_using_test_dataset.get_table_names()
+ assert 'sample' in tables
+ assert 'sample_one_row' in tables
+ assert 'sample_dml' in tables
+ assert len(tables) == 3
+
-def test_get_columns(inspector):
+def test_get_columns(inspector, inspector_using_test_dataset):
columns_without_schema = inspector.get_columns('test_pybigquery.sample')
columns_schema = inspector.get_columns('sample', 'test_pybigquery')
columns_queries = [columns_without_schema, columns_schema]
@@ -344,7 +411,22 @@ def test_get_columns(inspector):
assert col['default'] == sample_col['default']
assert col['type'].__class__.__name__ == sample_col['type'].__class__.__name__
+ columns_without_schema = inspector_using_test_dataset.get_columns('sample')
+ columns_schema = inspector_using_test_dataset.get_columns('sample', 'test_pybigquery')
+ columns_queries = [columns_without_schema, columns_schema]
+ for columns in columns_queries:
+ for i, col in enumerate(columns):
+ sample_col = SAMPLE_COLUMNS[i]
+ assert col['name'] == sample_col['name']
+ assert col['nullable'] == sample_col['nullable']
+ assert col['default'] == sample_col['default']
+ assert col['type'].__class__.__name__ == sample_col['type'].__class__.__name__
-def test_has_table(engine):
+
+def test_has_table(engine, engine_using_test_dataset):
assert engine.has_table('sample', 'test_pybigquery') is True
assert engine.has_table('test_pybigquery.sample') is True
+
+ assert engine_using_test_dataset.has_table('sample') is True
+ with pytest.raises(Exception):
+ assert engine_using_test_dataset.has_table('test_pybigquery.sample') is True
\ No newline at end of file