feat: add param merge_fields to upsert_on_duplicate

NightMarcher · Aug 24, 2023 · 186abe0 · 186abe0
1 parent 468a2a8
commit 186abe0
Show file tree

Hide file tree

Showing 7 changed files with 53 additions and 37 deletions.
diff --git a/README.md b/README.md
@@ -119,24 +119,24 @@ Generate sql and execute
 ```python
 await AccountMgr.upsert_on_duplicate(
     [
-        {"id": 7, "gender": 1, "name": "斉藤 修平", "locale": "ja_JP", "extend": {}},
-        {"id": 8, "gender": 1, "name": "Ojas Salvi", "locale": "en_IN", "extend": {}},
-        {"id": 9, "gender": 1, "name": "羊淑兰", "locale": "zh_CN", "extend": {}}
+        {"id": 10, "gender": 1, "name": "田中 知実", "locale": "ja_JP", "extend": {"rdm": 1}},
+        {"id": 11, "gender": 2, "name": "Tara Chadha", "locale": "en_IN", "extend": {"rdm": 10}},
+        {"id": 12, "gender": 2, "name": "吴磊", "locale": "zh_CN", "extend": {"rdm": 9}},
     ],
     insert_fields=["id", "gender", "name", "locale", "extend"],
-    upsert_fields=["name", "locale"],
-    using_values=False,
+    upsert_fields=["gender", "name"],
+    merge_fields=["extend"],
 )
 ```
 Generate sql and execute
 ```sql
     INSERT INTO account
       (id, gender, name, locale, extend)
     VALUES
-      (7, 1, '斉藤 修平', 'ja_JP', '{}'),
-      (8, 1, 'Ojas Salvi', 'en_IN', '{}'),
-      (9, 1, '羊淑兰', 'zh_CN', '{}')
-    AS `new_account` ON DUPLICATE KEY UPDATE name=`new_account`.name, locale=`new_account`.locale
+      (10, 1, '田中 知実', 'ja_JP', '{"rdm": 1}'),
+      (11, 2, 'Tara Chadha', 'en_IN', '{"rdm": 10}'),
+      (12, 2, '吴磊', 'zh_CN', '{"rdm": 9}')
+    AS `new_account` ON DUPLICATE KEY UPDATE gender=`new_account`.gender, name=`new_account`.name, extend=JSON_MERGE_PATCH(COALESCE(account.extend, '{}'), `new_account`.extend)
 ```
 
 ### **insert_into_select**

diff --git a/examples/service/routers/account.py b/examples/service/routers/account.py
@@ -160,12 +160,16 @@ async def bulk_upsert_view():
             "gender": faker.random_int(0, 2),
             "name": faker.name(),
             "locale": locale.value,
-            "extend": {},
+            "extend": {"rdm": faker.random_int(0, 10)},
         })
+    from json import dumps
+    print(dumps(dicts, ensure_ascii=False))
     row_cnt = await AccountMgr.upsert_on_duplicate(
         dicts,
         insert_fields=["id", "gender", "name", "locale", "extend"],
-        # upsert_fields=["name", "locale"],
+        upsert_fields=["gender", "name"],
+        merge_fields=["extend"],
+        using_values=True,
     )
     return {"row_cnt": row_cnt}
 

diff --git a/fastapi_esql/__init__.py b/fastapi_esql/__init__.py
@@ -22,7 +22,7 @@
     timing,
 )
 
-__version__ = "0.0.10"
+__version__ = "0.0.11"
 
 __all__ = [
     "QsParsingError",

diff --git a/fastapi_esql/orm/base_manager.py b/fastapi_esql/orm/base_manager.py
@@ -150,13 +150,15 @@ async def upsert_on_duplicate(
         dicts: List[Dict[str, Any]],
         insert_fields: List[str],
         upsert_fields: Optional[List[str]] = None,
+        merge_fields: Optional[List[str]] = None,
         using_values: bool = False,
     ):
         sql = SQLizer.upsert_on_duplicate(
             cls.table,
             dicts,
             insert_fields,
             upsert_fields,
+            merge_fields,
             using_values,
         )
         return await CursorHandler.sum_row_cnt(sql, cls.rw_conn, logger)

diff --git a/fastapi_esql/utils/sqlizer.py b/fastapi_esql/utils/sqlizer.py
@@ -196,6 +196,7 @@ def upsert_on_duplicate(
         dicts: List[Dict[str, Any]],
         insert_fields: List[str],
         upsert_fields: Optional[List[str]] = None,
+        merge_fields: Optional[List[str]] = None,
         using_values: bool = False,
     ) -> Optional[str]:
         if not all([table, dicts, insert_fields]):
@@ -208,13 +209,21 @@ def upsert_on_duplicate(
         # NOTE Beginning with MySQL 8.0.19, it is possible to use an alias for the row
         # https://dev.mysql.com/doc/refman/8.0/en/insert-on-duplicate.html
         on_duplicate = ""
-        if upsert_fields:
+        if upsert_fields or merge_fields:
+            upsert_fields = upsert_fields or []
+            merge_fields = merge_fields or []
             if using_values:
                 upserts = [f"{field}=VALUES({field})" for field in upsert_fields]
+                for mf in merge_fields:
+                    dict_obj = f"COALESCE({table}.{mf}, '{{}}')"
+                    upserts.append(f"{mf}=JSON_MERGE_PATCH({dict_obj}, VALUES({mf}))")
                 on_duplicate = f"ON DUPLICATE KEY UPDATE {', '.join(upserts)}"
             else:
                 new_table = f"`new_{table}`"
                 upserts = [f"{field}={new_table}.{field}" for field in upsert_fields]
+                for mf in merge_fields:
+                    dict_obj = f"COALESCE({table}.{mf}, '{{}}')"
+                    upserts.append(f"{mf}=JSON_MERGE_PATCH({dict_obj}, {new_table}.{mf})")
                 on_duplicate = f"AS {new_table} ON DUPLICATE KEY UPDATE {', '.join(upserts)}"
 
         sql = """

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastapi-efficient-sql"
-version = "0.0.10"
+version = "0.0.11"
 description = "Generate bulk DML SQL and execute them based on Tortoise ORM and mysql8.0+, and integrated with FastAPI."
 authors = ["BryanLee <[email protected]>"]
 keywords = ["sql", "fastapi", "tortoise-orm", "mysql8", "bulk-operation"]

diff --git a/tests/test_sqlizer.py b/tests/test_sqlizer.py
@@ -212,61 +212,62 @@ def test_upsert_on_duplicate(self):
         old_sql = SQLizer.upsert_on_duplicate(
             self.table,
             [
-                {"id": 7, "gender": 1, "name": "斉藤 修平", "locale": "ja_JP", "extend": {}},
-                {"id": 8, "gender": 1, "name": "Ojas Salvi", "locale": "en_IN", "extend": {}},
-                {"id": 9, "gender": 1, "name": "羊淑兰", "locale": "zh_CN", "extend": {}}
+                {"id": 10, "gender": 2, "name": "佐々木 美加子", "locale": "ja_JP", "extend": {"rdm": 6}},
+                {"id": 11, "gender": 0, "name": "Seher Bumb", "locale": "en_IN", "extend": {"rdm": 4}},
+                {"id": 12, "gender": 0, "name": "谢冬梅", "locale": "zh_CN", "extend": {"rdm": 6}},
             ],
             insert_fields=["id", "gender", "name", "locale", "extend"],
-            upsert_fields=["name", "locale"],
+            upsert_fields=["gender", "name"],
+            merge_fields=["extend"],
             using_values=True,
         )
         assert old_sql == """
     INSERT INTO account
       (id, gender, name, locale, extend)
     VALUES
-      (7, 1, '斉藤 修平', 'ja_JP', '{}'),
-      (8, 1, 'Ojas Salvi', 'en_IN', '{}'),
-      (9, 1, '羊淑兰', 'zh_CN', '{}')
-    ON DUPLICATE KEY UPDATE name=VALUES(name), locale=VALUES(locale)
+      (10, 2, '佐々木 美加子', 'ja_JP', '{"rdm": 6}'),
+      (11, 0, 'Seher Bumb', 'en_IN', '{"rdm": 4}'),
+      (12, 0, '谢冬梅', 'zh_CN', '{"rdm": 6}')
+    ON DUPLICATE KEY UPDATE gender=VALUES(gender), name=VALUES(name), extend=JSON_MERGE_PATCH(COALESCE(account.extend, '{}'), VALUES(extend))
 """
 
         new_sql = SQLizer.upsert_on_duplicate(
             self.table,
             [
-                {"id": 7, "gender": 1, "name": "斉藤 修平", "locale": "ja_JP", "extend": {}},
-                {"id": 8, "gender": 1, "name": "Ojas Salvi", "locale": "en_IN", "extend": {}},
-                {"id": 9, "gender": 1, "name": "羊淑兰", "locale": "zh_CN", "extend": {}}
+                {"id": 10, "gender": 1, "name": "田中 知実", "locale": "ja_JP", "extend": {"rdm": 1}},
+                {"id": 11, "gender": 2, "name": "Tara Chadha", "locale": "en_IN", "extend": {"rdm": 10}},
+                {"id": 12, "gender": 2, "name": "吴磊", "locale": "zh_CN", "extend": {"rdm": 9}},
             ],
             insert_fields=["id", "gender", "name", "locale", "extend"],
-            upsert_fields=["name", "locale"],
-            using_values=False,
+            upsert_fields=["gender", "name"],
+            merge_fields=["extend"],
         )
         assert new_sql == """
     INSERT INTO account
       (id, gender, name, locale, extend)
     VALUES
-      (7, 1, '斉藤 修平', 'ja_JP', '{}'),
-      (8, 1, 'Ojas Salvi', 'en_IN', '{}'),
-      (9, 1, '羊淑兰', 'zh_CN', '{}')
-    AS `new_account` ON DUPLICATE KEY UPDATE name=`new_account`.name, locale=`new_account`.locale
+      (10, 1, '田中 知実', 'ja_JP', '{"rdm": 1}'),
+      (11, 2, 'Tara Chadha', 'en_IN', '{"rdm": 10}'),
+      (12, 2, '吴磊', 'zh_CN', '{"rdm": 9}')
+    AS `new_account` ON DUPLICATE KEY UPDATE gender=`new_account`.gender, name=`new_account`.name, extend=JSON_MERGE_PATCH(COALESCE(account.extend, '{}'), `new_account`.extend)
 """
 
         only_insert_sql = SQLizer.upsert_on_duplicate(
             self.table,
             [
-                {"id": 7, "gender": 1, "name": "斉藤 修平", "locale": "ja_JP", "extend": {}},
-                {"id": 8, "gender": 1, "name": "Ojas Salvi", "locale": "en_IN", "extend": {}},
-                {"id": 9, "gender": 1, "name": "羊淑兰", "locale": "zh_CN", "extend": {}}
+                {"id": 10, "gender": 2, "name": "池田 幹", "locale": "ja_JP", "extend": {"rdm": 9}},
+                {"id": 11, "gender": 0, "name": "Sana Mani", "locale": "en_IN", "extend": {"rdm": 0}},
+                {"id": 12, "gender": 0, "name": "刘柳", "locale": "zh_CN", "extend": {"rdm": 9}},
             ],
             insert_fields=["id", "gender", "name", "locale", "extend"],
         )
         assert only_insert_sql == """
     INSERT INTO account
       (id, gender, name, locale, extend)
     VALUES
-      (7, 1, '斉藤 修平', 'ja_JP', '{}'),
-      (8, 1, 'Ojas Salvi', 'en_IN', '{}'),
-      (9, 1, '羊淑兰', 'zh_CN', '{}')
+      (10, 2, '池田 幹', 'ja_JP', '{"rdm": 9}'),
+      (11, 0, 'Sana Mani', 'en_IN', '{"rdm": 0}'),
+      (12, 0, '刘柳', 'zh_CN', '{"rdm": 9}')
 """
 
     def test_insert_into_select(self):