diff --git a/README.md b/README.md index 799f5e8d..bd4d3c76 100644 --- a/README.md +++ b/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 diff --git a/composer_local_dev/cli.py b/composer_local_dev/cli.py index bc8cbda9..246e3b99 100644 --- a/composer_local_dev/cli.py +++ b/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,12 +387,16 @@ 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, ): @@ -377,7 +404,7 @@ def start( 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() diff --git a/composer_local_dev/constants.py b/composer_local_dev/constants.py index e470b078..3ebab7a9 100644 --- a/composer_local_dev/constants.py +++ b/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" diff --git a/composer_local_dev/docker_files/entrypoint.sh b/composer_local_dev/docker_files/entrypoint.sh index 66172407..f710a1e7 100755 --- a/composer_local_dev/docker_files/entrypoint.sh +++ b/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})" diff --git a/composer_local_dev/environment.py b/composer_local_dev/environment.py index 6326d29a..9c495f3c 100644 --- a/composer_local_dev/environment.py +++ b/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, diff --git a/tests/e2e/incompatible_requirements.txt b/tests/e2e/incompatible_requirements.txt index b1d81767..f6cfce96 100644 --- a/tests/e2e/incompatible_requirements.txt +++ b/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 \ No newline at end of file diff --git a/tests/e2e/test_enable_ssh.py b/tests/e2e/test_enable_ssh.py new file mode 100644 index 00000000..a3535de9 --- /dev/null +++ b/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}") diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index c40fee68..8f3fadb6 100644 --- a/tests/unit/test_cli.py +++ b/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( diff --git a/tests/unit/test_environment.py b/tests/unit/test_environment.py index ea771d8b..b2fc6f57 100644 --- a/tests/unit/test_environment.py +++ b/tests/unit/test_environment.py @@ -483,6 +483,8 @@ def test_from_image( @pytest.mark.parametrize( "database_engine", constants.DatabaseEngine.choices() ) + @pytest.mark.parametrize("ssh_port", [None, 2222]) + @pytest.mark.parametrize("enable_ssh", [False, True]) @mock.patch("composer_local_dev.environment.docker.from_env") @mock.patch("composer_local_dev.environment.assert_image_exists") def test_create_and_load_from_config( @@ -492,6 +494,8 @@ def test_create_and_load_from_config( pypi_packages, database_engine, port, + enable_ssh, + ssh_port, tmp_path, ): env_dir_path = tmp_path / ".compose" / "my_env" @@ -504,6 +508,8 @@ def test_create_and_load_from_config( dags_path=str(pathlib.Path(tmp_path)), dag_dir_list_interval=10, port=port, + enable_ssh=enable_ssh, + ssh_port=ssh_port, pypi_packages=pypi_packages, database_engine=database_engine, ) @@ -702,7 +708,9 @@ def test_create_docker_container( "DAGS_FOLDER": "/home/airflow/gcs/dags", "COMPOSER_LOCATION": default_env.location, "COMPOSER_ENVIRONMENT": "my_env", + "COMPOSER_CONTAINER_AIRFLOW_USER_PASSWORD": "airflow", "COMPOSER_CONTAINER_RUN_AS_HOST_USER": "False", + "COMPOSER_CONTAINER_ENABLE_SSH": "False", "COMPOSER_HOST_USER_NAME": f"{getpass.getuser()}", "COMPOSER_HOST_USER_ID": f"{os.getuid() if platform.system() != 'Windows' else ''}", "AIRFLOW_HOME": "/home/airflow/airflow", @@ -768,6 +776,8 @@ def test_create_docker_container_with_custom_limits( "COMPOSER_LOCATION": default_env.location, "COMPOSER_ENVIRONMENT": "my_env", "COMPOSER_CONTAINER_RUN_AS_HOST_USER": "False", + "COMPOSER_CONTAINER_ENABLE_SSH": "False", + "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 ''}", "AIRFLOW_HOME": "/home/airflow/airflow", @@ -1140,6 +1150,8 @@ def test_get_environment_variables(self, mocked_docker): "COMPOSER_LOCATION": "eu-west", "COMPOSER_ENVIRONMENT": "env_name", "AIRFLOW_HOME": "/home/airflow/airflow", + "COMPOSER_CONTAINER_AIRFLOW_USER_PASSWORD": "airflow", + "COMPOSER_CONTAINER_ENABLE_SSH": "False", "COMPOSER_CONTAINER_RUN_AS_HOST_USER": "False", "COMPOSER_HOST_USER_NAME": f"{getpass.getuser()}", "COMPOSER_HOST_USER_ID": f"{os.getuid() if platform.system() != 'Windows' else ''}", @@ -1199,6 +1211,8 @@ def test_get_environment_variables_airflow_3(self, mocked_docker): "COMPOSER_LOCATION": "eu-west", "COMPOSER_ENVIRONMENT": "env_name", "AIRFLOW_HOME": "/home/airflow/airflow", + "COMPOSER_CONTAINER_AIRFLOW_USER_PASSWORD": "airflow", + "COMPOSER_CONTAINER_ENABLE_SSH": "False", "COMPOSER_CONTAINER_RUN_AS_HOST_USER": "False", "COMPOSER_HOST_USER_NAME": f"{getpass.getuser()}", "COMPOSER_HOST_USER_ID": f"{os.getuid() if platform.system() != 'Windows' else ''}",