You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "amoghrajesh (via GitHub)" <gi...@apache.org> on 2025/02/10 09:22:16 UTC
[PR] AIP-72: Improving Operator Links to Prevent User Code Execution in Webserver [airflow]
amoghrajesh opened a new pull request, #46613:
URL: https://github.com/apache/airflow/pull/46613
<!--
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.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
closes: https://github.com/apache/airflow/issues/46609
## Why?
If you want to add further links to operators you can define them via a plugin or provider package. Extra links will be displayed in task details page in Grid view.
Currently, users define operator links in one of two ways:
1. Using plugins and defining operators in those plugins for defining extra links
2. Directly using custom airflow operators to define operator extra links.
This feature doesn't work with Airflow 3, so while porting this to task sdk so it can be used by users, we decided to simplify the extra links code-flow in a way to also reduce / decouple the user code that runs using `get_link` of a operator extra link in the airflow webserver to outside the context of the webserver.
## What?
To get this working, a simple interface change is made and the following things get easier:
1. While defining an extra link, add a "xcom_key" required field to it.
For example:
```
class FooBarLink(BaseOperatorLink):
name = "foo-bar"
xcom_key = "link-key"
def get_link(self, operator, *, ti_key):
return f"http://www.example.com"
```
2. At runtime, from the task runner, after the task runs (when its `finalize`ing), it requests to push an XCOM which contains the "entire" link for the extra link with the xcom key and task id.
Ex:
```
def finalize(ti: RuntimeTaskInstance, log: Logger):
for oe in ti.task.operator_extra_link_dict.values():
link, xcom_key = oe.get_link(operator=ti.task, ti_key=ti.id), oe.xcom_key # type: ignore[arg-type]
log.debug("Setting xcom for operator extra link", link=link, xcom_key=xcom_key)
_xcom_push(ti, key=xcom_key, value=link)
```
3. After the task has completed running, when the user goes to grid view, he can see the extra button.
The link in this button comes by reading from xcom instead of running user code
## How?
### Changes of note:
#### BaseOperatorLink Model
Introducing a new `xcom_key` datafield that carries the xcom_key with which the entire link will be pushed.
Introduced a new slim interface that we will use during deserialisation which can return the link from xcoms instead of user defined "get_link".
```
@attrs.define()
class GenericOperatorLink(LoggingMixin):
"""A generic operator link class that can retrieve link only using XCOMs. Used while deserializing operators."""
name: str
xcom_key: str
def get_link(self, ti: TaskInstance) -> str:
"""
Retrieve the link from the XComs.
:param ti: Task instance from which to retrieve the link
:return: link to external system, but by pulling it from XComs
"""
self.log.info("Retrieving link from XComs with key: %s for task id: %s", self.xcom_key, ti.task_id)
return ti.xcom_pull(key=self.xcom_key, task_ids=ti.task_id) # type: ignore
```
#### DAG serialisation/ deserialisation
We need something in the serdag to say what links a task has, which we already have but the format is overly complex.
Difference in formats (then vs now):
<meta charset="utf-8"><b style="font-weight:normal;" id="docs-internal-guid-aeba09dd-7fff-6e38-86db-19bab8077e4b"><br /><br /><div dir="ltr" style="margin-left:0pt;" align="left">
| Older | Newer
-- | -- | --
Encoded Operator | '_operator_extra_links': [{'tests.www.views.test_views_extra_links.RaiseErrorLink': {}}, {'tests.www.views.test_views_extra_links.NoResponseLink': {}}, {'tests.www.views.test_views_extra_links.FooBarLink': {}}, {'tests_common.test_utils.mock_operators.AirflowLink': {}}]} | '_operator_extra_links': {'raise_error': 'key', 'no_response': 'key', 'foo-bar': 'link-key', 'airflow': 'airflow_link_key'}}
Decoded Operator | _operator_links_source ={ 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink': { 'index': 0 }},list(_operator_links_source.items()) =[ ( 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink', {'index': 0} )]list(_operator_links_source.items())[0] =( 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink', { 'index': 0 }) | deserialised_operator = SerializedBaseOperator.deserialize_operator(enc_operator)deserialised_operator.operator_extra_linksOut[8]: [GenericOperatorLink(name='raise_error', xcom_key='key'), GenericOperatorLink(name='no_response', xcom_key='key'), GenericOperatorLink(name='foo-bar', xcom_key='link-key'), GenericOperatorLink(name='airflow', xcom_key='airflow_link_key')]
</div><br /></b>
Example of how its stored in the serialised dag now:
```
{"__version": 1, "dag": {"tags": ["extra_links"], "edge_info": {}, "task_group": {"_group_id": null, "group_display_name": "", "prefix_group_id": true, "tooltip": "", "ui_color": "CornflowerBlue", "ui_fgcolor": "#000", "children": {"task": ["operator", "task"]}, "upstream_group_ids": [], "downstream_group_ids": [], "upstream_task_ids": [], "downstream_task_ids": []}, "timetable": {"__type": "airflow.timetables.simple.NullTimetable", "__var": {}}, "fileloc": "/files/dags/dags/extra_links_custom_operator.py", "dag_id": "extra_links_custom_operator", "timezone": "UTC", "catchup": false, "relative_fileloc": "dags/extra_links_custom_operator.py", "_processor_dags_folder": "/files/dags", "tasks": [{"__var": {"pool": "default_pool", "task_type": "DummyTestOperator", "downstream_task_ids": [], "template_fields": [], "on_failure_fail_dagrun": false, "template_fields_renderers": {}, "task_id": "task", "is_teardown": false, "ui_color": "#fff", "weight_rule": "downstream", "ui_fgcolor": "#000
", "start_from_trigger": false, "_needs_expansion": false, "template_ext": [], "is_setup": false, "_task_module": "unusual_prefix_86147154bd7e391eb5735fc97db99fe835bcf5b0_extra_links_custom_operator", "_is_empty": false, "start_trigger_args": null, "_operator_extra_links": {"no_response": "key", "foo-bar": "link-key"}, "label": "task", "_operator_name": "DummyTestOperator"}, "__type": "operator"}], "dag_dependencies": [], "params": []}}
```
- During deserialisation, we create a new GenericOperatorLink type link for each link defined per operator in the serdag
#### Abstract Operator
- The definitions for utils for operator extra links have been moved to task sdk.
- The models/abstractoperator inherits from the task sdk abstarctoperator.
- The `get_extra_links` in task sdk now follows the new function signature
## Testing
1. Operator Links Defined at Operator Level
DAG:
```
from airflow.models import BaseOperatorLink, BaseOperator
from airflow.sdk import DAG
class NoResponseLink(BaseOperatorLink):
name = "no_response"
xcom_key = "key"
def get_link(self, operator, *, ti_key):
return None
class FooBarLink(BaseOperatorLink):
name = "foo-bar"
xcom_key = "link-key"
def get_link(self, operator, *, ti_key):
return f"http://www.example.com"
class DummyTestOperator(BaseOperator):
operator_extra_links = (
NoResponseLink(),
FooBarLink(),
)
def execute(self, context):
print("Hello from custom operator", self.operator_extra_links)
with DAG(
dag_id="extra_links_custom_operator",
schedule=None,
catchup=False,
tags=["extra_links"],
):
t = DummyTestOperator(
task_id="task"
)
```
The link elements:
```
<a target="_blank" class="chakra-link chakra-button c-h1717v" href="http://www.example.com">foo-bar</a>
<a class="chakra-link chakra-button c-h1717v" href="null">no_response</a>
```
XCOMS pushed:
<img width="1471" alt="image" src="https://github.com/user-attachments/assets/bf3bfcc4-485c-473d-922b-e7aa2fe31ad8" />
Logs:
<img width="1471" alt="image" src="https://github.com/user-attachments/assets/dd4b8747-be43-47c7-82e5-9e72d210d55b" />
Last few lines:
```
{"timestamp":"2025-02-10T09:15:06.800071","level":"debug","event":"Sending request","json":"{\"state\":\"success\",\"end_date\":\"2025-02-10T09:15:06.800015Z\",\"task_outlets\":[],\"outlet_events\":[],\"type\":\"SucceedTask\"}\n","logger":"task"}
{"timestamp":"2025-02-10T09:15:06.800145","level":"debug","event":"Plugins are already loaded. Skipping.","logger":"airflow.plugins_manager"}
{"timestamp":"2025-02-10T09:15:06.800175","level":"debug","event":"Initialize extra operators links plugins","logger":"airflow.plugins_manager"}
{"timestamp":"2025-02-10T09:15:06.800263","level":"debug","event":"Setting xcom for operator extra link","link":null,"xcom_key":"key","logger":"task"}
{"timestamp":"2025-02-10T09:15:06.800411","level":"debug","event":"Sending request","json":"{\"key\":\"key\",\"value\":\"null\",\"dag_id\":\"extra_links_custom_operator\",\"run_id\":\"manual__2025-02-10T09:15:05.441050+00:00\",\"task_id\":\"task\",\"map_index\":-1,\"mapped_length\":null,\"type\":\"SetXCom\"}\n","logger":"task"}
{"timestamp":"2025-02-10T09:15:06.800430","level":"debug","event":"Setting xcom for operator extra link","link":"http://www.example.com/","xcom_key":"link-key","logger":"task"}
{"timestamp":"2025-02-10T09:15:06.800457","level":"debug","event":"Sending request","json":"{\"key\":\"link-key\",\"value\":\"\\\"http://www.example.com/\\\"\",\"dag_id\":\"extra_links_custom_operator\",\"run_id\":\"manual__2025-02-10T09:15:05.441050+00:00\",\"task_id\":\"task\",\"map_index\":-1,\"mapped_length\":null,\"type\":\"SetXCom\"}\n","logger":"task"}
```
2. Operator Links Defined as plugins
Plugin:
```
from airflow.plugins_manager import AirflowPlugin
from airflow.models.baseoperatorlink import BaseOperatorLink
from airflow.providers.standard.operators.python import PythonOperator
# define the extra link
class HTTPDocsLink(BaseOperatorLink):
# name the link button in the UI
name = "HTTP docs"
xcom_key = "key"
# add the button to one or more operators
operators = [PythonOperator]
# provide the link
def get_link(self, operator, *, ti_key=None):
return "https://developer.mozilla.org/en-US/docs/Web/HTTP"
# define the plugin class
class AirflowExtraLinkPlugin(AirflowPlugin):
name = "extra_link_plugin"
operator_extra_links = [
HTTPDocsLink(),
]
```
DAG:
```
from airflow.models.dag import DAG
from pendulum import datetime
from airflow.providers.standard.operators.python import PythonOperator
def handler():
print("Hello from the python operator!!")
with DAG(
dag_id="extra_links_plugin",
start_date=datetime(2022, 11, 1),
schedule=None,
catchup=False,
tags=["extra_links"],
):
call_api_simple = PythonOperator(
task_id="call_api_simple",
python_callable=handler,
)
```
<img width="1471" alt="image" src="https://github.com/user-attachments/assets/89bf6e7b-58d9-4801-a8f3-35f7b258e78b" />
The link elements:
```
<a target="_blank" class="chakra-link chakra-button c-h1717v" href="https://developer.mozilla.org/en-US/docs/Web/HTTP">HTTP docs</a>
```
XComs pushed:
<img width="1471" alt="image" src="https://github.com/user-attachments/assets/615f195d-8a63-4d7a-9f80-c6893631f2ba" />
Logs:
<img width="1471" alt="image" src="https://github.com/user-attachments/assets/96f8d4ba-eee6-4ec0-8ed9-adf6efab8b90" />
Last few logs lines:
```
{"timestamp":"2025-02-10T09:19:37.131169","level":"debug","event":"Sending request","json":"{\"state\":\"success\",\"end_date\":\"2025-02-10T09:19:37.131068Z\",\"task_outlets\":[],\"outlet_events\":[],\"type\":\"SucceedTask\"}\n","logger":"task"}
{"timestamp":"2025-02-10T09:19:37.131332","level":"debug","event":"Plugins are already loaded. Skipping.","logger":"airflow.plugins_manager"}
{"timestamp":"2025-02-10T09:19:37.131368","level":"debug","event":"Initialize extra operators links plugins","logger":"airflow.plugins_manager"}
{"timestamp":"2025-02-10T09:19:37.131524","level":"debug","event":"Setting xcom for operator extra link","link":"[https://developer.mozilla.org/en-US/docs/Web/HTTP","xcom_key":"key","logger":"task"](https://developer.mozilla.org/en-US/docs/Web/HTTP%22,%22xcom_key%22:%22key%22,%22logger%22:%22task%22)}
{"timestamp":"2025-02-10T09:19:37.131681","level":"debug","event":"Sending request","json":"{\"key\":\"key\",\"value\":\"\\\"https://developer.mozilla.org/en-US/docs/Web/HTTP\\\"\",\"dag_id\":\"extra_links_plugin\",\"run_id\":\"manual__2025-02-10T09:19:36.380827+00:00\",\"task_id\":\"call_api_simple\",\"map_index\":-1,\"mapped_length\":null,\"type\":\"SetXCom\"}\n","logger":"task"}
{"timestamp":"2025-02-10T09:19:37.877771Z","level":"info","event":"Hello from the python operator!!","chan":"stdout","logger":"task"}
```
## Whats pending?
- [ ] Adding newsfragments (significant / behaviour change?)
- [ ] Testing with SDK baseoperator
- [ ] Moving baseoperatorlink model to sdk definitions
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [newsfragments](https://github.com/apache/airflow/tree/main/newsfragments).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950746647
##########
airflow/models/baseoperatorlink.py:
##########
@@ -20,14 +20,47 @@
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING, ClassVar
-import attr
+import attrs
+
+from airflow.models.xcom import BaseXCom
+from airflow.utils.log.logging_mixin import LoggingMixin
if TYPE_CHECKING:
from airflow.models.baseoperator import BaseOperator
from airflow.models.taskinstancekey import TaskInstanceKey
-@attr.s(auto_attribs=True)
+@attrs.define()
+class GenericOperatorLink(LoggingMixin):
+ """A generic operator link class that can retrieve link only using XCOMs. Used while deserializing operators."""
Review Comment:
Handled here https://github.com/apache/airflow/pull/46613/commits/1852c6219b6997fa7d12160a85f2f8a5d7d0b2a5
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1951040973
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -754,6 +754,12 @@ def finalize(ti: RuntimeTaskInstance, state: TerminalTIState, log: Logger):
)
# TODO: Run task failure callbacks here
+ # Pushing xcom for each operator extra links defined
+ for oe in ti.task.operator_extra_link_dict.values():
+ link, xcom_key = oe.get_link(operator=ti.task, ti_key=ti.id), oe.xcom_key # type: ignore[arg-type]
+ log.debug("Setting xcom for operator extra link", link=link, xcom_key=xcom_key)
+ _xcom_push(ti, key=xcom_key, value=link)
Review Comment:
Done in last commit: a2a0cc559fee08f03fbcdc6b4b5c96b0171b3656
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1948907944
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1520,76 +1523,43 @@ def _deserialize_operator_extra_links(cls, encoded_op_links: list) -> dict[str,
raise AirflowException("Can't load plugins")
op_predefined_extra_links = {}
- for _operator_links_source in encoded_op_links:
- # Get the key, value pair as Tuple where key is OperatorLink ClassName
- # and value is the dictionary containing the arguments passed to the OperatorLink
+ for item in encoded_op_links.items():
+ # Get the name and xcom_key of the encoded operator and use it to create a GenericOperatorLink object
+ # during deserialization.
#
- # Example of a single iteration:
- #
- # _operator_links_source =
- # {
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink': {
- # 'index': 0
- # }
- # },
- #
- # list(_operator_links_source.items()) =
- # [
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {'index': 0}
- # )
- # ]
- #
- # list(_operator_links_source.items())[0] =
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {
- # 'index': 0
- # }
- # )
-
- _operator_link_class_path, data = next(iter(_operator_links_source.items()))
- if _operator_link_class_path in get_operator_extra_links():
- single_op_link_class = import_string(_operator_link_class_path)
- elif _operator_link_class_path in plugins_manager.registered_operator_link_classes:
- single_op_link_class = plugins_manager.registered_operator_link_classes[
- _operator_link_class_path
- ]
- else:
- log.error("Operator Link class %r not registered", _operator_link_class_path)
- return {}
-
- op_link_parameters = {param: cls.deserialize(value) for param, value in data.items()}
- op_predefined_extra_link: BaseOperatorLink = single_op_link_class(**op_link_parameters)
-
+ # Example:
+ # enc_operator['_operator_extra_links'] =
+ # {
+ # 'airflow': 'airflow_link_key',
+ # 'foo-bar': 'link-key',
+ # 'no_response': 'key',
+ # 'raise_error': 'key'
+ # }
+
+ name, xcom_key = item
Review Comment:
```suggestion
```
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1520,76 +1523,43 @@ def _deserialize_operator_extra_links(cls, encoded_op_links: list) -> dict[str,
raise AirflowException("Can't load plugins")
op_predefined_extra_links = {}
- for _operator_links_source in encoded_op_links:
- # Get the key, value pair as Tuple where key is OperatorLink ClassName
- # and value is the dictionary containing the arguments passed to the OperatorLink
+ for item in encoded_op_links.items():
+ # Get the name and xcom_key of the encoded operator and use it to create a GenericOperatorLink object
+ # during deserialization.
#
- # Example of a single iteration:
- #
- # _operator_links_source =
- # {
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink': {
- # 'index': 0
- # }
- # },
- #
- # list(_operator_links_source.items()) =
- # [
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {'index': 0}
- # )
- # ]
- #
- # list(_operator_links_source.items())[0] =
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {
- # 'index': 0
- # }
- # )
-
- _operator_link_class_path, data = next(iter(_operator_links_source.items()))
- if _operator_link_class_path in get_operator_extra_links():
- single_op_link_class = import_string(_operator_link_class_path)
- elif _operator_link_class_path in plugins_manager.registered_operator_link_classes:
- single_op_link_class = plugins_manager.registered_operator_link_classes[
- _operator_link_class_path
- ]
- else:
- log.error("Operator Link class %r not registered", _operator_link_class_path)
- return {}
-
- op_link_parameters = {param: cls.deserialize(value) for param, value in data.items()}
- op_predefined_extra_link: BaseOperatorLink = single_op_link_class(**op_link_parameters)
-
+ # Example:
+ # enc_operator['_operator_extra_links'] =
+ # {
+ # 'airflow': 'airflow_link_key',
+ # 'foo-bar': 'link-key',
+ # 'no_response': 'key',
+ # 'raise_error': 'key'
+ # }
+
+ name, xcom_key = item
+ op_predefined_extra_link = GenericOperatorLink(name=name, xcom_key=xcom_key)
op_predefined_extra_links.update({op_predefined_extra_link.name: op_predefined_extra_link})
Review Comment:
I think this new link class is going to break the API server -- it's going to call `link.get_link(operator=tasl, ti_key=ti.key)` which isn't defined on GenericOperatorLink.
Either we need to make that new GOL class have the same interface (and call `XCom.pull()` directly) or change the class/type hint for extra links. I think the former is better, as that way we can still keep global plugin links working with the same interface
##########
task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py:
##########
@@ -136,6 +140,81 @@ class AbstractOperator(Templater, DAGNode):
)
)
+ def unmap(self, resolve) -> BaseOperator:
+ """
+ Get the "normal" operator from current abstract operator.
+
+ MappedOperator uses this to unmap itself based on the map index. A non-
+ mapped operator (i.e. BaseOperator subclass) simply returns itself.
+
+ :meta private:
+ """
+ raise NotImplementedError()
+
+ @cached_property
+ def global_operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all global extra links."""
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.global_operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ return {link.name: link for link in plugins_manager.global_operator_extra_links}
+
+ @cached_property
+ def extra_links(self) -> list[str]:
+ return sorted(set(self.operator_extra_link_dict).union(self.global_operator_extra_link_dict))
+
+ @cached_property
+ def operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all extra links for the operator."""
+ op_extra_links_from_plugin: dict[str, Any] = {}
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ for ope in plugins_manager.operator_extra_links:
+ if ope.operators and self.operator_class in ope.operators:
+ op_extra_links_from_plugin.update({ope.name: ope})
+
+ operator_extra_links_all = {link.name: link for link in self.operator_extra_links}
+ # Extra links defined in Plugins overrides operator links defined in operator
+ operator_extra_links_all.update(op_extra_links_from_plugin)
+
+ return operator_extra_links_all
+
+ def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None:
+ """
+ For an operator, gets the URLs that the ``extra_links`` entry points to.
+
+ :meta private:
+
+ :raise ValueError: The error message of a ValueError will be passed on through to
+ the fronted to show up as a tooltip on the disabled link.
+ :param ti: The TaskInstance for the URL being searched for.
+ :param link_name: The name of the link we're looking for the URL for. Should be
+ one of the options specified in ``extra_links``.
+ """
+ link = self.operator_extra_link_dict.get(link_name)
+ if not link:
+ link = self.global_operator_extra_link_dict.get(link_name)
+ if not link:
+ return None
+
+ # TODO: replace with isinstance checks when operator links ported to sdk
+ try:
+ el = link.get_link(ti=ti)
+ # Temporary workaround till https://github.com/apache/airflow/issues/46513 is handled.
+ return el.strip('"')
Review Comment:
This doesn't make sense -- we aren't operating on XCom values here -- but the return of the actual OperatorLink class I thought?
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1520,76 +1523,43 @@ def _deserialize_operator_extra_links(cls, encoded_op_links: list) -> dict[str,
raise AirflowException("Can't load plugins")
op_predefined_extra_links = {}
- for _operator_links_source in encoded_op_links:
- # Get the key, value pair as Tuple where key is OperatorLink ClassName
- # and value is the dictionary containing the arguments passed to the OperatorLink
+ for item in encoded_op_links.items():
Review Comment:
Nit of sorts.
```suggestion
for name, xcom_key in encoded_op_links.items():
```
##########
newsfragments/46613.feature.rst:
##########
@@ -0,0 +1 @@
+Operator Links interface changed to not run user code in Airflow Webserver The Operator Extra links, which can be defined either via plugins or custom operators now do not execute any user code in the Airflow Webserver, but instead push the "full" links to XCom backend and the value is again fetched from the XCom backend when viewing task details in grid view.
Review Comment:
Can you give some more info/short example code about what users who have a custom link class written have to do?
##########
task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py:
##########
@@ -136,6 +140,81 @@ class AbstractOperator(Templater, DAGNode):
)
)
+ def unmap(self, resolve) -> BaseOperator:
+ """
+ Get the "normal" operator from current abstract operator.
+
+ MappedOperator uses this to unmap itself based on the map index. A non-
+ mapped operator (i.e. BaseOperator subclass) simply returns itself.
+
+ :meta private:
+ """
+ raise NotImplementedError()
+
+ @cached_property
+ def global_operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all global extra links."""
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.global_operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ return {link.name: link for link in plugins_manager.global_operator_extra_links}
Review Comment:
I wonder if we should keep global operator links defined /loaded from the webserver instead of in the task sdk
(I think yes, that way we don't need to use plugin manager in task SDK.
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1520,76 +1523,43 @@ def _deserialize_operator_extra_links(cls, encoded_op_links: list) -> dict[str,
raise AirflowException("Can't load plugins")
op_predefined_extra_links = {}
- for _operator_links_source in encoded_op_links:
- # Get the key, value pair as Tuple where key is OperatorLink ClassName
- # and value is the dictionary containing the arguments passed to the OperatorLink
+ for item in encoded_op_links.items():
+ # Get the name and xcom_key of the encoded operator and use it to create a GenericOperatorLink object
+ # during deserialization.
#
- # Example of a single iteration:
- #
- # _operator_links_source =
- # {
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink': {
- # 'index': 0
- # }
- # },
- #
- # list(_operator_links_source.items()) =
- # [
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {'index': 0}
- # )
- # ]
- #
- # list(_operator_links_source.items())[0] =
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {
- # 'index': 0
- # }
- # )
-
- _operator_link_class_path, data = next(iter(_operator_links_source.items()))
- if _operator_link_class_path in get_operator_extra_links():
- single_op_link_class = import_string(_operator_link_class_path)
- elif _operator_link_class_path in plugins_manager.registered_operator_link_classes:
- single_op_link_class = plugins_manager.registered_operator_link_classes[
- _operator_link_class_path
- ]
- else:
- log.error("Operator Link class %r not registered", _operator_link_class_path)
- return {}
-
- op_link_parameters = {param: cls.deserialize(value) for param, value in data.items()}
- op_predefined_extra_link: BaseOperatorLink = single_op_link_class(**op_link_parameters)
-
+ # Example:
+ # enc_operator['_operator_extra_links'] =
+ # {
+ # 'airflow': 'airflow_link_key',
+ # 'foo-bar': 'link-key',
+ # 'no_response': 'key',
+ # 'raise_error': 'key'
+ # }
+
+ name, xcom_key = item
+ op_predefined_extra_link = GenericOperatorLink(name=name, xcom_key=xcom_key)
op_predefined_extra_links.update({op_predefined_extra_link.name: op_predefined_extra_link})
return op_predefined_extra_links
@classmethod
- def _serialize_operator_extra_links(cls, operator_extra_links: Iterable[BaseOperatorLink]):
+ def _serialize_operator_extra_links(
+ cls, operator_extra_links: Iterable[BaseOperatorLink]
+ ) -> dict[str, str]:
"""
Serialize Operator Links.
- Store the import path of the OperatorLink and the arguments passed to it.
+ Store the "name" of the link mapped with the xcom_key which can be later used to retrieve this
+ operator extra link from XComs.
For example:
- ``[{'airflow.providers.google.cloud.links.bigquery.BigQueryDatasetLink': {}}]``
+ ``{'link-name-1': 'xcom-key-1'}``
:param operator_extra_links: Operator Link
:return: Serialized Operator Link
"""
- serialize_operator_extra_links = []
+ serialize_operator_extra_links = {}
for operator_extra_link in operator_extra_links:
- op_link_arguments = {
- param: cls.serialize(value) for param, value in attrs.asdict(operator_extra_link).items()
- }
-
- module_path = (
- f"{operator_extra_link.__class__.__module__}.{operator_extra_link.__class__.__name__}"
- )
- serialize_operator_extra_links.append({module_path: op_link_arguments})
-
+ serialize_operator_extra_links[operator_extra_link.name] = operator_extra_link.xcom_key
return serialize_operator_extra_links
Review Comment:
We can simplify this whole function:
```python
return {
link.name: link.xcom_key for link in operator_extra_links
}
```
##########
task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py:
##########
@@ -136,6 +140,81 @@ class AbstractOperator(Templater, DAGNode):
)
)
+ def unmap(self, resolve) -> BaseOperator:
+ """
+ Get the "normal" operator from current abstract operator.
+
+ MappedOperator uses this to unmap itself based on the map index. A non-
+ mapped operator (i.e. BaseOperator subclass) simply returns itself.
+
+ :meta private:
+ """
+ raise NotImplementedError()
+
+ @cached_property
+ def global_operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all global extra links."""
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.global_operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ return {link.name: link for link in plugins_manager.global_operator_extra_links}
+
+ @cached_property
+ def extra_links(self) -> list[str]:
+ return sorted(set(self.operator_extra_link_dict).union(self.global_operator_extra_link_dict))
+
+ @cached_property
+ def operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all extra links for the operator."""
+ op_extra_links_from_plugin: dict[str, Any] = {}
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ for ope in plugins_manager.operator_extra_links:
+ if ope.operators and self.operator_class in ope.operators:
+ op_extra_links_from_plugin.update({ope.name: ope})
+
+ operator_extra_links_all = {link.name: link for link in self.operator_extra_links}
+ # Extra links defined in Plugins overrides operator links defined in operator
+ operator_extra_links_all.update(op_extra_links_from_plugin)
+
+ return operator_extra_links_all
+
+ def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None:
+ """
+ For an operator, gets the URLs that the ``extra_links`` entry points to.
+
+ :meta private:
+
+ :raise ValueError: The error message of a ValueError will be passed on through to
+ the fronted to show up as a tooltip on the disabled link.
+ :param ti: The TaskInstance for the URL being searched for.
+ :param link_name: The name of the link we're looking for the URL for. Should be
+ one of the options specified in ``extra_links``.
+ """
+ link = self.operator_extra_link_dict.get(link_name)
+ if not link:
+ link = self.global_operator_extra_link_dict.get(link_name)
+ if not link:
+ return None
+
+ # TODO: replace with isinstance checks when operator links ported to sdk
+ try:
+ el = link.get_link(ti=ti)
Review Comment:
In the TaskSDK we should never call this form -- the one that takes the `ti` is on `GenericOperatorLink` which is used only in the Webserver, which shouldn't be called from this path -- or to put it another way: we shouldn't have Webserver specific code here in the TaskSDK.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #46613:
URL: https://github.com/apache/airflow/pull/46613#issuecomment-2650654829
Phew yeah that was hard. The providers tests were failing due to this interface change.
Providers needn't have to deal with whats happening at the deser part at all. Removed that portion from providers tests.
(Providers were validating if deser changes the operator extra links class)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1951042478
##########
task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py:
##########
@@ -136,6 +138,39 @@ class AbstractOperator(Templater, DAGNode):
)
)
+ @cached_property
+ def operator_extra_link_dict(self) -> dict[str, Any]:
Review Comment:
Yep it should be. I just ported it over the same way, so missed it.
Handled in latest commit
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #46613:
URL: https://github.com/apache/airflow/pull/46613#issuecomment-2647691470
Should something like this pass:
```
@pytest.mark.usefixtures("clear_all_logger_handlers")
def test_extra_operator_links_logs_error_for_non_registered_extra_links(self):
"""
Assert OperatorLinks not registered via Plugins and if it is not an inbuilt Operator Link,
it can still deserialize the DAG (does not error) but just logs an error.
We test NOT using caplog as this is flaky, we check that the task after deserialize
is missing the extra links.
"""
class TaskStateLink(BaseOperatorLink):
"""OperatorLink not registered via Plugins nor a built-in OperatorLink"""
name = "My Link"
def get_link(self, operator, *, ti_key):
return "https://www.google.com"
class MyOperator(BaseOperator):
"""Just a EmptyOperator using above defined Extra Operator Link"""
operator_extra_links = [TaskStateLink()]
def execute(self, context: Context):
pass
with DAG(dag_id="simple_dag", schedule=None, start_date=datetime(2019, 8, 1)) as dag:
MyOperator(task_id="blah")
serialized_dag = SerializedDAG.to_dict(dag)
sdag = SerializedDAG.from_dict(serialized_dag)
assert sdag.task_dict["blah"].operator_extra_links == []
```
I do not really understand what:
```
Assert OperatorLinks not registered via Plugins and if it is not an inbuilt Operator Link,
it can still deserialize the DAG (does not error) but just logs an error.
We test NOT using caplog as this is flaky, we check that the task after deserialize
is missing the extra links.
```
means.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1949376999
##########
task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py:
##########
@@ -136,6 +140,81 @@ class AbstractOperator(Templater, DAGNode):
)
)
+ def unmap(self, resolve) -> BaseOperator:
+ """
+ Get the "normal" operator from current abstract operator.
+
+ MappedOperator uses this to unmap itself based on the map index. A non-
+ mapped operator (i.e. BaseOperator subclass) simply returns itself.
+
+ :meta private:
+ """
+ raise NotImplementedError()
+
+ @cached_property
+ def global_operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all global extra links."""
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.global_operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ return {link.name: link for link in plugins_manager.global_operator_extra_links}
+
+ @cached_property
+ def extra_links(self) -> list[str]:
+ return sorted(set(self.operator_extra_link_dict).union(self.global_operator_extra_link_dict))
+
+ @cached_property
+ def operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all extra links for the operator."""
+ op_extra_links_from_plugin: dict[str, Any] = {}
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ for ope in plugins_manager.operator_extra_links:
+ if ope.operators and self.operator_class in ope.operators:
+ op_extra_links_from_plugin.update({ope.name: ope})
+
+ operator_extra_links_all = {link.name: link for link in self.operator_extra_links}
+ # Extra links defined in Plugins overrides operator links defined in operator
+ operator_extra_links_all.update(op_extra_links_from_plugin)
+
+ return operator_extra_links_all
+
+ def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None:
+ """
+ For an operator, gets the URLs that the ``extra_links`` entry points to.
+
+ :meta private:
+
+ :raise ValueError: The error message of a ValueError will be passed on through to
+ the fronted to show up as a tooltip on the disabled link.
+ :param ti: The TaskInstance for the URL being searched for.
+ :param link_name: The name of the link we're looking for the URL for. Should be
+ one of the options specified in ``extra_links``.
+ """
+ link = self.operator_extra_link_dict.get(link_name)
+ if not link:
+ link = self.global_operator_extra_link_dict.get(link_name)
+ if not link:
+ return None
+
+ # TODO: replace with isinstance checks when operator links ported to sdk
+ try:
+ el = link.get_link(ti=ti)
+ # Temporary workaround till https://github.com/apache/airflow/issues/46513 is handled.
+ return el.strip('"')
Review Comment:
The return of the actual operator link class is pushed to xcom during task runtime and when its pushed to XCOM, it gets those extra string quotes. For example, this link:
```
class FooBarLink(BaseOperatorLink):
name = "foo-bar"
def get_link(self, operator, *, ti_key):
return f"http://www.example.com"
```
Pushes xcom which looks like in the table:
```
"\"\\\"http://www.example.com\\\"\""
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950518337
##########
airflow/models/baseoperatorlink.py:
##########
@@ -20,14 +20,47 @@
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING, ClassVar
-import attr
+import attrs
+
+from airflow.models.xcom import BaseXCom
+from airflow.utils.log.logging_mixin import LoggingMixin
if TYPE_CHECKING:
from airflow.models.baseoperator import BaseOperator
from airflow.models.taskinstancekey import TaskInstanceKey
-@attr.s(auto_attribs=True)
+@attrs.define()
+class GenericOperatorLink(LoggingMixin):
+ """A generic operator link class that can retrieve link only using XCOMs. Used while deserializing operators."""
Review Comment:
I like `XComOperatorLink` here -> describes the purpose better. I will make this change
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #46613:
URL: https://github.com/apache/airflow/pull/46613#issuecomment-2651130424
Update:
- Migrated over all the functions using plugins manager from sdk/abstractoperator over to the serialisation module.\
- Handled all nits
- Fixed all tests (hopefully)
Whats left:
- Adding examples to the newsfragment (for which I need to extend it to support over 1 line)
- Test an e2e scenario 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1952769974
##########
airflow/models/baseoperatorlink.py:
##########
@@ -20,14 +20,47 @@
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING, ClassVar
-import attr
+import attrs
+
+from airflow.models.xcom import BaseXCom
+from airflow.utils.log.logging_mixin import LoggingMixin
if TYPE_CHECKING:
from airflow.models.baseoperator import BaseOperator
from airflow.models.taskinstancekey import TaskInstanceKey
-@attr.s(auto_attribs=True)
+@attrs.define()
+class GenericOperatorLink(LoggingMixin):
+ """A generic operator link class that can retrieve link only using XCOMs. Used while deserializing operators."""
+
+ name: str
+ xcom_key: str
+
+ def get_link(self, operator: BaseOperator, *, ti_key: TaskInstanceKey) -> str:
+ """
+ Retrieve the link from the XComs.
+
+ :param operator: The Airflow operator object this link is associated to.
+ :param ti_key: TaskInstance ID to return link for.
+ :return: link to external system, but by pulling it from XComs
+ """
+ self.log.info(
+ "Attempting to retrieve link from XComs with key: %s for task id: %s", self.xcom_key, ti_key
+ )
+ value = BaseXCom.get_one(key=self.xcom_key, run_id=ti_key.run_id)
+ if not value:
+ self.log.debug(
+ "No link with name: %s present in XCom as key: %s, returning empty link",
+ self.name,
+ self.xcom_key,
+ )
+ return ""
Review Comment:
Will just create an immediate follow up for this one
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1952810459
##########
newsfragments/46613.feature.rst:
##########
@@ -0,0 +1 @@
+Operator Links interface changed to not run user code in Airflow Webserver The Operator Extra links, which can be defined either via plugins or custom operators now do not execute any user code in the Airflow Webserver, but instead push the "full" links to XCom backend and the value is again fetched from the XCom backend when viewing task details in grid view.
Review Comment:
Tracking via this: https://github.com/apache/airflow/pull/46681
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950764591
##########
airflow/models/baseoperatorlink.py:
##########
@@ -20,14 +20,47 @@
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING, ClassVar
-import attr
+import attrs
+
+from airflow.models.xcom import BaseXCom
+from airflow.utils.log.logging_mixin import LoggingMixin
if TYPE_CHECKING:
from airflow.models.baseoperator import BaseOperator
from airflow.models.taskinstancekey import TaskInstanceKey
-@attr.s(auto_attribs=True)
+@attrs.define()
+class GenericOperatorLink(LoggingMixin):
+ """A generic operator link class that can retrieve link only using XCOMs. Used while deserializing operators."""
+
+ name: str
+ xcom_key: str
+
+ def get_link(self, operator: BaseOperator, *, ti_key: TaskInstanceKey) -> str:
+ """
+ Retrieve the link from the XComs.
+
+ :param operator: The Airflow operator object this link is associated to.
+ :param ti_key: TaskInstance ID to return link for.
+ :return: link to external system, but by pulling it from XComs
+ """
+ self.log.info(
+ "Attempting to retrieve link from XComs with key: %s for task id: %s", self.xcom_key, ti_key
+ )
+ value = BaseXCom.get_one(key=self.xcom_key, run_id=ti_key.run_id)
+ if not value:
+ self.log.debug(
+ "No link with name: %s present in XCom as key: %s, returning empty link",
+ self.name,
+ self.xcom_key,
+ )
+ return ""
Review Comment:
I'd probably say that this specific case should fail/except here, but maybe catch that in the API layer? (so that you can view an old TI even if the link xcom has been deleted by a db clean process?)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950781499
##########
task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py:
##########
@@ -136,6 +138,39 @@ class AbstractOperator(Templater, DAGNode):
)
)
+ @cached_property
+ def operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all extra links for the operator."""
+ op_extra_links_from_plugin: dict[str, Any] = {}
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ for ope in plugins_manager.operator_extra_links:
+ if ope.operators and self.operator_class in ope.operators:
+ op_extra_links_from_plugin.update({ope.name: ope})
Review Comment:
I still think this should be handled in deserialization, not at task run time, but very open to discussion on this.
WDYT?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950411699
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1520,76 +1523,43 @@ def _deserialize_operator_extra_links(cls, encoded_op_links: list) -> dict[str,
raise AirflowException("Can't load plugins")
op_predefined_extra_links = {}
- for _operator_links_source in encoded_op_links:
- # Get the key, value pair as Tuple where key is OperatorLink ClassName
- # and value is the dictionary containing the arguments passed to the OperatorLink
+ for item in encoded_op_links.items():
+ # Get the name and xcom_key of the encoded operator and use it to create a GenericOperatorLink object
+ # during deserialization.
#
- # Example of a single iteration:
- #
- # _operator_links_source =
- # {
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink': {
- # 'index': 0
- # }
- # },
- #
- # list(_operator_links_source.items()) =
- # [
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {'index': 0}
- # )
- # ]
- #
- # list(_operator_links_source.items())[0] =
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {
- # 'index': 0
- # }
- # )
-
- _operator_link_class_path, data = next(iter(_operator_links_source.items()))
- if _operator_link_class_path in get_operator_extra_links():
- single_op_link_class = import_string(_operator_link_class_path)
- elif _operator_link_class_path in plugins_manager.registered_operator_link_classes:
- single_op_link_class = plugins_manager.registered_operator_link_classes[
- _operator_link_class_path
- ]
- else:
- log.error("Operator Link class %r not registered", _operator_link_class_path)
- return {}
-
- op_link_parameters = {param: cls.deserialize(value) for param, value in data.items()}
- op_predefined_extra_link: BaseOperatorLink = single_op_link_class(**op_link_parameters)
-
+ # Example:
+ # enc_operator['_operator_extra_links'] =
+ # {
+ # 'airflow': 'airflow_link_key',
+ # 'foo-bar': 'link-key',
+ # 'no_response': 'key',
+ # 'raise_error': 'key'
+ # }
+
+ name, xcom_key = item
Review Comment:
Done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #46613:
URL: https://github.com/apache/airflow/pull/46613#issuecomment-2653790979
Last few tests to go, almost there!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #46613:
URL: https://github.com/apache/airflow/pull/46613#issuecomment-2656960162
Added some additional information about the design et al in the last self 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950782639
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -754,6 +754,12 @@ def finalize(ti: RuntimeTaskInstance, state: TerminalTIState, log: Logger):
)
# TODO: Run task failure callbacks here
+ # Pushing xcom for each operator extra links defined
+ for oe in ti.task.operator_extra_link_dict.values():
+ link, xcom_key = oe.get_link(operator=ti.task, ti_key=ti.id), oe.xcom_key # type: ignore[arg-type]
+ log.debug("Setting xcom for operator extra link", link=link, xcom_key=xcom_key)
+ _xcom_push(ti, key=xcom_key, value=link)
Review Comment:
Nit: lets move this to before the listeners (they _could_ timeout, and this should be quick and happen before hand)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950780678
##########
task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py:
##########
@@ -136,6 +138,39 @@ class AbstractOperator(Templater, DAGNode):
)
)
+ @cached_property
+ def operator_extra_link_dict(self) -> dict[str, Any]:
Review Comment:
Could/should this be `-> dict[str, BaseOperatorLink]` instead?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950384710
##########
airflow/models/baseoperatorlink.py:
##########
@@ -20,14 +20,47 @@
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING, ClassVar
-import attr
+import attrs
Review Comment:
The new library for attr
##########
airflow/models/baseoperatorlink.py:
##########
@@ -20,14 +20,47 @@
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING, ClassVar
-import attr
+import attrs
+
+from airflow.models.xcom import BaseXCom
+from airflow.utils.log.logging_mixin import LoggingMixin
if TYPE_CHECKING:
from airflow.models.baseoperator import BaseOperator
from airflow.models.taskinstancekey import TaskInstanceKey
-@attr.s(auto_attribs=True)
+@attrs.define()
+class GenericOperatorLink(LoggingMixin):
+ """A generic operator link class that can retrieve link only using XCOMs. Used while deserializing operators."""
+
+ name: str
+ xcom_key: str
+
+ def get_link(self, operator: BaseOperator, *, ti_key: TaskInstanceKey) -> str:
+ """
+ Retrieve the link from the XComs.
+
+ :param operator: The Airflow operator object this link is associated to.
+ :param ti_key: TaskInstance ID to return link for.
+ :return: link to external system, but by pulling it from XComs
+ """
+ self.log.info(
+ "Attempting to retrieve link from XComs with key: %s for task id: %s", self.xcom_key, ti_key
+ )
+ value = BaseXCom.get_one(key=self.xcom_key, run_id=ti_key.run_id)
+ if not value:
+ self.log.debug(
+ "No link with name: %s present in XCom as key: %s, returning empty link",
+ self.name,
+ self.xcom_key,
+ )
+ return ""
Review Comment:
Should we be this liberal here or just fail? cc: @ashb
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1166,6 +1166,25 @@ def __init__(self, *args, **kwargs):
self.template_fields = BaseOperator.template_fields
self.operator_extra_links = BaseOperator.operator_extra_links
+ def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None:
+ """
+ For an operator, gets the URLs that the ``extra_links`` entry points to.
+
+ :meta private:
+
+ :raise ValueError: The error message of a ValueError will be passed on through to
+ the fronted to show up as a tooltip on the disabled link.
+ :param ti: The TaskInstance for the URL being searched for.
+ :param link_name: The name of the link we're looking for the URL for. Should be
+ one of the options specified in ``extra_links``.
+ """
+ link = self.operator_extra_link_dict.get(link_name)
+ if not link:
+ link = self.global_operator_extra_link_dict.get(link_name)
+ if not link:
+ return None
+ return link.get_link(self.unmap(None), ti_key=ti.key)
Review Comment:
Can this stay this way or does it have to stay like the previous definition:
```
link: BaseOperatorLink | None = self.operator_extra_link_dict.get(link_name)
if not link:
link = self.global_operator_extra_link_dict.get(link_name)
if not link:
return None
parameters = inspect.signature(link.get_link).parameters
old_signature = all(name != "ti_key" for name, p in parameters.items() if p.kind != p.VAR_KEYWORD)
if old_signature:
return link.get_link(self.unmap(None), ti.dag_run.logical_date) # type: ignore[misc]
return link.get_link(self.unmap(None), ti_key=ti.key)
```
##########
newsfragments/46613.feature.rst:
##########
@@ -0,0 +1 @@
+Operator Links interface changed to not run user code in Airflow Webserver The Operator Extra links, which can be defined either via plugins or custom operators now do not execute any user code in the Airflow Webserver, but instead push the "full" links to XCom backend and the value is again fetched from the XCom backend when viewing task details in grid view.
Review Comment:
Yeah can do that. Right now, there's a precommit limitation that the feature.rst can only be a single line long. Let me push a fix for that.
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1520,76 +1523,43 @@ def _deserialize_operator_extra_links(cls, encoded_op_links: list) -> dict[str,
raise AirflowException("Can't load plugins")
op_predefined_extra_links = {}
- for _operator_links_source in encoded_op_links:
- # Get the key, value pair as Tuple where key is OperatorLink ClassName
- # and value is the dictionary containing the arguments passed to the OperatorLink
+ for item in encoded_op_links.items():
+ # Get the name and xcom_key of the encoded operator and use it to create a GenericOperatorLink object
+ # during deserialization.
#
- # Example of a single iteration:
- #
- # _operator_links_source =
- # {
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink': {
- # 'index': 0
- # }
- # },
- #
- # list(_operator_links_source.items()) =
- # [
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {'index': 0}
- # )
- # ]
- #
- # list(_operator_links_source.items())[0] =
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {
- # 'index': 0
- # }
- # )
-
- _operator_link_class_path, data = next(iter(_operator_links_source.items()))
- if _operator_link_class_path in get_operator_extra_links():
- single_op_link_class = import_string(_operator_link_class_path)
- elif _operator_link_class_path in plugins_manager.registered_operator_link_classes:
- single_op_link_class = plugins_manager.registered_operator_link_classes[
- _operator_link_class_path
- ]
- else:
- log.error("Operator Link class %r not registered", _operator_link_class_path)
- return {}
-
- op_link_parameters = {param: cls.deserialize(value) for param, value in data.items()}
- op_predefined_extra_link: BaseOperatorLink = single_op_link_class(**op_link_parameters)
-
+ # Example:
+ # enc_operator['_operator_extra_links'] =
+ # {
+ # 'airflow': 'airflow_link_key',
+ # 'foo-bar': 'link-key',
+ # 'no_response': 'key',
+ # 'raise_error': 'key'
+ # }
+
+ name, xcom_key = item
+ op_predefined_extra_link = GenericOperatorLink(name=name, xcom_key=xcom_key)
op_predefined_extra_links.update({op_predefined_extra_link.name: op_predefined_extra_link})
Review Comment:
I have followed the former here, please take a look
##########
airflow/models/abstractoperator.py:
##########
@@ -157,64 +153,6 @@ def priority_weight_total(self) -> int:
)
)
- @cached_property
- def operator_extra_link_dict(self) -> dict[str, Any]:
- """Returns dictionary of all extra links for the operator."""
- op_extra_links_from_plugin: dict[str, Any] = {}
- from airflow import plugins_manager
-
- plugins_manager.initialize_extra_operators_links_plugins()
- if plugins_manager.operator_extra_links is None:
- raise AirflowException("Can't load operators")
- for ope in plugins_manager.operator_extra_links:
- if ope.operators and self.operator_class in ope.operators:
- op_extra_links_from_plugin.update({ope.name: ope})
-
- operator_extra_links_all = {link.name: link for link in self.operator_extra_links}
- # Extra links defined in Plugins overrides operator links defined in operator
- operator_extra_links_all.update(op_extra_links_from_plugin)
-
- return operator_extra_links_all
-
- @cached_property
- def global_operator_extra_link_dict(self) -> dict[str, Any]:
- """Returns dictionary of all global extra links."""
- from airflow import plugins_manager
-
- plugins_manager.initialize_extra_operators_links_plugins()
- if plugins_manager.global_operator_extra_links is None:
- raise AirflowException("Can't load operators")
- return {link.name: link for link in plugins_manager.global_operator_extra_links}
-
- @cached_property
- def extra_links(self) -> list[str]:
- return sorted(set(self.operator_extra_link_dict).union(self.global_operator_extra_link_dict))
-
- def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None:
- """
- For an operator, gets the URLs that the ``extra_links`` entry points to.
-
- :meta private:
-
- :raise ValueError: The error message of a ValueError will be passed on through to
- the fronted to show up as a tooltip on the disabled link.
- :param ti: The TaskInstance for the URL being searched for.
- :param link_name: The name of the link we're looking for the URL for. Should be
- one of the options specified in ``extra_links``.
- """
- link: BaseOperatorLink | None = self.operator_extra_link_dict.get(link_name)
- if not link:
- link = self.global_operator_extra_link_dict.get(link_name)
- if not link:
- return None
-
- parameters = inspect.signature(link.get_link).parameters
- old_signature = all(name != "ti_key" for name, p in parameters.items() if p.kind != p.VAR_KEYWORD)
-
- if old_signature:
- return link.get_link(self.unmap(None), ti.dag_run.logical_date) # type: ignore[misc]
- return link.get_link(self.unmap(None), ti_key=ti.key)
-
Review Comment:
Moved these out into relevant places:
1. `get_extra_links` goes to SerialisedBaseOperator as its only called by the webserver and webserver shouldnt have to interact with abstractoperator
2. Rest of the stuff goes into abstractoperator of SDK, so it can be used within the SDK too and it only makes sense cos this class inherits that one
##########
newsfragments/46613.feature.rst:
##########
@@ -0,0 +1 @@
+Operator Links interface changed to not run user code in Airflow Webserver The Operator Extra links, which can be defined either via plugins or custom operators now do not execute any user code in the Airflow Webserver, but instead push the "full" links to XCom backend and the value is again fetched from the XCom backend when viewing task details in grid view.
Review Comment:
Yeah can do that. Right now, there's a precommit limitation that the feature.rst can only be a single line long. Let me push a fix for that.
##########
task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py:
##########
@@ -136,6 +138,39 @@ class AbstractOperator(Templater, DAGNode):
)
)
+ @cached_property
+ def operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all extra links for the operator."""
+ op_extra_links_from_plugin: dict[str, Any] = {}
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ for ope in plugins_manager.operator_extra_links:
+ if ope.operators and self.operator_class in ope.operators:
+ op_extra_links_from_plugin.update({ope.name: ope})
+
+ operator_extra_links_all = {link.name: link for link in self.operator_extra_links}
+ # Extra links defined in Plugins overrides operator links defined in operator
+ operator_extra_links_all.update(op_extra_links_from_plugin)
+
+ return operator_extra_links_all
+
+ @cached_property
+ def global_operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all global extra links."""
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.global_operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ return {link.name: link for link in plugins_manager.global_operator_extra_links}
+
+ @cached_property
+ def extra_links(self) -> list[str]:
+ return sorted(set(self.operator_extra_link_dict).union(self.global_operator_extra_link_dict))
+
Review Comment:
Moved from models/abstractoperator to here
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1951045116
##########
task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py:
##########
@@ -136,6 +140,81 @@ class AbstractOperator(Templater, DAGNode):
)
)
+ def unmap(self, resolve) -> BaseOperator:
+ """
+ Get the "normal" operator from current abstract operator.
+
+ MappedOperator uses this to unmap itself based on the map index. A non-
+ mapped operator (i.e. BaseOperator subclass) simply returns itself.
+
+ :meta private:
+ """
+ raise NotImplementedError()
+
+ @cached_property
+ def global_operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all global extra links."""
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.global_operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ return {link.name: link for link in plugins_manager.global_operator_extra_links}
+
+ @cached_property
+ def extra_links(self) -> list[str]:
+ return sorted(set(self.operator_extra_link_dict).union(self.global_operator_extra_link_dict))
+
+ @cached_property
+ def operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all extra links for the operator."""
+ op_extra_links_from_plugin: dict[str, Any] = {}
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ for ope in plugins_manager.operator_extra_links:
+ if ope.operators and self.operator_class in ope.operators:
+ op_extra_links_from_plugin.update({ope.name: ope})
+
+ operator_extra_links_all = {link.name: link for link in self.operator_extra_links}
+ # Extra links defined in Plugins overrides operator links defined in operator
+ operator_extra_links_all.update(op_extra_links_from_plugin)
+
+ return operator_extra_links_all
+
+ def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None:
+ """
+ For an operator, gets the URLs that the ``extra_links`` entry points to.
+
+ :meta private:
+
+ :raise ValueError: The error message of a ValueError will be passed on through to
+ the fronted to show up as a tooltip on the disabled link.
+ :param ti: The TaskInstance for the URL being searched for.
+ :param link_name: The name of the link we're looking for the URL for. Should be
+ one of the options specified in ``extra_links``.
+ """
+ link = self.operator_extra_link_dict.get(link_name)
+ if not link:
+ link = self.global_operator_extra_link_dict.get(link_name)
+ if not link:
+ return None
+
+ # TODO: replace with isinstance checks when operator links ported to sdk
+ try:
+ el = link.get_link(ti=ti)
Review Comment:
Outdated comment, it has been better handled, resolving
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh merged PR #46613:
URL: https://github.com/apache/airflow/pull/46613
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1948767468
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -717,7 +717,11 @@ def _push_xcom_if_needed(result: Any, ti: RuntimeTaskInstance):
_xcom_push(ti, "return_value", result, mapped_length=mapped_length)
-def finalize(log: Logger): ...
+def finalize(ti: RuntimeTaskInstance, log: Logger):
+ for oe in ti.task.operator_extra_link_dict.values():
+ link, xcom_key = oe.get_link(operator=ti.task, ti_key=ti.id), oe.xcom_key # type: ignore[arg-type]
+ log.debug("Setting xcom for operator extra link", link=link, xcom_key=xcom_key)
+ _xcom_push(ti, key=xcom_key, value=link)
Review Comment:
Question on previous behaviuor? Where links available for failed tasks? I guess they were (in which case this is the right place 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1949370470
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1520,76 +1523,43 @@ def _deserialize_operator_extra_links(cls, encoded_op_links: list) -> dict[str,
raise AirflowException("Can't load plugins")
op_predefined_extra_links = {}
- for _operator_links_source in encoded_op_links:
- # Get the key, value pair as Tuple where key is OperatorLink ClassName
- # and value is the dictionary containing the arguments passed to the OperatorLink
+ for item in encoded_op_links.items():
+ # Get the name and xcom_key of the encoded operator and use it to create a GenericOperatorLink object
+ # during deserialization.
#
- # Example of a single iteration:
- #
- # _operator_links_source =
- # {
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink': {
- # 'index': 0
- # }
- # },
- #
- # list(_operator_links_source.items()) =
- # [
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {'index': 0}
- # )
- # ]
- #
- # list(_operator_links_source.items())[0] =
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {
- # 'index': 0
- # }
- # )
-
- _operator_link_class_path, data = next(iter(_operator_links_source.items()))
- if _operator_link_class_path in get_operator_extra_links():
- single_op_link_class = import_string(_operator_link_class_path)
- elif _operator_link_class_path in plugins_manager.registered_operator_link_classes:
- single_op_link_class = plugins_manager.registered_operator_link_classes[
- _operator_link_class_path
- ]
- else:
- log.error("Operator Link class %r not registered", _operator_link_class_path)
- return {}
-
- op_link_parameters = {param: cls.deserialize(value) for param, value in data.items()}
- op_predefined_extra_link: BaseOperatorLink = single_op_link_class(**op_link_parameters)
-
+ # Example:
+ # enc_operator['_operator_extra_links'] =
+ # {
+ # 'airflow': 'airflow_link_key',
+ # 'foo-bar': 'link-key',
+ # 'no_response': 'key',
+ # 'raise_error': 'key'
+ # }
+
+ name, xcom_key = item
+ op_predefined_extra_link = GenericOperatorLink(name=name, xcom_key=xcom_key)
op_predefined_extra_links.update({op_predefined_extra_link.name: op_predefined_extra_link})
return op_predefined_extra_links
@classmethod
- def _serialize_operator_extra_links(cls, operator_extra_links: Iterable[BaseOperatorLink]):
+ def _serialize_operator_extra_links(
+ cls, operator_extra_links: Iterable[BaseOperatorLink]
+ ) -> dict[str, str]:
"""
Serialize Operator Links.
- Store the import path of the OperatorLink and the arguments passed to it.
+ Store the "name" of the link mapped with the xcom_key which can be later used to retrieve this
+ operator extra link from XComs.
For example:
- ``[{'airflow.providers.google.cloud.links.bigquery.BigQueryDatasetLink': {}}]``
+ ``{'link-name-1': 'xcom-key-1'}``
:param operator_extra_links: Operator Link
:return: Serialized Operator Link
"""
- serialize_operator_extra_links = []
+ serialize_operator_extra_links = {}
for operator_extra_link in operator_extra_links:
- op_link_arguments = {
- param: cls.serialize(value) for param, value in attrs.asdict(operator_extra_link).items()
- }
-
- module_path = (
- f"{operator_extra_link.__class__.__module__}.{operator_extra_link.__class__.__name__}"
- )
- serialize_operator_extra_links.append({module_path: op_link_arguments})
-
+ serialize_operator_extra_links[operator_extra_link.name] = operator_extra_link.xcom_key
return serialize_operator_extra_links
Review Comment:
Cool
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1520,76 +1523,43 @@ def _deserialize_operator_extra_links(cls, encoded_op_links: list) -> dict[str,
raise AirflowException("Can't load plugins")
op_predefined_extra_links = {}
- for _operator_links_source in encoded_op_links:
- # Get the key, value pair as Tuple where key is OperatorLink ClassName
- # and value is the dictionary containing the arguments passed to the OperatorLink
+ for item in encoded_op_links.items():
+ # Get the name and xcom_key of the encoded operator and use it to create a GenericOperatorLink object
+ # during deserialization.
#
- # Example of a single iteration:
- #
- # _operator_links_source =
- # {
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink': {
- # 'index': 0
- # }
- # },
- #
- # list(_operator_links_source.items()) =
- # [
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {'index': 0}
- # )
- # ]
- #
- # list(_operator_links_source.items())[0] =
- # (
- # 'airflow.providers.google.cloud.operators.bigquery.BigQueryConsoleIndexableLink',
- # {
- # 'index': 0
- # }
- # )
-
- _operator_link_class_path, data = next(iter(_operator_links_source.items()))
- if _operator_link_class_path in get_operator_extra_links():
- single_op_link_class = import_string(_operator_link_class_path)
- elif _operator_link_class_path in plugins_manager.registered_operator_link_classes:
- single_op_link_class = plugins_manager.registered_operator_link_classes[
- _operator_link_class_path
- ]
- else:
- log.error("Operator Link class %r not registered", _operator_link_class_path)
- return {}
-
- op_link_parameters = {param: cls.deserialize(value) for param, value in data.items()}
- op_predefined_extra_link: BaseOperatorLink = single_op_link_class(**op_link_parameters)
-
+ # Example:
+ # enc_operator['_operator_extra_links'] =
+ # {
+ # 'airflow': 'airflow_link_key',
+ # 'foo-bar': 'link-key',
+ # 'no_response': 'key',
+ # 'raise_error': 'key'
+ # }
+
+ name, xcom_key = item
+ op_predefined_extra_link = GenericOperatorLink(name=name, xcom_key=xcom_key)
op_predefined_extra_links.update({op_predefined_extra_link.name: op_predefined_extra_link})
return op_predefined_extra_links
@classmethod
- def _serialize_operator_extra_links(cls, operator_extra_links: Iterable[BaseOperatorLink]):
+ def _serialize_operator_extra_links(
+ cls, operator_extra_links: Iterable[BaseOperatorLink]
+ ) -> dict[str, str]:
"""
Serialize Operator Links.
- Store the import path of the OperatorLink and the arguments passed to it.
+ Store the "name" of the link mapped with the xcom_key which can be later used to retrieve this
+ operator extra link from XComs.
For example:
- ``[{'airflow.providers.google.cloud.links.bigquery.BigQueryDatasetLink': {}}]``
+ ``{'link-name-1': 'xcom-key-1'}``
:param operator_extra_links: Operator Link
:return: Serialized Operator Link
"""
- serialize_operator_extra_links = []
+ serialize_operator_extra_links = {}
for operator_extra_link in operator_extra_links:
- op_link_arguments = {
- param: cls.serialize(value) for param, value in attrs.asdict(operator_extra_link).items()
- }
-
- module_path = (
- f"{operator_extra_link.__class__.__module__}.{operator_extra_link.__class__.__name__}"
- )
- serialize_operator_extra_links.append({module_path: op_link_arguments})
-
+ serialize_operator_extra_links[operator_extra_link.name] = operator_extra_link.xcom_key
return serialize_operator_extra_links
Review Comment:
Cool, handled
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950510196
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1166,6 +1166,25 @@ def __init__(self, *args, **kwargs):
self.template_fields = BaseOperator.template_fields
self.operator_extra_links = BaseOperator.operator_extra_links
+ def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None:
+ """
+ For an operator, gets the URLs that the ``extra_links`` entry points to.
+
+ :meta private:
+
+ :raise ValueError: The error message of a ValueError will be passed on through to
+ the fronted to show up as a tooltip on the disabled link.
+ :param ti: The TaskInstance for the URL being searched for.
+ :param link_name: The name of the link we're looking for the URL for. Should be
+ one of the options specified in ``extra_links``.
+ """
+ link = self.operator_extra_link_dict.get(link_name)
+ if not link:
+ link = self.global_operator_extra_link_dict.get(link_name)
+ if not link:
+ return None
+ return link.get_link(self.unmap(None), ti_key=ti.key)
Review Comment:
Thanks a lot @uranusjr!! I was going to propose we remove the old signature anyways, thats helpful
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950498632
##########
airflow/models/baseoperatorlink.py:
##########
@@ -20,14 +20,47 @@
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING, ClassVar
-import attr
+import attrs
+
+from airflow.models.xcom import BaseXCom
+from airflow.utils.log.logging_mixin import LoggingMixin
if TYPE_CHECKING:
from airflow.models.baseoperator import BaseOperator
from airflow.models.taskinstancekey import TaskInstanceKey
-@attr.s(auto_attribs=True)
+@attrs.define()
+class GenericOperatorLink(LoggingMixin):
+ """A generic operator link class that can retrieve link only using XCOMs. Used while deserializing operators."""
Review Comment:
What does “generic” mean here?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #46613:
URL: https://github.com/apache/airflow/pull/46613#issuecomment-2653959565
Yeah finally!! 🥇
@ashb could you take a look again at this one?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1951029722
##########
task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py:
##########
@@ -136,6 +140,81 @@ class AbstractOperator(Templater, DAGNode):
)
)
+ def unmap(self, resolve) -> BaseOperator:
+ """
+ Get the "normal" operator from current abstract operator.
+
+ MappedOperator uses this to unmap itself based on the map index. A non-
+ mapped operator (i.e. BaseOperator subclass) simply returns itself.
+
+ :meta private:
+ """
+ raise NotImplementedError()
+
+ @cached_property
+ def global_operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all global extra links."""
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.global_operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ return {link.name: link for link in plugins_manager.global_operator_extra_links}
Review Comment:
Not a bad idea at all. Let me track this in the latest related comment
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #46613:
URL: https://github.com/apache/airflow/pull/46613#issuecomment-2652866890
A lot of core tests are failing, working on fixing them
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1949369352
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1520,76 +1523,43 @@ def _deserialize_operator_extra_links(cls, encoded_op_links: list) -> dict[str,
raise AirflowException("Can't load plugins")
op_predefined_extra_links = {}
- for _operator_links_source in encoded_op_links:
- # Get the key, value pair as Tuple where key is OperatorLink ClassName
- # and value is the dictionary containing the arguments passed to the OperatorLink
+ for item in encoded_op_links.items():
Review Comment:
Done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950511410
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1166,6 +1166,25 @@ def __init__(self, *args, **kwargs):
self.template_fields = BaseOperator.template_fields
self.operator_extra_links = BaseOperator.operator_extra_links
+ def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None:
+ """
+ For an operator, gets the URLs that the ``extra_links`` entry points to.
+
+ :meta private:
+
+ :raise ValueError: The error message of a ValueError will be passed on through to
+ the fronted to show up as a tooltip on the disabled link.
+ :param ti: The TaskInstance for the URL being searched for.
+ :param link_name: The name of the link we're looking for the URL for. Should be
+ one of the options specified in ``extra_links``.
+ """
+ link = self.operator_extra_link_dict.get(link_name)
+ if not link:
+ link = self.global_operator_extra_link_dict.get(link_name)
+ if not link:
+ return None
+ return link.get_link(self.unmap(None), ti_key=ti.key)
Review Comment:
Resolving based on your comment, irrelevant now
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950516013
##########
airflow/models/baseoperatorlink.py:
##########
@@ -20,14 +20,47 @@
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING, ClassVar
-import attr
+import attrs
+
+from airflow.models.xcom import BaseXCom
+from airflow.utils.log.logging_mixin import LoggingMixin
if TYPE_CHECKING:
from airflow.models.baseoperator import BaseOperator
from airflow.models.taskinstancekey import TaskInstanceKey
-@attr.s(auto_attribs=True)
+@attrs.define()
+class GenericOperatorLink(LoggingMixin):
+ """A generic operator link class that can retrieve link only using XCOMs. Used while deserializing operators."""
Review Comment:
Hmm that naming makes me wonder now.
Its more of a slim interface that the webserver APIs will call to fetch operator extra links, but the implementation is such that it will only happen through an xcom pull..
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950501227
##########
airflow/serialization/serialized_objects.py:
##########
@@ -1166,6 +1166,25 @@ def __init__(self, *args, **kwargs):
self.template_fields = BaseOperator.template_fields
self.operator_extra_links = BaseOperator.operator_extra_links
+ def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None:
+ """
+ For an operator, gets the URLs that the ``extra_links`` entry points to.
+
+ :meta private:
+
+ :raise ValueError: The error message of a ValueError will be passed on through to
+ the fronted to show up as a tooltip on the disabled link.
+ :param ti: The TaskInstance for the URL being searched for.
+ :param link_name: The name of the link we're looking for the URL for. Should be
+ one of the options specified in ``extra_links``.
+ """
+ link = self.operator_extra_link_dict.get(link_name)
+ if not link:
+ link = self.global_operator_extra_link_dict.get(link_name)
+ if not link:
+ return None
+ return link.get_link(self.unmap(None), ti_key=ti.key)
Review Comment:
See #46415
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950778051
##########
providers/amazon/tests/provider_tests/amazon/aws/links/test_base_aws.py:
##########
@@ -194,14 +194,10 @@ def assert_extra_link_url(
)
error_msg = f"{self.full_qualname!r} should be preserved after execution"
- assert ti.task.get_extra_links(ti, self.link_class.name) == expected_url, error_msg
- serialized_dag = self.dag_maker.get_serialized_data()
- deserialized_dag = SerializedDAG.from_dict(serialized_dag)
- deserialized_task = deserialized_dag.task_dict[self.task_id]
-
- error_msg = f"{self.full_qualname!r} should be preserved in deserialized tasks after execution"
- assert deserialized_task.get_extra_links(ti, self.link_class.name) == expected_url, error_msg
+ # This is the right thing to test in the provider as providers shouldn't influence ti.task
+ # This is also the source of the link, if all's fine here, the link will be pushed to xcom and fetched too
Review Comment:
```suggestion
```
This is correct, but to me doesn't serve value as a long term code comment.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1951032794
##########
task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py:
##########
@@ -136,6 +138,39 @@ class AbstractOperator(Templater, DAGNode):
)
)
+ @cached_property
+ def operator_extra_link_dict(self) -> dict[str, Any]:
+ """Returns dictionary of all extra links for the operator."""
+ op_extra_links_from_plugin: dict[str, Any] = {}
+ from airflow import plugins_manager
+
+ plugins_manager.initialize_extra_operators_links_plugins()
+ if plugins_manager.operator_extra_links is None:
+ raise AirflowException("Can't load operators")
+ for ope in plugins_manager.operator_extra_links:
+ if ope.operators and self.operator_class in ope.operators:
+ op_extra_links_from_plugin.update({ope.name: ope})
Review Comment:
We should yes, that will reduce the unwanted dependency between plugins manager and the task sdk. We should keep it a minimal slim package anyways.
Its not just global_operator_extra_link_dict, there's also operator_extra_link_dict and extra_links which I will move to their right places in the next commit
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1950498632
##########
airflow/models/baseoperatorlink.py:
##########
@@ -20,14 +20,47 @@
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING, ClassVar
-import attr
+import attrs
+
+from airflow.models.xcom import BaseXCom
+from airflow.utils.log.logging_mixin import LoggingMixin
if TYPE_CHECKING:
from airflow.models.baseoperator import BaseOperator
from airflow.models.taskinstancekey import TaskInstanceKey
-@attr.s(auto_attribs=True)
+@attrs.define()
+class GenericOperatorLink(LoggingMixin):
+ """A generic operator link class that can retrieve link only using XCOMs. Used while deserializing operators."""
Review Comment:
What does “generic” mean here? It only uses XCom and can be used by any operator?
Maybe something like XComOperatorLink would be more descriptive?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1951040451
##########
providers/amazon/tests/provider_tests/amazon/aws/links/test_base_aws.py:
##########
@@ -194,14 +194,10 @@ def assert_extra_link_url(
)
error_msg = f"{self.full_qualname!r} should be preserved after execution"
- assert ti.task.get_extra_links(ti, self.link_class.name) == expected_url, error_msg
- serialized_dag = self.dag_maker.get_serialized_data()
- deserialized_dag = SerializedDAG.from_dict(serialized_dag)
- deserialized_task = deserialized_dag.task_dict[self.task_id]
-
- error_msg = f"{self.full_qualname!r} should be preserved in deserialized tasks after execution"
- assert deserialized_task.get_extra_links(ti, self.link_class.name) == expected_url, error_msg
+ # This is the right thing to test in the provider as providers shouldn't influence ti.task
+ # This is also the source of the link, if all's fine here, the link will be pushed to xcom and fetched too
Review Comment:
Yeah sure removing
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1952776243
##########
newsfragments/46613.feature.rst:
##########
@@ -0,0 +1 @@
+Operator Links interface changed to not run user code in Airflow Webserver The Operator Extra links, which can be defined either via plugins or custom operators now do not execute any user code in the Airflow Webserver, but instead push the "full" links to XCom backend and the value is again fetched from the XCom backend when viewing task details in grid view.
Review Comment:
I will take this one in a follow up, creating a PR to allow more than 1 line in feature newsfragments.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1948810412
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -717,7 +717,11 @@ def _push_xcom_if_needed(result: Any, ti: RuntimeTaskInstance):
_xcom_push(ti, "return_value", result, mapped_length=mapped_length)
-def finalize(log: Logger): ...
+def finalize(ti: RuntimeTaskInstance, log: Logger):
+ for oe in ti.task.operator_extra_link_dict.values():
+ link, xcom_key = oe.get_link(operator=ti.task, ti_key=ti.id), oe.xcom_key # type: ignore[arg-type]
+ log.debug("Setting xcom for operator extra link", link=link, xcom_key=xcom_key)
+ _xcom_push(ti, key=xcom_key, value=link)
Review Comment:
Precisely, thats the behaviour. And also the reason why I kept this in finalize.
There's a use case, a link to "how to write a dag effectively" if a dag times out due to being poorly written :)
##########
task_sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -717,7 +717,11 @@ def _push_xcom_if_needed(result: Any, ti: RuntimeTaskInstance):
_xcom_push(ti, "return_value", result, mapped_length=mapped_length)
-def finalize(log: Logger): ...
+def finalize(ti: RuntimeTaskInstance, log: Logger):
+ for oe in ti.task.operator_extra_link_dict.values():
+ link, xcom_key = oe.get_link(operator=ti.task, ti_key=ti.id), oe.xcom_key # type: ignore[arg-type]
+ log.debug("Setting xcom for operator extra link", link=link, xcom_key=xcom_key)
+ _xcom_push(ti, key=xcom_key, value=link)
Review Comment:
Precisely, thats the behaviour. And also the reason why I kept this in finalize.
Here's a use case, a link to "how to write a dag effectively" if a dag times out due to being poorly written :)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-72: Improving Operator Links Interface to Prevent User Code Execution in Webserver [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #46613:
URL: https://github.com/apache/airflow/pull/46613#discussion_r1954725637
##########
tests_common/pytest_plugin.py:
##########
@@ -1606,6 +1607,18 @@ def _disable_redact(request: pytest.FixtureRequest, mocker):
return
+@pytest.fixture(autouse=True)
+def _mock_plugins(request: pytest.FixtureRequest):
+ """Disable redacted text in tests, except specific."""
+ if mark := next(request.node.iter_markers("mock_plugin_manager"), None):
+ from tests_common.test_utils.mock_plugins import mock_plugin_manager
+
+ with mock_plugin_manager(**mark.kwargs):
+ yield
+ return
+ yield
Review Comment:
This is important to introduce because I was running into a scenario where the `mock_plugin_manager` as a wrapper was running before the dag was serialising, leading to some issues like the test plugins loading when I didnt wish for them too.
##########
tests/api_fastapi/core_api/routes/public/test_extra_links.py:
##########
@@ -43,7 +39,33 @@
pytestmark = pytest.mark.db_test
-class TestExtraLinks:
+class GoogleLink(BaseOperatorLink):
+ name = "Google"
+
+ def get_link(self, operator, ti_key):
+ return "https://www.google.com"
Review Comment:
New style, not date time.
##########
tests/api_fastapi/core_api/routes/public/test_extra_links.py:
##########
@@ -158,15 +181,33 @@ def test_should_respond_200_missing_xcom(self, test_client):
assert response.status_code == 200
assert response.json() == {"Google Custom": None}
- @mock_plugin_manager(plugins=[])
- def test_should_respond_200_multiple_links(self, test_client):
+ def test_should_respond_200_multiple_links(self, test_client, session):
XCom.set(
key="search_query",
value=["TEST_LINK_VALUE_1", "TEST_LINK_VALUE_2"],
task_id=self.task_multiple_links,
dag_id=self.dag.dag_id,
run_id=self.dag_run_id,
+ session=session,
+ )
+ XCom.set(
+ key="bigquery_1",
+ value="https://console.cloud.google.com/bigquery?j=TEST_LINK_VALUE_1",
+ task_id=self.task_multiple_links,
+ dag_id=self.dag_id,
+ run_id=self.dag_run_id,
+ session=session,
+ )
+ XCom.set(
+ key="bigquery_2",
+ value="https://console.cloud.google.com/bigquery?j=TEST_LINK_VALUE_2",
+ task_id=self.task_multiple_links,
+ dag_id=self.dag_id,
+ run_id=self.dag_run_id,
+ session=session,
Review Comment:
Actually setting these, so it can be fetched from xcom instead of mocking.
##########
tests/api_fastapi/core_api/routes/public/test_extra_links.py:
##########
@@ -186,48 +226,19 @@ def test_should_respond_200_multiple_links_missing_xcom(self, test_client):
assert response.status_code == 200
assert response.json() == {"BigQuery Console #1": None, "BigQuery Console #2": None}
+ @pytest.mark.mock_plugin_manager(plugins=[AirflowPluginWithOperatorLinks])
Review Comment:
This fixture will override global one
##########
tests/api_fastapi/core_api/routes/public/test_extra_links.py:
##########
@@ -43,7 +39,33 @@
pytestmark = pytest.mark.db_test
-class TestExtraLinks:
+class GoogleLink(BaseOperatorLink):
+ name = "Google"
+
+ def get_link(self, operator, ti_key):
+ return "https://www.google.com"
+
+
+class S3LogLink(BaseOperatorLink):
+ name = "S3"
+ operators = [CustomOperator]
+
+ def get_link(self, operator, ti_key):
+ return f"https://s3.amazonaws.com/airflow-logs/{operator.dag_id}/{operator.task_id}/"
+
+
+class AirflowPluginWithOperatorLinks(AirflowPlugin):
+ name = "test_plugin"
+ global_operator_extra_links = [
+ GoogleLink(),
+ ]
+ operator_extra_links = [
+ S3LogLink(),
+ ]
+
+
+@pytest.mark.mock_plugin_manager(plugins=[])
Review Comment:
This is one such example of that scenario mentioned above.
##########
airflow/models/abstractoperator.py:
##########
@@ -157,64 +153,6 @@ def priority_weight_total(self) -> int:
)
)
- @cached_property
- def operator_extra_link_dict(self) -> dict[str, Any]:
- """Returns dictionary of all extra links for the operator."""
- op_extra_links_from_plugin: dict[str, Any] = {}
- from airflow import plugins_manager
-
- plugins_manager.initialize_extra_operators_links_plugins()
- if plugins_manager.operator_extra_links is None:
- raise AirflowException("Can't load operators")
- for ope in plugins_manager.operator_extra_links:
- if ope.operators and self.operator_class in ope.operators:
- op_extra_links_from_plugin.update({ope.name: ope})
-
- operator_extra_links_all = {link.name: link for link in self.operator_extra_links}
- # Extra links defined in Plugins overrides operator links defined in operator
- operator_extra_links_all.update(op_extra_links_from_plugin)
-
- return operator_extra_links_all
-
- @cached_property
- def global_operator_extra_link_dict(self) -> dict[str, Any]:
- """Returns dictionary of all global extra links."""
- from airflow import plugins_manager
-
- plugins_manager.initialize_extra_operators_links_plugins()
- if plugins_manager.global_operator_extra_links is None:
- raise AirflowException("Can't load operators")
- return {link.name: link for link in plugins_manager.global_operator_extra_links}
-
- @cached_property
- def extra_links(self) -> list[str]:
- return sorted(set(self.operator_extra_link_dict).union(self.global_operator_extra_link_dict))
-
- def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None:
- """
- For an operator, gets the URLs that the ``extra_links`` entry points to.
-
- :meta private:
-
- :raise ValueError: The error message of a ValueError will be passed on through to
- the fronted to show up as a tooltip on the disabled link.
- :param ti: The TaskInstance for the URL being searched for.
- :param link_name: The name of the link we're looking for the URL for. Should be
- one of the options specified in ``extra_links``.
- """
- link: BaseOperatorLink | None = self.operator_extra_link_dict.get(link_name)
- if not link:
- link = self.global_operator_extra_link_dict.get(link_name)
- if not link:
- return None
-
- parameters = inspect.signature(link.get_link).parameters
- old_signature = all(name != "ti_key" for name, p in parameters.items() if p.kind != p.VAR_KEYWORD)
-
- if old_signature:
- return link.get_link(self.unmap(None), ti.dag_run.logical_date) # type: ignore[misc]
- return link.get_link(self.unmap(None), ti_key=ti.key)
Review Comment:
Webserver should not be dealing with real abstarctoperator objects. It does and should deal with SerialisedBaseOperator objects. Hence moving this there and nuking this one!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org