diff --git a/CHANGELOG.md b/CHANGELOG.md index 1f0e1b955..c399a87a9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## Unreleased +## Fixed + +- Google Big Query: get project_id from connector config whatever auth mode (JWT/GoogleCreds). + ### [4.9.1] 2023-09-22 ## Fixed diff --git a/tests/google_big_query/test_google_big_query.py b/tests/google_big_query/test_google_big_query.py index 64851b99d..a0c72fe43 100644 --- a/tests/google_big_query/test_google_big_query.py +++ b/tests/google_big_query/test_google_big_query.py @@ -771,6 +771,37 @@ def mock_available_schs(): == 'THE_JWT_project_id' ) + # Now let say the JWT is not set or bad, or the project_id is also bad or + # not set, we need to fallback on project_id provided by the GoogleCreds + assert ( + GoogleBigQueryDataSource(query=',', name='MyGBQ-WITH-JWT', domain='foo').get_form( + GoogleBigQueryConnector( + name='MyGBQ', + credentials=GoogleCredentials( + type='my_type', + project_id='THE_GOOGLE_CREDS_project_id', + private_key_id='my_private_key_id', + private_key='my_private_key', + client_email='my_client_email@email.com', + client_id='my_client_id', + auth_uri='https://accounts.google.com/o/oauth2/auth', + token_uri='https://oauth2.googleapis.com/token', + auth_provider_x509_cert_url='https://www.googleapis.com/oauth2/v1/certs', + client_x509_cert_url='https://www.googleapis.com/robot/v1/metadata/x509/pika.com', + ), + jwt_credentials=JWTCredentials( + project_id='THE_JWT_project_id', + jwt_token='', # the jwt is empty + ), + scopes=[ + 'https://www.googleapis.com/auth/bigquery', + ], + ), + {}, + )['properties']['database']['default'] + == 'THE_GOOGLE_CREDS_project_id' # because the jwt_token value is not good or missing + ) + @pytest.mark.parametrize( 'input_query, expected_output', @@ -852,3 +883,34 @@ def test_optional_fields_validator_for_google_creds_or_jwt(): } # no error raised _ = GoogleBigQueryConnector(name='something', jwt_credentials=valid_credentials) + + +def test_get_project_id(_fixture_credentials: GoogleCredentials) -> None: + connector = GoogleBigQueryConnector( + name='MyGBQ', + credentials=_fixture_credentials, + jwt_credentials=JWTCredentials(jwt_token='token', project_id='123'), + scopes=[ + 'https://www.googleapis.com/auth/bigquery', + 'https://www.googleapis.com/auth/drive', + ], + ) + + # test get project id with jwt credentials + assert connector._get_project_id() == '123' + + # test get project id with credentials + connector.jwt_credentials = None + connector.credentials = _fixture_credentials + connector.credentials.project_id = '456' + assert connector._get_project_id() == '456' + + # On ProjectId missing... + connector = GoogleBigQueryConnector( + name='MyGBQ', + credentials=None, + jwt_credentials=JWTCredentials(jwt_token='', project_id=''), + scopes=[], + ) + with pytest.raises(GoogleClientCreationError): + connector._get_project_id() 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 0a37a2519..4025b73cc 100644 --- a/toucan_connectors/google_big_query/google_big_query_connector.py +++ b/toucan_connectors/google_big_query/google_big_query_connector.py @@ -102,13 +102,7 @@ def get_form(cls, connector: 'GoogleBigQueryConnector', current_config: dict[str __base__=cls, ).schema() - 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 + schema['properties']['database']['default'] = connector._get_project_id() return schema @@ -320,7 +314,7 @@ def _bigquery_client(self) -> bigquery.Client: # We try to instantiate the bigquery.Client with the given jwt-token _session = CustomRequestSession(self.jwt_credentials.jwt_token) client = GoogleBigQueryConnector._http_connect( - http_session=_session, project_id=self.jwt_credentials.project_id + http_session=_session, project_id=self._get_project_id() ) _LOGGER.info('bigqueryClient created with the JWT provided !') @@ -332,6 +326,22 @@ def _bigquery_client(self) -> bigquery.Client: # or we fallback on default google-credentials return self._bigquery_client_with_google_creds() + def _get_project_id(self) -> str: + """We need an util in other to check either jwt_creds are well set or + not for the configuration validation, because self.jwt_credentials can + be not None but its value are either empty or None...""" + if ( + self.jwt_credentials + and self.jwt_credentials.jwt_token + and self.jwt_credentials.project_id + ): + return self.jwt_credentials.project_id + elif self.credentials and self.credentials.project_id: + return self.credentials.project_id + else: + # project-id is missing and it's mandatory + raise GoogleClientCreationError + def _get_bigquery_client(self) -> bigquery.Client: with suppress(Exception): if bigquery_client := self._bigquery_client: