Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
mobuchowski merged PR #66342: URL: https://github.com/apache/airflow/pull/66342 -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
VladaZakharova commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4533307370 > @VladaZakharova Can you mark PR as non-draft? I'll ping Maciej for re-review Done! thanks a lot :) -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
kacpermuda commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4533297080 @VladaZakharova Can you mark PR as non-draft? I'll ping Maciej for re-review -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
VladaZakharova commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4532798109 hey there! can you please check changes again? :) -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
kacpermuda commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4486603986 I think this is ready to go out of draft, tests are passing, code looks good, one final review and merge from Maciej is needed. -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
VladaZakharova commented on code in PR #66342:
URL: https://github.com/apache/airflow/pull/66342#discussion_r3259253645
##
providers/openlineage/src/airflow/providers/openlineage/token_provider.py:
##
@@ -0,0 +1,137 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook
+
+AIRFLOW_CONNECTION_API_KEY_AUTH_TYPE = "airflow_connection_api_key"
+_DEFAULT_EXTRA_KEYS = ("apiKey", "api_key", "apikey", "token", "access_token")
+
+
+class OpenLineageAirflowConnectionAuthError(AirflowException):
+"""Raised when OpenLineage API key auth cannot be resolved from an Airflow
connection."""
+
+
+class OpenLineageAirflowConnectionConfigError(AirflowException):
+"""Raised when OpenLineage config cannot be resolved from an Airflow
connection."""
+
+
+class AirflowConnectionConfigProvider:
+"""
+Resolve OpenLineage client configuration from an Airflow connection.
+
+The connection extra contains the full OpenLineage client config, for
example
+``{"transport": {"type": "console"}}``.
+"""
+
+def __init__(self, conn_id: str) -> None:
+if not conn_id:
+raise OpenLineageAirflowConnectionConfigError(
+"OpenLineage connection config requires a non-empty connection
ID."
+)
+self.conn_id = conn_id
+
+def get_config(self) -> dict[str, Any]:
+connection = BaseHook.get_connection(self.conn_id)
+extra = connection.extra_dejson
+config = self._get_config_from_extra(extra)
+if config is not None:
+return config
+
+raise OpenLineageAirflowConnectionConfigError(
+"OpenLineage connection config could not find configuration in
connection "
+f"`{self.conn_id}`. Expected OpenLineage config with `transport`
in connection extra."
+)
+
+def _get_config_from_extra(self, extra: dict[str, Any]) -> dict[str, Any]
| None:
+if "transport" in extra:
+return self._validate_config(extra)
+
+return None
Review Comment:
yes, that's right
--
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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
VladaZakharova commented on code in PR #66342: URL: https://github.com/apache/airflow/pull/66342#discussion_r3259251588 ## providers/openlineage/tests/unit/openlineage/test_token_provider.py: ## @@ -0,0 +1,141 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from airflow.providers.common.compat.sdk import BaseHook, Connection +from airflow.providers.openlineage.token_provider import ( +AIRFLOW_CONNECTION_API_KEY_AUTH_TYPE, +AirflowConnectionConfigProvider, +AirflowConnectionTokenProvider, +OpenLineageAirflowConnectionAuthError, +OpenLineageAirflowConnectionConfigError, +resolve_airflow_connection_auth, +) + + [email protected](BaseHook, "get_connection") +def test_get_api_key_from_connection_password(mock_get_connection): +mock_get_connection.return_value = Connection( +conn_id="openlineage_default", conn_type="http", password="api-key" +) + +provider = AirflowConnectionTokenProvider({"conn_id": "openlineage_default"}) + +assert provider.get_api_key() == "api-key" + + [email protected](BaseHook, "get_connection") +def test_get_api_key_from_default_connection_id(mock_get_connection): +mock_get_connection.return_value = Connection( +conn_id="openlineage_default", conn_type="http", password="api-key" +) + +provider = AirflowConnectionTokenProvider({}, default_conn_id="openlineage_default") + +assert provider.get_api_key() == "api-key" + + [email protected](BaseHook, "get_connection") +def test_get_api_key_from_connection_extra(mock_get_connection): +mock_get_connection.return_value = Connection( +conn_id="openlineage_default", conn_type="http", extra='{"api_key": "api-key-from-extra"}' +) + +provider = AirflowConnectionTokenProvider({"conn_id": "openlineage_default"}) + +assert provider.get_api_key() == "api-key-from-extra" + + +def test_missing_conn_id_raises_custom_exception(): +with pytest.raises(OpenLineageAirflowConnectionAuthError, match="requires a non-empty `conn_id`"): +AirflowConnectionTokenProvider({}) + + [email protected](BaseHook, "get_connection") +def test_missing_token_raises_custom_exception(mock_get_connection): +mock_get_connection.return_value = Connection(conn_id="openlineage_default", conn_type="http") + +provider = AirflowConnectionTokenProvider({"conn_id": "openlineage_default"}) + +with pytest.raises(OpenLineageAirflowConnectionAuthError, match="could not find a token"): +provider.get_api_key() + + [email protected](BaseHook, "get_connection") +def test_resolve_connection_auth_in_composite_transport(mock_get_connection): Review Comment: Updated 👍 -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
potiuk commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4476912860 @VladaZakharova — Removing the `ready for maintainer review` label and converting back to draft. **CI has failed** since the label was added (`4` failing checks). The label's contract is that the PR is ready for maintainer review — a regression like this means the PR temporarily isn't. Check the failing checks, fix the regression, push, then mark "Ready for review" again to re-enter the queue. See the [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria). No rush. --- _Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this [two-stage triage process](https://github.com/apache/airflow/blob/main/contributing-docs/25_maintainer_pr_triage.md#why-the-first-pass-is-automated) so that our maintainers' limited time is spent where it matters most: the conversation with you._ -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
kacpermuda commented on code in PR #66342: URL: https://github.com/apache/airflow/pull/66342#discussion_r3257769716 ## providers/openlineage/tests/unit/openlineage/test_token_provider.py: ## @@ -0,0 +1,141 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from airflow.providers.common.compat.sdk import BaseHook, Connection +from airflow.providers.openlineage.token_provider import ( +AIRFLOW_CONNECTION_API_KEY_AUTH_TYPE, +AirflowConnectionConfigProvider, +AirflowConnectionTokenProvider, +OpenLineageAirflowConnectionAuthError, +OpenLineageAirflowConnectionConfigError, +resolve_airflow_connection_auth, +) + + [email protected](BaseHook, "get_connection") +def test_get_api_key_from_connection_password(mock_get_connection): +mock_get_connection.return_value = Connection( +conn_id="openlineage_default", conn_type="http", password="api-key" +) + +provider = AirflowConnectionTokenProvider({"conn_id": "openlineage_default"}) + +assert provider.get_api_key() == "api-key" + + [email protected](BaseHook, "get_connection") +def test_get_api_key_from_default_connection_id(mock_get_connection): +mock_get_connection.return_value = Connection( +conn_id="openlineage_default", conn_type="http", password="api-key" +) + +provider = AirflowConnectionTokenProvider({}, default_conn_id="openlineage_default") + +assert provider.get_api_key() == "api-key" + + [email protected](BaseHook, "get_connection") +def test_get_api_key_from_connection_extra(mock_get_connection): +mock_get_connection.return_value = Connection( +conn_id="openlineage_default", conn_type="http", extra='{"api_key": "api-key-from-extra"}' +) + +provider = AirflowConnectionTokenProvider({"conn_id": "openlineage_default"}) + +assert provider.get_api_key() == "api-key-from-extra" + + +def test_missing_conn_id_raises_custom_exception(): +with pytest.raises(OpenLineageAirflowConnectionAuthError, match="requires a non-empty `conn_id`"): +AirflowConnectionTokenProvider({}) + + [email protected](BaseHook, "get_connection") +def test_missing_token_raises_custom_exception(mock_get_connection): +mock_get_connection.return_value = Connection(conn_id="openlineage_default", conn_type="http") + +provider = AirflowConnectionTokenProvider({"conn_id": "openlineage_default"}) + +with pytest.raises(OpenLineageAirflowConnectionAuthError, match="could not find a token"): +provider.get_api_key() + + [email protected](BaseHook, "get_connection") +def test_resolve_connection_auth_in_composite_transport(mock_get_connection): Review Comment: Can we add one more tests here for nested composite transport? So composite(transports=[http, composite(transports=[http, console])]), and make sure that it also works as expected? ## providers/openlineage/src/airflow/providers/openlineage/token_provider.py: ## @@ -0,0 +1,137 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from typing import Any + +from airflow.providers.common.compat.sdk import AirflowException, BaseHook + +AIRFLOW_CONNECTION_API_KEY_AUTH_TYPE = "airflow_connection_api_key" +_DEFAULT_EXTRA_KEYS = ("apiKey", "api_key", "apikey", "token", "access_token") + + +class OpenLineageAirflowConnectionAuthError(AirflowException): +"""Raised when OpenLineage API key auth
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
VladaZakharova commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4475903460 > Hey @VladaZakharova , I see some comments being resolved or marked as addressed, but see no recent commits. Is all the relevant code pushed as intended? not really, i forgot to push my changes :D -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
kacpermuda commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4475760198 Hey @VladaZakharova , I see some comments being resolved or marked as addressed, but see no recent commits. Is all the relevant code pushed as intended? -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
VladaZakharova commented on code in PR #66342:
URL: https://github.com/apache/airflow/pull/66342#discussion_r3257016652
##
providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py:
##
@@ -103,25 +108,52 @@ def get_or_create_openlineage_client(self) ->
OpenLineageClient:
return self._client
def get_openlineage_config(self) -> dict | None:
-# First, try to read from YAML file
+# First, try to read from Airflow connection
+openlineage_config_conn_id = conf.config_conn_id()
+if openlineage_config_conn_id:
+config =
AirflowConnectionConfigProvider(openlineage_config_conn_id).get_config()
+self._resolve_airflow_connection_auth(config=config,
config_conn_id=openlineage_config_conn_id)
+return config
+self.log.debug("OpenLineage config_conn_id configuration not found.")
+
+# Second, try to read from YAML file
openlineage_config_path = conf.config_path(check_legacy_env_var=False)
if openlineage_config_path:
-config = self._read_yaml_config(openlineage_config_path)
-return config
+yaml_config = self._read_yaml_config(openlineage_config_path)
+if yaml_config is None:
+return None
+self._resolve_airflow_connection_auth(yaml_config)
Review Comment:
I think we’re okay here. The resolver is called for each config source, but
it only does anything when it finds an auth block with:
{"type": "airflow_connection_api_key"}
For normal OpenLineage config, like regular api_key auth, it does not read
any Airflow connection and should behave the same as before. I added a test for
that case to make sure we don’t accidentally change it later.
--
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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
VladaZakharova commented on code in PR #66342:
URL: https://github.com/apache/airflow/pull/66342#discussion_r3257007947
##
providers/openlineage/src/airflow/providers/openlineage/token_provider.py:
##
@@ -0,0 +1,126 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook
+
+AIRFLOW_CONNECTION_API_KEY_AUTH_TYPE = "airflow_connection_api_key"
+OPENLINEAGE_CONFIG_EXTRA_KEY = "openlineage_config"
+_DEFAULT_EXTRA_KEYS = ("apiKey", "api_key", "apikey", "token", "access_token")
+
+
+class OpenLineageAirflowConnectionAuthError(AirflowException):
+"""Raised when OpenLineage API key auth cannot be resolved from an Airflow
connection."""
+
+
+class OpenLineageAirflowConnectionConfigError(AirflowException):
+"""Raised when OpenLineage config cannot be resolved from an Airflow
connection."""
Review Comment:
I agree that a dedicated openlineage connection type would be nicer for
users. I’m just not sure we should add it in this PR, because it feels like a
separate feature from loading the config from a connection.
For now I updated the docs and provider metadata to say that config_conn_id
should point to a Generic Airflow connection. That at least gives users a clear
choice instead of “pick any connection type”. We can add a proper OpenLineage
connection type later in a separate PR. WDYT?
##
providers/openlineage/src/airflow/providers/openlineage/token_provider.py:
##
@@ -0,0 +1,126 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook
+
+AIRFLOW_CONNECTION_API_KEY_AUTH_TYPE = "airflow_connection_api_key"
+OPENLINEAGE_CONFIG_EXTRA_KEY = "openlineage_config"
+_DEFAULT_EXTRA_KEYS = ("apiKey", "api_key", "apikey", "token", "access_token")
+
+
+class OpenLineageAirflowConnectionAuthError(AirflowException):
+"""Raised when OpenLineage API key auth cannot be resolved from an Airflow
connection."""
+
+
+class OpenLineageAirflowConnectionConfigError(AirflowException):
+"""Raised when OpenLineage config cannot be resolved from an Airflow
connection."""
+
+
+class AirflowConnectionConfigProvider:
+"""
+Resolve OpenLineage client configuration from an Airflow connection.
+
+The connection extra can contain the full OpenLineage client config, for
example
+``{"transport": {"type": "console"}}``. For convenience, it can also
contain only the transport
+config, for example ``{"type": "console"}``.
+"""
+
+def __init__(self, conn_id: str) -> None:
+if not conn_id:
+raise OpenLineageAirflowConnectionConfigError(
+"OpenLineage connection config requires a non-empty connection
ID."
+)
+self.conn_id = conn_id
+
+def get_config(self) -> dict[str, Any]:
+connection = BaseHook.get_connection(self.conn_id)
+extra = connection.extra_dejson
+config = self._get_config_from_extra(extra)
+if config is not None:
+return config
+
+raise OpenLineageAirflowConnectionConfigError(
+"OpenLineage connection config could not find configuration in
connection "
+f"`{self.conn_id}`. Expected full OpenLineage config or transport
config in connection extra."
+)
+
+def
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
kacpermuda commented on code in PR #66342:
URL: https://github.com/apache/airflow/pull/66342#discussion_r3242982959
##
providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py:
##
@@ -103,25 +108,52 @@ def get_or_create_openlineage_client(self) ->
OpenLineageClient:
return self._client
def get_openlineage_config(self) -> dict | None:
-# First, try to read from YAML file
+# First, try to read from Airflow connection
+openlineage_config_conn_id = conf.config_conn_id()
+if openlineage_config_conn_id:
+config =
AirflowConnectionConfigProvider(openlineage_config_conn_id).get_config()
+self._resolve_airflow_connection_auth(config=config,
config_conn_id=openlineage_config_conn_id)
+return config
+self.log.debug("OpenLineage config_conn_id configuration not found.")
+
+# Second, try to read from YAML file
openlineage_config_path = conf.config_path(check_legacy_env_var=False)
if openlineage_config_path:
-config = self._read_yaml_config(openlineage_config_path)
-return config
+yaml_config = self._read_yaml_config(openlineage_config_path)
+if yaml_config is None:
+return None
+self._resolve_airflow_connection_auth(yaml_config)
Review Comment:
Should every resolve be guarded and actually fire only if there is an
airflow connection defined? I think it fires every time now. We need to make
sure that for users not setting up this conn, nothing changes.
##
providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py:
##
@@ -103,25 +108,52 @@ def get_or_create_openlineage_client(self) ->
OpenLineageClient:
return self._client
def get_openlineage_config(self) -> dict | None:
-# First, try to read from YAML file
+# First, try to read from Airflow connection
+openlineage_config_conn_id = conf.config_conn_id()
+if openlineage_config_conn_id:
+config =
AirflowConnectionConfigProvider(openlineage_config_conn_id).get_config()
+self._resolve_airflow_connection_auth(config=config,
config_conn_id=openlineage_config_conn_id)
+return config
+self.log.debug("OpenLineage config_conn_id configuration not found.")
+
+# Second, try to read from YAML file
openlineage_config_path = conf.config_path(check_legacy_env_var=False)
if openlineage_config_path:
-config = self._read_yaml_config(openlineage_config_path)
-return config
+yaml_config = self._read_yaml_config(openlineage_config_path)
+if yaml_config is None:
+return None
+self._resolve_airflow_connection_auth(yaml_config)
+return yaml_config
self.log.debug("OpenLineage config_path configuration not found.")
-# Second, try to get transport config
+# Third, try to get transport config
transport_config = conf.transport()
if not transport_config:
self.log.debug("OpenLineage transport configuration not found.")
return None
-return {"transport": transport_config}
+config = {"transport": transport_config}
+self._resolve_airflow_connection_auth(config)
+return config
@staticmethod
-def _read_yaml_config(path: str) -> dict | None:
+def _read_yaml_config(path: str) -> dict[str, Any] | None:
with open(path) as config_file:
return yaml.safe_load(config_file)
+@staticmethod
+def _resolve_airflow_connection_auth(
Review Comment:
Maybe this can live on some class in the new airflow token provider module?
or even as standalone function there? It's very specific, maybe not needed on
adapter ?
##
providers/openlineage/src/airflow/providers/openlineage/token_provider.py:
##
@@ -0,0 +1,126 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook
+
+AIRFLOW_CONNECTION_API_KEY_AUTH_TYPE = "airfl
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
kacpermuda commented on code in PR #66342:
URL: https://github.com/apache/airflow/pull/66342#discussion_r3242976451
##
providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py:
##
@@ -103,25 +108,52 @@ def get_or_create_openlineage_client(self) ->
OpenLineageClient:
return self._client
def get_openlineage_config(self) -> dict | None:
-# First, try to read from YAML file
+# First, try to read from Airflow connection
+openlineage_config_conn_id = conf.config_conn_id()
+if openlineage_config_conn_id:
+config =
AirflowConnectionConfigProvider(openlineage_config_conn_id).get_config()
+self._resolve_airflow_connection_auth(config=config,
config_conn_id=openlineage_config_conn_id)
+return config
+self.log.debug("OpenLineage config_conn_id configuration not found.")
+
+# Second, try to read from YAML file
openlineage_config_path = conf.config_path(check_legacy_env_var=False)
if openlineage_config_path:
-config = self._read_yaml_config(openlineage_config_path)
-return config
+yaml_config = self._read_yaml_config(openlineage_config_path)
+if yaml_config is None:
+return None
+self._resolve_airflow_connection_auth(yaml_config)
+return yaml_config
self.log.debug("OpenLineage config_path configuration not found.")
-# Second, try to get transport config
+# Third, try to get transport config
transport_config = conf.transport()
if not transport_config:
self.log.debug("OpenLineage transport configuration not found.")
return None
-return {"transport": transport_config}
+config = {"transport": transport_config}
+self._resolve_airflow_connection_auth(config)
+return config
@staticmethod
-def _read_yaml_config(path: str) -> dict | None:
+def _read_yaml_config(path: str) -> dict[str, Any] | None:
with open(path) as config_file:
return yaml.safe_load(config_file)
+@staticmethod
+def _resolve_airflow_connection_auth(
+config: dict[str, Any] | None, config_conn_id: str | None = None
+) -> None:
+if not isinstance(config, dict):
+return
+
+for key, value in config.items():
+if isinstance(value, dict) and value.get("type") ==
AIRFLOW_CONNECTION_API_KEY_AUTH_TYPE:
+provider = AirflowConnectionTokenProvider(value,
default_conn_id=config_conn_id)
+config[key] = {"type": "api_key", "apiKey":
provider.get_api_key()}
+else:
+OpenLineageAdapter._resolve_airflow_connection_auth(value,
config_conn_id=config_conn_id)
Review Comment:
Yup, composite transports can also be nested, so this is a good call.
--
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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
mobuchowski commented on code in PR #66342:
URL: https://github.com/apache/airflow/pull/66342#discussion_r3242630804
##
providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py:
##
@@ -103,25 +108,52 @@ def get_or_create_openlineage_client(self) ->
OpenLineageClient:
return self._client
def get_openlineage_config(self) -> dict | None:
-# First, try to read from YAML file
+# First, try to read from Airflow connection
+openlineage_config_conn_id = conf.config_conn_id()
+if openlineage_config_conn_id:
+config =
AirflowConnectionConfigProvider(openlineage_config_conn_id).get_config()
+self._resolve_airflow_connection_auth(config=config,
config_conn_id=openlineage_config_conn_id)
+return config
+self.log.debug("OpenLineage config_conn_id configuration not found.")
+
+# Second, try to read from YAML file
openlineage_config_path = conf.config_path(check_legacy_env_var=False)
if openlineage_config_path:
-config = self._read_yaml_config(openlineage_config_path)
-return config
+yaml_config = self._read_yaml_config(openlineage_config_path)
+if yaml_config is None:
+return None
+self._resolve_airflow_connection_auth(yaml_config)
+return yaml_config
self.log.debug("OpenLineage config_path configuration not found.")
-# Second, try to get transport config
+# Third, try to get transport config
transport_config = conf.transport()
if not transport_config:
self.log.debug("OpenLineage transport configuration not found.")
return None
-return {"transport": transport_config}
+config = {"transport": transport_config}
+self._resolve_airflow_connection_auth(config)
+return config
@staticmethod
-def _read_yaml_config(path: str) -> dict | None:
+def _read_yaml_config(path: str) -> dict[str, Any] | None:
with open(path) as config_file:
return yaml.safe_load(config_file)
+@staticmethod
+def _resolve_airflow_connection_auth(
+config: dict[str, Any] | None, config_conn_id: str | None = None
+) -> None:
+if not isinstance(config, dict):
+return
+
+for key, value in config.items():
+if isinstance(value, dict) and value.get("type") ==
AIRFLOW_CONNECTION_API_KEY_AUTH_TYPE:
+provider = AirflowConnectionTokenProvider(value,
default_conn_id=config_conn_id)
+config[key] = {"type": "api_key", "apiKey":
provider.get_api_key()}
+else:
+OpenLineageAdapter._resolve_airflow_connection_auth(value,
config_conn_id=config_conn_id)
Review Comment:
I think this won't work with composite transport, for example:
```
{"type": "composite", "transports": [{"auth": {"type":
"airflow_connection_api_key", ...}}]})
```
Can we make it work with it, something like
```python
@staticmethod
def _resolve_airflow_connection_auth(
config: dict[str, Any] | None, config_conn_id: str | None = None
) -> None:
if not isinstance(config, dict):
return
for key, value in config.items():
if isinstance(value, dict) and key == "auth" and value.get("type")
== AIRFLOW_CONNECTION_API_KEY_AUTH_TYPE:
provider = AirflowConnectionTokenProvider(value,
default_conn_id=config_conn_id)
config[key] = {"type": "api_key", "apiKey":
provider.get_api_key()}
elif key == "transports" and isinstance(value, list):
for item in value:
OpenLineageAdapter._resolve_airflow_connection_auth(item,
config_conn_id=config_conn_id)
else:
OpenLineageAdapter._resolve_airflow_connection_auth(value,
config_conn_id=config_conn_id)
```
##
providers/openlineage/src/airflow/providers/openlineage/conf.py:
##
@@ -54,6 +54,12 @@ def config_path(check_legacy_env_var: bool = True) -> str:
return option
+@cache
+def config_conn_id() -> str:
+"""[openlineage] config_conn_id."""
+return conf.get(_CONFIG_SECTION, "config_conn_id", fallback="")
Review Comment:
We need to add new configuration options to `provider.yaml` too - so the
docs are properly generated.
##
providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py:
##
@@ -103,25 +108,52 @@ def get_or_create_openlineage_client(self) ->
OpenLineageClient:
return self._client
def get_openlineage_config(self) -> dict | None:
-# First, try to read from YAML file
+# First, try to read from Airflow connection
+openlineage_config_conn_id = conf.config_conn_id()
+if openlineage_config_conn_id:
+config =
AirflowConnectionConfigProvider(openlineage_config_conn_
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
VladaZakharova commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4443035618 hi there! I tried to address your ideas, please check them when possible :) -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
mobuchowski commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4389342166 @VladaZakharova Sorry I was going to reply but totally lost track somewhere. Yeah, IMO it would be great if we could cover the whole transport config. An alternative would be to cover the subset - for example HTTP transport - in a way that would not conflict later, so use the same JSON structure. -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
kacpermuda commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4386661337 > should i wait for @mobuchowski response? or i can go with implementation? Go ahead, I'll be happy to review the PR, then we can ask Maciej for merge as he has the power to do it. -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
VladaZakharova commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4386453986 > I see the value in using Airflow connections for OpenLineage configuration, but I'd suggest expanding the scope beyond just auth tokens for HTTP transport (that are solving a very specific use case). > > Instead of a narrow solution, what if we allow storing the entire OpenLineage config dict (f.e. what can be put in the yaml file, or even just transport config) in an Airflow connection, that we can add to `get_openlineage_config` as yet another config source checked. (It should probably have precedence over yaml file or env vars, but that's something we can decide later) > > This approach would: > > * Work for any transport type, not just HTTP > * Support composite transports (e.g., two HTTP transports with different auth) > * Handle any OL config field, not just auth tokens > * Be more flexible for future use cases > > The current auth-token-only solution would miss users with composite transports or other config needs. Could you expand the scope to cover the full OL config dict? > > cc @mobuchowski hi there! okay, this sounds reasonable, i think it is worth trying should i wait for @mobuchowski response? or i can go with implementation? -- 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]
Re: [PR] Openlineage: read API key auth from Airflow connection [airflow]
kacpermuda commented on PR #66342: URL: https://github.com/apache/airflow/pull/66342#issuecomment-4371300412 I see the value in using Airflow connections for OpenLineage configuration, but I'd suggest expanding the scope beyond just auth tokens for HTTP transport (that are solving a very specific use case). Instead of a narrow solution, what if we allow storing the entire OpenLineage config dict (f.e. what can be put in the yaml file, or even just transport config) in an Airflow connection, that we can add to `get_openlineage_config` as yet another config source checked. (It should probably have precedence over yaml file or env vars, but that's something we can decide later) This approach would: - Work for any transport type, not just HTTP - Support composite transports (e.g., two HTTP transports with different auth) - Handle any OL config field, not just auth tokens - Be more flexible for future use cases The current auth-token-only solution would miss users with composite transports or other config needs. Could you expand the scope to cover the full OL config dict? cc @mobuchowski -- 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]
