Add HeadersProvider which provides needed gRPC connection headers (#5)

* Initial commit

Signed-off-by: teodordelibasic-db <[email protected]>

* Add header_provider.py

Signed-off-by: teodordelibasic-db <[email protected]>

* Update NEXT_CHANGELOG.md

Signed-off-by: teodordelibasic-db <[email protected]>

* Address comments

Signed-off-by: teodordelibasic-db <[email protected]>

---------

Signed-off-by: teodordelibasic-db <[email protected]>

Author: teodordelibasic-db

NEXT_CHANGELOG.md

@@ -4,10 +4,20 @@
 
 ### New Features and Improvements
 
+- Added `HeadersProvider` abstraction for flexible authentication strategies
+- Implemented `OAuthHeadersProvider` for OAuth 2.0 Client Credentials flow (default authentication method used by `create_stream()`)
+
 ### Bug Fixes
 
 ### Documentation
 
+- Added Azure workspace and endpoint URL examples
+
 ### Internal Changes
 
 ### API Changes
+
+- Added `HeadersProvider` abstract base class for custom header strategies
+- Added `OAuthHeadersProvider` class for OAuth 2.0 authentication with Databricks OIDC endpoint
+- Added `create_stream_with_headers_provider` method to `ZerobusSdk` and `aio.ZerobusSdk` for custom authentication header providers
+  - **Note**: Custom headers providers must include both `authorization` and `x-databricks-zerobus-table-name` headers

README.md

@@ -24,6 +24,7 @@ The Databricks Zerobus Ingest SDK for Python provides a high-performance client
 - [Usage Examples](#usage-examples)
   - [Blocking Ingestion](#blocking-ingestion)
   - [Non-Blocking Ingestion](#non-blocking-ingestion)
+- [Authentication](#authentication)
 - [Configuration](#configuration)
 - [Error Handling](#error-handling)
 - [API Reference](#api-reference)
@@ -479,6 +480,80 @@ async def main():
 asyncio.run(main())
 ```
 
+## Authentication
+
+The SDK uses OAuth 2.0 Client Credentials for authentication:
+
+```python
+from zerobus.sdk.sync import ZerobusSdk
+from zerobus.sdk.shared import TableProperties
+import record_pb2
+
+sdk = ZerobusSdk(server_endpoint, workspace_url)
+table_properties = TableProperties(table_name, record_pb2.AirQuality.DESCRIPTOR)
+
+# Create stream with OAuth authentication
+stream = sdk.create_stream(client_id, client_secret, table_properties)
+```
+
+The SDK automatically fetches access tokens and includes these headers:
+- `"authorization": "Bearer <oauth_token>"` - Obtained via OAuth 2.0 Client Credentials flow
+- `"x-databricks-zerobus-table-name": "<table_name>"` - The fully qualified table name
+
+### Advanced: Custom Headers
+
+For advanced use cases where you need to provide custom headers (e.g., for future authentication methods or additional metadata), you can implement a custom `HeadersProvider`:
+
+```python
+from zerobus.sdk.shared.headers_provider import HeadersProvider
+
+class CustomHeadersProvider(HeadersProvider):
+    """
+    Custom headers provider for advanced use cases.
+
+    Note: Currently, OAuth 2.0 Client Credentials (via create_stream())
+    is the standard authentication method. Use this only if you have
+    specific requirements for custom headers.
+    
+    IMPORTANT: Custom headers providers MUST include both:
+    - "authorization": "Bearer <token>" header for authentication
+    - "x-databricks-zerobus-table-name": "<table_name>" header for routing
+    """
+
+    def __init__(self, token: str, table_name: str):
+        self.token = token
+        self.table_name = table_name
+
+    def get_headers(self):
+        """
+        Return headers for gRPC metadata.
+
+        Returns:
+            List of (header_name, header_value) tuples
+        """
+        return [
+            ("authorization", f"Bearer {self.token}"),
+            ("x-databricks-zerobus-table-name", self.table_name),
+            ("x-custom-header", "custom-value"),  # Optional: additional custom headers
+        ]
+
+# Use the custom provider
+custom_provider = CustomHeadersProvider(
+    token="your-token",
+    table_name=table_properties.table_name
+)
+stream = sdk.create_stream_with_headers_provider(
+    custom_provider,
+    table_properties
+)
+```
+
+**Potential use cases for custom headers:**
+- Integration with existing token management systems
+- Additional metadata headers for request tracking
+- Future authentication methods
+- Special routing or service mesh requirements
+
 ## Configuration
 
 ### Stream Configuration Options
@@ -565,7 +640,22 @@ def create_stream(
     options: StreamConfigurationOptions = None
 ) -> ZerobusStream
 ```
-Creates a new ingestion stream. Returns a `ZerobusStream` instance.
+Creates a new ingestion stream using OAuth 2.0 Client Credentials authentication.
+
+Automatically includes these headers:
+- `"authorization": "Bearer <oauth_token>"` (fetched via OAuth 2.0 Client Credentials flow)
+- `"x-databricks-zerobus-table-name": "<table_name>"`
+
+Returns a `ZerobusStream` instance.
+
+```python
+def create_stream_with_headers_provider(
+    headers_provider: HeadersProvider,
+    table_properties: TableProperties,
+    options: StreamConfigurationOptions = None
+) -> ZerobusStream
+```
+Creates a new ingestion stream using a custom headers provider. For advanced use cases only where custom headers are required. Returns a `ZerobusStream` instance.
 
 ---
 
@@ -586,7 +676,22 @@ async def create_stream(
     options: StreamConfigurationOptions = None
 ) -> ZerobusStream
 ```
-Creates a new ingestion stream. Returns a `ZerobusStream` instance.
+Creates a new ingestion stream using OAuth 2.0 Client Credentials authentication.
+
+Automatically includes these headers:
+- `"authorization": "Bearer <oauth_token>"` (fetched via OAuth 2.0 Client Credentials flow)
+- `"x-databricks-zerobus-table-name": "<table_name>"`
+
+Returns a `ZerobusStream` instance.
+
+```python
+async def create_stream_with_headers_provider(
+    headers_provider: HeadersProvider,
+    table_properties: TableProperties,
+    options: StreamConfigurationOptions = None
+) -> ZerobusStream
+```
+Creates a new ingestion stream using a custom headers provider. For advanced use cases only where custom headers are required. Returns a `ZerobusStream` instance.
 
 ---
 
@@ -683,6 +788,62 @@ Returns the protobuf message descriptor.
 
 ---
 
+### HeadersProvider
+
+Abstract base class for providing headers to gRPC streams. For advanced use cases only.
+
+**Abstract Method:**
+
+```python
+@abstractmethod
+def get_headers(self) -> List[Tuple[str, str]]
+```
+Returns headers for gRPC metadata as a list of (header_name, header_value) tuples.
+
+**Built-in Implementation:**
+
+#### OAuthHeadersProvider
+
+OAuth 2.0 Client Credentials flow headers provider (used internally by `create_stream()`).
+
+```python
+OAuthHeadersProvider(
+    workspace_id: str,
+    workspace_url: str,
+    table_name: str,
+    client_id: str,
+    client_secret: str
+)
+```
+
+Returns these headers:
+- `"authorization": "Bearer <oauth_token>"` (fetched via OAuth 2.0)
+- `"x-databricks-zerobus-table-name": "<table_name>"`
+
+**Custom Implementation (Advanced):**
+
+For advanced use cases requiring custom headers, extend the `HeadersProvider` class:
+
+```python
+from zerobus.sdk.shared.headers_provider import HeadersProvider
+
+class MyCustomProvider(HeadersProvider):
+    def get_headers(self):
+        return [
+            ("authorization", "Bearer my-token"),  # Required
+            ("x-databricks-zerobus-table-name", "catalog.schema.table"),  # Required
+            ("x-custom-header", "value"),  # Optional: additional custom headers
+        ]
+```
+
+**Important:** Custom headers providers MUST include both required headers:
+- `"authorization"`: Bearer token for authentication
+- `"x-databricks-zerobus-table-name"`: Fully qualified table name for routing
+
+Note: Most users should use `create_stream()` with OAuth credentials rather than implementing a custom provider.
+
+---
+
 ### StreamConfigurationOptions
 
 Configuration options for stream behavior.

examples/async_example.py

@@ -6,6 +6,10 @@
 Use Case: Best for applications already using asyncio, async web frameworks (FastAPI, aiohttp),
 or when integrating ingestion with other asynchronous operations in an event loop.
 
+Authentication:
+  - Uses OAuth 2.0 Client Credentials (standard method)
+  - Includes example of custom headers provider for advanced use cases
+
 Note: Both sync and async APIs provide the same throughput and durability guarantees.
 Choose based on your application's architecture, not performance requirements.
 """
@@ -20,6 +24,7 @@
 
 from zerobus.sdk.aio import ZerobusSdk
 from zerobus.sdk.shared import StreamConfigurationOptions, TableProperties
+from zerobus.sdk.shared.headers_provider import HeadersProvider
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
@@ -56,6 +61,31 @@ def create_sample_record(index):
     return record_pb2.AirQuality(device_name=f"sensor-{index % 10}", temp=20 + (index % 15), humidity=50 + (index % 40))
 
 
+class CustomHeadersProvider(HeadersProvider):
+    """
+    Example custom headers provider for advanced use cases.
+
+    Note: OAuth 2.0 Client Credentials (via create_stream()) is the standard
+    authentication method. Use this only if you have specific requirements
+    for custom headers (e.g., custom metadata, existing token management, etc.).
+    """
+
+    def __init__(self, custom_token: str):
+        self.custom_token = custom_token
+
+    def get_headers(self):
+        """
+        Return custom headers for gRPC metadata.
+
+        Returns:
+            List of (header_name, header_value) tuples
+        """
+        return [
+            ("authorization", f"Bearer {self.custom_token}"),
+            ("x-custom-header", "custom-value"),
+        ]
+
+
 def create_ack_callback():
     """
     Creates an acknowledgment callback that logs progress.
@@ -111,8 +141,19 @@ async def main():
         table_properties = TableProperties(TABLE_NAME, record_pb2.AirQuality.DESCRIPTOR)
         logger.info(f"✓ Table properties configured for: {TABLE_NAME}")
 
-        # Step 4: Create a stream
+        # Step 4: Create a stream with OAuth 2.0 authentication
+        #
+        # Standard method: OAuth 2.0 Client Credentials
+        # The SDK automatically includes these headers:
+        #   - "authorization": "Bearer <oauth_token>" (fetched via OAuth 2.0 Client Credentials flow)
+        #   - "x-databricks-zerobus-table-name": "<table_name>"
         stream = await sdk.create_stream(CLIENT_ID, CLIENT_SECRET, table_properties, options)
+
+        # Advanced: Custom headers provider (for special use cases only)
+        # Uncomment to use custom headers instead of OAuth:
+        # custom_provider = CustomHeadersProvider(custom_token="your-custom-token")
+        # stream = await sdk.create_stream_with_headers_provider(custom_provider, table_properties, options)
+
         logger.info(f"✓ Stream created: {stream.stream_id}")
 
         # Step 5: Ingest records asynchronously

examples/sync_example.py

@@ -5,6 +5,10 @@
 
 Use Case: Best for applications that don't use asyncio or prefer blocking I/O patterns.
 
+Authentication:
+  - Uses OAuth 2.0 Client Credentials (standard method)
+  - Includes example of custom headers provider for advanced use cases
+
 Note: Both sync and async APIs provide the same throughput and durability guarantees.
 Choose based on your application's architecture, not performance requirements.
 """
@@ -17,6 +21,7 @@
 import record_pb2
 
 from zerobus.sdk.shared import StreamConfigurationOptions, TableProperties
+from zerobus.sdk.shared.headers_provider import HeadersProvider
 from zerobus.sdk.sync import ZerobusSdk
 
 # Configure logging
@@ -54,6 +59,31 @@ def create_sample_record(index):
     return record_pb2.AirQuality(device_name=f"sensor-{index % 10}", temp=20 + (index % 15), humidity=50 + (index % 40))
 
 
+class CustomHeadersProvider(HeadersProvider):
+    """
+    Example custom headers provider for advanced use cases.
+
+    Note: OAuth 2.0 Client Credentials (via create_stream()) is the standard
+    authentication method. Use this only if you have specific requirements
+    for custom headers (e.g., custom metadata, existing token management, etc.).
+    """
+
+    def __init__(self, custom_token: str):
+        self.custom_token = custom_token
+
+    def get_headers(self):
+        """
+        Return custom headers for gRPC metadata.
+
+        Returns:
+            List of (header_name, header_value) tuples
+        """
+        return [
+            ("authorization", f"Bearer {self.custom_token}"),
+            ("x-custom-header", "custom-value"),
+        ]
+
+
 def main():
     print("Starting synchronous ingestion example...")
     print("=" * 60)
@@ -93,8 +123,19 @@ def main():
         )
         logger.info("✓ Stream configuration created")
 
-        # Step 4: Create a stream
+        # Step 4: Create a stream with OAuth 2.0 authentication
+        #
+        # Standard method: OAuth 2.0 Client Credentials
+        # The SDK automatically includes these headers:
+        #   - "authorization": "Bearer <oauth_token>" (fetched via OAuth 2.0 Client Credentials flow)
+        #   - "x-databricks-zerobus-table-name": "<table_name>"
         stream = sdk.create_stream(CLIENT_ID, CLIENT_SECRET, table_properties, options)
+
+        # Advanced: Custom headers provider (for special use cases only)
+        # Uncomment to use custom headers instead of OAuth:
+        # custom_provider = CustomHeadersProvider(custom_token="your-custom-token")
+        # stream = sdk.create_stream_with_headers_provider(custom_provider, table_properties, options)
+
         logger.info(f"✓ Stream created: {stream.stream_id}")
 
         # Step 5: Ingest records synchronously

tests/mock_grpc.py

@@ -471,7 +471,7 @@ async def wrapper(test_instance, *args, **kwargs):
         # Run for async SDK
         logging.info(f"\nRunning test '{test_func.__name__}' for Async SDK")
         async_sdk_manager = SdkManager(ZerobusSdkAsync)
-        with patch("zerobus.sdk.aio.zerobus_sdk.get_zerobus_token", return_value="mock_token"):
+        with patch("zerobus.sdk.shared.headers_provider.get_zerobus_token", return_value="mock_token"):
             await test_func(test_instance, async_sdk_manager, *args, **kwargs)
         logging.info(f"Test '{test_func.__name__}' PASSED for Async SDK")
 
@@ -482,7 +482,7 @@ async def wrapper(test_instance, *args, **kwargs):
         loop = asyncio.get_running_loop()
 
         def run_sync_test_in_isolation():
-            with patch("zerobus.sdk.sync.zerobus_sdk.get_zerobus_token", return_value="mock_token"):
+            with patch("zerobus.sdk.shared.headers_provider.get_zerobus_token", return_value="mock_token"):
                 asyncio.run(coro)
 
         await loop.run_in_executor(None, run_sync_test_in_isolation)

tests/test_sdk.py

@@ -51,7 +51,7 @@ def create_ephemeral_stream(generator, **kwargs):
 
         with (
             patch("grpc.secure_channel", return_value=mock_channel),
-            patch("zerobus.sdk.sync.zerobus_sdk.get_zerobus_token", return_value="mock_token"),
+            patch("zerobus.sdk.shared.headers_provider.get_zerobus_token", return_value="mock_token"),
         ):
             sdk_handle = ZerobusSdk(SERVER_ENDPOINT, unity_catalog_url="https://test.unity.catalog.url")
             try: