@@ -97,6 +97,7 @@ class ConnectionImpl implements Connection {
97
97
private final Logger logger = Logger .getLogger (this .getClass ().getName ());
98
98
private BigQueryReadClient bqReadClient ;
99
99
private static final long EXECUTOR_TIMEOUT_SEC = 10 ;
100
+ private static final long BIGQUERY_TIMEOUT_SEC = 10 ;
100
101
private BlockingQueue <AbstractList <FieldValue >>
101
102
bufferFvl ; // initialized lazily iff we end up using the tabledata.list end point
102
103
private BlockingQueue <BigQueryResultImpl .Row >
@@ -148,8 +149,15 @@ public synchronized boolean close() throws BigQuerySQLException {
148
149
flagEndOfStream (); // an End of Stream flag in the buffer so that the `ResultSet.next()` stops
149
150
// advancing the cursor
150
151
queryTaskExecutor .shutdownNow ();
152
+ boolean isBqReadClientTerminated = true ;
151
153
try {
152
- if (queryTaskExecutor .awaitTermination (EXECUTOR_TIMEOUT_SEC , TimeUnit .SECONDS )) {
154
+ if (bqReadClient != null ) {
155
+ bqReadClient .shutdownNow ();
156
+ isBqReadClientTerminated =
157
+ bqReadClient .awaitTermination (BIGQUERY_TIMEOUT_SEC , TimeUnit .SECONDS );
158
+ }
159
+ if (queryTaskExecutor .awaitTermination (EXECUTOR_TIMEOUT_SEC , TimeUnit .SECONDS )
160
+ && isBqReadClientTerminated ) {
153
161
return true ;
154
162
} // else queryTaskExecutor.isShutdown() will be returned outside this try block
155
163
} catch (InterruptedException e ) {
@@ -159,7 +167,9 @@ public synchronized boolean close() throws BigQuerySQLException {
159
167
e ); // Logging InterruptedException instead of throwing the exception back, close method
160
168
// will return queryTaskExecutor.isShutdown()
161
169
}
162
- return queryTaskExecutor .isShutdown (); // check if the executor has been shutdown
170
+
171
+ return queryTaskExecutor .isShutdown ()
172
+ && isBqReadClientTerminated ; // check if the executor has been shutdown
163
173
}
164
174
165
175
/**
@@ -992,7 +1002,6 @@ BigQueryResult highThroughPutRead(
992
1002
// DO a regex check using order by and use multiple streams
993
1003
;
994
1004
ReadSession readSession = bqReadClient .createReadSession (builder .build ());
995
-
996
1005
bufferRow = new LinkedBlockingDeque <>(getBufferSize ());
997
1006
Map <String , Integer > arrowNameToIndex = new HashMap <>();
998
1007
// deserialize and populate the buffer async, so that the client isn't blocked
@@ -1050,6 +1059,7 @@ private void processArrowStreamAsync(
1050
1059
"\n " + Thread .currentThread ().getName () + " Interrupted @ markLast" ,
1051
1060
e );
1052
1061
}
1062
+ bqReadClient .shutdownNow (); // Shutdown the read client
1053
1063
queryTaskExecutor .shutdownNow (); // Shutdown the thread pool
1054
1064
}
1055
1065
};
0 commit comments