From 3e78effbe6d34eeb50235a8bf0faa923dada92d1 Mon Sep 17 00:00:00 2001 From: KonstantAnxiety <58992437+KonstantAnxiety@users.noreply.github.com> Date: Fri, 6 Sep 2024 18:02:21 +0300 Subject: [PATCH] fix(connectors): BI-5777 handle google api errors & bad creds error in BQ connector (#601) * fix(connectors) BI-5777 handle google api errors & bad creds error in BQ connector * remove debug print (shame on me) --- lib/dl_api_lib/dl_api_lib/error_handling.py | 1 + .../dl_connector_bigquery/core/adapters.py | 11 +++++- .../dl_connector_bigquery/core/sa_types.py | 2 +- .../core/type_transformer.py | 1 + .../ext/api/base.py | 18 ++++++++++ .../ext/api/test_connection.py | 36 ++++++++++++++++++- lib/dl_core/dl_core/exc.py | 5 +++ 7 files changed, 71 insertions(+), 3 deletions(-) diff --git a/lib/dl_api_lib/dl_api_lib/error_handling.py b/lib/dl_api_lib/dl_api_lib/error_handling.py index 7ccb9e6f9..3a8b800a2 100644 --- a/lib/dl_api_lib/dl_api_lib/error_handling.py +++ b/lib/dl_api_lib/dl_api_lib/error_handling.py @@ -54,6 +54,7 @@ common_exc.QueryConstructorError: status.BAD_REQUEST, common_exc.ResultRowCountLimitExceeded: status.BAD_REQUEST, common_exc.TableNameNotConfiguredError: status.BAD_REQUEST, + common_exc.MalformedCredentialsError: status.BAD_REQUEST, common_exc.NotAvailableError: status.BAD_REQUEST, common_exc.InvalidFieldError: status.BAD_REQUEST, common_exc.FieldNotFound: status.BAD_REQUEST, diff --git a/lib/dl_connector_bigquery/dl_connector_bigquery/core/adapters.py b/lib/dl_connector_bigquery/dl_connector_bigquery/core/adapters.py index ed9fc58b1..c031b8d26 100644 --- a/lib/dl_connector_bigquery/dl_connector_bigquery/core/adapters.py +++ b/lib/dl_connector_bigquery/dl_connector_bigquery/core/adapters.py @@ -7,13 +7,16 @@ Tuple, ) +from google.api_core.exceptions import GoogleAPIError from google.auth.credentials import Credentials as BQCredentials from google.cloud.bigquery import Client as BQClient import google.oauth2.service_account as g_service_account import sqlalchemy as sa import sqlalchemy_bigquery._types as bq_types +from dl_core import exc from dl_core.connection_executors.adapters.adapters_base_sa_classic import ( + LOGGER, BaseClassicAdapter, BaseConnLineConstructor, ) @@ -43,6 +46,8 @@ def _get_dsn_params( class BigQueryDefaultAdapter(BaseClassicAdapter[BigQueryConnTargetDTO]): + EXTRA_EXC_CLS = (GoogleAPIError,) + conn_type = CONNECTION_TYPE_BIGQUERY dsn_template = "{dialect}://{project_id}" conn_line_constructor_type = BigQueryConnLineConstructor @@ -72,7 +77,11 @@ def get_engine_kwargs(self) -> dict: } def _get_bq_credentials(self) -> BQCredentials: - credentials_info = json.loads(base64.b64decode(self._target_dto.credentials)) + try: + credentials_info = json.loads(base64.b64decode(self._target_dto.credentials)) + except (json.JSONDecodeError, UnicodeDecodeError) as e: + LOGGER.info("Got malformed bigquery credentials", exc_info=True) + raise exc.MalformedCredentialsError() from e credentials = g_service_account.Credentials.from_service_account_info(credentials_info) return credentials diff --git a/lib/dl_connector_bigquery/dl_connector_bigquery/core/sa_types.py b/lib/dl_connector_bigquery/dl_connector_bigquery/core/sa_types.py index d5862e018..34162be34 100644 --- a/lib/dl_connector_bigquery/dl_connector_bigquery/core/sa_types.py +++ b/lib/dl_connector_bigquery/dl_connector_bigquery/core/sa_types.py @@ -8,7 +8,7 @@ SQLALCHEMY_BIGQUERY_TYPES = { (BACKEND_TYPE_BIGQUERY, make_native_type(bq_types.DATE)): simple_instantiator(bq_types.DATE), - (BACKEND_TYPE_BIGQUERY, make_native_type(bq_types.DATETIME)): simple_instantiator(bq_types.DATE), + (BACKEND_TYPE_BIGQUERY, make_native_type(bq_types.DATETIME)): simple_instantiator(bq_types.DATETIME), (BACKEND_TYPE_BIGQUERY, make_native_type(bq_types.STRING)): simple_instantiator(bq_types.STRING), (BACKEND_TYPE_BIGQUERY, make_native_type(bq_types.BOOLEAN)): simple_instantiator(bq_types.BOOLEAN), (BACKEND_TYPE_BIGQUERY, make_native_type(bq_types.INTEGER)): simple_instantiator(bq_types.INTEGER), diff --git a/lib/dl_connector_bigquery/dl_connector_bigquery/core/type_transformer.py b/lib/dl_connector_bigquery/dl_connector_bigquery/core/type_transformer.py index 57248c09c..31cbc0049 100644 --- a/lib/dl_connector_bigquery/dl_connector_bigquery/core/type_transformer.py +++ b/lib/dl_connector_bigquery/dl_connector_bigquery/core/type_transformer.py @@ -24,6 +24,7 @@ class BigQueryTypeTransformer(TypeTransformer): UserDataType.date: make_native_type(bq_types.DATE), UserDataType.genericdatetime: make_native_type(bq_types.DATETIME), UserDataType.string: make_native_type(bq_types.STRING), + UserDataType.uuid: make_native_type(bq_types.STRING), UserDataType.boolean: make_native_type(bq_types.BOOLEAN), UserDataType.integer: make_native_type(bq_types.INTEGER), UserDataType.float: make_native_type(bq_types.FLOAT), diff --git a/lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/base.py b/lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/base.py index 7a25d8d1e..eff39f08c 100644 --- a/lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/base.py +++ b/lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/base.py @@ -29,6 +29,24 @@ def connection_params(self, bq_secrets) -> dict: ) +class BigQueryConnectionTestMalformedCreds(BigQueryConnectionTestBase): + @pytest.fixture(scope="class") + def connection_params(self, bq_secrets) -> dict: + return dict( + project_id=bq_secrets.get_project_id(), + credentials=bq_secrets.get_creds() + "asdf", + ) + + +class BigQueryConnectionTestBadProjectId(BigQueryConnectionTestBase): + @pytest.fixture(scope="class") + def connection_params(self, bq_secrets) -> dict: + return dict( + project_id=bq_secrets.get_project_id() + "123", + credentials=bq_secrets.get_creds(), + ) + + class BigQueryDatasetTestBase(BigQueryConnectionTestBase, DatasetTestBase): @pytest.fixture(scope="class") def dataset_params(self, sample_table) -> dict: diff --git a/lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/test_connection.py b/lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/test_connection.py index 15f02e05b..d6a7b2b66 100644 --- a/lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/test_connection.py +++ b/lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/test_connection.py @@ -1,7 +1,41 @@ +from typing import Optional + +from dl_api_client.dsmaker.api.http_sync_base import SyncHttpClientBase +from dl_api_lib_testing.connection_base import ConnectionTestBase from dl_api_lib_testing.connector.connection_suite import DefaultConnectorConnectionTestSuite +from dl_testing.regulated_test import RegulatedTestCase -from dl_connector_bigquery_tests.ext.api.base import BigQueryConnectionTestBase +from dl_connector_bigquery_tests.ext.api.base import ( + BigQueryConnectionTestBadProjectId, + BigQueryConnectionTestBase, + BigQueryConnectionTestMalformedCreds, +) class TestBigQueryConnection(BigQueryConnectionTestBase, DefaultConnectorConnectionTestSuite): pass + + +class ErrorHandlingTestBase(BigQueryConnectionTestBase, ConnectionTestBase, RegulatedTestCase): + def test_connection_sources_error( + self, + control_api_sync_client: SyncHttpClientBase, + saved_connection_id: str, + bi_headers: Optional[dict[str, str]], + ) -> None: + resp = control_api_sync_client.get( + url=f"/api/v1/connections/{saved_connection_id}/info/sources", + headers=bi_headers, + ) + assert resp.status_code == 400, resp.json + resp_data = resp.json + assert "code" in resp_data, resp_data + assert "message" in resp_data, resp_data + + +class TestErrorHandlingMalformedCreds(BigQueryConnectionTestMalformedCreds, ErrorHandlingTestBase): + pass + + +class TestErrorHandlingBadProjectId(BigQueryConnectionTestBadProjectId, ErrorHandlingTestBase): + pass diff --git a/lib/dl_core/dl_core/exc.py b/lib/dl_core/dl_core/exc.py index c9ca79fa8..e8e1e0e77 100644 --- a/lib/dl_core/dl_core/exc.py +++ b/lib/dl_core/dl_core/exc.py @@ -56,6 +56,11 @@ class TableNameInvalidError(DataSourceConfigurationError): err_code = DataSourceConfigurationError.err_code + ["TABLE_NAME_INVALID"] +class MalformedCredentialsError(DataSourceConfigurationError): + err_code = DataSourceConfigurationError.err_code + ["MALFORMED_CREDS"] + default_message = "Malformed credentials" + + class DatasetConfigurationError(DLBaseException): err_code = DLBaseException.err_code + ["DS_CONFIG"]