kacpermuda commented on code in PR #47444:
URL: https://github.com/apache/airflow/pull/47444#discussion_r1989102284
##########
providers/openlineage/src/airflow/providers/openlineage/utils/spark.py:
##########
@@ -53,35 +53,66 @@ def
_get_parent_job_information_as_spark_properties(context: Context) -> dict:
def _get_transport_information_as_spark_properties() -> dict:
"""Retrieve transport information as Spark properties."""
- transport =
get_openlineage_listener().adapter.get_or_create_openlineage_client().transport
- if transport.kind != "http":
- log.info(
- "OpenLineage transport type `%s` does not support automatic "
- "injection of OpenLineage transport information into Spark
properties.",
- transport.kind,
- )
- return {}
-
- properties = {
- "spark.openlineage.transport.type": transport.kind,
- "spark.openlineage.transport.url": transport.url,
- "spark.openlineage.transport.endpoint": transport.endpoint,
- "spark.openlineage.transport.timeoutInMillis": str(
- int(transport.timeout * 1000) # convert to milliseconds, as
required by Spark integration
- ),
- }
- if transport.compression:
- properties["spark.openlineage.transport.compression"] =
str(transport.compression)
- if hasattr(transport.config.auth, "api_key") and
transport.config.auth.get_bearer():
- properties["spark.openlineage.transport.auth.type"] = "api_key"
- properties["spark.openlineage.transport.auth.apiKey"] =
transport.config.auth.get_bearer()
+ def _get_transport_information(tp) -> dict:
+ properties = {
+ "type": tp.kind,
+ "url": tp.url,
+ "endpoint": tp.endpoint,
+ "timeoutInMillis": str(
+ int(tp.timeout * 1000) # convert to milliseconds, as required
by Spark integration
+ ),
+ }
+ if hasattr(tp, "compression") and tp.compression:
+ properties["compression"] = str(tp.compression)
+
+ if hasattr(tp.config.auth, "api_key") and tp.config.auth.get_bearer():
+ properties["auth.type"] = "api_key"
+ properties["auth.apiKey"] = tp.config.auth.get_bearer()
+
+ if hasattr(tp.config, "custom_headers") and tp.config.custom_headers:
+ for key, value in tp.config.custom_headers.items():
+ properties[f"headers.{key}"] = value
+ return properties
+
+ def _format_transport(props: dict, transport: dict, name: str | None):
+ for key, value in transport.items():
+ if name:
+ props[f"spark.openlineage.transport.transports.{name}.{key}"]
= value
+ else:
+ props[f"spark.openlineage.transport.{key}"] = value
+ return props
- if hasattr(transport.config, "custom_headers") and
transport.config.custom_headers:
- for key, value in transport.config.custom_headers.items():
- properties[f"spark.openlineage.transport.headers.{key}"] = value
+ transport =
get_openlineage_listener().adapter.get_or_create_openlineage_client().transport
- return properties
+ if transport.kind == "composite":
+ http_transports = {}
+ for nested_transport in transport.transports:
+ if nested_transport.kind == "http":
+ http_transports[nested_transport.name] =
_get_transport_information(nested_transport)
+ if len(http_transports) == 0:
+ log.warning(
+ "OpenLineage transport type `composite` does not contain http
transport. Skipping "
+ "injection of OpenLineage transport information into Spark
properties.",
+ )
+ return {}
Review Comment:
```suggestion
for nested_transport in transport.transports:
if nested_transport.kind == "http":
http_transports[nested_transport.name] =
_get_transport_information(nested_transport)
else:
log.warning(
"OpenLineage transport type `composite` contain `%s`
transport that will not be injected "
" into Spark properties. ... # TODO",
nested_transport.kind
)
if len(http_transports) == 0:
log.info(
"OpenLineage transport type `composite` does not contain
http transport. Skipping "
"injection of OpenLineage transport information into Spark
properties.",
)
return {}
```
This is what i meant earlier about adding warning log. The other one can
stay on the info level.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
