dara-core 1.20.0a1__py3-none-any.whl → 1.20.1a1__py3-none-any.whl

dara/core/interactivity/any_data_variable.py

@@ -1,11 +1,138 @@
-from typing_extensions import TypeAlias
+"""
+Copyright 2023 Impulse Innovations Limited
 
-from dara.core.interactivity.any_variable import AnyVariable
 
-# re-export for backwards compatibility
-from .tabular_variable import *  # noqa: F403
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
 
-AnyDataVariable: TypeAlias = AnyVariable
-"""
-Deprecated alias. Tabular variables are now DerivedVariable or ServerVariable
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
 """
+
+import abc
+import io
+import os
+from typing import Any, Awaitable, Callable, Literal, Optional, TypedDict, Union, cast
+
+import pandas
+from fastapi import UploadFile
+from pydantic import ConfigDict
+
+from dara.core.base_definitions import CachedRegistryEntry, UploadResolverDef
+from dara.core.interactivity.any_variable import AnyVariable
+from dara.core.interactivity.filtering import FilterQuery
+from dara.core.internal.cache_store.cache_store import CacheStore
+from dara.core.internal.registry_lookup import RegistryLookup
+from dara.core.internal.utils import run_user_handler
+
+
+class AnyDataVariable(AnyVariable, abc.ABC):
+    """
+    AnyDataVariable represents any variable that is specifically designed to hold datasets (i.e. DataVariable, DerivedDataVariable)
+
+    :param uid: the unique identifier for this variable; if not provided a random one is generated
+    :param filters: a dictionary of filters to apply to the data
+    """
+
+    uid: str
+    filters: Optional[FilterQuery] = None
+
+    def __init__(self, uid: Optional[str] = None, **kwargs) -> None:
+        super().__init__(uid=uid, **kwargs)
+
+    def filter(self, filters: FilterQuery):
+        return self.copy(update={'filters': filters}, deep=True)
+
+
+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 DataVariableRegistryEntry(CachedRegistryEntry):
+    """
+    Registry entry for DataVariable.
+    """
+
+    type: Literal['plain', 'derived']
+    get_data: Callable[..., Awaitable[Any]]
+    """Handler to get the data from the data variable. Defaults to DataVariable.get_value for type=plain, and DerivedDataVariable.get_data for type=derived"""
+
+    get_total_count: Callable[..., Awaitable[int]]
+    """Handler to get the total number of rows in the data variable. Defaults to DataVariable.get_total_count for type=plain, and DerivedDataVariable.get_total_count for type=derived"""
+
+    get_schema: Callable[..., Awaitable[DataFrameSchema]]
+    """Handler to get the schema for data variable. Defaults to DataVariable.get_schema for type=plain, and DerivedDataVariable.get_schema for type=derived"""
+    model_config = ConfigDict(extra='forbid', arbitrary_types_allowed=True)
+
+
+async def upload(data: UploadFile, data_uid: Optional[str] = None, resolver_id: Optional[str] = None):
+    """
+    Handler for uploading data.
+
+    :param data: the file to upload
+    :param data_uid: optional uid of the data variable to upload to
+    :param resolver_id: optional id of the upload resolver to use, falls back to default handlers for csv/xlsx
+    """
+    from dara.core.interactivity.data_variable import DataVariable
+    from dara.core.internal.registries import (
+        data_variable_registry,
+        upload_resolver_registry,
+        utils_registry,
+    )
+
+    store: CacheStore = utils_registry.get('Store')
+    registry_mgr: RegistryLookup = utils_registry.get('RegistryLookup')
+
+    if data.filename is None:
+        raise ValueError('Filename not provided')
+
+    variable = None
+
+    _name, file_type = os.path.splitext(data.filename)
+
+    if data_uid is not None:
+        try:
+            variable = await registry_mgr.get(data_variable_registry, data_uid)
+        except KeyError:
+            raise ValueError(f'Data Variable {data_uid} does not exist')
+
+        if variable.type == 'derived':
+            raise ValueError('Cannot upload data to DerivedDataVariable')
+
+    content = cast(bytes, await data.read())
+
+    resolver = None
+
+    # If Id is provided, lookup the definition from registry
+    if resolver_id is not None:
+        resolver_def: UploadResolverDef = await registry_mgr.get(upload_resolver_registry, resolver_id)
+        resolver = resolver_def.resolver
+
+    if resolver:
+        content = await run_user_handler(handler=resolver, args=(content, data.filename))
+    # If resolver is not provided, follow roughly the cl_dataset_parser logic
+    elif file_type == '.xlsx':
+        file_object_xlsx = io.BytesIO(content)
+        content = pandas.read_excel(file_object_xlsx, index_col=None)
+        content.columns = content.columns.str.replace('Unnamed: *', 'column_', regex=True)   # type: ignore
+    else:
+        # default to csv
+        file_object_csv = io.StringIO(content.decode('utf-8'))
+        content = pandas.read_csv(file_object_csv, index_col=0)
+        content.columns = content.columns.str.replace('Unnamed: *', 'column_', regex=True)   # type: ignore
+
+    # If a data variable is provided, update it with the new content
+    if variable:
+        DataVariable.update_value(variable, store, content)

dara/core/interactivity/any_variable.py

@@ -30,8 +30,9 @@ from fastapi.encoders import jsonable_encoder
 from pydantic import ConfigDict
 
 from dara.core.auth.definitions import SESSION_ID, USER, UserData
-from dara.core.base_definitions import BaseTask, PendingTask
+from dara.core.base_definitions import BaseTask
 from dara.core.base_definitions import DaraBaseModel as BaseModel
+from dara.core.base_definitions import PendingTask
 from dara.core.interactivity.condition import Condition, Operator
 from dara.core.internal.cache_store import CacheStore
 from dara.core.internal.tasks import TaskManager
@@ -111,7 +112,7 @@ async def get_current_value(variable: dict, timeout: float = 3, raw: bool = Fals
         user_identity = None
 
         if current_user is not None:
-            user_identity = current_user.identity_id
+            user_identity = current_user.identity_id or current_user.identity_name
         elif isinstance(auth_config, BasicAuthConfig):
             # basic auth - assume it's the single existing user
             user_identity = list(auth_config.users.keys())[0]
@@ -203,6 +204,7 @@ async def get_current_value(variable: dict, timeout: float = 3, raw: bool = Fals
 
         for session, channels in session_channels.items():
             for ws in channels:
+
                 raw_result = raw_results[ws]
                 # Skip values from clients where the variable is not registered
                 if raw_result == NOT_REGISTERED:
@@ -236,12 +238,12 @@ async def get_current_value(variable: dict, timeout: float = 3, raw: bool = Fals
 
         # If we're returning multiple values, in Jupyter environments print an explainer
         try:
-            from IPython import get_ipython  # pyright: ignore[reportMissingImports]
+            from IPython import get_ipython
         except ImportError:
             pass
         else:
             if get_ipython() is not None:
-                from IPython.display import HTML, display  # pyright: ignore[reportMissingImports]
+                from IPython.display import HTML, display
 
                 display(
                     HTML(
@@ -270,7 +272,7 @@ async def get_current_value(variable: dict, timeout: float = 3, raw: bool = Fals
         return results
 
 
-class AnyVariable(BaseModel, abc.ABC):  # noqa: PLW1641 # we override equals to create conditions, otherwise we should define hash
+class AnyVariable(BaseModel, abc.ABC):
     """
     Base class for all variables. Used for typing to specify that any variable can be provided.
     """
@@ -280,6 +282,7 @@ class AnyVariable(BaseModel, abc.ABC):  # noqa: PLW1641 # we override equals to
     uid: str
 
     def __init__(self, uid: Optional[str] = None, **kwargs) -> None:
+
         new_uid = uid
         if new_uid is None:
             new_uid = str(uuid.uuid4())

dara/core/interactivity/data_variable.py

@@ -17,15 +17,40 @@ limitations under the License.
 
 from __future__ import annotations
 
-from typing import Optional
+import asyncio
+from typing import Optional, Union, cast
 
+from anyio.abc import TaskGroup
 from pandas import DataFrame
+from pandas.io.json._table_schema import build_table_schema
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    SerializerFunctionWrapHandler,
+    model_serializer,
+)
 
-from dara.core.base_definitions import Cache, CacheArgType
-from dara.core.interactivity.server_variable import ServerVariable
+from dara.core.base_definitions import BaseCachePolicy, Cache, CacheArgType
+from dara.core.interactivity.any_data_variable import (
+    AnyDataVariable,
+    DataFrameSchema,
+    DataVariableRegistryEntry,
+)
+from dara.core.interactivity.filtering import (
+    FilterQuery,
+    Pagination,
+    apply_filters,
+    coerce_to_filter_query,
+)
+from dara.core.internal.cache_store import CacheStore
+from dara.core.internal.hashing import hash_object
+from dara.core.internal.pandas_utils import append_index, df_convert_to_internal
+from dara.core.internal.utils import call_async
+from dara.core.internal.websocket import WebsocketManager
+from dara.core.logging import eng_logger
 
 
-class DataVariable(ServerVariable):
+class DataVariable(AnyDataVariable):
     """
     DataVariable represents a variable that is specifically designed to hold datasets.
 
@@ -44,6 +69,11 @@ class DataVariable(ServerVariable):
     >>> UpdateVariable(func=some_data_transformation, variable=session_data)
     """
 
+    uid: str
+    filters: Optional[FilterQuery] = None
+    cache: Optional[BaseCachePolicy] = None
+    model_config = ConfigDict(extra='forbid', arbitrary_types_allowed=True, use_enum_values=True)
+
     def __init__(
         self,
         data: Optional[DataFrame] = None,
@@ -59,11 +89,239 @@ class DataVariable(ServerVariable):
         :param filters: a dictionary of filters to apply to the data
         :param cache: how to cache the result; 'user' per user, 'session' per session, 'global' for all users
         """
-        cache_policy = Cache.Policy.from_arg(cache)
+        cache = Cache.Policy.from_arg(cache)
 
-        if data is not None and cache_policy.cache_type is not Cache.Type.GLOBAL:
+        if data is not None and cache.cache_type is not Cache.Type.GLOBAL:
             raise ValueError('Data cannot be cached per session or per user if provided upfront')
 
-        scope = 'global' if cache_policy.cache_type == Cache.Type.GLOBAL else 'user'
+        super().__init__(cache=cache, uid=uid, **kwargs)
+
+        # Register the variable with the dataset
+        from dara.core.internal.registries import data_variable_registry
+
+        var_entry = DataVariableRegistryEntry(
+            cache=cache,
+            uid=str(self.uid),
+            type='plain',
+            get_data=DataVariable.get_value,
+            get_total_count=DataVariable.get_total_count,
+            get_schema=DataVariable.get_schema,
+        )
+        data_variable_registry.register(
+            str(self.uid),
+            var_entry,
+        )
+
+        # Put the data entry into the store if not empty (so cache='global')
+        # We don't create an entry in a different case since session key will be global anyway at this point
+        if data is not None:
+            from dara.core.internal.registries import utils_registry
+
+            store: CacheStore = utils_registry.get('Store')
+            call_async(self._update, var_entry, store, data)
+
+    @staticmethod
+    def _get_cache_key(uid: str) -> str:
+        """
+        Get a unique cache key for the data variable.
+
+        :param uid: uid of the DataVariable
+        """
+        return f'data-{uid}'
+
+    @staticmethod
+    def _get_schema_cache_key(uid: str) -> str:
+        """
+        Get a unique cache key for the data variable's schema.
+
+        :param uid: uid of the DataVariable
+        """
+        return f'schema-{uid}'
+
+    @classmethod
+    def _get_count_cache_key(cls, uid: str, filters: Optional[Union[FilterQuery, dict]]) -> str:
+        return f'{cls._get_cache_key(uid)}_{hash_object(filters)}'
+
+    @classmethod
+    async def _update(cls, var_entry: DataVariableRegistryEntry, store: CacheStore, data: Optional[DataFrame]):
+        """
+        Internal helper which updates the data variable entry in store.
+
+        TODO: for now data is always kept in store, in the future depending on the size data might be cached on disk
+        """
+        await store.set(var_entry, key=cls._get_cache_key(var_entry.uid), value=DataStoreEntry(data=append_index(data)))
+
+    @classmethod
+    def update_value(cls, var_entry: DataVariableRegistryEntry, store: CacheStore, data: Optional[DataFrame]):
+        """
+        Update the data entry and notify all clients about the update.
+
+        :param var_entry: the entry in the registry
+        :param data: the data to update
+        :param store: store instance
+        """
+        from dara.core.internal.registries import utils_registry
+
+        ws_mgr: WebsocketManager = utils_registry.get('WebsocketManager')
+        task_group: TaskGroup = utils_registry.get('TaskGroup')
+
+        # Update store
+        task_group.start_soon(cls._update, var_entry, store, data)
+
+        # Broadcast the update to all clients
+        task_group.start_soon(
+            ws_mgr.broadcast,
+            {
+                'data_id': str(var_entry.uid),
+            },
+        )
+
+    @classmethod
+    async def get_value(
+        cls,
+        var_entry: DataVariableRegistryEntry,
+        store: CacheStore,
+        filters: Optional[Union[FilterQuery, dict]] = None,
+        pagination: Optional[Pagination] = None,
+        format_for_display: bool = False,
+    ) -> Optional[DataFrame]:
+        """
+        Get the value of this DataVariable.
+        """
+        _uid_short = f'{var_entry.uid[:3]}..{var_entry.uid[-3:]}'
+        eng_logger.info(
+            f'Data Variable {_uid_short} get_value',
+            {'uid': var_entry.uid, 'filters': filters, 'pagination': pagination},
+        )
+
+        cache_key = cls._get_cache_key(var_entry.uid)
+        entry = await store.get(var_entry, key=cache_key)
+
+        eng_logger.debug(
+            f'Data Variable {_uid_short}',
+            'retrieved from cache',
+            {
+                'uid': var_entry.uid,
+                'size': len(entry.data.index) if entry is not None and entry.data is not None else 0,
+            },
+        )
+
+        if entry is None:
+            await asyncio.gather(
+                store.set(var_entry, key=cls._get_count_cache_key(var_entry.uid, filters), value=0, pin=True),
+                store.set(var_entry, key=cls._get_schema_cache_key(var_entry.uid), value=None, pin=True),
+            )
+            return None
+
+        data = None
+
+        if entry.data is not None:
+            filtered_data, count = apply_filters(entry.data, coerce_to_filter_query(filters), pagination)
+            if format_for_display and filtered_data is not None:
+                filtered_data = filtered_data.copy()
+                for col in filtered_data.columns:
+                    if filtered_data[col].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
+                        filtered_data.loc[:, col] = filtered_data[col].apply(str)
+            data = filtered_data
+            # Store count for given filters and schema
+            await asyncio.gather(
+                store.set(var_entry, key=cls._get_count_cache_key(var_entry.uid, filters), value=count, pin=True),
+                store.set(
+                    var_entry,
+                    key=cls._get_schema_cache_key(var_entry.uid),
+                    value=build_table_schema(df_convert_to_internal(entry.data)),
+                    pin=True,
+                ),
+            )
+        else:
+            await asyncio.gather(
+                store.set(var_entry, key=cls._get_count_cache_key(var_entry.uid, filters), value=0, pin=True),
+                store.set(var_entry, key=cls._get_schema_cache_key(var_entry.uid), value=None, pin=True),
+            )
+
+        # TODO: once path is supported, stream&filter from disk
+        if entry.path:
+            raise NotImplementedError('DataVariable.get_value() does not support disk caching yet')
+
+        eng_logger.info(
+            f'Data Variable {_uid_short} returning filtered data',
+            {'uid': var_entry.uid, 'size': len(data.index) if data is not None else 0},
+        )
+
+        return data
+
+    @classmethod
+    async def get_total_count(
+        cls, var_entry: DataVariableRegistryEntry, store: CacheStore, filters: Optional[FilterQuery]
+    ):
+        """
+        Get total count of the data variable.
+
+        :param var_entry: variable entry
+        :param store: store
+        :param filters: filters to get count for
+        """
+        cache_key = cls._get_count_cache_key(var_entry.uid, filters)
+        entry = await store.get(var_entry, key=cache_key, unpin=True)
+
+        if entry is None:
+            raise ValueError('Requested count for filter setup which has not been performed yet')
+
+        return entry
+
+    @classmethod
+    async def get_schema(cls, var_entry: DataVariableRegistryEntry, store: CacheStore):
+        """
+        Get the schema of the data variable.
+
+        :param var_entry: variable entry
+        :param store: store
+        """
+        cache_key = cls._get_schema_cache_key(var_entry.uid)
+        entry = await store.get(var_entry, key=cache_key, unpin=True)
+
+        return cast(DataFrameSchema, entry)
+
+    def reset(self):
+        raise NotImplementedError('DataVariable cannot be reset')
+
+    def update(self, value: Optional[DataFrame]):
+        """
+        Create an action to update the value of this Variable to a provided value.
+
+        ```python
+        import pandas as pd
+        from dara.core import DataVariable
+        from dara.components import Button
+
+        data = DataVariable(pd.DataFrame({'a': [1, 2, 3]}))
+
+        Button(
+            'Empty Data',
+            onclick=data.update(None),
+        )
+
+        ```
+        """
+        from dara.core.interactivity.actions import UpdateVariableImpl
+
+        return UpdateVariableImpl(variable=self, value=value)
+
+    @model_serializer(mode='wrap')
+    def ser_model(self, nxt: SerializerFunctionWrapHandler) -> dict:
+        parent_dict = nxt(self)
+        if 'data' in parent_dict:
+            parent_dict.pop('data')   # make sure data is not included in the serialised dict
+        return {**parent_dict, '__typename': 'DataVariable', 'uid': str(parent_dict['uid'])}
+
+
+class DataStoreEntry(BaseModel):
+    """
+    Entry in the cache store for a DataVariable.
+    Can either be a DataFrame or a path to a file on disk.
+    """
 
-        super().__init__(scope=scope, uid=uid, default=data, **kwargs)
+    data: Optional[DataFrame] = None
+    path: Optional[str] = None
+    model_config = ConfigDict(extra='forbid', arbitrary_types_allowed=True)