fix: set the appropriate project-id on GoCred/JWT auth mode.

ToucanToco · Sep 22, 2023 · ad1be92 · ad1be92
1 parent 6564b76
commit ad1be92
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 3 deletions.
diff --git 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
@@ -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