@quillsql/node 0.3.6 → 0.3.8

package/src/db/BigQuery.ts

@@ -0,0 +1,201 @@
+import { Client } from "../models/Client";
+import { BigQuery } from "@google-cloud/bigquery";
+import { QuillQueryResults } from "./DatabaseHelper";
+import { capitalize, depluralize } from "../utils/textProcessing";
+
+export interface BigQueryConfig {
+  datasetName: string;
+  projectId: string;
+  credentials: any;
+}
+
+export function connectToBigQuery(config: BigQueryConfig) {
+  return new BigQuery(config);
+}
+
+export async function runQueryBigQuery(
+  sql: string,
+  bigQuery: BigQuery
+): Promise<QuillQueryResults> {
+  const rows = await bigQuery.query(sql);
+  if (!rows[0] || rows[0].length === 0) return { fields: [], rows: [] };
+  const typedRows = rows[0] as { [fieldName: string]: any }[];
+  const fields = Object.keys(typedRows[0]).map((name: string) => ({
+    name,
+    dataTypeID: 1043,
+  }));
+  fields.forEach((field) => {
+    typedRows.some((row) => {
+      if (row[field.name] === null) return false;
+      field.dataTypeID = inferType(row[field.name]);
+      return true;
+    });
+  });
+  return {
+    fields: fields,
+    rows: typedRows,
+  };
+}
+
+export async function getSchemaBigQuery(bigQuery: BigQuery): Promise<string[]> {
+  const [datasets] = await bigQuery.getDatasets();
+  const definedDatasets = datasets.map((dataset) => dataset.id);
+  const filtered: string[] = [];
+  definedDatasets.forEach((dataset) => {
+    if (dataset !== undefined) {
+      filtered.push(dataset);
+    }
+  });
+  return filtered;
+}
+
+export async function getTablesBySchemaBigQuery(
+  bigQuery: BigQuery,
+  schemaName: string
+): Promise<string[]> {
+  const sql = `SELECT table_name FROM ${schemaName}.INFORMATION_SCHEMA.TABLES WHERE table_type = 'BASE TABLE'`;
+  const rows = await bigQuery.query(sql);
+  return rows[0].map((row: any) => row.table_name);
+}
+
+export async function getColumnsByTableBigQuery(
+  bigQuery: BigQuery,
+  schemaName: string,
+  tableName: string
+): Promise<string[]> {
+  const sql = `SELECT column_name FROM ${schemaName}.INFORMATION_SCHEMA.COLUMNS WHERE table_name = '${tableName}'`;
+  const rows = await bigQuery.query(sql);
+  return rows[0].map((row: any) => row.column_name);
+}
+
+export function formatBigQueryConfig(connectionString: string): BigQueryConfig {
+  const jsonStartIndex = connectionString.indexOf("{");
+  if (jsonStartIndex === -1) {
+    throw new Error("Invalid input string. No JSON data found.");
+  }
+
+  const datasetName = connectionString.substring(0, jsonStartIndex).trim();
+  const jsonString = connectionString.substring(jsonStartIndex);
+
+  try {
+    const serviceAccount = JSON.parse(jsonString);
+
+    // Validate if required fields are present
+    if (!serviceAccount.project_id || !serviceAccount.private_key) {
+      throw new Error(
+        "Invalid service account JSON. Required fields are missing."
+      );
+    }
+
+    return {
+      datasetName,
+      projectId: serviceAccount.project_id,
+      credentials: serviceAccount,
+    };
+  } catch (error) {
+    throw new Error("Failed to parse JSON string: " + error);
+  }
+}
+
+export async function getForeignKeysBigQuery(
+  connection: BigQuery,
+  schemaName: string,
+  tableName: string,
+  primaryKey: string
+): Promise<string[]> {
+  const depluralizedTableName = depluralize(tableName);
+  let sql = `SELECT column_name FROM ${schemaName}.INFORMATION_SCHEMA.COLUMNS 
+    WHERE table_name != '${tableName}' 
+    and (column_name = '${primaryKey}' 
+      or column_name = '${depluralizedTableName}_${primaryKey}' 
+      or column_name = '${depluralizedTableName}${capitalize(primaryKey)}')`;
+  const results = await runQueryBigQuery(sql, connection);
+  let foreignKeysString = results.rows.map((key) => {
+    return key.column_name;
+  });
+  foreignKeysString = foreignKeysString.filter(
+    (key) => key !== "id" && key !== "_id_"
+  );
+  foreignKeysString = [...new Set(foreignKeysString)];
+  if (foreignKeysString.length === 0) {
+    sql = `SELECT column_name FROM ${schemaName}.INFORMATION_SCHEMA.COLUMNS 
+      WHERE table_name != '${tableName}' 
+      and (column_name like '${depluralizedTableName}%'
+      or column_name like '%_id' 
+      or column_name like '%Id' 
+      or column_name like '%_${primaryKey}' 
+      or column_name like '%${capitalize(primaryKey)}')`;
+    const results = await runQueryBigQuery(sql, connection);
+    foreignKeysString = results.rows.map((key) => {
+      return key.column_name;
+    });
+    foreignKeysString = [...new Set(foreignKeysString)];
+  }
+  return foreignKeysString;
+}
+
+export async function getSchemaColumnInfoBigQuery(
+  connection: BigQuery,
+  schemaName: string,
+  tableNames: string[]
+): Promise<
+  { tableName: string; columns: { columnName: string; dataTypeId: number }[] }[]
+> {
+  const allColumns = await Promise.all(
+    tableNames.map(async (tableName) => {
+      const query = `
+      SELECT column_name as columnName, data_type as dataType
+      FROM ${schemaName}.INFORMATION_SCHEMA.COLUMNS
+      WHERE table_name = '${tableName}'
+      ORDER BY ordinal_position;
+    `;
+      const results = await runQueryBigQuery(query, connection);
+      return {
+        tableName,
+        displayName: tableName,
+        columns: results.rows.map((row: any) => ({
+          columnName: row.columnName,
+          displayName: row.columnName,
+          dataTypeId: inferType(convertBigQueryTypeToPostgresOID(row.dataType)),
+          fieldType: row.dataType,
+        })),
+      };
+    })
+  );
+  return allColumns;
+}
+
+function convertBigQueryTypeToPostgresOID(type: string): number {
+  const typeToOidMap: { [key: string]: number } = {
+    VARCHAR: 1043,
+    INTEGER: 23,
+    FLOAT: 700,
+    TIMESTAMP: 1114,
+    DATE: 1082,
+  };
+
+  const postgresType = typeToOidMap[type.toUpperCase()] || "VARCHAR"; // Default to 'text' if the type is not recognized
+  return typeToOidMap[postgresType] || 1043; // Default to OID for 'text' if the type is not recognized
+}
+
+function inferType(elem: any) {
+  if (typeof elem === "number") {
+    // Check if the number is a float or an integer
+    return Number.isInteger(elem) ? 23 : 700; // 23: integer, 700: real
+  }
+  if (typeof elem === "string") {
+    // Attempt to infer date, time, and timestamp formats
+    // Date in YYYY-MM-DD format
+    if (/^\d{4}-\d{2}-\d{2}$/.test(elem)) return 1082; // date
+    // Timestamp in YYYY-MM-DDTHH:MM:SS[.fraction] format
+    if (/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?$/.test(elem)) return 1114; // timestamp without timezone
+    // Timestamp in YYYY-MM-DDTHH:MM:SS[.fraction]Z format
+    if (/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?Z$/.test(elem))
+      return 1184; // timestamp with timezone
+    // Time in HH:MM:SS format
+    if (/^\d{2}:\d{2}:\d{2}$/.test(elem)) return 1083; // time
+    return 1043; // varchar
+  }
+  // Add more specific cases or different data types as needed
+  return 1043; // default or unknown type
+}

package/src/db/{CachedPools.ts → CachedConnection.ts}

@@ -3,6 +3,13 @@ import { Mapable, CacheCredentials } from "../models/Cache";
 import { QuillConfig } from "../models/Quill";
 import { createClient } from "redis";
 import { isSuperset } from "../utils/Error";
+import {
+  DatabaseConnection,
+  DatabaseType,
+  connectToDatabase,
+  disconnectFromDatabase,
+  runQueryByDatabase,
+} from "./DatabaseHelper";
 
 class PgError extends Error {
   code?: string;
@@ -21,14 +28,20 @@ class PgError extends Error {
 /** The TTL for new cache entries (default: 1h) */
 const DEFAULT_CACHE_TTL = 24 * 60 * 60;
 
-export class CachedPool {
-  public pool: Pool;
+export class CachedConnection {
+  public databaseType: DatabaseType;
+  public pool: DatabaseConnection;
   public orgId: any;
   public ttl: number;
   public cache: Mapable | null;
 
-  constructor(config: any, cacheConfig: Partial<CacheCredentials> = {}) {
-    this.pool = new Pool(config);
+  constructor(
+    databaseType: DatabaseType,
+    config: any,
+    cacheConfig: Partial<CacheCredentials> = {}
+  ) {
+    this.databaseType = databaseType;
+    this.pool = connectToDatabase(databaseType, config);
     this.ttl = cacheConfig?.ttl ?? DEFAULT_CACHE_TTL;
     this.cache = this.getCache(cacheConfig);
   }
@@ -36,30 +49,21 @@ export class CachedPool {
   public async query(text: string, values?: any[]): Promise<any> {
     try {
       if (!this.cache) {
-        const results = await this.pool.query(text, values);
-        return {
-          fields: results.fields.map((field: any) => ({
-            name: field.name,
-            dataTypeID: field.dataTypeID,
-          })),
-          rows: results.rows,
-        };
+        return await runQueryByDatabase(this.databaseType, this.pool, text);
       }
       const key: string = `${this.orgId}:${text}`;
       const cachedResult: string | null = await this.cache.get(key);
       if (cachedResult) {
         return JSON.parse(cachedResult);
       } else {
-        const newResult: any = await this.pool.query(text, values);
+        const newResult = await runQueryByDatabase(
+          this.databaseType,
+          this.pool,
+          text
+        );
         const newResultString: string = JSON.stringify(newResult);
         await this.cache.set(key, newResultString, "EX", DEFAULT_CACHE_TTL);
-        return {
-          fields: newResult.fields.map((field: any) => ({
-            name: field.name,
-            dataTypeID: field.dataTypeID,
-          })),
-          rows: newResult.rows,
-        };
+        return newResult;
       }
     } catch (err) {
       if (isSuperset(err, PgError)) {
@@ -92,6 +96,6 @@ export class CachedPool {
   }
 
   async close() {
-    await this.pool.end();
+    await disconnectFromDatabase(this.databaseType, this.pool);
   }
 }

package/src/db/DatabaseHelper.ts

@@ -0,0 +1,340 @@
+import { Pool } from "pg";
+import snowflake from "snowflake-sdk";
+import { Client } from "../models/Client";
+import { Pool as MysqlPool } from "mysql2";
+import { BigQuery } from "@google-cloud/bigquery";
+import {
+  PostgresConnectionConfig,
+  connectToPostgres,
+  disconnectFromPostgres,
+  formatPostgresConfig,
+  getTablesBySchemaPostgres,
+  getColumnsByTablePostgres,
+  runQueryPostgres,
+  getForeignKeysPostgres,
+  getSchemaColumnInfoPostgress,
+  getSchemasPostgres,
+} from "./Postgres";
+import {
+  SnowflakeConnectionConfig,
+  connectToSnowflake,
+  disconnectFromSnowflake,
+  formatSnowflakeConfig,
+  getTablesBySchemaSnowflake,
+  getColumnsByTableSnowflake,
+  runQuerySnowflake,
+  getForeignKeysSnowflake,
+  getSchemaColumnInfoSnowflake,
+  getSchemasSnowflake,
+} from "./Snowflake";
+import {
+  formatBigQueryConfig,
+  BigQueryConfig,
+  connectToBigQuery,
+  runQueryBigQuery,
+  getTablesBySchemaBigQuery,
+  getColumnsByTableBigQuery,
+  getForeignKeysBigQuery,
+  getSchemaColumnInfoBigQuery,
+  getSchemaBigQuery,
+} from "./BigQuery";
+import {
+  MysqlConnectionConfig,
+  connectToMysql,
+  disconnectFromMysql,
+  formatMysqlConfig,
+  runQueryMysql,
+  getTablesBySchemaMysql,
+  getColumnsByTableMysql,
+  getForeignKeysMysql,
+  getSchemaColumnInfoMysql,
+  getSchemasMysql,
+} from "./Mysql";
+
+export enum DatabaseType {
+  postgres = "postgres",
+  postgresql = "PostgreSQL",
+  snowflake = "Snowflake",
+  bigquery = "BigQuery",
+  mysql = "MySQL",
+}
+
+// export all database connections types
+export type DatabaseConnection =
+  | Pool
+  | snowflake.Connection
+  | BigQuery
+  | MysqlPool;
+
+export interface QuillQueryResults {
+  fields: { name: string; dataTypeID: number }[];
+  rows: { [fieldName: string]: any }[];
+}
+
+export function getDatabaseCredentials(
+  databaseType: DatabaseType,
+  connectionString: string
+):
+  | PostgresConnectionConfig
+  | SnowflakeConnectionConfig
+  | BigQueryConfig
+  | MysqlConnectionConfig
+  | undefined {
+  switch (databaseType) {
+    case "postgres":
+      return formatPostgresConfig(connectionString);
+    case "PostgreSQL":
+      return formatPostgresConfig(connectionString);
+    case "Snowflake":
+      return formatSnowflakeConfig(connectionString);
+    case "BigQuery":
+      return formatBigQueryConfig(connectionString);
+    case "MySQL":
+      return formatMysqlConfig(connectionString);
+    default:
+      return undefined;
+  }
+}
+
+export function connectToDatabase(
+  databaseType: DatabaseType,
+  config:
+    | PostgresConnectionConfig
+    | SnowflakeConnectionConfig
+    | BigQueryConfig
+    | MysqlConnectionConfig
+): DatabaseConnection {
+  switch (databaseType) {
+    case "postgres":
+      return connectToPostgres(config as PostgresConnectionConfig);
+    case "PostgreSQL":
+      return connectToPostgres(config as PostgresConnectionConfig);
+    case "Snowflake":
+      return connectToSnowflake(config as SnowflakeConnectionConfig);
+    case "BigQuery":
+      return connectToBigQuery(config as BigQueryConfig);
+    case "MySQL":
+      return connectToMysql(config as MysqlConnectionConfig);
+    default:
+      return connectToPostgres(config as PostgresConnectionConfig);
+  }
+}
+
+export function runQueryByDatabase(
+  databaseType: DatabaseType,
+  connection: DatabaseConnection,
+  sql: string
+): Promise<QuillQueryResults> | undefined {
+  switch (databaseType) {
+    case "postgres":
+      return runQueryPostgres(sql, connection as Pool);
+    case "PostgreSQL":
+      return runQueryPostgres(sql, connection as Pool);
+    case "Snowflake":
+      return runQuerySnowflake(sql, connection as snowflake.Connection);
+    case "BigQuery":
+      return runQueryBigQuery(sql, connection as BigQuery);
+    case "MySQL":
+      return runQueryMysql(sql, connection as MysqlPool);
+    default:
+      return undefined;
+  }
+}
+
+export function disconnectFromDatabase(
+  databaseType: DatabaseType,
+  database: DatabaseConnection
+) {
+  switch (databaseType) {
+    case "postgres":
+      return disconnectFromPostgres(database as Pool);
+    case "PostgreSQL":
+      return disconnectFromPostgres(database as Pool);
+    case "Snowflake":
+      return disconnectFromSnowflake(database as snowflake.Connection);
+    case "BigQuery":
+      return; // BigQuery does not need to be disconnected
+    case "MySQL":
+      return disconnectFromMysql(database as MysqlPool);
+    default:
+      return undefined;
+  }
+}
+
+export async function getSchemasByDatabase(
+  databaseType: DatabaseType,
+  connection: DatabaseConnection
+): Promise<string[] | undefined> {
+  switch (databaseType) {
+    case "postgres":
+      return getSchemasPostgres(connection as Pool);
+    case "PostgreSQL":
+      return getSchemasPostgres(connection as Pool);
+    case "Snowflake":
+      return getSchemasSnowflake(connection as snowflake.Connection);
+    case "BigQuery":
+      return getSchemaBigQuery(connection as BigQuery);
+    case "MySQL":
+      return getSchemasMysql(connection as MysqlPool);
+    default:
+      return undefined;
+  }
+}
+
+// INFORMATION SCHEMA SELECTS
+export async function getTablesBySchemaByDatabase(
+  databaseType: DatabaseType,
+  connection: DatabaseConnection,
+  schemaName: string
+): Promise<string[] | undefined> {
+  switch (databaseType) {
+    case "postgres":
+      return getTablesBySchemaPostgres(connection as Pool, schemaName);
+    case "PostgreSQL":
+      return getTablesBySchemaPostgres(connection as Pool, schemaName);
+    case "Snowflake":
+      return getTablesBySchemaSnowflake(
+        connection as snowflake.Connection,
+        schemaName
+      );
+    case "BigQuery":
+      return getTablesBySchemaBigQuery(connection as BigQuery, schemaName);
+    case "MySQL":
+      return getTablesBySchemaMysql(connection as MysqlPool, schemaName);
+    default:
+      return undefined;
+  }
+}
+
+// INFORMATION SCHEMA SELECTS
+export async function getColumnsByTableByDatabase(
+  databaseType: DatabaseType,
+  connection: DatabaseConnection,
+  schemaName: string,
+  tableName: string
+): Promise<string[] | undefined> {
+  switch (databaseType) {
+    case "postgres":
+      return getColumnsByTablePostgres(
+        connection as Pool,
+        schemaName,
+        tableName
+      );
+    case "PostgreSQL":
+      return getColumnsByTablePostgres(
+        connection as Pool,
+        schemaName,
+        tableName
+      );
+    case "Snowflake":
+      return getColumnsByTableSnowflake(
+        connection as snowflake.Connection,
+        schemaName,
+        tableName
+      );
+    case "BigQuery":
+      return getColumnsByTableBigQuery(
+        connection as BigQuery,
+        schemaName,
+        tableName
+      );
+    case "MySQL":
+      return getColumnsByTableMysql(
+        connection as MysqlPool,
+        schemaName,
+        tableName
+      );
+    default:
+      return undefined;
+  }
+}
+
+export async function getForiegnKeysByDatabase(
+  databaseType: DatabaseType,
+  connection: DatabaseConnection,
+  schemaName: string,
+  tableName: string,
+  primaryKey: string
+): Promise<string[] | undefined> {
+  switch (databaseType) {
+    case "postgres":
+      return getForeignKeysPostgres(
+        connection as Pool,
+        schemaName,
+        tableName,
+        primaryKey
+      );
+    case "PostgreSQL":
+      return getForeignKeysPostgres(
+        connection as Pool,
+        schemaName,
+        tableName,
+        primaryKey
+      );
+    case "Snowflake":
+      return getForeignKeysSnowflake(
+        connection as snowflake.Connection,
+        schemaName,
+        tableName,
+        primaryKey
+      );
+    case "BigQuery":
+      return getForeignKeysBigQuery(
+        connection as BigQuery,
+        schemaName,
+        tableName,
+        primaryKey
+      );
+    case "MySQL":
+      return getForeignKeysMysql(
+        connection as MysqlPool,
+        schemaName,
+        tableName,
+        primaryKey
+      );
+    default:
+      return undefined;
+  }
+}
+
+export function getColumnInfoBySchemaByDatabase(
+  databaseType: DatabaseType,
+  connection: DatabaseConnection,
+  schemaName: string,
+  tables: string[]
+) {
+  switch (databaseType) {
+    case "postgres":
+      return getSchemaColumnInfoPostgress(
+        connection as Pool,
+        schemaName,
+        tables
+      );
+    case "PostgreSQL":
+      return getSchemaColumnInfoPostgress(
+        connection as Pool,
+        schemaName,
+        tables
+      );
+    case "Snowflake":
+      return getSchemaColumnInfoSnowflake(
+        connection as snowflake.Connection,
+        schemaName,
+        tables
+      );
+    case "BigQuery":
+      return getSchemaColumnInfoBigQuery(
+        connection as BigQuery,
+        schemaName,
+        tables
+      );
+    case "MySQL":
+      return getSchemaColumnInfoMysql(
+        connection as MysqlPool,
+        schemaName,
+        tables
+      );
+    default:
+      return undefined;
+  }
+}