bizon 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

bizon/{sources → connectors/sources}/periscope/src/source.py

@@ -15,8 +15,6 @@ from bizon.source.source import AbstractSource
 
 BASE_URL = "https://app.periscopedata.com"
 
-URL_OWNERS = f"{BASE_URL}/users/owners"
-URL_VIEWS = f"{BASE_URL}/login_state/sql_views"
 URL_DATABASES = f"{BASE_URL}/welcome/remaining_state/site_models"
 
 
@@ -39,6 +37,7 @@ class PeriscopeSourceConfig(SourceConfig):
     workspace_name: str = Field(..., description="Name of the workspace")
     client_site_id: int = Field(..., description="Client site ID")
     database_id: int = Field(..., description="Fetch charts connected to this Database ID")
+    x_csrf_token: str = Field(..., description="CSRF token for the requests")
 
 
 class PeriscopeSource(AbstractSource):
@@ -49,7 +48,14 @@ class PeriscopeSource(AbstractSource):
 
     @staticmethod
     def streams() -> List[str]:
-        return ["dashboards", "charts", "users", "databases", "views"]
+        return [
+            "charts",
+            "dashboards_metadata",
+            "dashboards",
+            "databases",
+            "users",
+            "views",
+        ]
 
     @staticmethod
     def get_config_class() -> AbstractSource:
@@ -86,8 +92,9 @@ class PeriscopeSource(AbstractSource):
                         "sec-fetch-dest": "empty",
                         "sec-fetch-mode": "cors",
                         "sec-fetch-site": "same-origin",
-                        "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/118.0.0.0 Safari/537.36",  # noqa
+                        "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/118.0.0.0 Safari/537.36",
                         "x-requested-with": "XMLHttpRequest",
+                        "x-csrf-token": self.config.x_csrf_token,
                     },
                 )
             )
@@ -97,10 +104,8 @@ class PeriscopeSource(AbstractSource):
     def get_total_records_count(self) -> int | None:
         return None
 
-    def get_dashboards(self, pagination: dict = None) -> SourceIteration:
-        response = self.session.call(method="GET", url=f"{BASE_URL}/login_state/dashboards", params=self.http_params)
-        records_json = response.json()["Dashboard"]
-
+    @staticmethod
+    def transform_response_to_source_iteration(records: List[dict]) -> SourceIteration:
         return SourceIteration(
             next_pagination=dict(),
             records=[
@@ -108,14 +113,69 @@ class PeriscopeSource(AbstractSource):
                     id=record["id"],
                     data=record,
                 )
-                for record in records_json
+                for record in records
             ],
         )
 
+    def get_dashboards(self, pagination: dict = None) -> SourceIteration:
+        response = self.session.call(
+            method="GET",
+            url=f"{BASE_URL}/login_state/dashboards",
+            params=self.http_params,
+        )
+        records_json = response.json()["Dashboard"]
+        return self.transform_response_to_source_iteration(records_json)
+
+    def get_dashboards_metadata(self, pagination: dict = None) -> SourceIteration:
+
+        params = {
+            "client_site_id": self.config.client_site_id,
+            "filters": [{"name": "typeFilter", "input": "Dashboard"}],
+            "limit": 2000,
+            "query_plan": None,
+        }
+
+        response = self.session.call(
+            method="POST",
+            url=f"{BASE_URL}/global_search/search",
+            json=params,
+        )
+        records_json = response.json()["results"]["data"]
+        return self.transform_response_to_source_iteration(records_json)
+
     def get_dashboard_ids(self) -> List[int]:
         source_iteration = self.get_dashboards()
         return [record.id for record in source_iteration.records]
 
+    def _extract_raw_text_from_textbox(self, data: dict) -> str:
+        raw_text = []
+
+        def clean_text(text: str):
+            """Strip Byte Order Mark (BOM) and other unwanted whitespace."""
+            return text.replace("\ufeff", "").strip()
+
+        def traverse_nodes(nodes):
+            for node in nodes:
+                if node["object"] == "text":
+                    for leaf in node["leaves"]:
+                        raw_text.append(clean_text(leaf["text"]))
+                elif node["type"] == "link" and "data" in node and "url" in node["data"]:
+                    link_text = []
+                    for leaf in node["nodes"][0]["leaves"]:  # Assume a single text node in link
+                        link_text.append(clean_text(leaf["text"]))
+                    # Format as Markdown link
+                    raw_text.append(f"[{''.join(link_text)}]({node['data']['url']})")
+                elif "nodes" in node:  # If there are nested nodes
+                    traverse_nodes(node["nodes"])
+
+        if not data["text_data"]:
+            return ""
+
+        # Start traversal from the root nodes
+        traverse_nodes(data["text_data"]["document"]["nodes"])
+
+        return " ".join(raw_text)
+
     def _get_charts(self, dashboard_id: int) -> List[dict]:
         MAXIMUM_ITERATION: int = 1000
         iter_count: int = 0
@@ -154,12 +214,37 @@ class PeriscopeSource(AbstractSource):
                 iter_count += 1
                 iter_charts = response.json().get("Widget")
 
+                iter_textboxes = response.json().get("TextBox")
+
                 for chart in iter_charts:
+                    # Only fetch charts connected to gorgias-growth-production
                     if str(chart.get("database_id")) == str(self.config.database_id):
                         if chart.get("id") not in charts_list:
+
                             charts_list.add(chart.get("id"))
+
+                            chart["raw_text"] = None
+
+                            # In case the chart is a textbox, we parse the raw text
+                            if chart.get("content_id"):
+                                text_box = list(
+                                    filter(
+                                        lambda x: x.get("id") == chart.get("content_id"),
+                                        iter_textboxes,
+                                    )
+                                )
+
+                                if not text_box:
+                                    logger.error(
+                                        f"Failed to fetch the textbox with id: {chart.get('content_id')} for chart with id: {chart.get('id')}"
+                                    )
+
+                                if text_box:
+                                    chart["raw_text"] = self._extract_raw_text_from_textbox(text_box[0])
+
                             dashboard_charts.append(chart)
-            except Exception:
+            except Exception as e:
+                logger.error(f"Failed to fetch the dashboard with id: {dashboard_id} with error: {e}")
                 continue
 
         return dashboard_charts
@@ -200,11 +285,50 @@ class PeriscopeSource(AbstractSource):
             ],
         )
 
+    def get_views(self, pagination: dict = None) -> SourceIteration:
+        response = self.session.call(
+            method="GET",
+            url=f"{BASE_URL}/login_state/sql_views",
+            params=self.http_params,
+        )
+        records_json = response.json()["SqlView"]
+        return self.transform_response_to_source_iteration(records_json)
+
+    def get_users(self, pagination: dict = None) -> SourceIteration:
+        response = self.session.call(
+            method="GET",
+            url=f"{BASE_URL}/users/owners",
+            params=self.http_params,
+        )
+        records_json = response.json()
+        return self.transform_response_to_source_iteration(records_json)
+
+    def get_databases(self, pagination: dict = None) -> SourceIteration:
+        response = self.session.call(
+            method="GET",
+            url=URL_DATABASES,
+            params=self.http_params,
+        )
+        records_json = response.json()["Database"]
+        return self.transform_response_to_source_iteration(records_json)
+
     def get(self, pagination: dict = None) -> SourceIteration:
-        if self.config.stream_name == "dashboards":
+        if self.config.stream == "dashboards":
             return self.get_dashboards(pagination)
 
-        if self.config.stream_name == "charts":
+        elif self.config.stream == "charts":
             return self.get_charts(pagination)
 
-        raise NotImplementedError(f"Stream {self.config.stream_name} not implemented for Periscope")
+        elif self.config.stream == "dashboards_metadata":
+            return self.get_dashboards_metadata(pagination)
+
+        elif self.config.stream == "views":
+            return self.get_views(pagination)
+
+        elif self.config.stream == "users":
+            return self.get_users(pagination)
+
+        elif self.config.stream == "databases":
+            return self.get_databases(pagination)
+
+        raise NotImplementedError(f"Stream {self.config.stream} not implemented for Periscope")

bizon/connectors/sources/periscope/tests/periscope_pipeline_dashboard.py

@@ -0,0 +1,9 @@
+import os
+
+from bizon.cli.utils import parse_from_yaml
+from bizon.engine.engine import RunnerFactory
+
+config = parse_from_yaml(os.path.abspath("bizon/connectors/sources/periscope/config/periscope_dashboards.yml"))
+
+runner = RunnerFactory.create_from_config_dict(config=config)
+runner.run()

bizon/connectors/sources/pokeapi/config/pokeapi_pokemon_to_json.example.yml

@@ -0,0 +1,19 @@
+name: pokemon to json unnested
+
+source:
+  name: pokeapi
+  stream: pokemon
+
+destination:
+  name: file
+  config:
+    filepath: pokemon.json
+    buffer_flush_timeout: 2
+
+    unnest: true
+
+    record_schema:
+      - name: name
+        type: string
+      - name: url
+        type: string

bizon/connectors/sources/pokeapi/config/pokeapi_pokemon_to_logger.example.yml

@@ -0,0 +1,10 @@
+name: pokemon to logger
+
+source:
+  name: pokeapi
+  stream: pokemon
+
+destination:
+  name: logger
+  config:
+    buffer_flush_timeout: 2

bizon/connectors/sources/pokeapi/src/source.py

@@ -0,0 +1,79 @@
+from enum import Enum
+from typing import Any, List, Tuple
+
+from requests.auth import AuthBase
+
+from bizon.source.config import SourceConfig
+from bizon.source.models import SourceIteration, SourceRecord
+from bizon.source.source import AbstractSource
+
+BASE_URL = "https://pokeapi.co/api/v2"
+
+
+# Define the streams that the source supports
+class PokeAPIStreams(str, Enum):
+    POKEMON = "pokemon"
+    BERRY = "berry"
+    ITEM = "item"
+
+
+# Define the config class for the source
+class PokeAPISourceConfig(SourceConfig):
+    stream: PokeAPIStreams
+
+
+class PeriscopeSource(AbstractSource):
+
+    def __init__(self, config: PokeAPISourceConfig):
+        super().__init__(config)
+        self.config: PokeAPISourceConfig = config
+
+    @property
+    def url_entity(self) -> str:
+        return f"{BASE_URL}/{self.config.stream}"
+
+    @staticmethod
+    def streams() -> List[str]:
+        return [item.value for item in PokeAPIStreams]
+
+    @staticmethod
+    def get_config_class() -> AbstractSource:
+        return PokeAPISourceConfig
+
+    def check_connection(self) -> Tuple[bool | Any | None]:
+        # Make a request to the base URL to check if the connection is successful
+        _ = self.session.get(self.url_entity)
+        return True, None
+
+    def get_authenticator(self) -> AuthBase:
+        # We return None because we don't need any authentication
+        return None
+
+    def get_total_records_count(self) -> int | None:
+        # Return the total number of records in the stream
+        response = self.session.get(self.url_entity)
+        return response.json().get("count")
+
+    def get_entity_list(self, pagination: dict = None) -> SourceIteration:
+        # If pagination is provided, use the next URL to get the next set of records
+        url = pagination.get("next") if pagination else self.url_entity
+        response = self.session.get(url)
+
+        data = response.json()
+
+        return SourceIteration(
+            next_pagination={"next": data.get("next")} if data.get("next") else {},
+            records=[
+                SourceRecord(
+                    id=record["name"],
+                    data=record,
+                )
+                for record in data["results"]
+            ],
+        )
+
+    def get(self, pagination: dict = None) -> SourceIteration:
+        if self.config.stream in [PokeAPIStreams.POKEMON, PokeAPIStreams.BERRY, PokeAPIStreams.ITEM]:
+            return self.get_entity_list(pagination)
+
+        raise NotImplementedError(f"Stream {self.config.stream} not implemented for PokeAPI source")

bizon/{destinations → destination}/buffer.py

@@ -18,6 +18,11 @@ class DestinationBuffer:
         self.pagination = {}
         self.modified_at: List[datetime] = [datetime.now(tz=UTC)]
 
+    @property
+    def is_empty(self) -> bool:
+        """Check if buffer is empty"""
+        return self.df_destination_records.height == 0
+
     @property
     def current_size(self) -> int:
         """Return buffer size"""

bizon/destination/config.py

@@ -0,0 +1,83 @@
+from abc import ABC
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+
+class DestinationTypes(str, Enum):
+    BIGQUERY = "bigquery"
+    BIGQUERY_STREAMING = "bigquery_streaming"
+    BIGQUERY_STREAMING_V2 = "bigquery_streaming_v2"
+    LOGGER = "logger"
+    FILE = "file"
+
+
+class DestinationColumn(BaseModel, ABC):
+    name: str = Field(..., description="Name of the column")
+    type: str = Field(..., description="Type of the column")
+    description: Optional[str] = Field(None, description="Description of the column")
+
+
+class RecordSchemaConfig(BaseModel):
+    # Forbid extra keys in the model
+    model_config = ConfigDict(extra="forbid")
+
+    destination_id: str = Field(..., description="Destination ID")
+    record_schema: list[DestinationColumn] = Field(..., description="Record schema")
+
+
+class AbstractDestinationDetailsConfig(BaseModel):
+
+    # Forbid extra keys in the model
+    model_config = ConfigDict(extra="forbid")
+
+    buffer_size: int = Field(
+        default=50,
+        description="Buffer size in Mb for the destination. Set to 0 to disable and write directly to the destination.",
+    )
+
+    buffer_flush_timeout: int = Field(
+        default=600,
+        description="Maximum time in seconds for buffering after which the records will be written to the destination. Set to 0 to deactivate the timeout buffer check.",  # noqa
+    )
+
+    max_concurrent_threads: int = Field(
+        default=10,
+        description="Maximum number of concurrent threads to use for writing to the destination.",
+    )
+
+    record_schemas: Optional[list[RecordSchemaConfig]] = Field(
+        default=None, description="Schemas for the records. Required if unnest is set to true."
+    )
+
+    unnest: bool = Field(
+        default=False,
+        description="Unnest the data before writing to the destination. Schema should be provided in the model_config.",
+    )
+
+    authentication: Optional[BaseModel] = Field(
+        description="Authentication configuration for the destination, if needed", default=None
+    )
+
+    destination_id: Optional[str] = Field(
+        description="Destination ID, identifier to use to store the records in the destination", default=None
+    )
+
+    @field_validator("unnest", mode="before")
+    def validate_record_schema_if_unnest(cls, value, values):
+        if bool(value) and not values.data.get("record_schemas", []):
+            raise ValueError("At least one `record_schemas` must be provided if `unnest` is set to True.")
+        return value
+
+
+class AbstractDestinationConfig(BaseModel):
+    # Forbid extra keys in the model
+    model_config = ConfigDict(extra="forbid")
+
+    name: DestinationTypes = Field(..., description="Name of the destination")
+    alias: str = Field(
+        ...,
+        description="Alias of the destination, used for tracking the system name (ie bigquery for bigquery_streaming)",
+    )
+    config: AbstractDestinationDetailsConfig = Field(..., description="Configuration for the destination")

bizon/{destinations → destination}/destination.py

@@ -10,6 +10,8 @@ from pydantic import BaseModel, Field
 from bizon.common.models import SyncMetadata
 from bizon.engine.backend.backend import AbstractBackend
 from bizon.engine.backend.models import JobStatus
+from bizon.monitoring.monitor import AbstractMonitor
+from bizon.source.callback import AbstractSourceCallback
 from bizon.source.config import SourceSyncModes
 
 from .buffer import DestinationBuffer
@@ -43,13 +45,42 @@ class DestinationIteration(BaseModel):
 
 class AbstractDestination(ABC):
 
-    def __init__(self, sync_metadata: SyncMetadata, config: AbstractDestinationDetailsConfig, backend: AbstractBackend):
+    def __init__(
+        self,
+        sync_metadata: SyncMetadata,
+        config: AbstractDestinationDetailsConfig,
+        backend: AbstractBackend,
+        source_callback: AbstractSourceCallback,
+        monitor: AbstractMonitor,
+    ):
         self.sync_metadata = sync_metadata
         self.config = config
         self.backend = backend
+        self.monitor = monitor
         self.buffer = DestinationBuffer(
             buffer_size=self.config.buffer_size, buffer_flush_timeout=self.config.buffer_flush_timeout
         )
+        self.source_callback = source_callback
+        self.destination_id = config.destination_id
+
+        self._record_schemas = None
+        self._clustering_keys = None
+
+    @property
+    def record_schemas(self):
+        if self._record_schemas is None and self.config.record_schemas:
+            self._record_schemas = {
+                schema.destination_id: schema.record_schema for schema in self.config.record_schemas
+            }
+        return self._record_schemas
+
+    @property
+    def clustering_keys(self):
+        if self._clustering_keys is None and self.config.record_schemas:
+            self._clustering_keys = {
+                schema.destination_id: schema.clustering_keys for schema in self.config.record_schemas
+            }
+        return self._clustering_keys
 
     @abstractmethod
     def check_connection(self) -> bool:
@@ -75,7 +106,7 @@ class AbstractDestination(ABC):
         )
 
         logger.info(
-            f"Writing in destination from source iteration {self.buffer.from_iteration} to {self.buffer.to_iteration}"
+            f"Writing in destination {self.destination_id} from source iteration {self.buffer.from_iteration} to {self.buffer.to_iteration}"
         )
 
         success, error_msg = self.write_records(df_destination_records=self.buffer.df_destination_records)
@@ -83,7 +114,9 @@ class AbstractDestination(ABC):
         if success:
             # We wrote records to destination so we keep it
             destination_iteration.records_written = self.buffer.df_destination_records.height
-            logger.info(f"Successfully wrote {destination_iteration.records_written} records to destination")
+            logger.info(
+                f"Successfully wrote {destination_iteration.records_written} records to destination {self.destination_id}"
+            )
 
         else:
             # We failed to write records to destination so we keep the error message
@@ -112,8 +145,8 @@ class AbstractDestination(ABC):
         # Last iteration, write all records to destination
         if last_iteration:
 
-            if self.buffer.df_destination_records.height == 0 and self.config.buffer_size == 0:
-                logger.warning("No records to write to destination, already written, buffer is empty.")
+            if self.buffer.df_destination_records.height == 0 and self.buffer.is_empty:
+                logger.info("No records to write to destination, already written, buffer is empty.")
                 return DestinationBufferStatus.RECORDS_WRITTEN
 
             logger.debug("Writing last iteration records to destination")
@@ -143,12 +176,12 @@ class AbstractDestination(ABC):
 
         # Don't write empty records to destination
         if df_destination_records.height == 0 and not last_iteration:
-            logger.warning("No records to write to destination. Check source and queue provider.")
+            logger.info("No records to write to destination. Check source and queue provider.")
             return DestinationBufferStatus.NO_RECORDS
 
         # Write records to destination if buffer size is 0 or streaming
         if self.buffer.buffer_size == 0:
-            logger.info("Writing records to destination.")
+            logger.info(f"Writing records to destination {self.destination_id}.")
             self.buffer.add_source_iteration_records_to_buffer(
                 iteration=iteration, df_destination_records=df_destination_records, pagination=pagination
             )
@@ -161,6 +194,14 @@ class AbstractDestination(ABC):
         logger.info(
             f"Buffer ripeness {round(self.buffer.ripeness / 60, 2)} min. Max ripeness {round(self.buffer.buffer_flush_timeout / 60, 2)} min."  # noqa
         )
+        logger.info(
+            f"Current records size to process: {round(df_destination_records.estimated_size(unit='b') / 1024 / 1024, 2)} Mb."
+        )
+
+        if df_destination_records.estimated_size(unit="b") > self.buffer.buffer_size:
+            raise ValueError(
+                f"Records size {round(df_destination_records.estimated_size(unit='b') / 1024 / 1024, 2)} Mb is greater than buffer size {round(self.buffer.buffer_size / 1024 / 1024, 2)} Mb. Please increase destination buffer_size or reduce batch_size from the source."
+            )
 
         # Write buffer to destination if buffer is ripe and create a new buffer for the new iteration
         if self.buffer.is_ripe:
@@ -245,26 +286,73 @@ class DestinationFactory:
         sync_metadata: SyncMetadata,
         config: AbstractDestinationConfig,
         backend: AbstractBackend,
+        source_callback: AbstractSourceCallback,
+        monitor: AbstractMonitor,
     ) -> AbstractDestination:
 
         if config.name == DestinationTypes.LOGGER:
-            from .logger.src.destination import LoggerDestination
+            from bizon.connectors.destinations.logger.src.destination import (
+                LoggerDestination,
+            )
 
-            return LoggerDestination(sync_metadata=sync_metadata, config=config.config, backend=backend)
+            return LoggerDestination(
+                sync_metadata=sync_metadata,
+                config=config.config,
+                backend=backend,
+                source_callback=source_callback,
+                monitor=monitor,
+            )
 
         elif config.name == DestinationTypes.BIGQUERY:
-            from .bigquery.src.destination import BigQueryDestination
+            from bizon.connectors.destinations.bigquery.src.destination import (
+                BigQueryDestination,
+            )
 
-            return BigQueryDestination(sync_metadata=sync_metadata, config=config.config, backend=backend)
+            return BigQueryDestination(
+                sync_metadata=sync_metadata,
+                config=config.config,
+                backend=backend,
+                source_callback=source_callback,
+                monitor=monitor,
+            )
 
         elif config.name == DestinationTypes.BIGQUERY_STREAMING:
-            from .bigquery_streaming.src.destination import BigQueryStreamingDestination
+            from bizon.connectors.destinations.bigquery_streaming.src.destination import (
+                BigQueryStreamingDestination,
+            )
+
+            return BigQueryStreamingDestination(
+                sync_metadata=sync_metadata,
+                config=config.config,
+                backend=backend,
+                source_callback=source_callback,
+                monitor=monitor,
+            )
 
-            return BigQueryStreamingDestination(sync_metadata=sync_metadata, config=config.config, backend=backend)
+        elif config.name == DestinationTypes.BIGQUERY_STREAMING_V2:
+            from bizon.connectors.destinations.bigquery_streaming_v2.src.destination import (
+                BigQueryStreamingV2Destination,
+            )
+
+            return BigQueryStreamingV2Destination(
+                sync_metadata=sync_metadata,
+                config=config.config,
+                backend=backend,
+                source_callback=source_callback,
+                monitor=monitor,
+            )
 
         elif config.name == DestinationTypes.FILE:
-            from .file.src.destination import FileDestination
+            from bizon.connectors.destinations.file.src.destination import (
+                FileDestination,
+            )
 
-            return FileDestination(sync_metadata=sync_metadata, config=config.config, backend=backend)
+            return FileDestination(
+                sync_metadata=sync_metadata,
+                config=config.config,
+                backend=backend,
+                source_callback=source_callback,
+                monitor=monitor,
+            )
 
         raise ValueError(f"Destination {config.name}" f"with params {config} not found")

bizon/engine/backend/adapters/sqlalchemy/backend.py

@@ -173,7 +173,9 @@ class SQLAlchemyBackend(AbstractBackend):
 
     def _execute(self, select: Select, session: Optional[Session] = None) -> Result:
         session = session or self.session
-        return session.execute(select)
+        result = session.execute(select)
+        session.commit()
+        return result
 
     #### STREAM JOB ####