fix(bigquery-jdbc): abort session when connection is closed - #14273
fix(bigquery-jdbc): abort session when connection is closed#14273Neenu1995 wants to merge 6 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces changes to automatically abort active BigQuery sessions when a connection is closed. Specifically, it adds an abortSession method in BigQueryConnection that executes CALL BQ.ABORT_SESSION(); and cleans up session-related connection properties. Corresponding unit and integration tests have been added to verify this behavior. The feedback suggests a minor improvement to the exception message when handling InterruptedException to more accurately reflect that the interruption occurred during session abortion.
| QueryJobConfiguration.newBuilder("CALL BQ.ABORT_SESSION();") | ||
| .setConnectionProperties(this.queryProperties) | ||
| .build(); | ||
| Job abortJob = this.bigQuery.create(JobInfo.of(abortSessionJobConfig)); |
There was a problem hiding this comment.
Can we use query API instead? It'd allow us to save 1 roundtrip to the backend
| QueryJobConfiguration.Builder transactionBeginJobConfig = | ||
| QueryJobConfiguration.newBuilder("BEGIN TRANSACTION;"); | ||
| try { | ||
| if (this.sessionInfoConnectionProperty != null) { |
There was a problem hiding this comment.
This if looks weird to me. We check sessionInfoConnectionProperty, but we call setConnectionproperties(this.queryProperties). queryProperties have more data.
Should we always do setConnectionProperties & do seCreateSession(true) only when sessionInfo is null?
| Job job = this.bigQuery.create(JobInfo.of(transactionBeginJobConfig.build())); | ||
| job = job.waitFor(); | ||
| Job transactionBeginJob = this.bigQuery.getJob(job.getJobId()); | ||
| if (this.sessionInfoConnectionProperty == null |
There was a problem hiding this comment.
If we don't call updateSessionInfo due to any of 3 != null checks, this is an error, we can't mark transactionStarted = true. We need to either throw custom error or just remove these checks.
| } | ||
| } | ||
|
|
||
| synchronized void updateSessionInfo(String sessionId) { |
There was a problem hiding this comment.
Session is immutable during connection lifetime, right? This method should reflect that imo
| boolean transactionStarted; | ||
| volatile ConnectionProperty sessionInfoConnectionProperty; | ||
| // isSessionCreatedByDriver is false by default. | ||
| boolean isSessionCreatedByDriver = false; |
There was a problem hiding this comment.
IIUC, this PR will abort the session on connection close when it is created by our driver.
And we are not creating a new conn prop (e.g. keepSessionAlive default false) that would let users NOT abort the sessions created by the driver when connection is closed?
3 participants
Fixes #13922