From a61e38abb0effafe46516b013d6e6127f995109c Mon Sep 17 00:00:00 2001
From: Aly Badr <aly.badr@rero.ch>
Date: Mon, 21 Oct 2019 07:51:14 +0200
Subject: [PATCH] cli: reserve a range of pids

* Assigns all unused pids the status NEW.
* Assigns all reserved pids the status RESERVED.
* Improves method to insert identifer pids and reset sequence.

Co-Authored-by: Aly Badr <aly.badr@rero.ch>
---
 rero_ils/modules/cli.py                       | 61 ++++++++++++++++++-
 rero_ils/modules/providers.py                 | 12 ++--
 .../organisations/test_organisations_api.py   |  2 +-
 3 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/rero_ils/modules/cli.py b/rero_ils/modules/cli.py
index 93cbc97979..2b3acf2426 100644
--- a/rero_ils/modules/cli.py
+++ b/rero_ils/modules/cli.py
@@ -38,7 +38,8 @@
 from flask.cli import with_appcontext
 from flask_security.confirmable import confirm_user
 from invenio_accounts.cli import commit, users
-from invenio_pidstore.models import PersistentIdentifier
+from invenio_db import db
+from invenio_pidstore.models import PersistentIdentifier, PIDStatus
 from invenio_records.api import Record
 from invenio_records_rest.utils import obj_or_import_string
 from invenio_search.cli import es_version_check
@@ -299,9 +300,9 @@ def create(infile, pid_type, schema, verbose, dbcommit, reindex, append):
             error_file.write(']')
 
     if append:
-        table = record_class.provider.identifier
+        identifier = record_class.provider.identifier
         try:
-            append_fixtures_new_identifiers(table, sorted(pids))
+            append_fixtures_new_identifiers(identifier, sorted(pids), pid_type)
         except Exception as err:
             pass
 
@@ -783,3 +784,57 @@ def marc21json(xml_file, json_file_ok, xml_file_error, parallel, chunk,
     if count_ko:
         click.secho('Records with errors: ', fg='red', nl=False)
         click.secho(str(count_ko))
+
+
+@utils.command('reserve_pid_range')
+@click.option('-p', '--pid_type', 'pid_type', default=None,
+              help='pid type of the resource')
+@click.option('-n', '--records_number', 'records_number', default=None,
+              help='Number of records to load')
+@click.option('-u', '--unused', 'unused', is_flag=True, default=False,
+              help='Set unused (gaps) pids status to NEW ')
+@with_appcontext
+def reserve_pid_range(pid_type, records_number, unused):
+    """Reserve a range of pids for future records loading.
+
+    reserved pids will have the status RESERVED.
+    - pid_type: the pid type of the resource as configured in config.py
+    - records_number: number of new records(with pids) to load.
+    - unused: set that the status of unused (gaps) pids to NEW.
+    """
+    click.secho('Reserving pids for loading "%s" records' %
+                pid_type, fg='green')
+    try:
+        records_number = int(records_number)
+    except ValueError:
+        raise ValueError('Parameter records_number must be integer.')
+
+    try:
+        record_class = obj_or_import_string(
+            current_app.config
+            .get('RECORDS_REST_ENDPOINTS')
+            .get(pid_type).get('record_class', Record))
+    except AttributeError:
+        raise AttributeError('Invalid pid type.')
+
+    identifier = record_class.provider.identifier
+    reserved_pids = []
+    for number in range(0, records_number):
+        pid = identifier.next()
+        reserved_pids.append(pid)
+        record_class.provider.create(pid_type, pid_value=pid,
+                                     status=PIDStatus.RESERVED)
+        db.session.commit()
+    click.secho(
+        ('reserved_pids range, from: {min} to: {max}').format(
+            min=min(reserved_pids), max=max(reserved_pids)
+        ))
+    if unused:
+        for pid in range(1, identifier.max()):
+            if not db.session.query(
+                    identifier.query.filter(identifier.recid == pid).exists()
+            ).scalar():
+                record_class.provider.create(pid_type, pid_value=pid,
+                                             status=PIDStatus.NEW)
+                db.session.add(identifier(recid=pid))
+            db.session.commit()
diff --git a/rero_ils/modules/providers.py b/rero_ils/modules/providers.py
index cccd50b94d..2dc0b3c3fe 100644
--- a/rero_ils/modules/providers.py
+++ b/rero_ils/modules/providers.py
@@ -21,17 +21,19 @@
 
 from invenio_db import db
 from invenio_pidstore.errors import PIDDoesNotExistError
-from invenio_pidstore.models import PIDStatus
+from invenio_pidstore.models import PersistentIdentifier, PIDStatus
 from invenio_pidstore.providers.base import BaseProvider
+from sqlalchemy import text
 
 
-def append_fixtures_new_identifiers(table, pids):
+def append_fixtures_new_identifiers(identifier, pids, pid_type):
     """Insert pids into the indentifier table and update its sequence."""
     for pid in pids:
-        data = table(recid=pid)
-        db.session.add(data)
+        db.session.add(identifier(recid=pid))
     db.session.commit()
-    table._set_sequence(int(max(pids))+1)
+    max_pid = PersistentIdentifier.query.filter_by(
+        pid_type=pid_type).order_by(text('pid_value desc')).first().id
+    identifier._set_sequence(max_pid)
 
 
 class Provider(BaseProvider):
diff --git a/tests/ui/organisations/test_organisations_api.py b/tests/ui/organisations/test_organisations_api.py
index 24ed620b24..4f11e30eb4 100644
--- a/tests/ui/organisations/test_organisations_api.py
+++ b/tests/ui/organisations/test_organisations_api.py
@@ -53,5 +53,5 @@ def test_organisation_create(app, db, org_martigny_data, org_sion_data):
     assert org.get('pid') == '2'
 
     identifier = Organisation.provider.identifier
-    append_fixtures_new_identifiers(identifier, ['1', '2'])
+    append_fixtures_new_identifiers(identifier, ['1', '2'], 'org')
     assert identifier.next() == identifier.max() == 3