Skip to content

Commit c9e27f4

Browse files
kravets-levkokravets-levko
authored
[PECO-618] Better handling of in-progress/subsequent action taken on closed session/operation (#129)
* Explicitly close sessions when closing connection Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Close operations when closing session Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Extract closable objects logic to a dedicated class Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Add error code to OperationStateError Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Handle closed/cancelled state in DBSQLOperation Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Convert DBSQLSession to async/await syntax Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Handle closed state in DBSQLSession Signed-off-by: Levko Kravets <levko.ne@gmail.com> * DBSQLOperation/CompleteOperationHelper: handle direct results close response only when operation is being closed Signed-off-by: Levko Kravets <levko.ne@gmail.com> * When all records were returned via direct results and operation was closed, it still attempted to fetch data Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Fix tests (wrong test conditions, running operations on closed session) Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Fix lint errors Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Fix after merge conflicts Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Tidy up code Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Add tests Signed-off-by: Levko Kravets <levko.ne@gmail.com> * Fix/update tests Signed-off-by: Levko Kravets <levko.ne@gmail.com> --------- Signed-off-by: Levko Kravets <levko.ne@gmail.com>
1 parent f265368 commit c9e27f4

10 files changed

Lines changed: 585 additions & 48 deletions

‎lib/DBSQLClient.ts‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import PlainHttpAuthentication from './connection/auth/PlainHttpAuthentication';
1818
import DatabricksOAuth from './connection/auth/DatabricksOAuth';
1919
import IDBSQLLogger, { LogLevel } from './contracts/IDBSQLLogger';
2020
import DBSQLLogger from './DBSQLLogger';
21+
import CloseableCollection from './utils/CloseableCollection';
2122

2223
function prependSlash(str: string): string {
2324
if (str.length > 0 && str.charAt(0) !== '/') {
@@ -52,6 +53,8 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient {
5253

5354
private readonly thrift = thrift;
5455

56+
private sessions = new CloseableCollection<DBSQLSession>();
57+
5558
constructor(options?: ClientOptions) {
5659
super();
5760
this.logger = options?.logger || new DBSQLLogger();
@@ -147,7 +150,11 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient {
147150
});
148151

149152
Status.assert(response.status);
150-
return new DBSQLSession(driver, definedOrError(response.sessionHandle), this.logger);
153+
const session = new DBSQLSession(driver, definedOrError(response.sessionHandle), {
154+
logger: this.logger,
155+
});
156+
this.sessions.add(session);
157+
return session;
151158
}
152159

153160
private async getClient() {
@@ -206,6 +213,8 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient {
206213
}
207214

208215
public async close(): Promise<void> {
216+
await this.sessions.closeAll();
217+
209218
this.client = null;
210219
this.authProvider = null;
211220
this.connectionOptions = null;

‎lib/DBSQLOperation/OperationStatusHelper.ts‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
import { TOperationHandle, TOperationState, TGetOperationStatusResp } from '../../thrift/TCLIService_types';
1+
import { TGetOperationStatusResp, TOperationHandle, TOperationState } from '../../thrift/TCLIService_types';
22
import HiveDriver from '../hive/HiveDriver';
33
import Status from '../dto/Status';
44
import { WaitUntilReadyOptions } from '../contracts/IOperation';
5-
import OperationStateError from '../errors/OperationStateError';
5+
import OperationStateError, { OperationStateErrorCode } from '../errors/OperationStateError';
66

77
async function delay(ms?: number): Promise<void> {
88
return new Promise((resolve) => {
@@ -92,16 +92,16 @@ export default class OperationStatusHelper {
9292
case TOperationState.FINISHED_STATE:
9393
return true;
9494
case TOperationState.CANCELED_STATE:
95-
throw new OperationStateError('The operation was canceled by a client', response);
95+
throw new OperationStateError(OperationStateErrorCode.Canceled, response);
9696
case TOperationState.CLOSED_STATE:
97-
throw new OperationStateError('The operation was closed by a client', response);
97+
throw new OperationStateError(OperationStateErrorCode.Closed, response);
9898
case TOperationState.ERROR_STATE:
99-
throw new OperationStateError('The operation failed due to an error', response);
99+
throw new OperationStateError(OperationStateErrorCode.Error, response);
100100
case TOperationState.TIMEDOUT_STATE:
101-
throw new OperationStateError('The operation is in a timed out state', response);
101+
throw new OperationStateError(OperationStateErrorCode.Timeout, response);
102102
case TOperationState.UKNOWN_STATE:
103103
default:
104-
throw new OperationStateError('The operation is in an unrecognized state', response);
104+
throw new OperationStateError(OperationStateErrorCode.Unknown, response);
105105
}
106106
}
107107

‎lib/DBSQLOperation/index.ts‎

Lines changed: 67 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
11
import { stringify, NIL, parse } from 'uuid';
2-
import IOperation, { FetchOptions, GetSchemaOptions, FinishedOptions } from '../contracts/IOperation';
2+
import IOperation, {
3+
FetchOptions,
4+
FinishedOptions,
5+
GetSchemaOptions,
6+
WaitUntilReadyOptions,
7+
} from '../contracts/IOperation';
38
import HiveDriver from '../hive/HiveDriver';
49
import {
510
TGetOperationStatusResp,
@@ -8,22 +13,28 @@ import {
813
TSparkDirectResults,
914
} from '../../thrift/TCLIService_types';
1015
import Status from '../dto/Status';
11-
1216
import OperationStatusHelper from './OperationStatusHelper';
1317
import SchemaHelper from './SchemaHelper';
1418
import FetchResultsHelper from './FetchResultsHelper';
1519
import CompleteOperationHelper from './CompleteOperationHelper';
1620
import IDBSQLLogger, { LogLevel } from '../contracts/IDBSQLLogger';
21+
import OperationStateError, { OperationStateErrorCode } from '../errors/OperationStateError';
1722

1823
const defaultMaxRows = 100000;
1924

25+
interface DBSQLOperationConstructorOptions {
26+
logger: IDBSQLLogger;
27+
}
28+
2029
export default class DBSQLOperation implements IOperation {
2130
private readonly driver: HiveDriver;
2231

2332
private readonly operationHandle: TOperationHandle;
2433

2534
private readonly logger: IDBSQLLogger;
2635

36+
public onClose?: () => void;
37+
2738
private readonly _status: OperationStatusHelper;
2839

2940
private readonly _schema: SchemaHelper;
@@ -35,7 +46,7 @@ export default class DBSQLOperation implements IOperation {
3546
constructor(
3647
driver: HiveDriver,
3748
operationHandle: TOperationHandle,
38-
logger: IDBSQLLogger,
49+
{ logger }: DBSQLOperationConstructorOptions,
3950
directResults?: TSparkDirectResults,
4051
) {
4152
this.driver = driver;
@@ -95,17 +106,21 @@ export default class DBSQLOperation implements IOperation {
95106
* const result = await queryOperation.fetchChunk({maxRows: 1000});
96107
*/
97108
public async fetchChunk(options?: FetchOptions): Promise<Array<object>> {
109+
await this.failIfClosed();
110+
98111
if (!this._status.hasResultSet) {
99112
return [];
100113
}
101114

102-
await this._status.waitUntilReady(options);
115+
await this.waitUntilReady(options);
103116

104117
const [resultHandler, data] = await Promise.all([
105118
this._schema.getResultHandler(),
106119
this._data.fetch(options?.maxRows || defaultMaxRows),
107120
]);
108121

122+
await this.failIfClosed();
123+
109124
const result = await resultHandler.getValue(data ? [data] : []);
110125
this.logger?.log(
111126
LogLevel.debug,
@@ -120,6 +135,7 @@ export default class DBSQLOperation implements IOperation {
120135
* @throws {StatusError}
121136
*/
122137
public async status(progress: boolean = false): Promise<TGetOperationStatusResp> {
138+
await this.failIfClosed();
123139
this.logger?.log(LogLevel.debug, `Fetching status for operation with id: ${this.getId()}`);
124140
return this._status.status(progress);
125141
}
@@ -129,21 +145,37 @@ export default class DBSQLOperation implements IOperation {
129145
* @throws {StatusError}
130146
*/
131147
public async cancel(): Promise<Status> {
148+
if (this._completeOperation.closed || this._completeOperation.cancelled) {
149+
return Status.success();
150+
}
151+
132152
this.logger?.log(LogLevel.debug, `Cancelling operation with id: ${this.getId()}`);
133-
return this._completeOperation.cancel();
153+
const result = this._completeOperation.cancel();
154+
155+
// Cancelled operation becomes unusable, similarly to being closed
156+
this.onClose?.();
157+
return result;
134158
}
135159

136160
/**
137161
* Closes operation
138162
* @throws {StatusError}
139163
*/
140164
public async close(): Promise<Status> {
165+
if (this._completeOperation.closed || this._completeOperation.cancelled) {
166+
return Status.success();
167+
}
168+
141169
this.logger?.log(LogLevel.debug, `Closing operation with id: ${this.getId()}`);
142-
return this._completeOperation.close();
170+
const result = await this._completeOperation.close();
171+
172+
this.onClose?.();
173+
return result;
143174
}
144175

145176
public async finished(options?: FinishedOptions): Promise<void> {
146-
await this._status.waitUntilReady(options);
177+
await this.failIfClosed();
178+
await this.waitUntilReady(options);
147179
}
148180

149181
public async hasMoreRows(): Promise<boolean> {
@@ -163,13 +195,40 @@ export default class DBSQLOperation implements IOperation {
163195
}
164196

165197
public async getSchema(options?: GetSchemaOptions): Promise<TTableSchema | null> {
198+
await this.failIfClosed();
199+
166200
if (!this._status.hasResultSet) {
167201
return null;
168202
}
169203

170-
await this._status.waitUntilReady(options);
204+
await this.waitUntilReady(options);
171205

172206
this.logger?.log(LogLevel.debug, `Fetching schema for operation with id: ${this.getId()}`);
173207
return this._schema.fetch();
174208
}
209+
210+
private async failIfClosed(): Promise<void> {
211+
if (this._completeOperation.closed) {
212+
throw new OperationStateError(OperationStateErrorCode.Closed);
213+
}
214+
if (this._completeOperation.cancelled) {
215+
throw new OperationStateError(OperationStateErrorCode.Canceled);
216+
}
217+
}
218+
219+
private async waitUntilReady(options?: WaitUntilReadyOptions) {
220+
try {
221+
await this._status.waitUntilReady(options);
222+
} catch (error) {
223+
if (error instanceof OperationStateError) {
224+
if (error.errorCode === OperationStateErrorCode.Canceled) {
225+
this._completeOperation.cancelled = true;
226+
}
227+
if (error.errorCode === OperationStateErrorCode.Closed) {
228+
this._completeOperation.closed = true;
229+
}
230+
}
231+
throw error;
232+
}
233+
}
175234
}

0 commit comments

Comments
 (0)