Add optional connector connect success and failure callbacks (#708)

Add optional connector connect success and failure callbacks

Author: tim-quix

quixstreams/processing/context.py

@@ -101,6 +101,7 @@ def on_partition_revoke(self, topic: str, partition: int):
         self.pausing_manager.revoke(topic=topic, partition=partition)
 
     def __enter__(self):
+        self.sink_manager.start_sinks()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):

quixstreams/sinks/__init__.py

@@ -1,9 +1,19 @@
-from .base import BaseSink, BatchingSink, SinkBackpressureError, SinkBatch, SinkManager
+from .base import (
+    BaseSink,
+    BatchingSink,
+    ClientConnectFailureCallback,
+    ClientConnectSuccessCallback,
+    SinkBackpressureError,
+    SinkBatch,
+    SinkManager,
+)
 
 __all__ = [
     "BaseSink",
     "BatchingSink",
     "SinkBackpressureError",
     "SinkBatch",
     "SinkManager",
+    "ClientConnectSuccessCallback",
+    "ClientConnectFailureCallback",
 ]

quixstreams/sinks/base/__init__.py

@@ -1,12 +1,19 @@
 from .batch import SinkBatch
 from .exceptions import SinkBackpressureError
 from .manager import SinkManager
-from .sink import BaseSink, BatchingSink
+from .sink import (
+    BaseSink,
+    BatchingSink,
+    ClientConnectFailureCallback,
+    ClientConnectSuccessCallback,
+)
 
 __all__ = (
     "SinkBatch",
     "SinkBackpressureError",
     "SinkManager",
     "BatchingSink",
     "BaseSink",
+    "ClientConnectSuccessCallback",
+    "ClientConnectFailureCallback",
 )

quixstreams/sinks/base/manager.py

@@ -14,6 +14,10 @@ def register(self, sink: BaseSink):
         if sink_id not in self._sinks:
             self._sinks[id(sink)] = sink
 
+    def start_sinks(self):
+        for sink in self.sinks:
+            sink.start()
+
     @property
     def sinks(self) -> List[BaseSink]:
         return list(self._sinks.values())

quixstreams/sinks/base/sink.py

@@ -1,13 +1,25 @@
 import abc
 import logging
-from typing import Any, Dict, Tuple
+from typing import Any, Callable, Dict, Optional, Tuple
 
 from quixstreams.models import HeadersTuples
 from quixstreams.sinks.base.batch import SinkBatch
 
 logger = logging.getLogger(__name__)
 
 
+ClientConnectSuccessCallback = Callable[[], None]
+ClientConnectFailureCallback = Callable[[Optional[Exception]], None]
+
+
+def _default_on_client_connect_success():
+    logger.info("CONNECTED!")
+
+
+def _default_on_client_connect_failure(exception: Exception):
+    logger.error(f"ERROR: Failed while connecting to client: {exception}")
+
+
 class BaseSink(abc.ABC):
     """
     This is a base class for all sinks.
@@ -17,6 +29,26 @@ class BaseSink(abc.ABC):
     Note that Sinks are currently in beta, and their design may change over time.
     """
 
+    def __init__(
+        self,
+        on_client_connect_success: Optional[ClientConnectSuccessCallback] = None,
+        on_client_connect_failure: Optional[ClientConnectFailureCallback] = None,
+    ):
+        """
+        :param on_client_connect_success: An optional callback made after successful
+            client authentication, primarily for additional logging.
+        :param on_client_connect_failure: An optional callback made after failed
+            client authentication (which should raise an Exception).
+            Callback should accept the raised Exception as an argument.
+            Callback must resolve (or propagate/re-raise) the Exception.
+        """
+        self._on_client_connect_success = (
+            on_client_connect_success or _default_on_client_connect_success
+        )
+        self._on_client_connect_failure = (
+            on_client_connect_failure or _default_on_client_connect_failure
+        )
+
     @abc.abstractmethod
     def flush(self, topic: str, partition: int):
         """
@@ -48,6 +80,24 @@ def add(
         on flush().
         """
 
+    @abc.abstractmethod
+    def setup(self):
+        """
+        When applicable, set up the client here along with any validation to affirm a
+        valid/successful authentication/connection.
+        """
+
+    def start(self):
+        """
+        Called as part of `Application.run()` to initialize the sink's client.
+        Allows using a callback pattern around the connection attempt.
+        """
+        try:
+            self.setup()
+            self._on_client_connect_success()
+        except Exception as e:
+            self._on_client_connect_failure(e)
+
     def on_paused(self, topic: str, partition: int):
         """
         This method is triggered when the sink is paused due to backpressure, when
@@ -60,7 +110,7 @@ def on_paused(self, topic: str, partition: int):
 class BatchingSink(BaseSink):
     """
     A base class for batching sinks, that need to accumulate the data first before
-    sending it to the external destinatios.
+    sending it to the external destinations.
 
     Examples: databases, objects stores, and other destinations where
     writing every message is not optimal.
@@ -73,7 +123,23 @@ class BatchingSink(BaseSink):
 
     _batches: Dict[Tuple[str, int], SinkBatch]
 
-    def __init__(self):
+    def __init__(
+        self,
+        on_client_connect_success: Optional[ClientConnectSuccessCallback] = None,
+        on_client_connect_failure: Optional[ClientConnectFailureCallback] = None,
+    ):
+        """
+        :param on_client_connect_success: An optional callback made after successful
+            client authentication, primarily for additional logging.
+        :param on_client_connect_failure: An optional callback made after failed
+            client authentication (which should raise an Exception).
+            Callback should accept the raised Exception as an argument.
+            Callback must resolve (or propagate/re-raise) the Exception.
+        """
+        super().__init__(
+            on_client_connect_success=on_client_connect_success,
+            on_client_connect_failure=on_client_connect_failure,
+        )
         self._batches = {}
 
     def __repr__(self):

quixstreams/sinks/community/bigquery.py

@@ -3,9 +3,7 @@
 import time
 from datetime import date, datetime
 from decimal import Decimal
-from typing import Any, Mapping
-
-from typing_extensions import Optional
+from typing import Any, Mapping, Optional
 
 try:
     from google.cloud import bigquery
@@ -19,7 +17,12 @@
 
 from quixstreams.exceptions import QuixException
 from quixstreams.models import HeadersTuples
-from quixstreams.sinks import BatchingSink, SinkBatch
+from quixstreams.sinks import (
+    BatchingSink,
+    ClientConnectFailureCallback,
+    ClientConnectSuccessCallback,
+    SinkBatch,
+)
 
 __all__ = ("BigQuerySink", "BigQuerySinkException")
 
@@ -62,6 +65,8 @@ def __init__(
         ddl_timeout: float = 10.0,
         insert_timeout: float = 10.0,
         retry_timeout: float = 30.0,
+        on_client_connect_success: Optional[ClientConnectSuccessCallback] = None,
+        on_client_connect_failure: Optional[ClientConnectFailureCallback] = None,
         **kwargs,
     ):
         """
@@ -101,10 +106,19 @@ def __init__(
         :param retry_timeout: a total timeout for each request to BigQuery API.
             During this timeout, a request can be retried according
             to the client's default retrying policy.
+        :param on_client_connect_success: An optional callback made after successful
+            client authentication, primarily for additional logging.
+        :param on_client_connect_failure: An optional callback made after failed
+            client authentication (which should raise an Exception).
+            Callback should accept the raised Exception as an argument.
+            Callback must resolve (or propagate/re-raise) the Exception.
         :param kwargs: Additional keyword arguments passed to `bigquery.Client`.
         """
+        super().__init__(
+            on_client_connect_success=on_client_connect_success,
+            on_client_connect_failure=on_client_connect_failure,
+        )
 
-        super().__init__()
         self.location = location
         self.project_id = project_id
         self.dataset_id = f"{self.project_id}.{dataset_id}"
@@ -123,12 +137,16 @@ def __init__(
                 scopes=["https://www.googleapis.com/auth/bigquery"],
             )
             kwargs["credentials"] = credentials
-
-        self._client = bigquery.Client(**kwargs)
-        logger.info("Successfully authenticated to BigQuery.")
-        if self.schema_auto_update:
-            # Initialize a table in BigQuery if it doesn't exist already
-            self._init_table()
+        self._client: Optional[bigquery.Client] = None
+        self._client_settings = kwargs
+
+    def setup(self):
+        if not self._client:
+            self._client = bigquery.Client(**self._client_settings)
+            logger.info("Successfully authenticated to BigQuery.")
+            if self.schema_auto_update:
+                # Initialize a table in BigQuery if it doesn't exist already
+                self._init_table()
 
     def write(self, batch: SinkBatch):
         rows = []

quixstreams/sinks/community/file/destinations/azure.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Optional
 
 from quixstreams.sinks import SinkBatch
 from quixstreams.sinks.community.file.destinations.base import Destination
@@ -51,17 +52,17 @@ def __init__(
         :raises AzureContainerAccessDeniedError: If access to the container is denied.
         """
         self._container = container
-        self._client = self._get_client(connection_string)
-        self._validate_container()
+        self._auth = connection_string
+        self._client: Optional[ContainerClient] = None
 
-    def _get_client(self, auth: str) -> ContainerClient:
+    def _get_client(self) -> ContainerClient:
         """
         Get an Azure file container client and validate the container exists.
 
         :param auth: Azure client authentication string.
         :return: An Azure ContainerClient
         """
-        storage_client = BlobServiceClient.from_connection_string(auth)
+        storage_client = BlobServiceClient.from_connection_string(self._auth)
         container_client = storage_client.get_container_client(self._container)
         return container_client
 
@@ -85,6 +86,11 @@ def _validate_container(self) -> None:
             raise
         raise AzureContainerNotFoundError(f"Container not found: {self._container}")
 
+    def setup(self):
+        if not self._client:
+            self._client = self._get_client()
+            self._validate_container()
+
     def write(self, data: bytes, batch: SinkBatch) -> None:
         """
         Write data to Azure.

quixstreams/sinks/community/file/destinations/base.py

@@ -10,7 +10,7 @@
 
 logger = logging.getLogger(__name__)
 
-_UNSAFE_CHARACTERS_REGEX = re.compile(r"[^a-zA-Z0-9 ._]")
+_UNSAFE_CHARACTERS_REGEX = re.compile(r"[^a-zA-Z0-9 ._/]")
 
 
 class Destination(ABC):
@@ -25,14 +25,33 @@ class Destination(ABC):
     _base_directory: str = ""
     _extension: str = ""
 
-    def set_directory(self, directory: str) -> None:
+    @abstractmethod
+    def setup(self):
+        """Authenticate and validate connection here"""
+        ...
+
+    @abstractmethod
+    def write(self, data: bytes, batch: SinkBatch) -> None:
+        """Write the serialized data to storage.
+
+        :param data: The serialized data to write.
+        :param batch: The batch information containing topic, partition and offset
+            details.
+        """
+        ...
+
+    def set_directory(
+        self,
+        directory: str,
+    ) -> None:
         """Configure the base directory for storing files.
 
         :param directory: The base directory path where files will be stored.
         :raises ValueError: If the directory path contains invalid characters.
-            Only alphanumeric characters (a-zA-Z0-9), spaces, dots, and
+            Only alphanumeric characters (a-zA-Z0-9), spaces, dots, slashes, and
             underscores are allowed.
         """
+        # TODO: logic around "/" for sinks that require special handling of them
         if _UNSAFE_CHARACTERS_REGEX.search(directory):
             raise ValueError(
                 f"Invalid characters in directory path: {directory}. "
@@ -50,16 +69,6 @@ def set_extension(self, format: Format) -> None:
         self._extension = format.file_extension
         logger.info("File extension set to '%s'", self._extension)
 
-    @abstractmethod
-    def write(self, data: bytes, batch: SinkBatch) -> None:
-        """Write the serialized data to storage.
-
-        :param data: The serialized data to write.
-        :param batch: The batch information containing topic, partition and offset
-            details.
-        """
-        ...
-
     def _path(self, batch: SinkBatch) -> Path:
         """Generate the full path where the batch data should be stored.
 

quixstreams/sinks/community/file/destinations/local.py

@@ -29,6 +29,9 @@ def __init__(self, append: bool = False) -> None:
         self._mode = "ab" if append else "wb"
         logger.debug("LocalDestination initialized with append=%s", append)
 
+    def setup(self):
+        return
+
     def set_extension(self, format: Format) -> None:
         """Set the file extension and validate append mode compatibility.