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

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

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:

dara/core/internal/store.py

@@ -17,7 +17,7 @@ limitations under the License.
 
 from typing import Any, Dict, List, Optional
 
-from dara.core.base_definitions import CacheType, PendingTask, PendingValue
+from dara.core.base_definitions import CacheType, PendingTask
 from dara.core.internal.utils import get_cache_scope
 
 
@@ -59,8 +59,6 @@ class Store:
         cache_key = get_cache_scope(cache_type)
         value = self._store.get(cache_key, {}).get(key)
 
-        if isinstance(value, PendingValue):
-            return await value.wait()
         if isinstance(value, PendingTask):
             return await value.run()
         return value
@@ -86,33 +84,8 @@ class Store:
         if self._store.get(cache_key) is None:
             self._store[cache_key] = {}
 
-        # If there is a PendingValue set for this key then trigger its resolution
-        if isinstance(self._store[cache_key].get(key), PendingValue):
-            if error is not None:
-                self._store[cache_key][key].error(error)
-            else:
-                self._store[cache_key][key].resolve(value)
-
         self._store[cache_key][key] = value
 
-    def set_pending_value(self, key: str, cache_type: Optional[CacheType] = CacheType.GLOBAL):
-        """
-        Set a pending state for a value in the store. This will trigger the async behavior of the get call if subsequent
-        requests ask for the same key. A future is created in the store, which all requests then listen for the
-        resolution of before returning.
-
-        :param key: the key to set as pending
-        :param cache_type: whether to pull the value from the specified cache specific store or the global one, defaults to
-                            the global one
-        """
-        cache_key = get_cache_scope(cache_type)
-        if self._store.get(cache_key) is None:
-            self._store[cache_key] = {}
-
-        pending_val = PendingValue()
-
-        self._store[cache_key][key] = pending_val
-
     def set_pending_task(self, key: str, pending_task: PendingTask, cache_type: Optional[CacheType] = CacheType.GLOBAL):
         """
         Store a pending task state for a given key in the store. This will trigger the async behavior of the get call if subsequent
@@ -158,7 +131,7 @@ class Store:
                 # Otherwise go through and remove any non-pending values
                 keys = list(cache_type_store.keys())
                 for key in keys:
-                    if not isinstance(cache_type_store[key], (PendingValue, PendingTask)):
+                    if not isinstance(cache_type_store[key], PendingTask):
                         cache_type_store.pop(key)
 
     def list(self, cache_type: Optional[CacheType] = CacheType.GLOBAL) -> List[str]: