bizon 0.0.8__py3-none-any.whl → 0.0.9__py3-none-any.whl

bizon/destinations/bigquery/src/config.py

@@ -15,6 +15,13 @@ class GCSBufferFormat(str, Enum):
     CSV = "csv"
 
 
+class TimePartitioning(str, Enum):
+    DAY = "DAY"
+    HOUR = "HOUR"
+    MONTH = "MONTH"
+    YEAR = "YEAR"
+
+
 class BigQueryAuthentication(BaseModel):
     service_account_key: str = Field(
         description="Service Account Key JSON string. If empty it will be infered",
@@ -31,6 +38,10 @@ class BigQueryConfigDetails(AbstractDestinationDetailsConfig):
     )
     gcs_buffer_bucket: str
     gcs_buffer_format: Optional[GCSBufferFormat] = GCSBufferFormat.PARQUET
+
+    time_partitioning: Optional[TimePartitioning] = Field(
+        default=TimePartitioning.DAY, description="BigQuery Time partitioning type"
+    )
     authentication: Optional[BigQueryAuthentication] = None
 
 

bizon/destinations/bigquery/src/destination.py

@@ -1,4 +1,5 @@
 import io
+import json
 import os
 import tempfile
 from typing import List, Tuple
@@ -9,7 +10,7 @@ import pyarrow as pa
 import pyarrow.parquet as pq
 from google.api_core.exceptions import NotFound
 from google.cloud import bigquery, storage
-from google.cloud.bigquery import DatasetReference
+from google.cloud.bigquery import DatasetReference, TimePartitioning
 from loguru import logger
 from pytz import UTC
 
@@ -70,7 +71,25 @@ class BigQueryDestination(AbstractDestination):
                 bigquery.SchemaField("_source_timestamp", "TIMESTAMP", mode="REQUIRED"),
                 bigquery.SchemaField("_source_data", "STRING", mode="NULLABLE"),
                 bigquery.SchemaField("_bizon_extracted_at", "TIMESTAMP", mode="REQUIRED"),
-                bigquery.SchemaField("_bizon_loaded_at", "TIMESTAMP", mode="REQUIRED"),
+                bigquery.SchemaField(
+                    "_bizon_loaded_at", "TIMESTAMP", mode="REQUIRED", default_value_expression="CURRENT_TIMESTAMP()"
+                ),
+                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"),
             ]
 
@@ -81,7 +100,9 @@ class BigQueryDestination(AbstractDestination):
                 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"),
+                bigquery.SchemaField(
+                    "_bizon_loaded_at", "TIMESTAMP", mode="REQUIRED", default_value_expression="CURRENT_TIMESTAMP()"
+                ),
                 bigquery.SchemaField("_bizon_id", "STRING", mode="REQUIRED"),
             ]
 
@@ -113,6 +134,10 @@ class BigQueryDestination(AbstractDestination):
                 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")
 
@@ -155,15 +180,26 @@ class BigQueryDestination(AbstractDestination):
     # 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]):
+
+        # We always partition by the loaded_at field
+        time_partitioning = TimePartitioning(field="_bizon_loaded_at", type_=self.config.time_partitioning)
+
         job_config = bigquery.LoadJobConfig(
             source_format=bigquery.SourceFormat.PARQUET,
             write_disposition=bigquery.WriteDisposition.WRITE_APPEND,
             schema=self.get_bigquery_schema(destination_records=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
         )
+
         load_job.result()
 
     def write_records(self, destination_records: List[DestinationRecord]) -> Tuple[bool, str]:

bizon/destinations/config.py

@@ -13,6 +13,7 @@ class DestinationTypes(str, Enum):
 class NormalizationType(str, Enum):
     TABULAR = "tabular"  # Parse key / value pairs to columns
     NONE = "none"  # No normalization, raw data is stored
+    DEBEZIUM = "debezium"  # Debezium normalization
 
 
 class NormalizationConfig(BaseModel):

bizon/destinations/models.py

@@ -27,6 +27,39 @@ class DestinationRecord(BaseModel):
             source_data=source_record.data,
         )
 
+    def to_dict_debezium(self, parquet: bool = False) -> dict:
+        """Return the record as a dict with Debezium data"""
+
+        # Extract keys from Debezium message key and unnest
+        parsed_debezium_keys = json.loads(self.source_data["_bizon_message_key"])
+
+        # Parse Debezium Operation and deleted record
+        if self.source_data.get("op") == "d":
+            parsed_source_data = {"__deleted": True, **self.source_data["before"]}
+        else:
+            parsed_source_data = {"__deleted": False, **self.source_data["after"]}
+
+        if parquet:
+            return {
+                **{k: str(v) for k, v in parsed_debezium_keys.items()},
+                "_bizon_id": self.bizon_id,
+                "_bizon_extracted_at": int(self.bizon_extracted_at.timestamp() * 1_000_000),
+                "_bizon_loaded_at": self.bizon_loaded_at.timestamp(),
+                "_source_record_id": self.source_record_id,
+                "_source_timestamp": int(self.source_timestamp.timestamp() * 1_000_000),
+                "_source_data": json.dumps(parsed_source_data),
+            }
+
+        return {
+            **{k: str(v) for k, v in parsed_debezium_keys.items()},
+            "_bizon_id": self.bizon_id,
+            "_bizon_extracted_at": self.bizon_extracted_at,
+            "_bizon_loaded_at": self.bizon_loaded_at,
+            "_source_record_id": self.source_record_id,
+            "_source_timestamp": self.source_timestamp,
+            "_source_data": json.dumps(parsed_source_data),
+        }
+
     def to_dict_raw_json_data(self, parquet: bool = False) -> str:
         """Return the record as a dict with raw JSON data"""
 

bizon/sources/kafka/src/source.py

@@ -45,6 +45,8 @@ class KafkaSourceConfig(SourceConfig):
     consumer_timeout: int = Field(10, description="Kafka consumer timeout in seconds")
     group_id: str = Field("bizon", description="Kafka group id")
 
+    max_consumer_threads: int = Field(16, description="Maximum number of threads for the consumer")
+
     nb_bytes_schema_id: Literal[4, 8] = Field(
         4, description="Number of bytes for the schema id. 4 is the default for majority of the cases"
     )
@@ -223,6 +225,7 @@ class KafkaSource(AbstractSource):
                     raise ValueError(f"Number of bytes for schema id {self.config.nb_bytes_schema_id} not supported")
 
                 data = self.decode(message.value(), schema)
+                data["_bizon_message_key"] = message.key().decode("utf-8")
 
                 # Get the source timestamp
                 if self.parse_timestamp:
@@ -261,7 +264,7 @@ class KafkaSource(AbstractSource):
 
         # Use ThreadPoolExecutor to parallelize reading partitions
         records = []
-        with ThreadPoolExecutor(max_workers=nb_partitions) as executor:
+        with ThreadPoolExecutor(max_workers=min(nb_partitions, self.config.max_consumer_threads)) as executor:
             futures = {executor.submit(self.read_partition, i, topic_offsets): i for i in range(nb_partitions)}
             for future in as_completed(futures):
                 partition_records = future.result()

bizon/sources/kafka/tests/kafka_pipeline.py

@@ -3,5 +3,7 @@ import os
 from bizon.engine.engine import RunnerFactory
 
 if __name__ == "__main__":
-    runner = RunnerFactory.create_from_yaml(filepath=os.path.abspath("bizon/sources/kafka/config/kafka.yml"))
+    runner = RunnerFactory.create_from_yaml(
+        filepath=os.path.abspath("bizon/sources/kafka/config/kafka_teams_users_eu_west1_c511.yml")
+    )
     runner.run()

{bizon-0.0.8.dist-info → bizon-0.0.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: bizon
-Version: 0.0.8
+Version: 0.0.9
 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

{bizon-0.0.8.dist-info → bizon-0.0.9.dist-info}/RECORD

@@ -6,16 +6,16 @@ bizon/common/errors/backoff.py,sha256=z7RkQt1Npdh0sfD3hBDaiWQKe4iqS6ewvT1Q4Fds5a
 bizon/common/errors/errors.py,sha256=mrYx1uE2kOuR2pEaB7ztK1l2m0E4V-_-hxq-DuILerY,682
 bizon/common/models.py,sha256=7_HKAxOyN9eK8hmqahzHhmK-TYVAuRtGOgf4iadE7FI,1751
 bizon/destinations/bigquery/config/bigquery.example.yml,sha256=mvKtFS_PUuekyMh9xssuwRfFwLtR-rVvpIy5xmF5__k,1261
-bizon/destinations/bigquery/src/config.py,sha256=HU04yJhFxnSbKi0R3B-xZCPT8CZ3l8KAUm-vRsFAT8w,1119
-bizon/destinations/bigquery/src/destination.py,sha256=m0lCvAuyM4FhbS29n4OuJI3g_7P4X_3d-erDrvb7igY,9138
+bizon/destinations/bigquery/src/config.py,sha256=QlD-FdBJ8Q6nKPrOf5q28lHnyFE8khT41dSR1s2meeM,1378
+bizon/destinations/bigquery/src/destination.py,sha256=tPxE0IpHbR4zDkW5HaiHkgeDRDY2AibIPzY9iftZ2Uc,11079
 bizon/destinations/buffer.py,sha256=bFYkaoge-3AyKfGolqsuB3PWWtdPt65Fllrz-3X_uMI,2594
-bizon/destinations/config.py,sha256=qncIW-rVJdViyNqYW9seBoBK0XMAbn45clC4CcHWrQw,1634
+bizon/destinations/config.py,sha256=jD4nkG-sg7mzJMFKLErQBkJu7ri0PMbCRVU3xIvFT7E,1686
 bizon/destinations/destination.py,sha256=VAyGPmowNimvK_joZj-6ESk2ezGxDZHnKCIpKRA-Vus,10995
 bizon/destinations/file/src/config.py,sha256=C4BBIKzBH5343iLGR3aCubAGjPo0b2LegsCLjb77uFA,513
 bizon/destinations/file/src/destination.py,sha256=1VCrVdtzAzwSKgYq0JUOc3r2cM7314dV-eIoAFhM_64,1003
 bizon/destinations/logger/src/config.py,sha256=AWY3R9q3ZjD3uQ_KBq8VcW60deKSIHe3qtgCKjdywKk,433
 bizon/destinations/logger/src/destination.py,sha256=xTt03F3AMI9KhQno2tGoCr3eacrO62qjnOlpeEHk6tQ,868
-bizon/destinations/models.py,sha256=FJw32LBF9mcniFDhss5lp_UT1cENTwRapaP6iHSUhD8,2111
+bizon/destinations/models.py,sha256=hK7yXMoOArLJ5sUS9kgljXMBaq2vqu1l_7u707yS1KM,3630
 bizon/engine/backend/adapters/sqlalchemy/backend.py,sha256=R0CztRGc3_6PdIIgbbrDYD2OJRNhq9PPmD6PYK7-fjk,15567
 bizon/engine/backend/adapters/sqlalchemy/config.py,sha256=K-FpE_-VHnTSAQOduouhXFVy43EkrKbeZLqr9_OfeMw,1846
 bizon/engine/backend/backend.py,sha256=Bodqoo5qJHV0H2zJJeGytaHGiNZmBjnLBxiRgq6M3kE,5844
@@ -76,16 +76,16 @@ bizon/sources/hubspot/src/hubspot_objects.py,sha256=EmABx9XD8q6g4Uc5mHLv5YYl5KcI
 bizon/sources/hubspot/src/models/hs_object.py,sha256=-Y20H3-nenJyySMlvM4TPttPz4O8qm3ArKP_I8pxsuo,1235
 bizon/sources/hubspot/tests/hubspot_pipeline.py,sha256=e6dCF5_MHMySkeiF6kKrSAuCa_48J22-ZeSCZSjrfUI,216
 bizon/sources/kafka/config/kafka.example.yml,sha256=ZyHBmSWZ_5WQaBr9WzD05PuE6vi3hhYgHh2VZ-IU-Iw,755
-bizon/sources/kafka/src/source.py,sha256=jbyk19RpXSS6CPbEhXa178OzcD-izk4biMr1EiOKBzk,11128
-bizon/sources/kafka/tests/kafka_pipeline.py,sha256=G22IyaStejMhh5zwyyKmLDBpzdigo1uT6qsIxcqVWpg,212
+bizon/sources/kafka/src/source.py,sha256=28Cn_m8DOzsUdgbq0sUm36I7hB0TWRF6xEzg7TcrPrc,11343
+bizon/sources/kafka/tests/kafka_pipeline.py,sha256=DrMHq96ZDiQ2lWmxEf_aX7HmBg_qNOsSFGTuGmuhly8,252
 bizon/sources/periscope/config/periscope_charts.example.yml,sha256=rpFDAWeU5oZ3UOiX0sSAgd1X5lv6t-s3iqiDPnRqutU,477
 bizon/sources/periscope/config/periscope_dashboards.example.yml,sha256=sN2iGGqCQCvrMXcwxNGq_dR7-KZ1KtYdXmNYKXlfEpg,481
 bizon/sources/periscope/src/source.py,sha256=AZM-HDDjdTWj8akeeofQ_-G8YlnNHEKi2mjEQSYwOvE,7638
 bizon/sources/periscope/tests/periscope_pipeline_charts.py,sha256=mU0JtfhS1KmWsS3iovGhGxK7iPVWiYzjBM_QfRL3ZQI,275
 bizon/sources/periscope/tests/periscope_pipeline_dashboard.py,sha256=vZKN7UfH-lQIWrnfjPqQFjZm28UIw2m9OSg4yS-Wckk,279
 bizon/utils.py,sha256=HXaPiyxpWKoy3XN5vSYOve1ezlFeOYin3aFqTjcabUQ,81
-bizon-0.0.8.dist-info/LICENSE,sha256=AW7SjYVT2bBnXOxgDxqy_e_JF8jDCFlMCaPCF11wFDI,1072
-bizon-0.0.8.dist-info/METADATA,sha256=XgDIPAHHLNd2__CX0K9gO1NJwg_UWsgYTEc073-6Uiw,5646
-bizon-0.0.8.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-bizon-0.0.8.dist-info/entry_points.txt,sha256=wtCd-6JswSY8lPWYSvOf7ASX1zfKgmgXtgg5XQS5274,44
-bizon-0.0.8.dist-info/RECORD,,
+bizon-0.0.9.dist-info/LICENSE,sha256=AW7SjYVT2bBnXOxgDxqy_e_JF8jDCFlMCaPCF11wFDI,1072
+bizon-0.0.9.dist-info/METADATA,sha256=POEtr3jEzvy8ogs2WvJ0rnlcpqFXcnZjLwQebLWxNnw,5646
+bizon-0.0.9.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+bizon-0.0.9.dist-info/entry_points.txt,sha256=wtCd-6JswSY8lPWYSvOf7ASX1zfKgmgXtgg5XQS5274,44
+bizon-0.0.9.dist-info/RECORD,,