bizon 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

bizon/alerting/alerts.py

@@ -7,7 +7,6 @@ from bizon.alerting.models import AlertingConfig, AlertMethod, LogLevel
 
 
 class AbstractAlert(ABC):
-
     def __init__(self, type: AlertMethod, config: AlertingConfig, log_levels: List[LogLevel] = [LogLevel.ERROR]):
         self.type = type
         self.config = config

bizon/common/models.py

@@ -1,9 +1,12 @@
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 from bizon.alerting.models import AlertingConfig
-from bizon.connectors.destinations.bigquery.src.config import BigQueryConfig
+from bizon.connectors.destinations.bigquery.src.config import (
+    BigQueryColumn,
+    BigQueryConfig,
+)
 from bizon.connectors.destinations.bigquery_streaming.src.config import (
     BigQueryStreamingConfig,
 )
@@ -18,8 +21,74 @@ from bizon.source.config import SourceConfig, SourceSyncModes
 from bizon.transform.config import TransformModel
 
 
-class BizonConfig(BaseModel):
+class StreamSourceConfig(BaseModel):
+    """Source-specific stream routing configuration.
+
+    Uses extra='allow' to support source-specific fields like:
+    - topic (Kafka)
+    - endpoint (API sources)
+    - channel (other streaming sources)
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    # Common field for stream identifier
+    name: Optional[str] = Field(None, description="Stream identifier within the source")
+
+    # Kafka-specific
+    topic: Optional[str] = Field(None, description="Kafka topic name")
+
+    # API-specific
+    endpoint: Optional[str] = Field(None, description="API endpoint path")
+
+
+class StreamDestinationConfig(BaseModel):
+    """Destination configuration for a stream.
+
+    Supports destination-specific schema definitions and options.
+    Uses extra='allow' to support destination-specific overrides.
+    """
 
+    model_config = ConfigDict(extra="allow")
+
+    # Universal destination identifier
+    table_id: str = Field(..., description="Full destination identifier (e.g., project.dataset.table)")
+
+    # BigQuery-specific schema (can be extended for other destinations)
+    record_schema: Optional[list[BigQueryColumn]] = Field(None, description="Schema for the destination records")
+    clustering_keys: Optional[list[str]] = Field(None, description="Clustering keys for the destination table")
+
+
+class StreamConfig(BaseModel):
+    """Configuration for a single stream.
+
+    Consolidates source stream routing and destination configuration in one place,
+    eliminating duplication of destination_id between source and destination configs.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    name: str = Field(..., description="Logical name for this stream")
+    source: StreamSourceConfig = Field(..., description="Source-specific routing configuration")
+    destination: StreamDestinationConfig = Field(
+        ..., description="Destination configuration including table and schema"
+    )
+
+    @field_validator("destination")
+    @classmethod
+    def validate_table_id_format(cls, v: StreamDestinationConfig) -> StreamDestinationConfig:
+        """Ensure table_id follows expected format for BigQuery-like destinations."""
+        if v.table_id:
+            parts = v.table_id.split(".")
+            if len(parts) != 3:
+                raise ValueError(
+                    f"table_id must be in format 'project.dataset.table', got: {v.table_id}. "
+                    f"Found {len(parts)} parts instead of 3."
+                )
+        return v
+
+
+class BizonConfig(BaseModel):
     # Forbid extra keys in the model
     model_config = ConfigDict(extra="forbid")
 
@@ -63,6 +132,115 @@ class BizonConfig(BaseModel):
         default=None,
     )
 
+    streams: Optional[list[StreamConfig]] = Field(
+        None,
+        description="Stream routing configuration (opt-in for multi-table streaming). "
+        "Consolidates source stream definitions with destination tables and schemas.",
+    )
+
+    @field_validator("streams")
+    @classmethod
+    def validate_streams_config(cls, v: Optional[list[StreamConfig]], info) -> Optional[list[StreamConfig]]:
+        """Validate streams configuration consistency."""
+        if not v:
+            return v
+
+        # Check for duplicate stream names
+        names = [s.name for s in v]
+        if len(names) != len(set(names)):
+            duplicates = [name for name in names if names.count(name) > 1]
+            raise ValueError(f"Duplicate stream names found in streams configuration: {set(duplicates)}")
+
+        # Check for duplicate table_ids
+        table_ids = [s.destination.table_id for s in v]
+        if len(table_ids) != len(set(table_ids)):
+            duplicates = [tid for tid in table_ids if table_ids.count(tid) > 1]
+            raise ValueError(f"Duplicate table_ids found in streams configuration: {set(duplicates)}")
+
+        # Validate that source sync_mode is 'stream' if streams config is used
+        source_config = info.data.get("source") if info.data else None
+        if source_config and source_config.sync_mode != SourceSyncModes.STREAM:
+            raise ValueError(
+                f"Configuration Error: 'streams' configuration requires source.sync_mode='stream'. "
+                f"Current sync_mode: {source_config.sync_mode}. "
+                f"Please update your config to use:\n"
+                f"  source:\n"
+                f"    sync_mode: stream"
+            )
+
+        return v
+
+    @model_validator(mode="before")
+    @classmethod
+    def inject_config_from_streams(cls, data: Any) -> Any:
+        """Inject source and destination config from streams.
+
+        This runs BEFORE field validation, enriching both source and destination
+        configs from the streams configuration. This allows:
+        1. Sources (like Kafka) to omit topics - they're extracted from streams
+        2. Destinations with unnest=true to work without duplicate record_schemas
+
+        This is source-agnostic: each source type can extract what it needs from streams.
+        """
+        if not isinstance(data, dict):
+            return data
+
+        streams = data.get("streams")
+        if not streams:
+            return data
+
+        source = data.get("source")
+        if source and isinstance(source, dict):
+            source_name = source.get("name")
+
+            # Kafka: inject topics from streams
+            if source_name == "kafka":
+                # Check if topics is missing, None, or empty list
+                if not source.get("topics") or source.get("topics") == []:
+                    topics = []
+                    for stream in streams:
+                        if isinstance(stream, dict):
+                            stream_src = stream.get("source", {})
+                            stream_dest = stream.get("destination", {})
+                            if stream_src.get("topic"):
+                                topics.append(
+                                    {
+                                        "name": stream_src.get("topic"),
+                                        "destination_id": stream_dest.get("table_id", ""),
+                                    }
+                                )
+                    if topics:
+                        source["topics"] = topics
+
+        destination = data.get("destination")
+        if not destination or not isinstance(destination, dict):
+            return data
+
+        destination_config = destination.get("config")
+        if not destination_config or not isinstance(destination_config, dict):
+            return data
+
+        # Only inject if record_schemas is not already set or is empty
+        if not destination_config.get("record_schemas"):
+            # Build record_schemas from streams
+            record_schemas = []
+            for stream in streams:
+                if isinstance(stream, dict):
+                    stream_dest = stream.get("destination", {})
+                    if stream_dest.get("record_schema"):
+                        record_schema_config = {
+                            "destination_id": stream_dest.get("table_id"),
+                            "record_schema": stream_dest.get("record_schema"),
+                            "clustering_keys": stream_dest.get("clustering_keys"),
+                        }
+                        record_schemas.append(record_schema_config)
+
+            # Inject into destination config
+            if record_schemas:
+                destination_config["record_schemas"] = record_schemas
+
+        return data
+
 
 class SyncMetadata(BaseModel):
     """Model which stores general metadata around a sync.
@@ -75,6 +253,7 @@ class SyncMetadata(BaseModel):
     stream_name: str
     sync_mode: SourceSyncModes
     destination_name: str
+    destination_alias: str
 
     @classmethod
     def from_bizon_config(cls, job_id: str, config: BizonConfig) -> "SyncMetadata":
@@ -85,4 +264,5 @@ class SyncMetadata(BaseModel):
             stream_name=config.source.stream,
             sync_mode=config.source.sync_mode,
             destination_name=config.destination.name,
+            destination_alias=config.destination.alias,
         )

bizon/connectors/destinations/bigquery/src/config.py

@@ -98,7 +98,6 @@ class BigQueryRecordSchemaConfig(BaseModel):
 
 
 class BigQueryConfigDetails(AbstractDestinationDetailsConfig):
-
     # Table details
     project_id: str = Field(..., description="BigQuery Project ID")
     dataset_id: str = Field(..., description="BigQuery Dataset ID")
@@ -123,5 +122,6 @@ class BigQueryConfigDetails(AbstractDestinationDetailsConfig):
 
 class BigQueryConfig(AbstractDestinationConfig):
     name: Literal[DestinationTypes.BIGQUERY]
+    alias: str = "bigquery"
     buffer_size: Optional[int] = 400
     config: BigQueryConfigDetails

bizon/connectors/destinations/bigquery/src/destination.py

@@ -14,6 +14,7 @@ from loguru import logger
 from bizon.common.models import SyncMetadata
 from bizon.destination.destination import AbstractDestination
 from bizon.engine.backend.backend import AbstractBackend
+from bizon.monitoring.monitor import AbstractMonitor
 from bizon.source.config import SourceSyncModes
 from bizon.source.source import AbstractSourceCallback
 
@@ -21,15 +22,15 @@ from .config import BigQueryColumn, BigQueryConfigDetails
 
 
 class BigQueryDestination(AbstractDestination):
-
     def __init__(
         self,
         sync_metadata: SyncMetadata,
         config: BigQueryConfigDetails,
         backend: AbstractBackend,
         source_callback: AbstractSourceCallback,
+        monitor: AbstractMonitor,
     ):
-        super().__init__(sync_metadata, config, backend, source_callback)
+        super().__init__(sync_metadata, config, backend, source_callback, monitor)
         self.config: BigQueryConfigDetails = config
 
         if config.authentication and config.authentication.service_account_key:
@@ -54,7 +55,6 @@ class BigQueryDestination(AbstractDestination):
 
     @property
     def temp_table_id(self) -> str:
-
         if self.sync_metadata.sync_mode == SourceSyncModes.FULL_REFRESH:
             return f"{self.table_id}_temp"
 
@@ -65,7 +65,6 @@ class BigQueryDestination(AbstractDestination):
             return f"{self.table_id}"
 
     def get_bigquery_schema(self, df_destination_records: pl.DataFrame) -> List[bigquery.SchemaField]:
-
         # Case we unnest the data
         if self.config.unnest:
             return [
@@ -111,9 +110,7 @@ class BigQueryDestination(AbstractDestination):
     # https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.dbapi.DataError
 
     def convert_and_upload_to_buffer(self, df_destination_records: pl.DataFrame) -> str:
-
         if self.buffer_format == "parquet":
-
             # Upload the Parquet file to GCS
             file_name = f"{self.sync_metadata.source_name}/{self.sync_metadata.stream_name}/{str(uuid4())}.parquet"
 
@@ -151,7 +148,6 @@ class BigQueryDestination(AbstractDestination):
         )
 
     def load_to_bigquery(self, gcs_file: str, df_destination_records: pl.DataFrame):
-
         # We always partition by the loaded_at field
         time_partitioning = TimePartitioning(field="_bizon_loaded_at", type_=self.config.time_partitioning)
 
@@ -169,7 +165,6 @@ class BigQueryDestination(AbstractDestination):
         assert result.state == "DONE", f"Job failed with state {result.state} with error {result.error_result}"
 
     def write_records(self, df_destination_records: pl.DataFrame) -> Tuple[bool, str]:
-
         # Rename fields to match BigQuery schema
         df_destination_records = df_destination_records.rename(
             {
@@ -199,7 +194,17 @@ class BigQueryDestination(AbstractDestination):
     def finalize(self):
         if self.sync_metadata.sync_mode == SourceSyncModes.FULL_REFRESH:
             logger.info(f"Loading temp table {self.temp_table_id} data into {self.table_id} ...")
-            self.bq_client.query(f"CREATE OR REPLACE TABLE {self.table_id} AS SELECT * FROM {self.temp_table_id}")
+            query = f"CREATE OR REPLACE TABLE {self.table_id} AS SELECT * FROM {self.temp_table_id}"
+            result = self.bq_client.query(query)
+            bq_result = result.result()  # Waits for the job to completew
+            logger.info(f"BigQuery CREATE OR REPLACE query result: {bq_result}")
+            # Check if the destination table exists by fetching it; raise if it doesn't exist
+            try:
+                self.bq_client.get_table(self.table_id)
+            except NotFound:
+                logger.error(f"Table {self.table_id} not found")
+                raise Exception(f"Table {self.table_id} not found")
+            # Cleanup
             logger.info(f"Deleting temp table {self.temp_table_id} ...")
             self.bq_client.delete_table(self.temp_table_id, not_found_ok=True)
             return True

bizon/connectors/destinations/bigquery_streaming/config/bigquery_streaming.example.yml

@@ -0,0 +1,74 @@
+# BigQuery Streaming Destination Configuration
+# Uses the BigQuery Storage Write API for low-latency inserts
+#
+# Use this destination when:
+# - You need near real-time data loading
+# - Low latency is more important than cost optimization
+# - Working with streaming/continuous data sources
+#
+# Requirements:
+# - Service account with bigquery.dataEditor role
+# - Dataset must already exist
+
+name: source_to_bigquery_streaming
+
+source:
+  name: <YOUR_SOURCE>
+  stream: <YOUR_STREAM>
+  authentication:
+    type: api_key
+    params:
+      token: <YOUR_API_KEY>
+
+destination:
+  name: bigquery_streaming
+  config:
+    # GCP Project ID
+    project_id: <YOUR_GCP_PROJECT>
+
+    # BigQuery dataset (must exist)
+    dataset_id: <YOUR_DATASET>
+
+    # Dataset location (US, EU, etc.)
+    dataset_location: US
+
+    # Time partitioning (optional)
+    time_partitioning:
+      type: DAY  # Options: DAY, HOUR, MONTH, YEAR
+      field: _bizon_loaded_at
+
+    # Max rows per streaming request (max 10000)
+    bq_max_rows_per_request: 5000
+
+    # Buffer settings
+    buffer_size: 50           # MB before flushing
+    buffer_flush_timeout: 300  # Seconds before forcing flush
+
+    # Authentication (optional - uses ADC if not provided)
+    # authentication:
+    #   service_account_key: |
+    #     {
+    #       "type": "service_account",
+    #       "project_id": "<YOUR_GCP_PROJECT>",
+    #       ...
+    #     }
+
+    # Schema definition for unnesting (optional)
+    # Required if unnest: true
+    # unnest: true
+    # record_schemas:
+    #   - destination_id: my_table
+    #     record_schema:
+    #       - name: id
+    #         type: STRING
+    #         mode: REQUIRED
+    #       - name: created_at
+    #         type: TIMESTAMP
+    #         mode: NULLABLE
+
+engine:
+  backend:
+    type: bigquery
+    database: <YOUR_GCP_PROJECT>
+    schema: bizon_state
+    syncCursorInDBEvery: 10

bizon/connectors/destinations/bigquery_streaming/src/config.py

@@ -41,16 +41,17 @@ class BigQueryStreamingConfigDetails(AbstractDestinationDetailsConfig):
         description="BigQuery Time partitioning type",
     )
     authentication: Optional[BigQueryAuthentication] = None
-    bq_max_rows_per_request: Optional[int] = Field(30000, description="Max rows per buffer streaming request.")
+    bq_max_rows_per_request: Optional[int] = Field(
+        5000,
+        description="Max rows per buffer streaming request. Must not exceed 10000.",
+        le=10000,
+    )
     record_schemas: Optional[list[BigQueryRecordSchemaConfig]] = Field(
         default=None, description="Schema for the records. Required if unnest is set to true."
     )
-    use_legacy_streaming_api: bool = Field(
-        default=False,
-        description="[DEPRECATED] Use the legacy streaming API. This is required for some older BigQuery versions.",
-    )
 
 
 class BigQueryStreamingConfig(AbstractDestinationConfig):
     name: Literal[DestinationTypes.BIGQUERY_STREAMING]
+    alias: str = "bigquery"
     config: BigQueryStreamingConfigDetails

bizon/connectors/destinations/bigquery_streaming/src/destination.py

@@ -36,15 +36,14 @@ from bizon.connectors.destinations.bigquery.src.config import (
 )
 from bizon.destination.destination import AbstractDestination
 from bizon.engine.backend.backend import AbstractBackend
+from bizon.monitoring.monitor import AbstractMonitor
 from bizon.source.callback import AbstractSourceCallback
 
 from .config import BigQueryStreamingConfigDetails
 
 
 class BigQueryStreamingDestination(AbstractDestination):
-
     # Add constants for limits
-    MAX_ROWS_PER_REQUEST = 5000  # 5000 (max is 10000)
     MAX_REQUEST_SIZE_BYTES = 5 * 1024 * 1024  # 5 MB (max is 10MB)
     MAX_ROW_SIZE_BYTES = 0.9 * 1024 * 1024  # 1 MB
 
@@ -54,8 +53,9 @@ class BigQueryStreamingDestination(AbstractDestination):
         config: BigQueryStreamingConfigDetails,
         backend: AbstractBackend,
         source_callback: AbstractSourceCallback,
+        monitor: AbstractMonitor,
     ):  # type: ignore
-        super().__init__(sync_metadata, config, backend, source_callback)
+        super().__init__(sync_metadata, config, backend, source_callback, monitor)
         self.config: BigQueryStreamingConfigDetails = config
 
         if config.authentication and config.authentication.service_account_key:
@@ -77,7 +77,6 @@ class BigQueryStreamingDestination(AbstractDestination):
         return self.destination_id or f"{self.project_id}.{self.dataset_id}.{tabled_id}"
 
     def get_bigquery_schema(self) -> List[bigquery.SchemaField]:
-
         if self.config.unnest:
             if len(list(self.record_schemas.keys())) == 1:
                 self.destination_id = list(self.record_schemas.keys())[0]
@@ -168,7 +167,6 @@ class BigQueryStreamingDestination(AbstractDestination):
         Safe cast record values to the correct type for BigQuery.
         """
         for col in self.record_schemas[self.destination_id]:
-
             # Handle dicts as strings
             if col.type in [BigQueryColumnType.STRING, BigQueryColumnType.JSON]:
                 if isinstance(row[col.name], dict) or isinstance(row[col.name], list):
@@ -222,7 +220,7 @@ class BigQueryStreamingDestination(AbstractDestination):
         try:
             # Handle streaming batch
             if batch.get("stream_batch") and len(batch["stream_batch"]) > 0:
-                return self.bq_client.insert_rows_json(
+                self.bq_client.insert_rows_json(
                     table,
                     batch["stream_batch"],
                     row_ids=[None] * len(batch["stream_batch"]),
@@ -245,6 +243,10 @@ class BigQueryStreamingDestination(AbstractDestination):
                 if load_job.state != "DONE":
                     raise Exception(f"Failed to load rows to BigQuery: {load_job.errors}")
 
+                self.monitor.track_large_records_synced(
+                    num_records=len(batch["json_batch"]), extra_tags={"destination_id": self.destination_id}
+                )
+
         except Exception as e:
             logger.error(f"Error inserting batch: {str(e)}, type: {type(e)}")
             raise
@@ -347,10 +349,12 @@ class BigQueryStreamingDestination(AbstractDestination):
 
             # If adding this item would exceed either limit, yield current batch and start new one
             if (
-                len(current_batch) >= self.MAX_ROWS_PER_REQUEST
+                len(current_batch) >= self.bq_max_rows_per_request
                 or current_batch_size + item_size > self.MAX_REQUEST_SIZE_BYTES
             ):
-                logger.debug(f"Yielding batch of {len(current_batch)} rows, size: {current_batch_size/1024/1024:.2f}MB")
+                logger.debug(
+                    f"Yielding batch of {len(current_batch)} rows, size: {current_batch_size / 1024 / 1024:.2f}MB"
+                )
                 yield {"stream_batch": current_batch, "json_batch": large_rows}
                 current_batch = []
                 current_batch_size = 0
@@ -366,7 +370,7 @@ class BigQueryStreamingDestination(AbstractDestination):
         # Yield the last batch
         if current_batch:
             logger.debug(
-                f"Yielding streaming batch of {len(current_batch)} rows, size: {current_batch_size/1024/1024:.2f}MB"
+                f"Yielding streaming batch of {len(current_batch)} rows, size: {current_batch_size / 1024 / 1024:.2f}MB"
             )
             logger.debug(f"Yielding large rows batch of {len(large_rows)} rows")
             yield {"stream_batch": current_batch, "json_batch": large_rows}

bizon/connectors/destinations/bigquery_streaming_v2/config/bigquery_streaming_v2.example.yml

@@ -0,0 +1,79 @@
+# BigQuery Streaming V2 Destination Configuration
+# Uses the BigQuery Storage Write API (v2) for improved streaming performance
+#
+# Use this destination when:
+# - You need the latest BigQuery streaming features
+# - Working with high-volume streaming data
+# - Require better error handling and retry logic
+#
+# Differences from v1:
+# - Improved batching and retry logic
+# - Better handling of schema evolution
+# - Enhanced error reporting
+#
+# Requirements:
+# - Service account with bigquery.dataEditor role
+# - Dataset must already exist
+
+name: source_to_bigquery_streaming_v2
+
+source:
+  name: <YOUR_SOURCE>
+  stream: <YOUR_STREAM>
+  authentication:
+    type: api_key
+    params:
+      token: <YOUR_API_KEY>
+
+destination:
+  name: bigquery_streaming_v2
+  config:
+    # GCP Project ID
+    project_id: <YOUR_GCP_PROJECT>
+
+    # BigQuery dataset (must exist)
+    dataset_id: <YOUR_DATASET>
+
+    # Dataset location (US, EU, etc.)
+    dataset_location: US
+
+    # Time partitioning (optional)
+    time_partitioning:
+      type: DAY  # Options: DAY, HOUR, MONTH, YEAR
+      field: _bizon_loaded_at
+
+    # Max rows per streaming request (max 10000)
+    bq_max_rows_per_request: 5000
+
+    # Buffer settings
+    buffer_size: 50           # MB before flushing
+    buffer_flush_timeout: 300  # Seconds before forcing flush
+
+    # Authentication (optional - uses ADC if not provided)
+    # authentication:
+    #   service_account_key: |
+    #     {
+    #       "type": "service_account",
+    #       "project_id": "<YOUR_GCP_PROJECT>",
+    #       ...
+    #     }
+
+    # Schema definition for unnesting (optional)
+    # Required if unnest: true
+    # unnest: true
+    # record_schemas:
+    #   - destination_id: my_table
+    #     record_schema:
+    #       - name: id
+    #         type: STRING
+    #         mode: REQUIRED
+    #       - name: created_at
+    #         type: TIMESTAMP
+    #         mode: NULLABLE
+
+engine:
+  backend:
+    type: bigquery
+    database: <YOUR_GCP_PROJECT>
+    schema: bizon_state
+    syncCursorInDBEvery: 10

bizon/connectors/destinations/bigquery_streaming_v2/src/config.py

@@ -41,7 +41,11 @@ class BigQueryStreamingV2ConfigDetails(AbstractDestinationDetailsConfig):
         description="BigQuery Time partitioning type",
     )
     authentication: Optional[BigQueryAuthentication] = None
-    bq_max_rows_per_request: Optional[int] = Field(30000, description="Max rows per buffer streaming request.")
+    bq_max_rows_per_request: Optional[int] = Field(
+        5000,
+        description="Max rows per buffer streaming request. Must not exceed 10000.",
+        le=10000,
+    )
     record_schemas: Optional[list[BigQueryRecordSchemaConfig]] = Field(
         default=None, description="Schema for the records. Required if unnest is set to true."
     )
@@ -49,4 +53,5 @@ class BigQueryStreamingV2ConfigDetails(AbstractDestinationDetailsConfig):
 
 class BigQueryStreamingV2Config(AbstractDestinationConfig):
     name: Literal[DestinationTypes.BIGQUERY_STREAMING_V2]
+    alias: str = "bigquery"
     config: BigQueryStreamingV2ConfigDetails