thestage 0.6.1__py3-none-any.whl → 0.6.3__py3-none-any.whl

thestage/services/container/container_service.py

@@ -1,374 +1,374 @@
-from pathlib import Path
-from typing import List, Tuple, Optional, Dict
-
-import typer
-from thestage.entities.container import DockerContainerEntity
-from thestage.services.clients.thestage_api.dtos.container_param_request import DockerContainerActionRequestDto
-from thestage.services.clients.thestage_api.dtos.enums.container_pending_action import DockerContainerAction
-from thestage.services.clients.thestage_api.dtos.enums.container_status import DockerContainerStatus
-from thestage.entities.enums.shell_type import ShellType
-from thestage.services.clients.thestage_api.dtos.paginated_entity_list import PaginatedEntityList
-from thestage.services.container.mapper.container_mapper import ContainerMapper
-from thestage.services.filesystem_service import FileSystemService
-from thestage.services.remote_server_service import RemoteServerService
-from thestage.i18n.translation import __
-from thestage.services.abstract_service import AbstractService
-from thestage.services.clients.thestage_api.dtos.container_response import DockerContainerDto
-from thestage.helpers.error_handler import error_handler
-from thestage.services.clients.thestage_api.api_client import TheStageApiClient
-from thestage.services.config_provider.config_provider import ConfigProvider
-
-
-class ContainerService(AbstractService):
-
-    __thestage_api_client: TheStageApiClient = None
-    __config_provider: ConfigProvider = None
-
-    def __init__(
-            self,
-            thestage_api_client: TheStageApiClient,
-            config_provider: ConfigProvider,
-            remote_server_service: RemoteServerService,
-            file_system_service: FileSystemService,
-    ):
-        self.__config_provider = config_provider
-        self.__thestage_api_client = thestage_api_client
-        self.__remote_server_service = remote_server_service
-        self.__file_system_service = file_system_service
-
-
-    @error_handler()
-    def print_container_list(
-            self,
-            row: int,
-            page: int,
-            project_uid: Optional[str],
-            statuses: List[str],
-    ):
-        container_status_map = self.__thestage_api_client.get_container_business_status_map()
-
-        if not statuses:
-            statuses = ({key: container_status_map[key] for key in [
-                DockerContainerStatus.RUNNING,
-                DockerContainerStatus.STARTING,
-            ]}).values()
-
-        if "all" in statuses:
-            statuses = container_status_map.values()
-
-        for input_status_item in statuses:
-            if input_status_item not in container_status_map.values():
-                typer.echo(__("'%invalid_status%' is not one of %valid_statuses%", {
-                    'invalid_status': input_status_item,
-                    'valid_statuses': str(list(container_status_map.values()))
-                }))
-                raise typer.Exit(1)
-
-        typer.echo(__(
-            "Listing containers with the following statuses: %statuses%. To list all containers, use --status all",
-            placeholders={
-                'statuses': ', '.join([input_status_item for input_status_item in statuses])
-            }))
-
-        backend_statuses: List[str] = [key for key, value in container_status_map.items() if value in statuses]
-
-        project_id: Optional[int] = None
-        if project_uid:
-            project = self.__thestage_api_client.get_project_by_slug(slug=project_uid)
-            project_id = project.id
-
-        self.print(
-            func_get_data=self.get_list,
-            func_special_params={
-                'statuses': backend_statuses,
-                'project_id': project_id,
-            },
-            mapper=ContainerMapper(),
-            headers=list(map(lambda x: x.alias, DockerContainerEntity.model_fields.values())),
-            row=row,
-            page=page,
-            max_col_width=[15, 20, 25],
-            show_index="never",
-        )
-
-
-    @error_handler()
-    def get_list(
-            self,
-            statuses: List[str],
-            row: int = 5,
-            page: int = 1,
-            project_id: Optional[int] = None,
-    ) -> PaginatedEntityList[DockerContainerDto]:
-
-        list = self.__thestage_api_client.get_container_list(
-            statuses=statuses,
-            page=page,
-            limit=row,
-            project_id=project_id,
-        )
-
-        return list
-
-    @error_handler()
-    def get_container(
-            self,
-            container_id: Optional[int] = None,
-            container_slug: Optional[str] = None,
-    ) -> Optional[DockerContainerDto]:
-        return self.__thestage_api_client.get_container(
-            container_id=container_id,
-            container_slug=container_slug,
-        )
-
-    def get_server_auth(
-            self,
-            container: DockerContainerDto,
-            username_param: Optional[str],
-            private_key_path_override: Optional[str],
-    ) -> Tuple[str, str, Optional[str]]:
-        username = None
-        if container.instance_rented:
-            username = container.instance_rented.host_username
-            ip_address = container.instance_rented.ip_address
-        elif container.selfhosted_instance:
-            ip_address = container.selfhosted_instance.ip_address
-        else:
-            typer.echo(__("Neither rented nor self-hosted server instance found to connect to"))
-            raise typer.Exit(1)
-
-        if username_param:
-            username = username_param
-
-        if not username:
-            username = 'root'
-            typer.echo(__("No remote server username provided, using 'root' as username"))
-
-        private_key_path = private_key_path_override
-        if not private_key_path:
-            private_key_path = self.__config_provider.get_valid_private_key_path_by_ip_address(ip_address)
-            if private_key_path:
-                typer.echo(f'Using configured private key for {ip_address}: {private_key_path}')
-            else:
-                typer.echo(f'Using SSH agent to connect to {ip_address}')
-        else:
-            self.__config_provider.update_remote_server_config_entry(ip_address, Path(private_key_path))
-            typer.echo(f'Updated private key path for {ip_address}: {private_key_path}')
-
-        return username, ip_address, private_key_path
-
-    @error_handler()
-    def connect_to_container(
-            self,
-            container_uid: str,
-            username: Optional[str],
-            input_ssh_key_path: Optional[str],
-    ):
-        container: Optional[DockerContainerDto] = self.get_container(
-            container_slug=container_uid,
-        )
-
-        if not container:
-            typer.echo(f"Container with UID '{container_uid}' not found")
-            raise typer.Exit(1)
-
-        self.check_if_container_running(
-            container=container
-        )
-
-        if not container.system_name:
-            typer.echo(__("Unable to connect to container: container system_name is missing"))
-            raise typer.Exit(1)
-
-        starting_directory: str = '/'
-        workspace_mappings = {v for v in container.mappings.directory_mappings.values() if v.startswith('/workspace/') or v == '/workspace'}
-        if len(workspace_mappings) > 0:
-            starting_directory = '/workspace'
-
-        inference_mappings = {v for v in container.mappings.directory_mappings.values() if v.startswith('/opt/') or v == '/opt'}
-        if len(inference_mappings) > 0:
-            starting_directory = '/opt/project'
-
-        username, ip_address, private_key_path = self.get_server_auth(
-            container=container,
-            username_param=username,
-            private_key_path_override=input_ssh_key_path
-        )
-
-        shell: Optional[ShellType] = self.__remote_server_service.get_shell_from_container(
-            ip_address=ip_address,
-            username=username,
-            container_name=container.system_name,
-            private_key_path=private_key_path
-        )
-
-        if not shell:
-            typer.echo(f"Failed to start shell (bash, sh) in container: ensure user '{username}' has Docker access and compatible shell is available")
-            raise typer.Exit(1)
-
-        self.__remote_server_service.connect_to_container(
-            ip_address=ip_address,
-            username=username,
-            docker_name=container.system_name,
-            starting_directory=starting_directory,
-            shell=shell,
-            private_key_path=private_key_path
-        )
-
-    @error_handler()
-    def check_if_container_stopped(
-            self,
-            container: DockerContainerDto,
-    ) -> DockerContainerDto:
-        if container.frontend_status.status_key not in [
-            DockerContainerStatus.STOPPED.value,
-        ]:
-            typer.echo(__(f'Container is not stopped (status: \'{container.frontend_status.status_translation}\')'))
-            raise typer.Exit(1)
-
-        return container
-
-    @error_handler()
-    def check_if_container_running(
-            self,
-            container: DockerContainerDto,
-    ):
-        if container.frontend_status.status_key not in [
-            DockerContainerStatus.RUNNING.value,
-            DockerContainerStatus.BUSY.value,
-        ]:
-            typer.echo(__(f'Container is not running (status: \'{container.frontend_status.status_translation}\')'))
-            raise typer.Exit(1)
-
-
-    @staticmethod
-    def _get_new_path_from_mapping(
-            directory_mapping: Dict[str, str],
-            destination_path: str,
-    ) -> Tuple[Optional[str], Optional[str]]:
-
-        instance_path: Optional[str] = None
-        container_path: Optional[str] = None
-
-        for instance_mapping, container_mapping in directory_mapping.items():
-            if destination_path.startswith(f"{container_mapping}/") or destination_path == container_mapping:
-                instance_path = destination_path.replace(container_mapping, instance_mapping)
-                container_path = destination_path
-                # dont break, check all mapping list
-
-        if instance_path and container_path:
-            return instance_path, container_path
-        else:
-            return None, None
-
-
-    @error_handler()
-    def put_file_to_container(
-            self,
-            container: DockerContainerDto,
-            src_path: str,
-            copy_only_folder_contents: bool,
-            destination_path: Optional[str] = None,
-            username_param: Optional[str] = None,
-    ):
-        if not self.__file_system_service.check_if_path_exist(file=src_path):
-            typer.echo(__("File not found at specified path"))
-            raise typer.Exit(1)
-
-        username, ip_address, private_key_path = self.get_server_auth(
-            container=container,
-            username_param=username_param,
-            private_key_path_override=None
-        )
-
-        if not container.mappings or not container.mappings.directory_mappings:
-            typer.echo(__("Mapping folders not found"))
-            raise typer.Exit(1)
-
-        instance_path, container_path = self._get_new_path_from_mapping(
-            directory_mapping=container.mappings.directory_mappings,
-            destination_path=destination_path,
-        )
-
-        if not instance_path and not container_path:
-            typer.echo(__("Cannot find matching container volume mapping for specified file path"))
-            raise typer.Exit(1)
-
-        self.__remote_server_service.upload_data_to_container(
-            ip_address=ip_address,
-            username=username,
-            src_path=src_path,
-            dest_path=destination_path,
-            instance_path=instance_path,
-            container_path=container_path,
-            copy_only_folder_contents=copy_only_folder_contents,
-            private_key_path=private_key_path,
-        )
-
-    @error_handler()
-    def get_file_from_container(
-            self,
-            container: DockerContainerDto,
-            src_path: str,
-            copy_only_folder_contents: bool,
-            destination_path: Optional[str] = None,
-            username_param: Optional[str] = None,
-    ):
-        username, ip_address, private_key_path = self.get_server_auth(
-            container=container,
-            username_param=username_param,
-            private_key_path_override=None,
-        )
-
-        if not container.mappings or not container.mappings.directory_mappings:
-            typer.echo(__("Mapping folders not found"))
-            raise typer.Exit(1)
-
-        instance_path, container_path = self._get_new_path_from_mapping(
-            directory_mapping=container.mappings.directory_mappings,
-            destination_path=src_path,
-        )
-
-        if not instance_path and not container_path:
-            typer.echo(__("Cannot find matching container volume mapping for specified file path"))
-            raise typer.Exit(1)
-
-        self.__remote_server_service.download_data_from_container(
-            ip_address=ip_address,
-            username=username,
-            dest_path=destination_path,
-            instance_path=instance_path,
-            copy_only_folder_contents=copy_only_folder_contents,
-            private_key_path=private_key_path,
-        )
-
-
-    @error_handler()
-    def request_docker_container_action(
-            self,
-            container_uid: str,
-            action: DockerContainerAction,
-    ):
-        container: Optional[DockerContainerDto] = self.get_container(
-            container_slug=container_uid,
-        )
-        if not container:
-            typer.echo(f"Container with unique ID '{container_uid}' not found")
-            raise typer.Exit(1)
-
-        if action == DockerContainerAction.START:
-            self.check_if_container_stopped(container=container)
-
-        if action in [DockerContainerAction.STOP, DockerContainerAction.RESTART]:
-            self.check_if_container_running(container=container)
-
-        request_params = DockerContainerActionRequestDto(
-            dockerContainerId=container.id,
-            action=action,
-        )
-        result = self.__thestage_api_client.container_action(
-            request_param=request_params,
-        )
-
-        if result.is_success:
-            typer.echo(f'Docker container action scheduled: {action.value}')
+from pathlib import Path
+from typing import List, Tuple, Optional, Dict
+
+import typer
+from thestage.entities.container import DockerContainerEntity
+from thestage.services.clients.thestage_api.dtos.container_param_request import DockerContainerActionRequestDto
+from thestage.services.clients.thestage_api.dtos.enums.container_pending_action import DockerContainerAction
+from thestage.services.clients.thestage_api.dtos.enums.container_status import DockerContainerStatus
+from thestage.entities.enums.shell_type import ShellType
+from thestage.services.clients.thestage_api.dtos.paginated_entity_list import PaginatedEntityList
+from thestage.services.container.mapper.container_mapper import ContainerMapper
+from thestage.services.filesystem_service import FileSystemService
+from thestage.services.remote_server_service import RemoteServerService
+from thestage.i18n.translation import __
+from thestage.services.abstract_service import AbstractService
+from thestage.services.clients.thestage_api.dtos.container_response import DockerContainerDto
+from thestage.helpers.error_handler import error_handler
+from thestage.services.clients.thestage_api.api_client import TheStageApiClient
+from thestage.services.config_provider.config_provider import ConfigProvider
+
+
+class ContainerService(AbstractService):
+
+    __thestage_api_client: TheStageApiClient = None
+    __config_provider: ConfigProvider = None
+
+    def __init__(
+            self,
+            thestage_api_client: TheStageApiClient,
+            config_provider: ConfigProvider,
+            remote_server_service: RemoteServerService,
+            file_system_service: FileSystemService,
+    ):
+        self.__config_provider = config_provider
+        self.__thestage_api_client = thestage_api_client
+        self.__remote_server_service = remote_server_service
+        self.__file_system_service = file_system_service
+
+
+    @error_handler()
+    def print_container_list(
+            self,
+            row: int,
+            page: int,
+            project_uid: Optional[str],
+            statuses: List[str],
+    ):
+        container_status_map = self.__thestage_api_client.get_container_business_status_map()
+
+        if not statuses:
+            statuses = ({key: container_status_map[key] for key in [
+                DockerContainerStatus.RUNNING,
+                DockerContainerStatus.STARTING,
+            ]}).values()
+
+        if "all" in statuses:
+            statuses = container_status_map.values()
+
+        for input_status_item in statuses:
+            if input_status_item not in container_status_map.values():
+                typer.echo(__("'%invalid_status%' is not one of %valid_statuses%", {
+                    'invalid_status': input_status_item,
+                    'valid_statuses': str(list(container_status_map.values()))
+                }))
+                raise typer.Exit(1)
+
+        typer.echo(__(
+            "Listing containers with the following statuses: %statuses%. To list all containers, use --status all",
+            placeholders={
+                'statuses': ', '.join([input_status_item for input_status_item in statuses])
+            }))
+
+        backend_statuses: List[str] = [key for key, value in container_status_map.items() if value in statuses]
+
+        project_id: Optional[int] = None
+        if project_uid:
+            project = self.__thestage_api_client.get_project_by_slug(slug=project_uid)
+            project_id = project.id
+
+        self.print(
+            func_get_data=self.get_list,
+            func_special_params={
+                'statuses': backend_statuses,
+                'project_id': project_id,
+            },
+            mapper=ContainerMapper(),
+            headers=list(map(lambda x: x.alias, DockerContainerEntity.model_fields.values())),
+            row=row,
+            page=page,
+            max_col_width=[15, 20, 25],
+            show_index="never",
+        )
+
+
+    @error_handler()
+    def get_list(
+            self,
+            statuses: List[str],
+            row: int = 5,
+            page: int = 1,
+            project_id: Optional[int] = None,
+    ) -> PaginatedEntityList[DockerContainerDto]:
+
+        list = self.__thestage_api_client.get_container_list(
+            statuses=statuses,
+            page=page,
+            limit=row,
+            project_id=project_id,
+        )
+
+        return list
+
+    @error_handler()
+    def get_container(
+            self,
+            container_id: Optional[int] = None,
+            container_slug: Optional[str] = None,
+    ) -> Optional[DockerContainerDto]:
+        return self.__thestage_api_client.get_container(
+            container_id=container_id,
+            container_slug=container_slug,
+        )
+
+    def get_server_auth(
+            self,
+            container: DockerContainerDto,
+            username_param: Optional[str],
+            private_key_path_override: Optional[str],
+    ) -> Tuple[str, str, Optional[str]]:
+        username = None
+        if container.instance_rented:
+            username = container.instance_rented.host_username
+            ip_address = container.instance_rented.ip_address
+        elif container.selfhosted_instance:
+            ip_address = container.selfhosted_instance.ip_address
+        else:
+            typer.echo(__("Neither rented nor self-hosted server instance found to connect to"))
+            raise typer.Exit(1)
+
+        if username_param:
+            username = username_param
+
+        if not username:
+            username = 'root'
+            typer.echo(__("No remote server username provided, using 'root' as username"))
+
+        private_key_path = private_key_path_override
+        if not private_key_path:
+            private_key_path = self.__config_provider.get_valid_private_key_path_by_ip_address(ip_address)
+            if private_key_path:
+                typer.echo(f'Using configured private key for {ip_address}: {private_key_path}')
+            else:
+                typer.echo(f'Using SSH agent to connect to {ip_address}')
+        else:
+            self.__config_provider.update_remote_server_config_entry(ip_address, Path(private_key_path))
+            typer.echo(f'Updated private key path for {ip_address}: {private_key_path}')
+
+        return username, ip_address, private_key_path
+
+    @error_handler()
+    def connect_to_container(
+            self,
+            container_uid: str,
+            username: Optional[str],
+            input_ssh_key_path: Optional[str],
+    ):
+        container: Optional[DockerContainerDto] = self.get_container(
+            container_slug=container_uid,
+        )
+
+        if not container:
+            typer.echo(f"Container with UID '{container_uid}' not found")
+            raise typer.Exit(1)
+
+        self.check_if_container_running(
+            container=container
+        )
+
+        if not container.system_name:
+            typer.echo(__("Unable to connect to container: container system_name is missing"))
+            raise typer.Exit(1)
+
+        starting_directory: str = '/'
+        workspace_mappings = {v for v in container.mappings.directory_mappings.values() if v.startswith('/workspace/') or v == '/workspace'}
+        if len(workspace_mappings) > 0:
+            starting_directory = '/workspace'
+
+        inference_mappings = {v for v in container.mappings.directory_mappings.values() if v.startswith('/opt/') or v == '/opt'}
+        if len(inference_mappings) > 0:
+            starting_directory = '/opt/project'
+
+        username, ip_address, private_key_path = self.get_server_auth(
+            container=container,
+            username_param=username,
+            private_key_path_override=input_ssh_key_path
+        )
+
+        shell: Optional[ShellType] = self.__remote_server_service.get_shell_from_container(
+            ip_address=ip_address,
+            username=username,
+            container_name=container.system_name,
+            private_key_path=private_key_path
+        )
+
+        if not shell:
+            typer.echo(f"Failed to start shell (bash, sh) in container: ensure user '{username}' has Docker access and compatible shell is available")
+            raise typer.Exit(1)
+
+        self.__remote_server_service.connect_to_container(
+            ip_address=ip_address,
+            username=username,
+            docker_name=container.system_name,
+            starting_directory=starting_directory,
+            shell=shell,
+            private_key_path=private_key_path
+        )
+
+    @error_handler()
+    def check_if_container_stopped(
+            self,
+            container: DockerContainerDto,
+    ) -> DockerContainerDto:
+        if container.frontend_status.status_key not in [
+            DockerContainerStatus.STOPPED.value,
+        ]:
+            typer.echo(__(f'Container is not stopped (status: \'{container.frontend_status.status_translation}\')'))
+            raise typer.Exit(1)
+
+        return container
+
+    @error_handler()
+    def check_if_container_running(
+            self,
+            container: DockerContainerDto,
+    ):
+        if container.frontend_status.status_key not in [
+            DockerContainerStatus.RUNNING.value,
+            DockerContainerStatus.BUSY.value,
+        ]:
+            typer.echo(__(f'Container is not running (status: \'{container.frontend_status.status_translation}\')'))
+            raise typer.Exit(1)
+
+
+    @staticmethod
+    def _get_new_path_from_mapping(
+            directory_mapping: Dict[str, str],
+            destination_path: str,
+    ) -> Tuple[Optional[str], Optional[str]]:
+
+        instance_path: Optional[str] = None
+        container_path: Optional[str] = None
+
+        for instance_mapping, container_mapping in directory_mapping.items():
+            if destination_path.startswith(f"{container_mapping}/") or destination_path == container_mapping:
+                instance_path = destination_path.replace(container_mapping, instance_mapping)
+                container_path = destination_path
+                # dont break, check all mapping list
+
+        if instance_path and container_path:
+            return instance_path, container_path
+        else:
+            return None, None
+
+
+    @error_handler()
+    def put_file_to_container(
+            self,
+            container: DockerContainerDto,
+            src_path: str,
+            copy_only_folder_contents: bool,
+            destination_path: Optional[str] = None,
+            username_param: Optional[str] = None,
+    ):
+        if not self.__file_system_service.check_if_path_exist(file=src_path):
+            typer.echo(__("File not found at specified path"))
+            raise typer.Exit(1)
+
+        username, ip_address, private_key_path = self.get_server_auth(
+            container=container,
+            username_param=username_param,
+            private_key_path_override=None
+        )
+
+        if not container.mappings or not container.mappings.directory_mappings:
+            typer.echo(__("Mapping folders not found"))
+            raise typer.Exit(1)
+
+        instance_path, container_path = self._get_new_path_from_mapping(
+            directory_mapping=container.mappings.directory_mappings,
+            destination_path=destination_path,
+        )
+
+        if not instance_path and not container_path:
+            typer.echo(__("Cannot find matching container volume mapping for specified file path"))
+            raise typer.Exit(1)
+
+        self.__remote_server_service.upload_data_to_container(
+            ip_address=ip_address,
+            username=username,
+            src_path=src_path,
+            dest_path=destination_path,
+            instance_path=instance_path,
+            container_path=container_path,
+            copy_only_folder_contents=copy_only_folder_contents,
+            private_key_path=private_key_path,
+        )
+
+    @error_handler()
+    def get_file_from_container(
+            self,
+            container: DockerContainerDto,
+            src_path: str,
+            copy_only_folder_contents: bool,
+            destination_path: Optional[str] = None,
+            username_param: Optional[str] = None,
+    ):
+        username, ip_address, private_key_path = self.get_server_auth(
+            container=container,
+            username_param=username_param,
+            private_key_path_override=None,
+        )
+
+        if not container.mappings or not container.mappings.directory_mappings:
+            typer.echo(__("Mapping folders not found"))
+            raise typer.Exit(1)
+
+        instance_path, container_path = self._get_new_path_from_mapping(
+            directory_mapping=container.mappings.directory_mappings,
+            destination_path=src_path,
+        )
+
+        if not instance_path and not container_path:
+            typer.echo(__("Cannot find matching container volume mapping for specified file path"))
+            raise typer.Exit(1)
+
+        self.__remote_server_service.download_data_from_container(
+            ip_address=ip_address,
+            username=username,
+            dest_path=destination_path,
+            instance_path=instance_path,
+            copy_only_folder_contents=copy_only_folder_contents,
+            private_key_path=private_key_path,
+        )
+
+
+    @error_handler()
+    def request_docker_container_action(
+            self,
+            container_uid: str,
+            action: DockerContainerAction,
+    ):
+        container: Optional[DockerContainerDto] = self.get_container(
+            container_slug=container_uid,
+        )
+        if not container:
+            typer.echo(f"Container with unique ID '{container_uid}' not found")
+            raise typer.Exit(1)
+
+        if action == DockerContainerAction.START:
+            self.check_if_container_stopped(container=container)
+
+        if action in [DockerContainerAction.STOP, DockerContainerAction.RESTART]:
+            self.check_if_container_running(container=container)
+
+        request_params = DockerContainerActionRequestDto(
+            dockerContainerId=container.id,
+            action=action,
+        )
+        result = self.__thestage_api_client.container_action(
+            request_param=request_params,
+        )
+
+        if result.is_success:
+            typer.echo(f'Docker container action scheduled: {action.value}')