streamlit-octostar-utils 0.5.2.dev1__tar.gz → 0.5.3__tar.gz

{streamlit_octostar_utils-0.5.2.dev1 → streamlit_octostar_utils-0.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.5.2.dev1
+Version: 0.5.3
 Summary: 
 License: MIT
 License-File: LICENSE

{streamlit_octostar_utils-0.5.2.dev1 → streamlit_octostar_utils-0.5.3}/pyproject.toml

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

{streamlit_octostar_utils-0.5.2.dev1 → streamlit_octostar_utils-0.5.3}/streamlit_octostar_utils/api_crafter/celery.py

@@ -48,6 +48,65 @@ class RedisFileLock:
         return self.lock.__exit__(exc_type, exc_val, exc_tb)
 
 
+def _install_queue_limit_patch(queue_limits):
+    """Monkey-patch kombu's Redis Channel._put for atomic queue limits.
+
+    The original _put does a bare LPUSH.  The patched version wraps
+    limited queues in a Lua script (LLEN check + LPUSH) so the
+    "is there room?" test and the "enqueue" are a single atomic
+    Redis operation.  Queues without a configured limit are
+    forwarded to the original _put unchanged.
+    """
+    from kombu.transport.redis import Channel
+    from kombu.utils.json import dumps as kombu_dumps
+
+    _ATOMIC_ENQUEUE_LUA = """
+local total = 0
+for i = 1, #KEYS do
+    total = total + redis.call('LLEN', KEYS[i])
+end
+if total >= tonumber(ARGV[1]) then
+    return 0
+end
+redis.call('LPUSH', KEYS[1], ARGV[2])
+return 1
+"""
+
+    if not hasattr(Channel, "_queue_limits"):
+        Channel._queue_limits = {}
+    Channel._queue_limits.update(queue_limits)
+
+    if getattr(Channel, "_queue_limit_patched", False):
+        return
+
+    _original_put = Channel._put
+
+    def _put_with_limit(self, queue, message, **kwargs):
+        max_tasks = Channel._queue_limits.get(queue)
+        if max_tasks is not None:
+            pri = self._get_message_priority(message, reverse=False)
+            target_key = self._q_for_pri(queue, pri)
+            all_keys = [self._q_for_pri(queue, p) for p in self.priority_steps]
+            if target_key in all_keys:
+                all_keys.remove(target_key)
+            all_keys.insert(0, target_key)
+            with self.conn_or_acquire() as client:
+                accepted = client.eval(
+                    _ATOMIC_ENQUEUE_LUA, len(all_keys), *all_keys,
+                    max_tasks, kombu_dumps(message),
+                )
+                if not accepted:
+                    from streamlit_octostar_utils.api_crafter.celery import CeleryExecutor
+                    raise CeleryExecutor.QueueFullException(
+                        f"Queue '{queue}' has reached its limit of {max_tasks} tasks!"
+                    )
+        else:
+            _original_put(self, queue, message, **kwargs)
+
+    Channel._put = _put_with_limit
+    Channel._queue_limit_patched = True
+
+
 class CeleryQueueConfig:
     def __init__(
         self,
@@ -201,10 +260,11 @@ class CeleryExecutor(object):
         self.get_thread_pool = None
         self.set_thread_pool = None
         self.io_thread_pool = None
-        self.queue_semaphores = {
-            k: threading.Semaphore(v.max_tasks_in_queue) if v.max_tasks_in_queue else None
-            for k, v in self.queue_config.items()
-        }
+
+        # Atomic queue-depth limits
+        queue_limits = {k: v.max_tasks_in_queue for k, v in self.queue_config.items() if v.max_tasks_in_queue}
+        if queue_limits:
+            _install_queue_limit_patch(queue_limits)
 
         # Folder setup
         self.base_folder = Path(base_folder).resolve()
@@ -720,17 +780,6 @@ class CeleryExecutor(object):
             self.preload_on_worker_init()
             self.app.conf.dev_preload = True
 
-        def _check_queue_llen(queue_name):
-            if self._queue_stalled.get(queue_name, False):
-                raise CeleryExecutor.QueueStalledException(
-                    f"Queue '{queue_name}' is stalled. Service temporarily unavailable."
-                )
-            if self.redis_client.llen(queue_name) >= self.queue_config[queue_name].max_tasks_in_queue:
-                raise CeleryExecutor.QueueFullException(
-                    f"Queue '{queue_name}' has reached its limit of "
-                    f"{self.queue_config[queue_name].max_tasks_in_queue} tasks!"
-                )
-
         def _write_task_data(in_folder, task_args, task_kwargs, task_id):
             serialized_data = CelerySerialized(
                 folder=in_folder,
@@ -754,22 +803,12 @@ class CeleryExecutor(object):
         queue_name = getattr(task_fn, "queue", queue_name)
         queue_name = options.get("queue", queue_name)
 
-        sem = self.queue_semaphores.get(queue_name)
-        acquired = False
-        if sem is not None:
-            if not sem.acquire(blocking=False):
-                raise CeleryExecutor.QueueFullException(
-                    f"Queue '{queue_name}' has reached its limit of "
-                    f"{self.queue_config[queue_name].max_tasks_in_queue} tasks!"
-                )
-            acquired = True
+        if self._queue_stalled.get(queue_name, False):
+            raise CeleryExecutor.QueueStalledException(
+                f"Queue '{queue_name}' is stalled. Service temporarily unavailable."
+            )
 
         try:
-            if acquired:
-                await asyncio.get_running_loop().run_in_executor(
-                    self.set_thread_pool, _check_queue_llen, queue_name
-                )
-
             if part is not None:
                 await self._write_task_data_with_part(
                     task_id, args, kwargs, part
@@ -800,9 +839,6 @@ class CeleryExecutor(object):
             except Exception:
                 pass
             raise
-        finally:
-            if acquired:
-                sem.release()
         return task_id
 
     async def _write_task_data_with_part(self, task_id, args, kwargs, part):

{streamlit_octostar_utils-0.5.2.dev1 → streamlit_octostar_utils-0.5.3}/streamlit_octostar_utils/api_crafter/nifi.py

@@ -17,6 +17,7 @@ from starlette.exceptions import HTTPException as StarletteHTTPException
 
 from octostar.utils.workspace import upsert_entities
 from octostar.utils.ontology import fetch_ontology_data
+from octostar.utils.workspace.permissions import get_permissions, PermissionLevel
 from octostar.utils.pipeline import update_processing_status
 
 from octostar.client import make_client
@@ -84,6 +85,7 @@ class NifiEntityModel(BaseModel):
         is_temporary: bool = False
         exception: dict = Field(default_factory=dict)
         last_processor_name: Optional[str] = None
+        fallback_os_workspace: Optional[str] = None
 
     class RecordModel(BaseModel):
         model_config = ConfigDict(extra="allow")
@@ -473,6 +475,7 @@ class NifiContextManager(object):
     def __init__(self, json_data, lazy_sync=True):
         if not json_data:
             raise ValueError("Nifi context manager received list of 0 entities")
+        self.permissions = {}
         self.in_batches = None
         self.out_entities = None
         self.nonlazy_sync_ids = set()
@@ -549,6 +552,16 @@ class NifiContextManager(object):
     def __enter__(self):
         return self
 
+    def get_workspaces_permissions(self, workspace_ids):
+        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)
+        permissions = {}
+        for k in workspace_ids:
+            permissions[k] = self.permissions.get(k, PermissionLevel.NONE)
+        return permissions
+
     def request_entity_sync(
         self,
         entity,
@@ -957,7 +970,29 @@ class NifiEntity(object):
 
     @property
     def write_os_workspace(self):
-        return self.record.get("os_workspace")
+        permissions = self.context.get_workspaces_permissions(
+            [
+                e
+                for e in [
+                    self.record.get("os_workspace"),
+                    self.request.get("fallback_os_workspace"),
+                ]
+                if e
+            ]
+        )
+        if (
+            self.record.get("os_workspace")
+            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)
+            >= PermissionLevel.WRITE
+        ):
+            return self.request["fallback_os_workspace"]
+        else:
+            return None
 
     @property
     def label(self):
@@ -1175,6 +1210,7 @@ class NifiEntity(object):
             "is_temporary": True,
             "exception": {},
             "last_processor_name": None,
+            "fallback_os_workspace": self.request["fallback_os_workspace"],
         }
         child_entity = NifiEntity(
             self.context,
@@ -1400,7 +1436,6 @@ class NifiEntity(object):
         os_entity_uid=None,
         os_relationship_uid=None,
         os_entity_type=FRAGMENT_ENTITY_NAME,
-        os_parent_uid=None,
         previous_fragment_uid=None,
         previous_fragment_relationship_uid=None,
         previous_fragment_relationship=PREVIOUS_FRAGMENT_RELATIONSHIP,
@@ -1413,7 +1448,7 @@ class NifiEntity(object):
         fields = {
             **{k: v for k, v in self.record.items() if k.startswith("fragment") and v is not None},
             **fields,
-            "os_parent_uid": os_parent_uid or self.record["os_entity_uid"],
+            "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,