fix(connectors): BI-5777 handle google api errors & bad creds error i…

…n BQ connector (#601) * fix(connectors) BI-5777 handle google api errors & bad creds error in BQ connector * remove debug print (shame on me)
datalens-tech · Sep 6, 2024 · 3e78eff · 3e78eff
1 parent d467d40
commit 3e78eff
Show file tree

Hide file tree

Showing 7 changed files with 71 additions and 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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"]