Return empty primary-key metadata for missing objects in Thrift

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<Array<object>> {
+    return [];
+  }
+
+  public async hasMore(): Promise<boolean> {
+    return false;
+  }
+
+  public async waitUntilReady(): Promise<void> {
+    return undefined;
+  }
+
+  public async status(): Promise<OperationStatus> {
+    return { state: OperationState.Succeeded, hasResultSet: true };
+  }
+
+  public async getResultMetadata(): Promise<ResultMetadata> {
+    return { schema: this.schema, resultFormat: ResultFormat.ColumnBased, isStagingOperation: false };
+  }
+
+  public async cancel(): Promise<Status> {
+    return Status.success();
+  }
+
+  public async close(): Promise<Status> {
+    return Status.success();
+  }
+}

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<void> {
@@ -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,18 +101,25 @@ export default class ThriftOperationBackend implements IOperationBackend {
 
   private resultHandler?: ResultSlicer<any>;
 
-  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);
 
     if (directResults?.operationStatus) {
       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;

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<IOperationBackend> {

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),
+      ),
+    )
+  );
+}

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;