databricks_sqlwarehouse.ts

import { DBSQLClient } from "@databricks/sql";
import IDBSQLSession from "@databricks/sql/dist/contracts/IDBSQLSession";

interface DataBricksWareHouseConnect {
    token: string;
    host: string;
    path: string;
}

class DataBricksManager {
    private databricksWarehouseClient: DBSQLClient;
    private warehouseConnOptions: DataBricksWareHouseConnect;
    private warehouseCatalog: string;
    private warehouseSchema: string;

    constructor() {
        this.databricksWarehouseClient = new DBSQLClient();
    }

    public init = async (): Promise<void> => {
        this.warehouseConnOptions = {
            token: process.env.DATABRICKS_PAT,
            host: process.env.SERVER_HOSTNAME,
            path: process.env.HTTP_PATH,
        };
        this.warehouseCatalog = process.env.CATALOG;
        this.warehouseSchema = process.env.SCHEMA;
    };

    public async executeQuery(query: string): Promise<object[]> {
        try {
            await this.databricksWarehouseClient.connect(
                this.warehouseConnOptions
            );
            const session: IDBSQLSession =
                await this.databricksWarehouseClient.openSession();
            await session.executeStatement(
                `USE CATALOG ${this.warehouseCatalog}`
            );
            await session.executeStatement(`USE ${this.warehouseSchema}`);
            const executeOptions = {
                runAsync: true,
                maxRows: 10000, // This option enables the direct results feature.
            };
            const queryOperation = await session.executeStatement(
                query,
                executeOptions
            );
            const result = await queryOperation.fetchAll();
            await queryOperation.close();
            await session.close();
            return result;
        } catch (err) {
            console.error(err);
            throw err;
        }
    }

    public async close(): Promise<void> {
        await this.databricksWarehouseClient.close();
    }
}

export const databricksManager = new DataBricksManager();