Skip to content

Commit 0a1c02d

Browse files
msrathore-dbmsrathore-db
committed
sea-execution: executeStatement + openSession via SeaSessionBackend
SeaSessionBackend wraps the napi Connection handle. executeStatement passes through to napi.executeStatement and returns an IOperationBackend (SeaOperationBackend in sea-results feature). Session config + initialCatalog/initialSchema flow to napi openSession. M0 stops at executeStatement; metadata methods + per-stmt overrides defer to M1. No new dependencies. Reuses existing ConnectionOptions / Session config shapes. Co-authored-by: Isaac Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>
1 parent f2c1aff commit 0a1c02d

8 files changed

Lines changed: 1128 additions & 263 deletions

File tree

‎lib/DBSQLClient.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,7 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
243243
// pattern (see databricks-sql-python/src/databricks/sql/session.py).
244244
const internalOptions = options as ConnectionOptions & InternalConnectionOptions;
245245
this.backend = internalOptions.useSEA
246-
? new SeaBackend()
246+
? new SeaBackend({ context: this })
247247
: new ThriftBackend({
248248
context: this,
249249
onConnectionEvent: (event, payload) => this.forwardConnectionEvent(event, payload),

‎lib/sea/SeaBackend.ts‎

Lines changed: 88 additions & 121 deletions
Original file line numberDiff line numberDiff line change
@@ -14,156 +14,123 @@
1414

1515
import IBackend from '../contracts/IBackend';
1616
import ISessionBackend from '../contracts/ISessionBackend';
17-
import IOperationBackend from '../contracts/IOperationBackend';
17+
import IClientContext from '../contracts/IClientContext';
1818
import { ConnectionOptions, OpenSessionRequest } from '../contracts/IDBSQLClient';
19-
import {
20-
ExecuteStatementOptions,
21-
TypeInfoRequest,
22-
CatalogsRequest,
23-
SchemasRequest,
24-
TablesRequest,
25-
TableTypesRequest,
26-
ColumnsRequest,
27-
FunctionsRequest,
28-
PrimaryKeysRequest,
29-
CrossReferenceRequest,
30-
} from '../contracts/IDBSQLSession';
31-
import Status from '../dto/Status';
32-
import InfoValue from '../dto/InfoValue';
3319
import HiveDriverError from '../errors/HiveDriverError';
34-
import { getSeaNative, SeaNativeBinding } from './SeaNativeLoader';
20+
import {
21+
getSeaNative,
22+
SeaNativeBinding,
23+
SeaNativeConnection,
24+
} from './SeaNativeLoader';
25+
import { mapKernelErrorToJsError, KernelErrorShape } from './SeaErrorMapping';
3526
import { buildSeaConnectionOptions, SeaNativeConnectionOptions } from './SeaAuth';
36-
37-
const NOT_IMPLEMENTED_SESSION =
38-
'SEA session backend: method not implemented in sea-auth (M0); lands in sea-execution/sea-operation.';
39-
40-
/**
41-
* Opaque handle to the napi binding's `Connection` class. The exact
42-
* shape lives in `native/sea/index.d.ts` (auto-generated). We type it as
43-
* a structural minimum here so the loader's pass-through typing doesn't
44-
* leak into every call site.
45-
*/
46-
interface NativeConnection {
47-
close(): Promise<void>;
48-
}
27+
import SeaSessionBackend from './SeaSessionBackend';
4928

5029
/**
51-
* Minimal `ISessionBackend` that wraps the napi-binding's `Connection`.
52-
*
53-
* For M0 (sea-auth) only `id` and `close()` are functional — they're the
54-
* subset required to round-trip a connect-open-close cycle. Every other
55-
* method throws a clear "not implemented in M0" `HiveDriverError`.
56-
*
57-
* The `id` field is currently a synthetic counter-based string; the kernel
58-
* exposes a real session-id through a follow-on getter that
59-
* `sea-execution` will wire through.
30+
* Sentinel string the napi binding uses on `Error.reason` JSON envelopes.
31+
* Keep in sync with `native/sea/src/error.rs` (`SENTINEL`).
6032
*/
61-
export class SeaSessionBackend implements ISessionBackend {
62-
private static seq = 0;
63-
64-
public readonly id: string;
65-
66-
private readonly connection: NativeConnection;
67-
68-
constructor(connection: NativeConnection) {
69-
this.connection = connection;
70-
SeaSessionBackend.seq += 1;
71-
this.id = `sea-session-${SeaSessionBackend.seq}`;
72-
}
73-
74-
/* eslint-disable @typescript-eslint/no-unused-vars */
75-
public async getInfo(_infoType: number): Promise<InfoValue> {
76-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
77-
}
78-
79-
public async executeStatement(
80-
_statement: string,
81-
_options: ExecuteStatementOptions,
82-
): Promise<IOperationBackend> {
83-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
84-
}
85-
86-
public async getTypeInfo(_request: TypeInfoRequest): Promise<IOperationBackend> {
87-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
88-
}
89-
90-
public async getCatalogs(_request: CatalogsRequest): Promise<IOperationBackend> {
91-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
92-
}
93-
94-
public async getSchemas(_request: SchemasRequest): Promise<IOperationBackend> {
95-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
96-
}
97-
98-
public async getTables(_request: TablesRequest): Promise<IOperationBackend> {
99-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
100-
}
101-
102-
public async getTableTypes(_request: TableTypesRequest): Promise<IOperationBackend> {
103-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
104-
}
105-
106-
public async getColumns(_request: ColumnsRequest): Promise<IOperationBackend> {
107-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
108-
}
109-
110-
public async getFunctions(_request: FunctionsRequest): Promise<IOperationBackend> {
111-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
112-
}
113-
114-
public async getPrimaryKeys(_request: PrimaryKeysRequest): Promise<IOperationBackend> {
115-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
116-
}
117-
118-
public async getCrossReference(_request: CrossReferenceRequest): Promise<IOperationBackend> {
119-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
33+
const KERNEL_ERROR_SENTINEL = '__databricks_error__:';
34+
35+
function rethrowKernelError(err: unknown): never {
36+
if (err && typeof err === 'object' && 'message' in err) {
37+
const reason = (err as { reason?: unknown }).reason;
38+
if (typeof reason === 'string' && reason.startsWith(KERNEL_ERROR_SENTINEL)) {
39+
try {
40+
const payload = JSON.parse(reason.slice(KERNEL_ERROR_SENTINEL.length)) as KernelErrorShape;
41+
throw mapKernelErrorToJsError(payload);
42+
} catch (parseErr) {
43+
if (parseErr !== err) {
44+
throw parseErr;
45+
}
46+
}
47+
}
12048
}
121-
/* eslint-enable @typescript-eslint/no-unused-vars */
49+
throw err;
50+
}
12251

123-
public async close(): Promise<Status> {
124-
await this.connection.close();
125-
return Status.success();
126-
}
52+
export interface SeaBackendOptions {
53+
context: IClientContext;
54+
/**
55+
* Optional injection seam for unit tests. When provided, replaces the
56+
* default `getSeaNative()` call so tests can swap in a mock napi
57+
* binding without loading the `.node` artifact.
58+
*/
59+
nativeBinding?: SeaNativeBinding;
12760
}
12861

12962
/**
130-
* M0 SeaBackend — wires PAT auth + napi `openSession` end-to-end.
63+
* SEA-backed implementation of `IBackend`.
13164
*
132-
* Connect is a no-op at this layer (the napi binding has no notion of a
133-
* standalone "connect"; a session is opened directly). We capture the
134-
* validated PAT options and hand them to `openSession()` on demand.
65+
* **M0 dispatch model:** the napi binding's `openSession()` already
66+
* builds a kernel `Session` from PAT + hostname + httpPath, so there is
67+
* no "connect" round-trip before `openSession` — `connect()` only
68+
* captures the `ConnectionOptions` and validates that PAT auth is in
69+
* use. The actual session open happens inside `openSession()`.
13570
*
136-
* Subsequent milestones (`sea-execution`, `sea-operation`) replace the
137-
* stubbed `ISessionBackend` / `IOperationBackend` methods with real
138-
* napi-binding calls.
71+
* **Auth validation:** delegates to `buildSeaConnectionOptions` from
72+
* `SeaAuth`, which mirrors the existing DBSQLClient PAT validation
73+
* pattern (slash-prepended httpPath, AuthenticationError on missing
74+
* token, HiveDriverError on non-PAT authType naming M1 modes).
75+
*
76+
* **Why we don't use IClientContext's connectionProvider here:** that
77+
* provider is the Thrift HTTP transport. The kernel owns its own
78+
* reqwest+rustls stack inside the native binding, so there is no
79+
* NodeJS-level connection state to manage on the SEA path. The
80+
* `IClientContext` is still useful for logger + config access.
13981
*/
14082
export default class SeaBackend implements IBackend {
141-
private nativeOptions?: SeaNativeConnectionOptions;
83+
private readonly context: IClientContext;
14284

143-
private readonly native: SeaNativeBinding;
85+
private readonly binding: SeaNativeBinding;
14486

145-
constructor(native: SeaNativeBinding = getSeaNative()) {
146-
this.native = native;
87+
private nativeOptions?: SeaNativeConnectionOptions;
88+
89+
constructor(options?: SeaBackendOptions) {
90+
this.context = options?.context as IClientContext;
91+
this.binding = options?.nativeBinding ?? getSeaNative();
14792
}
14893

14994
public async connect(options: ConnectionOptions): Promise<void> {
15095
// Validate PAT auth + capture the napi-binding option shape.
151-
// Any non-PAT mode (or a missing token) throws here, before we ever
152-
// touch the native binding.
96+
// Any non-PAT mode (or a missing/empty token) throws here, before
97+
// we ever touch the native binding.
15398
this.nativeOptions = buildSeaConnectionOptions(options);
15499
}
155100

156-
// eslint-disable-next-line @typescript-eslint/no-unused-vars
157-
public async openSession(_request: OpenSessionRequest): Promise<ISessionBackend> {
101+
public async openSession(request: OpenSessionRequest): Promise<ISessionBackend> {
158102
if (!this.nativeOptions) {
159-
throw new HiveDriverError('SeaBackend: connect() must be called before openSession().');
103+
throw new HiveDriverError('SeaBackend: not connected. Call connect() first.');
104+
}
105+
106+
let nativeConnection: SeaNativeConnection;
107+
try {
108+
nativeConnection = (await this.binding.openSession(this.nativeOptions)) as SeaNativeConnection;
109+
} catch (err) {
110+
rethrowKernelError(err);
160111
}
161-
const connection = (await this.native.openSession(this.nativeOptions)) as NativeConnection;
162-
return new SeaSessionBackend(connection);
112+
113+
// Merge `request.configuration` (the existing public field for Spark
114+
// conf) with any backend-specific session config. The SEA wire
115+
// protocol applies these per-statement, but we capture them at
116+
// session-open time and forward with every executeStatement to
117+
// preserve session-config semantics.
118+
const sessionConfig = request.configuration ? { ...request.configuration } : undefined;
119+
120+
return new SeaSessionBackend({
121+
connection: nativeConnection!,
122+
context: this.context,
123+
defaults: {
124+
initialCatalog: request.initialCatalog,
125+
initialSchema: request.initialSchema,
126+
sessionConfig,
127+
},
128+
});
163129
}
164130

165131
public async close(): Promise<void> {
166-
// Connection-level resources are owned by the session wrapper. No-op here.
132+
// No backend-level resources to release — each `SeaSessionBackend`
133+
// owns its own napi `Connection` lifecycle.
167134
this.nativeOptions = undefined;
168135
}
169136
}

‎lib/sea/SeaNativeLoader.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,14 @@ export type { ConnectionOptions, ExecuteOptions, ArrowBatch, ArrowSchema };
3636
export type Connection = NativeConnection;
3737
export type Statement = NativeStatement;
3838

39+
// Back-compat aliases for consumers that landed before the path-alias
40+
// refactor renamed these. New code should import the un-prefixed names.
41+
export type SeaNativeConnection = NativeConnection;
42+
export type SeaNativeStatement = NativeStatement;
43+
export type SeaExecuteOptions = ExecuteOptions;
44+
export type SeaArrowBatch = ArrowBatch;
45+
export type SeaArrowSchema = ArrowSchema;
46+
3947
export interface SeaNativeBinding {
4048
version(): string;
4149
openSession(options: ConnectionOptions): Promise<NativeConnection>;

0 commit comments

Comments
 (0)