streamlit-octostar-utils 0.5.0.dev5__tar.gz → 0.5.0.dev7__tar.gz

{streamlit_octostar_utils-0.5.0.dev5 → streamlit_octostar_utils-0.5.0.dev7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.5.0.dev5
+Version: 0.5.0.dev7
 Summary: 
 License: MIT
 License-File: LICENSE

{streamlit_octostar_utils-0.5.0.dev5 → streamlit_octostar_utils-0.5.0.dev7}/pyproject.toml

@@ -5,7 +5,7 @@ include = '\.pyi?$'
 
 [tool.poetry]
 name = "streamlit-octostar-utils"
-version = "0.5.0-dev.5"
+version = "0.5.0-dev.7"
 description = ""
 license = "MIT"
 authors = ["Octostar"]

{streamlit_octostar_utils-0.5.0.dev5 → streamlit_octostar_utils-0.5.0.dev7}/streamlit_octostar_utils/api_crafter/nifi.py

@@ -45,6 +45,7 @@ TAG_RELATIONSHIP = RelationshipName(name="has_tag", type="mtm")
 FRAGMENT_RELATIONSHIP = RelationshipName(name="is_fragment_of", type="otm")
 PREVIOUS_FRAGMENT_RELATIONSHIP = RelationshipName(name="has_previous_fragment", type="otm")
 NEXT_FRAGMENT_RELATIONSHIP = RelationshipName(name="has_next_fragment", type="otm")
+SOURCE_FRAGMENT_ENTITY_RELATIONSHIP = RelationshipName(name="is_child_fragment_of", type="otm")
 OS_RESERVED_FIELDS = [
     "os_entity_uid", "entity_id", "entity_type", "os_concept",
     "entity_label", "os_created_at", "os_created_by",
@@ -212,6 +213,7 @@ class NifiFragmenter(object):
             raise ValueError("Cannot have more than 100k entities for fragmentation")
         identifier = str(uuid.uuid4())
         root_uid = fragments[0].record["os_entity_uid"]
+        root_type = fragments[0].record["os_concept"]
         for i, entity in enumerate(fragments):
             travel_dict(entity.request["nifi_attributes"], fragmenter_keylist.split("."), "w")(
                 {"identifier": identifier, "count": count, "index": i}
@@ -225,7 +227,7 @@ class NifiFragmenter(object):
                 "fragments_stack"
             ]
             travel_dict(entity.request["config"]["fragment"], fragmenter_keylist.split("."), "w")(
-                {"identifier": identifier, "count": count, "index": i, "root_uid": root_uid}
+                {"identifier": identifier, "count": count, "index": i, "root_uid": root_uid, "root_type": root_type}
             )
 
     def push_defragment_strategy(fragment, defragmenter_config):
@@ -424,30 +426,34 @@ class NifiFragmenter(object):
         )
 
     @staticmethod
-    def resolve_source_entity_uid(entity, fragment_root_source=None) -> str:
-        """Resolve the source_entity_uid to use for child fragments.
+    def resolve_source_entity(entity, fragment_root_source=None) -> tuple:
+        """Resolve the source entity (recursive originator) for child fragments.
 
         When fragment_root_source is set (a fragment name or stack index),
-        the UID is looked up via entity.get_fragment_root_uid() -- this is
-        necessary when the current entity is a clone that may not be persisted.
-        Otherwise falls back to the entity's own UID.
+        the root is looked up via entity.get_fragment_root().  When it is
+        None and the entity is already fragmented, the oldest fragmentation
+        level's root is used (the original non-fragment ancestor).  Otherwise
+        falls back to the entity itself.
 
         Args:
             entity: A NifiEntity or NifiEntityProxy.
-            fragment_root_source: None to use the entity's own UID, an int
+            fragment_root_source: None for automatic resolution, an int
                 to index into the fragments stack, or a string fragmenter
                 keylist name.
 
         Returns:
-            The resolved source entity UID string.
+            (source_entity_uid, source_entity_type) tuple.
         """
-        if fragment_root_source is None:
-            return entity.record["os_entity_uid"]
-        try:
-            idx = int(fragment_root_source)
-            return entity.get_fragment_root_uid(idx)
-        except (ValueError, TypeError):
-            return entity.get_fragment_root_uid(fragment_root_source)
+        if fragment_root_source is not None:
+            try:
+                idx = int(fragment_root_source)
+                return entity.get_fragment_root(idx)
+            except (ValueError, TypeError):
+                return entity.get_fragment_root(fragment_root_source)
+        stack = entity.request.get("config", {}).get("fragment", {}).get("fragments_stack", [])
+        if stack:
+            return entity.get_fragment_root(-1)
+        return (entity.record["os_entity_uid"], entity.record["os_concept"])
 
 
 class NifiEntityBatch(object):
@@ -519,7 +525,7 @@ class NifiContextManager(object):
             )
         entities = sorted(
             entities,
-            key=lambda x: string_to_datetime(x.record.get("os_last_updated_at")),
+            key=lambda x: string_to_datetime(x.request.get("entity_timestamp")),
         )
         entities = list({e.record["entity_id"]: e for e in entities}.values())
         entities = [
@@ -592,7 +598,7 @@ class NifiContextManager(object):
             all_entities.extend(_process_entity(entity, processor_name))
         all_entities = sorted(
             all_entities,
-            key=lambda x: string_to_datetime(x.record.get("os_last_updated_at")),
+            key=lambda x: string_to_datetime(x.request.get("entity_timestamp")),
         )
         self.out_entities = list({e.record["entity_id"]: e for e in all_entities}.values())
         self.sync_entities()
@@ -616,34 +622,68 @@ class NifiContextManager(object):
     @contextmanager
     def reindex_lock(self, entities: List[Union[dict, "NifiEntity"]], timeout: int = 7200):
         if not entities:
-            yield
+            yield True
             return
-        entities = [e if isinstance(e, dict) else e.record for e in entities]
+
+        import logging
+        _lock_logger = logging.getLogger(__name__)
+
+        records = []
+        for e in entities:
+            if isinstance(e, dict):
+                records.append(
+                    (e, e.get("entity_timestamp"))
+                )
+            else:
+                records.append(
+                    (e.record, e.request.get("entity_timestamp") if e.request else None)
+                )
+
         long_expiry = (datetime.now(timezone.utc) + timedelta(seconds=timeout)).strftime("%Y-%m-%dT%H:%M:%SZ")
         statuses = [
             {
-                "entity_id": e["os_entity_uid"],
-                "entity_type": e["os_concept"],
+                "entity_id": rec["os_entity_uid"],
+                "entity_type": rec["os_concept"],
+                "do_not_update_before": ts,
                 "do_not_reindex_before": long_expiry,
-                "do_not_update_before": long_expiry,
             }
-            for e in entities
+            for rec, ts in records
+            if ts is not None
         ]
-        update_processing_status.sync(statuses, client=self.client)
+        if len(statuses) < len(records):
+            missing = [rec["os_entity_uid"] for rec, ts in records if ts is None]
+            _lock_logger.warning(
+                f"reindex_lock: {len(missing)} entities have no entity_timestamp, cannot set lock: {missing}"
+            )
+        if not statuses:
+            yield True
+            return
+
+        result = update_processing_status.sync(statuses, client=self.client)
+        skipped_ids = result.get("skipped_entity_ids", []) if isinstance(result, dict) else []
+        if skipped_ids:
+            _lock_logger.warning(
+                f"reindex_lock: {len(skipped_ids)} entities have a newer pipeline run, aborting sync: {skipped_ids}"
+            )
+            yield False
+            return
+
         try:
-            yield
+            yield True
         finally:
             short_expiry = (datetime.now(timezone.utc) + timedelta(seconds=1)).strftime("%Y-%m-%dT%H:%M:%SZ")
             statuses = [
                 {
-                    "entity_id": e["os_entity_uid"],
-                    "entity_type": e["os_concept"],
+                    "entity_id": rec["os_entity_uid"],
+                    "entity_type": rec["os_concept"],
+                    "do_not_update_before": ts,
                     "do_not_reindex_before": short_expiry,
-                    "do_not_update_before": short_expiry,
                 }
-                for e in entities
+                for rec, ts in records
+                if ts is not None
             ]
-            update_processing_status.sync(statuses, client=self.client)
+            if statuses:
+                update_processing_status.sync(statuses, client=self.client)
 
     def sync_entities(self):
         if not self.lazy_sync:
@@ -658,14 +698,17 @@ class NifiContextManager(object):
         fetch_rel_entities, fetch_concept_rels = self._find_relationships_to_fetch(entities)
 
         all_entities_to_modify = attachables_to_write + [e[0] for e in entities_to_upsert]
-        with self.reindex_lock(all_entities_to_modify):
+        with self.reindex_lock(all_entities_to_modify) as locked:
+            if not locked:
+                return
             self._sync_write_attachables(attachables_to_write)
             self._sync_upsert_entities(entities_to_upsert)
         self._sync_fetch_relationships(entities, fetch_rel_entities, fetch_concept_rels)
 
+        now_ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
         for entity in all_entities_to_modify:
             entity.request["is_temporary"] = False
-            entity.request["entity_timestamp"] = entity.record["os_last_updated_at"]
+            entity.request["entity_timestamp"] = now_ts
 
         for entity in entities:
             entity.sync_params = {}
@@ -712,7 +755,7 @@ class NifiContextManager(object):
                 file.record["entity_id"] = file.record["os_entity_uid"]
                 file.record["entity_type"] = file.record["os_concept"]
                 file.record["entity_label"] = file.label
-            file.request["entity_timestamp"] = file.record["os_last_updated_at"]
+            file.request["entity_timestamp"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
             file.request["is_temporary"] = False
 
     def _sync_upsert_entities(self, entities_to_upsert):
@@ -736,7 +779,7 @@ class NifiContextManager(object):
             entity.record["entity_id"] = entity.record["os_entity_uid"]
             entity.record["entity_type"] = entity.record["os_concept"]
             entity.record["entity_label"] = entity.label
-            entity.request["entity_timestamp"] = entity.record["os_last_updated_at"]
+            entity.request["entity_timestamp"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
             entity.request["is_temporary"] = False
 
     def _find_relationships_to_fetch(self, entities) -> dict:
@@ -1052,7 +1095,17 @@ class NifiEntity(object):
             )
         return NifiFragmenter.get_fragment_info(self, key).get("index", 0) == 0
 
-    def get_fragment_root_uid(self, fragment_name_or_idx) -> str:
+    def get_fragment_root(self, fragment_name_or_idx) -> tuple:
+        """Return (root_uid, root_type) for a given fragmentation level.
+
+        Args:
+            fragment_name_or_idx: An int to index into fragments_stack
+                (e.g. -1 for the oldest fragmentation), or a string
+                fragmenter keylist name.
+
+        Returns:
+            (root_uid, root_type) tuple for the root entity at that level.
+        """
         fragment_config = self.request.get("config", {}).get("fragment", {})
         fragments_stack = fragment_config.get("fragments_stack", [])
         if isinstance(fragment_name_or_idx, int):
@@ -1077,7 +1130,10 @@ class NifiEntity(object):
             raise KeyError(
                 f"No root_uid found in fragment config for '{fragment_key}'"
             )
-        return fragment_data["root_uid"]
+        return (
+            fragment_data["root_uid"],
+            fragment_data.get("root_type", "os_fragment"),
+        )
 
     def to_json(self):
         if self.drop_on_output:
@@ -1111,7 +1167,7 @@ class NifiEntity(object):
         proxy_otm_children = [c.record for c in proxy_otm_children]
         full_entity_children = sorted(
             full_entity_children,
-            key=lambda x: string_to_datetime(x.record.get("os_last_updated_at")),
+            key=lambda x: string_to_datetime(x.request.get("entity_timestamp")),
         )
         full_entity_children = list({c.uid: c.to_json() for c in full_entity_children}.values())
         children = full_entity_children + proxy_entity_children + proxy_otm_children
@@ -1144,7 +1200,7 @@ class NifiEntity(object):
                 "relationships": ont_relationships,
                 "label_keys": ont_label_keys,
             },
-            "entity_timestamp": None,
+            "entity_timestamp": self.request.get("entity_timestamp"),
             "sync_params": {
                 NifiContextManager.SyncFlag.UPSERT_ENTITY_ALL.name: True,
             },
@@ -1371,6 +1427,7 @@ class NifiEntity(object):
         self,
         os_workspace,
         source_entity_uid,
+        source_entity_type,
         filename,
         filetype,
         file: Union[Contents, bytes],
@@ -1385,9 +1442,12 @@ class NifiEntity(object):
         next_fragment_uid=None,
         next_fragment_relationship_uid=None,
         next_fragment_relationship=NEXT_FRAGMENT_RELATIONSHIP,
+        source_entity_relationship_uid=None,
+        source_entity_relationship=SOURCE_FRAGMENT_ENTITY_RELATIONSHIP,
     ):
         fields = {
             **fields,
+            "os_parent_uid": self.record["os_entity_uid"],
             "source_entity_uid": source_entity_uid,
             "previous_entity_uid": previous_fragment_uid,
             "next_entity_uid": next_fragment_uid,
@@ -1403,9 +1463,14 @@ class NifiEntity(object):
             os_relationship_uid,
             os_entity_type,
         )
-        prev_rel, next_rel = None, None
         fragment_uid = child_entity.record["os_entity_uid"]
         fragment_type = child_entity.record["os_concept"]
+        source_rel = self._add_relationship(
+            os_workspace, fragment_uid, fragment_type,
+            source_entity_uid, source_entity_type,
+            source_entity_relationship, {}, source_entity_relationship_uid,
+        )
+        prev_rel, next_rel = None, None
         if previous_fragment_uid:
             prev_rel = self._add_relationship(
                 os_workspace, fragment_uid, fragment_type,
@@ -1418,7 +1483,7 @@ class NifiEntity(object):
                 next_fragment_uid, os_entity_type,
                 next_fragment_relationship, {}, next_fragment_relationship_uid,
             )
-        return child_entity, child_rel, prev_rel, next_rel
+        return child_entity, child_rel, source_rel, prev_rel, next_rel
 
     def add_tag(self, os_workspace, name, group, order, color, fields={}):
         return self.add_child_entity(
@@ -1445,12 +1510,6 @@ class NifiEntity(object):
         to_entity.annotations = recursive_update_dict(to_entity.annotations, annotations_to_propagate, merge_method)
 
 
-def more_recent_than(record_a, record_b):
-    return string_to_datetime(record_a.get("os_last_updated_at")) >= string_to_datetime(
-        record_b.get("os_last_updated_at")
-    )
-
-
 class NifiRoute(Route):
     def __init__(self, app, tasks_routes, processor_name, celery_executor, router=None):
         self.app = app