From a3f230a139737b79ef21612f8d97f3c5b15811a5 Mon Sep 17 00:00:00 2001 From: Luka Peschke Date: Tue, 23 Jan 2024 16:32:28 +0100 Subject: [PATCH] feat(gbq): Implement a simple status check validating the private key format (#1470) Signed-off-by: Luka Peschke --- CHANGELOG.md | 4 ++ tests/conftest.py | 7 +++ .../google_big_query/test_google_big_query.py | 52 ++++++++++++++----- tests/utils/test_pem.py | 6 --- .../google_big_query_connector.py | 12 ++++- 5 files changed, 61 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 06e29c122..fdbd72c78 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index 1bfd69723..a027eb59f 100644 --- a/tests/conftest.py +++ b/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(): """ + + +@pytest.fixture +def sanitized_pem_key() -> str: + with open(join(dirname(__file__), 'utils', 'fixtures', 'sanitized_pem_key.pem'), 'r') as f: + return f.read() diff --git a/tests/google_big_query/test_google_big_query.py b/tests/google_big_query/test_google_big_query.py index 37845abe4..5559d69ba 100644 --- a/tests/google_big_query/test_google_big_query.py +++ b/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 diff --git a/tests/utils/test_pem.py b/tests/utils/test_pem.py index 7c2817266..00d292539 100644 --- a/tests/utils/test_pem.py +++ b/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: diff --git a/toucan_connectors/google_big_query/google_big_query_connector.py b/toucan_connectors/google_big_query/google_big_query_connector.py index 8404a7522..844d143c3 100644 --- a/toucan_connectors/google_big_query/google_big_query_connector.py +++ b/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))