thestage 0.6.4__py3-none-any.whl → 0.6.6__py3-none-any.whl

thestage/services/connect/connect_service.py

@@ -3,21 +3,18 @@ import typer
 
 from thestage.cli_command import CliCommand
 from thestage.cli_command_helper import check_command_permission
+from thestage.color_scheme.color_scheme import ColorScheme
 from thestage.i18n.translation import __
 from thestage.services.clients.thestage_api.core.http_client_exception import HttpClientException
-from thestage.services.clients.thestage_api.dtos.enums.container_status import DockerContainerStatus
-from thestage.services.clients.thestage_api.dtos.enums.selfhosted_status import SelfhostedBusinessStatus
 from thestage.services.clients.thestage_api.dtos.enums.instance_rented_status import InstanceRentedBusinessStatus
 from thestage.services.abstract_service import AbstractService
 from thestage.helpers.error_handler import error_handler
 from thestage.services.clients.thestage_api.api_client import TheStageApiClient
-from thestage.services.clients.thestage_api.dtos.enums.task_status import TaskStatus
 from thestage.services.clients.thestage_api.dtos.instance_rented_response import InstanceRentedDto
 from thestage.services.container.container_service import ContainerService
 from thestage.services.instance.instance_service import InstanceService
 from thestage.services.logging.logging_service import LoggingService
-from thestage.services.task.dto.task_dto import TaskDto
-from thestage.helpers.logger.app_logger import app_logger
+from rich import print
 
 
 class ConnectService(AbstractService):
@@ -44,107 +41,113 @@ class ConnectService(AbstractService):
     @error_handler()
     def connect_to_entity(
             self,
-            uid: str,
+            input_entity_identifier: str,
             username: Optional[str],
             private_key_path: Optional[str],
     ):
-        try:
-            instance_selfhosted = self.__thestage_api_client.get_selfhosted_instance(instance_slug=uid)
-        except HttpClientException as e:
-            app_logger.warn(f"get_selfhosted_instance: code {e.get_status_code()}")
-            instance_selfhosted = None
-
-        try:
-            instance_rented = self.__thestage_api_client.get_rented_instance(instance_slug=uid)
-        except HttpClientException as e:
-            app_logger.warn(f"get_rented_instance: code {e.get_status_code()}")
-            instance_rented = None
-
-        try:
-            container = self.__thestage_api_client.get_container(container_slug=uid,)
-        except HttpClientException as e:
-            app_logger.warn(f"get_container: code {e.get_status_code()}")
-            container = None
-
-        task: Optional[TaskDto] = None
-        if uid.isdigit():
-            try:
-                task_view_response = self.__thestage_api_client.get_task(task_id=int(uid))
-            except HttpClientException as e:
-                app_logger.warn(f"get_task error: code {e.get_status_code()}")
-                task_view_response = None
-            if task_view_response and task_view_response.task:
-                task = task_view_response.task
-
-        rented_exists = int(instance_rented is not None and instance_rented.frontend_status.status_key == InstanceRentedBusinessStatus.ONLINE)
-        selfhosted_exists = int(instance_selfhosted is not None)
-        container_exists = int(container is not None)
-        task_exists = int(task is not None)
-
-        rented_presence = int(rented_exists and instance_rented.frontend_status.status_key == InstanceRentedBusinessStatus.ONLINE)
-        selfhosted_presence = int(selfhosted_exists and instance_selfhosted.frontend_status.status_key == SelfhostedBusinessStatus.RUNNING)
-        container_presence = int(container_exists and (container.frontend_status.status_key == DockerContainerStatus.RUNNING or container.frontend_status.status_key == DockerContainerStatus.BUSY))
-        task_presence = int(task_exists and task.frontend_status.status_key in [TaskStatus.RUNNING, TaskStatus.SCHEDULED])
-
-        if rented_exists:
-            typer.echo(__("Found a rented instance with the provided UID in status: '%rented_status%'", {"rented_status": instance_rented.frontend_status.status_translation}))
-
-        if selfhosted_exists:
-            typer.echo(__("Found a self-hosted instance with the provided UID in status: '%selfhosted_status%'", {"selfhosted_status": instance_selfhosted.frontend_status.status_translation}))
-
-        if container_exists:
-            typer.echo(__("Found a docker container with the provided UID in status: '%container_status%'", {"container_status": container.frontend_status.status_translation}))
-
-        if task_exists:
-            typer.echo(__("Found a task with the provided ID in status: '%task_status%'", {"task_status": task.frontend_status.status_translation}))
-
-        if (rented_presence + selfhosted_presence + container_presence + task_presence) > 1:
+        resolved_options = self.__thestage_api_client.resolve_connect_options(entity_identifier=input_entity_identifier)
+        entities_available_for_connect_count = 0
+        task_presence = False
+        container_presence = False
+        rented_presence = False
+        selfhosted_presence = False
+        resolved_entity_public_id = None
+
+        if resolved_options.taskMatchData:
+            for task_item in resolved_options.taskMatchData:
+                message = f"Found a task with with matching {task_item.matchedField} in status: '{task_item.frontendStatus.status_translation}' (ID: {task_item.publicId})"
+                line_color = ColorScheme.SUCCESS if task_item.canConnect else 'default'
+                print(f"[{line_color}]{message}[{line_color}]")
+                if task_item.canConnect:
+                    task_presence = True
+                    resolved_entity_public_id = task_item.publicId
+                    entities_available_for_connect_count += 1
+
+        if resolved_options.dockerContainerMatchData:
+            for container_item in resolved_options.dockerContainerMatchData:
+                message = f"Found a container with matching {container_item.matchedField} in status: '{container_item.frontendStatus.status_translation}' (ID: {container_item.publicId})"
+                line_color = ColorScheme.SUCCESS if container_item.canConnect else 'default'
+                print(f"[{line_color}]{message}[{line_color}]")
+                if container_item.canConnect:
+                    container_presence = True
+                    resolved_entity_public_id = container_item.publicId
+                    entities_available_for_connect_count += 1
+
+        if resolved_options.instanceRentedMatchData:
+            for instance_rented_item in resolved_options.instanceRentedMatchData:
+                message = f"Found a rented instance with matching {instance_rented_item.matchedField} in status: '{instance_rented_item.frontendStatus.status_translation}'  (ID: {instance_rented_item.publicId})"
+                line_color = ColorScheme.SUCCESS if instance_rented_item.canConnect else 'default'
+                print(f"[{line_color}]{message}[{line_color}]")
+
+                if instance_rented_item.canConnect:
+                    rented_presence = True
+                    resolved_entity_public_id = instance_rented_item.publicId
+                    entities_available_for_connect_count += 1
+
+        if resolved_options.selfhostedInstanceMatchData:
+            for selfhosted_item in resolved_options.selfhostedInstanceMatchData:
+                message = f"Found a self-hosted instance with matching {selfhosted_item.matchedField} in status: '{selfhosted_item.frontendStatus.status_translation}' (ID: {selfhosted_item.publicId})"
+                line_color = ColorScheme.SUCCESS if selfhosted_item.canConnect else 'default'
+                print(f"[{line_color}]{message}[{line_color}]")
+
+                if selfhosted_item.canConnect:
+                    selfhosted_presence = True
+                    resolved_entity_public_id = selfhosted_item.publicId
+                    entities_available_for_connect_count += 1
+
+        if entities_available_for_connect_count > 1:
             typer.echo("Provided identifier caused ambiguity")
             typer.echo("Consider running a dedicated command to connect to the entity you need")
             raise typer.Exit(code=1)
 
-        if (rented_presence + selfhosted_presence + container_presence + task_presence) == 0:
+        if entities_available_for_connect_count == 0:
             typer.echo("There is nothing to connect to with the provided identifier")
             raise typer.Exit(code=1)
 
         if rented_presence:
             check_command_permission(CliCommand.INSTANCE_RENTED_CONNECT)
-            typer.echo("Connecting to rented instance...")
+            typer.echo(f"Connecting to rented instance '{resolved_entity_public_id}'...")
             self.__instance_service.connect_to_rented_instance(
-                instance_rented_slug=uid,
+                instance_rented_public_id=resolved_entity_public_id,
+                instance_rented_slug=None,
                 input_ssh_key_path=private_key_path
             )
 
         if container_presence:
             check_command_permission(CliCommand.CONTAINER_CONNECT)
-            typer.echo("Connecting to docker container...")
+            typer.echo(f"Connecting to docker container '{resolved_entity_public_id}'...")
             self.__container_service.connect_to_container(
-                container_uid=uid,
+                container_public_id=resolved_entity_public_id,
+                container_slug=None,
                 username=username,
                 input_ssh_key_path=private_key_path
             )
 
         if selfhosted_presence:
             check_command_permission(CliCommand.INSTANCE_SELF_HOSTED_CONNECT)
-            typer.echo("Connecting to self-hosted instance...")
+            typer.echo(f"Connecting to self-hosted instance '{resolved_entity_public_id}'...")
 
             self.__instance_service.connect_to_selfhosted_instance(
-                selfhosted_instance_slug=uid,
+                selfhosted_instance_public_id=resolved_entity_public_id,
+                selfhosted_instance_slug=None,
                 username=username,
                 input_ssh_key_path=private_key_path
             )
 
         if task_presence:
-            typer.echo(__("Connecting to task..."))
-            self.__logging_service.stream_task_logs_with_controls(task_id=int(uid))
+            typer.echo(f"Connecting to task '{resolved_entity_public_id}'...")
+            self.__logging_service.stream_task_logs_with_controls(task_public_id=resolved_entity_public_id)
 
 
     @error_handler()
-    def upload_ssh_key(self, public_key_contents: str, instance_slug: Optional[str]):
+    def upload_ssh_key(self, public_key_contents: str, instance_public_id: Optional[str], instance_slug: Optional[str]):
         instance_rented: Optional[InstanceRentedDto] = None
-        if instance_slug:
+        if instance_slug or instance_public_id:
             try:
-                instance_rented = self.__thestage_api_client.get_rented_instance(instance_slug=instance_slug)
+                instance_rented = self.__thestage_api_client.get_rented_instance(
+                    instance_public_id=instance_public_id,
+                    instance_slug=instance_slug
+                )
             except HttpClientException as e:
                 instance_rented = None
 
@@ -159,7 +162,7 @@ class ConnectService(AbstractService):
             public_key=public_key_contents
         )
 
-        ssh_key_pair_id = is_user_already_has_key_response.sshKeyPairId
+        ssh_key_pair_public_id = is_user_already_has_key_response.sshKeyPairPublicId
         is_adding_key_to_user = not is_user_already_has_key_response.isUserHasPublicKey
 
         if is_adding_key_to_user and not note_to_send:
@@ -179,12 +182,12 @@ class ConnectService(AbstractService):
                 note=note_to_send
             )
             typer.echo(f"Public key '{note_to_send}' added to your profile")
-            ssh_key_pair_id = add_ssh_key_to_user_response.sshKeyPairId
+            ssh_key_pair_public_id = add_ssh_key_to_user_response.sshKeyPairPublicId
 
         if instance_rented:
             self.__thestage_api_client.add_public_ssh_key_to_instance_rented(
-                instance_rented_id=instance_rented.id,
-                ssh_key_pair_id=ssh_key_pair_id
+                instance_rented_public_id=instance_rented.public_id,
+                ssh_key_pair_public_id=ssh_key_pair_public_id
             )
 
             if instance_rented.frontend_status.status_key != InstanceRentedBusinessStatus.ONLINE:

thestage/services/container/container_service.py

@@ -42,7 +42,8 @@ class ContainerService(AbstractService):
             self,
             row: int,
             page: int,
-            project_uid: Optional[str],
+            project_public_id: Optional[str],
+            project_slug: Optional[str],
             statuses: List[str],
     ):
         container_status_map = self.__thestage_api_client.get_container_business_status_map()
@@ -72,22 +73,18 @@ class ContainerService(AbstractService):
 
         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,
+                'project_slug': project_slug,
+                'project_public_id': project_public_id,
             },
             mapper=ContainerMapper(),
             headers=list(map(lambda x: x.alias, DockerContainerEntity.model_fields.values())),
             row=row,
             page=page,
-            max_col_width=[15, 20, 25],
+            max_col_width=[35, 20, 25],
             show_index="never",
         )
 
@@ -98,26 +95,29 @@ class ContainerService(AbstractService):
             statuses: List[str],
             row: int = 5,
             page: int = 1,
-            project_id: Optional[int] = None,
+            project_public_id: Optional[str] = None,
+            project_slug: Optional[str] = None,
     ) -> PaginatedEntityList[DockerContainerDto]:
 
         list = self.__thestage_api_client.get_container_list(
             statuses=statuses,
             page=page,
             limit=row,
-            project_id=project_id,
+            project_public_id=project_public_id,
+            project_slug=project_slug
         )
 
         return list
 
+    # TODO delete this proxy method
     @error_handler()
     def get_container(
             self,
-            container_id: Optional[int] = None,
+            container_public_id: Optional[str] = None,
             container_slug: Optional[str] = None,
     ) -> Optional[DockerContainerDto]:
         return self.__thestage_api_client.get_container(
-            container_id=container_id,
+            container_public_id=container_public_id,
             container_slug=container_slug,
         )
 
@@ -160,16 +160,18 @@ class ContainerService(AbstractService):
     @error_handler()
     def connect_to_container(
             self,
-            container_uid: str,
+            container_public_id: Optional[str],
+            container_slug: Optional[str],
             username: Optional[str],
             input_ssh_key_path: Optional[str],
     ):
         container: Optional[DockerContainerDto] = self.get_container(
-            container_slug=container_uid,
+            container_public_id=container_public_id,
+            container_slug=container_slug,
         )
 
         if not container:
-            typer.echo(f"Container with UID '{container_uid}' not found")
+            typer.echo(f"Container not found")
             raise typer.Exit(1)
 
         self.check_if_container_running(
@@ -346,14 +348,16 @@ class ContainerService(AbstractService):
     @error_handler()
     def request_docker_container_action(
             self,
-            container_uid: str,
+            container_public_id: Optional[str],
+            container_slug: Optional[str],
             action: DockerContainerAction,
     ):
         container: Optional[DockerContainerDto] = self.get_container(
-            container_slug=container_uid,
+            container_public_id=container_public_id,
+            container_slug=container_slug,
         )
         if not container:
-            typer.echo(f"Container with unique ID '{container_uid}' not found")
+            typer.echo(f"Container not found")
             raise typer.Exit(1)
 
         if action == DockerContainerAction.START:
@@ -363,7 +367,7 @@ class ContainerService(AbstractService):
             self.check_if_container_running(container=container)
 
         request_params = DockerContainerActionRequestDto(
-            dockerContainerId=container.id,
+            dockerContainerPublicId=container.public_id,
             action=action,
         )
         result = self.__thestage_api_client.container_action(

thestage/services/container/mapper/container_mapper.py

@@ -22,8 +22,9 @@ class ContainerMapper(AbstractMapper):
 
         return DockerContainerEntity(
             status=item.frontend_status.status_translation if item.frontend_status else '',
+            public_id=item.public_id or '',
             slug=item.slug or '',
-            title=item.title or '',
+            project_slug=item.project.slug if item.project else '',
             instance_type=instance_type,
             instance_slug=instance_slug,
             docker_image=item.docker_image or '',

thestage/services/instance/instance_service.py

@@ -35,21 +35,6 @@ class InstanceService(AbstractService):
         self.__remote_server_service = remote_server_service
         self.__config_provider = config_provider
 
-    def get_rented_instance(
-            self,
-            instance_slug: str,
-    ) -> Optional[InstanceRentedDto]:
-        return self.__thestage_api_client.get_rented_instance(
-            instance_slug=instance_slug,
-        )
-
-    def get_self_hosted_instance(
-            self,
-            instance_slug: str,
-    ) -> Optional[SelfHostedInstanceDto]:
-        return self.__thestage_api_client.get_selfhosted_instance(
-            instance_slug=instance_slug,
-        )
 
     @error_handler()
     def check_instance_status_to_connect(
@@ -116,10 +101,14 @@ class InstanceService(AbstractService):
     @error_handler()
     def connect_to_rented_instance(
             self,
-            instance_rented_slug: str,
+            instance_rented_public_id: Optional[str],
+            instance_rented_slug: Optional[str],
             input_ssh_key_path: Optional[str]
     ):
-        instance = self.get_rented_instance(instance_slug=instance_rented_slug)
+        instance = self.__thestage_api_client.get_rented_instance(
+            instance_public_id=instance_rented_public_id,
+            instance_slug=instance_rented_slug,
+        )
 
         if instance:
             self.check_instance_status_to_connect(
@@ -151,7 +140,8 @@ class InstanceService(AbstractService):
     @error_handler()
     def connect_to_selfhosted_instance(
             self,
-            selfhosted_instance_slug: str,
+            selfhosted_instance_public_id: Optional[str],
+            selfhosted_instance_slug: Optional[str],
             username: str,
             input_ssh_key_path: Optional[str],
     ):
@@ -159,7 +149,10 @@ class InstanceService(AbstractService):
             username = 'root'
             typer.echo(__("No remote server username provided, using 'root' as username"))
 
-        instance = self.get_self_hosted_instance(instance_slug=selfhosted_instance_slug)
+        instance = self.__thestage_api_client.get_selfhosted_instance(
+            instance_public_id=selfhosted_instance_public_id,
+            instance_slug=selfhosted_instance_slug,
+        )
 
         if instance:
             self.check_selfhosted_status_to_connect(
@@ -184,7 +177,7 @@ class InstanceService(AbstractService):
             if input_ssh_key_path:
                 self.__config_provider.update_remote_server_config_entry(ip_address=instance.ip_address, ssh_key_path=Path(input_ssh_key_path))
         else:
-            typer.echo(__("Server instance not found: %instance_item%", {'instance_item': selfhosted_instance_slug}))
+            typer.echo("Self-hosted instance not found")
 
 
     @error_handler()

thestage/services/instance/mapper/instance_mapper.py

@@ -13,12 +13,10 @@ class InstanceMapper(AbstractMapper):
 
         return RentedInstanceEntity(
             slug=item.slug if item.slug else '',
-            title=item.title if item.title else '',
+            public_id=item.public_id if item.public_id else '',
             cpu_type=item.cpu_type if item.cpu_type else '',
             gpu_type=item.gpu_type if item.gpu_type else '',
             cpu_cores=str(item.cpu_cores) if item.cpu_cores else '',
             ip_address=item.ip_address if item.ip_address else '',
             status=item.frontend_status.status_translation if item.frontend_status else '',
-            created_at=str(item.created_at.strftime("%Y-%m-%d %H:%M:%S")) if item.created_at else '',
-            updated_at=str(item.updated_at.strftime("%Y-%m-%d %H:%M:%S")) if item.updated_at else '',
         )

thestage/services/instance/mapper/selfhosted_mapper.py

@@ -3,6 +3,7 @@ from typing import Optional, Any, Tuple
 
 from thestage.entities.rented_instance import RentedInstanceEntity
 from thestage.entities.self_hosted_instance import SelfHostedInstanceEntity
+from thestage.services.clients.thestage_api.dtos.enums.gpu_name import InstanceGpuType
 from thestage.services.clients.thestage_api.dtos.selfhosted_instance_response import SelfHostedInstanceDto
 from thestage.services.abstract_mapper import AbstractMapper
 
@@ -18,16 +19,14 @@ class SelfHostedMapper(AbstractMapper):
             gpus = [item.type.value for item in item.detected_gpus.gpus]
 
         if len(gpus) == 0:
-            gpus = ['NO_GPU']
+            gpus = [InstanceGpuType.NO_GPU]
 
         return SelfHostedInstanceEntity(
             slug=item.slug,
-            title=item.title,
+            public_id=item.public_id,
             cpu_type=item.cpu_type,
             cpu_cores=item.cpu_cores,
             gpu_type=', '.join(gpus),
             ip_address=item.ip_address,
             status=item.frontend_status.status_translation if item.frontend_status else None,
-            created_at=item.created_at.strftime("%Y-%m-%d %H:%M:%S") if item.created_at else '',
-            updated_at=item.updated_at.strftime("%Y-%m-%d %H:%M:%S") if item.updated_at else '',
         )