forge-sql-orm 2.1.9 → 2.1.11

package/src/lib/drizzle/extensions/additionalActions.ts

@@ -8,7 +8,7 @@ import {
 import { SelectedFields } from "drizzle-orm/mysql-core/query-builders/select.types";
 import { applyFromDriverTransform, ForgeSqlOrmOptions, mapSelectFieldsWithAlias } from "../../..";
 import { MySqlSelectBase, MySqlSelectBuilder } from "drizzle-orm/mysql-core";
-import type { MySqlTable } from "drizzle-orm/mysql-core/table";
+import { MySqlTable } from "drizzle-orm/mysql-core/table";
 import {
   MySqlDeleteBase,
   MySqlInsertBuilder,
@@ -26,9 +26,10 @@ import { isSQLWrapper, SQLWrapper } from "drizzle-orm/sql/sql";
 import type { MySqlQueryResultKind } from "drizzle-orm/mysql-core/session";
 import { getTableColumns, Query, SQL } from "drizzle-orm";
 import { MySqlDialect } from "drizzle-orm/mysql-core/dialect";
-import type {
+import {
   GetSelectTableName,
   GetSelectTableSelection,
+  SelectResultField,
 } from "drizzle-orm/query-builders/select.types";
 
 // ============================================================================
@@ -142,12 +143,19 @@ export type SelectAllFromAliasedType = <T extends MySqlTable>(
   table: T,
 ) => MySqlSelectBase<
   GetSelectTableName<T>,
-  T["_"]["columns"] extends undefined ? GetSelectTableSelection<T> : T["_"]["columns"],
-  T["_"]["columns"] extends undefined ? "single" : "partial",
+  GetSelectTableSelection<T>,
+  "single",
   MySqlRemotePreparedQueryHKT,
   GetSelectTableName<T> extends string ? Record<string & GetSelectTableName<T>, "not-null"> : {},
   false,
   never,
+  {
+    [K in keyof {
+      [Key in keyof GetSelectTableSelection<T>]: SelectResultField<GetSelectTableSelection<T>[Key]>;
+    }]: {
+      [Key in keyof GetSelectTableSelection<T>]: SelectResultField<GetSelectTableSelection<T>[Key]>;
+    }[K];
+  }[],
   any
 >;
 
@@ -158,12 +166,19 @@ export type SelectAllDistinctFromAliasedType = <T extends MySqlTable>(
   table: T,
 ) => MySqlSelectBase<
   GetSelectTableName<T>,
-  T["_"]["columns"] extends undefined ? GetSelectTableSelection<T> : T["_"]["columns"],
-  T["_"]["columns"] extends undefined ? "single" : "partial",
+  GetSelectTableSelection<T>,
+  "single",
   MySqlRemotePreparedQueryHKT,
   GetSelectTableName<T> extends string ? Record<string & GetSelectTableName<T>, "not-null"> : {},
   false,
   never,
+  {
+    [K in keyof {
+      [Key in keyof GetSelectTableSelection<T>]: SelectResultField<GetSelectTableSelection<T>[Key]>;
+    }]: {
+      [Key in keyof GetSelectTableSelection<T>]: SelectResultField<GetSelectTableSelection<T>[Key]>;
+    }[K];
+  }[],
   any
 >;
 
@@ -175,12 +190,19 @@ export type SelectAllFromCacheableAliasedType = <T extends MySqlTable>(
   cacheTtl?: number,
 ) => MySqlSelectBase<
   GetSelectTableName<T>,
-  T["_"]["columns"] extends undefined ? GetSelectTableSelection<T> : T["_"]["columns"],
-  T["_"]["columns"] extends undefined ? "single" : "partial",
+  GetSelectTableSelection<T>,
+  "single",
   MySqlRemotePreparedQueryHKT,
   GetSelectTableName<T> extends string ? Record<string & GetSelectTableName<T>, "not-null"> : {},
   false,
   never,
+  {
+    [K in keyof {
+      [Key in keyof GetSelectTableSelection<T>]: SelectResultField<GetSelectTableSelection<T>[Key]>;
+    }]: {
+      [Key in keyof GetSelectTableSelection<T>]: SelectResultField<GetSelectTableSelection<T>[Key]>;
+    }[K];
+  }[],
   any
 >;
 
@@ -192,12 +214,19 @@ export type SelectAllDistinctFromCacheableAliasedType = <T extends MySqlTable>(
   cacheTtl?: number,
 ) => MySqlSelectBase<
   GetSelectTableName<T>,
-  T["_"]["columns"] extends undefined ? GetSelectTableSelection<T> : T["_"]["columns"],
-  T["_"]["columns"] extends undefined ? "single" : "partial",
+  GetSelectTableSelection<T>,
+  "single",
   MySqlRemotePreparedQueryHKT,
   GetSelectTableName<T> extends string ? Record<string & GetSelectTableName<T>, "not-null"> : {},
   false,
   never,
+  {
+    [K in keyof {
+      [Key in keyof GetSelectTableSelection<T>]: SelectResultField<GetSelectTableSelection<T>[Key]>;
+    }]: {
+      [Key in keyof GetSelectTableSelection<T>]: SelectResultField<GetSelectTableSelection<T>[Key]>;
+    }[K];
+  }[],
   any
 >;
 
@@ -774,8 +803,52 @@ export function patchDbWithSelectAliased(
    * const users = await db.selectFrom(userTable).where(eq(userTable.id, 1));
    * ```
    */
-  db.selectFrom = function <T extends MySqlTable>(table: T) {
-    return db.selectAliased(getTableColumns(table)).from(table);
+  db.selectFrom = function <T extends MySqlTable>(
+    table: T,
+  ): MySqlSelectBase<
+    GetSelectTableName<T>,
+    GetSelectTableSelection<T>,
+    "single",
+    MySqlRemotePreparedQueryHKT,
+    GetSelectTableName<T> extends string ? Record<string & GetSelectTableName<T>, "not-null"> : {},
+    false,
+    never,
+    {
+      [K in keyof {
+        [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+          GetSelectTableSelection<T>[Key]
+        >;
+      }]: {
+        [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+          GetSelectTableSelection<T>[Key]
+        >;
+      }[K];
+    }[],
+    any
+  > {
+    return db.selectAliased(getTableColumns(table)).from(table) as unknown as MySqlSelectBase<
+      GetSelectTableName<T>,
+      GetSelectTableSelection<T>,
+      "single",
+      MySqlRemotePreparedQueryHKT,
+      GetSelectTableName<T> extends string
+        ? Record<string & GetSelectTableName<T>, "not-null">
+        : {},
+      false,
+      never,
+      {
+        [K in keyof {
+          [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+            GetSelectTableSelection<T>[Key]
+          >;
+        }]: {
+          [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+            GetSelectTableSelection<T>[Key]
+          >;
+        }[K];
+      }[],
+      any
+    >;
   };
 
   /**
@@ -790,8 +863,55 @@ export function patchDbWithSelectAliased(
    * const users = await db.selectFromCacheable(userTable, 300).where(eq(userTable.id, 1));
    * ```
    */
-  db.selectFromCacheable = function <T extends MySqlTable>(table: T, cacheTtl?: number) {
-    return db.selectAliasedCacheable(getTableColumns(table), cacheTtl).from(table);
+  db.selectFromCacheable = function <T extends MySqlTable>(
+    table: T,
+    cacheTtl?: number,
+  ): MySqlSelectBase<
+    GetSelectTableName<T>,
+    GetSelectTableSelection<T>,
+    "single",
+    MySqlRemotePreparedQueryHKT,
+    GetSelectTableName<T> extends string ? Record<string & GetSelectTableName<T>, "not-null"> : {},
+    false,
+    never,
+    {
+      [K in keyof {
+        [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+          GetSelectTableSelection<T>[Key]
+        >;
+      }]: {
+        [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+          GetSelectTableSelection<T>[Key]
+        >;
+      }[K];
+    }[],
+    any
+  > {
+    return db
+      .selectAliasedCacheable(getTableColumns(table), cacheTtl)
+      .from(table) as unknown as MySqlSelectBase<
+      GetSelectTableName<T>,
+      GetSelectTableSelection<T>,
+      "single",
+      MySqlRemotePreparedQueryHKT,
+      GetSelectTableName<T> extends string
+        ? Record<string & GetSelectTableName<T>, "not-null">
+        : {},
+      false,
+      never,
+      {
+        [K in keyof {
+          [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+            GetSelectTableSelection<T>[Key]
+          >;
+        }]: {
+          [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+            GetSelectTableSelection<T>[Key]
+          >;
+        }[K];
+      }[],
+      any
+    >;
   };
 
   /**
@@ -805,8 +925,54 @@ export function patchDbWithSelectAliased(
    * const uniqueUsers = await db.selectDistinctFrom(userTable).where(eq(userTable.status, 'active'));
    * ```
    */
-  db.selectDistinctFrom = function <T extends MySqlTable>(table: T) {
-    return db.selectAliasedDistinct(getTableColumns(table)).from(table);
+  db.selectDistinctFrom = function <T extends MySqlTable>(
+    table: T,
+  ): MySqlSelectBase<
+    GetSelectTableName<T>,
+    GetSelectTableSelection<T>,
+    "single",
+    MySqlRemotePreparedQueryHKT,
+    GetSelectTableName<T> extends string ? Record<string & GetSelectTableName<T>, "not-null"> : {},
+    false,
+    never,
+    {
+      [K in keyof {
+        [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+          GetSelectTableSelection<T>[Key]
+        >;
+      }]: {
+        [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+          GetSelectTableSelection<T>[Key]
+        >;
+      }[K];
+    }[],
+    any
+  > {
+    return db
+      .selectAliasedDistinct(getTableColumns(table))
+      .from(table) as unknown as MySqlSelectBase<
+      GetSelectTableName<T>,
+      GetSelectTableSelection<T>,
+      "single",
+      MySqlRemotePreparedQueryHKT,
+      GetSelectTableName<T> extends string
+        ? Record<string & GetSelectTableName<T>, "not-null">
+        : {},
+      false,
+      never,
+      {
+        [K in keyof {
+          [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+            GetSelectTableSelection<T>[Key]
+          >;
+        }]: {
+          [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+            GetSelectTableSelection<T>[Key]
+          >;
+        }[K];
+      }[],
+      any
+    >;
   };
 
   /**
@@ -821,8 +987,55 @@ export function patchDbWithSelectAliased(
    * const uniqueUsers = await db.selectDistinctFromCacheable(userTable, 300).where(eq(userTable.status, 'active'));
    * ```
    */
-  db.selectDistinctFromCacheable = function <T extends MySqlTable>(table: T, cacheTtl?: number) {
-    return db.selectAliasedDistinctCacheable(getTableColumns(table), cacheTtl).from(table);
+  db.selectDistinctFromCacheable = function <T extends MySqlTable>(
+    table: T,
+    cacheTtl?: number,
+  ): MySqlSelectBase<
+    GetSelectTableName<T>,
+    GetSelectTableSelection<T>,
+    "single",
+    MySqlRemotePreparedQueryHKT,
+    GetSelectTableName<T> extends string ? Record<string & GetSelectTableName<T>, "not-null"> : {},
+    false,
+    never,
+    {
+      [K in keyof {
+        [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+          GetSelectTableSelection<T>[Key]
+        >;
+      }]: {
+        [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+          GetSelectTableSelection<T>[Key]
+        >;
+      }[K];
+    }[],
+    any
+  > {
+    return db
+      .selectAliasedDistinctCacheable(getTableColumns(table), cacheTtl)
+      .from(table) as unknown as MySqlSelectBase<
+      GetSelectTableName<T>,
+      GetSelectTableSelection<T>,
+      "single",
+      MySqlRemotePreparedQueryHKT,
+      GetSelectTableName<T> extends string
+        ? Record<string & GetSelectTableName<T>, "not-null">
+        : {},
+      false,
+      never,
+      {
+        [K in keyof {
+          [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+            GetSelectTableSelection<T>[Key]
+          >;
+        }]: {
+          [Key in keyof GetSelectTableSelection<T>]: SelectResultField<
+            GetSelectTableSelection<T>[Key]
+          >;
+        }[K];
+      }[],
+      any
+    >;
   };
 
   // ============================================================================

package/src/utils/cacheUtils.ts

@@ -332,7 +332,9 @@ export async function getFromCache<T>(
   }
 
   try {
-    const cacheResult = await kvs.entity<CacheEntity>(options.cacheEntityName).get(key);
+    const cacheResult = (await kvs.entity<CacheEntity>(options.cacheEntityName).get(key)) as
+      | CacheEntity
+      | undefined;
 
     if (
       cacheResult &&

package/src/utils/forgeDriver.ts

@@ -1,6 +1,10 @@
 import { sql, UpdateQueryResponse } from "@forge/sql";
 import { saveMetaDataToContext } from "./metadataContextUtils";
 import { getOperationType } from "./requestTypeContextUtils";
+import {withTimeout} from "./sqlUtils";
+
+const timeoutMs =10000;
+const timeoutMessage = `Atlassian @forge/sql did not return a response within ${timeoutMs}ms (${timeoutMs / 1000} seconds), so the request is blocked. Possible causes: slow query, network issues, or exceeding Forge SQL limits.`;
 
 /**
  * Metadata structure for Forge SQL query results.
@@ -62,36 +66,6 @@ export function isUpdateQueryResponse(obj: unknown): obj is UpdateQueryResponse
   );
 }
 
-/**
- * Executes a promise with a timeout.
- *
- * @param promise - The promise to execute
- * @param timeoutMs - Timeout in milliseconds (default: 10000ms)
- * @returns Promise that resolves with the result or rejects on timeout
- * @throws {Error} When the operation times out
- */
-async function withTimeout<T>(promise: Promise<T>, timeoutMs: number = 10000): Promise<T> {
-  let timeoutId: ReturnType<typeof setTimeout> | undefined;
-
-  const timeoutPromise = new Promise<never>((_, reject) => {
-    timeoutId = setTimeout(() => {
-      reject(
-        new Error(
-          `Atlassian @forge/sql did not return a response within ${timeoutMs}ms (${timeoutMs / 1000} seconds), so the request is blocked. Possible causes: slow query, network issues, or exceeding Forge SQL limits.`,
-        ),
-      );
-    }, timeoutMs);
-  });
-
-  try {
-    return await Promise.race([promise, timeoutPromise]);
-  } finally {
-    if (timeoutId) {
-      clearTimeout(timeoutId);
-    }
-  }
-}
-
 function inlineParams(sql: string, params: unknown[]): string {
   let i = 0;
   return sql.replace(/\?/g, () => {
@@ -124,7 +98,7 @@ async function processDDLResult(method: QueryMethod, result: any): Promise<Forge
 
   if (Array.isArray(result.rows)) {
     if (method === "execute") {
-      return { rows: result.rows };
+      return { rows: [result.rows] };
     } else {
       const rows = (result.rows as any[]).map((r) => Object.values(r as Record<string, unknown>));
       return { rows };
@@ -148,17 +122,12 @@ async function processExecuteMethod(query: string, params: unknown[]): Promise<F
     sqlStatement.bindParams(...params);
   }
 
-  const result = await withTimeout(sqlStatement.execute());
+  const result = await withTimeout(sqlStatement.execute(), timeoutMessage, timeoutMs);
   await saveMetaDataToContext(result.metadata as ForgeSQLMetadata);
-  if (!result?.rows) {
+  if (!result.rows) {
     return { rows: [[]] };
   }
 
-  if (isUpdateQueryResponse(result.rows)) {
-    const oneRow = result.rows as any;
-    return { rows: [oneRow] };
-  }
-
   return { rows: [result.rows] };
 }
 
@@ -176,10 +145,10 @@ async function processAllMethod(query: string, params: unknown[]): Promise<Forge
     await sqlStatement.bindParams(...params);
   }
 
-  const result = (await withTimeout(sqlStatement.execute())) as ForgeSQLResult;
+  const result = (await withTimeout(sqlStatement.execute(), timeoutMessage, timeoutMs)) as ForgeSQLResult;
   await saveMetaDataToContext(result.metadata);
 
-  if (!result?.rows) {
+  if (!result.rows) {
     return { rows: [] };
   }
 
@@ -217,10 +186,9 @@ export const forgeDriver = async (
   method: QueryMethod,
 ): Promise<ForgeDriverResult> => {
   const operationType = await getOperationType();
-
   // Handle DDL operations
   if (operationType === "DDL") {
-    const result = await withTimeout(sql.executeDDL(inlineParams(query, params)));
+    const result = await withTimeout(sql.executeDDL(inlineParams(query, params)), timeoutMessage, timeoutMs);
     return await processDDLResult(method, result);
   }
 

package/src/utils/forgeDriverProxy.ts

@@ -1,11 +1,33 @@
 import { forgeDriver } from "./forgeDriver";
 import { injectSqlHints, SqlHints } from "./sqlHints";
+import { ForgeSqlOperation } from "../core/ForgeSQLQueryBuilder";
+import { printQueriesWithPlan } from "./sqlUtils";
 
 /**
- * Creates a proxy for the forgeDriver that injects SQL hints
+ * Error codes and constants for query analysis
+ */
+const QUERY_ERROR_CODES = {
+  TIMEOUT: "SQL_QUERY_TIMEOUT",
+  OUT_OF_MEMORY_ERRNO: 8175,
+} as const;
+
+/**
+ * Delay to wait for CLUSTER_STATEMENTS_SUMMARY to be populated
+ */
+const STATEMENTS_SUMMARY_DELAY_MS = 200;
+
+/**
+ * Creates a proxy for the forgeDriver that injects SQL hints and handles query analysis
+ * @param forgeSqlOperation - The ForgeSQL operation instance
+ * @param options - SQL hints to inject
+ * @param logRawSqlQuery - Whether to log raw SQL queries
  * @returns A proxied version of the forgeDriver
  */
-export function createForgeDriverProxy(options?: SqlHints, logRawSqlQuery?: boolean) {
+export function createForgeDriverProxy(
+  forgeSqlOperation: ForgeSqlOperation,
+  options?: SqlHints,
+  logRawSqlQuery?: boolean,
+) {
   return async (
     query: string,
     params: any[],
@@ -20,16 +42,46 @@ export function createForgeDriverProxy(options?: SqlHints, logRawSqlQuery?: bool
 
     if (options && logRawSqlQuery && modifiedQuery !== query) {
       // eslint-disable-next-line no-console
-      console.debug("injected Hints: " + modifiedQuery);
+      console.debug(`SQL Hints injected: ${modifiedQuery}`);
     }
+
+    const queryStartTime = Date.now();
+
     try {
-      // Call the original forgeDriver with the modified query
+      // Execute the query with injected hints
       return await forgeDriver(modifiedQuery, params, method);
-    } catch (error) {
+    } catch (error: any) {
+      // Check if this is a timeout or out-of-memory error that we want to analyze
+      const isTimeoutError = error.code === QUERY_ERROR_CODES.TIMEOUT;
+      const isOutOfMemoryError =
+        error?.context?.debug?.errno === QUERY_ERROR_CODES.OUT_OF_MEMORY_ERRNO;
+
+      if (isTimeoutError || isOutOfMemoryError) {
+        if (isTimeoutError) {
+          // eslint-disable-next-line no-console
+          console.error(` TIMEOUT detected - Query exceeded time limit`);
+        } else {
+          // eslint-disable-next-line no-console
+          console.error(`OUT OF MEMORY detected - Query exceeded memory limit`);
+        }
+
+        // Wait for CLUSTER_STATEMENTS_SUMMARY to be populated with our failed query data
+        await new Promise((resolve) => setTimeout(resolve, STATEMENTS_SUMMARY_DELAY_MS));
+
+        const queryEndTime = Date.now();
+        const queryDuration = queryEndTime - queryStartTime;
+
+        // Analyze the failed query using CLUSTER_STATEMENTS_SUMMARY
+        await printQueriesWithPlan(forgeSqlOperation, queryDuration);
+      }
+
+      // Log SQL error details if requested
       if (logRawSqlQuery) {
         // eslint-disable-next-line no-console
-        console.debug("SQL Error:", JSON.stringify(error));
+        console.debug(`SQL Error Details:`, JSON.stringify(error, null, 2));
       }
+
+      // Re-throw the original error
       throw error;
     }
   };

package/src/utils/metadataContextUtils.ts

@@ -1,19 +1,34 @@
 import { AsyncLocalStorage } from "node:async_hooks";
 import { ForgeSQLMetadata } from "./forgeDriver";
+import { ForgeSqlOperation } from "../core/ForgeSQLQueryBuilder";
+import { printQueriesWithPlan } from "./sqlUtils";
 
 export type MetadataQueryContext = {
   totalDbExecutionTime: number;
   totalResponseSize: number;
-  lastMetadata?: ForgeSQLMetadata;
+  beginTime: Date;
+  printQueriesWithPlan: () => Promise<void>;
+  forgeSQLORM: ForgeSqlOperation;
 };
 export const metadataQueryContext = new AsyncLocalStorage<MetadataQueryContext>();
 
-export async function saveMetaDataToContext(metadata: ForgeSQLMetadata): Promise<void> {
+export async function saveMetaDataToContext(metadata?: ForgeSQLMetadata): Promise<void> {
   const context = metadataQueryContext.getStore();
-  if (context && metadata) {
-    context.totalResponseSize += metadata.responseSize;
-    context.totalDbExecutionTime += metadata.dbExecutionTime;
-    context.lastMetadata = metadata;
+  if (context) {
+    context.printQueriesWithPlan = async () => {
+      if (process.env.NODE_ENV !== "test") {
+        await new Promise((r) => setTimeout(r, 200));
+      }
+       await printQueriesWithPlan(
+        context.forgeSQLORM,
+        Date.now() - context.beginTime.getTime(),
+      );
+    };
+    if (metadata) {
+      context.totalResponseSize += metadata.responseSize;
+      context.totalDbExecutionTime += metadata.dbExecutionTime;
+    }
+    // Log the results to console
   }
 }