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

bizon/connectors/sources/kafka/src/source.py

@@ -1,7 +1,8 @@
 import traceback
+from collections.abc import Mapping
 from datetime import datetime
-from functools import lru_cache
-from typing import Any, List, Mapping, Tuple
+from functools import cache
+from typing import Any, List, Tuple
 
 import orjson
 from avro.schema import Schema, parse
@@ -12,6 +13,7 @@ from confluent_kafka import (
     Message,
     TopicPartition,
 )
+from confluent_kafka.cimpl import KafkaException as CimplKafkaException
 from loguru import logger
 from pydantic import BaseModel
 from pytz import UTC
@@ -27,12 +29,13 @@ from .config import KafkaSourceConfig, MessageEncoding, SchemaRegistryType
 from .decode import (
     Hashabledict,
     decode_avro_message,
-    get_header_bytes,
     parse_global_id_from_serialized_message,
 )
 
 
 class SchemaNotFound(Exception):
+    """Schema not found in the Schema Registry"""
+
     pass
 
 
@@ -57,13 +60,25 @@ class TopicOffsets(BaseModel):
         return sum([partition.last for partition in self.partitions.values()])
 
 
-class KafkaSource(AbstractSource):
+def on_error(err: KafkaError):
+    # Fires for client-level errors (incl. DNS resolve failures)
+    if err.fatal():
+        logger.error(f"Kafka client error: {err} | fatal={err.fatal()} retriable={err.retriable()}")
+        raise KafkaException(err)
+    else:
+        logger.warning(f"Kafka client error: {err} | fatal={err.fatal()} retriable={err.retriable()}")
+
 
+class KafkaSource(AbstractSource):
     def __init__(self, config: KafkaSourceConfig):
         super().__init__(config)
 
         self.config: KafkaSourceConfig = config
 
+        # Ensure topics is always a list (not None)
+        if self.config.topics is None:
+            self.config.topics = []
+
         # Kafka consumer configuration.
         if self.config.authentication.type == AuthType.BASIC:
             self.config.consumer_config["sasl.mechanisms"] = "PLAIN"
@@ -74,11 +89,58 @@ class KafkaSource(AbstractSource):
         self.config.consumer_config["group.id"] = self.config.group_id
         self.config.consumer_config["bootstrap.servers"] = self.config.bootstrap_servers
 
+        # Set the error callback
+        self.config.consumer_config["error_cb"] = on_error
+
         # Consumer instance
-        self.consumer = Consumer(self.config.consumer_config, logger=logger)
+        self.consumer = Consumer(self.config.consumer_config)
 
+        # Map topic_name to destination_id
         self.topic_map = {topic.name: topic.destination_id for topic in self.config.topics}
 
+    def set_streams_config(self, streams: list) -> None:
+        """Configure Kafka topics from streams config.
+
+        This method enriches self.config.topics from the streams configuration,
+        ensuring that subsequent source instantiations (e.g., in init_job) have
+        access to the topics without duplication in the YAML config.
+
+        When a top-level 'streams' configuration is present, this method:
+        1. Extracts Kafka topics from streams (topic field)
+        2. Builds TopicConfig objects with destination_id from streams
+        3. Populates self.config.topics if empty (modifies bizon_config.source in-place)
+        4. Updates topic_map for record routing
+
+        Args:
+            streams: List of StreamConfig objects from BizonConfig.streams
+        """
+        from .config import TopicConfig
+
+        # Extract topics from streams
+        topics_from_streams = []
+        streams_map = {}
+
+        for stream in streams:
+            if hasattr(stream.source, "topic") and stream.source.topic:
+                topic_name = stream.source.topic
+                streams_map[topic_name] = stream
+
+                # Build TopicConfig from stream
+                topic_config = TopicConfig(name=topic_name, destination_id=stream.destination.table_id)
+                topics_from_streams.append(topic_config)
+
+        # Populate self.config.topics from streams (modifies bizon_config.source in-place)
+        # This ensures check_connection() and subsequent source instantiations have topics
+        if not self.config.topics and topics_from_streams:
+            self.config.topics = topics_from_streams
+            logger.info(f"Kafka: Populated {len(topics_from_streams)} topics from streams config")
+            for topic_config in topics_from_streams:
+                logger.info(f"  - Topic: {topic_config.name} -> {topic_config.destination_id}")
+
+        # Update topic_map with destination table_ids from streams
+        for topic, stream_config in streams_map.items():
+            self.topic_map[topic] = stream_config.destination.table_id
+
     @staticmethod
     def streams() -> List[str]:
         return ["topic"]
@@ -99,24 +161,52 @@ class KafkaSource(AbstractSource):
     def check_connection(self) -> Tuple[bool | Any | None]:
         """Check the connection to the Kafka source"""
 
-        logger.info(f"Found: {len(self.consumer.list_topics().topics)} topics")
+        # Validate that topics have been configured
+        if not self.config.topics:
+            error_msg = (
+                "No topics configured. Either provide topics in source config or use streams configuration. "
+                "If using streams config, ensure set_streams_config() is called before check_connection()."
+            )
+            logger.error(error_msg)
+            return False, error_msg
+
+        try:
+            # Use a short timeout to avoid hanging on connection issues
+            cluster_metadata = self.consumer.list_topics(timeout=self.config.consumer_timeout)
+            topics = cluster_metadata.topics
+
+            logger.info(f"Found: {len(topics)} topics")
+
+            config_topics = [topic.name for topic in self.config.topics]
 
-        topics = self.consumer.list_topics().topics
+            # Display consumer config
+            # We ignore the key sasl.password and sasl.username
+            consumer_config = self.config.consumer_config.copy()
+            consumer_config.pop("sasl.password", None)
+            consumer_config.pop("sasl.username", None)
+            logger.info(f"Consumer config: {consumer_config}")
 
-        config_topics = [topic.name for topic in self.config.topics]
+            for topic in config_topics:
+                if topic not in topics:
+                    logger.error(f"Topic {topic} not found, available topics: {topics.keys()}")
+                    return False, f"Topic {topic} not found"
 
-        for topic in config_topics:
-            if topic not in topics:
-                logger.error(f"Topic {topic} not found, available topics: {topics.keys()}")
-                return False, f"Topic {topic} not found"
+                logger.info(f"Topic {topic} has {len(topics[topic].partitions)} partitions")
 
-            logger.info(f"Topic {topic} has {len(topics[topic].partitions)} partitions")
+            return True, None
 
-        return True, None
+        except KafkaException as e:
+            error_msg = f"Kafka connection failed: {e}"
+            logger.error(error_msg)
+            return False, error_msg
+        except Exception as e:
+            error_msg = f"Connection check failed: {e}"
+            logger.error(error_msg)
+            return False, error_msg
 
     def get_number_of_partitions(self, topic: str) -> int:
         """Get the number of partitions for the topic"""
-        return len(self.consumer.list_topics().topics[topic].partitions)
+        return len(self.consumer.list_topics(timeout=self.config.consumer_timeout).topics[topic].partitions)
 
     def get_offset_partitions(self, topic: str) -> TopicOffsets:
         """Get the offsets for each partition of the topic"""
@@ -124,7 +214,9 @@ class KafkaSource(AbstractSource):
         partitions: Mapping[int, OffsetPartition] = {}
 
         for i in range(self.get_number_of_partitions(topic)):
-            offsets = self.consumer.get_watermark_offsets(TopicPartition(topic, i))
+            offsets = self.consumer.get_watermark_offsets(
+                TopicPartition(topic, i), timeout=self.config.consumer_timeout
+            )
             partitions[i] = OffsetPartition(first=offsets[0], last=offsets[1])
 
         return TopicOffsets(name=topic, partitions=partitions)
@@ -137,7 +229,7 @@ class KafkaSource(AbstractSource):
             total_records += self.get_offset_partitions(topic).total_offset
         return total_records
 
-    @lru_cache(maxsize=None)
+    @cache
     def get_schema_from_registry(self, global_id: int) -> Tuple[Hashabledict, Schema]:
         """Get the schema from the registry, return a hashable dict and an avro schema object"""
 
@@ -174,40 +266,40 @@ class KafkaSource(AbstractSource):
         else:
             raise ValueError(f"Schema registry type {self.config.authentication.schema_registry_type} not supported")
 
-    def decode_avro(self, message: Message) -> dict:
-        # Get the header bytes and the global id from the message
-        header_message_bytes = get_header_bytes(
-            nb_bytes_schema_id=self.config.nb_bytes_schema_id, message=message.value()
-        )
-        global_id = parse_global_id_from_serialized_message(
-            nb_bytes_schema_id=self.config.nb_bytes_schema_id,
-            header_message_bytes=header_message_bytes,
+    def decode_avro(self, message: Message) -> Tuple[dict, dict]:
+        """Decode the message as avro and return the parsed message and the schema"""
+        global_id, nb_bytes_schema_id = parse_global_id_from_serialized_message(
+            message=message.value(),
         )
 
         try:
             hashable_dict_schema, avro_schema = self.get_schema_from_registry(global_id=global_id)
         except SchemaNotFound as e:
             logger.error(
-                (
-                    f"Message on topic {message.topic()} partition {message.partition()} at offset {message.offset()} has a  SchemaID of {global_id} which is not found in Registry."
-                    f"message value: {message.value()}."
-                )
+                f"Message on topic {message.topic()} partition {message.partition()} at offset {message.offset()} has a  SchemaID of {global_id} which is not found in Registry."
+                f"message value: {message.value()}."
             )
             logger.error(traceback.format_exc())
             raise e
 
-        return decode_avro_message(
-            message=message,
-            nb_bytes_schema_id=self.config.nb_bytes_schema_id,
-            hashable_dict_schema=hashable_dict_schema,
-            avro_schema=avro_schema,
+        return (
+            decode_avro_message(
+                message_value=message.value(),
+                nb_bytes_schema_id=nb_bytes_schema_id,
+                avro_schema=avro_schema,
+            ),
+            hashable_dict_schema,
         )
 
-    def decode_utf_8(self, message: Message):
+    def decode_utf_8(self, message: Message) -> Tuple[dict, dict]:
+        """Decode the message as utf-8 and return the parsed message and the schema"""
         # Decode the message as utf-8
-        return orjson.loads(message.value().decode("utf-8"))
+        return orjson.loads(message.value().decode("utf-8")), {}
 
-    def decode(self, message):
+    def decode(self, message) -> Tuple[dict, dict]:
+        """Decode the message based on the encoding type
+        Returns parsed message and the schema
+        """
         if self.config.message_encoding == MessageEncoding.AVRO:
             return self.decode_avro(message)
 
@@ -218,50 +310,64 @@ class KafkaSource(AbstractSource):
             raise ValueError(f"Message encoding {self.config.message_encoding} not supported")
 
     def parse_encoded_messages(self, encoded_messages: list) -> List[SourceRecord]:
+        """Parse the encoded Kafka messages and return a list of SourceRecord"""
 
         records = []
 
         for message in encoded_messages:
+            MESSAGE_LOG_METADATA = (
+                f"Message for topic {message.topic()} partition {message.partition()} and offset {message.offset()}"
+            )
 
             if message.error():
                 # If the message is too large, we skip it and update the offset
                 if message.error().code() == KafkaError.MSG_SIZE_TOO_LARGE:
-                    logger.warning(
-                        (
-                            f"Message for topic {message.topic()} partition {message.partition()} and offset {message.offset()} has been skipped. "
-                            f"Raised MSG_SIZE_TOO_LARGE, we suppose the message does not exist. Double-check in Confluent Cloud."
-                        )
+                    logger.error(
+                        f"{MESSAGE_LOG_METADATA} is too large. "
+                        "Raised MSG_SIZE_TOO_LARGE, if manually setting the offset, the message might not exist. Double-check in Confluent Cloud."
                     )
-                    continue
 
-                logger.error(
-                    (
-                        f"Error while consuming message for topic {message.topic()} partition {message.partition()} and offset {message.offset()}: "
-                        f"{message.error()}"
-                    )
-                )
+                logger.error(f"{MESSAGE_LOG_METADATA}: {message.error()}")
                 raise KafkaException(message.error())
 
             # We skip tombstone messages
             if self.config.skip_message_empty_value and not message.value():
-                logger.debug(
-                    f"Message for topic {message.topic()} partition {message.partition()} and offset {message.offset()} is empty, skipping."
-                )
+                logger.debug(f"{MESSAGE_LOG_METADATA} is empty, skipping.")
                 continue
 
+            # Parse message keys
+            if message.key():
+                try:
+                    message_keys = orjson.loads(message.key().decode("utf-8"))
+                except orjson.JSONDecodeError as e:
+                    # We skip messages with invalid keys
+                    if self.config.skip_message_invalid_keys:
+                        logger.warning(f"{MESSAGE_LOG_METADATA} has an invalid key={message.key()}, skipping.")
+                        # Skip the message
+                        continue
+
+                    logger.error(
+                        f"{MESSAGE_LOG_METADATA}: Error while parsing message key: {e}, raw key: {message.key()}"
+                    )
+                    raise e
+            else:
+                message_keys = {}
+
             # Decode the message
             try:
+                decoded_message, hashable_dict_schema = self.decode(message)
 
                 data = {
                     "topic": message.topic(),
                     "offset": message.offset(),
                     "partition": message.partition(),
                     "timestamp": message.timestamp()[1],
-                    "keys": orjson.loads(message.key().decode("utf-8")) if message.key() else {},
+                    "keys": message_keys,
                     "headers": (
                         {key: value.decode("utf-8") for key, value in message.headers()} if message.headers() else {}
                     ),
-                    "value": self.decode(message),
+                    "value": decoded_message,
+                    "schema": hashable_dict_schema,
                 }
 
                 records.append(
@@ -275,17 +381,27 @@ class KafkaSource(AbstractSource):
 
             except Exception as e:
                 logger.error(
-                    (
-                        f"Error while decoding message for topic {message.topic()} on partition {message.partition()}: {e} at offset {message.offset()} "
-                        f"with value: {message.value()} and key: {message.key()}"
-                    )
+                    f"{MESSAGE_LOG_METADATA}: Error while decoding message: {e} "
+                    f"with value: {message.value()} and key: {message.key()}"
                 )
-                # Try to parse error message from the message
+
+                # Try to parse error message from the message value
                 try:
                     message_raw_text = message.value().decode("utf-8")
                     logger.error(f"Parsed Kafka value: {message_raw_text}")
                 except UnicodeDecodeError:
-                    logger.error("Message is not a valid UTF-8 string")
+                    logger.error("Message value is not a valid UTF-8 string")
+
+                # Try to parse error message from the message headers
+                if message.headers():
+                    try:
+                        headers_dict = {key: value.decode("utf-8") for key, value in message.headers()}
+                        logger.error(f"Parsed Kafka headers: {headers_dict}")
+                    except UnicodeDecodeError as header_error:
+                        logger.error(f"Some message headers are not valid UTF-8 strings: {header_error}")
+                        logger.error(f"Raw message headers: {list(message.headers())}")
+                else:
+                    logger.error("Message headers are None or empty")
 
                 logger.error(traceback.format_exc())
                 raise e
@@ -358,4 +474,9 @@ class KafkaSource(AbstractSource):
             return self.read_topics_manually(pagination)
 
     def commit(self):
-        self.consumer.commit(asynchronous=False)
+        """Commit the offsets of the consumer"""
+        try:
+            self.consumer.commit(asynchronous=False)
+        except CimplKafkaException as e:
+            logger.error(f"Kafka exception occurred during commit: {e}")
+            logger.info("Gracefully exiting without committing offsets due to Kafka exception")

bizon/connectors/sources/kafka/tests/kafka_pipeline.py

@@ -3,5 +3,5 @@ import os
 from bizon.engine.engine import RunnerFactory
 
 if __name__ == "__main__":
-    runner = RunnerFactory.create_from_yaml(filepath=os.path.abspath("test-pipeline-streaming.yml"))
+    runner = RunnerFactory.create_from_yaml(filepath=os.path.abspath("test-pipeline-streaming-v2.yml"))
     runner.run()

bizon/connectors/sources/notion/config/api_key.example.yml

@@ -0,0 +1,35 @@
+# Notion Source Configuration
+# This example shows how to configure the Notion source connector
+
+source:
+  name: notion
+  stream: pages  # Options: databases, data_sources, pages, blocks, users
+  authentication:
+    type: api_key
+    params:
+      token: secret_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx  # Your Notion integration token
+
+  # List of database IDs to fetch data from
+  # Find the ID in the database URL: notion.so/{workspace}/{database_id}?v=...
+  database_ids:
+    - "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"
+    - "yyyyyyyy-yyyy-yyyy-yyyy-yyyyyyyyyyyy"
+
+  # List of specific page IDs to fetch (optional)
+  # Find the ID in the page URL: notion.so/{page_id}
+  page_ids:
+    - "zzzzzzzz-zzzz-zzzz-zzzz-zzzzzzzzzzzz"
+
+  # Whether to fetch nested blocks recursively (default: true)
+  # Only applies to blocks stream
+  fetch_blocks_recursively: true
+
+  # Number of results per API call (1-100, default: 100)
+  page_size: 100
+
+destination:
+  name: bigquery
+  config:
+    project_id: my-project
+    dataset_id: notion_data
+    # ... other destination config

bizon/connectors/sources/notion/src/config.py

@@ -0,0 +1,59 @@
+from enum import Enum
+from typing import Any, Dict, List
+
+from pydantic import Field
+
+from bizon.source.config import SourceConfig
+
+
+class NotionStreams(str, Enum):
+    DATABASES = "databases"
+    DATA_SOURCES = "data_sources"
+    PAGES = "pages"
+    BLOCKS = "blocks"
+    BLOCKS_MARKDOWN = "blocks_markdown"
+    USERS = "users"
+    # Streams that fetch all accessible content (no database_ids/page_ids required)
+    ALL_PAGES = "all_pages"
+    ALL_DATABASES = "all_databases"
+    ALL_DATA_SOURCES = "all_data_sources"
+    ALL_BLOCKS_MARKDOWN = "all_blocks_markdown"
+
+
+class NotionSourceConfig(SourceConfig):
+    stream: NotionStreams
+
+    database_ids: List[str] = Field(
+        default_factory=list,
+        description="List of Notion database IDs to fetch. Required for databases, data_sources, pages, and blocks streams.",
+    )
+    page_ids: List[str] = Field(
+        default_factory=list,
+        description="List of Notion page IDs to fetch. Used for pages and blocks streams.",
+    )
+    fetch_blocks_recursively: bool = Field(
+        default=True,
+        description="Whether to fetch nested blocks recursively. Only applies to blocks stream.",
+    )
+    max_recursion_depth: int = Field(
+        default=5,
+        ge=1,
+        le=100,
+        description="Maximum nesting depth for recursive block fetching. Prevents infinite loops.",
+    )
+    page_size: int = Field(
+        default=100,
+        ge=1,
+        le=100,
+        description="Number of results per page (max 100)",
+    )
+    max_workers: int = Field(
+        default=3,
+        ge=1,
+        le=10,
+        description="Number of concurrent workers for fetching blocks. Keep low to respect rate limits.",
+    )
+    database_filters: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Map of database_id -> Notion filter object. Filters are passed directly to Notion API.",
+    )