From 64797f3bbf0f2408cd83fe2577b1a1dc38f151e7 Mon Sep 17 00:00:00 2001
From: Sergey Motornyuk <sergey.motornyuk@linkdigital.com.au>
Date: Thu, 7 Dec 2023 22:20:50 +0200
Subject: [PATCH] chore: allow immediate index

---
 ckanext/federated_index/logic/action.py | 39 ++++++++++++++++++++++---
 ckanext/federated_index/logic/auth.py   |  7 +++++
 ckanext/federated_index/logic/schema.py | 11 +++++++
 ckanext/federated_index/storage.py      |  4 ++-
 4 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/ckanext/federated_index/logic/action.py b/ckanext/federated_index/logic/action.py
index 9e9761a..92761e0 100644
--- a/ckanext/federated_index/logic/action.py
+++ b/ckanext/federated_index/logic/action.py
@@ -55,6 +55,10 @@ def federated_index_profile_refresh(
 
     for pkg in profile.fetch_packages(payload):
         db.add(pkg["id"], pkg)
+        if data_dict["index"]:
+            tk.get_action("federated_index_profile_index")(
+                tk.fresh_context(context), {"profile": profile, "ids": [pkg["id"]]}
+            )
 
     return {
         "profile": profile.id,
@@ -125,8 +129,9 @@ def federated_index_profile_index(
 
         try:
             model.Session.flush()
+
         except IntegrityError:
-            log.exception("Cannot index package %s", pkg_dict["name"])
+            log.exception("Cannot index package %s", pkg_dict["id"])
             model.Session.rollback()
             continue
 
@@ -136,10 +141,13 @@ def federated_index_profile_index(
         try:
             package_index.remove_dict(pkg_dict)
             package_index.update_dict(pkg_dict, True)
-        except (search.SearchIndexError, TypeError):
-            log.exception("Cannot index package %s", pkg_dict["name"])
+
+        except (search.SearchIndexError, TypeError, tk.ObjectNotFound):
+            log.exception("Cannot index package %s", pkg_dict["id"])
+
         else:
-            log.debug("Successfully indexed package %s", pkg_dict["name"])
+            log.debug("Successfully indexed package %s", pkg_dict["id"])
+
         finally:
             model.Session.rollback()
 
@@ -178,3 +186,26 @@ def federated_index_profile_clear(
         "profile": data_dict["profile"].id,
         "count": resp.hits,
     }
+
+
+@validate(schema.profile_refresh)
+def federated_index_profile_remove(
+    context: Any,
+    data_dict: dict[str, Any],
+) -> dict[str, Any]:
+    """Remove stored data for the given profile.
+
+    Args:
+        profile(str|Profile): name of the profile or Profile instance
+    """
+    tk.check_access("federated_index_profile_remove", context, data_dict)
+    profile: shared.Profile = data_dict["profile"]
+
+    db = storage.get_storage(profile)
+    count = db.count()
+    db.reset()
+
+    return {
+        "profile": profile.id,
+        "count": count,
+    }
diff --git a/ckanext/federated_index/logic/auth.py b/ckanext/federated_index/logic/auth.py
index 9d843dd..b0c3c1e 100644
--- a/ckanext/federated_index/logic/auth.py
+++ b/ckanext/federated_index/logic/auth.py
@@ -38,3 +38,10 @@ def federated_index_profile_clear(
     data_dict: dict[str, Any],
 ) -> Any:
     return authz.is_authorized("federated_index_access", context, data_dict)
+
+
+def federated_index_profile_remove(
+    context: Any,
+    data_dict: dict[str, Any],
+) -> Any:
+    return authz.is_authorized("federated_index_access", context, data_dict)
diff --git a/ckanext/federated_index/logic/schema.py b/ckanext/federated_index/logic/schema.py
index 657534b..9aa9e97 100644
--- a/ckanext/federated_index/logic/schema.py
+++ b/ckanext/federated_index/logic/schema.py
@@ -18,6 +18,7 @@ def profile_refresh(
         "reset": [boolean_validator],
         "search_payload": [default("{}"), convert_to_json_if_string, dict_only],
         "since_last_refresh": [boolean_validator],
+        "index": [boolean_validator],
     }
 
 
@@ -56,3 +57,13 @@ def profile_clear(
     return {
         "profile": [not_empty, federated_index_profile],
     }
+
+
+@validator_args
+def profile_remove(
+    not_empty: types.Validator,
+    federated_index_profile: types.Validator,
+) -> types.Schema:
+    return {
+        "profile": [not_empty, federated_index_profile],
+    }
diff --git a/ckanext/federated_index/storage.py b/ckanext/federated_index/storage.py
index 359ded2..ab99194 100644
--- a/ckanext/federated_index/storage.py
+++ b/ckanext/federated_index/storage.py
@@ -114,6 +114,7 @@ def count(self):
     def reset(self):
         stmt = sa.delete(Record).where(Record.profile_id == self.profile.id)
         model.Session.execute(stmt)
+        model.Session.commit()
 
     def scan(
         self, offset: int = 0, limit: int | None = None
@@ -132,4 +133,5 @@ def scan(
             yield pkg
 
     def get(self, id: str) -> dict[str, Any] | None:
-        return Record.get(id, self.profile.id)
+        if record := Record.get(id, self.profile.id):
+            return record.data