openlineage-prefect 0.0.1b3__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.
@@ -0,0 +1,81 @@
1
+ Metadata-Version: 2.3
2
+ Name: openlineage-prefect
3
+ Version: 0.0.1b3
4
+ Summary: OpenLineage integration with Prefect
5
+ Keywords: openlineage,prefect
6
+ Author: OpenLineage
7
+ License: Apache-2.0
8
+ Classifier: Development Status :: 4 - Beta
9
+ Classifier: Intended Audience :: Developers
10
+ Classifier: License :: OSI Approved :: Apache Software License
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: Programming Language :: Python :: 3.12
13
+ Requires-Dist: prefect>=3.8.5
14
+ Requires-Dist: openlineage-python>=1.53.0
15
+ Requires-Dist: mypy ; extra == 'dev'
16
+ Requires-Dist: openlineage-python ; extra == 'dev'
17
+ Requires-Dist: pillow ; extra == 'dev'
18
+ Requires-Dist: pre-commit ; extra == 'dev'
19
+ Requires-Dist: pytest>=8.3 ; extra == 'dev'
20
+ Requires-Dist: pytest-asyncio ; extra == 'dev'
21
+ Requires-Dist: pytest-cov ; extra == 'dev'
22
+ Requires-Dist: pytest-env ; extra == 'dev'
23
+ Requires-Dist: pytest-timeout>=2.3.1 ; extra == 'dev'
24
+ Requires-Dist: pytest-xdist ; extra == 'dev'
25
+ Requires-Python: >=3.12
26
+ Project-URL: Homepage, https://openlineage.io/
27
+ Project-URL: Source, https://github.com/OpenLineage/OpenLineage/tree/main/integration/prefect
28
+ Provides-Extra: dev
29
+ Description-Content-Type: text/markdown
30
+
31
+ # Openlineage Prefect
32
+
33
+ ## Basic Configuration
34
+
35
+ At a minimum, define a namespace, Prefect API URL and transport consisting of a type, URL and endpoint. All three can be supplied using environment variables.
36
+
37
+ For example:
38
+
39
+ ```sh
40
+ export OPENLINEAGE_NAMESPACE='prefect_test' &&
41
+ export OPENLINEAGE__TRANSPORT__TYPE='http' &&
42
+ export OPENLINEAGE__TRANSPORT__URL='http://lineageconsumer.com:5000' &&
43
+ export OPENLINEAGE__TRANSPORT__ENDPOINT='/api/v1/lineage' &&
44
+ export PREFECT_API_URL='http://prefecthost.com:4200/api'
45
+ ```
46
+
47
+ For more details about OpenLineage transport options and how to configure them, consult the [OpenLineage Python Client Documentation](https://openlineage.io/docs/client/python/).
48
+
49
+ ## Execution
50
+
51
+ Import the `openlineage_prefect` package and execute `collect_and_process_runs()` asynchronously in its own process.
52
+
53
+ For example:
54
+
55
+ ```py
56
+ import asyncio
57
+ import openlineage_prefect
58
+ from openlineage_prefect.prefect_adapter.listener import PrefectOpenLineageListener
59
+
60
+ async def main():
61
+ await PrefectOpenLineageListener().collect_and_process_runs()
62
+
63
+ if __name__ == "__main__":
64
+ asyncio.run(main())
65
+ ```
66
+
67
+ ## Datasets
68
+
69
+ The integration looks for datasets in Prefect Artifacts. To attach input and output datasets to job runs, use `create_table_artifact()` from the Artifact library. Provide a namespace, typically the dataset URI, and name to the adapter via an artifact's `table` and `description`, respectively. Distinguish the type of dataset by appending `_output` or `_input` to the description.
70
+
71
+ For example:
72
+
73
+ ```py
74
+ ol_table = [{"database_uri":"duckdb:///customers_db", "table":"customers"}]
75
+
76
+ create_table_artifact(
77
+ key="upstream-insert",
78
+ table=ol_table,
79
+ description="ol-dataset_output"
80
+ )
81
+ ```
@@ -0,0 +1,51 @@
1
+ # Openlineage Prefect
2
+
3
+ ## Basic Configuration
4
+
5
+ At a minimum, define a namespace, Prefect API URL and transport consisting of a type, URL and endpoint. All three can be supplied using environment variables.
6
+
7
+ For example:
8
+
9
+ ```sh
10
+ export OPENLINEAGE_NAMESPACE='prefect_test' &&
11
+ export OPENLINEAGE__TRANSPORT__TYPE='http' &&
12
+ export OPENLINEAGE__TRANSPORT__URL='http://lineageconsumer.com:5000' &&
13
+ export OPENLINEAGE__TRANSPORT__ENDPOINT='/api/v1/lineage' &&
14
+ export PREFECT_API_URL='http://prefecthost.com:4200/api'
15
+ ```
16
+
17
+ For more details about OpenLineage transport options and how to configure them, consult the [OpenLineage Python Client Documentation](https://openlineage.io/docs/client/python/).
18
+
19
+ ## Execution
20
+
21
+ Import the `openlineage_prefect` package and execute `collect_and_process_runs()` asynchronously in its own process.
22
+
23
+ For example:
24
+
25
+ ```py
26
+ import asyncio
27
+ import openlineage_prefect
28
+ from openlineage_prefect.prefect_adapter.listener import PrefectOpenLineageListener
29
+
30
+ async def main():
31
+ await PrefectOpenLineageListener().collect_and_process_runs()
32
+
33
+ if __name__ == "__main__":
34
+ asyncio.run(main())
35
+ ```
36
+
37
+ ## Datasets
38
+
39
+ The integration looks for datasets in Prefect Artifacts. To attach input and output datasets to job runs, use `create_table_artifact()` from the Artifact library. Provide a namespace, typically the dataset URI, and name to the adapter via an artifact's `table` and `description`, respectively. Distinguish the type of dataset by appending `_output` or `_input` to the description.
40
+
41
+ For example:
42
+
43
+ ```py
44
+ ol_table = [{"database_uri":"duckdb:///customers_db", "table":"customers"}]
45
+
46
+ create_table_artifact(
47
+ key="upstream-insert",
48
+ table=ol_table,
49
+ description="ol-dataset_output"
50
+ )
51
+ ```
@@ -0,0 +1,62 @@
1
+ [build-system]
2
+ requires = ["uv_build >= 0.12.10, <0.13.0"]
3
+ build-backend = "uv_build"
4
+
5
+ [project]
6
+ name = "openlineage-prefect"
7
+ version = "0.0.1b3"
8
+ description = "OpenLineage integration with Prefect"
9
+ readme = "README.md"
10
+ requires-python = ">=3.12"
11
+ keywords = [
12
+ "openlineage",
13
+ "prefect",
14
+ ]
15
+ classifiers = [
16
+ "Development Status :: 4 - Beta",
17
+ "Intended Audience :: Developers",
18
+ "License :: OSI Approved :: Apache Software License",
19
+ "Programming Language :: Python :: 3",
20
+ "Programming Language :: Python :: 3.12",
21
+ ]
22
+ dependencies = [
23
+ "prefect>=3.8.5",
24
+ "openlineage-python>=1.53.0",
25
+ ]
26
+
27
+ [[project.authors]]
28
+ name = "OpenLineage"
29
+
30
+ [project.license]
31
+ text = "Apache-2.0"
32
+
33
+ [project.optional-dependencies]
34
+ dev = [
35
+ "mypy",
36
+ "openlineage-python",
37
+ "pillow",
38
+ "pre-commit",
39
+ "pytest>=8.3",
40
+ "pytest-asyncio",
41
+ "pytest-cov",
42
+ "pytest-env",
43
+ "pytest-timeout>=2.3.1",
44
+ "pytest-xdist",
45
+ ]
46
+
47
+ [project.urls]
48
+ Homepage = "https://openlineage.io/"
49
+ Source = "https://github.com/OpenLineage/OpenLineage/tree/main/integration/prefect"
50
+
51
+ [tool.bumpversion]
52
+ current_version = "0.0.1b3"
53
+ commit = false
54
+ tag = false
55
+ parse = '(?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(?P<rc>.*)'
56
+ serialize = [
57
+ "{major}.{minor}.{patch}{rc}",
58
+ "{major}.{minor}.{patch}",
59
+ ]
60
+
61
+ [tool.bumpversion.files]
62
+ filename = "pyproject.toml"
@@ -0,0 +1,56 @@
1
+ [build-system]
2
+ requires = ["uv_build >= 0.12.10, <0.13.0"]
3
+ build-backend = "uv_build"
4
+
5
+ [project]
6
+ name = "openlineage-prefect"
7
+ version = "0.0.1b3"
8
+ description = "OpenLineage integration with Prefect"
9
+ readme = "README.md"
10
+ authors = [{name = "OpenLineage"}]
11
+ license = {text = "Apache-2.0"}
12
+ requires-python = ">=3.12"
13
+ keywords = ["openlineage", "prefect"]
14
+ classifiers = [
15
+ "Development Status :: 4 - Beta",
16
+ "Intended Audience :: Developers",
17
+ "License :: OSI Approved :: Apache Software License",
18
+ "Programming Language :: Python :: 3",
19
+ "Programming Language :: Python :: 3.12",
20
+ ]
21
+
22
+ dependencies = [
23
+ "prefect>=3.8.5",
24
+ "openlineage-python>=1.53.0",
25
+ ]
26
+
27
+ [project.optional-dependencies]
28
+ dev = [
29
+ "mypy",
30
+ "openlineage-python",
31
+ "pillow",
32
+ "pre-commit",
33
+ "pytest>=8.3",
34
+ "pytest-asyncio",
35
+ "pytest-cov",
36
+ "pytest-env",
37
+ "pytest-timeout>=2.3.1",
38
+ "pytest-xdist",
39
+ ]
40
+
41
+ [project.urls]
42
+ Homepage = "https://openlineage.io/"
43
+ Source = "https://github.com/OpenLineage/OpenLineage/tree/main/integration/prefect"
44
+
45
+ [tool.bumpversion]
46
+ current_version = "0.0.1b3"
47
+ commit = false
48
+ tag = false
49
+ parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)(?P<rc>.*)"
50
+ serialize = [
51
+ "{major}.{minor}.{patch}{rc}",
52
+ "{major}.{minor}.{patch}"
53
+ ]
54
+
55
+ [tool.bumpversion.files]
56
+ filename = "pyproject.toml"
@@ -0,0 +1,2 @@
1
+ # Copyright 2018-2026 contributors to the OpenLineage project
2
+ # SPDX-License-Identifier: Apache-2.0
@@ -0,0 +1,190 @@
1
+ # Copyright 2018-2026 contributors to the OpenLineage project
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ import logging
5
+ from datetime import datetime
6
+
7
+ from .facets.run_facets import PrefectDeploymentRunFacet
8
+ from openlineage.client import OpenLineageClient
9
+ from openlineage.client.event_v2 import Dataset
10
+ from openlineage.client.facet import (
11
+ JobTypeJobFacet,
12
+ NominalTimeRunFacet,
13
+ ParentRunFacet,
14
+ )
15
+ from openlineage.client.facet_v2 import job_dependencies_run, processing_engine_run
16
+ from openlineage.client.run import Job, Run, RunEvent, RunState
17
+
18
+ PRODUCER: str = "https://github.com/OpenLineage/openlineage/integration/prefect"
19
+
20
+ logger: logging.Logger = logging.getLogger(__name__)
21
+
22
+
23
+ class PrefectOpenLineageAdapter:
24
+ def __init__(self, client: OpenLineageClient | None = None):
25
+ self.client = client or OpenLineageClient()
26
+
27
+ def create_and_emit_flow_event(
28
+ self,
29
+ run_id: str,
30
+ event_type: str,
31
+ event_time: datetime,
32
+ flow_name: str | None = None,
33
+ flow_namespace: str | None = None,
34
+ prefect_version: str | None = None,
35
+ deployment_id: str | None = None,
36
+ deployment_created: str | None = None,
37
+ deployment_updated: str | None = None,
38
+ deployment_name: str | None = None,
39
+ ) -> RunEvent:
40
+ """Create and emit a flow-level OpenLineage event."""
41
+
42
+ match event_type:
43
+ case "START":
44
+ event_type = RunState.START
45
+ case "COMPLETE":
46
+ event_type = RunState.COMPLETE
47
+ case "FAILED":
48
+ event_type = RunState.FAIL
49
+
50
+ if deployment_id:
51
+ run_facets = {
52
+ "prefectDeployment": PrefectDeploymentRunFacet(
53
+ deploymentId=deployment_id,
54
+ created=deployment_created,
55
+ updated=deployment_updated,
56
+ name=deployment_name,
57
+ ),
58
+ "processingEngine": processing_engine_run.ProcessingEngineRunFacet(
59
+ version=prefect_version, name="Prefect"
60
+ ),
61
+ }
62
+ else:
63
+ run_facets = {
64
+ "processingEngine": processing_engine_run.ProcessingEngineRunFacet(
65
+ version=prefect_version, name="Prefect"
66
+ )
67
+ }
68
+
69
+ job_facets = {
70
+ "jobType": JobTypeJobFacet(
71
+ processingType="BATCH", integration="Prefect", jobType="FLOW"
72
+ )
73
+ }
74
+
75
+ run_event = RunEvent(
76
+ eventType=event_type,
77
+ eventTime=event_time.isoformat(),
78
+ run=Run(run_id, run_facets),
79
+ job=Job(flow_namespace, flow_name, job_facets),
80
+ producer=PRODUCER,
81
+ )
82
+
83
+ try:
84
+ self.client.emit(run_event)
85
+ logger.info("Emitted OpenLineage event successfully.")
86
+ except Exception:
87
+ logger.exception("OpenLineage event not sent.")
88
+
89
+ def create_and_emit_task_event(
90
+ self,
91
+ run_id: str,
92
+ event_type: str,
93
+ event_time: datetime,
94
+ expectedevent_time: datetime | None = None,
95
+ flow_run_id: str | None = None,
96
+ flow_name: str | None = None,
97
+ task_name: str | None = None,
98
+ namespace: str | None = None,
99
+ job_deps: list | None = None,
100
+ prefect_version: str | None = None,
101
+ deployment_id: str | None = None,
102
+ deployment_created: str | None = None,
103
+ deployment_updated: str | None = None,
104
+ deployment_name: str | None = None,
105
+ input_datasets: list | None = None,
106
+ output_datasets: list | None = None,
107
+ ) -> RunEvent:
108
+ """Create and emit a task-level OpenLineage event."""
109
+
110
+ match event_type:
111
+ case "START":
112
+ event_type = RunState.START
113
+ case "COMPLETE":
114
+ event_type = RunState.COMPLETE
115
+ case "FAILED":
116
+ event_type = RunState.FAIL
117
+
118
+ if deployment_id:
119
+ run_facets = {
120
+ "nominalTime": NominalTimeRunFacet(nominalStartTime=expectedevent_time),
121
+ "processingEngine": processing_engine_run.ProcessingEngineRunFacet(
122
+ version=prefect_version, name="Prefect"
123
+ ),
124
+ "parentRun": ParentRunFacet(
125
+ run={"run_id": flow_run_id},
126
+ job={"namespace": namespace, "name": flow_name},
127
+ ),
128
+ "prefectDeployment": PrefectDeploymentRunFacet(
129
+ deploymentId=deployment_id,
130
+ created=deployment_created,
131
+ updated=deployment_updated,
132
+ name=deployment_name,
133
+ ),
134
+ }
135
+ else:
136
+ run_facets = {
137
+ "nominalTime": NominalTimeRunFacet(nominalStartTime=expectedevent_time),
138
+ "processingEngine": processing_engine_run.ProcessingEngineRunFacet(
139
+ version=prefect_version, name="Prefect"
140
+ ),
141
+ "parentRun": ParentRunFacet(
142
+ run={"runId": flow_run_id},
143
+ job={"namespace": namespace, "name": flow_name},
144
+ ),
145
+ }
146
+
147
+ if job_deps:
148
+ upstream_jobs = [
149
+ job_dependencies_run.JobDependency(
150
+ job=job_dependencies_run.JobIdentifier(
151
+ namespace=dep["namespace"], name=dep["name"]
152
+ )
153
+ )
154
+ for dep in job_deps
155
+ ]
156
+ run_facets["jobDependencies"] = (
157
+ job_dependencies_run.JobDependenciesRunFacet(upstream=upstream_jobs)
158
+ )
159
+
160
+ job_facets = {
161
+ "jobType": JobTypeJobFacet(
162
+ processingType="BATCH", integration="Prefect", jobType="TASK"
163
+ )
164
+ }
165
+
166
+ inputs = [
167
+ Dataset(namespace=dataset["uri"], name=dataset["table"])
168
+ for dataset in input_datasets
169
+ ]
170
+ outputs = [
171
+ Dataset(namespace=dataset["uri"], name=dataset["table"])
172
+ for dataset in output_datasets
173
+ ]
174
+
175
+ run_event = RunEvent(
176
+ eventType=event_type,
177
+ eventTime=event_time.isoformat(),
178
+ run=Run(run_id, run_facets),
179
+ job=Job(namespace, task_name, job_facets),
180
+ producer=PRODUCER,
181
+ inputs=inputs,
182
+ outputs=outputs,
183
+ )
184
+
185
+ try:
186
+ self.client.emit(run_event)
187
+ logger.info("Emitted OpenLineage event successfully.")
188
+ except Exception:
189
+ logger.exception("OpenLineage event not sent.")
190
+
@@ -0,0 +1,46 @@
1
+ {
2
+ "$schema": "https://json-schema.org/draft/2020-12/schema",
3
+ "$id": "https://github.com/OpenLineage/openlineage/integration/facets/PrefectDeploymentRunFacet.json",
4
+ "$defs": {
5
+ "PrefectDeploymentRunFacet": {
6
+ "allOf": [
7
+ {
8
+ "$ref": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/RunFacet"
9
+ },
10
+ {
11
+ "type": "object",
12
+ "properties": {
13
+ "deploymentId": {
14
+ "description": "The ID of the Prefect deployment",
15
+ "type": "string"
16
+ },
17
+ "created": {
18
+ "description": "An [ISO-8601](https://en.wikipedia.org/wiki/ISO_8601) timestamp representing when the deployment was created",
19
+ "type": "string",
20
+ "format": "date-time",
21
+ "example": "2020-12-17T03:00:00.000Z"
22
+ },
23
+ "updated": {
24
+ "description": "An [ISO-8601](https://en.wikipedia.org/wiki/ISO_8601) timestamp representing when the deployment was last updated",
25
+ "type": "string",
26
+ "format": "date-time",
27
+ "example": "2020-12-17T04:00:00.000Z"
28
+ },
29
+ "name": {
30
+ "description": "The name of the Prefect deployment",
31
+ "type": "string"
32
+ }
33
+ },
34
+ "required": ["deploymentId", "created", "updated", "name"]
35
+ }
36
+ ],
37
+ "type": "object"
38
+ }
39
+ },
40
+ "type": "object",
41
+ "properties": {
42
+ "PrefectDeploymentRunFacet": {
43
+ "$ref": "#/$defs/PrefectDeploymentRunFacet"
44
+ }
45
+ }
46
+ }
@@ -0,0 +1,30 @@
1
+ # Copyright 2018-2026 contributors to the OpenLineage project
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ import attr
5
+ from openlineage.client.facet import BaseFacet
6
+
7
+
8
+ @attr.define
9
+ class PrefectDeploymentRunFacet(BaseFacet):
10
+ deploymentId: str
11
+ created: str
12
+ updated: str
13
+ name: str
14
+
15
+ def __init__(self, deploymentId, created, updated, name):
16
+ super().__init__()
17
+ self.deploymentId = deploymentId
18
+ self.created = created
19
+ self.updated = updated
20
+ self.name = name
21
+
22
+ @staticmethod
23
+ def _get_schema() -> str:
24
+ return "https://raw.githubusercontent.com/OpenLineage/openlineage/integration/prefect/facets/PrefectDeploymentRunFacet.json"
25
+
26
+ @staticmethod
27
+ def _get_producer() -> str:
28
+ return (
29
+ "https://github.com/OpenLineage/openlineage/integration/prefect"
30
+ )
@@ -0,0 +1,382 @@
1
+ # Copyright 2018-2026 contributors to the OpenLineage project
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ import ast
5
+ import asyncio
6
+ import logging
7
+ import os
8
+ from datetime import datetime
9
+
10
+ from .adapter import PrefectOpenLineageAdapter
11
+ from openlineage.client.uuid import generate_static_uuid
12
+
13
+ from prefect import client
14
+ from prefect.client.orchestration import get_client
15
+ from prefect.events.clients import get_events_subscriber
16
+ from prefect.events.filters import EventFilter, EventNameFilter
17
+ from prefect.events.schemas.events import Event
18
+ from prefect.exceptions import ObjectNotFound, PrefectHTTPStatusError
19
+
20
+ JOB_NAMESPACE: str = os.environ.get("OPENLINEAGE_NAMESPACE", "default")
21
+
22
+ logger: logging.Logger = logging.getLogger(__name__)
23
+
24
+
25
+ class PrefectOpenLineageListener:
26
+ def __init__(
27
+ self,
28
+ client: client = None,
29
+ adapter: PrefectOpenLineageAdapter = None,
30
+ ):
31
+ self.client = client or get_client()
32
+ self.ol_adapter = adapter or PrefectOpenLineageAdapter()
33
+
34
+ def build_run_id(
35
+ self, execution_time: datetime, run_name: str, namespace: str
36
+ ) -> str:
37
+ """
38
+ Build a deterministic UUID for the OpenLineage run based on the execution
39
+ time, run name, and namespace.
40
+ """
41
+
42
+ return str(
43
+ generate_static_uuid(
44
+ instant=execution_time,
45
+ data=f"{namespace}.{run_name}".encode(),
46
+ )
47
+ )
48
+
49
+ async def get_deployment_and_flow_info(self, flow_run_id: str) -> tuple:
50
+ flow_run = await self.client.read_flow_run(flow_run_id)
51
+ flow_id = flow_run.flow_id
52
+ flow = await self.client.read_flow(flow_id)
53
+ flow_name = flow.name
54
+
55
+ try:
56
+ deployment = await self.client.read_deployment(flow_run.deployment_id)
57
+ try:
58
+ ns = deployment.job_variables["env"]["OPENLINEAGE_NAMESPACE"]
59
+ except KeyError:
60
+ ns = JOB_NAMESPACE
61
+ logger.info(
62
+ "OPENLINEAGE_NAMESPACE deployment variable not found. Using \
63
+ OPENLINEAGE_NAMESPACE env variable."
64
+ )
65
+ if JOB_NAMESPACE == "default":
66
+ logger.info(
67
+ "OPENLINEAGE_NAMESPACE env variable not set. Namespace will be 'default.'"
68
+ )
69
+ return (
70
+ str(deployment.id),
71
+ flow_run.start_time,
72
+ deployment.created.isoformat(),
73
+ deployment.updated.isoformat(),
74
+ deployment.name,
75
+ ns,
76
+ flow_name,
77
+ )
78
+
79
+ except AttributeError:
80
+ logger.info("Deployment not found for flow run: %s", flow_run_id)
81
+ ns = JOB_NAMESPACE
82
+ logger.info(
83
+ "OPENLINEAGE_NAMESPACE deployment variable not found. Using \
84
+ OPENLINEAGE_NAMESPACE env variable."
85
+ )
86
+ if JOB_NAMESPACE == "default":
87
+ logger.info(
88
+ "OPENLINEAGE_NAMESPACE env variable not set. Namespace will be 'default.'"
89
+ )
90
+
91
+ return (None, flow_run.start_time, None, None, None, ns, flow_name)
92
+
93
+ async def get_prefect_version(self) -> str | None:
94
+ """Retrieve the Prefect version from the Prefect API."""
95
+
96
+ try:
97
+ response = await self.client._client.get("/admin/version")
98
+ return response.json()
99
+ except TypeError:
100
+ logger.info(
101
+ "Cannot get the Prefect version. Did you set the PREFECT_API_URL?"
102
+ )
103
+
104
+ async def get_flow_ns(self, flow_run_id: str) -> str:
105
+ """
106
+ Looks for OPENLINEAGE_NAMESPACE job env variable in deployment.
107
+ """
108
+ flow_run = await self.client.read_flow_run(flow_run_id)
109
+ try:
110
+ deployment = await self.client.read_deployment(flow_run.deployment_id)
111
+ try:
112
+ ns = deployment.job_variables["env"]["OPENLINEAGE_NAMESPACE"]
113
+ except KeyError:
114
+ ns = JOB_NAMESPACE
115
+ logger.info(
116
+ "OPENLINEAGE_NAMESPACE deployment variable not found. Using \
117
+ OPENLINEAGE_NAMESPACE env variable."
118
+ )
119
+ if JOB_NAMESPACE == "default":
120
+ logger.info(
121
+ "OPENLINEAGE_NAMESPACE env variable not found. Namespace will be 'default.'"
122
+ )
123
+ return ns
124
+ except AttributeError:
125
+ logger.info("Deployment not found for flow run: %s", flow_run_id)
126
+ logger.info(
127
+ "OPENLINEAGE_NAMESPACE deployment variable not found. Using \
128
+ OPENLINEAGE_NAMESPACE env variable."
129
+ )
130
+ if JOB_NAMESPACE == "default":
131
+ logger.info(
132
+ "OPENLINEAGE_NAMESPACE env variable not found. Namespace will be 'default.'"
133
+ )
134
+ return JOB_NAMESPACE
135
+
136
+ async def get_job_ns(self, task_run_id: str) -> str:
137
+ """Look for the OPENLINEAGE_NAMESPACE job env variable in the parent deployment."""
138
+
139
+ task_run = await self.client.read_task_run(task_run_id)
140
+ return await self.get_flow_ns(task_run.flow_run_id)
141
+
142
+ async def get_flow_run_start_time(self, flow_run_id: str) -> datetime:
143
+ """Retrieve the start time of a flow run."""
144
+
145
+ flow_run = await self.client.read_flow_run(flow_run_id)
146
+ return flow_run.start_time
147
+
148
+ async def get_artifacts_by_task_run(self, run_id: str) -> list[dict]:
149
+ """Retrieve artifacts associated with a given task run ID."""
150
+ payload = {"artifacts": {"task_run_id": {"any_": [run_id]}}}
151
+ response = await self.client._client.post("/artifacts/filter", json=payload)
152
+ if response.status_code == 200:
153
+ dataset_info = []
154
+ artifacts = response.json()
155
+ for artifact in artifacts:
156
+ if artifact["description"] and "ol-dataset" in artifact["description"]:
157
+ dataset_type = artifact["description"].split("_")[-1].lower()
158
+ data_list = ast.literal_eval(artifact["data"])
159
+ uri = data_list[0]["database_uri"].lower()
160
+ table = data_list[0]["table"].lower()
161
+ dataset_info.append(
162
+ {"uri": uri, "table": table, "dataset_type": dataset_type}
163
+ )
164
+ return dataset_info
165
+ logger.info("No datasets found for task run.")
166
+ return []
167
+
168
+ async def get_parent_runs(
169
+ self, payload: dict, prefect_task_run_id: str
170
+ ) -> list[dict]:
171
+ """Retrieve the parent runs for a given task run."""
172
+
173
+ try:
174
+ parent_runs = []
175
+ task_parents = payload["task_run"]["task_inputs"]["__parents__"]
176
+ for parent in task_parents:
177
+ task_run_id: str | None = (
178
+ parent["id"] if parent["input_type"] == "task_run" else None
179
+ )
180
+ if task_run_id:
181
+ parent_namespace: dict = await self.get_job_ns(task_run_id)
182
+ parent_run = await self.client.read_task_run(task_run_id)
183
+ parent_name = parent_run.name[0:-4]
184
+ parent_run_id = self.build_run_id(
185
+ parent_run.start_time, parent_name, parent_namespace
186
+ )
187
+ parent_runs.append(
188
+ {
189
+ "name": parent_name,
190
+ "namespace": parent_namespace,
191
+ "id": parent_run_id,
192
+ }
193
+ )
194
+ return parent_runs
195
+ except KeyError:
196
+ logger.info("No task parents found for %s", prefect_task_run_id)
197
+ return []
198
+
199
+ async def collect_and_process_flow_runs(
200
+ self, prefect_version: str, event: Event, event_state: str
201
+ ) -> None:
202
+ """Retrieve the flow runs for a given event and emit OpenLineage events."""
203
+
204
+ for res in event.related:
205
+ if res["prefect.resource.role"] == "flow":
206
+ flow_name = res["prefect.resource.name"]
207
+ prefect_flow_run_id = event.resource.id.split(".")[-1]
208
+ event_time: datetime = datetime.fromisoformat(
209
+ event.resource["prefect.state-timestamp"]
210
+ )
211
+ deployment_and_flow_info = await self.get_deployment_and_flow_info(
212
+ prefect_flow_run_id
213
+ )
214
+ deployment_id = deployment_and_flow_info[0]
215
+ start_time: datetime = deployment_and_flow_info[1]
216
+ deployment_created = deployment_and_flow_info[2]
217
+ deployment_updated = deployment_and_flow_info[3]
218
+ deployment_name = deployment_and_flow_info[4]
219
+ flow_namespace = deployment_and_flow_info[5]
220
+ try:
221
+ ol_flow_run_id: str = self.build_run_id(
222
+ start_time, flow_name, flow_namespace
223
+ )
224
+ except AttributeError:
225
+ logger.info(
226
+ "No Prefect run found for %s. OpenLineage event will not be emitted.",
227
+ prefect_flow_run_id,
228
+ )
229
+ continue
230
+
231
+ self.ol_adapter.create_and_emit_flow_event(
232
+ run_id=ol_flow_run_id,
233
+ event_type=event_state,
234
+ event_time=event_time,
235
+ flow_name=flow_name,
236
+ flow_namespace=flow_namespace,
237
+ prefect_version=prefect_version,
238
+ deployment_id=deployment_id,
239
+ deployment_created=deployment_created,
240
+ deployment_updated=deployment_updated,
241
+ deployment_name=deployment_name,
242
+ )
243
+
244
+ async def collect_and_process_task_runs(
245
+ self, prefect_version: str, event: Event, event_state: str
246
+ ) -> None:
247
+ """Retrieve the task runs for a given event and emit OpenLineage events."""
248
+
249
+ task_name = event.resource.name[0:-4]
250
+ event_time = datetime.fromisoformat(event.resource["prefect.state-timestamp"])
251
+ expected_start_time = event.payload["task_run"]["expected_start_time"]
252
+ prefect_task_run_id = event.resource.id.split(".")[-1]
253
+ try:
254
+ task_run = await self.client.read_task_run(prefect_task_run_id)
255
+ namespace = await self.get_job_ns(prefect_task_run_id)
256
+
257
+ # Skip task runs without a start time
258
+ if task_run.start_time:
259
+ run_start_time = task_run.start_time
260
+ else:
261
+ logger.warning(
262
+ "No start time found for task run %s. An event will not be emitted.",
263
+ prefect_task_run_id,
264
+ )
265
+ return
266
+
267
+ ol_task_run_id: str = self.build_run_id(
268
+ run_start_time, task_name, namespace
269
+ )
270
+
271
+ # Get datasets from Prefect Artifacts
272
+ datasets = await self.get_artifacts_by_task_run(prefect_task_run_id)
273
+ input_datasets = [
274
+ dataset for dataset in datasets if dataset["dataset_type"] == "input"
275
+ ]
276
+ output_datasets = [
277
+ dataset for dataset in datasets if dataset["dataset_type"] == "output"
278
+ ]
279
+
280
+ # Get job dependencies (Prefect "parents") info for JobDependenciesRunFacet
281
+ parent_runs = await self.get_parent_runs(event.payload, prefect_task_run_id)
282
+
283
+ # Get flow run info for ParentRunFacet
284
+ for res in event.related:
285
+ if res["prefect.resource.role"] == "flow-run":
286
+ flow_run_id = res["prefect.resource.id"].split(".")[-1]
287
+ deployment_and_flow_info = await self.get_deployment_and_flow_info(
288
+ flow_run_id
289
+ )
290
+ deployment_id = deployment_and_flow_info[0]
291
+ deployment_created = deployment_and_flow_info[2]
292
+ deployment_updated = deployment_and_flow_info[3]
293
+ deployment_name = deployment_and_flow_info[4]
294
+ flow_name = deployment_and_flow_info[6]
295
+ flow_start_time = deployment_and_flow_info[1]
296
+ try:
297
+ ol_flow_run_id: str = self.build_run_id(
298
+ flow_start_time, flow_name, namespace
299
+ )
300
+ except AttributeError:
301
+ logger.info(
302
+ "No Prefect run found for %s. ParentRunFacet will not be included.",
303
+ flow_run_id,
304
+ )
305
+ continue
306
+
307
+ self.ol_adapter.create_and_emit_task_event(
308
+ run_id=ol_task_run_id,
309
+ event_type=event_state,
310
+ event_time=event_time,
311
+ expectedevent_time=expected_start_time,
312
+ flow_run_id=ol_flow_run_id,
313
+ flow_name=flow_name,
314
+ task_name=task_name,
315
+ namespace=namespace,
316
+ job_deps=parent_runs,
317
+ prefect_version=prefect_version,
318
+ deployment_id=deployment_id,
319
+ deployment_created=deployment_created,
320
+ deployment_updated=deployment_updated,
321
+ deployment_name=deployment_name,
322
+ input_datasets=input_datasets,
323
+ output_datasets=output_datasets,
324
+ )
325
+ except (PrefectHTTPStatusError, ObjectNotFound):
326
+ logger.info(
327
+ "No Prefect run found for %s, will not attempt to create OpenLineage event.",
328
+ prefect_task_run_id,
329
+ )
330
+
331
+ async def collect_and_process_runs(self) -> None:
332
+ """Collect and process Prefect events."""
333
+
334
+ try:
335
+ os.environ.get("PREFECT_API_URL")
336
+ except TypeError:
337
+ logger.warning("PREFECT_API_URL not set. Prefect events will not be received.")
338
+ return
339
+
340
+ filter_criteria = EventFilter(
341
+ event=EventNameFilter(
342
+ prefix=[
343
+ "prefect.task-run.",
344
+ "prefect.flow-run.",
345
+ "prefect.asset.materialization.",
346
+ ]
347
+ )
348
+ )
349
+
350
+ async with get_events_subscriber(filter=filter_criteria) as subscriber, self.client:
351
+ prefect_version = await self.get_prefect_version()
352
+
353
+ async for event in subscriber:
354
+ entity_type = event.event.split(".")[1]
355
+ prefect_state = event.event.split(".")[-1]
356
+
357
+ if prefect_state in ["Running", "Completed", "Failed"]:
358
+ match prefect_state:
359
+ case "Running":
360
+ event_state = "START"
361
+ case "Completed":
362
+ event_state = "COMPLETE"
363
+ case "Failed":
364
+ event_state = "FAILED"
365
+
366
+ if entity_type == "flow-run":
367
+ await self.collect_and_process_flow_runs(
368
+ prefect_version, event, event_state
369
+ )
370
+
371
+ if entity_type == "task-run":
372
+ await self.collect_and_process_task_runs(
373
+ prefect_version, event, event_state
374
+ )
375
+
376
+
377
+ async def main():
378
+ await PrefectOpenLineageListener().collect_and_process_runs()
379
+
380
+
381
+ if __name__ == "__main__":
382
+ asyncio.run(main())