bizon 0.0.3.dev1__tar.gz → 0.0.4.dev2__tar.gz

{bizon-0.0.3.dev1 → bizon-0.0.4.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: bizon
-Version: 0.0.3.dev1
+Version: 0.0.4.dev2
 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
@@ -16,6 +16,7 @@ Provides-Extra: kafka
 Provides-Extra: postgres
 Provides-Extra: rabbitmq
 Requires-Dist: backoff (>=2.2.1,<3.0.0)
+Requires-Dist: click (>=8.1.7,<9.0.0)
 Requires-Dist: dpath (>=2.2.0,<3.0.0)
 Requires-Dist: faker (>=26.0.0,<27.0.0)
 Requires-Dist: google-cloud-bigquery (>=3.25.0,<4.0.0) ; extra == "bigquery"
@@ -32,6 +33,7 @@ Requires-Dist: pyarrow (>=16.1.0,<17.0.0)
 Requires-Dist: pydantic (>=2.8.2,<3.0.0)
 Requires-Dist: pydantic-extra-types (>=2.9.0,<3.0.0)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
+Requires-Dist: pytz (>=2024.2,<2025.0)
 Requires-Dist: pyyaml (>=6.0.1,<7.0.0)
 Requires-Dist: requests (>=2.28.2,<3.0.0)
 Requires-Dist: sqlalchemy (>=2.0.32,<3.0.0)
@@ -61,11 +63,10 @@ pip install bizon
 ```
 
 ## Usage
-```python
-from yaml import safe_load
-from bizon.engine.runner import RunnerFactory
 
-yaml_config = """
+Create a file named `config.yml` in your working directory with the following content:
+
+```yaml
 source:
   source_name: dummy
   stream_name: creatures
@@ -78,13 +79,13 @@ destination:
   name: logger
   config:
     dummy: dummy
-"""
-
-config = safe_load(yaml_config)
-runner = RunnerFactory.create_from_config_dict(config=config)
-runner.run()
 ```
 
+Run the pipeline with the following command:
+
+```bash
+bizon run config.yml
+```
 ## Backend configuration
 
 Backend is the interface used by Bizon to store its state. It can be configured in the `backend` section of the configuration file. The following backends are supported:

{bizon-0.0.3.dev1 → bizon-0.0.4.dev2}/README.md

@@ -21,11 +21,10 @@ pip install bizon
 ```
 
 ## Usage
-```python
-from yaml import safe_load
-from bizon.engine.runner import RunnerFactory
 
-yaml_config = """
+Create a file named `config.yml` in your working directory with the following content:
+
+```yaml
 source:
   source_name: dummy
   stream_name: creatures
@@ -38,13 +37,13 @@ destination:
   name: logger
   config:
     dummy: dummy
-"""
-
-config = safe_load(yaml_config)
-runner = RunnerFactory.create_from_config_dict(config=config)
-runner.run()
 ```
 
+Run the pipeline with the following command:
+
+```bash
+bizon run config.yml
+```
 ## Backend configuration
 
 Backend is the interface used by Bizon to store its state. It can be configured in the `backend` section of the configuration file. The following backends are supported:

bizon-0.0.4.dev2/bizon/cli/main.py

@@ -0,0 +1,119 @@
+import click
+
+from bizon.engine.engine import RunnerFactory
+from bizon.source.discover import discover_all_sources
+
+from .utils import (
+    parse_from_yaml,
+    set_custom_source_path_in_config,
+    set_debug_mode,
+    set_runner_in_config,
+)
+
+
+@click.group()
+def cli():
+    """Bizon CLI."""
+    pass
+
+
+# Create a 'destination' group under 'bizon'
+@cli.group()
+def source():
+    """Subcommands for handling sources."""
+    pass
+
+
+@source.command()
+def list():
+    """List available sources."""
+
+    click.echo("Retrieving available sources...")
+    sources = discover_all_sources()
+
+    click.echo("Available sources:")
+    for source_name, source_model in sources.items():
+        if not source_model.available_streams:
+            click.echo(
+                f"{source_name} - NOT AVAILABLE, run 'pip install bizon[{source_name}]' to install missing dependencies."
+            )
+        else:
+            click.echo(f"{source_name} - {source_model.available_streams}")
+
+
+# Create a 'destination' group under 'bizon'
+@cli.group()
+def stream():
+    """Subcommands for handling streams."""
+    pass
+
+
+@stream.command()
+@click.argument("source_name", type=click.STRING)
+def list(source_name: str):
+    """List available streams for a source."""
+    sources = discover_all_sources()
+    source_model = sources.get(source_name)
+    if not source_model:
+        click.echo(f"Source {source_name} not found.")
+        return
+
+    click.echo(f"Available streams for {source_name}:")
+    for stream in source_model.streams:
+        stream_mode = "[Supports incremental]" if stream.supports_incremental else "[Full refresh only]"
+        click.echo(f"{stream_mode} - {stream.name}")
+
+
+# Create a 'destination' group under 'bizon'
+@cli.group()
+def destination():
+    """Subcommands for handling destinations."""
+    pass
+
+
+@cli.command()
+@click.argument("filename", type=click.Path(exists=True))
+@click.option(
+    "--custom-source",
+    required=False,
+    type=click.Path(exists=True),
+    help="Custom Python file implementing a Bizon source.",
+)
+@click.option(
+    "--runner",
+    required=False,
+    type=click.Choice(["thread", "process"]),
+    default="thread",
+    show_default=True,
+    help="Runner type to use. Thread or Process.",
+)
+@click.option(
+    "--debug",
+    required=False,
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Enable debug mode.",
+)
+def run(filename: str, custom_source: str, runner: str, debug, 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)
+
+    # Override source_file_path param in config
+    set_custom_source_path_in_config(config=config, custom_source=ctx.get_parameter_source("custom-source"))
+
+    # Override runner param in config
+    set_runner_in_config(config=config, runner=runner)
+
+    runner = RunnerFactory.create_from_config_dict(config=config)
+    runner.run()
+
+
+if __name__ == "__main__":
+    cli()

bizon-0.0.4.dev2/bizon/cli/utils.py

@@ -0,0 +1,31 @@
+import os
+
+import yaml
+
+
+def parse_from_yaml(path_to_yaml) -> dict:
+    with open(path_to_yaml) as f:
+        config = yaml.safe_load(f)
+    return config
+
+
+def set_debug_mode(debug: bool):
+    # Set Log Level to DEBUG
+    if debug:
+        os.environ["LOGURU_LEVEL"] = "DEBUG"
+
+
+def set_custom_source_path_in_config(config: dict, custom_source: str):
+    if custom_source:
+        config["source"]["source_file_path"] = custom_source
+
+
+def set_runner_in_config(config: dict, runner: str):
+    if runner:
+        if "engine" not in config:
+            config["engine"] = {}
+
+        if "runner" not in config["engine"]:
+            config["engine"]["runner"] = {"type": runner}
+
+        config["engine"]["runner"]["type"] = runner

{bizon-0.0.3.dev1 → bizon-0.0.4.dev2}/bizon/destinations/bigquery/tests/test_bigquery_client.py

@@ -6,7 +6,7 @@ from random import randint
 import pytest
 from faker import Faker
 
-from bizon.cli.parser import parse_from_yaml
+from bizon.cli.utils import parse_from_yaml
 from bizon.destinations.bigquery.src.config import BigQueryConfig
 from bizon.destinations.bigquery.src.destination import BigQueryDestination
 from bizon.source.config import SourceConfig

{bizon-0.0.3.dev1 → bizon-0.0.4.dev2}/bizon/destinations/destination.py

@@ -154,6 +154,11 @@ class AbstractDestination(ABC):
     ) -> bool:
         """Write records to destination and update the cursor for the given iteration"""
 
+        # 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()

{bizon-0.0.3.dev1 → bizon-0.0.4.dev2}/bizon/engine/backend/adapters/sqlalchemy/backend.py

@@ -29,7 +29,8 @@ class SQLAlchemyBackend(AbstractBackend):
 
     def __init__(self, config: Union[PostgresConfigDetails, SQLiteConfigDetails], type: BackendTypes):
         super().__init__(config, type)
-        self.engine: Engine = self._get_engine()
+
+        self._engine = None
 
         self.config: Union[
             PostgresConfigDetails,
@@ -40,9 +41,10 @@ class SQLAlchemyBackend(AbstractBackend):
     @property
     def session(self) -> Generator[Session, None, None]:
         """yields a SQLAlchemy connection"""
+        engine = self.get_engine()
         session_ = scoped_session(
             sessionmaker(
-                bind=self.engine,
+                bind=engine,
                 expire_on_commit=False,
             )
         )
@@ -51,14 +53,19 @@ class SQLAlchemyBackend(AbstractBackend):
 
         session_.close()
 
+        # For SQLite, we need to dispose the engine after each operation to prevent database lock
+        self.handle_dispose_sqlite(engine)
+
     def _get_engine_bigquery(self) -> Engine:
         if hasattr(self.config, "service_account_key") and self.config.service_account_key:
             return create_engine(
-                f"bigquery://{self.config.database}/{self.config.schema}",
+                f"bigquery://{self.config.database}/{self.config.schema_name}",
                 echo=self.config.echoEngine,
                 credentials_info=self.config.service_account_key,
             )
-        return create_engine(f"bigquery://{self.config.database}/{self.config.schema}", echo=self.config.echoEngine)
+        return create_engine(
+            f"bigquery://{self.config.database}/{self.config.schema_name}", echo=self.config.echoEngine
+        )
 
     def _get_engine_postgres(self) -> Engine:
         return create_engine(
@@ -75,14 +82,15 @@ class SQLAlchemyBackend(AbstractBackend):
         if self.type == BackendTypes.POSTGRES:
             return self._get_engine_postgres()
 
-        # SQLite in a file
+        # SQLite in a file, ok for small tests
         if self.type == BackendTypes.SQLITE:
             return create_engine(
-                f"sqlite:///{self.config.database}.db",
+                f"sqlite:///{self.config.database}.sqlite3",
                 echo=self.config.echoEngine,
+                connect_args={"check_same_thread": False},
             )
 
-        # SQLite in memory
+        # ONLY FOR UNIT TESTS: SQLite in memory
         if self.type == BackendTypes.SQLITE_IN_MEMORY:
             return create_engine(
                 f"sqlite:///:memory:",
@@ -97,22 +105,44 @@ class SQLAlchemyBackend(AbstractBackend):
             logger.warning("SQLite does not support schemas")
             return True
 
-        with self.engine.connect() as connection:
-            if not inspect(connection).has_schema(self.config.schema):
+        engine = self.get_engine()
+
+        with engine.connect() as connection:
+            if not inspect(connection).has_schema(self.config.schema_name):
                 logger.error(
-                    f"Schema or dataset {self.config.schema} does not exist in the database, you need to create it first."
+                    f"Schema or dataset {self.config.schema_name} does not exist in the database, you need to create it first."
                 )
                 raise Exception(
-                    f"Schema or dataset {self.config.schema} does not exist in the database, you need to create it first."
+                    f"Schema or dataset {self.config.schema_name} does not exist in the database, you need to create it first."
                 )
 
+        self.handle_dispose_sqlite(engine)
+
+    def handle_dispose_sqlite(self, engine: Engine):
+        if self.type == BackendTypes.SQLITE:
+            engine.dispose()
+
+    def get_engine(self) -> Engine:
+        """Return the SQLAlchemy engine"""
+        if self.type == BackendTypes.SQLITE:
+            return self._get_engine()
+
+        if not self._engine:
+            self._engine = self._get_engine()
+
+        return self._engine
+
     #### INIT DATABASE ####
 
     def create_all_tables(self):
-        Base.metadata.create_all(self.engine)
+        engine = self.get_engine()
+        Base.metadata.create_all(engine)
+        self.handle_dispose_sqlite(engine)
 
     def drop_all_tables(self):
-        Base.metadata.drop_all(self.engine)
+        engine = self.get_engine()
+        Base.metadata.drop_all(engine)
+        self.handle_dispose_sqlite(engine)
 
     def check_prerequisites(self) -> bool:
         """Check if the database contains the necessary tables, return True if entities are present
@@ -125,19 +155,23 @@ class SQLAlchemyBackend(AbstractBackend):
 
         all_entities_exist = True
 
+        engine = self.get_engine()
+
         # Check if TABLE_STREAM_INFO exists, otherwise create it
-        if not inspect(self.engine).has_table(TABLE_STREAM_INFO):
+        if not inspect(engine).has_table(TABLE_STREAM_INFO):
             all_entities_exist = False
             logger.info(f"Table {TABLE_STREAM_INFO} does not exist in the database, we will create it")
 
-        if not inspect(self.engine).has_table(TABLE_SOURCE_CURSOR):
+        if not inspect(engine).has_table(TABLE_SOURCE_CURSOR):
             all_entities_exist = False
             logger.info(f"Table {TABLE_SOURCE_CURSOR} does not exist in the database, we will create it")
 
-        if not inspect(self.engine).has_table(TABLE_DESTINATION_CURSOR):
+        if not inspect(engine).has_table(TABLE_DESTINATION_CURSOR):
             all_entities_exist = False
             logger.info(f"Table {TABLE_DESTINATION_CURSOR} does not exist in the database, we will create it")
 
+        self.handle_dispose_sqlite(engine)
+
         return all_entities_exist
 
     def _add_and_commit(self, obj, session: Optional[Session] = None):
@@ -158,7 +192,7 @@ class SQLAlchemyBackend(AbstractBackend):
         source_name: str,
         stream_name: str,
         total_records_to_fetch: Optional[int] = None,
-        job_status: JobStatus = JobStatus.NOT_STARTED,
+        job_status: JobStatus = JobStatus.STARTED,
         session: Optional[Session] = None,
     ) -> StreamJob:
         """Create new StreamJob record in dbt and return its ID"""
@@ -228,7 +262,7 @@ class SQLAlchemyBackend(AbstractBackend):
         iteration: int,
         rows_fetched: int,
         next_pagination: dict,
-        cursor_status: CursorStatus = CursorStatus.NOT_STARTED,
+        cursor_status: CursorStatus = CursorStatus.STARTED,
         error_message: Optional[str] = None,
         session: Optional[Session] = None,
     ) -> str:

{bizon-0.0.3.dev1 → bizon-0.0.4.dev2}/bizon/engine/backend/adapters/sqlalchemy/config.py

@@ -55,14 +55,18 @@ class SQLiteInMemoryConfig(AbstractBackendConfig):
 
 ## BIGQUERY ##
 class BigQueryConfigDetails(SQLAlchemyConfigDetails):
+
     database: str = Field(
         description="GCP Project name",
         default=...,
     )
-    schema: str = Field(
+
+    schema_name: str = Field(
         description="BigQuery Dataset name",
         default=...,
+        alias="schema",
     )
+
     service_account_key: str = Field(
         description="Service Account Key JSON string. If empty it will be infered",
         default="",

{bizon-0.0.3.dev1 → bizon-0.0.4.dev2}/bizon/engine/backend/backend.py

@@ -35,7 +35,7 @@ class AbstractBackend(ABC):
         source_name: str,
         stream_name: str,
         total_records_to_fetch: Optional[int] = None,
-        job_status: JobStatus = JobStatus.NOT_STARTED,
+        job_status: JobStatus = JobStatus.STARTED,
         session: Optional[Session] = None,
     ) -> StreamJob:
         """Create new StreamJob record in db and return it"""
@@ -69,7 +69,7 @@ class AbstractBackend(ABC):
         iteration: int,
         rows_fetched: int,
         next_pagination: dict,
-        cursor_status: CursorStatus = CursorStatus.NOT_STARTED,
+        cursor_status: CursorStatus = CursorStatus.STARTED,
         error_message: Optional[str] = None,
         session: Optional[Session] = None,
     ) -> str:
@@ -161,8 +161,10 @@ class BackendFactory:
 
             return SQLAlchemyBackend(config=config.config, type=config.type)
 
+        # ONLY FOR UNIT TESTS
         elif config.type == BackendTypes.SQLITE_IN_MEMORY:
             from .adapters.sqlalchemy.backend import SQLAlchemyBackend
 
             return SQLAlchemyBackend(config=config.config, type=config.type)
+
         raise ValueError(f"Unsupported backend type: {config.type}")

{bizon-0.0.3.dev1 → bizon-0.0.4.dev2}/bizon/engine/backend/config.py

@@ -15,7 +15,7 @@ class BackendTypes(str, Enum):
 
 class AbstractBackendConfigDetails(BaseModel, ABC):
     database: str = Field(..., description="Database name")
-    schema: str = Field(..., description="Schema name")
+    schema_name: str = Field(..., description="Schema name", alias="schema")
     syncCursorInDBEvery: int = Field(10, description="Number of iterations before syncing the cursor in the database")
 
 

{bizon-0.0.3.dev1 → bizon-0.0.4.dev2}/bizon/engine/backend/models.py

@@ -16,7 +16,7 @@ def generate_uuid():
 
 
 class JobStatus(str, Enum):
-    NOT_STARTED = "not_started"
+    STARTED = "started"
     RUNNING = "running"
     SUCCESS = "success"
     FAILED = "failed"
@@ -25,7 +25,7 @@ class JobStatus(str, Enum):
 
 
 class CursorStatus(str, Enum):
-    NOT_STARTED = "not_started"
+    STARTED = "started"
     PULLING = "pulling"
     PULLED = "pulled"
     WRITTEN_IN_DESTINATION = "written_in_destination"
@@ -50,7 +50,7 @@ class StreamJob(Base):
         DateTime, nullable=False, default=datetime.now(tz=UTC), doc="Timestamp when the job was created"
     )
     updated_at = Column(DateTime, nullable=True, default=None, doc="Timestamp when the job was last updated")
-    status = Column(String(100), default=JobStatus.NOT_STARTED, doc="Status of the job")
+    status = Column(String(100), default=JobStatus.STARTED, doc="Status of the job")
     error_message = Column(String(255), nullable=True, doc="Error message if the job failed", default=None)
 
     source_cursor = relationship("SourceCursor", cascade="all, delete")
@@ -71,7 +71,7 @@ class SourceCursor(Base):
     rows_fetched = Column(Integer, default=0)
     next_pagination = Column(String, nullable=True)
     attempt = Column(Integer, default=0, doc="Number of attempts to pull the data for this cursor")
-    status = Column(String(100), default=CursorStatus.NOT_STARTED, doc="Status of the cursor")
+    status = Column(String(100), default=CursorStatus.STARTED, doc="Status of the cursor")
     error_message = Column(
         String(500), nullable=True, doc="Error message if pulling failed for this cursor", default=None
     )
@@ -94,7 +94,7 @@ class DestinationCursor(Base):
     to_source_iteration = Column(Integer, default=0)
     rows_written = Column(Integer, default=0)
     attempt = Column(Integer, default=0, doc="Number of attempts to pull the data for this cursor")
-    status = Column(String(100), default=CursorStatus.NOT_STARTED, doc="Status of the cursor")
+    status = Column(String(100), default=CursorStatus.STARTED, doc="Status of the cursor")
     error_message = Column(
         String(500), nullable=True, doc="Error message if pulling failed for this cursor", default=None
     )

{bizon-0.0.3.dev1 → bizon-0.0.4.dev2}/bizon/engine/config.py

@@ -19,7 +19,7 @@ from .queue.adapters.python_queue.config import (
 )
 from .queue.adapters.rabbitmq.config import RabbitMQConfig
 from .queue.config import QueueTypes
-from .runners.config import RunnerConfig, RunnerTypes, ThreadsConfig
+from .runners.config import RunnerConfig, RunnerFuturesConfig, RunnerTypes
 
 
 class EngineConfig(BaseModel):
@@ -34,10 +34,10 @@ class EngineConfig(BaseModel):
         BigQuerySQLAlchemyConfig,
     ] = Field(
         description="Configuration for the backend",
-        default=SQLiteInMemoryConfig(
-            type=BackendTypes.SQLITE_IN_MEMORY,
+        default=SQLiteSQLAlchemyConfig(
+            type=BackendTypes.SQLITE,
             config=SQLiteConfigDetails(
-                database="NOT_USED_IN_SQLITE",
+                database="bizon",
                 schema="NOT_USED_IN_SQLITE",
             ),
             syncCursorInDBEvery=2,
@@ -54,7 +54,7 @@ class EngineConfig(BaseModel):
         default=PythonQueueConfig(
             type=QueueTypes.PYTHON_QUEUE,
             config=PythonQueueConfigDetails(
-                queue=PythonQueueQueueConfig(max_size=1000),
+                queue=PythonQueueQueueConfig(max_size=0),
                 consumer=PythonQueueConsumerConfig(poll_interval=2),
             ),
         ),
@@ -64,10 +64,8 @@ class EngineConfig(BaseModel):
     runner: RunnerConfig = Field(
         description="Runner to use for the pipeline",
         default=RunnerConfig(
-            type=RunnerTypes.THREADS,
-            config=ThreadsConfig(
-                consumer_start_delay=2,
-            ),
+            type=RunnerTypes.THREAD,
+            config=RunnerFuturesConfig(),
             log_level="INFO",
         ),
     )

bizon-0.0.4.dev2/bizon/engine/engine.py

@@ -0,0 +1,31 @@
+from loguru import logger
+
+from bizon.cli.utils import parse_from_yaml
+from bizon.common.models import BizonConfig
+
+from .config import RunnerTypes
+from .runners.runner import AbstractRunner
+
+
+class RunnerFactory:
+    @staticmethod
+    def create_from_config_dict(config: dict) -> AbstractRunner:
+
+        bizon_config = BizonConfig.model_validate(obj=config)
+
+        if bizon_config.engine.runner.type == RunnerTypes.THREAD:
+            from .runners.adapters.thread import ThreadRunner
+
+            return ThreadRunner(config=config)
+
+        if bizon_config.engine.runner.type == RunnerTypes.PROCESS:
+            from .runners.adapters.process import ProcessRunner
+
+            return ProcessRunner(config=config)
+
+        raise ValueError(f"Runner type {bizon_config.engine.runner.type} is not supported")
+
+    @staticmethod
+    def create_from_yaml(filepath: str) -> AbstractRunner:
+        config = parse_from_yaml(filepath)
+        return RunnerFactory.create_from_config_dict(config)

bizon-0.0.4.dev2/bizon/engine/pipeline/consumer.py

@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+
+from bizon.destinations.destination import AbstractDestination
+from bizon.engine.queue.config import AbstractQueueConfig
+
+
+class AbstractQueueConsumer(ABC):
+    def __init__(self, config: AbstractQueueConfig, destination: AbstractDestination):
+        self.config = config
+        self.destination = destination
+
+    @abstractmethod
+    def run(self):
+        pass

bizon-0.0.4.dev2/bizon/engine/pipeline/models.py

@@ -0,0 +1,10 @@
+from enum import Enum
+
+
+class PipelineReturnStatus(Enum):
+    """Producer error types"""
+
+    SUCCESS = "success"
+    QUEUE_ERROR = "queue_error"
+    SOURCE_ERROR = "source_error"
+    BACKEND_ERROR = "backend_error"