From 777dca6a81d37463ea5ccf2aeabaf8b05343dedd Mon Sep 17 00:00:00 2001 From: Yury Dzerin Date: Wed, 2 Feb 2022 12:37:21 +0100 Subject: [PATCH 1/2] Implement engine restart method --- src/firebolt/common/urls.py | 1 + src/firebolt/model/engine.py | 73 ++++++++++++++++++++++--------- tests/unit/service/conftest.py | 2 +- tests/unit/service/test_engine.py | 37 ++++++++++++++++ 4 files changed, 92 insertions(+), 21 deletions(-) diff --git a/src/firebolt/common/urls.py b/src/firebolt/common/urls.py index af64b53a5b4..07a29d74d89 100644 --- a/src/firebolt/common/urls.py +++ b/src/firebolt/common/urls.py @@ -10,6 +10,7 @@ ACCOUNT_ENGINE_URL = "/core/v1/accounts/{account_id}/engines/{engine_id}" ACCOUNT_ENGINE_START_URL = ACCOUNT_ENGINE_URL + ":start" +ACCOUNT_ENGINE_RESTART_URL = ACCOUNT_ENGINE_URL + ":restart" ACCOUNT_ENGINE_STOP_URL = ACCOUNT_ENGINE_URL + ":stop" ACCOUNT_ENGINES_URL = "/core/v1/accounts/{account_id}/engines" ACCOUNT_ENGINE_BY_NAME_URL = ACCOUNT_ENGINES_URL + ":getIdByName" diff --git a/src/firebolt/model/engine.py b/src/firebolt/model/engine.py index fd10d74b391..e91edd9d3ea 100644 --- a/src/firebolt/model/engine.py +++ b/src/firebolt/model/engine.py @@ -10,6 +10,7 @@ from firebolt.common.exception import NoAttachedDatabaseError from firebolt.common.urls import ( + ACCOUNT_ENGINE_RESTART_URL, ACCOUNT_ENGINE_START_URL, ACCOUNT_ENGINE_STOP_URL, ACCOUNT_ENGINE_URL, @@ -251,7 +252,7 @@ def start( f"Engine (engine_id={engine.engine_id}, name={engine.name}) stopped." ) - engine = self._send_start() + engine = self._send_engine_request(ACCOUNT_ENGINE_START_URL) logger.info( f"Starting Engine (engine_id={engine.engine_id}, name={engine.name})" ) @@ -277,16 +278,6 @@ def start( return engine - def _send_start(self) -> Engine: - response = self._service.client.post( - url=ACCOUNT_ENGINE_START_URL.format( - account_id=self._service.account_id, engine_id=self.engine_id - ) - ) - return Engine.parse_obj_with_service( - obj=response.json()["engine"], engine_service=self._service - ) - @check_attached_to_database def stop( self, wait_for_stop: bool = False, wait_timeout_seconds: int = 3600 @@ -294,17 +285,9 @@ def stop( """Stop an Engine running on Firebolt.""" timeout_time = time.time() + wait_timeout_seconds - response = self._service.client.post( - url=ACCOUNT_ENGINE_STOP_URL.format( - account_id=self._service.account_id, engine_id=self.engine_id - ) - ) + engine = self._send_engine_request(ACCOUNT_ENGINE_STOP_URL) logger.info(f"Stopping Engine (engine_id={self.engine_id}, name={self.name})") - engine = Engine.parse_obj_with_service( - obj=response.json()["engine"], engine_service=self._service - ) - while wait_for_stop and engine.current_status_summary not in { EngineStatusSummary.ENGINE_STATUS_SUMMARY_STOPPED, EngineStatusSummary.ENGINE_STATUS_SUMMARY_FAILED, @@ -320,6 +303,46 @@ def stop( return engine + @check_attached_to_database + def restart( + self, + wait_for_startup: bool = True, + wait_timeout_seconds: int = 3600, + ) -> Engine: + """ + Restart an engine. + + Args: + wait_for_startup: + If True, wait for startup to complete. + If false, return immediately after requesting startup. + wait_timeout_seconds: + Number of seconds to wait for startup to complete + before raising a TimeoutError. + + Returns: + The updated Engine from Firebolt. + """ + timeout_time = time.time() + wait_timeout_seconds + + engine = self._send_engine_request(ACCOUNT_ENGINE_RESTART_URL) + logger.info(f"Stopping Engine (engine_id={self.engine_id}, name={self.name})") + + while wait_for_startup and engine.current_status_summary not in { + EngineStatusSummary.ENGINE_STATUS_SUMMARY_RUNNING, + EngineStatusSummary.ENGINE_STATUS_SUMMARY_FAILED, + }: + wait( + seconds=5, + timeout_time=timeout_time, + error_message=f"Could not restart engine within {wait_timeout_seconds} seconds.", # noqa: E501 + verbose=False, + ) + + engine = engine.get_latest() + + return engine + def delete(self) -> Engine: """Delete an Engine from Firebolt.""" response = self._service.client.delete( @@ -331,3 +354,13 @@ def delete(self) -> Engine: return Engine.parse_obj_with_service( obj=response.json()["engine"], engine_service=self._service ) + + def _send_engine_request(self, url: str) -> Engine: + response = self._service.client.post( + url=url.format( + account_id=self._service.account_id, engine_id=self.engine_id + ) + ) + return Engine.parse_obj_with_service( + obj=response.json()["engine"], engine_service=self._service + ) diff --git a/tests/unit/service/conftest.py b/tests/unit/service/conftest.py index 48622fc80eb..44b78a3d7dd 100644 --- a/tests/unit/service/conftest.py +++ b/tests/unit/service/conftest.py @@ -214,7 +214,7 @@ def do_mock( request: httpx.Request = None, **kwargs, ) -> Response: - assert request.url == engine_url + # assert request.url == engine_url return Response( status_code=httpx.codes.OK, json={"engine": mock_engine.dict()}, diff --git a/tests/unit/service/test_engine.py b/tests/unit/service/test_engine.py index 8a105a59f6c..b57f6e9bede 100644 --- a/tests/unit/service/test_engine.py +++ b/tests/unit/service/test_engine.py @@ -265,3 +265,40 @@ def test_attach_to_database( engine.attach_to_database(database=database) assert engine.database == database + + +def test_engine_restart( + httpx_mock: HTTPXMock, + auth_callback: Callable, + auth_url: str, + provider_callback: Callable, + provider_url: str, + settings: Settings, + mock_engine: Engine, + account_id_callback: Callable, + account_id_url: str, + engine_callback: Callable, + account_engine_url: str, + bindings_callback: Callable, + bindings_url: str, + database_callback: Callable, + database_url: str, +): + httpx_mock.add_callback(auth_callback, url=auth_url) + httpx_mock.add_callback(provider_callback, url=provider_url) + + httpx_mock.add_callback(account_id_callback, url=account_id_url) + httpx_mock.add_callback(auth_callback, url=auth_url) + + httpx_mock.add_callback( + engine_callback, url=f"{account_engine_url}:restart", method="POST" + ) + httpx_mock.add_callback(bindings_callback, url=bindings_url) + httpx_mock.add_callback(database_callback, url=database_url) + + manager = ResourceManager(settings=settings) + + mock_engine._service = manager.engines + engine = mock_engine.restart(wait_for_startup=False) + + assert engine.name == mock_engine.name From 8c8e8bcfe686c9476fb921d52682f6ca65a706b1 Mon Sep 17 00:00:00 2001 From: Yury Dzerin Date: Mon, 7 Feb 2022 14:07:00 +0100 Subject: [PATCH 2/2] add removed assert --- tests/unit/service/conftest.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/service/conftest.py b/tests/unit/service/conftest.py index 44b78a3d7dd..5ed8b1d80c5 100644 --- a/tests/unit/service/conftest.py +++ b/tests/unit/service/conftest.py @@ -1,5 +1,6 @@ import json from typing import Callable, List +from urllib.parse import urlparse import httpx import pytest @@ -214,7 +215,7 @@ def do_mock( request: httpx.Request = None, **kwargs, ) -> Response: - # assert request.url == engine_url + assert urlparse(engine_url).path in request.url.path return Response( status_code=httpx.codes.OK, json={"engine": mock_engine.dict()},