Skip to content

Commit 3a315a9

Browse files
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 50f436e commit 3a315a9

8 files changed

Lines changed: 1138 additions & 302 deletions

File tree

lib/DBSQLClient.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -232,14 +232,14 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
232232
// doesn't ship in the public `.d.ts`. Mirrors Python's `kwargs.get("use_sea")`
233233
// pattern (see databricks-sql-python/src/databricks/sql/session.py).
234234
const internalOptions = options as ConnectionOptions & InternalConnectionOptions;
235-
236235
if (internalOptions.useSEA) {
237236
// The SEA backend authenticates inside the native binding; the
238237
// Thrift auth/connection providers are never read on this path, so
239238
// we don't build them (avoids validating the PAT twice and
240239
// constructing a throwaway OAuth provider for an OAuth+useSEA call).
240+
// The backend reads logger/config off the IClientContext it's given.
241241
this.logger.log(LogLevel.info, 'Connecting via the SEA (native) backend');
242-
this.backend = new SeaBackend(undefined, this.logger);
242+
this.backend = new SeaBackend({ context: this });
243243
} else {
244244
this.authProvider = this.createAuthProvider(options, authProvider);
245245
this.connectionProvider = this.createConnectionProvider(options);

lib/sea/SeaBackend.ts

Lines changed: 88 additions & 148 deletions
Original file line numberDiff line numberDiff line change
@@ -14,183 +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 IDBSQLLogger, { LogLevel } from '../contracts/IDBSQLLogger';
35-
import { getSeaNative, SeaNativeBinding } from './SeaNativeLoader';
20+
import {
21+
getSeaNative,
22+
SeaNativeBinding,
23+
SeaNativeConnection,
24+
} from './SeaNativeLoader';
25+
import { mapKernelErrorToJsError, KernelErrorShape } from './SeaErrorMapping';
3626
import { buildSeaConnectionOptions, SeaNativeConnectionOptions } from './SeaAuth';
37-
38-
const NOT_IMPLEMENTED_SESSION =
39-
'SEA session backend: method not implemented in sea-auth (M0); lands in sea-execution/sea-operation.';
40-
41-
/**
42-
* Opaque handle to the napi binding's `Connection` class. The exact
43-
* shape lives in `native/sea/index.d.ts` (auto-generated). We type it as
44-
* a structural minimum here so the loader's pass-through typing doesn't
45-
* leak into every call site.
46-
*/
47-
interface NativeConnection {
48-
/** Server-issued session id (kernel `Connection.sessionId` getter). */
49-
readonly sessionId: string;
50-
close(): Promise<void>;
51-
}
27+
import SeaSessionBackend from './SeaSessionBackend';
5228

5329
/**
54-
* Minimal `ISessionBackend` that wraps the napi-binding's `Connection`.
55-
*
56-
* For M0 (sea-auth) only `id` and `close()` are functional — they're the
57-
* subset required to round-trip a connect-open-close cycle. Every other
58-
* method throws a clear "not implemented in M0" `HiveDriverError`.
59-
*
60-
* `id` is the server-issued session id read straight off the kernel
61-
* `Connection` (its `sessionId` getter, readable even after close()), so
62-
* the value logged by `DBSQLSession` correlates with kernel / server logs
63-
* rather than being a process-local synthetic counter.
30+
* Sentinel string the napi binding uses on `Error.reason` JSON envelopes.
31+
* Keep in sync with `native/sea/src/error.rs` (`SENTINEL`).
6432
*/
65-
export class SeaSessionBackend implements ISessionBackend {
66-
public readonly id: string;
67-
68-
private readonly connection: NativeConnection;
69-
70-
private readonly logger?: IDBSQLLogger;
71-
72-
constructor(connection: NativeConnection, logger?: IDBSQLLogger) {
73-
this.connection = connection;
74-
this.logger = logger;
75-
this.id = connection.sessionId;
76-
}
77-
78-
/* eslint-disable @typescript-eslint/no-unused-vars */
79-
public async getInfo(_infoType: number): Promise<InfoValue> {
80-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
81-
}
82-
83-
public async executeStatement(
84-
_statement: string,
85-
_options: ExecuteStatementOptions,
86-
): Promise<IOperationBackend> {
87-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
88-
}
89-
90-
public async getTypeInfo(_request: TypeInfoRequest): Promise<IOperationBackend> {
91-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
92-
}
93-
94-
public async getCatalogs(_request: CatalogsRequest): Promise<IOperationBackend> {
95-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
96-
}
97-
98-
public async getSchemas(_request: SchemasRequest): Promise<IOperationBackend> {
99-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
100-
}
101-
102-
public async getTables(_request: TablesRequest): Promise<IOperationBackend> {
103-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
104-
}
105-
106-
public async getTableTypes(_request: TableTypesRequest): Promise<IOperationBackend> {
107-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
108-
}
109-
110-
public async getColumns(_request: ColumnsRequest): Promise<IOperationBackend> {
111-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
112-
}
113-
114-
public async getFunctions(_request: FunctionsRequest): Promise<IOperationBackend> {
115-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
116-
}
117-
118-
public async getPrimaryKeys(_request: PrimaryKeysRequest): Promise<IOperationBackend> {
119-
throw new HiveDriverError(NOT_IMPLEMENTED_SESSION);
120-
}
121-
122-
public async getCrossReference(_request: CrossReferenceRequest): Promise<IOperationBackend> {
123-
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+
}
12448
}
125-
/* eslint-enable @typescript-eslint/no-unused-vars */
49+
throw err;
50+
}
12651

127-
public async close(): Promise<Status> {
128-
this.logger?.log(LogLevel.debug, `SEA session closing with id: ${this.id}`);
129-
await this.connection.close();
130-
return Status.success();
131-
}
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;
13260
}
13361

13462
/**
135-
* M0 SeaBackend — wires PAT auth + napi `openSession` end-to-end.
63+
* SEA-backed implementation of `IBackend`.
64+
*
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()`.
13670
*
137-
* Connect is a no-op at this layer (the napi binding has no notion of a
138-
* standalone "connect"; a session is opened directly). We capture the
139-
* validated PAT options and hand them to `openSession()` on demand.
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).
14075
*
141-
* Subsequent milestones (`sea-execution`, `sea-operation`) replace the
142-
* stubbed `ISessionBackend` / `IOperationBackend` methods with real
143-
* napi-binding calls.
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.
14481
*/
14582
export default class SeaBackend implements IBackend {
146-
private nativeOptions?: SeaNativeConnectionOptions;
83+
private readonly context: IClientContext;
14784

148-
private readonly injectedNative?: SeaNativeBinding;
85+
private readonly binding: SeaNativeBinding;
14986

150-
private cachedNative?: SeaNativeBinding;
151-
152-
private readonly logger?: IDBSQLLogger;
153-
154-
// `native` is injectable (tests pass a fake); production leaves it
155-
// undefined and the binding is resolved lazily on first use so that
156-
// constructing a SeaBackend never throws on a platform without the
157-
// optional `.node` — the clearer auth/option validation in connect()
158-
// runs first.
159-
constructor(native?: SeaNativeBinding, logger?: IDBSQLLogger) {
160-
this.injectedNative = native;
161-
this.logger = logger;
162-
}
87+
private nativeOptions?: SeaNativeConnectionOptions;
16388

164-
private get native(): SeaNativeBinding {
165-
if (!this.cachedNative) {
166-
this.cachedNative = this.injectedNative ?? getSeaNative();
167-
}
168-
return this.cachedNative;
89+
constructor(options?: SeaBackendOptions) {
90+
this.context = options?.context as IClientContext;
91+
this.binding = options?.nativeBinding ?? getSeaNative();
16992
}
17093

17194
public async connect(options: ConnectionOptions): Promise<void> {
172-
// Validate PAT auth + capture the napi-binding option shape. Any
173-
// non-PAT mode (or a missing token) throws here, before we ever touch
174-
// the native binding. NOTE: unlike Thrift, this performs no network
175-
// round-trip — the session is opened lazily in openSession(), so a
176-
// resolved connect() does not by itself prove the endpoint is
177-
// reachable or the credential is valid.
95+
// Validate PAT auth + capture the napi-binding option shape.
96+
// Any non-PAT mode (or a missing/empty token) throws here, before
97+
// we ever touch the native binding.
17898
this.nativeOptions = buildSeaConnectionOptions(options);
17999
}
180100

181-
// eslint-disable-next-line @typescript-eslint/no-unused-vars
182-
public async openSession(_request: OpenSessionRequest): Promise<ISessionBackend> {
101+
public async openSession(request: OpenSessionRequest): Promise<ISessionBackend> {
183102
if (!this.nativeOptions) {
184-
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);
185111
}
186-
const connection = (await this.native.openSession(this.nativeOptions)) as NativeConnection;
187-
const session = new SeaSessionBackend(connection, this.logger);
188-
this.logger?.log(LogLevel.info, `SEA session opened with id: ${session.id}`);
189-
return session;
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+
});
190129
}
191130

192131
public async close(): Promise<void> {
193-
// 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.
194134
this.nativeOptions = undefined;
195135
}
196136
}

lib/sea/SeaNativeLoader.ts

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,23 @@ export type SeaArrowSchema = NativeArrowSchema;
4646
export type SeaConnection = NativeConnection;
4747
export type SeaStatement = NativeStatement;
4848

49+
// Back-compat aliases for the downstream SEA stack branches that landed
50+
// against the pre-rename loader. The merged kernel (@databricks/sql-kernel)
51+
// moved per-statement catalog/schema/sessionConfig to session-level
52+
// `openSession`, so `ExecuteOptions` no longer exists on the binding;
53+
// `SeaExecuteOptions` is kept as a deprecated shim describing the old
54+
// per-statement shape so the stack keeps compiling. Per-statement options
55+
// are now applied at session creation — see native/sea/README.md.
56+
export type SeaNativeConnection = NativeConnection;
57+
export type SeaNativeStatement = NativeStatement;
58+
export type SeaNativeConnectionOptions = NativeConnectionOptions;
59+
/** @deprecated per-statement options moved to session-level `openSession`. */
60+
export interface SeaExecuteOptions {
61+
initialCatalog?: string;
62+
initialSchema?: string;
63+
sessionConfig?: Record<string, string>;
64+
}
65+
4966
/**
5067
* The full native binding surface, derived from the generated module
5168
* so it can never drift from the `.d.ts` contract: when the kernel

0 commit comments

Comments
 (0)