diff --git a/pybigquery/_helpers.py b/pybigquery/_helpers.py new file mode 100644 index 00000000..35f7e4ab --- /dev/null +++ b/pybigquery/_helpers.py @@ -0,0 +1,60 @@ +# Copyright 2021 The PyBigQuery Authors +# +# Use of this source code is governed by an MIT-style +# license that can be found in the LICENSE file or at +# https://opensource.org/licenses/MIT. + +from google.api_core import client_info +import google.auth +from google.cloud import bigquery +from google.oauth2 import service_account +import sqlalchemy + + +USER_AGENT_TEMPLATE = "sqlalchemy/{}" +SCOPES = ( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/drive", +) + + +def google_client_info(): + user_agent = USER_AGENT_TEMPLATE.format(sqlalchemy.__version__) + return client_info.ClientInfo(user_agent=user_agent) + + +def create_bigquery_client( + credentials_info=None, + credentials_path=None, + default_query_job_config=None, + location=None, + project_id=None, +): + default_project = None + + if credentials_path: + credentials = service_account.Credentials.from_service_account_file( + credentials_path + ) + credentials = credentials.with_scopes(SCOPES) + default_project = credentials.project + elif credentials_info: + credentials = service_account.Credentials.from_service_account_info( + credentials_info + ) + credentials = credentials.with_scopes(SCOPES) + default_project = credentials.project + else: + credentials, default_project = google.auth.default(scopes=SCOPES) + + if project_id is None: + project_id = default_project + + return bigquery.Client( + client_info=google_client_info(), + project=project_id, + credentials=credentials, + location=location, + default_query_job_config=default_query_job_config, + ) diff --git a/pybigquery/api.py b/pybigquery/api.py index ae886612..f8e1894e 100644 --- a/pybigquery/api.py +++ b/pybigquery/api.py @@ -22,19 +22,18 @@ from __future__ import absolute_import from __future__ import unicode_literals -from google.cloud.bigquery import Client, QueryJobConfig +from google.cloud.bigquery import QueryJobConfig + +from pybigquery import _helpers class ApiClient(object): 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, location=self.location - ) - else: - self.client = Client(location=self.location) + self.client = _helpers.create_bigquery_client( + credentials_path=credentials_path, location=location + ) def dry_run_query(self, query): job_config = QueryJobConfig() diff --git a/pybigquery/sqlalchemy_bigquery.py b/pybigquery/sqlalchemy_bigquery.py index ff83f319..e7caddd8 100644 --- a/pybigquery/sqlalchemy_bigquery.py +++ b/pybigquery/sqlalchemy_bigquery.py @@ -25,11 +25,9 @@ import operator from google import auth -from google.cloud import bigquery from google.cloud.bigquery import dbapi from google.cloud.bigquery.schema import SchemaField from google.cloud.bigquery.table import TableReference -from google.oauth2 import service_account from google.api_core.exceptions import NotFound from sqlalchemy.exc import NoSuchTableError from sqlalchemy import types, util @@ -46,6 +44,7 @@ import re from .parse_url import parse_url +from pybigquery import _helpers FIELD_ILLEGAL_CHARACTERS = re.compile(r"[^\w]+") @@ -342,30 +341,6 @@ def _add_default_dataset_to_job_config(job_config, project_id, dataset_id): job_config.default_dataset = "{}.{}".format(project_id, dataset_id) - def _create_client_from_credentials( - self, credentials, default_query_job_config, project_id - ): - if project_id is None: - project_id = credentials.project_id - - scopes = ( - "https://www.googleapis.com/auth/bigquery", - "https://www.googleapis.com/auth/cloud-platform", - "https://www.googleapis.com/auth/drive", - ) - credentials = credentials.with_scopes(scopes) - - self._add_default_dataset_to_job_config( - default_query_job_config, project_id, self.dataset_id - ) - - return bigquery.Client( - project=project_id, - credentials=credentials, - location=self.location, - default_query_job_config=default_query_job_config, - ) - def create_connect_args(self, url): ( project_id, @@ -380,34 +355,16 @@ def create_connect_args(self, url): self.location = location or self.location self.credentials_path = credentials_path or self.credentials_path self.dataset_id = dataset_id - - if self.credentials_path: - credentials = service_account.Credentials.from_service_account_file( - self.credentials_path - ) - client = self._create_client_from_credentials( - credentials, default_query_job_config, project_id - ) - - elif self.credentials_info: - credentials = service_account.Credentials.from_service_account_info( - self.credentials_info - ) - client = self._create_client_from_credentials( - credentials, default_query_job_config, project_id - ) - - else: - self._add_default_dataset_to_job_config( - default_query_job_config, project_id, dataset_id - ) - - client = bigquery.Client( - project=project_id, - location=self.location, - default_query_job_config=default_query_job_config, - ) - + self._add_default_dataset_to_job_config( + default_query_job_config, project_id, dataset_id + ) + client = _helpers.create_bigquery_client( + credentials_path=self.credentials_path, + credentials_info=self.credentials_info, + project_id=project_id, + location=self.location, + default_query_job_config=default_query_job_config, + ) return ([client], {}) def _json_deserializer(self, row): diff --git a/setup.py b/setup.py index f1f87a53..dc0f5b42 100644 --- a/setup.py +++ b/setup.py @@ -65,7 +65,8 @@ def readme(): platforms="Posix; MacOS X; Windows", install_requires=[ "sqlalchemy>=1.1.9,<1.4.0dev", - "google-cloud-bigquery>=1.6.0", + "google-auth>=1.2.0,<2.0dev", + "google-cloud-bigquery>=1.12.0", "future", ], python_requires=">=3.6, <3.10", diff --git a/testing/constraints-3.6.txt b/testing/constraints-3.6.txt index 264a6efd..ab72cf88 100644 --- a/testing/constraints-3.6.txt +++ b/testing/constraints-3.6.txt @@ -5,4 +5,5 @@ # # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", sqlalchemy==1.1.9 -google-cloud-bigquery==1.6.0 +google-auth==1.2.0 +google-cloud-bigquery==1.12.0 diff --git a/tests/unit/test_helpers.py b/tests/unit/test_helpers.py new file mode 100644 index 00000000..38be0453 --- /dev/null +++ b/tests/unit/test_helpers.py @@ -0,0 +1,138 @@ +# Copyright 2021 The PyBigQuery Authors +# +# Use of this source code is governed by an MIT-style +# license that can be found in the LICENSE file or at +# https://opensource.org/licenses/MIT. + +from unittest import mock + +import google.auth +import google.auth.credentials +from google.oauth2 import service_account +import pytest + + +class AnonymousCredentialsWithProject(google.auth.credentials.AnonymousCredentials): + """Fake credentials to trick isinstance""" + + def __init__(self, project): + super().__init__() + self.project = project + + def with_scopes(self, scopes): + return self + + +@pytest.fixture(scope="session") +def module_under_test(): + from pybigquery import _helpers + + return _helpers + + +def test_create_bigquery_client_with_credentials_path(monkeypatch, module_under_test): + mock_service_account = mock.create_autospec(service_account.Credentials) + mock_service_account.from_service_account_file.return_value = AnonymousCredentialsWithProject( + "service-account-project" + ) + monkeypatch.setattr(service_account, "Credentials", mock_service_account) + + bqclient = module_under_test.create_bigquery_client( + credentials_path="path/to/key.json", + ) + + assert bqclient.project == "service-account-project" + + +def test_create_bigquery_client_with_credentials_path_respects_project( + monkeypatch, module_under_test +): + """Test that project_id is used, even when there is a default project. + + https://github.com/googleapis/python-bigquery-sqlalchemy/issues/48 + """ + mock_service_account = mock.create_autospec(service_account.Credentials) + mock_service_account.from_service_account_file.return_value = AnonymousCredentialsWithProject( + "service-account-project" + ) + monkeypatch.setattr(service_account, "Credentials", mock_service_account) + + bqclient = module_under_test.create_bigquery_client( + credentials_path="path/to/key.json", project_id="connection-url-project", + ) + + assert bqclient.project == "connection-url-project" + + +def test_create_bigquery_client_with_credentials_info(monkeypatch, module_under_test): + mock_service_account = mock.create_autospec(service_account.Credentials) + mock_service_account.from_service_account_info.return_value = AnonymousCredentialsWithProject( + "service-account-project" + ) + monkeypatch.setattr(service_account, "Credentials", mock_service_account) + + bqclient = module_under_test.create_bigquery_client( + credentials_info={ + "type": "service_account", + "project_id": "service-account-project", + }, + ) + + assert bqclient.project == "service-account-project" + + +def test_create_bigquery_client_with_credentials_info_respects_project( + monkeypatch, module_under_test +): + """Test that project_id is used, even when there is a default project. + + https://github.com/googleapis/python-bigquery-sqlalchemy/issues/48 + """ + mock_service_account = mock.create_autospec(service_account.Credentials) + mock_service_account.from_service_account_info.return_value = AnonymousCredentialsWithProject( + "service-account-project" + ) + monkeypatch.setattr(service_account, "Credentials", mock_service_account) + + bqclient = module_under_test.create_bigquery_client( + credentials_info={ + "type": "service_account", + "project_id": "service-account-project", + }, + project_id="connection-url-project", + ) + + assert bqclient.project == "connection-url-project" + + +def test_create_bigquery_client_with_default_credentials( + monkeypatch, module_under_test +): + def mock_default_credentials(*args, **kwargs): + return (google.auth.credentials.AnonymousCredentials(), "default-project") + + monkeypatch.setattr(google.auth, "default", mock_default_credentials) + + bqclient = module_under_test.create_bigquery_client() + + assert bqclient.project == "default-project" + + +def test_create_bigquery_client_with_default_credentials_respects_project( + monkeypatch, module_under_test +): + """Test that project_id is used, even when there is a default project. + + https://github.com/googleapis/python-bigquery-sqlalchemy/issues/48 + """ + + def mock_default_credentials(*args, **kwargs): + return (google.auth.credentials.AnonymousCredentials(), "default-project") + + monkeypatch.setattr(google.auth, "default", mock_default_credentials) + + bqclient = module_under_test.create_bigquery_client( + project_id="connection-url-project", + ) + + assert bqclient.project == "connection-url-project"