rara-tools 0.0.4__py3-none-any.whl → 0.0.8__py3-none-any.whl

rara_tools/constants/digitizer.py

@@ -0,0 +1,13 @@
+class StatusKeys:
+    CLEAN_UP = "digitizer_clean_up"
+    ELASTICSEARCH_UPLOAD = "digitizer_elasticsearch_upload"
+    UPLOAD = "s3_upload"
+    DOWNLOAD = "digitizer_s3_download"
+    OCR = "digitizer_ocr"
+
+
+class Queue:
+    IO = "io"
+    DOWNLOAD = "download"
+    FINISH = "finish"
+    OCR = "ocr"

rara_tools/constants/general.py

@@ -0,0 +1,10 @@
+class Status:
+    FAILED = "FAILED"
+    PENDING = "PENDING"
+    RUNNING = "RUNNING"
+    COMPLETED = "COMPLETED"
+    RETRYING = "RETRYING"
+
+
+class Queue:
+    CORE = "core"

rara_tools/elastic.py

@@ -1,6 +1,10 @@
-from typing import Dict, Optional, List
+from typing import Dict, Optional, Any, Iterator
+
+import elasticsearch_dsl
+from elastic_transport import ObjectApiResponse
 from elasticsearch import Elasticsearch
-from elasticsearch_dsl import Search
+from elasticsearch.helpers import bulk
+from elasticsearch_dsl import Index
 
 from .decorators import _elastic_connection
 
@@ -8,11 +12,54 @@ from .decorators import _elastic_connection
 class KataElastic:
     """A class to manage all required Elasticsearch operations for Kata.
     """
+
+    TYPE_MAPPING = {
+        "keyword": elasticsearch_dsl.Keyword,
+        "text": elasticsearch_dsl.Text,
+        "float": elasticsearch_dsl.Float,
+        "integer": elasticsearch_dsl.Integer,
+        "date": elasticsearch_dsl.Date,
+    }
+
+    DEFAULT_MAPPING = {
+        "text": "keyword",
+        "parent_id": "keyword",
+        "text_quality": "float",
+        "n_chars": "integer",
+        "n_words": "integer",
+        "language": "keyword",
+        "end_page": "integer",
+        "start_page": "integer",
+        "sequence_nr": "integer",
+        "section_title": "keyword",
+        "section_type": "keyword",
+        "section_meta": "keyword",
+    }
+
     def __init__(self, elasticsearch_url: str, timeout: Optional[int] = None):
         self.timeout = timeout
         self.elasticsearch_url = elasticsearch_url
         self.elasticsearch = Elasticsearch(self.elasticsearch_url, request_timeout=self.timeout)
 
+    def _produce_rollover_index(self, index_prefix: str, rollover_limit: int) -> str:
+        indices = self.elasticsearch.indices.get(index=f"{index_prefix}-*", expand_wildcards="open")
+        sorted_indices = sorted([(k, v["settings"]["index"]["creation_date"]) for k, v in indices.items()], key=lambda x: x[1], reverse=True)
+        sorted_indices = [i[0] for i in sorted_indices]
+
+        # new index name if none exist
+        if not len(sorted_indices):
+            last_index_name = f"{index_prefix}-0"
+            last_index_count = 0
+        else:
+            last_index_name = sorted_indices[0]
+            last_index_count = self.elasticsearch.count(index=last_index_name)["count"]
+        # check the size of the last index of the pipeline
+        if last_index_count >= rollover_limit:
+            new_index_number = int(last_index_name[-1]) + 1
+            last_index_name = f"{index_prefix}-{new_index_number}"
+
+        return last_index_name
+
     @_elastic_connection
     def check(self) -> bool:
         """Checks Elasticsearch connection.
@@ -22,25 +69,42 @@ class KataElastic:
             return True
         return False
 
+    def generate_mapping(self, schema: dict | None = None) -> dict:
+        mapping_dsl = elasticsearch_dsl.Mapping()
+        mapping = schema or self.DEFAULT_MAPPING
+        for field_name, field_type in mapping.items():
+            if field_type in self.TYPE_MAPPING:
+                # We instantiate the class stored in the type mapping.
+                mapping_dsl.field(field_name, self.TYPE_MAPPING[field_type]())
+        return mapping_dsl.to_dict()
+
+    @_elastic_connection
+    def add_mapping(self, index_name: str, schema: dict):
+        index = Index(name=index_name)
+        return index.put_mapping(body=schema, using=self.elasticsearch)
+
     @_elastic_connection
     def create_index(
             self,
             index: str,
             shards: int = 3,
             replicas: int = 1,
-            settings: Optional[dict] = None
-    ) -> Dict:
+            settings: Optional[dict] = None,
+    ) -> Dict | None:
         """Creates empty index.
         :param: index str: Name of the index to create.
         :param: shards int: Number of shards for the index.
         :param: replicas int: Number of replicas of the index.
         :param: settings dict: Overwrite settings for the index.
         """
-        body = settings or {
-            "number_of_shards": shards,
-            "number_of_replicas": replicas,
-        }
-        return self.elasticsearch.indices.create(index=index, settings=body)
+
+        index_exists = self.elasticsearch.indices.exists(index=index).body
+        if index_exists is False:
+            setting_body = settings or {
+                "number_of_shards": shards,
+                "number_of_replicas": replicas,
+            }
+            return self.elasticsearch.indices.create(index=index, settings=setting_body)
 
     @_elastic_connection
     def delete_index(self, index: str, ignore: Optional[bool] = True) -> Dict:
@@ -49,11 +113,11 @@ class KataElastic:
         :param: ignore bool: Ignore errors because of closed/deleted index.
         :return: Dict of Elastic's acknowledgement of the action.
         """
-        response = self.elasticsearch.indices.delete(index=index, ignore_unavailable=ignore)
+        response = self.elasticsearch.indices.delete(index=index, ignore_unavailable=ignore, expand_wildcards="open")
         return response
 
     @_elastic_connection
-    def delete_document(self, index: str, document_id: str) -> Dict:
+    def delete_document(self, index: str, document_id: str) -> ObjectApiResponse[Any]:
         """Deletes document fom index.
         :param: document_id str: ID of the document to be deleted.
         :param: index str: Index where the document is to be found.
@@ -63,6 +127,25 @@ class KataElastic:
         response = self.elasticsearch.delete(id=document_id, index=index)
         return response
 
+    @_elastic_connection
+    def bulk_index(
+            self,
+            documents: Iterator[dict],
+            index_prefix: str,
+            rollover_limit: int,
+            refresh="false",
+            create_index: bool = True
+    ) -> (int, int):
+        last_index_name = self._produce_rollover_index(index_prefix, rollover_limit)
+        if create_index:
+            response = self.create_index(index=last_index_name)
+            response = self.add_mapping(index_name=last_index_name, schema=self.generate_mapping())
+            pass
+
+        actions = [{"_index": last_index_name, "_source": document} for document in documents]
+        successful_count, error_count = bulk(actions=actions, client=self.elasticsearch, max_retries=3, refresh=refresh)
+        return successful_count, error_count
+
     @_elastic_connection
     def index_document(self, index: str, body: dict, document_id: Optional[str] = None) -> Dict:
         """Indexes document.
@@ -78,15 +161,15 @@ class KataElastic:
         return indexed
 
     @_elastic_connection
-    def get_documents_by_key(self, index: str, document_key: str) -> List:
-        """This method is for retrieving all texts/pages of the original document.
-        :param: index str: Index to search the documents from.
-        :param: document_key str: parent_id field that connects pages of document together.
-        :return: List of matching documents. 
-        """
-        s = Search(using=self.elasticsearch, index=index)
-        docs = s.query("match", parent_id=document_key).execute()
-        return docs
+    def get_documents_by_key(self, index: str, document_key: str, sort_fields=("start_page", "end_page", "sequence_nr",)):
+        index = f"{index}-*"
+        s = elasticsearch_dsl.Search(using=self.elasticsearch, index=index)
+        s = s.query("match", parent_id=document_key).sort(*sort_fields)
+        # Since scan doesn't allow for sorting, we do it manually after fetching the documents.
+        documents = sorted(
+            s.scan(), key=lambda doc: [getattr(doc, field) for field in sort_fields]
+        )
+        return documents
 
     def __str__(self) -> str:
-         return self.elasticsearch_url
+        return self.elasticsearch_url

{rara_tools-0.0.4.dist-info → rara_tools-0.0.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: rara-tools
-Version: 0.0.4
+Version: 0.0.8
 Summary: Tools to support Kata's work.
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10

rara_tools-0.0.8.dist-info/RECORD

@@ -0,0 +1,13 @@
+rara_tools/decorators.py,sha256=rYDk5CEHhCZvqeFaHku8qLMv7G7NTMWppHwLg3ZeVj4,2186
+rara_tools/elastic.py,sha256=nNlCmoyKfCkM_2r1jtbpSpUn4S8IrLOKak17QwhNSvs,7146
+rara_tools/exceptions.py,sha256=FtuHG-2snaEfADA25HjjutGNQzNo6sTdSfqk9VrzOuE,374
+rara_tools/s3.py,sha256=eqMiOKbjXvXY04JJV68gmOU-4DUnwEaeYdhjQSI6crU,4440
+rara_tools/task_reporter.py,sha256=WCcZts9dAUokPc4vbrG3-lNAFLnWaMgE3b3iaUB7mr8,3256
+rara_tools/constants/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+rara_tools/constants/digitizer.py,sha256=gJ3jOMwuZfKcLqgOAxTyB266VYsskLabJiMUiSz3xX4,297
+rara_tools/constants/general.py,sha256=E9Jaw-YxocS_tOZw9QBoxO3e9KK5EMbLoM0R7D4Iflw,171
+rara_tools-0.0.8.dist-info/LICENSE.md,sha256=hkZVnIZll7e_KNEQzeY94Y9tlzVL8iVZBTMBvDykksU,35142
+rara_tools-0.0.8.dist-info/METADATA,sha256=TMrOrd_YtH83jCAzbNpBrHlcN7ta6VQwYBD_HqH3unM,3820
+rara_tools-0.0.8.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+rara_tools-0.0.8.dist-info/top_level.txt,sha256=JwfB5b8BAtW5OFKRln2AQ_WElTRyIBM4nO0FKN1cupY,11
+rara_tools-0.0.8.dist-info/RECORD,,

rara_tools-0.0.4.dist-info/RECORD

@@ -1,10 +0,0 @@
-rara_tools/decorators.py,sha256=rYDk5CEHhCZvqeFaHku8qLMv7G7NTMWppHwLg3ZeVj4,2186
-rara_tools/elastic.py,sha256=MVqai6wDQlDQeHQzAKsRpxOchI29y3W1UiridgfH6d4,3718
-rara_tools/exceptions.py,sha256=FtuHG-2snaEfADA25HjjutGNQzNo6sTdSfqk9VrzOuE,374
-rara_tools/s3.py,sha256=eqMiOKbjXvXY04JJV68gmOU-4DUnwEaeYdhjQSI6crU,4440
-rara_tools/task_reporter.py,sha256=WCcZts9dAUokPc4vbrG3-lNAFLnWaMgE3b3iaUB7mr8,3256
-rara_tools-0.0.4.dist-info/LICENSE.md,sha256=hkZVnIZll7e_KNEQzeY94Y9tlzVL8iVZBTMBvDykksU,35142
-rara_tools-0.0.4.dist-info/METADATA,sha256=onb2qPn7IXknjCM09yxKfuEZYJlTlUwetCrS2ZXlPio,3820
-rara_tools-0.0.4.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-rara_tools-0.0.4.dist-info/top_level.txt,sha256=JwfB5b8BAtW5OFKRln2AQ_WElTRyIBM4nO0FKN1cupY,11
-rara_tools-0.0.4.dist-info/RECORD,,