Skip to content

Commit 8ce183c

Browse files
Neenu1995Neenu1995
authored
fix(bqjdbc): support Service Account Impersonation in ADC (googleapis#12967)
b/506258278
1 parent 1721e7c commit 8ce183c

5 files changed

Lines changed: 77 additions & 26 deletions

File tree

‎java-bigquery/google-cloud-bigquery-jdbc/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -307,6 +307,11 @@
307307
<artifactId>mockito-core</artifactId>
308308
<scope>test</scope>
309309
</dependency>
310+
<dependency>
311+
<groupId>org.mockito</groupId>
312+
<artifactId>mockito-inline</artifactId>
313+
<scope>test</scope>
314+
</dependency>
310315
<dependency>
311316
<groupId>org.mockito</groupId>
312317
<artifactId>mockito-junit-jupiter</artifactId>

‎java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtility.java‎

Lines changed: 16 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -230,26 +230,22 @@ static Map<String, String> parseOAuthProperties(DataSource ds, String callerClas
230230
break;
231231
}
232232

233-
if (authType == AuthType.GOOGLE_SERVICE_ACCOUNT
234-
|| authType == AuthType.GOOGLE_USER_ACCOUNT
235-
|| authType == AuthType.PRE_GENERATED_TOKEN) {
236-
oauthProperties.put(
237-
BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_EMAIL_PROPERTY_NAME,
238-
ds.getOAuthSAImpersonationEmail());
239-
oauthProperties.put(
240-
BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_CHAIN_PROPERTY_NAME,
241-
ds.getOAuthSAImpersonationChain());
242-
oauthProperties.put(
243-
BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_SCOPES_PROPERTY_NAME,
244-
ds.getOAuthSAImpersonationScopes() != null
245-
? ds.getOAuthSAImpersonationScopes()
246-
: BIGQUERY_SCOPE);
247-
oauthProperties.put(
248-
BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_TOKEN_LIFETIME_PROPERTY_NAME,
249-
ds.getOAuthSAImpersonationTokenLifetime() != null
250-
? ds.getOAuthSAImpersonationTokenLifetime()
251-
: BigQueryJdbcUrlUtility.DEFAULT_OAUTH_SA_IMPERSONATION_TOKEN_LIFETIME_VALUE);
252-
}
233+
oauthProperties.put(
234+
BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_EMAIL_PROPERTY_NAME,
235+
ds.getOAuthSAImpersonationEmail());
236+
oauthProperties.put(
237+
BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_CHAIN_PROPERTY_NAME,
238+
ds.getOAuthSAImpersonationChain());
239+
oauthProperties.put(
240+
BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_SCOPES_PROPERTY_NAME,
241+
ds.getOAuthSAImpersonationScopes() != null
242+
? ds.getOAuthSAImpersonationScopes()
243+
: BIGQUERY_SCOPE);
244+
oauthProperties.put(
245+
BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_TOKEN_LIFETIME_PROPERTY_NAME,
246+
ds.getOAuthSAImpersonationTokenLifetime() != null
247+
? ds.getOAuthSAImpersonationTokenLifetime()
248+
: BigQueryJdbcUrlUtility.DEFAULT_OAUTH_SA_IMPERSONATION_TOKEN_LIFETIME_VALUE);
253249
return oauthProperties;
254250
}
255251

@@ -284,12 +280,9 @@ static GoogleCredentials getCredentials(
284280
getPreGeneratedTokensCredentials(authProperties, overrideProperties, callerClassName);
285281
break;
286282
case APPLICATION_DEFAULT_CREDENTIALS:
287-
// This auth method doesn't support service account impersonation
288-
289283
credentials = getApplicationDefaultCredentials(callerClassName);
290284
break;
291285
case EXTERNAL_ACCOUNT_AUTH:
292-
// This auth method doesn't support service account impersonation
293286
credentials = getExternalAccountAuthCredentials(authProperties, callerClassName);
294287
break;
295288
default:

‎java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryCallableStatementTest.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -914,7 +914,7 @@ public void testSetDateCalParamByName() throws SQLException {
914914
Calendar expectedCal = mock(Calendar.class);
915915

916916
doReturn(1L).when(expectedDate).getTime();
917-
doReturn(1L).when(expectedCal).getTime();
917+
doReturn(new java.util.Date(1L)).when(expectedCal).getTime();
918918
doReturn(1L).when(expectedCal).getTimeInMillis();
919919
statement.setDate(PARAM_KEY, expectedDate, expectedCal);
920920
Date actual = statement.getDate(PARAM_KEY);
@@ -1033,7 +1033,7 @@ public void testSetTimeCalParamByName() throws SQLException {
10331033
Calendar expectedCal = mock(Calendar.class);
10341034

10351035
doReturn(1L).when(expectedTime).getTime();
1036-
doReturn(1L).when(expectedCal).getTime();
1036+
doReturn(new java.util.Date(1L)).when(expectedCal).getTime();
10371037
doReturn(1L).when(expectedCal).getTimeInMillis();
10381038
statement.setTime(PARAM_KEY, expectedTime, expectedCal);
10391039
Time actual = statement.getTime(PARAM_KEY);
@@ -1062,7 +1062,7 @@ public void testSetTimestampCalParamByName() throws SQLException {
10621062
Calendar expectedCal = mock(Calendar.class);
10631063

10641064
doReturn(1L).when(expectedTimestamp).getTime();
1065-
doReturn(1L).when(expectedCal).getTime();
1065+
doReturn(new java.util.Date(1L)).when(expectedCal).getTime();
10661066
doReturn(1L).when(expectedCal).getTimeInMillis();
10671067
statement.setTimestamp(PARAM_KEY, expectedTimestamp, expectedCal);
10681068
Timestamp actual = statement.getTimestamp(PARAM_KEY);

‎java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtilityTest.java‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -377,6 +377,48 @@ public void testParseUserImpersonationNonDefault() {
377377
result.get(BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_TOKEN_LIFETIME_PROPERTY_NAME));
378378
}
379379

380+
@Test
381+
public void testParseUserImpersonationForADC() {
382+
Map<String, String> result =
383+
BigQueryJdbcOAuthUtility.parseOAuthProperties(
384+
DataSource.fromUrl(
385+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
386+
+ "OAuthType=3;ProjectId=MyBigQueryProject;"
387+
+ "[email protected];"),
388+
"");
389+
390+
assertEquals("APPLICATION_DEFAULT_CREDENTIALS", result.get("OAuthType"));
391+
assertEquals(
392+
"[email protected]",
393+
result.get(BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_EMAIL_PROPERTY_NAME));
394+
}
395+
396+
@Test
397+
public void testGetServiceAccountImpersonatedCredentialsForADC() throws Exception {
398+
GoogleCredentials dummySourceCredentials = GoogleCredentials.newBuilder().build();
399+
400+
try (org.mockito.MockedStatic<GoogleCredentials> mockedCreds =
401+
org.mockito.Mockito.mockStatic(GoogleCredentials.class)) {
402+
mockedCreds.when(GoogleCredentials::getApplicationDefault).thenReturn(dummySourceCredentials);
403+
404+
Map<String, String> authProperties =
405+
BigQueryJdbcOAuthUtility.parseOAuthProperties(
406+
DataSource.fromUrl(
407+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
408+
+ "OAuthType=3;ProjectId=MyBigQueryProject;"
409+
+ "[email protected];"),
410+
"");
411+
412+
GoogleCredentials credentials =
413+
BigQueryJdbcOAuthUtility.getCredentials(
414+
authProperties, java.util.Collections.EMPTY_MAP, false, null);
415+
416+
assertThat(credentials).isInstanceOf(ImpersonatedCredentials.class);
417+
assertThat(((ImpersonatedCredentials) credentials).getSourceCredentials())
418+
.isEqualTo(dummySourceCredentials);
419+
}
420+
}
421+
380422
@Test
381423
public void testGetServiceAccountImpersonatedCredentials() {
382424
Map<String, String> authProperties =

‎java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -371,4 +371,15 @@ public void testServiceAccountAuthenticationWithChainedImpersonation()
371371
.toString();
372372
validateConnection(connection_uri);
373373
}
374+
375+
@Test
376+
public void testADCAuthenticationWithImpersonation() throws IOException, SQLException {
377+
final JsonObject authJson = getAuthJson();
378+
379+
String connection_uri =
380+
getBaseUri(3, authJson.get("project_id").getAsString())
381+
.append("ServiceAccountImpersonationEmail", authJson.get("client_email").getAsString())
382+
.toString();
383+
validateConnection(connection_uri);
384+
}
374385
}

0 commit comments

Comments
 (0)