dara-core 1.19.0__py3-none-any.whl → 1.20.0a1__py3-none-any.whl

dara/core/internal/multi_resource_lock.py

@@ -0,0 +1,70 @@
+from collections import Counter
+from contextlib import asynccontextmanager
+
+import anyio
+
+
+class MultiResourceLock:
+    """
+    A class that manages multiple named locks for concurrent access to shared resources.
+
+    This class allows for acquiring and releasing locks on named resources, ensuring
+    that only one task can access a specific resource at a time. It automatically
+    creates locks for new resources and cleans them up when they're no longer in use.
+
+    :reentrant:
+        If True a task can acquire the same resource more than once; every
+        subsequent acquire of an already-held lock is a no-op.  If False the
+        second attempt raises ``RuntimeError``.
+    """
+
+    def __init__(self):
+        self._locks: dict[str, anyio.Lock] = {}
+        self._waiters = Counter[str]()
+        self._cleanup_lock = anyio.Lock()
+
+    def is_locked(self, resource_name: str) -> bool:
+        """
+        Check if a lock for the specified resource is currently held.
+
+        :param resource_name (str): The name of the resource to check.
+        :return: True if the lock is held, False otherwise.
+        """
+        return resource_name in self._locks and self._locks[resource_name].locked()
+
+    @asynccontextmanager
+    async def acquire(self, resource_name: str):
+        """
+        Acquire a lock for the specified resource.
+
+        This method is an async context manager that acquires a lock for the given
+        resource name. If the lock doesn't exist, it creates one. It also keeps
+        track of waiters to ensure proper cleanup when the resource is no longer in use.
+
+        :param resource_name (str): The name of the resource to lock.
+
+        Usage:
+        ```python
+        async with multi_lock.acquire_lock("resource_a"):
+            # Critical section for "resource_a"
+            ...
+        ```
+
+        Note:
+            The lock is automatically released when exiting the context manager.
+        """
+
+        async with self._cleanup_lock:
+            if resource_name not in self._locks:
+                self._locks[resource_name] = anyio.Lock()
+            self._waiters[resource_name] += 1
+
+        try:
+            async with self._locks[resource_name]:
+                yield
+        finally:
+            async with self._cleanup_lock:
+                self._waiters[resource_name] -= 1
+                if self._waiters[resource_name] <= 0:
+                    del self._waiters[resource_name]
+                    del self._locks[resource_name]

dara/core/internal/normalization.py

@@ -31,7 +31,6 @@ from typing import (
 from typing_extensions import TypedDict, TypeGuard
 
 from dara.core.base_definitions import DaraBaseModel as BaseModel
-from dara.core.internal.hashing import hash_object
 
 JsonLike = Union[Mapping, List]
 
@@ -81,10 +80,6 @@ def _get_identifier(obj: Referrable) -> str:
         nested = ','.join(cast(List[str], obj['nested']))
         identifier = f'{identifier}:{nested}'
 
-    if _is_referrable_with_filters(obj):
-        filter_hash = hash_object(obj['filters'])
-        identifier = f'{identifier}:{filter_hash}'
-
     return identifier
 
 

dara/core/internal/pandas_utils.py

@@ -15,13 +15,24 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from typing import Optional, TypeVar, cast
+import json
+import uuid
+from typing import Any, Literal, Optional, TypeVar, Union, cast, overload
 
-from pandas import DataFrame, MultiIndex
+from pandas import DataFrame, MultiIndex, Series
+from typing_extensions import TypedDict, TypeGuard
 
 INDEX = '__index__'
 
 
+@overload
+def append_index(df: DataFrame) -> DataFrame: ...
+
+
+@overload
+def append_index(df: None) -> None: ...
+
+
 def append_index(df: Optional[DataFrame]) -> Optional[DataFrame]:
     """
     Add a numerical index column to the dataframe
@@ -65,6 +76,12 @@ def df_convert_to_internal(original_df: DataFrame) -> DataFrame:
     if any(isinstance(c, str) and c.startswith('__col__') for c in df.columns):
         return df
 
+    # Apply display transformations to the DataFrame
+    format_for_display(df)
+
+    # Append index to match the way we process the original DataFrame
+    df = cast(DataFrame, append_index(df))
+
     # Handle hierarchical columns: [(A, B), (A, C)] -> ['A_B', 'A_C']
     if isinstance(df.columns, MultiIndex):
         df.columns = ['_'.join(col).strip() if col[0] != INDEX else INDEX for col in df.columns.values]
@@ -90,4 +107,89 @@ def df_convert_to_internal(original_df: DataFrame) -> DataFrame:
 
 
 def df_to_json(df: DataFrame) -> str:
-    return df_convert_to_internal(df).to_json(orient='records') or ''
+    return df_convert_to_internal(df).to_json(orient='records', date_unit='ns') or ''
+
+
+def format_for_display(df: DataFrame) -> None:
+    """
+    Apply transformations to a DataFrame to make it suitable for display.
+    Not: this does NOT make a copy of the DataFrame
+    """
+    for col in df.columns:
+        column_data = df[col]
+        if isinstance(column_data, DataFrame):
+            # Handle duplicate column names - format each column in the sub-DataFrame
+            for sub_col in column_data.columns:
+                if isinstance(column_data[sub_col], Series) and column_data[sub_col].dtype == 'object':
+                    column_data.loc[:, sub_col] = column_data[sub_col].apply(str)
+        elif column_data.dtype == 'object':
+            # We need to convert all values to string to avoid issues with
+            # displaying data in the Table component, for example when
+            # displaying datetime and number objects in the same column
+            df.loc[:, col] = column_data.apply(str)
+
+
+class FieldType(TypedDict):
+    name: Union[str, tuple[str, ...]]
+    type: Literal['integer', 'number', 'boolean', 'datetime', 'duration', 'any', 'str']
+
+
+class DataFrameSchema(TypedDict):
+    fields: list[FieldType]
+    primaryKey: list[str]
+
+
+class DataResponse(TypedDict):
+    data: Optional[DataFrame]
+    count: int
+    schema: Optional[DataFrameSchema]
+
+
+def is_data_response(response: Any) -> TypeGuard[DataResponse]:
+    has_shape = isinstance(response, dict) and 'data' in response and 'count' in response
+    if not has_shape:
+        return False
+    return response['data'] is None or isinstance(response['data'], DataFrame)
+
+
+def data_response_to_json(response: DataResponse) -> str:
+    """
+    Serialize a DataResponse to JSON.
+
+    json.dumps() custom serializers only accept value->value mappings, whereas `to_json` on pandas returns a string directly.
+    To avoid double serialization, we first insert a placeholder string and then replace it with the actual serialized JSON.
+    """
+    placeholder = str(uuid.uuid4())
+
+    def _custom_serializer(obj: Any) -> Any:
+        if isinstance(obj, DataFrame):
+            return placeholder
+        raise TypeError(f'Object of type {type(obj)} is not JSON serializable')
+
+    result = json.dumps(response, default=_custom_serializer)
+    result = result.replace(
+        f'"{placeholder}"', df_to_json(response['data']) if response['data'] is not None else 'null'
+    )
+    return result
+
+
+def build_data_response(data: DataFrame, count: int) -> DataResponse:
+    data_internal = df_convert_to_internal(data)
+    schema = get_schema(data_internal)
+
+    return DataResponse(data=data, count=count, schema=schema)
+
+
+def get_schema(df: DataFrame):
+    from pandas.io.json._table_schema import build_table_schema
+
+    raw_schema = build_table_schema(df)
+
+    for field_data in cast(list, raw_schema['fields']):
+        if field_data.get('type') == 'datetime':
+            # for datetime fields we need to know the resolution, so we get the actual e.g. `datetime64[ns]` string
+            column_name = field_data.get('name')
+            dtype_str = str(df[column_name].dtype)
+            field_data['type'] = dtype_str
+
+    return cast(DataFrameSchema, raw_schema)

dara/core/internal/pool/definitions.py

@@ -95,7 +95,7 @@ class TaskDefinition:
     def __await__(self):
         """Await the underlying event, then return or raise the result"""
         yield from self.event.wait().__await__()
-        if isinstance(self.result, Exception):
+        if isinstance(self.result, BaseException):
             raise self.result
         return self.result
 

dara/core/internal/pool/task_pool.py

@@ -158,7 +158,7 @@ class TaskPool:
 
         task = self.tasks.pop(task_uid)
         if not task.event.is_set():
-            task.result = Exception('Task cancelled')
+            task.result = anyio.get_cancelled_exc_class()()
             task.event.set()
 
         # Task in progress, stop the worker
@@ -499,11 +499,14 @@ class TaskPool:
             while self.status not in (PoolStatus.ERROR, PoolStatus.STOPPED):
                 await anyio.sleep(0.1)
 
-                self._handle_excess_workers()
-                self._handle_orphaned_workers()
-                self._handle_dead_workers()
-                self._create_workers()
-                await self._process_next_worker_message()
+                try:
+                    self._handle_excess_workers()
+                    self._handle_orphaned_workers()
+                    self._handle_dead_workers()
+                    self._create_workers()
+                    await self._process_next_worker_message()
+                except Exception as e:
+                    dev_logger.error('Error in task pool', e)
         finally:
             self.loop_stopped.set()
 

dara/core/internal/pool/utils.py

@@ -27,6 +27,8 @@ from typing import Any, Callable, Optional, Tuple
 import anyio
 from tblib import Traceback
 
+from dara.core.logging import dev_logger
+
 
 class SubprocessException:
     """
@@ -131,20 +133,23 @@ async def stop_process_async(process: BaseProcess, timeout: float = 3):
     # Terminate and wait for it to shutdown
     process.terminate()
 
-    # mimic process.join() in an async way to not block
-    await wait_while(process.is_alive, timeout)
-
-    # If it's still alive
-    if process.is_alive():
-        try:
-            os.kill(process.pid, signal.SIGKILL)
-            await wait_while(process.is_alive, timeout)
-        except OSError as e:
-            raise RuntimeError(f'Unable to terminate subprocess with PID {process.pid}') from e
-
-    # If it's still alive raise an exception
-    if process.is_alive():
-        raise RuntimeError(f'Unable to terminate subprocess with PID {process.pid}')
+    try:
+        # mimic process.join() in an async way to not block
+        await wait_while(process.is_alive, timeout)
+
+        # If it's still alive
+        if process.is_alive():
+            try:
+                os.kill(process.pid, signal.SIGKILL)
+                await wait_while(process.is_alive, timeout)
+            except OSError as e:
+                raise RuntimeError(f'Unable to terminate subprocess with PID {process.pid}') from e
+
+        # If it's still alive raise an exception
+        if process.is_alive():
+            raise RuntimeError(f'Unable to terminate subprocess with PID {process.pid}')
+    except Exception as e:
+        dev_logger.error('Error stopping process', e)
 
 
 def stop_process(process: BaseProcess, timeout: float = 3):

dara/core/internal/registries.py

@@ -27,11 +27,11 @@ from dara.core.definitions import (
     EndpointConfiguration,
     Template,
 )
-from dara.core.interactivity.data_variable import DataVariableRegistryEntry
 from dara.core.interactivity.derived_variable import (
     DerivedVariableRegistryEntry,
     LatestValueRegistryEntry,
 )
+from dara.core.interactivity.server_variable import ServerVariableRegistryEntry
 from dara.core.internal.download import DownloadDataEntry
 from dara.core.internal.registry import Registry, RegistryType
 from dara.core.internal.websocket import CustomClientMessagePayload
@@ -44,7 +44,8 @@ upload_resolver_registry = Registry[UploadResolverDef](
 )  # functions for upload resolvers requiring backend calls
 component_registry = Registry[ComponentTypeAnnotation](RegistryType.COMPONENTS, CORE_COMPONENTS)
 config_registry = Registry[EndpointConfiguration](RegistryType.ENDPOINT_CONFIG)
-data_variable_registry = Registry[DataVariableRegistryEntry](RegistryType.DATA_VARIABLE, allow_duplicates=False)
+server_variable_registry = Registry[ServerVariableRegistryEntry](RegistryType.SERVER_VARIABLE, allow_duplicates=False)
+"""map of server variable uid -> server variable entry"""
 derived_variable_registry = Registry[DerivedVariableRegistryEntry](
     RegistryType.DERIVED_VARIABLE, allow_duplicates=False
 )

dara/core/internal/registry.py

@@ -32,7 +32,7 @@ class RegistryType(str, Enum):
     DOWNLOAD = 'Download'
     COMPONENTS = 'Components'
     ENDPOINT_CONFIG = 'Endpoint Configuration'
-    DATA_VARIABLE = 'DataVariable'
+    SERVER_VARIABLE = 'ServerVariable'
     DERIVED_VARIABLE = 'DerivedVariable'
     LAST_VALUE = 'LatestValue'
     TEMPLATE = 'Template'

dara/core/internal/registry_lookup.py

@@ -16,7 +16,7 @@ limitations under the License.
 """
 
 from collections.abc import Coroutine
-from typing import Callable, Dict, Literal, Union
+from typing import Callable, Dict, Literal, TypeVar, Union
 
 from dara.core.internal.registry import Registry, RegistryType
 from dara.core.internal.utils import async_dedupe
@@ -24,8 +24,8 @@ from dara.core.internal.utils import async_dedupe
 RegistryLookupKey = Literal[
     RegistryType.ACTION,
     RegistryType.COMPONENTS,
-    RegistryType.DATA_VARIABLE,
     RegistryType.DERIVED_VARIABLE,
+    RegistryType.SERVER_VARIABLE,
     RegistryType.STATIC_KWARGS,
     RegistryType.UPLOAD_RESOLVER,
     RegistryType.BACKEND_STORE,
@@ -33,6 +33,8 @@ RegistryLookupKey = Literal[
 ]
 CustomRegistryLookup = Dict[RegistryLookupKey, Callable[[str], Coroutine]]
 
+RegistryType = TypeVar('RegistryType')
+
 
 class RegistryLookup:
     """
@@ -45,7 +47,7 @@ class RegistryLookup:
         self.handlers = handlers
 
     @async_dedupe
-    async def get(self, registry: Registry, uid: str):
+    async def get(self, registry: Registry[RegistryType], uid: str) -> RegistryType:
         """
         Get the entry from registry by uid.
         If uid is not in registry and it has a external handler that defined, will execute the handler

dara/core/internal/routing.py

@@ -16,7 +16,6 @@ limitations under the License.
 """
 
 import inspect
-import json
 import os
 from collections.abc import Mapping
 from functools import wraps
@@ -24,7 +23,6 @@ from importlib.metadata import version
 from typing import Any, Callable, Dict, List, Optional
 
 import anyio
-import pandas
 from fastapi import (
     APIRouter,
     Body,
@@ -40,26 +38,28 @@ from fastapi.responses import StreamingResponse
 from pandas import DataFrame
 from pydantic import BaseModel
 from starlette.background import BackgroundTask
+from starlette.status import HTTP_415_UNSUPPORTED_MEDIA_TYPE
 
 from dara.core.auth.routes import verify_session
-from dara.core.base_definitions import ActionResolverDef, BaseTask, UploadResolverDef
+from dara.core.base_definitions import ActionResolverDef, BaseTask, NonTabularDataError, UploadResolverDef
 from dara.core.configuration import Configuration
-from dara.core.interactivity.any_data_variable import DataVariableRegistryEntry, upload
+from dara.core.interactivity.any_data_variable import upload
 from dara.core.interactivity.filtering import FilterQuery, Pagination
+from dara.core.interactivity.server_variable import ServerVariable
 from dara.core.internal.cache_store import CacheStore
 from dara.core.internal.download import DownloadRegistryEntry
 from dara.core.internal.execute_action import CURRENT_ACTION_ID
 from dara.core.internal.normalization import NormalizedPayload, denormalize, normalize
-from dara.core.internal.pandas_utils import df_to_json
+from dara.core.internal.pandas_utils import data_response_to_json, df_to_json, is_data_response
 from dara.core.internal.registries import (
     action_def_registry,
     action_registry,
     backend_store_registry,
     component_registry,
-    data_variable_registry,
     derived_variable_registry,
     download_code_registry,
     latest_value_registry,
+    server_variable_registry,
     static_kwargs_registry,
     template_registry,
     upload_resolver_registry,
@@ -307,135 +307,65 @@ def create_router(config: Configuration):
         except KeyError as err:
             raise ValueError(f'Could not find latest value for derived variable with uid: {uid}') from err
 
-    class DataVariableRequestBody(BaseModel):
+    class TabularRequestBody(BaseModel):
         filters: Optional[FilterQuery] = None
-        cache_key: Optional[str] = None
-        ws_channel: Optional[str] = None
+        ws_channel: str
+        dv_values: Optional[NormalizedPayload[List[Any]]] = None
+        """DerivedVariable values if variable is a DerivedVariable"""
+        force_key: Optional[str] = None
+        """Optional force key if variable is a DerivedVariable and a recalculation is forced"""
 
-    @core_api_router.post('/data-variable/{uid}', dependencies=[Depends(verify_session)])
-    async def get_data_variable(
+    @core_api_router.post('/tabular-variable/{uid}', dependencies=[Depends(verify_session)])
+    async def get_tabular_variable(
         uid: str,
-        body: DataVariableRequestBody,
+        body: TabularRequestBody,
         offset: Optional[int] = None,
         limit: Optional[int] = None,
         order_by: Optional[str] = None,
         index: Optional[str] = None,
     ):
+        """
+        Generic endpoint for getting tabular data from a variable.
+        Supports ServerVariables and DerivedVariables.
+        """
+        WS_CHANNEL.set(body.ws_channel)
+
         try:
-            store: CacheStore = utils_registry.get('Store')
-            task_mgr: TaskManager = utils_registry.get('TaskManager')
+            pagination = Pagination(offset=offset, limit=limit, orderBy=order_by, index=index)
             registry_mgr: RegistryLookup = utils_registry.get('RegistryLookup')
-            data_variable_entry: DataVariableRegistryEntry = await registry_mgr.get(data_variable_registry, uid)
-
-            data = None
-            WS_CHANNEL.set(body.ws_channel)
-
-            if data_variable_entry.type == 'derived':
-                if body.cache_key is None:
-                    raise HTTPException(
-                        status_code=400,
-                        detail='Cache key is required for derived data variables',
-                    )
-
-                if body.ws_channel is None:
-                    raise HTTPException(
-                        status_code=400,
-                        detail='Websocket channel is required for derived data variables',
-                    )
-
-                derived_variable_entry = await registry_mgr.get(derived_variable_registry, uid)
-
-                data = await data_variable_entry.get_data(
-                    derived_variable_entry,
-                    data_variable_entry,
-                    body.cache_key,
-                    store,
-                    body.filters,
-                    Pagination(offset=offset, limit=limit, orderBy=order_by, index=index),
-                    format_for_display=True,
-                )
-                if isinstance(data, BaseTask):
-                    await task_mgr.run_task(data, body.ws_channel)
-                    return {'task_id': data.task_id}
-            elif data_variable_entry.type == 'plain':
-                data = await data_variable_entry.get_data(
-                    data_variable_entry,
-                    store,
-                    body.filters,
-                    Pagination(offset=offset, limit=limit, orderBy=order_by, index=index),
-                    format_for_display=True,
-                )
 
-            dev_logger.debug(
-                f'DataVariable {data_variable_entry.uid[:3]}..{data_variable_entry.uid[-3:]}',
-                'return value',
-                {
-                    'value': data.describe() if isinstance(data, pandas.DataFrame) else None,
-                    'uid': uid,
-                },  # type: ignore
-            )
-
-            if data is None:
-                return None
-
-            # Explicitly convert to JSON to avoid implicit serialization;
-            # return as records as that makes more sense in a JSON structure
-            return Response(
-                content=df_to_json(data) if isinstance(data, pandas.DataFrame) else data,
-                media_type='application/json',
-            )  # type: ignore
-        except ValueError as e:
-            raise HTTPException(status_code=400, detail=str(e)) from e
+            # ServerVariable
+            if body.dv_values is None:
+                server_variable_entry = await registry_mgr.get(server_variable_registry, uid)
+                data_response = await ServerVariable.get_tabular_data(server_variable_entry, body.filters, pagination)
+                return Response(data_response_to_json(data_response), media_type='application/json')
 
-    class DataVariableCountRequestBody(BaseModel):
-        cache_key: Optional[str] = None
-        filters: Optional[FilterQuery] = None
-
-    @core_api_router.post('/data-variable/{uid}/count', dependencies=[Depends(verify_session)])
-    async def get_data_variable_count(uid: str, body: Optional[DataVariableCountRequestBody] = None):
-        try:
+            # DerivedVariable
             store: CacheStore = utils_registry.get('Store')
-            registry_mgr: RegistryLookup = utils_registry.get('RegistryLookup')
-            variable_def = await registry_mgr.get(data_variable_registry, uid)
+            task_mgr: TaskManager = utils_registry.get('TaskManager')
+            variable_def = await registry_mgr.get(derived_variable_registry, uid)
+            values = denormalize(body.dv_values.data, body.dv_values.lookup)
 
-            if variable_def.type == 'plain':
-                return await variable_def.get_total_count(
-                    variable_def, store, body.filters if body is not None else None
-                )
+            result = await variable_def.get_tabular_data(
+                variable_def, store, task_mgr, values, body.force_key, pagination, body.filters
+            )
 
-            if body is None or body.cache_key is None:
-                raise HTTPException(
-                    status_code=400,
-                    detail="Cache key is required when requesting DerivedDataVariable's count",
-                )
+            if isinstance(result, BaseTask):
+                await task_mgr.run_task(result, body.ws_channel)
+                return {'task_id': result.task_id}
 
-            return await variable_def.get_total_count(variable_def, store, body.cache_key, body.filters)
-        except ValueError as e:
-            raise HTTPException(status_code=400, detail=str(e)) from e
+            return Response(data_response_to_json(result), media_type='application/json')
+        except NonTabularDataError as e:
+            raise HTTPException(status_code=HTTP_415_UNSUPPORTED_MEDIA_TYPE, detail=str(e)) from e
 
-    @core_api_router.get('/data-variable/{uid}/schema', dependencies=[Depends(verify_session)])
-    async def get_data_variable_schema(uid: str, cache_key: Optional[str] = None):
-        try:
-            store: CacheStore = utils_registry.get('Store')
-            registry_mgr: RegistryLookup = utils_registry.get('RegistryLookup')
-            data_def = await registry_mgr.get(data_variable_registry, uid)
-
-            if data_def.type == 'plain':
-                return await data_def.get_schema(data_def, store)
-
-            if cache_key is None:
-                raise HTTPException(
-                    status_code=400,
-                    detail='Cache key is required when requesting DerivedDataVariable schema',
-                )
-
-            # Use the other registry for derived variables
-            derived_ref = await registry_mgr.get(derived_variable_registry, uid)
-            data = await data_def.get_schema(derived_ref, store, cache_key)
-            content = json.dumps(jsonable_encoder(data)) if isinstance(data, dict) else data
-            return Response(content=content, media_type='application/json')
-        except ValueError as e:
-            raise HTTPException(status_code=400, detail=str(e)) from e
+    @core_api_router.get('/server-variable/{uid}/sequence', dependencies=[Depends(verify_session)])
+    async def get_server_variable_sequence(
+        uid: str,
+    ):
+        registry_mgr: RegistryLookup = utils_registry.get('RegistryLookup')
+        server_variable_entry = await registry_mgr.get(server_variable_registry, uid)
+        seq_num = await ServerVariable.get_sequence_number(server_variable_entry)
+        return {'sequence_number': seq_num}
 
     @core_api_router.post('/data/upload', dependencies=[Depends(verify_session)])
     async def upload_data(
@@ -473,7 +403,6 @@ def create_router(config: Configuration):
         values: NormalizedPayload[List[Any]]
         force_key: Optional[str] = None
         ws_channel: str
-        is_data_variable: Optional[bool] = False
 
     @core_api_router.post('/derived-variable/{uid}', dependencies=[Depends(verify_session)])
     async def get_derived_variable(uid: str, body: DerivedStateRequestBody):
@@ -553,9 +482,11 @@ def create_router(config: Configuration):
                 {'value': res},
             )
 
-            # Serialize dataframes correctly
+            # Serialize dataframes correctly, either direct or as a DataResponse
             if isinstance(res, DataFrame):
-                return Response(df_to_json(res))
+                return Response(df_to_json(res), media_type='application/json')
+            elif is_data_response(res):
+                return Response(data_response_to_json(res), media_type='application/json')
 
             return res
         except Exception as err: