From 36441177a63bfc01931239e61afe93a342fbd1b4 Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Tue, 7 Mar 2023 15:18:10 +0200 Subject: [PATCH 1/4] Implement REST client (subset of databricks-sdk-go) Signed-off-by: Levko Kravets --- lib/rest/RestClient.ts | 76 +++++++++++++++ lib/rest/Types.ts | 165 +++++++++++++++++++++++++++++++ package-lock.json | 216 ++++++++++++++++++++++++++++++++++++++++- package.json | 2 + 4 files changed, 454 insertions(+), 5 deletions(-) create mode 100644 lib/rest/RestClient.ts create mode 100644 lib/rest/Types.ts diff --git a/lib/rest/RestClient.ts b/lib/rest/RestClient.ts new file mode 100644 index 00000000..302f08ba --- /dev/null +++ b/lib/rest/RestClient.ts @@ -0,0 +1,76 @@ +import fetch from 'node-fetch'; + +import { + ExecuteStatementRequest, + ExecuteStatementResponse, + CancelExecutionRequest, + GetStatementRequest, + GetStatementResponse, + GetStatementResultChunkNRequest, + ResultData, +} from './Types'; + +export interface RestClientOptions { + host: string; + headers?: Record; +} + +enum HTTPMethod { + GET = 'GET', + POST = 'POST', +} + +export default class RestClient { + private options: RestClientOptions; + + private async doRequest(method: string, path: string, payload: P): Promise { + const { host, headers } = this.options; + const response = await fetch(`https://${host}${path}`, { + method, + headers: { + Accept: 'application/json', + 'Content-Type': 'application/json', + ...headers, + }, + body: method !== HTTPMethod.GET ? JSON.stringify(payload) : undefined, + }); + const result = await response.json(); + return result as R; + } + + constructor(options: RestClientOptions) { + this.options = options; + } + + public executeStatement(request: ExecuteStatementRequest) { + return this.doRequest( + HTTPMethod.POST, + '/api/2.0/sql/statements/', + request, + ); + } + + public cancelExecution(request: CancelExecutionRequest) { + return this.doRequest( + HTTPMethod.POST, + `/api/2.0/sql/statements/${request.statement_id}/cancel`, + request, + ); + } + + public getStatement(request: GetStatementRequest): Promise { + return this.doRequest( + HTTPMethod.GET, + `/api/2.0/sql/statements/${request.statement_id}`, + request, + ); + } + + public getStatementResultChunkN(request: GetStatementResultChunkNRequest): Promise { + return this.doRequest( + HTTPMethod.GET, + `/api/2.0/sql/statements/${request.statement_id}/result/chunks/${request.chunk_index}`, + request, + ); + } +} diff --git a/lib/rest/Types.ts b/lib/rest/Types.ts new file mode 100644 index 00000000..3139b50f --- /dev/null +++ b/lib/rest/Types.ts @@ -0,0 +1,165 @@ +export enum Disposition { + ExternalLinks = 'EXTERNAL_LINKS', + Inline = 'INLINE', +} + +export enum Format { + ArrowStream = 'ARROW_STREAM', + JsonArray = 'JSON_ARRAY', +} + +export enum TimeoutAction { + Cancel = 'CANCEL', + Continue = 'CONTINUE', +} + +export interface ExecuteStatementRequest { + catalog: string; + disposition: Disposition; + format: Format; + on_wait_timeout: TimeoutAction; + schema: string; + statement: string; + wait_timeout: string; + warehouse_id: string; +} + +export interface ExecuteStatementResponse { + manifest: ResultManifest; + result: ResultData; + statement_id: string; + status: StatementStatus; +} + +export enum StatementState { + Canceled = 'CANCELED', + Closed = 'CLOSED', + Failed = 'FAILED', + Pending = 'PENDING', + Running = 'RUNNING', + Succeeded = 'SUCCEEDED', +} + +export interface StatementStatus { + error: ServiceError; + state: StatementState; +} + +export interface ServiceError { + error_code: ServiceErrorCode; + message: string; +} + +export enum ServiceErrorCode { + Aborted = 'ABORTED', + AlreadyExists = 'ALREADY_EXISTS', + BadRequest = 'BAD_REQUEST', + Cancelled = 'CANCELLED', + DeadlineExceeded = 'DEADLINE_EXCEEDED', + InternalError = 'INTERNAL_ERROR', + IoError = 'IO_ERROR', + NotFound = 'NOT_FOUND', + ResourceExhausted = 'RESOURCE_EXHAUSTED', + ServiceUnderMaintenance = 'SERVICE_UNDER_MAINTENANCE', + TemporarilyUnavailable = 'TEMPORARILY_UNAVAILABLE', + Unauthenticated = 'UNAUTHENTICATED', + Unknown = 'UNKNOWN', + WorkspaceTemporarilyUnavailable = 'WORKSPACE_TEMPORARILY_UNAVAILABLE', +} + +export interface ResultData { + byte_count: number; // int64 + chunk_index: number; + data_array: string[][]; + external_links: ExternalLink[]; + next_chunk_index: number; + next_chunk_internal_link: string; + row_count: number; // int64 + row_offset: number; // int64 +} + +interface ExternalLink { + byte_count: number; // int64 + chunk_index: number; + expiration: string; + external_link: string; + next_chunk_index: number; + next_chunk_internal_link: string; + row_count: number; // int64 + row_offset: number; // int64 +} + +export interface ResultManifest { + chunks: ChunkInfo[]; + format: Format; + schema: ResultSchema; + total_byte_count: number; // int64 + total_chunk_count: number; + total_row_count: number; // int64 +} + +export interface ChunkInfo { + byte_count: number; // int64 + chunk_index: number; + next_chunk_index: number; + next_chunk_internal_link: string; + row_count: number; // int64 + row_offset: number; // int64 +} + +export interface ResultSchema { + column_count: number; + columns: ColumnInfo[]; +} + +export interface ColumnInfo { + name: string; + position: number; + type_interval_type: string; + type_name: ColumnInfoTypeName; + type_precision: number; + type_scale: number; + type_text: string; +} + +export enum ColumnInfoTypeName { + Array = 'ARRAY', + Binary = 'BINARY', + Boolean = 'BOOLEAN', + Byte = 'BYTE', + Char = 'CHAR', + Date = 'DATE', + Decimal = 'DECIMAL', + Double = 'DOUBLE', + Float = 'FLOAT', + Int = 'INT', + Interval = 'INTERVAL', + Long = 'LONG', + Map = 'MAP', + Null = 'NULL', + Short = 'SHORT', + String = 'STRING', + Struct = 'STRUCT', + Timestamp = 'TIMESTAMP', + UserDefinedType = 'USER_DEFINED_TYPE', +} + +export interface CancelExecutionRequest { + statement_id: string; +} + +export interface GetStatementRequest { + statement_id: string; +} + +export interface GetStatementResponse { + manifest: ResultManifest; + result: ResultData; + statement_id: string; + status: StatementStatus; +} + +export interface GetStatementResultChunkNRequest { + chunk_index: number; + statement_id: string; +} diff --git a/package-lock.json b/package-lock.json index 4afe3daa..9ec22a1b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,16 +1,20 @@ { "name": "@databricks/sql", - "version": "1.1.0", + "version": "1.1.1", "lockfileVersion": 2, "requires": true, "packages": { "": { "name": "@databricks/sql", - "version": "1.0.0", + "version": "1.1.1", + "bundleDependencies": [ + "thrift" + ], "hasInstallScript": true, "license": "Apache 2.0", "dependencies": { "commander": "^9.3.0", + "node-fetch": "^2.6.9", "node-int64": "^0.4.0", "patch-package": "^6.5.0", "thrift": "^0.16.0", @@ -19,6 +23,7 @@ }, "devDependencies": { "@types/node": "^18.0.0", + "@types/node-fetch": "^2.6.2", "@types/node-int64": "^0.4.29", "@types/thrift": "^0.10.11", "@types/uuid": "^8.3.4", @@ -792,6 +797,16 @@ "integrity": "sha512-NcKK6Ts+9LqdHJaW6HQmgr7dT/i3GOHG+pt6BiWv++5SnjtRd4NXeiuN2kA153SjhXPR/AhHIPHPbrsbpUVOww==", "dev": true }, + "node_modules/@types/node-fetch": { + "version": "2.6.2", + "resolved": "https://registry.npmjs.org/@types/node-fetch/-/node-fetch-2.6.2.tgz", + "integrity": "sha512-DHqhlq5jeESLy19TYhLakJ07kNumXWjcDdxXsLUMJZ6ue8VZJj4kLPQVE/2mdHh3xZziNF1xppu5lwmS53HR+A==", + "dev": true, + "dependencies": { + "@types/node": "*", + "form-data": "^3.0.0" + } + }, "node_modules/@types/node-int64": { "version": "0.4.29", "resolved": "https://registry.npmjs.org/@types/node-int64/-/node-int64-0.4.29.tgz", @@ -1238,7 +1253,14 @@ "node_modules/async-limiter": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/async-limiter/-/async-limiter-1.0.1.tgz", - "integrity": "sha512-csOlWGAcRFJaI6m+F2WKdnMKr4HhdhFVBk0H/QbJFMCr+uO2kwohwXQPxw/9OCxp05r5ghVBFSyioixx3gfkNQ==" + "integrity": "sha512-csOlWGAcRFJaI6m+F2WKdnMKr4HhdhFVBk0H/QbJFMCr+uO2kwohwXQPxw/9OCxp05r5ghVBFSyioixx3gfkNQ==", + "inBundle": true + }, + "node_modules/asynckit": { + "version": "0.4.0", + "resolved": "https://registry.npmjs.org/asynckit/-/asynckit-0.4.0.tgz", + "integrity": "sha512-Oei9OH4tRh0YqU3GxhX79dM/mwVgvbZJaSNaRk+bshkj0S5cfHcgYakreBjrHwatXKbz+IoIdYLxrKim2MjW0Q==", + "dev": true }, "node_modules/axe-core": { "version": "4.4.3", @@ -1292,7 +1314,8 @@ "node_modules/browser-or-node": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/browser-or-node/-/browser-or-node-1.3.0.tgz", - "integrity": "sha512-0F2z/VSnLbmEeBcUrSuDH5l0HxTXdQQzLjkmBR4cYfvg1zJrKSlmIZFqyFR8oX0NrwPhy3c3HQ6i3OxMbew4Tg==" + "integrity": "sha512-0F2z/VSnLbmEeBcUrSuDH5l0HxTXdQQzLjkmBR4cYfvg1zJrKSlmIZFqyFR8oX0NrwPhy3c3HQ6i3OxMbew4Tg==", + "inBundle": true }, "node_modules/browser-stdout": { "version": "1.3.1", @@ -1552,6 +1575,18 @@ "text-hex": "1.0.x" } }, + "node_modules/combined-stream": { + "version": "1.0.8", + "resolved": "https://registry.npmjs.org/combined-stream/-/combined-stream-1.0.8.tgz", + "integrity": "sha512-FQN4MRfuJeHf7cBbBMJFXhKSDq+2kAArBlmRBvcvFE5BB1HZKXtSFASDhdlz9zOYwxh8lDdnvmMOe/+5cdoEdg==", + "dev": true, + "dependencies": { + "delayed-stream": "~1.0.0" + }, + "engines": { + "node": ">= 0.8" + } + }, "node_modules/commander": { "version": "9.4.0", "resolved": "https://registry.npmjs.org/commander/-/commander-9.4.0.tgz", @@ -1689,6 +1724,15 @@ "url": "https://github.com/sponsors/ljharb" } }, + "node_modules/delayed-stream": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/delayed-stream/-/delayed-stream-1.0.0.tgz", + "integrity": "sha512-ZySD7Nf91aLB0RxL4KGrKHBXl7Eds1DAmEdcoVawXnLD7SDhpNgtuII2aAkg7a7QS41jxPSZ17p4VdGnMHk3MQ==", + "dev": true, + "engines": { + "node": ">=0.4.0" + } + }, "node_modules/diff": { "version": "5.0.0", "resolved": "https://registry.npmjs.org/diff/-/diff-5.0.0.tgz", @@ -2487,6 +2531,20 @@ "node": ">=8.0.0" } }, + "node_modules/form-data": { + "version": "3.0.1", + "resolved": "https://registry.npmjs.org/form-data/-/form-data-3.0.1.tgz", + "integrity": "sha512-RHkBKtLWUVwd7SqRIvCZMEvAMoGUp0XU+seQiZejj0COz3RI3hWP4sCv3gZWWLjJTd7rGwcsF5eKZGii0r/hbg==", + "dev": true, + "dependencies": { + "asynckit": "^0.4.0", + "combined-stream": "^1.0.8", + "mime-types": "^2.1.12" + }, + "engines": { + "node": ">= 6" + } + }, "node_modules/fromentries": { "version": "1.3.2", "resolved": "https://registry.npmjs.org/fromentries/-/fromentries-1.3.2.tgz", @@ -3217,6 +3275,7 @@ "version": "4.0.1", "resolved": "https://registry.npmjs.org/isomorphic-ws/-/isomorphic-ws-4.0.1.tgz", "integrity": "sha512-BhBvN2MBpWTaSHdWRb/bwdZJ1WaehQ2L1KngkCkfLUGF0mAWAT1sQUQacEmQ0jXkFw/czDXPNQSL5u2/Krsz1w==", + "inBundle": true, "peerDependencies": { "ws": "*" } @@ -3594,6 +3653,27 @@ "node": ">=8.6" } }, + "node_modules/mime-db": { + "version": "1.52.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.52.0.tgz", + "integrity": "sha512-sPU4uV7dYlvtWJxwwxHD0PuihVNiE7TyAbQ5SWxDCB9mUYvOgroQOwYQQOKPJ8CIbE+1ETVlOoK1UC2nU3gYvg==", + "dev": true, + "engines": { + "node": ">= 0.6" + } + }, + "node_modules/mime-types": { + "version": "2.1.35", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-2.1.35.tgz", + "integrity": "sha512-ZDY+bPm5zTTF+YpCrAU9nK0UgICYPT0QtT1NZWFv4s++TNkcgVaT0g6+4R2uI4MjQjzysHB1zxuWL50hzaeXiw==", + "dev": true, + "dependencies": { + "mime-db": "1.52.0" + }, + "engines": { + "node": ">= 0.6" + } + }, "node_modules/minimatch": { "version": "3.1.2", "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-3.1.2.tgz", @@ -3734,10 +3814,30 @@ "path-to-regexp": "^1.7.0" } }, + "node_modules/node-fetch": { + "version": "2.6.9", + "resolved": "https://registry.npmjs.org/node-fetch/-/node-fetch-2.6.9.tgz", + "integrity": "sha512-DJm/CJkZkRjKKj4Zi4BsKVZh3ValV5IR5s7LVZnW+6YMh0W1BfNA8XSs6DLMGYlId5F3KnA70uu2qepcR08Qqg==", + "dependencies": { + "whatwg-url": "^5.0.0" + }, + "engines": { + "node": "4.x || >=6.0.0" + }, + "peerDependencies": { + "encoding": "^0.1.0" + }, + "peerDependenciesMeta": { + "encoding": { + "optional": true + } + } + }, "node_modules/node-int64": { "version": "0.4.0", "resolved": "https://registry.npmjs.org/node-int64/-/node-int64-0.4.0.tgz", - "integrity": "sha512-O5lz91xSOeoXP6DulyHfllpq+Eg00MWitZIbtPfoSEvqIHdl5gfcY6hYzDWnj0qD5tz52PI08u9qUvSVeUBeHw==" + "integrity": "sha512-O5lz91xSOeoXP6DulyHfllpq+Eg00MWitZIbtPfoSEvqIHdl5gfcY6hYzDWnj0qD5tz52PI08u9qUvSVeUBeHw==", + "inBundle": true }, "node_modules/node-preload": { "version": "0.2.1", @@ -4482,6 +4582,7 @@ "version": "1.5.1", "resolved": "https://registry.npmjs.org/q/-/q-1.5.1.tgz", "integrity": "sha512-kV/CThkXo6xyFEZUugw/+pIOywXcDbFYgSct5cT3gqlbkBE1SJdwy6UQoZvodiWF/ckQLZyDE/Bu1M6gVu5lVw==", + "inBundle": true, "engines": { "node": ">=0.6.0", "teleport": ">=0.2.0" @@ -5022,6 +5123,7 @@ "version": "0.16.0", "resolved": "https://registry.npmjs.org/thrift/-/thrift-0.16.0.tgz", "integrity": "sha512-W8DpGyTPlIaK3f+e1XOCLxefaUWXtrOXAaVIDbfYhmVyriYeAKgsBVFNJUV1F9SQ2SPt2sG44AZQxSGwGj/3VA==", + "inBundle": true, "dependencies": { "browser-or-node": "^1.2.1", "isomorphic-ws": "^4.0.1", @@ -5064,6 +5166,11 @@ "node": ">=8.0" } }, + "node_modules/tr46": { + "version": "0.0.3", + "resolved": "https://registry.npmjs.org/tr46/-/tr46-0.0.3.tgz", + "integrity": "sha512-N3WMsuqV66lT30CrXNbEjx4GEwlow3v6rr4mCcv6prnfwhS01rkgyFdjPNBYd9br7LpXV1+Emh01fHnq2Gdgrw==" + }, "node_modules/triple-beam": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/triple-beam/-/triple-beam-1.3.0.tgz", @@ -5255,6 +5362,20 @@ "integrity": "sha512-l8lCEmLcLYZh4nbunNZvQCJc5pv7+RCwa8q/LdUx8u7lsWvPDKmpodJAJNwkAhJC//dFY48KuIEmjtd4RViDrA==", "dev": true }, + "node_modules/webidl-conversions": { + "version": "3.0.1", + "resolved": "https://registry.npmjs.org/webidl-conversions/-/webidl-conversions-3.0.1.tgz", + "integrity": "sha512-2JAn3z8AR6rjK8Sm8orRC0h/bcl/DqL7tRPdGZ4I1CjdF+EaMLmYxBHyXuKL849eucPFhvBoxMsflfOb8kxaeQ==" + }, + "node_modules/whatwg-url": { + "version": "5.0.0", + "resolved": "https://registry.npmjs.org/whatwg-url/-/whatwg-url-5.0.0.tgz", + "integrity": "sha512-saE57nupxk6v3HY35+jzBwYa0rKSy0XR8JSxZPwgLr7ys0IBzhGviA1/TUGJLmSVqs8pb9AnvICXEuOHLprYTw==", + "dependencies": { + "tr46": "~0.0.3", + "webidl-conversions": "^3.0.0" + } + }, "node_modules/which": { "version": "2.0.2", "resolved": "https://registry.npmjs.org/which/-/which-2.0.2.tgz", @@ -5379,6 +5500,7 @@ "version": "5.2.3", "resolved": "https://registry.npmjs.org/ws/-/ws-5.2.3.tgz", "integrity": "sha512-jZArVERrMsKUatIdnLzqvcfydI85dvd/Fp1u/VOpfdDWQ4c9qWXe+VIeAbQ5FrDwciAkr+lzofXLz3Kuf26AOA==", + "inBundle": true, "dependencies": { "async-limiter": "~1.0.0" } @@ -6080,6 +6202,16 @@ "integrity": "sha512-NcKK6Ts+9LqdHJaW6HQmgr7dT/i3GOHG+pt6BiWv++5SnjtRd4NXeiuN2kA153SjhXPR/AhHIPHPbrsbpUVOww==", "dev": true }, + "@types/node-fetch": { + "version": "2.6.2", + "resolved": "https://registry.npmjs.org/@types/node-fetch/-/node-fetch-2.6.2.tgz", + "integrity": "sha512-DHqhlq5jeESLy19TYhLakJ07kNumXWjcDdxXsLUMJZ6ue8VZJj4kLPQVE/2mdHh3xZziNF1xppu5lwmS53HR+A==", + "dev": true, + "requires": { + "@types/node": "*", + "form-data": "^3.0.0" + } + }, "@types/node-int64": { "version": "0.4.29", "resolved": "https://registry.npmjs.org/@types/node-int64/-/node-int64-0.4.29.tgz", @@ -6379,6 +6511,12 @@ "resolved": "https://registry.npmjs.org/async-limiter/-/async-limiter-1.0.1.tgz", "integrity": "sha512-csOlWGAcRFJaI6m+F2WKdnMKr4HhdhFVBk0H/QbJFMCr+uO2kwohwXQPxw/9OCxp05r5ghVBFSyioixx3gfkNQ==" }, + "asynckit": { + "version": "0.4.0", + "resolved": "https://registry.npmjs.org/asynckit/-/asynckit-0.4.0.tgz", + "integrity": "sha512-Oei9OH4tRh0YqU3GxhX79dM/mwVgvbZJaSNaRk+bshkj0S5cfHcgYakreBjrHwatXKbz+IoIdYLxrKim2MjW0Q==", + "dev": true + }, "axe-core": { "version": "4.4.3", "resolved": "https://registry.npmjs.org/axe-core/-/axe-core-4.4.3.tgz", @@ -6616,6 +6754,15 @@ "text-hex": "1.0.x" } }, + "combined-stream": { + "version": "1.0.8", + "resolved": "https://registry.npmjs.org/combined-stream/-/combined-stream-1.0.8.tgz", + "integrity": "sha512-FQN4MRfuJeHf7cBbBMJFXhKSDq+2kAArBlmRBvcvFE5BB1HZKXtSFASDhdlz9zOYwxh8lDdnvmMOe/+5cdoEdg==", + "dev": true, + "requires": { + "delayed-stream": "~1.0.0" + } + }, "commander": { "version": "9.4.0", "resolved": "https://registry.npmjs.org/commander/-/commander-9.4.0.tgz", @@ -6719,6 +6866,12 @@ "object-keys": "^1.1.1" } }, + "delayed-stream": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/delayed-stream/-/delayed-stream-1.0.0.tgz", + "integrity": "sha512-ZySD7Nf91aLB0RxL4KGrKHBXl7Eds1DAmEdcoVawXnLD7SDhpNgtuII2aAkg7a7QS41jxPSZ17p4VdGnMHk3MQ==", + "dev": true + }, "diff": { "version": "5.0.0", "resolved": "https://registry.npmjs.org/diff/-/diff-5.0.0.tgz", @@ -7336,6 +7489,17 @@ "signal-exit": "^3.0.2" } }, + "form-data": { + "version": "3.0.1", + "resolved": "https://registry.npmjs.org/form-data/-/form-data-3.0.1.tgz", + "integrity": "sha512-RHkBKtLWUVwd7SqRIvCZMEvAMoGUp0XU+seQiZejj0COz3RI3hWP4sCv3gZWWLjJTd7rGwcsF5eKZGii0r/hbg==", + "dev": true, + "requires": { + "asynckit": "^0.4.0", + "combined-stream": "^1.0.8", + "mime-types": "^2.1.12" + } + }, "fromentries": { "version": "1.3.2", "resolved": "https://registry.npmjs.org/fromentries/-/fromentries-1.3.2.tgz", @@ -8131,6 +8295,21 @@ "picomatch": "^2.3.1" } }, + "mime-db": { + "version": "1.52.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.52.0.tgz", + "integrity": "sha512-sPU4uV7dYlvtWJxwwxHD0PuihVNiE7TyAbQ5SWxDCB9mUYvOgroQOwYQQOKPJ8CIbE+1ETVlOoK1UC2nU3gYvg==", + "dev": true + }, + "mime-types": { + "version": "2.1.35", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-2.1.35.tgz", + "integrity": "sha512-ZDY+bPm5zTTF+YpCrAU9nK0UgICYPT0QtT1NZWFv4s++TNkcgVaT0g6+4R2uI4MjQjzysHB1zxuWL50hzaeXiw==", + "dev": true, + "requires": { + "mime-db": "1.52.0" + } + }, "minimatch": { "version": "3.1.2", "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-3.1.2.tgz", @@ -8244,6 +8423,14 @@ "path-to-regexp": "^1.7.0" } }, + "node-fetch": { + "version": "2.6.9", + "resolved": "https://registry.npmjs.org/node-fetch/-/node-fetch-2.6.9.tgz", + "integrity": "sha512-DJm/CJkZkRjKKj4Zi4BsKVZh3ValV5IR5s7LVZnW+6YMh0W1BfNA8XSs6DLMGYlId5F3KnA70uu2qepcR08Qqg==", + "requires": { + "whatwg-url": "^5.0.0" + } + }, "node-int64": { "version": "0.4.0", "resolved": "https://registry.npmjs.org/node-int64/-/node-int64-0.4.0.tgz", @@ -9217,6 +9404,11 @@ "is-number": "^7.0.0" } }, + "tr46": { + "version": "0.0.3", + "resolved": "https://registry.npmjs.org/tr46/-/tr46-0.0.3.tgz", + "integrity": "sha512-N3WMsuqV66lT30CrXNbEjx4GEwlow3v6rr4mCcv6prnfwhS01rkgyFdjPNBYd9br7LpXV1+Emh01fHnq2Gdgrw==" + }, "triple-beam": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/triple-beam/-/triple-beam-1.3.0.tgz", @@ -9354,6 +9546,20 @@ "integrity": "sha512-l8lCEmLcLYZh4nbunNZvQCJc5pv7+RCwa8q/LdUx8u7lsWvPDKmpodJAJNwkAhJC//dFY48KuIEmjtd4RViDrA==", "dev": true }, + "webidl-conversions": { + "version": "3.0.1", + "resolved": "https://registry.npmjs.org/webidl-conversions/-/webidl-conversions-3.0.1.tgz", + "integrity": "sha512-2JAn3z8AR6rjK8Sm8orRC0h/bcl/DqL7tRPdGZ4I1CjdF+EaMLmYxBHyXuKL849eucPFhvBoxMsflfOb8kxaeQ==" + }, + "whatwg-url": { + "version": "5.0.0", + "resolved": "https://registry.npmjs.org/whatwg-url/-/whatwg-url-5.0.0.tgz", + "integrity": "sha512-saE57nupxk6v3HY35+jzBwYa0rKSy0XR8JSxZPwgLr7ys0IBzhGviA1/TUGJLmSVqs8pb9AnvICXEuOHLprYTw==", + "requires": { + "tr46": "~0.0.3", + "webidl-conversions": "^3.0.0" + } + }, "which": { "version": "2.0.2", "resolved": "https://registry.npmjs.org/which/-/which-2.0.2.tgz", diff --git a/package.json b/package.json index 709214fc..27988d6f 100644 --- a/package.json +++ b/package.json @@ -49,6 +49,7 @@ "license": "Apache 2.0", "devDependencies": { "@types/node": "^18.0.0", + "@types/node-fetch": "^2.6.2", "@types/node-int64": "^0.4.29", "@types/thrift": "^0.10.11", "@types/uuid": "^8.3.4", @@ -71,6 +72,7 @@ }, "dependencies": { "commander": "^9.3.0", + "node-fetch": "^2.6.9", "node-int64": "^0.4.0", "patch-package": "^6.5.0", "thrift": "^0.16.0", From 906bedb86a7a49a2cce0a57aaced7cccb169c754 Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Tue, 7 Mar 2023 17:11:23 +0200 Subject: [PATCH 2/4] Implement Exec API facade so it can be used with existing Thrift-based codebase (PoC) Signed-off-by: Levko Kravets --- lib/DBSQLClient.ts | 25 +- lib/DBSQLOperation/CompleteOperationHelper.ts | 6 +- lib/DBSQLOperation/FetchResultsHelper.ts | 6 +- lib/DBSQLOperation/OperationStatusHelper.ts | 6 +- lib/DBSQLOperation/SchemaHelper.ts | 6 +- lib/DBSQLOperation/index.ts | 6 +- lib/DBSQLSession.ts | 5 +- lib/rest/RestClient.ts | 5 + lib/rest/RestDriver.ts | 408 ++++++++++++++++++ lib/rest/Types.ts | 8 +- 10 files changed, 461 insertions(+), 20 deletions(-) create mode 100644 lib/rest/RestDriver.ts diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts index a40c740e..e500f3b7 100644 --- a/lib/DBSQLClient.ts +++ b/lib/DBSQLClient.ts @@ -21,6 +21,9 @@ import PlainHttpAuthentication from './connection/auth/PlainHttpAuthentication'; import IDBSQLLogger, { LogLevel } from './contracts/IDBSQLLogger'; import DBSQLLogger from './DBSQLLogger'; +import RestClient from './rest/RestClient'; +import RestDriver from './rest/RestDriver'; + function prependSlash(str: string): string { if (str.length > 0 && str.charAt(0) !== '/') { return `/${str}`; @@ -42,7 +45,8 @@ function getInitialNamespaceOptions(catalogName?: string, schemaName?: string) { } export default class DBSQLClient extends EventEmitter implements IDBSQLClient { - private client: TCLIService.Client | null; + private client: RestClient | null; + private driver: RestDriver | null; private connection: IThriftConnection | null; @@ -63,6 +67,7 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient { this.statusFactory = new StatusFactory(); this.logger = options?.logger || new DBSQLLogger(); this.client = null; + this.driver = null; this.connection = null; this.logger.log(LogLevel.info, 'Created DBSQLClient'); } @@ -99,7 +104,19 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient { this.connection = await this.connectionProvider.connect(this.getConnectionOptions(options), this.authProvider); - this.client = this.thrift.createClient(TCLIService, this.connection.getConnection()); + // this.client = this.thrift.createClient(TCLIService, this.connection.getConnection()); + this.client = new RestClient({ + host: options.host, + warehouseId: + options.path + .split('/') + .filter((s) => s !== '') + .pop() || '', // take last path fragment + headers: { + Authorization: `Basic ${Buffer.from(`token:${options.token}`).toString('base64')}`, + }, + }); + this.driver = new RestDriver(this.client); this.connection.getConnection().on('error', (error: Error) => { // Error.stack already contains error type and message, so log stack if available, @@ -141,11 +158,11 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient { * const session = await client.openSession(); */ openSession(request: OpenSessionRequest = {}): Promise { - if (!this.connection?.isConnected()) { + if (!this.connection?.isConnected() || !this.driver) { return Promise.reject(new HiveDriverError('DBSQLClient: connection is lost')); } - const driver = new HiveDriver(this.getClient()); + const driver = this.driver; return driver .openSession({ diff --git a/lib/DBSQLOperation/CompleteOperationHelper.ts b/lib/DBSQLOperation/CompleteOperationHelper.ts index d24b24be..7e1a0875 100644 --- a/lib/DBSQLOperation/CompleteOperationHelper.ts +++ b/lib/DBSQLOperation/CompleteOperationHelper.ts @@ -3,8 +3,10 @@ import HiveDriver from '../hive/HiveDriver'; import StatusFactory from '../factory/StatusFactory'; import Status from '../dto/Status'; +import RestDriver from '../rest/RestDriver'; + export default class CompleteOperationHelper { - private driver: HiveDriver; + private driver: RestDriver; private operationHandle: TOperationHandle; @@ -14,7 +16,7 @@ export default class CompleteOperationHelper { cancelled: boolean = false; - constructor(driver: HiveDriver, operationHandle: TOperationHandle, closeOperation?: TCloseOperationResp) { + constructor(driver: RestDriver, operationHandle: TOperationHandle, closeOperation?: TCloseOperationResp) { this.driver = driver; this.operationHandle = operationHandle; diff --git a/lib/DBSQLOperation/FetchResultsHelper.ts b/lib/DBSQLOperation/FetchResultsHelper.ts index 13dfc6cc..58fb33b8 100644 --- a/lib/DBSQLOperation/FetchResultsHelper.ts +++ b/lib/DBSQLOperation/FetchResultsHelper.ts @@ -10,6 +10,8 @@ import { ColumnCode, FetchType, Int64 } from '../hive/Types'; import HiveDriver from '../hive/HiveDriver'; import StatusFactory from '../factory/StatusFactory'; +import RestDriver from '../rest/RestDriver'; + function checkIfOperationHasMoreRows(response: TFetchResultsResp): boolean { if (response.hasMoreRows) { return true; @@ -37,7 +39,7 @@ function checkIfOperationHasMoreRows(response: TFetchResultsResp): boolean { } export default class FetchResultsHelper { - private driver: HiveDriver; + private driver: RestDriver; private operationHandle: TOperationHandle; @@ -50,7 +52,7 @@ export default class FetchResultsHelper { hasMoreRows: boolean = false; constructor( - driver: HiveDriver, + driver: RestDriver, operationHandle: TOperationHandle, prefetchedResults: Array, ) { diff --git a/lib/DBSQLOperation/OperationStatusHelper.ts b/lib/DBSQLOperation/OperationStatusHelper.ts index 45aa74ec..f3ee3dab 100644 --- a/lib/DBSQLOperation/OperationStatusHelper.ts +++ b/lib/DBSQLOperation/OperationStatusHelper.ts @@ -4,6 +4,8 @@ import StatusFactory from '../factory/StatusFactory'; import { WaitUntilReadyOptions } from '../contracts/IOperation'; import OperationStateError from '../errors/OperationStateError'; +import RestDriver from '../rest/RestDriver'; + async function delay(ms?: number): Promise { return new Promise((resolve) => { setTimeout(() => { @@ -13,7 +15,7 @@ async function delay(ms?: number): Promise { } export default class OperationStatusHelper { - private driver: HiveDriver; + private driver: RestDriver; private operationHandle: TOperationHandle; @@ -27,7 +29,7 @@ export default class OperationStatusHelper { hasResultSet: boolean = false; - constructor(driver: HiveDriver, operationHandle: TOperationHandle, operationStatus?: TGetOperationStatusResp) { + constructor(driver: RestDriver, operationHandle: TOperationHandle, operationStatus?: TGetOperationStatusResp) { this.driver = driver; this.operationHandle = operationHandle; this.hasResultSet = operationHandle.hasResultSet; diff --git a/lib/DBSQLOperation/SchemaHelper.ts b/lib/DBSQLOperation/SchemaHelper.ts index 33d66f00..31b99bd3 100644 --- a/lib/DBSQLOperation/SchemaHelper.ts +++ b/lib/DBSQLOperation/SchemaHelper.ts @@ -6,8 +6,10 @@ import JsonResult from '../result/JsonResult'; import HiveDriverError from '../errors/HiveDriverError'; import { definedOrError } from '../utils'; +import RestDriver from '../rest/RestDriver'; + export default class SchemaHelper { - private driver: HiveDriver; + private driver: RestDriver; private operationHandle: TOperationHandle; @@ -15,7 +17,7 @@ export default class SchemaHelper { private metadata?: TGetResultSetMetadataResp; - constructor(driver: HiveDriver, operationHandle: TOperationHandle, metadata?: TGetResultSetMetadataResp) { + constructor(driver: RestDriver, operationHandle: TOperationHandle, metadata?: TGetResultSetMetadataResp) { this.driver = driver; this.operationHandle = operationHandle; this.metadata = metadata; diff --git a/lib/DBSQLOperation/index.ts b/lib/DBSQLOperation/index.ts index 92f286f0..355775a1 100644 --- a/lib/DBSQLOperation/index.ts +++ b/lib/DBSQLOperation/index.ts @@ -15,10 +15,12 @@ import FetchResultsHelper from './FetchResultsHelper'; import CompleteOperationHelper from './CompleteOperationHelper'; import IDBSQLLogger, { LogLevel } from '../contracts/IDBSQLLogger'; +import RestDriver from '../rest/RestDriver'; + const defaultMaxRows = 100000; export default class DBSQLOperation implements IOperation { - private driver: HiveDriver; + private driver: RestDriver; private operationHandle: TOperationHandle; @@ -33,7 +35,7 @@ export default class DBSQLOperation implements IOperation { private _completeOperation: CompleteOperationHelper; constructor( - driver: HiveDriver, + driver: RestDriver, operationHandle: TOperationHandle, logger: IDBSQLLogger, directResults?: TSparkDirectResults, diff --git a/lib/DBSQLSession.ts b/lib/DBSQLSession.ts index 6c31f19a..899671ed 100644 --- a/lib/DBSQLSession.ts +++ b/lib/DBSQLSession.ts @@ -21,6 +21,7 @@ import StatusFactory from './factory/StatusFactory'; import InfoValue from './dto/InfoValue'; import { definedOrError } from './utils'; import IDBSQLLogger, { LogLevel } from './contracts/IDBSQLLogger'; +import RestDriver from './rest/RestDriver'; const defaultMaxRows = 100000; @@ -43,7 +44,7 @@ function getDirectResultsOptions(maxRows: number | null = defaultMaxRows) { } export default class DBSQLSession implements IDBSQLSession { - private driver: HiveDriver; + private driver: RestDriver; private sessionHandle: TSessionHandle; @@ -51,7 +52,7 @@ export default class DBSQLSession implements IDBSQLSession { private logger: IDBSQLLogger; - constructor(driver: HiveDriver, sessionHandle: TSessionHandle, logger: IDBSQLLogger) { + constructor(driver: RestDriver, sessionHandle: TSessionHandle, logger: IDBSQLLogger) { this.driver = driver; this.sessionHandle = sessionHandle; this.statusFactory = new StatusFactory(); diff --git a/lib/rest/RestClient.ts b/lib/rest/RestClient.ts index 302f08ba..0e67fb47 100644 --- a/lib/rest/RestClient.ts +++ b/lib/rest/RestClient.ts @@ -12,6 +12,7 @@ import { export interface RestClientOptions { host: string; + warehouseId: string; headers?: Record; } @@ -42,6 +43,10 @@ export default class RestClient { this.options = options; } + public getWarehouseId() { + return this.options.warehouseId; + } + public executeStatement(request: ExecuteStatementRequest) { return this.doRequest( HTTPMethod.POST, diff --git a/lib/rest/RestDriver.ts b/lib/rest/RestDriver.ts new file mode 100644 index 00000000..6f45022f --- /dev/null +++ b/lib/rest/RestDriver.ts @@ -0,0 +1,408 @@ +import Int64 from 'node-int64'; + +import { + TCancelDelegationTokenReq, + TCancelDelegationTokenResp, + TCancelOperationReq, + TCancelOperationResp, + TCloseOperationReq, + TCloseOperationResp, + TCloseSessionReq, + TCloseSessionResp, + TColumn, + TColumnDesc, + TExecuteStatementReq, + TExecuteStatementResp, + TFetchResultsReq, + TFetchResultsResp, + TGetCatalogsReq, + TGetCatalogsResp, + TGetColumnsReq, + TGetColumnsResp, + TGetCrossReferenceReq, + TGetCrossReferenceResp, + TGetDelegationTokenReq, + TGetDelegationTokenResp, + TGetFunctionsReq, + TGetFunctionsResp, + TGetInfoReq, + TGetInfoResp, + TGetOperationStatusReq, + TGetOperationStatusResp, + TGetPrimaryKeysReq, + TGetPrimaryKeysResp, + TGetResultSetMetadataReq, + TGetResultSetMetadataResp, + TGetSchemasReq, + TGetSchemasResp, + TGetTablesReq, + TGetTablesResp, + TGetTableTypesReq, + TGetTableTypesResp, + TGetTypeInfoReq, + TGetTypeInfoResp, + THandleIdentifier, + TOpenSessionReq, + TOpenSessionResp, + TOperationHandle, + TOperationState, + TOperationType, + TProtocolVersion, + TRenewDelegationTokenReq, + TRenewDelegationTokenResp, + TRowSet, + TSessionHandle, + TSparkRowSetType, + TStatus, + TStatusCode, + TTypeId, +} from '../../thrift/TCLIService_types'; + +import RestClient from './RestClient'; +import { + ColumnInfoTypeName, + Disposition, + ExecuteStatementResponse, + Format, + ResultData, + ResultSchema, + StatementState, + StatementStatus, + TimeoutAction, +} from './Types'; + +class NotImplementedError extends Error { + constructor() { + super('RestDriver: Method not implemented'); + } +} + +function restOperationStateToThriftOperationState(state: StatementState): TOperationState { + switch (state) { + case StatementState.Canceled: + return TOperationState.CANCELED_STATE; + case StatementState.Closed: + return TOperationState.CLOSED_STATE; + case StatementState.Failed: + return TOperationState.ERROR_STATE; + case StatementState.Pending: + return TOperationState.PENDING_STATE; + case StatementState.Running: + return TOperationState.RUNNING_STATE; + case StatementState.Succeeded: + return TOperationState.FINISHED_STATE; + default: + return TOperationState.UKNOWN_STATE; + } +} + +function restTypeNameToThriftTypeId(typeName: ColumnInfoTypeName): TTypeId { + switch (typeName) { + case ColumnInfoTypeName.Array: + return TTypeId.ARRAY_TYPE; + case ColumnInfoTypeName.Binary: + return TTypeId.BINARY_TYPE; + case ColumnInfoTypeName.Boolean: + return TTypeId.BOOLEAN_TYPE; + case ColumnInfoTypeName.Byte: + return TTypeId.TINYINT_TYPE; // ?? + case ColumnInfoTypeName.Char: + return TTypeId.CHAR_TYPE; + case ColumnInfoTypeName.Date: + return TTypeId.DATE_TYPE; + case ColumnInfoTypeName.Decimal: + return TTypeId.DECIMAL_TYPE; + case ColumnInfoTypeName.Double: + return TTypeId.DOUBLE_TYPE; + case ColumnInfoTypeName.Float: + return TTypeId.FLOAT_TYPE; + case ColumnInfoTypeName.Int: + return TTypeId.INT_TYPE; + case ColumnInfoTypeName.Interval: + return TTypeId.INTERVAL_DAY_TIME_TYPE; // ?? + case ColumnInfoTypeName.Long: + return TTypeId.INT_TYPE; // ?? + case ColumnInfoTypeName.Map: + return TTypeId.MAP_TYPE; + case ColumnInfoTypeName.Null: + return TTypeId.NULL_TYPE; + case ColumnInfoTypeName.Short: + return TTypeId.SMALLINT_TYPE; // ?? + case ColumnInfoTypeName.String: + return TTypeId.STRING_TYPE; + case ColumnInfoTypeName.Struct: + return TTypeId.STRUCT_TYPE; + case ColumnInfoTypeName.Timestamp: + return TTypeId.TIMESTAMP_TYPE; + case ColumnInfoTypeName.UserDefinedType: + return TTypeId.USER_DEFINED_TYPE; + default: + return TTypeId.NULL_TYPE; // ?? + } +} + +function restJsonResultToThriftColumnar(schema: ResultSchema, result: ResultData): TRowSet { + const columns: TColumn[] = schema.columns.map(() => ({ + stringVal: { values: [], nulls: Buffer.alloc(0) }, + })); + + if (result.data_array) { + result.data_array.forEach((row) => { + for (let i = 0; i < row.length; i++) { + columns[i].stringVal?.values.push(row[i]); + } + }); + } + + return { + startRowOffset: new Int64(0), + rows: [], + columns, + columnCount: columns.length, + }; +} + +export default class RestDriver { + private client: RestClient; + + // REST API doesn't support sessions; we emulate it, but only single session is supported in this PoC + private session?: TOpenSessionReq = undefined; + + // In this PoC we don't support concurrent queries, so just store currently processed query id to simplify things + private currentStatement?: ExecuteStatementResponse = undefined; + private currentResultChunk?: ResultData = undefined; + + constructor(client: RestClient) { + this.client = client; + } + + private processStatementStatus(status: StatementStatus) { + switch (status.state) { + case StatementState.Canceled: + case StatementState.Closed: + this.currentStatement = undefined; + return; + case StatementState.Failed: + throw new Error(`${status.error.error_code}: ${status.error.message}`); + default: + return; + } + } + + async openSession(request: TOpenSessionReq): Promise { + this.session = request; + return new TOpenSessionResp({ + status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), + serverProtocolVersion: request.client_protocol || TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V6, + sessionHandle: new TSessionHandle({ + sessionId: new THandleIdentifier({ + guid: Buffer.alloc(16), + secret: Buffer.alloc(0), + }), + }), + }); + } + + async closeSession(request: TCloseSessionReq): Promise { + this.session = undefined; + return new TCloseSessionResp({ + status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), + }); + } + + async executeStatement(request: TExecuteStatementReq): Promise { + this.currentStatement = await this.client.executeStatement({ + catalog: this.session?.initialNamespace?.catalogName, + schema: this.session?.initialNamespace?.schemaName, + disposition: Disposition.Inline, + format: Format.JsonArray, + on_wait_timeout: TimeoutAction.Continue, + wait_timeout: '5s', + warehouse_id: this.client.getWarehouseId(), + statement: request.statement, + }); + + this.processStatementStatus(this.currentStatement.status); + + if (!this.currentStatement) { + return new TExecuteStatementResp({ + status: new TStatus({ statusCode: TStatusCode.INVALID_HANDLE_STATUS }), + }); + } + + return new TExecuteStatementResp({ + status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), + operationHandle: new TOperationHandle({ + operationId: new THandleIdentifier({ + guid: Buffer.alloc(16), + secret: Buffer.alloc(0), + }), + operationType: TOperationType.EXECUTE_STATEMENT, + hasResultSet: this.currentStatement?.result?.chunk_index !== undefined, + }), + }); + } + + async getResultSetMetadata(request: TGetResultSetMetadataReq): Promise { + if (!this.currentStatement) { + return new TGetResultSetMetadataResp({ + status: new TStatus({ statusCode: TStatusCode.ERROR_STATUS, errorMessage: 'Invalid handle' }), + }); + } + + const columns: TColumnDesc[] = []; + this.currentStatement.manifest?.schema?.columns?.forEach((column) => { + columns.push({ + columnName: column.name, + position: column.position + 1, // thrift columns are 1-based, rest columns are 0-based + typeDesc: { + types: [ + { + primitiveEntry: { + type: restTypeNameToThriftTypeId(column.type_name), + }, + }, + ], + }, + }); + }); + + return new TGetResultSetMetadataResp({ + status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), + resultFormat: TSparkRowSetType.COLUMN_BASED_SET, + schema: { columns }, + }); + } + + async fetchResults(request: TFetchResultsReq): Promise { + if (!this.currentStatement) { + return new TFetchResultsResp({ + status: new TStatus({ statusCode: TStatusCode.ERROR_STATUS, errorMessage: 'Invalid handle' }), + }); + } + + if (!this.currentResultChunk) { + this.currentResultChunk = this.currentStatement.result; + } else { + this.currentResultChunk = await this.client.getStatementResultChunkN({ + statement_id: this.currentStatement.statement_id, + chunk_index: this.currentResultChunk.next_chunk_index, + }); + } + + if (!this.currentResultChunk) { + return new TFetchResultsResp({ + status: new TStatus({ statusCode: TStatusCode.ERROR_STATUS, errorMessage: 'No more data' }), + }); + } + + const schema = this.currentStatement.manifest?.schema || { + column_count: 0, + columns: [], + }; + + return new TFetchResultsResp({ + status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), + hasMoreRows: this.currentResultChunk?.next_chunk_index !== undefined, + results: restJsonResultToThriftColumnar(schema, this.currentResultChunk), + }); + } + + async getInfo(request: TGetInfoReq): Promise { + throw new NotImplementedError(); + } + + async getTypeInfo(request: TGetTypeInfoReq): Promise { + throw new NotImplementedError(); + } + + async getCatalogs(request: TGetCatalogsReq): Promise { + throw new NotImplementedError(); + } + + async getSchemas(request: TGetSchemasReq): Promise { + throw new NotImplementedError(); + } + + async getTables(request: TGetTablesReq): Promise { + throw new NotImplementedError(); + } + + async getTableTypes(request: TGetTableTypesReq): Promise { + throw new NotImplementedError(); + } + + async getColumns(request: TGetColumnsReq): Promise { + throw new NotImplementedError(); + } + + async getFunctions(request: TGetFunctionsReq): Promise { + throw new NotImplementedError(); + } + + async getPrimaryKeys(request: TGetPrimaryKeysReq): Promise { + throw new NotImplementedError(); + } + + async getCrossReference(request: TGetCrossReferenceReq): Promise { + throw new NotImplementedError(); + } + + async getOperationStatus(request: TGetOperationStatusReq): Promise { + if (this.currentStatement && !this.currentStatement.manifest) { + const response = await this.client.getStatement({ statement_id: this.currentStatement.statement_id }); + this.currentStatement.status = response.status; + this.currentStatement.manifest = response.manifest; + this.currentStatement.result = response.result; + + console.log(JSON.stringify(response, null, 2)); + + this.processStatementStatus(this.currentStatement.status); + } + + if (!this.currentStatement) { + return new TGetOperationStatusResp({ + status: new TStatus({ statusCode: TStatusCode.ERROR_STATUS, errorMessage: 'Invalid handle' }), + }); + } + + return new TGetOperationStatusResp({ + status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), + operationState: restOperationStateToThriftOperationState(this.currentStatement.status.state), + errorCode: 0, + errorMessage: `${this.currentStatement.status.error?.error_code}: ${this.currentStatement.status.error?.message}`, + hasResultSet: Boolean(this.currentStatement.result), + }); + } + + async cancelOperation(request: TCancelOperationReq): Promise { + if (this.currentStatement) { + await this.client.cancelExecution({ + statement_id: this.currentStatement.statement_id, + }); + this.currentStatement = undefined; + } + return new TCancelOperationResp({ + status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), + }); + } + + async closeOperation(request: TCloseOperationReq): Promise { + this.currentStatement = undefined; + return new TCloseOperationResp({ + status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), + }); + } + + async getDelegationToken(request: TGetDelegationTokenReq): Promise { + throw new NotImplementedError(); + } + + async cancelDelegationToken(request: TCancelDelegationTokenReq): Promise { + throw new NotImplementedError(); + } + + async renewDelegationToken(request: TRenewDelegationTokenReq): Promise { + throw new NotImplementedError(); + } +} diff --git a/lib/rest/Types.ts b/lib/rest/Types.ts index 3139b50f..f807130e 100644 --- a/lib/rest/Types.ts +++ b/lib/rest/Types.ts @@ -14,19 +14,19 @@ export enum TimeoutAction { } export interface ExecuteStatementRequest { - catalog: string; + catalog?: string; disposition: Disposition; format: Format; on_wait_timeout: TimeoutAction; - schema: string; + schema?: string; statement: string; wait_timeout: string; warehouse_id: string; } export interface ExecuteStatementResponse { - manifest: ResultManifest; - result: ResultData; + manifest?: ResultManifest; + result?: ResultData; statement_id: string; status: StatementStatus; } From dda8ffed1c706af39a068cbb48a556918e06d90e Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Tue, 14 Mar 2023 19:15:28 +0200 Subject: [PATCH 3/4] Switch to Arrow-based results --- lib/DBSQLOperation/FetchResultsHelper.ts | 24 +----- lib/rest/RestClient.ts | 20 ++++- lib/rest/RestDriver.ts | 94 +++++++++++++----------- lib/rest/Types.ts | 6 +- 4 files changed, 73 insertions(+), 71 deletions(-) diff --git a/lib/DBSQLOperation/FetchResultsHelper.ts b/lib/DBSQLOperation/FetchResultsHelper.ts index 58fb33b8..d0a99567 100644 --- a/lib/DBSQLOperation/FetchResultsHelper.ts +++ b/lib/DBSQLOperation/FetchResultsHelper.ts @@ -13,29 +13,7 @@ import StatusFactory from '../factory/StatusFactory'; import RestDriver from '../rest/RestDriver'; function checkIfOperationHasMoreRows(response: TFetchResultsResp): boolean { - if (response.hasMoreRows) { - return true; - } - - const columns = response.results?.columns || []; - - if (columns.length === 0) { - return false; - } - - const column: TColumn = columns[0]; - - const columnValue = - column[ColumnCode.binaryVal] || - column[ColumnCode.boolVal] || - column[ColumnCode.byteVal] || - column[ColumnCode.doubleVal] || - column[ColumnCode.i16Val] || - column[ColumnCode.i32Val] || - column[ColumnCode.i64Val] || - column[ColumnCode.stringVal]; - - return (columnValue?.values?.length || 0) > 0; + return response.hasMoreRows || false; } export default class FetchResultsHelper { diff --git a/lib/rest/RestClient.ts b/lib/rest/RestClient.ts index 0e67fb47..7f046a4c 100644 --- a/lib/rest/RestClient.ts +++ b/lib/rest/RestClient.ts @@ -22,7 +22,7 @@ enum HTTPMethod { } export default class RestClient { - private options: RestClientOptions; + private readonly options: RestClientOptions; private async doRequest(method: string, path: string, payload: P): Promise { const { host, headers } = this.options; @@ -78,4 +78,22 @@ export default class RestClient { request, ); } + + public getStatementResultChunk(internalLink: string): Promise { + return this.doRequest( + HTTPMethod.GET, + internalLink, + undefined, + ); + } + + public async fetchExternalLink(url: string): Promise { + const { host, headers } = this.options; + const response = await fetch(url, { + method: HTTPMethod.GET, + headers: {}, + }); + const result = await response.arrayBuffer(); + return Buffer.from(result); + } } diff --git a/lib/rest/RestDriver.ts b/lib/rest/RestDriver.ts index 6f45022f..d9932c50 100644 --- a/lib/rest/RestDriver.ts +++ b/lib/rest/RestDriver.ts @@ -141,27 +141,6 @@ function restTypeNameToThriftTypeId(typeName: ColumnInfoTypeName): TTypeId { } } -function restJsonResultToThriftColumnar(schema: ResultSchema, result: ResultData): TRowSet { - const columns: TColumn[] = schema.columns.map(() => ({ - stringVal: { values: [], nulls: Buffer.alloc(0) }, - })); - - if (result.data_array) { - result.data_array.forEach((row) => { - for (let i = 0; i < row.length; i++) { - columns[i].stringVal?.values.push(row[i]); - } - }); - } - - return { - startRowOffset: new Int64(0), - rows: [], - columns, - columnCount: columns.length, - }; -} - export default class RestDriver { private client: RestClient; @@ -170,7 +149,9 @@ export default class RestDriver { // In this PoC we don't support concurrent queries, so just store currently processed query id to simplify things private currentStatement?: ExecuteStatementResponse = undefined; - private currentResultChunk?: ResultData = undefined; + + private resultLinks: Array = []; + private nextChunkLinks: Array = []; constructor(client: RestClient) { this.client = client; @@ -189,6 +170,21 @@ export default class RestDriver { } } + private processResultData(result?: ResultData) { + if (!result) return; + + console.log(result); + + result.external_links?.forEach((link) => { + if (link.external_link) { + this.resultLinks.push(link.external_link); + } + if (link.next_chunk_internal_link) { + this.nextChunkLinks.push(link.next_chunk_internal_link); + } + }); + } + async openSession(request: TOpenSessionReq): Promise { this.session = request; return new TOpenSessionResp({ @@ -214,10 +210,10 @@ export default class RestDriver { this.currentStatement = await this.client.executeStatement({ catalog: this.session?.initialNamespace?.catalogName, schema: this.session?.initialNamespace?.schemaName, - disposition: Disposition.Inline, - format: Format.JsonArray, + disposition: Disposition.ExternalLinks, + format: Format.ArrowStream, on_wait_timeout: TimeoutAction.Continue, - wait_timeout: '5s', + wait_timeout: '30s', warehouse_id: this.client.getWarehouseId(), statement: request.statement, }); @@ -230,6 +226,8 @@ export default class RestDriver { }); } + this.processResultData(this.currentStatement?.result); + return new TExecuteStatementResp({ status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), operationHandle: new TOperationHandle({ @@ -238,7 +236,7 @@ export default class RestDriver { secret: Buffer.alloc(0), }), operationType: TOperationType.EXECUTE_STATEMENT, - hasResultSet: this.currentStatement?.result?.chunk_index !== undefined, + hasResultSet: this.resultLinks.length > 0 || this.nextChunkLinks.length > 0, }), }); } @@ -269,8 +267,9 @@ export default class RestDriver { return new TGetResultSetMetadataResp({ status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), - resultFormat: TSparkRowSetType.COLUMN_BASED_SET, + resultFormat: TSparkRowSetType.ARROW_BASED_SET, schema: { columns }, + arrowSchema: Buffer.alloc(0), }); } @@ -281,30 +280,37 @@ export default class RestDriver { }); } - if (!this.currentResultChunk) { - this.currentResultChunk = this.currentStatement.result; - } else { - this.currentResultChunk = await this.client.getStatementResultChunkN({ - statement_id: this.currentStatement.statement_id, - chunk_index: this.currentResultChunk.next_chunk_index, - }); + console.log(this.resultLinks.length, this.nextChunkLinks.length); + + if (this.resultLinks.length === 0) { + const nextChunkLink = this.nextChunkLinks.pop(); + if (nextChunkLink) { + const result = await this.client.getStatementResultChunk(nextChunkLink); + this.processResultData(result); + } } - if (!this.currentResultChunk) { + const resultLink = this.resultLinks.pop(); + if (resultLink) { + const arrowData = await this.client.fetchExternalLink(resultLink); return new TFetchResultsResp({ - status: new TStatus({ statusCode: TStatusCode.ERROR_STATUS, errorMessage: 'No more data' }), + status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), + hasMoreRows: this.resultLinks.length > 0 || this.nextChunkLinks.length > 0, + results: { + startRowOffset: new Int64(0), + rows: [], + arrowBatches: [ + { + batch: arrowData, + rowCount: new Int64(0), + }, + ], + }, }); } - const schema = this.currentStatement.manifest?.schema || { - column_count: 0, - columns: [], - }; - return new TFetchResultsResp({ - status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), - hasMoreRows: this.currentResultChunk?.next_chunk_index !== undefined, - results: restJsonResultToThriftColumnar(schema, this.currentResultChunk), + status: new TStatus({ statusCode: TStatusCode.ERROR_STATUS, errorMessage: 'No more data' }), }); } diff --git a/lib/rest/Types.ts b/lib/rest/Types.ts index f807130e..f5e2e77e 100644 --- a/lib/rest/Types.ts +++ b/lib/rest/Types.ts @@ -71,9 +71,9 @@ export interface ResultData { byte_count: number; // int64 chunk_index: number; data_array: string[][]; - external_links: ExternalLink[]; - next_chunk_index: number; - next_chunk_internal_link: string; + external_links?: ExternalLink[]; + next_chunk_index?: number; + next_chunk_internal_link?: string; row_count: number; // int64 row_offset: number; // int64 } From a4c382d50c9e538bd684baa7ce517b5d900e120c Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Tue, 14 Mar 2023 20:11:05 +0200 Subject: [PATCH 4/4] Properly wait for warehouse start and query execution complete Signed-off-by: Levko Kravets --- lib/rest/RestClient.ts | 6 +----- lib/rest/RestDriver.ts | 27 ++++++++++++++++----------- lib/rest/Types.ts | 4 ++-- 3 files changed, 19 insertions(+), 18 deletions(-) diff --git a/lib/rest/RestClient.ts b/lib/rest/RestClient.ts index 7f046a4c..f03f1bbf 100644 --- a/lib/rest/RestClient.ts +++ b/lib/rest/RestClient.ts @@ -80,11 +80,7 @@ export default class RestClient { } public getStatementResultChunk(internalLink: string): Promise { - return this.doRequest( - HTTPMethod.GET, - internalLink, - undefined, - ); + return this.doRequest(HTTPMethod.GET, internalLink, undefined); } public async fetchExternalLink(url: string): Promise { diff --git a/lib/rest/RestDriver.ts b/lib/rest/RestDriver.ts index d9932c50..832f37c5 100644 --- a/lib/rest/RestDriver.ts +++ b/lib/rest/RestDriver.ts @@ -173,8 +173,6 @@ export default class RestDriver { private processResultData(result?: ResultData) { if (!result) return; - console.log(result); - result.external_links?.forEach((link) => { if (link.external_link) { this.resultLinks.push(link.external_link); @@ -185,6 +183,16 @@ export default class RestDriver { }); } + private isPendingOrRunning() { + const state = this.currentStatement?.status?.state; + return state === StatementState.Pending || state === StatementState.Running; + } + + private hasResultSet() { + if (this.isPendingOrRunning()) return true; + return this.resultLinks.length > 0 || this.nextChunkLinks.length > 0; + } + async openSession(request: TOpenSessionReq): Promise { this.session = request; return new TOpenSessionResp({ @@ -213,7 +221,7 @@ export default class RestDriver { disposition: Disposition.ExternalLinks, format: Format.ArrowStream, on_wait_timeout: TimeoutAction.Continue, - wait_timeout: '30s', + wait_timeout: '5s', warehouse_id: this.client.getWarehouseId(), statement: request.statement, }); @@ -236,7 +244,7 @@ export default class RestDriver { secret: Buffer.alloc(0), }), operationType: TOperationType.EXECUTE_STATEMENT, - hasResultSet: this.resultLinks.length > 0 || this.nextChunkLinks.length > 0, + hasResultSet: this.hasResultSet(), }), }); } @@ -280,8 +288,6 @@ export default class RestDriver { }); } - console.log(this.resultLinks.length, this.nextChunkLinks.length); - if (this.resultLinks.length === 0) { const nextChunkLink = this.nextChunkLinks.pop(); if (nextChunkLink) { @@ -295,7 +301,7 @@ export default class RestDriver { const arrowData = await this.client.fetchExternalLink(resultLink); return new TFetchResultsResp({ status: new TStatus({ statusCode: TStatusCode.SUCCESS_STATUS }), - hasMoreRows: this.resultLinks.length > 0 || this.nextChunkLinks.length > 0, + hasMoreRows: this.hasResultSet(), results: { startRowOffset: new Int64(0), rows: [], @@ -355,15 +361,14 @@ export default class RestDriver { } async getOperationStatus(request: TGetOperationStatusReq): Promise { - if (this.currentStatement && !this.currentStatement.manifest) { + if (this.currentStatement && this.isPendingOrRunning()) { const response = await this.client.getStatement({ statement_id: this.currentStatement.statement_id }); this.currentStatement.status = response.status; this.currentStatement.manifest = response.manifest; this.currentStatement.result = response.result; - console.log(JSON.stringify(response, null, 2)); - this.processStatementStatus(this.currentStatement.status); + this.processResultData(this.currentStatement?.result); } if (!this.currentStatement) { @@ -377,7 +382,7 @@ export default class RestDriver { operationState: restOperationStateToThriftOperationState(this.currentStatement.status.state), errorCode: 0, errorMessage: `${this.currentStatement.status.error?.error_code}: ${this.currentStatement.status.error?.message}`, - hasResultSet: Boolean(this.currentStatement.result), + hasResultSet: this.hasResultSet(), }); } diff --git a/lib/rest/Types.ts b/lib/rest/Types.ts index f5e2e77e..a283b807 100644 --- a/lib/rest/Types.ts +++ b/lib/rest/Types.ts @@ -153,8 +153,8 @@ export interface GetStatementRequest { } export interface GetStatementResponse { - manifest: ResultManifest; - result: ResultData; + manifest?: ResultManifest; + result?: ResultData; statement_id: string; status: StatementStatus; }