feat(gbq): Implement a simple status check validating the private key…

… format (#1470)

Signed-off-by: Luka Peschke <luka.peschke@toucantoco.com>

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog (Pypi package)
 
+### Changed
+
+- Google Big Query: A simple status check that validates the private key's format has been implemented
+
 ## [3.23.25] 2024-01-17
 
 ### Fixed

tests/conftest.py

@@ -2,6 +2,7 @@
 import time
 from contextlib import suppress
 from os import environ, getenv, path
+from os.path import dirname, join
 from typing import Any
 
 import pytest
@@ -208,3 +209,9 @@ def xml_response():
     </ListOfLanguagesByCodeResponse>
   </soap:Body>
 </soap:Envelope>"""
+
+
+@pytest.fixture
+def sanitized_pem_key() -> str:
+    with open(join(dirname(__file__), 'utils', 'fixtures', 'sanitized_pem_key.pem'), 'r') as f:
+        return f.read()

tests/google_big_query/test_google_big_query.py

@@ -10,7 +10,8 @@
 from google.cloud.bigquery.table import RowIterator
 from google.oauth2.service_account import Credentials
 from pandas.util.testing import assert_frame_equal  # <-- for testing dataframes
-from pytest_mock import MockFixture
+from pydantic import SecretStr
+from pytest_mock import MockerFixture, MockFixture
 
 from toucan_connectors.google_big_query.google_big_query_connector import (
     GoogleBigQueryConnector,
@@ -21,7 +22,7 @@
 
 
 @pytest.fixture
-def _fixture_credentials():
+def fixture_credentials() -> GoogleCredentials:
     my_credentials = GoogleCredentials(
         type='my_type',
         project_id='my_project_id',
@@ -38,7 +39,7 @@ def _fixture_credentials():
 
 
 @pytest.fixture
-def _fixture_scope():
+def fixture_scope():
     scopes = [
         'https://www.googleapis.com/auth/bigquery',
         'https://www.googleapis.com/auth/drive',
@@ -112,9 +113,9 @@ def test_prepare_parameters_empty():
 
 @patch('google.cloud.bigquery.Client', autospec=True)
 @patch('cryptography.hazmat.primitives.serialization.load_pem_private_key')
-def test_connect(load_pem_private_key, client, _fixture_credentials, _fixture_scope):
+def test_connect(load_pem_private_key, client, fixture_credentials, fixture_scope):
     credentials = GoogleBigQueryConnector._get_google_credentials(
-        _fixture_credentials, _fixture_scope
+        fixture_credentials, fixture_scope
     )
     assert isinstance(credentials, Credentials)
     connection = GoogleBigQueryConnector._connect(credentials)
@@ -157,10 +158,10 @@ def test_execute_error(client, execute, result, to_dataframe):
     'toucan_connectors.google_big_query.google_big_query_connector.GoogleBigQueryConnector._execute_query',
     return_value=pandas.DataFrame({'a': [1, 1], 'b': [2, 2]}),
 )
-def test_retrieve_data(execute, connect, credentials, _fixture_credentials):
+def test_retrieve_data(execute, connect, credentials, fixture_credentials):
     connector = GoogleBigQueryConnector(
         name='MyGBQ',
-        credentials=_fixture_credentials,
+        credentials=fixture_credentials,
         scopes=[
             'https://www.googleapis.com/auth/bigquery',
             'https://www.googleapis.com/auth/drive',
@@ -176,7 +177,7 @@ def test_retrieve_data(execute, connect, credentials, _fixture_credentials):
     assert_frame_equal(pandas.DataFrame({'a': [1, 1], 'b': [2, 2]}), result)
 
 
-def test_get_model(mocker: MockFixture, _fixture_credentials) -> None:
+def test_get_model(mocker: MockFixture, fixture_credentials) -> None:
     class FakeResponse:
         def __init__(self) -> None:
             ...
@@ -283,7 +284,7 @@ def to_dataframe(self) -> Generator[Any, Any, Any]:
     )
     connector = GoogleBigQueryConnector(
         name='MyGBQ',
-        credentials=_fixture_credentials,
+        credentials=fixture_credentials,
         scopes=[
             'https://www.googleapis.com/auth/bigquery',
             'https://www.googleapis.com/auth/drive',
@@ -482,7 +483,7 @@ def to_dataframe(self) -> Generator[Any, Any, Any]:
     )
 
 
-def test_get_model_multi_location(mocker: MockFixture, _fixture_credentials) -> None:
+def test_get_model_multi_location(mocker: MockFixture, fixture_credentials) -> None:
     fake_resp_1 = mocker.MagicMock()
     fake_resp_1.to_dataframe.return_value = pd.DataFrame(
         [
@@ -550,7 +551,7 @@ def test_get_model_multi_location(mocker: MockFixture, _fixture_credentials) ->
     )
     connector = GoogleBigQueryConnector(
         name='MyGBQ',
-        credentials=_fixture_credentials,
+        credentials=fixture_credentials,
         scopes=[
             'https://www.googleapis.com/auth/bigquery',
             'https://www.googleapis.com/auth/drive',
@@ -662,7 +663,7 @@ def test_get_model_multi_location(mocker: MockFixture, _fixture_credentials) ->
     assert mocked_query.call_args_list[2][1] == {'location': 'Toulouse'}
 
 
-def test_get_form(mocker: MockFixture, _fixture_credentials: MockFixture) -> None:
+def test_get_form(mocker: MockerFixture, fixture_credentials: GoogleCredentials) -> None:
     def mock_available_schs():
         return ['ok', 'test']
 
@@ -675,7 +676,7 @@ def mock_available_schs():
         GoogleBigQueryDataSource(query=',', name='MyGBQ', domain='foo').get_form(
             GoogleBigQueryConnector(
                 name='MyGBQ',
-                credentials=_fixture_credentials,
+                credentials=fixture_credentials,
                 scopes=[
                     'https://www.googleapis.com/auth/bigquery',
                     'https://www.googleapis.com/auth/drive',
@@ -685,3 +686,28 @@ def mock_available_schs():
         )['properties']['database']['default']
         == 'my_project_id'
     )
+
+
+def test_get_status(
+    mocker: MockerFixture, fixture_credentials: GoogleCredentials, sanitized_pem_key: str
+) -> None:
+    connector = GoogleBigQueryConnector(
+        name='MyGBQ',
+        credentials=fixture_credentials,
+        scopes=[
+            'https://www.googleapis.com/auth/bigquery',
+            'https://www.googleapis.com/auth/drive',
+        ],
+    )
+
+    status = connector.get_status()
+    assert status.status is False
+    assert status.details == [('Private key validity', False)]
+    assert status.error is not None
+    assert 'Could not deserialize key data' in status.error
+
+    connector.credentials.private_key = SecretStr(sanitized_pem_key)
+    status = connector.get_status()
+    assert status.status is True
+    assert status.details == [('Private key validity', True)]
+    assert status.error is None

tests/utils/test_pem.py

@@ -11,12 +11,6 @@ def pem_key_with_spaces() -> str:
         return f.read()
 
 
-@pytest.fixture
-def sanitized_pem_key() -> str:
-    with open(join(dirname(__file__), 'fixtures', 'sanitized_pem_key.pem'), 'r') as f:
-        return f.read()
-
-
 @pytest.fixture
 def pem_bundle_with_spaces() -> str:
     with open(join(dirname(__file__), 'fixtures', 'pem_bundle_with_spaces.pem'), 'r') as f:

toucan_connectors/google_big_query/google_big_query_connector.py

@@ -13,7 +13,7 @@
 from google.oauth2.service_account import Credentials
 from pydantic import Field, create_model
 
-from toucan_connectors.common import sanitize_query
+from toucan_connectors.common import ConnectorStatus, sanitize_query
 from toucan_connectors.google_credentials import GoogleCredentials, get_google_oauth2_credentials
 from toucan_connectors.toucan_connector import (
     DiscoverableConnector,
@@ -77,6 +77,9 @@ def _define_query_param(name: str, value: Any) -> BigQueryParam:
         return bigquery_helpers.scalar_to_query_parameter(value=value, name=name)
 
 
+_KEY_CHECK_NAME = 'Private key validity'
+
+
 class GoogleBigQueryConnector(ToucanConnector, DiscoverableConnector):
     data_source_model: GoogleBigQueryDataSource
 
@@ -349,3 +352,10 @@ def get_model(
     ) -> list[TableInfo]:
         """Retrieves the database tree structure using current connection"""
         return self._get_project_structure(db_name, schema_name)
+
+    def get_status(self) -> ConnectorStatus:
+        try:
+            get_google_oauth2_credentials(self.credentials)
+            return ConnectorStatus(status=True, details=[(_KEY_CHECK_NAME, True)], error=None)
+        except Exception as exc:
+            return ConnectorStatus(status=False, details=[(_KEY_CHECK_NAME, False)], error=str(exc))