diff --git a/invenio_drafts_resources/services/records/service.py b/invenio_drafts_resources/services/records/service.py
index 925690f..820a621 100644
--- a/invenio_drafts_resources/services/records/service.py
+++ b/invenio_drafts_resources/services/records/service.py
@@ -258,10 +258,11 @@ def read_latest(self, identity, id_, expand=False):
 
     @unit_of_work()
     def update_draft(
-        self, identity, id_, data, revision_id=None, uow=None, expand=False
+        self, identity, id_, data, revision_id=None, uow=None, expand=False, draft=None
     ):
         """Replace a draft."""
-        draft = self.draft_cls.pid.resolve(id_, registered_only=False)
+        if draft:
+            draft = self.draft_cls.pid.resolve(id_, registered_only=False)
 
         self.check_revision_id(draft, revision_id)