streamlit-octostar-utils 0.2.12a1__tar.gz → 0.2.13__tar.gz

{streamlit_octostar_utils-0.2.12a1 → streamlit_octostar_utils-0.2.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.2.12a1
+Version: 0.2.13
 Summary: 
 License: MIT
 License-File: LICENSE

{streamlit_octostar_utils-0.2.12a1 → streamlit_octostar_utils-0.2.13}/pyproject.toml

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

{streamlit_octostar_utils-0.2.12a1 → streamlit_octostar_utils-0.2.13}/streamlit_octostar_utils/api_crafter/nifi.py

@@ -163,12 +163,16 @@ class NifiEntityProxy(object):
                     if child_entity.uid == uid_to_search:
                         found_entity = child_entity
                     else:
-                        found_entity = _recursive_search_expanded_proxy(child_entity._proxy, uid_to_search)
+                        found_entity = _recursive_search_expanded_proxy(
+                            child_entity._proxy, uid_to_search
+                        )
                 if found_entity:
                     return found_entity
 
         if not self._proxy:
-            main_entities = itertools.chain(*[b.entities for b in self.context.in_batches])
+            main_entities = itertools.chain(
+                *[b.entities for b in self.context.in_batches]
+            )
             main_entities = {e.record["entity_id"]: e for e in main_entities}
             if main_entities.get(self.uid):
                 self._proxy = main_entities.get(self.uid)
@@ -179,7 +183,9 @@ class NifiEntityProxy(object):
                     self._proxy = found_entity._proxy
                     return self._proxy
             ## TODO: Try to get the entity from the database with query_ontology()
-            raise AttributeError(f"Cannot find children with UUID {self.uid}! It may exist in the database?")
+            raise AttributeError(
+                f"Cannot find children with UUID {self.uid}! It may exist in the database?"
+            )
 
     def __getattr__(self, name):
         if name in self.__dict__:
@@ -213,24 +219,28 @@ class NifiFragmenter(object):
             raise ValueError("Must have at least 2 entities for fragmentation")
         identifier = str(uuid.uuid4())
         for i, entity in enumerate(fragments):
-            travel_dict(entity.request["nifi_attributes"], fragmenter_keylist.split("."), "w")(
-                {"identifier": identifier, "count": count, "index": i}
-            )
+            travel_dict(
+                entity.request["nifi_attributes"], fragmenter_keylist.split("."), "w"
+            )({"identifier": identifier, "count": count, "index": i})
             if "fragment" not in entity.request["config"]:
                 entity.request["config"]["fragment"] = {}
             if "fragments_stack" not in entity.request["config"]["fragment"]:
                 entity.request["config"]["fragment"]["fragments_stack"] = []
-            entity.request["config"]["fragment"]["fragments_stack"].insert(0, fragmenter_keylist)
-            entity.request["nifi_attributes"]["fragments_stack"] = entity.request["config"]["fragment"][
-                "fragments_stack"
-            ]
-            travel_dict(entity.request["config"]["fragment"], fragmenter_keylist.split("."), "w")(
-                {"identifier": identifier, "count": count, "index": i}
+            entity.request["config"]["fragment"]["fragments_stack"].insert(
+                0, fragmenter_keylist
             )
+            entity.request["nifi_attributes"]["fragments_stack"] = entity.request[
+                "config"
+            ]["fragment"]["fragments_stack"]
+            travel_dict(
+                entity.request["config"]["fragment"], fragmenter_keylist.split("."), "w"
+            )({"identifier": identifier, "count": count, "index": i})
 
     def push_defragment_strategy(fragment, defragmenter_config):
         pointer = fragment.request["config"]
-        last_fragmenter_keylist = fragment.request["config"]["fragment"]["fragments_stack"][0]
+        last_fragmenter_keylist = fragment.request["config"]["fragment"][
+            "fragments_stack"
+        ][0]
         for k in ("fragment." + last_fragmenter_keylist).split("."):
             if not pointer.get(k):
                 pointer[k] = {}
@@ -249,7 +259,7 @@ class NifiEntityBatch(object):
 
 class NifiContextManager(object):
     HEADLESS_PROCESSOR_NAME = "headless"
-
+    
     class SyncFlag(Enum):
         UPSERT_ENTITY_ALL = 0  # bool
         UPSERT_ENTITY_SPECIFIC_FIELDS = 1  # 'fields': list of record fields
@@ -270,7 +280,9 @@ class NifiContextManager(object):
     @property
     def ontology(self):
         if not self._ontology:
-            self._ontology = fetch_ontology_data.sync(ontology_name=self.ontology_name, client=self.client)
+            self._ontology = fetch_ontology_data.sync(
+                ontology_name=self.ontology_name, client=self.client
+            )
         return self._ontology
 
     def _config_get(entity, keylist):
@@ -339,7 +351,9 @@ class NifiContextManager(object):
         return self
 
     def get_workspaces_permissions(self, workspace_ids):
-        permissions_to_fetch = list(set(workspace_ids).difference(set(list(self.permissions.keys()))))
+        permissions_to_fetch = list(
+            set(workspace_ids).difference(set(list(self.permissions.keys())))
+        )
         if permissions_to_fetch:
             permissions = get_permissions.sync(permissions_to_fetch, client=self.client)
             self.permissions.update(permissions)
@@ -369,13 +383,18 @@ class NifiContextManager(object):
             entities.append(entity)
             for child_entity in entity.children_entities:
                 if not child_entity.drop_on_output:
-                    if child_entity.output_as_independent or child_entity.output_as_child:
+                    if (
+                        child_entity.output_as_independent
+                        or child_entity.output_as_child
+                    ):
                         if processor_name != NifiContextManager.HEADLESS_PROCESSOR_NAME:
                             child_entity.request["last_processor_name"] = processor_name
                     if child_entity.output_as_independent:
                         if not child_entity._proxy:
                             child_entity.fetch_proxy()
-                        entities.extend(_process_entity(child_entity._proxy, processor_name))
+                        entities.extend(
+                            _process_entity(child_entity._proxy, processor_name)
+                        )
             return entities
 
         entities = itertools.chain(*[b.entities for b in entity_batches])
@@ -386,7 +405,9 @@ class NifiContextManager(object):
             all_entities,
             key=lambda x: string_to_datetime(x.record.get("os_last_updated_at")),
         )
-        self.out_entities = list({e.record["entity_id"]: e for e in all_entities}.values())
+        self.out_entities = list(
+            {e.record["entity_id"]: e for e in all_entities}.values()
+        )
         self.sync_entities()
         return [entity for entity in self.jsonify(self.out_entities)["content"]]
 
@@ -394,19 +415,23 @@ class NifiContextManager(object):
         error_response = DefaultErrorRoute.format_error(exc)
         entity.request["exception"]["code"] = error_response.status_code
         entity.request["exception"]["body"] = json.loads(error_response.body)["message"]
-        travel_dict(entity.request["nifi_attributes"], ["invokehttp", "response", "body"], "w")(
-            entity.request["exception"]["body"]
-        )
-        travel_dict(entity.request["nifi_attributes"], ["invokehttp", "response", "code"], "w")(
-            entity.request["exception"]["code"]
-        )
+        travel_dict(
+            entity.request["nifi_attributes"], ["invokehttp", "response", "body"], "w"
+        )(entity.request["exception"]["body"])
+        travel_dict(
+            entity.request["nifi_attributes"], ["invokehttp", "response", "code"], "w"
+        )(entity.request["exception"]["code"])
         entity.request["nifi_attributes"]["raised_exc"] = True
 
     def sync_entities(self):
         if not self.lazy_sync:
             entities = self.out_entities
         else:
-            entities = [e for e in self.out_entities if e.record["entity_id"] in self.nonlazy_sync_ids]
+            entities = [
+                e
+                for e in self.out_entities
+                if e.record["entity_id"] in self.nonlazy_sync_ids
+            ]
         if not entities:
             return
         reserved_fields = [
@@ -436,16 +461,18 @@ class NifiContextManager(object):
         for entity in entities:
             if entity.sync_params.get(NifiContextManager.SyncFlag.FETCH_RELATIONSHIPS):
                 concept_name = entity.record["entity_type"]
-                rels_to_fetch = entity.sync_params.get(NifiContextManager.SyncFlag.FETCH_RELATIONSHIPS, [])
+                rels_to_fetch = entity.sync_params.get(
+                    NifiContextManager.SyncFlag.FETCH_RELATIONSHIPS, []
+                )
                 for rel in rels_to_fetch:
                     if rel not in fetch_relationships_entities:
                         fetch_relationships_entities[rel] = []
                     fetch_relationships_entities[rel].append(entity)
                 if concept_name not in fetch_concept_relationships:
                     fetch_concept_relationships[concept_name] = set()
-                fetch_concept_relationships[concept_name] = fetch_concept_relationships[concept_name].union(
-                    set(rels_to_fetch)
-                )
+                fetch_concept_relationships[concept_name] = fetch_concept_relationships[
+                    concept_name
+                ].union(set(rels_to_fetch))
         for k in fetch_concept_relationships.keys():
             fetch_concept_relationships[k] = list(fetch_concept_relationships[k])
         # UPSERT ENTITIES
@@ -491,7 +518,7 @@ class NifiContextManager(object):
             file.request["is_temporary"] = False
             file.request["entity_timestamp"] = file.record["os_last_updated_at"]
         # FETCH RELATIONSHIPS
-        """
+        '''
         if fetch_relationships_entities:
             relationship_mappings_info = relationship_mappings.sync_detailed(
                 client=self.client
@@ -547,7 +574,7 @@ class NifiContextManager(object):
                             child_rel.request["entity_timestamp"] = rel.get(
                                 "os_last_updated_at"
                             )
-        """
+        '''
         # CLEAN SYNC PARAMS
         for entity in entities:
             entity.sync_params = {}
@@ -556,10 +583,15 @@ class NifiContextManager(object):
         for entity in entities:
             fields = set()
 
-            if entity.sync_params.get(NifiContextManager.SyncFlag.UPSERT_ENTITY_ALL) or entity.request["is_temporary"]:
+            if (
+                    entity.sync_params.get(NifiContextManager.SyncFlag.UPSERT_ENTITY_ALL)
+                    or entity.request["is_temporary"]
+            ):
                 fields = fields.union(set(list(entity.record.keys())))
 
-            if entity.sync_params.get(NifiContextManager.SyncFlag.UPSERT_ENTITY_SPECIFIC_FIELDS):
+            if entity.sync_params.get(
+                    NifiContextManager.SyncFlag.UPSERT_ENTITY_SPECIFIC_FIELDS
+            ):
                 fields = fields.union(
                     set(
                         entity.sync_params.get(
@@ -570,7 +602,9 @@ class NifiContextManager(object):
                     )
                 )
             if fields:
-                entities_to_upsert.append((entity, [f for f in list(fields) if f not in reserved_fields]))
+                entities_to_upsert.append(
+                    (entity, [f for f in list(fields) if f not in reserved_fields])
+                )
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_val is not None:
@@ -584,7 +618,9 @@ class NifiContextManager(object):
             for entity in entities:
                 if isinstance(entity, NifiEntityProxy):
                     children.extend(entity.children_entities)
-                    children.extend(_recursive_collect_proxies(entity.children_entities))
+                    children.extend(
+                        _recursive_collect_proxies(entity.children_entities)
+                    )
             return children
 
         all_proxies = _recursive_collect_proxies(entities)
@@ -597,21 +633,15 @@ class NifiContextManager(object):
 
 
 class NifiEntity(object):
-    def __init__(self, context, request, record, annotations, all_independent_uids, children=[], contents=None):
+    def __init__(
+        self, context, request, record, annotations, all_independent_uids, children=[], contents=None
+    ):
         self.context = context
         self.request = request
         self.record = record
         self.annotations = annotations
-        assert (
-            self.record.get("os_entity_uid")
-            and self.record.get("entity_id")
-            and self.record["os_entity_uid"] == self.record["entity_id"]
-        )
-        assert (
-            self.record.get("os_concept")
-            and self.record.get("entity_type")
-            and self.record["os_concept"] == self.record["entity_type"]
-        )
+        assert self.record.get("os_entity_uid") and self.record.get("entity_id") and self.record["os_entity_uid"] == self.record["entity_id"]
+        assert self.record.get("os_concept") and self.record.get("entity_type") and self.record["os_concept"] == self.record["entity_type"]
         if "entity_label" not in self.record:
             self.record["entity_label"] = self.label
         children = [c for c in children if isinstance(c, (str, dict))]
@@ -631,7 +661,9 @@ class NifiEntity(object):
         child_types = [c["entity_type"] for c in proxy_entity_children] + [
             c["record"]["entity_type"] for c in full_entity_children
         ]
-        output_as_child = [False] * len(proxy_entity_children) + [True] * len(full_entity_children)
+        output_as_child = [False] * len(proxy_entity_children) + [True] * len(
+            full_entity_children
+        )
         output_as_independent = [uid in all_independent_uids for uid in child_uids]
         full_entity_children = [
             NifiEntity(
@@ -645,7 +677,9 @@ class NifiEntity(object):
             )
             for c in full_entity_children
         ]
-        proxy_otm_children = [NifiOTMRelationshipProxy(**otm_child) for otm_child in proxy_otm_children]
+        proxy_otm_children = [
+            NifiOTMRelationshipProxy(**otm_child) for otm_child in proxy_otm_children
+        ]
         child_proxies = [None] * len(proxy_entity_children) + full_entity_children
         self.children = [
             NifiEntityProxy(
@@ -673,18 +707,21 @@ class NifiEntity(object):
 
     @property
     def sync_params(self):
-        return {NifiContextManager.SyncFlag[k]: v for k, v in (self.request.get("sync_params") or {}).items()}
+        return {
+            NifiContextManager.SyncFlag[k]: v
+            for k, v in (self.request.get("sync_params") or {}).items()
+        }
 
     @sync_params.setter
     def sync_params(self, new_params):
         self.request["sync_params"] = {
-            (k.name if isinstance(k, NifiContextManager.SyncFlag) else k): v for k, v in new_params.items()
+            (k.name if isinstance(k, NifiContextManager.SyncFlag) else k): v
+            for k, v in new_params.items()
         }
 
     @property
     def metadata(self):
         return self.annotations
-
     @metadata.setter
     def metadata(self, new_metadata):
         self.annotations = new_metadata
@@ -707,7 +744,9 @@ class NifiEntity(object):
         contents_pointer = deepcopy(self.request["contents_pointer"])
         ptr_location = contents_pointer.get("location")
         if ptr_location == "attachment" and not contents_pointer.get("pointer"):
-            contents_pointer["pointer"] = f"{self.record['os_workspace']}/{self.record['os_entity_uid']}"
+            contents_pointer["pointer"] = (
+                f"{self.record['os_workspace']}/{self.record['os_entity_uid']}"
+            )
         return self.request["contents_pointer"]
 
     @contents_pointer.setter
@@ -727,7 +766,9 @@ class NifiEntity(object):
         return list(
             filter(
                 lambda x: isinstance(x, NifiOTMRelationshipProxy)
-                or is_child_concept_fn(x.entity_type, RELATIONSHIP_ENTITY_NAME, self.context.ontology),
+                or is_child_concept_fn(
+                    x.entity_type, RELATIONSHIP_ENTITY_NAME, self.context.ontology
+                ),
                 self.children,
             )
         )
@@ -746,12 +787,18 @@ class NifiEntity(object):
         )
         if (
             self.record.get("os_workspace")
-            and (permissions.get(self.record.get("os_workspace")) or PermissionLevel.NONE) >= PermissionLevel.WRITE
+            and (
+                permissions.get(self.record.get("os_workspace")) or PermissionLevel.NONE
+            )
+            >= PermissionLevel.WRITE
         ):
             return self.record["os_workspace"]
         elif (
             self.request.get("fallback_os_workspace")
-            and (permissions.get(self.request.get("fallback_os_workspace")) or PermissionLevel.NONE)
+            and (
+                permissions.get(self.request.get("fallback_os_workspace"))
+                or PermissionLevel.NONE
+            )
             >= PermissionLevel.WRITE
         ):
             return self.request["fallback_os_workspace"]
@@ -760,7 +807,7 @@ class NifiEntity(object):
 
     @property
     def label(self):
-        label_keys = self.request["ontology"]["label_keys"]
+        label_keys = self.request["ontology_info"]["label_keys"]
         label = " ".join([(self.record.get(field) or "") for field in label_keys]).strip()
         if not label:
             label = None
@@ -785,10 +832,14 @@ class NifiEntity(object):
         entity_type = None
         if isinstance(self, NifiEntityProxy):
             entity_type = self.entity_type
-            return entity_type == type or is_child_concept_fn(entity_type, type, self.context.ontology)
+            return entity_type == type or is_child_concept_fn(
+                entity_type, type, self.context.ontology
+            )
         else:
             entity_type = self.record["entity_type"]
-            return entity_type == type or type in self.request["ontology_info"]["parents"]
+            return (
+                entity_type == type or type in self.request["ontology_info"]["parents"]
+            )
 
     def is_fragmented(self) -> bool:
         return bool(self.request["config"].get("fragment", {}).get("fragments_stack"))
@@ -804,7 +855,7 @@ class NifiEntity(object):
                     if _is_sub_fragment_recursive(value):
                         return True
             return False
-
+        
         if not self.is_fragmented():
             return True
         fragment = entity.request.get("config", {}).get("fragment", {})
@@ -831,7 +882,10 @@ class NifiEntity(object):
                 else:
                     proxy_entity_children.append(child)
         proxy_entity_children = list({c.uid: c for c in proxy_entity_children}.values())
-        proxy_entity_children = [{"entity_id": c.uid, "entity_type": c.entity_type} for c in proxy_entity_children]
+        proxy_entity_children = [
+            {"entity_id": c.uid, "entity_type": c.entity_type}
+            for c in proxy_entity_children
+        ]
         proxy_otm_children = list(
             {
                 c.record["os_entity_uid_from"]
@@ -847,7 +901,9 @@ class NifiEntity(object):
             full_entity_children,
             key=lambda x: string_to_datetime(x.record.get("os_last_updated_at")),
         )
-        full_entity_children = list({c.uid: c.to_json() for c in full_entity_children}.values())
+        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
         return {
             "request": self.request,
@@ -862,9 +918,9 @@ class NifiEntity(object):
         random_id = str(uuid.uuid4())
         username = self.jwt_data["username"]
         if entity_type == self.record["entity_type"]:
-            ont_parents = self.request["ontology"]["parents"]
-            ont_relationships = self.request["ontology"]["relationships"]
-            ont_label_keys = self.request["ontology"]["label_keys"]
+            ont_parents = self.request["ontology_info"]["parents"]
+            ont_relationships = self.request["ontology_info"]["relationships"]
+            ont_label_keys = self.request["ontology_info"]["label_keys"]
         else:
             ont_parents = self.context.ontology["concepts"][entity_type]["parents"]
             ont_relationships = self.context.ontology["concepts"][entity_type]["relationships"]
@@ -931,7 +987,11 @@ class NifiEntity(object):
     ):
         return self._add_entity(
             os_relationship_workspace,
-            (LOCAL_RELATIONSHIP_ENTITY_NAME if os_relationship_workspace else RELATIONSHIP_ENTITY_NAME),
+            (
+                LOCAL_RELATIONSHIP_ENTITY_NAME
+                if os_relationship_workspace
+                else RELATIONSHIP_ENTITY_NAME
+            ),
             {
                 **relationship_fields,
                 "os_entity_uid_from": os_entity_uid_from,
@@ -1013,7 +1073,9 @@ class NifiEntity(object):
             os_relationship_type,
         )
         child_entity._contents = file
-        child_entity.request["contents_pointer"] = NifiEntityModel.RequestModel.ContentsPointerModel(location="local")
+        child_entity.request["contents_pointer"] = (
+            NifiEntityModel.RequestModel.ContentsPointerModel(location="local")
+        )
         return child_entity, child_rel
 
     def add_tag(self, os_workspace, name, group, order, color):
@@ -1026,20 +1088,21 @@ class NifiEntity(object):
         )
 
     def add_metadata(
-        self,
-        json,
-        merge_method: Callable[[Any, Any], Any],
-        recurse: Union[bool, int] = False,
+        self, json, merge_method: Callable[[Any, Any], Any], recurse: Union[bool, int] = False,
     ):
         if not self.metadata:
             self.metadata = {}
-        self.metadata = recursive_update_dict(self.metadata, json, merge_method, recurse)
+        self.metadata = recursive_update_dict(
+            self.metadata, json, merge_method, recurse
+        )
 
     def propagate_metadata(self, to_entity, fields=None, merge_method=lambda _, v2: v2):
         metadata_to_propagate = deepcopy(self.metadata)
         if fields:
             metadata_to_propagate = {k: v for k, v in self.metadata if k in fields}
-        to_entity.metadata = recursive_update_dict(to_entity.metadata, metadata_to_propagate, merge_method)
+        to_entity.metadata = recursive_update_dict(
+            to_entity.metadata, metadata_to_propagate, merge_method
+        )
 
 
 def more_recent_than(record_a, record_b):
@@ -1098,10 +1161,19 @@ class NifiRoute(Route):
             query_params = request.query_params
             processor_suffix = query_params["processor_suffix"]
             body = await request.json()
-            processor_name = "processor." + self.processor_name + "." + op.replace("-", "_") + "." + processor_suffix
+            processor_name = (
+                "processor."
+                + self.processor_name
+                + "."
+                + op.replace("-", "_")
+                + "."
+                + processor_suffix
+            )
             if op not in self.endpoints.keys():
                 raise StarletteHTTPException(401, f"Route {op} is forbidden for NiFi.")
-            task_id = await self.celery_executor.send_task(self.endpoints[op], args=[body, processor_name])
+            task_id = await self.celery_executor.send_task(
+                self.endpoints[op], args=[body, processor_name]
+            )
             return task_id
 
     @staticmethod

{streamlit_octostar_utils-0.2.12a1 → streamlit_octostar_utils-0.2.13}/streamlit_octostar_utils/nlp/language.py

@@ -1,6 +1,6 @@
 import re
 import py3langid as langid
-import iso639 as languages
+from iso639 import Lang
 
 
 def detect_language(text, min_confidence=None):
@@ -11,5 +11,11 @@ def detect_language(text, min_confidence=None):
     if min_confidence and confidence < min_confidence:
         return None, confidence
     detected_lang = re.sub("[^A-Za-z]", "", detected_lang).lower()
-    detected_lang = languages.to_name(detected_lang).lower()
+    detected_lang = Lang(detected_lang).name.lower()
     return detected_lang, confidence
+
+def to_name(alpha2):
+    return Lang(alpha2).name.lower()
+
+def to_alpha2(name):
+    return Lang(name).pt1

{streamlit_octostar_utils-0.2.12a1 → streamlit_octostar_utils-0.2.13}/streamlit_octostar_utils/ontology/inheritance.py

@@ -12,4 +12,4 @@ def get_label_keys(type, ontology):
             if not label_key:
                 continue
             label_keys[label_key] = None
-    return list(label_keys.keys())
+    return list(label_keys.keys())

streamlit_octostar_utils-0.2.13/streamlit_octostar_utils/ontology/relationships.py

@@ -0,0 +1,28 @@
+from .inheritance import is_child_concept
+
+
+def get_relationships_between_concepts(source, target, ontology):
+    ontology_rels = {k["relationship_name"]: k for k in ontology["relationships"]}
+    from_rels = ontology["concepts"][source]["relationships"]
+    from_rels = [ontology_rels[r] for r in from_rels]
+    from_rels = [
+        r
+        for r in from_rels
+        if is_child_concept(
+            target,
+            r["target_concept"],
+            {
+                "concepts": ontology["concepts"],
+                "relationships": list(ontology_rels.values()),
+            },
+        )
+    ]
+    return from_rels
+
+
+def invert_relationships(rels, ontology):
+    ontology_rels = {k["relationship_name"]: k for k in ontology["relationships"]}
+    inverses = []
+    for rel in rels:
+        inverses.append(ontology_rels[rel]["inverse_name"])
+    return inverses

streamlit_octostar_utils-0.2.12a1/streamlit_octostar_utils/ontology/expand_entities.py

@@ -1,594 +0,0 @@
-from itertools import groupby, chain
-import time
-import uuid
-from octostar.utils.ontology import multiquery_ontology
-from octostar.utils.exceptions import StopAsyncIterationWithResult
-
-from .inheritance import is_child_concept
-
-REL_FETCHED_FIELDS = [
-    "os_relationship_name",
-    "os_entity_uid_from",
-    "os_entity_uid_to",
-    "os_entity_uid",
-    "os_entity_type_from",
-    "os_entity_type_to",
-    "os_workspace",
-]
-TARGETS_FETCHED_FIELDS = ["os_entity_uid", "os_concept", "entity_label", "os_workspace"]
-
-
-class ExecutionMetrics:
-    def __init__(self):
-        self.n_queries = 0
-        self.n_relationships = 0
-        self.n_target_entities = 0
-        self.exec_time = 0.0
-        self.exec_times = {}
-        self.n_cancelled_queries = 0
-        self.relationship_names = set()
-        self.timeout = 0
-
-    def print_metrics(self):
-        print("Execution Metrics:")
-        print(f"  Execution Time: {self.exec_time}")
-        print(f"  Per-Type Execution Times: {self.exec_times}")
-        print(f"  Number of Queries: {self.n_queries}")
-        print(f"  Number of Cancelled Queries: {self.n_cancelled_queries}")
-        print(f"  Timeout per Query: {self.timeout}")
-        print(f"  Number of Relationships: {self.n_relationships}")
-        print(f"  Number of Target Entities: {self.n_target_entities}")
-        print(f"  Relationship Names: {self.relationship_names}")
-
-
-async def _get_stream_result(stream):
-    result = None
-    try:
-        async for _ in stream:
-            pass
-    except StopAsyncIterationWithResult as e:
-        result = e.value
-    return result
-
-
-def _left_join(left_list, right_list, left_keys, right_keys):
-    left_keys = left_keys if isinstance(left_keys, (list, tuple)) else [left_keys]
-    right_keys = right_keys if isinstance(right_keys, (list, tuple)) else [right_keys]
-
-    def _make_composite_key(item, keys):
-        return tuple(item[key] for key in keys)
-
-    left_list.sort(key=lambda item: _make_composite_key(item, left_keys))
-    right_list.sort(key=lambda item: _make_composite_key(item, right_keys))
-    left_groups = groupby(
-        left_list, key=lambda item: _make_composite_key(item, left_keys)
-    )
-    right_groups = groupby(
-        right_list, key=lambda item: _make_composite_key(item, right_keys)
-    )
-    right_dict = {key: list(group) for key, group in right_groups}
-    result = []
-    for left_key_val, left_items in left_groups:
-        left_items = list(left_items)
-        associated_rights = right_dict.get(left_key_val, [])
-        for left_item in left_items:
-            result.append((left_item, associated_rights))
-    return result
-
-
-async def expand_otm(
-    expanded_entities,
-    entities_by_concept,
-    ontology_rels,
-    concepts_to_otm_rels,
-    metrics,
-    client,
-    timeout,
-    limit,
-):
-    start_time = time.time()
-    otm_queries = []
-    target_fields = ",".join(
-        ["`" + rel_field + "`" for rel_field in TARGETS_FETCHED_FIELDS]
-    )
-    for concept_name, rels in concepts_to_otm_rels.items():
-        for rel_name in rels:
-            rel = ontology_rels[rel_name]
-            source_properties = rel["source_properties"].split(",")
-            source_properties_values = [
-                [entity[p] for p in source_properties]
-                for entity in entities_by_concept[concept_name]
-            ]
-            source_properties_values = [
-                values
-                for values in source_properties_values
-                if all(value is not None for value in values)
-            ]
-            if source_properties_values:
-                target_prop_names = rel["target_properties"].split(",")
-                target_prop_names = ",".join(
-                    ["`" + prop + "`" for prop in target_prop_names]
-                )
-                source_properties_values = ",".join(
-                    [
-                        "(" + ",".join(["'" + value + "'" for value in values]) + ")"
-                        for values in source_properties_values
-                    ]
-                )
-                if target_prop_names and source_properties_values:
-                    query = f"""SELECT {target_fields} FROM `etimbr`.`{rel['target_concept']}` WHERE ({target_prop_names}) IN ({source_properties_values}) LIMIT {limit}"""
-                    otm_queries.append(
-                        {
-                            "concept_name": concept_name,
-                            "query": query,
-                            "relationship_name": rel_name,
-                            "source_properties": rel["source_properties"].split(","),
-                            "target_properties": rel["target_properties"].split(","),
-                        }
-                    )
-    otm_queries = {str(i): otm_queries[i] for i in range(len(otm_queries))}
-    otm_stream = multiquery_ontology.streaming(
-        sql_queries={
-            query_id: query_data["query"]
-            for query_id, query_data in otm_queries.items()
-        },
-        client=client,
-        timeout=timeout,
-    )
-    metrics.n_queries += len(otm_queries)
-    all_results = await _get_stream_result(otm_stream)
-    for query_id, data in all_results.items():
-        if data == None:
-            metrics.n_cancelled_queries += 1
-        if not data:
-            continue
-        query_data = otm_queries[query_id]
-        data = _left_join(
-            [
-                e[0]
-                for e in expanded_entities.values()
-                if e[0]["os_concept"] == query_data["concept_name"]
-            ],
-            data,
-            query_data["source_properties"],
-            query_data["target_properties"],
-        )
-        data = {
-            e[0]["os_entity_uid"]: (
-                e[0],
-                {
-                    r["os_entity_uid"]: (
-                        {"os_relationship_name": query_data["relationship_name"]},
-                        {r["os_entity_uid"]: r},
-                    )
-                    for r in e[1]
-                },
-            )
-            for e in data
-        }
-        for entity_id, entities in data.items():
-            if entity_id not in expanded_entities:
-                expanded_entities[entity_id] = [entities[0], {}]
-            for sub_entity_id, sub_entities in entities[1].items():
-                if sub_entity_id not in expanded_entities[entity_id][1]:
-                    expanded_entities[entity_id][1][sub_entity_id] = [
-                        sub_entities[0],
-                        {},
-                    ]
-                expanded_entities[entity_id][1][sub_entity_id][1].update(
-                    sub_entities[1]
-                )
-    metrics.exec_times["otm"] = time.time() - start_time
-    return expanded_entities
-
-
-async def expand_mtm_mixed(
-    expanded_entities,
-    entities_by_concept,
-    ontology,
-    concepts_to_mixed_rels,
-    ontology_rels,
-    metrics,
-    client,
-    timeout,
-    limit,
-):
-    start_time = time.time()
-    mixed_queries = []
-    match_patterns = {}
-    rel_fields = ",".join(["`" + rel_field + "`" for rel_field in REL_FETCHED_FIELDS])
-    for concept_name, rels in concepts_to_mixed_rels.items():
-        if concept_name not in match_patterns:
-            match_patterns[concept_name] = []
-        match_patterns[concept_name].extend(
-            [entity["os_entity_uid"] for entity in entities_by_concept[concept_name]]
-        )
-    for concept_name, entity_ids in match_patterns.items():
-        rels = concepts_to_mixed_rels[concept_name]
-        entity_ids = ",".join(["'" + uid + "'" for uid in entity_ids])
-        for rel_name in rels:
-            rel = ontology_rels[rel_name]
-            target_concept = rel["target_concept"]
-            relationship_fields = ",".join(
-                [
-                    f"`{rel_name}[{target_concept}]_"
-                    + prop
-                    + "` AS "
-                    + "`rel__"
-                    + prop
-                    + "`"
-                    for prop in REL_FETCHED_FIELDS
-                ]
-            )
-            target_fields = ",".join(
-                [
-                    f"`{rel_name}[{target_concept}]."
-                    + prop
-                    + "` AS "
-                    + "`tgt__"
-                    + prop
-                    + "`"
-                    for prop in TARGETS_FETCHED_FIELDS
-                ]
-            )
-            all_fields = (
-                f"{relationship_fields}, {target_fields}, `os_entity_uid`".strip(", ")
-            )
-            if entity_ids:
-                mixed_queries.append(
-                    {
-                        "query": f"SELECT {all_fields} FROM `dtimbr`.`{concept_name}` WHERE `os_entity_uid` IN ({entity_ids}) LIMIT {limit}"
-                    }
-                )
-    mixed_queries = {str(i): mixed_queries[i] for i in range(len(mixed_queries))}
-    metrics.n_queries += len(mixed_queries)
-    local_stream = multiquery_ontology.streaming(
-        sql_queries={
-            query_id: query_data["query"]
-            for query_id, query_data in mixed_queries.items()
-        },
-        client=client,
-        timeout=timeout / 2.0,
-    )
-    all_results = await _get_stream_result(local_stream)
-    for (
-        _,
-        data,
-    ) in (
-        all_results.items()
-    ):  ## TO BE TESTED (data is always empty at the moment due to timbr bug)
-        if data == None:
-            metrics.n_cancelled_queries += 1
-        if not data:
-            continue
-        data = _left_join(
-            [e[0] for e in expanded_entities.values()],
-            data,
-            "os_entity_uid",
-            "os_entity_uid",
-        )
-        for elem in data:
-            elem["#rel__os_entity_uid"] = elem["rel__os_entity_uid"] or (
-                "temp-" + str(uuid.uuid4())
-            )
-        data = {
-            e[0]["os_entity_uid"]: (
-                e[0],
-                {
-                    rt["#rel__os_entity_uid"]: (
-                        {k[5:]: v for k, v in rt.items() if k.startswith("rel__")},
-                        {
-                            rt["tgt__os_entity_uid"]: {
-                                k[5:]: v for k, v in rt.items() if k.startswith("tgt__")
-                            }
-                        },
-                    )
-                    for rt in e[1]
-                },
-            )
-            for e in data
-        }
-        data = {
-            e[0]["os_entity_uid"]: (
-                e[0],
-                {r["os_entity_uid"]: ({}, {r["os_entity_uid"]: r}) for r in e[1]},
-            )
-            for e in data
-        }
-        for entity_id, entities in data.items():
-            if entity_id not in expanded_entities:
-                expanded_entities[entity_id] = [entities[0], {}]
-            for sub_entity_id, sub_entities in entities[1].items():
-                if sub_entity_id not in expanded_entities[entity_id][1]:
-                    expanded_entities[entity_id][1][sub_entity_id] = [
-                        sub_entities[0],
-                        {},
-                    ]
-                expanded_entities[entity_id][1][sub_entity_id][1].update(
-                    sub_entities[1]
-                )
-    metrics.exec_times["mtm_mixed"] = time.time() - start_time
-    return expanded_entities
-
-
-async def expand_mtm_local(
-    expanded_entities,
-    entities_by_concept,
-    ontology_rels,
-    concepts_to_local_rels,
-    metrics,
-    client,
-    timeout,
-    limit,
-):
-    start_time = time.time()
-    local_queries = {"from": [], "to": []}
-    match_patterns = {}
-    target_fields = ",".join(
-        ["`" + rel_field + "`" for rel_field in TARGETS_FETCHED_FIELDS]
-    )
-    rel_fields = ",".join(["`" + rel_field + "`" for rel_field in REL_FETCHED_FIELDS])
-    for concept_name in concepts_to_local_rels.keys():
-        if concept_name not in match_patterns:
-            match_patterns[concept_name] = []
-        match_patterns[concept_name].extend(
-            [entity["os_entity_uid"] for entity in entities_by_concept[concept_name]]
-        )
-    for concept_name, entity_ids in match_patterns.items():
-        rel_names = concepts_to_local_rels[concept_name]
-        inverse_names = [
-            ontology_rels[rel_name]["inverse_name"] for rel_name in rel_names
-        ]
-        rel_names = ",".join(["'" + rel_name + "'" for rel_name in rel_names])
-        inverse_names = ",".join(["'" + rel_name + "'" for rel_name in inverse_names])
-        entity_ids = ",".join(["'" + uid + "'" for uid in entity_ids])
-        if rel_names and inverse_names and entity_ids:
-            local_queries["from"].append(
-                {
-                    "query": f"SELECT {rel_fields} FROM `timbr`.`os_workspace_relationship` WHERE `os_relationship_name` IN ({rel_names}) AND `os_entity_uid_from` IN ({entity_ids}) LIMIT {limit}"
-                }
-            )
-            local_queries["to"].append(
-                {
-                    "query": f"SELECT {rel_fields} FROM `timbr`.`os_workspace_relationship` WHERE `os_relationship_name` IN ({inverse_names}) AND `os_entity_uid_to` IN ({entity_ids}) LIMIT {limit}"
-                }
-            )
-    local_queries = {
-        **{
-            "from_" + str(i): local_queries["from"][i]
-            for i in range(len(local_queries["from"]))
-        },
-        **{
-            "to_" + str(i): local_queries["to"][i]
-            for i in range(len(local_queries["to"]))
-        },
-    }
-    metrics.n_queries += len(local_queries)
-    local_stream = multiquery_ontology.streaming(
-        sql_queries={
-            query_id: query_data["query"]
-            for query_id, query_data in local_queries.items()
-        },
-        client=client,
-        timeout=timeout / 2.0,
-    )
-    all_results = await _get_stream_result(local_stream)
-    middle_entities = []
-    for query_id, data in all_results.items():
-        if data == None:
-            metrics.n_cancelled_queries += 1
-        if not data:
-            continue
-        query_id = query_id.split("_")
-        data = _left_join(
-            [e[0] for e in expanded_entities.values()],
-            data,
-            "os_entity_uid",
-            "os_entity_uid_" + query_id[0],
-        )
-        inverse_direction = "to" if query_id[0] == "from" else "from"
-        data = {
-            e[0]["os_entity_uid"]: (e[0], {r["os_entity_uid"]: (r, {}) for r in e[1]})
-            for e in data
-        }
-        middle_entities.extend(
-            [(inverse_direction, r) for e in data.values() for r in e[1].values()]
-        )
-        for entity_id, entities in data.items():
-            if entity_id not in expanded_entities:
-                expanded_entities[entity_id] = [entities[0], {}]
-            for sub_entity_id, sub_entities in entities[1].items():
-                if sub_entity_id not in expanded_entities[entity_id][1]:
-                    expanded_entities[entity_id][1][sub_entity_id] = [
-                        sub_entities[0],
-                        {},
-                    ]
-                expanded_entities[entity_id][1][sub_entity_id][1].update(
-                    sub_entities[1]
-                )
-    middle_entities = [
-        {
-            "entity_id": r[1][0]["os_entity_uid_" + r[0]],
-            "concept_name": r[1][0]["os_entity_type_" + r[0]],
-            "direction": r[0],
-            "relationship": r[1],
-        }
-        for r in middle_entities
-    ]
-    middle_entities = sorted(middle_entities, key=lambda x: x["concept_name"])
-    middle_entities = groupby(middle_entities, key=lambda x: x["concept_name"])
-    middle_entities = {e[0]: list(e[1]) for e in middle_entities}
-    local_queries_2 = []
-    for concept_name, entities in middle_entities.items():
-        entity_ids = ",".join(["'" + e["entity_id"] + "'" for e in entities])
-        if entity_ids:
-            local_queries_2.append(
-                {
-                    "query": f"SELECT {target_fields} FROM `timbr`.`{concept_name}` WHERE `os_entity_uid` IN ({entity_ids}) LIMIT {limit}"
-                }
-            )
-    local_queries_2 = {str(i): local_queries_2[i] for i in range(len(local_queries_2))}
-    metrics.n_queries += len(local_queries_2)
-    local_stream_2 = multiquery_ontology.streaming(
-        sql_queries={
-            query_id: query_data["query"]
-            for query_id, query_data in local_queries_2.items()
-        },
-        client=client,
-        timeout=timeout / 2.0,
-    )
-    all_results = await _get_stream_result(local_stream_2)
-    middle_entities = list(chain(*middle_entities.values()))
-    for query_id, data in all_results.items():
-        if data == None:
-            metrics.n_cancelled_queries += 1
-        if not data:
-            continue
-        data = _left_join(middle_entities, data, "entity_id", "os_entity_uid")
-        for entry in data:
-            relationship_targets = entry[0]["relationship"]
-            for target_entity in entry[1]:
-                relationship_targets[1][target_entity["os_entity_uid"]] = target_entity
-        if data:
-            for rel in data:
-                targets = rel[0]["relationship"][1]
-                rel_uid = rel[0]["relationship"][0]["os_entity_uid"]
-                for entity in expanded_entities.values():
-                    entity_rels = entity[1]
-                    if rel_uid in entity_rels:
-                        entity_rels[rel_uid][1].update(targets)
-    metrics.exec_times["mtm_local"] = time.time() - start_time
-    return expanded_entities
-
-
-async def expand_entities(
-    entities,
-    ontology,
-    relationship_mappings_info,
-    client,
-    relationship_names_by_entity_type=None,
-    batch_size=10000,
-    avg_limit_per_entity=20,
-    timeout=10.0,
-):
-    metrics = ExecutionMetrics()
-    metrics.timeout = timeout
-    start_time = time.time()
-    ordered_entities = sorted(entities, key=lambda x: x["entity_type"])
-    ordered_entities = list({e["os_entity_uid"]: e for e in entities}.values())
-    entity_batches = [
-        ordered_entities[i : i + batch_size]
-        for i in range(0, len(ordered_entities), batch_size)
-    ]
-    expanded_entities = {}
-    for batch in entity_batches:
-        new_entities = await expand_entities_batch(
-            batch,
-            ontology,
-            relationship_mappings_info,
-            relationship_names_by_entity_type,
-            metrics,
-            client,
-            avg_limit_per_entity,
-            timeout,
-        )
-        expanded_entities.update(new_entities)
-    metrics.exec_time = time.time() - start_time
-    for _, entity_data in expanded_entities.items():
-        for rel in entity_data[1].values():
-            if rel[0]:
-                metrics.n_relationships += 1
-                metrics.relationship_names.add(rel[0]["os_relationship_name"])
-            metrics.n_target_entities += len(rel[1])
-    metrics.relationship_names = list(metrics.relationship_names)
-    expanded_entities = [
-        (e[0], [(r[0], list(r[1].values())[0]) for r in e[1].values()])
-        for e in expanded_entities.values()
-    ]
-    return expanded_entities, metrics
-
-
-async def expand_entities_batch(
-    entities,
-    ontology,
-    relationship_mappings,
-    relationship_names_by_entity_type,
-    metrics,
-    client,
-    avg_limit_per_entity=20,
-    timeout=10.0,
-):
-    limit = avg_limit_per_entity * len(entities)
-    expanded_entities = {e["os_entity_uid"]: [e, {}] for e in entities}
-    entities_by_concept = sorted(entities, key=lambda x: x["entity_type"])
-    entities_by_concept = groupby(entities_by_concept, key=lambda x: x["entity_type"])
-    entities_by_concept = {e[0]: list(e[1]) for e in entities_by_concept}
-    ontology_rels = {r["relationship_name"]: r for r in ontology["relationships"]}
-    concepts_to_rels = {
-        cn: c["relationships"] for cn, c in ontology["concepts"].items()
-    }
-    concepts_to_mixed_rels = {}
-    concepts_to_otm_rels = {}
-    concepts_to_local_rels = {}
-    for concept_name in entities_by_concept.keys():
-        relationship_names = None
-        if relationship_names_by_entity_type:
-            relationship_names = set()
-            for cn_name in relationship_names_by_entity_type.keys():
-                if is_child_concept(concept_name, cn_name, ontology):
-                    relationship_names = relationship_names.union(
-                        set(relationship_names_by_entity_type[cn_name])
-                    )
-        rels = concepts_to_rels[concept_name]
-        if relationship_names:
-            rels = [r for r in rels if r in relationship_names]
-        filtered_mtm_rels = []
-        filtered_otm_rels = []
-        filtered_local_mtm_rels = []
-        for rel in rels:
-            if not ontology_rels[rel]["is_mtm"]:
-                filtered_otm_rels.append(rel)
-            elif rel in relationship_mappings["unmapped"]:
-                continue
-            elif rel in relationship_mappings["local_only"]:
-                filtered_local_mtm_rels.append(rel)
-            else:
-                filtered_mtm_rels.append(rel)
-        concepts_to_mixed_rels[concept_name] = filtered_mtm_rels
-        concepts_to_otm_rels[concept_name] = filtered_otm_rels
-        concepts_to_local_rels[concept_name] = filtered_local_mtm_rels
-    # OTM QUERIES: fetch target entities directly from the target tables with an IN statement
-    expanded_entities = await expand_otm(
-        expanded_entities,
-        entities_by_concept,
-        ontology_rels,
-        concepts_to_otm_rels,
-        metrics,
-        client,
-        timeout,
-        limit,
-    )
-    # LOCAL MTM QUERIES: fetch relationships directly from the os_workspace_relationship table, then fetch target entities
-    expanded_entities = await expand_mtm_local(
-        expanded_entities,
-        entities_by_concept,
-        ontology_rels,
-        concepts_to_local_rels,
-        metrics,
-        client,
-        timeout,
-        limit,
-    )
-    # MIXED MTM QUERIES: query via timbr to make the most of JOINs
-    expanded_entities = await expand_mtm_mixed(
-        expanded_entities,
-        entities_by_concept,
-        ontology,
-        concepts_to_mixed_rels,
-        ontology_rels,
-        metrics,
-        client,
-        timeout,
-        limit,
-    )
-    return expanded_entities