openlineage-python 1.22.0__tar.gz → 1.23.0__tar.gz

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: openlineage-python
-Version: 1.22.0
+Version: 1.23.0
 Summary: OpenLineage Python Client
 Author-email: OpenLineage <info@openlineage.io>
 Keywords: openlineage

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/openlineage/client/client.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+import json
 import logging
 import os
 import warnings
@@ -16,6 +17,8 @@ if TYPE_CHECKING:
     from requests import Session
     from requests.adapters import HTTPAdapter
 
+import contextlib
+
 from openlineage.client import event_v2
 from openlineage.client.run import DatasetEvent, JobEvent, RunEvent
 from openlineage.client.transport import Transport, TransportFactory, get_default_factory
@@ -40,6 +43,9 @@ _T = TypeVar("_T", bound="OpenLineageClient")
 
 
 class OpenLineageClient:
+    DYNAMIC_ENV_VARS_PREFIX = "OPENLINEAGE__"
+    DEFAULT_URL_TRANSPORT_NAME = "default_http"
+
     def __init__(  # noqa: PLR0913
         self,
         url: str | None = None,
@@ -64,6 +70,8 @@ class OpenLineageClient:
         # reload yaml each time config is referred to.
         self._config: dict[str, dict[str, str]] | None = None
 
+        self._alias_env_vars()
+
         self.transport = self._resolve_transport(
             url=url, options=options, session=session, transport=transport, factory=factory
         )
@@ -131,7 +139,7 @@ class OpenLineageClient:
                 self._config = {}
         return self._config
 
-    def _resolve_transport(self, **kwargs: Any) -> Transport:
+    def _resolve_transport(self, **kwargs: Any) -> Transport:  # noqa: PLR0911
         """
         Resolves the transport mechanism based on the provided arguments or environment settings.
 
@@ -166,11 +174,16 @@ class OpenLineageClient:
                 url=kwargs["url"], options=kwargs.get("options"), session=kwargs.get("session")
             )
 
-        # 5. Check HTTP transport initialization with env variables
+        # 5. Check transport initialization with env variables
+        if config := self._load_config_from_env_variables():
+            factory = kwargs.get("factory") or get_default_factory()
+            return factory.create(config["transport"])
+
+        # 6. Check HTTP transport initialization with env variables
         if os.environ.get("OPENLINEAGE_URL"):
             return self._http_transport_from_env_variables()
 
-        # 6. If all else fails, print events to console
+        # 7. If all else fails, print events to console
         from openlineage.client.transport.console import ConsoleConfig, ConsoleTransport
 
         log.warning("Couldn't find any OpenLineage transport configuration; will print events to console.")
@@ -245,3 +258,68 @@ class OpenLineageClient:
                 session=session,
             ),
         )
+
+    def _alias_env_vars(self) -> None:
+        default_transport_name = self.DEFAULT_URL_TRANSPORT_NAME.upper()
+        if url := os.environ.get("OPENLINEAGE_URL"):
+            if any(
+                k.startswith(f"OPENLINEAGE__TRANSPORT__TRANSPORTS__{default_transport_name}")
+                for k in os.environ
+            ):
+                log.warning(
+                    "%s already found in environment variables, skipping aliasing OPENLINEAGE_URL",
+                    default_transport_name,
+                )
+                return
+            os.environ[f"OPENLINEAGE__TRANSPORT__TRANSPORTS__{default_transport_name}__TYPE"] = "http"
+            os.environ[f"OPENLINEAGE__TRANSPORT__TRANSPORTS__{default_transport_name}__URL"] = url
+            if api_key := os.environ.get("OPENLINEAGE_API_KEY"):
+                os.environ[
+                    f"OPENLINEAGE__TRANSPORT__TRANSPORTS__{default_transport_name}__AUTH"
+                ] = json.dumps(
+                    {
+                        "type": "api_key",
+                        "apiKey": api_key,
+                    }
+                )
+            if endpoint := os.environ.get("OPENLINEAGE_ENDPOINT"):
+                os.environ[
+                    f"OPENLINEAGE__TRANSPORT__TRANSPORTS__{default_transport_name}__ENDPOINT"
+                ] = endpoint
+
+    @classmethod
+    def _load_config_from_env_variables(cls) -> dict[str, Any] | None:
+        config: dict[str, Any] = {}
+
+        # get os.environ.items only starting with OPENLINEAGE_ prefix and reverse sort
+        # to make sure that top-level keys have precedence
+        env_vars = sorted(
+            filter(lambda k: k[0].startswith(cls.DYNAMIC_ENV_VARS_PREFIX), os.environ.items()), reverse=True
+        )
+
+        for env_key, env_value in env_vars:
+            keys = env_key[len(cls.DYNAMIC_ENV_VARS_PREFIX) :].split("__")
+
+            # Parse value (try to parse as JSON, otherwise lowercase the value)
+            with contextlib.suppress(json.JSONDecodeError):
+                env_value = json.loads(env_value)  # noqa: PLW2901
+
+            cls._insert_into_config(config, keys, env_value)
+
+        if (transport_config := config.get("transport")) is None or transport_config.get("type") is None:
+            return None
+
+        return config
+
+    @staticmethod
+    def _insert_into_config(config: dict[str, Any], key_path: list[str], value: str) -> None:
+        keys = [key.lower() for key in key_path]
+
+        current = config
+        for key in keys[:-1]:
+            if key not in current:
+                current[key] = {}
+            current = current[key]
+
+        # Overwrite if key already exists
+        current[keys[-1]] = value

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/openlineage/client/constants.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-__version__ = "1.22.0"
+__version__ = "1.23.0"
 
 DEFAULT_TIMEOUT_MS = 5000
 DEFAULT_NAMESPACE_NAME = "default"

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/openlineage/client/generated/column_lineage_dataset.py

@@ -15,9 +15,15 @@ class ColumnLineageDatasetFacet(DatasetFacet):
     fields: dict[str, Fields]
     """Column level lineage that maps output fields into input fields used to evaluate them."""
 
+    dataset: list[InputField] | None = attr.field(factory=list)
+    """
+    Column level lineage that affects the whole dataset. This includes filtering, sorting, grouping
+    (aggregates), joining, window functions, etc.
+    """
+
     @staticmethod
     def _get_schema() -> str:
-        return "https://openlineage.io/spec/facets/1-1-0/ColumnLineageDatasetFacet.json#/$defs/ColumnLineageDatasetFacet"
+        return "https://openlineage.io/spec/facets/1-2-0/ColumnLineageDatasetFacet.json#/$defs/ColumnLineageDatasetFacet"
 
 
 @attr.define
@@ -35,6 +41,8 @@ class Fields(RedactMixin):
 
 @attr.define
 class InputField(RedactMixin):
+    """Represents a single dependency on some field (column)."""
+
     namespace: str
     """The input dataset namespace"""
 
@@ -47,6 +55,10 @@ class InputField(RedactMixin):
     transformations: list[Transformation] | None = attr.field(factory=list)
     _skip_redact: ClassVar[list[str]] = ["namespace", "name", "field"]
 
+    @staticmethod
+    def _get_schema() -> str:
+        return "https://openlineage.io/spec/facets/1-2-0/ColumnLineageDatasetFacet.json#/$defs/InputField"
+
 
 @attr.define
 class Transformation(RedactMixin):

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/openlineage/client/generated/datasource_dataset.py

@@ -21,6 +21,8 @@ class DatasourceDatasetFacet(DatasetFacet):
 
     @uri.validator
     def uri_check(self, attribute: str, value: str) -> None:  # noqa: ARG002
+        if value is None:
+            return
         from urllib.parse import urlparse
 
         urlparse(value)

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/openlineage/client/generated/nominal_time_run.py

@@ -39,6 +39,8 @@ class NominalTimeRunFacet(RunFacet):
 
     @nominalEndTime.validator
     def nominalendtime_check(self, attribute: str, value: str) -> None:  # noqa: ARG002
+        if value is None:
+            return
         from dateutil import parser
 
         parser.isoparse(value)

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/openlineage/client/transport/composite.py

@@ -35,7 +35,7 @@ class CompositeConfig(Config):
             transport will halt the emission process for subsequent transports.
     """
 
-    transports: list[dict[str, Any]] = attr.ib()
+    transports: list[dict[str, Any]] | dict[str, dict[str, Any]] = attr.ib()
     continue_on_failure: bool = attr.ib(default=True)
 
     @classmethod
@@ -63,7 +63,12 @@ class CompositeTransport(Transport):
         from openlineage.client.transport import get_default_factory
 
         transports = []
-        for transport_config in self.config.transports:
+        config_transports = self.config.transports
+        if isinstance(config_transports, dict):
+            config_transports = [
+                {**config, "name": name} for name, config in config_transports.items() if config
+            ]
+        for transport_config in config_transports:
             transports.append(get_default_factory().create(transport_config))
         return transports
 

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/openlineage/client/transport/factory.py

@@ -50,6 +50,8 @@ class DefaultTransportFactory(TransportFactory):
             msg = "You need to pass transport type in config."
             raise TypeError(msg) from None
 
+        transport_name = config.pop("name", None)
+
         transport_class_type_or_str = self.transports.get(transport_type, transport_type)
 
         if isinstance(transport_class_type_or_str, str):
@@ -68,4 +70,7 @@ class DefaultTransportFactory(TransportFactory):
             msg = f"Config {config_class} has to be class, and subclass of Config"
             raise TypeError(msg)
 
-        return transport_class(config_class.from_dict(config))  # type: ignore[call-arg]
+        transport: Transport = transport_class(config_class.from_dict(config))  # type: ignore[call-arg]
+        if transport_name and not transport.name:
+            transport.name = transport_name
+        return transport

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/openlineage/client/transport/http.py

@@ -87,6 +87,8 @@ class HttpConfig(Config):
     session: Session | None = attr.ib(default=None)
     # not set by TransportFactory
     adapter: HTTPAdapter | None = attr.ib(default=None)
+    # custom headers support
+    custom_headers: dict[str, str] = attr.ib(factory=dict)
 
     @classmethod
     def from_dict(cls, params: dict[str, Any]) -> HttpConfig:
@@ -150,6 +152,7 @@ class HttpTransport(Transport):
             self.session.headers["Content-Type"] = "application/json"
             auth_headers = self._auth_headers(config.auth)
             self.session.headers.update(auth_headers)
+            self.session.headers.update(config.custom_headers)
         self.timeout = config.timeout
         self.verify = config.verify
         self.compression = config.compression
@@ -168,6 +171,9 @@ class HttpTransport(Transport):
         http_client.HTTPConnection.debuglevel = 0
         body, headers = self._prepare_request(Serde.to_json(event))
 
+        # Update headers with custom headers from the config
+        headers.update(self.config.custom_headers)
+
         if self.session:
             resp = self.session.post(
                 url=urljoin(self.url, self.endpoint),
@@ -179,6 +185,7 @@ class HttpTransport(Transport):
         else:
             headers["Content-Type"] = "application/json"
             headers.update(self._auth_headers(self.config.auth))
+            headers.update(self.config.custom_headers)
             with Session() as session:
                 resp = session.post(
                     url=urljoin(self.url, self.endpoint),

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/openlineage/client/transport/kafka.py

@@ -41,6 +41,9 @@ class KafkaConfig(Config):
 
     @classmethod
     def from_dict(cls: type[_T], params: dict[str, Any]) -> _T:
+        # alias message_key to messageKey
+        if message_key := params.pop("message_key", None):
+            params["messageKey"] = params.get("messageKey") or message_key
         if "config" not in params:
             msg = "kafka `config` not passed to KafkaConfig"
             raise RuntimeError(msg)

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/openlineage/client/transport/transport.py

@@ -35,6 +35,7 @@ class Config:
 
 class Transport:
     kind: str | None = None
+    name: str | None = None
     config_class: type[Config] = Config
 
     def emit(self, event: Event) -> Any:

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/pyproject.toml

@@ -6,7 +6,7 @@ requires = [
 
 [project]
 name = "openlineage-python"
-version = "1.22.0"
+version = "1.23.0"
 description = "OpenLineage Python Client"
 readme = "README.md"
 keywords = [

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/tests/test_client.py

@@ -20,13 +20,16 @@ from openlineage.client.run import (
     RunEvent,
     RunState,
 )
-from openlineage.client.transport.http import HttpTransport
+from openlineage.client.transport.composite import CompositeTransport
+from openlineage.client.transport.console import ConsoleTransport
+from openlineage.client.transport.http import ApiKeyTokenProvider, HttpTransport, TokenProvider
 from openlineage.client.transport.noop import NoopTransport
 from openlineage.client.uuid import generate_new_uuid
 
 if TYPE_CHECKING:
     from pathlib import Path
 
+    from openlineage.client.transport.kafka import KafkaTransport
     from pytest_mock import MockerFixture
 
 
@@ -389,3 +392,193 @@ def test_http_transport_from_url_no_options() -> None:
     assert transport.timeout == timeout
     assert transport.verify is True
     assert transport.config.auth.api_key == "xxx"
+
+
+@patch.dict(
+    os.environ, {"OPENLINEAGE_URL": "http://example.com", "OPENLINEAGE__TRANSPORT__TYPE": "composite"}
+)
+def test_composite_transport_with_aliased_url() -> None:
+    transport: CompositeTransport = OpenLineageClient().transport
+    assert transport.kind == CompositeTransport.kind
+    assert len(transport.transports) == 1
+    assert transport.transports[0].kind == HttpTransport.kind
+    assert isinstance(transport.transports[0].config.auth, TokenProvider)
+
+
+@patch.dict(
+    os.environ,
+    {
+        "OPENLINEAGE_URL": "http://example.com",
+        "OPENLINEAGE_API_KEY": "random_key",
+        "OPENLINEAGE__TRANSPORT__TYPE": "composite",
+    },
+)
+def test_composite_transport_with_aliased_url_and_api_key() -> None:
+    transport: CompositeTransport = OpenLineageClient().transport
+    assert transport.kind == CompositeTransport.kind
+    assert len(transport.transports) == 1
+    assert transport.transports[0].kind == HttpTransport.kind
+    assert isinstance(transport.transports[0].config.auth, ApiKeyTokenProvider)
+    assert transport.transports[0].config.auth.api_key == "random_key"
+
+
+@patch.dict(
+    os.environ,
+    {
+        "OPENLINEAGE_URL": "http://example.com",
+        "OPENLINEAGE__TRANSPORT__TYPE": "composite",
+        "OPENLINEAGE__TRANSPORT__TRANSPORTS__ANOTHER__TYPE": "console",
+    },
+)
+def test_composite_transport_with_aliased_url_and_second_transport() -> None:
+    transport: CompositeTransport = OpenLineageClient().transport
+    assert transport.kind == CompositeTransport.kind
+    assert len(transport.transports) == 2  # noqa: PLR2004
+    assert transport.transports[0].kind == HttpTransport.kind
+
+
+@patch.dict(
+    os.environ,
+    {
+        "OPENLINEAGE_URL": "http://example.com",
+        "OPENLINEAGE__TRANSPORT__TYPE": "composite",
+        "OPENLINEAGE__TRANSPORT__TRANSPORTS__DEFAULT_HTTP": "{}",
+        "OPENLINEAGE__TRANSPORT__TRANSPORTS__ANOTHER__TYPE": "console",
+    },
+)
+def test_composite_transport_with_aliased_url_and_overriden_alias() -> None:
+    transport: CompositeTransport = OpenLineageClient().transport
+    assert transport.kind == CompositeTransport.kind
+    assert len(transport.transports) == 1
+    assert transport.transports[0].kind == ConsoleTransport.kind
+
+
+@patch.dict(
+    os.environ,
+    {
+        "OPENLINEAGE_URL": "http://example.com",
+        "OPENLINEAGE__TRANSPORT__TYPE": "composite",
+        "OPENLINEAGE__TRANSPORT__TRANSPORTS__DEFAULT_HTTP__TYPE": "console",
+        "OPENLINEAGE__TRANSPORT__TRANSPORTS__ANOTHER__TYPE": "console",
+    },
+)
+def test_openlineage_url_does_not_alias_when_transport_exists() -> None:
+    transport: CompositeTransport = OpenLineageClient().transport
+    assert transport.kind == CompositeTransport.kind
+    assert len(transport.transports) == 2  # noqa: PLR2004
+    assert transport.transports[0].kind == transport.transports[1].kind == ConsoleTransport.kind
+
+
+@patch.dict(
+    os.environ,
+    {
+        "OPENLINEAGE__TRANSPORT__TYPE": "kafka",
+        "OPENLINEAGE__TRANSPORT__TOPIC": "my_topic",
+        "OPENLINEAGE__TRANSPORT__CONFIG": '{"bootstrap.servers": "localhost:9092,another.host:9092", "acks": "all", "retries": 3}',  # noqa: E501
+        "OPENLINEAGE__TRANSPORT__FLUSH": "true",
+        "OPENLINEAGE__TRANSPORT__MESSAGE_KEY": "some-value",
+    },
+)
+def test_kafka_transport_configured_with_aliased_message_key() -> None:
+    transport: KafkaTransport = OpenLineageClient().transport
+    assert transport.message_key == "some-value"
+    assert transport.flush is True
+    assert transport.kafka_config.config == {
+        "bootstrap.servers": "localhost:9092,another.host:9092",
+        "acks": "all",
+        "retries": 3,
+    }
+
+
+class TestOpenLineageConfigLoader:
+    @pytest.mark.parametrize(
+        ("env_vars", "expected_config"),
+        [
+            (
+                {
+                    "OPENLINEAGE__TRANSPORT__TYPE": "http",
+                    "OPENLINEAGE__TRANSPORT__URL": "http://localhost:5050",
+                    "OPENLINEAGE__TRANSPORT__AUTH__API_KEY": "random_token",
+                },
+                {
+                    "transport": {
+                        "type": "http",
+                        "url": "http://localhost:5050",
+                        "auth": {"api_key": "random_token"},
+                    }
+                },
+            ),
+            (
+                {
+                    "OPENLINEAGE__TRANSPORT__TYPE": "composite",
+                    "OPENLINEAGE__TRANSPORT__TRANSPORTS__FIRST__TYPE": "http",
+                    "OPENLINEAGE__TRANSPORT__TRANSPORTS__FIRST__URL": "http://localhost:5050",
+                    "OPENLINEAGE__TRANSPORT__TRANSPORTS__SECOND__TYPE": "console",
+                },
+                {
+                    "transport": {
+                        "type": "composite",
+                        "transports": {
+                            "first": {"type": "http", "url": "http://localhost:5050"},
+                            "second": {"type": "console"},
+                        },
+                    }
+                },
+            ),
+            (
+                {"OPENLINEAGE__TRANSPORT": '{"type": "console"}', "OPENLINEAGE__TRANSPORT__TYPE": "http"},
+                {"transport": {"type": "console"}},
+            ),
+            (
+                {
+                    "OPENLINEAGE__TRANSPORT__TYPE": "kafka",
+                    "OPENLINEAGE__TRANSPORT__PROPERTIES": '{"key.serializer": "org.apache.kafka.common.serialization.StringSerializer"}',  # noqa: E501
+                },
+                {
+                    "transport": {
+                        "type": "kafka",
+                        "properties": {
+                            "key.serializer": "org.apache.kafka.common.serialization.StringSerializer"
+                        },
+                    }
+                },
+            ),
+            (
+                {
+                    "OPENLINEAGE__TRANSPORT__TYPE": "test",
+                    "OPENLINEAGE__TRANSPORT__MY_NAME__LIST": '["first", "SeCond"]',
+                },
+                {"transport": {"my_name": {"list": ["first", "SeCond"]}, "type": "test"}},
+            ),
+            (
+                {
+                    "OPENLINEAGE__TRANSPORT__TYPE": "http",
+                    "OPENLINEAGE__TRANSPORT__URL": "http://localhost:5050",
+                    "OPENLINEAGE__TRANSPORT__AUTH__API_KEY": "random_token",
+                    "OPENLINEAGE__TRANSPORT__AUTH__COMPRESSION": "gzip",
+                },
+                {
+                    "transport": {
+                        "type": "http",
+                        "url": "http://localhost:5050",
+                        "auth": {"api_key": "random_token", "compression": "gzip"},
+                    }
+                },
+            ),
+            (
+                {
+                    "OPENLINEAGE__TRANSPORT": '{"type": "console"}',
+                    "OPENLINEAGE__TRANSPORT__TYPE": "http",
+                    "OPENLINEAGE__TRANSPORT__URL": "http://localhost:5050",
+                    "OPENLINEAGE__TRANSPORT__AUTH__API_KEY": "random_token",
+                    "OPENLINEAGE__TRANSPORT__AUTH__COMPRESSION": "gzip",
+                },
+                {"transport": {"type": "console"}},
+            ),
+        ],
+    )
+    @patch.dict(os.environ, {})
+    def test_config_loader(self, env_vars, expected_config):
+        with patch.dict(os.environ, env_vars):
+            config = OpenLineageClient._load_config_from_env_variables()  # noqa: SLF001
+            assert config == expected_config

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/tests/test_composite.py

@@ -67,6 +67,16 @@ def test_composite_transport_create_transports():
     assert transport.transports[0].kind == "fake"
 
 
+def test_composite_transport_from_dict_config():
+    config = CompositeConfig(
+        transports={"fake_trans": {"type": "tests.transport.FakeTransport"}}, continue_on_failure=True
+    )
+    transport = CompositeTransport(config)
+    assert len(transport.transports) == 1
+    assert transport.transports[0].kind == "fake"
+    assert transport.transports[0].name == "fake_trans"
+
+
 @mock.patch("openlineage.client.transport.get_default_factory")
 def test_transports(mock_factory):
     mock_transport1 = MagicMock(spec=Transport)

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/tests/test_facet_v2.py

@@ -43,6 +43,11 @@ def test_set_producer_from_argument():
     assert facet._producer == "http://another.producer"  # noqa: SLF001
 
 
+def test_optional_attributed_not_validated():
+    """Don't pass optional value with validator."""
+    nominal_time_run.NominalTimeRunFacet(nominalStartTime="2020-12-17T03:00:00.001Z")
+
+
 def test_custom_facet() -> None:
     session = mock.MagicMock()
     client = OpenLineageClient(url="http://example.com", session=session)

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/tests/test_http.py

@@ -4,12 +4,14 @@ from __future__ import annotations
 
 import datetime
 import gzip
+import os
 from typing import TYPE_CHECKING
+from unittest.mock import MagicMock, patch
 
 from openlineage.client import OpenLineageClient
 from openlineage.client.run import Job, Run, RunEvent, RunState
 from openlineage.client.serde import Serde
-from openlineage.client.transport.http import HttpCompression, HttpConfig, HttpTransport
+from openlineage.client.transport.http import ApiKeyTokenProvider, HttpCompression, HttpConfig, HttpTransport
 from openlineage.client.uuid import generate_new_uuid
 from requests import Session
 
@@ -163,3 +165,119 @@ def test_http_config_configs_session() -> None:
         assert not hasattr(config.auth, "api_key")
         assert config.session is s
         assert config.adapter is None
+
+
+@patch("requests.Session.post")
+def test_http_transport_custom_headers_applied(mock_post):
+    custom_headers = {"X-Custom-Header": "CustomValue", "X-Another-Header": "AnotherValue"}
+
+    config = HttpConfig(url="http://example.com", custom_headers=custom_headers)
+
+    transport = HttpTransport(config)
+    mock_event = MagicMock()
+
+    with patch("openlineage.client.serde.Serde.to_json", return_value='{"mock": "event"}'):
+        transport.emit(mock_event)
+
+    mock_post.assert_called_once()
+    _, kwargs = mock_post.call_args
+    headers = kwargs["headers"]
+
+    assert headers["X-Custom-Header"] == "CustomValue"
+    assert headers["X-Another-Header"] == "AnotherValue"
+
+
+@patch("requests.Session.post")
+def test_http_transport_auth_and_custom_headers_applied(mock_post):
+    custom_headers = {"X-Custom-Header": "CustomValue"}
+    auth_token = "Bearer test_token"  # noqa: S105
+
+    # Set up config with an ApiKeyTokenProvider
+    config = HttpConfig(
+        url="http://example.com",
+        auth=ApiKeyTokenProvider({"api_key": "test_token"}),
+        custom_headers=custom_headers,
+    )
+
+    transport = HttpTransport(config)
+    mock_event = MagicMock()
+
+    with patch("openlineage.client.serde.Serde.to_json", return_value='{"mock": "event"}'):
+        transport.emit(mock_event)
+
+    mock_post.assert_called_once()
+    _, kwargs = mock_post.call_args
+    headers = kwargs["headers"]
+
+    assert headers["Authorization"] == auth_token
+    assert headers["X-Custom-Header"] == "CustomValue"
+
+
+@patch("requests.Session.post")
+def test_http_transport_no_custom_headers(mock_post):
+    config = HttpConfig(url="http://example.com")
+    transport = HttpTransport(config)
+    mock_event = MagicMock()
+
+    with patch("openlineage.client.serde.Serde.to_json", return_value='{"mock": "event"}'):
+        transport.emit(mock_event)
+
+    mock_post.assert_called_once()
+    _, kwargs = mock_post.call_args
+    headers = kwargs["headers"]
+
+    # Only the Content-Type header should be set if no custom headers
+    assert headers["Content-Type"] == "application/json"
+    assert "X-Custom-Header" not in headers
+
+
+@patch("requests.Session.post")
+def test_http_transport_compression_with_custom_headers(mock_post):
+    custom_headers = {"X-Custom-Header": "CustomValue"}
+
+    config = HttpConfig(
+        url="http://example.com", compression=HttpCompression.GZIP, custom_headers=custom_headers
+    )
+    transport = HttpTransport(config)
+    mock_event = MagicMock()
+
+    with patch("openlineage.client.serde.Serde.to_json", return_value='{"mock": "event"}'), patch(
+        "gzip.compress", return_value=b"compressed_data"
+    ):
+        transport.emit(mock_event)
+
+    mock_post.assert_called_once()
+    _, kwargs = mock_post.call_args
+    headers = kwargs["headers"]
+    data = kwargs["data"]
+
+    assert headers["X-Custom-Header"] == "CustomValue"
+    assert headers["Content-Encoding"] == "gzip"
+    assert data == b"compressed_data"
+
+
+@patch("requests.Session.post")
+@patch.dict(
+    os.environ,
+    {
+        "OPENLINEAGE__TRANSPORT__TYPE": "http",
+        "OPENLINEAGE__TRANSPORT__URL": "http://example.com",
+        "OPENLINEAGE__TRANSPORT__CUSTOM_HEADERS__CUSTOM_HEADER": "FIRST",
+        "OPENLINEAGE__TRANSPORT__CUSTOM_HEADERS__ANOTHER_HEADER": "second",
+    },
+)
+def test_http_transport_with_custom_headers_from_env_vars(mock_post):
+    transport = OpenLineageClient().transport
+    mock_event = MagicMock()
+
+    with patch("openlineage.client.serde.Serde.to_json", return_value='{"mock": "event"}'), patch(
+        "gzip.compress", return_value=b"compressed_data"
+    ):
+        transport.emit(mock_event)
+
+    mock_post.assert_called_once()
+    _, kwargs = mock_post.call_args
+    headers = kwargs["headers"]
+
+    assert headers["custom_header"] == "FIRST"
+    assert headers["another_header"] == "second"

{openlineage_python-1.22.0 → openlineage_python-1.23.0}/tests/test_kafka.py

@@ -106,6 +106,22 @@ def test_kafka_loads_partial_config_with_defaults() -> None:
     assert config.flush is True
 
 
+def test_kafka_config_accepts_aliased_message_key() -> None:
+    config = KafkaConfig.from_dict(
+        {
+            "type": "kafka",
+            "topic": "my_topic",
+            "message_key": "some-value",
+            "flush": True,
+            "config": {"bootstrap.servers": "localhost:9092,another.host:9092", "acks": "all", "retries": 3},
+        }
+    )
+
+    assert config.messageKey == "some-value"
+    assert config.topic == "my_topic"
+    assert config.flush is True
+
+
 def test_kafka_load_config_fails_on_no_config() -> None:
     with pytest.raises(TypeError):
         KafkaConfig.from_dict(