feat: add python remote debug support

README.md

@@ -228,6 +228,16 @@ composer-dev create example-local-environment \
   --plugins-path example_directory/plugins
 ```
 
+## Remote debugging
+To facilitate remote debugging in the local Airflow environment, use the `--enable-ssh` flag.  
+This flag allows SSH access into the container, enabling you to debug your DAGs directly.  
+For SSH connection, use the `airflow` user. The default password for this user is `airflow`.  
+If you need to change this default password, set the `COMPOSER_CONTAINER_AIRFLOW_USER_PASSWORD` environment variable
+in the variables.env file.
+
+Use the `--ssh-port` option to change which host port is mapped to the container's `ssh` port.
+The default host port used is 10022.
+
 ## Enable the container user to access mounted files and directories from the host
 
 By default, the Composer container runs as the user `airflow` with UID 999. The user needs to have access the files and

composer_local_dev/cli.py

@@ -48,7 +48,7 @@
         },
         {
             "name": "Environment options",
-            "options": ["--web-server-port", "--dags-path", "--plugins-path"],
+            "options": ["--web-server-port", "--dags-path", "--plugins-path", "--enable-ssh", "--ssh-port"],
         },
         {
             "name": "Container Memory and CPUs limit",
@@ -178,6 +178,21 @@ def cli():
     metavar="PORT",
 )
 
+option_enable_ssh = click.option(
+    "--enable-ssh",
+    is_flag=True,
+    default=False,
+    help="Enable SSH daemon in the environment.",
+    metavar="ENABLE_SSH",
+)
+
+option_ssh_port = click.option(
+    "--ssh-port",
+    type=click.IntRange(min=0, max=65535),
+    help="Port used by SSH daemon",
+    show_default="read from the configuration file",
+    metavar="SSHD_PORT",
+)
 
 def _complete_environment(ctx, param, incomplete):
     env_dirs = files.get_environment_directories()
@@ -246,6 +261,8 @@ def _complete_environment(ctx, param, incomplete):
 )
 @option_location
 @option_port
+@option_enable_ssh
+@option_ssh_port
 @click.option(
     "--dags-path",
     help="Path to DAGs folder. If it does not exist, it will be created.",
@@ -279,6 +296,8 @@ def create(
     project: Optional[str],
     location: str,
     web_server_port: Optional[int],
+    enable_ssh: Optional[bool],
+    ssh_port: Optional[int],
     environment: str,
     verbose: bool,
     debug: bool,
@@ -339,6 +358,8 @@ def create(
             location=location,
             env_dir_path=env_dir,
             web_server_port=web_server_port,
+            enable_ssh=enable_ssh,
+            ssh_port=ssh_port,
             dags_path=dags_path,
             plugins_path=plugins_path,
             database_engine=database_engine,
@@ -352,6 +373,8 @@ def create(
             location=location,
             env_dir_path=env_dir,
             port=web_server_port,
+            enable_ssh=enable_ssh,
+            ssh_port=ssh_port,
             dags_path=dags_path,
             plugins_path=plugins_path,
             database_engine=database_engine,
@@ -364,20 +387,24 @@ def create(
 @cli.command()
 @optional_environment
 @option_port
+@option_enable_ssh
+@option_ssh_port
 @verbose_mode
 @debug_mode
 @errors.catch_exceptions()
 def start(
     environment: Optional[str],
     web_server_port: Optional[int],
+    enable_ssh: Optional[bool],
+    ssh_port: Optional[int],
     verbose: bool,
     debug: bool,
 ):
     """Start Composer environment."""
     utils.setup_logging(verbose, debug)
     env_path = files.resolve_environment_path(environment)
     env = composer_environment.Environment.load_from_config(
-        env_path, web_server_port
+        env_path, web_server_port,enable_ssh,ssh_port
     )
     console.get_console().print(f"Starting {env.name} composer environment...")
     env.start()
@@ -404,12 +431,16 @@ def stop(environment: Optional[str], verbose: bool, debug: bool):
 @cli.command()
 @optional_environment
 @option_port
+@option_enable_ssh
+@option_ssh_port
 @verbose_mode
 @debug_mode
 @errors.catch_exceptions()
 def restart(
     environment: Optional[str],
     web_server_port: Optional[int],
+    enable_ssh: Optional[bool],
+    ssh_port: Optional[int],
     verbose: bool,
     debug: bool,
 ):
@@ -422,7 +453,7 @@ def restart(
     utils.setup_logging(verbose, debug)
     env_path = files.resolve_environment_path(environment)
     env = composer_environment.Environment.load_from_config(
-        env_path, web_server_port
+        env_path, web_server_port,enable_ssh,ssh_port
     )
     env.restart()
 

composer_local_dev/constants.py

@@ -28,6 +28,9 @@
 CONFLICT_ERROR_CODE = 409
 SERVER_ERROR_CODE = 500
 
+DEFAULT_ENABLE_SSH = False
+DEFAULT_SSH_PORT = 10022
+
 
 class ContainerStatus(str, enum.Enum):
     RUNNING = "running"

composer_local_dev/docker_files/entrypoint.sh

@@ -102,6 +102,24 @@ create_user() {
   fi
 }
 
+install_and_run_sshd() {
+  if ! command -v /usr/sbin/sshd &> /dev/null
+  then
+    echo "Installing sshd"
+    sudo bash -c 'cat << EOF > /etc/apt/sources.list.d/default.list
+deb http://archive.ubuntu.com/ubuntu/ noble main restricted
+deb http://archive.ubuntu.com/ubuntu/ noble-updates main restricted
+EOF'
+    sudo apt-get -qq update > /dev/null 2>&1
+    sudo apt-get -qqy install openssh-server > /dev/null 2>&1
+    sudo chown airflow:airflow /etc/ssh/ssh_host_*_key
+    sudo mkdir /run/sshd
+    echo "airflow:${COMPOSER_CONTAINER_AIRFLOW_USER_PASSWORD}" | sudo chpasswd
+  fi
+  echo "Starting sshd"
+  sudo /usr/sbin/sshd
+}
+
 main() {
   sudo chown airflow:airflow airflow
   sudo chmod +x $run_as_user
@@ -116,6 +134,10 @@ main() {
     sudo chown -R airflow:airflow "$FAST_API_DIR"
   fi
 
+  if [ "${COMPOSER_CONTAINER_ENABLE_SSHD}" = "True" ]; then
+    install_and_run_sshd
+  fi
+
   if [ "${COMPOSER_CONTAINER_RUN_AS_HOST_USER}" = "True" ]; then
     create_user "${COMPOSER_HOST_USER_NAME}" "${COMPOSER_HOST_USER_ID}" || true
     echo "Running Airflow as user ${COMPOSER_HOST_USER_NAME}(${COMPOSER_HOST_USER_ID})"

composer_local_dev/environment.py

@@ -358,7 +358,13 @@ def get_environments_status(
 
 
 class EnvironmentConfig:
-    def __init__(self, env_dir_path: pathlib.Path, port: Optional[int]):
+    def __init__(
+        self,
+        env_dir_path: pathlib.Path,
+        port: Optional[int],
+        enable_ssh: Optional[bool] = None,
+        ssh_port: Optional[int] = None,
+    ):
         self.env_dir_path = env_dir_path
         self.config = self.load_configuration_from_file()
         self.project_id = self.get_str_param("composer_project_id")
@@ -383,6 +389,16 @@ def __init__(self, env_dir_path: pathlib.Path, port: Optional[int]):
         )
         self.database_engine = self.get_str_param("database_engine")
 
+        self.enable_ssh = self.config.get("enable_ssh", constants.DEFAULT_ENABLE_SSH)
+        self.ssh_port = self.config.get("ssh_port", constants.DEFAULT_SSH_PORT)
+
+        if enable_ssh is not None:
+            self.enable_ssh = enable_ssh
+
+        if ssh_port is not None:
+            self.ssh_port = self.parse_int_param("ssh_port", allowed_range=(0, 65535))
+
+
     def load_configuration_from_file(self) -> Dict:
         """
         Load environment configuration from json file.
@@ -462,6 +478,8 @@ def __init__(
         memory_limit: Optional[str] = None,
         cpu_count: Optional[int] = None,
         port: Optional[int] = None,
+        enable_ssh: Optional[bool] = False,
+        ssh_port: Optional[int] = None,
         pypi_packages: Optional[Dict] = None,
         environment_vars: Optional[Dict] = None,
     ):
@@ -494,6 +512,8 @@ def __init__(
             self.database_engine == constants.DatabaseEngine.sqlite3
         )
         self.port: int = port if port is not None else 8080
+        self.enable_ssh: bool = enable_ssh if enable_ssh is not None else constants.DEFAULT_ENABLE_SSH
+        self.ssh_port: int = ssh_port if ssh_port is not None else constants.DEFAULT_SSH_PORT
         self.pypi_packages = (
             pypi_packages if pypi_packages is not None else dict()
         )
@@ -548,9 +568,15 @@ def assert_valid_environment_configuration(
                 )
 
     @classmethod
-    def load_from_config(cls, env_dir_path: pathlib.Path, port: Optional[int]):
+    def load_from_config(
+        cls,
+        env_dir_path: pathlib.Path,
+        port: Optional[int],
+        enable_ssh: Optional[bool] = None,
+        ssh_port: Optional[int] = None,
+    ):
         """Create local environment using 'config.json' configuration file."""
-        config = EnvironmentConfig(env_dir_path, port)
+        config = EnvironmentConfig(env_dir_path, port, enable_ssh, ssh_port)
         environment_vars = load_environment_variables(env_dir_path)
         Environment.assert_valid_environment_configuration(
             config, environment_vars
@@ -568,6 +594,8 @@ def load_from_config(cls, env_dir_path: pathlib.Path, port: Optional[int]):
             database_engine=config.database_engine,
             memory_limit=config.memory_limit,
             cpu_count=config.cpu_count,
+            enable_ssh=config.enable_ssh,
+            ssh_port=config.ssh_port,
             environment_vars=environment_vars,
         )
 
@@ -584,6 +612,8 @@ def from_source_environment(
         database_engine: str,
         memory_limit: Optional[str] = None,
         cpu_count: Optional[int] = None,
+        enable_ssh: Optional[bool] = False,
+        ssh_port: Optional[int] = None,
     ):
         """
         Create Environment using configuration retrieved from Composer
@@ -608,6 +638,8 @@ def from_source_environment(
             plugins_path=plugins_path,
             dag_dir_list_interval=10,
             port=web_server_port,
+            enable_ssh=enable_ssh,
+            ssh_port=ssh_port,
             pypi_packages=pypi_packages,
             environment_vars=env_variables,
             database_engine=database_engine,
@@ -666,7 +698,7 @@ def get_environment_variables_for_image_version(self) -> Dict:
         return env_vars
 
     def get_default_environment_variables(
-        self, default_db_variables: Dict[str, str]
+            self, default_db_variables: Dict[str, str], enable_ssh: bool = False
     ) -> Dict:
         """Return environment variables that will be set inside container."""
         return {
@@ -680,6 +712,8 @@ def get_default_environment_variables(
             # By default, the container runs as the user `airflow` with UID 999. Set
             # this env variable to "True" to make it run as the current host user.
             "COMPOSER_CONTAINER_RUN_AS_HOST_USER": "False",
+            "COMPOSER_CONTAINER_ENABLE_SSH": str(enable_ssh),
+            "COMPOSER_CONTAINER_AIRFLOW_USER_PASSWORD": "airflow",
             "COMPOSER_HOST_USER_NAME": f"{getpass.getuser()}",
             "COMPOSER_HOST_USER_ID": f"{os.getuid() if platform.system() != 'Windows' else ''}",
             "COMPOSER_ENVIRONMENT": self.name,
@@ -725,6 +759,8 @@ def write_environment_config_to_config_file(self):
             "database_engine": self.database_engine,
             "memory_limit": self.container_memory_limit,
             "cpu_count": self.container_cpu_count,
+            "enable_ssh": bool(self.enable_ssh),
+            "ssh_port": int(self.ssh_port),
         }
         with open(self.env_dir_path / "config.json", "w") as fp:
             json.dump(config, fp, indent=4)
@@ -804,7 +840,8 @@ def create_docker_container(self):
             db_mounts,
         )
         db_vars = db_extras["env_vars"]
-        default_vars = self.get_default_environment_variables(db_vars)
+        default_vars = self.get_default_environment_variables(db_vars, self.enable_ssh)
+
         env_vars = {**default_vars, **self.environment_vars}
         if (
             platform.system() == "Windows"
@@ -825,6 +862,10 @@ def create_docker_container(self):
         cpu_count = (
             self.container_cpu_count or constants.DOCKER_CONTAINER_CPU_COUNT
         )
+
+        if env_vars["COMPOSER_CONTAINER_ENABLE_SSH"] == "True":
+            ports["22/tcp"] = self.ssh_port
+
         try:
             container = self.create_container(
                 image=self.image_tag,

tests/e2e/incompatible_requirements.txt

@@ -1,4 +1,4 @@
 protobuf==5.29.6
 google-cloud-dlp==3.9.2
-apache-airflow-providers-google==8.6.0
+apache-airflow-providers-google==8.10.0
 proto-plus==1.20.5

tests/e2e/test_enable_ssh.py

@@ -0,0 +1,40 @@
+import pathlib
+
+import pytest
+
+from tests.e2e import (
+    assert_example_dag_listed,
+    assert_example_dag_succeeded,
+    run_app,
+)
+
+
+@pytest.mark.e2e
+def test_enable_ssh_airflow(composer_image_version, valid_project_id, env_name):
+    dags_dir = pathlib.Path(__file__).parent / "example_dag"
+    run_app(
+        f"create --from-image-version {composer_image_version} "
+        f"-p {valid_project_id} --dags-path {dags_dir} "
+        f"--enable-ssh {env_name}"
+    )
+    run_app(f"start {env_name}")
+    assert_example_dag_listed()
+    assert_example_dag_succeeded(env_name, airflow_major_version=2)
+    run_app(f"stop {env_name}")
+
+
+@pytest.mark.e2e
+def test_enable_ssh_airflow_3(
+    composer_image_version_airflow_3, valid_project_id, env_name
+):
+    dags_dir = pathlib.Path(__file__).parent / "example_dag"
+    run_app(
+        f"create --from-image-version {composer_image_version_airflow_3} "
+        f"-p {valid_project_id} --dags-path {dags_dir} "
+        f"--enable-ssh {env_name}"
+    )
+    # Copy requirements.txt with already satisfied deps to our environment
+    run_app(f"start {env_name}")
+    assert_example_dag_listed()
+    assert_example_dag_succeeded(env_name, airflow_major_version=3)
+    run_app(f"stop {env_name}")

tests/unit/test_cli.py

@@ -182,15 +182,15 @@ def mocked_resolve_env(self, env_path):
         ) as mock_check:
             yield mock_check
 
-    def assert_env_loaded(self, mocked_env, env_path, port=None):
-        mocked_env.load_from_config.assert_called_with(env_path, port)
+    def assert_env_loaded(self, mocked_env, env_path, port=None, enable_ssh=False, ssh_port=None):
+        mocked_env.load_from_config.assert_called_with(env_path, port, enable_ssh, ssh_port)
 
-    def assert_run_command(self, command, mocked_env, env_path, port=None):
+    def assert_run_command(self, command, mocked_env, env_path, port=None, enable_ssh=False, ssh_port=None):
         run_composer_and_assert_exit_code(
             command,
             exit_code=0,
         )
-        self.assert_env_loaded(mocked_env, env_path, port)
+        self.assert_env_loaded(mocked_env, env_path, port, enable_ssh, ssh_port)
 
     @pytest.mark.parametrize("command", ["start", "restart"])
     def test_start_command(