Fix running Composer as host user

composer_local_dev/docker_files/entrypoint.sh

@@ -1,4 +1,4 @@
-#!/bin/sh"
+#!/bin/sh
 
 # Copyright 2022 Google LLC
 #
@@ -17,27 +17,29 @@
 set -xe
 
 init_airflow() {
-  sudo chown airflow:airflow airflow
+  # The run command is defined in the main() function below
+  # If COMPOSER_HOST_USER_NAME is True, run will execute the command as the local user
+  # Otherwise, the command will be executed as is
 
-  mkdir -p ${AIRFLOW__CORE__DAGS_FOLDER}
-  mkdir -p ${AIRFLOW__CORE__PLUGINS_FOLDER}
-  mkdir -p ${AIRFLOW__CORE__DATA_FOLDER}
+  run mkdir -p ${AIRFLOW__CORE__DAGS_FOLDER}
+  run mkdir -p ${AIRFLOW__CORE__PLUGINS_FOLDER}
+  run mkdir -p ${AIRFLOW__CORE__DATA_FOLDER}
 
   # That file exists in Composer < 1.19.2 and is responsible for linking name
   # `python` to python3 exec, in Composer >= 1.19.2 name `python` is already
   # linked to python3 and file no longer exist.
   if [ -f /var/local/setup_python_command.sh ]; then
-      /var/local/setup_python_command.sh
+      run /var/local/setup_python_command.sh
   fi
 
-  pip3 install --upgrade -r composer_requirements.txt
-  pip3 check
+  run pip3 install --upgrade -r composer_requirements.txt
+  run pip3 check
 
-  airflow db init
+  run airflow db init
 
   # Allow non-authenticated access to UI for Airflow 2.*
   if ! grep -Fxq "AUTH_ROLE_PUBLIC = 'Admin'" /home/airflow/airflow/webserver_config.py; then
-    echo "AUTH_ROLE_PUBLIC = 'Admin'" >> /home/airflow/airflow/webserver_config.py
+    run sh -c "echo \"AUTH_ROLE_PUBLIC = 'Admin'\" >> /home/airflow/airflow/webserver_config.py"
   fi
 }
 
@@ -59,29 +61,46 @@ create_user() {
   sudo find /var -user "${old_user_id}" -exec chown -h "${user_name}" {} \;
 }
 
-run_airflow_as_host_user() {
-  create_user "${COMPOSER_HOST_USER_NAME}" "${COMPOSER_HOST_USER_ID}"
+prepare_running_airflow_as_host_user() {
+  # Do not recreate user if it already exists
+  create_user "${COMPOSER_HOST_USER_NAME}" "${COMPOSER_HOST_USER_ID}" || true
+
+  # Define run command
+  sudo cat >/usr/local/bin/run <<EOF
+#!/bin/sh
+
+sudo -E -u ${COMPOSER_HOST_USER_NAME} env ENV=${ENV} PATH=${PATH} "\$@"
+EOF
+
   echo "Running Airflow as user ${COMPOSER_HOST_USER_NAME}(${COMPOSER_HOST_USER_ID})"
-  sudo -E -u "${COMPOSER_HOST_USER_NAME}" env PATH=${PATH} airflow scheduler &
-  sudo -E -u "${COMPOSER_HOST_USER_NAME}" env PATH=${PATH} airflow triggerer &
-  exec sudo -E -u "${COMPOSER_HOST_USER_NAME}" env PATH=${PATH} airflow webserver
 }
 
-run_airflow_as_airflow_user() {
+prepare_running_airflow_as_airflow_user() {
   echo "Running Airflow as user airflow(999)"
-  airflow scheduler &
-  airflow triggerer &
-  exec airflow webserver
+
+  # Define run command
+  sudo cat >/usr/local/bin/run <<EOF
+#!/bin/sh
+
+exec "\$@"
+EOF
 }
 
 main() {
-  init_airflow
+  sudo chown airflow:airflow airflow
 
   if [ "${COMPOSER_CONTAINER_RUN_AS_HOST_USER}" = "True" ]; then
-    run_airflow_as_host_user
+    prepare_running_airflow_as_host_user
   else
-    run_airflow_as_airflow_user
+    prepare_running_airflow_as_airflow_user
   fi
+
+  sudo chmod +x /usr/local/bin/run
+
+  init_airflow
+  run airflow scheduler &
+  run airflow triggerer &
+  exec run airflow webserver
 }
 
 main "$@"

composer_local_dev/environment.py

@@ -858,6 +858,7 @@ def run_airflow_command(self, command: List) -> None:
         """
         container = self.get_container(assert_running=True)
         command.insert(0, "airflow")
+        command.insert(0, "run")
         result = container.exec_run(cmd=command)
         console.get_console().print(result.output.decode())