forge-sql-orm 2.0.8 → 2.0.9

package/src/core/ForgeSQLQueryBuilder.ts

@@ -1,6 +1,5 @@
 import { UpdateQueryResponse } from "@forge/sql";
 import { SqlParameters } from "@forge/sql/out/sql-statement";
-import { MySql2Database } from "drizzle-orm/mysql2/driver";
 import {
   AnyMySqlSelectQueryBuilder,
   AnyMySqlTable,
@@ -14,7 +13,7 @@ import {
 import { InferInsertModel, SQL } from "drizzle-orm";
 import moment from "moment/moment";
 import { parseDateTime } from "../utils/sqlUtils";
-import { MySql2PreparedQueryHKT } from "drizzle-orm/mysql2/index";
+import { MySqlRemoteDatabase, MySqlRemotePreparedQueryHKT } from "drizzle-orm/mysql-proxy/index";
 
 // ============= Core Types =============
 
@@ -45,15 +44,15 @@ export interface QueryBuilderForgeSql {
    * Creates a new query builder for the given entity.
    * @returns {MySql2Database<Record<string, unknown>>} The Drizzle database instance for building queries
    */
-  getDrizzleQueryBuilder(): MySql2Database<Record<string, unknown>>;
+  getDrizzleQueryBuilder(): MySqlRemoteDatabase<Record<string, unknown>>;
 
   select<TSelection extends SelectedFields>(
     fields: TSelection,
-  ): MySqlSelectBuilder<TSelection, MySql2PreparedQueryHKT>;
+  ): MySqlSelectBuilder<TSelection, MySqlRemotePreparedQueryHKT>;
 
   selectDistinct<TSelection extends SelectedFields>(
     fields: TSelection,
-  ): MySqlSelectBuilder<TSelection, MySql2PreparedQueryHKT>;
+  ): MySqlSelectBuilder<TSelection, MySqlRemotePreparedQueryHKT>;
 }
 
 // ============= CRUD Operations =============

package/src/utils/forgeDriver.ts

@@ -1,28 +1,39 @@
-import { sql } from "@forge/sql";
-import { AnyMySql2Connection } from "drizzle-orm/mysql2/driver";
+import { sql, UpdateQueryResponse } from "@forge/sql";
 
 interface ForgeSQLResult {
   rows: Record<string, unknown>[] | Record<string, unknown>;
 }
 
-export const forgeDriver = {
-  query: async (query: { sql: string }, params?: unknown[]) => {
-    try {
-      const sqlStatement = await sql.prepare<unknown>(query.sql);
+export const forgeDriver = async (
+  query: string,
+  params: any[],
+  method: "all" | "execute",
+): Promise<{
+  rows: any[];
+  insertId?: number;
+  affectedRows?: number;
+}> => {
+  try {
+    if (method == "execute") {
+      const sqlStatement = sql.prepare<UpdateQueryResponse>(query);
+      if (params) {
+        sqlStatement.bindParams(...params);
+      }
+      const updateQueryResponseResults = await sqlStatement.execute();
+      let result = updateQueryResponseResults.rows as any;
+      return { ...result, rows: [result] };
+    } else {
+      const sqlStatement = await sql.prepare<unknown>(query);
       if (params) {
         await sqlStatement.bindParams(...params);
       }
       const result = (await sqlStatement.execute()) as ForgeSQLResult;
       let rows;
-      if (Array.isArray(result.rows)) {
-        rows = [result.rows.map((r) => Object.values(r as Record<string, unknown>))];
-      } else {
-        rows = [result.rows as Record<string, unknown>];
-      }
-      return rows;
-    } catch (error) {
-      console.error("SQL Error:", JSON.stringify(error));
-      throw error;
+      rows = (result.rows as any[]).map((r) => Object.values(r as Record<string, unknown>));
+      return { rows: rows };
     }
-  },
-} as unknown as AnyMySql2Connection;
+  } catch (error) {
+    console.error("SQL Error:", JSON.stringify(error));
+    throw error;
+  }
+};