From d6ea77dda8a840122c94eea448aabfa2752ef312 Mon Sep 17 00:00:00 2001
From: Jakub Knejzlik <jakub.knejzlik@gmail.com>
Date: Fri, 22 Nov 2024 12:37:50 +0100
Subject: [PATCH] Add support for expresison values in update statement

---
 src/Mutation-serialization.test.ts | 14 ++++++++++---
 src/Mutation.test.ts               | 10 ++++++++--
 src/Mutation.ts                    | 31 ++++++++++++++++++++---------
 src/stories/1_Query.stories.mdx    | 32 ++++++++++++++++++++++++++++++
 4 files changed, 73 insertions(+), 14 deletions(-)

diff --git a/src/Mutation-serialization.test.ts b/src/Mutation-serialization.test.ts
index 8cecbcd..74b804d 100644
--- a/src/Mutation-serialization.test.ts
+++ b/src/Mutation-serialization.test.ts
@@ -1,4 +1,5 @@
 import { Conditions } from "./Condition";
+import { Fn } from "./Function";
 import { MutationBase } from "./Mutation";
 import { Q } from "./Query";
 
@@ -19,10 +20,17 @@ describe("Mutation serialization and deserialization", () => {
 
   it("should handle round-trip JSON serialization and deserialization for an update", () => {
     const originalQuery = Q.update("table")
-      .set({ foo: 123, bar: "baz" })
+      .set({
+        foo: 123,
+        bar: "baz",
+        total: Q.expr(Fn.multiply("amount", "price")),
+      })
       .where(Conditions.equal("foo", 123));
-    const jsonStr = JSON.stringify(originalQuery.toJSON());
-    const deserializedQuery = MutationBase.deserialize(jsonStr);
+    const ser = originalQuery.serialize();
+    const deserializedQuery = MutationBase.deserialize(ser);
+    expect(deserializedQuery.toSQL()).toEqual(
+      'UPDATE `table` SET `foo` = 123, `bar` = "baz", `total` = (`amount` * `price`) WHERE `foo` = 123'
+    );
     expect(deserializedQuery.toSQL()).toEqual(originalQuery.toSQL());
   });
 });
diff --git a/src/Mutation.test.ts b/src/Mutation.test.ts
index 0fcf202..4855fb8 100644
--- a/src/Mutation.test.ts
+++ b/src/Mutation.test.ts
@@ -1,4 +1,5 @@
 import { Cond } from "./Condition";
+import { Fn } from "./Function";
 import { Q } from "./Query";
 
 describe("Mutation builder SQL", () => {
@@ -21,10 +22,15 @@ describe("Mutation builder SQL", () => {
   it("should return SQL for update", () => {
     expect(
       Q.update("users")
-        .set({ name: "John Doe", age: 42, isActive: true })
+        .set({
+          name: "John Doe",
+          age: 42,
+          isActive: true,
+          total: Fn.multiply("amount", "price"),
+        })
         .toSQL()
     ).toEqual(
-      'UPDATE `users` SET `name` = "John Doe", `age` = 42, `isActive` = true'
+      'UPDATE `users` SET `name` = "John Doe", `age` = 42, `isActive` = true, `total` = (`amount` * `price`)'
     );
   });
 });
diff --git a/src/Mutation.ts b/src/Mutation.ts
index 8a0a78c..85b860b 100644
--- a/src/Mutation.ts
+++ b/src/Mutation.ts
@@ -1,6 +1,7 @@
 import { Condition } from "./Condition";
+import { Expression, ExpressionValue } from "./Expression";
 import { ISQLFlavor } from "./Flavor";
-import { Table } from "./Query";
+import { Q, Table } from "./Query";
 import { MySQLFlavor } from "./flavors/mysql";
 import {
   IMetadata,
@@ -10,7 +11,8 @@ import {
   OperationType,
 } from "./interfaces";
 
-type RowRecord = Record<string, any>;
+type RowRecord = Record<string, Expression>;
+type RowRecordInput = Record<string, Expression | any>;
 
 export class MutationBase {
   protected _table: Table;
@@ -113,7 +115,7 @@ export class InsertMutation
     return clone;
   }
 
-  values(values: RowRecord[]): this {
+  values(values: RowRecordInput[]): this {
     const clone = this.clone();
     clone._values = [...clone._values, ...values];
     return clone;
@@ -173,9 +175,13 @@ export class UpdateMutation
     return clone;
   }
 
-  set(values: RowRecord): this {
+  set(values: RowRecordInput): this {
     const clone = this.clone();
-    clone._values = { ...clone._values, ...values };
+    const _values = {};
+    for (const [key, value] of Object.entries(values)) {
+      _values[key] = value instanceof Expression ? value : Q.exprValue(value);
+    }
+    clone._values = { ...clone._values, ..._values };
     return clone;
   }
 
@@ -192,8 +198,7 @@ export class UpdateMutation
       this._values
     )
       .map(
-        ([key, value]) =>
-          `${flavor.escapeColumn(key)} = ${flavor.escapeValue(value)}`
+        ([key, value]) => `${flavor.escapeColumn(key)} = ${value.toSQL(flavor)}`
       )
       .join(", ")}`;
     if (this._where.length) {
@@ -209,17 +214,25 @@ export class UpdateMutation
   }
 
   toJSON() {
+    const _values = {};
+    for (const [key, value] of Object.entries(this._values)) {
+      _values[key] = value.serialize();
+    }
     return {
       type: OperationType.UPDATE,
       table: this._table.toJSON(),
-      values: this._values,
+      values: _values,
       where: this._where.map((condition) => condition.toJSON()),
     };
   }
 
   static fromJSON({ table, values, where }: any): UpdateMutation {
     const updateMutation = new UpdateMutation(table.source, table.alias);
-    updateMutation._values = values;
+    const _values = {};
+    for (const [key, value] of Object.entries(values)) {
+      _values[key] = Expression.deserialize(value as ExpressionValue);
+    }
+    updateMutation._values = _values;
     updateMutation._where = where.map((condition: any) =>
       Condition.fromJSON(condition)
     );
diff --git a/src/stories/1_Query.stories.mdx b/src/stories/1_Query.stories.mdx
index 9ca59c1..d9cd3c3 100644
--- a/src/stories/1_Query.stories.mdx
+++ b/src/stories/1_Query.stories.mdx
@@ -279,6 +279,38 @@ Q.select().from(Q.select().from('table').where(Cond.equal('field', 'blah')),'sub
 `}
 />
 
+## Mutations
+
+You can also generate insert/update/delete statements
+
+### Insert
+
+<QueryPreview
+  code={`
+Q.insert('table').values([{foo:'blah'}]);
+`}
+/>
+
+### Update
+
+<QueryPreview
+  code={`
+Q.update('table')
+  .set({ 
+    foo: 'blah', 
+    total: Fn.divide(Fn.multiply('amount','price'),Q.raw(2)) 
+  }).where(Cond.equal('name','john'));
+`}
+/>
+
+### Delete
+
+<QueryPreview
+  code={`
+Q.delete('table').where(Cond.equal('id',123));
+`}
+/>
+
 ## Examples
 
 ### Basic Select Query