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

bizon/engine/engine.py

@@ -1,4 +1,4 @@
-from loguru import logger
+from os import getenv
 
 from bizon.cli.utils import parse_from_yaml
 from bizon.common.models import BizonConfig
@@ -7,10 +7,24 @@ from .config import RunnerTypes
 from .runner.runner import AbstractRunner
 
 
+def replace_env_variables_in_config(config: dict) -> dict:
+    """Replace templated secrets with actual values from environment variables"""
+    for key, value in config.items():
+        if isinstance(value, dict):
+            config[key] = replace_env_variables_in_config(value)
+        elif isinstance(value, str):
+            if value.startswith("BIZON_ENV_"):
+                config[key] = getenv(value)
+    return config
+
+
 class RunnerFactory:
     @staticmethod
     def create_from_config_dict(config: dict) -> AbstractRunner:
 
+        # Replace env variables in config
+        config = replace_env_variables_in_config(config=config)
+
         bizon_config = BizonConfig.model_validate(obj=config)
 
         if bizon_config.engine.runner.type == RunnerTypes.THREAD:
@@ -23,6 +37,11 @@ class RunnerFactory:
 
             return ProcessRunner(config=config)
 
+        if bizon_config.engine.runner.type == RunnerTypes.STREAM:
+            from .runner.adapters.streaming import StreamingRunner
+
+            return StreamingRunner(config=config)
+
         raise ValueError(f"Runner type {bizon_config.engine.runner.type} is not supported")
 
     @staticmethod

bizon/engine/pipeline/consumer.py

@@ -1,15 +1,83 @@
+import multiprocessing
+import multiprocessing.synchronize
+import threading
+import traceback
 from abc import ABC, abstractmethod
-from enum import Enum
+from typing import Union
 
-from bizon.destinations.destination import AbstractDestination
-from bizon.engine.queue.config import AbstractQueueConfig
+from loguru import logger
+
+from bizon.destination.destination import AbstractDestination
+from bizon.engine.pipeline.models import PipelineReturnStatus
+from bizon.engine.queue.config import (
+    QUEUE_TERMINATION,
+    AbstractQueueConfig,
+    QueueMessage,
+)
+from bizon.monitoring.monitor import AbstractMonitor
+from bizon.transform.transform import Transform
 
 
 class AbstractQueueConsumer(ABC):
-    def __init__(self, config: AbstractQueueConfig, destination: AbstractDestination):
+    def __init__(
+        self,
+        config: AbstractQueueConfig,
+        destination: AbstractDestination,
+        transform: Transform,
+        monitor: AbstractMonitor,
+    ):
         self.config = config
         self.destination = destination
+        self.transform = transform
+        self.monitor = monitor
 
     @abstractmethod
-    def run(self):
+    def run(self, stop_event: Union[multiprocessing.synchronize.Event, threading.Event]) -> PipelineReturnStatus:
         pass
+
+    def process_queue_message(self, queue_message: QueueMessage) -> PipelineReturnStatus:
+
+        # Apply the transformation
+        try:
+            df_source_records = self.transform.apply_transforms(df_source_records=queue_message.df_source_records)
+        except Exception as e:
+            logger.error(f"Error applying transformation: {e}")
+            logger.error(traceback.format_exc())
+            self.monitor.track_pipeline_status(PipelineReturnStatus.TRANSFORM_ERROR)
+            return PipelineReturnStatus.TRANSFORM_ERROR
+
+        # Handle last iteration
+        try:
+            if queue_message.signal == QUEUE_TERMINATION:
+                logger.info("Received termination signal, waiting for destination to close gracefully ...")
+                self.destination.write_records_and_update_cursor(
+                    df_source_records=df_source_records,
+                    iteration=queue_message.iteration,
+                    extracted_at=queue_message.extracted_at,
+                    pagination=queue_message.pagination,
+                    last_iteration=True,
+                )
+                self.monitor.track_pipeline_status(PipelineReturnStatus.SUCCESS)
+                return PipelineReturnStatus.SUCCESS
+
+        except Exception as e:
+            logger.error(f"Error writing records to destination: {e}")
+            self.monitor.track_pipeline_status(PipelineReturnStatus.DESTINATION_ERROR)
+            return PipelineReturnStatus.DESTINATION_ERROR
+
+        # Write the records to the destination
+        try:
+            self.destination.write_records_and_update_cursor(
+                df_source_records=df_source_records,
+                iteration=queue_message.iteration,
+                extracted_at=queue_message.extracted_at,
+                pagination=queue_message.pagination,
+            )
+            return PipelineReturnStatus.RUNNING
+
+        except Exception as e:
+            logger.error(f"Error writing records to destination: {e}")
+            self.monitor.track_pipeline_status(PipelineReturnStatus.DESTINATION_ERROR)
+            return PipelineReturnStatus.DESTINATION_ERROR
+
+        raise RuntimeError("Should not reach this point")

bizon/engine/pipeline/models.py

@@ -1,10 +1,15 @@
 from enum import Enum
 
 
-class PipelineReturnStatus(Enum):
+class PipelineReturnStatus(str, Enum):
     """Producer error types"""
 
-    SUCCESS = "success"
+    BACKEND_ERROR = "backend_error"
+    DESTINATION_ERROR = "destination_error"
+    KILLED_BY_RUNNER = "killed_by_runner"
     QUEUE_ERROR = "queue_error"
+    RUNNING = "running"
     SOURCE_ERROR = "source_error"
-    BACKEND_ERROR = "backend_error"
+    SUCCESS = "success"
+    TRANSFORM_ERROR = "transform_error"
+    STREAM_ERROR = "stream_error"

bizon/engine/pipeline/producer.py

@@ -1,8 +1,11 @@
 import ast
+import multiprocessing
+import multiprocessing.synchronize
+import threading
 import traceback
 from datetime import datetime
 from time import sleep
-from typing import Tuple
+from typing import Tuple, Union
 
 from loguru import logger
 from pytz import UTC
@@ -28,7 +31,7 @@ class Producer:
 
     @property
     def name(self) -> str:
-        return f"producer-{self.source.config.source_name}-{self.source.config.stream_name}"
+        return f"producer-{self.source.config.name}-{self.source.config.stream}"
 
     def get_or_create_cursor(self, job_id: str, session=None) -> Cursor:
         """Get or create a cursor for the current stream, return the cursor"""
@@ -47,8 +50,8 @@ class Producer:
 
             # Initialize the recovery from the DestinationCursor
             cursor = Cursor.from_db(
-                source_name=self.source.config.source_name,
-                stream_name=self.source.config.stream_name,
+                source_name=self.source.config.name,
+                stream_name=self.source.config.stream,
                 job_id=job_id,
                 total_records=job.total_records_to_fetch,
                 iteration=cursor_from_db.to_source_iteration + 1,
@@ -60,8 +63,8 @@ class Producer:
             total_records = self.source.get_total_records_count()
             # Initialize the cursor
             cursor = Cursor(
-                source_name=self.source.config.source_name,
-                stream_name=self.source.config.stream_name,
+                source_name=self.source.config.name,
+                stream_name=self.source.config.stream,
                 job_id=job_id,
                 total_records=total_records,
             )
@@ -99,7 +102,9 @@ class Producer:
 
         return False, queue_size, approximate_nb_records_in_queue
 
-    def run(self, job_id: int):
+    def run(
+        self, job_id: int, stop_event: Union[multiprocessing.synchronize.Event, threading.Event]
+    ) -> PipelineReturnStatus:
 
         return_value: PipelineReturnStatus = PipelineReturnStatus.SUCCESS
 
@@ -128,6 +133,10 @@ class Producer:
 
         while not cursor.is_finished:
 
+            if stop_event.is_set():
+                logger.info("Stop event is set, terminating producer ...")
+                return PipelineReturnStatus.KILLED_BY_RUNNER
+
             timestamp_start_iteration = datetime.now(tz=UTC)
 
             # Handle the case where last cursor already reach max_iterations
@@ -142,8 +151,8 @@ class Producer:
                     self.backend.create_source_cursor(
                         job_id=job_id,
                         name=self.bizon_config.name,
-                        source_name=self.source.config.source_name,
-                        stream_name=self.source.config.stream_name,
+                        source_name=self.source.config.name,
+                        stream_name=self.source.config.stream,
                         iteration=cursor.iteration,
                         rows_fetched=cursor.rows_fetched,
                         next_pagination=cursor.pagination,

bizon/engine/queue/adapters/kafka/consumer.py

@@ -3,9 +3,9 @@ import json
 from kafka import KafkaConsumer
 from loguru import logger
 
-from bizon.destinations.destination import AbstractDestination
+from bizon.destination.destination import AbstractDestination
 from bizon.engine.pipeline.consumer import AbstractQueueConsumer
-from bizon.engine.queue.queue import QUEUE_TERMINATION, QueueMessage
+from bizon.engine.queue.config import QUEUE_TERMINATION, QueueMessage
 
 from .config import KafkaConfigDetails
 

bizon/engine/queue/adapters/kafka/queue.py

@@ -4,8 +4,9 @@ from typing import Union
 from kafka import KafkaProducer
 from loguru import logger
 
-from bizon.destinations.destination import AbstractDestination
-from bizon.engine.queue.queue import QUEUE_TERMINATION, AbstractQueue, QueueMessage
+from bizon.destination.destination import AbstractDestination
+from bizon.engine.queue.config import QUEUE_TERMINATION, QueueMessage
+from bizon.engine.queue.queue import AbstractQueue
 
 from .config import KafkaConfigDetails
 from .consumer import KafkaConsumer_

bizon/engine/queue/adapters/python_queue/consumer.py

@@ -1,36 +1,53 @@
+import multiprocessing
+import multiprocessing.synchronize
+import threading
+from typing import Union
+
 from loguru import logger
 
-from bizon.destinations.destination import AbstractDestination
+from bizon.destination.destination import AbstractDestination
 from bizon.engine.pipeline.consumer import AbstractQueueConsumer
-from bizon.engine.queue.queue import QUEUE_TERMINATION, AbstractQueue, QueueMessage
+from bizon.engine.pipeline.models import PipelineReturnStatus
+from bizon.engine.queue.config import QueueMessage
+from bizon.engine.queue.queue import AbstractQueue
+from bizon.monitoring.monitor import AbstractMonitor
+from bizon.transform.transform import Transform
 
 from .config import PythonQueueConfig
 
 
 class PythonQueueConsumer(AbstractQueueConsumer):
-    def __init__(self, config: PythonQueueConfig, queue: AbstractQueue, destination: AbstractDestination):
-        super().__init__(config, destination=destination)
+    def __init__(
+        self,
+        config: PythonQueueConfig,
+        queue: AbstractQueue,
+        destination: AbstractDestination,
+        transform: Transform,
+        monitor: AbstractMonitor,
+    ):
+        super().__init__(
+            config,
+            destination=destination,
+            transform=transform,
+            monitor=monitor,
+        )
         self.queue = queue
+        self.monitor.track_pipeline_status(PipelineReturnStatus.RUNNING)
+
+    def run(self, stop_event: Union[threading.Event, multiprocessing.synchronize.Event]) -> PipelineReturnStatus:
 
-    def run(self) -> None:
         while True:
+            # Handle kill signal from the runner
+            if stop_event.is_set():
+                logger.info("Stop event is set, closing consumer ...")
+                self.monitor.track_pipeline_status(PipelineReturnStatus.KILLED_BY_RUNNER)
+                return PipelineReturnStatus.KILLED_BY_RUNNER
+
+            # Retrieve the message from the queue
             queue_message: QueueMessage = self.queue.get()
 
-            if queue_message.signal == QUEUE_TERMINATION:
-                logger.info("Received termination signal, waiting for destination to close gracefully ...")
-                self.destination.write_records_and_update_cursor(
-                    df_source_records=queue_message.df_source_records,
-                    iteration=queue_message.iteration,
-                    extracted_at=queue_message.extracted_at,
-                    pagination=queue_message.pagination,
-                    last_iteration=True,
-                )
-                break
-
-            self.destination.write_records_and_update_cursor(
-                df_source_records=queue_message.df_source_records,
-                iteration=queue_message.iteration,
-                extracted_at=queue_message.extracted_at,
-                pagination=queue_message.pagination,
-            )
-            self.queue.task_done()
+            status = self.process_queue_message(queue_message)
+
+            if status != PipelineReturnStatus.RUNNING:
+                self.queue.task_done()
+                return status

bizon/engine/queue/adapters/python_queue/queue.py

@@ -5,14 +5,13 @@ from typing import Union
 
 from loguru import logger
 
-from bizon.destinations.destination import AbstractDestination
-from bizon.engine.queue.queue import (
-    QUEUE_TERMINATION,
-    AbstractQueue,
-    AbstractQueueConsumer,
-    QueueMessage,
-)
+from bizon.destination.destination import AbstractDestination
+from bizon.engine.queue.config import QUEUE_TERMINATION, QueueMessage
+from bizon.engine.queue.queue import AbstractQueue, AbstractQueueConsumer
+from bizon.monitoring.monitor import AbstractMonitor
+from bizon.source.callback import AbstractSourceCallback
 from bizon.source.models import SourceIteration
+from bizon.transform.transform import Transform
 
 from .config import PythonQueueConfigDetails
 from .consumer import PythonQueueConsumer
@@ -31,8 +30,19 @@ class PythonQueue(AbstractQueue):
         # No connection to establish for PythonQueue
         pass
 
-    def get_consumer(self, destination: AbstractDestination) -> AbstractQueueConsumer:
-        return PythonQueueConsumer(config=self.config, queue=self.queue, destination=destination)
+    def get_consumer(
+        self,
+        destination: AbstractDestination,
+        transform: Transform,
+        monitor: AbstractMonitor,
+    ) -> AbstractQueueConsumer:
+        return PythonQueueConsumer(
+            config=self.config,
+            queue=self.queue,
+            destination=destination,
+            transform=transform,
+            monitor=monitor,
+        )
 
     def put_queue_message(self, queue_message: QueueMessage):
         if not self.queue.full():

bizon/engine/queue/adapters/rabbitmq/consumer.py

@@ -2,12 +2,9 @@ import pika
 import pika.connection
 from loguru import logger
 
-from bizon.destinations.destination import AbstractDestination
-from bizon.engine.queue.queue import (
-    QUEUE_TERMINATION,
-    AbstractQueueConsumer,
-    QueueMessage,
-)
+from bizon.destination.destination import AbstractDestination
+from bizon.engine.queue.config import QUEUE_TERMINATION
+from bizon.engine.queue.queue import AbstractQueueConsumer, QueueMessage
 
 from .config import RabbitMQConfigDetails
 

bizon/engine/queue/adapters/rabbitmq/queue.py

@@ -3,9 +3,10 @@ from typing import Union
 import pika
 from loguru import logger
 
-from bizon.destinations.destination import AbstractDestination
+from bizon.destination.destination import AbstractDestination
 from bizon.engine.pipeline.consumer import AbstractQueueConsumer
-from bizon.engine.queue.queue import QUEUE_TERMINATION, AbstractQueue, QueueMessage
+from bizon.engine.queue.config import QUEUE_TERMINATION, QueueMessage
+from bizon.engine.queue.queue import AbstractQueue
 
 from .config import RabbitMQConfigDetails
 from .consumer import RabbitMQConsumer

bizon/engine/queue/config.py

@@ -1,7 +1,23 @@
 from abc import ABC
+from dataclasses import dataclass
+from datetime import datetime
 from enum import Enum
+from typing import Optional
 
+import polars as pl
 from pydantic import BaseModel, ConfigDict, Field
+from pytz import UTC
+
+QUEUE_TERMINATION = "TERMINATION"
+
+
+@dataclass
+class QueueMessage:
+    iteration: int
+    df_source_records: pl.DataFrame
+    extracted_at: datetime = datetime.now(tz=UTC)
+    pagination: Optional[dict] = None
+    signal: Optional[str] = None
 
 
 class QueueTypes(str, Enum):

bizon/engine/queue/queue.py

@@ -1,28 +1,23 @@
 import json
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
 from datetime import datetime
-from typing import Optional, Union
+from typing import Union
 
 import polars as pl
 from pytz import UTC
 
-from bizon.destinations.destination import AbstractDestination
+from bizon.destination.destination import AbstractDestination
 from bizon.engine.pipeline.consumer import AbstractQueueConsumer
+from bizon.monitoring.monitor import AbstractMonitor
 from bizon.source.models import SourceIteration, source_record_schema
+from bizon.transform.transform import Transform
 
-from .config import AbastractQueueConfigDetails, AbstractQueueConfig, QueueTypes
-
-QUEUE_TERMINATION = "TERMINATION"
-
-
-@dataclass
-class QueueMessage:
-    iteration: int
-    df_source_records: pl.DataFrame
-    extracted_at: datetime = datetime.now(tz=UTC)
-    pagination: Optional[dict] = None
-    signal: Optional[str] = None
+from .config import (
+    AbastractQueueConfigDetails,
+    AbstractQueueConfig,
+    QueueMessage,
+    QueueTypes,
+)
 
 
 class AbstractQueue(ABC):
@@ -35,7 +30,12 @@ class AbstractQueue(ABC):
         pass
 
     @abstractmethod
-    def get_consumer(self, destination: AbstractDestination) -> AbstractQueueConsumer:
+    def get_consumer(
+        self,
+        destination: AbstractDestination,
+        transform: Transform,
+        monitor: AbstractMonitor,
+    ) -> AbstractQueueConsumer:
         pass
 
     @abstractmethod
@@ -71,6 +71,7 @@ class AbstractQueue(ABC):
                 "id": [record.id for record in source_iteration.records],
                 "data": [json.dumps(record.data, ensure_ascii=False) for record in source_iteration.records],
                 "timestamp": [record.timestamp for record in source_iteration.records],
+                "destination_id": [record.destination_id for record in source_iteration.records],
             },
             schema=source_record_schema,
         )

bizon/engine/runner/adapters/process.py

@@ -1,5 +1,6 @@
 import concurrent.futures
 import time
+import traceback
 
 from loguru import logger
 
@@ -50,6 +51,7 @@ class ProcessRunner(AbstractRunner):
             future_consumer = executor.submit(
                 AbstractRunner.instanciate_and_run_consumer,
                 self.bizon_config,
+                self.config,
                 job.id,
                 **extra_kwargs,
             )
@@ -68,8 +70,19 @@ class ProcessRunner(AbstractRunner):
                 result_producer = future_producer.result()
                 logger.info(f"Producer process stopped running with result: {result_producer}")
 
+                if result_producer.SUCCESS:
+                    logger.info("Producer thread has finished successfully, will wait for consumer to finish ...")
+                else:
+                    logger.error("Producer thread failed, stopping consumer ...")
+                    executor.shutdown(wait=False)
+
             if not future_consumer.running():
-                result_consumer = future_consumer.result()
-                logger.info(f"Consumer process stopped running with result: {result_consumer}")
+                try:
+                    future_consumer.result()
+                except Exception as e:
+                    logger.error(f"Consumer thread stopped running with error {e}")
+                    logger.error(traceback.format_exc())
+                finally:
+                    executor.shutdown(wait=False)
 
         return True

bizon/engine/runner/adapters/streaming.py

@@ -0,0 +1,121 @@
+import os
+import time
+from datetime import datetime
+from typing import List
+
+import polars as pl
+import simplejson as json
+from loguru import logger
+from pytz import UTC
+
+from bizon.common.models import BizonConfig, SyncMetadata
+from bizon.destination.models import transform_to_df_destination_records
+from bizon.engine.pipeline.models import PipelineReturnStatus
+from bizon.engine.runner.config import RunnerStatus
+from bizon.engine.runner.runner import AbstractRunner
+from bizon.source.models import SourceRecord, source_record_schema
+
+
+class StreamingRunner(AbstractRunner):
+    def __init__(self, config: BizonConfig):
+        super().__init__(config)
+
+    @staticmethod
+    def convert_source_records(records: List[SourceRecord]) -> pl.DataFrame:
+        return pl.DataFrame(
+            {
+                "id": [record.id for record in records],
+                "data": [json.dumps(record.data, ensure_ascii=False) for record in records],
+                "timestamp": [record.timestamp for record in records],
+                "destination_id": [record.destination_id for record in records],
+            },
+            schema=source_record_schema,
+        )
+
+    @staticmethod
+    def convert_to_destination_records(df_source_records: pl.DataFrame, extracted_at: datetime) -> pl.DataFrame:
+        return transform_to_df_destination_records(df_source_records=df_source_records, extracted_at=extracted_at)
+
+    def run(self) -> RunnerStatus:
+        job = self.init_job(bizon_config=self.bizon_config, config=self.config)
+        backend = self.get_backend(bizon_config=self.bizon_config)
+        source = self.get_source(bizon_config=self.bizon_config, config=self.config)
+
+        sync_metadata = SyncMetadata.from_bizon_config(job_id=job.id, config=self.bizon_config)
+        monitor = self.get_monitoring_client(sync_metadata=sync_metadata, bizon_config=self.bizon_config)
+
+        destination = self.get_destination(
+            bizon_config=self.bizon_config,
+            backend=backend,
+            job_id=job.id,
+            source_callback=None,
+            monitor=monitor,
+        )
+
+        transform = self.get_transform(bizon_config=self.bizon_config)
+
+        destination.buffer.buffer_size = 0  # force buffer to be flushed immediately
+        iteration = 0
+
+        while True:
+
+            if source.config.max_iterations and iteration > source.config.max_iterations:
+                logger.info(f"Max iterations {source.config.max_iterations} reached, terminating stream ...")
+                break
+
+            with monitor.trace(operation_name="bizon.stream.iteration"):
+                source_iteration = source.get()
+
+                destination_id_indexed_records = {}
+
+                if len(source_iteration.records) == 0:
+                    logger.info("No new records found, stopping iteration")
+                    time.sleep(2)
+                    monitor.track_pipeline_status(PipelineReturnStatus.SUCCESS)
+                    iteration += 1
+                    continue
+
+                for record in source_iteration.records:
+                    if destination_id_indexed_records.get(record.destination_id):
+                        destination_id_indexed_records[record.destination_id].append(record)
+                    else:
+                        destination_id_indexed_records[record.destination_id] = [record]
+
+                for destination_id, records in destination_id_indexed_records.items():
+                    df_source_records = StreamingRunner.convert_source_records(records)
+
+                    dsm_headers = monitor.track_source_iteration(records=records)
+
+                    # Apply transformation
+                    df_source_records = transform.apply_transforms(df_source_records=df_source_records)
+
+                    df_destination_records = StreamingRunner.convert_to_destination_records(
+                        df_source_records, datetime.now(tz=UTC)
+                    )
+                    # Override destination_id
+                    destination.destination_id = destination_id
+                    destination.write_or_buffer_records(
+                        df_destination_records=df_destination_records,
+                        iteration=iteration,
+                        pagination=None,
+                    )
+                    monitor.track_records_synced(
+                        num_records=len(df_destination_records),
+                        destination_id=destination_id,
+                        extra_tags={"destination_id": destination_id},
+                        headers=dsm_headers,
+                    )
+
+                if os.getenv("ENVIRONMENT") == "production":
+                    try:
+                        source.commit()
+                    except Exception as e:
+                        logger.error(f"Error committing source: {e}")
+                        monitor.track_pipeline_status(PipelineReturnStatus.ERROR)
+                        return RunnerStatus(stream=PipelineReturnStatus.ERROR)
+
+                iteration += 1
+
+                monitor.track_pipeline_status(PipelineReturnStatus.SUCCESS)
+
+        return RunnerStatus(stream=PipelineReturnStatus.SUCCESS)  # return when max iterations is reached