From 35511debb8ea6fed788d44cf85ccff4dac2d740f Mon Sep 17 00:00:00 2001 From: Omri Bromberg Date: Wed, 12 May 2021 23:50:30 +0300 Subject: [PATCH] add and fix tests, reformatting --- pybigquery/parse_url.py | 26 +++++++++++++++++++++++--- pybigquery/sqlalchemy_bigquery.py | 6 ++++-- tests/unit/fauxdbi.py | 2 +- tests/unit/test_parse_url.py | 26 ++++++++++++++++++++++++-- 4 files changed, 52 insertions(+), 8 deletions(-) diff --git a/pybigquery/parse_url.py b/pybigquery/parse_url.py index e9de35fb..aee1b02b 100644 --- a/pybigquery/parse_url.py +++ b/pybigquery/parse_url.py @@ -91,7 +91,10 @@ def parse_url(url): # noqa: C901 try: list_tables_page_size = int(str_list_tables_page_size) except ValueError: - raise ValueError("invalid int in url query list_tables_page_size: " + str_list_tables_page_size) + raise ValueError( + "invalid int in url query list_tables_page_size: " + + str_list_tables_page_size + ) # if only these "non-config" values were present, the dict will now be empty if not query: @@ -105,9 +108,18 @@ def parse_url(url): # noqa: C901 arraysize, credentials_path, QueryJobConfig(), + list_tables_page_size, ) else: - return project_id, location, dataset_id, arraysize, credentials_path, None + return ( + project_id, + location, + dataset_id, + arraysize, + credentials_path, + None, + list_tables_page_size, + ) job_config = QueryJobConfig() @@ -247,4 +259,12 @@ def parse_url(url): # noqa: C901 "invalid write_disposition in url query: " + query["write_disposition"] ) - return project_id, location, dataset_id, arraysize, credentials_path, job_config, list_tables_page_size + return ( + project_id, + location, + dataset_id, + arraysize, + credentials_path, + job_config, + list_tables_page_size, + ) diff --git a/pybigquery/sqlalchemy_bigquery.py b/pybigquery/sqlalchemy_bigquery.py index 204936d1..6ae02b9a 100644 --- a/pybigquery/sqlalchemy_bigquery.py +++ b/pybigquery/sqlalchemy_bigquery.py @@ -603,7 +603,7 @@ def create_connect_args(self, url): arraysize, credentials_path, default_query_job_config, - list_tables_page_size + list_tables_page_size, ) = parse_url(url) self.arraysize = arraysize or self.arraysize @@ -648,7 +648,9 @@ def _get_table_or_view_names(self, connection, table_type, schema=None): continue try: - tables = client.list_tables(dataset.reference, max_results=self.list_tables_page_size) + tables = client.list_tables( + dataset.reference, max_results=self.list_tables_page_size + ) for table in tables: if table_type == table.table_type: result.append(get_table_name(table)) diff --git a/tests/unit/fauxdbi.py b/tests/unit/fauxdbi.py index 70cbb8aa..083b87a0 100644 --- a/tests/unit/fauxdbi.py +++ b/tests/unit/fauxdbi.py @@ -415,7 +415,7 @@ def list_datasets(self): google.cloud.bigquery.Dataset("myproject.yourdataset"), ] - def list_tables(self, dataset): + def list_tables(self, dataset, max_results): with contextlib.closing(self.connection.connection.cursor()) as cursor: cursor.execute("select * from sqlite_master") return [ diff --git a/tests/unit/test_parse_url.py b/tests/unit/test_parse_url.py index bf9f8855..1cb94158 100644 --- a/tests/unit/test_parse_url.py +++ b/tests/unit/test_parse_url.py @@ -50,6 +50,7 @@ def url_with_everything(): "?credentials_path=/some/path/to.json" "&location=some-location" "&arraysize=1000" + "&list_tables_page_size=5000" "&clustering_fields=a,b,c" "&create_disposition=CREATE_IF_NEEDED" "&destination=different-project.different-dataset.table" @@ -72,12 +73,14 @@ def test_basic(url_with_everything): arraysize, credentials_path, job_config, + list_tables_page_size, ) = parse_url(url_with_everything) assert project_id == "some-project" assert location == "some-location" assert dataset_id == "some-dataset" assert arraysize == 1000 + assert list_tables_page_size == 5000 assert credentials_path == "/some/path/to.json" assert isinstance(job_config, QueryJobConfig) @@ -134,6 +137,7 @@ def test_all_values(url_with_everything, param, value, default): "param, value", [ ("arraysize", "not-int"), + ("list_tables_page_size", "not-int"), ("create_disposition", "not-attribute"), ("destination", "not.fully-qualified"), ("dry_run", "not-bool"), @@ -165,7 +169,15 @@ def test_empty_with_non_config(): "bigquery:///?location=some-location&arraysize=1000&credentials_path=/some/path/to.json" ) ) - project_id, location, dataset_id, arraysize, credentials_path, job_config = url + ( + project_id, + location, + dataset_id, + arraysize, + credentials_path, + job_config, + list_tables_page_size, + ) = url assert project_id is None assert location == "some-location" @@ -173,17 +185,27 @@ def test_empty_with_non_config(): assert arraysize == 1000 assert credentials_path == "/some/path/to.json" assert job_config is None + assert list_tables_page_size 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, + job_config, + list_tables_page_size, + ) = url assert project_id is None assert location is None assert dataset_id == "some-dataset" assert arraysize is None assert credentials_path is None + assert list_tables_page_size 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