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

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

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

{openlineage_python-1.22.0 → openlineage_python-1.24.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
@@ -9,16 +10,28 @@ from typing import TYPE_CHECKING, Any, TypeVar, Union, cast
 
 import attr
 import yaml
-from openlineage.client.filter import Filter, create_filter
+from openlineage.client.filter import Filter, FilterConfig, create_filter
 from openlineage.client.serde import Serde
+from openlineage.client.utils import deep_merge_dicts
 
 if TYPE_CHECKING:
     from requests import Session
     from requests.adapters import HTTPAdapter
 
+import contextlib
+
 from openlineage.client import event_v2
+from openlineage.client.facets import FacetsConfig
+from openlineage.client.generated.environment_variables_run import (
+    EnvironmentVariable,
+    EnvironmentVariablesRunFacet,
+)
 from openlineage.client.run import DatasetEvent, JobEvent, RunEvent
-from openlineage.client.transport import Transport, TransportFactory, get_default_factory
+from openlineage.client.transport import (
+    Transport,
+    TransportFactory,
+    get_default_factory,
+)
 from openlineage.client.transport.http import HttpConfig, HttpTransport, create_token_provider
 from openlineage.client.transport.noop import NoopConfig, NoopTransport
 
@@ -35,11 +48,32 @@ class OpenLineageClientOptions:
     adapter: HTTPAdapter = attr.ib(default=None)
 
 
+@attr.s
+class OpenLineageConfig:
+    transport: dict[str, Any] | None = attr.ib(factory=dict)
+    facets: FacetsConfig = attr.ib(factory=FacetsConfig)
+    filters: list[FilterConfig] = attr.ib(factory=list)
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> OpenLineageConfig:
+        config = cls()
+        if "transport" in params:
+            config.transport = params["transport"]
+        if "facets" in params:
+            config.facets = FacetsConfig(**params["facets"])
+        if "filters" in params:
+            config.filters = [FilterConfig(**filter_config) for filter_config in params["filters"]]
+        return config
+
+
 log = logging.getLogger(__name__)
 _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,
@@ -47,6 +81,8 @@ class OpenLineageClient:
         session: Session | None = None,
         transport: Transport | None = None,
         factory: TransportFactory | None = None,
+        *,
+        config: dict[str, str] | None = None,
     ) -> None:
         # Set parent's logging level if environment variable is present
         custom_logging_level = os.getenv("OPENLINEAGE_CLIENT_LOGGING", None)
@@ -62,7 +98,11 @@ class OpenLineageClient:
 
         # Make config ellipsis - as a guard value to not try to
         # reload yaml each time config is referred to.
-        self._config: dict[str, dict[str, str]] | None = None
+        self._config: OpenLineageConfig | None = None
+
+        self.user_defined_config: dict[str, str] | None = config
+
+        self._alias_env_vars()
 
         self.transport = self._resolve_transport(
             url=url, options=options, session=session, transport=transport, factory=factory
@@ -70,7 +110,7 @@ class OpenLineageClient:
         log.info("OpenLineageClient will use `%s` transport", self.transport.kind)
 
         self._filters: list[Filter] = []
-        for conf in self.config.get("filters", []):
+        for conf in self.config.filters:
             _filter = create_filter(conf)
             if _filter:
                 self._filters.append(_filter)
@@ -86,7 +126,7 @@ class OpenLineageClient:
 
     @classmethod
     def from_dict(cls: type[_T], config: dict[str, str]) -> _T:
-        return cls(transport=get_default_factory().create(config=config))
+        return cls(config=config)
 
     def filter_event(
         self,
@@ -114,6 +154,8 @@ class OpenLineageClient:
             log.debug("OpenLineage event has been filtered out and will not be emitted.")
             return
 
+        event = self.add_environment_facets(event)
+
         if log.isEnabledFor(logging.DEBUG):
             val = Serde.to_json(event).encode("utf-8")
             log.debug("OpenLineageClient will emit event %s", val)
@@ -121,14 +163,36 @@ class OpenLineageClient:
         log.debug("OpenLineage event successfully emitted.")
 
     @property
-    def config(self) -> dict[str, Any]:
-        """Content of OpenLineage YAML config file."""
+    def config(self) -> OpenLineageConfig:
+        """
+        Retrieves the OpenLineage configuration.
+
+        This property method returns the content of the OpenLineage YAML config file.
+        The configuration is determined by merging sources in the following order of precedence:
+        1. User-defined configuration passed to the client constructor.
+        2. YAML config file located in one of the following paths:
+        - Path specified by the `OPENLINEAGE_CONFIG` environment variable.
+        - Current working directory.
+        - `$HOME/.openlineage`.
+        3. Environment variables with the `OPENLINEAGE__` prefix.
+        If the configuration is not already loaded, it will be constructed by merging the above sources.
+        In case of a TypeError during the parsing of the configuration, a ValueError will be raised indicating
+        that the structure of the config does not match the expected format.
+        """
         if self._config is None:
-            config_path = self._find_yaml_config_path()
-            if config_path:
-                self._config = self._get_config_file_content(config_path)
-            else:
-                self._config = {}
+            config_dict: dict[str, Any] = {}
+            if self.user_defined_config:
+                config_dict = self.user_defined_config
+            if config_path := self._find_yaml_config_path():
+                config_dict = deep_merge_dicts(self._get_config_file_content(config_path), config_dict)
+            if config_from_env_vars := self._load_config_from_env_variables():
+                config_dict = deep_merge_dicts(config_from_env_vars, config_dict)
+            try:
+                self._config = OpenLineageConfig.from_dict(config_dict)
+            except TypeError as e:
+                # raise exception that structure of the config does not match
+                msg = "Failed to parse OpenLineage config."
+                raise ValueError(msg) from e
         return self._config
 
     def _resolve_transport(self, **kwargs: Any) -> Transport:
@@ -156,9 +220,9 @@ class OpenLineageClient:
             return cast(Transport, kwargs["transport"])
 
         # 3. Check if transport configuration is provided in YAML config file
-        if self.config.get("transport"):
+        if self.config.transport and self.config.transport.get("type"):
             factory = kwargs.get("factory") or get_default_factory()
-            return factory.create(self.config["transport"])
+            return factory.create(self.config.transport)
 
         # 4. Check legacy HTTP transport initialization with url and options
         if kwargs.get("url"):
@@ -245,3 +309,90 @@ 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)
+
+        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
+
+    def add_environment_facets(self, event: Event) -> Event:
+        """
+        Adds environment variables as facets to the event object.
+        """
+        if isinstance(event, RunEvent) and (env_vars := self._collect_environment_variables()):
+            event.run.facets["environmentVariables"] = EnvironmentVariablesRunFacet(
+                environmentVariables=[
+                    EnvironmentVariable(name=name, value=value) for name, value in env_vars.items()
+                ]
+            )
+        return event
+
+    def _collect_environment_variables(self) -> dict[str, str]:
+        """
+        Collects and returns a dictionary of relevant environment variables.
+        """
+        filtered_vars = {k: v for k, v in os.environ.items() if k in self.config.facets.environment_variables}
+        missing_vars = set(self.config.facets.environment_variables) - set(filtered_vars)
+        if missing_vars:
+            log.warning(
+                "The following environment variables are missing: %s when adding to OpenLineage event",
+                missing_vars,
+            )
+        return filtered_vars

{openlineage_python-1.22.0 → openlineage_python-1.24.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.24.0"
 
 DEFAULT_TIMEOUT_MS = 5000
 DEFAULT_NAMESPACE_NAME = "default"

{openlineage_python-1.22.0 → openlineage_python-1.24.0}/openlineage/client/facet_v2.py

@@ -9,6 +9,7 @@ from openlineage.client.generated import (
     datasource_dataset,
     documentation_dataset,
     documentation_job,
+    environment_variables_run,
     error_message_run,
     external_query_run,
     extraction_error_run,
@@ -54,6 +55,7 @@ __all__ = [
     "datasource_dataset",
     "documentation_dataset",
     "documentation_job",
+    "environment_variables_run",
     "error_message_run",
     "external_query_run",
     "extraction_error_run",

openlineage_python-1.24.0/openlineage/client/facets.py

@@ -0,0 +1,10 @@
+# Copyright 2018-2024 contributors to the OpenLineage project
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import attr
+
+
+@attr.s
+class FacetsConfig:
+    environment_variables: list[str] = attr.ib(factory=list)

{openlineage_python-1.22.0 → openlineage_python-1.24.0}/openlineage/client/filter.py

@@ -5,12 +5,20 @@ from __future__ import annotations
 import re
 import typing
 
+import attr
 from openlineage.client.event_v2 import RunEvent as RunEvent_v2
 from openlineage.client.run import RunEvent
 
 RunEventType = typing.Union[RunEvent, RunEvent_v2]
 
 
+@attr.s
+class FilterConfig:
+    type: str | None = attr.ib(default=None)
+    match: str | None = attr.ib(default=None)
+    regex: str | None = attr.ib(default=None)
+
+
 class Filter:
     def filter_event(self, event: RunEventType) -> RunEventType | None:
         ...
@@ -36,12 +44,12 @@ class RegexFilter(Filter):
         return event
 
 
-def create_filter(conf: dict[str, str]) -> Filter | None:
-    if "type" not in conf:
+def create_filter(conf: FilterConfig) -> Filter | None:
+    if not conf.type:
         return None
     # Switch in 3.10 🙂
-    if conf["type"] == "exact":
-        return ExactMatchFilter(match=conf["match"])
-    if conf["type"] == "regex":
-        return RegexFilter(regex=conf["regex"])
+    if conf.type == "exact" and conf.match:
+        return ExactMatchFilter(match=conf.match)
+    if conf.type == "regex" and conf.regex:
+        return RegexFilter(regex=conf.regex)
     return None

{openlineage_python-1.22.0 → openlineage_python-1.24.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.24.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.24.0/openlineage/client/generated/environment_variables_run.py

@@ -0,0 +1,35 @@
+# Copyright 2018-2024 contributors to the OpenLineage project
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+import attr
+from openlineage.client.generated.base import RunFacet
+from openlineage.client.utils import RedactMixin
+
+
+@attr.define
+class EnvironmentVariable(RedactMixin):
+    name: str
+    """The name of the environment variable."""
+
+    value: str
+    """The value of the environment variable."""
+
+    _skip_redact: ClassVar[list[str]] = ["name"]
+
+    @staticmethod
+    def _get_schema() -> str:
+        return "https://openlineage.io/spec/facets/1-0-0/EnvironmentVariablesRunFacet.json#/$defs/EnvironmentVariable"
+
+
+@attr.define
+class EnvironmentVariablesRunFacet(RunFacet):
+    environmentVariables: list[EnvironmentVariable]  # noqa: N815
+    """The environment variables for the run."""
+
+    @staticmethod
+    def _get_schema() -> str:
+        return "https://openlineage.io/spec/facets/1-0-0/EnvironmentVariablesRunFacet.json#/$defs/EnvironmentVariablesRunFacet"

{openlineage_python-1.22.0 → openlineage_python-1.24.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.24.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.24.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.24.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.24.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.24.0}/openlineage/client/transport/transport.py

@@ -20,6 +20,8 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, TypeVar
 
+import attr
+
 if TYPE_CHECKING:
     from openlineage.client.client import Event
 
@@ -27,6 +29,7 @@ if TYPE_CHECKING:
 _T = TypeVar("_T", bound="Config")
 
 
+@attr.s
 class Config:
     @classmethod
     def from_dict(cls: type[_T], params: dict[str, Any]) -> _T:  # noqa: ARG003
@@ -35,6 +38,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.24.0}/openlineage/client/utils.py

@@ -36,6 +36,25 @@ def get_only_specified_fields(clazz: type[Any], params: dict[str, Any]) -> dict[
     return {key: value for key, value in params.items() if key in field_keys}
 
 
+def deep_merge_dicts(dict1: dict[Any, Any], dict2: dict[Any, Any]) -> dict[Any, Any]:
+    """Deep merges two dictionaries.
+
+    This function merges two dictionaries while handling nested dictionaries.
+    For keys that exist in both dictionaries, the values from dict2 take precedence.
+    If a key exists in both dictionaries and the values are dictionaries themselves,
+    they are merged recursively.
+    This function merges only dictionaries. If key is of different type, e.g. list
+    it does not work properly.
+    """
+    merged = dict1.copy()
+    for k, v in dict2.items():
+        if k in merged and isinstance(v, dict):
+            merged[k] = deep_merge_dicts(merged.get(k, {}), v)
+        else:
+            merged[k] = v
+    return merged
+
+
 class RedactMixin:
     _skip_redact: ClassVar[list[str]] = []
 

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

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

{openlineage_python-1.22.0 → openlineage_python-1.24.0}/redact_fields.yml

@@ -92,6 +92,13 @@
   classes:
     - class_name: DocumentationJobFacet
       redact_fields: []
+- module: environment_variables_run
+  classes:
+    - class_name: EnvironmentVariable
+      redact_fields:
+        - name
+    - class_name: EnvironmentVariablesRunFacet
+      redact_fields: []
 - module: error_message_run
   classes:
     - class_name: ErrorMessageRunFacet