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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (105) hide show
  1. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/PKG-INFO +5 -1
  2. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/client.py +15 -7
  3. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/constants.py +1 -1
  4. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/parent_run.py +42 -1
  5. openlineage_python-1.31.0/openlineage/client/generator/base.py +257 -0
  6. openlineage_python-1.31.0/openlineage/client/generator/cli.py +53 -0
  7. openlineage_python-1.31.0/openlineage/client/generator/generate.py +100 -0
  8. openlineage_python-1.31.0/openlineage/client/generator/templates/Enum.jinja2 +22 -0
  9. openlineage_python-1.31.0/openlineage/client/generator/templates/dataclass.jinja2 +196 -0
  10. openlineage_python-1.31.0/openlineage/client/generator/templates/facet_v2.jinja2 +25 -0
  11. openlineage_python-1.31.0/openlineage/client/generator/templates/root.jinja2 +1 -0
  12. openlineage_python-1.31.0/openlineage/client/generator/templates/validators.jinja2 +43 -0
  13. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/transport/kafka.py +40 -13
  14. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/pyproject.toml +10 -1
  15. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/redact_fields.yml +11 -0
  16. openlineage_python-1.31.0/tests/__init__.py +2 -0
  17. openlineage_python-1.31.0/tests/generator/__init__.py +2 -0
  18. openlineage_python-1.31.0/tests/generator/test_base.py +148 -0
  19. openlineage_python-1.31.0/tests/generator/test_cli.py +53 -0
  20. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/test_client.py +127 -79
  21. openlineage_python-1.31.0/tests/test_kafka.py +477 -0
  22. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tox.ini +1 -0
  23. openlineage_python-1.30.0/tests/test_kafka.py +0 -338
  24. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/.gitignore +0 -0
  25. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/README.md +0 -0
  26. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/docs/Makefile +0 -0
  27. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/docs/conf.py +0 -0
  28. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/docs/convert.py +0 -0
  29. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/docs/index.rst +0 -0
  30. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/docs/make.bat +0 -0
  31. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/docs/source/openlineage.client.rst +0 -0
  32. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/__init__.py +0 -0
  33. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/event_v2.py +0 -0
  34. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/facet.py +0 -0
  35. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/facet_v2.py +0 -0
  36. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/facets.py +0 -0
  37. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/filter.py +0 -0
  38. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/base.py +0 -0
  39. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/column_lineage_dataset.py +0 -0
  40. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/data_quality_assertions_dataset.py +0 -0
  41. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/data_quality_metrics_input_dataset.py +0 -0
  42. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/dataset_type_dataset.py +0 -0
  43. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/dataset_version_dataset.py +0 -0
  44. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/datasource_dataset.py +0 -0
  45. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/documentation_dataset.py +0 -0
  46. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/documentation_job.py +0 -0
  47. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/environment_variables_run.py +0 -0
  48. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/error_message_run.py +0 -0
  49. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/external_query_run.py +0 -0
  50. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/extraction_error_run.py +0 -0
  51. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/input_statistics_input_dataset.py +0 -0
  52. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/job_type_job.py +0 -0
  53. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/lifecycle_state_change_dataset.py +0 -0
  54. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/nominal_time_run.py +0 -0
  55. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/output_statistics_output_dataset.py +0 -0
  56. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/ownership_dataset.py +0 -0
  57. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/ownership_job.py +0 -0
  58. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/processing_engine_run.py +0 -0
  59. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/schema_dataset.py +0 -0
  60. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/source_code_job.py +0 -0
  61. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/source_code_location_job.py +0 -0
  62. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/sql_job.py +0 -0
  63. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/storage_dataset.py +0 -0
  64. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/symlinks_dataset.py +0 -0
  65. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/tags_dataset.py +0 -0
  66. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/tags_job.py +0 -0
  67. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/generated/tags_run.py +0 -0
  68. openlineage_python-1.30.0/tests/__init__.py → openlineage_python-1.31.0/openlineage/client/generator/header.py +0 -0
  69. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/py.typed +0 -0
  70. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/run.py +0 -0
  71. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/serde.py +0 -0
  72. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/tags.py +0 -0
  73. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/transport/__init__.py +0 -0
  74. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/transport/composite.py +0 -0
  75. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/transport/console.py +0 -0
  76. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/transport/factory.py +0 -0
  77. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/transport/file.py +1 -1
  78. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/transport/http.py +0 -0
  79. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/transport/msk_iam.py +0 -0
  80. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/transport/noop.py +0 -0
  81. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/transport/transport.py +0 -0
  82. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/utils.py +0 -0
  83. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/openlineage/client/uuid.py +0 -0
  84. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/config/config.yml +0 -0
  85. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/config/exact_filter.yml +0 -0
  86. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/config/http.yml +0 -0
  87. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/config/openlineage.yml +0 -0
  88. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/config/regex_filter.yml +0 -0
  89. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/conftest.py +0 -0
  90. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/example_full_event.json +0 -0
  91. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/nominal_time_without_end.json +0 -0
  92. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/serde_example_dataset_event.json +0 -0
  93. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/serde_example_job_event.json +0 -0
  94. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/serde_example_run_event.json +0 -0
  95. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/test_composite.py +0 -0
  96. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/test_events.py +0 -0
  97. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/test_facet.py +0 -0
  98. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/test_facet_v2.py +0 -0
  99. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/test_factory.py +0 -0
  100. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/test_file.py +0 -0
  101. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/test_http.py +0 -0
  102. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/test_msk_iam.py +0 -0
  103. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/test_utils.py +0 -0
  104. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/test_uuid.py +0 -0
  105. {openlineage_python-1.30.0 → openlineage_python-1.31.0}/tests/transport.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: openlineage-python
3
- Version: 1.30.0
3
+ Version: 1.31.0
4
4
  Summary: OpenLineage Python Client
5
5
  Author-email: OpenLineage <info@openlineage.io>
6
6
  Keywords: openlineage
@@ -23,6 +23,10 @@ Requires-Dist: beautifulsoup4>=4.12.0; extra == 'docs'
23
23
  Requires-Dist: pydata-sphinx-theme>=0.13; extra == 'docs'
24
24
  Requires-Dist: sphinx-autodoc-typehints>=2.0.0; extra == 'docs'
25
25
  Requires-Dist: sphinx>=7.0; extra == 'docs'
26
+ Provides-Extra: generator
27
+ Requires-Dist: click; extra == 'generator'
28
+ Requires-Dist: datamodel-code-generator==0.25.4; extra == 'generator'
29
+ Requires-Dist: ruff==0.3.5; extra == 'generator'
26
30
  Provides-Extra: kafka
27
31
  Requires-Dist: confluent-kafka>=2.1.1; extra == 'kafka'
28
32
  Provides-Extra: msk-iam
@@ -163,6 +163,11 @@ class OpenLineageClient:
163
163
  if type(event) not in get_args(Event):
164
164
  msg = "`emit` only accepts RunEvent, DatasetEvent, JobEvent classes"
165
165
  raise ValueError(msg)
166
+
167
+ if log.isEnabledFor(logging.DEBUG):
168
+ val = Serde.to_json(event).encode("utf-8")
169
+ log.debug("OpenLineageClient will *try* to emit event %s", val)
170
+
166
171
  if not self.transport:
167
172
  log.error("Tried to emit OpenLineage event, but transport is not configured.")
168
173
  return
@@ -176,9 +181,6 @@ class OpenLineageClient:
176
181
  event = self.add_environment_facets(event)
177
182
  event = self.update_event_tags_facets(event)
178
183
 
179
- if log.isEnabledFor(logging.DEBUG):
180
- val = Serde.to_json(event).encode("utf-8")
181
- log.debug("OpenLineageClient will emit event %s", val)
182
184
  self.transport.emit(event)
183
185
  log.debug("OpenLineage event successfully emitted.")
184
186
 
@@ -420,15 +422,21 @@ class OpenLineageClient:
420
422
  """
421
423
  Creates or updates job and run tag facets based on user-supplied environment variables
422
424
  """
425
+ run_event_types = (RunEvent, event_v2.RunEvent)
426
+ run_and_job_event_types = (RunEvent, event_v2.RunEvent, JobEvent, event_v2.JobEvent)
423
427
  tags_job = self.config.tags.job
424
- if (isinstance(event, (JobEvent, RunEvent))) and tags_job:
428
+ if isinstance(event, run_and_job_event_types) and tags_job:
429
+ # Ensure facets exists
430
+ event.job.facets = {} if not event.job.facets else event.job.facets
425
431
  tags_facet = event.job.facets.get("tags", TagsJobFacet())
426
- event.job.facets["tags"] = self._update_tag_facet(tags_facet, tags_job) # type: ignore [arg-type]
432
+ event.job.facets["tags"] = self._update_tag_facet(tags_facet, tags_job) # type: ignore [arg-type, assignment]
427
433
 
428
434
  tags_run = self.config.tags.run
429
- if (isinstance(event, RunEvent)) and tags_run:
435
+ if isinstance(event, run_event_types) and tags_run:
436
+ # Ensure facets exists
437
+ event.run.facets = {} if not event.run.facets else event.run.facets
430
438
  tags_facet = event.run.facets.get("tags", TagsRunFacet())
431
- event.run.facets["tags"] = self._update_tag_facet(tags_facet, tags_run) # type: ignore [arg-type]
439
+ event.run.facets["tags"] = self._update_tag_facet(tags_facet, tags_run) # type: ignore [arg-type, assignment]
432
440
 
433
441
  return event
434
442
 
@@ -2,7 +2,7 @@
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  from __future__ import annotations
4
4
 
5
- __version__ = "1.30.0"
5
+ __version__ = "1.31.0"
6
6
 
7
7
  DEFAULT_TIMEOUT_MS = 5000
8
8
  DEFAULT_NAMESPACE_NAME = "default"
@@ -28,11 +28,12 @@ class ParentRunFacet(RunFacet):
28
28
 
29
29
  run: Run
30
30
  job: Job
31
+ root: Root | None = attr.field(default=None)
31
32
  _additional_skip_redact: ClassVar[list[str]] = ["job", "run"]
32
33
 
33
34
  @staticmethod
34
35
  def _get_schema() -> str:
35
- return "https://openlineage.io/spec/facets/1-0-1/ParentRunFacet.json#/$defs/ParentRunFacet"
36
+ return "https://openlineage.io/spec/facets/1-1-0/ParentRunFacet.json#/$defs/ParentRunFacet"
36
37
 
37
38
  @classmethod
38
39
  def create(cls, runId: str, namespace: str, name: str) -> ParentRunFacet: # noqa: N803
@@ -46,6 +47,46 @@ class ParentRunFacet(RunFacet):
46
47
  return cls(run=Run(runId=runId), job=Job(namespace=namespace, name=name))
47
48
 
48
49
 
50
+ @attr.define
51
+ class Root(RedactMixin):
52
+ run: RootRun
53
+ job: RootJob
54
+ _skip_redact: ClassVar[list[str]] = ["run", "job"]
55
+
56
+
57
+ @attr.define
58
+ class RootJob(RedactMixin):
59
+ namespace: str
60
+ """The namespace containing root job"""
61
+
62
+ name: str
63
+ """The unique name containing root job within that namespace"""
64
+
65
+ _skip_redact: ClassVar[list[str]] = ["namespace", "name"]
66
+
67
+ @staticmethod
68
+ def _get_schema() -> str:
69
+ return "https://openlineage.io/spec/facets/1-1-0/ParentRunFacet.json#/$defs/RootJob"
70
+
71
+
72
+ @attr.define
73
+ class RootRun(RedactMixin):
74
+ runId: str = attr.field() # noqa: N815
75
+ """The globally unique ID of the root run associated with the root job."""
76
+
77
+ _skip_redact: ClassVar[list[str]] = ["runId"]
78
+
79
+ @staticmethod
80
+ def _get_schema() -> str:
81
+ return "https://openlineage.io/spec/facets/1-1-0/ParentRunFacet.json#/$defs/RootRun"
82
+
83
+ @runId.validator
84
+ def runid_check(self, attribute: str, value: str) -> None: # noqa: ARG002
85
+ from uuid import UUID
86
+
87
+ UUID(value)
88
+
89
+
49
90
  @attr.define
50
91
  class Run(RedactMixin):
51
92
  runId: str = attr.field() # noqa: N815
@@ -0,0 +1,257 @@
1
+ # Copyright 2018-2025 contributors to the OpenLineage project
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ from __future__ import annotations
4
+
5
+ import copy
6
+ import json
7
+ import logging
8
+ import os
9
+ import pathlib
10
+ import re
11
+ import subprocess
12
+ import tempfile
13
+ from typing import Any
14
+
15
+ from datamodel_code_generator import DataModelType, PythonVersion
16
+ from datamodel_code_generator.imports import Import
17
+ from datamodel_code_generator.model import get_data_model_types
18
+ from datamodel_code_generator.model import pydantic as pydantic_model
19
+ from datamodel_code_generator.parser.jsonschema import JsonSchemaParser
20
+ from datamodel_code_generator.types import Types
21
+
22
+ log = logging.getLogger(__name__)
23
+
24
+
25
+ def camel_to_snake(string: str) -> str:
26
+ """Convert camel_case string to SnakeCase"""
27
+ _under_score_1 = re.compile(r"([^_])([A-Z][a-z]+)")
28
+ _under_score_2 = re.compile("([a-z0-9])([A-Z])")
29
+ subbed = _under_score_1.sub(r"\1_\2", string)
30
+ return _under_score_2.sub(r"\1_\2", subbed).lower()
31
+
32
+
33
+ # import attrs instead of dataclass
34
+ DATA_MODELS_TYPES = get_data_model_types(
35
+ DataModelType.DataclassesDataclass, target_python_version=PythonVersion.PY_38
36
+ )
37
+ NEW_MODEL = DATA_MODELS_TYPES.data_model
38
+ NEW_MODEL.DEFAULT_IMPORTS = (Import.from_full_path("attr"),)
39
+
40
+ # custom code
41
+ SET_PRODUCER_CODE = """
42
+ PRODUCER = DEFAULT_PRODUCER
43
+
44
+ def set_producer(producer: str) -> None:
45
+ global PRODUCER # noqa: PLW0603
46
+ PRODUCER = producer
47
+ """
48
+
49
+ # locations definitions
50
+ FILE_LOCATION = pathlib.Path(__file__).resolve().parent
51
+ TEMPLATES_LOCATION = FILE_LOCATION / "templates"
52
+
53
+ # Global dictionaries intended to be imported by other modules
54
+ # Contains schema URLs and base IDs parsed from specification files
55
+ SCHEMA_URLS: dict[str, str] = {}
56
+ BASE_IDS: dict[str, str] = {}
57
+
58
+
59
+ def deep_merge_dicts(dict1: dict[str, Any], dict2: dict[str, Any]) -> dict[str, Any]:
60
+ """Deep merges two dictionaries.
61
+
62
+ This function merges two dictionaries while handling nested dictionaries.
63
+ For keys that exist in both dictionaries, the values from dict2 take precedence.
64
+ If a key exists in both dictionaries and the values are dictionaries themselves,
65
+ they are merged recursively.
66
+ This function merges only dictionaries. If key is of different type, e.g. list
67
+ it does not work properly.
68
+ """
69
+ merged = dict1.copy()
70
+ for k, v in dict2.items():
71
+ if k in merged and isinstance(v, dict):
72
+ merged[k] = deep_merge_dicts(merged.get(k, {}), v)
73
+ else:
74
+ merged[k] = v
75
+ return merged
76
+
77
+
78
+ def parse_additional_data(spec: dict[str, Any], file_name: str) -> None:
79
+ """Parse additional data from spec files.
80
+
81
+ Parses:
82
+ * schema URLs
83
+ * base IDs
84
+
85
+ Updates the module-level SCHEMA_URLS and BASE_IDS dictionaries
86
+ that can be imported by other modules.
87
+ """
88
+ base_id = spec["$id"]
89
+ for name, _ in spec["$defs"].items():
90
+ SCHEMA_URLS[name] = f"{base_id}#/$defs/{name}"
91
+ BASE_IDS[file_name] = spec["$id"]
92
+
93
+
94
+ def load_specs(base_spec_location: pathlib.Path, facets_spec_location: pathlib.Path) -> list[pathlib.Path]:
95
+ """Load base `OpenLineage.json` and other facets' spec files"""
96
+ locations = []
97
+ if facets_spec_location.is_dir():
98
+ facets_spec_files = sorted(pathlib.Path(os.path.abspath(facets_spec_location)).glob("*.json"))
99
+ else:
100
+ facets_spec_files = [pathlib.Path(os.path.abspath(facets_spec_location))]
101
+
102
+ file_specs = [
103
+ base_spec_location.resolve(), # Base spec must be the first element
104
+ *facets_spec_files,
105
+ ]
106
+ for file_spec in file_specs:
107
+ spec = json.loads(file_spec.read_text())
108
+ parse_additional_data(spec, file_spec.name)
109
+ locations.append(file_spec)
110
+
111
+ return locations
112
+
113
+
114
+ def parse_and_generate(
115
+ locations: list[pathlib.Path], extra_template_data: dict[str, Any] | None = None
116
+ ) -> dict[str, Any]:
117
+ """Parse and generate data models from a given specification."""
118
+ temporary_locations = []
119
+
120
+ current_dir = pathlib.Path.cwd()
121
+ with tempfile.TemporaryDirectory() as tmp:
122
+ tmp_directory = pathlib.Path(tmp).resolve()
123
+ for location in locations:
124
+ tmp_location = (tmp_directory / location.name).resolve()
125
+ tmp_location.write_text(location.read_text())
126
+ temporary_locations.append(tmp_location)
127
+
128
+ os.chdir(tmp_directory)
129
+ # first parse OpenLineage.json
130
+ parser = JsonSchemaParser(
131
+ source=temporary_locations[:1],
132
+ data_model_type=NEW_MODEL,
133
+ data_model_root_type=DATA_MODELS_TYPES.root_model,
134
+ data_model_field_type=DATA_MODELS_TYPES.field_model,
135
+ data_type_manager_type=DATA_MODELS_TYPES.data_type_manager,
136
+ dump_resolve_reference_action=DATA_MODELS_TYPES.dump_resolve_reference_action,
137
+ special_field_name_prefix="",
138
+ use_schema_description=True,
139
+ field_constraints=True,
140
+ use_union_operator=True,
141
+ use_standard_collections=True,
142
+ base_class="openlineage.client.utils.RedactMixin",
143
+ class_name="ClassToBeSkipped",
144
+ use_field_description=True,
145
+ use_double_quotes=True,
146
+ keep_model_order=True,
147
+ custom_template_dir=TEMPLATES_LOCATION,
148
+ extra_template_data=extra_template_data, # type: ignore[arg-type]
149
+ additional_imports=[
150
+ "typing.ClassVar",
151
+ "typing.Any",
152
+ "typing.cast",
153
+ "openlineage.client.constants.DEFAULT_PRODUCER",
154
+ ],
155
+ )
156
+
157
+ # keep information about uuid and date-time formats
158
+ # this is going to be changed back to str type hint in jinja template
159
+ data_type_manager: pydantic_model.DataTypeManager = parser.data_type_manager # type: ignore[assignment]
160
+ uuid_type = copy.deepcopy(data_type_manager.type_map[Types.uuid])
161
+ uuid_type.type = "uuid"
162
+ data_type_manager.type_map[Types.uuid] = uuid_type
163
+
164
+ date_time_type = copy.deepcopy(data_type_manager.type_map[Types.date_time])
165
+ date_time_type.type = "date-time"
166
+ data_type_manager.type_map[Types.date_time] = date_time_type
167
+
168
+ uri_type = copy.deepcopy(data_type_manager.type_map[Types.date_time])
169
+ uri_type.type = "uri"
170
+ data_type_manager.type_map[Types.uri] = uri_type
171
+
172
+ parser.parse(format_=False)
173
+
174
+ # parse rest of spec
175
+ parser.source = temporary_locations[1:]
176
+
177
+ # change paths so that parser sees base objects as local
178
+ parser.model_resolver.references = {
179
+ k.replace("OpenLineage.json", BASE_IDS["OpenLineage.json"]): v
180
+ for k, v in parser.model_resolver.references.items()
181
+ }
182
+
183
+ output = parser.parse(format_=False)
184
+
185
+ # go back to original directory
186
+ os.chdir(current_dir)
187
+
188
+ return output # type: ignore[return-value]
189
+
190
+
191
+ def separate_imports(code: str) -> tuple[str, str]:
192
+ """Separates a Python script code (as string) into imports and the rest."""
193
+ imports_section = []
194
+ rest_of_code = []
195
+ in_import = False
196
+ for line in code.splitlines():
197
+ if line.startswith(("import ", "from ")):
198
+ in_import = True
199
+ elif in_import and line.strip() == "":
200
+ in_import = False
201
+ if in_import:
202
+ imports_section.append(line)
203
+ else:
204
+ rest_of_code.append(line)
205
+ return ("\n".join(imports_section), "\n".join(rest_of_code))
206
+
207
+
208
+ def format_and_save_output(
209
+ output: str,
210
+ location: pathlib.Path,
211
+ add_set_producer_code: bool | None = False,
212
+ header: str | None = None,
213
+ ) -> None:
214
+ """Adjust and format generated file."""
215
+
216
+ # save temporary file to the same directory that output
217
+ # so that ruff receives same rules
218
+ with tempfile.NamedTemporaryFile(
219
+ "w", prefix=location.stem.lower(), suffix=".py", dir=location.parent, delete=False
220
+ ) as tmpfile:
221
+ if header:
222
+ tmpfile.write(header + "\n")
223
+ output = output.replace("from .OpenLineage", "from openlineage.client.generated.base")
224
+ imports_section, rest_of_code = separate_imports(output)
225
+ tmpfile.write(imports_section)
226
+ if add_set_producer_code:
227
+ tmpfile.write(SET_PRODUCER_CODE)
228
+ tmpfile.write(rest_of_code)
229
+ tmpfile.flush()
230
+
231
+ # run ruff lint
232
+ with subprocess.Popen(
233
+ args=["ruff", "check", tmpfile.name, "--fix"],
234
+ stdout=subprocess.DEVNULL,
235
+ stderr=subprocess.DEVNULL,
236
+ close_fds=True,
237
+ ) as lint_process:
238
+ if lint_process.returncode:
239
+ log.warning("Ruff lint failed: %s", lint_process.returncode)
240
+
241
+ # run ruff format
242
+ with subprocess.Popen(
243
+ args=["ruff", "format", tmpfile.name],
244
+ stdout=subprocess.DEVNULL,
245
+ stderr=subprocess.DEVNULL,
246
+ close_fds=True,
247
+ ) as format_process:
248
+ if format_process.returncode:
249
+ log.warning("Ruff lint failed: %s", format_process.returncode)
250
+
251
+ # move file to output location
252
+ if location.name == "open_lineage.py":
253
+ location = location.with_name("base.py")
254
+ os.rename(tmpfile.name, location)
255
+
256
+ if lint_process.returncode or format_process.returncode:
257
+ log.warning("%s failed on ruff.", location)
@@ -0,0 +1,53 @@
1
+ # Copyright 2018-2025 contributors to the OpenLineage project
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ from __future__ import annotations
4
+
5
+ import pathlib
6
+ import tempfile
7
+
8
+ import click
9
+ import requests
10
+ from openlineage.client.generator.base import (
11
+ camel_to_snake,
12
+ format_and_save_output,
13
+ load_specs,
14
+ parse_and_generate,
15
+ )
16
+
17
+ BASE_SPEC_URL = "https://openlineage.io/spec/2-0-2/OpenLineage.json"
18
+
19
+
20
+ def get_base_spec() -> pathlib.Path:
21
+ """Get the base spec from the OpenLineage repository."""
22
+ response = requests.get(BASE_SPEC_URL, timeout=10)
23
+ response.raise_for_status() # Raise an HTTPError for bad responses
24
+ tmpfile = pathlib.Path(tempfile.mkdtemp()) / "OpenLineage.json"
25
+ tmpfile.write_text(response.text)
26
+ return tmpfile
27
+
28
+
29
+ @click.command()
30
+ @click.argument("facets_spec_location", type=pathlib.Path)
31
+ @click.option("--output-location", type=pathlib.Path, default=None)
32
+ def main(facets_spec_location: pathlib.Path, output_location: pathlib.Path | None) -> None:
33
+ base_spec_location = get_base_spec()
34
+ locations = load_specs(base_spec_location=base_spec_location, facets_spec_location=facets_spec_location)
35
+ results = parse_and_generate(locations)
36
+ if not output_location:
37
+ output_location = pathlib.Path(tempfile.mkdtemp())
38
+ modules = {
39
+ output_location.joinpath(*[camel_to_snake(n.replace("Facet", "")) for n in name]): result.body
40
+ for name, result in sorted(results.items())
41
+ }
42
+ for path, body in modules.items():
43
+ if path and not path.parent.exists():
44
+ path.parent.mkdir(parents=True)
45
+ if body and path.name not in ("open_lineage.py", "__init__.py"):
46
+ format_and_save_output(body, path, path.name == "open_lineage.py")
47
+ if path.name not in ("open_lineage.py", "__init__.py"):
48
+ with open(path) as file:
49
+ print(f"Generator would generate file {path.name}:\r\n\r\n{file.read()}")
50
+
51
+
52
+ if __name__ == "__main__":
53
+ main()
@@ -0,0 +1,100 @@
1
+ # Copyright 2018-2025 contributors to the OpenLineage project
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ from __future__ import annotations
4
+
5
+ import pathlib
6
+ import pkgutil
7
+ import shutil
8
+ from collections import defaultdict
9
+ from typing import Any
10
+
11
+ import click
12
+ import jinja2
13
+ import yaml
14
+ from base import ( # type: ignore[import-not-found]
15
+ SCHEMA_URLS,
16
+ camel_to_snake,
17
+ deep_merge_dicts,
18
+ format_and_save_output,
19
+ load_specs,
20
+ parse_and_generate,
21
+ )
22
+
23
+ # locations definitions
24
+ FILE_LOCATION = pathlib.Path(__file__).resolve().parent
25
+ PYTHON_CLIENT_LOCATION = FILE_LOCATION.parent.parent.parent
26
+ BASE_LOCATION = PYTHON_CLIENT_LOCATION.parent.parent
27
+ BASE_SPEC_LOCATION = BASE_LOCATION / "spec" / "OpenLineage.json"
28
+ FACETS_SPEC_LOCATION = BASE_LOCATION / "spec" / "facets"
29
+ DEFAULT_OUTPUT_LOCATION = PYTHON_CLIENT_LOCATION / "openlineage" / "client" / "generated"
30
+ TEMPLATES_LOCATION = FILE_LOCATION / "templates"
31
+
32
+ HEADER = (FILE_LOCATION / "header.py").resolve().read_text()
33
+
34
+ # structures to customize code generation
35
+ REDACT_FIELDS = yaml.safe_load((PYTHON_CLIENT_LOCATION / "redact_fields.yml").read_text())
36
+
37
+
38
+ def get_redact_fields(module_name: str) -> dict[str, dict[str, list[str]]]:
39
+ module_entry = next((entry for entry in REDACT_FIELDS if entry == module_name), None)
40
+ if not module_entry:
41
+ msg = f"{module_name} should be present in redact_fields"
42
+ raise ValueError(msg)
43
+ return {class_name: {"redact_fields": fields} for class_name, fields in module_entry}
44
+
45
+
46
+ def generate_facet_v2_module(module_location: pathlib.Path) -> None:
47
+ modules = [name for _, name, _ in pkgutil.iter_modules([str(module_location)]) if name != "base"]
48
+
49
+ facet_v2_template = (TEMPLATES_LOCATION / "facet_v2.jinja2").read_text()
50
+ output = jinja2.Environment(autoescape=True).from_string(facet_v2_template).render(facets_modules=modules)
51
+ format_and_save_output(
52
+ output=output,
53
+ location=PYTHON_CLIENT_LOCATION / "openlineage" / "client" / "facet_v2.py",
54
+ header=HEADER,
55
+ )
56
+
57
+
58
+ @click.command()
59
+ @click.option("--output-location", type=pathlib.Path, default=DEFAULT_OUTPUT_LOCATION)
60
+ def main(output_location: pathlib.Path) -> None:
61
+ locations = load_specs(base_spec_location=BASE_SPEC_LOCATION, facets_spec_location=FACETS_SPEC_LOCATION)
62
+
63
+ extra_schema_urls = {obj_name: {"_schemaURL": schema_url} for obj_name, schema_url in SCHEMA_URLS.items()}
64
+
65
+ extra_redact_fields: dict[str, dict[str, list[dict[str, Any]]]] = defaultdict(lambda: {"redactions": []})
66
+
67
+ for module_entry in REDACT_FIELDS:
68
+ for clazz in module_entry["classes"]:
69
+ class_name = clazz["class_name"]
70
+ # defaultdict automatically creates the key-value pair if not found
71
+ extra_redact_fields[class_name]["redactions"].append(
72
+ {
73
+ "fields": clazz["redact_fields"],
74
+ "module_name": module_entry["module"],
75
+ }
76
+ )
77
+
78
+ extra_template_data = defaultdict(dict, deep_merge_dicts(extra_redact_fields, extra_schema_urls))
79
+
80
+ results = parse_and_generate(locations, extra_template_data)
81
+
82
+ if output_location.exists():
83
+ shutil.rmtree(output_location)
84
+ for name, result in sorted(results.items()):
85
+ path = output_location.joinpath(*[camel_to_snake(n.replace("Facet", "")) for n in name])
86
+ if path and not path.parent.exists():
87
+ path.parent.mkdir(parents=True)
88
+ if result.body:
89
+ format_and_save_output(
90
+ output=result.body,
91
+ location=path,
92
+ add_set_producer_code=path.name == "open_lineage.py",
93
+ header=HEADER,
94
+ )
95
+
96
+ generate_facet_v2_module(DEFAULT_OUTPUT_LOCATION)
97
+
98
+
99
+ if __name__ == "__main__":
100
+ main()
@@ -0,0 +1,22 @@
1
+ {#- this template adds indentation and wrapping for docstrings -#}
2
+ {% for decorator in decorators -%}
3
+ {{ decorator }}
4
+ {% endfor -%}
5
+ class {{ class_name }}({{ base_class }}):
6
+ {%- if description -%}
7
+ {% if description | length > 100 %}
8
+ """
9
+ {{ description | wordwrap(100) | indent(4)}}
10
+ """
11
+ {%- else %}
12
+ """{{ description }}"""
13
+ {% endif -%}
14
+ {%- endif %}
15
+ {%- for field in fields %}
16
+ {{ field.name }} = {{ field.default }}
17
+ {%- if field.docstring %}
18
+ """
19
+ {{ field.docstring | indent(4) }}
20
+ """
21
+ {%- endif %}
22
+ {%- endfor -%}