bizon 0.0.4.dev2__tar.gz → 0.0.4.dev4__tar.gz

{bizon-0.0.4.dev2 → bizon-0.0.4.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: bizon
-Version: 0.0.4.dev2
+Version: 0.0.4.dev4
 Summary: Extract and load your data reliably from API Clients with native fault-tolerant and checkpointing mechanism.
 Author: Antoine Balliet
 Author-email: antoine.balliet@gmail.com
@@ -45,7 +45,8 @@ Extract and load your largest data streams with a framework you can trust for bi
 
 ## Features
 - **Natively fault-tolerant**: Bizon uses a checkpointing mechanism to keep track of the progress and recover from the last checkpoint.
-- **Queue system agnostic**: Bizon is agnostic of the queuing system, you can use any queuing system like Python Queue, Kafka or Redpanda. Thanks to the `bizon.queue.Queue` interface, adapters can be written for any queuing system.
+- **High throughput**: Bizon is designed to handle high throughput and can process billions of records.
+- **Queue system agnostic**: Bizon is agnostic of the queuing system, you can use any queuing system among Python Queue, RabbitMQ, Kafka or Redpanda. Thanks to the `bizon.engine.queue.Queue` interface, adapters can be written for any queuing system.
 - **Pipeline metrics**: Bizon provides exhaustive pipeline metrics and implement OpenTelemetry for tracing. You can monitor:
     - ETAs for completion
     - Number of records processed
@@ -64,6 +65,14 @@ pip install bizon
 
 ## Usage
 
+### List available sources and streams
+```bash
+bizon source list
+bizon stream list <source_name>
+```
+
+### Create a pipeline
+
 Create a file named `config.yml` in your working directory with the following content:
 
 ```yaml
@@ -97,7 +106,8 @@ Backend is the interface used by Bizon to store its state. It can be configured
 
 Queue is the interface used by Bizon to exchange data between `Source` and `Destination`. It can be configured in the `queue` section of the configuration file. The following queues are supported:
 - `python_queue`: Python Queue, useful for testing and development.
-- `kafka`: Apache Kafka, for production use and high throughput.
+- `rabbitmq`: RabbitMQ, for production use and high throughput.
+- `kafka`: Apache Kafka, for production use and high throughput and strong persistence.
 
 ## Start syncing your data 🚀
 

{bizon-0.0.4.dev2 → bizon-0.0.4.dev4}/README.md

@@ -3,7 +3,8 @@ Extract and load your largest data streams with a framework you can trust for bi
 
 ## Features
 - **Natively fault-tolerant**: Bizon uses a checkpointing mechanism to keep track of the progress and recover from the last checkpoint.
-- **Queue system agnostic**: Bizon is agnostic of the queuing system, you can use any queuing system like Python Queue, Kafka or Redpanda. Thanks to the `bizon.queue.Queue` interface, adapters can be written for any queuing system.
+- **High throughput**: Bizon is designed to handle high throughput and can process billions of records.
+- **Queue system agnostic**: Bizon is agnostic of the queuing system, you can use any queuing system among Python Queue, RabbitMQ, Kafka or Redpanda. Thanks to the `bizon.engine.queue.Queue` interface, adapters can be written for any queuing system.
 - **Pipeline metrics**: Bizon provides exhaustive pipeline metrics and implement OpenTelemetry for tracing. You can monitor:
     - ETAs for completion
     - Number of records processed
@@ -22,6 +23,14 @@ pip install bizon
 
 ## Usage
 
+### List available sources and streams
+```bash
+bizon source list
+bizon stream list <source_name>
+```
+
+### Create a pipeline
+
 Create a file named `config.yml` in your working directory with the following content:
 
 ```yaml
@@ -55,7 +64,8 @@ Backend is the interface used by Bizon to store its state. It can be configured
 
 Queue is the interface used by Bizon to exchange data between `Source` and `Destination`. It can be configured in the `queue` section of the configuration file. The following queues are supported:
 - `python_queue`: Python Queue, useful for testing and development.
-- `kafka`: Apache Kafka, for production use and high throughput.
+- `rabbitmq`: RabbitMQ, for production use and high throughput.
+- `kafka`: Apache Kafka, for production use and high throughput and strong persistence.
 
 ## Start syncing your data 🚀
 

bizon-0.0.4.dev4/bizon/__main__.py

@@ -0,0 +1,4 @@
+from bizon.cli.main import cli
+
+if __name__ == "__main__":
+    cli()

{bizon-0.0.4.dev2 → bizon-0.0.4.dev4}/bizon/cli/main.py

@@ -1,12 +1,13 @@
 import click
 
 from bizon.engine.engine import RunnerFactory
+from bizon.engine.runner.config import LoggerLevel
 from bizon.source.discover import discover_all_sources
 
 from .utils import (
     parse_from_yaml,
     set_custom_source_path_in_config,
-    set_debug_mode,
+    set_log_level,
     set_runner_in_config,
 )
 
@@ -50,7 +51,7 @@ def stream():
 
 @stream.command()
 @click.argument("source_name", type=click.STRING)
-def list(source_name: str):
+def list(source_name: str):  # noqa
     """List available streams for a source."""
     sources = discover_all_sources()
     source_model = sources.get(source_name)
@@ -88,25 +89,29 @@ def destination():
     help="Runner type to use. Thread or Process.",
 )
 @click.option(
-    "--debug",
+    "--log-level",
     required=False,
-    is_flag=True,
+    type=click.Choice([level.name for level in LoggerLevel]),
     show_default=True,
-    default=False,
-    help="Enable debug mode.",
+    help="Log level to use.",
 )
-def run(filename: str, custom_source: str, runner: str, debug, help="Run a bizon pipeline from a YAML file."):
+def run(
+    filename: str,
+    custom_source: str,
+    runner: str,
+    log_level: LoggerLevel,
+    help="Run a bizon pipeline from a YAML file.",
+):
     """Run a bizon pipeline from a YAML file."""
-    ctx = click.get_current_context()
 
     # Parse config from YAML file as a dictionary
     config = parse_from_yaml(filename)
 
     # Set debug mode
-    set_debug_mode(debug)
+    set_log_level(config=config, level=log_level)
 
     # Override source_file_path param in config
-    set_custom_source_path_in_config(config=config, custom_source=ctx.get_parameter_source("custom-source"))
+    set_custom_source_path_in_config(config=config, custom_source=custom_source)
 
     # Override runner param in config
     set_runner_in_config(config=config, runner=runner)
@@ -114,6 +119,8 @@ def run(filename: str, custom_source: str, runner: str, debug, help="Run a bizon
     runner = RunnerFactory.create_from_config_dict(config=config)
     runner.run()
 
+    click.echo("Pipeline finished.")
+
 
 if __name__ == "__main__":
     cli()

{bizon-0.0.4.dev2 → bizon-0.0.4.dev4}/bizon/cli/utils.py

@@ -1,5 +1,3 @@
-import os
-
 import yaml
 
 
@@ -9,10 +7,17 @@ def parse_from_yaml(path_to_yaml) -> dict:
     return config
 
 
-def set_debug_mode(debug: bool):
+# TODO: Refacto
+def set_log_level(config: dict, level: str):
     # Set Log Level to DEBUG
-    if debug:
-        os.environ["LOGURU_LEVEL"] = "DEBUG"
+    if level:
+        if "engine" not in config:
+            config["engine"] = {}
+
+        if "runner" not in config["engine"]:
+            config["engine"]["runner"] = {"log_level": level}
+
+        config["engine"]["runner"]["log_level"] = level
 
 
 def set_custom_source_path_in_config(config: dict, custom_source: str):
@@ -20,6 +25,7 @@ def set_custom_source_path_in_config(config: dict, custom_source: str):
         config["source"]["source_file_path"] = custom_source
 
 
+# TODO: Refacto
 def set_runner_in_config(config: dict, runner: str):
     if runner:
         if "engine" not in config:

{bizon-0.0.4.dev2 → bizon-0.0.4.dev4}/bizon/destinations/bigquery/config/bigquery.example.yml

@@ -13,6 +13,8 @@ destination:
   # Must have the bigquery.dataEditor and storage.objectUser on the supplied dataset and bucket
   name: bigquery
   config:
+    buffer_size: 10 # in Mb
+    buffer_flush_timeout: 300 # in seconds
     dataset_id: bizon_test
     dataset_location: US
     project_id: my-gcp-project-id

bizon-0.0.4.dev4/bizon/destinations/buffer.py

@@ -0,0 +1,76 @@
+import sys
+from datetime import datetime
+from typing import List
+
+from bizon.destinations.models import DestinationRecord
+
+
+class DestinationBuffer:
+
+    def __init__(self, buffer_size: int, buffer_flush_timeout: int) -> None:
+        self.buffer_size = buffer_size * 1024 * 1024  # Convert to bytes
+        self.buffer_flush_timeout = buffer_flush_timeout
+        self.records: List[DestinationRecord] = []
+        self._iterations: List[int] = []
+        self.pagination = {}
+        self.modified_at: List[datetime] = [datetime.utcnow()]
+
+    @property
+    def current_size(self) -> int:
+        """Return buffer size"""
+        return sys.getsizeof(self.records)
+
+    @property
+    def buffer_free_space_pct(self) -> float:
+        """Return free space in buffer in percentage"""
+        return round((self.buffer_free_space / self.buffer_size) * 100, 3)
+
+    @property
+    def from_iteration(self) -> int:
+        """Return the smallest iteration in buffer"""
+        if not self._iterations:
+            raise ValueError("Buffer is empty")
+        return min(self._iterations)
+
+    @property
+    def to_iteration(self) -> int:
+        """Return the largest iteration in buffer"""
+        if not self._iterations:
+            raise ValueError("Buffer is empty")
+        return max(self._iterations)
+
+    @property
+    def buffer_free_space(self) -> int:
+        """Return free space for records in buffer"""
+        assert self.current_size <= self.buffer_size, "Buffer size exceeded"
+        return self.buffer_size - self.current_size
+
+    @property
+    def ripeness(self) -> float:
+        """Return buffer ripeness"""
+        if self.buffer_flush_timeout == 0:
+            return 0
+        return round((max(self.modified_at) - min(self.modified_at)).seconds, 2)
+
+    @property
+    def is_ripe(self) -> bool:
+        """Check if buffer is ripe for flushing based on the timeout"""
+        if self.buffer_flush_timeout == 0:
+            return False
+        return (max(self.modified_at) - min(self.modified_at)).seconds >= self.buffer_flush_timeout
+
+    def flush(self):
+        """Flush buffer"""
+        self.records = []
+        self._iterations = []
+        self.pagination = {}
+        self.modified_at = []
+
+    def add_source_iteration_records_to_buffer(
+        self, iteration: int, records: List[DestinationRecord], pagination: dict = None
+    ):
+        """Add records for the given iteration to buffer"""
+        self.records.extend(records)
+        self._iterations.append(iteration)
+        self.pagination = pagination
+        self.modified_at.append(datetime.utcnow())

{bizon-0.0.4.dev2 → bizon-0.0.4.dev4}/bizon/destinations/config.py

@@ -20,7 +20,14 @@ class NormalizationConfig(BaseModel):
 
 
 class AbstractDestinationDetailsConfig(BaseModel):
-    buffer_size: int = Field(default=2000, description="Buffer size for the destination")
+    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.",
+    )
     normalization: Optional[NormalizationConfig] = Field(
         description="Normalization configuration, by default no normalization is applied",
         default=NormalizationConfig(type=NormalizationType.NONE),

bizon-0.0.4.dev4/bizon/destinations/destination.py

@@ -0,0 +1,266 @@
+import sys
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import List, Optional, Tuple, Union
+
+from loguru import logger
+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.source.config import SourceSyncModes
+from bizon.source.models import SourceRecord
+
+from .buffer import DestinationBuffer
+from .config import (
+    AbstractDestinationConfig,
+    AbstractDestinationDetailsConfig,
+    DestinationTypes,
+)
+from .models import DestinationRecord
+
+
+class DestinationBufferStatus(str, Enum):
+    """Destination buffer status"""
+
+    RECORDS_WRITTEN = "RECORDS_WRITTEN"
+    RECORDS_WRITTEN_THEN_BUFFERED = "RECORDS_WRITTEN_THEN_BUFFERED"
+    RECORDS_BUFFERED = "RECORDS_BUFFERED"
+    NO_RECORDS = "NO_RECORDS"
+
+
+class DestinationIteration(BaseModel):
+    success: bool = Field(..., description="Success status of the iteration")
+    error_message: Optional[str] = Field(None, description="Error message if iteration failed")
+    records_written: int = Field(0, description="Number of records written to the destination")
+    from_source_iteration: Optional[int] = Field(None, description="From source iteration identifier buffer starts")
+    to_source_iteration: Optional[int] = Field(
+        None, description="To source iteration identifier buffer ends, inclusive"
+    )
+    pagination: Optional[dict] = Field(None, description="Source pagination for next interation recovery purposes")
+
+
+class AbstractDestination(ABC):
+
+    def __init__(self, sync_metadata: SyncMetadata, config: AbstractDestinationDetailsConfig, backend: AbstractBackend):
+        self.sync_metadata = sync_metadata
+        self.config = config
+        self.backend = backend
+        self.buffer = DestinationBuffer(
+            buffer_size=self.config.buffer_size, buffer_flush_timeout=self.config.buffer_flush_timeout
+        )
+
+    @abstractmethod
+    def check_connection(self) -> bool:
+        pass
+
+    @abstractmethod
+    def delete_table(self) -> bool:
+        """Delete table in destination"""
+        pass
+
+    @abstractmethod
+    def write_records(self, destination_records: List[DestinationRecord]) -> Tuple[bool, Union[str, None]]:
+        """Write records to destination for the given iteration and return success status and error message"""
+        pass
+
+    @staticmethod
+    def convert_source_records_to_destination_records(
+        source_records: List[SourceRecord], extracted_at: datetime
+    ) -> List[DestinationRecord]:
+        """Convert source records to destination records"""
+        return [
+            DestinationRecord.from_source_record(source_record=source_record, extracted_at=extracted_at)
+            for source_record in source_records
+        ]
+
+    def prepare_destination(self):
+        """Prepare destination before writing records"""
+        # Delete table if full refresh
+        if self.sync_metadata.sync_mode == SourceSyncModes.FULL_REFRESH:
+            self.delete_table()
+
+    def buffer_flush_handler(self, session=None) -> DestinationIteration:
+        # TODO: Add retry strategy
+
+        # Initialize destination iteration
+        destination_iteration = DestinationIteration(
+            success=False,
+            records_written=0,
+            pagination=self.buffer.pagination,
+        )
+
+        success, error_msg = self.write_records(destination_records=self.buffer.records)
+
+        if success:
+            # We wrote records to destination so we keep it
+            destination_iteration.records_written = len(self.buffer.records)
+
+        else:
+            # We failed to write records to destination so we keep the error message
+            destination_iteration.error_message = error_msg
+
+        destination_iteration.success = success
+        destination_iteration.from_source_iteration = self.buffer.from_iteration
+        destination_iteration.to_source_iteration = self.buffer.to_iteration
+        destination_iteration.pagination = self.buffer.pagination
+
+        # Update destination cursor
+        self.create_cursors(destination_iteration=destination_iteration)
+
+        return destination_iteration
+
+    def write_or_buffer_records(
+        self,
+        destination_records: List[DestinationRecord],
+        iteration: int,
+        last_iteration: bool = False,
+        session=None,
+        pagination: dict = None,
+    ) -> DestinationBufferStatus:
+        """Write records to destination or buffer them for the given iteration"""
+
+        # Last iteration, write all records to destination
+        if last_iteration:
+            logger.debug("Writing last iteration records to destination")
+            assert len(destination_records) == 0, "Last iteration should not have any records"
+            destination_iteration = self.buffer_flush_handler(session=session)
+
+            if destination_iteration.success:
+                # Update job status to success
+                self.backend.update_stream_job_status(
+                    job_id=self.sync_metadata.job_id,
+                    job_status=JobStatus.SUCCEEDED,
+                    session=session,
+                )
+
+            self.buffer.flush()
+            return DestinationBufferStatus.RECORDS_WRITTEN
+
+        # Don't write empty records to destination
+        if len(destination_records) == 0 and not last_iteration:
+            logger.warning("No records to write to destination. Check source and queue provider.")
+            return DestinationBufferStatus.NO_RECORDS
+
+        # Write records to destination if buffer size is 0
+        if self.buffer.buffer_size == 0:
+            logger.info("Writing last iteration records to destination")
+            self.buffer.add_source_iteration_records_to_buffer(
+                iteration=iteration, records=destination_records, pagination=pagination
+            )
+            self.buffer_flush_handler(session=session)
+            self.buffer.flush()
+            return DestinationBufferStatus.RECORDS_WRITTEN
+
+        logger.debug(f"Buffer free space {self.buffer.buffer_free_space_pct}%")
+        logger.debug(f"Buffer current size {self.buffer.current_size} bytes")
+        logger.info(
+            f"Buffer ripeness {self.buffer.ripeness / 60} min. Max ripeness {self.buffer.buffer_flush_timeout / 60} min."  # noqa
+        )
+
+        # Write buffer to destination if buffer is ripe and create a new buffer for the new iteration
+        if self.buffer.is_ripe:
+            logger.info(
+                f"Buffer is ripe (buffering for longer than buffer_flush_timeout: {self.buffer.buffer_flush_timeout} seconds), writing buffer to destination"  # noqa
+            )
+            self.buffer_flush_handler(session=session)
+            self.buffer.flush()
+            self.buffer.add_source_iteration_records_to_buffer(
+                iteration=iteration, records=destination_records, pagination=pagination
+            )
+            return DestinationBufferStatus.RECORDS_WRITTEN_THEN_BUFFERED
+
+        # Buffer can hold all records from this iteration
+        elif self.buffer.buffer_free_space >= sys.getsizeof(destination_records):
+            self.buffer.add_source_iteration_records_to_buffer(
+                iteration=iteration, records=destination_records, pagination=pagination
+            )
+            return DestinationBufferStatus.RECORDS_BUFFERED
+
+        # Buffer can contain some records from this iteration
+        # For now we will write all records to destination and then buffer the remaining records
+        else:
+            self.buffer_flush_handler(session=session)
+            self.buffer.flush()
+            self.buffer.add_source_iteration_records_to_buffer(
+                iteration=iteration, records=destination_records, pagination=pagination
+            )
+            return DestinationBufferStatus.RECORDS_WRITTEN_THEN_BUFFERED
+
+    def create_cursors(self, destination_iteration: DestinationIteration):
+        self.backend.create_destination_cursor(
+            job_id=self.sync_metadata.job_id,
+            source_name=self.sync_metadata.source_name,
+            stream_name=self.sync_metadata.stream_name,
+            destination_name=self.sync_metadata.destination_name,
+            from_source_iteration=destination_iteration.from_source_iteration,
+            to_source_iteration=destination_iteration.to_source_iteration,
+            rows_written=destination_iteration.records_written,
+            pagination=destination_iteration.pagination,
+            success=destination_iteration.success,
+        )
+
+    def write_records_and_update_cursor(
+        self,
+        source_records: List[SourceRecord],
+        iteration: int,
+        extracted_at: datetime,
+        last_iteration: bool = False,
+        pagination: dict = None,
+    ) -> bool:
+        """
+        Write records to destination and update the cursor for the given iteration.
+        Stores the source pagination for recovery purposes.
+        """
+
+        # Case when producer failed to fetch data from first iteration
+        if iteration == 0 and len(source_records) == 0:
+            logger.warning("Source failed to fetch data from the first iteration, no records will be written.")
+            return False
+
+        # Prepare destination
+        if iteration == 0:
+            self.prepare_destination()
+
+        destination_records = self.convert_source_records_to_destination_records(
+            source_records=source_records,
+            extracted_at=extracted_at,
+        )
+
+        # Buffer records otherwise write to destination
+        self.write_or_buffer_records(
+            destination_records=destination_records,
+            iteration=iteration,
+            last_iteration=last_iteration,
+            pagination=pagination,
+        )
+
+        return True
+
+
+class DestinationFactory:
+    @staticmethod
+    def get_destination(
+        sync_metadata: SyncMetadata,
+        config: AbstractDestinationConfig,
+        backend: AbstractBackend,
+    ) -> AbstractDestination:
+
+        if config.name == DestinationTypes.LOGGER:
+            from .logger.src.destination import LoggerDestination
+
+            return LoggerDestination(sync_metadata=sync_metadata, config=config.config, backend=backend)
+
+        elif config.name == DestinationTypes.BIGQUERY:
+            from .bigquery.src.destination import BigQueryDestination
+
+            return BigQueryDestination(sync_metadata=sync_metadata, config=config.config, backend=backend)
+
+        elif config.name == DestinationTypes.FILE:
+            from .file.src.destination import FileDestination
+
+            return FileDestination(sync_metadata=sync_metadata, config=config.config, backend=backend)
+
+        raise ValueError(f"Destination {config.name}" f"with params {config} not found")

{bizon-0.0.4.dev2 → bizon-0.0.4.dev4}/bizon/destinations/file/src/destination.py

@@ -1,4 +1,3 @@
-import json
 from typing import List, Tuple
 
 from loguru import logger
@@ -26,6 +25,5 @@ class FileDestination(AbstractDestination):
     def write_records(self, destination_records: List[DestinationRecord]) -> Tuple[bool, str]:
         with open(self.config.filepath, "a") as f:
             for record in destination_records:
-                json.dump(record.model_dump_json(), f)
-                f.write("\n")
+                f.write(record.model_dump_json() + "\n")
         return True, ""

{bizon-0.0.4.dev2 → bizon-0.0.4.dev4}/bizon/destinations/models.py

@@ -1,6 +1,5 @@
 import json
 from datetime import datetime
-from typing import Optional
 from uuid import uuid4
 
 from pydantic import BaseModel, Field
@@ -9,14 +8,6 @@ from pytz import UTC
 from bizon.source.models import SourceRecord
 
 
-class DestinationIteration(BaseModel):
-    success: bool = Field(..., description="Success status of the iteration")
-    error_message: Optional[str] = Field(None, description="Error message if iteration failed")
-    records_written: int = Field(0, description="Number of records written to the destination")
-    from_source_iteration: Optional[int] = Field(None, description="From source iteration identifier buffer starts")
-    to_source_iteration: Optional[int] = Field(None, description="To source iteration identifier buffer ends")
-
-
 class DestinationRecord(BaseModel):
     bizon_id: str = Field(..., description="Bizon unique identifier of the record")
     bizon_extracted_at: datetime = Field(..., description="Datetime when the record was extracted")