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

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

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

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

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

{splight_lib-5.4.4 → 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.4 → splight_lib-5.5.0.dev1}/splight_lib/client/datalake/remote_client.py

@@ -26,7 +26,7 @@ from splight_lib.stringcase import camelcase
 logger = get_splight_logger()
 
 
-class RemoteDatalakeClient(AbstractDatalakeClient):
+class SyncRemoteDatalakeClient(AbstractDatalakeClient):
     _PREFIX = "/data"
 
     def __init__(
@@ -198,7 +198,7 @@ class RemoteDatalakeClient(AbstractDatalakeClient):
         return df
 
 
-class BufferedRemoteDatalakeClient(RemoteDatalakeClient):
+class BufferedAsyncRemoteDatalakeClient(SyncRemoteDatalakeClient):
     _PREFIX = "data"
 
     def __init__(
@@ -239,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,
@@ -277,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.4 → 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.4 → splight_lib-5.5.0.dev1}/splight_lib/models/base.py

@@ -170,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.4 → 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.4 → 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.4/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