chore: handle Kafka connection error, add client_options in BQ

.pre-commit-config.yaml

@@ -1,6 +1,3 @@
-default_language_version:
-  python: python3.8
-
 repos:
   - repo: https://github.com/psf/black
     rev: 24.4.2

bizon/common/models.py

@@ -75,6 +75,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 +86,5 @@ def from_bizon_config(cls, job_id: str, config: BizonConfig) -> "SyncMetadata":
             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

@@ -123,5 +123,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 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
 
@@ -28,8 +29,9 @@ def __init__(
         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:

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,6 +36,7 @@
 )
 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
@@ -44,7 +45,6 @@
 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 +54,9 @@ def __init__(
         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:
@@ -222,7 +223,7 @@ def _insert_batch(self, table, batch):
         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 +246,10 @@ def _insert_batch(self, table, batch):
                 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,7 +352,7 @@ def batch(self, iterable):
 
             # 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")

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

@@ -41,12 +41,17 @@ 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."
     )
 
 
 class BigQueryStreamingV2Config(AbstractDestinationConfig):
     name: Literal[DestinationTypes.BIGQUERY_STREAMING_V2]
+    alias: str = "bigquery"
     config: BigQueryStreamingV2ConfigDetails