From ad1be92583502d5eea34cbb69695bda403be6679 Mon Sep 17 00:00:00 2001 From: sanix-darker Date: Fri, 22 Sep 2023 12:19:05 +0200 Subject: [PATCH] fix: set the appropriate project-id on GoCred/JWT auth mode. --- .../google_big_query/test_google_big_query.py | 23 +++++++++++++++++-- .../google_big_query_connector.py | 9 +++++++- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/tests/google_big_query/test_google_big_query.py b/tests/google_big_query/test_google_big_query.py index 7667fed6f..285ed06f0 100644 --- a/tests/google_big_query/test_google_big_query.py +++ b/tests/google_big_query/test_google_big_query.py @@ -48,7 +48,7 @@ def _fixture_credentials() -> GoogleCredentials: @pytest.fixture def _jwt_fixture_credentials() -> JWTCredentials: my_credentials = JWTCredentials( - project_id='my_project_id', + project_id='THE_JWT_project_id', jwt_token='valid-jwt', ) return my_credentials @@ -731,7 +731,11 @@ 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: MockFixture, + _fixture_credentials: MockFixture, + _jwt_fixture_credentials: MockFixture +) -> None: def mock_available_schs(): return ['ok', 'test'] @@ -755,6 +759,21 @@ def mock_available_schs(): == 'my_project_id' ) + assert ( + GoogleBigQueryDataSource(query=',', name='MyGBQ-WITH-JWT', domain='foo').get_form( + GoogleBigQueryConnector( + name='MyGBQ', + jwt_credentials=_jwt_fixture_credentials, + scopes=[ + 'https://www.googleapis.com/auth/bigquery', + ], + ), + {}, + )['properties']['database']['default'] + == 'THE_JWT_project_id' + ) + + @pytest.mark.parametrize( 'input_query, expected_output', 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 c32dd744d..0a37a2519 100644 --- a/toucan_connectors/google_big_query/google_big_query_connector.py +++ b/toucan_connectors/google_big_query/google_big_query_connector.py @@ -101,7 +101,14 @@ def get_form(cls, connector: 'GoogleBigQueryConnector', current_config: dict[str db_schema=strlist_to_enum('db_schema', connector._available_schs), __base__=cls, ).schema() - schema['properties']['database']['default'] = connector.credentials.project_id + + project_id = '' + if connector.jwt_credentials: + project_id = connector.jwt_credentials.project_id + elif connector.credentials: + project_id = connector.credentials.project_id + + schema['properties']['database']['default'] = project_id return schema