From b76134f4c06c0844c9df0e6fc0ce0dfea611f869 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Sat, 24 Oct 2026 18:22:24 +0000 Subject: [PATCH] Fix Thrift primary-key metadata for missing objects Signed-off-by: Vu Anh Phung --- lib/thrift-backend/EmptyOperationBackend.ts | 44 +++ lib/thrift-backend/ThriftOperationBackend.ts | 44 ++- lib/thrift-backend/ThriftSessionBackend.ts | 24 +- lib/thrift-backend/metadataErrors.ts | 17 + lib/thrift-backend/primaryKeysSchema.ts | 14 + .../ThriftPrimaryKeysMetadata.test.ts | 325 ++++++++++++++++++ 6 files changed, 464 insertions(+), 4 deletions(-) create mode 100644 lib/thrift-backend/EmptyOperationBackend.ts create mode 100644 lib/thrift-backend/metadataErrors.ts create mode 100644 lib/thrift-backend/primaryKeysSchema.ts create mode 100644 tests/unit/thrift-backend/ThriftPrimaryKeysMetadata.test.ts diff --git a/lib/thrift-backend/EmptyOperationBackend.ts b/lib/thrift-backend/EmptyOperationBackend.ts new file mode 100644 index 00000000..daa9c13c --- /dev/null +++ b/lib/thrift-backend/EmptyOperationBackend.ts @@ -0,0 +1,44 @@ +import { v4 } from 'uuid'; +import { TTableSchema } from '../../thrift/TCLIService_types'; +import IOperationBackend from '../contracts/IOperationBackend'; +import { OperationState, OperationStatus } from '../contracts/OperationStatus'; +import { ResultFormat, ResultMetadata } from '../contracts/ResultMetadata'; +import Status from '../dto/Status'; + +export default class EmptyOperationBackend implements IOperationBackend { + public readonly id = v4(); + + constructor(private readonly schema: TTableSchema) {} + + public hasResultSet(): boolean { + return true; + } + + public async fetchChunk(): Promise> { + return []; + } + + public async hasMore(): Promise { + return false; + } + + public async waitUntilReady(): Promise { + return undefined; + } + + public async status(): Promise { + return { state: OperationState.Succeeded, hasResultSet: true }; + } + + public async getResultMetadata(): Promise { + return { schema: this.schema, resultFormat: ResultFormat.ColumnBased, isStagingOperation: false }; + } + + public async cancel(): Promise { + return Status.success(); + } + + public async close(): Promise { + return Status.success(); + } +} diff --git a/lib/thrift-backend/ThriftOperationBackend.ts b/lib/thrift-backend/ThriftOperationBackend.ts index 1111584e..58c95288 100644 --- a/lib/thrift-backend/ThriftOperationBackend.ts +++ b/lib/thrift-backend/ThriftOperationBackend.ts @@ -7,6 +7,8 @@ import { TSparkRowSetType, TCloseOperationResp, TOperationState, + TStatusCode, + TTableSchema, } from '../../thrift/TCLIService_types'; import IOperationBackend, { IOperationBackendWaitOptions } from '../contracts/IOperationBackend'; import IClientContext from '../contracts/IClientContext'; @@ -24,12 +26,15 @@ import ArrowResultConverter from '../result/ArrowResultConverter'; import ResultSlicer from '../result/ResultSlicer'; import { definedOrError } from '../utils'; import HiveDriverError from '../errors/HiveDriverError'; +import isObjectNotFound from './metadataErrors'; +import { synthesizeThriftResultSetMetadata } from './wireSynthesis'; interface ThriftOperationBackendOptions { handle: TOperationHandle; directResults?: TSparkDirectResults; context: IClientContext; functionCatalog?: string | null; + emptyResultOnNotFound?: TTableSchema; } async function delay(ms?: number): Promise { @@ -78,6 +83,8 @@ export default class ThriftOperationBackend implements IOperationBackend { private readonly functionCatalog?: string | null; + private readonly emptyResultOnNotFound?: TTableSchema; + private readonly operationHandle: TOperationHandle; private readonly _data: RowSetProvider; @@ -94,10 +101,18 @@ export default class ThriftOperationBackend implements IOperationBackend { private resultHandler?: ResultSlicer; - constructor({ handle, directResults, context, functionCatalog }: ThriftOperationBackendOptions) { + constructor({ + handle, + directResults, + context, + functionCatalog, + emptyResultOnNotFound, + }: ThriftOperationBackendOptions) { this.operationHandle = handle; this.context = context; this.functionCatalog = functionCatalog; + this.emptyResultOnNotFound = emptyResultOnNotFound; + this.metadata = directResults?.resultSetMetadata; const useOnlyPrefetchedResults = Boolean(directResults?.closeOperation); @@ -105,7 +120,6 @@ export default class ThriftOperationBackend implements IOperationBackend { this.processOperationStatusResponse(directResults.operationStatus); } - this.metadata = directResults?.resultSetMetadata; this._data = new RowSetProvider( this.context, this.operationHandle, @@ -130,7 +144,7 @@ export default class ThriftOperationBackend implements IOperationBackend { } public hasResultSet(): boolean { - return Boolean(this.operationHandle.hasResultSet); + return Boolean(this.emptyResultOnNotFound || this.operationHandle.hasResultSet); } public async fetchChunk({ @@ -377,6 +391,30 @@ export default class ThriftOperationBackend implements IOperationBackend { } private processOperationStatusResponse(response: TGetOperationStatusResp) { + const isError = + response.operationState === TOperationState.ERROR_STATE || + (response.operationState == null && response.status.statusCode === TStatusCode.ERROR_STATUS); + if ( + this.emptyResultOnNotFound && + isError && + response.status.statusCode !== TStatusCode.INVALID_HANDLE_STATUS && + (isObjectNotFound(response) || isObjectNotFound(response.status)) + ) { + this.metadata = synthesizeThriftResultSetMetadata({ + schema: this.emptyResultOnNotFound, + resultFormat: ResultFormat.ColumnBased, + isStagingOperation: false, + }); + this.resultHandler = new ResultSlicer(this.context, { + fetchNext: async () => [], + hasMore: async () => false, + }); + response = { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + operationState: TOperationState.FINISHED_STATE, + hasResultSet: true, + }; + } Status.assert(response.status); this.state = response.operationState ?? this.state; diff --git a/lib/thrift-backend/ThriftSessionBackend.ts b/lib/thrift-backend/ThriftSessionBackend.ts index 1c63b356..3c2ef266 100644 --- a/lib/thrift-backend/ThriftSessionBackend.ts +++ b/lib/thrift-backend/ThriftSessionBackend.ts @@ -9,6 +9,8 @@ import { TSparkParameter, TProtocolVersion, TExecuteStatementReq, + TStatusCode, + TOperationState, } from '../../thrift/TCLIService_types'; import ISessionBackend from '../contracts/ISessionBackend'; import IOperationBackend from '../contracts/IOperationBackend'; @@ -33,6 +35,9 @@ import { DBSQLParameter, DBSQLParameterValue } from '../DBSQLParameter'; import { LogLevel } from '../contracts/IDBSQLLogger'; import ThriftOperationBackend from './ThriftOperationBackend'; import escapeCatalogName from '../utils/escapeCatalogName'; +import EmptyOperationBackend from './EmptyOperationBackend'; +import isObjectNotFound from './metadataErrors'; +import primaryKeysSchema from './primaryKeysSchema'; interface OperationResponseShape { status: TStatus; @@ -313,7 +318,24 @@ export default class ThriftSessionBackend implements ISessionBackend { runAsync: this.getRunAsyncForMetadataOperations(), ...getDirectResultsOptions(request.maxRows, this.context.getConfig()), }); - return this.createOperationBackend(response); + const notFound = response.status.statusCode === TStatusCode.ERROR_STATUS && isObjectNotFound(response.status); + if (notFound && !response.operationHandle) { + return new EmptyOperationBackend(primaryKeysSchema); + } + if (!notFound) { + Status.assert(response.status); + } + return new ThriftOperationBackend({ + handle: definedOrError(response.operationHandle), + context: this.context, + directResults: notFound + ? { + ...response.directResults, + operationStatus: { status: response.status, operationState: TOperationState.ERROR_STATE }, + } + : response.directResults, + emptyResultOnNotFound: primaryKeysSchema, + }); } public async getCrossReference(request: CrossReferenceRequest): Promise { diff --git a/lib/thrift-backend/metadataErrors.ts b/lib/thrift-backend/metadataErrors.ts new file mode 100644 index 00000000..288294a6 --- /dev/null +++ b/lib/thrift-backend/metadataErrors.ts @@ -0,0 +1,17 @@ +interface MetadataError { + sqlState?: string | null; + errorMessage?: string | null; + displayMessage?: string | null; +} + +export default function isObjectNotFound(error: MetadataError): boolean { + // Match Databricks JDBC's metadata not-found classification. + return ( + error.sqlState === '42704' || + [error.errorMessage, error.displayMessage].some((message) => + ['NO_SUCH_CATALOG_EXCEPTION', 'SCHEMA_NOT_FOUND', 'TABLE_OR_VIEW_NOT_FOUND', 'INVALID_PARAMETER_VALUE'].some( + (errorClass) => message?.includes(errorClass), + ), + ) + ); +} diff --git a/lib/thrift-backend/primaryKeysSchema.ts b/lib/thrift-backend/primaryKeysSchema.ts new file mode 100644 index 00000000..45defb63 --- /dev/null +++ b/lib/thrift-backend/primaryKeysSchema.ts @@ -0,0 +1,14 @@ +import { TTableSchema, TTypeId } from '../../thrift/TCLIService_types'; + +const primaryKeysSchema: TTableSchema = { + columns: ['TABLE_CAT', 'TABLE_SCHEM', 'TABLE_NAME', 'COLUMN_NAME', 'KEY_SEQ', 'PK_NAME'].map((columnName, index) => ({ + columnName, + typeDesc: { + types: [{ primitiveEntry: { type: columnName === 'KEY_SEQ' ? TTypeId.INT_TYPE : TTypeId.STRING_TYPE } }], + }, + position: index + 1, + comment: '', + })), +}; + +export default primaryKeysSchema; diff --git a/tests/unit/thrift-backend/ThriftPrimaryKeysMetadata.test.ts b/tests/unit/thrift-backend/ThriftPrimaryKeysMetadata.test.ts new file mode 100644 index 00000000..ac14f366 --- /dev/null +++ b/tests/unit/thrift-backend/ThriftPrimaryKeysMetadata.test.ts @@ -0,0 +1,325 @@ +import assert from 'node:assert/strict'; +import { expect } from 'chai'; +import Int64 from 'node-int64'; +import sinon from 'sinon'; +import { + TGetOperationStatusResp, + TOperationState, + TOperationType, + TSparkRowSetType, + TStatusCode, + TTypeId, +} from '../../../thrift/TCLIService_types'; +import IOperation from '../../../lib/contracts/IOperation'; +import DBSQLOperation from '../../../lib/DBSQLOperation'; +import { PrimaryKeysRequest } from '../../../lib/contracts/IDBSQLSession'; +import OperationStateError, { OperationStateErrorCode } from '../../../lib/errors/OperationStateError'; +import StatusError from '../../../lib/errors/StatusError'; +import ClientContextStub from '../.stubs/ClientContextStub'; +import { createSessionForTest } from '../.stubs/createSessionForTest'; + +const request = { schemaName: 'missing_schema', tableName: 'table' }; +const expectedColumns = [ + ['TABLE_CAT', TTypeId.STRING_TYPE], + ['TABLE_SCHEM', TTypeId.STRING_TYPE], + ['TABLE_NAME', TTypeId.STRING_TYPE], + ['COLUMN_NAME', TTypeId.STRING_TYPE], + ['KEY_SEQ', TTypeId.INT_TYPE], + ['PK_NAME', TTypeId.STRING_TYPE], +]; + +function createContext(): ClientContextStub { + const context = new ClientContextStub(); + context.driver.getPrimaryKeysResp.operationHandle!.hasResultSet = true; + context.driver.getPrimaryKeysResp.operationHandle!.operationType = TOperationType.GET_FUNCTIONS; + return context; +} + +function createSession(context: ClientContextStub) { + return createSessionForTest({ + handle: { sessionId: { guid: Buffer.alloc(16), secret: Buffer.alloc(16) } }, + context, + }); +} + +function missingSchemaStatus(): TGetOperationStatusResp { + return { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + operationState: TOperationState.ERROR_STATE, + sqlState: '42704', + errorMessage: '[SCHEMA_NOT_FOUND] The schema cannot be found', + hasResultSet: false, + }; +} + +async function expectEmptyPrimaryKeys(operation: IOperation): Promise { + const schema = await operation.getSchema(); + expect( + schema?.columns.map((column) => [column.columnName, column.typeDesc.types[0].primitiveEntry?.type]), + ).to.deep.equal(expectedColumns); + expect(schema?.columns.map((column) => column.position)).to.deep.equal([1, 2, 3, 4, 5, 6]); + await operation.finished(); + expect(await operation.fetchAll()).to.deep.equal([]); + expect(await operation.hasMoreRows()).to.be.false; + const metadata = await (operation as DBSQLOperation).getMetadata(); + expect(metadata.status.statusCode).to.equal(TStatusCode.SUCCESS_STATUS); + expect(metadata.resultFormat).to.equal(TSparkRowSetType.COLUMN_BASED_SET); + const status = await operation.status(); + expect(status.operationState).to.equal(TOperationState.FINISHED_STATE); + expect(status.status.statusCode).to.equal(TStatusCode.SUCCESS_STATUS); + expect(status.hasResultSet).to.be.true; + expect(status.sqlState).to.be.undefined; + expect(status.errorMessage).to.be.undefined; +} + +describe('Thrift getPrimaryKeys missing objects', () => { + afterEach(() => sinon.restore()); + + const notFoundErrors = [ + { sqlState: '42704', errorMessage: 'Object is missing' }, + { sqlState: 'HY000', errorMessage: '[NO_SUCH_CATALOG_EXCEPTION] Missing catalog' }, + { sqlState: 'HY000', errorMessage: '[SCHEMA_NOT_FOUND] Missing schema' }, + { sqlState: '42P01', errorMessage: '[TABLE_OR_VIEW_NOT_FOUND] Missing table' }, + { sqlState: '22023', errorMessage: '[INVALID_PARAMETER_VALUE] Invalid catalog' }, + ]; + + for (const error of notFoundErrors) { + for (const delivery of ['initial', 'direct', 'poll']) { + it(`returns a typed empty result for ${error.errorMessage} via ${delivery}`, async () => { + const context = createContext(); + const driver = sinon.spy(context.driver); + const failure = { ...missingSchemaStatus(), ...error }; + if (delivery === 'initial') { + context.driver.getPrimaryKeysResp = { + status: { statusCode: TStatusCode.ERROR_STATUS, ...error }, + }; + } else if (delivery === 'direct') { + context.driver.getPrimaryKeysResp.directResults = { + operationStatus: failure, + resultSetMetadata: context.driver.getResultSetMetadataResp, + resultSet: context.driver.fetchResultsResp, + }; + } else { + context.driver.getOperationStatusResp = failure; + } + const operation = await createSession(context).getPrimaryKeys(request); + + await expectEmptyPrimaryKeys(operation); + expect(driver.getResultSetMetadata.called).to.be.false; + expect(driver.fetchResults.called).to.be.false; + expect(driver.getOperationStatus.callCount).to.equal(delivery === 'poll' ? 1 : 0); + await operation.close(); + await operation.close(); + expect(driver.closeOperation.callCount).to.equal(delivery === 'initial' ? 0 : 1); + }); + } + } + + it('closes a real handle returned with an initial error', async () => { + const context = createContext(); + context.driver.getPrimaryKeysResp.status = { + statusCode: TStatusCode.ERROR_STATUS, + sqlState: '42704', + }; + const closeOperation = sinon.spy(context.driver, 'closeOperation'); + const operation = await createSession(context).getPrimaryKeys(request); + + await expectEmptyPrimaryKeys(operation); + await operation.close(); + expect( + closeOperation.calledOnceWithExactly({ + operationHandle: context.driver.getPrimaryKeysResp.operationHandle!, + }), + ).to.be.true; + }); + + it('uses a prefetched close response after recovering direct results', async () => { + const context = createContext(); + context.driver.getPrimaryKeysResp.directResults = { + operationStatus: missingSchemaStatus(), + closeOperation: context.driver.closeOperationResp, + }; + const closeOperation = sinon.spy(context.driver, 'closeOperation'); + const operation = await createSession(context).getPrimaryKeys(request); + + await expectEmptyPrimaryKeys(operation); + expect((await operation.close()).isSuccess).to.be.true; + expect(closeOperation.called).to.be.false; + }); + + it('polls asynchronously and reports recovered success to callbacks', async () => { + const context = createContext(); + const poll = sinon.stub(context.driver, 'getOperationStatus').resolves(missingSchemaStatus()); + poll.onFirstCall().resolves({ + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + operationState: TOperationState.RUNNING_STATE, + }); + const operation = await createSession(context).getPrimaryKeys(request); + expect(poll.called).to.be.false; + const callback = sinon.spy(); + + await operation.finished({ callback }); + expect(callback.args.map(([status]) => status.operationState)).to.deep.equal([ + TOperationState.RUNNING_STATE, + TOperationState.FINISHED_STATE, + ]); + await expectEmptyPrimaryKeys(operation); + }); + + it('recovers a not-found RPC status without an operation state', async () => { + const context = createContext(); + context.driver.getOperationStatusResp = { + status: { statusCode: TStatusCode.ERROR_STATUS, errorMessage: '[SCHEMA_NOT_FOUND] Missing schema' }, + }; + const operation = await createSession(context).getPrimaryKeys(request); + + await expectEmptyPrimaryKeys(operation); + }); + + it('waits for metadata even when the initial handle has no result set', async () => { + const context = createContext(); + context.driver.getPrimaryKeysResp.operationHandle!.hasResultSet = false; + context.driver.getOperationStatusResp = missingSchemaStatus(); + const operation = await createSession(context).getPrimaryKeys(request); + + await expectEmptyPrimaryKeys(operation); + }); + + it('does not recover invalid-handle RPC errors', async () => { + const context = createContext(); + context.driver.getOperationStatusResp = { + ...missingSchemaStatus(), + status: { statusCode: TStatusCode.INVALID_HANDLE_STATUS, errorMessage: 'Invalid handle' }, + }; + const operation = await createSession(context).getPrimaryKeys(request); + + await assert.rejects(operation.finished(), (error) => { + expect(error).to.be.instanceOf(StatusError); + expect(error).to.have.property('message', 'Invalid handle'); + return true; + }); + }); + + it('recognizes a not-found error in displayMessage', async () => { + const context = createContext(); + context.driver.getOperationStatusResp = { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + operationState: TOperationState.ERROR_STATE, + displayMessage: '[SCHEMA_NOT_FOUND] Missing schema', + }; + const operation = await createSession(context).getPrimaryKeys(request); + + await expectEmptyPrimaryKeys(operation); + }); + + for (const sqlState of ['42501', '28000', 'HYT00', '22023']) { + it(`preserves unrelated failures with SQLSTATE ${sqlState}`, async () => { + const context = createContext(); + const errorStatus = { statusCode: TStatusCode.ERROR_STATUS, sqlState, errorMessage: 'Unrelated failure' }; + context.driver.getPrimaryKeysResp.status = errorStatus; + const session = createSession(context); + await assert.rejects(session.getPrimaryKeys(request), (error) => { + expect(error).to.be.instanceOf(StatusError); + expect(error).to.have.property('sqlState', sqlState); + return true; + }); + context.driver.getPrimaryKeysResp.status = { statusCode: TStatusCode.SUCCESS_STATUS }; + const failure = { ...missingSchemaStatus(), sqlState, errorMessage: 'Unrelated failure' }; + context.driver.getOperationStatusResp = failure; + const operation = await session.getPrimaryKeys(request); + await assert.rejects(operation.fetchAll(), (error) => { + expect(error).to.be.instanceOf(OperationStateError); + expect(error).to.have.property('response', failure); + return true; + }); + }); + } + + for (const [operationState, errorCode] of [ + [TOperationState.CANCELED_STATE, OperationStateErrorCode.Canceled], + [TOperationState.CLOSED_STATE, OperationStateErrorCode.Closed], + [TOperationState.TIMEDOUT_STATE, OperationStateErrorCode.Timeout], + ] as const) { + it(`preserves ${TOperationState[operationState]} even with a not-found message`, async () => { + const context = createContext(); + context.driver.getOperationStatusResp = { ...missingSchemaStatus(), operationState }; + const operation = await createSession(context).getPrimaryKeys(request); + + await assert.rejects(operation.finished(), (error) => { + expect(error).to.be.instanceOf(OperationStateError); + expect(error).to.have.property('errorCode', errorCode); + return true; + }); + }); + } + + for (const method of ['executeStatement', 'getTables']) { + it(`continues throwing missing-object errors for ${method}`, async () => { + const context = createContext(); + context.driver.getOperationStatusResp = missingSchemaStatus(); + const session = createSession(context); + const operation = await (method === 'executeStatement' + ? session.executeStatement('SELECT * FROM missing_schema.table') + : session.getTables(request)); + + await assert.rejects(operation.finished(), OperationStateError); + }); + } + + for (const catalogName of [undefined, null, '', 'catalog']) { + it(`preserves real primary keys with catalog=${JSON.stringify(catalogName)}`, async () => { + const context = createContext(); + context.driver.getResultSetMetadataResp.schema = { + columns: expectedColumns.map(([name, type], index) => ({ + columnName: name === 'KEY_SEQ' ? 'KEQ_SEQ' : String(name), + typeDesc: { types: [{ primitiveEntry: { type: type as TTypeId } }] }, + position: index + 1, + })), + }; + const session = createSession(context); + const operation = await session.getPrimaryKeys({ ...request, catalogName } as PrimaryKeysRequest); + + // The successful server response is authoritative, including its wire metadata. + const metadata = context.driver.getResultSetMetadataResp; + const fetch = sinon.stub(context.driver, 'fetchResults').resolves({ status: metadata.status }); + fetch.onFirstCall().resolves({ + status: metadata.status, + results: { + startRowOffset: new Int64(0), + rows: [], + columns: [ + { stringVal: { values: ['catalog'], nulls: Buffer.alloc(0) } }, + { stringVal: { values: ['schema'], nulls: Buffer.alloc(0) } }, + { stringVal: { values: ['table'], nulls: Buffer.alloc(0) } }, + { stringVal: { values: ['id'], nulls: Buffer.alloc(0) } }, + { i32Val: { values: [1], nulls: Buffer.alloc(0) } }, + { stringVal: { values: ['pk_table'], nulls: Buffer.alloc(0) } }, + ], + }, + }); + + expect(await operation.getSchema()).to.equal(metadata.schema); + expect(await operation.fetchAll()).to.deep.equal([ + { + TABLE_CAT: 'catalog', + TABLE_SCHEM: 'schema', + TABLE_NAME: 'table', + COLUMN_NAME: 'id', + KEQ_SEQ: 1, + PK_NAME: 'pk_table', + }, + ]); + expect(await (operation as DBSQLOperation).getMetadata()).to.equal(metadata); + expect(context.driver.getPrimaryKeysReq?.catalogName).to.equal(catalogName); + }); + } + + it('preserves a successful empty server response', async () => { + const context = createContext(); + context.driver.fetchResultsResp.results = { startRowOffset: new Int64(0), rows: [], columns: [] }; + const operation = await createSession(context).getPrimaryKeys(request); + + expect(await operation.fetchAll()).to.deep.equal([]); + expect(await operation.getSchema()).to.equal(context.driver.getResultSetMetadataResp.schema); + }); +});