cli: add possibility to specify a vanity PID for new drafts

* note: this is *not* recommended, but can come in handy if the
        assignment of a specific PID is a requirement

invenio_utilities_tuw/cli/drafts.py

@@ -18,7 +18,13 @@ from flask.cli import with_appcontext
 from invenio_files_rest.models import ObjectVersion
 
 from ..utils import get_draft_file_service, get_record_service
-from .options import option_as_user, option_owners, option_pid_type, option_pid_value
+from .options import (
+    option_as_user,
+    option_owners,
+    option_pid_type,
+    option_pid_value,
+    option_vanity_pid,
+)
 from .utils import (
     convert_to_recid,
     create_record_from_metadata,
@@ -69,8 +75,9 @@ def list_drafts(user):
     help="publish the draft after creation (default: false)",
 )
 @option_owners
+@option_vanity_pid
 @with_appcontext
-def create_draft(metadata_path, publish, user, owners):
+def create_draft(metadata_path, publish, user, owners, vanity_pid):
     """Create a new record draft with the specified metadata.
 
     The specified metadata path can either point to a JSON file containing the metadata,
@@ -91,7 +98,7 @@ def create_draft(metadata_path, publish, user, owners):
         if owners:
             metadata = set_record_owners(metadata, owners)
 
-        draft = create_record_from_metadata(metadata, identity)
+        draft = create_record_from_metadata(metadata, identity, vanity_pid=vanity_pid)
         recid = draft["id"]
 
     elif isdir(metadata_path):

invenio_utilities_tuw/cli/options.py

@@ -58,6 +58,15 @@ option_owners = click.option(
     help="email address of the record owner to set (can be specified multiple times)",
 )
 
+option_vanity_pid = click.option(
+    "--vanity-pid",
+    "-V",
+    "vanity_pid",
+    metavar="PID_VALUE",
+    required=False,
+    help="vanity PID, to assign to the object (not recommended)",
+)
+
 # user management options
 
 option_only_list_active_users = click.option(

invenio_utilities_tuw/cli/utils.py

@@ -14,6 +14,7 @@ from flask_principal import Identity
 from invenio_access import any_user
 from invenio_access.utils import get_identity
 from invenio_accounts import current_accounts
+from invenio_db import db
 from invenio_pidstore.models import PersistentIdentifier
 
 from ..utils import get_record_service
@@ -31,10 +32,36 @@ def read_metadata(metadata_file_path):
     return metadata
 
 
-def create_record_from_metadata(metadata, identity):
+def create_record_from_metadata(
+    metadata, identity, vanity_pid=None, vanity_pid_type="recid"
+):
     """Create a draft from the specified metadata."""
     service = get_record_service()
+
+    if vanity_pid is not None:
+        # check if the vanity PID is already taken, before doing anything stupid
+        count = PersistentIdentifier.query.filter_by(
+            pid_value=vanity_pid, pid_type=vanity_pid_type
+        ).count()
+
+        if count > 0:
+            raise Exception(
+                "PID '{}:{}' is already taken".format(vanity_pid_type, vanity_pid)
+            )
+
     draft = service.create(identity=identity, data=metadata)
+    actual_draft = draft._record if hasattr(draft, "_record") else draft
+
+    if vanity_pid:
+        # service.update_draft() is called to update the IDs in the record's metadata
+        # (via record.commit()), re-index the record, and commit the db session
+        if service.indexer:
+            service.indexer.delete(actual_draft)
+
+        actual_draft.pid.pid_value = vanity_pid
+        db.session.commit()
+        draft = service.update_draft(vanity_pid, identity=identity, data=metadata)
+
     return draft