octostar-python-client 0.1.759__py3-none-any.whl

octostar/utils/workspace/upsert_relationships.py

@@ -0,0 +1,194 @@
+from typing import Dict, Union, List, Any, TypedDict, Optional
+import uuid
+import os
+import itertools
+import logging
+
+_logger = logging.getLogger(__name__)
+
+from ..ontology import query_ontology
+from . import upsert_entities
+from ...client import Client, get_default_client
+
+
+class EntityBase(TypedDict):
+    entity_type: str
+    os_entity_uid: str
+
+
+def sync(
+    entities_from: List[EntityBase],
+    entities_to: List[EntityBase],
+    os_relationship_names: Union[str, List[str]],
+    os_relationship_workspaces: Union[str, List[str]],
+    relationship_fields: Optional[List[Dict[str, Any]]] = None,
+    os_relationship_uids: Optional[List[Optional[str]]] = None,
+    allow_multi_cardinality: bool = True,
+    ontology_name: str = None,
+    client: Client = None,
+):
+    """
+    # Create or update a set of local relationships between pairs of entities (sources and targets)
+
+    ## Arguments
+    - `entities_from`: The source entities, as a list of dictionaries
+        Each dictionary must contain:
+        - `os_workspace`: The workspace ID belonging to the entity. This will also be the workspace ID of
+            the corresponding relationship
+        - `entity_type`: The concept name for the entity
+        - `os_entity_uid`: The ID for the entity. Only necessary in case of an update
+    - `entities_to`: The target entities, as a list of dictionaries. These entities are aligned with the
+        source entities to create relationships between them. For these entities, os_workspace can be None
+        if the target is a global entity
+    - `os_relationship_name`: A single string used as the name of all relationships, or a list of relationship names
+    - `os_relationship_workspaces`: A single string used as the workspace ID of all relationships, or a list of workspace IDs (one per relationship)
+    - `relationship_fields`: A list of dictionaries, each with the additional properties for a relationship
+    - `os_relationship_uids`: The list of IDs for the relationships. Only necessary in case of an update
+    - `allow_multi_cardinality`: If True, allows multiple relationships from the same source, target, and relationship name. Otherwise, it updates them
+        Note this uniqueness is only applied per-workspace
+    - `ontology_name`: The name of the ontology
+    - `client`: The Client with which to connect to Octostar. If None, the default one is used
+
+    ## Returns
+    The list of created/updated relationship records
+
+    ## Raises
+    - `AssertionError`: If the length of the input lists do not match or if some source entities
+        do not have a workspace ID
+    - `NotImplementedError`: If the operation is not supported for the ontology or for the given data
+    - `ApiValidationError`: The server rejected one or more
+        relationships. The per-entity breakdown is exposed as
+        structured attributes on the exception (see the class
+        docstring for the shape of each):
+        - `entity_errors`: ``list[{"entity_id": str, "error": str}]``
+        - `invalid_fields`: ``dict[entity_id, list[field_name]]`` —
+          drop these fields and re-upsert to recover
+        - `reserved_keyword_entity_types`: ``set[str]`` of entity-type
+          tokens flagged as reserved keywords
+        Subclass of `ApiConnectionError` for backward compatibility.
+    - `ApiConnectionError`: If the transport layer failed (network error,
+        non-2xx response).
+    - `ValueError`: If relationship uniqueness is enforced when the relationship is already not unique
+    """
+    if not client:
+        client = get_default_client()
+    if not ontology_name:
+        ontology_name = client.ontology
+    if ontology_name != client.ontology:
+        raise NotImplementedError(
+            "This operation is currently only supported on the current ontology!"
+        )
+    assert len(entities_from) > 0
+    assert len(entities_from) == len(entities_to)
+    if isinstance(os_relationship_workspaces, str):
+        os_relationship_workspaces = [
+            os_relationship_workspaces for _ in range(len(entities_from))
+        ]
+    assert len(os_relationship_workspaces) == len(entities_from)
+    if not os_relationship_uids:
+        os_relationship_uids = [str(uuid.uuid4()) for _ in range(len(entities_from))]
+    for i in range(len(os_relationship_uids)):
+        if not os_relationship_uids[i]:
+            os_relationship_uids[i] = str(uuid.uuid4())
+    assert len(os_relationship_uids) == len(entities_from)
+    if isinstance(os_relationship_names, str):
+        os_relationship_names = [
+            os_relationship_names for _ in range(len(entities_from))
+        ]
+    assert len(entities_from) == len(os_relationship_names)
+    if not relationship_fields:
+        relationship_fields = [dict() for _ in range(len(entities_from))]
+    for i in range(len(relationship_fields)):
+        if not relationship_fields[i]:
+            relationship_fields[i] = dict()
+    assert len(relationship_fields) == len(entities_from)
+    # ensure each relationship exists uniquely, fetch the IDs in that case
+    if not allow_multi_cardinality:
+        check_existence_table = [
+            (
+                os_relationship_workspaces[i],
+                os_relationship_names[i],
+                entities_from[i]["os_entity_uid"],
+                entities_from[i]["entity_type"],
+                entities_to[i]["os_entity_uid"],
+                entities_to[i]["entity_type"],
+            )
+            for i in range(len(entities_from))
+        ]
+        check_existence_subquery = ""
+        for i in range(len(check_existence_table)):
+            elem = check_existence_table[i]
+            if i == 0:
+                check_existence_subquery = f"""SELECT '{i}' AS n,
+'{elem[0]}' AS os_workspace,
+'{elem[1]}' AS os_relationship_name,
+'{elem[2]}' AS os_entity_uid_from,
+'{elem[3]}' AS os_entity_type_from,
+'{elem[4]}' AS os_entity_uid_to,
+'{elem[5]}' AS os_entity_type_to """
+            else:
+                check_existence_subquery += f"""SELECT '{i}', '{elem[0]}', '{elem[1]}', '{elem[2]}', '{elem[3]}', '{elem[4]}', '{elem[5]}' """
+            if i < len(check_existence_table) - 1:
+                check_existence_subquery += " UNION ALL "
+        check_existence_query = (
+            """SELECT
+os_entity_uid,
+v.n AS n,
+v.os_relationship_name AS os_relationship_name,
+v.os_workspace AS os_workspace,
+v.os_entity_uid_from AS os_entity_uid_from,
+v.os_entity_type_from AS os_entity_type_from,
+v.os_entity_uid_to AS os_entity_uid_to,
+v.os_entity_type_to AS os_entity_type_to
+FROM dtimbr.os_workspace_relationship AS o RIGHT JOIN ("""
+            + check_existence_subquery
+            + """
+) AS v ON
+v.os_workspace = o.os_workspace AND v.os_relationship_name = o.os_relationship_name
+AND v.os_entity_uid_from = o.os_entity_uid_from AND
+v.os_entity_type_from = o.os_entity_type_from AND v.os_entity_uid_to = o.os_entity_uid_to
+AND v.os_entity_type_to = o.os_entity_type_to"""
+        )
+        result = query_ontology.sync(check_existence_query, client=client)
+        result = list(sorted(result, key=lambda x: x["n"]))
+        for i, result_i in itertools.groupby(result, key=lambda x: x["n"]):
+            result_i = list(result_i)
+            if result_i:
+                if len(result_i) > 1:
+                    raise ValueError(
+                        "Cannot enforce cardinality = 1 when cardinality is already greater than 1!"
+                    )
+                os_relationship_uids[int(i)] = result_i[0]["os_entity_uid"]
+    entities = [
+        {
+            "entity_type": "os_workspace_relationship",
+            "os_entity_uid": os_relationship_uids[i],
+            "fields": {
+                **relationship_fields[i],
+                "os_entity_uid_from": entities_from[i]["os_entity_uid"],
+                "os_entity_type_from": entities_from[i]["entity_type"],
+                "os_entity_uid_to": entities_to[i]["os_entity_uid"],
+                "os_entity_type_to": entities_to[i]["entity_type"],
+                "os_relationship_name": os_relationship_names[i],
+            },
+        }
+        for i in range(len(entities_from))
+    ]
+    return upsert_entities.sync(os_relationship_workspaces, entities, client=client)
+
+
+def asyncio(
+    entities_from: List[EntityBase],
+    entities_to: List[EntityBase],
+    os_relationship_names: Union[str, List[str]],
+    os_relationship_workspaces: Union[str, List[str]],
+    relationship_fields: Optional[List[Dict[str, Any]]] = None,
+    os_relationship_uids: Optional[List[Optional[str]]] = None,
+    allow_multi_cardinality: bool = False,
+    ontology_name: str = None,
+    client: Client = None,
+):
+    """
+    # NOT IMPLEMENTED
+    """
+    raise NotImplementedError()

octostar/utils/workspace/write_attachment.py

@@ -0,0 +1,263 @@
+from typing import Dict, Optional, Union
+from io import BytesIO
+import json
+import uuid
+import logging
+import httpx
+import aiofiles
+from urllib.parse import quote
+
+_logger = logging.getLogger(__name__)
+
+from ...client import Client, get_default_client
+from ..commons import network_retry_strategy
+from ..exceptions import ApiConnectionError
+
+
+DEFAULT_TIMEOUT = 120
+
+
+def _entity_url(client: Client, os_workspace: str) -> str:
+    """Build the entity creation endpoint URL."""
+    base = client.get_base_url_v1()
+    return f"{base}/api/v2/entities/{quote(os_workspace, safe='')}"
+
+
+def _create_entity_sync(
+    os_workspace: str,
+    entity_fields: dict,
+    file=None,
+    content_type: str = None,
+    client: "Client" = None,
+) -> dict:
+    """Create or upsert an entity via the multipart endpoint.
+
+    Sends a multipart/form-data request with an ``entity`` JSON part
+    and an optional ``file`` binary part.
+    """
+    url = _entity_url(client, os_workspace)
+    headers = dict(client.get_headers())
+
+    entity_json = json.dumps(entity_fields).encode("utf-8")
+
+    should_close = False
+    file_obj = None
+    if file is not None:
+        if isinstance(file, str):
+            file_obj = open(file, "rb")
+            should_close = True
+        elif isinstance(file, bytes):
+            file_obj = BytesIO(file)
+        else:
+            file_obj = file
+
+    try:
+        response = None
+        try:
+            for attempt in network_retry_strategy():
+                with attempt:
+                    parts = [("entity", ("entity", entity_json, "application/json"))]
+                    if file_obj is not None:
+                        if hasattr(file_obj, "seek"):
+                            file_obj.seek(0)
+                        parts.append(
+                            (
+                                "file",
+                                (
+                                    "file",
+                                    file_obj,
+                                    content_type or "application/octet-stream",
+                                ),
+                            )
+                        )
+                    with httpx.Client(timeout=DEFAULT_TIMEOUT) as http_client:
+                        response = http_client.post(url, files=parts, headers=headers)
+                        response.raise_for_status()
+        except Exception:
+            raise ApiConnectionError("write_attachment", response, client)
+    finally:
+        if should_close and file_obj:
+            file_obj.close()
+
+    return response.json()
+
+
+async def _create_entity_async(
+    os_workspace: str,
+    entity_fields: dict,
+    file=None,
+    content_type: str = None,
+    client: "Client" = None,
+) -> dict:
+    """Create or upsert an entity via the multipart endpoint (async).
+
+    Reads file content into memory for async compatibility, then sends a
+    multipart/form-data request with an ``entity`` JSON part and an optional
+    ``file`` binary part.
+    """
+    url = _entity_url(client, os_workspace)
+    headers = dict(client.get_headers())
+
+    entity_json = json.dumps(entity_fields).encode("utf-8")
+
+    file_content = None
+    if file is not None:
+        if isinstance(file, str):
+            async with aiofiles.open(file, "rb") as f:
+                file_content = await f.read()
+        elif isinstance(file, bytes):
+            file_content = file
+        else:
+            file_content = file.read() if hasattr(file, "read") else file
+
+    response = None
+    try:
+        for attempt in network_retry_strategy():
+            with attempt:
+                parts = [("entity", ("entity", entity_json, "application/json"))]
+                if file_content is not None:
+                    parts.append(
+                        (
+                            "file",
+                            (
+                                "file",
+                                file_content,
+                                content_type or "application/octet-stream",
+                            ),
+                        )
+                    )
+                async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as http_client:
+                    response = await http_client.post(url, files=parts, headers=headers)
+                    response.raise_for_status()
+    except Exception:
+        raise ApiConnectionError("write_attachment", response, client)
+
+    return response.json()
+
+
+def sync(
+    os_workspace: str,
+    os_entity_uid: Union[str, None],
+    entity_type: str,
+    filetype: str,
+    file: Union[str, bytes, BytesIO],
+    fields: Optional[Dict] = None,
+    ontology_name: str = None,
+    client: Client = None,
+):
+    """
+    # Write an attachment to a workspace for an attachable entity
+
+    This includes creating or updating the local entity record as well as writing
+    the attachment contents to storage so that they can be retrieved with
+    read_attachment(). Note that using upsert_entity() only allows to save a record
+    about an entity, not the attachment of said entity.
+
+    ## Arguments
+    - `os_workspace`: The workspace ID the entity should belong to
+    - `os_entity_uid`: The entity ID
+    - `entity_type`: The ontological entity type. Ensure this inherits from 'os_attachable'
+    - `filetype`: The MIME type for the attachment
+    - `file`: The attachment, either an UTF-8 string or a bytes-like
+    - `fields`: Optional dict of additional entity fields to include on the entity.
+    - `ontology_name`: The name of the ontology
+    - `client`: The Client with which to connect to Octostar. If None, the default one is used
+
+    ## Returns
+    The record of the written entity
+
+    ## Raises
+    - `NotImplementedError`: If the operation is not supported for the ontology
+    - `ConnectionError`: If multiple (ambiguous) S3 urls are returned for this file
+    - `ApiConnectionError`: If the operation was unsuccessful on the server
+    """
+    if not client:
+        client = get_default_client()
+    if not ontology_name:
+        ontology_name = client.ontology
+    if ontology_name != client.ontology:
+        raise NotImplementedError(
+            "This operation is currently only supported on the current ontology!"
+        )
+    if not os_entity_uid:
+        os_entity_uid = str(uuid.uuid4())
+
+    entity_fields = {**(fields or {})}
+    entity_fields.update(
+        {
+            "entity_id": os_entity_uid,
+            "entity_type": entity_type,
+            "os_item_content_type": filetype,
+        }
+    )
+    return _create_entity_sync(
+        os_workspace=os_workspace,
+        entity_fields=entity_fields,
+        file=file,
+        content_type=filetype,
+        client=client,
+    )
+
+
+async def asyncio(
+    os_workspace: str,
+    os_entity_uid: Union[str, None],
+    entity_type: str,
+    filetype: str,
+    file: Union[str, bytes, BytesIO],
+    fields: Optional[Dict] = None,
+    ontology_name: str = None,
+    client: Client = None,
+):
+    """
+    # Write an attachment to a workspace for an attachable entity (async)
+
+    This includes creating or updating the local entity record as well as writing
+    the attachment contents to storage so that they can be retrieved with
+    read_attachment(). Note that using upsert_entity() only allows to save a record
+    about an entity, not the attachment of said entity.
+
+    ## Arguments
+    - `os_workspace`: The workspace ID the entity should belong to
+    - `os_entity_uid`: The entity ID
+    - `entity_type`: The ontological entity type. Ensure this inherits from 'os_attachable'
+    - `filetype`: The MIME type for the attachment
+    - `file`: The attachment, either an UTF-8 string or a bytes-like
+    - `fields`: Optional dict of additional entity fields to include on the entity.
+    - `ontology_name`: The name of the ontology
+    - `client`: The Client with which to connect to Octostar. If None, the default one is used
+
+    ## Returns
+    The record of the written entity
+
+    ## Raises
+    - `NotImplementedError`: If the operation is not supported for the ontology
+    - `ConnectionError`: If multiple (ambiguous) S3 urls are returned for this file
+    - `ApiConnectionError`: If the operation was unsuccessful on the server
+    """
+    if not client:
+        client = get_default_client()
+    if not ontology_name:
+        ontology_name = client.ontology
+    if ontology_name != client.ontology:
+        raise NotImplementedError(
+            "This operation is currently only supported on the current ontology!"
+        )
+    if not os_entity_uid:
+        os_entity_uid = str(uuid.uuid4())
+
+    entity_fields = {**(fields or {})}
+    entity_fields.update(
+        {
+            "entity_id": os_entity_uid,
+            "entity_type": entity_type,
+            "os_item_content_type": filetype,
+        }
+    )
+    return await _create_entity_async(
+        os_workspace=os_workspace,
+        entity_fields=entity_fields,
+        file=file,
+        content_type=filetype,
+        client=client,
+    )