Fix DockerOperator xcom

airflow/providers/docker/example_dags/example_docker_xcom.py

@@ -0,0 +1,86 @@
+#
+# 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.
+# pylint: disable=missing-function-docstring
+"""
+This sample writes and pull from xcom using DockerOperator.
+The following operators are being used: DockerOperator &
+BashOperator.
+"""
+from datetime import timedelta
+
+from airflow import DAG
+from airflow.operators.bash_operator import BashOperator
+from airflow.operators.docker_operator import DockerOperator
+from airflow.utils.dates import days_ago
+
+with DAG(
+    dag_id="example_docker_xcom",
+    schedule_interval="@daily",
+    start_date=days_ago(5),
+    catchup=True,
+    max_active_runs=2,
+    default_args={
+        "owner": "airflow",
+        "retries": 2,
+        "retry_delay": timedelta(minutes=5),
+    },
+) as dag:
+
+    write_xcom_last = DockerOperator(
+        task_id="write_xcom_last",
+        image="python:3-slim",
+        api_version="auto",
+        command="""python -c "import logging;[
+            print(i) for i in range(10)];logging.warning('this is a warning')" """,
+        do_xcom_push=True,
+        xcom_all=False,
+        auto_remove=True,
+        network_mode="bridge",
+    )
+
+    write_xcom_all = DockerOperator(
+        task_id="write_xcom_all",
+        image="python:3-slim",
+        api_version="auto",
+        command="""python -c "import logging;[
+            print(i) for i in range(10)];logging.warning('this is a warning')" """,
+        do_xcom_push=True,
+        xcom_all=True,
+        auto_remove=True,
+        network_mode="bridge",
+    )
+
+    pull_xcom_last = DockerOperator(
+        task_id="pull_xcom",
+        image="ubuntu:20.04",
+        api_version="auto",
+        command="""echo {{ ti.xcom_pull(task_ids="write_xcom_last") }}""",
+        do_xcom_push=True,
+        xcom_all=False,
+        auto_remove=True,
+        network_mode="bridge",
+    )
+
+    pull_xcom_all = BashOperator(
+        task_id="pull_xcom_all",
+        bash_command="""echo '{{ ti.xcom_pull(task_ids="write_xcom_all") }}' """,
+        do_xcom_push=True,
+    )
+
+    write_xcom_last >> pull_xcom_last
+    write_xcom_all >> pull_xcom_all

airflow/providers/docker/operators/docker.py

@@ -19,6 +19,7 @@
 Implements Docker operator
 """
 import ast
+from collections import deque
 from tempfile import TemporaryDirectory
 from typing import Dict, Iterable, List, Optional, Union
 
@@ -127,10 +128,10 @@ class DockerOperator(BaseOperator):
     :type cap_add: list[str]
     """
 
-    template_fields = ('command', 'environment', 'container_name')
+    template_fields = ("command", "environment", "container_name")
     template_ext = (
-        '.sh',
-        '.bash',
+        ".sh",
+        ".bash",
     )
 
     # pylint: disable=too-many-arguments,too-many-locals
@@ -143,7 +144,7 @@ def __init__(
         command: Optional[Union[str, List[str]]] = None,
         container_name: Optional[str] = None,
         cpus: float = 1.0,
-        docker_url: str = 'unix://var/run/docker.sock',
+        docker_url: str = "unix://var/run/docker.sock",
         environment: Optional[Dict] = None,
         private_environment: Optional[Dict] = None,
         force_pull: bool = False,
@@ -225,10 +226,10 @@ def _run_image(self) -> Optional[str]:
         """
         Run a Docker container with the provided image
         """
-        self.log.info('Starting docker container from image %s', self.image)
+        self.log.info("Starting docker container from image %s", self.image)
 
-        with TemporaryDirectory(prefix='airflowtmp', dir=self.host_tmp_dir) as host_tmp_dir:
-            self.volumes.append('{0}:{1}'.format(host_tmp_dir, self.tmp_dir))
+        with TemporaryDirectory(prefix="airflowtmp", dir=self.host_tmp_dir) as host_tmp_dir:
+            self.volumes.append("{0}:{1}".format(host_tmp_dir, self.tmp_dir))
 
             if not self.cli:
                 raise Exception("The 'cli' should be initialized before!")
@@ -256,29 +257,34 @@ def _run_image(self) -> Optional[str]:
 
             lines = self.cli.attach(container=self.container['Id'], stdout=True, stderr=True, stream=True)
 
-            self.cli.start(self.container['Id'])
+            def gen_output(stdout=False, stderr=False):
+                return (
+                    templ.decode("utf-8").strip() if hasattr(templ, "decode") else str(templ)
+                    for templ in self.cli.logs(
+                        self.container["Id"], stream=True, stdout=stdout, stderr=stderr
+                    )
+                )
 
-            line = ''
-            for line in lines:
-                line = line.strip()
-                if hasattr(line, 'decode'):
-                    # Note that lines returned can also be byte sequences so we have to handle decode here
-                    line = line.decode('utf-8')
+            for line in gen_output(stdout=True, stderr=True):
                 self.log.info(line)
 
-            result = self.cli.wait(self.container['Id'])
-            if result['StatusCode'] != 0:
-                raise AirflowException('docker container failed: ' + repr(result))
-
-            # duplicated conditional logic because of expensive operation
-            ret = None
+            return_value = None
             if self.do_xcom_push:
-                ret = self.cli.logs(container=self.container['Id']) if self.xcom_all else line.encode('utf-8')
+                lines = gen_output(stdout=True)
+                if self.xcom_all:
+                    return_value = "\n".join(lines)
+                else:
+                    line_deque = deque(lines, maxlen=1)
+                    return_value = line_deque.pop()
+
+            result = self.cli.wait(self.container["Id"])
+            if result["StatusCode"] != 0:
+                raise AirflowException("docker container failed: " + repr(result))
 
             if self.auto_remove:
-                self.cli.remove_container(self.container['Id'])
+                self.cli.remove_container(self.container["Id"])
 
-            return ret
+            return return_value
 
     def execute(self, context) -> Optional[str]:
         self.cli = self._get_cli()
@@ -287,12 +293,12 @@ def execute(self, context) -> Optional[str]:
 
         # Pull the docker image if `force_pull` is set or image does not exist locally
         if self.force_pull or not self.cli.images(name=self.image):
-            self.log.info('Pulling docker image %s', self.image)
+            self.log.info("Pulling docker image %s", self.image)
             for output in self.cli.pull(self.image, stream=True, decode=True):
                 if isinstance(output, str):
                     self.log.info("%s", output)
-                if isinstance(output, dict) and 'status' in output:
-                    self.log.info("%s", output['status'])
+                if isinstance(output, dict) and "status" in output:
+                    self.log.info("%s", output["status"])
 
         self.environment['AIRFLOW_TMP_DIR'] = self.tmp_dir
         return self._run_image()
@@ -311,16 +317,16 @@ def get_command(self) -> Union[List[str], str]:
         :return: the command (or commands)
         :rtype: str | List[str]
         """
-        if isinstance(self.command, str) and self.command.strip().find('[') == 0:
+        if isinstance(self.command, str) and self.command.strip().find("[") == 0:
             commands = ast.literal_eval(self.command)
         else:
             commands = self.command
         return commands
 
     def on_kill(self) -> None:
         if self.cli is not None:
-            self.log.info('Stopping docker container')
-            self.cli.stop(self.container['Id'])
+            self.log.info("Stopping docker container")
+            self.cli.stop(self.container["Id"])
 
     def __get_tls_config(self) -> Optional[tls.TLSConfig]:
         tls_config = None
@@ -334,5 +340,5 @@ def __get_tls_config(self) -> Optional[tls.TLSConfig]:
                 ssl_version=self.tls_ssl_version,  # noqa
                 assert_hostname=self.tls_hostname,
             )
-            self.docker_url = self.docker_url.replace('tcp://', 'https://')
+            self.docker_url = self.docker_url.replace("tcp://", "https://")
         return tls_config