@databricks/appkit 0.1.5 → 0.3.0

package/dist/connectors/lakebase/client.js.map

@@ -1 +1 @@
-{"version":3,"file":"client.js","names":[],"sources":["../../../src/connectors/lakebase/client.ts"],"sourcesContent":["import { randomUUID } from \"node:crypto\";\nimport type { WorkspaceClient } from \"@databricks/sdk-experimental\";\nimport { ApiClient, Config } from \"@databricks/sdk-experimental\";\nimport pg from \"pg\";\nimport {\n  type Counter,\n  type Histogram,\n  SpanStatusCode,\n  TelemetryManager,\n  type TelemetryProvider,\n} from \"@/telemetry\";\nimport { deepMerge } from \"../../utils\";\nimport { lakebaseDefaults } from \"./defaults\";\nimport type {\n  LakebaseConfig,\n  LakebaseConnectionConfig,\n  LakebaseCredentials,\n} from \"./types\";\n\n/**\n * Enterprise-grade connector for Databricks Lakebase\n * @example Simplest - everything from env/context\n * ```typescript\n * const connector = new LakebaseConnector();\n * await connector.query('SELECT * FROM users');\n * ```\n *\n * @example With explicit connection string\n * ```typescript\n * const connector = new LakebaseConnector({\n *   connectionString: 'postgresql://...'\n * });\n * ```\n */\nexport class LakebaseConnector {\n  private readonly name: string = \"lakebase\";\n  private readonly CACHE_BUFFER_MS = 2 * 60 * 1000;\n  private readonly config: LakebaseConfig;\n  private readonly connectionConfig: LakebaseConnectionConfig;\n  private pool: pg.Pool | null = null;\n  private credentials: LakebaseCredentials | null = null;\n\n  // telemetry\n  private readonly telemetry: TelemetryProvider;\n  private readonly telemetryMetrics: {\n    queryCount: Counter;\n    queryDuration: Histogram;\n  };\n\n  constructor(userConfig?: Partial<LakebaseConfig>) {\n    this.config = deepMerge(lakebaseDefaults, userConfig);\n    this.connectionConfig = this.parseConnectionConfig();\n\n    this.telemetry = TelemetryManager.getProvider(\n      this.name,\n      this.config.telemetry,\n    );\n    this.telemetryMetrics = {\n      queryCount: this.telemetry\n        .getMeter()\n        .createCounter(\"lakebase.query.count\", {\n          description: \"Total number of queries executed\",\n          unit: \"1\",\n        }),\n      queryDuration: this.telemetry\n        .getMeter()\n        .createHistogram(\"lakebase.query.duration\", {\n          description: \"Duration of queries executed\",\n          unit: \"ms\",\n        }),\n    };\n\n    // validate configuration\n    if (this.config.maxPoolSize < 1) {\n      throw new Error(\"maxPoolSize must be at least 1\");\n    }\n  }\n\n  /**\n   * Execute a SQL query\n   *\n   * @example\n   * ```typescript\n   * const users = await connector.query('SELECT * FROM users');\n   * const user = await connector.query('SELECT * FROM users WHERE id = $1', [123]);\n   * ```\n   */\n  async query<T extends pg.QueryResultRow>(\n    sql: string,\n    params?: any[],\n    retryCount: number = 0,\n  ): Promise<pg.QueryResult<T>> {\n    const startTime = Date.now();\n\n    return this.telemetry.startActiveSpan(\n      \"lakebase.query\",\n      {\n        attributes: {\n          \"db.system\": \"lakebase\",\n          \"db.statement\": sql.substring(0, 500),\n          \"db.retry_count\": retryCount,\n        },\n      },\n      async (span) => {\n        try {\n          const pool = await this.getPool();\n          const result = await pool.query<T>(sql, params);\n          span.setAttribute(\"db.rows_affected\", result.rowCount ?? 0);\n          span.setStatus({ code: SpanStatusCode.OK });\n          return result;\n        } catch (error) {\n          // retry on auth failure\n          if (this.isAuthError(error)) {\n            span.addEvent(\"auth_error_retry\");\n            await this.rotateCredentials();\n            const newPool = await this.getPool();\n            const result = await newPool.query<T>(sql, params);\n            span.setAttribute(\"db.rows_affected\", result.rowCount ?? 0);\n            span.setStatus({ code: SpanStatusCode.OK });\n            return result;\n          }\n\n          // retry on transient errors, but only once\n          if (this.isTransientError(error) && retryCount < 1) {\n            span.addEvent(\"transient_error_retry\");\n            await new Promise((resolve) => setTimeout(resolve, 100));\n            return await this.query<T>(sql, params, retryCount + 1);\n          }\n\n          span.recordException(error as Error);\n          span.setStatus({ code: SpanStatusCode.ERROR });\n\n          throw error;\n        } finally {\n          const duration = Date.now() - startTime;\n          this.telemetryMetrics.queryCount.add(1);\n          this.telemetryMetrics.queryDuration.record(duration);\n          span.end();\n        }\n      },\n    );\n  }\n\n  /**\n   * Execute a transaction\n   *\n   * COMMIT and ROLLBACK are automatically managed by the transaction function.\n   *\n   * @param callback - Callback function to execute within the transaction context\n   * @example\n   * ```typescript\n   * await connector.transaction(async (client) => {\n   *   await client.query('INSERT INTO accounts (name) VALUES ($1)', ['Alice']);\n   *   await client.query('INSERT INTO logs (action) VALUES ($1)', ['Created Alice']);\n   * });\n   * ```\n   */\n  async transaction<T>(\n    callback: (client: pg.PoolClient) => Promise<T>,\n    retryCount: number = 0,\n  ): Promise<T> {\n    const startTime = Date.now();\n    return this.telemetry.startActiveSpan(\n      \"lakebase.transaction\",\n      {\n        attributes: {\n          \"db.system\": \"lakebase\",\n          \"db.retry_count\": retryCount,\n        },\n      },\n      async (span) => {\n        const pool = await this.getPool();\n        const client = await pool.connect();\n        try {\n          await client.query(\"BEGIN\");\n          const result = await callback(client);\n          await client.query(\"COMMIT\");\n          span.setStatus({ code: SpanStatusCode.OK });\n          return result;\n        } catch (error) {\n          try {\n            await client.query(\"ROLLBACK\");\n          } catch {}\n          // retry on auth failure\n          if (this.isAuthError(error)) {\n            span.addEvent(\"auth_error_retry\");\n            client.release();\n            await this.rotateCredentials();\n            const newPool = await this.getPool();\n            const retryClient = await newPool.connect();\n            try {\n              await client.query(\"BEGIN\");\n              const result = await callback(retryClient);\n              await client.query(\"COMMIT\");\n              span.setStatus({ code: SpanStatusCode.OK });\n              return result;\n            } catch (retryError) {\n              try {\n                await retryClient.query(\"ROLLBACK\");\n              } catch {}\n              throw retryError;\n            } finally {\n              retryClient.release();\n            }\n          }\n\n          // retry on transient errors, but only once\n          if (this.isTransientError(error) && retryCount < 1) {\n            span.addEvent(\"transaction_error_retry\");\n            client.release();\n            await new Promise((resolve) => setTimeout(resolve, 100));\n            return await this.transaction<T>(callback, retryCount + 1);\n          }\n          span.recordException(error as Error);\n          span.setStatus({ code: SpanStatusCode.ERROR });\n          throw error;\n        } finally {\n          client.release();\n          const duration = Date.now() - startTime;\n          this.telemetryMetrics.queryCount.add(1);\n          this.telemetryMetrics.queryDuration.record(duration);\n          span.end();\n        }\n      },\n    );\n  }\n\n  /** Check if database connection is healthy */\n  async healthCheck(): Promise<boolean> {\n    return this.telemetry.startActiveSpan(\n      \"lakebase.healthCheck\",\n      {},\n      async (span) => {\n        try {\n          const result = await this.query<{ result: number }>(\n            \"SELECT 1 as result\",\n          );\n          const healthy = result.rows[0]?.result === 1;\n          span.setAttribute(\"db.healthy\", healthy);\n          span.setStatus({ code: SpanStatusCode.OK });\n          return healthy;\n        } catch {\n          span.setAttribute(\"db.healthy\", false);\n          span.setStatus({ code: SpanStatusCode.ERROR });\n          return false;\n        } finally {\n          span.end();\n        }\n      },\n    );\n  }\n\n  /** Close connection pool (call on shutdown) */\n  async close(): Promise<void> {\n    if (this.pool) {\n      await this.pool.end().catch((error: unknown) => {\n        console.error(\"Error closing connection pool:\", error);\n      });\n      this.pool = null;\n    }\n    this.credentials = null;\n  }\n\n  /** Setup graceful shutdown to close connection pools */\n  shutdown(): void {\n    process.on(\"SIGTERM\", () => this.close());\n    process.on(\"SIGINT\", () => this.close());\n    this.close();\n  }\n\n  /** Get Databricks workspace client - from config or execution context */\n  private getWorkspaceClient(): WorkspaceClient {\n    if (this.config.workspaceClient) {\n      return this.config.workspaceClient;\n    }\n\n    try {\n      const { getWorkspaceClient: getClient } = require(\"../../context\");\n      const client = getClient();\n\n      // cache it for subsequent calls\n      this.config.workspaceClient = client;\n      return client;\n    } catch (_error) {\n      throw new Error(\n        \"Databricks workspace client not available. Either pass it in config or ensure ServiceContext is initialized.\",\n      );\n    }\n  }\n\n  /** Get or create connection pool */\n  private async getPool(): Promise<pg.Pool> {\n    if (!this.connectionConfig) {\n      throw new Error(\n        \"Lakebase connection not configured. \" +\n          \"Set PGHOST, PGDATABASE, PGAPPNAME env vars, provide a connectionString, or pass explicit config.\",\n      );\n    }\n\n    if (!this.pool) {\n      const creds = await this.getCredentials();\n      this.pool = this.createPool(creds);\n    }\n    return this.pool;\n  }\n\n  /** Create PostgreSQL pool */\n  private createPool(credentials: {\n    username: string;\n    password: string;\n  }): pg.Pool {\n    const { host, database, port, sslMode } = this.connectionConfig;\n\n    const pool = new pg.Pool({\n      host,\n      port,\n      database,\n      user: credentials.username,\n      password: credentials.password,\n      max: this.config.maxPoolSize,\n      idleTimeoutMillis: this.config.idleTimeoutMs,\n      connectionTimeoutMillis: this.config.connectionTimeoutMs,\n      ssl: sslMode === \"require\" ? { rejectUnauthorized: true } : false,\n    });\n\n    pool.on(\"error\", (error: Error & { code?: string }) => {\n      console.error(\"Connection pool error:\", error.message, {\n        code: error.code,\n      });\n    });\n\n    return pool;\n  }\n\n  /** Get or fetch credentials with caching */\n  private async getCredentials(): Promise<{\n    username: string;\n    password: string;\n  }> {\n    const now = Date.now();\n\n    // return cached if still valid\n    if (\n      this.credentials &&\n      now < this.credentials.expiresAt - this.CACHE_BUFFER_MS\n    ) {\n      return this.credentials;\n    }\n\n    // fetch new credentials\n    const username = await this.fetchUsername();\n    const { token, expiresAt } = await this.fetchPassword();\n\n    this.credentials = {\n      username,\n      password: token,\n      expiresAt,\n    };\n\n    return { username, password: token };\n  }\n\n  /** Rotate credentials and recreate pool */\n  private async rotateCredentials(): Promise<void> {\n    // clear cached credentials\n    this.credentials = null;\n\n    if (this.pool) {\n      const oldPool = this.pool;\n      this.pool = null;\n      oldPool.end().catch((error: unknown) => {\n        console.error(\n          \"Error closing old connection pool during rotation:\",\n          error,\n        );\n      });\n    }\n  }\n\n  /** Fetch username from Databricks */\n  private async fetchUsername(): Promise<string> {\n    const workspaceClient = this.getWorkspaceClient();\n    const user = await workspaceClient.currentUser.me();\n    if (!user.userName) {\n      throw new Error(\"Failed to get current user from Databricks workspace\");\n    }\n    return user.userName;\n  }\n\n  /** Fetch password (OAuth token) from Databricks */\n  private async fetchPassword(): Promise<{ token: string; expiresAt: number }> {\n    const workspaceClient = this.getWorkspaceClient();\n    const config = new Config({ host: workspaceClient.config.host });\n    const apiClient = new ApiClient(config);\n\n    if (!this.connectionConfig.appName) {\n      throw new Error(`Database app name not found in connection config`);\n    }\n\n    const credentials = await apiClient.request({\n      path: `/api/2.0/database/credentials`,\n      method: \"POST\",\n      headers: new Headers(),\n      raw: false,\n      payload: {\n        instance_names: [this.connectionConfig.appName],\n        request_id: randomUUID(),\n      },\n    });\n\n    if (!this.validateCredentials(credentials)) {\n      throw new Error(\n        `Failed to generate credentials for instance: ${this.connectionConfig.appName}`,\n      );\n    }\n\n    const expiresAt = new Date(credentials.expiration_time).getTime();\n\n    return { token: credentials.token, expiresAt };\n  }\n\n  /** Check if error is auth failure */\n  private isAuthError(error: unknown): boolean {\n    return (\n      typeof error === \"object\" &&\n      error !== null &&\n      \"code\" in error &&\n      (error as any).code === \"28P01\"\n    );\n  }\n\n  /** Check if error is transient */\n  private isTransientError(error: unknown): boolean {\n    if (typeof error !== \"object\" || error === null || !(\"code\" in error)) {\n      return false;\n    }\n\n    const code = (error as any).code;\n    return (\n      code === \"ECONNRESET\" ||\n      code === \"ECONNREFUSED\" ||\n      code === \"ETIMEDOUT\" ||\n      code === \"57P01\" || // admin_shutdown\n      code === \"57P03\" || // cannot_connect_now\n      code === \"08006\" || // connection_failure\n      code === \"08003\" || // connection_does_not_exist\n      code === \"08000\" // connection_exception\n    );\n  }\n\n  /** Type guard for credentials */\n  private validateCredentials(\n    value: unknown,\n  ): value is { token: string; expiration_time: string } {\n    if (typeof value !== \"object\" || value === null) {\n      return false;\n    }\n\n    const credentials = value as { token: string; expiration_time: string };\n    return (\n      \"token\" in credentials &&\n      typeof credentials.token === \"string\" &&\n      \"expiration_time\" in credentials &&\n      typeof credentials.expiration_time === \"string\" &&\n      new Date(credentials.expiration_time).getTime() > Date.now()\n    );\n  }\n\n  /** Parse connection configuration from config or environment */\n  private parseConnectionConfig(): LakebaseConnectionConfig {\n    if (this.config.connectionString) {\n      return this.parseConnectionString(this.config.connectionString);\n    }\n\n    // get connection from config\n    if (this.config.host && this.config.database && this.config.appName) {\n      return {\n        host: this.config.host,\n        database: this.config.database,\n        port: this.config.port ?? 5432,\n        sslMode: this.config.sslMode ?? \"require\",\n        appName: this.config.appName,\n      };\n    }\n\n    // get connection from environment variables\n    const pgHost = process.env.PGHOST;\n    const pgDatabase = process.env.PGDATABASE;\n    const pgAppName = process.env.PGAPPNAME;\n    if (!pgHost || !pgDatabase || !pgAppName) {\n      throw new Error(\n        \"Lakebase connection not configured. Required env vars: PGHOST, PGDATABASE, PGAPPNAME. \" +\n          \"Optional: PGPORT (default: 5432), PGSSLMODE (default: require).\",\n      );\n    }\n    const pgPort = process.env.PGPORT;\n    const port = pgPort ? parseInt(pgPort, 10) : 5432;\n\n    if (Number.isNaN(port)) {\n      throw new Error(`Invalid port: ${pgPort}. Must be a number.`);\n    }\n\n    const pgSSLMode = process.env.PGSSLMODE;\n    const sslMode =\n      (pgSSLMode as \"require\" | \"disable\" | \"prefer\") || \"require\";\n\n    return {\n      host: pgHost,\n      database: pgDatabase,\n      port,\n      sslMode,\n      appName: pgAppName,\n    };\n  }\n\n  private parseConnectionString(\n    connectionString: string,\n  ): LakebaseConnectionConfig {\n    const url = new URL(connectionString);\n    const appName = url.searchParams.get(\"appName\");\n    if (!appName) {\n      throw new Error(\"Connection string must include appName parameter\");\n    }\n\n    return {\n      host: url.hostname,\n      database: url.pathname.slice(1), // remove leading slash\n      port: url.port ? parseInt(url.port, 10) : 5432,\n      sslMode:\n        (url.searchParams.get(\"sslmode\") as \"require\" | \"disable\" | \"prefer\") ??\n        \"require\",\n      appName: appName,\n    };\n  }\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;AAkCA,IAAa,oBAAb,MAA+B;CAe7B,YAAY,YAAsC;cAdlB;yBACG,MAAS;cAGb;qBACmB;AAUhD,OAAK,SAAS,UAAU,kBAAkB,WAAW;AACrD,OAAK,mBAAmB,KAAK,uBAAuB;AAEpD,OAAK,YAAY,iBAAiB,YAChC,KAAK,MACL,KAAK,OAAO,UACb;AACD,OAAK,mBAAmB;GACtB,YAAY,KAAK,UACd,UAAU,CACV,cAAc,wBAAwB;IACrC,aAAa;IACb,MAAM;IACP,CAAC;GACJ,eAAe,KAAK,UACjB,UAAU,CACV,gBAAgB,2BAA2B;IAC1C,aAAa;IACb,MAAM;IACP,CAAC;GACL;AAGD,MAAI,KAAK,OAAO,cAAc,EAC5B,OAAM,IAAI,MAAM,iCAAiC;;;;;;;;;;;CAarD,MAAM,MACJ,KACA,QACA,aAAqB,GACO;EAC5B,MAAM,YAAY,KAAK,KAAK;AAE5B,SAAO,KAAK,UAAU,gBACpB,kBACA,EACE,YAAY;GACV,aAAa;GACb,gBAAgB,IAAI,UAAU,GAAG,IAAI;GACrC,kBAAkB;GACnB,EACF,EACD,OAAO,SAAS;AACd,OAAI;IAEF,MAAM,SAAS,OADF,MAAM,KAAK,SAAS,EACP,MAAS,KAAK,OAAO;AAC/C,SAAK,aAAa,oBAAoB,OAAO,YAAY,EAAE;AAC3D,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;YACA,OAAO;AAEd,QAAI,KAAK,YAAY,MAAM,EAAE;AAC3B,UAAK,SAAS,mBAAmB;AACjC,WAAM,KAAK,mBAAmB;KAE9B,MAAM,SAAS,OADC,MAAM,KAAK,SAAS,EACP,MAAS,KAAK,OAAO;AAClD,UAAK,aAAa,oBAAoB,OAAO,YAAY,EAAE;AAC3D,UAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,YAAO;;AAIT,QAAI,KAAK,iBAAiB,MAAM,IAAI,aAAa,GAAG;AAClD,UAAK,SAAS,wBAAwB;AACtC,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,IAAI,CAAC;AACxD,YAAO,MAAM,KAAK,MAAS,KAAK,QAAQ,aAAa,EAAE;;AAGzD,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAE9C,UAAM;aACE;IACR,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,WAAW,IAAI,EAAE;AACvC,SAAK,iBAAiB,cAAc,OAAO,SAAS;AACpD,SAAK,KAAK;;IAGf;;;;;;;;;;;;;;;;CAiBH,MAAM,YACJ,UACA,aAAqB,GACT;EACZ,MAAM,YAAY,KAAK,KAAK;AAC5B,SAAO,KAAK,UAAU,gBACpB,wBACA,EACE,YAAY;GACV,aAAa;GACb,kBAAkB;GACnB,EACF,EACD,OAAO,SAAS;GAEd,MAAM,SAAS,OADF,MAAM,KAAK,SAAS,EACP,SAAS;AACnC,OAAI;AACF,UAAM,OAAO,MAAM,QAAQ;IAC3B,MAAM,SAAS,MAAM,SAAS,OAAO;AACrC,UAAM,OAAO,MAAM,SAAS;AAC5B,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;YACA,OAAO;AACd,QAAI;AACF,WAAM,OAAO,MAAM,WAAW;YACxB;AAER,QAAI,KAAK,YAAY,MAAM,EAAE;AAC3B,UAAK,SAAS,mBAAmB;AACjC,YAAO,SAAS;AAChB,WAAM,KAAK,mBAAmB;KAE9B,MAAM,cAAc,OADJ,MAAM,KAAK,SAAS,EACF,SAAS;AAC3C,SAAI;AACF,YAAM,OAAO,MAAM,QAAQ;MAC3B,MAAM,SAAS,MAAM,SAAS,YAAY;AAC1C,YAAM,OAAO,MAAM,SAAS;AAC5B,WAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,aAAO;cACA,YAAY;AACnB,UAAI;AACF,aAAM,YAAY,MAAM,WAAW;cAC7B;AACR,YAAM;eACE;AACR,kBAAY,SAAS;;;AAKzB,QAAI,KAAK,iBAAiB,MAAM,IAAI,aAAa,GAAG;AAClD,UAAK,SAAS,0BAA0B;AACxC,YAAO,SAAS;AAChB,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,IAAI,CAAC;AACxD,YAAO,MAAM,KAAK,YAAe,UAAU,aAAa,EAAE;;AAE5D,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,UAAM;aACE;AACR,WAAO,SAAS;IAChB,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,WAAW,IAAI,EAAE;AACvC,SAAK,iBAAiB,cAAc,OAAO,SAAS;AACpD,SAAK,KAAK;;IAGf;;;CAIH,MAAM,cAAgC;AACpC,SAAO,KAAK,UAAU,gBACpB,wBACA,EAAE,EACF,OAAO,SAAS;AACd,OAAI;IAIF,MAAM,WAHS,MAAM,KAAK,MACxB,qBACD,EACsB,KAAK,IAAI,WAAW;AAC3C,SAAK,aAAa,cAAc,QAAQ;AACxC,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;WACD;AACN,SAAK,aAAa,cAAc,MAAM;AACtC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,WAAO;aACC;AACR,SAAK,KAAK;;IAGf;;;CAIH,MAAM,QAAuB;AAC3B,MAAI,KAAK,MAAM;AACb,SAAM,KAAK,KAAK,KAAK,CAAC,OAAO,UAAmB;AAC9C,YAAQ,MAAM,kCAAkC,MAAM;KACtD;AACF,QAAK,OAAO;;AAEd,OAAK,cAAc;;;CAIrB,WAAiB;AACf,UAAQ,GAAG,iBAAiB,KAAK,OAAO,CAAC;AACzC,UAAQ,GAAG,gBAAgB,KAAK,OAAO,CAAC;AACxC,OAAK,OAAO;;;CAId,AAAQ,qBAAsC;AAC5C,MAAI,KAAK,OAAO,gBACd,QAAO,KAAK,OAAO;AAGrB,MAAI;GACF,MAAM,EAAE,oBAAoB;GAC5B,MAAM,SAAS,WAAW;AAG1B,QAAK,OAAO,kBAAkB;AAC9B,UAAO;WACA,QAAQ;AACf,SAAM,IAAI,MACR,+GACD;;;;CAKL,MAAc,UAA4B;AACxC,MAAI,CAAC,KAAK,iBACR,OAAM,IAAI,MACR,uIAED;AAGH,MAAI,CAAC,KAAK,MAAM;GACd,MAAM,QAAQ,MAAM,KAAK,gBAAgB;AACzC,QAAK,OAAO,KAAK,WAAW,MAAM;;AAEpC,SAAO,KAAK;;;CAId,AAAQ,WAAW,aAGP;EACV,MAAM,EAAE,MAAM,UAAU,MAAM,YAAY,KAAK;EAE/C,MAAM,OAAO,IAAI,GAAG,KAAK;GACvB;GACA;GACA;GACA,MAAM,YAAY;GAClB,UAAU,YAAY;GACtB,KAAK,KAAK,OAAO;GACjB,mBAAmB,KAAK,OAAO;GAC/B,yBAAyB,KAAK,OAAO;GACrC,KAAK,YAAY,YAAY,EAAE,oBAAoB,MAAM,GAAG;GAC7D,CAAC;AAEF,OAAK,GAAG,UAAU,UAAqC;AACrD,WAAQ,MAAM,0BAA0B,MAAM,SAAS,EACrD,MAAM,MAAM,MACb,CAAC;IACF;AAEF,SAAO;;;CAIT,MAAc,iBAGX;EACD,MAAM,MAAM,KAAK,KAAK;AAGtB,MACE,KAAK,eACL,MAAM,KAAK,YAAY,YAAY,KAAK,gBAExC,QAAO,KAAK;EAId,MAAM,WAAW,MAAM,KAAK,eAAe;EAC3C,MAAM,EAAE,OAAO,cAAc,MAAM,KAAK,eAAe;AAEvD,OAAK,cAAc;GACjB;GACA,UAAU;GACV;GACD;AAED,SAAO;GAAE;GAAU,UAAU;GAAO;;;CAItC,MAAc,oBAAmC;AAE/C,OAAK,cAAc;AAEnB,MAAI,KAAK,MAAM;GACb,MAAM,UAAU,KAAK;AACrB,QAAK,OAAO;AACZ,WAAQ,KAAK,CAAC,OAAO,UAAmB;AACtC,YAAQ,MACN,sDACA,MACD;KACD;;;;CAKN,MAAc,gBAAiC;EAE7C,MAAM,OAAO,MADW,KAAK,oBAAoB,CACd,YAAY,IAAI;AACnD,MAAI,CAAC,KAAK,SACR,OAAM,IAAI,MAAM,uDAAuD;AAEzE,SAAO,KAAK;;;CAId,MAAc,gBAA+D;EAG3E,MAAM,YAAY,IAAI,UADP,IAAI,OAAO,EAAE,MADJ,KAAK,oBAAoB,CACC,OAAO,MAAM,CAAC,CACzB;AAEvC,MAAI,CAAC,KAAK,iBAAiB,QACzB,OAAM,IAAI,MAAM,mDAAmD;EAGrE,MAAM,cAAc,MAAM,UAAU,QAAQ;GAC1C,MAAM;GACN,QAAQ;GACR,SAAS,IAAI,SAAS;GACtB,KAAK;GACL,SAAS;IACP,gBAAgB,CAAC,KAAK,iBAAiB,QAAQ;IAC/C,YAAY,YAAY;IACzB;GACF,CAAC;AAEF,MAAI,CAAC,KAAK,oBAAoB,YAAY,CACxC,OAAM,IAAI,MACR,gDAAgD,KAAK,iBAAiB,UACvE;EAGH,MAAM,YAAY,IAAI,KAAK,YAAY,gBAAgB,CAAC,SAAS;AAEjE,SAAO;GAAE,OAAO,YAAY;GAAO;GAAW;;;CAIhD,AAAQ,YAAY,OAAyB;AAC3C,SACE,OAAO,UAAU,YACjB,UAAU,QACV,UAAU,SACT,MAAc,SAAS;;;CAK5B,AAAQ,iBAAiB,OAAyB;AAChD,MAAI,OAAO,UAAU,YAAY,UAAU,QAAQ,EAAE,UAAU,OAC7D,QAAO;EAGT,MAAM,OAAQ,MAAc;AAC5B,SACE,SAAS,gBACT,SAAS,kBACT,SAAS,eACT,SAAS,WACT,SAAS,WACT,SAAS,WACT,SAAS,WACT,SAAS;;;CAKb,AAAQ,oBACN,OACqD;AACrD,MAAI,OAAO,UAAU,YAAY,UAAU,KACzC,QAAO;EAGT,MAAM,cAAc;AACpB,SACE,WAAW,eACX,OAAO,YAAY,UAAU,YAC7B,qBAAqB,eACrB,OAAO,YAAY,oBAAoB,YACvC,IAAI,KAAK,YAAY,gBAAgB,CAAC,SAAS,GAAG,KAAK,KAAK;;;CAKhE,AAAQ,wBAAkD;AACxD,MAAI,KAAK,OAAO,iBACd,QAAO,KAAK,sBAAsB,KAAK,OAAO,iBAAiB;AAIjE,MAAI,KAAK,OAAO,QAAQ,KAAK,OAAO,YAAY,KAAK,OAAO,QAC1D,QAAO;GACL,MAAM,KAAK,OAAO;GAClB,UAAU,KAAK,OAAO;GACtB,MAAM,KAAK,OAAO,QAAQ;GAC1B,SAAS,KAAK,OAAO,WAAW;GAChC,SAAS,KAAK,OAAO;GACtB;EAIH,MAAM,SAAS,QAAQ,IAAI;EAC3B,MAAM,aAAa,QAAQ,IAAI;EAC/B,MAAM,YAAY,QAAQ,IAAI;AAC9B,MAAI,CAAC,UAAU,CAAC,cAAc,CAAC,UAC7B,OAAM,IAAI,MACR,wJAED;EAEH,MAAM,SAAS,QAAQ,IAAI;EAC3B,MAAM,OAAO,SAAS,SAAS,QAAQ,GAAG,GAAG;AAE7C,MAAI,OAAO,MAAM,KAAK,CACpB,OAAM,IAAI,MAAM,iBAAiB,OAAO,qBAAqB;AAO/D,SAAO;GACL,MAAM;GACN,UAAU;GACV;GACA,SARgB,QAAQ,IAAI,aAEuB;GAOnD,SAAS;GACV;;CAGH,AAAQ,sBACN,kBAC0B;EAC1B,MAAM,MAAM,IAAI,IAAI,iBAAiB;EACrC,MAAM,UAAU,IAAI,aAAa,IAAI,UAAU;AAC/C,MAAI,CAAC,QACH,OAAM,IAAI,MAAM,mDAAmD;AAGrE,SAAO;GACL,MAAM,IAAI;GACV,UAAU,IAAI,SAAS,MAAM,EAAE;GAC/B,MAAM,IAAI,OAAO,SAAS,IAAI,MAAM,GAAG,GAAG;GAC1C,SACG,IAAI,aAAa,IAAI,UAAU,IAChC;GACO;GACV"}
+{"version":3,"file":"client.js","names":[],"sources":["../../../src/connectors/lakebase/client.ts"],"sourcesContent":["import { randomUUID } from \"node:crypto\";\nimport type { WorkspaceClient } from \"@databricks/sdk-experimental\";\nimport { ApiClient, Config } from \"@databricks/sdk-experimental\";\nimport pg from \"pg\";\nimport {\n  type Counter,\n  type Histogram,\n  SpanStatusCode,\n  TelemetryManager,\n  type TelemetryProvider,\n} from \"@/telemetry\";\nimport {\n  AppKitError,\n  AuthenticationError,\n  ConfigurationError,\n  ConnectionError,\n  ValidationError,\n} from \"../../errors\";\nimport { createLogger } from \"../../logging/logger\";\nimport { deepMerge } from \"../../utils\";\nimport { lakebaseDefaults } from \"./defaults\";\nimport type {\n  LakebaseConfig,\n  LakebaseConnectionConfig,\n  LakebaseCredentials,\n} from \"./types\";\n\nconst logger = createLogger(\"connectors:lakebase\");\n\n/**\n * Enterprise-grade connector for Databricks Lakebase\n * @example Simplest - everything from env/context\n * ```typescript\n * const connector = new LakebaseConnector();\n * await connector.query('SELECT * FROM users');\n * ```\n *\n * @example With explicit connection string\n * ```typescript\n * const connector = new LakebaseConnector({\n *   connectionString: 'postgresql://...'\n * });\n * ```\n */\nexport class LakebaseConnector {\n  private readonly name: string = \"lakebase\";\n  private readonly CACHE_BUFFER_MS = 2 * 60 * 1000;\n  private readonly config: LakebaseConfig;\n  private readonly connectionConfig: LakebaseConnectionConfig;\n  private pool: pg.Pool | null = null;\n  private credentials: LakebaseCredentials | null = null;\n\n  // telemetry\n  private readonly telemetry: TelemetryProvider;\n  private readonly telemetryMetrics: {\n    queryCount: Counter;\n    queryDuration: Histogram;\n  };\n\n  constructor(userConfig?: Partial<LakebaseConfig>) {\n    this.config = deepMerge(lakebaseDefaults, userConfig);\n    this.connectionConfig = this.parseConnectionConfig();\n\n    this.telemetry = TelemetryManager.getProvider(\n      this.name,\n      this.config.telemetry,\n    );\n    this.telemetryMetrics = {\n      queryCount: this.telemetry\n        .getMeter()\n        .createCounter(\"lakebase.query.count\", {\n          description: \"Total number of queries executed\",\n          unit: \"1\",\n        }),\n      queryDuration: this.telemetry\n        .getMeter()\n        .createHistogram(\"lakebase.query.duration\", {\n          description: \"Duration of queries executed\",\n          unit: \"ms\",\n        }),\n    };\n\n    // validate configuration\n    if (this.config.maxPoolSize < 1) {\n      throw ValidationError.invalidValue(\n        \"maxPoolSize\",\n        this.config.maxPoolSize,\n        \"at least 1\",\n      );\n    }\n  }\n\n  /**\n   * Execute a SQL query\n   *\n   * @example\n   * ```typescript\n   * const users = await connector.query('SELECT * FROM users');\n   * const user = await connector.query('SELECT * FROM users WHERE id = $1', [123]);\n   * ```\n   */\n  async query<T extends pg.QueryResultRow>(\n    sql: string,\n    params?: any[],\n    retryCount: number = 0,\n  ): Promise<pg.QueryResult<T>> {\n    const startTime = Date.now();\n\n    return this.telemetry.startActiveSpan(\n      \"lakebase.query\",\n      {\n        attributes: {\n          \"db.system\": \"lakebase\",\n          \"db.statement\": sql.substring(0, 500),\n          \"db.retry_count\": retryCount,\n        },\n      },\n      async (span) => {\n        try {\n          const pool = await this.getPool();\n          const result = await pool.query<T>(sql, params);\n          span.setAttribute(\"db.rows_affected\", result.rowCount ?? 0);\n          span.setStatus({ code: SpanStatusCode.OK });\n          return result;\n        } catch (error) {\n          // retry on auth failure\n          if (this.isAuthError(error)) {\n            span.addEvent(\"auth_error_retry\");\n            await this.rotateCredentials();\n            const newPool = await this.getPool();\n            const result = await newPool.query<T>(sql, params);\n            span.setAttribute(\"db.rows_affected\", result.rowCount ?? 0);\n            span.setStatus({ code: SpanStatusCode.OK });\n            return result;\n          }\n\n          // retry on transient errors, but only once\n          if (this.isTransientError(error) && retryCount < 1) {\n            span.addEvent(\"transient_error_retry\");\n            await new Promise((resolve) => setTimeout(resolve, 100));\n            return await this.query<T>(sql, params, retryCount + 1);\n          }\n\n          span.recordException(error as Error);\n          span.setStatus({ code: SpanStatusCode.ERROR });\n\n          if (error instanceof AppKitError) {\n            throw error;\n          }\n          throw ConnectionError.queryFailed(error as Error);\n        } finally {\n          const duration = Date.now() - startTime;\n          this.telemetryMetrics.queryCount.add(1);\n          this.telemetryMetrics.queryDuration.record(duration);\n          span.end();\n        }\n      },\n    );\n  }\n\n  /**\n   * Execute a transaction\n   *\n   * COMMIT and ROLLBACK are automatically managed by the transaction function.\n   *\n   * @param callback - Callback function to execute within the transaction context\n   * @example\n   * ```typescript\n   * await connector.transaction(async (client) => {\n   *   await client.query('INSERT INTO accounts (name) VALUES ($1)', ['Alice']);\n   *   await client.query('INSERT INTO logs (action) VALUES ($1)', ['Created Alice']);\n   * });\n   * ```\n   */\n  async transaction<T>(\n    callback: (client: pg.PoolClient) => Promise<T>,\n    retryCount: number = 0,\n  ): Promise<T> {\n    const startTime = Date.now();\n    return this.telemetry.startActiveSpan(\n      \"lakebase.transaction\",\n      {\n        attributes: {\n          \"db.system\": \"lakebase\",\n          \"db.retry_count\": retryCount,\n        },\n      },\n      async (span) => {\n        const pool = await this.getPool();\n        const client = await pool.connect();\n        try {\n          await client.query(\"BEGIN\");\n          const result = await callback(client);\n          await client.query(\"COMMIT\");\n          span.setStatus({ code: SpanStatusCode.OK });\n          return result;\n        } catch (error) {\n          try {\n            await client.query(\"ROLLBACK\");\n          } catch {}\n          // retry on auth failure\n          if (this.isAuthError(error)) {\n            span.addEvent(\"auth_error_retry\");\n            client.release();\n            await this.rotateCredentials();\n            const newPool = await this.getPool();\n            const retryClient = await newPool.connect();\n            try {\n              await client.query(\"BEGIN\");\n              const result = await callback(retryClient);\n              await client.query(\"COMMIT\");\n              span.setStatus({ code: SpanStatusCode.OK });\n              return result;\n            } catch (retryError) {\n              try {\n                await retryClient.query(\"ROLLBACK\");\n              } catch {}\n              throw retryError;\n            } finally {\n              retryClient.release();\n            }\n          }\n\n          // retry on transient errors, but only once\n          if (this.isTransientError(error) && retryCount < 1) {\n            span.addEvent(\"transaction_error_retry\");\n            client.release();\n            await new Promise((resolve) => setTimeout(resolve, 100));\n            return await this.transaction<T>(callback, retryCount + 1);\n          }\n          span.recordException(error as Error);\n          span.setStatus({ code: SpanStatusCode.ERROR });\n\n          if (error instanceof AppKitError) {\n            throw error;\n          }\n          throw ConnectionError.transactionFailed(error as Error);\n        } finally {\n          client.release();\n          const duration = Date.now() - startTime;\n          this.telemetryMetrics.queryCount.add(1);\n          this.telemetryMetrics.queryDuration.record(duration);\n          span.end();\n        }\n      },\n    );\n  }\n\n  /** Check if database connection is healthy */\n  async healthCheck(): Promise<boolean> {\n    return this.telemetry.startActiveSpan(\n      \"lakebase.healthCheck\",\n      {},\n      async (span) => {\n        try {\n          const result = await this.query<{ result: number }>(\n            \"SELECT 1 as result\",\n          );\n          const healthy = result.rows[0]?.result === 1;\n          span.setAttribute(\"db.healthy\", healthy);\n          span.setStatus({ code: SpanStatusCode.OK });\n          return healthy;\n        } catch {\n          span.setAttribute(\"db.healthy\", false);\n          span.setStatus({ code: SpanStatusCode.ERROR });\n          return false;\n        } finally {\n          span.end();\n        }\n      },\n    );\n  }\n\n  /** Close connection pool (call on shutdown) */\n  async close(): Promise<void> {\n    if (this.pool) {\n      await this.pool.end().catch((error: unknown) => {\n        logger.error(\"Error closing connection pool: %O\", error);\n      });\n      this.pool = null;\n    }\n    this.credentials = null;\n  }\n\n  /** Setup graceful shutdown to close connection pools */\n  shutdown(): void {\n    process.on(\"SIGTERM\", () => this.close());\n    process.on(\"SIGINT\", () => this.close());\n    this.close();\n  }\n\n  /** Get Databricks workspace client - from config or execution context */\n  private getWorkspaceClient(): WorkspaceClient {\n    if (this.config.workspaceClient) {\n      return this.config.workspaceClient;\n    }\n\n    try {\n      const { getWorkspaceClient: getClient } = require(\"../../context\");\n      const client = getClient();\n\n      // cache it for subsequent calls\n      this.config.workspaceClient = client;\n      return client;\n    } catch (_error) {\n      throw ConnectionError.clientUnavailable(\n        \"Databricks workspace client\",\n        \"Either pass it in config or ensure ServiceContext is initialized\",\n      );\n    }\n  }\n\n  /** Get or create connection pool */\n  private async getPool(): Promise<pg.Pool> {\n    if (!this.connectionConfig) {\n      throw ConfigurationError.invalidConnection(\n        \"Lakebase\",\n        \"Set PGHOST, PGDATABASE, PGAPPNAME env vars, provide a connectionString, or pass explicit config\",\n      );\n    }\n\n    if (!this.pool) {\n      const creds = await this.getCredentials();\n      this.pool = this.createPool(creds);\n    }\n    return this.pool;\n  }\n\n  /** Create PostgreSQL pool */\n  private createPool(credentials: {\n    username: string;\n    password: string;\n  }): pg.Pool {\n    const { host, database, port, sslMode } = this.connectionConfig;\n\n    const pool = new pg.Pool({\n      host,\n      port,\n      database,\n      user: credentials.username,\n      password: credentials.password,\n      max: this.config.maxPoolSize,\n      idleTimeoutMillis: this.config.idleTimeoutMs,\n      connectionTimeoutMillis: this.config.connectionTimeoutMs,\n      ssl: sslMode === \"require\" ? { rejectUnauthorized: true } : false,\n    });\n\n    pool.on(\"error\", (error: Error & { code?: string }) => {\n      logger.error(\n        \"Connection pool error: %s (code: %s)\",\n        error.message,\n        error.code,\n      );\n    });\n\n    return pool;\n  }\n\n  /** Get or fetch credentials with caching */\n  private async getCredentials(): Promise<{\n    username: string;\n    password: string;\n  }> {\n    const now = Date.now();\n\n    // return cached if still valid\n    if (\n      this.credentials &&\n      now < this.credentials.expiresAt - this.CACHE_BUFFER_MS\n    ) {\n      return this.credentials;\n    }\n\n    // fetch new credentials\n    const username = await this.fetchUsername();\n    const { token, expiresAt } = await this.fetchPassword();\n\n    this.credentials = {\n      username,\n      password: token,\n      expiresAt,\n    };\n\n    return { username, password: token };\n  }\n\n  /** Rotate credentials and recreate pool */\n  private async rotateCredentials(): Promise<void> {\n    // clear cached credentials\n    this.credentials = null;\n\n    if (this.pool) {\n      const oldPool = this.pool;\n      this.pool = null;\n      oldPool.end().catch((error: unknown) => {\n        logger.error(\n          \"Error closing old connection pool during rotation: %O\",\n          error,\n        );\n      });\n    }\n  }\n\n  /** Fetch username from Databricks */\n  private async fetchUsername(): Promise<string> {\n    const workspaceClient = this.getWorkspaceClient();\n    const user = await workspaceClient.currentUser.me();\n    if (!user.userName) {\n      throw AuthenticationError.userLookupFailed();\n    }\n    return user.userName;\n  }\n\n  /** Fetch password (OAuth token) from Databricks */\n  private async fetchPassword(): Promise<{ token: string; expiresAt: number }> {\n    const workspaceClient = this.getWorkspaceClient();\n    const config = new Config({ host: workspaceClient.config.host });\n    const apiClient = new ApiClient(config);\n\n    if (!this.connectionConfig.appName) {\n      throw ConfigurationError.resourceNotFound(\"Database app name\");\n    }\n\n    const credentials = await apiClient.request({\n      path: `/api/2.0/database/credentials`,\n      method: \"POST\",\n      headers: new Headers(),\n      raw: false,\n      payload: {\n        instance_names: [this.connectionConfig.appName],\n        request_id: randomUUID(),\n      },\n    });\n\n    if (!this.validateCredentials(credentials)) {\n      throw AuthenticationError.credentialsFailed(\n        this.connectionConfig.appName,\n      );\n    }\n\n    const expiresAt = new Date(credentials.expiration_time).getTime();\n\n    return { token: credentials.token, expiresAt };\n  }\n\n  /** Check if error is auth failure */\n  private isAuthError(error: unknown): boolean {\n    return (\n      typeof error === \"object\" &&\n      error !== null &&\n      \"code\" in error &&\n      (error as any).code === \"28P01\"\n    );\n  }\n\n  /** Check if error is transient */\n  private isTransientError(error: unknown): boolean {\n    if (typeof error !== \"object\" || error === null || !(\"code\" in error)) {\n      return false;\n    }\n\n    const code = (error as any).code;\n    return (\n      code === \"ECONNRESET\" ||\n      code === \"ECONNREFUSED\" ||\n      code === \"ETIMEDOUT\" ||\n      code === \"57P01\" || // admin_shutdown\n      code === \"57P03\" || // cannot_connect_now\n      code === \"08006\" || // connection_failure\n      code === \"08003\" || // connection_does_not_exist\n      code === \"08000\" // connection_exception\n    );\n  }\n\n  /** Type guard for credentials */\n  private validateCredentials(\n    value: unknown,\n  ): value is { token: string; expiration_time: string } {\n    if (typeof value !== \"object\" || value === null) {\n      return false;\n    }\n\n    const credentials = value as { token: string; expiration_time: string };\n    return (\n      \"token\" in credentials &&\n      typeof credentials.token === \"string\" &&\n      \"expiration_time\" in credentials &&\n      typeof credentials.expiration_time === \"string\" &&\n      new Date(credentials.expiration_time).getTime() > Date.now()\n    );\n  }\n\n  /** Parse connection configuration from config or environment */\n  private parseConnectionConfig(): LakebaseConnectionConfig {\n    if (this.config.connectionString) {\n      return this.parseConnectionString(this.config.connectionString);\n    }\n\n    // get connection from config\n    if (this.config.host && this.config.database && this.config.appName) {\n      return {\n        host: this.config.host,\n        database: this.config.database,\n        port: this.config.port ?? 5432,\n        sslMode: this.config.sslMode ?? \"require\",\n        appName: this.config.appName,\n      };\n    }\n\n    // get connection from environment variables\n    const pgHost = process.env.PGHOST;\n    const pgDatabase = process.env.PGDATABASE;\n    const pgAppName = process.env.PGAPPNAME;\n    if (!pgHost || !pgDatabase || !pgAppName) {\n      throw ConfigurationError.invalidConnection(\n        \"Lakebase\",\n        \"Required env vars: PGHOST, PGDATABASE, PGAPPNAME. Optional: PGPORT (default: 5432), PGSSLMODE (default: require)\",\n      );\n    }\n    const pgPort = process.env.PGPORT;\n    const port = pgPort ? parseInt(pgPort, 10) : 5432;\n\n    if (Number.isNaN(port)) {\n      throw ValidationError.invalidValue(\"port\", pgPort, \"a number\");\n    }\n\n    const pgSSLMode = process.env.PGSSLMODE;\n    const sslMode =\n      (pgSSLMode as \"require\" | \"disable\" | \"prefer\") || \"require\";\n\n    return {\n      host: pgHost,\n      database: pgDatabase,\n      port,\n      sslMode,\n      appName: pgAppName,\n    };\n  }\n\n  private parseConnectionString(\n    connectionString: string,\n  ): LakebaseConnectionConfig {\n    const url = new URL(connectionString);\n    const appName = url.searchParams.get(\"appName\");\n    if (!appName) {\n      throw ConfigurationError.missingConnectionParam(\"appName\");\n    }\n\n    return {\n      host: url.hostname,\n      database: url.pathname.slice(1), // remove leading slash\n      port: url.port ? parseInt(url.port, 10) : 5432,\n      sslMode:\n        (url.searchParams.get(\"sslmode\") as \"require\" | \"disable\" | \"prefer\") ??\n        \"require\",\n      appName: appName,\n    };\n  }\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;aAiBsB;AAUtB,MAAM,SAAS,aAAa,sBAAsB;;;;;;;;;;;;;;;;AAiBlD,IAAa,oBAAb,MAA+B;CAe7B,YAAY,YAAsC;cAdlB;yBACG,MAAS;cAGb;qBACmB;AAUhD,OAAK,SAAS,UAAU,kBAAkB,WAAW;AACrD,OAAK,mBAAmB,KAAK,uBAAuB;AAEpD,OAAK,YAAY,iBAAiB,YAChC,KAAK,MACL,KAAK,OAAO,UACb;AACD,OAAK,mBAAmB;GACtB,YAAY,KAAK,UACd,UAAU,CACV,cAAc,wBAAwB;IACrC,aAAa;IACb,MAAM;IACP,CAAC;GACJ,eAAe,KAAK,UACjB,UAAU,CACV,gBAAgB,2BAA2B;IAC1C,aAAa;IACb,MAAM;IACP,CAAC;GACL;AAGD,MAAI,KAAK,OAAO,cAAc,EAC5B,OAAM,gBAAgB,aACpB,eACA,KAAK,OAAO,aACZ,aACD;;;;;;;;;;;CAaL,MAAM,MACJ,KACA,QACA,aAAqB,GACO;EAC5B,MAAM,YAAY,KAAK,KAAK;AAE5B,SAAO,KAAK,UAAU,gBACpB,kBACA,EACE,YAAY;GACV,aAAa;GACb,gBAAgB,IAAI,UAAU,GAAG,IAAI;GACrC,kBAAkB;GACnB,EACF,EACD,OAAO,SAAS;AACd,OAAI;IAEF,MAAM,SAAS,OADF,MAAM,KAAK,SAAS,EACP,MAAS,KAAK,OAAO;AAC/C,SAAK,aAAa,oBAAoB,OAAO,YAAY,EAAE;AAC3D,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;YACA,OAAO;AAEd,QAAI,KAAK,YAAY,MAAM,EAAE;AAC3B,UAAK,SAAS,mBAAmB;AACjC,WAAM,KAAK,mBAAmB;KAE9B,MAAM,SAAS,OADC,MAAM,KAAK,SAAS,EACP,MAAS,KAAK,OAAO;AAClD,UAAK,aAAa,oBAAoB,OAAO,YAAY,EAAE;AAC3D,UAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,YAAO;;AAIT,QAAI,KAAK,iBAAiB,MAAM,IAAI,aAAa,GAAG;AAClD,UAAK,SAAS,wBAAwB;AACtC,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,IAAI,CAAC;AACxD,YAAO,MAAM,KAAK,MAAS,KAAK,QAAQ,aAAa,EAAE;;AAGzD,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAE9C,QAAI,iBAAiB,YACnB,OAAM;AAER,UAAM,gBAAgB,YAAY,MAAe;aACzC;IACR,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,WAAW,IAAI,EAAE;AACvC,SAAK,iBAAiB,cAAc,OAAO,SAAS;AACpD,SAAK,KAAK;;IAGf;;;;;;;;;;;;;;;;CAiBH,MAAM,YACJ,UACA,aAAqB,GACT;EACZ,MAAM,YAAY,KAAK,KAAK;AAC5B,SAAO,KAAK,UAAU,gBACpB,wBACA,EACE,YAAY;GACV,aAAa;GACb,kBAAkB;GACnB,EACF,EACD,OAAO,SAAS;GAEd,MAAM,SAAS,OADF,MAAM,KAAK,SAAS,EACP,SAAS;AACnC,OAAI;AACF,UAAM,OAAO,MAAM,QAAQ;IAC3B,MAAM,SAAS,MAAM,SAAS,OAAO;AACrC,UAAM,OAAO,MAAM,SAAS;AAC5B,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;YACA,OAAO;AACd,QAAI;AACF,WAAM,OAAO,MAAM,WAAW;YACxB;AAER,QAAI,KAAK,YAAY,MAAM,EAAE;AAC3B,UAAK,SAAS,mBAAmB;AACjC,YAAO,SAAS;AAChB,WAAM,KAAK,mBAAmB;KAE9B,MAAM,cAAc,OADJ,MAAM,KAAK,SAAS,EACF,SAAS;AAC3C,SAAI;AACF,YAAM,OAAO,MAAM,QAAQ;MAC3B,MAAM,SAAS,MAAM,SAAS,YAAY;AAC1C,YAAM,OAAO,MAAM,SAAS;AAC5B,WAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,aAAO;cACA,YAAY;AACnB,UAAI;AACF,aAAM,YAAY,MAAM,WAAW;cAC7B;AACR,YAAM;eACE;AACR,kBAAY,SAAS;;;AAKzB,QAAI,KAAK,iBAAiB,MAAM,IAAI,aAAa,GAAG;AAClD,UAAK,SAAS,0BAA0B;AACxC,YAAO,SAAS;AAChB,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,IAAI,CAAC;AACxD,YAAO,MAAM,KAAK,YAAe,UAAU,aAAa,EAAE;;AAE5D,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAE9C,QAAI,iBAAiB,YACnB,OAAM;AAER,UAAM,gBAAgB,kBAAkB,MAAe;aAC/C;AACR,WAAO,SAAS;IAChB,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,WAAW,IAAI,EAAE;AACvC,SAAK,iBAAiB,cAAc,OAAO,SAAS;AACpD,SAAK,KAAK;;IAGf;;;CAIH,MAAM,cAAgC;AACpC,SAAO,KAAK,UAAU,gBACpB,wBACA,EAAE,EACF,OAAO,SAAS;AACd,OAAI;IAIF,MAAM,WAHS,MAAM,KAAK,MACxB,qBACD,EACsB,KAAK,IAAI,WAAW;AAC3C,SAAK,aAAa,cAAc,QAAQ;AACxC,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;WACD;AACN,SAAK,aAAa,cAAc,MAAM;AACtC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,WAAO;aACC;AACR,SAAK,KAAK;;IAGf;;;CAIH,MAAM,QAAuB;AAC3B,MAAI,KAAK,MAAM;AACb,SAAM,KAAK,KAAK,KAAK,CAAC,OAAO,UAAmB;AAC9C,WAAO,MAAM,qCAAqC,MAAM;KACxD;AACF,QAAK,OAAO;;AAEd,OAAK,cAAc;;;CAIrB,WAAiB;AACf,UAAQ,GAAG,iBAAiB,KAAK,OAAO,CAAC;AACzC,UAAQ,GAAG,gBAAgB,KAAK,OAAO,CAAC;AACxC,OAAK,OAAO;;;CAId,AAAQ,qBAAsC;AAC5C,MAAI,KAAK,OAAO,gBACd,QAAO,KAAK,OAAO;AAGrB,MAAI;GACF,MAAM,EAAE,oBAAoB;GAC5B,MAAM,SAAS,WAAW;AAG1B,QAAK,OAAO,kBAAkB;AAC9B,UAAO;WACA,QAAQ;AACf,SAAM,gBAAgB,kBACpB,+BACA,mEACD;;;;CAKL,MAAc,UAA4B;AACxC,MAAI,CAAC,KAAK,iBACR,OAAM,mBAAmB,kBACvB,YACA,kGACD;AAGH,MAAI,CAAC,KAAK,MAAM;GACd,MAAM,QAAQ,MAAM,KAAK,gBAAgB;AACzC,QAAK,OAAO,KAAK,WAAW,MAAM;;AAEpC,SAAO,KAAK;;;CAId,AAAQ,WAAW,aAGP;EACV,MAAM,EAAE,MAAM,UAAU,MAAM,YAAY,KAAK;EAE/C,MAAM,OAAO,IAAI,GAAG,KAAK;GACvB;GACA;GACA;GACA,MAAM,YAAY;GAClB,UAAU,YAAY;GACtB,KAAK,KAAK,OAAO;GACjB,mBAAmB,KAAK,OAAO;GAC/B,yBAAyB,KAAK,OAAO;GACrC,KAAK,YAAY,YAAY,EAAE,oBAAoB,MAAM,GAAG;GAC7D,CAAC;AAEF,OAAK,GAAG,UAAU,UAAqC;AACrD,UAAO,MACL,wCACA,MAAM,SACN,MAAM,KACP;IACD;AAEF,SAAO;;;CAIT,MAAc,iBAGX;EACD,MAAM,MAAM,KAAK,KAAK;AAGtB,MACE,KAAK,eACL,MAAM,KAAK,YAAY,YAAY,KAAK,gBAExC,QAAO,KAAK;EAId,MAAM,WAAW,MAAM,KAAK,eAAe;EAC3C,MAAM,EAAE,OAAO,cAAc,MAAM,KAAK,eAAe;AAEvD,OAAK,cAAc;GACjB;GACA,UAAU;GACV;GACD;AAED,SAAO;GAAE;GAAU,UAAU;GAAO;;;CAItC,MAAc,oBAAmC;AAE/C,OAAK,cAAc;AAEnB,MAAI,KAAK,MAAM;GACb,MAAM,UAAU,KAAK;AACrB,QAAK,OAAO;AACZ,WAAQ,KAAK,CAAC,OAAO,UAAmB;AACtC,WAAO,MACL,yDACA,MACD;KACD;;;;CAKN,MAAc,gBAAiC;EAE7C,MAAM,OAAO,MADW,KAAK,oBAAoB,CACd,YAAY,IAAI;AACnD,MAAI,CAAC,KAAK,SACR,OAAM,oBAAoB,kBAAkB;AAE9C,SAAO,KAAK;;;CAId,MAAc,gBAA+D;EAG3E,MAAM,YAAY,IAAI,UADP,IAAI,OAAO,EAAE,MADJ,KAAK,oBAAoB,CACC,OAAO,MAAM,CAAC,CACzB;AAEvC,MAAI,CAAC,KAAK,iBAAiB,QACzB,OAAM,mBAAmB,iBAAiB,oBAAoB;EAGhE,MAAM,cAAc,MAAM,UAAU,QAAQ;GAC1C,MAAM;GACN,QAAQ;GACR,SAAS,IAAI,SAAS;GACtB,KAAK;GACL,SAAS;IACP,gBAAgB,CAAC,KAAK,iBAAiB,QAAQ;IAC/C,YAAY,YAAY;IACzB;GACF,CAAC;AAEF,MAAI,CAAC,KAAK,oBAAoB,YAAY,CACxC,OAAM,oBAAoB,kBACxB,KAAK,iBAAiB,QACvB;EAGH,MAAM,YAAY,IAAI,KAAK,YAAY,gBAAgB,CAAC,SAAS;AAEjE,SAAO;GAAE,OAAO,YAAY;GAAO;GAAW;;;CAIhD,AAAQ,YAAY,OAAyB;AAC3C,SACE,OAAO,UAAU,YACjB,UAAU,QACV,UAAU,SACT,MAAc,SAAS;;;CAK5B,AAAQ,iBAAiB,OAAyB;AAChD,MAAI,OAAO,UAAU,YAAY,UAAU,QAAQ,EAAE,UAAU,OAC7D,QAAO;EAGT,MAAM,OAAQ,MAAc;AAC5B,SACE,SAAS,gBACT,SAAS,kBACT,SAAS,eACT,SAAS,WACT,SAAS,WACT,SAAS,WACT,SAAS,WACT,SAAS;;;CAKb,AAAQ,oBACN,OACqD;AACrD,MAAI,OAAO,UAAU,YAAY,UAAU,KACzC,QAAO;EAGT,MAAM,cAAc;AACpB,SACE,WAAW,eACX,OAAO,YAAY,UAAU,YAC7B,qBAAqB,eACrB,OAAO,YAAY,oBAAoB,YACvC,IAAI,KAAK,YAAY,gBAAgB,CAAC,SAAS,GAAG,KAAK,KAAK;;;CAKhE,AAAQ,wBAAkD;AACxD,MAAI,KAAK,OAAO,iBACd,QAAO,KAAK,sBAAsB,KAAK,OAAO,iBAAiB;AAIjE,MAAI,KAAK,OAAO,QAAQ,KAAK,OAAO,YAAY,KAAK,OAAO,QAC1D,QAAO;GACL,MAAM,KAAK,OAAO;GAClB,UAAU,KAAK,OAAO;GACtB,MAAM,KAAK,OAAO,QAAQ;GAC1B,SAAS,KAAK,OAAO,WAAW;GAChC,SAAS,KAAK,OAAO;GACtB;EAIH,MAAM,SAAS,QAAQ,IAAI;EAC3B,MAAM,aAAa,QAAQ,IAAI;EAC/B,MAAM,YAAY,QAAQ,IAAI;AAC9B,MAAI,CAAC,UAAU,CAAC,cAAc,CAAC,UAC7B,OAAM,mBAAmB,kBACvB,YACA,mHACD;EAEH,MAAM,SAAS,QAAQ,IAAI;EAC3B,MAAM,OAAO,SAAS,SAAS,QAAQ,GAAG,GAAG;AAE7C,MAAI,OAAO,MAAM,KAAK,CACpB,OAAM,gBAAgB,aAAa,QAAQ,QAAQ,WAAW;AAOhE,SAAO;GACL,MAAM;GACN,UAAU;GACV;GACA,SARgB,QAAQ,IAAI,aAEuB;GAOnD,SAAS;GACV;;CAGH,AAAQ,sBACN,kBAC0B;EAC1B,MAAM,MAAM,IAAI,IAAI,iBAAiB;EACrC,MAAM,UAAU,IAAI,aAAa,IAAI,UAAU;AAC/C,MAAI,CAAC,QACH,OAAM,mBAAmB,uBAAuB,UAAU;AAG5D,SAAO;GACL,MAAM,IAAI;GACV,UAAU,IAAI,SAAS,MAAM,EAAE;GAC/B,MAAM,IAAI,OAAO,SAAS,IAAI,MAAM,GAAG,GAAG;GAC1C,SACG,IAAI,aAAa,IAAI,UAAU,IAChC;GACO;GACV"}

package/dist/connectors/sql-warehouse/client.js

@@ -1,10 +1,18 @@
+import { createLogger } from "../../logging/logger.js";
 import { TelemetryManager } from "../../telemetry/telemetry-manager.js";
 import { SpanKind, SpanStatusCode } from "../../telemetry/index.js";
+import { AppKitError } from "../../errors/base.js";
+import { ConnectionError } from "../../errors/connection.js";
+import { ExecutionError } from "../../errors/execution.js";
+import { ValidationError } from "../../errors/validation.js";
+import { init_errors } from "../../errors/index.js";
 import { ArrowStreamProcessor } from "../../stream/arrow-stream-processor.js";
 import { executeStatementDefaults } from "./defaults.js";
 import { Context } from "@databricks/sdk-experimental";
 
 //#region src/connectors/sql-warehouse/client.ts
+init_errors();
+const logger = createLogger("connectors:sql-warehouse");
 var SQLWarehouseConnector = class {
 	constructor(config) {
 		this.name = "sql-warehouse";
@@ -37,6 +45,7 @@ var SQLWarehouseConnector = class {
 	async executeStatement(workspaceClient, input, signal) {
 		const startTime = Date.now();
 		let success = false;
+		if (signal?.aborted) throw ExecutionError.canceled();
 		return this.telemetry.startActiveSpan("sql.query", {
 			kind: SpanKind.CLIENT,
 			attributes: {
@@ -48,9 +57,24 @@ var SQLWarehouseConnector = class {
 				"db.has_parameters": !!input.parameters
 			}
 		}, async (span) => {
+			let abortHandler;
+			let isAborted = false;
+			if (signal) {
+				abortHandler = () => {
+					if (!span.isRecording()) return;
+					isAborted = true;
+					span.setAttribute("cancelled", true);
+					span.setStatus({
+						code: SpanStatusCode.ERROR,
+						message: "Query cancelled by client"
+					});
+					span.end();
+				};
+				signal.addEventListener("abort", abortHandler, { once: true });
+			}
 			try {
-				if (!input.statement) throw new Error("Statement is required: Please provide a SQL statement to execute");
-				if (!input.warehouse_id) throw new Error("Warehouse ID is required: Please provide a warehouse_id to execute the statement");
+				if (!input.statement) throw ValidationError.missingField("statement");
+				if (!input.warehouse_id) throw ValidationError.missingField("warehouse_id");
 				const body = {
 					statement: input.statement,
 					parameters: input.parameters,
@@ -66,7 +90,7 @@ var SQLWarehouseConnector = class {
 				};
 				span.addEvent("statement.submitting", { "db.warehouse_id": input.warehouse_id });
 				const response = await workspaceClient.statementExecution.executeStatement(body, this._createContext(signal));
-				if (!response) throw new Error("No response received from SQL Warehouse API");
+				if (!response) throw ConnectionError.apiFailure("SQL Warehouse");
 				const status = response.status;
 				const statementId = response.statement_id;
 				span.setAttribute("db.statement_id", statementId);
@@ -84,27 +108,37 @@ var SQLWarehouseConnector = class {
 					case "SUCCEEDED":
 						result = this._transformDataArray(response);
 						break;
-					case "FAILED": throw new Error(`Statement failed: ${status.error?.message || "Unknown error"}`);
-					case "CANCELED": throw new Error("Statement was canceled");
-					case "CLOSED": throw new Error("Statement execution completed but results are no longer available (CLOSED state)");
-					default: throw new Error(`Unknown statement state: ${status?.state}`);
+					case "FAILED": throw ExecutionError.statementFailed(status.error?.message);
+					case "CANCELED": throw ExecutionError.canceled();
+					case "CLOSED": throw ExecutionError.resultsClosed();
+					default: throw ExecutionError.unknownState(String(status?.state ?? "unknown"));
 				}
 				const resultData = result.result;
-				if (resultData?.data) span.setAttribute("db.result.row_count", resultData.data.length);
-				else if (resultData?.data_array) span.setAttribute("db.result.row_count", resultData.data_array.length);
+				const rowCount = resultData?.data?.length ?? resultData?.data_array?.length ?? 0;
+				if (rowCount > 0) span.setAttribute("db.result.row_count", rowCount);
+				const duration = Date.now() - startTime;
+				logger.event()?.setContext("sql-warehouse", {
+					warehouse_id: input.warehouse_id,
+					rows_returned: rowCount,
+					query_duration_ms: duration
+				});
 				success = true;
-				span.setStatus({ code: SpanStatusCode.OK });
+				if (!isAborted) span.setStatus({ code: SpanStatusCode.OK });
 				return result;
 			} catch (error) {
-				span.recordException(error);
-				span.setStatus({
-					code: SpanStatusCode.ERROR,
-					message: error instanceof Error ? error.message : String(error)
-				});
-				throw error;
+				if (!isAborted) {
+					span.recordException(error);
+					span.setStatus({
+						code: SpanStatusCode.ERROR,
+						message: error instanceof Error ? error.message : String(error)
+					});
+				}
+				if (error instanceof AppKitError) throw error;
+				throw ExecutionError.statementFailed(error instanceof Error ? error.message : String(error));
 			} finally {
+				if (abortHandler && signal) signal.removeEventListener("abort", abortHandler);
 				const duration = Date.now() - startTime;
-				span.end();
+				if (!isAborted) span.end();
 				const attributes = {
 					"db.warehouse_id": input.warehouse_id,
 					"db.catalog": input.catalog ?? "",
@@ -135,13 +169,13 @@ var SQLWarehouseConnector = class {
 					span.setAttribute("db.polling.current_attempt", pollCount);
 					const elapsedTime = Date.now() - startTime;
 					if (elapsedTime > timeout) {
-						const error = /* @__PURE__ */ new Error(`Statement polling timeout exceeded after ${timeout}ms (elapsed: ${elapsedTime}ms)`);
+						const error = ExecutionError.statementFailed(`Polling timeout exceeded after ${timeout}ms (elapsed: ${elapsedTime}ms)`);
 						span.recordException(error);
 						span.setStatus({ code: SpanStatusCode.ERROR });
 						throw error;
 					}
 					if (signal?.aborted) {
-						const error = /* @__PURE__ */ new Error("Request aborted");
+						const error = ExecutionError.canceled();
 						span.recordException(error);
 						span.setStatus({ code: SpanStatusCode.ERROR });
 						throw error;
@@ -152,7 +186,7 @@ var SQLWarehouseConnector = class {
 						"poll.elapsed_ms": elapsedTime
 					});
 					const response = await workspaceClient.statementExecution.getStatement({ statement_id: statementId }, this._createContext(signal));
-					if (!response) throw new Error("No response received from SQL Warehouse API");
+					if (!response) throw ConnectionError.apiFailure("SQL Warehouse");
 					const status = response.status;
 					span.addEvent("polling.status_check", {
 						"db.status": status?.state,
@@ -170,10 +204,10 @@ var SQLWarehouseConnector = class {
 							});
 							span.setStatus({ code: SpanStatusCode.OK });
 							return this._transformDataArray(response);
-						case "FAILED": throw new Error(`Statement failed: ${status.error?.message || "Unknown error"}`);
-						case "CANCELED": throw new Error("Statement was canceled");
-						case "CLOSED": throw new Error("Statement execution completed but results are no longer available (CLOSED state)");
-						default: throw new Error(`Unknown statement state: ${status?.state}`);
+						case "FAILED": throw ExecutionError.statementFailed(status.error?.message);
+						case "CANCELED": throw ExecutionError.canceled();
+						case "CLOSED": throw ExecutionError.resultsClosed();
+						default: throw ExecutionError.unknownState(String(status?.state ?? "unknown"));
 					}
 					await new Promise((resolve) => setTimeout(resolve, delay));
 					delay = Math.min(delay * 2, maxDelayBetweenPolls);
@@ -184,7 +218,8 @@ var SQLWarehouseConnector = class {
 					code: SpanStatusCode.ERROR,
 					message: error instanceof Error ? error.message : String(error)
 				});
-				throw error;
+				if (error instanceof AppKitError) throw error;
+				throw ExecutionError.statementFailed(error instanceof Error ? error.message : String(error));
 			} finally {
 				span.end();
 			}
@@ -242,7 +277,7 @@ var SQLWarehouseConnector = class {
 				const response = await workspaceClient.statementExecution.getStatement({ statement_id: jobId }, this._createContext(signal));
 				const chunks = response.result?.external_links;
 				const schema = response.manifest?.schema;
-				if (!chunks || !schema) throw new Error("No chunks or schema found in response");
+				if (!chunks || !schema) throw ExecutionError.missingData("chunks or schema");
 				span.setAttribute("arrow.chunk_count", chunks.length);
 				const result = await this.arrowProcessor.processChunks(chunks, schema, signal);
 				span.setAttribute("arrow.data_size_bytes", result.data.length);
@@ -252,6 +287,10 @@ var SQLWarehouseConnector = class {
 					operation: "arrow.getData",
 					status: "success"
 				});
+				logger.event()?.setContext("sql-warehouse", {
+					arrow_data_size_bytes: result.data.length,
+					arrow_job_id: jobId
+				});
 				return result;
 			} catch (error) {
 				span.setStatus({
@@ -264,8 +303,9 @@ var SQLWarehouseConnector = class {
 					operation: "arrow.getData",
 					status: "error"
 				});
-				console.error(`Failed Arrow job: ${jobId}`, error);
-				throw error;
+				logger.error("Failed Arrow job: %s %O", jobId, error);
+				if (error instanceof AppKitError) throw error;
+				throw ExecutionError.statementFailed(error instanceof Error ? error.message : String(error));
 			}
 		});
 	}

package/dist/connectors/sql-warehouse/client.js.map

@@ -1 +1 @@
-{"version":3,"file":"client.js","names":[],"sources":["../../../src/connectors/sql-warehouse/client.ts"],"sourcesContent":["import {\n  Context,\n  type sql,\n  type WorkspaceClient,\n} from \"@databricks/sdk-experimental\";\nimport { ArrowStreamProcessor } from \"../../stream/arrow-stream-processor\";\nimport type { TelemetryOptions } from \"shared\";\nimport type { TelemetryProvider } from \"../../telemetry\";\nimport {\n  type Counter,\n  type Histogram,\n  type Span,\n  SpanKind,\n  SpanStatusCode,\n  TelemetryManager,\n} from \"../../telemetry\";\nimport { executeStatementDefaults } from \"./defaults\";\n\nexport interface SQLWarehouseConfig {\n  timeout?: number;\n  telemetry?: TelemetryOptions;\n}\n\nexport class SQLWarehouseConnector {\n  private readonly name = \"sql-warehouse\";\n\n  private config: SQLWarehouseConfig;\n\n  // Lazy-initialized: only created when Arrow format is used\n  private _arrowProcessor: ArrowStreamProcessor | null = null;\n  // telemetry\n  private readonly telemetry: TelemetryProvider;\n  private readonly telemetryMetrics: {\n    queryCount: Counter;\n    queryDuration: Histogram;\n  };\n\n  constructor(config: SQLWarehouseConfig) {\n    this.config = config;\n\n    this.telemetry = TelemetryManager.getProvider(\n      this.name,\n      this.config.telemetry,\n    );\n    this.telemetryMetrics = {\n      queryCount: this.telemetry.getMeter().createCounter(\"query.count\", {\n        description: \"Total number of queries executed\",\n        unit: \"1\",\n      }),\n      queryDuration: this.telemetry\n        .getMeter()\n        .createHistogram(\"query.duration\", {\n          description: \"Duration of queries executed\",\n          unit: \"ms\",\n        }),\n    };\n  }\n\n  /**\n   * Lazily initializes and returns the ArrowStreamProcessor.\n   * Only created on first Arrow format query to avoid unnecessary allocation.\n   */\n  private get arrowProcessor(): ArrowStreamProcessor {\n    if (!this._arrowProcessor) {\n      this._arrowProcessor = new ArrowStreamProcessor({\n        timeout: this.config.timeout || executeStatementDefaults.timeout,\n        maxConcurrentDownloads:\n          ArrowStreamProcessor.DEFAULT_MAX_CONCURRENT_DOWNLOADS,\n        retries: ArrowStreamProcessor.DEFAULT_RETRIES,\n      });\n    }\n    return this._arrowProcessor;\n  }\n\n  async executeStatement(\n    workspaceClient: WorkspaceClient,\n    input: sql.ExecuteStatementRequest,\n    signal?: AbortSignal,\n  ) {\n    const startTime = Date.now();\n    let success = false;\n\n    return this.telemetry.startActiveSpan(\n      \"sql.query\",\n      {\n        kind: SpanKind.CLIENT,\n        attributes: {\n          \"db.system\": \"databricks\",\n          \"db.warehouse_id\": input.warehouse_id || \"\",\n          \"db.catalog\": input.catalog ?? \"\",\n          \"db.schema\": input.schema ?? \"\",\n          \"db.statement\": input.statement?.substring(0, 500) || \"\",\n          \"db.has_parameters\": !!input.parameters,\n        },\n      },\n      async (span: Span) => {\n        try {\n          // validate required fields\n          if (!input.statement) {\n            throw new Error(\n              \"Statement is required: Please provide a SQL statement to execute\",\n            );\n          }\n\n          if (!input.warehouse_id) {\n            throw new Error(\n              \"Warehouse ID is required: Please provide a warehouse_id to execute the statement\",\n            );\n          }\n\n          const body: sql.ExecuteStatementRequest = {\n            statement: input.statement,\n            parameters: input.parameters,\n            warehouse_id: input.warehouse_id,\n            catalog: input.catalog,\n            schema: input.schema,\n            wait_timeout:\n              input.wait_timeout || executeStatementDefaults.wait_timeout,\n            disposition:\n              input.disposition || executeStatementDefaults.disposition,\n            format: input.format || executeStatementDefaults.format,\n            byte_limit: input.byte_limit,\n            row_limit: input.row_limit,\n            on_wait_timeout:\n              input.on_wait_timeout || executeStatementDefaults.on_wait_timeout,\n          };\n\n          span.addEvent(\"statement.submitting\", {\n            \"db.warehouse_id\": input.warehouse_id,\n          });\n\n          const response =\n            await workspaceClient.statementExecution.executeStatement(\n              body,\n              this._createContext(signal),\n            );\n\n          if (!response) {\n            throw new Error(\"No response received from SQL Warehouse API\");\n          }\n          const status = response.status;\n          const statementId = response.statement_id as string;\n\n          span.setAttribute(\"db.statement_id\", statementId);\n          span.addEvent(\"statement.submitted\", {\n            \"db.statement_id\": response.statement_id,\n            \"db.status\": status?.state,\n          });\n\n          let result:\n            | sql.StatementResponse\n            | { result: { statement_id: string; status: sql.StatementStatus } };\n\n          switch (status?.state) {\n            case \"RUNNING\":\n            case \"PENDING\":\n              span.addEvent(\"statement.polling_started\", {\n                \"db.status\": response.status?.state,\n              });\n              result = await this._pollForStatementResult(\n                workspaceClient,\n                statementId,\n                this.config.timeout,\n                signal,\n              );\n              break;\n            case \"SUCCEEDED\":\n              result = this._transformDataArray(response);\n              break;\n            case \"FAILED\":\n              throw new Error(\n                `Statement failed: ${status.error?.message || \"Unknown error\"}`,\n              );\n            case \"CANCELED\":\n              throw new Error(\"Statement was canceled\");\n            case \"CLOSED\":\n              throw new Error(\n                \"Statement execution completed but results are no longer available (CLOSED state)\",\n              );\n            default:\n              throw new Error(`Unknown statement state: ${status?.state}`);\n          }\n\n          const resultData = result.result as any;\n          if (resultData?.data) {\n            span.setAttribute(\"db.result.row_count\", resultData.data.length);\n          } else if (resultData?.data_array) {\n            span.setAttribute(\n              \"db.result.row_count\",\n              resultData.data_array.length,\n            );\n          }\n\n          success = true;\n          span.setStatus({ code: SpanStatusCode.OK });\n          return result;\n        } catch (error) {\n          span.recordException(error as Error);\n          span.setStatus({\n            code: SpanStatusCode.ERROR,\n            message: error instanceof Error ? error.message : String(error),\n          });\n          throw error;\n        } finally {\n          const duration = Date.now() - startTime;\n          span.end();\n\n          const attributes = {\n            \"db.warehouse_id\": input.warehouse_id,\n            \"db.catalog\": input.catalog ?? \"\",\n            \"db.schema\": input.schema ?? \"\",\n            \"db.statement\": input.statement?.substring(0, 500) || \"\",\n            success: success.toString(),\n          };\n\n          this.telemetryMetrics.queryCount.add(1, attributes);\n          this.telemetryMetrics.queryDuration.record(duration, attributes);\n        }\n      },\n      { name: this.name, includePrefix: true },\n    );\n  }\n\n  private async _pollForStatementResult(\n    workspaceClient: WorkspaceClient,\n    statementId: string,\n    timeout = executeStatementDefaults.timeout,\n    signal?: AbortSignal,\n  ) {\n    return this.telemetry.startActiveSpan(\n      \"sql.poll\",\n      {\n        attributes: {\n          \"db.statement_id\": statementId,\n          \"db.polling.timeout\": timeout,\n        },\n      },\n      async (span: Span) => {\n        try {\n          const startTime = Date.now();\n          let delay = 1000;\n          const maxDelayBetweenPolls = 5000; // max 5 seconds between polls\n          let pollCount = 0;\n\n          while (true) {\n            pollCount++;\n            span.setAttribute(\"db.polling.current_attempt\", pollCount);\n\n            // check if timeout exceeded\n            const elapsedTime = Date.now() - startTime;\n            if (elapsedTime > timeout) {\n              const error = new Error(\n                `Statement polling timeout exceeded after ${timeout}ms (elapsed: ${elapsedTime}ms)`,\n              );\n              span.recordException(error);\n              span.setStatus({ code: SpanStatusCode.ERROR });\n              throw error;\n            }\n\n            if (signal?.aborted) {\n              const error = new Error(\"Request aborted\");\n              span.recordException(error);\n              span.setStatus({ code: SpanStatusCode.ERROR });\n              throw error;\n            }\n\n            span.addEvent(\"polling.attempt\", {\n              \"poll.attempt\": pollCount,\n              \"poll.delay_ms\": delay,\n              \"poll.elapsed_ms\": elapsedTime,\n            });\n\n            const response =\n              await workspaceClient.statementExecution.getStatement(\n                {\n                  statement_id: statementId,\n                },\n                this._createContext(signal),\n              );\n            if (!response) {\n              throw new Error(\"No response received from SQL Warehouse API\");\n            }\n\n            const status = response.status;\n\n            span.addEvent(\"polling.status_check\", {\n              \"db.status\": status?.state,\n              \"poll.attempt\": pollCount,\n            });\n\n            switch (status?.state) {\n              case \"PENDING\":\n              case \"RUNNING\":\n                // continue polling\n                break;\n              case \"SUCCEEDED\":\n                span.setAttribute(\"db.polling.attempts\", pollCount);\n                span.setAttribute(\"db.polling.total_duration_ms\", elapsedTime);\n                span.addEvent(\"polling.completed\", {\n                  \"poll.attempts\": pollCount,\n                  \"poll.duration_ms\": elapsedTime,\n                });\n                span.setStatus({ code: SpanStatusCode.OK });\n                return this._transformDataArray(response);\n              case \"FAILED\":\n                throw new Error(\n                  `Statement failed: ${\n                    status.error?.message || \"Unknown error\"\n                  }`,\n                );\n              case \"CANCELED\":\n                throw new Error(\"Statement was canceled\");\n              case \"CLOSED\":\n                throw new Error(\n                  \"Statement execution completed but results are no longer available (CLOSED state)\",\n                );\n              default:\n                throw new Error(`Unknown statement state: ${status?.state}`);\n            }\n\n            // continue polling after delay\n            await new Promise((resolve) => setTimeout(resolve, delay));\n            delay = Math.min(delay * 2, maxDelayBetweenPolls);\n          }\n        } catch (error) {\n          span.recordException(error as Error);\n          span.setStatus({\n            code: SpanStatusCode.ERROR,\n            message: error instanceof Error ? error.message : String(error),\n          });\n          throw error;\n        } finally {\n          span.end();\n        }\n      },\n      { name: this.name, includePrefix: true },\n    );\n  }\n\n  private _transformDataArray(response: sql.StatementResponse) {\n    if (response.manifest?.format === \"ARROW_STREAM\") {\n      return this.updateWithArrowStatus(response);\n    }\n\n    if (!response.result?.data_array || !response.manifest?.schema?.columns) {\n      return response;\n    }\n\n    const columns = response.manifest.schema.columns;\n\n    const transformedData = response.result.data_array.map((row) => {\n      const obj: Record<string, unknown> = {};\n      row.forEach((value, index) => {\n        const column = columns[index];\n        const columnName = column?.name || `column_${index}`;\n\n        // attempt to parse JSON strings for string columns\n        if (\n          column?.type_name === \"STRING\" &&\n          typeof value === \"string\" &&\n          value &&\n          (value[0] === \"{\" || value[0] === \"[\")\n        ) {\n          try {\n            obj[columnName] = JSON.parse(value);\n          } catch {\n            // if parsing fails, keep as string\n            obj[columnName] = value;\n          }\n        } else {\n          obj[columnName] = value;\n        }\n      });\n      return obj;\n    });\n\n    // remove data_array\n    const { data_array: _data_array, ...restResult } = response.result;\n    return {\n      ...response,\n      result: {\n        ...restResult,\n        data: transformedData,\n      },\n    };\n  }\n\n  private updateWithArrowStatus(response: sql.StatementResponse): {\n    result: { statement_id: string; status: sql.StatementStatus };\n  } {\n    return {\n      result: {\n        statement_id: response.statement_id as string,\n        status: {\n          state: response.status?.state,\n          error: response.status?.error,\n        } as sql.StatementStatus,\n      },\n    };\n  }\n\n  async getArrowData(\n    workspaceClient: WorkspaceClient,\n    jobId: string,\n    signal?: AbortSignal,\n  ): Promise<ReturnType<typeof this.arrowProcessor.processChunks>> {\n    const startTime = Date.now();\n\n    return this.telemetry.startActiveSpan(\n      \"arrow.getData\",\n      {\n        kind: SpanKind.CLIENT,\n        attributes: {\n          \"db.system\": \"databricks\",\n          \"arrow.job_id\": jobId,\n        },\n      },\n      async (span: Span) => {\n        try {\n          const response =\n            await workspaceClient.statementExecution.getStatement(\n              { statement_id: jobId },\n              this._createContext(signal),\n            );\n\n          const chunks = response.result?.external_links;\n          const schema = response.manifest?.schema;\n\n          if (!chunks || !schema) {\n            throw new Error(\"No chunks or schema found in response\");\n          }\n\n          span.setAttribute(\"arrow.chunk_count\", chunks.length);\n\n          const result = await this.arrowProcessor.processChunks(\n            chunks,\n            schema,\n            signal,\n          );\n\n          span.setAttribute(\"arrow.data_size_bytes\", result.data.length);\n          span.setStatus({ code: SpanStatusCode.OK });\n\n          const duration = Date.now() - startTime;\n          this.telemetryMetrics.queryDuration.record(duration, {\n            operation: \"arrow.getData\",\n            status: \"success\",\n          });\n\n          return result;\n        } catch (error) {\n          span.setStatus({\n            code: SpanStatusCode.ERROR,\n            message: error instanceof Error ? error.message : \"Unknown error\",\n          });\n          span.recordException(error as Error);\n\n          const duration = Date.now() - startTime;\n          this.telemetryMetrics.queryDuration.record(duration, {\n            operation: \"arrow.getData\",\n            status: \"error\",\n          });\n\n          console.error(`Failed Arrow job: ${jobId}`, error);\n          throw error;\n        }\n      },\n    );\n  }\n\n  // create context for cancellation token\n  private _createContext(signal?: AbortSignal) {\n    return new Context({\n      cancellationToken: {\n        isCancellationRequested: signal?.aborted ?? false,\n        onCancellationRequested: (cb: () => void) => {\n          signal?.addEventListener(\"abort\", cb, { once: true });\n        },\n      },\n    });\n  }\n}\n"],"mappings":";;;;;;;AAuBA,IAAa,wBAAb,MAAmC;CAcjC,YAAY,QAA4B;cAbhB;yBAK+B;AASrD,OAAK,SAAS;AAEd,OAAK,YAAY,iBAAiB,YAChC,KAAK,MACL,KAAK,OAAO,UACb;AACD,OAAK,mBAAmB;GACtB,YAAY,KAAK,UAAU,UAAU,CAAC,cAAc,eAAe;IACjE,aAAa;IACb,MAAM;IACP,CAAC;GACF,eAAe,KAAK,UACjB,UAAU,CACV,gBAAgB,kBAAkB;IACjC,aAAa;IACb,MAAM;IACP,CAAC;GACL;;;;;;CAOH,IAAY,iBAAuC;AACjD,MAAI,CAAC,KAAK,gBACR,MAAK,kBAAkB,IAAI,qBAAqB;GAC9C,SAAS,KAAK,OAAO,WAAW,yBAAyB;GACzD,wBACE,qBAAqB;GACvB,SAAS,qBAAqB;GAC/B,CAAC;AAEJ,SAAO,KAAK;;CAGd,MAAM,iBACJ,iBACA,OACA,QACA;EACA,MAAM,YAAY,KAAK,KAAK;EAC5B,IAAI,UAAU;AAEd,SAAO,KAAK,UAAU,gBACpB,aACA;GACE,MAAM,SAAS;GACf,YAAY;IACV,aAAa;IACb,mBAAmB,MAAM,gBAAgB;IACzC,cAAc,MAAM,WAAW;IAC/B,aAAa,MAAM,UAAU;IAC7B,gBAAgB,MAAM,WAAW,UAAU,GAAG,IAAI,IAAI;IACtD,qBAAqB,CAAC,CAAC,MAAM;IAC9B;GACF,EACD,OAAO,SAAe;AACpB,OAAI;AAEF,QAAI,CAAC,MAAM,UACT,OAAM,IAAI,MACR,mEACD;AAGH,QAAI,CAAC,MAAM,aACT,OAAM,IAAI,MACR,mFACD;IAGH,MAAM,OAAoC;KACxC,WAAW,MAAM;KACjB,YAAY,MAAM;KAClB,cAAc,MAAM;KACpB,SAAS,MAAM;KACf,QAAQ,MAAM;KACd,cACE,MAAM,gBAAgB,yBAAyB;KACjD,aACE,MAAM,eAAe,yBAAyB;KAChD,QAAQ,MAAM,UAAU,yBAAyB;KACjD,YAAY,MAAM;KAClB,WAAW,MAAM;KACjB,iBACE,MAAM,mBAAmB,yBAAyB;KACrD;AAED,SAAK,SAAS,wBAAwB,EACpC,mBAAmB,MAAM,cAC1B,CAAC;IAEF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,iBACvC,MACA,KAAK,eAAe,OAAO,CAC5B;AAEH,QAAI,CAAC,SACH,OAAM,IAAI,MAAM,8CAA8C;IAEhE,MAAM,SAAS,SAAS;IACxB,MAAM,cAAc,SAAS;AAE7B,SAAK,aAAa,mBAAmB,YAAY;AACjD,SAAK,SAAS,uBAAuB;KACnC,mBAAmB,SAAS;KAC5B,aAAa,QAAQ;KACtB,CAAC;IAEF,IAAI;AAIJ,YAAQ,QAAQ,OAAhB;KACE,KAAK;KACL,KAAK;AACH,WAAK,SAAS,6BAA6B,EACzC,aAAa,SAAS,QAAQ,OAC/B,CAAC;AACF,eAAS,MAAM,KAAK,wBAClB,iBACA,aACA,KAAK,OAAO,SACZ,OACD;AACD;KACF,KAAK;AACH,eAAS,KAAK,oBAAoB,SAAS;AAC3C;KACF,KAAK,SACH,OAAM,IAAI,MACR,qBAAqB,OAAO,OAAO,WAAW,kBAC/C;KACH,KAAK,WACH,OAAM,IAAI,MAAM,yBAAyB;KAC3C,KAAK,SACH,OAAM,IAAI,MACR,mFACD;KACH,QACE,OAAM,IAAI,MAAM,4BAA4B,QAAQ,QAAQ;;IAGhE,MAAM,aAAa,OAAO;AAC1B,QAAI,YAAY,KACd,MAAK,aAAa,uBAAuB,WAAW,KAAK,OAAO;aACvD,YAAY,WACrB,MAAK,aACH,uBACA,WAAW,WAAW,OACvB;AAGH,cAAU;AACV,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;YACA,OAAO;AACd,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU;KACb,MAAM,eAAe;KACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM;KAChE,CAAC;AACF,UAAM;aACE;IACR,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,KAAK;IAEV,MAAM,aAAa;KACjB,mBAAmB,MAAM;KACzB,cAAc,MAAM,WAAW;KAC/B,aAAa,MAAM,UAAU;KAC7B,gBAAgB,MAAM,WAAW,UAAU,GAAG,IAAI,IAAI;KACtD,SAAS,QAAQ,UAAU;KAC5B;AAED,SAAK,iBAAiB,WAAW,IAAI,GAAG,WAAW;AACnD,SAAK,iBAAiB,cAAc,OAAO,UAAU,WAAW;;KAGpE;GAAE,MAAM,KAAK;GAAM,eAAe;GAAM,CACzC;;CAGH,MAAc,wBACZ,iBACA,aACA,UAAU,yBAAyB,SACnC,QACA;AACA,SAAO,KAAK,UAAU,gBACpB,YACA,EACE,YAAY;GACV,mBAAmB;GACnB,sBAAsB;GACvB,EACF,EACD,OAAO,SAAe;AACpB,OAAI;IACF,MAAM,YAAY,KAAK,KAAK;IAC5B,IAAI,QAAQ;IACZ,MAAM,uBAAuB;IAC7B,IAAI,YAAY;AAEhB,WAAO,MAAM;AACX;AACA,UAAK,aAAa,8BAA8B,UAAU;KAG1D,MAAM,cAAc,KAAK,KAAK,GAAG;AACjC,SAAI,cAAc,SAAS;MACzB,MAAM,wBAAQ,IAAI,MAChB,4CAA4C,QAAQ,eAAe,YAAY,KAChF;AACD,WAAK,gBAAgB,MAAM;AAC3B,WAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,YAAM;;AAGR,SAAI,QAAQ,SAAS;MACnB,MAAM,wBAAQ,IAAI,MAAM,kBAAkB;AAC1C,WAAK,gBAAgB,MAAM;AAC3B,WAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,YAAM;;AAGR,UAAK,SAAS,mBAAmB;MAC/B,gBAAgB;MAChB,iBAAiB;MACjB,mBAAmB;MACpB,CAAC;KAEF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,aACvC,EACE,cAAc,aACf,EACD,KAAK,eAAe,OAAO,CAC5B;AACH,SAAI,CAAC,SACH,OAAM,IAAI,MAAM,8CAA8C;KAGhE,MAAM,SAAS,SAAS;AAExB,UAAK,SAAS,wBAAwB;MACpC,aAAa,QAAQ;MACrB,gBAAgB;MACjB,CAAC;AAEF,aAAQ,QAAQ,OAAhB;MACE,KAAK;MACL,KAAK,UAEH;MACF,KAAK;AACH,YAAK,aAAa,uBAAuB,UAAU;AACnD,YAAK,aAAa,gCAAgC,YAAY;AAC9D,YAAK,SAAS,qBAAqB;QACjC,iBAAiB;QACjB,oBAAoB;QACrB,CAAC;AACF,YAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,cAAO,KAAK,oBAAoB,SAAS;MAC3C,KAAK,SACH,OAAM,IAAI,MACR,qBACE,OAAO,OAAO,WAAW,kBAE5B;MACH,KAAK,WACH,OAAM,IAAI,MAAM,yBAAyB;MAC3C,KAAK,SACH,OAAM,IAAI,MACR,mFACD;MACH,QACE,OAAM,IAAI,MAAM,4BAA4B,QAAQ,QAAQ;;AAIhE,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,MAAM,CAAC;AAC1D,aAAQ,KAAK,IAAI,QAAQ,GAAG,qBAAqB;;YAE5C,OAAO;AACd,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU;KACb,MAAM,eAAe;KACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM;KAChE,CAAC;AACF,UAAM;aACE;AACR,SAAK,KAAK;;KAGd;GAAE,MAAM,KAAK;GAAM,eAAe;GAAM,CACzC;;CAGH,AAAQ,oBAAoB,UAAiC;AAC3D,MAAI,SAAS,UAAU,WAAW,eAChC,QAAO,KAAK,sBAAsB,SAAS;AAG7C,MAAI,CAAC,SAAS,QAAQ,cAAc,CAAC,SAAS,UAAU,QAAQ,QAC9D,QAAO;EAGT,MAAM,UAAU,SAAS,SAAS,OAAO;EAEzC,MAAM,kBAAkB,SAAS,OAAO,WAAW,KAAK,QAAQ;GAC9D,MAAM,MAA+B,EAAE;AACvC,OAAI,SAAS,OAAO,UAAU;IAC5B,MAAM,SAAS,QAAQ;IACvB,MAAM,aAAa,QAAQ,QAAQ,UAAU;AAG7C,QACE,QAAQ,cAAc,YACtB,OAAO,UAAU,YACjB,UACC,MAAM,OAAO,OAAO,MAAM,OAAO,KAElC,KAAI;AACF,SAAI,cAAc,KAAK,MAAM,MAAM;YAC7B;AAEN,SAAI,cAAc;;QAGpB,KAAI,cAAc;KAEpB;AACF,UAAO;IACP;EAGF,MAAM,EAAE,YAAY,aAAa,GAAG,eAAe,SAAS;AAC5D,SAAO;GACL,GAAG;GACH,QAAQ;IACN,GAAG;IACH,MAAM;IACP;GACF;;CAGH,AAAQ,sBAAsB,UAE5B;AACA,SAAO,EACL,QAAQ;GACN,cAAc,SAAS;GACvB,QAAQ;IACN,OAAO,SAAS,QAAQ;IACxB,OAAO,SAAS,QAAQ;IACzB;GACF,EACF;;CAGH,MAAM,aACJ,iBACA,OACA,QAC+D;EAC/D,MAAM,YAAY,KAAK,KAAK;AAE5B,SAAO,KAAK,UAAU,gBACpB,iBACA;GACE,MAAM,SAAS;GACf,YAAY;IACV,aAAa;IACb,gBAAgB;IACjB;GACF,EACD,OAAO,SAAe;AACpB,OAAI;IACF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,aACvC,EAAE,cAAc,OAAO,EACvB,KAAK,eAAe,OAAO,CAC5B;IAEH,MAAM,SAAS,SAAS,QAAQ;IAChC,MAAM,SAAS,SAAS,UAAU;AAElC,QAAI,CAAC,UAAU,CAAC,OACd,OAAM,IAAI,MAAM,wCAAwC;AAG1D,SAAK,aAAa,qBAAqB,OAAO,OAAO;IAErD,MAAM,SAAS,MAAM,KAAK,eAAe,cACvC,QACA,QACA,OACD;AAED,SAAK,aAAa,yBAAyB,OAAO,KAAK,OAAO;AAC9D,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;IAE3C,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,cAAc,OAAO,UAAU;KACnD,WAAW;KACX,QAAQ;KACT,CAAC;AAEF,WAAO;YACA,OAAO;AACd,SAAK,UAAU;KACb,MAAM,eAAe;KACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU;KACnD,CAAC;AACF,SAAK,gBAAgB,MAAe;IAEpC,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,cAAc,OAAO,UAAU;KACnD,WAAW;KACX,QAAQ;KACT,CAAC;AAEF,YAAQ,MAAM,qBAAqB,SAAS,MAAM;AAClD,UAAM;;IAGX;;CAIH,AAAQ,eAAe,QAAsB;AAC3C,SAAO,IAAI,QAAQ,EACjB,mBAAmB;GACjB,yBAAyB,QAAQ,WAAW;GAC5C,0BAA0B,OAAmB;AAC3C,YAAQ,iBAAiB,SAAS,IAAI,EAAE,MAAM,MAAM,CAAC;;GAExD,EACF,CAAC"}
+{"version":3,"file":"client.js","names":[],"sources":["../../../src/connectors/sql-warehouse/client.ts"],"sourcesContent":["import {\n  Context,\n  type sql,\n  type WorkspaceClient,\n} from \"@databricks/sdk-experimental\";\nimport type { TelemetryOptions } from \"shared\";\nimport {\n  AppKitError,\n  ConnectionError,\n  ExecutionError,\n  ValidationError,\n} from \"../../errors\";\nimport { createLogger } from \"../../logging/logger\";\nimport { ArrowStreamProcessor } from \"../../stream/arrow-stream-processor\";\nimport type { TelemetryProvider } from \"../../telemetry\";\nimport {\n  type Counter,\n  type Histogram,\n  type Span,\n  SpanKind,\n  SpanStatusCode,\n  TelemetryManager,\n} from \"../../telemetry\";\nimport { executeStatementDefaults } from \"./defaults\";\n\nconst logger = createLogger(\"connectors:sql-warehouse\");\n\nexport interface SQLWarehouseConfig {\n  timeout?: number;\n  telemetry?: TelemetryOptions;\n}\n\nexport class SQLWarehouseConnector {\n  private readonly name = \"sql-warehouse\";\n\n  private config: SQLWarehouseConfig;\n\n  // Lazy-initialized: only created when Arrow format is used\n  private _arrowProcessor: ArrowStreamProcessor | null = null;\n  // telemetry\n  private readonly telemetry: TelemetryProvider;\n  private readonly telemetryMetrics: {\n    queryCount: Counter;\n    queryDuration: Histogram;\n  };\n\n  constructor(config: SQLWarehouseConfig) {\n    this.config = config;\n\n    this.telemetry = TelemetryManager.getProvider(\n      this.name,\n      this.config.telemetry,\n    );\n    this.telemetryMetrics = {\n      queryCount: this.telemetry.getMeter().createCounter(\"query.count\", {\n        description: \"Total number of queries executed\",\n        unit: \"1\",\n      }),\n      queryDuration: this.telemetry\n        .getMeter()\n        .createHistogram(\"query.duration\", {\n          description: \"Duration of queries executed\",\n          unit: \"ms\",\n        }),\n    };\n  }\n\n  /**\n   * Lazily initializes and returns the ArrowStreamProcessor.\n   * Only created on first Arrow format query to avoid unnecessary allocation.\n   */\n  private get arrowProcessor(): ArrowStreamProcessor {\n    if (!this._arrowProcessor) {\n      this._arrowProcessor = new ArrowStreamProcessor({\n        timeout: this.config.timeout || executeStatementDefaults.timeout,\n        maxConcurrentDownloads:\n          ArrowStreamProcessor.DEFAULT_MAX_CONCURRENT_DOWNLOADS,\n        retries: ArrowStreamProcessor.DEFAULT_RETRIES,\n      });\n    }\n    return this._arrowProcessor;\n  }\n\n  async executeStatement(\n    workspaceClient: WorkspaceClient,\n    input: sql.ExecuteStatementRequest,\n    signal?: AbortSignal,\n  ) {\n    const startTime = Date.now();\n    let success = false;\n\n    // if signal is aborted, throw an error\n    if (signal?.aborted) {\n      throw ExecutionError.canceled();\n    }\n\n    return this.telemetry.startActiveSpan(\n      \"sql.query\",\n      {\n        kind: SpanKind.CLIENT,\n        attributes: {\n          \"db.system\": \"databricks\",\n          \"db.warehouse_id\": input.warehouse_id || \"\",\n          \"db.catalog\": input.catalog ?? \"\",\n          \"db.schema\": input.schema ?? \"\",\n          \"db.statement\": input.statement?.substring(0, 500) || \"\",\n          \"db.has_parameters\": !!input.parameters,\n        },\n      },\n      async (span: Span) => {\n        let abortHandler: (() => void) | undefined;\n        let isAborted = false;\n\n        if (signal) {\n          abortHandler = () => {\n            // abort span if not recording\n            if (!span.isRecording()) return;\n            isAborted = true;\n            span.setAttribute(\"cancelled\", true);\n            span.setStatus({\n              code: SpanStatusCode.ERROR,\n              message: \"Query cancelled by client\",\n            });\n            span.end();\n          };\n          signal.addEventListener(\"abort\", abortHandler, { once: true });\n        }\n\n        try {\n          // validate required fields\n          if (!input.statement) {\n            throw ValidationError.missingField(\"statement\");\n          }\n\n          if (!input.warehouse_id) {\n            throw ValidationError.missingField(\"warehouse_id\");\n          }\n\n          const body: sql.ExecuteStatementRequest = {\n            statement: input.statement,\n            parameters: input.parameters,\n            warehouse_id: input.warehouse_id,\n            catalog: input.catalog,\n            schema: input.schema,\n            wait_timeout:\n              input.wait_timeout || executeStatementDefaults.wait_timeout,\n            disposition:\n              input.disposition || executeStatementDefaults.disposition,\n            format: input.format || executeStatementDefaults.format,\n            byte_limit: input.byte_limit,\n            row_limit: input.row_limit,\n            on_wait_timeout:\n              input.on_wait_timeout || executeStatementDefaults.on_wait_timeout,\n          };\n\n          span.addEvent(\"statement.submitting\", {\n            \"db.warehouse_id\": input.warehouse_id,\n          });\n\n          const response =\n            await workspaceClient.statementExecution.executeStatement(\n              body,\n              this._createContext(signal),\n            );\n\n          if (!response) {\n            throw ConnectionError.apiFailure(\"SQL Warehouse\");\n          }\n          const status = response.status;\n          const statementId = response.statement_id as string;\n\n          span.setAttribute(\"db.statement_id\", statementId);\n          span.addEvent(\"statement.submitted\", {\n            \"db.statement_id\": response.statement_id,\n            \"db.status\": status?.state,\n          });\n\n          let result:\n            | sql.StatementResponse\n            | { result: { statement_id: string; status: sql.StatementStatus } };\n\n          switch (status?.state) {\n            case \"RUNNING\":\n            case \"PENDING\":\n              span.addEvent(\"statement.polling_started\", {\n                \"db.status\": response.status?.state,\n              });\n              result = await this._pollForStatementResult(\n                workspaceClient,\n                statementId,\n                this.config.timeout,\n                signal,\n              );\n              break;\n            case \"SUCCEEDED\":\n              result = this._transformDataArray(response);\n              break;\n            case \"FAILED\":\n              throw ExecutionError.statementFailed(status.error?.message);\n            case \"CANCELED\":\n              throw ExecutionError.canceled();\n            case \"CLOSED\":\n              throw ExecutionError.resultsClosed();\n            default:\n              throw ExecutionError.unknownState(\n                String(status?.state ?? \"unknown\"),\n              );\n          }\n\n          const resultData = result.result as any;\n          const rowCount =\n            resultData?.data?.length ?? resultData?.data_array?.length ?? 0;\n\n          if (rowCount > 0) {\n            span.setAttribute(\"db.result.row_count\", rowCount);\n          }\n\n          const duration = Date.now() - startTime;\n          logger.event()?.setContext(\"sql-warehouse\", {\n            warehouse_id: input.warehouse_id,\n            rows_returned: rowCount,\n            query_duration_ms: duration,\n          });\n\n          success = true;\n          // only set success status if not aborted\n          if (!isAborted) {\n            span.setStatus({ code: SpanStatusCode.OK });\n          }\n          return result;\n        } catch (error) {\n          // only record error if not already handled by abort\n          if (!isAborted) {\n            span.recordException(error as Error);\n            span.setStatus({\n              code: SpanStatusCode.ERROR,\n              message: error instanceof Error ? error.message : String(error),\n            });\n          }\n\n          if (error instanceof AppKitError) {\n            throw error;\n          }\n          throw ExecutionError.statementFailed(\n            error instanceof Error ? error.message : String(error),\n          );\n        } finally {\n          // remove abort handler\n          if (abortHandler && signal) {\n            signal.removeEventListener(\"abort\", abortHandler);\n          }\n\n          const duration = Date.now() - startTime;\n\n          // end span if not already ended by abort handler\n          if (!isAborted) {\n            span.end();\n          }\n\n          const attributes = {\n            \"db.warehouse_id\": input.warehouse_id,\n            \"db.catalog\": input.catalog ?? \"\",\n            \"db.schema\": input.schema ?? \"\",\n            \"db.statement\": input.statement?.substring(0, 500) || \"\",\n            success: success.toString(),\n          };\n\n          this.telemetryMetrics.queryCount.add(1, attributes);\n          this.telemetryMetrics.queryDuration.record(duration, attributes);\n        }\n      },\n      { name: this.name, includePrefix: true },\n    );\n  }\n\n  private async _pollForStatementResult(\n    workspaceClient: WorkspaceClient,\n    statementId: string,\n    timeout = executeStatementDefaults.timeout,\n    signal?: AbortSignal,\n  ) {\n    return this.telemetry.startActiveSpan(\n      \"sql.poll\",\n      {\n        attributes: {\n          \"db.statement_id\": statementId,\n          \"db.polling.timeout\": timeout,\n        },\n      },\n      async (span: Span) => {\n        try {\n          const startTime = Date.now();\n          let delay = 1000;\n          const maxDelayBetweenPolls = 5000; // max 5 seconds between polls\n          let pollCount = 0;\n\n          while (true) {\n            pollCount++;\n            span.setAttribute(\"db.polling.current_attempt\", pollCount);\n\n            // check if timeout exceeded\n            const elapsedTime = Date.now() - startTime;\n            if (elapsedTime > timeout) {\n              const error = ExecutionError.statementFailed(\n                `Polling timeout exceeded after ${timeout}ms (elapsed: ${elapsedTime}ms)`,\n              );\n              span.recordException(error);\n              span.setStatus({ code: SpanStatusCode.ERROR });\n              throw error;\n            }\n\n            if (signal?.aborted) {\n              const error = ExecutionError.canceled();\n              span.recordException(error);\n              span.setStatus({ code: SpanStatusCode.ERROR });\n              throw error;\n            }\n\n            span.addEvent(\"polling.attempt\", {\n              \"poll.attempt\": pollCount,\n              \"poll.delay_ms\": delay,\n              \"poll.elapsed_ms\": elapsedTime,\n            });\n\n            const response =\n              await workspaceClient.statementExecution.getStatement(\n                {\n                  statement_id: statementId,\n                },\n                this._createContext(signal),\n              );\n            if (!response) {\n              throw ConnectionError.apiFailure(\"SQL Warehouse\");\n            }\n\n            const status = response.status;\n\n            span.addEvent(\"polling.status_check\", {\n              \"db.status\": status?.state,\n              \"poll.attempt\": pollCount,\n            });\n\n            switch (status?.state) {\n              case \"PENDING\":\n              case \"RUNNING\":\n                // continue polling\n                break;\n              case \"SUCCEEDED\":\n                span.setAttribute(\"db.polling.attempts\", pollCount);\n                span.setAttribute(\"db.polling.total_duration_ms\", elapsedTime);\n                span.addEvent(\"polling.completed\", {\n                  \"poll.attempts\": pollCount,\n                  \"poll.duration_ms\": elapsedTime,\n                });\n                span.setStatus({ code: SpanStatusCode.OK });\n                return this._transformDataArray(response);\n              case \"FAILED\":\n                throw ExecutionError.statementFailed(status.error?.message);\n              case \"CANCELED\":\n                throw ExecutionError.canceled();\n              case \"CLOSED\":\n                throw ExecutionError.resultsClosed();\n              default:\n                throw ExecutionError.unknownState(\n                  String(status?.state ?? \"unknown\"),\n                );\n            }\n\n            // continue polling after delay\n            await new Promise((resolve) => setTimeout(resolve, delay));\n            delay = Math.min(delay * 2, maxDelayBetweenPolls);\n          }\n        } catch (error) {\n          span.recordException(error as Error);\n          span.setStatus({\n            code: SpanStatusCode.ERROR,\n            message: error instanceof Error ? error.message : String(error),\n          });\n\n          if (error instanceof AppKitError) {\n            throw error;\n          }\n          throw ExecutionError.statementFailed(\n            error instanceof Error ? error.message : String(error),\n          );\n        } finally {\n          span.end();\n        }\n      },\n      { name: this.name, includePrefix: true },\n    );\n  }\n\n  private _transformDataArray(response: sql.StatementResponse) {\n    if (response.manifest?.format === \"ARROW_STREAM\") {\n      return this.updateWithArrowStatus(response);\n    }\n\n    if (!response.result?.data_array || !response.manifest?.schema?.columns) {\n      return response;\n    }\n\n    const columns = response.manifest.schema.columns;\n\n    const transformedData = response.result.data_array.map((row) => {\n      const obj: Record<string, unknown> = {};\n      row.forEach((value, index) => {\n        const column = columns[index];\n        const columnName = column?.name || `column_${index}`;\n\n        // attempt to parse JSON strings for string columns\n        if (\n          column?.type_name === \"STRING\" &&\n          typeof value === \"string\" &&\n          value &&\n          (value[0] === \"{\" || value[0] === \"[\")\n        ) {\n          try {\n            obj[columnName] = JSON.parse(value);\n          } catch {\n            // if parsing fails, keep as string\n            obj[columnName] = value;\n          }\n        } else {\n          obj[columnName] = value;\n        }\n      });\n      return obj;\n    });\n\n    // remove data_array\n    const { data_array: _data_array, ...restResult } = response.result;\n    return {\n      ...response,\n      result: {\n        ...restResult,\n        data: transformedData,\n      },\n    };\n  }\n\n  private updateWithArrowStatus(response: sql.StatementResponse): {\n    result: { statement_id: string; status: sql.StatementStatus };\n  } {\n    return {\n      result: {\n        statement_id: response.statement_id as string,\n        status: {\n          state: response.status?.state,\n          error: response.status?.error,\n        } as sql.StatementStatus,\n      },\n    };\n  }\n\n  async getArrowData(\n    workspaceClient: WorkspaceClient,\n    jobId: string,\n    signal?: AbortSignal,\n  ): Promise<ReturnType<typeof this.arrowProcessor.processChunks>> {\n    const startTime = Date.now();\n\n    return this.telemetry.startActiveSpan(\n      \"arrow.getData\",\n      {\n        kind: SpanKind.CLIENT,\n        attributes: {\n          \"db.system\": \"databricks\",\n          \"arrow.job_id\": jobId,\n        },\n      },\n      async (span: Span) => {\n        try {\n          const response =\n            await workspaceClient.statementExecution.getStatement(\n              { statement_id: jobId },\n              this._createContext(signal),\n            );\n\n          const chunks = response.result?.external_links;\n          const schema = response.manifest?.schema;\n\n          if (!chunks || !schema) {\n            throw ExecutionError.missingData(\"chunks or schema\");\n          }\n\n          span.setAttribute(\"arrow.chunk_count\", chunks.length);\n\n          const result = await this.arrowProcessor.processChunks(\n            chunks,\n            schema,\n            signal,\n          );\n\n          span.setAttribute(\"arrow.data_size_bytes\", result.data.length);\n          span.setStatus({ code: SpanStatusCode.OK });\n\n          const duration = Date.now() - startTime;\n          this.telemetryMetrics.queryDuration.record(duration, {\n            operation: \"arrow.getData\",\n            status: \"success\",\n          });\n\n          logger.event()?.setContext(\"sql-warehouse\", {\n            arrow_data_size_bytes: result.data.length,\n            arrow_job_id: jobId,\n          });\n\n          return result;\n        } catch (error) {\n          span.setStatus({\n            code: SpanStatusCode.ERROR,\n            message: error instanceof Error ? error.message : \"Unknown error\",\n          });\n          span.recordException(error as Error);\n\n          const duration = Date.now() - startTime;\n          this.telemetryMetrics.queryDuration.record(duration, {\n            operation: \"arrow.getData\",\n            status: \"error\",\n          });\n\n          logger.error(\"Failed Arrow job: %s %O\", jobId, error);\n\n          if (error instanceof AppKitError) {\n            throw error;\n          }\n          throw ExecutionError.statementFailed(\n            error instanceof Error ? error.message : String(error),\n          );\n        }\n      },\n    );\n  }\n\n  // create context for cancellation token\n  private _createContext(signal?: AbortSignal) {\n    return new Context({\n      cancellationToken: {\n        isCancellationRequested: signal?.aborted ?? false,\n        onCancellationRequested: (cb: () => void) => {\n          signal?.addEventListener(\"abort\", cb, { once: true });\n        },\n      },\n    });\n  }\n}\n"],"mappings":";;;;;;;;;;;;;aAWsB;AActB,MAAM,SAAS,aAAa,2BAA2B;AAOvD,IAAa,wBAAb,MAAmC;CAcjC,YAAY,QAA4B;cAbhB;yBAK+B;AASrD,OAAK,SAAS;AAEd,OAAK,YAAY,iBAAiB,YAChC,KAAK,MACL,KAAK,OAAO,UACb;AACD,OAAK,mBAAmB;GACtB,YAAY,KAAK,UAAU,UAAU,CAAC,cAAc,eAAe;IACjE,aAAa;IACb,MAAM;IACP,CAAC;GACF,eAAe,KAAK,UACjB,UAAU,CACV,gBAAgB,kBAAkB;IACjC,aAAa;IACb,MAAM;IACP,CAAC;GACL;;;;;;CAOH,IAAY,iBAAuC;AACjD,MAAI,CAAC,KAAK,gBACR,MAAK,kBAAkB,IAAI,qBAAqB;GAC9C,SAAS,KAAK,OAAO,WAAW,yBAAyB;GACzD,wBACE,qBAAqB;GACvB,SAAS,qBAAqB;GAC/B,CAAC;AAEJ,SAAO,KAAK;;CAGd,MAAM,iBACJ,iBACA,OACA,QACA;EACA,MAAM,YAAY,KAAK,KAAK;EAC5B,IAAI,UAAU;AAGd,MAAI,QAAQ,QACV,OAAM,eAAe,UAAU;AAGjC,SAAO,KAAK,UAAU,gBACpB,aACA;GACE,MAAM,SAAS;GACf,YAAY;IACV,aAAa;IACb,mBAAmB,MAAM,gBAAgB;IACzC,cAAc,MAAM,WAAW;IAC/B,aAAa,MAAM,UAAU;IAC7B,gBAAgB,MAAM,WAAW,UAAU,GAAG,IAAI,IAAI;IACtD,qBAAqB,CAAC,CAAC,MAAM;IAC9B;GACF,EACD,OAAO,SAAe;GACpB,IAAI;GACJ,IAAI,YAAY;AAEhB,OAAI,QAAQ;AACV,yBAAqB;AAEnB,SAAI,CAAC,KAAK,aAAa,CAAE;AACzB,iBAAY;AACZ,UAAK,aAAa,aAAa,KAAK;AACpC,UAAK,UAAU;MACb,MAAM,eAAe;MACrB,SAAS;MACV,CAAC;AACF,UAAK,KAAK;;AAEZ,WAAO,iBAAiB,SAAS,cAAc,EAAE,MAAM,MAAM,CAAC;;AAGhE,OAAI;AAEF,QAAI,CAAC,MAAM,UACT,OAAM,gBAAgB,aAAa,YAAY;AAGjD,QAAI,CAAC,MAAM,aACT,OAAM,gBAAgB,aAAa,eAAe;IAGpD,MAAM,OAAoC;KACxC,WAAW,MAAM;KACjB,YAAY,MAAM;KAClB,cAAc,MAAM;KACpB,SAAS,MAAM;KACf,QAAQ,MAAM;KACd,cACE,MAAM,gBAAgB,yBAAyB;KACjD,aACE,MAAM,eAAe,yBAAyB;KAChD,QAAQ,MAAM,UAAU,yBAAyB;KACjD,YAAY,MAAM;KAClB,WAAW,MAAM;KACjB,iBACE,MAAM,mBAAmB,yBAAyB;KACrD;AAED,SAAK,SAAS,wBAAwB,EACpC,mBAAmB,MAAM,cAC1B,CAAC;IAEF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,iBACvC,MACA,KAAK,eAAe,OAAO,CAC5B;AAEH,QAAI,CAAC,SACH,OAAM,gBAAgB,WAAW,gBAAgB;IAEnD,MAAM,SAAS,SAAS;IACxB,MAAM,cAAc,SAAS;AAE7B,SAAK,aAAa,mBAAmB,YAAY;AACjD,SAAK,SAAS,uBAAuB;KACnC,mBAAmB,SAAS;KAC5B,aAAa,QAAQ;KACtB,CAAC;IAEF,IAAI;AAIJ,YAAQ,QAAQ,OAAhB;KACE,KAAK;KACL,KAAK;AACH,WAAK,SAAS,6BAA6B,EACzC,aAAa,SAAS,QAAQ,OAC/B,CAAC;AACF,eAAS,MAAM,KAAK,wBAClB,iBACA,aACA,KAAK,OAAO,SACZ,OACD;AACD;KACF,KAAK;AACH,eAAS,KAAK,oBAAoB,SAAS;AAC3C;KACF,KAAK,SACH,OAAM,eAAe,gBAAgB,OAAO,OAAO,QAAQ;KAC7D,KAAK,WACH,OAAM,eAAe,UAAU;KACjC,KAAK,SACH,OAAM,eAAe,eAAe;KACtC,QACE,OAAM,eAAe,aACnB,OAAO,QAAQ,SAAS,UAAU,CACnC;;IAGL,MAAM,aAAa,OAAO;IAC1B,MAAM,WACJ,YAAY,MAAM,UAAU,YAAY,YAAY,UAAU;AAEhE,QAAI,WAAW,EACb,MAAK,aAAa,uBAAuB,SAAS;IAGpD,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,WAAO,OAAO,EAAE,WAAW,iBAAiB;KAC1C,cAAc,MAAM;KACpB,eAAe;KACf,mBAAmB;KACpB,CAAC;AAEF,cAAU;AAEV,QAAI,CAAC,UACH,MAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAE7C,WAAO;YACA,OAAO;AAEd,QAAI,CAAC,WAAW;AACd,UAAK,gBAAgB,MAAe;AACpC,UAAK,UAAU;MACb,MAAM,eAAe;MACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM;MAChE,CAAC;;AAGJ,QAAI,iBAAiB,YACnB,OAAM;AAER,UAAM,eAAe,gBACnB,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM,CACvD;aACO;AAER,QAAI,gBAAgB,OAClB,QAAO,oBAAoB,SAAS,aAAa;IAGnD,MAAM,WAAW,KAAK,KAAK,GAAG;AAG9B,QAAI,CAAC,UACH,MAAK,KAAK;IAGZ,MAAM,aAAa;KACjB,mBAAmB,MAAM;KACzB,cAAc,MAAM,WAAW;KAC/B,aAAa,MAAM,UAAU;KAC7B,gBAAgB,MAAM,WAAW,UAAU,GAAG,IAAI,IAAI;KACtD,SAAS,QAAQ,UAAU;KAC5B;AAED,SAAK,iBAAiB,WAAW,IAAI,GAAG,WAAW;AACnD,SAAK,iBAAiB,cAAc,OAAO,UAAU,WAAW;;KAGpE;GAAE,MAAM,KAAK;GAAM,eAAe;GAAM,CACzC;;CAGH,MAAc,wBACZ,iBACA,aACA,UAAU,yBAAyB,SACnC,QACA;AACA,SAAO,KAAK,UAAU,gBACpB,YACA,EACE,YAAY;GACV,mBAAmB;GACnB,sBAAsB;GACvB,EACF,EACD,OAAO,SAAe;AACpB,OAAI;IACF,MAAM,YAAY,KAAK,KAAK;IAC5B,IAAI,QAAQ;IACZ,MAAM,uBAAuB;IAC7B,IAAI,YAAY;AAEhB,WAAO,MAAM;AACX;AACA,UAAK,aAAa,8BAA8B,UAAU;KAG1D,MAAM,cAAc,KAAK,KAAK,GAAG;AACjC,SAAI,cAAc,SAAS;MACzB,MAAM,QAAQ,eAAe,gBAC3B,kCAAkC,QAAQ,eAAe,YAAY,KACtE;AACD,WAAK,gBAAgB,MAAM;AAC3B,WAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,YAAM;;AAGR,SAAI,QAAQ,SAAS;MACnB,MAAM,QAAQ,eAAe,UAAU;AACvC,WAAK,gBAAgB,MAAM;AAC3B,WAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,YAAM;;AAGR,UAAK,SAAS,mBAAmB;MAC/B,gBAAgB;MAChB,iBAAiB;MACjB,mBAAmB;MACpB,CAAC;KAEF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,aACvC,EACE,cAAc,aACf,EACD,KAAK,eAAe,OAAO,CAC5B;AACH,SAAI,CAAC,SACH,OAAM,gBAAgB,WAAW,gBAAgB;KAGnD,MAAM,SAAS,SAAS;AAExB,UAAK,SAAS,wBAAwB;MACpC,aAAa,QAAQ;MACrB,gBAAgB;MACjB,CAAC;AAEF,aAAQ,QAAQ,OAAhB;MACE,KAAK;MACL,KAAK,UAEH;MACF,KAAK;AACH,YAAK,aAAa,uBAAuB,UAAU;AACnD,YAAK,aAAa,gCAAgC,YAAY;AAC9D,YAAK,SAAS,qBAAqB;QACjC,iBAAiB;QACjB,oBAAoB;QACrB,CAAC;AACF,YAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,cAAO,KAAK,oBAAoB,SAAS;MAC3C,KAAK,SACH,OAAM,eAAe,gBAAgB,OAAO,OAAO,QAAQ;MAC7D,KAAK,WACH,OAAM,eAAe,UAAU;MACjC,KAAK,SACH,OAAM,eAAe,eAAe;MACtC,QACE,OAAM,eAAe,aACnB,OAAO,QAAQ,SAAS,UAAU,CACnC;;AAIL,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,MAAM,CAAC;AAC1D,aAAQ,KAAK,IAAI,QAAQ,GAAG,qBAAqB;;YAE5C,OAAO;AACd,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU;KACb,MAAM,eAAe;KACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM;KAChE,CAAC;AAEF,QAAI,iBAAiB,YACnB,OAAM;AAER,UAAM,eAAe,gBACnB,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM,CACvD;aACO;AACR,SAAK,KAAK;;KAGd;GAAE,MAAM,KAAK;GAAM,eAAe;GAAM,CACzC;;CAGH,AAAQ,oBAAoB,UAAiC;AAC3D,MAAI,SAAS,UAAU,WAAW,eAChC,QAAO,KAAK,sBAAsB,SAAS;AAG7C,MAAI,CAAC,SAAS,QAAQ,cAAc,CAAC,SAAS,UAAU,QAAQ,QAC9D,QAAO;EAGT,MAAM,UAAU,SAAS,SAAS,OAAO;EAEzC,MAAM,kBAAkB,SAAS,OAAO,WAAW,KAAK,QAAQ;GAC9D,MAAM,MAA+B,EAAE;AACvC,OAAI,SAAS,OAAO,UAAU;IAC5B,MAAM,SAAS,QAAQ;IACvB,MAAM,aAAa,QAAQ,QAAQ,UAAU;AAG7C,QACE,QAAQ,cAAc,YACtB,OAAO,UAAU,YACjB,UACC,MAAM,OAAO,OAAO,MAAM,OAAO,KAElC,KAAI;AACF,SAAI,cAAc,KAAK,MAAM,MAAM;YAC7B;AAEN,SAAI,cAAc;;QAGpB,KAAI,cAAc;KAEpB;AACF,UAAO;IACP;EAGF,MAAM,EAAE,YAAY,aAAa,GAAG,eAAe,SAAS;AAC5D,SAAO;GACL,GAAG;GACH,QAAQ;IACN,GAAG;IACH,MAAM;IACP;GACF;;CAGH,AAAQ,sBAAsB,UAE5B;AACA,SAAO,EACL,QAAQ;GACN,cAAc,SAAS;GACvB,QAAQ;IACN,OAAO,SAAS,QAAQ;IACxB,OAAO,SAAS,QAAQ;IACzB;GACF,EACF;;CAGH,MAAM,aACJ,iBACA,OACA,QAC+D;EAC/D,MAAM,YAAY,KAAK,KAAK;AAE5B,SAAO,KAAK,UAAU,gBACpB,iBACA;GACE,MAAM,SAAS;GACf,YAAY;IACV,aAAa;IACb,gBAAgB;IACjB;GACF,EACD,OAAO,SAAe;AACpB,OAAI;IACF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,aACvC,EAAE,cAAc,OAAO,EACvB,KAAK,eAAe,OAAO,CAC5B;IAEH,MAAM,SAAS,SAAS,QAAQ;IAChC,MAAM,SAAS,SAAS,UAAU;AAElC,QAAI,CAAC,UAAU,CAAC,OACd,OAAM,eAAe,YAAY,mBAAmB;AAGtD,SAAK,aAAa,qBAAqB,OAAO,OAAO;IAErD,MAAM,SAAS,MAAM,KAAK,eAAe,cACvC,QACA,QACA,OACD;AAED,SAAK,aAAa,yBAAyB,OAAO,KAAK,OAAO;AAC9D,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;IAE3C,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,cAAc,OAAO,UAAU;KACnD,WAAW;KACX,QAAQ;KACT,CAAC;AAEF,WAAO,OAAO,EAAE,WAAW,iBAAiB;KAC1C,uBAAuB,OAAO,KAAK;KACnC,cAAc;KACf,CAAC;AAEF,WAAO;YACA,OAAO;AACd,SAAK,UAAU;KACb,MAAM,eAAe;KACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU;KACnD,CAAC;AACF,SAAK,gBAAgB,MAAe;IAEpC,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,cAAc,OAAO,UAAU;KACnD,WAAW;KACX,QAAQ;KACT,CAAC;AAEF,WAAO,MAAM,2BAA2B,OAAO,MAAM;AAErD,QAAI,iBAAiB,YACnB,OAAM;AAER,UAAM,eAAe,gBACnB,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM,CACvD;;IAGN;;CAIH,AAAQ,eAAe,QAAsB;AAC3C,SAAO,IAAI,QAAQ,EACjB,mBAAmB;GACjB,yBAAyB,QAAQ,WAAW;GAC5C,0BAA0B,OAAmB;AAC3C,YAAQ,iBAAiB,SAAS,IAAI,EAAE,MAAM,MAAM,CAAC;;GAExD,EACF,CAAC"}

package/dist/context/service-context.js

@@ -1,4 +1,8 @@
 import { __esmMin } from "../_virtual/rolldown_runtime.js";
+import { AuthenticationError } from "../errors/authentication.js";
+import { ConfigurationError } from "../errors/configuration.js";
+import { InitializationError } from "../errors/initialization.js";
+import { init_errors } from "../errors/index.js";
 import { name, version } from "../appkit/package.js";
 import { WorkspaceClient } from "@databricks/sdk-experimental";
 import { coerce } from "semver";
@@ -14,6 +18,7 @@ function getClientOptions() {
 }
 var ServiceContext;
 var init_service_context = __esmMin((() => {
+	init_errors();
 	ServiceContext = class ServiceContext {
 		static {
 			this.instance = null;
@@ -37,7 +42,7 @@ var init_service_context = __esmMin((() => {
 		* @throws Error if not initialized
 		*/
 		static get() {
-			if (!ServiceContext.instance) throw new Error("ServiceContext not initialized. Call ServiceContext.initialize() first.");
+			if (!ServiceContext.instance) throw InitializationError.notInitialized("ServiceContext", "Call ServiceContext.initialize() first");
 			return ServiceContext.instance;
 		}
 		/**
@@ -55,9 +60,9 @@ var init_service_context = __esmMin((() => {
 		* @throws Error if token is not provided
 		*/
 		static createUserContext(token, userId, userName) {
-			if (!token) throw new Error("User token is required to create user context");
+			if (!token) throw AuthenticationError.missingToken("user token");
 			const host = process.env.DATABRICKS_HOST;
-			if (!host) throw new Error("DATABRICKS_HOST environment variable is required for user context");
+			if (!host) throw ConfigurationError.missingEnvVar("DATABRICKS_HOST");
 			const serviceCtx = ServiceContext.get();
 			return {
 				client: new WorkspaceClient({
@@ -84,7 +89,7 @@ var init_service_context = __esmMin((() => {
 			const warehouseId = ServiceContext.getWarehouseId(client);
 			const workspaceId = ServiceContext.getWorkspaceId(client);
 			const currentUser = await client.currentUser.me();
-			if (!currentUser.id) throw new Error("Service user ID not found");
+			if (!currentUser.id) throw ConfigurationError.resourceNotFound("Service user ID");
 			return {
 				client,
 				serviceUserId: currentUser.id,
@@ -102,7 +107,7 @@ var init_service_context = __esmMin((() => {
 				query: {},
 				responseHeaders: ["x-databricks-org-id"]
 			});
-			if (!response["x-databricks-org-id"]) throw new Error("Workspace ID not found");
+			if (!response["x-databricks-org-id"]) throw ConfigurationError.resourceNotFound("Workspace ID");
 			return response["x-databricks-org-id"];
 		}
 		static async getWarehouseId(client) {
@@ -126,12 +131,12 @@ var init_service_context = __esmMin((() => {
 				const warehouses = (response.warehouses || []).sort((a, b) => {
 					return priorities[a.state] - priorities[b.state];
 				});
-				if (response.warehouses.length === 0) throw new Error("Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.");
+				if (response.warehouses.length === 0) throw ConfigurationError.resourceNotFound("Warehouse ID", "Please configure the DATABRICKS_WAREHOUSE_ID environment variable");
 				const firstWarehouse = warehouses[0];
-				if (firstWarehouse.state === "DELETED" || firstWarehouse.state === "DELETING" || !firstWarehouse.id) throw new Error("Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.");
+				if (firstWarehouse.state === "DELETED" || firstWarehouse.state === "DELETING" || !firstWarehouse.id) throw ConfigurationError.resourceNotFound("Warehouse ID", "Please configure the DATABRICKS_WAREHOUSE_ID environment variable");
 				return firstWarehouse.id;
 			}
-			throw new Error("Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.");
+			throw ConfigurationError.resourceNotFound("Warehouse ID", "Please configure the DATABRICKS_WAREHOUSE_ID environment variable");
 		}
 		/**
 		* Reset the service context. Only for testing purposes.

package/dist/context/service-context.js.map

@@ -1 +1 @@
-{"version":3,"file":"service-context.js","names":["productName","productVersion"],"sources":["../../src/context/service-context.ts"],"sourcesContent":["import {\n  type ClientOptions,\n  type sql,\n  WorkspaceClient,\n} from \"@databricks/sdk-experimental\";\nimport { coerce } from \"semver\";\nimport {\n  name as productName,\n  version as productVersion,\n} from \"../../package.json\";\nimport type { UserContext } from \"./user-context\";\n\n/**\n * Service context holds the service principal client and shared resources.\n * This is initialized once at app startup and shared across all requests.\n */\nexport interface ServiceContextState {\n  /** WorkspaceClient authenticated as the service principal */\n  client: WorkspaceClient;\n  /** The service principal's user ID */\n  serviceUserId: string;\n  /** Promise that resolves to the warehouse ID */\n  warehouseId: Promise<string>;\n  /** Promise that resolves to the workspace ID */\n  workspaceId: Promise<string>;\n}\n\nfunction getClientOptions(): ClientOptions {\n  const isDev = process.env.NODE_ENV === \"development\";\n  const semver = coerce(productVersion);\n  const normalizedVersion = (semver?.version ??\n    productVersion) as ClientOptions[\"productVersion\"];\n\n  return {\n    product: productName,\n    productVersion: normalizedVersion,\n    ...(isDev && { userAgentExtra: { mode: \"dev\" } }),\n  };\n}\n\n/**\n * ServiceContext is a singleton that manages the service principal's\n * WorkspaceClient and shared resources like warehouse/workspace IDs.\n *\n * It's initialized once at app startup and provides the foundation\n * for both service principal and user context execution.\n */\nexport class ServiceContext {\n  private static instance: ServiceContextState | null = null;\n  private static initPromise: Promise<ServiceContextState> | null = null;\n\n  /**\n   * Initialize the service context. Should be called once at app startup.\n   * Safe to call multiple times - will return the same instance.\n   */\n  static async initialize(): Promise<ServiceContextState> {\n    if (ServiceContext.instance) {\n      return ServiceContext.instance;\n    }\n\n    if (ServiceContext.initPromise) {\n      return ServiceContext.initPromise;\n    }\n\n    ServiceContext.initPromise = ServiceContext.createContext();\n    ServiceContext.instance = await ServiceContext.initPromise;\n    return ServiceContext.instance;\n  }\n\n  /**\n   * Get the initialized service context.\n   * @throws Error if not initialized\n   */\n  static get(): ServiceContextState {\n    if (!ServiceContext.instance) {\n      throw new Error(\n        \"ServiceContext not initialized. Call ServiceContext.initialize() first.\",\n      );\n    }\n    return ServiceContext.instance;\n  }\n\n  /**\n   * Check if the service context has been initialized.\n   */\n  static isInitialized(): boolean {\n    return ServiceContext.instance !== null;\n  }\n\n  /**\n   * Create a user context from request headers.\n   *\n   * @param token - The user's access token from x-forwarded-access-token header\n   * @param userId - The user's ID from x-forwarded-user header\n   * @param userName - Optional user name\n   * @throws Error if token is not provided\n   */\n  static createUserContext(\n    token: string,\n    userId: string,\n    userName?: string,\n  ): UserContext {\n    if (!token) {\n      throw new Error(\"User token is required to create user context\");\n    }\n\n    const host = process.env.DATABRICKS_HOST;\n    if (!host) {\n      throw new Error(\n        \"DATABRICKS_HOST environment variable is required for user context\",\n      );\n    }\n\n    const serviceCtx = ServiceContext.get();\n\n    // Create user client with the OAuth token from Databricks Apps\n    // Note: We use authType: \"pat\" because the token is passed as a Bearer token\n    // just like a PAT, even though it's technically an OAuth token\n    const userClient = new WorkspaceClient(\n      {\n        token,\n        host,\n        authType: \"pat\",\n      },\n      getClientOptions(),\n    );\n\n    return {\n      client: userClient,\n      userId,\n      userName,\n      warehouseId: serviceCtx.warehouseId,\n      workspaceId: serviceCtx.workspaceId,\n      isUserContext: true,\n    };\n  }\n\n  /**\n   * Get the client options for WorkspaceClient.\n   * Exposed for testing purposes.\n   */\n  static getClientOptions(): ClientOptions {\n    return getClientOptions();\n  }\n\n  private static async createContext(): Promise<ServiceContextState> {\n    const client = new WorkspaceClient({}, getClientOptions());\n\n    const warehouseId = ServiceContext.getWarehouseId(client);\n    const workspaceId = ServiceContext.getWorkspaceId(client);\n    const currentUser = await client.currentUser.me();\n\n    if (!currentUser.id) {\n      throw new Error(\"Service user ID not found\");\n    }\n\n    return {\n      client,\n      serviceUserId: currentUser.id,\n      warehouseId,\n      workspaceId,\n    };\n  }\n\n  private static async getWorkspaceId(\n    client: WorkspaceClient,\n  ): Promise<string> {\n    if (process.env.DATABRICKS_WORKSPACE_ID) {\n      return process.env.DATABRICKS_WORKSPACE_ID;\n    }\n\n    const response = (await client.apiClient.request({\n      path: \"/api/2.0/preview/scim/v2/Me\",\n      method: \"GET\",\n      headers: new Headers(),\n      raw: false,\n      query: {},\n      responseHeaders: [\"x-databricks-org-id\"],\n    })) as { \"x-databricks-org-id\": string };\n\n    if (!response[\"x-databricks-org-id\"]) {\n      throw new Error(\"Workspace ID not found\");\n    }\n\n    return response[\"x-databricks-org-id\"];\n  }\n\n  private static async getWarehouseId(\n    client: WorkspaceClient,\n  ): Promise<string> {\n    if (process.env.DATABRICKS_WAREHOUSE_ID) {\n      return process.env.DATABRICKS_WAREHOUSE_ID;\n    }\n\n    if (process.env.NODE_ENV === \"development\") {\n      const response = (await client.apiClient.request({\n        path: \"/api/2.0/sql/warehouses\",\n        method: \"GET\",\n        headers: new Headers(),\n        raw: false,\n        query: {\n          skip_cannot_use: \"true\",\n        },\n      })) as { warehouses: sql.EndpointInfo[] };\n\n      const priorities: Record<sql.State, number> = {\n        RUNNING: 0,\n        STOPPED: 1,\n        STARTING: 2,\n        STOPPING: 3,\n        DELETED: 99,\n        DELETING: 99,\n      };\n\n      const warehouses = (response.warehouses || []).sort((a, b) => {\n        return (\n          priorities[a.state as sql.State] - priorities[b.state as sql.State]\n        );\n      });\n\n      if (response.warehouses.length === 0) {\n        throw new Error(\n          \"Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.\",\n        );\n      }\n\n      const firstWarehouse = warehouses[0];\n      if (\n        firstWarehouse.state === \"DELETED\" ||\n        firstWarehouse.state === \"DELETING\" ||\n        !firstWarehouse.id\n      ) {\n        throw new Error(\n          \"Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.\",\n        );\n      }\n\n      return firstWarehouse.id;\n    }\n\n    throw new Error(\n      \"Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.\",\n    );\n  }\n\n  /**\n   * Reset the service context. Only for testing purposes.\n   */\n  static reset(): void {\n    ServiceContext.instance = null;\n    ServiceContext.initPromise = null;\n  }\n}\n"],"mappings":";;;;;;AA2BA,SAAS,mBAAkC;CACzC,MAAM,QAAQ,QAAQ,IAAI,aAAa;AAKvC,QAAO;EACL,SAASA;EACT,gBANa,OAAOC,QAAe,EACF,WACjCA;EAKA,GAAI,SAAS,EAAE,gBAAgB,EAAE,MAAM,OAAO,EAAE;EACjD;;;;CAUU,iBAAb,MAAa,eAAe;;mBAC4B;;;sBACY;;;;;;EAMlE,aAAa,aAA2C;AACtD,OAAI,eAAe,SACjB,QAAO,eAAe;AAGxB,OAAI,eAAe,YACjB,QAAO,eAAe;AAGxB,kBAAe,cAAc,eAAe,eAAe;AAC3D,kBAAe,WAAW,MAAM,eAAe;AAC/C,UAAO,eAAe;;;;;;EAOxB,OAAO,MAA2B;AAChC,OAAI,CAAC,eAAe,SAClB,OAAM,IAAI,MACR,0EACD;AAEH,UAAO,eAAe;;;;;EAMxB,OAAO,gBAAyB;AAC9B,UAAO,eAAe,aAAa;;;;;;;;;;EAWrC,OAAO,kBACL,OACA,QACA,UACa;AACb,OAAI,CAAC,MACH,OAAM,IAAI,MAAM,gDAAgD;GAGlE,MAAM,OAAO,QAAQ,IAAI;AACzB,OAAI,CAAC,KACH,OAAM,IAAI,MACR,oEACD;GAGH,MAAM,aAAa,eAAe,KAAK;AAcvC,UAAO;IACL,QAViB,IAAI,gBACrB;KACE;KACA;KACA,UAAU;KACX,EACD,kBAAkB,CACnB;IAIC;IACA;IACA,aAAa,WAAW;IACxB,aAAa,WAAW;IACxB,eAAe;IAChB;;;;;;EAOH,OAAO,mBAAkC;AACvC,UAAO,kBAAkB;;EAG3B,aAAqB,gBAA8C;GACjE,MAAM,SAAS,IAAI,gBAAgB,EAAE,EAAE,kBAAkB,CAAC;GAE1D,MAAM,cAAc,eAAe,eAAe,OAAO;GACzD,MAAM,cAAc,eAAe,eAAe,OAAO;GACzD,MAAM,cAAc,MAAM,OAAO,YAAY,IAAI;AAEjD,OAAI,CAAC,YAAY,GACf,OAAM,IAAI,MAAM,4BAA4B;AAG9C,UAAO;IACL;IACA,eAAe,YAAY;IAC3B;IACA;IACD;;EAGH,aAAqB,eACnB,QACiB;AACjB,OAAI,QAAQ,IAAI,wBACd,QAAO,QAAQ,IAAI;GAGrB,MAAM,WAAY,MAAM,OAAO,UAAU,QAAQ;IAC/C,MAAM;IACN,QAAQ;IACR,SAAS,IAAI,SAAS;IACtB,KAAK;IACL,OAAO,EAAE;IACT,iBAAiB,CAAC,sBAAsB;IACzC,CAAC;AAEF,OAAI,CAAC,SAAS,uBACZ,OAAM,IAAI,MAAM,yBAAyB;AAG3C,UAAO,SAAS;;EAGlB,aAAqB,eACnB,QACiB;AACjB,OAAI,QAAQ,IAAI,wBACd,QAAO,QAAQ,IAAI;AAGrB,OAAI,QAAQ,IAAI,aAAa,eAAe;IAC1C,MAAM,WAAY,MAAM,OAAO,UAAU,QAAQ;KAC/C,MAAM;KACN,QAAQ;KACR,SAAS,IAAI,SAAS;KACtB,KAAK;KACL,OAAO,EACL,iBAAiB,QAClB;KACF,CAAC;IAEF,MAAM,aAAwC;KAC5C,SAAS;KACT,SAAS;KACT,UAAU;KACV,UAAU;KACV,SAAS;KACT,UAAU;KACX;IAED,MAAM,cAAc,SAAS,cAAc,EAAE,EAAE,MAAM,GAAG,MAAM;AAC5D,YACE,WAAW,EAAE,SAAsB,WAAW,EAAE;MAElD;AAEF,QAAI,SAAS,WAAW,WAAW,EACjC,OAAM,IAAI,MACR,6FACD;IAGH,MAAM,iBAAiB,WAAW;AAClC,QACE,eAAe,UAAU,aACzB,eAAe,UAAU,cACzB,CAAC,eAAe,GAEhB,OAAM,IAAI,MACR,6FACD;AAGH,WAAO,eAAe;;AAGxB,SAAM,IAAI,MACR,6FACD;;;;;EAMH,OAAO,QAAc;AACnB,kBAAe,WAAW;AAC1B,kBAAe,cAAc"}
+{"version":3,"file":"service-context.js","names":["productName","productVersion"],"sources":["../../src/context/service-context.ts"],"sourcesContent":["import {\n  type ClientOptions,\n  type sql,\n  WorkspaceClient,\n} from \"@databricks/sdk-experimental\";\nimport { coerce } from \"semver\";\nimport {\n  name as productName,\n  version as productVersion,\n} from \"../../package.json\";\nimport {\n  AuthenticationError,\n  ConfigurationError,\n  InitializationError,\n} from \"../errors\";\nimport type { UserContext } from \"./user-context\";\n\n/**\n * Service context holds the service principal client and shared resources.\n * This is initialized once at app startup and shared across all requests.\n */\nexport interface ServiceContextState {\n  /** WorkspaceClient authenticated as the service principal */\n  client: WorkspaceClient;\n  /** The service principal's user ID */\n  serviceUserId: string;\n  /** Promise that resolves to the warehouse ID */\n  warehouseId: Promise<string>;\n  /** Promise that resolves to the workspace ID */\n  workspaceId: Promise<string>;\n}\n\nfunction getClientOptions(): ClientOptions {\n  const isDev = process.env.NODE_ENV === \"development\";\n  const semver = coerce(productVersion);\n  const normalizedVersion = (semver?.version ??\n    productVersion) as ClientOptions[\"productVersion\"];\n\n  return {\n    product: productName,\n    productVersion: normalizedVersion,\n    ...(isDev && { userAgentExtra: { mode: \"dev\" } }),\n  };\n}\n\n/**\n * ServiceContext is a singleton that manages the service principal's\n * WorkspaceClient and shared resources like warehouse/workspace IDs.\n *\n * It's initialized once at app startup and provides the foundation\n * for both service principal and user context execution.\n */\nexport class ServiceContext {\n  private static instance: ServiceContextState | null = null;\n  private static initPromise: Promise<ServiceContextState> | null = null;\n\n  /**\n   * Initialize the service context. Should be called once at app startup.\n   * Safe to call multiple times - will return the same instance.\n   */\n  static async initialize(): Promise<ServiceContextState> {\n    if (ServiceContext.instance) {\n      return ServiceContext.instance;\n    }\n\n    if (ServiceContext.initPromise) {\n      return ServiceContext.initPromise;\n    }\n\n    ServiceContext.initPromise = ServiceContext.createContext();\n    ServiceContext.instance = await ServiceContext.initPromise;\n    return ServiceContext.instance;\n  }\n\n  /**\n   * Get the initialized service context.\n   * @throws Error if not initialized\n   */\n  static get(): ServiceContextState {\n    if (!ServiceContext.instance) {\n      throw InitializationError.notInitialized(\n        \"ServiceContext\",\n        \"Call ServiceContext.initialize() first\",\n      );\n    }\n    return ServiceContext.instance;\n  }\n\n  /**\n   * Check if the service context has been initialized.\n   */\n  static isInitialized(): boolean {\n    return ServiceContext.instance !== null;\n  }\n\n  /**\n   * Create a user context from request headers.\n   *\n   * @param token - The user's access token from x-forwarded-access-token header\n   * @param userId - The user's ID from x-forwarded-user header\n   * @param userName - Optional user name\n   * @throws Error if token is not provided\n   */\n  static createUserContext(\n    token: string,\n    userId: string,\n    userName?: string,\n  ): UserContext {\n    if (!token) {\n      throw AuthenticationError.missingToken(\"user token\");\n    }\n\n    const host = process.env.DATABRICKS_HOST;\n    if (!host) {\n      throw ConfigurationError.missingEnvVar(\"DATABRICKS_HOST\");\n    }\n\n    const serviceCtx = ServiceContext.get();\n\n    // Create user client with the OAuth token from Databricks Apps\n    // Note: We use authType: \"pat\" because the token is passed as a Bearer token\n    // just like a PAT, even though it's technically an OAuth token\n    const userClient = new WorkspaceClient(\n      {\n        token,\n        host,\n        authType: \"pat\",\n      },\n      getClientOptions(),\n    );\n\n    return {\n      client: userClient,\n      userId,\n      userName,\n      warehouseId: serviceCtx.warehouseId,\n      workspaceId: serviceCtx.workspaceId,\n      isUserContext: true,\n    };\n  }\n\n  /**\n   * Get the client options for WorkspaceClient.\n   * Exposed for testing purposes.\n   */\n  static getClientOptions(): ClientOptions {\n    return getClientOptions();\n  }\n\n  private static async createContext(): Promise<ServiceContextState> {\n    const client = new WorkspaceClient({}, getClientOptions());\n\n    const warehouseId = ServiceContext.getWarehouseId(client);\n    const workspaceId = ServiceContext.getWorkspaceId(client);\n    const currentUser = await client.currentUser.me();\n\n    if (!currentUser.id) {\n      throw ConfigurationError.resourceNotFound(\"Service user ID\");\n    }\n\n    return {\n      client,\n      serviceUserId: currentUser.id,\n      warehouseId,\n      workspaceId,\n    };\n  }\n\n  private static async getWorkspaceId(\n    client: WorkspaceClient,\n  ): Promise<string> {\n    if (process.env.DATABRICKS_WORKSPACE_ID) {\n      return process.env.DATABRICKS_WORKSPACE_ID;\n    }\n\n    const response = (await client.apiClient.request({\n      path: \"/api/2.0/preview/scim/v2/Me\",\n      method: \"GET\",\n      headers: new Headers(),\n      raw: false,\n      query: {},\n      responseHeaders: [\"x-databricks-org-id\"],\n    })) as { \"x-databricks-org-id\": string };\n\n    if (!response[\"x-databricks-org-id\"]) {\n      throw ConfigurationError.resourceNotFound(\"Workspace ID\");\n    }\n\n    return response[\"x-databricks-org-id\"];\n  }\n\n  private static async getWarehouseId(\n    client: WorkspaceClient,\n  ): Promise<string> {\n    if (process.env.DATABRICKS_WAREHOUSE_ID) {\n      return process.env.DATABRICKS_WAREHOUSE_ID;\n    }\n\n    if (process.env.NODE_ENV === \"development\") {\n      const response = (await client.apiClient.request({\n        path: \"/api/2.0/sql/warehouses\",\n        method: \"GET\",\n        headers: new Headers(),\n        raw: false,\n        query: {\n          skip_cannot_use: \"true\",\n        },\n      })) as { warehouses: sql.EndpointInfo[] };\n\n      const priorities: Record<sql.State, number> = {\n        RUNNING: 0,\n        STOPPED: 1,\n        STARTING: 2,\n        STOPPING: 3,\n        DELETED: 99,\n        DELETING: 99,\n      };\n\n      const warehouses = (response.warehouses || []).sort((a, b) => {\n        return (\n          priorities[a.state as sql.State] - priorities[b.state as sql.State]\n        );\n      });\n\n      if (response.warehouses.length === 0) {\n        throw ConfigurationError.resourceNotFound(\n          \"Warehouse ID\",\n          \"Please configure the DATABRICKS_WAREHOUSE_ID environment variable\",\n        );\n      }\n\n      const firstWarehouse = warehouses[0];\n      if (\n        firstWarehouse.state === \"DELETED\" ||\n        firstWarehouse.state === \"DELETING\" ||\n        !firstWarehouse.id\n      ) {\n        throw ConfigurationError.resourceNotFound(\n          \"Warehouse ID\",\n          \"Please configure the DATABRICKS_WAREHOUSE_ID environment variable\",\n        );\n      }\n\n      return firstWarehouse.id;\n    }\n\n    throw ConfigurationError.resourceNotFound(\n      \"Warehouse ID\",\n      \"Please configure the DATABRICKS_WAREHOUSE_ID environment variable\",\n    );\n  }\n\n  /**\n   * Reset the service context. Only for testing purposes.\n   */\n  static reset(): void {\n    ServiceContext.instance = null;\n    ServiceContext.initPromise = null;\n  }\n}\n"],"mappings":";;;;;;;;;;AAgCA,SAAS,mBAAkC;CACzC,MAAM,QAAQ,QAAQ,IAAI,aAAa;AAKvC,QAAO;EACL,SAASA;EACT,gBANa,OAAOC,QAAe,EACF,WACjCA;EAKA,GAAI,SAAS,EAAE,gBAAgB,EAAE,MAAM,OAAO,EAAE;EACjD;;;;cA5BgB;CAsCN,iBAAb,MAAa,eAAe;;mBAC4B;;;sBACY;;;;;;EAMlE,aAAa,aAA2C;AACtD,OAAI,eAAe,SACjB,QAAO,eAAe;AAGxB,OAAI,eAAe,YACjB,QAAO,eAAe;AAGxB,kBAAe,cAAc,eAAe,eAAe;AAC3D,kBAAe,WAAW,MAAM,eAAe;AAC/C,UAAO,eAAe;;;;;;EAOxB,OAAO,MAA2B;AAChC,OAAI,CAAC,eAAe,SAClB,OAAM,oBAAoB,eACxB,kBACA,yCACD;AAEH,UAAO,eAAe;;;;;EAMxB,OAAO,gBAAyB;AAC9B,UAAO,eAAe,aAAa;;;;;;;;;;EAWrC,OAAO,kBACL,OACA,QACA,UACa;AACb,OAAI,CAAC,MACH,OAAM,oBAAoB,aAAa,aAAa;GAGtD,MAAM,OAAO,QAAQ,IAAI;AACzB,OAAI,CAAC,KACH,OAAM,mBAAmB,cAAc,kBAAkB;GAG3D,MAAM,aAAa,eAAe,KAAK;AAcvC,UAAO;IACL,QAViB,IAAI,gBACrB;KACE;KACA;KACA,UAAU;KACX,EACD,kBAAkB,CACnB;IAIC;IACA;IACA,aAAa,WAAW;IACxB,aAAa,WAAW;IACxB,eAAe;IAChB;;;;;;EAOH,OAAO,mBAAkC;AACvC,UAAO,kBAAkB;;EAG3B,aAAqB,gBAA8C;GACjE,MAAM,SAAS,IAAI,gBAAgB,EAAE,EAAE,kBAAkB,CAAC;GAE1D,MAAM,cAAc,eAAe,eAAe,OAAO;GACzD,MAAM,cAAc,eAAe,eAAe,OAAO;GACzD,MAAM,cAAc,MAAM,OAAO,YAAY,IAAI;AAEjD,OAAI,CAAC,YAAY,GACf,OAAM,mBAAmB,iBAAiB,kBAAkB;AAG9D,UAAO;IACL;IACA,eAAe,YAAY;IAC3B;IACA;IACD;;EAGH,aAAqB,eACnB,QACiB;AACjB,OAAI,QAAQ,IAAI,wBACd,QAAO,QAAQ,IAAI;GAGrB,MAAM,WAAY,MAAM,OAAO,UAAU,QAAQ;IAC/C,MAAM;IACN,QAAQ;IACR,SAAS,IAAI,SAAS;IACtB,KAAK;IACL,OAAO,EAAE;IACT,iBAAiB,CAAC,sBAAsB;IACzC,CAAC;AAEF,OAAI,CAAC,SAAS,uBACZ,OAAM,mBAAmB,iBAAiB,eAAe;AAG3D,UAAO,SAAS;;EAGlB,aAAqB,eACnB,QACiB;AACjB,OAAI,QAAQ,IAAI,wBACd,QAAO,QAAQ,IAAI;AAGrB,OAAI,QAAQ,IAAI,aAAa,eAAe;IAC1C,MAAM,WAAY,MAAM,OAAO,UAAU,QAAQ;KAC/C,MAAM;KACN,QAAQ;KACR,SAAS,IAAI,SAAS;KACtB,KAAK;KACL,OAAO,EACL,iBAAiB,QAClB;KACF,CAAC;IAEF,MAAM,aAAwC;KAC5C,SAAS;KACT,SAAS;KACT,UAAU;KACV,UAAU;KACV,SAAS;KACT,UAAU;KACX;IAED,MAAM,cAAc,SAAS,cAAc,EAAE,EAAE,MAAM,GAAG,MAAM;AAC5D,YACE,WAAW,EAAE,SAAsB,WAAW,EAAE;MAElD;AAEF,QAAI,SAAS,WAAW,WAAW,EACjC,OAAM,mBAAmB,iBACvB,gBACA,oEACD;IAGH,MAAM,iBAAiB,WAAW;AAClC,QACE,eAAe,UAAU,aACzB,eAAe,UAAU,cACzB,CAAC,eAAe,GAEhB,OAAM,mBAAmB,iBACvB,gBACA,oEACD;AAGH,WAAO,eAAe;;AAGxB,SAAM,mBAAmB,iBACvB,gBACA,oEACD;;;;;EAMH,OAAO,QAAc;AACnB,kBAAe,WAAW;AAC1B,kBAAe,cAAc"}

package/dist/errors/authentication.d.ts

@@ -0,0 +1,38 @@
+import { AppKitError } from "./base.js";
+
+//#region src/errors/authentication.d.ts
+
+/**
+ * Error thrown when authentication fails.
+ * Use for missing tokens, invalid credentials, or authorization failures.
+ *
+ * @example
+ * ```typescript
+ * throw new AuthenticationError("User token is required");
+ * throw new AuthenticationError("Failed to generate credentials", { cause: originalError });
+ * ```
+ */
+declare class AuthenticationError extends AppKitError {
+  readonly code = "AUTHENTICATION_ERROR";
+  readonly statusCode = 401;
+  readonly isRetryable = false;
+  /**
+   * Create an authentication error for missing token
+   */
+  static missingToken(tokenType?: string): AuthenticationError;
+  /**
+   * Create an authentication error for missing user identity
+   */
+  static missingUserId(): AuthenticationError;
+  /**
+   * Create an authentication error for credential generation failure
+   */
+  static credentialsFailed(instance: string, cause?: Error): AuthenticationError;
+  /**
+   * Create an authentication error for failed user lookup
+   */
+  static userLookupFailed(cause?: Error): AuthenticationError;
+}
+//#endregion
+export { AuthenticationError };
+//# sourceMappingURL=authentication.d.ts.map

package/dist/errors/authentication.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"authentication.d.ts","names":[],"sources":["../../src/errors/authentication.ts"],"sourcesContent":[],"mappings":";;;;;;AAYA;;;;;;;;AAAyC,cAA5B,mBAAA,SAA4B,WAAA,CAAA;EAAW,SAAA,IAAA,GAAA,sBAAA;;;;;;2CAQD;;;;0BASzB;;;;qDAYd,QACP;;;;kCAU6B,QAAQ"}