From 24b0fe7f537baaa9338680e62f56f063a65e4e09 Mon Sep 17 00:00:00 2001 From: vshvechko Date: Mon, 17 Jun 2024 17:55:15 +0300 Subject: [PATCH] fix(cli): fix cli commands for not responding arbitrary servers (M2-7021) --- .../workspaces/commands/arbitrary_server.py | 36 ++++++++++++++----- src/apps/workspaces/domain/workspace.py | 6 ++++ src/apps/workspaces/service/workspace.py | 8 +++-- 3 files changed, 39 insertions(+), 11 deletions(-) diff --git a/src/apps/workspaces/commands/arbitrary_server.py b/src/apps/workspaces/commands/arbitrary_server.py index 898e6810d3e..72011f92d3f 100644 --- a/src/apps/workspaces/commands/arbitrary_server.py +++ b/src/apps/workspaces/commands/arbitrary_server.py @@ -1,3 +1,4 @@ +import asyncio import http import io import uuid @@ -53,8 +54,12 @@ def print_data_table(data: WorkspaceArbitraryFields) -> None: print(table) +def error_msg(msg: str): + print(f"[bold red]Error: {msg}[/bold red]") + + def error(msg: str): - print(f"[bold red]Error: \n{msg}[/bold red]") + error_msg(msg) raise typer.Abort() @@ -180,8 +185,14 @@ async def show( if not data: print(f"[bold green]Arbitrary settings are not configured for {owner_email}[/bold green]") return - alembic_version = await get_version(data.database_uri) arbitrary_fields = WorkspaceArbitraryFields.from_orm(data) + try: + alembic_version = await get_version(data.database_uri) + except asyncio.TimeoutError: + alembic_version = "[bold red]ERROR: Timeout[/bold red]" + except Exception as e: + alembic_version = f"[bold red]ERROR: {e}[/bold red]" + output = WorkSpaceArbitraryConsoleOutput( **arbitrary_fields.dict(), email=owner_email, user_id=owner.id, alembic_version=alembic_version ) @@ -191,8 +202,13 @@ async def show( user_crud = UsersCRUD(session) for data in workspaces: user = await user_crud.get_by_id(data.user_id) - alembic_version = await get_version(data.database_uri) arbitrary_fields = WorkspaceArbitraryFields.from_orm(data) + try: + alembic_version = await get_version(data.database_uri) + except asyncio.TimeoutError: + alembic_version = "[bold red]ERROR: Timeout[/bold red]" + except Exception as e: + alembic_version = f"[bold red]ERROR: {e}[/bold red]" output = WorkSpaceArbitraryConsoleOutput( **arbitrary_fields.dict(), email=user.email_encrypted, @@ -243,7 +259,9 @@ async def ping(owner_email: str = typer.Argument(..., help="Workspace owner emai async with session_maker() as session: try: owner = await UsersCRUD(session).get_by_email(owner_email) - data = await WorkspaceService(session, owner.id).get_arbitrary_info_by_owner_id_if_use_arbitrary(owner.id) + data = await WorkspaceService(session, owner.id).get_arbitrary_info_by_owner_id_if_use_arbitrary( + owner.id, in_use_only=False + ) except (UserNotFound, UserIsDeletedError): error(f"User with email {owner_email} not found") except WorkspaceNotFoundError as e: @@ -256,9 +274,11 @@ async def ping(owner_email: str = typer.Argument(..., help="Workspace owner emai print("Check database availability.") await arb_session.execute("select current_date") print(f"[green]Database for user [bold]{owner_email}[/bold] is available.[/green]") + except asyncio.TimeoutError: + error_msg("Timeout error") except Exception as e: - error(str(e)) - print(f"Check bucket {data.storage_bucket} availability.") + error_msg(str(e)) + print(f'Check bucket "{data.storage_bucket}" availability.') storage = await select_storage(owner_id=owner.id, session=session) key = "mindlogger.txt" presigned_data = storage.generate_presigned_post(data.storage_bucket, key) @@ -272,8 +292,8 @@ async def ping(owner_email: str = typer.Argument(..., help="Workspace owner emai if response.status_code == http.HTTPStatus.NO_CONTENT: print(f"[green]Bucket {data.storage_bucket} for user {owner_email} is available.[/green]") else: - print(f"Can not upload test file to bucket {data.storage_bucket} for user {owner_email}") + error_msg("File upload error") print(response.content) except httpx.HTTPError as e: - print(f"Can not upload test file to bucket {data.storage_bucket} for user {owner_email}") + error_msg("File upload error") error(str(e)) diff --git a/src/apps/workspaces/domain/workspace.py b/src/apps/workspaces/domain/workspace.py index ba9ed9b81fc..e331a6dd493 100644 --- a/src/apps/workspaces/domain/workspace.py +++ b/src/apps/workspaces/domain/workspace.py @@ -309,6 +309,12 @@ class WorkSpaceArbitraryConsoleOutput(WorkspaceArbitraryFields): email: str alembic_version: str | None + @validator("use_arbitrary") + def format_arbitrary_usage(cls, value): + if value: + return "[green]True[/green]" + return "[red]False[/red]" + class WorkspaceArbitraryCreate(WorkspaceArbitraryFields): database_uri: str diff --git a/src/apps/workspaces/service/workspace.py b/src/apps/workspaces/service/workspace.py index 48926924868..25e78724d64 100644 --- a/src/apps/workspaces/service/workspace.py +++ b/src/apps/workspaces/service/workspace.py @@ -249,12 +249,14 @@ async def get_arbitrary_info_if_use_arbitrary(self, applet_id: uuid.UUID) -> Wor except ValidationError: return None - async def get_arbitrary_info_by_owner_id_if_use_arbitrary(self, owner_id: uuid.UUID) -> WorkspaceArbitrary | None: + async def get_arbitrary_info_by_owner_id_if_use_arbitrary( + self, owner_id: uuid.UUID, *, in_use_only=True + ) -> WorkspaceArbitrary | None: schema = await UserWorkspaceCRUD(self.session).get_by_user_id(owner_id) - if not schema or not schema.use_arbitrary or not schema.database_uri: + if not schema or (in_use_only and not schema.use_arbitrary) or not schema.database_uri: return None try: - return WorkspaceArbitrary.from_orm(schema) if schema else None + return WorkspaceArbitrary.from_orm(schema) except ValidationError: return None