bizon 0.0.11__py3-none-any.whl → 0.0.13__py3-none-any.whl

bizon/destinations/bigquery/src/destination.py

@@ -2,22 +2,19 @@ import io
 import json
 import os
 import tempfile
+import traceback
 from typing import List, Tuple
 from uuid import uuid4
 
-import pandas as pd
-import pyarrow as pa
-import pyarrow.parquet as pq
+import polars as pl
 from google.api_core.exceptions import NotFound
 from google.cloud import bigquery, storage
 from google.cloud.bigquery import DatasetReference, TimePartitioning
 from loguru import logger
-from pytz import UTC
 
 from bizon.common.models import SyncMetadata
 from bizon.destinations.config import NormalizationType
 from bizon.destinations.destination import AbstractDestination
-from bizon.destinations.models import DestinationRecord
 from bizon.engine.backend.backend import AbstractBackend
 from bizon.source.config import SourceSyncModes
 
@@ -62,7 +59,7 @@ class BigQueryDestination(AbstractDestination):
         elif self.sync_metadata.sync_mode == SourceSyncModes.STREAM:
             return f"{self.table_id}"
 
-    def get_bigquery_schema(self, destination_records: List[DestinationRecord]) -> List[bigquery.SchemaField]:
+    def get_bigquery_schema(self, df_destination_records: pl.DataFrame) -> List[bigquery.SchemaField]:
 
         # we keep raw data in the column source_data
         if self.config.normalization.type == NormalizationType.NONE:
@@ -77,26 +74,13 @@ class BigQueryDestination(AbstractDestination):
                 bigquery.SchemaField("_bizon_id", "STRING", mode="REQUIRED"),
             ]
 
-        elif self.config.normalization.type == NormalizationType.DEBEZIUM:
-            assert (
-                "_bizon_message_key" in destination_records[0].source_data
-            ), "Debezium records must have a '_bizon_message_key' key"
-            message_keys = json.loads(destination_records[0].source_data["_bizon_message_key"])
-            return [bigquery.SchemaField(key, "STRING", mode="NULLABLE") for key in message_keys] + [
-                bigquery.SchemaField("_source_data", "STRING", mode="NULLABLE"),
-                bigquery.SchemaField("_source_record_id", "STRING", mode="REQUIRED"),
-                bigquery.SchemaField("_source_timestamp", "TIMESTAMP", mode="REQUIRED"),
-                bigquery.SchemaField("_bizon_extracted_at", "TIMESTAMP", mode="REQUIRED"),
-                bigquery.SchemaField(
-                    "_bizon_loaded_at", "TIMESTAMP", mode="REQUIRED", default_value_expression="CURRENT_TIMESTAMP()"
-                ),
-                bigquery.SchemaField("_bizon_id", "STRING", mode="REQUIRED"),
-            ]
-
         # If normalization is tabular, we parse key / value pairs to columns
         elif self.config.normalization.type == NormalizationType.TABULAR:
-            first_record_keys = destination_records[0].source_data.keys()
-            return [bigquery.SchemaField(key, "STRING", mode="NULLABLE") for key in first_record_keys] + [
+
+            # We use the first record to infer the schema of tabular data (key / value pairs)
+            source_data_keys = list(json.loads(df_destination_records["source_data"][0]).keys())
+
+            return [bigquery.SchemaField(key, "STRING", mode="NULLABLE") for key in source_data_keys] + [
                 bigquery.SchemaField("_source_record_id", "STRING", mode="REQUIRED"),
                 bigquery.SchemaField("_source_timestamp", "TIMESTAMP", mode="REQUIRED"),
                 bigquery.SchemaField("_bizon_extracted_at", "TIMESTAMP", mode="REQUIRED"),
@@ -108,58 +92,6 @@ class BigQueryDestination(AbstractDestination):
 
         raise NotImplementedError(f"Normalization type {self.config.normalization.type} is not supported")
 
-    def get_batch_records_as_df(self, destination_records: List[DestinationRecord]) -> pd.DataFrame:
-
-        # We keep raw data in a column -> convert the SourceRecord to a DestinationRecord
-        if self.config.normalization.type == NormalizationType.NONE:
-            df = pd.DataFrame([record.to_dict_raw_json_data(parquet=True) for record in destination_records])
-            df["_bizon_loaded_at"] = pd.Timestamp.now(tz=UTC)
-
-        # If normalization is tabular, we can just convert the data to a DataFrame parsing first-level keys
-        elif self.config.normalization.type == NormalizationType.TABULAR:
-            list_data_dict = [record.source_data for record in destination_records]
-            df = pd.DataFrame(list_data_dict).astype(str)
-            df["_bizon_id"] = [uuid4().hex for _ in range(len(destination_records))]
-
-            df["_bizon_extracted_at"] = [
-                int(record.source_timestamp.timestamp() * 1_000_000) for record in destination_records
-            ]
-
-            df["_bizon_loaded_at"] = pd.Timestamp.now(tz=UTC)
-
-            df["_source_record_id"] = [record.source_record_id for record in destination_records]
-
-            # We need to convert the source datetime to a int timestamp
-            df["_source_timestamp"] = [
-                int(record.source_timestamp.timestamp() * 1_000_000) for record in destination_records
-            ]
-
-        elif self.config.normalization.type == NormalizationType.DEBEZIUM:
-            df = pd.DataFrame([record.to_dict_debezium(parquet=True) for record in destination_records])
-            df["_bizon_loaded_at"] = pd.Timestamp.now(tz=UTC)
-
-        else:
-            raise NotImplementedError(f"Normalization type {self.config.normalization.type} is not supported")
-
-        return df
-
-    def convert_and_upload_to_buffer(self, destination_records: List[DestinationRecord]):
-
-        df = self.get_batch_records_as_df(destination_records)
-
-        # Convert DataFrame to Parquet in-memory
-        if self.buffer_format == "parquet":
-            table = pa.Table.from_pandas(df)
-            buffer = io.BytesIO()
-            pq.write_table(table, buffer)
-            buffer.seek(0)
-
-            # Upload the Parquet file to GCS
-            file_name = f"{self.sync_metadata.source_name}/{self.sync_metadata.stream_name}/{str(uuid4())}.parquet"
-            blob = self.buffer_bucket.blob(file_name)
-            blob.upload_from_file(buffer, content_type="application/octet-stream")
-            return file_name
-
     def check_connection(self) -> bool:
         dataset_ref = DatasetReference(self.project_id, self.dataset_id)
 
@@ -179,7 +111,25 @@ class BigQueryDestination(AbstractDestination):
     # https://cloud.google.com/python/docs/reference/storage/latest/retry_timeout
     # https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.dbapi.DataError
 
-    def load_to_bigquery(self, gcs_file: str, destination_records: List[DestinationRecord]):
+    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"
+
+            with io.BytesIO() as stream:
+                df_destination_records.write_parquet(stream)
+                stream.seek(0)
+
+                blob = self.buffer_bucket.blob(file_name)
+                blob.upload_from_file(stream, content_type="application/octet-stream")
+
+            return file_name
+
+        raise NotImplementedError(f"Buffer format {self.buffer_format} is not supported")
+
+    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)
@@ -187,34 +137,41 @@ class BigQueryDestination(AbstractDestination):
         job_config = bigquery.LoadJobConfig(
             source_format=bigquery.SourceFormat.PARQUET,
             write_disposition=bigquery.WriteDisposition.WRITE_APPEND,
-            schema=self.get_bigquery_schema(destination_records=destination_records),
+            schema=self.get_bigquery_schema(df_destination_records=df_destination_records),
             time_partitioning=time_partitioning,
         )
 
-        if self.config.normalization.type == NormalizationType.DEBEZIUM:
-            job_config.clustering_fields = list(
-                json.loads(destination_records[0].source_data["_bizon_message_key"]).keys()
-            )
-
         load_job = self.bq_client.load_table_from_uri(
             f"gs://{self.buffer_bucket_name}/{gcs_file}", self.temp_table_id, job_config=job_config
         )
+        result = load_job.result()  # Waits for the job to complete
+        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(
+            {
+                # Bizon fields
+                "bizon_extracted_at": "_bizon_extracted_at",
+                "bizon_id": "_bizon_id",
+                "bizon_loaded_at": "_bizon_loaded_at",
+                # Source fields
+                "source_record_id": "_source_record_id",
+                "source_timestamp": "_source_timestamp",
+                "source_data": "_source_data",
+            },
+        )
 
-        load_job.result()
-
-    def write_records(self, destination_records: List[DestinationRecord]) -> Tuple[bool, str]:
-
-        # Here we can check if these IDs are already present in BigQuery
-        # Using SourceRecord.id values
-
-        gs_file_name = self.convert_and_upload_to_buffer(destination_records=destination_records)
+        gs_file_name = self.convert_and_upload_to_buffer(df_destination_records=df_destination_records)
 
         try:
-            self.load_to_bigquery(gs_file_name, destination_records=destination_records)
+            self.load_to_bigquery(gcs_file=gs_file_name, df_destination_records=df_destination_records)
             self.cleanup(gs_file_name)
         except Exception as e:
             self.cleanup(gs_file_name)
             logger.error(f"Error loading data to BigQuery: {e}")
+            logger.error(traceback.format_exc())
             return False, str(e)
         return True, ""
 

bizon/destinations/bigquery_streaming/src/config.py

@@ -10,11 +10,6 @@ from bizon.destinations.config import (
 )
 
 
-class GCSBufferFormat(str, Enum):
-    PARQUET = "parquet"
-    CSV = "csv"
-
-
 class TimePartitioning(str, Enum):
     DAY = "DAY"
     HOUR = "HOUR"
@@ -29,7 +24,7 @@ class BigQueryAuthentication(BaseModel):
     )
 
 
-class BigQueryConfigDetails(AbstractDestinationDetailsConfig):
+class BigQueryStreamingConfigDetails(AbstractDestinationDetailsConfig):
     project_id: str
     dataset_id: str
     dataset_location: Optional[str] = "US"
@@ -40,16 +35,9 @@ class BigQueryConfigDetails(AbstractDestinationDetailsConfig):
         default=TimePartitioning.DAY, description="BigQuery Time partitioning type"
     )
     authentication: Optional[BigQueryAuthentication] = None
-
-    buffer_size: int = Field(default=0, description="Buffer size in MB")
-
-    @field_validator("buffer_size", mode="after")
-    def validate_buffer_size(cls, value: int) -> int:
-        if value != 0:
-            raise ValueError("Buffer size must be 0, we directly stream to BigQuery")
-        return value
+    bq_max_rows_per_request: Optional[int] = Field(30000, description="Max rows per buffer streaming request.")
 
 
 class BigQueryStreamingConfig(AbstractDestinationConfig):
     name: Literal[DestinationTypes.BIGQUERY_STREAMING]
-    config: BigQueryConfigDetails
+    config: BigQueryStreamingConfigDetails

bizon/destinations/bigquery_streaming/src/destination.py

@@ -1,29 +1,33 @@
-import json
 import os
 import tempfile
-from typing import List, Tuple
+from concurrent.futures import ThreadPoolExecutor
+from typing import List, Tuple, Type
 
+import polars as pl
 from google.api_core.exceptions import NotFound
-from google.cloud import bigquery, bigquery_storage_v1, storage
+from google.cloud import bigquery, bigquery_storage_v1
 from google.cloud.bigquery import DatasetReference, TimePartitioning
-from google.cloud.bigquery_storage_v1.types import AppendRowsRequest, ProtoRows
-from loguru import logger
+from google.cloud.bigquery_storage_v1.types import (
+    AppendRowsRequest,
+    ProtoRows,
+    ProtoSchema,
+)
+from google.protobuf.message import Message
 
 from bizon.common.models import SyncMetadata
 from bizon.destinations.config import NormalizationType
 from bizon.destinations.destination import AbstractDestination
-from bizon.destinations.models import DestinationRecord
 from bizon.engine.backend.backend import AbstractBackend
 
-from .config import BigQueryConfigDetails
+from .config import BigQueryStreamingConfigDetails
 from .proto_utils import get_proto_schema_and_class
 
 
 class BigQueryStreamingDestination(AbstractDestination):
 
-    def __init__(self, sync_metadata: SyncMetadata, config: BigQueryConfigDetails, backend: AbstractBackend):
+    def __init__(self, sync_metadata: SyncMetadata, config: BigQueryStreamingConfigDetails, backend: AbstractBackend):
         super().__init__(sync_metadata, config, backend)
-        self.config: BigQueryConfigDetails = config
+        self.config: BigQueryStreamingConfigDetails = config
 
         if config.authentication and config.authentication.service_account_key:
             with tempfile.NamedTemporaryFile(delete=False) as temp:
@@ -34,16 +38,16 @@ class BigQueryStreamingDestination(AbstractDestination):
         self.project_id = config.project_id
         self.bq_client = bigquery.Client(project=self.project_id)
         self.bq_storage_client = bigquery_storage_v1.BigQueryWriteClient()
-        self.gcs_client = storage.Client(project=self.project_id)
         self.dataset_id = config.dataset_id
         self.dataset_location = config.dataset_location
+        self.bq_max_rows_per_request = config.bq_max_rows_per_request
 
     @property
     def table_id(self) -> str:
         tabled_id = self.config.table_id or f"{self.sync_metadata.source_name}_{self.sync_metadata.stream_name}"
         return f"{self.project_id}.{self.dataset_id}.{tabled_id}"
 
-    def get_bigquery_schema(self, destination_records: List[DestinationRecord]) -> List[bigquery.SchemaField]:
+    def get_bigquery_schema(self) -> List[bigquery.SchemaField]:
 
         # we keep raw data in the column source_data
         if self.config.normalization.type == NormalizationType.NONE:
@@ -58,35 +62,6 @@ class BigQueryStreamingDestination(AbstractDestination):
                 bigquery.SchemaField("_bizon_id", "STRING", mode="REQUIRED"),
             ]
 
-        elif self.config.normalization.type == NormalizationType.DEBEZIUM:
-            assert (
-                "_bizon_message_key" in destination_records[0].source_data
-            ), "Debezium records must have a '_bizon_message_key' key"
-            message_keys = json.loads(destination_records[0].source_data["_bizon_message_key"])
-            return [bigquery.SchemaField(key, "STRING", mode="NULLABLE") for key in message_keys] + [
-                bigquery.SchemaField("_source_data", "STRING", mode="NULLABLE"),
-                bigquery.SchemaField("_source_record_id", "STRING", mode="REQUIRED"),
-                bigquery.SchemaField("_source_timestamp", "TIMESTAMP", mode="REQUIRED"),
-                bigquery.SchemaField("_bizon_extracted_at", "TIMESTAMP", mode="REQUIRED"),
-                bigquery.SchemaField(
-                    "_bizon_loaded_at", "TIMESTAMP", mode="REQUIRED", default_value_expression="CURRENT_TIMESTAMP()"
-                ),
-                bigquery.SchemaField("_bizon_id", "STRING", mode="REQUIRED"),
-            ]
-
-        # If normalization is tabular, we parse key / value pairs to columns
-        elif self.config.normalization.type == NormalizationType.TABULAR:
-            first_record_keys = destination_records[0].source_data.keys()
-            return [bigquery.SchemaField(key, "STRING", mode="NULLABLE") for key in first_record_keys] + [
-                bigquery.SchemaField("_source_record_id", "STRING", mode="REQUIRED"),
-                bigquery.SchemaField("_source_timestamp", "TIMESTAMP", mode="REQUIRED"),
-                bigquery.SchemaField("_bizon_extracted_at", "TIMESTAMP", mode="REQUIRED"),
-                bigquery.SchemaField(
-                    "_bizon_loaded_at", "TIMESTAMP", mode="REQUIRED", default_value_expression="CURRENT_TIMESTAMP()"
-                ),
-                bigquery.SchemaField("_bizon_id", "STRING", mode="REQUIRED"),
-            ]
-
         raise NotImplementedError(f"Normalization type {self.config.normalization.type} is not supported")
 
     def check_connection(self) -> bool:
@@ -100,21 +75,45 @@ class BigQueryStreamingDestination(AbstractDestination):
             dataset = self.bq_client.create_dataset(dataset)
         return True
 
-    def load_to_bigquery_via_streaming(self, destination_records: List[DestinationRecord]) -> str:
+    def append_rows_to_stream(
+        self,
+        write_client: bigquery_storage_v1.BigQueryWriteClient,
+        stream_name: str,
+        proto_schema: ProtoSchema,
+        serialized_rows: List[bytes],
+    ):
+        request = AppendRowsRequest(
+            write_stream=stream_name,
+            proto_rows=AppendRowsRequest.ProtoData(
+                rows=ProtoRows(serialized_rows=serialized_rows),
+                writer_schema=proto_schema,
+            ),
+        )
+        response = write_client.append_rows(iter([request]))
+        return response.code().name
+
+    @staticmethod
+    def to_protobuf_serialization(TableRowClass: Type[Message], row: dict) -> bytes:
+        """Convert a row to a protobuf serialization"""
+        record = TableRowClass()
+        record._bizon_id = row["bizon_id"]
+        record._bizon_extracted_at = row["bizon_extracted_at"].strftime("%Y-%m-%d %H:%M:%S.%f")
+        record._bizon_loaded_at = row["bizon_loaded_at"].strftime("%Y-%m-%d %H:%M:%S.%f")
+        record._source_record_id = row["source_record_id"]
+        record._source_timestamp = row["source_timestamp"].strftime("%Y-%m-%d %H:%M:%S.%f")
+        record._source_data = row["source_data"]
+        return record.SerializeToString()
+
+    def load_to_bigquery_via_streaming(self, df_destination_records: pl.DataFrame) -> str:
+        # TODO: for now no clustering keys
         clustering_keys = []
 
-        if self.config.normalization.type == NormalizationType.DEBEZIUM:
-            clustering_keys = list(json.loads(destination_records[0].source_data["_bizon_message_key"]).keys())
-
         # Create table if it doesnt exist
-        schema = self.get_bigquery_schema(destination_records=destination_records)
+        schema = self.get_bigquery_schema()
         table = bigquery.Table(self.table_id, schema=schema)
         time_partitioning = TimePartitioning(field="_bizon_loaded_at", type_=self.config.time_partitioning)
         table.time_partitioning = time_partitioning
 
-        if clustering_keys:
-            table.clustering_fields = clustering_keys
-
         table = self.bq_client.create_table(table, exists_ok=True)
 
         # Create the stream
@@ -127,22 +126,29 @@ class BigQueryStreamingDestination(AbstractDestination):
         proto_schema, TableRow = get_proto_schema_and_class(clustering_keys)
 
         serialized_rows = [
-            record.to_protobuf_serialization(
-                TableRow, debezium=self.config.normalization.type == NormalizationType.DEBEZIUM
-            )
-            for record in destination_records
+            self.to_protobuf_serialization(TableRowClass=TableRow, row=row)
+            for row in df_destination_records.iter_rows(named=True)
         ]
 
-        request = AppendRowsRequest(
-            write_stream=stream_name,
-            proto_rows=AppendRowsRequest.ProtoData(
-                rows=ProtoRows(serialized_rows=serialized_rows),
-                writer_schema=proto_schema,
-            ),
-        )
-        response = write_client.append_rows(iter([request]))
-        assert response.code().name == "OK"
+        results = []
+        with ThreadPoolExecutor() as executor:
+            futures = [
+                executor.submit(self.append_rows_to_stream, write_client, stream_name, proto_schema, batch_rows)
+                for batch_rows in self.batch(serialized_rows)
+            ]
+            for future in futures:
+                results.append(future.result())
+
+        assert all([r == "OK" for r in results]) is True, "Failed to append rows to stream"
 
-    def write_records(self, destination_records: List[DestinationRecord]) -> Tuple[bool, str]:
-        self.load_to_bigquery_via_streaming(destination_records=destination_records)
+    def write_records(self, df_destination_records: pl.DataFrame) -> Tuple[bool, str]:
+        self.load_to_bigquery_via_streaming(df_destination_records=df_destination_records)
         return True, ""
+
+    def batch(self, iterable):
+        """
+        Yield successive batches of size `batch_size` from `iterable`.
+        """
+
+        for i in range(0, len(iterable), self.bq_max_rows_per_request):
+            yield iterable[i : i + self.bq_max_rows_per_request]  # noqa

bizon/destinations/buffer.py

@@ -1,8 +1,11 @@
-import sys
 from datetime import datetime
 from typing import List
 
-from bizon.destinations.models import DestinationRecord
+from loguru import logger
+from polars import DataFrame
+from pytz import UTC
+
+from .models import destination_record_schema
 
 
 class DestinationBuffer:
@@ -10,15 +13,15 @@ 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.df_destination_records: DataFrame = DataFrame(schema=destination_record_schema)
         self._iterations: List[int] = []
         self.pagination = {}
-        self.modified_at: List[datetime] = [datetime.utcnow()]
+        self.modified_at: List[datetime] = [datetime.now(tz=UTC)]
 
     @property
     def current_size(self) -> int:
         """Return buffer size"""
-        return sys.getsizeof(self.records)
+        return self.df_destination_records.estimated_size(unit="b")
 
     @property
     def buffer_free_space_pct(self) -> float:
@@ -61,16 +64,20 @@ class DestinationBuffer:
 
     def flush(self):
         """Flush buffer"""
-        self.records = []
+        self.df_destination_records = DataFrame(schema=destination_record_schema)
         self._iterations = []
         self.pagination = {}
         self.modified_at = []
 
     def add_source_iteration_records_to_buffer(
-        self, iteration: int, records: List[DestinationRecord], pagination: dict = None
+        self, iteration: int, df_destination_records: DataFrame, pagination: dict = None
     ):
         """Add records for the given iteration to buffer"""
-        self.records.extend(records)
+        self.df_destination_records.vstack(df_destination_records, in_place=True)
         self._iterations.append(iteration)
         self.pagination = pagination
-        self.modified_at.append(datetime.utcnow())
+        self.modified_at.append(datetime.now(tz=UTC))
+
+        logger.info(
+            f"Added {df_destination_records.height} records to buffer for iteration {iteration} - {self.df_destination_records.estimated_size(unit='mb')} MB"
+        )