cube-msfabric-driver/src/MSFabricDriver.ts

Cube.js MS Fabric Database Driver using msnodesqlv8

import { BaseDriver, DriverInterface } from "@cubejs-backend/base-driver";
import * as sql from "msnodesqlv8";
import { MSFabricQuery } from "./MSFabricQuery";
import { Connection } from "msnodesqlv8/types";

type DatabaseStructure = Record<
  string,
  Record<
    string,
    Array<{
      name: string;
      type: string;
      isNullable: boolean;
    }>
  >
>;

interface ShowTableRow {
  database: string;
  tableName: string;
  isTemporary: boolean;
}

interface ColumnInfo {
  name: string;
  type: string;
  normalizedType: string;
  isNullable: string;
}

// Updated SqlPool interface to match msnodesqlv8 types
interface SqlPool {
  query(sql: string, callback: (err: Error | null, rows: any[]) => void): void;
  query(
    sql: string,
    params: any[],
    callback: (err: Error | null, rows: any[]) => void
  ): void;
  close(): Promise<void>;
}

export interface MSFabricDriverConfig {
  connectionString: string;
  maxPoolSize?: number;
  acquireTimeout?: number;
}

export class MSFabricDriver extends BaseDriver implements DriverInterface {
  private pool: SqlPool | null = null;
  private readonly connectionString: string;
  private readonly maxPoolSize: number;
  private readonly acquireTimeout: number;

  constructor({
    connectionString,
    maxPoolSize = 8,
    acquireTimeout = 30000,
  }: MSFabricDriverConfig) {
    super();
    this.connectionString = connectionString;
    this.maxPoolSize = maxPoolSize;
    this.acquireTimeout = acquireTimeout;
  }

  public static dialectClass() {
    return MSFabricQuery;
  }

  public override async testConnection(): Promise<void> {
    try {
      await this.query("SELECT 1 as test");
      console.log("Connection test successful");
    } catch (error) {
      throw new Error(
        `Connection test failed: ${error instanceof Error ? error.message : String(error)}`
      );
    }
  }

  private async getConnection(): Promise<SqlPool> {
    if (!this.pool) {
      try {
        const connection = await new Promise<Connection>((resolve, reject) => {
          sql.open(this.connectionString, (err, conn: Connection) => {
            if (err) reject(err);
            else resolve(conn);
          });
        });

        // Create adapter object that implements SqlPool interface
        this.pool = {
          query(
            sql: string,
            paramsOrCallback:
              | any[]
              | ((err: Error | null, rows: any[]) => void),
            callback?: (err: Error | null, rows: any[]) => void
          ) {
            if (typeof paramsOrCallback === "function") {
              connection.query(sql, paramsOrCallback as any);
            } else {
              connection.query(sql, paramsOrCallback, callback as any);
            }
          },
          close: async () => {
            return new Promise((resolve) => {
              connection.close(() => resolve());
            });
          },
        };
      } catch (error) {
        throw new Error(
          `Failed to create connection pool: ${error instanceof Error ? error.message : String(error)}`
        );
      }
    }
    return this.pool;
  }

  public async query<T = unknown>(
    query: string,
    values: unknown[] = []
  ): Promise<T[]> {
    const pool = await this.getConnection();

    try {
      const result = await new Promise<T[]>((resolve, reject) => {
        const callback = (err: Error | null, rows: T[]) => {
          if (err) reject(err);
          else resolve(rows);
        };

        if (values.length > 0) {
          pool.query(query, values, callback);
        } else {
          pool.query(query, callback);
        }
      });

      // check if column data is type of date, convert it to iso string, use by moment.js - temporary solution
      // const transform = result.map((row: T) => {
      //   const transformedRow: any = {};
      //   for (const key in row) {
      //     if (Object.prototype.hasOwnProperty.call(row, key)) {
      //       const value = row[key];
      //       if (value instanceof Date) {
      //         transformedRow[key] = value.toISOString();
      //       } else {
      //         transformedRow[key] = value;
      //       }
      //     }
      //   }
      //   return transformedRow;
      // });

      // return transform;
      return result as T[];
    } catch (error) {
      throw new Error(
        `Query failed: ${error instanceof Error ? error.message : String(error)}`
      );
    }
  }

  public override async release(): Promise<void> {
    if (this.pool) {
      try {
        await this.pool.close();
        this.pool = null;
      } catch (error) {
        throw new Error(
          `Failed to close connection pool: ${error instanceof Error ? error.message : String(error)}`
        );
      }
    }
  }

  public override async tablesSchema(): Promise<DatabaseStructure> {
    const schema: DatabaseStructure = {};

    try {
      const tables = await this.query<ShowTableRow>(`
        SELECT 
          TABLE_CATALOG AS [database],
          TABLE_NAME AS tableName,
          CAST(0 as bit) as isTemporary
        FROM INFORMATION_SCHEMA.TABLES
        WHERE TABLE_TYPE = 'BASE TABLE'
      `);

      for (const table of tables) {
        if (!schema[table.database]) {
          schema[table.database] = {};
        }

        // Use parameterized query for security
        const columns = await this.query<ColumnInfo>(
          `SELECT 
            COLUMN_NAME as name,
            DATA_TYPE as type,
            CASE 
              WHEN DATA_TYPE IN ('datetime', 'datetime2', 'smalldatetime', 'datetimeoffset', 'date', 'time')
              THEN 'timestamp'
              ELSE UPPER(DATA_TYPE)
            END as normalizedType,
            IS_NULLABLE as isNullable
          FROM INFORMATION_SCHEMA.COLUMNS
          WHERE TABLE_CATALOG = ?
            AND TABLE_NAME = ?
          ORDER BY ORDINAL_POSITION`,
          [table.database, table.tableName]
        );

        schema[table.database][table.tableName] = columns.map((column) => ({
          name: column.name,
          type: column.normalizedType.toLowerCase(),
          isNullable: column.isNullable === "YES",
        }));
      }

      return schema;
    } catch (error) {
      throw new Error(
        `Failed to get schema: ${error instanceof Error ? error.message : String(error)}`
      );
    }
  }
}