It seems that the issue might be deeper in how the custom credentials class interacts with the BigQuery client and how token expiry and refresh are managed.You need try to streamline the implementation and ensure the credentials are managed correctly.
Ensure that the custom implementation of OAuth2Credentials correctly manages token refresh:
import com.google.auth.oauth2.AccessToken;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.auth.oauth2.OAuth2Credentials;
import java.io.IOException;
import java.util.Date;
import java.util.Map;
public class CustomCredentials extends OAuth2Credentials {
private String wifToken;
public CustomCredentials(String wifToken, AccessToken accessToken) {
super(accessToken);
this.wifToken = wifToken;
}
@Override
public AccessToken refreshAccessToken() throws IOException {
// Generate a new WIF token and then generate a new access token
this.wifToken = generateWifToken();
return generateAccessToken(this.wifToken);
}
private String generateWifToken() {
// Your logic to generate WIF token
return "newWifToken";
}
private AccessToken generateAccessToken(String wifToken) throws IOException {
// Your logic to generate Access Token using WIF token
Date newExpiryDate = new Date(System.currentTimeMillis() + TimeUnit.HOURS.toMillis(1)); // Set expiry time appropriately
return new AccessToken("newAccessToken", newExpiryDate);
}
@Override
public Map<String, List<String>> getRequestMetadata(URI uri) throws IOException {
if (shouldRefresh()) {
refresh();
}
return super.getRequestMetadata(uri);
}
private boolean shouldRefresh() {
return getAccessToken().getExpirationTime().getTime() - System.currentTimeMillis() < TimeUnit.MINUTES.toMillis(5);
}
@Override
public boolean hasRequestMetadata() {
return true;
}
@Override
public boolean hasRequestMetadataOnly() {
return true;
}
}
Next, create the BigQuery client using the custom credentials:
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.BigQueryOptions;
public class BigQueryService {
private BigQuery bigQuery;
public BigQueryService(String wifToken) throws IOException {
CustomCredentials customCredentials = new CustomCredentials(wifToken, generateInitialAccessToken(wifToken));
this.bigQuery = BigQueryOptions.newBuilder()
.setCredentials(customCredentials)
.build()
.getService();
}
private AccessToken generateInitialAccessToken(String wifToken) throws IOException {
// Your logic to generate initial Access Token using WIF token
Date newExpiryDate = new Date(System.currentTimeMillis() + TimeUnit.HOURS.toMillis(1)); // Set expiry time appropriately
return new AccessToken("initialAccessToken", newExpiryDate);
}
public BigQuery getBigQuery() {
return bigQuery;
}
}
- Ensure the refreshAccessToken() method correctly generates a new WIF token and access token.
- Implement the expiry check in both the constructor and getRequestMetadata() method to ensure the token is refreshed before it expires.
- Implement a retry mechanism in the application to handle transient errors.
- Add detailed logging to track token generation, refresh, and usage. This can help identify issues in the process.
- Ensure environment variables and configurations are set correctly in Cloud Run.
- If the application is multi-threaded, ensure thread-safety in your credentials handling.
- Thoroughly test the implementation with different record sizes and processing times to ensure stability.
public class Main {
public static void main(String[] args) throws IOException {
String wifToken = "yourInitialWifToken";
BigQueryService bigQueryService = new BigQueryService(wifToken);
BigQuery bigQuery = bigQueryService.getBigQuery();
// Your BigQuery operations
}
}