openlineage-python 1.30.0__tar.gz → 1.31.0__tar.gz

{openlineage_python-1.30.0 → openlineage_python-1.31.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openlineage-python
-Version: 1.30.0
+Version: 1.31.0
 Summary: OpenLineage Python Client
 Author-email: OpenLineage <info@openlineage.io>
 Keywords: openlineage
@@ -23,6 +23,10 @@ Requires-Dist: beautifulsoup4>=4.12.0; extra == 'docs'
 Requires-Dist: pydata-sphinx-theme>=0.13; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints>=2.0.0; extra == 'docs'
 Requires-Dist: sphinx>=7.0; extra == 'docs'
+Provides-Extra: generator
+Requires-Dist: click; extra == 'generator'
+Requires-Dist: datamodel-code-generator==0.25.4; extra == 'generator'
+Requires-Dist: ruff==0.3.5; extra == 'generator'
 Provides-Extra: kafka
 Requires-Dist: confluent-kafka>=2.1.1; extra == 'kafka'
 Provides-Extra: msk-iam

{openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/client.py

@@ -163,6 +163,11 @@ class OpenLineageClient:
         if type(event) not in get_args(Event):
             msg = "`emit` only accepts RunEvent, DatasetEvent, JobEvent classes"
             raise ValueError(msg)
+
+        if log.isEnabledFor(logging.DEBUG):
+            val = Serde.to_json(event).encode("utf-8")
+            log.debug("OpenLineageClient will *try* to emit event %s", val)
+
         if not self.transport:
             log.error("Tried to emit OpenLineage event, but transport is not configured.")
             return
@@ -176,9 +181,6 @@ class OpenLineageClient:
         event = self.add_environment_facets(event)
         event = self.update_event_tags_facets(event)
 
-        if log.isEnabledFor(logging.DEBUG):
-            val = Serde.to_json(event).encode("utf-8")
-            log.debug("OpenLineageClient will emit event %s", val)
         self.transport.emit(event)
         log.debug("OpenLineage event successfully emitted.")
 
@@ -420,15 +422,21 @@ class OpenLineageClient:
         """
         Creates or updates job and run tag facets based on user-supplied environment variables
         """
+        run_event_types = (RunEvent, event_v2.RunEvent)
+        run_and_job_event_types = (RunEvent, event_v2.RunEvent, JobEvent, event_v2.JobEvent)
         tags_job = self.config.tags.job
-        if (isinstance(event, (JobEvent, RunEvent))) and tags_job:
+        if isinstance(event, run_and_job_event_types) and tags_job:
+            # Ensure facets exists
+            event.job.facets = {} if not event.job.facets else event.job.facets
             tags_facet = event.job.facets.get("tags", TagsJobFacet())
-            event.job.facets["tags"] = self._update_tag_facet(tags_facet, tags_job)  # type: ignore [arg-type]
+            event.job.facets["tags"] = self._update_tag_facet(tags_facet, tags_job)  # type: ignore [arg-type, assignment]
 
         tags_run = self.config.tags.run
-        if (isinstance(event, RunEvent)) and tags_run:
+        if isinstance(event, run_event_types) and tags_run:
+            # Ensure facets exists
+            event.run.facets = {} if not event.run.facets else event.run.facets
             tags_facet = event.run.facets.get("tags", TagsRunFacet())
-            event.run.facets["tags"] = self._update_tag_facet(tags_facet, tags_run)  # type: ignore [arg-type]
+            event.run.facets["tags"] = self._update_tag_facet(tags_facet, tags_run)  # type: ignore [arg-type, assignment]
 
         return event
 

{openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/constants.py

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

{openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/parent_run.py

@@ -28,11 +28,12 @@ class ParentRunFacet(RunFacet):
 
     run: Run
     job: Job
+    root: Root | None = attr.field(default=None)
     _additional_skip_redact: ClassVar[list[str]] = ["job", "run"]
 
     @staticmethod
     def _get_schema() -> str:
-        return "https://openlineage.io/spec/facets/1-0-1/ParentRunFacet.json#/$defs/ParentRunFacet"
+        return "https://openlineage.io/spec/facets/1-1-0/ParentRunFacet.json#/$defs/ParentRunFacet"
 
     @classmethod
     def create(cls, runId: str, namespace: str, name: str) -> ParentRunFacet:  # noqa: N803
@@ -46,6 +47,46 @@ class ParentRunFacet(RunFacet):
         return cls(run=Run(runId=runId), job=Job(namespace=namespace, name=name))
 
 
+@attr.define
+class Root(RedactMixin):
+    run: RootRun
+    job: RootJob
+    _skip_redact: ClassVar[list[str]] = ["run", "job"]
+
+
+@attr.define
+class RootJob(RedactMixin):
+    namespace: str
+    """The namespace containing root job"""
+
+    name: str
+    """The unique name containing root job within that namespace"""
+
+    _skip_redact: ClassVar[list[str]] = ["namespace", "name"]
+
+    @staticmethod
+    def _get_schema() -> str:
+        return "https://openlineage.io/spec/facets/1-1-0/ParentRunFacet.json#/$defs/RootJob"
+
+
+@attr.define
+class RootRun(RedactMixin):
+    runId: str = attr.field()  # noqa: N815
+    """The globally unique ID of the root run associated with the root job."""
+
+    _skip_redact: ClassVar[list[str]] = ["runId"]
+
+    @staticmethod
+    def _get_schema() -> str:
+        return "https://openlineage.io/spec/facets/1-1-0/ParentRunFacet.json#/$defs/RootRun"
+
+    @runId.validator
+    def runid_check(self, attribute: str, value: str) -> None:  # noqa: ARG002
+        from uuid import UUID
+
+        UUID(value)
+
+
 @attr.define
 class Run(RedactMixin):
     runId: str = attr.field()  # noqa: N815

openlineage_python-1.31.0/openlineage/client/generator/base.py

@@ -0,0 +1,257 @@
+# Copyright 2018-2025 contributors to the OpenLineage project
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import copy
+import json
+import logging
+import os
+import pathlib
+import re
+import subprocess
+import tempfile
+from typing import Any
+
+from datamodel_code_generator import DataModelType, PythonVersion
+from datamodel_code_generator.imports import Import
+from datamodel_code_generator.model import get_data_model_types
+from datamodel_code_generator.model import pydantic as pydantic_model
+from datamodel_code_generator.parser.jsonschema import JsonSchemaParser
+from datamodel_code_generator.types import Types
+
+log = logging.getLogger(__name__)
+
+
+def camel_to_snake(string: str) -> str:
+    """Convert camel_case string to SnakeCase"""
+    _under_score_1 = re.compile(r"([^_])([A-Z][a-z]+)")
+    _under_score_2 = re.compile("([a-z0-9])([A-Z])")
+    subbed = _under_score_1.sub(r"\1_\2", string)
+    return _under_score_2.sub(r"\1_\2", subbed).lower()
+
+
+# import attrs instead of dataclass
+DATA_MODELS_TYPES = get_data_model_types(
+    DataModelType.DataclassesDataclass, target_python_version=PythonVersion.PY_38
+)
+NEW_MODEL = DATA_MODELS_TYPES.data_model
+NEW_MODEL.DEFAULT_IMPORTS = (Import.from_full_path("attr"),)
+
+# custom code
+SET_PRODUCER_CODE = """
+PRODUCER = DEFAULT_PRODUCER
+
+def set_producer(producer: str) -> None:
+    global PRODUCER  # noqa: PLW0603
+    PRODUCER = producer
+"""
+
+# locations definitions
+FILE_LOCATION = pathlib.Path(__file__).resolve().parent
+TEMPLATES_LOCATION = FILE_LOCATION / "templates"
+
+# Global dictionaries intended to be imported by other modules
+# Contains schema URLs and base IDs parsed from specification files
+SCHEMA_URLS: dict[str, str] = {}
+BASE_IDS: dict[str, str] = {}
+
+
+def deep_merge_dicts(dict1: dict[str, Any], dict2: dict[str, Any]) -> dict[str, 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
+
+
+def parse_additional_data(spec: dict[str, Any], file_name: str) -> None:
+    """Parse additional data from spec files.
+
+    Parses:
+        * schema URLs
+        * base IDs
+
+    Updates the module-level SCHEMA_URLS and BASE_IDS dictionaries
+    that can be imported by other modules.
+    """
+    base_id = spec["$id"]
+    for name, _ in spec["$defs"].items():
+        SCHEMA_URLS[name] = f"{base_id}#/$defs/{name}"
+        BASE_IDS[file_name] = spec["$id"]
+
+
+def load_specs(base_spec_location: pathlib.Path, facets_spec_location: pathlib.Path) -> list[pathlib.Path]:
+    """Load base `OpenLineage.json` and other facets' spec files"""
+    locations = []
+    if facets_spec_location.is_dir():
+        facets_spec_files = sorted(pathlib.Path(os.path.abspath(facets_spec_location)).glob("*.json"))
+    else:
+        facets_spec_files = [pathlib.Path(os.path.abspath(facets_spec_location))]
+
+    file_specs = [
+        base_spec_location.resolve(),  # Base spec must be the first element
+        *facets_spec_files,
+    ]
+    for file_spec in file_specs:
+        spec = json.loads(file_spec.read_text())
+        parse_additional_data(spec, file_spec.name)
+        locations.append(file_spec)
+
+    return locations
+
+
+def parse_and_generate(
+    locations: list[pathlib.Path], extra_template_data: dict[str, Any] | None = None
+) -> dict[str, Any]:
+    """Parse and generate data models from a given specification."""
+    temporary_locations = []
+
+    current_dir = pathlib.Path.cwd()
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp_directory = pathlib.Path(tmp).resolve()
+        for location in locations:
+            tmp_location = (tmp_directory / location.name).resolve()
+            tmp_location.write_text(location.read_text())
+            temporary_locations.append(tmp_location)
+
+        os.chdir(tmp_directory)
+        # first parse OpenLineage.json
+        parser = JsonSchemaParser(
+            source=temporary_locations[:1],
+            data_model_type=NEW_MODEL,
+            data_model_root_type=DATA_MODELS_TYPES.root_model,
+            data_model_field_type=DATA_MODELS_TYPES.field_model,
+            data_type_manager_type=DATA_MODELS_TYPES.data_type_manager,
+            dump_resolve_reference_action=DATA_MODELS_TYPES.dump_resolve_reference_action,
+            special_field_name_prefix="",
+            use_schema_description=True,
+            field_constraints=True,
+            use_union_operator=True,
+            use_standard_collections=True,
+            base_class="openlineage.client.utils.RedactMixin",
+            class_name="ClassToBeSkipped",
+            use_field_description=True,
+            use_double_quotes=True,
+            keep_model_order=True,
+            custom_template_dir=TEMPLATES_LOCATION,
+            extra_template_data=extra_template_data,  # type: ignore[arg-type]
+            additional_imports=[
+                "typing.ClassVar",
+                "typing.Any",
+                "typing.cast",
+                "openlineage.client.constants.DEFAULT_PRODUCER",
+            ],
+        )
+
+        # keep information about uuid and date-time formats
+        # this is going to be changed back to str type hint in jinja template
+        data_type_manager: pydantic_model.DataTypeManager = parser.data_type_manager  # type: ignore[assignment]
+        uuid_type = copy.deepcopy(data_type_manager.type_map[Types.uuid])
+        uuid_type.type = "uuid"
+        data_type_manager.type_map[Types.uuid] = uuid_type
+
+        date_time_type = copy.deepcopy(data_type_manager.type_map[Types.date_time])
+        date_time_type.type = "date-time"
+        data_type_manager.type_map[Types.date_time] = date_time_type
+
+        uri_type = copy.deepcopy(data_type_manager.type_map[Types.date_time])
+        uri_type.type = "uri"
+        data_type_manager.type_map[Types.uri] = uri_type
+
+        parser.parse(format_=False)
+
+        # parse rest of spec
+        parser.source = temporary_locations[1:]
+
+        # change paths so that parser sees base objects as local
+        parser.model_resolver.references = {
+            k.replace("OpenLineage.json", BASE_IDS["OpenLineage.json"]): v
+            for k, v in parser.model_resolver.references.items()
+        }
+
+        output = parser.parse(format_=False)
+
+    # go back to original directory
+    os.chdir(current_dir)
+
+    return output  # type: ignore[return-value]
+
+
+def separate_imports(code: str) -> tuple[str, str]:
+    """Separates a Python script code (as string) into imports and the rest."""
+    imports_section = []
+    rest_of_code = []
+    in_import = False
+    for line in code.splitlines():
+        if line.startswith(("import ", "from ")):
+            in_import = True
+        elif in_import and line.strip() == "":
+            in_import = False
+        if in_import:
+            imports_section.append(line)
+        else:
+            rest_of_code.append(line)
+    return ("\n".join(imports_section), "\n".join(rest_of_code))
+
+
+def format_and_save_output(
+    output: str,
+    location: pathlib.Path,
+    add_set_producer_code: bool | None = False,
+    header: str | None = None,
+) -> None:
+    """Adjust and format generated file."""
+
+    # save temporary file to the same directory that output
+    # so that ruff receives same rules
+    with tempfile.NamedTemporaryFile(
+        "w", prefix=location.stem.lower(), suffix=".py", dir=location.parent, delete=False
+    ) as tmpfile:
+        if header:
+            tmpfile.write(header + "\n")
+        output = output.replace("from .OpenLineage", "from openlineage.client.generated.base")
+        imports_section, rest_of_code = separate_imports(output)
+        tmpfile.write(imports_section)
+        if add_set_producer_code:
+            tmpfile.write(SET_PRODUCER_CODE)
+        tmpfile.write(rest_of_code)
+        tmpfile.flush()
+
+    # run ruff lint
+    with subprocess.Popen(
+        args=["ruff", "check", tmpfile.name, "--fix"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        close_fds=True,
+    ) as lint_process:
+        if lint_process.returncode:
+            log.warning("Ruff lint failed: %s", lint_process.returncode)
+
+    # run ruff format
+    with subprocess.Popen(
+        args=["ruff", "format", tmpfile.name],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        close_fds=True,
+    ) as format_process:
+        if format_process.returncode:
+            log.warning("Ruff lint failed: %s", format_process.returncode)
+
+    # move file to output location
+    if location.name == "open_lineage.py":
+        location = location.with_name("base.py")
+    os.rename(tmpfile.name, location)
+
+    if lint_process.returncode or format_process.returncode:
+        log.warning("%s failed on ruff.", location)

openlineage_python-1.31.0/openlineage/client/generator/cli.py

@@ -0,0 +1,53 @@
+# Copyright 2018-2025 contributors to the OpenLineage project
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import pathlib
+import tempfile
+
+import click
+import requests
+from openlineage.client.generator.base import (
+    camel_to_snake,
+    format_and_save_output,
+    load_specs,
+    parse_and_generate,
+)
+
+BASE_SPEC_URL = "https://openlineage.io/spec/2-0-2/OpenLineage.json"
+
+
+def get_base_spec() -> pathlib.Path:
+    """Get the base spec from the OpenLineage repository."""
+    response = requests.get(BASE_SPEC_URL, timeout=10)
+    response.raise_for_status()  # Raise an HTTPError for bad responses
+    tmpfile = pathlib.Path(tempfile.mkdtemp()) / "OpenLineage.json"
+    tmpfile.write_text(response.text)
+    return tmpfile
+
+
+@click.command()
+@click.argument("facets_spec_location", type=pathlib.Path)
+@click.option("--output-location", type=pathlib.Path, default=None)
+def main(facets_spec_location: pathlib.Path, output_location: pathlib.Path | None) -> None:
+    base_spec_location = get_base_spec()
+    locations = load_specs(base_spec_location=base_spec_location, facets_spec_location=facets_spec_location)
+    results = parse_and_generate(locations)
+    if not output_location:
+        output_location = pathlib.Path(tempfile.mkdtemp())
+    modules = {
+        output_location.joinpath(*[camel_to_snake(n.replace("Facet", "")) for n in name]): result.body
+        for name, result in sorted(results.items())
+    }
+    for path, body in modules.items():
+        if path and not path.parent.exists():
+            path.parent.mkdir(parents=True)
+        if body and path.name not in ("open_lineage.py", "__init__.py"):
+            format_and_save_output(body, path, path.name == "open_lineage.py")
+        if path.name not in ("open_lineage.py", "__init__.py"):
+            with open(path) as file:
+                print(f"Generator would generate file {path.name}:\r\n\r\n{file.read()}")
+
+
+if __name__ == "__main__":
+    main()

openlineage_python-1.31.0/openlineage/client/generator/generate.py

@@ -0,0 +1,100 @@
+# Copyright 2018-2025 contributors to the OpenLineage project
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import pathlib
+import pkgutil
+import shutil
+from collections import defaultdict
+from typing import Any
+
+import click
+import jinja2
+import yaml
+from base import (  # type: ignore[import-not-found]
+    SCHEMA_URLS,
+    camel_to_snake,
+    deep_merge_dicts,
+    format_and_save_output,
+    load_specs,
+    parse_and_generate,
+)
+
+# locations definitions
+FILE_LOCATION = pathlib.Path(__file__).resolve().parent
+PYTHON_CLIENT_LOCATION = FILE_LOCATION.parent.parent.parent
+BASE_LOCATION = PYTHON_CLIENT_LOCATION.parent.parent
+BASE_SPEC_LOCATION = BASE_LOCATION / "spec" / "OpenLineage.json"
+FACETS_SPEC_LOCATION = BASE_LOCATION / "spec" / "facets"
+DEFAULT_OUTPUT_LOCATION = PYTHON_CLIENT_LOCATION / "openlineage" / "client" / "generated"
+TEMPLATES_LOCATION = FILE_LOCATION / "templates"
+
+HEADER = (FILE_LOCATION / "header.py").resolve().read_text()
+
+# structures to customize code generation
+REDACT_FIELDS = yaml.safe_load((PYTHON_CLIENT_LOCATION / "redact_fields.yml").read_text())
+
+
+def get_redact_fields(module_name: str) -> dict[str, dict[str, list[str]]]:
+    module_entry = next((entry for entry in REDACT_FIELDS if entry == module_name), None)
+    if not module_entry:
+        msg = f"{module_name} should be present in redact_fields"
+        raise ValueError(msg)
+    return {class_name: {"redact_fields": fields} for class_name, fields in module_entry}
+
+
+def generate_facet_v2_module(module_location: pathlib.Path) -> None:
+    modules = [name for _, name, _ in pkgutil.iter_modules([str(module_location)]) if name != "base"]
+
+    facet_v2_template = (TEMPLATES_LOCATION / "facet_v2.jinja2").read_text()
+    output = jinja2.Environment(autoescape=True).from_string(facet_v2_template).render(facets_modules=modules)
+    format_and_save_output(
+        output=output,
+        location=PYTHON_CLIENT_LOCATION / "openlineage" / "client" / "facet_v2.py",
+        header=HEADER,
+    )
+
+
+@click.command()
+@click.option("--output-location", type=pathlib.Path, default=DEFAULT_OUTPUT_LOCATION)
+def main(output_location: pathlib.Path) -> None:
+    locations = load_specs(base_spec_location=BASE_SPEC_LOCATION, facets_spec_location=FACETS_SPEC_LOCATION)
+
+    extra_schema_urls = {obj_name: {"_schemaURL": schema_url} for obj_name, schema_url in SCHEMA_URLS.items()}
+
+    extra_redact_fields: dict[str, dict[str, list[dict[str, Any]]]] = defaultdict(lambda: {"redactions": []})
+
+    for module_entry in REDACT_FIELDS:
+        for clazz in module_entry["classes"]:
+            class_name = clazz["class_name"]
+            # defaultdict automatically creates the key-value pair if not found
+            extra_redact_fields[class_name]["redactions"].append(
+                {
+                    "fields": clazz["redact_fields"],
+                    "module_name": module_entry["module"],
+                }
+            )
+
+    extra_template_data = defaultdict(dict, deep_merge_dicts(extra_redact_fields, extra_schema_urls))
+
+    results = parse_and_generate(locations, extra_template_data)
+
+    if output_location.exists():
+        shutil.rmtree(output_location)
+    for name, result in sorted(results.items()):
+        path = output_location.joinpath(*[camel_to_snake(n.replace("Facet", "")) for n in name])
+        if path and not path.parent.exists():
+            path.parent.mkdir(parents=True)
+        if result.body:
+            format_and_save_output(
+                output=result.body,
+                location=path,
+                add_set_producer_code=path.name == "open_lineage.py",
+                header=HEADER,
+            )
+
+    generate_facet_v2_module(DEFAULT_OUTPUT_LOCATION)
+
+
+if __name__ == "__main__":
+    main()

openlineage_python-1.31.0/openlineage/client/generator/templates/Enum.jinja2

@@ -0,0 +1,22 @@
+{#- this template adds indentation and wrapping for docstrings -#}
+{% for decorator in decorators -%}
+{{ decorator }}
+{% endfor -%}
+class {{ class_name }}({{ base_class }}):
+{%- if description -%}
+    {% if description | length > 100 %}
+    """
+    {{ description | wordwrap(100) | indent(4)}}
+    """
+    {%- else %}
+    """{{ description }}"""
+    {% endif -%}
+{%- endif %}
+{%- for field in fields %}
+    {{ field.name }} = {{ field.default }}
+    {%- if field.docstring %}
+    """
+    {{ field.docstring | indent(4) }}
+    """
+    {%- endif %}
+{%- endfor -%}