Skip to content

Commit

Permalink
Revert "Revert "Remove OAuth Token From Query Params""
Browse files Browse the repository at this point in the history
  • Loading branch information
shivanesabharwal authored Nov 23, 2023
1 parent 8f9b83b commit c2dcf5d
Show file tree
Hide file tree
Showing 4 changed files with 67 additions and 19 deletions.
32 changes: 31 additions & 1 deletion src/main/java/net/starschema/clouddb/jdbc/BQSupportFuncts.java
Original file line number Diff line number Diff line change
Expand Up @@ -646,6 +646,36 @@ static QueryResponse runSyncQuery(
Map<String, String> labels,
boolean useQueryCache)
throws IOException {
return getSyncQuery(
bigquery,
projectId,
querySql,
dataSet,
useLegacySql,
maxBillingBytes,
queryTimeoutMs,
maxResults,
labels,
useQueryCache)
.execute();
}

/*
* Gets a query as specified, but does not execute it.
* Introduced for assertions on the property of the query.
* */
static Bigquery.Jobs.Query getSyncQuery(
Bigquery bigquery,
String projectId,
String querySql,
String dataSet,
Boolean useLegacySql,
Long maxBillingBytes,
Long queryTimeoutMs,
Long maxResults,
Map<String, String> labels,
boolean useQueryCache)
throws IOException {
QueryRequest qr =
new QueryRequest()
.setLabels(labels)
Expand All @@ -661,7 +691,7 @@ static QueryResponse runSyncQuery(
qr.setMaxResults(maxResults);
}

return bigquery.jobs().query(projectId, qr).execute();
return bigquery.jobs().query(projectId, qr);
}

/**
Expand Down
15 changes: 0 additions & 15 deletions src/main/java/net/starschema/clouddb/jdbc/Oauth2Bigquery.java
Original file line number Diff line number Diff line change
Expand Up @@ -119,9 +119,6 @@ private static Bigquery.Builder createBqBuilderForCredential(
if (userAgent != null) {
requestInitializer.setUserAgent(userAgent);
}
if (oauthToken != null) {
requestInitializer.setOauthToken(oauthToken);
}

bqBuilder.setBigqueryRequestInitializer(requestInitializer);
}
Expand Down Expand Up @@ -506,20 +503,11 @@ public boolean handleResponse(
private static class BigQueryRequestUserAgentInitializer extends BigqueryRequestInitializer {

String userAgent = null;
String oauthToken = null;

public void setUserAgent(String userAgent) {
this.userAgent = userAgent;
}

public void setOauthToken(String oauthToken) {
this.oauthToken = oauthToken;
}

public String getOauthToken() {
return this.oauthToken;
}

@Override
public void initializeBigqueryRequest(BigqueryRequest<?> request) throws IOException {
if (userAgent != null) {
Expand All @@ -529,9 +517,6 @@ public void initializeBigqueryRequest(BigqueryRequest<?> request) throws IOExcep

request.setRequestHeaders(currentHeaders);
}
if (oauthToken != null) {
request.setOauthToken(oauthToken);
}
}
}
}
34 changes: 34 additions & 0 deletions src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import com.google.api.client.testing.http.MockHttpTransport;
import com.google.api.client.testing.http.MockLowLevelHttpRequest;
import com.google.api.client.testing.http.MockLowLevelHttpResponse;
import com.google.api.services.bigquery.Bigquery.Jobs.Query;
import com.google.common.collect.ImmutableMap;
import java.io.IOException;
import java.nio.file.Files;
Expand Down Expand Up @@ -172,6 +173,39 @@ public void canConnectWithOAuthAccessToken()
stmt.executeQuery("SELECT * FROM orders limit 1");
}

@Test
public void oAuthAccessTokenOnlyInHeader()
throws SQLException, IOException, GeneralSecurityException {
// generate access token from service account credentials
Properties serviceProps = getProperties("/protectedaccount.properties");
String accessToken =
Oauth2Bigquery.generateAccessToken(
serviceProps.getProperty("user"),
serviceProps.getProperty("path"),
serviceProps.getProperty("password"),
null);

Properties oauthProps = getProperties("/oauthaccount.properties");
oauthProps.setProperty("oauthaccesstoken", accessToken);
String url = BQSupportFuncts.constructUrlFromPropertiesFile(oauthProps, true, null);
BQConnection bqConn = new BQConnection(url, new Properties());
BQStatement stmt = new BQStatement(oauthProps.getProperty("projectid"), bqConn);
Query query =
BQSupportFuncts.getSyncQuery(
bqConn.getBigquery(),
oauthProps.getProperty("projectid"),
"SELECT * FROM orders limit 1",
bqConn.getDataSet(),
bqConn.getUseLegacySql(),
null,
stmt.getSyncTimeoutMillis(),
(long) stmt.getMaxRows(),
stmt.getAllLabels(),
bqConn.getUseQueryCache());
String oAuthToken = query.getOauthToken();
Assert.assertTrue(oAuthToken == null);
}

@Test
public void unauthorizedResponseForInvalidOAuthAccessToken()
throws SQLException, IOException, GeneralSecurityException {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -75,9 +75,8 @@ public void Connect() throws Exception {
PreparedStatementTests.con =
DriverManager.getConnection(
BQSupportFuncts.constructUrlFromPropertiesFile(
BQSupportFuncts.readFromPropFile(
"src/test/resources/installedaccount1.properties")),
BQSupportFuncts.readFromPropFile("src/test/resources/installedaccount1.properties"));
BQSupportFuncts.readFromPropFile("installedaccount1.properties")),
BQSupportFuncts.readFromPropFile("installedaccount1.properties"));
} catch (Exception e) {
e.printStackTrace();
}
Expand Down

0 comments on commit c2dcf5d

Please sign in to comment.