kaxil commented on a change in pull request #8651: URL: https://github.com/apache/airflow/pull/8651#discussion_r421157338
########## File path: airflow/task/context/current.py ########## @@ -0,0 +1,64 @@ +# 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. + +import contextlib +import logging +from typing import Any, Dict + +from airflow.exceptions import AirflowException + +_CURRENT_CONTEXT = [] +log = logging.getLogger(__name__) + + [email protected] +def set_current_context(context): Review comment: Can we add TypeHint for `context` ########## File path: airflow/config_templates/config.yml ########## @@ -366,6 +366,17 @@ type: string example: "path.to.CustomXCom" default: "airflow.models.xcom.BaseXCom" + - name: additional_execute_contextmanager + description: | + Custom user function that returns a context manager. Syntax is "package.method". + Context is entered when operator starts executing task. __enter__() will be called + before the operator's execute method, and __exit__() shortly after. + Function's signature should accept two positional parameters - task instance + and execution context Review comment: ```suggestion Context is entered when operator starts executing task. ``__enter__()`` will be called before the operator's execute method, and ``__exit__()`` shortly after. Function's signature should accept two positional parameters - task instance and execution context ``` I think you can add what it is equivalent to, to make it clear: ```python @my_context_manager def execute(context): ... ``` ########## File path: airflow/models/taskinstance.py ########## @@ -1110,6 +1116,26 @@ def signal_handler(signum, frame): session.merge(self) session.commit() + def get_additional_execution_contextmanager(self, execution_context): + """ + Retrieves the user defined execution context callback from the configuration, + and validates that it is indeed a context manager + :param execution_context: the current execution context to be passed to user ctx Review comment: A blank line is needed for correct Sphinx doc rendering ```suggestion and validates that it is indeed a context manager :param execution_context: the current execution context to be passed to user ctx ``` ########## File path: docs/howto/use-additional-execute-contextmanager.rst ########## @@ -0,0 +1,47 @@ + .. 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. + + + +Defining Additional Execute Context Manager +=============================== Review comment: ```suggestion Defining Additional Execute Context Manager =========================================== ``` ########## File path: docs/howto/use-additional-execute-contextmanager.rst ########## @@ -0,0 +1,47 @@ + .. 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. + + + +Defining Additional Execute Context Manager +=============================== + +Creating new context manager +---------------------------- + +Users can create their own exeuction context manager to allow context management on a higher level. Review comment: ```suggestion Users can create their own execution context manager to allow context management on a higher level. ``` ########## File path: docs/howto/use-additional-execute-contextmanager.rst ########## @@ -0,0 +1,47 @@ + .. 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. + + + +Defining Additional Execute Context Manager +=============================== + +Creating new context manager +---------------------------- + +Users can create their own exeuction context manager to allow context management on a higher level. +To do so, one must define a new context manager in one of their files: + + .. code-block:: python + + import contextlib + + + @contextlib.contextmanager + def example_user_context_manager(task_instance, execution_context): + """ + An example of a context manager that a user may provide. + """ + in_context = True + try: + yield in_context + finally: + in_context = False + +Notice the context manager's signature, it receives two parameters: +1. task_instance - the executing task instance object (can also be retrieved from execution context via ``"ti"`` key. Review comment: ```suggestion 1. ``task_instance`` - the executing task instance object (can also be retrieved from execution context via ``"ti"`` key. ``` ########## File path: docs/howto/use-additional-execute-contextmanager.rst ########## @@ -0,0 +1,47 @@ + .. 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. + + + +Defining Additional Execute Context Manager +=============================== + +Creating new context manager +---------------------------- + +Users can create their own exeuction context manager to allow context management on a higher level. +To do so, one must define a new context manager in one of their files: + + .. code-block:: python + + import contextlib + + + @contextlib.contextmanager + def example_user_context_manager(task_instance, execution_context): + """ + An example of a context manager that a user may provide. + """ + in_context = True + try: + yield in_context + finally: + in_context = False + +Notice the context manager's signature, it receives two parameters: +1. task_instance - the executing task instance object (can also be retrieved from execution context via ``"ti"`` key. +2. execution_context - the execution context that is provided to an operator's ``execute`` function. Review comment: ```suggestion 2. ``execution_context`` - the execution context that is provided to an operator's ``execute`` function. ``` ########## File path: airflow/config_templates/config.yml ########## @@ -366,6 +366,17 @@ type: string example: "path.to.CustomXCom" default: "airflow.models.xcom.BaseXCom" + - name: additional_execute_contextmanager + description: | + Custom user function that returns a context manager. Syntax is "package.method". + Context is entered when operator starts executing task. __enter__() will be called + before the operator's execute method, and __exit__() shortly after. + Function's signature should accept two positional parameters - task instance + and execution context Review comment: if not here, we should add it to `docs/howto/use-additional-execute-contextmanager.rst` ########## File path: tests/task/context/test_additional_execute_contextmanager.py ########## @@ -0,0 +1,134 @@ +# 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. + +import contextlib +import os +import sys +from datetime import datetime, timedelta +from unittest import mock +from unittest.mock import MagicMock, Mock + +import pytest + +from airflow import DAG, configuration +from airflow.exceptions import AirflowException +from airflow.models import TaskInstance +from airflow.models.baseoperator import BaseOperator + +DEFAULT_ARGS = { + "owner": "test", + "depends_on_past": True, + "start_date": datetime.utcnow(), Review comment: Using `utcnow()` is discouraged ########## File path: tests/task/context/test_current_context.py ########## @@ -0,0 +1,98 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from datetime import datetime, timedelta + +import pytest + +from airflow import DAG +from airflow.exceptions import AirflowException +from airflow.models.baseoperator import BaseOperator +from airflow.operators.python import PythonOperator +from airflow.task.context.current import get_current_context, set_current_context + +DEFAULT_ARGS = { + "owner": "test", + "depends_on_past": True, + "start_date": datetime.utcnow(), Review comment: same 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. For queries about this service, please contact Infrastructure at: [email protected]
