Merge pull request #436 from pieces-app/fix-auth-commands

fix auth commands

Author: bishoy-at-pieces

src/pieces/_vendor/pieces_os_client/wrapper/basic_identifier/user.py

@@ -1,5 +1,5 @@
 from typing import TYPE_CHECKING, Optional
-
+from threading import Thread
 from .basic import Basic
 
 from pieces._vendor.pieces_os_client.models.allocation_status_enum import AllocationStatusEnum
@@ -53,20 +53,27 @@ def _on_login_connect(self):
 			thread: The thread handling the login process.
 			timeout: The maximum time to wait for the login process.
 		"""
-		self.connect(True)
+		self.connect()
 
 	def login(self, connect_after_login=True, timeout=120):
 		"""
 		Logs the user into the OS and optionally connects to the cloud.
 
 		Args:
-			connect_after_login: A flag indicating if the user should connect to the cloud after login (default is True).
-			timeout: The maximum time to wait for the login process (default is 120 seconds).
+				connect_after_login: A flag indicating if the user should connect to the cloud after login (default is True).
+				timeout: The maximum time to wait for the login process (default is 120 seconds).
 		"""
-		thread = self.pieces_client.os_api.sign_into_os(async_req=True)
+		result = {}
+
+		def target():
+			result['user'] = self.pieces_client.os_api.sign_into_os()
+
+		thread = Thread(target=target)
+		thread.start()
+		thread.join(timeout)
+
 		if connect_after_login:
-			user = thread.get(timeout)
-			self.user_profile = user
+			self.user_profile = result.get('user')
 			self._on_login_connect()
 
 	def logout(self):
@@ -84,8 +91,20 @@ def connect(self, async_req = False):
 		"""
 		if not self.user_profile:
 			raise PermissionError("You must be logged in to use this feature")
-		self.on_user_callback(self.user_profile, True)  # Set the connecting to cloud bool to true
-		self.pieces_client.allocations_api.allocations_connect_new_cloud(self.user_profile,async_req=async_req)
+		self.on_user_callback(
+			self.user_profile, True
+		)  # Set the connecting to cloud bool to true
+		if async_req:
+			thread = Thread(
+				target=self.pieces_client.allocations_api.allocations_connect_new_cloud,
+				args=(self.user_profile,),
+			)
+			thread.start()
+		else:
+			self.pieces_client.allocations_api.allocations_connect_new_cloud(
+				self.user_profile
+			)
+
 
 	def disconnect(self):
 		"""

src/pieces/_vendor/pieces_os_client/wrapper/client.py

@@ -188,7 +188,7 @@ def model_id(self):
     def model_name(self, model):
         models = self.get_models()
         if model not in models:
-            raise ValueError("Not a vaild model name, the available models are"
+            raise ValueError("Not a valid model name, the available models are"
                              f"{', '.join(models.keys())}")
         self._model_name = model
         self._model_id = models[model]
@@ -282,15 +282,6 @@ def pieces_os_installer(self, callback: Callable[[DownloadModel], None]) -> PosI
         """
         return PosInstaller(callback, self.app_name)
 
-    def pool(self, api_call, args):
-        """
-            call the api async without stopping the main thread
-            Create thread pool on first request
-            avoids instantiating unused threadpool for blocking clients.
-            return the ThreadPool created
-        """
-        return self.api_client.pool.apply_async(api_call, args)
-
 
 # Register the function to be called on exit
 atexit.register(PiecesClient.close)

src/pieces/command_interface/auth_commands.py

@@ -1,4 +1,7 @@
 import argparse
+from pieces._vendor.pieces_os_client.models.allocation_status_enum import (
+    AllocationStatusEnum,
+)
 from pieces.base_command import BaseCommand
 from pieces.urls import URLs
 from pieces.settings import Settings
@@ -40,9 +43,13 @@ def execute(self, **kwargs) -> int:
             Settings.pieces_client.user_api.user_snapshot().user
         )
         if Settings.pieces_client.user.user_profile:
+            status = Settings.pieces_client.user.cloud_status or AllocationStatusEnum.DISCONNECTED
             Settings.logger.print(
-                f"Signed in as {Settings.pieces_client.user.name}\nemail: {Settings.pieces_client.user.email}"
+                f"Signed in as {Settings.pieces_client.user.name}\nEmail: {Settings.pieces_client.user.email}\nCloud status: {status.value.title()}"
             )
+            if status == AllocationStatusEnum.DISCONNECTED:
+                Settings.logger.print("Connecting to the Pieces Cloud...")
+                Settings.pieces_client.user.connect()
             return 0
         try:
             Settings.pieces_client.user.login(True)

src/pieces/core/assets_command.py

@@ -47,9 +47,11 @@ def wrapper(*args, **kwargs):
         try:
             if AssetsCommands.current_asset is None:
                 raise ValueError("No material selected")
-            AssetsCommands.current_asset.asset  # Check if the current asset is vaild
+            AssetsCommands.current_asset.asset  # Check if the current asset is valid
         except (ValueError, NotFoundException):
             ListCommand.list_assets()
+        if AssetsCommands.current_asset is None:
+            return
         return func(asset=AssetsCommands.current_asset, *args, **kwargs)
 
     return wrapper