splight-lib 5.4.3__tar.gz → 5.5.0.dev1__tar.gz

{splight_lib-5.4.3 → splight_lib-5.5.0.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: splight-lib
-Version: 5.4.3
+Version: 5.5.0.dev1
 Summary: Splight Library
 Author: Splight Dev
 Author-email: dev@splight-ae.com

{splight_lib-5.4.3 → splight_lib-5.5.0.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "splight-lib"
-version = "5.4.3"
+version = "5.5.0.dev1"
 description = "Splight Library"
 authors = ["Splight Dev <dev@splight-ae.com>"]
 readme = "README.md"

{splight_lib-5.4.3 → splight_lib-5.5.0.dev1}/splight_lib/client/datalake/__init__.py

@@ -1,9 +1,9 @@
 from splight_lib.client.datalake.builder import DatalakeClientBuilder
 from splight_lib.client.datalake.local_client import LocalDatalakeClient
-from splight_lib.client.datalake.remote_client import RemoteDatalakeClient
+from splight_lib.client.datalake.remote_client import SyncRemoteDatalakeClient
 
 __all__ = [
     DatalakeClientBuilder,
-    RemoteDatalakeClient,
+    SyncRemoteDatalakeClient,
     LocalDatalakeClient,
 ]

splight_lib-5.5.0.dev1/splight_lib/client/datalake/builder.py

@@ -0,0 +1,26 @@
+from typing import Any, Dict
+
+from splight_lib.client.datalake.abstract import AbstractDatalakeClient
+from splight_lib.client.datalake.local_client import LocalDatalakeClient
+from splight_lib.client.datalake.remote_client import (
+    BufferedAsyncRemoteDatalakeClient,
+    BufferedSyncRemoteDataClient,
+    SyncRemoteDatalakeClient,
+)
+from splight_lib.settings import DatalakeClientType
+
+DL_CLIENT_TYPE_MAP = {
+    DatalakeClientType.BUFFERED_ASYNC: BufferedAsyncRemoteDatalakeClient,
+    DatalakeClientType.BUFFERED_SYNC: BufferedSyncRemoteDataClient,
+    DatalakeClientType.SYNC: SyncRemoteDatalakeClient,
+    DatalakeClientType.LOCAL: LocalDatalakeClient,
+}
+
+
+class DatalakeClientBuilder:
+    @staticmethod
+    def build(
+        dl_client_type: DatalakeClientType = DatalakeClientType.BUFFERED_ASYNC,
+        parameters: Dict[str, Any] = {},
+    ) -> AbstractDatalakeClient:
+        return DL_CLIENT_TYPE_MAP[dl_client_type](**parameters)

{splight_lib-5.4.3 → splight_lib-5.5.0.dev1}/splight_lib/client/datalake/remote_client.py

@@ -1,6 +1,7 @@
+from datetime import datetime
 from threading import Lock, Thread
 from time import sleep
-from typing import Dict, List, Union
+from typing import Dict, List, Optional, Union
 
 import pandas as pd
 from furl import furl
@@ -25,7 +26,7 @@ from splight_lib.stringcase import camelcase
 logger = get_splight_logger()
 
 
-class RemoteDatalakeClient(AbstractDatalakeClient):
+class SyncRemoteDatalakeClient(AbstractDatalakeClient):
     _PREFIX = "/data"
 
     def __init__(
@@ -85,12 +86,18 @@ class RemoteDatalakeClient(AbstractDatalakeClient):
         collection: str = DEFAULT_COLLECTION,
         sort_field: str = DEFAULT_SORT_FIELD,
         sort_direction: int = -1,
+        from_timestamp: Optional[datetime] = None,
+        to_timestamp: Optional[datetime] = None,
+        extra_pipeline: List[Dict] = [],
+        aggregation_query: Dict = {},
         limit: int = 10000,
-        **filters,
     ) -> List[Dict]:
         # POST /data/read
         url = self._base_url / f"{self._PREFIX}/read"
+        pipeline = [{"$match": match}, *extra_pipeline]
         data_request = DataRequest(
+            from_timestmap=from_timestamp,
+            to_timestamp=to_timestamp,
             collection=collection,
             sort_field=sort_field,
             sort_direction=sort_direction,
@@ -99,11 +106,11 @@ class RemoteDatalakeClient(AbstractDatalakeClient):
                 {
                     "ref_id": "output",
                     "type": "QUERY",
-                    "pipeline": [{"$match": match}],
+                    "pipeline": pipeline,
                     "expression": None,
                 }
             ],
-            **filters,
+            aggregation_query=aggregation_query,
         )
         response = self._restclient.post(url, json=data_request.dict())
         response.raise_for_status()
@@ -116,12 +123,18 @@ class RemoteDatalakeClient(AbstractDatalakeClient):
         collection: str = DEFAULT_COLLECTION,
         sort_field: str = DEFAULT_SORT_FIELD,
         sort_direction: int = -1,
+        from_timestamp: Optional[datetime] = None,
+        to_timestamp: Optional[datetime] = None,
+        extra_pipeline: List[Dict] = [],
+        aggregation_query: Dict = {},
         limit: int = 10000,
-        **filters,
     ) -> List[Dict]:
         # POST /data/read
         url = self._base_url / f"{self._PREFIX}/read"
+        pipeline = [{"$match": match}, *extra_pipeline]
         data_request = DataRequest(
+            from_timestmap=from_timestamp,
+            to_timestamp=to_timestamp,
             collection=collection,
             sort_field=sort_field,
             sort_direction=sort_direction,
@@ -130,15 +143,11 @@ class RemoteDatalakeClient(AbstractDatalakeClient):
                 {
                     "ref_id": "output",
                     "type": "QUERY",
-                    "pipeline": [
-                        {
-                            "$match": match,
-                        }
-                    ],
+                    "pipeline": pipeline,
                     "expression": None,
                 }
             ],
-            **filters,
+            aggregation_query=aggregation_query,
         )
         response = await self._restclient.async_post(
             url, json=data_request.dict()
@@ -161,11 +170,23 @@ class RemoteDatalakeClient(AbstractDatalakeClient):
         collection: str = DEFAULT_COLLECTION,
         sort_field: str = DEFAULT_SORT_FIELD,
         sort_direction: int = -1,
+        from_timestamp: Optional[datetime] = None,
+        to_timestamp: Optional[datetime] = None,
+        extra_pipeline: List[Dict] = [],
+        aggregation_query: Dict = {},
         limit: int = 10000,
-        **filters,
+        **kwargs,
     ) -> pd.DataFrame:
         data = self._get(
-            match, collection, sort_field, sort_direction, limit, **filters
+            match,
+            collection,
+            sort_field,
+            sort_direction,
+            from_timestamp,
+            to_timestamp,
+            extra_pipeline,
+            aggregation_query,
+            limit,
         )
         df = pd.DataFrame(data)
         if not df.empty:
@@ -177,7 +198,7 @@ class RemoteDatalakeClient(AbstractDatalakeClient):
         return df
 
 
-class BufferedRemoteDatalakeClient(RemoteDatalakeClient):
+class BufferedAsyncRemoteDatalakeClient(SyncRemoteDatalakeClient):
     _PREFIX = "data"
 
     def __init__(
@@ -218,9 +239,9 @@ class BufferedRemoteDatalakeClient(RemoteDatalakeClient):
                 buffer_size, buffer_timeout
             ),
         }
+        self._lock = Lock()
         self._flush_thread = Thread(target=self._flusher, daemon=True)
         self._flush_thread.start()
-        self._lock = Lock()
         logger.debug(
             "Buffered Remote datalake client initialized.",
             tags=LogTags.DATALAKE,
@@ -256,22 +277,111 @@ class BufferedRemoteDatalakeClient(RemoteDatalakeClient):
     def _flush_buffer(
         self, collection: str, buffer: DatalakeDocumentBuffer
     ) -> None:
-        if buffer.should_flush():
-            try:
+        with self._lock:
+            if buffer.should_flush():
+                try:
+                    logger.debug(
+                        "Flushing datalake buffer with %s elements",
+                        len(buffer.data),
+                    )
+                    self._send_documents(collection, buffer.data)
+                    buffer.reset()
+                except Exception:
+                    logger.error("Unable to save documents", exc_info=True)
+
+    @retry(SPLIGHT_REQUEST_EXCEPTIONS, tries=3, delay=2, jitter=1)
+    def _send_documents(self, collection: str, docs: List[Dict]) -> List[Dict]:
+        url = self._base_url / f"{self._PREFIX}/write"
+        collection = camelcase(collection)
+        data = {
+            "collection": collection,
+            "records": docs,
+        }
+        response = self._restclient.post(url, json=data)
+        response.raise_for_status()
+        return docs
+
+
+class BufferedSyncRemoteDataClient(SyncRemoteDatalakeClient):
+    _PREFIX = "data"
+
+    def __init__(
+        self,
+        base_url: str,
+        access_id: str,
+        secret_key: str,
+        buffer_size: int = DL_BUFFER_SIZE,
+        buffer_timeout: float = DL_BUFFER_TIMEOUT,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(
+            base_url,
+            access_id,
+            secret_key,
+            buffer_size,
+            buffer_timeout,
+            *args,
+            **kwargs,
+        )
+        self._base_url = furl(base_url)
+        token = SplightAuthToken(
+            access_key=access_id,
+            secret_key=secret_key,
+        )
+        self._restclient = SplightRestClient()
+        self._restclient.update_headers(token.header)
+
+        logger.debug(
+            "Initializing buffer with size %s and timeout %s",
+            buffer_size,
+            buffer_timeout,
+        )
+        self._data_buffers = {
+            "default": DatalakeDocumentBuffer(buffer_size, buffer_timeout),
+            "routine_evaluations": DatalakeDocumentBuffer(
+                buffer_size, buffer_timeout
+            ),
+        }
+        self._lock = Lock()
+        logger.debug(
+            "Synchronous Buffered Remote datalake client initialized.",
+            tags=LogTags.DATALAKE,
+        )
+
+    def save(
+        self, collection: str, instances: Union[List[Dict], Dict]
+    ) -> List[Dict]:
+        logger.debug("Saving documents in datalake")
+        instances = instances if isinstance(instances, List) else [instances]
+        if collection not in self._data_buffers:
+            raise InvalidCollectionName(collection)
+        buffer = self._data_buffers[collection]
+        with self._lock:
+            if buffer.should_flush():
                 logger.debug(
                     "Flushing datalake buffer with %s elements",
                     len(buffer.data),
                 )
                 self._send_documents(collection, buffer.data)
                 buffer.reset()
-            except Exception:
-                logger.error("Unable to save documents", exc_info=True)
+            buffer.add_documents(instances)
+        return instances
+
+    def save_dataframe(
+        self, dataframe: pd.DataFrame, collection: str = DEFAULT_COLLECTION
+    ) -> None:
+        logger.debug("Saving dataframe in datalake")
+        dataframe["timestamp"] = dataframe["timestamp"].apply(
+            lambda x: x.tz_localize(tz="UTC").isoformat()
+        )
+        instances = dataframe.to_dict("records")
+        self.save(collection, instances)
 
     @retry(SPLIGHT_REQUEST_EXCEPTIONS, tries=3, delay=2, jitter=1)
     def _send_documents(self, collection: str, docs: List[Dict]) -> List[Dict]:
         url = self._base_url / f"{self._PREFIX}/write"
         collection = camelcase(collection)
-        collection = camelcase(collection)
         data = {
             "collection": collection,
             "records": docs,

{splight_lib-5.4.3 → splight_lib-5.5.0.dev1}/splight_lib/client/tests/test_datalake.py

@@ -2,7 +2,7 @@ import os
 
 from pytest_mock import MockerFixture
 
-from splight_lib.client.datalake import RemoteDatalakeClient  # noqa E402
+from splight_lib.client.datalake import SyncRemoteDatalakeClient  # noqa E402
 from splight_lib.client.datalake.remote_client import (  # noqa E402
     SplightRestClient,
 )
@@ -34,7 +34,7 @@ def test_initialization(mocker: MockerFixture):
         "update_headers",
         return_value=None,
     )
-    _ = RemoteDatalakeClient(
+    _ = SyncRemoteDatalakeClient(
         base_url=base_url,
         access_id=access_id,
         secret_key=secret_key,
@@ -49,7 +49,7 @@ def test_save(mocker: MockerFixture):
     secret_key = os.getenv("SECRET_KEY")
     access_id = os.getenv("ACCESS_ID")
 
-    client = RemoteDatalakeClient(
+    client = SyncRemoteDatalakeClient(
         base_url=base_url,
         access_id=access_id,
         secret_key=secret_key,

{splight_lib-5.4.3 → splight_lib-5.5.0.dev1}/splight_lib/models/base.py

@@ -1,9 +1,9 @@
 import json
 from datetime import datetime, timezone
-from typing import ClassVar, Dict, List, Optional
+from typing import ClassVar, Dict, List
 
 import pandas as pd
-from pydantic import BaseModel, Field, PrivateAttr, field_validator
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
 
 from splight_lib.client.database import DatabaseClientBuilder
 from splight_lib.client.database.abstract import AbstractDatabaseClient
@@ -75,14 +75,10 @@ class SplightDatalakeBaseModel(BaseModel):
     timestamp: datetime = Field(
         default_factory=lambda: datetime.now(timezone.utc)
     )
-    instance_id: str = ""
-    instance_type: str = ""
     _collection_name: ClassVar[str] = "DatalakeModel"
     _dl_client: AbstractDatalakeClient = PrivateAttr()
 
-    @field_validator("instance_id", "instance_type", mode="before")
-    def convert_to_string(cls, value: Optional[str]):
-        return value if value else ""
+    model_config = ConfigDict(extra="ignore")
 
     @classmethod
     def get(
@@ -174,8 +170,7 @@ class SplightDatalakeBaseModel(BaseModel):
     @staticmethod
     def __get_datalake_client() -> AbstractDatalakeClient:
         db_client = DatalakeClientBuilder.build(
-            local=settings.LOCAL_ENVIRONMENT,
-            use_buffer=settings.USE_BUFFER,
+            dl_client_type=settings.DL_CLIENT_TYPE,
             parameters={
                 "path": settings.CURRENT_DIR,
                 "base_url": settings.SPLIGHT_PLATFORM_API_HOST,

{splight_lib-5.4.3 → splight_lib-5.5.0.dev1}/splight_lib/models/native.py

@@ -23,21 +23,24 @@ class NativeOutput(SplightDatalakeBaseModel):
     def get(
         cls, asset: str, attribute: str, **params: Dict
     ) -> List["NativeOutput"]:
-        params["output_format"] = cls._output_format
         return super().get(asset, attribute, **params)
 
     @classmethod
-    async def async_get(cls, **params: Dict) -> List["NativeOutput"]:
-        params["output_format"] = cls._output_format
-        return await super().async_get(**params)
+    async def async_get(
+        cls, asset: str, attribute: str, **params: Dict
+    ) -> List["NativeOutput"]:
+        return await super().async_get(asset, attribute, **params)
 
     @classmethod
     def get_dataframe(cls, asset: str, attribute: str, **params: Dict):
-        params["output_format"] = cls._output_format
-        return super().get_dataframe(asset, attribute, **params)
+        df = super().get_dataframe(asset, attribute, **params)
+        df["output_format"] = cls._output_format
+        return df
 
     @classmethod
-    def latest(cls, asset: str, attribute: str, expiration: timedelta = None):
+    def latest(
+        cls, asset: str, attribute: str, expiration: Optional[timedelta] = None
+    ):
         from_timestamp = None
         if expiration:
             from_timestamp = datetime.now(timezone.utc) - expiration

{splight_lib-5.4.3 → splight_lib-5.5.0.dev1}/splight_lib/models/pipeline.py

@@ -12,7 +12,7 @@ from splight_lib.stringcase import camelcase
 
 def get_datalake_client() -> AbstractDatalakeClient:
     client = DatalakeClientBuilder.build(
-        local=settings.LOCAL_ENVIRONMENT,
+        dl_client_type=settings.DL_CLIENT_TYPE,
         parameters={
             "path": settings.CURRENT_DIR,
             "base_url": settings.SPLIGHT_PLATFORM_API_HOST,

{splight_lib-5.4.3 → splight_lib-5.5.0.dev1}/splight_lib/settings.py

@@ -1,4 +1,5 @@
 import os
+from enum import Enum
 from typing import Any, Dict, Optional, Tuple, Type
 
 import yaml
@@ -11,6 +12,13 @@ from splight_lib.constants import DL_BUFFER_SIZE, DL_BUFFER_TIMEOUT
 SPLIGHT_HOME = os.path.join(os.path.expanduser("~"), ".splight")
 
 
+class DatalakeClientType(str, Enum):
+    LOCAL = "local"
+    SYNC = "sync"
+    BUFFERED_SYNC = "buffered_sync"
+    BUFFERED_ASYNC = "buffered_async"
+
+
 class Singleton:
     def __new__(cls, *args, **kw):
         if not hasattr(cls, "_instance"):
@@ -58,12 +66,13 @@ class SplightSettings(BaseSettings, Singleton):
     SPLIGHT_PLATFORM_API_HOST: str = "https://api.splight-ai.com"
 
     # Parameters for local environment
+    # TODO: to deprecate this and its effects
     LOCAL_ENVIRONMENT: bool = False
     CURRENT_DIR: Optional[str] = None
 
-    # Parameters for Buffered Datalake Client
+    # Parameters for the datalake client
     # Review if is better to use another class for only the DL Client settings
-    USE_BUFFER: bool = True
+    DL_CLIENT_TYPE: DatalakeClientType = DatalakeClientType.BUFFERED_ASYNC
     DL_BUFFER_SIZE: int = DL_BUFFER_SIZE
     DL_BUFFER_TIMEOUT: float = DL_BUFFER_TIMEOUT  # seconds
 

splight_lib-5.4.3/splight_lib/client/datalake/builder.py

@@ -1,24 +0,0 @@
-from typing import Any, Dict
-
-from splight_lib.client.datalake.abstract import AbstractDatalakeClient
-from splight_lib.client.datalake.local_client import LocalDatalakeClient
-from splight_lib.client.datalake.remote_client import (
-    BufferedRemoteDatalakeClient,
-    RemoteDatalakeClient,
-)
-
-
-class DatalakeClientBuilder:
-    @staticmethod
-    def build(
-        local: bool = False,
-        use_buffer: bool = True,
-        parameters: Dict[str, Any] = {},
-    ) -> AbstractDatalakeClient:
-        if local:
-            dl_client = LocalDatalakeClient(**parameters)
-        elif use_buffer:
-            dl_client = BufferedRemoteDatalakeClient(**parameters)
-        else:
-            dl_client = RemoteDatalakeClient(**parameters)
-        return dl_client