From f4f6f138849d70de4099f54cccf376c2d4b2e540 Mon Sep 17 00:00:00 2001
From: Gavin Haynes <haynesgt@users.noreply.github.com>
Date: Wed, 24 Apr 2024 08:35:09 -0700
Subject: [PATCH] Add support for script as a dict in update (#1562)

---
 elasticsearch_dsl/_async/document.py           | 10 ++++++++--
 elasticsearch_dsl/_sync/document.py            | 10 ++++++++--
 tests/test_integration/_async/test_document.py | 18 ++++++++++++++++++
 tests/test_integration/_sync/test_document.py  | 18 ++++++++++++++++++
 4 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/elasticsearch_dsl/_async/document.py b/elasticsearch_dsl/_async/document.py
index e3cfa38d..23d3c282 100644
--- a/elasticsearch_dsl/_async/document.py
+++ b/elasticsearch_dsl/_async/document.py
@@ -252,6 +252,8 @@ async def update(
         :arg doc_as_upsert:  Instead of sending a partial doc plus an upsert
             doc, setting doc_as_upsert to true will use the contents of doc as
             the upsert value
+        :arg script: the source code of the script as a string, or a dictionary
+            with script attributes to update.
         :arg return_doc_meta: set to ``True`` to return all metadata from the
             index API call instead of only the operation result
 
@@ -268,11 +270,15 @@ async def update(
                 body["upsert"] = upsert
 
             if script:
-                script = {"source": script}
+                if isinstance(script, str):
+                    script = {"source": script}
             else:
                 script = {"id": script_id}
 
-            script["params"] = fields
+            if "params" not in script:
+                script["params"] = fields
+            else:
+                script["params"].update(fields)
 
             body["script"] = script
             body["scripted_upsert"] = scripted_upsert
diff --git a/elasticsearch_dsl/_sync/document.py b/elasticsearch_dsl/_sync/document.py
index 850de9a4..7567ff1a 100644
--- a/elasticsearch_dsl/_sync/document.py
+++ b/elasticsearch_dsl/_sync/document.py
@@ -250,6 +250,8 @@ def update(
         :arg doc_as_upsert:  Instead of sending a partial doc plus an upsert
             doc, setting doc_as_upsert to true will use the contents of doc as
             the upsert value
+        :arg script: the source code of the script as a string, or a dictionary
+            with script attributes to update.
         :arg return_doc_meta: set to ``True`` to return all metadata from the
             index API call instead of only the operation result
 
@@ -266,11 +268,15 @@ def update(
                 body["upsert"] = upsert
 
             if script:
-                script = {"source": script}
+                if isinstance(script, str):
+                    script = {"source": script}
             else:
                 script = {"id": script_id}
 
-            script["params"] = fields
+            if "params" not in script:
+                script["params"] = fields
+            else:
+                script["params"].update(fields)
 
             body["script"] = script
             body["scripted_upsert"] = scripted_upsert
diff --git a/tests/test_integration/_async/test_document.py b/tests/test_integration/_async/test_document.py
index 1c74335b..008f71dc 100644
--- a/tests/test_integration/_async/test_document.py
+++ b/tests/test_integration/_async/test_document.py
@@ -241,6 +241,24 @@ async def test_update_script(async_write_client):
     assert w.views == 47
 
 
+@pytest.mark.asyncio
+async def test_update_script_with_dict(async_write_client):
+    await Wiki.init()
+    w = Wiki(owner=User(name="Honza Kral"), _id="elasticsearch-py", views=42)
+    await w.save()
+
+    await w.update(
+        script={
+            "source": "ctx._source.views += params.inc1 + params.inc2",
+            "params": {"inc1": 2},
+            "lang": "painless",
+        },
+        inc2=3,
+    )
+    w = await Wiki.get(id="elasticsearch-py")
+    assert w.views == 47
+
+
 @pytest.mark.asyncio
 async def test_update_retry_on_conflict(async_write_client):
     await Wiki.init()
diff --git a/tests/test_integration/_sync/test_document.py b/tests/test_integration/_sync/test_document.py
index c5675fd7..ea508b95 100644
--- a/tests/test_integration/_sync/test_document.py
+++ b/tests/test_integration/_sync/test_document.py
@@ -241,6 +241,24 @@ def test_update_script(write_client):
     assert w.views == 47
 
 
+@pytest.mark.sync
+def test_update_script_with_dict(write_client):
+    Wiki.init()
+    w = Wiki(owner=User(name="Honza Kral"), _id="elasticsearch-py", views=42)
+    w.save()
+
+    w.update(
+        script={
+            "source": "ctx._source.views += params.inc1 + params.inc2",
+            "params": {"inc1": 2},
+            "lang": "painless",
+        },
+        inc2=3,
+    )
+    w = Wiki.get(id="elasticsearch-py")
+    assert w.views == 47
+
+
 @pytest.mark.sync
 def test_update_retry_on_conflict(write_client):
     Wiki.init()